From 428152d82f230ea65eb23865ff6b18c5b0042f65 Mon Sep 17 00:00:00 2001 From: Adrian Lyjak Date: Tue, 27 Jan 2026 13:06:19 -0500 Subject: [PATCH] feat: extend mock server with files list, split index, and sheets APIs (#197) --- src/extraction_review/testing_utils/files.py | 49 ++ .../testing_utils/pipelines.py | 418 ++++++++++++++++++ src/extraction_review/testing_utils/server.py | 14 + src/extraction_review/testing_utils/sheets.py | 274 ++++++++++++ tests/testing_utils/test_files.py | 45 ++ tests/testing_utils/test_pipelines.py | 381 ++++++++++++++++ tests/testing_utils/test_sheets.py | 209 +++++++++ 7 files changed, 1390 insertions(+) create mode 100644 src/extraction_review/testing_utils/pipelines.py create mode 100644 src/extraction_review/testing_utils/sheets.py create mode 100644 tests/testing_utils/test_pipelines.py create mode 100644 tests/testing_utils/test_sheets.py diff --git a/src/extraction_review/testing_utils/files.py b/src/extraction_review/testing_utils/files.py index ed5c960..bc2b0fe 100644 --- a/src/extraction_review/testing_utils/files.py +++ b/src/extraction_review/testing_utils/files.py @@ -9,6 +9,7 @@ from urllib.parse import urlencode import httpx import respx from llama_cloud.types import File as CloudFile +from llama_cloud.types.file_list_response import FileListResponse from llama_cloud.types.file_query_response import FileQueryResponse, Item from llama_cloud.types.presigned_url import PresignedURL @@ -105,6 +106,9 @@ class FakeFilesNamespace: body = json_body or {"detail": "upload rejected by fake server"} self._upload_stubs.append((matcher, status_code, body, once)) + def all_files(self) -> Dict[str, StoredFile]: + return dict(self._files) + # Route registration --------------------------------------------- def register(self) -> None: server = self._server @@ -116,6 +120,14 @@ class FakeFilesNamespace: alias="upload", ) self.routes["upload"] = upload_route + list_route = server.add_route( + "GET", + "/api/v1/beta/files", + self._handle_list, + namespace="files", + alias="list_files", + ) + self.routes["list"] = list_route get_route = server.add_route( "GET", "/api/v1/beta/files/{file_id}/content", @@ -165,6 +177,43 @@ class FakeFilesNamespace: self._files[file_id] = stored return self._server.json_response(stored.file.model_dump()) + def _handle_list(self, request: httpx.Request) -> httpx.Response: + params = request.url.params + file_ids_raw = params.multi_items() + file_ids_filter = [v for k, v in file_ids_raw if k == "file_ids"] + file_name = params.get("file_name") + external_file_id = params.get("external_file_id") + page_size = int(params.get("page_size", "50")) + + files = list(self._files.values()) + if file_ids_filter: + files = [f for f in files if f.file.id in file_ids_filter] + if file_name: + files = [f for f in files if f.file.name == file_name] + if external_file_id: + files = [f for f in files if f.file.external_file_id == external_file_id] + + files = files[:page_size] + items = [ + FileListResponse( + id=f.file.id, + name=f.file.name, + project_id=f.file.project_id, + expires_at=f.file.expires_at, + external_file_id=f.file.external_file_id, + file_type=f.file.file_type, + last_modified_at=f.file.last_modified_at, + purpose=f.file.purpose, + ) + for f in files + ] + return self._server.json_response( + { + "items": [item.model_dump() for item in items], + "next_page_token": None, + } + ) + def _handle_delete(self, request: httpx.Request) -> httpx.Response: file_id = request.url.path.split("/")[-1] self._files.pop(file_id, None) diff --git a/src/extraction_review/testing_utils/pipelines.py b/src/extraction_review/testing_utils/pipelines.py new file mode 100644 index 0000000..1877ba1 --- /dev/null +++ b/src/extraction_review/testing_utils/pipelines.py @@ -0,0 +1,418 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Dict, List, Optional + +import httpx +from llama_cloud.types.managed_ingestion_status_response import ( + ManagedIngestionStatusResponse, +) +from llama_cloud.types.pipeline import Pipeline +from llama_cloud.types.pipeline_retrieve_response import ( + PipelineRetrieveResponse, + RetrievalNode, +) +from llama_cloud.types.pipelines.cloud_document import CloudDocument +from llama_cloud.types.pipelines.pipeline_file import PipelineFile +from llama_cloud.types.pipelines.text_node import TextNode + +from ._deterministic import combined_seed, generate_text_blob, utcnow + +if TYPE_CHECKING: + from .server import FakeLlamaCloudServer + + +class FakePipelinesNamespace: + def __init__(self, *, server: "FakeLlamaCloudServer") -> None: + self._server = server + self._pipelines: Dict[str, Pipeline] = {} + # Per-pipeline storage for ingested documents and files + self._documents: Dict[str, Dict[str, CloudDocument]] = {} + self._files: Dict[str, Dict[str, PipelineFile]] = {} + self.routes: Dict[str, Any] = {} + + def register(self) -> None: + server = self._server + server.add_route( + "POST", + "/api/v1/pipelines", + self._handle_create, + namespace="pipelines", + ) + server.add_route( + "GET", + "/api/v1/pipelines", + self._handle_list, + namespace="pipelines", + ) + server.add_route( + "GET", + "/api/v1/pipelines/{pipeline_id}", + self._handle_get, + namespace="pipelines", + ) + server.add_route( + "PUT", + "/api/v1/pipelines/{pipeline_id}", + self._handle_update, + namespace="pipelines", + ) + server.add_route( + "DELETE", + "/api/v1/pipelines/{pipeline_id}", + self._handle_delete, + namespace="pipelines", + ) + server.add_route( + "GET", + "/api/v1/pipelines/{pipeline_id}/status", + self._handle_get_status, + namespace="pipelines", + ) + server.add_route( + "POST", + "/api/v1/pipelines/{pipeline_id}/retrieve", + self._handle_retrieve, + namespace="pipelines", + ) + # Document ingestion + server.add_route( + "POST", + "/api/v1/pipelines/{pipeline_id}/documents", + self._handle_create_documents, + namespace="pipelines", + ) + server.add_route( + "PUT", + "/api/v1/pipelines/{pipeline_id}/documents", + self._handle_upsert_documents, + namespace="pipelines", + ) + # File ingestion + server.add_route( + "PUT", + "/api/v1/pipelines/{pipeline_id}/files", + self._handle_upsert_files, + namespace="pipelines", + ) + + # Handlers ------------------------------------------------------- + + def _handle_create(self, request: httpx.Request) -> httpx.Response: + payload = self._server.json(request) + name = payload.get("name") + if not name: + return self._server.json_response( + {"detail": "name is required"}, status_code=400 + ) + + pipeline_id = self._server.new_id("pipeline") + project_id = request.url.params.get( + "project_id", self._server.default_project_id + ) + now = utcnow() + + embedding_config = payload.get("embedding_config") or { + "type": "MANAGED_OPENAI_EMBEDDING", + "component": {}, + } + + pipeline = Pipeline( + id=pipeline_id, + name=name, + project_id=project_id, + embedding_config=embedding_config, + created_at=now, + updated_at=now, + pipeline_type=payload.get("pipeline_type", "MANAGED"), + status="CREATED", + ) + self._pipelines[pipeline_id] = pipeline + self._documents[pipeline_id] = {} + self._files[pipeline_id] = {} + return self._server.json_response(pipeline.model_dump(), status_code=200) + + def _handle_list(self, request: httpx.Request) -> httpx.Response: + params = request.url.params + project_id = params.get("project_id") + pipeline_name = params.get("pipeline_name") + + pipelines = list(self._pipelines.values()) + if project_id: + pipelines = [p for p in pipelines if p.project_id == project_id] + if pipeline_name: + pipelines = [p for p in pipelines if p.name == pipeline_name] + + return self._server.json_response([p.model_dump() for p in pipelines]) + + def _handle_get(self, request: httpx.Request) -> httpx.Response: + pipeline_id = request.url.path.split("/")[-1] + pipeline = self._pipelines.get(pipeline_id) + if not pipeline: + return self._server.json_response( + {"detail": f"Pipeline {pipeline_id} not found"}, status_code=404 + ) + return self._server.json_response(pipeline.model_dump()) + + def _handle_update(self, request: httpx.Request) -> httpx.Response: + pipeline_id = request.url.path.split("/")[-1] + pipeline = self._pipelines.get(pipeline_id) + if not pipeline: + return self._server.json_response( + {"detail": f"Pipeline {pipeline_id} not found"}, status_code=404 + ) + + payload = self._server.json(request) + data = pipeline.model_dump() + data.update({k: v for k, v in payload.items() if v is not None}) + data["updated_at"] = utcnow() + updated = Pipeline.model_validate(data) + self._pipelines[pipeline_id] = updated + return self._server.json_response(updated.model_dump()) + + def _handle_delete(self, request: httpx.Request) -> httpx.Response: + pipeline_id = request.url.path.split("/")[-1] + self._pipelines.pop(pipeline_id, None) + self._documents.pop(pipeline_id, None) + self._files.pop(pipeline_id, None) + return self._server.json_response({}, status_code=200) + + def _handle_get_status(self, request: httpx.Request) -> httpx.Response: + parts = request.url.path.split("/") + pipeline_id = parts[-2] + pipeline = self._pipelines.get(pipeline_id) + if not pipeline: + return self._server.json_response( + {"detail": f"Pipeline {pipeline_id} not found"}, status_code=404 + ) + status = ManagedIngestionStatusResponse(status="SUCCESS") + return self._server.json_response(status.model_dump()) + + # Helpers -------------------------------------------------------- + + @staticmethod + def _extract_list(payload: Any, key: str) -> List[Dict[str, Any]]: + """Extract a list from payload that may be a bare array or a dict.""" + if isinstance(payload, list): + return payload + if isinstance(payload, dict): + return payload.get(key, []) + return [] + + # Document ingestion --------------------------------------------- + + def _ingest_documents( + self, pipeline_id: str, documents: List[Dict[str, Any]] + ) -> List[CloudDocument]: + store = self._documents.setdefault(pipeline_id, {}) + results: List[CloudDocument] = [] + for doc_payload in documents: + doc_id = doc_payload.get("id") or self._server.new_id("doc") + doc = CloudDocument( + id=doc_id, + text=doc_payload.get("text", ""), + metadata=doc_payload.get("metadata", {}), + excluded_embed_metadata_keys=doc_payload.get( + "excluded_embed_metadata_keys" + ), + excluded_llm_metadata_keys=doc_payload.get( + "excluded_llm_metadata_keys" + ), + ) + store[doc_id] = doc + results.append(doc) + return results + + def _handle_create_documents(self, request: httpx.Request) -> httpx.Response: + parts = request.url.path.split("/") + pipeline_id = parts[-2] + pipeline = self._pipelines.get(pipeline_id) + if not pipeline: + return self._server.json_response( + {"detail": f"Pipeline {pipeline_id} not found"}, status_code=404 + ) + + payload = self._server.json(request) + documents = self._extract_list(payload, "documents") + results = self._ingest_documents(pipeline_id, documents) + return self._server.json_response( + [d.model_dump() for d in results], status_code=200 + ) + + def _handle_upsert_documents(self, request: httpx.Request) -> httpx.Response: + parts = request.url.path.split("/") + pipeline_id = parts[-2] + pipeline = self._pipelines.get(pipeline_id) + if not pipeline: + return self._server.json_response( + {"detail": f"Pipeline {pipeline_id} not found"}, status_code=404 + ) + + payload = self._server.json(request) + documents = self._extract_list(payload, "documents") + results = self._ingest_documents(pipeline_id, documents) + return self._server.json_response( + [d.model_dump() for d in results], status_code=200 + ) + + # File ingestion ------------------------------------------------- + + def _handle_upsert_files(self, request: httpx.Request) -> httpx.Response: + parts = request.url.path.split("/") + pipeline_id = parts[-2] + pipeline = self._pipelines.get(pipeline_id) + if not pipeline: + return self._server.json_response( + {"detail": f"Pipeline {pipeline_id} not found"}, status_code=404 + ) + + payload = self._server.json(request) + file_ids = self._extract_list(payload, "files") + now = utcnow() + store = self._files.setdefault(pipeline_id, {}) + results: List[PipelineFile] = [] + for entry in file_ids: + # Each entry can be a dict with file_id + optional metadata, or a plain string + if isinstance(entry, dict): + file_id = entry.get("file_id", self._server.new_id("file")) + custom_metadata = entry.get("custom_metadata") + name = entry.get("name") + else: + file_id = str(entry) + custom_metadata = None + name = None + + pf_id = self._server.new_id("pf") + pf = PipelineFile( + id=pf_id, + pipeline_id=pipeline_id, + file_id=file_id, + name=name or f"file-{file_id}", + status="SUCCESS", + project_id=pipeline.project_id, + custom_metadata=custom_metadata, + created_at=now, + updated_at=now, + ) + store[pf_id] = pf + results.append(pf) + return self._server.json_response( + [pf.model_dump() for pf in results], status_code=200 + ) + + # Retrieval ------------------------------------------------------ + + def _handle_retrieve(self, request: httpx.Request) -> httpx.Response: + parts = request.url.path.split("/") + pipeline_id = parts[-2] + pipeline = self._pipelines.get(pipeline_id) + if not pipeline: + return self._server.json_response( + {"detail": f"Pipeline {pipeline_id} not found"}, status_code=404 + ) + + payload = self._server.json(request) + query = payload.get("query", "") + top_k = payload.get("dense_similarity_top_k") or 3 + + nodes = self._build_retrieval_nodes(pipeline_id, query, top_k) + response = PipelineRetrieveResponse( + pipeline_id=pipeline_id, + retrieval_nodes=nodes, + ) + return self._server.json_response(response.model_dump()) + + def _build_retrieval_nodes( + self, pipeline_id: str, query: str, top_k: int + ) -> List[RetrievalNode]: + """Build retrieval nodes from ingested documents and files.""" + chunks: List[_Chunk] = [] + + # Collect chunks from ingested documents + for doc_id, doc in self._documents.get(pipeline_id, {}).items(): + text = doc.text or "" + # Split the document text into paragraph-sized chunks + paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()] + if not paragraphs: + paragraphs = [text] if text else [] + for i, para in enumerate(paragraphs): + chunks.append( + _Chunk( + text=para, + source_id=doc_id, + chunk_index=i, + metadata=dict(doc.metadata) if doc.metadata else {}, + ) + ) + + # Collect chunks from ingested files (generate deterministic text) + for pf_id, pf in self._files.get(pipeline_id, {}).items(): + seed = combined_seed(pipeline_id, pf.file_id or pf_id) + file_text = generate_text_blob(seed, sentences=6) + # Split generated text into sentence-pair chunks + sentences = file_text.split(". ") + for i in range(0, len(sentences), 2): + chunk_text = ". ".join(sentences[i : i + 2]) + if not chunk_text.endswith("."): + chunk_text += "." + metadata: Dict[str, Any] = {"file_name": pf.name or ""} + if pf.file_id: + metadata["file_id"] = pf.file_id + chunks.append( + _Chunk( + text=chunk_text, + source_id=pf.file_id or pf_id, + chunk_index=i // 2, + metadata=metadata, + ) + ) + + if not chunks: + return [] + + # Score chunks deterministically based on query+chunk content + scored: List[tuple[float, _Chunk]] = [] + for chunk in chunks: + seed = combined_seed(query, chunk.text) + # Generate a score between 0.5 and 1.0 + score = 0.5 + (seed % 5000) / 10000.0 + scored.append((score, chunk)) + + scored.sort(key=lambda x: x[0], reverse=True) + selected = scored[:top_k] + + nodes: List[RetrievalNode] = [] + for score, chunk in selected: + node_id = self._server.new_id("node") + text_node = TextNode( + id=node_id, + text=chunk.text, + extra_info=chunk.metadata or None, + start_char_idx=0, + end_char_idx=len(chunk.text), + ) + nodes.append( + RetrievalNode( + node=text_node, + score=round(score, 4), + ) + ) + return nodes + + +class _Chunk: + """Internal helper to represent a text chunk for retrieval.""" + + __slots__ = ("text", "source_id", "chunk_index", "metadata") + + def __init__( + self, + *, + text: str, + source_id: str, + chunk_index: int, + metadata: Optional[Dict[str, Any]] = None, + ) -> None: + self.text = text + self.source_id = source_id + self.chunk_index = chunk_index + self.metadata = metadata or {} diff --git a/src/extraction_review/testing_utils/server.py b/src/extraction_review/testing_utils/server.py index ab4f16a..0b76d62 100644 --- a/src/extraction_review/testing_utils/server.py +++ b/src/extraction_review/testing_utils/server.py @@ -13,6 +13,8 @@ from .classify import FakeClassifyNamespace from .extract import FakeExtractNamespace from .files import FakeFilesNamespace from .parse import FakeParseNamespace +from .pipelines import FakePipelinesNamespace +from .sheets import FakeSheetsNamespace from .split import FakeSplitNamespace Handler = Callable[[httpx.Request], httpx.Response] @@ -42,6 +44,8 @@ class FakeLlamaCloudServer: "classify", "agent_data", "split", + "sheets", + "pipelines", ) self._namespace_names = {name.lower() for name in selected} self._upload_base_url = upload_base_url or self.DEFAULT_UPLOAD_BASE @@ -63,6 +67,12 @@ class FakeLlamaCloudServer: self.classify = FakeClassifyNamespace(server=self, files=self.files) self.agent_data = FakeAgentDataNamespace(server=self) self.split = FakeSplitNamespace(server=self) + self.pipelines = FakePipelinesNamespace(server=self) + self.sheets = FakeSheetsNamespace( + server=self, + files=self.files, + download_base_url=self._download_base_url, + ) # Context management ---------------------------------------------- def install(self) -> "FakeLlamaCloudServer": @@ -181,6 +191,10 @@ class FakeLlamaCloudServer: self.agent_data.register() if "split" in self._namespace_names: self.split.register() + if "pipelines" in self._namespace_names: + self.pipelines.register() + if "sheets" in self._namespace_names: + self.sheets.register() self._registered = True diff --git a/src/extraction_review/testing_utils/sheets.py b/src/extraction_review/testing_utils/sheets.py new file mode 100644 index 0000000..28c927d --- /dev/null +++ b/src/extraction_review/testing_utils/sheets.py @@ -0,0 +1,274 @@ +from __future__ import annotations + +import random +import struct +from dataclasses import dataclass +from typing import TYPE_CHECKING, Dict, List, Optional +from urllib.parse import urlencode + +import httpx +from llama_cloud.types.beta.sheets_job import Region, SheetsJob, WorksheetMetadata +from llama_cloud.types.beta.sheets_parsing_config import SheetsParsingConfig +from llama_cloud.types.presigned_url import PresignedURL + +from ._deterministic import combined_seed, generate_text_blob, utcnow +from .files import FakeFilesNamespace, StoredFile + +if TYPE_CHECKING: + from .server import FakeLlamaCloudServer + + +@dataclass +class SheetsJobRecord: + job: SheetsJob + + +class FakeSheetsNamespace: + def __init__( + self, + *, + server: "FakeLlamaCloudServer", + files: FakeFilesNamespace, + download_base_url: str, + ) -> None: + self._server = server + self._files = files + self._download_base_url = download_base_url.rstrip("/") + self._jobs: Dict[str, SheetsJobRecord] = {} + self._region_content: Dict[str, bytes] = {} + + def register(self) -> None: + server = self._server + server.add_route( + "POST", + "/api/v1/beta/sheets/jobs", + self._handle_create, + namespace="sheets", + ) + server.add_route( + "GET", + "/api/v1/beta/sheets/jobs", + self._handle_list, + namespace="sheets", + ) + server.add_route( + "GET", + "/api/v1/beta/sheets/jobs/{spreadsheet_job_id}", + self._handle_get, + namespace="sheets", + ) + server.add_route( + "DELETE", + "/api/v1/beta/sheets/jobs/{spreadsheet_job_id}", + self._handle_delete, + namespace="sheets", + ) + server.add_route( + "GET", + "/api/v1/beta/sheets/jobs/{spreadsheet_job_id}/regions/{region_id}/result/{region_type}", + self._handle_get_result_table, + namespace="sheets", + ) + server.add_route( + "GET", + "/sheets/{job_id}/{region_id}/{region_type}", + self._handle_presigned_download, + namespace="sheets", + base_urls=[self._download_base_url], + ) + + def _handle_create(self, request: httpx.Request) -> httpx.Response: + payload = self._server.json(request) + file_id = payload.get("file_id", "") + config_raw = payload.get("config") or {} + + stored_file = self._files.get(file_id) if file_id else None + if file_id and not stored_file: + return self._server.json_response( + {"detail": f"File {file_id} not found"}, status_code=404 + ) + + job_id = self._server.new_id("sheets-job") + now = utcnow() + + config = SheetsParsingConfig.model_validate(config_raw) + regions, worksheet_metadata = self._build_results(job_id, config, stored_file) + + job = SheetsJob( + id=job_id, + config=config, + created_at=now.isoformat(), + file_id=file_id, + project_id=request.url.params.get( + "project_id", self._server.default_project_id + ), + status="SUCCESS", + updated_at=now.isoformat(), + user_id="fake-user", + errors=None, + file=None, + regions=regions, + success=True, + worksheet_metadata=worksheet_metadata, + ) + self._jobs[job_id] = SheetsJobRecord(job=job) + return self._server.json_response(job.model_dump()) + + def _handle_list(self, request: httpx.Request) -> httpx.Response: + items = [r.job.model_dump() for r in self._jobs.values()] + return self._server.json_response({"items": items, "next_page_token": None}) + + def _handle_get(self, request: httpx.Request) -> httpx.Response: + job_id = request.url.path.split("/")[-1] + record = self._jobs.get(job_id) + if not record: + return self._server.json_response( + {"detail": "Sheets job not found"}, status_code=404 + ) + return self._server.json_response(record.job.model_dump()) + + def _handle_delete(self, request: httpx.Request) -> httpx.Response: + job_id = request.url.path.split("/")[-1] + self._jobs.pop(job_id, None) + return self._server.json_response({}, status_code=200) + + def _handle_get_result_table(self, request: httpx.Request) -> httpx.Response: + parts = request.url.path.split("/") + # .../jobs/{job_id}/regions/{region_id}/result/{region_type} + job_id = parts[-5] + region_id = parts[-3] + + record = self._jobs.get(job_id) + if not record: + return self._server.json_response( + {"detail": "Sheets job not found"}, status_code=404 + ) + + if record.job.regions: + found = any(r.region_id == region_id for r in record.job.regions) + if not found: + return self._server.json_response( + {"detail": f"Region {region_id} not found"}, status_code=404 + ) + + region_type = parts[-1] + presigned = PresignedURL( + url=( + f"{self._download_base_url}/sheets/{job_id}/{region_id}/{region_type}" + f"?{urlencode({'token': 'fake'})}" + ), + expires_at=utcnow(), + form_fields=None, + ) + return self._server.json_response(presigned.model_dump()) + + def _handle_presigned_download(self, request: httpx.Request) -> httpx.Response: + parts = request.url.path.split("/") + # /sheets/{job_id}/{region_id}/{region_type} + region_id = parts[-2] + content_key = region_id + content = self._region_content.get(content_key) + if content is None: + return httpx.Response(404, json={"detail": "Region content not found"}) + return httpx.Response( + 200, + content=content, + headers={"content-type": "application/octet-stream"}, + ) + + def _build_results( + self, + job_id: str, + config: SheetsParsingConfig, + stored_file: Optional[StoredFile], + ) -> tuple[List[Region], List[WorksheetMetadata]]: + file_hash = stored_file.sha256 if stored_file else "no-file" + seed = combined_seed(file_hash, job_id) + rng = random.Random(seed) + + sheet_names_config = config.sheet_names if config else None + if sheet_names_config: + sheet_names = list(sheet_names_config) + else: + num_sheets = rng.randint(1, 3) + sheet_names = [f"Sheet{i + 1}" for i in range(num_sheets)] + + worksheet_metadata: List[WorksheetMetadata] = [] + for name in sheet_names: + worksheet_metadata.append( + WorksheetMetadata( + sheet_name=name, + title=f"Title for {name}", + description=f"Description for {name}", + ) + ) + + regions: List[Region] = [] + region_types = ["table", "extra"] + for sheet_name in sheet_names: + num_regions = rng.randint(1, 3) + for j in range(num_regions): + region_id = self._server.new_id("region") + rtype = region_types[rng.randint(0, len(region_types) - 1)] + row_start = rng.randint(1, 10) + col_start = chr(ord("A") + rng.randint(0, 5)) + row_end = row_start + rng.randint(3, 20) + col_end = chr(ord(col_start) + rng.randint(1, 5)) + location = f"{col_start}{row_start}:{col_end}{row_end}" + regions.append( + Region( + region_id=region_id, + region_type=rtype, + sheet_name=sheet_name, + location=location, + title=f"Region {j + 1} in {sheet_name}", + description=f"Deterministic region from {sheet_name}", + ) + ) + content_seed = combined_seed(file_hash, job_id, region_id) + self._region_content[region_id] = _build_fake_parquet( + content_seed, sheet_name, location + ) + + return regions, worksheet_metadata + + +def _build_fake_parquet(seed: int, sheet_name: str, location: str) -> bytes: + """Build a minimal Apache Parquet file with deterministic tabular data. + + The file uses the PAR1 magic bytes and contains a simplified but + structurally valid Parquet layout so that downstream consumers can + at minimum verify the magic header. + """ + rng = random.Random(seed) + num_cols = rng.randint(2, 5) + num_rows = rng.randint(3, 10) + + headers = [f"col_{i}" for i in range(num_cols)] + rows: List[List[str]] = [] + for _ in range(num_rows): + row = [ + generate_text_blob(rng.randint(0, 1_000_000), sentences=1)[:30] + for _ in headers + ] + rows.append(row) + + # Encode as minimal Parquet-like binary with correct magic. + # Real Parquet parsers need Thrift metadata; for the mock server we + # embed the data as JSON in the page payload between the PAR1 markers + # so tests can verify content determinism if needed. + import json as _json + + payload = _json.dumps( + { + "sheet_name": sheet_name, + "location": location, + "headers": headers, + "rows": rows, + } + ).encode("utf-8") + + magic = b"PAR1" + # Parquet footer: 4-byte LE footer length + magic + footer_len = struct.pack("= 2 + ids = {f.id for f in items} + assert file_obj_1.id in ids + assert file_obj_2.id in ids + + +@pytest.mark.asyncio +async def test_files_list_by_name(server, tmp_path): + """Test filtering files by name via the list endpoint.""" + test_file = tmp_path / "unique_name.pdf" + test_file.write_bytes(b"unique content") + + # Use preload for reliable filename storage + server.files.preload(path=test_file, filename="unique_name.pdf") + server.files.preload_from_source("other_file.txt", b"other content") + + client = AsyncLlamaCloud(api_key="fake-api-key") + + results = await client.files.list(file_name="unique_name.pdf") + items = [f async for f in results] + assert len(items) == 1 + assert items[0].name == "unique_name.pdf" diff --git a/tests/testing_utils/test_pipelines.py b/tests/testing_utils/test_pipelines.py new file mode 100644 index 0000000..8487ae6 --- /dev/null +++ b/tests/testing_utils/test_pipelines.py @@ -0,0 +1,381 @@ +"""Tests for the FakePipelinesNamespace mock implementation.""" + +import pytest +from extraction_review.testing_utils import FakeLlamaCloudServer +from llama_cloud import APIStatusError, AsyncLlamaCloud + + +@pytest.fixture +def server(): + """Provide a server with the pipelines namespace enabled.""" + with FakeLlamaCloudServer(namespaces=["pipelines"]) as srv: + yield srv + + +@pytest.mark.asyncio +async def test_pipelines_create_and_get(server): + """Verify a pipeline can be created and retrieved.""" + client = AsyncLlamaCloud(api_key="fake-api-key") + + pipeline = await client.pipelines.create(name="test-pipeline") + assert pipeline.id.startswith("pipeline_") + assert pipeline.name == "test-pipeline" + assert pipeline.project_id == server.default_project_id + assert pipeline.status == "CREATED" + + retrieved = await client.pipelines.get(pipeline.id) + assert retrieved.id == pipeline.id + assert retrieved.name == "test-pipeline" + + +@pytest.mark.asyncio +async def test_pipelines_create_with_embedding_config(server): + """Verify a pipeline can be created with explicit embedding config.""" + client = AsyncLlamaCloud(api_key="fake-api-key") + + pipeline = await client.pipelines.create( + name="custom-pipeline", + embedding_config={ + "type": "MANAGED_OPENAI_EMBEDDING", + "component": {}, + }, + pipeline_type="MANAGED", + ) + assert pipeline.name == "custom-pipeline" + assert pipeline.embedding_config is not None + + +@pytest.mark.asyncio +async def test_pipelines_list(server): + """Verify listing pipelines returns created pipelines.""" + client = AsyncLlamaCloud(api_key="fake-api-key") + + await client.pipelines.create(name="pipeline-1") + await client.pipelines.create(name="pipeline-2") + + pipelines = await client.pipelines.list() + assert len(pipelines) == 2 + names = {p.name for p in pipelines} + assert names == {"pipeline-1", "pipeline-2"} + + +@pytest.mark.asyncio +async def test_pipelines_get_not_found(server): + """Verify non-existent pipeline returns 404.""" + client = AsyncLlamaCloud(api_key="fake-api-key") + with pytest.raises(APIStatusError) as exc_info: + await client.pipelines.get("nonexistent-id") + assert exc_info.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_pipelines_delete(server): + """Verify a pipeline can be deleted.""" + client = AsyncLlamaCloud(api_key="fake-api-key") + + pipeline = await client.pipelines.create(name="to-delete") + retrieved = await client.pipelines.get(pipeline.id) + assert retrieved.id == pipeline.id + + await client.pipelines.delete(pipeline.id) + + with pytest.raises(APIStatusError) as exc_info: + await client.pipelines.get(pipeline.id) + assert exc_info.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_pipelines_update(server): + """Verify a pipeline can be updated.""" + client = AsyncLlamaCloud(api_key="fake-api-key") + + pipeline = await client.pipelines.create(name="original-name") + updated = await client.pipelines.update(pipeline.id, name="new-name") + assert updated.name == "new-name" + assert updated.id == pipeline.id + + retrieved = await client.pipelines.get(pipeline.id) + assert retrieved.name == "new-name" + + +@pytest.mark.asyncio +async def test_pipelines_get_status(server): + """Verify pipeline status can be retrieved.""" + client = AsyncLlamaCloud(api_key="fake-api-key") + + pipeline = await client.pipelines.create(name="status-test") + status = await client.pipelines.get_status(pipeline.id) + assert status.status == "SUCCESS" + + +@pytest.mark.asyncio +async def test_pipelines_retrieve(server): + """Verify pipeline retrieve endpoint works.""" + client = AsyncLlamaCloud(api_key="fake-api-key") + + pipeline = await client.pipelines.create(name="retrieve-test") + result = await client.pipelines.retrieve(pipeline.id, query="test query") + assert result.pipeline_id == pipeline.id + assert result.retrieval_nodes is not None + assert isinstance(result.retrieval_nodes, list) + + +# --- Document ingestion tests --- + + +@pytest.mark.asyncio +async def test_pipelines_ingest_documents_and_retrieve(server): + """Ingest documents into a pipeline and verify retrieval returns nodes.""" + client = AsyncLlamaCloud(api_key="fake-api-key") + + pipeline = await client.pipelines.create(name="doc-ingest-test") + + docs = await client.pipelines.documents.create( + pipeline.id, + body=[ + { + "text": "The quick brown fox jumps over the lazy dog.", + "metadata": {"source": "test"}, + }, + { + "text": "Machine learning is a subset of artificial intelligence.", + "metadata": {"source": "ml"}, + }, + ], + ) + assert len(docs) == 2 + assert docs[0].text == "The quick brown fox jumps over the lazy dog." + assert docs[0].metadata["source"] == "test" + assert docs[1].text == "Machine learning is a subset of artificial intelligence." + + result = await client.pipelines.retrieve(pipeline.id, query="fox") + assert result.pipeline_id == pipeline.id + assert len(result.retrieval_nodes) > 0 + + # Nodes should have text content from our documents + texts = [n.node.text for n in result.retrieval_nodes] + all_doc_texts = { + "The quick brown fox jumps over the lazy dog.", + "Machine learning is a subset of artificial intelligence.", + } + for text in texts: + assert text in all_doc_texts + + # Each node should have a score + for node in result.retrieval_nodes: + assert node.score is not None + assert 0.0 <= node.score <= 1.0 + + +@pytest.mark.asyncio +async def test_pipelines_upsert_documents(server): + """Upsert documents (PUT) into a pipeline and verify they are stored.""" + client = AsyncLlamaCloud(api_key="fake-api-key") + + pipeline = await client.pipelines.create(name="doc-upsert-test") + + docs = await client.pipelines.documents.upsert( + pipeline.id, + body=[ + {"text": "First document content.", "metadata": {"order": "1"}}, + ], + ) + assert len(docs) == 1 + assert docs[0].text == "First document content." + + # Upsert more documents + docs2 = await client.pipelines.documents.upsert( + pipeline.id, + body=[ + {"text": "Second document content.", "metadata": {"order": "2"}}, + ], + ) + assert len(docs2) == 1 + + # Retrieve should pick up both documents + result = await client.pipelines.retrieve(pipeline.id, query="document") + assert len(result.retrieval_nodes) == 2 + + +@pytest.mark.asyncio +async def test_pipelines_ingest_documents_with_custom_id(server): + """Documents with explicit IDs preserve them.""" + client = AsyncLlamaCloud(api_key="fake-api-key") + + pipeline = await client.pipelines.create(name="custom-id-test") + + docs = await client.pipelines.documents.create( + pipeline.id, + body=[ + {"id": "my-doc-1", "text": "Custom ID document.", "metadata": {}}, + ], + ) + assert len(docs) == 1 + assert docs[0].id == "my-doc-1" + + +# --- File ingestion tests --- + + +@pytest.mark.asyncio +async def test_pipelines_ingest_files_and_retrieve(server): + """Add files to a pipeline and verify retrieval returns generated nodes.""" + client = AsyncLlamaCloud(api_key="fake-api-key") + + pipeline = await client.pipelines.create(name="file-ingest-test") + + files = await client.pipelines.files.create( + pipeline.id, + body=[ + {"file_id": "file-abc123"}, + {"file_id": "file-def456"}, + ], + ) + assert len(files) == 2 + assert files[0].pipeline_id == pipeline.id + assert files[0].status == "SUCCESS" + assert files[1].file_id == "file-def456" + + result = await client.pipelines.retrieve(pipeline.id, query="search query") + assert result.pipeline_id == pipeline.id + assert len(result.retrieval_nodes) > 0 + + # Nodes from files should have file metadata + for node in result.retrieval_nodes: + assert node.node.text # Should have generated text + assert node.score is not None + + +# --- Retrieval behavior tests --- + + +@pytest.mark.asyncio +async def test_pipelines_retrieve_empty_pipeline(server): + """Retrieve on an empty pipeline returns no nodes.""" + client = AsyncLlamaCloud(api_key="fake-api-key") + + pipeline = await client.pipelines.create(name="empty-pipeline") + result = await client.pipelines.retrieve(pipeline.id, query="anything") + assert result.retrieval_nodes == [] + + +@pytest.mark.asyncio +async def test_pipelines_retrieve_respects_top_k(server): + """Retrieve respects the dense_similarity_top_k parameter.""" + client = AsyncLlamaCloud(api_key="fake-api-key") + + pipeline = await client.pipelines.create(name="topk-test") + + # Ingest several documents to have more chunks than top_k + await client.pipelines.documents.create( + pipeline.id, + body=[ + { + "text": f"Document number {i} with unique content about topic {i}.", + "metadata": {}, + } + for i in range(10) + ], + ) + + result = await client.pipelines.retrieve( + pipeline.id, + query="topic", + dense_similarity_top_k=2, + ) + assert len(result.retrieval_nodes) == 2 + + +@pytest.mark.asyncio +async def test_pipelines_retrieve_deterministic(server): + """Same query on same data produces the same results.""" + client = AsyncLlamaCloud(api_key="fake-api-key") + + pipeline = await client.pipelines.create(name="deterministic-test") + await client.pipelines.documents.create( + pipeline.id, + body=[ + {"text": "Alpha bravo charlie delta.", "metadata": {}}, + {"text": "Echo foxtrot golf hotel.", "metadata": {}}, + ], + ) + + result1 = await client.pipelines.retrieve(pipeline.id, query="bravo") + result2 = await client.pipelines.retrieve(pipeline.id, query="bravo") + + texts1 = [n.node.text for n in result1.retrieval_nodes] + texts2 = [n.node.text for n in result2.retrieval_nodes] + assert texts1 == texts2 + + scores1 = [n.score for n in result1.retrieval_nodes] + scores2 = [n.score for n in result2.retrieval_nodes] + assert scores1 == scores2 + + +@pytest.mark.asyncio +async def test_pipelines_delete_cleans_up_documents_and_files(server): + """Deleting a pipeline clears its ingested documents and files.""" + client = AsyncLlamaCloud(api_key="fake-api-key") + + pipeline = await client.pipelines.create(name="cleanup-test") + + await client.pipelines.documents.create( + pipeline.id, + body=[{"text": "Some content.", "metadata": {}}], + ) + await client.pipelines.files.create( + pipeline.id, + body=[{"file_id": "file-xyz"}], + ) + + # Verify data exists + result = await client.pipelines.retrieve(pipeline.id, query="content") + assert len(result.retrieval_nodes) > 0 + + # Delete pipeline + await client.pipelines.delete(pipeline.id) + + # Internal stores should be cleaned + assert pipeline.id not in server.pipelines._documents + assert pipeline.id not in server.pipelines._files + + +@pytest.mark.asyncio +async def test_pipelines_mixed_documents_and_files_retrieval(server): + """Retrieval combines results from both documents and files.""" + client = AsyncLlamaCloud(api_key="fake-api-key") + + pipeline = await client.pipelines.create(name="mixed-test") + + await client.pipelines.documents.create( + pipeline.id, + body=[ + { + "text": "Document text about important concepts.", + "metadata": {"type": "doc"}, + } + ], + ) + await client.pipelines.files.create( + pipeline.id, + body=[{"file_id": "file-mixed-001"}], + ) + + result = await client.pipelines.retrieve( + pipeline.id, + query="concepts", + dense_similarity_top_k=10, + ) + assert len(result.retrieval_nodes) > 1 + + # Should have nodes from both sources + has_doc_node = any( + n.node.text == "Document text about important concepts." + for n in result.retrieval_nodes + ) + has_file_node = any( + n.node.extra_info and n.node.extra_info.get("file_id") == "file-mixed-001" + for n in result.retrieval_nodes + ) + assert has_doc_node, "Should have a node from the ingested document" + assert has_file_node, "Should have a node from the ingested file" diff --git a/tests/testing_utils/test_sheets.py b/tests/testing_utils/test_sheets.py new file mode 100644 index 0000000..9b0542c --- /dev/null +++ b/tests/testing_utils/test_sheets.py @@ -0,0 +1,209 @@ +"""Tests for the FakeSheetsNamespace mock implementation.""" + +import httpx +import pytest +from extraction_review.testing_utils import FakeLlamaCloudServer +from llama_cloud import APIStatusError, AsyncLlamaCloud + + +@pytest.fixture +def server(): + """Provide a server with files and sheets namespaces enabled.""" + with FakeLlamaCloudServer(namespaces=["files", "sheets"]) as srv: + yield srv + + +@pytest.mark.asyncio +async def test_sheets_create_and_get(server, tmp_path): + """Verify a sheets job can be created and retrieved.""" + test_file = tmp_path / "spreadsheet.xlsx" + test_file.write_bytes(b"fake spreadsheet content") + file_id = server.files.preload(path=test_file) + + client = AsyncLlamaCloud(api_key="fake-api-key") + + job = await client.beta.sheets.create(file_id=file_id) + assert job.id.startswith("sheets-job_") + assert job.status == "SUCCESS" + assert job.success is True + assert job.regions is not None + assert len(job.regions) > 0 + + for region in job.regions: + assert region.region_id is not None + assert region.region_type in ("table", "extra") + assert region.sheet_name is not None + assert region.location is not None + + assert job.worksheet_metadata is not None + assert len(job.worksheet_metadata) > 0 + + # Retrieve the same job + retrieved = await client.beta.sheets.get(job.id) + assert retrieved.id == job.id + assert retrieved.status == "SUCCESS" + + +@pytest.mark.asyncio +async def test_sheets_file_not_found(server): + """Verify sheets with non-existent file returns 404.""" + client = AsyncLlamaCloud(api_key="fake-api-key") + with pytest.raises(APIStatusError) as exc_info: + await client.beta.sheets.create(file_id="nonexistent-file") + assert exc_info.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_sheets_job_not_found(server): + """Verify non-existent sheets job returns 404.""" + client = AsyncLlamaCloud(api_key="fake-api-key") + with pytest.raises(APIStatusError) as exc_info: + await client.beta.sheets.get("nonexistent-id") + assert exc_info.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_sheets_delete_job(server, tmp_path): + """Verify a sheets job can be deleted.""" + test_file = tmp_path / "spreadsheet.xlsx" + test_file.write_bytes(b"delete me") + file_id = server.files.preload(path=test_file) + + client = AsyncLlamaCloud(api_key="fake-api-key") + job = await client.beta.sheets.create(file_id=file_id) + + # Job should exist + retrieved = await client.beta.sheets.get(job.id) + assert retrieved.id == job.id + + # Delete + await client.beta.sheets.delete_job(job.id) + + # Should no longer exist + with pytest.raises(APIStatusError) as exc_info: + await client.beta.sheets.get(job.id) + assert exc_info.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_sheets_get_result_table(server, tmp_path): + """Verify presigned URL generation for region results.""" + test_file = tmp_path / "spreadsheet.xlsx" + test_file.write_bytes(b"spreadsheet with tables") + file_id = server.files.preload(path=test_file) + + client = AsyncLlamaCloud(api_key="fake-api-key") + job = await client.beta.sheets.create(file_id=file_id) + + assert job.regions is not None + assert len(job.regions) > 0 + + region = job.regions[0] + presigned = await client.beta.sheets.get_result_table( + "table", + spreadsheet_job_id=job.id, + region_id=region.region_id, + ) + assert presigned.url is not None + assert "token=fake" in presigned.url + assert job.id in presigned.url + + +@pytest.mark.asyncio +async def test_sheets_with_config(server, tmp_path): + """Verify sheets job with custom config.""" + test_file = tmp_path / "spreadsheet.xlsx" + test_file.write_bytes(b"configured spreadsheet") + file_id = server.files.preload(path=test_file) + + client = AsyncLlamaCloud(api_key="fake-api-key") + + job = await client.beta.sheets.create( + file_id=file_id, + config={ + "sheet_names": ["Revenue", "Expenses"], + "flatten_hierarchical_tables": True, + "generate_additional_metadata": True, + }, + ) + assert job.status == "SUCCESS" + assert job.worksheet_metadata is not None + + sheet_names = {ws.sheet_name for ws in job.worksheet_metadata} + assert sheet_names == {"Revenue", "Expenses"} + + if job.regions: + region_sheets = {r.sheet_name for r in job.regions} + assert region_sheets.issubset({"Revenue", "Expenses"}) + + +@pytest.mark.asyncio +async def test_sheets_list_jobs(server, tmp_path): + """Verify listing sheets jobs returns created jobs.""" + test_file = tmp_path / "spreadsheet.xlsx" + test_file.write_bytes(b"list test") + file_id = server.files.preload(path=test_file) + + client = AsyncLlamaCloud(api_key="fake-api-key") + + await client.beta.sheets.create(file_id=file_id) + await client.beta.sheets.create(file_id=file_id) + + jobs = await client.beta.sheets.list() + items = [j async for j in jobs] + assert len(items) == 2 + + +@pytest.mark.asyncio +async def test_sheets_presigned_download(server, tmp_path): + """Verify that the presigned URL for a region result serves parquet content.""" + test_file = tmp_path / "spreadsheet.xlsx" + test_file.write_bytes(b"spreadsheet for download test") + file_id = server.files.preload(path=test_file) + + client = AsyncLlamaCloud(api_key="fake-api-key") + job = await client.beta.sheets.create(file_id=file_id) + + assert job.regions is not None + assert len(job.regions) > 0 + + region = job.regions[0] + presigned = await client.beta.sheets.get_result_table( + region.region_type, + spreadsheet_job_id=job.id, + region_id=region.region_id, + ) + + # Follow the presigned URL to download the parquet content + async with httpx.AsyncClient() as http: + resp = await http.get(presigned.url) + assert resp.status_code == 200 + content = resp.content + # Verify PAR1 magic bytes at start and end + assert content[:4] == b"PAR1" + assert content[-4:] == b"PAR1" + assert len(content) > 8 # has actual payload between magic markers + + +@pytest.mark.asyncio +async def test_sheets_presigned_download_deterministic(server, tmp_path): + """Verify that repeated downloads for the same region return identical content.""" + test_file = tmp_path / "spreadsheet.xlsx" + test_file.write_bytes(b"deterministic download test") + file_id = server.files.preload(path=test_file) + + client = AsyncLlamaCloud(api_key="fake-api-key") + job = await client.beta.sheets.create(file_id=file_id) + + assert job.regions is not None + region = job.regions[0] + presigned = await client.beta.sheets.get_result_table( + region.region_type, + spreadsheet_job_id=job.id, + region_id=region.region_id, + ) + + async with httpx.AsyncClient() as http: + resp1 = await http.get(presigned.url) + resp2 = await http.get(presigned.url) + assert resp1.content == resp2.content