mirror of
https://github.com/run-llama/template-workflow-extract-basic.git
synced 2026-07-19 18:53:47 -04:00
feat: extend mock server with files list, split index, and sheets APIs (#197)
This commit is contained in:
Adrian Lyjak
@@ -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)
|
||||
|
||||
@@ -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 {}
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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("<I", len(payload))
|
||||
return magic + payload + footer_len + magic
|
||||
@@ -115,3 +115,48 @@ async def test_files_query_by_id(server, tmp_path):
|
||||
assert len(response.items) == 1
|
||||
assert response.total_size == 1
|
||||
assert response.items[0].id == file_obj_1.id
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_files_list(server, tmp_path):
|
||||
"""Test that you can list files using the GET /files endpoint."""
|
||||
test_file_1 = tmp_path / "file_a.txt"
|
||||
test_file_1.write_bytes(b"content a")
|
||||
test_file_2 = tmp_path / "file_b.txt"
|
||||
test_file_2.write_bytes(b"content b")
|
||||
|
||||
client = AsyncLlamaCloud(api_key="fake-api-key")
|
||||
file_obj_1 = await client.files.create(
|
||||
file=test_file_1,
|
||||
purpose="extract",
|
||||
)
|
||||
file_obj_2 = await client.files.create(
|
||||
file=test_file_2,
|
||||
purpose="extract",
|
||||
)
|
||||
|
||||
# List all files
|
||||
all_files = await client.files.list()
|
||||
items = [f async for f in all_files]
|
||||
assert len(items) >= 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"
|
||||
|
||||
@@ -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"
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user