diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index d8d9dbb38..6b349a09b 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -31,5 +31,5 @@ jobs: cache: pip - run: pip install -r requirements.txt pytest - if: matrix.agent-frameworks == 'with' - run: pip install openai-agents claude-agent-sdk + run: pip install openai-agents claude-agent-sdk anthropic - run: python -m pytest -q diff --git a/.gitignore b/.gitignore index 5193735ca..b5c223b31 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,4 @@ __pycache__ logs/ .pageindex/ dist/ +*.doc_id diff --git a/README.md b/README.md index 5ce0ca5e6..27e3084a4 100644 --- a/README.md +++ b/README.md @@ -173,9 +173,10 @@ python3 run_pageindex.py --pdf_path /path/to/your/document.pdf
Optional parameters
-You can customize the processing with additional optional arguments: +You can customize the processing with additional optional arguments (the structure-tuning flags below require --mode standard): ``` +--mode Processing mode: flash (default) or standard --model LLM model to use (default: gpt-4o-2024-11-20) --toc-check-pages Pages to check for table of contents (default: 20) --max-pages-per-node Max pages per node (default: 10) @@ -199,13 +200,13 @@ python3 run_pageindex.py --md_path /path/to/your/document.md
> ### ⚡ PageIndex Flash *(preview)* -> **PageIndex Flash** ([`pageindex/flash`](pageindex/flash)) generates tree structures from PDFs in seconds. Structure extraction is purely heuristic-based, no LLM needed. LLM is only used to generate node summaries. +> **PageIndex Flash** ([`pageindex/flash`](pageindex/flash)) generates tree structures from PDFs in seconds. Structure extraction is purely heuristic-based, no LLM needed. An LLM is used only for node summaries and the optimization's expansion pass. > > ```bash -> python3 run_pageindex.py --flash --pdf_path /path/to/your/document.pdf +> python3 run_pageindex.py --mode flash --pdf_path /path/to/your/document.pdf > ``` > -> Add `--optimize` to refine the tree structure for more efficient retrieval (with an LLM expansion pass). +> Tree optimization for retrieval (a deterministic merge, then an LLM expansion pass) is on by default; pass `--optimize off` to disable. ## 🚀 Agentic Vectorless RAG: An Example diff --git a/examples/agentic_vectorless_rag_demo.py b/examples/agentic_vectorless_rag_demo.py index 4fe5f179f..ac0db360a 100644 --- a/examples/agentic_vectorless_rag_demo.py +++ b/examples/agentic_vectorless_rag_demo.py @@ -6,20 +6,21 @@ chunking, PageIndex builds a hierarchical tree index and uses agentic LLM reasoning for human-like, context-aware retrieval. -Agent tools: - - get_document() — document metadata (status, page count, etc.) - - get_document_structure() — tree structure index of a document - - get_page_content() — retrieve text content of specific pages +The agent tools come straight from the SDK — ``client.as_openai_tools()`` +exposes the PageIndex tool contract (browse_documents, get_document, +get_document_structure, get_page_content) and ``client.agent_instructions()`` +provides the retrieval playbook, so the whole agent is a few lines. Swap +``PageIndexLocalClient()`` for ``PageIndexCloudClient(api_key=...)`` and the +same code runs against the cloud. Steps: 1 — Index a PDF locally and view its tree structure index 2 — View document metadata 3 — Ask a question (agent reasons over the index and auto-calls tools) -Requirements: pip install openai-agents; OPENAI_API_KEY in the environment. +Requirements: pip install "pageindex[openai]"; OPENAI_API_KEY in the environment. """ import sys -import json import asyncio import concurrent.futures from pathlib import Path @@ -27,62 +28,30 @@ sys.path.insert(0, str(Path(__file__).parent.parent)) -from agents import Agent, Runner, function_tool, set_tracing_disabled -from agents.model_settings import ModelSettings +from agents import Agent, Runner, set_tracing_disabled from agents.stream_events import RawResponsesStreamEvent, RunItemStreamEvent from openai.types.responses import ResponseTextDeltaEvent, ResponseReasoningSummaryTextDeltaEvent -from pageindex import PageIndexClient +from pageindex import PageIndexAPIError, PageIndexLocalClient import pageindex.utils as utils PDF_URL = "https://arxiv.org/pdf/2603.15031" _EXAMPLES_DIR = Path(__file__).parent PDF_PATH = _EXAMPLES_DIR / "documents" / "attention-residuals.pdf" +DOC_ID_PATH = _EXAMPLES_DIR / "documents" / "attention-residuals.doc_id" STORAGE_PATH = _EXAMPLES_DIR / ".pageindex" -AGENT_SYSTEM_PROMPT = """ -You are PageIndex, a document QA assistant. -TOOL USE: -- Call get_document() first to confirm status and page count. -- Call get_document_structure() to identify relevant page ranges. -- Call get_page_content(pages="5-7") with tight ranges; never fetch the whole document. -- Before each tool call, output one short sentence explaining the reason. -Answer based only on tool output. Be concise. -""" - -def query_agent(client: PageIndexClient, doc_id: str, prompt: str, verbose: bool = False) -> str: +def query_agent(client: PageIndexLocalClient, doc_id: str, prompt: str, verbose: bool = False) -> str: """Run a document QA agent using the OpenAI Agents SDK. Streams text output token-by-token and returns the full answer string. Tool calls are always printed; verbose=True also prints arguments and output previews. """ - - @function_tool - def get_document() -> str: - """Get document metadata: status, page count, name, and description.""" - return json.dumps(client.get_document(doc_id)) - - @function_tool - def get_document_structure() -> str: - """Get the document's full tree structure (without text) to find relevant sections.""" - return json.dumps(client.get_document_structure(doc_id), ensure_ascii=False) - - @function_tool - def get_page_content(pages: str) -> str: - """ - Get the text content of specific pages. - Use tight ranges: e.g. '5-7' for pages 5 to 7, '3,8' for pages 3 and 8, '12' for page 12. - """ - return json.dumps(client.get_page_content(doc_id, pages), ensure_ascii=False) - agent = Agent( - name="PageIndex", - instructions=AGENT_SYSTEM_PROMPT, - tools=[get_document, get_document_structure, get_page_content], - model=getattr(client, "retrieve_model", None), - # model_settings=ModelSettings(reasoning={"effort": "low", "summary": "auto"}), # Uncomment to enable reasoning + **client.openai_agent_config(doc_id=doc_id), + # model_settings=ModelSettings(reasoning={"effort": "low", "summary": "auto"}), # from agents.model_settings import ModelSettings ) async def _run(): @@ -152,21 +121,33 @@ async def _run(): print("Download complete.\n") # Setup: local mode — no PageIndex API key needed, your LLM key does the work - client = PageIndexClient(storage_path=str(STORAGE_PATH)) + client = PageIndexLocalClient(storage_path=str(STORAGE_PATH)) # Step 1: Index PDF and view tree structure print("=" * 60) print("Step 1: Index PDF and view tree structure") print("=" * 60) - doc_id = next( - (doc["id"] for doc in client.list_documents(limit=100)["documents"] - if doc["name"] == PDF_PATH.name), - None, - ) + doc_id = None + if DOC_ID_PATH.exists(): + cached = DOC_ID_PATH.read_text().strip() + try: + client.get_document(cached) + doc_id = cached + except PageIndexAPIError: + DOC_ID_PATH.unlink() + if doc_id is None: + # The .doc_id cache is gitignored — on a fresh clone with an + # existing store, find the already-indexed copy by name instead of + # re-indexing it. + doc_id = next( + (doc["id"] for doc in client.list_documents(limit=100)["documents"] + if doc["name"] == PDF_PATH.name), None) if doc_id: + DOC_ID_PATH.write_text(doc_id) print(f"\nLoaded cached doc_id: {doc_id}") else: - doc_id = client.submit_document(str(PDF_PATH))["doc_id"] + doc_id = client.submit_document(str(PDF_PATH), wait=True)["doc_id"] + DOC_ID_PATH.write_text(doc_id) print(f"\nIndexed. doc_id: {doc_id}") print("\nTree Structure (top-level sections):") structure = client.get_tree(doc_id, node_summary=True)["result"] diff --git a/pageindex/__init__.py b/pageindex/__init__.py index 3513668a2..8a2383014 100644 --- a/pageindex/__init__.py +++ b/pageindex/__init__.py @@ -18,26 +18,36 @@ ] _LAZY = { + "page_index": ".page_index_classic", + "page_index_main": ".page_index_classic", "page_index_flash": ".flash", "optimize_tree": ".tree_optimize", "md_to_tree": ".page_index_md", } -_SUBMODULES = {"client", "cloud_api", "errors", "flash", "local_api", - "local_store", "page_index_classic", "page_index_md", "tree_optimize", - "utils"} - +_SUBMODULES = {"agent_tools", "client", "cloud_api", "errors", "flash", + "integrations", "local_api", "local_chat", "local_store", + "mcp_bridge", "page_index_classic", "page_index_md", + "tree_optimize", "utils"} def __getattr__(name): if name.startswith("_"): + # Dunder probes (copy, pickle, inspect) are the frequent unknown + # names — they must not trigger the classic import below. raise AttributeError(f"module {__name__!r} has no attribute {name!r}") import importlib if name in _SUBMODULES: return importlib.import_module(f".{name}", __name__) - module = importlib.import_module(_LAZY.get(name, ".page_index_classic"), __name__) + # Pre-0.2.10 compat: unknown names fall through to the classic module, + # whose public surface (ConfigLoader, count_tokens, ...) resolved as + # package attributes. A non-underscore typo pays one classic import + # before its AttributeError — not worth an allowlist. + module = importlib.import_module(_LAZY.get(name, ".page_index_classic"), + __name__) try: value = getattr(module, name) except AttributeError: - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None + raise AttributeError( + f"module {__name__!r} has no attribute {name!r}") from None globals()[name] = value return value diff --git a/pageindex/_version.py b/pageindex/_version.py new file mode 100644 index 000000000..da5c00c2c --- /dev/null +++ b/pageindex/_version.py @@ -0,0 +1,10 @@ +"""Installed-package version, shared by every surface that reports it upstream.""" +from __future__ import annotations + + +def sdk_version() -> str: + try: + from importlib.metadata import version + return version("pageindex") + except Exception: + return "0.0.0" diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py new file mode 100644 index 000000000..e00eb3ca9 --- /dev/null +++ b/pageindex/agent_tools.py @@ -0,0 +1,1633 @@ +"""Agent tools: the cloud MCP tool contract, executed against a PageIndexClient. + +Tool names and the surviving input-schema structure match the PageIndex +cloud MCP server — the local surface hides the documented cloud-only +parameters — so agent prompts port across the cloud MCP connection and +this in-process layer. Only the tools that exist in every mode are +registered (no folders, search_documents, or get_document_image), and the +guidance strings (tool descriptions) adapt to the local surface the same +way the agent instructions do — they never teach capabilities that only +exist on the cloud. + +Tools never raise for any invocation their signatures accept: every +outcome, including errors, is returned as the same JSON envelope the cloud +emits ({"success": true, ...} / {"error": ...}). Arguments outside a pruned +local signature fail at the Python call boundary; the call_tool path +answers them with the guided error envelope instead. +""" +from __future__ import annotations + +import copy +import difflib +import inspect +import json +import re +import threading +import time +import weakref +from typing import Any, Callable, Optional + +from .errors import PageIndexAPIError + +TOOL_RESPONSE_CHAR_LIMIT = 100_000 +STRUCTURE_FIRST_PAGE_THRESHOLD = 20 + +_CHAR_BUDGET = int(TOOL_RESPONSE_CHAR_LIMIT * 0.95) +_PAGES_SPEC_RE = re.compile(r"^(\d+(-\d+)?)(,\s*\d+(-\d+)?)*$") +_MAX_REQUESTED_PAGES = 10_000 +_SIMILAR_NAMES_LIMIT = 3 +_TOOL_WAIT_TIMEOUT = 180.0 # "up to 3 minutes", per the wait_for_completion schema +_TOOL_WAIT_INTERVAL = 5.0 + +_DOC_NAME_DESCRIPTION = ( + 'Copy the `name` field verbatim from a browse_documents() or ' + 'search_documents() response (case-sensitive, include extension). ' + 'Example: "Q3 Report.pdf". If the response shows two documents with the ' + 'same name, pass `folder_id` alongside to disambiguate.' +) +_FOLDER_ID_DISAMBIGUATOR_DESCRIPTION = ( + 'Disambiguator for same-name documents. Copy the `folder_id` from the ' + 'intended browse/search result; use "root" for root-level documents, or ' + '"shared-with-me"/"following" for the read-only folders at the library ' + 'root; omit if `doc_name` is unique. Copy any folder_id verbatim from a ' + 'browse_documents()/get_folder_structure() response, never construct one.' +) +_WAIT_FOR_COMPLETION_DESCRIPTION = ( + "If true and document is processing, automatically wait up to 3 minutes " + "until completed. Reduces repeated tool calls." +) + +#: Tool names, descriptions, and parameter schemas, identical to the cloud +#: MCP server's tools/list. +TOOL_CONTRACT: dict[str, dict[str, Any]] = { + "browse_documents": { + "annotations": {"readOnlyHint": True, "openWorldHint": False}, + "description": ( + "Primary document retrieval tool. After orienting with " + "get_folder_structure() (when available), use this for all " + "document-related questions. The bare call returns root-level " + "sub-folders and documents; pass folder_id to drill into a " + 'sub-folder level by level. Use sort="relevance" + query for ' + "semantic ranking. Do NOT jump to search_documents() first — it " + "is an escalation path, only after " + 'browse_documents(sort="relevance") has failed.' + ), + "schema": { + "type": "object", + "properties": { + "folder_id": { + "type": "string", + "default": "root", + "description": ( + 'Folder scope (default "root"). Pass a specific folder ' + 'ID to scope into that folder, or "root" to reference ' + "the library root. The read-only \"shared-with-me\" and " + '"following" folders live at the library root — pass ' + "one of those ids to browse them. Copy any folder_id " + "verbatim from a browse/tree response, never construct " + "one. Combine with `recursive` to control breadth." + ), + }, + "recursive": { + "type": "boolean", + "default": False, + "description": ( + "Whether to include documents from descendant folders. " + "When false (default), returns the direct contents of " + "folder_id along with its sub-folders — prefer this for " + "level-by-level exploration so you retain folder " + "hierarchy context. When true, flattens all descendant " + "documents into one list and omits sub-folders — use " + "only when a non-recursive browse of the target folder " + "returned no relevant results and you need to widen the " + "scope, or the user explicitly requests a flat listing." + ), + }, + "sort": { + "type": "string", + "enum": ["time", "relevance"], + "default": "time", + "description": ( + 'Sort order. "time" (default) sorts by upload date ' + '(newest first); "relevance" orders documents by ' + "semantic relevance to `query`. Relevance also works " + "inside the read-only shared folders — pass their " + "folder_id — but at the library root it ranks only " + "your own documents." + ), + }, + "query": { + "type": "string", + "description": ( + "Search query for relevance ranking. Required when " + 'sort="relevance"; must be omitted when sort="time".' + ), + }, + "offset": { + "type": "integer", + "minimum": 0, + "maximum": 9007199254740991, + "default": 0, + "description": ( + "Zero-based pagination offset. Pass the value of " + "`next_offset` from the previous response to fetch the " + "next page." + ), + }, + "limit": { + "type": "number", + "minimum": 1, + "maximum": 50, + "default": 10, + "description": ( + "Number of documents to return per page (1-50, " + "default 10)" + ), + }, + }, + "required": [], + }, + }, + "get_document": { + "annotations": {"readOnlyHint": True, "openWorldHint": False}, + "description": ( + "Check a document's processing status and metadata. `status` is " + 'one of "pending", "queued", "processing", "completed", or ' + '"failed" — call this before `get_document_structure()` or ' + "`get_page_content()` to confirm the document is ready." + ), + "schema": { + "type": "object", + "properties": { + "doc_name": { + "type": "string", + "minLength": 1, + "description": _DOC_NAME_DESCRIPTION, + }, + "folder_id": { + "anyOf": [{"type": "string"}, {"type": "null"}], + "description": _FOLDER_ID_DISAMBIGUATOR_DESCRIPTION, + }, + "wait_for_completion": { + "type": "boolean", + "default": False, + "description": _WAIT_FOR_COMPLETION_DESCRIPTION, + }, + }, + "required": ["doc_name"], + }, + }, + "get_document_structure": { + "annotations": {"readOnlyHint": True, "openWorldHint": False}, + "description": ( + "Extract a document's hierarchical outline (headers, sections, " + f"page references). REQUIRED for documents over " + f"{STRUCTURE_FIRST_PAGE_THRESHOLD} pages — call this first to " + "locate relevant sections, then pass their page numbers to " + "`get_page_content()`. Use the `part` parameter to iterate large " + "outlines until `pagination.has_more` is false." + ), + "schema": { + "type": "object", + "properties": { + "doc_name": { + "type": "string", + "minLength": 1, + "description": _DOC_NAME_DESCRIPTION, + }, + "folder_id": { + "anyOf": [{"type": "string"}, {"type": "null"}], + "description": _FOLDER_ID_DISAMBIGUATOR_DESCRIPTION, + }, + "part": { + "type": "integer", + "minimum": 1, + "maximum": 9007199254740991, + "default": 1, + "description": ( + "Part number for pagination (1-based, default 1). For " + "large outlines, increment until the response's " + "`pagination.has_more` becomes false." + ), + }, + "wait_for_completion": { + "type": "boolean", + "default": False, + "description": _WAIT_FOR_COMPLETION_DESCRIPTION, + }, + }, + "required": ["doc_name"], + }, + }, + "get_page_content": { + "annotations": {"readOnlyHint": True, "openWorldHint": False}, + "description": ( + "Extract page content from a processed document. Use tight, " + "targeted page ranges — never the whole document at once. For " + f"documents over {STRUCTURE_FIRST_PAGE_THRESHOLD} pages, call " + "`get_document_structure()` first to pick relevant sections. " + "Embedded image paths in the response feed into " + "`get_document_image()`." + ), + "schema": { + "type": "object", + "properties": { + "doc_name": { + "type": "string", + "minLength": 1, + "description": _DOC_NAME_DESCRIPTION, + }, + "folder_id": { + "anyOf": [{"type": "string"}, {"type": "null"}], + "description": _FOLDER_ID_DISAMBIGUATOR_DESCRIPTION, + }, + "pages": { + "type": "string", + "minLength": 1, + "pattern": r"^(\d+(-\d+)?)(,\s*\d+(-\d+)?)*$", + "description": ( + 'Page specification: "5", "3,7,10", "5-10", or ' + '"1-3,7,9-12"' + ), + }, + "wait_for_completion": { + "type": "boolean", + "default": False, + "description": _WAIT_FOR_COMPLETION_DESCRIPTION, + }, + }, + "required": ["doc_name", "pages"], + }, + }, + "remove_document": { + "annotations": {"readOnlyHint": False, "destructiveHint": True, + "idempotentHint": True, "openWorldHint": False}, + "description": ( + "Permanently delete documents and all associated data. Only invoke " + "when the user explicitly names the documents AND confirms " + "deletion. Returns `results` — one entry per requested document: " + '`{ doc_name, status: "deleted" | "not_found" | "failed", ' + "error? }`. Inspect each entry for per-document failures. This " + "action is irreversible." + ), + "schema": { + "type": "object", + "properties": { + "doc_names": { + "type": "array", + "items": {"type": "string", "minLength": 1}, + "minItems": 1, + "maxItems": 10, + "description": ( + "Array of document names to delete. Each name must be " + "copied verbatim from the `name` field of a " + "browse_documents() or search_documents() response " + "(case-sensitive, include extension). Example: " + '["Q3 Report.pdf", "draft.pdf"]. Max 10 per call.' + ), + }, + "folder_id": { + "anyOf": [{"type": "string"}, {"type": "null"}], + "description": _FOLDER_ID_DISAMBIGUATOR_DESCRIPTION, + }, + }, + "required": ["doc_names"], + }, + }, +} + +_READ_TOOLS = ("browse_documents", "get_document", "get_document_structure", + "get_page_content") +_MANAGEMENT_TOOLS = ("remove_document",) + + +# ── response envelopes ── + +_ToolResult = tuple[dict, bool] + + +def _success(data: dict[str, Any], next_steps: dict[str, Any]) -> tuple[dict, bool]: + return {"success": True, **data, "next_steps": next_steps}, False + + +def _failure(error: str, details: Optional[dict[str, Any]], + next_steps: dict[str, Any], error_code: Optional[str] = None, + ) -> tuple[dict, bool]: + payload: dict[str, Any] = {"error": error} + if error_code: + payload["errorCode"] = error_code + if details: + payload.update(details) + payload["next_steps"] = next_steps + return payload, True + + +def _dumps(payload: dict[str, Any]) -> str: + return json.dumps(payload, ensure_ascii=False) + + +# ── document listing / name resolution ── + +def _all_documents(client) -> list[dict[str, Any]]: + """Every document the client can list, newest first (both modes list + newest-first; paging preserves that order).""" + documents: list[dict[str, Any]] = [] + offset = 0 + while True: + page = client.list_documents(limit=100, offset=offset) + batch = page.get("documents") or [] + documents.extend(batch) + # Advance by what actually arrived — stepping by the requested + # limit skips documents whenever a server caps its page size. + offset += len(batch) + total = page.get("total") + # An empty page is the reliable terminator; `total` (absent or + # None on some backends) only saves the final empty-page request. + if not batch or (isinstance(total, int) and offset >= total): + return documents + + +def _normalize_created_at(value: Any) -> str: + """Emit the cloud tool format (ISO-8601 UTC with 'Z', millisecond + precision) from either mode's createdAt string.""" + if not isinstance(value, str) or not value: + return "" + try: + from datetime import datetime, timezone + parsed = datetime.fromisoformat(value.replace("Z", "+00:00")) + if parsed.tzinfo is None: + parsed = parsed.replace(tzinfo=timezone.utc) + parsed = parsed.astimezone(timezone.utc) + return parsed.isoformat(timespec="milliseconds").replace("+00:00", "Z") + except ValueError: + return value + + +def _flat_metadata(value: Any) -> Optional[dict[str, Any]]: + """User-facing string|number|boolean metadata fields only, or None.""" + if not isinstance(value, dict): + return None + flat = {key: val for key, val in value.items() + if isinstance(val, (str, int, float, bool))} + return flat or None + + +def _scope_documents(documents: list[dict[str, Any]], + allowed_ids: Optional[frozenset]) -> list[dict[str, Any]]: + if allowed_ids is None: + return documents + return [doc for doc in documents if doc.get("id") in allowed_ids] + + +def _resolve_document( + client, doc_name: str, + documents: Optional[list[dict[str, Any]]] = None, + allowed_ids: Optional[frozenset] = None, +) -> "tuple[Optional[dict[str, Any]], Optional[_ToolResult]]": + """Resolve doc_name to a list entry. Same-name duplicates resolve to the + newest match. Returns (entry, None) or (None, error_payload_pair).""" + if documents is None: + documents = _all_documents(client) + documents = _scope_documents(documents, allowed_ids) + matches = [doc for doc in documents if doc.get("name") == doc_name] + if matches: + return max(matches, key=lambda d: d.get("createdAt") or ""), None + names = [str(doc.get("name")) for doc in documents if doc.get("name")] + similar = difflib.get_close_matches(doc_name, names, n=_SIMILAR_NAMES_LIMIT, + cutoff=0.5) + message = ( + "Document not found. Did you mean: " + + ", ".join(f'"{name}"' for name in similar) + "?" + if similar else "Document not found or you do not have access to it" + ) + return None, _failure( + message, + {"doc_name": doc_name, "similar_files": similar}, + { + "summary": "The requested document does not exist or is not accessible", + "options": [ + "Verify the document name is correct", + "Use browse_documents() to see your recent documents", + "Check if the document was deleted", + ], + }, + "NOT_FOUND", + ) + + +def _refetch_entry(client, doc_id: str) -> Optional[dict[str, Any]]: + try: + return client.get_document(doc_id) + except PageIndexAPIError: + return None + + +def _await_completion(client, entry: dict[str, Any], wait: bool) -> dict[str, Any]: + """Re-poll a processing document for up to 3 minutes when wait is set.""" + doc_id = entry.get("id") + if not wait or not doc_id or entry.get("status") in ("completed", "failed"): + return entry + deadline = time.monotonic() + _TOOL_WAIT_TIMEOUT + current = entry + while time.monotonic() < deadline: + time.sleep(_TOOL_WAIT_INTERVAL) + refreshed = _refetch_entry(client, doc_id) + if refreshed is None: + return current + if refreshed.get("metadata") is None: + # Status refetches omit (or null out) custom metadata; keep the + # listing's copy. + refreshed["metadata"] = current.get("metadata") + current = {**current, **refreshed} + if current.get("status") in ("completed", "failed"): + return current + return current + + +def _not_ready_error(doc_name: str, status: Any, operation: str, + timed_out: bool) -> tuple[dict, bool]: + if status == "failed": + return _failure( + f"Document processing failed. Current status: {status}", + {"doc_name": doc_name}, + { + "summary": "Document processing has failed", + "options": [ + "Index the document again with " + "PageIndexClient.submit_document()", + "Use browse_documents() to work with other documents", + ], + }, + "INVALID_INPUT", + ) + if timed_out: + return _failure( + f"Document is still processing. Current status: {status}", + {"doc_name": doc_name}, + { + "summary": "Document processing timeout", + "options": [ + "Try again later when processing is complete", + "Check status with get_document()", + ], + }, + "INVALID_INPUT", + ) + return _failure( + f"Document is not ready for {operation}. Current status: {status}", + {"doc_name": doc_name}, + { + "summary": "Document is still processing", + "options": [ + "Wait for document processing to complete", + "Check status with browse_documents() or get_document()", + ], + }, + "INVALID_INPUT", + ) + + +def _folder_unsupported(param: str) -> tuple[dict, bool]: + return _failure( + f"Folders are not supported in local mode yet — omit {param}.", + None, + { + "summary": "This local library does not have folders yet", + "options": ["Retry the call without a folder_id", + "Use browse_documents() to list the library root", + "Folders are available on PageIndex cloud (PageIndexCloudClient with an API key)"], + }, + "INVALID_INPUT", + ) + + +# ── page spec handling ── + +def _parse_page_spec( + pages: str, doc_name: str, +) -> "tuple[Optional[list[int]], Optional[_ToolResult]]": + """Expand '1-3,7' into a sorted, deduplicated page list, or an error.""" + invalid = _failure( + "Invalid page specification format", + {"doc_name": doc_name}, + { + "summary": "Failed to parse the pages parameter", + "options": [ + 'Use valid formats: "5", "3,7,10", "5-10", or "1-3,7,9-12"', + "Ensure page numbers are positive integers", + ], + }, + "INVALID_INPUT", + ) + if not isinstance(pages, str) or not _PAGES_SPEC_RE.match(pages.strip()): + return None, invalid + too_many = _failure( + f"Too many pages requested (over {_MAX_REQUESTED_PAGES})", + {"doc_name": doc_name}, + { + "summary": "The page specification spans too many pages", + "options": [ + "Request a narrower page range", + "The response holds only a few pages per call - page through with several smaller requests", + ], + }, + "INVALID_INPUT", + ) + expanded: set[int] = set() + for part in pages.split(","): + part = part.strip() + if "-" in part: + start, end = (int(x) for x in part.split("-", 1)) + if start > end: + return None, invalid + else: + start = end = int(part) + # Bound each part arithmetically before materializing it: a spec like + # "1-1000000000" would otherwise expand to billions of integers + # inside the caller's process. The cap is on distinct pages, so + # overlapping parts (a parent section plus its children) don't + # double-count. + if end - start + 1 > _MAX_REQUESTED_PAGES: + return None, too_many + expanded.update(range(start, end + 1)) + if len(expanded) > _MAX_REQUESTED_PAGES: + return None, too_many + if any(page < 1 for page in expanded): + return None, _failure( + "Invalid page numbers. Page numbers must be positive integers", + {"doc_name": doc_name}, + { + "summary": "Invalid page numbers provided", + "options": [ + "Page numbers must be positive integers (>= 1)", + "Check the page specification format", + ], + }, + "INVALID_INPUT", + ) + return sorted(expanded), None + + +def _format_page_spec(pages: list[int]) -> str: + """Compress [1,2,3,5] into '1-3,5'.""" + if not pages: + return "" + ordered = sorted(set(pages)) + ranges = [] + start = prev = ordered[0] + for page in ordered[1:]: + if page == prev + 1: + prev = page + continue + ranges.append(f"{start}" if start == prev else f"{start}-{prev}") + start = prev = page + ranges.append(f"{start}" if start == prev else f"{start}-{prev}") + return ",".join(ranges) + + +# ── structure formatting / splitting ── + +_STRUCTURE_KEY_ORDER = ("title", "node_id", "start_index", "end_index", + "page_index", "prefix_summary", "summary", "nodes") + + +def _format_structure(node: Any) -> Any: + """Drop node text and normalize key order, recursively.""" + if isinstance(node, list): + return [_format_structure(item) for item in node] + if isinstance(node, dict): + stripped = {key: value for key, value in node.items() if key != "text"} + if "nodes" in stripped: + stripped["nodes"] = _format_structure(stripped["nodes"]) + ordered = {key: stripped[key] for key in _STRUCTURE_KEY_ORDER + if key in stripped} + ordered.update({key: value for key, value in stripped.items() + if key not in ordered}) + return ordered + return node + + +def _serialized_size(value: Any) -> int: + return len(json.dumps(value, ensure_ascii=False)) + + +def _split_structure(structure: Any, budget: int) -> list[Any]: + """Split a formatted structure into chunks of at most ~budget serialized + chars. The paginated response shape matches the cloud tool (its chunk + type admits node-or-list); chunk boundaries are implementation-defined. + An unsplit structure keeps its natural shape; once split, every chunk + is a list of nodes — the `structure` field must not change JSON type + between parts of one paginated response.""" + if _serialized_size(structure) <= budget: + return [structure] + nodes = structure if isinstance(structure, list) else [structure] + chunks: list[Any] = [] + group: list[Any] = [] + group_size = 0 + for node in nodes: + size = _serialized_size(node) + if size > budget: + if group: + chunks.append(group) + group, group_size = [], 0 + chunks.extend([part] + for part in _split_oversized_node(node, budget)) + continue + if group and group_size + size > budget: + chunks.append(group) + group, group_size = [], 0 + group.append(node) + group_size += size + if group: + chunks.append(group) + return chunks or [structure] + + +def _split_oversized_node(node: Any, budget: int) -> list[Any]: + children = node.get("nodes") if isinstance(node, dict) else None + if not children: + return [node] + shell = {key: value for key, value in node.items() if key != "nodes"} + shell_size = _serialized_size(shell) + child_budget = max(budget - shell_size, budget // 2) + parts = [] + for chunk in _split_structure(children, child_budget): + # A recursive result is either the unsplit children (natural shape) + # or always-list chunks; normalize for the shell's "nodes". + parts.append({**shell, + "nodes": chunk if isinstance(chunk, list) else [chunk]}) + return parts + + +# ── tool implementations (client-backed; mode-blind) ── + +def _browse_documents(client, folder_id: str = "root", recursive: bool = False, + sort: str = "time", query: Optional[str] = None, + offset: int = 0, limit: int = 10, + _allowed_ids: Optional[frozenset] = None) -> tuple[dict, bool]: + if folder_id != "root": + return _folder_unsupported("folder_id") + if sort not in ("time", "relevance"): + return _failure( + 'Invalid sort mode — only the default "time" sort is available ' + "in local mode.", None, + {"summary": "Invalid sort mode", + "options": ['Use sort="time" (newest first) or omit sort', + "Semantic ranking is available on PageIndex cloud (PageIndexCloudClient with an API key)"]}, + "INVALID_INPUT", + ) + if sort == "relevance" or query: + # Semantic ranking is a cloud capability; like folders, it is not + # imitated here. + return _failure( + "Relevance ranking is not supported in local mode yet — use " + "the default time sort.", None, + {"summary": "This local library does not have semantic ranking yet", + "options": ["Retry without sort/query and match the returned names and descriptions against the intent yourself", + "Page through the full library with `offset: next_offset`", + "Semantic ranking is available on PageIndex cloud (PageIndexCloudClient with an API key)"]}, + "INVALID_INPUT", + ) + try: + offset = max(int(offset), 0) + limit = min(max(int(limit), 1), 50) + except (TypeError, ValueError): + return _failure("offset and limit must be numbers", None, + {"summary": "Invalid pagination parameters", + "options": ["Pass integer offset and limit values"]}, + "INVALID_INPUT") + + if _allowed_ids is None: + listing = client.list_documents(limit=limit, offset=offset) + window = listing.get("documents") or [] + total = listing.get("total") + else: + scoped = _scope_documents(_all_documents(client), _allowed_ids) + window, total = scoped[offset:offset + limit], len(scoped) + window_end = offset + len(window) + has_more = bool(window) and (window_end < total if isinstance(total, int) + else len(window) == limit) + next_offset = window_end if has_more else None + + page_has_processing = False + page_has_failed = False + items = [] + for doc in window: + status = doc.get("status") or "unknown" + if status == "failed": + page_has_failed = True + elif status != "completed": + page_has_processing = True + item = { + "name": doc.get("name") or "Unknown Document", + "description": doc.get("description") or "No description provided", + "status": status, + "created_at": _normalize_created_at(doc.get("createdAt")), + } + metadata = _flat_metadata(doc.get("metadata")) + if metadata is not None: + item["metadata"] = metadata + items.append(item) + + data: dict[str, Any] = { + "documents": items, + "sort": sort, + "next_offset": next_offset, + "has_more": has_more, + } + if not recursive: + data["folders"] = [] + + if not items and offset == 0: + next_steps = { + "summary": "Nothing to show", + "options": ["Nothing here. Index documents with " + "PageIndexClient.submit_document() to get started."], + "auto_retry": "Index a document with " + "PageIndexClient.submit_document() to get started", + } + return _success(data, next_steps) + + options = [] + if items: + options.append("Use get_document() with a document name to view details") + options.append( + "Results returned ≠ correct results. Verify these documents match " + "the user's actual intent (topic, time period, document type) " + "before proceeding." + + (" If they do not match, page through the rest of the library." + if has_more else "") + + " Do NOT use general knowledge as a substitute." + ) + if page_has_processing: + options.append("Some documents on this page are still processing. " + "Use get_document() to check individual status.") + if page_has_failed: + options.append("Some documents on this page failed processing. " + "Use get_document() to see error details.") + if has_more: + options.append("Use browse_documents() with `offset: next_offset` to " + "load more documents") + summary = (f"Showing {len(items)} document(s)" + + (" (more available)" if has_more else "") + if items else "Nothing to show") + return _success(data, {"summary": summary, "options": options}) + + +def _get_document(client, doc_name: str, folder_id: Optional[str] = None, + wait_for_completion: bool = False, + _allowed_ids: Optional[frozenset] = None) -> tuple[dict, bool]: + if folder_id not in (None, "root"): + return _folder_unsupported("folder_id") + entry, error = _resolve_document(client, doc_name, allowed_ids=_allowed_ids) + if error is not None: + return error + assert entry is not None + entry = _await_completion(client, entry, wait_for_completion) + + status = entry.get("status") or "unknown" + is_processing = status not in ("completed", "failed") + is_ready = status == "completed" + page_num = entry.get("pageNum") or 0 + name = entry.get("name") or "Unknown Document" + + suggestions: list[str] = [] + if is_processing: + suggestions.append("Document is still processing. Processing status " + "can be checked later.") + elif is_ready: + suggestions.append("Document is ready for analysis.") + if page_num > 0: + if page_num <= 5: + suggestions.extend([ + f"This is a short document with {page_num} pages.", + f'First explore structure: get_document_structure(doc_name: "{name}")', + f'Then extract all content: get_page_content(doc_name: "{name}", pages: "1-{page_num}")', + ]) + elif page_num <= STRUCTURE_FIRST_PAGE_THRESHOLD: + suggestions.extend([ + f"This document has {page_num} pages.", + f'First explore structure: get_document_structure(doc_name: "{name}")', + f'Then extract key pages: get_page_content(doc_name: "{name}", pages: "1,5,10")', + ]) + else: + suggestions.extend([ + f"This is a large document with {page_num} pages.", + f'First explore structure: get_document_structure(doc_name: "{name}")', + f'Then target specific sections: get_page_content(doc_name: "{name}", pages: "1-3")', + ]) + else: + suggestions.append("Document processing failed. Index the document " + "again with PageIndexClient.submit_document().") + + data: dict[str, Any] = { + "name": name, + "description": entry.get("description") or "No description provided", + "status": status, + "created_at": _normalize_created_at(entry.get("createdAt")), + "page_count": page_num or None, + "folder_id": entry.get("folderId"), + } + metadata = _flat_metadata(entry.get("metadata")) + if metadata is not None: + data["metadata"] = metadata + + return _success(data, { + "summary": ("Document is ready for analysis and querying." if is_ready + else "Document is still being processed." if is_processing + else "Document processing has failed."), + "options": suggestions, + **({"auto_retry": "Document processing status can be monitored periodically"} + if is_processing else {}), + }) + + +def _get_document_structure(client, doc_name: str, + folder_id: Optional[str] = None, part: int = 1, + wait_for_completion: bool = False, + _allowed_ids: Optional[frozenset] = None) -> tuple[dict, bool]: + if folder_id not in (None, "root"): + return _folder_unsupported("folder_id") + entry, error = _resolve_document(client, doc_name, allowed_ids=_allowed_ids) + if error is not None: + return error + assert entry is not None + waited = wait_for_completion and entry.get("status") not in ("completed", "failed") + entry = _await_completion(client, entry, wait_for_completion) + if entry.get("status") != "completed": + return _not_ready_error(doc_name, entry.get("status"), + "structure retrieval", + waited and entry.get("status") != "failed") + + try: + raw_tree = getattr(getattr(client, "_api", None), "raw_tree", None) + tree = raw_tree(entry["id"]) if raw_tree is not None else None + if tree is None: + tree = client.get_tree(entry["id"], node_summary=True).get("result") + except PageIndexAPIError as exc: + return _failure( + f"Failed to retrieve document structure: {exc}", + {"doc_name": doc_name}, + { + "summary": "Failed to retrieve document structure due to an error", + "options": [ + "The document may not exist or is not accessible", + "Check if the document name is correct", + "Try again in a few moments", + ], + }, + "INTERNAL_ERROR", + ) + if tree is None: + return _failure( + "Structure not available for this document", + {"doc_name": doc_name}, + { + "summary": "Structure not available for this document", + "options": [ + "The document may not have been processed correctly or structure extraction may have failed", + "Try processing the document again if possible", + ], + }, + "INTERNAL_ERROR", + ) + + formatted = _format_structure(tree) + chunks = _split_structure(formatted, _CHAR_BUDGET) + total_parts = max(1, len(chunks)) + try: + requested_part = int(part) + except (TypeError, ValueError): + requested_part = 1 + current = min(max(requested_part, 1), total_parts) + + if total_parts == 1: + return _success( + {"doc_name": doc_name, "structure": chunks[0]}, + { + "summary": "Document structure retrieved successfully.", + "options": [ + "Use get_page_content() to extract specific content from pages", + ], + }, + ) + + next_steps = ( + { + "summary": f"Showing part {current} of {total_parts}.", + "options": [ + f"Request next part with part: {current + 1}", + f"Jump to last part with part: {total_parts}", + "Proceed to get_page_content() for specific sections", + ], + } + if current < total_parts else + { + "summary": "All parts retrieved for current pagination.", + "options": [ + "Use get_page_content() to extract specific content from pages", + ], + } + ) + return _success( + { + "doc_name": doc_name, + "total_parts": total_parts, + "structure": chunks[current - 1], + "pagination": { + "part": current, + "total_parts": total_parts, + "has_more": current < total_parts, + }, + }, + next_steps, + ) + + +def _get_page_content(client, doc_name: str, pages: str, + folder_id: Optional[str] = None, + wait_for_completion: bool = False, + _allowed_ids: Optional[frozenset] = None) -> tuple[dict, bool]: + if folder_id not in (None, "root"): + return _folder_unsupported("folder_id") + entry, error = _resolve_document(client, doc_name, allowed_ids=_allowed_ids) + if error is not None: + return error + assert entry is not None + waited = wait_for_completion and entry.get("status") not in ("completed", "failed") + entry = _await_completion(client, entry, wait_for_completion) + if entry.get("status") != "completed": + return _not_ready_error(doc_name, entry.get("status"), + "page content retrieval", + waited and entry.get("status") != "failed") + + requested, error = _parse_page_spec(pages, doc_name) + if error is not None: + return error + assert requested is not None + + try: + page_data = client.get_ocr(entry["id"], format="page").get("result") or [] + except PageIndexAPIError as exc: + return _failure( + f"Failed to retrieve page content: {exc}", + {"doc_name": doc_name}, + { + "summary": "Unable to retrieve page content due to a service issue.", + "options": [ + "Verify the document name is correct using browse_documents()", + "Check if the document processing is complete with get_document()", + "Ensure the requested page numbers are valid", + ], + "auto_retry": "This may be a temporary issue - you can try " + "the request again", + }, + "INTERNAL_ERROR", + ) + + by_index = {item["page_index"]: item for item in page_data + if isinstance(item, dict) + and isinstance(item.get("page_index"), int)} + max_page = max(by_index, default=0) + + out_of_range = [page for page in requested if page > max_page] + valid_pages = [page for page in requested if page <= max_page] + if out_of_range and not valid_pages: + return _failure( + f"All requested pages are out of range. Document has {max_page} " + f"pages, but you requested pages: {_format_page_spec(out_of_range)}", + { + "doc_name": doc_name, + "max_pages": max_page, + "requested_pages": _format_page_spec(out_of_range), + }, + { + "summary": "All requested pages are out of range for this document", + "options": [ + f"Request pages between 1 and {max_page}", + "Use get_document() to check document page count", + ], + }, + "INVALID_INPUT", + ) + + content = [] + included: list[int] = [] + remaining: list[int] = [] + budget = _CHAR_BUDGET + for page in valid_pages: + item = by_index.get(page) + markdown = item.get("markdown") if item else None + text = (markdown if isinstance(markdown, str) + else f"Page {page} content not available") + if not included or budget - len(text) >= 0: + content.append({"page": page, "text": text}) + included.append(page) + budget -= len(text) + else: + remaining.append(page) + + options = [ + "Use get_document_structure() to understand document organization", + "Request additional pages as needed", + ] + if remaining: + options.insert(0, f"For remaining pages, request: {_format_page_spec(remaining)}") + if out_of_range: + options.insert(0, f"Document has {max_page} pages total - request " + f"pages 1-{max_page}") + if remaining or out_of_range: + parts = [f"Retrieved {len(included)} of {len(requested)} " + "requested pages."] + if remaining: + parts.append(f"Pages {_format_page_spec(remaining)} were " + "omitted due to response size limits.") + if out_of_range: + parts.append(f"Pages {_format_page_spec(out_of_range)} " + "were out of range.") + summary = " ".join(parts) + else: + summary = (f"Successfully retrieved content for {len(content)} " + f"page{'' if len(content) == 1 else 's'}.") + return _success( + { + "doc_name": doc_name, + "total_pages": max_page, + "requested_pages": _format_page_spec(requested), + "returned_pages": _format_page_spec(included), + "content": content, + }, + {"summary": summary, "options": options}, + ) + + +def _remove_document(client, doc_names: list[str], + folder_id: Optional[str] = None, + _allowed_ids: Optional[frozenset] = None) -> tuple[dict, bool]: + if folder_id not in (None, "root"): + return _folder_unsupported("folder_id") + if not isinstance(doc_names, list) or not doc_names: + return _failure("At least one document name is required", None, + {"summary": "No document names provided", + "options": ["Pass doc_names as a non-empty array"]}, + "INVALID_INPUT") + # Validate every element before deleting anything: a rejection envelope + # must mean nothing was destroyed. + if not all(isinstance(name, str) and name.strip() for name in doc_names): + return _failure( + "doc_names must be an array of non-empty document name strings", + None, + {"summary": "Invalid document names", + "options": ["Copy each name verbatim from a browse_documents() " + "response"]}, + "INVALID_INPUT") + doc_names = list(dict.fromkeys(doc_names)) + if len(doc_names) > 10: + return _failure("Maximum 10 documents can be deleted at once", None, + {"summary": "Too many documents in one call", + "options": ["Delete at most 10 documents per call"]}, + "INVALID_INPUT") + documents = _all_documents(client) + results = [] + for doc_name in doc_names: + entry, error = _resolve_document(client, doc_name, documents=documents, + allowed_ids=_allowed_ids) + if error is not None or entry is None: + results.append({"doc_name": doc_name, "status": "not_found"}) + continue + try: + client.delete_document(entry["id"]) + results.append({"doc_name": doc_name, "status": "deleted"}) + except Exception as exc: + # Any escape here (OSError, transport errors) would discard the + # entries for documents already irreversibly deleted. + results.append({"doc_name": doc_name, "status": "failed", + "error": str(exc)}) + deleted = sum(1 for item in results if item["status"] == "deleted") + return _success( + {"results": results}, + { + "summary": f"Deleted {deleted} of {len(doc_names)} document(s).", + "options": ["Use browse_documents() to review the remaining library"], + }, + ) + + +_IMPLEMENTATIONS: dict[str, Callable[..., tuple[dict, bool]]] = { + "browse_documents": _browse_documents, + "get_document": _get_document, + "get_document_structure": _get_document_structure, + "get_page_content": _get_page_content, + "remove_document": _remove_document, +} + + +def tool_names(include_management: bool = False) -> tuple[str, ...]: + return _READ_TOOLS + (_MANAGEMENT_TOOLS if include_management else ()) + + +def _coerce_bool_args(name: str, kwargs: dict[str, Any]) -> None: + """Models routinely send booleans as JSON strings ("false"); the bare + truthiness tests downstream would read those as True.""" + properties = TOOL_CONTRACT.get(name, {}).get("schema", {}).get( + "properties", {}) + for key, spec in properties.items(): + value = kwargs.get(key) + if spec.get("type") == "boolean" and isinstance(value, str): + kwargs[key] = value.strip().lower() not in ("false", "no", "0", "") + + +def call_tool(client, name: str, arguments: dict[str, Any], + doc_ids=None) -> tuple[str, bool]: + """Run one contract tool; returns (envelope_json, is_error). Never raises + for tool-level failures — unexpected exceptions become error envelopes. + ``doc_ids`` restricts every document lookup to that allowlist (the local + chat surfaces' doc_id scope).""" + implementation = _IMPLEMENTATIONS.get(name) + if implementation is None: + payload, _ = _failure( + f"Unknown tool: {name}", + {"tool_name": name, "available_tools": list(_IMPLEMENTATIONS)}, + {"summary": "Tool not found", + "options": [f"Available tools: {', '.join(_IMPLEMENTATIONS)}"]}, + "INVALID_INPUT", + ) + return _dumps(payload), True + if arguments is not None and not isinstance(arguments, dict): + payload, is_error = _failure( + f"Invalid arguments for {name}: expected a JSON object, got " + f"{type(arguments).__name__}", None, + {"summary": "Invalid tool arguments", + "options": [f"Pass {name}() arguments as a JSON object of its " + "parameters"]}, + "INVALID_INPUT", + ) + return _dumps(payload), is_error + # Underscore-prefixed keys are the SDK's private channel (the scope + # below), never model arguments. None ≡ omitted (the contract's + # "omit if ..." semantics, same as the cloud bridge invoker). + kwargs = {key: value for key, value in (arguments or {}).items() + if not key.startswith("_") and value is not None} + _coerce_bool_args(name, kwargs) + try: + if doc_ids is not None: + ids = [doc_ids] if isinstance(doc_ids, str) else doc_ids + kwargs["_allowed_ids"] = frozenset(str(one_id) for one_id in ids) + bound = inspect.signature(implementation).bind(client, **kwargs) + except TypeError as exc: + payload, is_error = _failure( + f"Invalid arguments for {name}: {exc}", None, + {"summary": "Invalid tool arguments", + "options": [f"Check the {name}() parameter names and types"]}, + "INVALID_INPUT", + ) + return _dumps(payload), is_error + try: + payload, is_error = implementation(*bound.args, **bound.kwargs) + except Exception as exc: # tool calls must never raise into the agent loop + payload, is_error = _failure( + f"{name} failed: {exc}", None, + {"summary": "Unexpected error while running the tool", + "options": ["Try the request again"], + "auto_retry": "This is likely a temporary issue - you can try " + "the request again"}, + "INTERNAL_ERROR", + ) + return _dumps(payload), is_error + + +# ── plain-function materialization (the `client.agent_tools()` surface) ── + +def _tool_docstring(description: str, properties: dict[str, Any]) -> str: + lines = [description, "", "Args:"] + for param, spec in properties.items(): + lines.append(f" {param}: {spec.get('description', '')}") + return "\n".join(lines) + + +_LOCAL_HIDDEN_PARAMS: dict[str, tuple[str, ...]] = { + "browse_documents": ("folder_id", "recursive", "sort", "query"), + "get_document": ("folder_id",), + "get_document_structure": ("folder_id",), + "get_page_content": ("folder_id",), + "remove_document": ("folder_id",), +} + +_LOCAL_DOC_NAME_DESCRIPTION = ( + 'Copy the `name` field verbatim from a browse_documents() response ' + '(case-sensitive, include extension). Example: "Q3 Report.pdf". ' + "Document names are unique in a local library." +) + +_LOCAL_DESCRIPTIONS: dict[str, str] = { + "browse_documents": ( + "Primary document retrieval tool — first choice for any " + "document-related question. Lists your documents newest first with " + "names and descriptions; match them against the user's intent and " + "page through with `offset: next_offset` (limit up to 50) while " + "`has_more` is true. " + 'Folder browsing and semantic ranking (sort="relevance") are not ' + "supported in local mode yet — they work on PageIndex cloud." + ), + "get_page_content": TOOL_CONTRACT["get_page_content"]["description"] + .replace(" Embedded image paths in the response feed into " + "`get_document_image()`.", ""), +} + +_LOCAL_PARAM_DESCRIPTIONS: dict[tuple[str, str], str] = { + ("get_document", "doc_name"): _LOCAL_DOC_NAME_DESCRIPTION, + ("get_document_structure", "doc_name"): _LOCAL_DOC_NAME_DESCRIPTION, + ("get_page_content", "doc_name"): _LOCAL_DOC_NAME_DESCRIPTION, + ("remove_document", "doc_names"): ( + "Array of document names to delete. Each name must be copied " + "verbatim from the `name` field of a browse_documents() response " + '(case-sensitive, include extension). Example: ["Q3 Report.pdf", ' + '"draft.pdf"]. Max 10 per call.' + ), +} + + +def _local_description(name: str) -> str: + return _LOCAL_DESCRIPTIONS.get(name) or TOOL_CONTRACT[name]["description"] + + +def _local_schema(name: str) -> dict[str, Any]: + schema = copy.deepcopy(TOOL_CONTRACT[name]["schema"]) + for param in _LOCAL_HIDDEN_PARAMS.get(name, ()): + schema["properties"].pop(param, None) + for (tool_name, param), text in _LOCAL_PARAM_DESCRIPTIONS.items(): + if tool_name == name and param in schema["properties"]: + schema["properties"][param]["description"] = text + return schema + + +def _docstring(name: str) -> str: + return _tool_docstring(_local_description(name), + _local_schema(name)["properties"]) + + +_SCHEMA_TYPE_MAP = {"string": str, "integer": int, "number": float, + "boolean": bool, "array": list, "object": dict} + + +def _annotation_for(spec: dict) -> Any: + schema_type = spec.get("type") + if schema_type is None and isinstance(spec.get("anyOf"), list): + # Nullable unions arrive as anyOf: [{type: string}, {type: null}]. + options = [option for option in spec["anyOf"] + if isinstance(option, dict) and option.get("type")] + schema_type = [option["type"] for option in options] + # `items` lives on the array option, not the union shell. + spec = next((option for option in options + if option["type"] == "array"), spec) + nullable = False + if isinstance(schema_type, list): + nullable = "null" in schema_type + bases = [t for t in schema_type if t != "null"] + schema_type = bases[0] if bases else None + base = _SCHEMA_TYPE_MAP.get(schema_type or "", Any) + if base is list: + # Strict function calling rejects arrays whose item type was lost + # in the annotation round-trip; parameterize when it is known. + item_type = (spec["items"].get("type") + if isinstance(spec.get("items"), dict) else None) + element = (_SCHEMA_TYPE_MAP.get(item_type) + if isinstance(item_type, str) else None) + if element is not None: + base = list[element] + return Optional[base] if nullable else base + + +def _bridge_invoker(bridge, name: str) -> "Callable[[dict], tuple[str, bool]]": + """One cloud tool call proxied over MCP: None-valued arguments are + dropped (None ≡ omitted, matching the contract's "omit if ..." + semantics) and failures are contained in the error envelope. Returns + (envelope_text, is_error), like call_tool.""" + def _invoke(arguments: dict[str, Any]) -> tuple[str, bool]: + try: + arguments = {key: value for key, value in arguments.items() + if value is not None} + return bridge.call_tool(name, arguments) + except Exception as exc: + payload, _ = _failure( + f"{name} failed: {exc}", None, + {"summary": "Unexpected error while running the tool", + "options": ["Try the request again"], + "auto_retry": "This is likely a temporary issue - you can " + "try the request again"}, + "INTERNAL_ERROR", + ) + return _dumps(payload), True + return _invoke + + +def _make_bridge_function(bridge, meta: dict) -> Callable[..., str]: + """One plain function for a cloud tool: real signature and docstring from + the server's schema, invocation proxied over MCP, errors contained.""" + import keyword + + name = str(meta.get("name") or "") + schema = meta.get("inputSchema") or {} + properties: dict[str, Any] = schema.get("properties") or {} + required = set(schema.get("required") or []) + _invoke = _bridge_invoker(bridge, name) + + params_usable = all(param.isidentifier() and not keyword.iskeyword(param) + and param != "_invoke" + for param in properties) + if not params_usable: + def proxy(**kwargs: Any) -> str: + return _invoke(kwargs)[0] + else: + ordered = ([p for p in properties if p in required] + + [p for p in properties if p not in required]) + rendered = ", ".join( + p if p in required else f"{p}={properties[p].get('default')!r}" + for p in ordered + ) + args_literal = "{" + ", ".join(f"'{p}': {p}" for p in ordered) + "}" + namespace: dict[str, Any] = {"_invoke": _invoke} + exec(f"def _synthesized({rendered}):\n" + f" return _invoke({args_literal})[0]", namespace) + proxy = namespace["_synthesized"] + annotations: dict[str, Any] = {} + for p in ordered: + annotation = _annotation_for(properties[p]) + if p not in required and "default" not in properties[p]: + # Absent-but-non-nullable params must admit None, or strict + # schemas force the model to always send a value. + annotation = Optional[annotation] + annotations[p] = annotation + annotations["return"] = str + proxy.__annotations__ = annotations + proxy.__name__ = proxy.__qualname__ = name or "tool" + proxy.__doc__ = _tool_docstring(meta.get("description") or "", properties) + return proxy + + +_BRIDGES: "weakref.WeakKeyDictionary" = weakref.WeakKeyDictionary() +_BRIDGES_LOCK = threading.Lock() + + +def _cloud_bridge(client): + """One bridge per client: tool discovery and instructions share a single + MCP session. Weak-keyed off the instance so clients stay picklable; the + lock closes the check-then-set race under concurrent first calls.""" + with _BRIDGES_LOCK: + bridge = _BRIDGES.get(client) + if bridge is None: + from .mcp_bridge import McpBridge + bridge = McpBridge( + f"{client.BASE_URL}/mcp", + {"Authorization": f"Bearer {client.api_key}"}, + ) + _BRIDGES[client] = bridge + return bridge + + +def _read_only_tools(tools_meta: list[dict]) -> list[dict]: + """The management gate for consumers without a framework permission + layer: only tools the server marks read-only, guarded against a server + annotation regression silently disabling every tool.""" + filtered = [meta for meta in tools_meta + if (meta.get("annotations") or {}).get("readOnlyHint") is True] + if tools_meta and not filtered: + raise PageIndexAPIError( + "The MCP server returned tools but none are annotated " + "read-only — a server annotation regression would otherwise " + "silently disable every tool. Pass include_management=True " + "to expose the unfiltered list." + ) + return filtered + + +def _build_cloud_agent_tools(client, include_management: bool) -> list[Callable[..., str]]: + bridge = _cloud_bridge(client) + tools_meta = bridge.list_tools() + if not include_management: + tools_meta = _read_only_tools(tools_meta) + return [_make_bridge_function(bridge, meta) for meta in tools_meta] + + +def _require_local_scope(client, doc_ids) -> None: + """The allowlist is enforced in-process; cloud lookups run server-side, + so accepting doc_ids there would be advisory-only — refuse loudly.""" + if doc_ids is not None and getattr(client, "api_key", None): + raise PageIndexAPIError( + "doc_ids scoping applies to local tools only — cloud calls " + "are scoped server-side." + ) + + +def _tool_specs(client, include_management: bool = False, doc_ids=None, + ) -> "list[tuple[str, str, dict, Callable[[dict], tuple[str, bool]]]]": + """(name, description, schema, invoke) per tool, for adapters that take + the wire schema verbatim. ``invoke`` returns (envelope_text, is_error). + Schemas are copies (frameworks keep the dict by reference). ``doc_ids`` + is the local chat scope; cloud scoping is server-side.""" + _require_local_scope(client, doc_ids) + if getattr(client, "api_key", None): + bridge = _cloud_bridge(client) + tools_meta = bridge.list_tools() + if not include_management: + tools_meta = _read_only_tools(tools_meta) + return [(str(meta.get("name") or "tool"), + meta.get("description") or "", + copy.deepcopy(meta.get("inputSchema")) + or {"type": "object", "properties": {}}, + _bridge_invoker(bridge, str(meta.get("name") or "tool"))) + for meta in tools_meta] + + def local_invoke(name: str) -> "Callable[[dict], tuple[str, bool]]": + def invoke(arguments: dict) -> tuple[str, bool]: + return call_tool(client, name, arguments, doc_ids=doc_ids) + return invoke + + return [(name, _local_description(name), _local_schema(name), + local_invoke(name)) + for name in tool_names(include_management)] + + +def build_agent_tools(client, include_management: bool = False) -> list[Callable[..., str]]: + """Plain synchronous functions bound to `client`. + + Cloud: one function per tool of the live cloud MCP tool set, signatures + synthesized from the server's schemas, calls proxied over MCP. Local: + the built-in contract tools over the local store. Every function returns + the JSON envelope as a string and never raises for arguments its + signature accepts (cloud-only parameters are absent from the local + signatures; the call_tool path answers them with the guided envelope). + """ + if getattr(client, "api_key", None): + return _build_cloud_agent_tools(client, include_management) + + def browse_documents(offset: int = 0, limit: int = 10) -> str: + return call_tool(client, "browse_documents", { + "offset": offset, "limit": limit, + })[0] + + def get_document(doc_name: str, wait_for_completion: bool = False) -> str: + return call_tool(client, "get_document", { + "doc_name": doc_name, + "wait_for_completion": wait_for_completion, + })[0] + + def get_document_structure(doc_name: str, part: int = 1, + wait_for_completion: bool = False) -> str: + return call_tool(client, "get_document_structure", { + "doc_name": doc_name, "part": part, + "wait_for_completion": wait_for_completion, + })[0] + + def get_page_content(doc_name: str, pages: str, + wait_for_completion: bool = False) -> str: + return call_tool(client, "get_page_content", { + "doc_name": doc_name, "pages": pages, + "wait_for_completion": wait_for_completion, + })[0] + + def remove_document(doc_names: list[str]) -> str: + return call_tool(client, "remove_document", { + "doc_names": doc_names, + })[0] + + functions = { + "browse_documents": browse_documents, + "get_document": get_document, + "get_document_structure": get_document_structure, + "get_page_content": get_page_content, + "remove_document": remove_document, + } + tools = [] + for name in tool_names(include_management): + function = functions[name] + function.__doc__ = _docstring(name) + tools.append(function) + return tools + + +# ── agent instructions ── + +_INSTRUCTIONS_HEADER = ( + "PageIndex by Vectify AI is a document platform for uploading and " + "managing long PDFs (research papers, financial reports, legal docs, " + "textbooks, etc.)." +) + +_READING_WORKFLOW = f"""\ +READING WORKFLOW: +- For documents over {STRUCTURE_FIRST_PAGE_THRESHOLD} pages: call get_document_structure() first to locate relevant sections, then get_page_content() with targeted page ranges. +- For small documents ({STRUCTURE_FIRST_PAGE_THRESHOLD} pages or fewer): call get_page_content() directly.""" + +_TOOL_USAGE_RULES = """\ +TOOL USAGE RULES: +- Invoke a tool only when all required parameters are present or clearly inferable. Never invent placeholder values. +- If a tool returns an error, present the provided next_steps/options to the user instead of retrying blindly.""" + +_DISCOVERY = """\ +DOCUMENT DISCOVERY: +- browse_documents() — DEFAULT discovery tool, first choice for any document-related question. It lists your documents newest first with names and descriptions; match them against the user's intent, and page through with `offset: next_offset` while has_more is true.""" + +_DECISION = """\ +DECISION: +- "What do I have / list / recent" → browse_documents() +- ANY question that needs a document to answer (including "find THE paper about Y") → browse_documents(), then pick the documents whose name/description matches the question""" + +_AFTER_DISCOVERY = """\ +- Skip discovery ONLY for questions with NO possible document connection (e.g., "capital of France"). +- After discovery: 1 match or 1 clearly best match → proceed to read and answer without asking. Multiple equally relevant → ask user to pick. +- Results returned ≠ correct results. If the returned documents do not clearly match the user's intent (e.g., wrong topic, wrong time period, wrong document type), treat it the same as "not found" and continue the PERSISTENCE protocol below.""" + +_PERSISTENCE = """\ +PERSISTENCE (before concluding the target document is not in the library): +This protocol applies both when results are empty AND when results are returned but none match the user's intent. Do NOT give up after a single discovery attempt. Follow these steps in order: +1. browse_documents() and compare every returned name/description against the user's intent +2. Page through the ENTIRE library with `limit: 50` and `offset: next_offset` until has_more is false — MANDATORY, must be completed before concluding "not found" +3. Re-scan for loose matches: synonyms, abbreviations, and partial titles in names/descriptions can identify the target +Only after ALL three steps have been tried may you conclude the document is not in the library. Do NOT fall back to general knowledge — if the user's question references their own documents, exhaust every discovery path first.""" + +AGENT_INSTRUCTIONS = "\n\n".join([ + _INSTRUCTIONS_HEADER, + _READING_WORKFLOW, + _TOOL_USAGE_RULES, + _DISCOVERY, + _DECISION, + _AFTER_DISCOVERY, + _PERSISTENCE, +]) + + +def _base_instructions(client) -> str: + """Cloud: the live instructions the MCP server serves for this key's + tool set. Local: the built-in subset instructions.""" + if not getattr(client, "api_key", None): + return AGENT_INSTRUCTIONS + instructions = _cloud_bridge(client).instructions() + if not isinstance(instructions, str) or not instructions.strip(): + raise PageIndexAPIError( + "The MCP server returned no agent instructions — refusing to " + "substitute the SDK's local-subset guidance, which does not " + "cover the cloud tool set." + ) + return instructions + + +def doc_targeting_block(client, doc_id, scoped: bool = False) -> Optional[str]: + """The doc_id targeting text: names, metadata, and the directive to work + within those documents. Shared by agent_instructions and the local chat + surfaces (a leading conversation item on the OpenAI surfaces, a system + block on messages()). Raises when a doc_id's name is shadowed by a newer + same-name document — the name-addressed tools could not reach it. With + ``scoped`` (surfaces whose tools resolve names inside the doc_id + allowlist) only a same-name duplicate within the targeted set + shadows.""" + if doc_id is None: + return None + doc_ids = [doc_id] if isinstance(doc_id, str) else list(doc_id) + if not doc_ids: + return None + details = [client.get_document(one_id) for one_id in doc_ids] + listing = _all_documents(client) + documents = ([{**detail, "id": one_id} + for one_id, detail in zip(doc_ids, details)] + if scoped else listing) + for one_id, detail in zip(doc_ids, details): + entry, _ = _resolve_document(client, str(detail.get("name")), + documents=documents) + if entry is not None and entry.get("id") != one_id: + raise PageIndexAPIError( + f'Document "{detail.get("name")}" (doc_id: {one_id}) is ' + "shadowed by a newer document with the same name (doc_id: " + f'{entry.get("id")}). The tools address documents by name ' + "and would read the newer one. Rename or remove the " + "duplicate, or pass the newer doc_id." + ) + by_id = {doc.get("id"): doc for doc in listing} + for one_id, detail in zip(doc_ids, details): + if detail.get("metadata") is None: + tags = _flat_metadata(by_id.get(one_id, {}).get("metadata")) + if tags is not None: + detail["metadata"] = tags + context = json.dumps(details, ensure_ascii=False) + if len(details) == 1: + return ( + f"The user has specified document: {details[0].get('name')}\n" + f"Document metadata: {context}\n" + "Use this document's name to retrieve its content with " + "get_document_structure() and get_page_content()." + ) + names = ", ".join(str(item.get("name")) for item in details) + return ( + f"The user has specified documents: {names}\n" + f"Documents metadata: {context}\n" + "Use these documents' names to retrieve their content with " + "get_document_structure() and get_page_content()." + ) + + +def build_agent_instructions(client, doc_id=None, scoped: bool = False) -> str: + """Orchestration guidance for document QA agents; with doc_id, appends + the target documents and directs the agent to work within them.""" + base = _base_instructions(client) + block = doc_targeting_block(client, doc_id, scoped=scoped) + return base if block is None else base + "\n\n" + block diff --git a/pageindex/client.py b/pageindex/client.py index 158c9b6f7..009b7cf41 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -1,23 +1,36 @@ """PageIndex SDK client: the 0.2.x cloud surface, now with a local mode.""" from __future__ import annotations -from typing import Any, Iterator, Optional, Union +import os +import time +import warnings +from typing import Any, Callable, Iterator, Optional, Union from .errors import PageIndexAPIError def _parse_pages(pages: str) -> list[int]: - result = [] + result: set[int] = set() + too_many = (f"Page specification '{pages}' spans more than " + "10000 pages; request a narrower range") for part in pages.split(","): part = part.strip() if "-" in part: start, end = (int(x) for x in part.split("-", 1)) if start > end: raise ValueError(f"Invalid range '{part}': start must be <= end") - result.extend(range(start, end + 1)) else: - result.append(int(part)) - return sorted(set(result)) + start = end = int(part) + # Bound each part arithmetically before materializing it — a spec + # like "1-999999999" would otherwise expand to a billion integers. + # The cap is on distinct pages, so overlapping parts (a parent + # section plus its children) don't double-count. + if end - start + 1 > 10_000: + raise ValueError(too_many) + result.update(range(start, end + 1)) + if len(result) > 10_000: + raise ValueError(too_many) + return sorted(result) def _normalize_retrieve_model(model: str) -> str: @@ -47,10 +60,12 @@ class PageIndexClient: trees. Defaults to the packaged config (see pageindex/config.yaml). summary_model (str, optional): Local mode only — LLM used for node summaries and document descriptions. - retrieve_model (str, optional): Local mode only — exposed as - ``client.retrieve_model`` (the agent demo reads it); the SDK - itself consumes it once agent-based local chat lands in a - later release. + retrieve_model (str, optional): Local mode only — the model the + local chat surfaces (``chat_completions``, ``responses``) + default to, exposed as ``client.retrieve_model``. + ``provider/model`` names route through LiteLLM; for an + OpenAI-compatible server that itself serves slashed model ids + (vLLM, TGI), prefix ``openai/`` (e.g. ``openai/Qwen/...``). storage_path (str, optional): Local mode only — directory where indexed documents are stored. Defaults to ``./.pageindex``. @@ -62,10 +77,9 @@ class PageIndexClient: instead of inferring it from api_key. Local mode differences (all documented per method): indexing is - synchronous, only PDFs are supported, and ``chat_completions`` (until - agent-based local chat lands in a later release) / folders / - ``beta_headers`` / the deprecated retrieval API (``submit_query``, - ``get_retrieval``) are cloud-only. + synchronous, only PDFs are supported, and folders / ``beta_headers`` / + the deprecated retrieval API (``submit_query``, ``get_retrieval``) are + cloud-only. """ BASE_URL = "https://api.pageindex.ai" @@ -126,39 +140,94 @@ def submit_document( beta_headers: Optional[list[str]] = None, folder_id: Optional[str] = None, metadata: Optional[dict] = None, + wait: bool = False, ) -> dict[str, Any]: """ - Submit a PDF document for processing. Returns {'doc_id': ...}. - - Cloud: uploads the file; processing is asynchronous — poll - ``is_retrieval_ready(doc_id)`` before retrieving. - - Local: indexes the document in this call (it blocks while your LLM - builds the tree — minutes for a standard index of a long document), - then stores it under ``storage_path``. Pass ``mode="flash"`` to build - the tree with PageIndex Flash (layout-based extraction, no LLM calls - for the structure; node summaries and the document description still - use ``summary_model``). ``beta_headers`` and ``folder_id`` are + Submit a PDF document for processing. Returns {'doc_id': ..., 'name': ...}. + + Cloud: uploads the file; processing is asynchronous. Pass + ``wait=True`` to block until the document is ready, or poll + ``get_document(doc_id)['status']`` yourself. + + Local: indexes the document in this call and stores it under + ``storage_path``. Defaults to Flash indexing: layout-based extraction, + refined for retrieval (a deterministic merge, then an LLM expansion + pass); node summaries, the expansion pass, and the document + description use ``summary_model``. Pass ``mode="standard"`` for a + full LLM-built tree (slower). ``beta_headers`` and ``folder_id`` are cloud-only. Args: file_path (str): Path to the PDF file. - mode (str, optional): Processing mode. Local mode supports - "standard" and "flash"; omit it for standard indexing. Cloud - modes are passed through (e.g. "mcp"). + mode (str, optional): Processing mode. Local defaults to "flash"; + pass "standard" for a full LLM-built tree. Cloud modes are + passed through (e.g. "mcp"). beta_headers (list[str], optional): Cloud-only beta feature headers. folder_id (str, optional): Cloud-only folder (workspace) ID. metadata (dict, optional): Your own JSON-serializable tags for the document; returned in get_tree/get_ocr responses and list_documents entries (both modes). + wait (bool): Return only once the document is ready for use. + Cloud: polls status until "completed" (raises on "failed" or + after 30 minutes). Local: indexing is synchronous already, so + this changes nothing. Leave False to submit many documents + concurrently and poll afterwards. Returns: - dict: {'doc_id': ...} + dict: {'doc_id': ..., 'name': ...}. 'name' is the stored document + name: a taken name gains a numeric suffix (name_1..name_99) + and a UserWarning is emitted. Older cloud servers omit 'name'. """ - return self._api.submit_document( + result = self._api.submit_document( file_path=file_path, mode=mode, beta_headers=beta_headers, folder_id=folder_id, metadata=metadata, ) + stored = result.get("name") + if stored and stored != os.path.basename(file_path): + warnings.warn( + f'Document "{os.path.basename(file_path)}" was stored as ' + f'"{stored}".', + stacklevel=2, + ) + if wait: + self._wait_until_ready(result["doc_id"]) + return result + + def _wait_until_ready(self, doc_id: str, timeout: float = 1800.0) -> None: + import requests + interval = 2.0 + deadline = time.monotonic() + timeout + poll_failures = 0 + while True: + try: + status = self.get_document(doc_id).get("status") + poll_failures = 0 + except (PageIndexAPIError, requests.RequestException) as exc: + # Tolerate transient poll failures; a 30-minute wait should + # not die on one 502 or dropped connection. + poll_failures += 1 + if poll_failures >= 3: + raise PageIndexAPIError( + f"Could not poll document status (doc_id: {doc_id}): " + f"{exc}. Processing continues in the cloud — poll " + "get_document(doc_id) for status." + ) from exc + status = None + if status == "completed": + return + if status == "failed": + raise PageIndexAPIError( + f"Document processing failed (doc_id: {doc_id})." + ) + if time.monotonic() >= deadline: + raise PageIndexAPIError( + f"Timed out after {int(timeout)}s waiting for document " + f"processing (doc_id: {doc_id}, last status: {status}). " + "Processing continues in the cloud — poll " + "get_document(doc_id) for status." + ) + time.sleep(interval) + interval = min(interval * 1.5, 15.0) # ---------- OCR FUNCTIONALITY ---------- @@ -256,11 +325,11 @@ def submit_query(self, doc_id: str, query: str, thinking: bool = False) -> dict[ Cloud-only: the cloud API marks this endpoint deprecated in favor of chat completions, so local mode does not implement it — raises - PageIndexAPIError. Use ``chat_completions`` (cloud) instead. + PageIndexAPIError. Use ``chat_completions`` instead. """ return self._require_cloud( "submit_query is cloud-only — the retrieval API is deprecated in " - "favor of chat completions; use chat_completions in cloud mode." + "favor of chat completions; use chat_completions instead." ).submit_query(doc_id=doc_id, query=query, thinking=thinking) def get_retrieval(self, retrieval_id: str) -> dict[str, Any]: @@ -269,55 +338,219 @@ def get_retrieval(self, retrieval_id: str) -> dict[str, Any]: Cloud-only: the cloud API marks this endpoint deprecated in favor of chat completions, so local mode does not implement it — raises - PageIndexAPIError. Use ``chat_completions`` (cloud) instead. + PageIndexAPIError. Use ``chat_completions`` instead. """ return self._require_cloud( "get_retrieval is cloud-only — the retrieval API is deprecated in " - "favor of chat completions; use chat_completions in cloud mode." + "favor of chat completions; use chat_completions instead." ).get_retrieval(retrieval_id=retrieval_id) # ---------- CHAT COMPLETIONS ---------- def chat_completions( self, - messages: list[dict[str, str]], + messages: Union[str, list[dict[str, str]]], stream: bool = False, doc_id: Optional[Union[str, list[str]]] = None, temperature: Optional[float] = None, stream_metadata: bool = False, enable_citations: bool = False, + model: Optional[str] = None, + max_turns: Optional[int] = None, ) -> Union[dict[str, Any], Iterator[str], Iterator[dict[str, Any]]]: """ - PageIndex Chat Completions, scoped to specific PageIndex documents. + PageIndex Chat Completions: document QA in one call. + + Cloud: the hosted chat endpoint. Local: a managed document-QA agent + run over the local tools against your own LLM backend's + /chat/completions (requires ``pageindex[openai]``; the OpenAI SDK's + usual env config — OPENAI_API_KEY, OPENAI_BASE_URL — selects the + backend, so any OpenAI-compatible server works; a ``/`` in the + model name means LiteLLM provider routing, so prefix ``openai/`` + when the backend itself serves slashed ids, e.g. + ``openai/Qwen/...`` on vLLM). The non-stream + response carries the final answer only; streaming yields the + agent's visible text as it is produced, including narration before + tool calls. ``finish_reason`` reports loop completion ("stop") — + the engine does not surface per-turn backend finish reasons. For + the tool-use process and prompt-cache round-trip use + ``responses()`` or ``messages()``. Args: - messages: Conversation messages with 'role' and 'content' keys. + messages: Conversation messages with 'role' and 'content' keys, + or a bare query string (it becomes a single user message). + Local also accepts system/developer messages — their content + is appended to the managed system prompt. stream: Enable streaming responses. doc_id: Document ID or list of IDs to scope the conversation. - temperature: Sampling temperature (0.0-1.0). + Keep it identical across a conversation's calls — the + targeting block it adds is re-set each call and is part + of the cached prompt prefix. + temperature: Sampling temperature, passed through to the model. stream_metadata: With stream=True, yield chunk dicts instead of text pieces. - enable_citations: Enable citation instructions in responses. + enable_citations: Cloud-only — local mode raises (citations need + block-level OCR data local mode does not store). + model: Local only — backend model name (defaults to + ``retrieve_model``). The cloud endpoint selects its own. + max_turns: Local only — cap on agent turns per call. Returns: - stream=False: complete response dict ({'id', 'object', 'created', 'choices', 'usage'}) - stream=True, stream_metadata=False: iterator of text chunks - stream=True, stream_metadata=True: iterator of chunk dicts - - Local: not yet supported — raises PageIndexAPIError. Agent-based - local chat arrives in a later release. """ - return self._require_cloud( - "chat_completions is not yet supported in local mode — it arrives " - "in a later release. Create the client with an api_key to use " - "cloud chat." - ).chat_completions( + if isinstance(messages, str): + if not messages.strip(): + raise PageIndexAPIError( + "messages must be a non-empty string or a list of " + "message dicts.") + messages = [{"role": "user", "content": messages}] + from .cloud_api import CloudAPI + if not isinstance(self._api, CloudAPI): + from .local_chat import run_chat_completions + return run_chat_completions( + self, messages, stream=stream, doc_id=doc_id, + temperature=temperature, stream_metadata=stream_metadata, + enable_citations=enable_citations, model=model, + max_turns=max_turns, + ) + if model is not None or max_turns is not None: + raise PageIndexAPIError( + "model and max_turns are local-mode parameters — the cloud " + "chat endpoint selects its own model." + ) + return self._api.chat_completions( messages=messages, stream=stream, doc_id=doc_id, temperature=temperature, stream_metadata=stream_metadata, enable_citations=enable_citations, ) + def responses( + self, + input: Union[str, list[dict[str, Any]]], + model: Optional[str] = None, + stream: bool = False, + doc_id: Optional[Union[str, list[str]]] = None, + instructions: Optional[str] = None, + temperature: Optional[float] = None, + top_p: Optional[float] = None, + max_turns: Optional[int] = None, + ) -> Union[dict[str, Any], Iterator[dict[str, Any]]]: + """ + Document QA over the OpenAI Responses protocol — the agentic surface. + + Local only for now. Drives your backend's /responses end to end (no + translation layer). The envelope is official Responses shape — + ``output`` carries the model-produced items and parses with the + openai SDK types — and the whole process transcript (including the + tool outputs the SDK executed) rides in the extra ``items`` field. + Append the returned ``items`` to your next call's ``input`` verbatim + to keep provider prompt-cache prefix continuity and the agent's + memory of what it already read. + + Requires ``pageindex[openai]`` and a backend that supports the + Responses API; backends that only speak chat.completions should use + ``chat_completions()``. Provider-prefixed models (``anthropic/…``) + route through LiteLLM's chat.completions adapter and are therefore + refused here — use ``chat_completions()`` or ``messages()`` for + those. + + Args: + input: A user message string, or a list of Responses input items + (round-trip prior ``items`` here). + model: Backend model name (defaults to ``retrieve_model``). + stream: Yield Responses stream events as dicts — one logical + response per call: per-turn backend lifecycle events are + collapsed, sequence numbers are reassigned monotonically, + and ``output_index`` is re-based onto the single logical + ``output``. The single final event is the terminal + ``response.*`` for the run's status; its ``response`` + carries the tool outputs in ``items``. + doc_id: Document ID or list of IDs to scope the conversation. + Keep it identical across a conversation's calls — the + targeting block it adds is re-set each call and is part + of the cached prompt prefix. + instructions: Appended to the managed system prompt. + temperature / top_p: Passed through to the model. + max_turns: Cap on agent turns per call. + """ + from .cloud_api import CloudAPI + if isinstance(self._api, CloudAPI): + raise PageIndexAPIError( + "responses is not available on PageIndex cloud yet — it is " + "a local-mode surface for now." + ) + from .local_chat import run_responses + return run_responses( + self, input, model=model, stream=stream, doc_id=doc_id, + instructions=instructions, temperature=temperature, top_p=top_p, + max_turns=max_turns, + ) + + def messages( + self, + messages: Union[str, list[dict[str, Any]]], + model: str, + max_tokens: Optional[int] = None, + stream: bool = False, + doc_id: Optional[Union[str, list[str]]] = None, + system: Optional[Union[str, list[dict[str, Any]]]] = None, + temperature: Optional[float] = None, + top_p: Optional[float] = None, + top_k: Optional[int] = None, + stop_sequences: Optional[list[str]] = None, + max_turns: Optional[int] = None, + ) -> Union[dict[str, Any], Iterator[Any]]: + """ + Document QA over the Anthropic Messages protocol — Claude-native. + + Local only for now. Drives Anthropic's /v1/messages via the + Anthropic SDK's own tool runner (requires ``pageindex[anthropic]``; + ANTHROPIC_API_KEY selects the backend). ``tool_use``/``tool_result`` + round-trip is the format's native behavior: the response is the + final message envelope with cross-turn aggregated ``usage`` plus a + ``messages`` field — the full new turn sequence, valid for verbatim + append to your history. The managed system prompt carries a + ``cache_control`` breakpoint. + + Args: + messages: Native Messages-format history (including prior + tool_use/tool_result blocks on round-trip), or a bare query + string (it becomes a single user message). + model: Required — there is no cross-vendor default to guess. + max_tokens: Per-turn output budget the Messages API requires on + the wire; the default is resolved per model (8192, or 4096 + for the claude-3 generation whose ceiling is lower) so the + simple call needs only a question. Passed through. + stream: Yield the Anthropic SDK's event stream across turns + (its native event objects, including SDK-synthesized + convenience events), one message sequence per turn. + doc_id: Document ID or list of IDs to scope the conversation. + Keep it identical across a conversation's calls — the + targeting block it adds is re-set each call. + system: Appended after the managed system blocks. + temperature / top_p / top_k / stop_sequences: Passed through. + max_turns: Cap on agent turns per call (default 10, like the + OpenAI surfaces). A truncated run reports + ``stop_reason: "tool_use"`` and its ``messages`` remain + valid for continuation. + """ + from .cloud_api import CloudAPI + if isinstance(self._api, CloudAPI): + raise PageIndexAPIError( + "messages is not available on PageIndex cloud yet — it is " + "a local-mode surface for now." + ) + from .local_chat import run_messages + return run_messages( + self, messages, model=model, max_tokens=max_tokens, + stream=stream, doc_id=doc_id, system=system, + temperature=temperature, top_p=top_p, top_k=top_k, + stop_sequences=stop_sequences, max_turns=max_turns, + ) + # ---------- DOCUMENT MANAGEMENT ---------- def get_document(self, doc_id: str) -> dict[str, Any]: @@ -365,6 +598,332 @@ def list_documents( """ return self._api.list_documents(limit=limit, offset=offset, folder_id=folder_id) + # ---------- AGENT INTEGRATION ---------- + + def agent_tools(self, include_management: bool = False) -> list[Callable[..., str]]: + """ + Plain functions for any agent framework (LangChain, PydanticAI, ...). + For the OpenAI / Claude Agent SDKs, prefer ``as_openai_tools()`` / + ``as_claude_mcp()``. + + Cloud: the full cloud tool set, discovered live from the PageIndex + MCP server when this method is called — one function per tool, + signature and docstring synthesized from the server's schemas, calls + executed from your process over MCP. Raises PageIndexAPIError if the + server cannot be reached. Local: the built-in tools over the local + store (``browse_documents``, ``get_document``, + ``get_document_structure``, ``get_page_content``). + + Each function takes JSON-serializable arguments, returns a JSON + string, and reports failures inside that JSON instead of raising. + + Args: + include_management (bool): Also expose tools that modify the + library. Local: adds ``remove_document``. Cloud: by default + only tools the server marks read-only are exposed; True + exposes the server's complete list (upload, delete, ...). + """ + from .agent_tools import build_agent_tools + return build_agent_tools(self, include_management) + + def as_openai_tools(self, include_management: bool = False, + hosted: bool = False, + doc_id: Optional[Union[str, list[str]]] = None) -> list: + """ + Tools for the OpenAI Agents SDK — pass to ``Agent(tools=...)`` + (or ``openai_agent_config()`` for all the Agent slots in one + call). + + Cloud (default): the full live read tool set (search, folders, + images — as enabled for your key) as plain function tools, + discovered from the PageIndex MCP server and executed from your + process — works with any model backend. Binary tool results + (e.g. ``get_document_image``) arrive as text placeholder stubs + on this in-process path. Pass ``hosted=True`` to + hand the connection to OpenAI instead: one hosted MCP tool, tool + calls executed server-side (lowest latency; requires an + OpenAI-hosted model on the Responses API). The framework's own + ``MCPServerStreamableHttp`` — ``params={"url": + f"{BASE_URL}/mcp?tools=read", "headers": {"Authorization": + "Bearer "}}`` (drop ``?tools=read`` for + the full tool set) — is the async-native alternative for its + ``mcp_servers=`` slot. + + Local: the in-process tools, any model backend; ``hosted`` does + not apply. + + Requires ``openai-agents`` (``pip install 'pageindex[openai]'``), + imported only when this method is called. + + Args: + include_management (bool): Also expose tools that modify the + library (delete, upload). Default off: the in-process + cloud default serves only server-annotated read-only + tools, and ``hosted=True`` connects OpenAI to the + read-only endpoint (``/mcp?tools=read``) instead. + hosted (bool): Cloud only — hand the MCP connection to OpenAI + for server-side tool execution (OpenAI models only). + doc_id: Local only — restrict the tools to this document ID + (or list of IDs), enforced at the tool layer: out-of-scope + lookups return NOT_FOUND. Raises on cloud, where scoping + is server-side. + """ + from .integrations.openai_agents import build_openai_tools + return build_openai_tools(self, include_management, hosted, + doc_ids=doc_id) + + def _local_doc_scope(self, doc_id): + """doc_id for the tool layer: passed through locally (structural + allowlist), dropped on cloud where scoping is server-side and the + config helpers keep prompt-level targeting.""" + if not getattr(self, "api_key", None): + return doc_id + if doc_id is not None and not doc_id: + # Cloud has no tool-layer allowlist to make an empty scope mean + # "nothing"; dropping it would silently mean "everything". + raise PageIndexAPIError( + "doc_id is empty. Pass one or more document IDs, or omit " + "doc_id to give the agent the whole library.") + return None + + def openai_agent_config( + self, + doc_id: Optional[Union[str, list[str]]] = None, + include_management: bool = False, + model: Optional[str] = None, + ) -> dict[str, Any]: + """ + Document QA ``Agent`` kwargs for the OpenAI Agents SDK in one + call:: + + agent = Agent(**client.openai_agent_config()) + + Sugar over the explicit form — ``agent_instructions`` (with + ``doc_id`` targeting) as the instructions and + ``as_openai_tools`` as the tools; local clients also carry their + configured ``retrieve_model`` (cloud omits ``model`` so the + framework default applies). To customize further, switch to + those methods directly. + + Args: + doc_id: Document ID or list of IDs to target, as in + ``agent_instructions``. Local: also enforced at the tool + layer, not just prompted. Cloud: prompt-level targeting + (tool scoping is server-side). + include_management (bool): Also expose tools that modify the + library. + model: Backend model name; overrides the local default. + """ + from .agent_tools import build_agent_instructions + scope = self._local_doc_scope(doc_id) + config: dict[str, Any] = { + "name": "PageIndex", + "instructions": build_agent_instructions(self, doc_id, + scoped=scope is not None), + "tools": self.as_openai_tools(include_management, doc_id=scope), + } + model = model or getattr(self, "retrieve_model", None) + if model: + config["model"] = model + return config + + def as_anthropic_tools(self, include_management: bool = False, + asynchronous: bool = False, + doc_id: Optional[Union[str, list[str]]] = None, + ) -> list: + """ + Runnable tools for the Anthropic SDK's tool runner — pass to + ``client.beta.messages.tool_runner(tools=...)`` (or + ``anthropic_runner_config()`` for the whole setup in one call). + The default flavor is for the sync ``Anthropic`` client; pass + ``asynchronous=True`` for ``AsyncAnthropic``. For a manual + ``messages.create`` loop, serialize with + ``[tool.to_dict() for tool in ...]``. + + Cloud: the full live read tool set (search, folders, images — as + enabled for your key), discovered from the PageIndex MCP server + and executed from your process; the server's input schemas pass + through verbatim (MCP and the Messages API share the schema + shape), and binary tool results (e.g. ``get_document_image``) + arrive as text placeholder stubs on this in-process path. The + server-side alternative is the Messages API's beta + MCP connector — ``mcp_servers=[{"type": "url", "name": + "pageindex", "url": f"{BASE_URL}/mcp?tools=read", + "authorization_token": }]`` (drop + ``?tools=read`` for the full tool set) — with no client-side + tools involved. Local: the in-process tools — the same set + ``messages()`` runs internally. + + Requires ``anthropic>=0.108.0`` + (``pip install 'pageindex[anthropic]'``), imported only when this + method is called. + + Args: + include_management (bool): Also expose tools that modify the + library. Local: adds ``remove_document``. Cloud: by default + only tools the server marks read-only are exposed; True + exposes the server's complete list (upload, delete, ...). + asynchronous (bool): Build ``beta_async_tool`` runnables for + ``AsyncAnthropic`` (each tool call runs in a worker + thread, keeping blocking I/O off your event loop). The + sync and async runners each accept only their own flavor. + doc_id: Local only — restrict the tools to this document ID + (or list of IDs), enforced at the tool layer: out-of-scope + lookups return NOT_FOUND. Raises on cloud, where scoping + is server-side. + """ + from .integrations.anthropic_sdk import build_anthropic_tools + return build_anthropic_tools(self, include_management, asynchronous, + doc_ids=doc_id) + + def anthropic_runner_config( + self, + model: str, + doc_id: Optional[Union[str, list[str]]] = None, + include_management: bool = False, + asynchronous: bool = False, + max_tokens: Optional[int] = None, + max_turns: Optional[int] = None, + ) -> dict[str, Any]: + """ + Document QA ``tool_runner`` kwargs for the Anthropic SDK in one + call — only your ``messages`` remain:: + + runner = anthropic_client.beta.messages.tool_runner( + **client.anthropic_runner_config(model="claude-sonnet-4-5"), + messages=[{"role": "user", "content": "..."}], + ) + + Sugar over the explicit form — ``agent_instructions`` (with + ``doc_id`` targeting) as the system prompt and + ``as_anthropic_tools`` as the tools — plus the same defaults + ``messages()`` applies: a per-model ``max_tokens`` and a + ``max_iterations`` bound of 10. To customize further, switch to + those methods directly. + + Args: + model: Backend model name (also resolves the ``max_tokens`` + default). + doc_id: Document ID or list of IDs to target, as in + ``agent_instructions``. Local: also enforced at the tool + layer, not just prompted. Cloud: prompt-level targeting + (tool scoping is server-side). + include_management (bool): Also expose tools that modify the + library. + asynchronous (bool): Build async runnables for + ``AsyncAnthropic``. + max_tokens: Per-turn output budget; default resolved per + model. + max_turns: Agent-loop bound; default 10. + """ + from .agent_tools import build_agent_instructions + from .local_chat import _default_max_tokens + scope = self._local_doc_scope(doc_id) + return { + "model": model, + "max_tokens": (max_tokens if max_tokens is not None + else _default_max_tokens(model)), + "system": build_agent_instructions(self, doc_id, + scoped=scope is not None), + "tools": self.as_anthropic_tools(include_management, asynchronous, + doc_id=scope), + "max_iterations": max_turns if max_turns is not None else 10, + } + + def as_claude_mcp(self, include_management: bool = False, + doc_id: Optional[Union[str, list[str]]] = None): + """ + ``mcp_servers`` entry for the Claude Agent SDK. + + Cloud: returns the remote PageIndex MCP config. + ``include_management`` picks the endpoint, so the URL itself is + the gate — the default connects to the read-only endpoint + (``/mcp?tools=read``: the server registers only read-only tools), + ``True`` connects to the full tool set. Local: returns an + in-process SDK MCP server exposing the agent tools, gated the + same way at registration (requires ``claude-agent-sdk``; + ``pip install 'pageindex[claude]'``). ``doc_id`` (local only) + restricts those tools to that document ID (or list), enforced at + the tool layer; it raises on cloud, where scoping is server-side. + + Cloud hosts that surface MCP server instructions receive the same + guidance ``agent_instructions()`` returns natively — passing both + duplicates the text (harmless). ``system_prompt`` stays the + recommended channel: it is guaranteed delivery, carries ``doc_id`` + targeting, and is the only channel local mode has. + + Usage (or ``claude_agent_config()`` for all three slots in one + call):: + + options = ClaudeAgentOptions( + system_prompt=client.agent_instructions(), + mcp_servers={"pageindex": client.as_claude_mcp()}, + # Pre-approval only — the server itself is already gated. + allowed_tools=["mcp__pageindex"], + ) + """ + from .integrations.claude_agent_sdk import build_claude_mcp + return build_claude_mcp(self, include_management, doc_ids=doc_id) + + def claude_agent_config( + self, + doc_id: Optional[Union[str, list[str]]] = None, + include_management: bool = False, + server_name: str = "pageindex", + ) -> dict[str, Any]: + """ + Document QA ``ClaudeAgentOptions`` kwargs in one call:: + + options = ClaudeAgentOptions(**client.claude_agent_config()) + + Sugar over the explicit form — the managed system prompt + (``agent_instructions``) and the server entry (``as_claude_mcp``, + itself the tool gate) with its ``allowed_tools`` pre-approval, + one ``include_management`` and ``server_name`` applied + everywhere. To customize (your own system prompt, extra + servers), switch to those methods directly. + + Args: + doc_id: Document ID or list of IDs to target, as in + ``agent_instructions``. Local: also enforced at the tool + layer, not just prompted. Cloud: prompt-level targeting + (tool scoping is server-side). + include_management (bool): Also allow tools that modify the + library. + server_name (str): Key the server is registered under. + """ + from .agent_tools import build_agent_instructions + scope = self._local_doc_scope(doc_id) + return { + "system_prompt": build_agent_instructions(self, doc_id, + scoped=scope is not None), + "mcp_servers": {server_name: self.as_claude_mcp( + include_management, doc_id=scope)}, + # Pre-approval only — the server itself is already gated (the + # read-only endpoint on cloud, the registered set locally). + "allowed_tools": [f"mcp__{server_name}"], + } + + def agent_instructions(self, doc_id: Optional[Union[str, list[str]]] = None) -> str: + """ + Orchestration guidance for document QA agents — pass as the agent's + system prompt (or append to your own). + + Cloud: the live instructions the PageIndex MCP server serves for + your key's tool set, fetched over the same session as + ``agent_tools()`` — server-side guidance updates arrive without an + SDK release. Raises PageIndexAPIError if the server cannot be + reached. Local: the built-in guidance for the in-process tools. + + With ``doc_id`` (str or list, same shape as ``chat_completions``), + appends the target documents' names and metadata and directs the + agent to work within them. Raises PageIndexAPIError if a doc_id does + not exist, or if its name is shadowed by a newer same-name document + (the name-addressed tools could not reach it). + """ + from .agent_tools import build_agent_instructions + return build_agent_instructions(self, doc_id) + # ---------- FOLDER MANAGEMENT ---------- def create_folder( diff --git a/pageindex/cloud_api.py b/pageindex/cloud_api.py index b7597cc9a..ab7a9c885 100644 --- a/pageindex/cloud_api.py +++ b/pageindex/cloud_api.py @@ -55,7 +55,9 @@ def submit_document( returned in get_tree/get_ocr responses and list_documents entries. Defaults to None. Returns: - dict: {'doc_id': ...} + dict: {'doc_id': ...} — plus 'name', the stored document name + (a taken name gains a numeric suffix), when the server + returns it. """ data = {'if_retrieval': True} if mode is not None: diff --git a/pageindex/flash/README.md b/pageindex/flash/README.md index 99d236181..0d23a3c35 100644 --- a/pageindex/flash/README.md +++ b/pageindex/flash/README.md @@ -11,9 +11,9 @@ an LLM. ```python from pageindex.flash import page_index_flash -tree = page_index_flash("paper.pdf") -tree = page_index_flash("paper.pdf", summary=False) # tree structure only, no LLM -tree = page_index_flash("paper.pdf", optimize=True) # refined tree for retrieval +tree = page_index_flash("paper.pdf") # optimized tree + summaries +tree = page_index_flash("paper.pdf", summary=False, optimize=False) # raw tree only, no LLM +tree = page_index_flash("paper.pdf", optimize="merge") # deterministic merge, no LLM expand ``` Takes a file path or an `io.BytesIO` stream and returns the tree as a dict. @@ -22,12 +22,11 @@ Summaries are on by default and need an LLM API key. ### Command line ```bash -python3 run_pageindex.py --pdf_path document.pdf --flash -python3 run_pageindex.py --pdf_path document.pdf --flash --no-summary # tree structure only, no LLM -python3 run_pageindex.py --pdf_path document.pdf --flash --optimize # refined tree for retrieval +python3 run_pageindex.py --mode flash --pdf_path document.pdf # optimized tree + summaries +python3 run_pageindex.py --mode flash --pdf_path document.pdf --no-summary --optimize off # raw tree only, no LLM ``` -Writes the tree to `results/_structure_flash.json`. +Writes the tree to `results/_structure.json`. ## Output diff --git a/pageindex/flash/api.py b/pageindex/flash/api.py index 74656d162..bf62d9657 100644 --- a/pageindex/flash/api.py +++ b/pageindex/flash/api.py @@ -96,15 +96,24 @@ def _optimize(structure, page_texts, do_expand, model): def page_index_flash(pdf, summary=True, summary_model=None, - optimize=False, optimize_expand=True, + optimize: str | bool = "full", optimize_expand=None, optimize_model=None, summary_concurrency=None, use_embedded_toc=True) -> dict: - """Build a PageIndex tree structure from a PDF using layout statistics, without an LLM. Args: pdf: path to a PDF file (``str`` or ``pathlib.Path``) or an in-memory binary stream (``io.BytesIO``). summary: if True, generate LLM summaries for each node (requires ``summary_model``). summary_model: the LLM model identifier to use for summary generation. optimize: if True, refine the tree for search cost before summaries: a deterministic merge collapses subtrees whose structure does not beat a linear scan, keeping the removed titles on the parent as ``key_items``, then an LLM pass expands oversized sections. Without it the extracted tree is returned unchanged. optimize_expand: if False, run the merge but skip the LLM expansion. optimize_model: the LLM model for expand (defaults to the summary model). summary_concurrency: maximum simultaneous summary model calls; None uses the library default. use_embedded_toc: if True, consume the PDF's embedded bookmarks when trustworthy: deep bookmarks become the frame and the detected sections they lack are grafted back in after noise filtering, coarse ones become the chapter frame with detected nodes re-hung under them (deeper sparse entries are filled in when the page text confirms them, and garbled extracted titles are repaired from the bookmark strings), garbage ones are ignored; adds a ``toc_source`` key to the result. On by default; pass False for the pure detected structure. Returns: dict with keys ``doc_name``, ``doc_title``, ``structure`` (a list of nested ``{"title", "start_index", "end_index", "nodes"}`` dicts; page indexes are 1-based) and ``has_abstract_or_references_section`` (True when a top-level entry is an abstract or references heading). With ``optimize`` an ``optimize`` key reports merge/expand counts and before/after search-cost metrics. """ + """Build a PageIndex tree structure from a PDF using layout statistics, without an LLM. Args: pdf: path to a PDF file (``str`` or ``pathlib.Path``) or an in-memory binary stream (``io.BytesIO``). summary: if True, generate LLM summaries for each node (requires ``summary_model``). summary_model: the LLM model identifier to use for summary generation. optimize: ``"full"`` for merge + LLM expand, ``"merge"`` for deterministic merge only, ``False`` to disable. ``True`` is accepted as ``"full"`` for backward compatibility. optimize_model: the LLM model for expand (defaults to the summary model). summary_concurrency: maximum simultaneous summary model calls; None uses the library default. use_embedded_toc: if True, consume the PDF's embedded bookmarks when trustworthy: deep bookmarks become the frame and the detected sections they lack are grafted back in after noise filtering, coarse ones become the chapter frame with detected nodes re-hung under them (deeper sparse entries are filled in when the page text confirms them, and garbled extracted titles are repaired from the bookmark strings), garbage ones are ignored; adds a ``toc_source`` key to the result. On by default; pass False for the pure detected structure. Returns: dict with keys ``doc_name``, ``doc_title``, ``structure`` (a list of nested ``{"title", "start_index", "end_index", "nodes"}`` dicts; page indexes are 1-based) and ``has_abstract_or_references_section`` (True when a top-level entry is an abstract or references heading). With ``optimize`` an ``optimize`` key reports merge/expand counts and before/after search-cost metrics. """ + if optimize is True: + optimize = "full" + if not optimize: + optimize = False + elif optimize not in ("full", "merge"): + raise ValueError( + f"optimize must be 'full', 'merge', or False, got {optimize!r}") + if optimize_expand is not None and optimize: + optimize = "full" if optimize_expand else "merge" result = extract_toc(_validate_pdf(pdf), use_embedded_toc=use_embedded_toc) structure = result.get("structure", []) if optimize and structure: result["optimize"] = _optimize(structure, result.get("page_texts") or [], - optimize_expand, + optimize == "full", optimize_model or summary_model) if summary and structure: import asyncio diff --git a/pageindex/integrations/__init__.py b/pageindex/integrations/__init__.py new file mode 100644 index 000000000..e42ccf64c --- /dev/null +++ b/pageindex/integrations/__init__.py @@ -0,0 +1,5 @@ +"""Framework adapters for the agent tools layer. + +These modules import their target frameworks lazily, at call time — the +frameworks are never required to install or import pageindex. +""" diff --git a/pageindex/integrations/anthropic_sdk.py b/pageindex/integrations/anthropic_sdk.py new file mode 100644 index 000000000..089b0809f --- /dev/null +++ b/pageindex/integrations/anthropic_sdk.py @@ -0,0 +1,59 @@ +"""Anthropic SDK adapter for the tool runner's tools=... slot. + +Cloud clients get one runnable tool per live cloud MCP tool — the server's +input schemas pass through verbatim (MCP inputSchema and Messages API +input_schema are the same shape), calls proxied over MCP. Local clients get +the in-process tools — the same set messages() runs internally. Failed +calls raise ToolError so the runner emits the tool_result with +``is_error: true`` and the envelope as its content. +""" +from __future__ import annotations + +import asyncio +from typing import Any + +from ..errors import PageIndexAPIError + + +def build_anthropic_tools(client, include_management: bool = False, + asynchronous: bool = False, doc_ids=None) -> list: + try: + from anthropic import beta_async_tool, beta_tool + from anthropic.lib.tools import ToolError + except ImportError as exc: + raise PageIndexAPIError( + "as_anthropic_tools requires the Anthropic SDK tool runner " + "(anthropic>=0.108.0) — pip install -U anthropic (or pip install " + "'pageindex[anthropic]')." + ) from exc + from ..agent_tools import _tool_specs + + def wrap(name, description, schema, invoke): + """One runnable tool in the caller's flavor: the sync runner and the + async runner each accept only their own kind, and the async variant + moves the blocking bridge/store call into a worker thread so it + never blocks the caller's event loop.""" + def run(kwargs: dict) -> str: + text, is_error = invoke(kwargs) + if is_error: + raise ToolError(text) + return text + + if asynchronous: + async def _afn(**kwargs: Any) -> str: + return await asyncio.to_thread(run, kwargs) + + _afn.__name__ = name + return beta_async_tool(_afn, name=name, description=description, + input_schema=schema) + + def _fn(**kwargs: Any) -> str: + return run(kwargs) + + _fn.__name__ = name + return beta_tool(_fn, name=name, description=description, + input_schema=schema) + + return [wrap(*spec) + for spec in _tool_specs(client, include_management, + doc_ids=doc_ids)] diff --git a/pageindex/integrations/claude_agent_sdk.py b/pageindex/integrations/claude_agent_sdk.py new file mode 100644 index 000000000..8c76cb434 --- /dev/null +++ b/pageindex/integrations/claude_agent_sdk.py @@ -0,0 +1,67 @@ +"""Claude Agent SDK adapter: one value for the mcp_servers slot. + +Cloud clients get the remote PageIndex MCP config — the framework connects +directly, and include_management picks the endpoint (the read-only +``?tools=read`` URL by default); local clients get an in-process SDK MCP +server over the same tool contract, gated the same way at registration. +""" +from __future__ import annotations + +import asyncio +from typing import Any + +from .._version import sdk_version +from ..errors import PageIndexAPIError + + +def build_claude_mcp(client, include_management: bool = False, doc_ids=None): + from ..agent_tools import _require_local_scope + _require_local_scope(client, doc_ids) + if getattr(client, "api_key", None): + # include_management picks the endpoint — the URL itself is the + # gate (?tools=read serves only readOnlyHint-annotated tools). + suffix = "" if include_management else "?tools=read" + return { + "type": "http", + "url": f"{client.BASE_URL}/mcp{suffix}", + "headers": {"Authorization": f"Bearer {client.api_key}"}, + } + + try: + from claude_agent_sdk import create_sdk_mcp_server, tool + except ImportError as exc: + raise PageIndexAPIError( + "as_claude_mcp in local mode requires the Claude Agent SDK — " + "pip install claude-agent-sdk (or pip install 'pageindex[claude]')." + ) from exc + from ..agent_tools import (TOOL_CONTRACT, _local_description, + _local_schema, call_tool, tool_names) + + def make_handler(name: str): + async def handler(arguments: dict[str, Any]) -> dict[str, Any]: + text, is_error = await asyncio.to_thread( + call_tool, client, name, arguments or {}, doc_ids + ) + result: dict[str, Any] = {"content": [{"type": "text", "text": text}]} + if is_error: + result["is_error"] = True + return result + return handler + + def tool_kwargs(name: str) -> dict: + annotations = TOOL_CONTRACT[name].get("annotations") + if not annotations: + return {} + try: + from claude_agent_sdk import ToolAnnotations + except ImportError: + return {} + return {"annotations": ToolAnnotations(**annotations)} + + tools = [ + tool(name, _local_description(name), + _local_schema(name), **tool_kwargs(name))(make_handler(name)) + for name in tool_names(include_management) + ] + return create_sdk_mcp_server(name="pageindex", version=sdk_version(), + tools=tools) diff --git a/pageindex/integrations/openai_agents.py b/pageindex/integrations/openai_agents.py new file mode 100644 index 000000000..36c062d2f --- /dev/null +++ b/pageindex/integrations/openai_agents.py @@ -0,0 +1,78 @@ +"""OpenAI Agents SDK adapter for the Agent(tools=...) slot. + +Cloud clients default to the live read tool set as plain FunctionTools via +the MCP bridge; pass hosted=True to use a single HostedMCPTool instead +(the model connects to the PageIndex cloud MCP server from OpenAI's side — +the read-only ``?tools=read`` endpoint by default). Local clients get the +in-process tools wrapped as FunctionTools. Tools are built as FunctionTool +directly so the contract/server JSON schema goes to the model verbatim — +function_tool() would regenerate it from a Python signature, dropping +items/enum/pattern/bounds and rejecting object-typed parameters. +""" +from __future__ import annotations + +import asyncio +import json +from typing import Any + +from ..errors import PageIndexAPIError + + +def build_openai_tools(client, include_management: bool = False, + hosted: bool = False, doc_ids=None) -> list: + try: + from agents import FunctionTool, HostedMCPTool + except ImportError as exc: + raise PageIndexAPIError( + "as_openai_tools requires the OpenAI Agents SDK — " + "pip install openai-agents (or pip install 'pageindex[openai]')." + ) from exc + from ..agent_tools import (_dumps, _failure, _require_local_scope, + _tool_specs) + _require_local_scope(client, doc_ids) + if getattr(client, "api_key", None) and hosted: + # include_management picks the endpoint — the URL itself is the + # gate (?tools=read serves only readOnlyHint-annotated tools), so + # nothing needs the Responses API approval flow. + suffix = "" if include_management else "?tools=read" + return [HostedMCPTool(tool_config={ + "type": "mcp", + "server_label": "pageindex", + "server_url": f"{client.BASE_URL}/mcp{suffix}", + "headers": {"Authorization": f"Bearer {client.api_key}"}, + "require_approval": "never", + })] + + def wrap(name, description, schema, invoke): + async def on_invoke_tool(ctx: Any, args_json: str) -> str: + # strict_json_schema is off, so the provider never validates the + # payload; a malformed or non-object argument string must come + # back as the guided error envelope — raising here aborts the + # caller's whole run (hand-built FunctionTools have no + # failure_error_function to hand the error back to the model). + try: + parsed = json.loads(args_json) if args_json else {} + except ValueError: + parsed = None + if not isinstance(parsed, dict): + payload, _ = _failure( + f"Invalid arguments for {name}: expected a JSON object, " + f"got: {(args_json or '')[:200]!r}", None, + {"summary": "Malformed tool arguments", + "options": [f"Re-send the {name} call with a JSON " + "object of its parameters"]}, + "INVALID_INPUT") + return _dumps(payload) + arguments = {key: value for key, value in parsed.items() + if value is not None} + text, _ = await asyncio.to_thread(invoke, arguments) + return text + + return FunctionTool(name=name, description=description, + params_json_schema=schema, + on_invoke_tool=on_invoke_tool, + strict_json_schema=False) + + return [wrap(*spec) + for spec in _tool_specs(client, include_management, + doc_ids=doc_ids)] diff --git a/pageindex/local_api.py b/pageindex/local_api.py index 0e9f682c8..9a82c48f9 100644 --- a/pageindex/local_api.py +++ b/pageindex/local_api.py @@ -75,8 +75,10 @@ def submit_document( if mode not in (None, "standard", "flash"): raise PageIndexAPIError( f"Failed to submit document: unknown local processing mode {mode!r}. " - "Supported: None or 'standard' for standard indexing, or 'flash'." + "Supported: 'flash' (default) or 'standard'." ) + if mode is None: + mode = "flash" file_path = os.path.abspath(os.path.expanduser(str(file_path))) if not os.path.isfile(file_path): raise FileNotFoundError(f"No such file: {file_path}") @@ -97,6 +99,7 @@ def submit_document( raise PageIndexAPIError( "Failed to submit document: PDF has no content. All pages are blank." ) + self._unique_doc_name(os.path.basename(file_path)) try: if mode == "flash": @@ -115,21 +118,37 @@ def submit_document( doc_id = "pi-" + uuid.uuid4().hex meta = { "id": doc_id, - "name": os.path.basename(file_path), + "name": self._unique_doc_name(os.path.basename(file_path)), "description": description, "status": "completed", "createdAt": _now_iso(), "pageNum": len(page_texts), "folderId": None, "metadata": metadata, - "mode": mode or "standard", + "mode": mode, } pages = [{"page_index": i + 1, "markdown": text} for i, text in enumerate(page_texts)] from .utils import remove_fields self._store.save_document( doc_id, meta, remove_fields(structure, fields=["text"]), pages) - return {"doc_id": doc_id} + return {"doc_id": doc_id, "name": meta["name"]} + + def _unique_doc_name(self, name: str) -> str: + """Mirror the cloud upload: a taken name gets _1.._99 appended, + beyond that the submit is rejected.""" + taken = {meta.get("name") for meta in self._store.list_metas()} + if name not in taken: + return name + base, ext = os.path.splitext(name) + for num in range(1, 100): + candidate = f"{base}_{num}{ext}" + if candidate not in taken: + return candidate + raise PageIndexAPIError( + "Failed to submit document: Too many files with similar names. " + "Please use a different file name." + ) @staticmethod def _extract_page_texts(file_path: str) -> list[str]: @@ -163,8 +182,16 @@ def _index_flash(self, file_path: str, page_texts: list[str]) -> tuple[list, str from .flash import page_index_flash from .utils import (add_node_text, create_clean_structure_for_description, generate_doc_description, write_node_id) + import litellm + env = litellm.validate_environment(self._summary_model) + if not env["keys_in_environment"]: + raise PageIndexAPIError( + f"Failed to submit document: missing API key for " + f"{self._summary_model}: {', '.join(env['missing_keys'])}") result = page_index_flash(file_path, summary=True, - summary_model=self._summary_model) + summary_model=self._summary_model, + optimize="full", + optimize_model=self._summary_model) structure = result.get("structure", []) if not structure: raise PageIndexAPIError( @@ -190,6 +217,11 @@ def _load_tree_with_text(self, doc_id: str, error_prefix: str) -> list: add_node_text(structure, pdf_pages) return structure + def raw_tree(self, doc_id: str) -> list | None: + """Stored tree verbatim — keeps start_index/end_index, which + get_tree's cloud wire shape renames and drops.""" + return self._store.get_tree(doc_id) + def get_tree(self, doc_id: str, node_summary: bool = False, include_text: bool = True) -> dict[str, Any]: meta = self._require_doc(doc_id, "Failed to get tree result") diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py new file mode 100644 index 000000000..df890ebb6 --- /dev/null +++ b/pageindex/local_chat.py @@ -0,0 +1,848 @@ +"""Managed local chat: document-QA agents over the local tools. + +Three methods, three backend protocols, routed 1:1: ``chat_completions`` +drives the backend's /chat/completions (any OpenAI-compatible backend, +final answer only), ``responses`` drives /responses (official-shape +envelope; the full process transcript rides in ``items`` — round-trip it +for provider prompt-cache continuation and agent memory), ``messages`` +drives Anthropic's /v1/messages via the SDK's own tool runner +(tool_use/tool_result round-trip is the format's native behavior). + +Content passes through untouched — the caller's messages, the model's +answers, tool outputs. Native stop reasons pass through on ``messages``; +the OpenAI engine's abstraction does not surface per-turn finish reasons, +so ``chat_completions`` reports loop completion as ``"stop"``, while +``responses`` reports the backend's terminal ``status`` where the wire +surfaces one (recorded at the transport layer — the framework discards +it). The SDK owns gatekeeping (structural validation), table-setting +(managed instructions, tools, doc targeting), tool execution, and billing +(usage aggregation, envelope ids). +""" +from __future__ import annotations + +import asyncio +import concurrent.futures +import hashlib +import json +import queue +import threading +import time +import uuid +from typing import Any, Iterator, Optional, Union + +from .agent_tools import AGENT_INSTRUCTIONS, doc_targeting_block +from .errors import PageIndexAPIError + +CHAT_HEADER = ( + "You are PageIndex by Vectify AI, a document-focused assistant. " + "Be concise, never use emojis, and do not expose tool names." +) + + +# ── shared: prompt, doc targeting, validation, sync bridges ── + +def _managed_instructions(extra_system: list[str]) -> str: + return "\n\n".join([CHAT_HEADER, AGENT_INSTRUCTIONS, *extra_system]) + + +def _doc_block(client, doc_id) -> Optional[str]: + if doc_id is None: + return None + if not isinstance(doc_id, (str, list)): + raise PageIndexAPIError("doc_id must be a string or a list of " + "strings.") + doc_ids = [doc_id] if isinstance(doc_id, str) else list(doc_id) + missing = [] + for one_id in doc_ids: + try: + client.get_document(one_id) + except PageIndexAPIError: + missing.append(str(one_id)) + if missing: + raise PageIndexAPIError( + "Documents not found or access denied: " + ", ".join(missing) + ) + # scoped: the chat surfaces also pass doc_id into the tool layer, so + # name resolution happens inside the allowlist — only a duplicate name + # within the targeted set shadows. + return doc_targeting_block(client, doc_id, scoped=True) + + +def _system_text(content: Any) -> str: + """Text of a system/developer message: a string, or text parts joined.""" + if isinstance(content, str): + return content + if isinstance(content, list): + texts = [part.get("text") for part in content + if isinstance(part, dict) and isinstance(part.get("text"), str)] + if texts: + return "\n".join(texts) + raise PageIndexAPIError( + "system message content must be a string or a list of text parts." + ) + + +def _split_chat_messages(messages) -> "tuple[list[str], list[dict]]": + """Validate the chat_completions surface's messages: system/developer + content joins the managed instructions; user/assistant history passes + through. Tool-history round-trips belong to responses()/messages().""" + if not isinstance(messages, list) or not messages: + raise PageIndexAPIError("messages must be a non-empty list.") + system_texts: list[str] = [] + history: list[dict] = [] + for message in messages: + if not isinstance(message, dict) or "role" not in message: + raise PageIndexAPIError( + "Each message must be a dict with 'role' and 'content'.") + role = message["role"] + if role in ("system", "developer"): + system_texts.append(_system_text(message.get("content"))) + elif role in ("user", "assistant"): + content = message.get("content") + if not isinstance(content, str): + raise PageIndexAPIError( + "chat_completions content must be a string; for " + "structured items use responses() or messages()." + ) + history.append({"role": role, "content": content}) + else: + raise PageIndexAPIError( + f"Unsupported role for chat_completions: {role!r}. Tool " + "history round-trips belong to responses() or messages()." + ) + if not history: + raise PageIndexAPIError("messages must contain a user or assistant " + "message.") + return system_texts, history + + +def _run_sync(coro): + try: + asyncio.get_running_loop() + except RuntimeError: + has_loop = False + else: + has_loop = True + if not has_loop: + return asyncio.run(coro) + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: + return pool.submit(asyncio.run, coro).result() + + +_SENTINEL = object() + + +def _stream_sync(agen_factory) -> Iterator[Any]: + """Drive an async generator from a background thread; yield synchronously. + + Closing the iterator cancels the run between items: the pump stops, and + the async generator's cleanup cancels the underlying agent task, so no + further model turns or tool executions start. An in-flight backend + request cannot be aborted mid-turn. + """ + items: "queue.Queue[Any]" = queue.Queue(maxsize=32) + cancelled = threading.Event() + + def deliver(item) -> bool: + while not cancelled.is_set(): + try: + items.put(item, timeout=0.1) + return True + except queue.Full: + continue + return False + + def pump(): + async def consume(): + agen = agen_factory() + + async def drain(): + async for item in agen: + if not deliver(item): + break + + # The watchdog lets cancellation land even while drain() is + # awaiting the backend — a plain async-for would only notice + # between items. + task = asyncio.ensure_future(drain()) + try: + while not task.done(): + if cancelled.is_set(): + task.cancel() + break + await asyncio.sleep(0.05) + try: + await task + except asyncio.CancelledError: + pass + finally: + await agen.aclose() + + try: + asyncio.run(consume()) + except BaseException as exc: # re-raised on the consumer thread + deliver(exc) + return + deliver(_SENTINEL) + + threading.Thread(target=pump, daemon=True).start() + try: + while True: + item = items.get() + if item is _SENTINEL: + return + if isinstance(item, BaseException): + raise item + yield item + finally: + cancelled.set() + + +# ── OpenAI engine (chat_completions / responses) ── + +def _require_openai_agents(method: str) -> None: + try: + import agents # noqa: F401 + except ImportError as exc: + raise PageIndexAPIError( + f"{method} in local mode requires the OpenAI Agents SDK — " + "pip install openai-agents (or pip install 'pageindex[openai]')." + ) from exc + + +def _openai_model(protocol: str, model_name: str): + """The backend protocol driver — the seam tests replace with a fake. + + ``litellm//`` (the client's normalized retrieve_model + form) and bare ``/`` paths drive the provider through + LiteLLM — chat.completions only, so the responses protocol refuses them + instead of silently downgrading; a first segment LiteLLM does not know + (a HuggingFace repo id like ``Qwen/...``) is refused with the + ``openai/`` escape instead of failing inside LiteLLM at request time; + an ``openai/`` prefix strips to the OpenAI SDK; bare names go to the + OpenAI SDK as-is.""" + if "/" in model_name and not model_name.startswith("openai/"): + if protocol == "responses": + raise PageIndexAPIError( + f"responses() cannot drive " + f"'{model_name.removeprefix('litellm/')}': provider-prefixed " + "models route through LiteLLM, which speaks chat.completions, " + "not the Responses API. Use chat_completions() (or messages() " + "for Anthropic models), or point OPENAI_BASE_URL at a " + "Responses-capable backend and use a bare or " + "'openai/'-prefixed model name." + ) + try: + from agents.extensions.models.litellm_model import LitellmModel + import litellm + except ImportError: + raise PageIndexAPIError( + f"'{model_name}' routes through LiteLLM, but litellm is not " + "installed. Run: pip install 'litellm>=1.30'" + ) + wire = model_name.removeprefix("litellm/") + providers = getattr(litellm, "provider_list", None) + if providers and wire.split("/", 1)[0] not in providers: + raise PageIndexAPIError( + f"'{wire}' routes through LiteLLM, but " + f"'{wire.split('/', 1)[0]}' is not a LiteLLM provider. For an " + "OpenAI-compatible server (vLLM, TGI, Ollama) serving this " + f"model id, use 'openai/{wire}' and point OPENAI_BASE_URL " + "at the server." + ) + return LitellmModel(wire) + import openai + model_name = model_name.removeprefix("openai/") + try: + backend = openai.AsyncOpenAI() + except openai.OpenAIError as exc: + raise PageIndexAPIError( + f"The OpenAI backend is not configured: {exc}") from exc + if protocol == "chat": + from agents.models.openai_chatcompletions import ( + OpenAIChatCompletionsModel) + return OpenAIChatCompletionsModel(model_name, backend) + from agents.models.openai_responses import OpenAIResponsesModel + return OpenAIResponsesModel(model_name, openai_client=backend) + + +def _reported_model(model_name: str) -> str: + """The name the provider actually serves — routing prefixes stripped.""" + return model_name.removeprefix("litellm/").removeprefix("openai/") + + +def _openai_agent(client, protocol: str, model_name: str, instructions: str, + temperature, top_p, doc_ids=None): + from agents import Agent, ModelSettings + from .integrations.openai_agents import build_openai_tools + return Agent( + name="PageIndex", + instructions=instructions, + tools=build_openai_tools(client, doc_ids=doc_ids), + model=_openai_model(protocol, model_name), + model_settings=ModelSettings(temperature=temperature, top_p=top_p), + ) + + +def _validate_max_turns(max_turns) -> None: + if max_turns is not None and (not isinstance(max_turns, int) + or max_turns < 1): + raise PageIndexAPIError("max_turns must be a positive integer.") + + +def _conversation_group_id(model_name: str, instructions: str, items) -> str: + """Stable per-conversation cache-routing key: openai-agents hashes + RunConfig.group_id into the OpenAI prompt_cache_key, and without one it + stamps every run with a fresh key, tagging a round-tripped prefix as a + different cache group. Keyed on the prefix identity — model, + instructions, first conversation item — so a conversation's + continuations share one route without pooling unrelated conversations. + Callers pass the conversation's own items, never the SDK-prepended + doc-targeting block: that block is byte-identical for every + conversation about a document and would pool them all under one key.""" + seed = json.dumps([model_name, instructions, + items[0] if items else None], + sort_keys=True, default=str) + return "pageindex-" + hashlib.sha256(seed.encode()).hexdigest()[:16] + + +def _run_kwargs(max_turns, group_id: str) -> dict: + # No traces — the caller opted into QA, not telemetry. + from agents import RunConfig + kwargs: dict = {"run_config": RunConfig(tracing_disabled=True, + group_id=group_id)} + if max_turns is not None: + kwargs["max_turns"] = max_turns + return kwargs + + +def _record_response_status(agent, recorded: dict) -> None: + """Capture each turn's terminal Response status at the transport client: + openai-agents' non-streaming path discards Response.status, so a final + turn truncated at the output cap would otherwise report as a clean + completion. No-op for backends without an OpenAI responses resource + (the streaming path records from lifecycle events instead).""" + responses = getattr(getattr(getattr(agent, "model", None), "_client", None), + "responses", None) + create = getattr(responses, "create", None) + if create is None: + return + + async def recording_create(*args, **kwargs): + response = await create(*args, **kwargs) + if getattr(response, "status", None): + recorded["status"] = response.status + for field in ("incomplete_details", "error"): + value = getattr(response, field, None) + recorded[field] = (value.model_dump(mode="json") + if hasattr(value, "model_dump") else value) + return response + + responses.create = recording_create + + +async def _aclose_backend(agent) -> None: + """Close the per-call AsyncOpenAI client before its event loop ends — + otherwise httpx tears down pooled connections on a closed loop and + emits 'Task exception was never retrieved' noise.""" + backend = getattr(getattr(agent, "model", None), "_client", None) + close = getattr(backend, "close", None) + if close is not None: + try: + await close() + except Exception: + pass + + +async def _run_closing(agent, coro): + try: + return await coro + finally: + await _aclose_backend(agent) + + +def _wrap_max_turns(max_turns) -> PageIndexAPIError: + limit = max_turns if max_turns is not None else "the default limit" + return PageIndexAPIError( + f"The agent did not finish within max_turns ({limit}). Raise " + "max_turns, or narrow the question." + ) + + +def _usage_sums(raw_responses) -> "tuple[int, int, int, int, int]": + prompt = completion = cached = cache_write = reasoning = 0 + for r in raw_responses: + prompt += r.usage.input_tokens + completion += r.usage.output_tokens + details = getattr(r.usage, "input_tokens_details", None) + cached += getattr(details, "cached_tokens", 0) or 0 + cache_write += getattr(details, "cache_write_tokens", 0) or 0 + details = getattr(r.usage, "output_tokens_details", None) + reasoning += getattr(details, "reasoning_tokens", 0) or 0 + return prompt, completion, cached, cache_write, reasoning + + +def _openai_usage(raw_responses) -> dict: + """Cross-turn sums, chat.completions dialect.""" + prompt, completion, cached, _, reasoning = _usage_sums(raw_responses) + return {"prompt_tokens": prompt, "completion_tokens": completion, + "total_tokens": prompt + completion, + "prompt_tokens_details": {"cached_tokens": cached}, + "completion_tokens_details": {"reasoning_tokens": reasoning}} + + +def _responses_usage(raw_responses) -> dict: + """Cross-turn sums, Responses dialect.""" + prompt, completion, cached, cache_write, reasoning = ( + _usage_sums(raw_responses)) + return {"input_tokens": prompt, + "input_tokens_details": {"cached_tokens": cached, + "cache_write_tokens": cache_write}, + "output_tokens": completion, + "output_tokens_details": {"reasoning_tokens": reasoning}, + "total_tokens": prompt + completion} + + +def run_chat_completions(client, messages, stream: bool = False, + doc_id=None, temperature: Optional[float] = None, + stream_metadata: bool = False, + enable_citations: bool = False, + model: Optional[str] = None, + max_turns: Optional[int] = None, + ) -> Union[dict, Iterator[str], Iterator[dict]]: + if enable_citations: + raise PageIndexAPIError( + "enable_citations is cloud-only — citations need block-level OCR " + "data that local mode does not store." + ) + _require_openai_agents("chat_completions") + _validate_max_turns(max_turns) + system_texts, history = _split_chat_messages(messages) + block = _doc_block(client, doc_id) + items = ([{"role": "user", "content": block}] if block else []) + history + model_name = model or client.retrieve_model + reported_model = _reported_model(model_name) + managed = _managed_instructions(system_texts) + agent = _openai_agent(client, "chat", model_name, managed, + temperature, None, doc_ids=doc_id) + run_kwargs = _run_kwargs(max_turns, + _conversation_group_id(model_name, managed, + history)) + import openai + from agents import Runner + from agents.exceptions import AgentsException, MaxTurnsExceeded + if not stream: + try: + result = _run_sync(_run_closing(agent, + Runner.run(agent, input=items, **run_kwargs))) + except MaxTurnsExceeded as exc: + raise _wrap_max_turns(max_turns) from exc + except AgentsException as exc: + raise PageIndexAPIError( + f"The agent backend failed: {exc}") from exc + except openai.OpenAIError as exc: + raise PageIndexAPIError( + f"The model backend failed: {exc}") from exc + return { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": int(time.time()), + "model": reported_model, + "choices": [{ + "index": 0, + "message": {"role": "assistant", + "content": result.final_output or ""}, + "finish_reason": "stop", + }], + "usage": _openai_usage(result.raw_responses), + } + + chat_id = f"chatcmpl-{uuid.uuid4().hex}" + created = int(time.time()) + + def chunk(delta: dict, finish=None) -> dict: + return { + "id": chat_id, "object": "chat.completion.chunk", + "created": created, "model": reported_model, + "choices": [{"index": 0, "delta": delta, + "finish_reason": finish}], + } + + async def agen(): + from openai.types.responses import ResponseTextDeltaEvent + streamed = Runner.run_streamed(agent, input=items, **run_kwargs) + completed = False + # First yield inside the try: a consumer that stops on the opening + # chunk must still tear the run down via the finally below. + try: + yield chunk({"role": "assistant", "content": ""}) + async for event in streamed.stream_events(): + if (event.type == "raw_response_event" + and isinstance(event.data, ResponseTextDeltaEvent)): + yield chunk({"content": event.data.delta}) + completed = True + except MaxTurnsExceeded as exc: + raise _wrap_max_turns(max_turns) from exc + except AgentsException as exc: + raise PageIndexAPIError( + f"The agent backend failed: {exc}") from exc + except openai.OpenAIError as exc: + raise PageIndexAPIError( + f"The model backend failed: {exc}") from exc + finally: + if not completed and hasattr(streamed, "cancel"): + streamed.cancel() # abandoned/failed: stop the agent task + await _aclose_backend(agent) + yield chunk({}, finish="stop") + yield { + "id": chat_id, "object": "chat.completion.chunk", + "created": created, "model": reported_model, "choices": [], + "usage": _openai_usage(streamed.raw_responses), + } + + if stream_metadata: + return _stream_sync(agen) + return (piece["choices"][0]["delta"]["content"] + for piece in _stream_sync(agen) + if piece.get("choices") + and "content" in piece["choices"][0]["delta"] + and piece["choices"][0]["delta"]["content"]) + + +def run_responses(client, input, model: Optional[str] = None, + stream: bool = False, doc_id=None, + instructions: Optional[str] = None, + temperature: Optional[float] = None, + top_p: Optional[float] = None, + max_turns: Optional[int] = None, + ) -> Union[dict, Iterator[dict]]: + _require_openai_agents("responses") + _validate_max_turns(max_turns) + if isinstance(input, str) and input.strip(): + items = [{"role": "user", "content": input}] + elif (isinstance(input, list) and input + and all(isinstance(item, dict) for item in input)): + items = list(input) + else: + raise PageIndexAPIError("input must be a non-empty string or list " + "of item dicts.") + block = _doc_block(client, doc_id) + conversation = items + if block: + items = [{"role": "user", "content": block}] + items + extra = [instructions] if instructions else [] + model_name = model or client.retrieve_model + managed = _managed_instructions(extra) + agent = _openai_agent(client, "responses", model_name, managed, + temperature, top_p, doc_ids=doc_id) + run_kwargs = _run_kwargs(max_turns, + _conversation_group_id(model_name, managed, + conversation)) + recorded: dict = {} + import openai + from agents import Runner + from agents.exceptions import AgentsException, MaxTurnsExceeded + + def envelope(transcript: list, raw_responses) -> dict: + return { + "id": f"resp_{uuid.uuid4().hex}", + "object": "response", + "created_at": int(time.time()), + "model": _reported_model(model_name), + "status": recorded.get("status") or "completed", + "output": [item for item in transcript + if item.get("type") != "function_call_output"], + "items": transcript, + "usage": _responses_usage(raw_responses), + "instructions": managed, + "tools": [{"type": "function", "name": tool.name, + "description": tool.description, + "parameters": tool.params_json_schema, + "strict": getattr(tool, "strict_json_schema", True)} + for tool in agent.tools], + "tool_choice": "auto", + "parallel_tool_calls": True, + "temperature": temperature, + "top_p": top_p, + "max_output_tokens": None, + "error": recorded.get("error"), + "incomplete_details": recorded.get("incomplete_details"), + "metadata": None, + } + + if not stream: + _record_response_status(agent, recorded) + try: + result = _run_sync(_run_closing(agent, + Runner.run(agent, input=[dict(item) for item in items], + **run_kwargs))) + except MaxTurnsExceeded as exc: + raise _wrap_max_turns(max_turns) from exc + except AgentsException as exc: + raise PageIndexAPIError( + f"The agent backend failed: {exc}") from exc + except openai.OpenAIError as exc: + raise PageIndexAPIError( + f"The model backend failed: {exc}") from exc + transcript = result.to_input_list()[len(items):] + return envelope(transcript, result.raw_responses) + + lifecycle = {"response.created", "response.in_progress", + "response.completed", "response.failed", + "response.incomplete", "response.queued"} + + async def agen(): + streamed = Runner.run_streamed(agent, + input=[dict(item) for item in items], + **run_kwargs) + sequence = 0 + # output_index addresses an item's position in the logical + # response.output (the final envelope's list). Backend events + # carry per-turn indexes that restart at 0 each turn, so they are + # re-based by the count of items already committed by prior turns. + output_offset = 0 + completed = False + try: + async for event in streamed.stream_events(): + if event.type == "raw_response_event": + data = event.data.model_dump(exclude_unset=True) + if data.get("type") in lifecycle: + if data["type"] in ("response.completed", + "response.incomplete", + "response.failed"): + # Per-turn terminal state; the last turn's wins + # and feeds the final envelope below. + state = data.get("response") or {} + for field in ("status", "incomplete_details", + "error"): + recorded[field] = state.get(field) + output_offset += len(state.get("output") or []) + continue + if isinstance(data.get("output_index"), int): + data["output_index"] += output_offset + sequence += 1 + data["sequence_number"] = sequence + yield data + completed = True + except MaxTurnsExceeded as exc: + raise _wrap_max_turns(max_turns) from exc + except AgentsException as exc: + if recorded.get("status") not in ("failed", "incomplete"): + raise PageIndexAPIError( + f"The agent backend failed: {exc}") from exc + completed = True + except openai.OpenAIError as exc: + raise PageIndexAPIError( + f"The model backend failed: {exc}") from exc + finally: + if not completed and hasattr(streamed, "cancel"): + streamed.cancel() # abandoned/failed: stop the agent task + await _aclose_backend(agent) + transcript = streamed.to_input_list()[len(items):] + sequence += 1 + status = recorded.get("status") or "completed" + terminal = {"incomplete": "response.incomplete", + "failed": "response.failed"}.get(status, + "response.completed") + yield {"type": terminal, "sequence_number": sequence, + "response": envelope(transcript, streamed.raw_responses)} + + return _stream_sync(agen) + + +# ── Anthropic engine (messages) ── + +def _require_anthropic() -> None: + try: + import anthropic # noqa: F401 + except ImportError as exc: + raise PageIndexAPIError( + "messages in local mode requires the Anthropic SDK — " + "pip install anthropic (or pip install 'pageindex[anthropic]')." + ) from exc + try: + from anthropic import beta_tool # noqa: F401 + from anthropic.lib.tools import ToolError # noqa: F401 + except ImportError as exc: + raise PageIndexAPIError( + "messages in local mode requires anthropic >= 0.108.0 (the tool " + "runner with ToolError) — pip install -U anthropic." + ) from exc + + +def _anthropic_client(): + """The backend client — the seam tests replace with a fake transport.""" + import anthropic + return anthropic.Anthropic() + + +def _anthropic_system(extra_system, block: Optional[str]) -> list[dict]: + """System blocks: cache_control marks the stable managed prefix only + (the API allows 4 breakpoints total — the varying doc block and caller + blocks must not consume the budget); the doc block and caller system + content follow as their own blocks.""" + blocks = [{"type": "text", + "text": CHAT_HEADER + "\n\n" + AGENT_INSTRUCTIONS, + "cache_control": {"type": "ephemeral"}}] + if block: + blocks.append({"type": "text", "text": block}) + if extra_system is None: + return blocks + if isinstance(extra_system, str): + if extra_system.strip(): + blocks.append({"type": "text", "text": extra_system}) + return blocks + if isinstance(extra_system, list): + return blocks + list(extra_system) + raise PageIndexAPIError("system must be a string or a list of blocks.") + + +def _dump_block(block) -> Any: + """A content block as a plain JSON dict, minus SDK-internal fields the + API rejects (ParsedBetaTextBlock.__api_exclude__, e.g. parsed_output).""" + if hasattr(block, "model_dump"): + exclude = getattr(type(block), "__api_exclude__", None) + return block.model_dump(mode="json", + exclude=set(exclude) if exclude else None) + return block + + +def _dump_message(message) -> dict: + message = dict(message) + content = message.get("content") + if isinstance(content, list): + message["content"] = [_dump_block(item) for item in content] + return message + + +def _anthropic_usage(turns, final_usage: dict) -> dict: + """The final turn's native usage dict with the token counters replaced + by cross-turn sums (None-safe); all other native fields survive.""" + totals = dict(final_usage) + for field in ("input_tokens", "output_tokens", + "cache_creation_input_tokens", "cache_read_input_tokens"): + values = [getattr(turn.usage, field, None) for turn in turns] + counted = [value for value in values if isinstance(value, int)] + if counted: + totals[field] = sum(counted) + return totals + + +_CLAUDE_4096_MODELS = ("claude-3-opus", "claude-3-sonnet", "claude-3-haiku", + "claude-3-5-sonnet-20240620") + + +def _default_max_tokens(model: str) -> int: + """The wire-required per-turn budget when the caller sets none: 8192, + except the claude-3 generation whose output ceiling is 4096.""" + return 4096 if model.startswith(_CLAUDE_4096_MODELS) else 8192 + + +def run_messages(client, messages, model: str, + max_tokens: Optional[int] = None, + stream: bool = False, doc_id=None, system=None, + temperature: Optional[float] = None, + top_p: Optional[float] = None, + top_k: Optional[int] = None, + stop_sequences: Optional[list[str]] = None, + max_turns: Optional[int] = None, + ) -> Union[dict, Iterator[Any]]: + from .integrations.anthropic_sdk import build_anthropic_tools + + _require_anthropic() + import anthropic + _validate_max_turns(max_turns) + if isinstance(messages, str) and messages.strip(): + messages = [{"role": "user", "content": messages}] + if (not isinstance(messages, list) or not messages + or not all(isinstance(message, dict) for message in messages)): + raise PageIndexAPIError("messages must be a non-empty string or a " + "list of message dicts.") + block = _doc_block(client, doc_id) + prepared = [dict(message) for message in messages] + passthrough = {key: value for key, value in { + "temperature": temperature, "top_p": top_p, "top_k": top_k, + "stop_sequences": stop_sequences, + }.items() if value is not None} + runner = _anthropic_client().beta.messages.tool_runner( + max_tokens=(max_tokens if max_tokens is not None + else _default_max_tokens(model)), + messages=prepared, + model=model, + tools=build_anthropic_tools(client, doc_ids=doc_id), + system=_anthropic_system(system, block), + stream=stream, + # Bounded like the OpenAI surfaces (their framework default is 10). + max_iterations=max_turns if max_turns is not None else 10, + **passthrough, + ) + + if stream: + def events() -> Iterator[Any]: + try: + for turn_stream in runner: + for event in turn_stream: + yield event + except anthropic.AnthropicError as exc: + raise PageIndexAPIError( + f"The model backend failed: {exc}") from exc + return events() + + try: + turns = [turn for turn in runner] + except anthropic.AnthropicError as exc: + raise PageIndexAPIError( + f"The model backend failed: {exc}") from exc + if not turns: + raise PageIndexAPIError("The model returned no response.") + captured: dict = {} + + def capture(params): + captured.update(params) + return params + + runner.set_messages_params(capture) + if not captured.get("messages"): + # The conversation is read back through a mutator; if a vendor + # change stops it delivering params, the envelope would silently + # lose the tool turns — fail loudly instead. + raise PageIndexAPIError( + "Could not read the conversation back from the anthropic tool " + "runner — the installed anthropic version is incompatible with " + "this pageindex release." + ) + conversation = list(captured["messages"]) + final = turns[-1] + envelope = final.model_dump(mode="json") + envelope["content"] = [_dump_block(item) for item in final.content] + envelope["usage"] = _anthropic_usage(turns, envelope.get("usage") or {}) + # The full turn sequence (assistant tool_use + user tool_result + final), + # valid for verbatim append to the caller's history. The runner appends + # a turn to its params only when it executed tools from it — content + # carried tool_use blocks and the turn was not a refusal. stop_reason + # alone cannot tell: a max_tokens turn with complete tool_use blocks + # still executes. Whether final's tool_use ids already sit in the + # history is the ground truth for "already appended". + new_messages = [_dump_message(message) + for message in conversation[len(prepared):]] + final_blocks = [_dump_block(item) for item in final.content] + final_ids = {block["id"] for block in final_blocks + if block.get("type") == "tool_use"} + history_ids = {block.get("id") + for message in new_messages + if (message.get("role") == "assistant" + and isinstance(message.get("content"), list)) + for block in message["content"] + if (isinstance(block, dict) + and block.get("type") == "tool_use")} + if not final_ids or not final_ids <= history_ids: + # Unexecuted tool_use blocks (refusal turns) have no tool_result, + # so they cannot enter an appendable history — strip them, as the + # SDK itself does when it rebuilds params around such a turn. + appendable = [block for block in final_blocks + if block.get("type") != "tool_use"] + if appendable: + new_messages = new_messages + [ + {"role": "assistant", "content": appendable}] + envelope["messages"] = new_messages + return envelope diff --git a/pageindex/mcp_bridge.py b/pageindex/mcp_bridge.py new file mode 100644 index 000000000..7d8d153c2 --- /dev/null +++ b/pageindex/mcp_bridge.py @@ -0,0 +1,210 @@ +"""Minimal MCP client (streamable HTTP) for the PageIndex cloud MCP server. + +Backs the cloud branches of ``client.agent_tools()`` and +``client.agent_instructions()``: ``tools/list`` discovers the live tool set, +``tools/call`` executes a tool, and the ``initialize`` handshake carries the +server's agent instructions. Synchronous, requests-only. +Works against both stateful and stateless servers: a session id returned by +``initialize`` is echoed back, and a session-carrying request rejected with +HTTP 404 (the spec's expired-session status) re-initializes once and +retries; a 400 is an ordinary bad request and is never replayed. +""" +from __future__ import annotations + +import json +import threading +from typing import Any, Optional + +import requests + +from ._version import sdk_version +from .errors import PageIndexAPIError + +_PROTOCOL_VERSION = "2025-06-18" +_TIMEOUT = (10, 240) # tools may wait server-side (wait_for_completion: 3 min) + + +def _parse_sse(text: str) -> list[dict]: + """JSON-RPC messages out of a text/event-stream body.""" + messages = [] + text = text.replace("\r\n", "\n").replace("\r", "\n") + for block in text.split("\n\n"): + data_lines = [line[5:].removeprefix(" ") for line in block.splitlines() + if line.startswith("data:")] + if not data_lines: + continue + try: + messages.append(json.loads("\n".join(data_lines))) + except ValueError: + continue + return messages + + +class McpBridge: + def __init__(self, url: str, headers: dict[str, str]): + self._url = url + self._auth_headers = dict(headers) + self._session_id: Optional[str] = None + self._protocol_version: Optional[str] = None + self._instructions: Optional[str] = None + self._initialized = False + self._lock = threading.RLock() + self._next_id = 0 + + # ── JSON-RPC over streamable HTTP ── + + def _post(self, payload: dict, session_id: Optional[str] = None, + protocol_version: Optional[str] = None) -> requests.Response: + headers = { + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + **self._auth_headers, + } + if session_id: + headers["Mcp-Session-Id"] = session_id + if protocol_version: + headers["MCP-Protocol-Version"] = protocol_version + try: + return requests.post(self._url, json=payload, headers=headers, + timeout=_TIMEOUT) + except requests.RequestException as exc: + raise PageIndexAPIError( + f"Could not reach the PageIndex MCP server: {exc}" + ) from exc + + def _extract_result(self, response: requests.Response, request_id: int) -> Any: + content_type = response.headers.get("Content-Type", "") + if "text/event-stream" in content_type: + # SSE is UTF-8 by spec; requests guesses latin-1 for charset-less + # text/* and would mojibake every non-ASCII character. + messages = _parse_sse(response.content.decode("utf-8", + errors="replace")) + else: + try: + messages = [response.json()] + except ValueError as exc: + raise PageIndexAPIError( + f"MCP server returned a non-JSON response " + f"(HTTP {response.status_code})." + ) from exc + # Strict id correlation only — accepting any result-bearing message + # would return a stale or mis-correlated reply as this call's. + reply = next((m for m in messages if m.get("id") == request_id), None) + if reply is None: + raise PageIndexAPIError( + "MCP server response contained no reply matching the request." + ) + if "error" in reply: + error = reply["error"] or {} + raise PageIndexAPIError( + f"MCP error {error.get('code')}: {error.get('message')}" + ) + return reply.get("result") + + def _request(self, method: str, params: Optional[dict] = None, + _retry: bool = True) -> Any: + self._ensure_initialized() + with self._lock: + self._next_id += 1 + request_id = self._next_id + session_id = self._session_id + protocol_version = self._protocol_version + payload: dict[str, Any] = {"jsonrpc": "2.0", "id": request_id, + "method": method} + if params is not None: + payload["params"] = params + response = self._post(payload, session_id, protocol_version) + if response.status_code == 404 and session_id and _retry: + # Session expired (stateful servers; the spec's 404): the server + # refused the request at session validation, so replaying it is + # safe. 400 is an ordinary bad request — replaying one would + # re-run side effects. Reset only if no other thread has already + # re-initialized, then retry once on the fresh session. + with self._lock: + if self._session_id == session_id: + self._initialized = False + self._session_id = None + self._protocol_version = None + return self._request(method, params, _retry=False) + if response.status_code >= 400: + raise PageIndexAPIError( + f"MCP request failed: HTTP {response.status_code} " + f"({response.text[:200]})" + ) + return self._extract_result(response, request_id) + + def _ensure_initialized(self) -> None: + with self._lock: + if self._initialized: + return + self._next_id += 1 + request_id = self._next_id + response = self._post({ + "jsonrpc": "2.0", "id": request_id, "method": "initialize", + "params": { + "protocolVersion": _PROTOCOL_VERSION, + "capabilities": {}, + "clientInfo": {"name": "pageindex-python-sdk", + "version": sdk_version()}, + }, + }) + if response.status_code >= 400: + raise PageIndexAPIError( + f"Could not connect to the PageIndex MCP server: HTTP " + f"{response.status_code} ({response.text[:200]}). Check " + "your API key." + ) + result = self._extract_result(response, request_id) or {} + self._session_id = response.headers.get("Mcp-Session-Id") + self._protocol_version = result.get("protocolVersion", + _PROTOCOL_VERSION) + self._instructions = result.get("instructions") + self._initialized = True + # Sent inside the lock so no concurrent thread can slip a + # request between the handshake and this notification. + try: + self._post({"jsonrpc": "2.0", + "method": "notifications/initialized"}, + self._session_id, self._protocol_version) + except PageIndexAPIError: + pass # advisory; a server that required it fails the next request + + # ── public surface ── + + def instructions(self) -> Optional[str]: + """The server's agent instructions from the initialize handshake.""" + self._ensure_initialized() + return self._instructions + + def list_tools(self) -> list[dict]: + tools: list[dict] = [] + cursor: Optional[str] = None + while True: + params = {"cursor": cursor} if cursor else {} + result = self._request("tools/list", params) or {} + tools.extend(result.get("tools") or []) + cursor = result.get("nextCursor") + if not cursor: + return tools + + def call_tool(self, name: str, arguments: dict[str, Any]) -> "tuple[str, bool]": + """Returns (text, is_error) — is_error is the server's MCP isError + marking, which callers must carry to their framework's own error + channel.""" + result = self._request("tools/call", + {"name": name, "arguments": arguments}) or {} + is_error = bool(result.get("isError")) + texts = [] + for block in result.get("content") or []: + if isinstance(block, dict) and block.get("type") == "text": + texts.append(block.get("text", "")) + elif isinstance(block, dict) and isinstance(block.get("data"), str): + # Base64 payloads (image/audio) become a metadata stub — + # dumped verbatim they hand the model the raw blob. Revisit + # if tool results ever pass through as real multimodal input. + kind = block.get("mimeType") or block.get("type") or "binary" + size_kb = max(1, len(block["data"]) * 3 // 4096) + texts.append(f"[{kind} content omitted: ~{size_kb} KB]") + else: + texts.append(json.dumps(block, ensure_ascii=False)) + return "\n".join(texts), is_error diff --git a/pyproject.toml b/pyproject.toml index deac66be3..c316451b6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "pageindex" -version = "0.2.9" +version = "0.2.10" description = "Python SDK for PageIndex — reasoning-based, vectorless document retrieval, cloud and local" readme = "README.md" license = "MIT" @@ -10,9 +10,6 @@ classifiers = [ "Intended Audience :: Developers", "License :: OSI Approved :: MIT License", "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.7", - "Programming Language :: Python :: 3.8", - "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", @@ -28,7 +25,7 @@ include = [ exclude = ["pageindex/flash/assets"] [tool.poetry.dependencies] -python = ">=3.7" +python = ">=3.10" requests = ">=2.28.0" openai = ">=1.70.0" litellm = ">=1.84.0" @@ -38,6 +35,17 @@ sortedcontainers = ">=2.4.0" regex = ">=2024.0.0" python-dotenv = ">=1.0.0" pyyaml = ">=6.0" +# Older releases break string prompts with SDK MCP servers (#597, #780). +claude-agent-sdk = { version = ">=0.1.53", optional = true } +# Older releases crash on current openai before the request is sent. +openai-agents = { version = ">=0.18.1", optional = true } +# Older releases execute a refusal turn's tool_use blocks. +anthropic = { version = ">=0.108.0", optional = true } + +[tool.poetry.extras] +claude = ["claude-agent-sdk"] +openai = ["openai-agents"] +anthropic = ["anthropic"] [tool.poetry.group.dev.dependencies] pytest = ">=7.0" diff --git a/run_pageindex.py b/run_pageindex.py index 452f08174..80c01f16f 100644 --- a/run_pageindex.py +++ b/run_pageindex.py @@ -10,15 +10,18 @@ parser = argparse.ArgumentParser(description='Process PDF or Markdown document and generate structure') parser.add_argument('--pdf_path', type=str, help='Path to the PDF file') parser.add_argument('--md_path', type=str, help='Path to the Markdown file') - parser.add_argument('--flash', action='store_true', help='Use PageIndex Flash (with --pdf_path)') + parser.add_argument('--mode', choices=['flash', 'standard'], default='flash', + help='Processing mode (default: flash)') + parser.add_argument('--flash', action='store_true', default=False, + help=argparse.SUPPRESS) parser.add_argument('--embedded-toc', action=argparse.BooleanOptionalAction, default=None, - help='Use the PDF\'s embedded bookmarks when trustworthy (default: on with --flash)') + help='Use the PDF\'s embedded bookmarks when trustworthy (default: on in flash mode)') parser.add_argument('--summary', action=argparse.BooleanOptionalAction, default=None, - help='Generate node summaries with an LLM (default: on with --flash)') - parser.add_argument('--optimize', nargs='?', const='full', choices=['full', 'merge'], + help='Generate node summaries with an LLM (default: on in flash mode)') + parser.add_argument('--optimize', nargs='?', const='full', choices=['full', 'merge', 'off'], default=None, - help='Refine the tree for search cost: a deterministic merge, then an ' - 'LLM expansion pass; pass `merge` to run the merge alone (PDF only)') + help='Refine the tree for search cost (default: full in flash mode). ' + '`merge` for deterministic merge only; `off` to disable') parser.add_argument('--model', type=str, default=None, help='Model to use (overrides config.yaml)') parser.add_argument('--summary-model', type=str, default=None, @@ -48,18 +51,32 @@ parser.add_argument('--summary-token-threshold', type=int, default=200, help='Token threshold for generating summaries (markdown only)') args = parser.parse_args() - + if args.flash: + args.mode = 'flash' + # Validate that exactly one file type is specified if not args.pdf_path and not args.md_path: raise ValueError("Either --pdf_path or --md_path must be specified") if args.pdf_path and args.md_path: raise ValueError("Only one of --pdf_path or --md_path can be specified") - if args.optimize and not (args.pdf_path and args.flash): - raise ValueError("--optimize requires --flash with --pdf_path") - if args.embedded_toc is not None and not (args.pdf_path and args.flash): - raise ValueError("--embedded-toc requires --flash with --pdf_path") - if args.summary is not None and not (args.pdf_path and args.flash): - raise ValueError("--summary requires --flash with --pdf_path") + if args.optimize in ('full', 'merge') and not (args.pdf_path and args.mode == 'flash'): + raise ValueError("--optimize requires Flash mode with --pdf_path") + if args.optimize is None: + args.optimize = 'full' if args.mode == 'flash' else 'off' + if args.embedded_toc is not None and not (args.pdf_path and args.mode == 'flash'): + raise ValueError("--embedded-toc requires Flash mode with --pdf_path") + if args.summary is not None and not (args.pdf_path and args.mode == 'flash'): + raise ValueError("--summary requires Flash mode with --pdf_path") + if args.pdf_path and args.mode == 'flash': + for flag, value in (('--toc-check-pages', args.toc_check_pages), + ('--max-pages-per-node', args.max_pages_per_node), + ('--max-tokens-per-node', args.max_tokens_per_node), + ('--if-add-node-id', args.if_add_node_id), + ('--if-add-node-summary', args.if_add_node_summary), + ('--if-add-doc-description', args.if_add_doc_description), + ('--if-add-node-text', args.if_add_node_text)): + if value is not None: + raise ValueError(f"{flag} is not supported in flash mode; use --mode standard") if args.pdf_path: # Validate PDF file @@ -68,22 +85,23 @@ if not os.path.isfile(args.pdf_path): raise ValueError(f"PDF file not found: {args.pdf_path}") - if args.flash: + if args.mode == 'flash': from pageindex.flash import page_index_flash - if args.optimize == 'full': - from pageindex.tree_optimize import default_model - from pageindex.utils import _is_openai_model - expand_model = args.model or default_model() - if _is_openai_model(expand_model) and not os.getenv("OPENAI_API_KEY"): - raise SystemExit(f"OPENAI_API_KEY is not set (expand model: {expand_model}).") + summary_model = args.summary_model or args.model + will_summarize = args.summary if args.summary is not None else True + if summary_model and (will_summarize or args.optimize == 'full'): + import litellm + env = litellm.validate_environment(summary_model) + if not env["keys_in_environment"]: + raise SystemExit( + f"Missing API key for {summary_model}: {', '.join(env['missing_keys'])}") toc_with_page_number = page_index_flash( args.pdf_path, - optimize=args.optimize is not None, - optimize_expand=args.optimize == 'full', - optimize_model=args.model, - summary_model=args.summary_model or args.model, + optimize=args.optimize if args.optimize != 'off' else False, + optimize_model=summary_model, + summary_model=summary_model, use_embedded_toc=args.embedded_toc if args.embedded_toc is not None else True, - summary=args.summary if args.summary is not None else True, + summary=will_summarize, ) if 'optimize' in toc_with_page_number: o = toc_with_page_number['optimize'] @@ -110,7 +128,7 @@ # Save results pdf_name = os.path.splitext(os.path.basename(args.pdf_path))[0] - suffix = '_structure_flash' if args.flash else '_structure' + suffix = '_structure' output_dir = './results' output_file = f'{output_dir}/{pdf_name}{suffix}.json' os.makedirs(output_dir, exist_ok=True) diff --git a/tests/data/cloud_mcp_contract.json b/tests/data/cloud_mcp_contract.json new file mode 100644 index 000000000..25711a6ba --- /dev/null +++ b/tests/data/cloud_mcp_contract.json @@ -0,0 +1,215 @@ +{ + "_provenance": "Frozen copy of the PageIndex cloud MCP server's tool contract (names, input schemas, descriptions, and annotations as served via tools/list). The parity test asserts pageindex.agent_tools.TOOL_CONTRACT matches this file; update both together only when the cloud contract changes.", + "tools": { + "browse_documents": { + "annotations": { + "readOnlyHint": true, + "openWorldHint": false + }, + "description": "Primary document retrieval tool. After orienting with get_folder_structure() (when available), use this for all document-related questions. The bare call returns root-level sub-folders and documents; pass folder_id to drill into a sub-folder level by level. Use sort=\"relevance\" + query for semantic ranking. Do NOT jump to search_documents() first — it is an escalation path, only after browse_documents(sort=\"relevance\") has failed.", + "schema": { + "type": "object", + "properties": { + "folder_id": { + "type": "string", + "default": "root", + "description": "Folder scope (default \"root\"). Pass a specific folder ID to scope into that folder, or \"root\" to reference the library root. The read-only \"shared-with-me\" and \"following\" folders live at the library root — pass one of those ids to browse them. Copy any folder_id verbatim from a browse/tree response, never construct one. Combine with `recursive` to control breadth." + }, + "recursive": { + "type": "boolean", + "default": false, + "description": "Whether to include documents from descendant folders. When false (default), returns the direct contents of folder_id along with its sub-folders — prefer this for level-by-level exploration so you retain folder hierarchy context. When true, flattens all descendant documents into one list and omits sub-folders — use only when a non-recursive browse of the target folder returned no relevant results and you need to widen the scope, or the user explicitly requests a flat listing." + }, + "sort": { + "type": "string", + "enum": [ + "time", + "relevance" + ], + "default": "time", + "description": "Sort order. \"time\" (default) sorts by upload date (newest first); \"relevance\" orders documents by semantic relevance to `query`. Relevance also works inside the read-only shared folders — pass their folder_id — but at the library root it ranks only your own documents." + }, + "query": { + "type": "string", + "description": "Search query for relevance ranking. Required when sort=\"relevance\"; must be omitted when sort=\"time\"." + }, + "offset": { + "type": "integer", + "minimum": 0, + "maximum": 9007199254740991, + "default": 0, + "description": "Zero-based pagination offset. Pass the value of `next_offset` from the previous response to fetch the next page." + }, + "limit": { + "type": "number", + "minimum": 1, + "maximum": 50, + "default": 10, + "description": "Number of documents to return per page (1-50, default 10)" + } + }, + "required": [] + } + }, + "get_document": { + "annotations": { + "readOnlyHint": true, + "openWorldHint": false + }, + "description": "Check a document's processing status and metadata. `status` is one of \"pending\", \"queued\", \"processing\", \"completed\", or \"failed\" — call this before `get_document_structure()` or `get_page_content()` to confirm the document is ready.", + "schema": { + "type": "object", + "properties": { + "doc_name": { + "type": "string", + "minLength": 1, + "description": "Copy the `name` field verbatim from a browse_documents() or search_documents() response (case-sensitive, include extension). Example: \"Q3 Report.pdf\". If the response shows two documents with the same name, pass `folder_id` alongside to disambiguate." + }, + "folder_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Disambiguator for same-name documents. Copy the `folder_id` from the intended browse/search result; use \"root\" for root-level documents, or \"shared-with-me\"/\"following\" for the read-only folders at the library root; omit if `doc_name` is unique. Copy any folder_id verbatim from a browse_documents()/get_folder_structure() response, never construct one." + }, + "wait_for_completion": { + "type": "boolean", + "default": false, + "description": "If true and document is processing, automatically wait up to 3 minutes until completed. Reduces repeated tool calls." + } + }, + "required": [ + "doc_name" + ] + } + }, + "get_document_structure": { + "annotations": { + "readOnlyHint": true, + "openWorldHint": false + }, + "description": "Extract a document's hierarchical outline (headers, sections, page references). REQUIRED for documents over 20 pages — call this first to locate relevant sections, then pass their page numbers to `get_page_content()`. Use the `part` parameter to iterate large outlines until `pagination.has_more` is false.", + "schema": { + "type": "object", + "properties": { + "doc_name": { + "type": "string", + "minLength": 1, + "description": "Copy the `name` field verbatim from a browse_documents() or search_documents() response (case-sensitive, include extension). Example: \"Q3 Report.pdf\". If the response shows two documents with the same name, pass `folder_id` alongside to disambiguate." + }, + "folder_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Disambiguator for same-name documents. Copy the `folder_id` from the intended browse/search result; use \"root\" for root-level documents, or \"shared-with-me\"/\"following\" for the read-only folders at the library root; omit if `doc_name` is unique. Copy any folder_id verbatim from a browse_documents()/get_folder_structure() response, never construct one." + }, + "part": { + "type": "integer", + "minimum": 1, + "maximum": 9007199254740991, + "default": 1, + "description": "Part number for pagination (1-based, default 1). For large outlines, increment until the response's `pagination.has_more` becomes false." + }, + "wait_for_completion": { + "type": "boolean", + "default": false, + "description": "If true and document is processing, automatically wait up to 3 minutes until completed. Reduces repeated tool calls." + } + }, + "required": [ + "doc_name" + ] + } + }, + "get_page_content": { + "annotations": { + "readOnlyHint": true, + "openWorldHint": false + }, + "description": "Extract page content from a processed document. Use tight, targeted page ranges — never the whole document at once. For documents over 20 pages, call `get_document_structure()` first to pick relevant sections. Embedded image paths in the response feed into `get_document_image()`.", + "schema": { + "type": "object", + "properties": { + "doc_name": { + "type": "string", + "minLength": 1, + "description": "Copy the `name` field verbatim from a browse_documents() or search_documents() response (case-sensitive, include extension). Example: \"Q3 Report.pdf\". If the response shows two documents with the same name, pass `folder_id` alongside to disambiguate." + }, + "folder_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Disambiguator for same-name documents. Copy the `folder_id` from the intended browse/search result; use \"root\" for root-level documents, or \"shared-with-me\"/\"following\" for the read-only folders at the library root; omit if `doc_name` is unique. Copy any folder_id verbatim from a browse_documents()/get_folder_structure() response, never construct one." + }, + "pages": { + "type": "string", + "minLength": 1, + "pattern": "^(\\d+(-\\d+)?)(,\\s*\\d+(-\\d+)?)*$", + "description": "Page specification: \"5\", \"3,7,10\", \"5-10\", or \"1-3,7,9-12\"" + }, + "wait_for_completion": { + "type": "boolean", + "default": false, + "description": "If true and document is processing, automatically wait up to 3 minutes until completed. Reduces repeated tool calls." + } + }, + "required": [ + "doc_name", + "pages" + ] + } + }, + "remove_document": { + "annotations": { + "readOnlyHint": false, + "destructiveHint": true, + "idempotentHint": true, + "openWorldHint": false + }, + "description": "Permanently delete documents and all associated data. Only invoke when the user explicitly names the documents AND confirms deletion. Returns `results` — one entry per requested document: `{ doc_name, status: \"deleted\" | \"not_found\" | \"failed\", error? }`. Inspect each entry for per-document failures. This action is irreversible.", + "schema": { + "type": "object", + "properties": { + "doc_names": { + "type": "array", + "items": { + "type": "string", + "minLength": 1 + }, + "minItems": 1, + "maxItems": 10, + "description": "Array of document names to delete. Each name must be copied verbatim from the `name` field of a browse_documents() or search_documents() response (case-sensitive, include extension). Example: [\"Q3 Report.pdf\", \"draft.pdf\"]. Max 10 per call." + }, + "folder_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Disambiguator for same-name documents. Copy the `folder_id` from the intended browse/search result; use \"root\" for root-level documents, or \"shared-with-me\"/\"following\" for the read-only folders at the library root; omit if `doc_name` is unique. Copy any folder_id verbatim from a browse_documents()/get_folder_structure() response, never construct one." + } + }, + "required": [ + "doc_names" + ] + } + } + } +} diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py new file mode 100644 index 000000000..5c7f7d368 --- /dev/null +++ b/tests/test_agent_tools.py @@ -0,0 +1,2212 @@ +"""Agent tools layer: cloud-contract parity and behavior against a seeded +local store (no LLM calls; one live parity test gated on PAGEINDEX_API_KEY).""" +import asyncio +import json +import os +import re +import sys +import time +import types +from pathlib import Path + +import pytest + +import pageindex.agent_tools as agent_tools_module +import pageindex.client as client_module +from pageindex import PageIndexAPIError, PageIndexCloudClient, PageIndexLocalClient +from pageindex.agent_tools import ( + AGENT_INSTRUCTIONS, + TOOL_CONTRACT, + call_tool, + tool_names, +) +from pageindex.local_store import DocStore + +SNAPSHOT_PATH = Path(__file__).parent / "data" / "cloud_mcp_contract.json" + + +def seed_doc(storage_path, doc_id, name, *, created_at="2026-08-01T10:00:00.123000", + description="A test document", metadata=None, tree=None, pages=None, + page_num=None): + pages = pages if pages is not None else [ + {"page_index": 1, "markdown": "Page one text about apples"}, + {"page_index": 2, "markdown": "Page two text about bananas"}, + ] + tree = tree if tree is not None else [{ + "title": "Doc", "node_id": "0000", "start_index": 1, "end_index": 2, + "summary": "root summary", "text": "ROOT TEXT", + "nodes": [ + {"title": "Intro", "node_id": "0001", "start_index": 1, + "end_index": 1, "summary": "intro summary", "text": "INTRO TEXT"}, + {"title": "Body", "node_id": "0002", "start_index": 2, + "end_index": 2, "summary": "body summary", "text": "BODY TEXT"}, + ], + }] + meta = { + "id": doc_id, "name": name, "description": description, + "status": "completed", "createdAt": created_at, + "pageNum": page_num if page_num is not None else len(pages), + "folderId": None, "metadata": metadata, "mode": "standard", + } + DocStore(storage_path).save_document(doc_id, meta, tree, pages) + return doc_id + + +@pytest.fixture +def store_path(tmp_path): + return str(tmp_path / "store") + + +@pytest.fixture +def client(store_path): + return PageIndexLocalClient(storage_path=store_path) + + +def run(client, name, **arguments): + text, is_error = call_tool(client, name, arguments) + return json.loads(text), is_error + + +# ── contract parity ── + +def test_contract_matches_snapshot(): + snapshot = json.loads(SNAPSHOT_PATH.read_text(encoding="utf-8")) + assert snapshot["tools"] == TOOL_CONTRACT + + +def test_tool_surface_and_docstrings(client): + import inspect + from pageindex.agent_tools import _LOCAL_HIDDEN_PARAMS, _local_schema + tools = client.agent_tools() + assert [tool.__name__ for tool in tools] == list(tool_names()) + with_management = client.agent_tools(include_management=True) + assert [tool.__name__ for tool in with_management][-1] == "remove_document" + for tool in with_management: + exposed = list(_local_schema(tool.__name__)["properties"]) + assert list(inspect.signature(tool).parameters) == exposed + for param in exposed: + assert param in tool.__doc__ + # Cloud-only params are hidden, not documented-then-retracted: + # strict-schema frameworks cannot express the dead-end calls at all. + # (The description may still mention them as cloud capabilities.) + args_section = tool.__doc__.split("Args:", 1)[1] + for hidden in _LOCAL_HIDDEN_PARAMS.get(tool.__name__, ()): + assert f"{hidden}:" not in args_section + docs = {tool.__name__: tool.__doc__ for tool in tools} + # Tools whose cloud description has no cloud-only content keep it + # verbatim; browse_documents serves the localized guidance. + assert docs["get_document"].startswith( + TOOL_CONTRACT["get_document"]["description"]) + assert docs["browse_documents"].startswith( + "Primary document retrieval tool") + + +def test_local_schema_structure_matches_contract(): + """The local surface is the contract minus the documented cloud-only + params; the surviving params' names, types, defaults, bounds, and + required stay byte-identical — localization may only touch description + strings.""" + import copy + from pageindex.agent_tools import _LOCAL_HIDDEN_PARAMS, _local_schema + + def stripped(schema, drop=()): + schema = copy.deepcopy(schema) + for param in drop: + schema["properties"].pop(param, None) + for spec in schema["properties"].values(): + spec.pop("description", None) + return schema + + for name, contract in TOOL_CONTRACT.items(): + hidden = _LOCAL_HIDDEN_PARAMS.get(name, ()) + assert not (set(hidden) & set(contract["schema"].get("required", []))), name + assert stripped(_local_schema(name)) == stripped(contract["schema"], + drop=hidden), name + + +def test_local_guidance_references_only_local_tools(client): + """Local descriptions must not send the agent to tools that are not + registered here (the cloud text names search_documents, + get_folder_structure, and get_document_image).""" + registered = set(tool_names(include_management=True)) + for tool in client.agent_tools(include_management=True): + named = set(re.findall(r"\b(\w+)\(", tool.__doc__)) + assert named <= registered, (tool.__name__, named - registered) + + +def test_local_guidance_points_cloud_only_capabilities_at_cloud(client): + tools = client.agent_tools(include_management=True) + browse = tools[0].__doc__ + assert "not supported in local mode yet" in browse + assert "PageIndex cloud" in browse + # Capability-phrase guard, all docstrings: cloud-only language must not + # drift back in via a contract refresh. browse alone keeps exactly one + # sort="relevance" mention — the sanctioned pointer to the cloud. + for tool in tools: + doc = tool.__doc__ + for phrase in ("shared-with-me", "sub-folder", "get_folder_structure", + "search_documents", "get_document_image"): + assert phrase not in doc, (tool.__name__, phrase) + expected = 1 if tool.__name__ == "browse_documents" else 0 + assert doc.count('sort="relevance"') == expected, tool.__name__ + + +# ── browse_documents ── + +def test_browse_documents_shape(client, store_path): + seed_doc(store_path, "pi-a", "older.pdf", created_at="2026-08-01T10:00:00.123000") + seed_doc(store_path, "pi-b", "newer.pdf", created_at="2026-08-02T10:00:00.456000", + metadata={"team": "research", "year": 2026, "nested": {"x": 1}}) + payload, is_error = run(client, "browse_documents") + assert not is_error + assert payload["success"] is True + assert payload["folders"] == [] + assert payload["has_more"] is False + assert payload["next_offset"] is None + names = [doc["name"] for doc in payload["documents"]] + assert names == ["newer.pdf", "older.pdf"] + newer = payload["documents"][0] + assert newer["status"] == "completed" + assert newer["created_at"] == "2026-08-02T10:00:00.456Z" + assert newer["metadata"] == {"team": "research", "year": 2026} + assert "folder_id" not in newer + assert "next_steps" in payload + + flat, _ = run(client, "browse_documents", recursive=True) + assert "folders" not in flat + + +def test_browse_documents_pagination(client, store_path): + for index in range(3): + seed_doc(store_path, f"pi-{index}", f"doc{index}.pdf", + created_at=f"2026-08-0{index + 1}T10:00:00.000000") + first, _ = run(client, "browse_documents", limit=2) + assert [d["name"] for d in first["documents"]] == ["doc2.pdf", "doc1.pdf"] + assert first["has_more"] is True and first["next_offset"] == 2 + assert "page through the rest" in json.dumps(first["next_steps"]) + second, _ = run(client, "browse_documents", limit=2, offset=2) + assert [d["name"] for d in second["documents"]] == ["doc0.pdf"] + assert second["has_more"] is False + # No paging advice when there is nothing left to page through. + assert "page through the rest" not in json.dumps(second["next_steps"]) + + +def test_browse_documents_relevance_unsupported(client, store_path): + """Semantic ranking is cloud-side; like folders, local answers with an + honest error instead of a keyword imitation.""" + seed_doc(store_path, "pi-a", "attention.pdf", + description="Transformers and attention mechanisms") + payload, is_error = run(client, "browse_documents", sort="relevance", + query="attention transformers") + assert is_error and payload["errorCode"] == "INVALID_INPUT" + assert "not supported in local mode" in payload["error"] + + stray_query, is_error = run(client, "browse_documents", query="x") + assert is_error and "not supported in local mode" in stray_query["error"] + bad_sort, is_error = run(client, "browse_documents", sort="banana") + assert is_error and bad_sort["errorCode"] == "INVALID_INPUT" + # The invalid-sort guidance must not prescribe the cloud-only value. + assert 'Use sort="relevance"' not in json.dumps(bad_sort) + assert "local mode" in bad_sort["error"] + + +def test_browse_documents_empty_and_folder_error(client): + payload, is_error = run(client, "browse_documents") + assert not is_error + assert payload["documents"] == [] + assert "submit_document" in json.dumps(payload) + + folder, is_error = run(client, "browse_documents", folder_id="folder-123") + assert is_error and folder["errorCode"] == "INVALID_INPUT" + + +# ── get_document ── + +def test_get_document(client, store_path): + seed_doc(store_path, "pi-a", "report.pdf", metadata={"team": "research"}) + payload, is_error = run(client, "get_document", doc_name="report.pdf") + assert not is_error + assert payload["name"] == "report.pdf" + assert payload["status"] == "completed" + assert payload["page_count"] == 2 + assert payload["folder_id"] is None + assert payload["created_at"].endswith("Z") + assert payload["metadata"] == {"team": "research"} + assert any("short document" in option + for option in payload["next_steps"]["options"]) + + +def test_get_document_not_found_suggests_similar(client, store_path): + seed_doc(store_path, "pi-a", "annual-report.pdf") + payload, is_error = run(client, "get_document", doc_name="anual-report.pdf") + assert is_error + assert payload["errorCode"] == "NOT_FOUND" + assert "annual-report.pdf" in payload["similar_files"] + assert "Did you mean" in payload["error"] + + +def test_get_document_duplicate_names_resolve_newest(client, store_path): + seed_doc(store_path, "pi-old", "same.pdf", description="old copy", + created_at="2026-08-01T10:00:00.000000") + seed_doc(store_path, "pi-new", "same.pdf", description="new copy", + created_at="2026-08-02T10:00:00.000000") + payload, _ = run(client, "get_document", doc_name="same.pdf") + assert payload["description"] == "new copy" + + +# ── get_document_structure ── + +def test_structure_strips_text_and_orders_keys(client, store_path): + seed_doc(store_path, "pi-a", "report.pdf") + payload, is_error = run(client, "get_document_structure", doc_name="report.pdf") + assert not is_error + assert payload["doc_name"] == "report.pdf" + assert "pagination" not in payload and "total_parts" not in payload + serialized = json.dumps(payload["structure"]) + assert "ROOT TEXT" not in serialized and "INTRO TEXT" not in serialized + # Cloud structure node shape: start_index/end_index/summary (live-verified). + root = payload["structure"][0] + assert list(root)[:4] == ["title", "node_id", "start_index", "end_index"] + assert root["summary"] == "root summary" + assert (root["start_index"], root["end_index"]) == (1, 2) + assert root["nodes"][0]["summary"] == "intro summary" + assert root["nodes"][0]["end_index"] == 1 + + +def test_structure_multipart_pagination(client, store_path): + big_tree = [{ + "title": f"Chapter {index}", "node_id": f"{index:04d}", + "start_index": index + 1, "end_index": index + 1, + "summary": "s" * 4000, "text": "T", + } for index in range(60)] + seed_doc(store_path, "pi-big", "big.pdf", tree=big_tree, + pages=[{"page_index": 1, "markdown": "x"}]) + first, _ = run(client, "get_document_structure", doc_name="big.pdf") + assert first["total_parts"] > 1 + assert first["pagination"] == { + "part": 1, "total_parts": first["total_parts"], "has_more": True, + } + titles = [] + for part in range(1, first["total_parts"] + 1): + payload, _ = run(client, "get_document_structure", doc_name="big.pdf", + part=part) + # Every part of one paginated response is a list — a consumer that + # iterates part 1 must not silently iterate dict keys on part 2. + assert isinstance(payload["structure"], list) + titles.extend(node["title"] for node in payload["structure"]) + assert payload["pagination"]["has_more"] == (part < first["total_parts"]) + assert titles == [f"Chapter {index}" for index in range(60)] + + clamped, _ = run(client, "get_document_structure", doc_name="big.pdf", + part=999) + assert clamped["pagination"]["part"] == first["total_parts"] + + +def test_split_structure_chunks_never_change_type(): + """A single-node group used to come out as a bare dict while its + sibling parts were lists — same response sequence, flipping JSON type.""" + from pageindex.agent_tools import _split_structure + small = {"title": "s", "node_id": "0001"} + big = {"title": "b", "node_id": "0002", + "nodes": [{"title": f"c{index}", "summary": "x" * 40} + for index in range(10)]} + chunks = _split_structure([small, small, big], 200) + assert len(chunks) > 1 + assert all(isinstance(chunk, list) for chunk in chunks) + # Unsplit structures keep their natural shape (cloud fallback parity). + assert _split_structure(small, 10_000) == [small] + assert _split_structure([small], 10_000) == [[small]] + + +# ── get_page_content ── + +def test_page_content(client, store_path): + seed_doc(store_path, "pi-a", "report.pdf") + payload, is_error = run(client, "get_page_content", doc_name="report.pdf", + pages="1-2") + assert not is_error + assert payload["total_pages"] == 2 + assert payload["requested_pages"] == "1-2" + assert payload["returned_pages"] == "1-2" + assert payload["content"] == [ + {"page": 1, "text": "Page one text about apples"}, + {"page": 2, "text": "Page two text about bananas"}, + ] + + +def test_page_content_out_of_range(client, store_path): + seed_doc(store_path, "pi-a", "report.pdf") + mixed, is_error = run(client, "get_page_content", doc_name="report.pdf", + pages="1,99") + assert not is_error + assert mixed["returned_pages"] == "1" + assert "out of range" in mixed["next_steps"]["summary"] + + all_out, is_error = run(client, "get_page_content", doc_name="report.pdf", + pages="99") + assert is_error and all_out["errorCode"] == "INVALID_INPUT" + assert all_out["max_pages"] == 2 + + +def test_out_of_range_pages_reported_as_ranges(client, store_path): + """Spans compress — enumerating them one by one buries the response.""" + seed_doc(store_path, "pi-a", "report.pdf") + partial, is_error = run(client, "get_page_content", doc_name="report.pdf", + pages="1,5-9") + assert not is_error + assert "Pages 5-9 were out of range" in partial["next_steps"]["summary"] + + spread, is_error = run(client, "get_page_content", doc_name="report.pdf", + pages="1,5,9") + assert not is_error + assert "Pages 5,9 were out of range" in spread["next_steps"]["summary"] + + all_out, is_error = run(client, "get_page_content", doc_name="report.pdf", + pages="5-9") + assert is_error + assert all_out["error"].endswith("you requested pages: 5-9") + assert all_out["requested_pages"] == "5-9" + + +@pytest.mark.parametrize("bad_spec", ["abc", "5-3", "1,,2", "-3", ""]) +def test_page_content_invalid_spec(client, store_path, bad_spec): + seed_doc(store_path, "pi-a", "report.pdf") + payload, is_error = run(client, "get_page_content", doc_name="report.pdf", + pages=bad_spec) + assert is_error and payload["errorCode"] == "INVALID_INPUT" + + +def test_page_content_zero_page_rejected(client, store_path): + seed_doc(store_path, "pi-a", "report.pdf") + payload, is_error = run(client, "get_page_content", doc_name="report.pdf", + pages="0") + assert is_error + assert "positive integers" in payload["error"] + + +def test_page_content_preserves_blank_pages(client, store_path): + pages = [ + {"page_index": 1, "markdown": ""}, + {"page_index": 2, "markdown": "content"}, + ] + seed_doc(store_path, "pi-a", "blanks.pdf", pages=pages) + payload, is_error = run(client, "get_page_content", doc_name="blanks.pdf", + pages="1-2") + assert not is_error + assert payload["content"][0] == {"page": 1, "text": ""} + assert payload["content"][1] == {"page": 2, "text": "content"} + + +def test_created_at_accepts_z_suffixed_input(client, store_path): + seed_doc(store_path, "pi-a", "cloudlike.pdf", + created_at="2026-08-01T10:00:00.123Z") + payload, _ = run(client, "browse_documents") + assert payload["documents"][0]["created_at"] == "2026-08-01T10:00:00.123Z" + + +def test_page_content_char_budget(client, store_path): + pages = [ + {"page_index": 1, "markdown": "x" * 96_000}, + {"page_index": 2, "markdown": "short"}, + ] + seed_doc(store_path, "pi-a", "huge.pdf", pages=pages) + payload, is_error = run(client, "get_page_content", doc_name="huge.pdf", + pages="1-2") + assert not is_error + assert payload["returned_pages"] == "1" + assert "size limits" in payload["next_steps"]["summary"] + assert any("For remaining pages, request: 2" in option + for option in payload["next_steps"]["options"]) + + +def test_page_content_reports_truncation_and_out_of_range_together( + client, store_path): + """Size truncation must not hide behind the out-of-range report (or + vice versa) — the agent otherwise believes it holds every in-range + page.""" + pages = [ + {"page_index": 1, "markdown": "x" * 96_000}, + {"page_index": 2, "markdown": "short"}, + ] + seed_doc(store_path, "pi-a", "huge.pdf", pages=pages) + payload, is_error = run(client, "get_page_content", doc_name="huge.pdf", + pages="1-2,99") + assert not is_error + assert payload["returned_pages"] == "1" + summary = payload["next_steps"]["summary"] + assert "size limits" in summary and "out of range" in summary + + +# ── remove_document (management-gated) ── + +def test_remove_document(client, store_path): + seed_doc(store_path, "pi-a", "report.pdf") + payload, is_error = run(client, "remove_document", + doc_names=["report.pdf", "ghost.pdf"]) + assert not is_error + assert payload["results"] == [ + {"doc_name": "report.pdf", "status": "deleted"}, + {"doc_name": "ghost.pdf", "status": "not_found"}, + ] + assert client.list_documents()["total"] == 0 + + +def test_remove_document_rejects_non_string_names_before_deleting(client, + store_path): + """A rejection envelope must mean nothing was destroyed — the bad + element is caught before the delete loop starts.""" + seed_doc(store_path, "pi-a", "report.pdf") + payload, is_error = run(client, "remove_document", + doc_names=["report.pdf", 123]) + assert is_error and payload["errorCode"] == "INVALID_INPUT" + assert client.list_documents()["total"] == 1 + + +def test_remove_document_partial_failure_keeps_results(client, store_path, + monkeypatch): + """A non-API error mid-batch must not discard the entries for documents + already irreversibly deleted — a generic INTERNAL_ERROR envelope would + tell the agent nothing was removed and to retry.""" + seed_doc(store_path, "pi-a", "a.pdf") + seed_doc(store_path, "pi-b", "b.pdf") + real = client.delete_document + + def flaky(doc_id): + if doc_id == "pi-b": + raise OSError(13, "Permission denied") + return real(doc_id) + + monkeypatch.setattr(client, "delete_document", flaky) + payload, is_error = run(client, "remove_document", + doc_names=["a.pdf", "b.pdf"]) + assert not is_error + assert payload["results"] == [ + {"doc_name": "a.pdf", "status": "deleted"}, + {"doc_name": "b.pdf", "status": "failed", + "error": "[Errno 13] Permission denied"}, + ] + + +def test_management_tools_hidden_by_default(client): + assert "remove_document" not in [t.__name__ for t in client.agent_tools()] + + +# ── doc_id scope (the local chat surfaces' allowlist) ── + +def test_call_tool_doc_scope_limits_every_lookup(client, store_path): + seed_doc(store_path, "pi-a", "report.pdf") + seed_doc(store_path, "pi-b", "payroll.pdf", + created_at="2026-08-02T10:00:00.123000") + + text, is_error = call_tool(client, "browse_documents", {}, + doc_ids=["pi-a"]) + browse = json.loads(text) + assert not is_error + assert [doc["name"] for doc in browse["documents"]] == ["report.pdf"] + assert browse["has_more"] is False + + text, is_error = call_tool(client, "get_page_content", + {"doc_name": "payroll.pdf", "pages": "1"}, + doc_ids="pi-a") + assert is_error and json.loads(text)["errorCode"] == "NOT_FOUND" + + text, is_error = call_tool(client, "get_document", + {"doc_name": "report.pdf"}, doc_ids="pi-a") + assert not is_error + + # An empty allowlist scopes to nothing — it must not read as "unscoped". + text, is_error = call_tool(client, "browse_documents", {}, doc_ids=[]) + assert not is_error and json.loads(text)["documents"] == [] + + +def test_call_tool_scope_channel_not_injectable(client, store_path): + """Model arguments cannot smuggle an allowlist: underscore keys are + stripped before binding.""" + seed_doc(store_path, "pi-a", "report.pdf") + text, is_error = call_tool(client, "browse_documents", + {"_allowed_ids": ["pi-none"]}) + assert not is_error + assert json.loads(text)["documents"] + + +# ── error containment ── + +def test_tools_never_raise(client, store_path, monkeypatch): + seed_doc(store_path, "pi-a", "report.pdf") + monkeypatch.setattr(client._api._store, "get_tree", + lambda *a, **k: (_ for _ in ()).throw(RuntimeError("boom"))) + payload, is_error = run(client, "get_document_structure", + doc_name="report.pdf") + assert is_error + assert "boom" in payload["error"] + + +def test_unknown_argument_becomes_error_envelope(client, store_path): + seed_doc(store_path, "pi-a", "report.pdf") + payload, is_error = run(client, "get_document", doc_name="report.pdf", + bogus=True) + assert is_error and payload["errorCode"] == "INVALID_INPUT" + + +def test_execution_type_error_is_internal_not_invalid_input(client, store_path, + monkeypatch): + """Only bind-time TypeErrors are argument errors; a TypeError raised + mid-execution must not masquerade as an input rejection.""" + seed_doc(store_path, "pi-a", "report.pdf") + monkeypatch.setattr(client._api._store, "get_tree", + lambda *a, **k: (_ for _ in ()).throw( + TypeError("wrong shape"))) + payload, is_error = run(client, "get_document_structure", + doc_name="report.pdf") + assert is_error and payload["errorCode"] == "INTERNAL_ERROR" + assert "wrong shape" in payload["error"] + + +def test_unknown_tool_envelope_uses_standard_formatting(client): + text, is_error = call_tool(client, "nope", {}) + assert is_error + assert text == json.dumps(json.loads(text), ensure_ascii=False) + + +# ── framework adapters ── + +def test_as_openai_tools_missing_dependency(client, monkeypatch): + monkeypatch.setitem(sys.modules, "agents", None) + with pytest.raises(PageIndexAPIError, match="openai-agents"): + client.as_openai_tools() + + +def test_as_openai_tools_local_in_process(client): + pytest.importorskip("agents") + tools = client.as_openai_tools() + assert [tool.name for tool in tools] == list(tool_names()) + + +def test_as_openai_tools_cloud_default_uses_bridge(monkeypatch): + pytest.importorskip("agents") + from agents import FunctionTool + import pageindex.mcp_bridge as mcp_bridge + monkeypatch.setattr(mcp_bridge, "McpBridge", _FakeBridge) + cloud = PageIndexCloudClient(api_key="pi-test-key") + tools = cloud.as_openai_tools() + assert all(isinstance(tool, FunctionTool) for tool in tools) + assert [tool.name for tool in tools] == ["search_documents", "get_document"] + + +def test_as_openai_tools_cloud_hosted_opt_in(): + pytest.importorskip("agents") + from agents import HostedMCPTool + cloud = PageIndexCloudClient(api_key="pi-test-key") + tools = cloud.as_openai_tools(hosted=True) + assert len(tools) == 1 + assert isinstance(tools[0], HostedMCPTool) + config = tools[0].tool_config + assert config["server_url"] == "https://api.pageindex.ai/mcp?tools=read" + assert config["headers"] == {"Authorization": "Bearer pi-test-key"} + assert config["server_label"] == "pageindex" + + +def test_as_openai_tools_local_ignores_hosted(client): + pytest.importorskip("agents") + assert ([tool.name for tool in client.as_openai_tools(hosted=True)] + == [tool.name for tool in client.as_openai_tools()] + == list(tool_names())) + + +def test_as_openai_tools_schemas_pass_through_verbatim(client): + """The contract schema goes to the model as-is — regenerating it from a + Python signature dropped items/enum/pattern/bounds.""" + pytest.importorskip("agents") + from pageindex.agent_tools import _local_schema + tools = {tool.name: tool + for tool in client.as_openai_tools(include_management=True)} + assert (tools["remove_document"].params_json_schema + == _local_schema("remove_document")) + pages = tools["get_page_content"].params_json_schema["properties"]["pages"] + assert pages["pattern"] and pages["minLength"] == 1 + assert all(tool.strict_json_schema is False for tool in tools.values()) + + +def test_as_openai_tools_invocation_runs_call_tool(client, store_path): + pytest.importorskip("agents") + seed_doc(store_path, "pi-a", "report.pdf") + tool = {t.name: t for t in client.as_openai_tools()}["get_document"] + out = asyncio.run(tool.on_invoke_tool( + None, json.dumps({"doc_name": "report.pdf", "folder_id": None}))) + payload = json.loads(out) + assert payload["success"] is True and payload["name"] == "report.pdf" + + +def test_as_openai_tools_malformed_args_answer_the_model(client, store_path): + """strict_json_schema is off, so a truncated or non-object argument + string is reachable; raising here aborted the caller's whole run — + the model must get the guided envelope back and retry instead.""" + pytest.importorskip("agents") + seed_doc(store_path, "pi-a", "report.pdf") + tool = {t.name: t for t in client.as_openai_tools()}["get_document"] + for bad in ('{not json', '[1, 2]', '"x"', 'null'): + out = asyncio.run(tool.on_invoke_tool(None, bad)) + payload = json.loads(out) + assert payload["errorCode"] == "INVALID_INPUT" + assert "JSON object" in payload["error"] + + +def test_as_openai_tools_cloud_object_params_survive(monkeypatch): + """An object-typed server parameter used to abort the whole build with + agents.exceptions.UserError; array items used to degrade to {}.""" + pytest.importorskip("agents") + import pageindex.mcp_bridge as mcp_bridge + + schema = { + "type": "object", + "properties": { + "filters": {"type": "object", "additionalProperties": False}, + "paths": {"type": "array", + "items": {"type": "string", "minLength": 1}}, + }, + "required": ["paths"], + } + + class _ObjBridge: + def __init__(self, url, headers): + pass + + def list_tools(self): + return [{"name": "get_document_image", + "description": "d", + "annotations": {"readOnlyHint": True}, + "inputSchema": schema}] + + def call_tool(self, name, arguments): + return json.dumps({"success": True}), False + + monkeypatch.setattr(mcp_bridge, "McpBridge", _ObjBridge) + cloud = PageIndexCloudClient(api_key="pi-test-key") + tools = cloud.as_openai_tools() + assert len(tools) == 1 + assert tools[0].params_json_schema == schema + assert tools[0].params_json_schema is not schema # copied, not aliased + + +def test_as_claude_mcp_cloud_needs_no_framework(monkeypatch): + monkeypatch.setitem(sys.modules, "claude_agent_sdk", None) + cloud = PageIndexCloudClient(api_key="pi-test-key") + # The URL is the gate: default → read-only endpoint, management opt-in + # → the full tool set. + assert cloud.as_claude_mcp() == { + "type": "http", + "url": "https://api.pageindex.ai/mcp?tools=read", + "headers": {"Authorization": "Bearer pi-test-key"}, + } + assert (cloud.as_claude_mcp(include_management=True)["url"] + == "https://api.pageindex.ai/mcp") + + +def test_as_claude_mcp_local_missing_dependency(client, monkeypatch): + monkeypatch.setitem(sys.modules, "claude_agent_sdk", None) + with pytest.raises(PageIndexAPIError, match="claude-agent-sdk"): + client.as_claude_mcp() + + +def test_as_claude_mcp_local_when_installed(client): + pytest.importorskip("claude_agent_sdk") + server = client.as_claude_mcp() + assert server is not None + if isinstance(server, dict): + assert server.get("type") != "http" + + +def test_claude_agent_config_is_sugar_over_the_explicit_form( + cloud_with_fake_bridge): + cloud, _ = cloud_with_fake_bridge + config = cloud.claude_agent_config() + assert config["system_prompt"] == "SERVER GUIDANCE" + server = config["mcp_servers"]["pageindex"] + assert server["type"] == "http" + assert server["url"] == "https://api.pageindex.ai/mcp?tools=read" + # Pre-approval only: the URL is the gate. + assert config["allowed_tools"] == ["mcp__pageindex"] + renamed = cloud.claude_agent_config(server_name="docs", + include_management=True) + assert set(renamed["mcp_servers"]) == {"docs"} + assert renamed["mcp_servers"]["docs"]["url"] == "https://api.pageindex.ai/mcp" + assert renamed["allowed_tools"] == ["mcp__docs"] + + +def test_claude_agent_config_local(client, store_path): + pytest.importorskip("claude_agent_sdk") + seed_doc(store_path, "pi-a", "report.pdf") + config = client.claude_agent_config(doc_id="pi-a") + assert "report.pdf" in config["system_prompt"] + assert config["allowed_tools"] == ["mcp__pageindex"] + + +def test_openai_agent_config_local(client, store_path): + pytest.importorskip("agents") + from agents import Agent + seed_doc(store_path, "pi-a", "report.pdf") + config = client.openai_agent_config(doc_id="pi-a") + assert config["name"] == "PageIndex" + assert "report.pdf" in config["instructions"] + assert [tool.name for tool in config["tools"]] == list(tool_names()) + assert config["model"] == client.retrieve_model + assert client.openai_agent_config(model="gpt-x")["model"] == "gpt-x" + assert Agent(**client.openai_agent_config()).name == "PageIndex" + + +def test_openai_agent_config_cloud_omits_model(cloud_with_fake_bridge): + pytest.importorskip("agents") + cloud, _ = cloud_with_fake_bridge + config = cloud.openai_agent_config() + assert "model" not in config + assert config["instructions"] == "SERVER GUIDANCE" + assert [tool.name for tool in config["tools"]] == ["search_documents", + "get_document"] + + +def test_anthropic_runner_config_shapes(client, store_path): + pytest.importorskip("anthropic") + import anthropic + from anthropic.lib.tools import BetaAsyncFunctionTool + seed_doc(store_path, "pi-a", "report.pdf") + config = client.anthropic_runner_config(model="claude-3-opus-20240229", + doc_id="pi-a") + assert config["max_tokens"] == 4096 + assert config["max_iterations"] == 10 + assert "report.pdf" in config["system"] + assert [tool.name for tool in config["tools"]] == list(tool_names()) + assert (client.anthropic_runner_config(model="claude-sonnet-4-5") + ["max_tokens"] == 8192) + override = client.anthropic_runner_config(model="claude-sonnet-4-5", + max_tokens=99, max_turns=3) + assert override["max_tokens"] == 99 and override["max_iterations"] == 3 + async_tools = client.anthropic_runner_config( + model="claude-sonnet-4-5", asynchronous=True)["tools"] + assert all(isinstance(tool, BetaAsyncFunctionTool) + for tool in async_tools) + # The kwargs must construct a real runner (construction is offline — + # requests start on iteration), pinning tool_runner's parameter names. + runner = anthropic.Anthropic(api_key="test").beta.messages.tool_runner( + **client.anthropic_runner_config(model="claude-sonnet-4-5"), + messages=[{"role": "user", "content": "q"}]) + assert runner is not None + + +def test_anthropic_runner_config_cloud(cloud_with_fake_bridge): + pytest.importorskip("anthropic") + cloud, _ = cloud_with_fake_bridge + config = cloud.anthropic_runner_config(model="claude-sonnet-4-5") + assert config["system"] == "SERVER GUIDANCE" + assert [tool.name for tool in config["tools"]] == ["search_documents", + "get_document"] + + +# ── config helpers: doc_id is structural in the tools, not just prompted ── + +def test_openai_agent_config_doc_scope_enforced_in_tools(client, store_path): + pytest.importorskip("agents") + seed_doc(store_path, "pi-a", "report.pdf") + seed_doc(store_path, "pi-b", "payroll.pdf", + created_at="2026-08-02T10:00:00.123000") + tools = {tool.name: tool + for tool in client.openai_agent_config(doc_id="pi-a")["tools"]} + out = asyncio.run(tools["get_page_content"].on_invoke_tool( + None, json.dumps({"doc_name": "payroll.pdf", "pages": "1"}))) + assert json.loads(out)["errorCode"] == "NOT_FOUND" + out = asyncio.run(tools["browse_documents"].on_invoke_tool(None, "{}")) + assert [doc["name"] + for doc in json.loads(out)["documents"]] == ["report.pdf"] + + +def test_anthropic_runner_config_doc_scope_enforced_in_tools(client, + store_path): + pytest.importorskip("anthropic") + from anthropic.lib.tools import ToolError + seed_doc(store_path, "pi-a", "report.pdf") + seed_doc(store_path, "pi-b", "payroll.pdf", + created_at="2026-08-02T10:00:00.123000") + config = client.anthropic_runner_config(model="claude-sonnet-4-5", + doc_id="pi-a") + tools = {tool.name: tool for tool in config["tools"]} + with pytest.raises(ToolError, match="NOT_FOUND"): + tools["get_page_content"].call({"doc_name": "payroll.pdf", + "pages": "1"}) + browse = json.loads(tools["browse_documents"].call({})) + assert [doc["name"] for doc in browse["documents"]] == ["report.pdf"] + + +def test_claude_agent_config_doc_scope_enforced_in_tools(client, store_path): + pytest.importorskip("claude_agent_sdk") + from mcp.types import CallToolRequest, CallToolRequestParams + seed_doc(store_path, "pi-a", "report.pdf") + seed_doc(store_path, "pi-b", "payroll.pdf", + created_at="2026-08-02T10:00:00.123000") + config = client.claude_agent_config(doc_id="pi-a") + server = config["mcp_servers"]["pageindex"] + handler = server["instance"].request_handlers[CallToolRequest] + result = asyncio.run(handler(CallToolRequest( + method="tools/call", + params=CallToolRequestParams( + name="get_page_content", + arguments={"doc_name": "payroll.pdf", "pages": "1"})))) + payload = json.loads(result.root.content[0].text) + assert payload["errorCode"] == "NOT_FOUND" + + +def test_openai_agent_config_scoped_shadow_check(client, store_path): + """The bundles' tools resolve names inside the allowlist, so a same-name + document outside the target set must not block — only an in-set + duplicate shadows.""" + pytest.importorskip("agents") + seed_doc(store_path, "pi-old", "report.pdf") + seed_doc(store_path, "pi-new", "report.pdf", + created_at="2026-08-02T10:00:00.123000") + config = client.openai_agent_config(doc_id="pi-old") + assert "report.pdf" in config["instructions"] + with pytest.raises(PageIndexAPIError, match="shadowed"): + client.openai_agent_config(doc_id=["pi-old", "pi-new"]) + + +def test_anthropic_runner_config_scoped_shadow_check(client, store_path): + pytest.importorskip("anthropic") + seed_doc(store_path, "pi-old", "report.pdf") + seed_doc(store_path, "pi-new", "report.pdf", + created_at="2026-08-02T10:00:00.123000") + config = client.anthropic_runner_config(model="claude-sonnet-4-5", + doc_id="pi-old") + assert "report.pdf" in config["system"] + + +def test_claude_agent_config_scoped_shadow_check(client, store_path): + pytest.importorskip("claude_agent_sdk") + seed_doc(store_path, "pi-old", "report.pdf") + seed_doc(store_path, "pi-new", "report.pdf", + created_at="2026-08-02T10:00:00.123000") + config = client.claude_agent_config(doc_id="pi-old") + assert "report.pdf" in config["system_prompt"] + + +def test_doc_scope_rejected_on_cloud_openai(): + pytest.importorskip("agents") + cloud = PageIndexCloudClient(api_key="pi-test-key") + with pytest.raises(PageIndexAPIError, match="server-side"): + cloud.as_openai_tools(doc_id="pi-a") + # The hosted branch returns before _tool_specs — it must reject too, + # not silently drop the allowlist. + with pytest.raises(PageIndexAPIError, match="server-side"): + cloud.as_openai_tools(hosted=True, doc_id="pi-a") + + +def test_doc_scope_rejected_on_cloud_anthropic(): + pytest.importorskip("anthropic") + cloud = PageIndexCloudClient(api_key="pi-test-key") + with pytest.raises(PageIndexAPIError, match="server-side"): + cloud.as_anthropic_tools(doc_id="pi-a") + + +def test_doc_scope_rejected_on_cloud_claude(): + cloud = PageIndexCloudClient(api_key="pi-test-key") + with pytest.raises(PageIndexAPIError, match="server-side"): + cloud.as_claude_mcp(doc_id="pi-a") + + +def test_as_anthropic_tools_missing_dependency(client, monkeypatch): + monkeypatch.setitem(sys.modules, "anthropic", None) + with pytest.raises(PageIndexAPIError, match="anthropic"): + client.as_anthropic_tools() + + +def test_as_anthropic_tools_local_in_process(client, store_path): + pytest.importorskip("anthropic") + from anthropic.lib.tools import BetaFunctionTool + from pageindex.agent_tools import _local_description, _local_schema + tools = client.as_anthropic_tools() + # The sync flavor is load-bearing: the sync runner (and messages()) + # rejects async tools and vice versa. + assert all(isinstance(tool, BetaFunctionTool) for tool in tools) + assert [tool.name for tool in tools] == list(tool_names()) + browse = {tool.name: tool for tool in tools}["browse_documents"] + assert browse.input_schema == _local_schema("browse_documents") + assert browse.description == _local_description("browse_documents") + seed_doc(store_path, "pi-a", "report.pdf") + assert "report.pdf" in browse.call({}) + + +def test_as_anthropic_tools_async_flavor(client, store_path): + pytest.importorskip("anthropic") + from anthropic.lib.tools import BetaAsyncFunctionTool + tools = client.as_anthropic_tools(asynchronous=True) + assert all(isinstance(tool, BetaAsyncFunctionTool) for tool in tools) + assert [tool.name for tool in tools] == list(tool_names()) + seed_doc(store_path, "pi-a", "report.pdf") + browse = {tool.name: tool for tool in tools}["browse_documents"] + assert "report.pdf" in asyncio.run(browse.call({})) + + +def test_as_anthropic_tools_local_management_opt_in(client): + pytest.importorskip("anthropic") + names = [tool.name + for tool in client.as_anthropic_tools(include_management=True)] + assert names == list(tool_names(include_management=True)) + assert "remove_document" in names + + +def test_as_anthropic_tools_local_failures_raise_toolerror(client, store_path): + """Error envelopes surface as ToolError so the runner marks the + tool_result is_error: true — a bare return would read as success.""" + pytest.importorskip("anthropic") + from anthropic.lib.tools import ToolError + seed_doc(store_path, "pi-a", "report.pdf") + tools = {tool.name: tool for tool in client.as_anthropic_tools()} + with pytest.raises(ToolError) as excinfo: + tools["get_document"].call({"doc_name": "ghost.pdf"}) + assert json.loads(excinfo.value.content)["errorCode"] == "NOT_FOUND" + assert "report.pdf" in tools["browse_documents"].call({}) + + +def test_as_anthropic_tools_cloud_iserror_raises_toolerror( + cloud_with_fake_bridge): + """The server's MCP isError marking must reach the runner's error + channel, not arrive as a successful tool_result.""" + pytest.importorskip("anthropic") + from anthropic.lib.tools import ToolError + cloud, created = cloud_with_fake_bridge + tools = cloud.as_anthropic_tools() + created["bridge"].call_tool = lambda name, arguments: ( + '{"error": "denied"}', True) + with pytest.raises(ToolError) as excinfo: + tools[0].call({"query": "q"}) + assert json.loads(excinfo.value.content)["error"] == "denied" + + +def test_as_anthropic_tools_cloud_schemas_pass_through(cloud_with_fake_bridge): + pytest.importorskip("anthropic") + cloud, created = cloud_with_fake_bridge + tools = cloud.as_anthropic_tools() + assert [tool.name for tool in tools] == ["search_documents", "get_document"] + bridge = created["bridge"] + assert tools[0].input_schema == bridge.tools[0]["inputSchema"] + # Equal but not aliased: beta_tool stores the dict by reference, so the + # builder must hand out copies of the bridge's cached metas. + assert tools[0].input_schema is not bridge.tools[0]["inputSchema"] + assert tools[0].description == bridge.tools[0]["description"] + # Calls route over the bridge; None-valued arguments mean "omitted". + out = tools[1].call({"doc_name": "x.pdf", "folder_id": None}) + assert bridge.calls == [("get_document", {"doc_name": "x.pdf"})] + assert json.loads(out)["success"] is True + + +def test_as_anthropic_tools_cloud_async_flavor(cloud_with_fake_bridge): + pytest.importorskip("anthropic") + from anthropic.lib.tools import BetaAsyncFunctionTool + cloud, created = cloud_with_fake_bridge + tools = cloud.as_anthropic_tools(asynchronous=True) + assert all(isinstance(tool, BetaAsyncFunctionTool) for tool in tools) + out = asyncio.run(tools[1].call({"doc_name": "x.pdf"})) + assert created["bridge"].calls == [("get_document", {"doc_name": "x.pdf"})] + assert json.loads(out)["success"] is True + + +def test_as_anthropic_tools_cloud_management_opt_in(cloud_with_fake_bridge): + pytest.importorskip("anthropic") + cloud, _ = cloud_with_fake_bridge + names = [tool.name + for tool in cloud.as_anthropic_tools(include_management=True)] + assert names == ["search_documents", "get_document", + "remove_document", "unannotated_tool"] + + +def test_as_anthropic_tools_cloud_contains_bridge_errors(cloud_with_fake_bridge): + """Bridge failures become error envelopes raised as ToolError — the + runner turns that into a tool_result with is_error: true and the + envelope as content.""" + pytest.importorskip("anthropic") + from anthropic.lib.tools import ToolError + cloud, created = cloud_with_fake_bridge + tools = cloud.as_anthropic_tools() + + def boom(name, arguments): + raise RuntimeError("bridge down") + + created["bridge"].call_tool = boom + with pytest.raises(ToolError) as excinfo: + tools[0].call({"query": "q"}) + payload = json.loads(excinfo.value.content) + assert payload["errorCode"] == "INTERNAL_ERROR" + assert "bridge down" in payload["error"] + + +def test_agent_tools_work_without_frameworks(client, store_path, monkeypatch): + monkeypatch.setitem(sys.modules, "agents", None) + monkeypatch.setitem(sys.modules, "claude_agent_sdk", None) + monkeypatch.setitem(sys.modules, "anthropic", None) + seed_doc(store_path, "pi-a", "report.pdf") + browse = client.agent_tools()[0] + assert "report.pdf" in browse() + + +# ── cloud agent_tools: MCP bridge ── + +class _FakeBridge: + def __init__(self, url, headers): + self.url = url + self.headers = headers + self.calls = [] + read_only = {"readOnlyHint": True, "openWorldHint": False} + self.tools = [ + { + "name": "search_documents", + "description": "ESCALATION tool — keyword search.", + "annotations": read_only, + "inputSchema": { + "type": "object", + "properties": { + "query": {"type": "string", "description": "Keyword query."}, + "limit": {"type": "number", "default": 10}, + }, + "required": ["query"], + }, + }, + { + "name": "get_document", + "description": "Check a document's status.", + "annotations": read_only, + "inputSchema": { + "type": "object", + "properties": { + "doc_name": {"type": "string"}, + "folder_id": {"anyOf": [{"type": "string"}, + {"type": "null"}]}, + }, + "required": ["doc_name"], + }, + }, + { + "name": "remove_document", + "description": "Permanently delete documents.", + "annotations": {"readOnlyHint": False, "destructiveHint": True}, + "inputSchema": { + "type": "object", + "properties": {"doc_names": {"type": "array"}}, + "required": ["doc_names"], + }, + }, + { + "name": "unannotated_tool", + "description": "A tool the server sent without annotations.", + "inputSchema": {"type": "object", "properties": {}, + "required": []}, + }, + ] + + def list_tools(self): + return self.tools + + def instructions(self): + return "SERVER GUIDANCE" + + def call_tool(self, name, arguments): + self.calls.append((name, arguments)) + return json.dumps({"success": True, "tool": name, + "args": arguments}), False + + +@pytest.fixture +def cloud_with_fake_bridge(monkeypatch): + import pageindex.mcp_bridge as mcp_bridge + created = {} + + def factory(url, headers): + created["bridge"] = _FakeBridge(url, headers) + return created["bridge"] + + monkeypatch.setattr(mcp_bridge, "McpBridge", factory) + return PageIndexCloudClient(api_key="pi-test-key"), created + + +def test_cloud_agent_tools_discover_live_tool_set(cloud_with_fake_bridge): + cloud, created = cloud_with_fake_bridge + tools = cloud.agent_tools() + bridge = created["bridge"] + assert bridge.url == "https://api.pageindex.ai/mcp" + assert bridge.headers == {"Authorization": "Bearer pi-test-key"} + # Default: only tools the server marks read-only; unannotated tools are + # treated as non-read-only. + assert [t.__name__ for t in tools] == ["search_documents", "get_document"] + assert "ESCALATION tool" in tools[0].__doc__ + + +def test_cloud_agent_tools_management_gate(cloud_with_fake_bridge): + cloud, _ = cloud_with_fake_bridge + names = [t.__name__ for t in cloud.agent_tools(include_management=True)] + assert names == ["search_documents", "get_document", "remove_document", + "unannotated_tool"] + + +def test_cloud_agent_tools_signatures_from_schema(cloud_with_fake_bridge): + import inspect + cloud, _ = cloud_with_fake_bridge + search, get_document = cloud.agent_tools() + params = inspect.signature(search).parameters + assert list(params) == ["query", "limit"] + assert params["query"].default is inspect.Parameter.empty + assert params["limit"].default == 10 + assert search.__annotations__["query"] is str + folder_param = inspect.signature(get_document).parameters["folder_id"] + assert folder_param.default is None + # The live server encodes nullables as anyOf; the annotation must still + # come out Optional[str], not Any. + from typing import Optional + assert get_document.__annotations__["folder_id"] == Optional[str] + + +def test_cloud_agent_tools_proxy_and_drop_none(cloud_with_fake_bridge): + cloud, created = cloud_with_fake_bridge + _, get_document = cloud.agent_tools() + result = json.loads(get_document("report.pdf")) + assert result["tool"] == "get_document" + assert result["args"] == {"doc_name": "report.pdf"} # folder_id=None dropped + assert created["bridge"].calls == [("get_document", {"doc_name": "report.pdf"})] + + +def test_cloud_agent_tools_null_description_survives(): + """A server may send description: null — .get(key, default) does not + apply the default to it, and agent_tools() died with a TypeError while + the _tool_specs path handled the same payload fine.""" + from pageindex.agent_tools import _make_bridge_function + + class _Bridge: + def call_tool(self, name, arguments): + return json.dumps({"success": True}), False + + tool = _make_bridge_function(_Bridge(), { + "name": "search_documents", + "description": None, + "inputSchema": {"type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"]}, + }) + assert tool.__name__ == "search_documents" + assert json.loads(tool(query="x"))["success"] is True + + +def test_cloud_agent_tools_call_errors_contained(cloud_with_fake_bridge): + cloud, created = cloud_with_fake_bridge + search, _ = cloud.agent_tools() + created["bridge"].call_tool = lambda *a, **k: (_ for _ in ()).throw( + RuntimeError("network down")) + payload = json.loads(search(query="x")) + assert payload["errorCode"] == "INTERNAL_ERROR" + assert "network down" in payload["error"] + + +def test_cloud_agent_tools_list_failure_raises(monkeypatch): + import pageindex.mcp_bridge as mcp_bridge + + class _DeadBridge: + def __init__(self, url, headers): + pass + + def list_tools(self): + raise PageIndexAPIError("Could not connect") + + monkeypatch.setattr(mcp_bridge, "McpBridge", _DeadBridge) + cloud = PageIndexCloudClient(api_key="pi-test-key") + with pytest.raises(PageIndexAPIError, match="Could not connect"): + cloud.agent_tools() + + +def test_mcp_bridge_protocol(monkeypatch): + import requests as requests_mod + from pageindex.mcp_bridge import McpBridge + import pageindex.mcp_bridge as mcp_bridge + + posts = [] + + class _Resp: + def __init__(self, status, body=None, headers=None, text=""): + self.status_code = status + self._body = body + self.headers = headers or {"Content-Type": "application/json"} + self.text = text or (json.dumps(body) if body else "") + self.content = self.text.encode("utf-8") + + def json(self): + if self._body is None: + raise ValueError("no body") + return self._body + + session_alive = {"first": True} + + def fake_post(url, json=None, headers=None, timeout=None): + posts.append({"payload": json, "headers": headers}) + method = json.get("method") + rid = json.get("id") + if method == "initialize": + return _Resp(200, {"jsonrpc": "2.0", "id": rid, + "result": {"protocolVersion": "2025-06-18", + "instructions": "SERVER GUIDANCE"}}, + {"Content-Type": "application/json", + "Mcp-Session-Id": "sess-1"}) + if method == "notifications/initialized": + return _Resp(202) + if method == "tools/list": + # SSE-framed response exercises the event-stream parser; the + # em-dash guards UTF-8 decoding (SSE is UTF-8 by spec). + body = {"jsonrpc": "2.0", "id": rid, + "result": {"tools": [{"name": "t1", + "description": "reads — never writes"}], + "nextCursor": None}} + import json as json_mod + return _Resp(200, None, + {"Content-Type": "text/event-stream"}, + f"event: message\ndata: {json_mod.dumps(body)}\n\n") + if method == "tools/call": + if session_alive["first"]: + session_alive["first"] = False + return _Resp(404, text="session expired") + return _Resp(200, {"jsonrpc": "2.0", "id": rid, "result": { + "content": [{"type": "text", "text": "hello"}, + {"type": "text", "text": "world"}]}}) + raise AssertionError(f"unexpected method {method}") + + # Replace the module's own `requests` binding — patching the shared + # requests module would leak the fake process-wide. + monkeypatch.setattr(mcp_bridge, "requests", types.SimpleNamespace( + post=fake_post, RequestException=requests_mod.RequestException)) + bridge = McpBridge("https://api.pageindex.ai/mcp", + {"Authorization": "Bearer k"}) + + tools = bridge.list_tools() + assert tools == [{"name": "t1", "description": "reads — never writes"}] + # Captured during the handshake — serving it must not post again. + posts_before = len(posts) + assert bridge.instructions() == "SERVER GUIDANCE" + assert len(posts) == posts_before + list_headers = posts[-1]["headers"] + assert list_headers["Mcp-Session-Id"] == "sess-1" + assert list_headers["MCP-Protocol-Version"] == "2025-06-18" + assert list_headers["Authorization"] == "Bearer k" + + # First tools/call 404s (expired session) → re-initialize → retry succeeds. + text, is_error = bridge.call_tool("t1", {"a": 1}) + assert (text, is_error) == ("hello\nworld", False) + methods = [p["payload"]["method"] for p in posts] + assert methods.count("initialize") == 2 + # The expired session's negotiated state must not leak into the new + # handshake. + reinit = [p for p in posts if p["payload"].get("method") == "initialize"][1] + assert "MCP-Protocol-Version" not in reinit["headers"] + assert "Mcp-Session-Id" not in reinit["headers"] + + +def test_mcp_bridge_400_is_an_error_not_session_expiry(monkeypatch): + """The spec's expired-session status is 404; a 400 is an ordinary bad + request — treating it as expiry replayed the rejected call (running a + management tool's side effect twice) behind a spurious re-initialize.""" + import requests as requests_mod + import pageindex.mcp_bridge as mcp_bridge + from pageindex.mcp_bridge import McpBridge + + posts = [] + + class _Resp: + def __init__(self, status, body=None, headers=None, text=""): + self.status_code = status + self._body = body + self.headers = headers or {"Content-Type": "application/json"} + self.text = text or (json.dumps(body) if body else "") + self.content = self.text.encode("utf-8") + + def json(self): + if self._body is None: + raise ValueError("no body") + return self._body + + def fake_post(url, json=None, headers=None, timeout=None): + posts.append(json.get("method")) + rid = json.get("id") + if json.get("method") == "initialize": + return _Resp(200, {"jsonrpc": "2.0", "id": rid, + "result": {"protocolVersion": "2025-06-18"}}, + {"Content-Type": "application/json", + "Mcp-Session-Id": "sess-1"}) + if json.get("method") == "notifications/initialized": + return _Resp(202) + return _Resp(400, text="unknown tool") + + monkeypatch.setattr(mcp_bridge, "requests", types.SimpleNamespace( + post=fake_post, RequestException=requests_mod.RequestException)) + bridge = McpBridge("https://api.pageindex.ai/mcp", + {"Authorization": "Bearer k"}) + with pytest.raises(PageIndexAPIError, match="HTTP 400"): + bridge.call_tool("nope", {}) + # Exactly one call attempt, no replay, no re-initialize; the live + # session survives for the next request. + assert posts.count("tools/call") == 1 + assert posts.count("initialize") == 1 + assert bridge._session_id == "sess-1" + + +def test_mcp_bridge_init_notification_bars_concurrent_requests(monkeypatch): + """No thread may send a request between the initialize handshake and + notifications/initialized — strict servers reject such requests with + HTTP 400, which the bridge never replays. The notification's fake + transport stalls to hold that window open; a racing thread would post + its tools/list inside it.""" + import threading + import requests as requests_mod + import pageindex.mcp_bridge as mcp_bridge + from pageindex.mcp_bridge import McpBridge + + events = [] + events_lock = threading.Lock() + in_notification = threading.Event() + + class _Resp: + def __init__(self, status, body=None): + self.status_code = status + self._body = body + self.headers = {"Content-Type": "application/json"} + self.text = json.dumps(body) if body else "" + self.content = self.text.encode("utf-8") + + def json(self): + if self._body is None: + raise ValueError("no body") + return self._body + + def fake_post(url, json=None, headers=None, timeout=None): + method = json.get("method") + with events_lock: + events.append(("start", method)) + if method == "notifications/initialized": + in_notification.set() + time.sleep(0.2) + rid = json.get("id") + if method == "initialize": + resp = _Resp(200, {"jsonrpc": "2.0", "id": rid, + "result": {"protocolVersion": "2025-06-18"}}) + elif method == "notifications/initialized": + resp = _Resp(202) + else: + resp = _Resp(200, {"jsonrpc": "2.0", "id": rid, + "result": {"tools": [], "nextCursor": None}}) + with events_lock: + events.append(("end", method)) + return resp + + monkeypatch.setattr(mcp_bridge, "requests", types.SimpleNamespace( + post=fake_post, RequestException=requests_mod.RequestException)) + bridge = McpBridge("https://api.pageindex.ai/mcp", + {"Authorization": "Bearer k"}) + + errors = [] + + def list_tools(): + try: + bridge.list_tools() + except BaseException as exc: + errors.append(exc) + + first = threading.Thread(target=list_tools) + first.start() + assert in_notification.wait(5), "handshake never reached the notification" + second = threading.Thread(target=list_tools) + second.start() + first.join(5) + second.join(5) + assert not first.is_alive() and not second.is_alive() + assert not errors + + notified = events.index(("end", "notifications/initialized")) + first_list = events.index(("start", "tools/list")) + assert notified < first_list, ( + f"tools/list overtook notifications/initialized: {events}") + assert events.count(("start", "initialize")) == 1 + + +def test_mcp_bridge_blob_blocks_become_stubs(): + """Non-text content used to be json.dumps'd wholesale, handing the + model the raw base64 payload of an image tool's response.""" + from pageindex.mcp_bridge import McpBridge + + bridge = McpBridge("https://api.pageindex.ai/mcp", {}) + blob = "A" * 8192 # ~6 KB decoded + bridge._request = lambda method, params: {"content": [ + {"type": "text", "text": "Page 3 of report.pdf"}, + {"type": "image", "mimeType": "image/png", "data": blob}, + ]} + text, is_error = bridge.call_tool("get_document_image", {}) + assert not is_error + assert "Page 3 of report.pdf" in text + assert "AAAA" not in text + assert "[image/png content omitted: ~6 KB]" in text + + +# ── review-round regressions ── + +def test_synth_optional_no_default_param_is_nullable(): + """A non-required, no-default schema param must annotate Optional, or + strict schemas force the model to always send a value (browse.query).""" + from pageindex.agent_tools import _make_bridge_function, TOOL_CONTRACT + from typing import get_args + + class _Bridge: + def call_tool(self, name, args): + return json.dumps(args), False + + meta = {"name": "browse_documents", + "description": "d", + "inputSchema": TOOL_CONTRACT["browse_documents"]["schema"]} + fn = _make_bridge_function(_Bridge(), meta) + assert type(None) in get_args(fn.__annotations__["query"]) + + +def test_synth_array_params_keep_their_item_type(): + """The schema→annotation round-trip flattened arrays to bare `list`; + function_tool then emits {"type": "array", "items": {}}, which strict + function calling rejects.""" + from typing import Optional + from pageindex.agent_tools import _make_bridge_function + + class _Bridge: + def call_tool(self, name, args): + return json.dumps(args), False + + meta = {"name": "remove_documents", "description": "d", + "inputSchema": { + "type": "object", + "properties": { + "doc_ids": {"type": "array", "items": {"type": "string"}}, + "tags": {"anyOf": [{"type": "array", + "items": {"type": "integer"}}, + {"type": "null"}]}, + "mixed": {"type": "array", + "items": {"type": ["string", "null"]}}, + }, + "required": ["doc_ids", "mixed"], + }} + fn = _make_bridge_function(_Bridge(), meta) + assert fn.__annotations__["doc_ids"] == list[str] + assert fn.__annotations__["tags"] == Optional[list[int]] + # A type-array in items (nullable elements) degrades to bare list — + # it must not crash the build on an unhashable dict key. + assert fn.__annotations__["mixed"] == list + + +def test_synth_escape_hatches(): + from pageindex.agent_tools import _make_bridge_function + + calls = [] + + class _Bridge: + def call_tool(self, name, args): + calls.append((name, args)) + return "ok", False + + # Tool named "_invoke" must not recurse into itself. + invoke_named = _make_bridge_function(_Bridge(), { + "name": "_invoke", "description": "d", + "inputSchema": {"type": "object", "properties": {"x": {"type": "string"}}, + "required": ["x"]}}) + assert invoke_named("v") == "ok" + assert calls[-1] == ("_invoke", {"x": "v"}) + + # Param named "dict" must not shadow the builtin. + dict_param = _make_bridge_function(_Bridge(), { + "name": "t", "description": "d", + "inputSchema": {"type": "object", "properties": {"dict": {"type": "string"}}, + "required": ["dict"]}}) + assert dict_param("v") == "ok" + assert calls[-1] == ("t", {"dict": "v"}) + + # Non-identifier tool name still gets a real signature. + import inspect + dashed = _make_bridge_function(_Bridge(), { + "name": "page-content.v2", "description": "d", + "inputSchema": {"type": "object", "properties": {"a": {"type": "string"}}, + "required": ["a"]}}) + assert dashed.__name__ == "page-content.v2" + assert list(inspect.signature(dashed).parameters) == ["a"] + assert dashed("v") == "ok" + + +def test_annotation_for_both_nullable_encodings(): + """Servers have emitted nullables as type-arrays and as anyOf unions; + both must map to Optional, not degrade to Any.""" + from typing import Optional + from pageindex.agent_tools import _annotation_for + assert _annotation_for({"type": "string"}) is str + assert _annotation_for({"type": ["string", "null"]}) == Optional[str] + assert (_annotation_for({"anyOf": [{"type": "string"}, {"type": "null"}]}) + == Optional[str]) + + +def test_cloud_agent_tools_empty_filter_raises(monkeypatch): + import pageindex.mcp_bridge as mcp_bridge + + class _AllWriteBridge: + def __init__(self, url, headers): + pass + + def list_tools(self): + return [{"name": "remove_document", + "annotations": {"readOnlyHint": False}, + "inputSchema": {"type": "object", "properties": {}}}] + + monkeypatch.setattr(mcp_bridge, "McpBridge", _AllWriteBridge) + cloud = PageIndexCloudClient(api_key="pi-test-key") + with pytest.raises(PageIndexAPIError, match="annotation"): + cloud.agent_tools() + assert len(cloud.agent_tools(include_management=True)) == 1 + + +def test_bridge_call_tool_surfaces_iserror(monkeypatch): + import requests as requests_mod + import pageindex.mcp_bridge as mcp_bridge + from pageindex.mcp_bridge import McpBridge + + class _Resp: + def __init__(self, status, body=None): + self.status_code = status + self._body = body + self.headers = {"Content-Type": "application/json"} + self.text = json.dumps(body) if body else "" + self.content = self.text.encode("utf-8") + + def json(self): + if self._body is None: + raise ValueError("no body") + return self._body + + def fake_post(url, json=None, headers=None, timeout=None): + method = json.get("method") + rid = json.get("id") + if method == "initialize": + return _Resp(200, {"jsonrpc": "2.0", "id": rid, "result": {}}) + if method == "notifications/initialized": + return _Resp(202) + return _Resp(200, {"jsonrpc": "2.0", "id": rid, "result": { + "isError": True, + "content": [{"type": "text", "text": '{"error": "denied"}'}]}}) + + monkeypatch.setattr(mcp_bridge, "requests", types.SimpleNamespace( + post=fake_post, RequestException=requests_mod.RequestException)) + bridge = McpBridge("https://api.pageindex.ai/mcp", {}) + assert bridge.call_tool("t", {}) == ('{"error": "denied"}', True) + + +def test_bridge_rejects_mismatched_reply_id(monkeypatch): + """A result-bearing message with the wrong id must not be returned as + this call's reply.""" + import requests as requests_mod + import pageindex.mcp_bridge as mcp_bridge + from pageindex.mcp_bridge import McpBridge + + class _Resp: + def __init__(self, status, body=None): + self.status_code = status + self._body = body + self.headers = {"Content-Type": "application/json"} + self.text = json.dumps(body) if body else "" + self.content = self.text.encode("utf-8") + + def json(self): + if self._body is None: + raise ValueError("no body") + return self._body + + def fake_post(url, json=None, headers=None, timeout=None): + method = json.get("method") + rid = json.get("id") + if method == "initialize": + return _Resp(200, {"jsonrpc": "2.0", "id": rid, "result": {}}) + if method == "notifications/initialized": + return _Resp(202) + return _Resp(200, {"jsonrpc": "2.0", "id": rid - 1, # stale reply + "result": {"content": [{"type": "text", + "text": "old"}]}}) + + monkeypatch.setattr(mcp_bridge, "requests", types.SimpleNamespace( + post=fake_post, RequestException=requests_mod.RequestException)) + bridge = McpBridge("https://api.pageindex.ai/mcp", {}) + with pytest.raises(PageIndexAPIError, match="no reply matching"): + bridge.call_tool("t", {}) + + +def test_sse_crlf_multi_message(): + from pageindex.mcp_bridge import _parse_sse + body = ('event: message\r\ndata: {"jsonrpc":"2.0","method":"notifications/progress"}\r\n\r\n' + 'event: message\r\ndata: {"jsonrpc":"2.0","id":7,"result":{"ok":true}}\r\n\r\n') + messages = _parse_sse(body) + assert len(messages) == 2 + assert messages[1]["result"] == {"ok": True} + + +def test_bridge_transport_error_is_pageindex_error(monkeypatch): + import requests as requests_mod + import pageindex.mcp_bridge as mcp_bridge + from pageindex.mcp_bridge import McpBridge + + def dead_post(*args, **kwargs): + raise requests_mod.ConnectionError("dns down") + + monkeypatch.setattr(mcp_bridge, "requests", types.SimpleNamespace( + post=dead_post, RequestException=requests_mod.RequestException)) + bridge = McpBridge("https://api.pageindex.ai/mcp", {}) + with pytest.raises(PageIndexAPIError, match="Could not reach"): + bridge.list_tools() + + +def test_await_completion_preserves_metadata_over_null_refetch(monkeypatch): + """A status refetch that nulls out metadata must not clobber the + listing's copy (setdefault is a no-op on an existing None value).""" + import pageindex.agent_tools as agent_tools_mod + monkeypatch.setattr(agent_tools_mod, "time", types.SimpleNamespace( + monotonic=time.monotonic, sleep=lambda seconds: None)) + + class _Client: + def get_document(self, doc_id): + return {"id": doc_id, "status": "completed", "metadata": None} + + entry = {"id": "pi-x", "status": "processing", + "metadata": {"team": "research"}} + merged = agent_tools_mod._await_completion(_Client(), entry, True) + assert merged["status"] == "completed" + assert merged["metadata"] == {"team": "research"} + + +def test_browse_time_sort_uses_native_pagination(client, store_path, monkeypatch): + """Time-sorted browsing must page through list_documents directly, not + fetch the whole library to slice one window.""" + for index in range(3): + seed_doc(store_path, f"pi-{index}", f"doc{index}.pdf", + created_at=f"2026-08-0{index + 1}T10:00:00.000000") + calls = [] + original = client.list_documents + + def spy(**kwargs): + calls.append(kwargs) + return original(**kwargs) + + monkeypatch.setattr(client, "list_documents", spy) + payload, is_error = run(client, "browse_documents", limit=2) + assert not is_error + assert calls == [{"limit": 2, "offset": 0}] + assert [d["name"] for d in payload["documents"]] == ["doc2.pdf", "doc1.pdf"] + assert payload["has_more"] is True and payload["next_offset"] == 2 + + +def test_all_documents_survives_short_pages_and_missing_total(): + """The full-library walk behind every name resolution must trust what + actually arrives: a server capping page size, omitting `total`, or + sending total: null silently truncated the library (or raised).""" + from pageindex.agent_tools import _all_documents + + docs = [{"id": f"pi-{index}"} for index in range(120)] + + def make_client(total_field, page_cap): + class _Client: + calls = 0 + + def list_documents(self, limit, offset): + type(self).calls += 1 + page = {"documents": docs[offset:offset + min(limit, + page_cap)]} + if total_field != "omit": + page["total"] = total_field + return page + return _Client() + + assert _all_documents(make_client(120, 50)) == docs # short pages + assert _all_documents(make_client("omit", 100)) == docs # no total + assert _all_documents(make_client(None, 100)) == docs # total: null + exact = make_client(120, 100) # well-behaved server: + assert _all_documents(exact) == docs + assert type(exact).calls == 2 # ...total still saves the empty page + + +def test_null_arguments_mean_omitted(client, store_path): + """Adapters that forward the model's null values verbatim (the Claude + MCP handler) used to trip parameter validation — None ≡ omitted is + enforced once, in call_tool.""" + seed_doc(store_path, "pi-a", "report.pdf") + payload, is_error = run(client, "browse_documents", folder_id=None, + sort=None, query=None) + assert not is_error + assert [doc["name"] for doc in payload["documents"]] == ["report.pdf"] + + +def test_page_spec_span_bomb_rejected(client, store_path): + """An absurd range must be rejected arithmetically, not expanded into + billions of integers in the caller's process.""" + seed_doc(store_path, "pi-a", "report.pdf") + payload, is_error = run(client, "get_page_content", doc_name="report.pdf", + pages="1-1000000000") + assert is_error and payload["errorCode"] == "INVALID_INPUT" + assert "Too many pages" in payload["error"] + + +def test_agent_instructions_shadowed_doc_id_raises(client, store_path): + seed_doc(store_path, "pi-old", "report.pdf", + created_at="2026-08-01T10:00:00.000000") + seed_doc(store_path, "pi-new", "report.pdf", + created_at="2026-08-02T10:00:00.000000") + with pytest.raises(PageIndexAPIError, match="shadowed"): + client.agent_instructions(doc_id="pi-old") + text = client.agent_instructions(doc_id="pi-new") + assert "report.pdf" in text + + +def test_wait_tolerates_transient_network_failures(fake_cloud_client, monkeypatch): + import requests as requests_mod + cloud = fake_cloud_client(["processing", "completed"]) + original = cloud._api.get_document + state = {"raised": False} + + def flaky(doc_id): + if not state["raised"]: + state["raised"] = True + raise requests_mod.ConnectionError("network blip") + return original(doc_id) + + monkeypatch.setattr(cloud._api, "get_document", flaky) + assert cloud.submit_document("x.pdf", wait=True) == {"doc_id": "pi-fake"} + + +def test_failed_document_status_message(client, store_path): + seed_doc(store_path, "pi-a", "broken.pdf") + import pageindex.agent_tools as agent_tools_mod + payload, is_error = agent_tools_mod._not_ready_error( + "broken.pdf", "failed", "structure retrieval", timed_out=False) + assert is_error + assert "failed" in payload["error"] + assert any("submit_document" in option + for option in payload["next_steps"]["options"]) + + +def test_hosted_gate_is_the_endpoint(): + """The URL is the gate on hosted mode too — no approval-flow gating, + the read-only endpoint simply has no write tools.""" + pytest.importorskip("agents") + cloud = PageIndexCloudClient(api_key="pi-test-key") + gated = cloud.as_openai_tools(hosted=True)[0].tool_config + assert gated["server_url"] == "https://api.pageindex.ai/mcp?tools=read" + assert gated["require_approval"] == "never" + open_config = cloud.as_openai_tools(hosted=True, + include_management=True)[0].tool_config + assert open_config["server_url"] == "https://api.pageindex.ai/mcp" + assert open_config["require_approval"] == "never" + + +def test_wait_tolerates_transient_poll_failures(fake_cloud_client, monkeypatch): + cloud = fake_cloud_client(["processing", "completed"]) + original = cloud._api.get_document + state = {"raised": False} + + def flaky(doc_id): + if not state["raised"]: + state["raised"] = True + raise PageIndexAPIError("502") + return original(doc_id) + + monkeypatch.setattr(cloud._api, "get_document", flaky) + assert cloud.submit_document("x.pdf", wait=True) == {"doc_id": "pi-fake"} + + +LIVE_KEY = os.getenv("PAGEINDEX_API_KEY") + + +@pytest.mark.skipif(not LIVE_KEY, reason="PAGEINDEX_API_KEY not set") +def test_live_cloud_contract_parity(): + """Real-drift detector: the frozen contract must match the live server + on every shared tool, including the annotations the gates rely on.""" + from pageindex.mcp_bridge import McpBridge + bridge = McpBridge("https://api.pageindex.ai/mcp", + {"Authorization": f"Bearer {LIVE_KEY}"}) + live = {t["name"]: t for t in bridge.list_tools()} + for name, ours in TOOL_CONTRACT.items(): + real = live.get(name) + assert real is not None, f"{name} missing from live tools/list" + assert real.get("description") == ours["description"], name + real_schema = real.get("inputSchema") or {} + real_props = real_schema.get("properties") or {} + assert set(real_props) == set(ours["schema"]["properties"]), name + # Full per-param equality: a drifted type, default, enum, or bound + # breaks calls just as surely as a renamed parameter. + for param, spec in ours["schema"]["properties"].items(): + assert real_props[param] == spec, (name, param) + assert (sorted(real_schema.get("required") or []) + == sorted(ours["schema"].get("required", []))), name + for key, value in (ours.get("annotations") or {}).items(): + assert (real.get("annotations") or {}).get(key) == value, (name, key) + + +@pytest.mark.skipif(not LIVE_KEY, reason="PAGEINDEX_API_KEY not set") +def test_live_cloud_envelope_field_parity(tmp_path): + """Response-envelope drift alarm: every field the local tools emit must + exist in the live cloud tool's response for the analogous call — a cloud + rename of a shared field (has_more, next_offset, content, ...) fails + here. Guidance wording is deliberately localized and not compared.""" + from pageindex.mcp_bridge import McpBridge + bridge = McpBridge("https://api.pageindex.ai/mcp", + {"Authorization": f"Bearer {LIVE_KEY}"}) + cloud_browse = json.loads( + bridge.call_tool("browse_documents", {"limit": 2})[0]) + assert cloud_browse.get("success") is True and cloud_browse["documents"] + doc_name = cloud_browse["documents"][0]["name"] + cloud = { + "browse_documents": cloud_browse, + "get_document": json.loads(bridge.call_tool( + "get_document", {"doc_name": doc_name})[0]), + "get_document_structure": json.loads(bridge.call_tool( + "get_document_structure", {"doc_name": doc_name})[0]), + "get_page_content": json.loads(bridge.call_tool( + "get_page_content", {"doc_name": doc_name, "pages": "1"})[0]), + } + + store = str(tmp_path / "store") + local_client = PageIndexLocalClient(storage_path=store) + seed_doc(store, "pi-parity", "parity.pdf") + local = { + "browse_documents": run(local_client, "browse_documents")[0], + "get_document": run(local_client, "get_document", + doc_name="parity.pdf")[0], + "get_document_structure": run(local_client, "get_document_structure", + doc_name="parity.pdf")[0], + "get_page_content": run(local_client, "get_page_content", + doc_name="parity.pdf", pages="1")[0], + } + + for name in cloud: + assert cloud[name].get("success") is True, name + missing = set(local[name]) - set(cloud[name]) + assert not missing, (name, missing) + assert (set(local[name]["next_steps"]) + <= set(cloud[name]["next_steps"]) | {"auto_retry"}), name + + local_doc = local["browse_documents"]["documents"][0] + cloud_doc = cloud_browse["documents"][0] + assert set(local_doc) - set(cloud_doc) <= {"metadata"} + + local_nodes = local["get_document_structure"]["structure"] + cloud_nodes = cloud["get_document_structure"]["structure"] + local_node = local_nodes[0] if isinstance(local_nodes, list) else local_nodes + cloud_node = cloud_nodes[0] if isinstance(cloud_nodes, list) else cloud_nodes + assert (set(local_node) + <= set(cloud_node) | {"page_index", "prefix_summary"}) + + assert (set(local["get_page_content"]["content"][0]) + <= set(cloud["get_page_content"]["content"][0])) + + +@pytest.mark.skipif(not LIVE_KEY, reason="PAGEINDEX_API_KEY not set") +def test_live_cloud_instructions_nonempty(): + """The empty-instructions guard raises for cloud clients; the real + server must actually serve instructions in its initialize result.""" + from pageindex.mcp_bridge import McpBridge + bridge = McpBridge("https://api.pageindex.ai/mcp", + {"Authorization": f"Bearer {LIVE_KEY}"}) + assert bridge.instructions() + + +# ── agent_instructions ── + +def test_agent_instructions_default(client): + text = client.agent_instructions() + assert text == AGENT_INSTRUCTIONS + assert "READING WORKFLOW" in text + assert "browse_documents" in text + assert "search_documents" not in text + assert "get_folder_structure" not in text + assert 'sort="relevance"' not in text # cloud-side capability + + +def test_agent_instructions_with_doc_id(client, store_path): + seed_doc(store_path, "pi-a", "report.pdf") + text = client.agent_instructions(doc_id="pi-a") + assert text.startswith(AGENT_INSTRUCTIONS) + assert "The user has specified document: report.pdf" in text + + seed_doc(store_path, "pi-b", "other.pdf") + multi = client.agent_instructions(doc_id=["pi-a", "pi-b"]) + assert "The user has specified documents: report.pdf, other.pdf" in multi + + with pytest.raises(PageIndexAPIError): + client.agent_instructions(doc_id="pi-missing") + + +def test_local_instructions_name_only_local_tools(): + """The local instructions are trimmed from the cloud server's; every + tool they name must exist in the local registry, or the trim drifted.""" + named = set(re.findall(r"\b(\w+)\(", AGENT_INSTRUCTIONS)) + assert named + assert named <= set(tool_names(include_management=True)) + + +def test_cloud_agent_instructions_served_live(monkeypatch): + """Cloud clients serve the server's live instructions from the MCP + initialize handshake — over the same bridge session as agent_tools().""" + import pageindex.mcp_bridge as mcp_bridge + created = [] + + class _Bridge(_FakeBridge): + def __init__(self, url, headers): + super().__init__(url, headers) + created.append(self) + + def instructions(self): + return "LIVE CLOUD GUIDANCE" + + monkeypatch.setattr(mcp_bridge, "McpBridge", _Bridge) + cloud = PageIndexCloudClient(api_key="pi-test-key") + cloud.agent_tools() + assert cloud.agent_instructions() == "LIVE CLOUD GUIDANCE" + assert len(created) == 1 + + +def test_cloud_bridge_cache_threadsafe_and_pickle_clean(monkeypatch): + """One bridge per client even under concurrent first calls, and the + bridge lives off the instance so cloud clients stay picklable.""" + import pickle + import threading + import time as time_mod + import pageindex.mcp_bridge as mcp_bridge + created = [] + + class _Bridge(_FakeBridge): + def __init__(self, url, headers): + time_mod.sleep(0.01) # widen the construction window + super().__init__(url, headers) + created.append(self) + + def instructions(self): + return "LIVE" + + monkeypatch.setattr(mcp_bridge, "McpBridge", _Bridge) + cloud = PageIndexCloudClient(api_key="pi-test-key") + workers = ([threading.Thread(target=cloud.agent_tools) for _ in range(4)] + + [threading.Thread(target=cloud.agent_instructions) + for _ in range(4)]) + for worker in workers: + worker.start() + for worker in workers: + worker.join() + assert len(created) == 1 + pickle.dumps(cloud) + + +def test_cloud_agent_instructions_blank_or_nonstring_raises(monkeypatch): + """Whitespace-only or non-string initialize.instructions must hit the + same honest error as a missing one — never a blank system prompt.""" + import pageindex.mcp_bridge as mcp_bridge + + for bad in (" \n\t ", {"not": "a string"}): + class _SilentBridge: + def __init__(self, url, headers): + pass + + def instructions(self, _value=bad): + return _value + + monkeypatch.setattr(mcp_bridge, "McpBridge", _SilentBridge) + cloud = PageIndexCloudClient(api_key="pi-test-key") + with pytest.raises(PageIndexAPIError, match="no agent instructions"): + cloud.agent_instructions() + + +def test_cloud_agent_instructions_empty_raises(monkeypatch): + """An empty server response must raise, not silently substitute the + subset guidance — same posture as the annotation-regression guard.""" + import pageindex.mcp_bridge as mcp_bridge + + class _SilentBridge: + def __init__(self, url, headers): + pass + + def instructions(self): + return None + + monkeypatch.setattr(mcp_bridge, "McpBridge", _SilentBridge) + cloud = PageIndexCloudClient(api_key="pi-test-key") + with pytest.raises(PageIndexAPIError, match="no agent instructions"): + cloud.agent_instructions() + + +# ── submit_document(wait=True) ── + +class _FakeCloudAPI: + def __init__(self, statuses): + self._statuses = list(statuses) + self.polls = 0 + + def submit_document(self, **kwargs): + return {"doc_id": "pi-fake"} + + def get_document(self, doc_id): + self.polls += 1 + status = (self._statuses.pop(0) if len(self._statuses) > 1 + else self._statuses[0]) + return {"id": doc_id, "status": status} + + +@pytest.fixture +def fake_cloud_client(tmp_path, monkeypatch): + monkeypatch.setattr(client_module, "time", types.SimpleNamespace( + monotonic=time.monotonic, sleep=lambda seconds: None)) + + def build(statuses): + cloud = PageIndexLocalClient(storage_path=str(tmp_path / "unused")) + cloud._api = _FakeCloudAPI(statuses) + return cloud + return build + + +def test_submit_wait_polls_until_completed(fake_cloud_client): + cloud = fake_cloud_client(["processing", "processing", "completed"]) + result = cloud.submit_document("whatever.pdf", wait=True) + assert result == {"doc_id": "pi-fake"} + assert cloud._api.polls == 3 + + +def test_submit_wait_raises_on_failed(fake_cloud_client): + cloud = fake_cloud_client(["processing", "failed"]) + with pytest.raises(PageIndexAPIError, match="failed"): + cloud.submit_document("whatever.pdf", wait=True) + + +def test_submit_wait_times_out(fake_cloud_client, monkeypatch): + clock = {"now": 0.0} + + def fake_monotonic(): + clock["now"] += 700.0 + return clock["now"] + + monkeypatch.setattr(client_module, "time", types.SimpleNamespace( + monotonic=fake_monotonic, sleep=lambda seconds: None)) + cloud = fake_cloud_client(["processing"]) + with pytest.raises(PageIndexAPIError, match="Timed out"): + cloud.submit_document("whatever.pdf", wait=True) + + +def test_submit_without_wait_does_not_poll(fake_cloud_client): + cloud = fake_cloud_client(["processing"]) + assert cloud.submit_document("whatever.pdf") == {"doc_id": "pi-fake"} + assert cloud._api.polls == 0 + + +def test_submit_warns_when_stored_name_differs(fake_cloud_client): + cloud = fake_cloud_client(["processing"]) + cloud._api.submit_document = lambda **kwargs: { + "doc_id": "pi-fake", "name": "whatever_1.pdf"} + with pytest.warns(UserWarning, match='stored as "whatever_1.pdf"'): + result = cloud.submit_document("docs/whatever.pdf") + assert result["name"] == "whatever_1.pdf" + + +def test_submit_wait_poll_error_carries_doc_id(fake_cloud_client, monkeypatch): + """A poll that dies on transient errors must keep the uploaded doc_id + recoverable, like the timeout and failed branches do.""" + cloud = fake_cloud_client(["processing"]) + + def boom(doc_id): + raise PageIndexAPIError("Failed to get document metadata: 502") + + monkeypatch.setattr(cloud, "get_document", boom) + with pytest.raises(PageIndexAPIError, match="pi-fake"): + cloud.submit_document("whatever.pdf", wait=True) + + +def test_config_helpers_reject_empty_doc_id_on_cloud(): + """An explicitly empty scope must not silently widen to the whole + library — cloud has no tool-layer allowlist to enforce it.""" + cloud = PageIndexCloudClient(api_key="pi-test-key") + with pytest.raises(PageIndexAPIError, match="doc_id is empty"): + cloud.openai_agent_config(doc_id=[]) + with pytest.raises(PageIndexAPIError, match="doc_id is empty"): + cloud.anthropic_runner_config(model="claude-sonnet-4-5", doc_id=[]) + with pytest.raises(PageIndexAPIError, match="doc_id is empty"): + cloud.claude_agent_config(doc_id=[]) + + +def test_call_tool_coerces_string_booleans(client, store_path, monkeypatch): + """Models routinely send booleans as JSON strings — "false" must not + read as True (a full wait_for_completion stall).""" + seed_doc(store_path, "pi-1", "a.pdf") + seen = {} + real = agent_tools_module._await_completion + + def spy(spy_client, entry, wait): + seen["wait"] = wait + return real(spy_client, entry, wait) + + monkeypatch.setattr(agent_tools_module, "_await_completion", spy) + run(client, "get_document", doc_name="a.pdf", wait_for_completion="false") + assert seen["wait"] is False + run(client, "get_document", doc_name="a.pdf", wait_for_completion="true") + assert seen["wait"] is True + + +def test_call_tool_rejects_non_object_arguments(client): + """A non-dict arguments value must come back as the guided envelope, + never raise into the agent loop.""" + for bad in ([1, 2], "doc_name=a.pdf"): + text, is_error = call_tool(client, "browse_documents", bad) + payload = json.loads(text) + assert is_error and payload["errorCode"] == "INVALID_INPUT" + + +def test_remove_document_repeated_name_deletes_once(client, store_path): + seed_doc(store_path, "pi-1", "a.pdf") + payload, is_error = run(client, "remove_document", + doc_names=["a.pdf", "a.pdf"]) + assert not is_error + assert payload["results"] == [{"doc_name": "a.pdf", "status": "deleted"}] + assert "1 of 1" in payload["next_steps"]["summary"] + + +def test_page_spec_cap_counts_distinct_pages(): + """Overlapping parts are normal tree output (a parent section plus its + children) — the cap is on the union, not the sum.""" + pages, error = agent_tools_module._parse_page_spec("1-5000,2000-9000", + "a.pdf") + assert error is None and pages is not None and len(pages) == 9000 + pages, error = agent_tools_module._parse_page_spec("1-10001", "a.pdf") + assert pages is None and "Too many pages" in error[0]["error"] + + +def test_browse_documents_pages_by_rows_returned(): + """A backend that caps its page size must not make the cursor skip + documents, and a null total must not crash (same guards as + _all_documents).""" + class _Capping: + def list_documents(self, limit, offset): + docs = [{"id": f"pi-{i}", "name": f"d{i}.pdf", + "status": "completed"} + for i in range(offset, min(offset + 5, 30))] + return {"documents": docs, "total": 30} + + payload, is_error = agent_tools_module._browse_documents(_Capping(), + limit=10) + assert not is_error + assert payload["has_more"] is True and payload["next_offset"] == 5 + + class _NullTotal: + def list_documents(self, limit, offset): + return {"documents": [{"name": "d.pdf", "status": "completed"}], + "total": None} + + payload, is_error = agent_tools_module._browse_documents(_NullTotal(), + limit=10) + assert not is_error + assert payload["has_more"] is False and payload["next_offset"] is None + + +def test_agent_instructions_carry_user_metadata(client, store_path): + """The targeting block promises names and metadata; local get_document + keeps the 7-key detail wire shape, so the tags come from the listing.""" + seed_doc(store_path, "pi-1", "report.pdf", + metadata={"quarter": "Q3", "year": 2025}) + text = client.agent_instructions(doc_id="pi-1") + assert '"quarter": "Q3"' in text and '"year": 2025' in text diff --git a/tests/test_client.py b/tests/test_client.py index 50b7f5178..aad4e3979 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -47,7 +47,7 @@ def fake_page_index_main(doc, opt=None, logger=None, page_list=None): "doc_description": "A test document.", "structure": json.loads(json.dumps(STRUCTURE))} monkeypatch.setattr(page_index_module, "page_index_main", fake_page_index_main) - return local_client.submit_document(sample_pdf)["doc_id"] + return local_client.submit_document(sample_pdf, mode="standard")["doc_id"] # ── constructor ── @@ -151,6 +151,16 @@ def test_get_page_content(local_client, indexed_doc): local_client.get_page_content(indexed_doc, "abc") +def test_get_page_content_span_bomb_rejected(local_client, indexed_doc): + """An absurd range must be rejected arithmetically, not expanded into + a billion integers in the caller's process (the tool layer already + refused; the public client method did not).""" + with pytest.raises(ValueError, match="spans more than 10000"): + local_client.get_page_content(indexed_doc, "1-1000001") + # At the bound itself the spec still parses. + assert local_client.get_page_content(indexed_doc, "5-10004") == [] + + def test_submit_does_not_create_cwd_logs(local_client, sample_pdf, tmp_path, monkeypatch): monkeypatch.chdir(tmp_path) def fake_page_index_main(doc, opt=None, logger=None, page_list=None): @@ -158,15 +168,60 @@ def fake_page_index_main(doc, opt=None, logger=None, page_list=None): return {"doc_name": "sample.pdf", "doc_description": None, "structure": json.loads(json.dumps(STRUCTURE))} monkeypatch.setattr(page_index_module, "page_index_main", fake_page_index_main) - local_client.submit_document(sample_pdf) + local_client.submit_document(sample_pdf, mode="standard") assert not (tmp_path / "logs").exists() +def test_submit_duplicate_name_gets_suffix(local_client, sample_pdf, monkeypatch): + """Mirror the cloud upload: a second submit of the same file name is + stored as name_1, not as a same-name duplicate.""" + def fake_page_index_main(doc, opt=None, logger=None, page_list=None): + return {"doc_name": "sample.pdf", "doc_description": "d", + "structure": json.loads(json.dumps(STRUCTURE))} + monkeypatch.setattr(page_index_module, "page_index_main", fake_page_index_main) + first = local_client.submit_document(sample_pdf, mode="standard") + assert first["name"] == "sample.pdf" + with pytest.warns(UserWarning, match='stored as "sample_1.pdf"'): + second = local_client.submit_document(sample_pdf, mode="standard") + assert second["name"] == "sample_1.pdf" + names = {d["id"]: d["name"] + for d in local_client.list_documents()["documents"]} + assert names[first["doc_id"]] == "sample.pdf" + assert names[second["doc_id"]] == "sample_1.pdf" + + +def test_submit_duplicate_name_exhaustion(local_client, monkeypatch): + api = local_client._api + metas = ([{"name": "x.pdf"}] + + [{"name": f"x_{num}.pdf"} for num in range(1, 100)]) + monkeypatch.setattr(api._store, "list_metas", lambda: metas) + with pytest.raises(PageIndexAPIError, match="Too many files"): + api._unique_doc_name("x.pdf") + + +def test_submit_name_exhaustion_rejects_before_indexing( + local_client, sample_pdf, monkeypatch, +): + api = local_client._api + metas = ([{"name": "sample.pdf"}] + + [{"name": f"sample_{num}.pdf"} for num in range(1, 100)]) + monkeypatch.setattr(api._store, "list_metas", lambda: metas) + monkeypatch.setattr( + page_index_module, "page_index_main", + lambda *args, **kwargs: pytest.fail( + "indexer ran despite name exhaustion"), + ) + with pytest.raises(PageIndexAPIError, match="Too many files"): + local_client.submit_document(sample_pdf, mode="standard") + + def test_submit_flash(local_client, sample_pdf, monkeypatch): calls = {} def fake_flash(pdf, summary=True, summary_model=None, **kwargs): calls["summary"] = summary calls["summary_model"] = summary_model + calls["optimize"] = kwargs.get("optimize") + calls["optimize_model"] = kwargs.get("optimize_model") return {"doc_name": "sample.pdf", "structure": [{"title": "Flash Root", "start_index": 1, "end_index": 2, "summary": "s", "nodes": []}]} @@ -174,13 +229,34 @@ def fake_flash(pdf, summary=True, summary_model=None, **kwargs): monkeypatch.setattr(pageindex.utils, "llm_completion", lambda model, prompt, **kw: "Flash description.") doc_id = local_client.submit_document(sample_pdf, mode="flash")["doc_id"] - assert calls == {"summary": True, "summary_model": local_client.summary_model} + assert calls == {"summary": True, "summary_model": local_client.summary_model, + "optimize": "full", + "optimize_model": local_client.summary_model} root = local_client.get_tree(doc_id)["result"][0] assert root["node_id"] == "0000" assert "Hello page one" in root["text"] assert local_client.get_document(doc_id)["description"] == "Flash description." +def test_submit_defaults_to_flash(local_client, sample_pdf, monkeypatch): + monkeypatch.setattr( + pageindex.flash, "page_index_flash", + lambda pdf, **kwargs: { + "doc_name": "sample.pdf", + "structure": [{"title": "Flash Root", "start_index": 1, + "end_index": 2, "summary": "s", "nodes": []}]}) + monkeypatch.setattr(pageindex.utils, "llm_completion", + lambda model, prompt, **kw: "Flash description.") + doc_id = local_client.submit_document(sample_pdf)["doc_id"] + assert local_client._api._store.get_meta(doc_id)["mode"] == "flash" + + +def test_page_index_flash_rejects_unknown_optimize(): + from pageindex.flash import page_index_flash + with pytest.raises(ValueError, match="optimize must be"): + page_index_flash("never-opened.pdf", optimize="off") + + def test_llm_completion_missing_key_raises_immediately(monkeypatch): import openai monkeypatch.delenv("OPENAI_API_KEY", raising=False) @@ -282,7 +358,7 @@ def test_submit_with_metadata(local_client, sample_pdf, monkeypatch): "doc_name": "sample.pdf", "doc_description": None, "structure": json.loads(json.dumps(STRUCTURE))}) tags = {"project": "alpha", "year": 2026} - doc_id = local_client.submit_document(sample_pdf, metadata=tags)["doc_id"] + doc_id = local_client.submit_document(sample_pdf, mode="standard", metadata=tags)["doc_id"] assert local_client.get_tree(doc_id)["metadata"] == tags assert local_client.get_ocr(doc_id)["metadata"] == tags assert local_client.list_documents()["documents"][0]["metadata"] == tags @@ -432,7 +508,8 @@ def test_torn_delete_never_lists_ghost(local_client, indexed_doc, tmp_path): def test_corrupt_doc_json_is_contained(local_client, indexed_doc, sample_pdf, tmp_path): - second = local_client.submit_document(sample_pdf)["doc_id"] + with pytest.warns(UserWarning): # same-name resubmit → stored as sample_1.pdf + second = local_client.submit_document(sample_pdf, mode="standard")["doc_id"] (tmp_path / "store" / "docs" / indexed_doc / "doc.json").write_text("{truncated") # manifest still holds a good copy of the meta — served consistently @@ -599,8 +676,12 @@ def test_retrieval_endpoints_cloud_only(local_client): local_client.get_retrieval("any") -def test_chat_completions_cloud_only(local_client): - with pytest.raises(PageIndexAPIError, match="not yet supported in local mode"): +def test_chat_completions_local_needs_agents_extra(local_client, monkeypatch): + """Local chat is implemented (see test_local_chat.py); without the + openai-agents extra it raises the actionable install error.""" + import sys + monkeypatch.setitem(sys.modules, "agents", None) + with pytest.raises(PageIndexAPIError, match="pageindex\\[openai\\]"): local_client.chat_completions( messages=[{"role": "user", "content": "q"}]) @@ -710,3 +791,21 @@ def test_cloud_chat_stream_parsing(cloud, monkeypatch): messages=[{"role": "user", "content": "q"}], stream=True, stream_metadata=True)) assert {"object": "chat.completion.citations", "citations": []} in chunks + + +def test_cloud_chat_accepts_query_string(cloud): + client, calls, fake = cloud + fake.payload = {"choices": [{"message": {"content": "ok"}}]} + client.chat_completions("What status?") + assert calls[-1]["json"]["messages"] == [ + {"role": "user", "content": "What status?"}] + with pytest.raises(PageIndexAPIError, match="non-empty string"): + client.chat_completions(" ") + + +def test_parse_pages_overlap_counts_union(): + from pageindex.client import _parse_pages + pages = _parse_pages("1-5000,2000-9000") + assert len(pages) == 9000 and pages[0] == 1 and pages[-1] == 9000 + with pytest.raises(ValueError, match="spans more than"): + _parse_pages("1-10001") diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py new file mode 100644 index 000000000..37553b27f --- /dev/null +++ b/tests/test_local_chat.py @@ -0,0 +1,1273 @@ +"""Local chat surfaces: three protocols over fake backends — no network, +no LLM keys. Tool execution runs for real against a seeded local store.""" +import asyncio +import json +import sys +import types + +import httpx # via the hard `openai` dependency +import pytest + +import pageindex.local_chat as local_chat +from pageindex import (PageIndexAPIError, PageIndexCloudClient, + PageIndexLocalClient) +from pageindex.local_chat import CHAT_HEADER +from pageindex.local_store import DocStore + + +def seed_doc(storage_path, doc_id, name): + pages = [{"page_index": 1, "markdown": "Page one text about apples"}] + tree = [{"title": "Doc", "node_id": "0000", "start_index": 1, + "end_index": 1, "summary": "root summary", "text": "ROOT"}] + meta = { + "id": doc_id, "name": name, "description": "A test document", + "status": "completed", "createdAt": "2026-08-01T10:00:00.123000", + "pageNum": 1, "folderId": None, "metadata": None, "mode": "standard", + } + DocStore(storage_path).save_document(doc_id, meta, tree, pages) + return doc_id + + +@pytest.fixture +def store_path(tmp_path): + return str(tmp_path / "store") + + +@pytest.fixture +def client(store_path): + return PageIndexLocalClient(storage_path=str(store_path)) + + +# ── OpenAI engine fakes (chat_completions / responses) ── +# Section-scoped skips: each engine's tests skip independently, so a +# machine with only one extra installed still covers the other surface. + +try: + import agents # noqa: F401 + _HAS_AGENTS = True +except ImportError: + _HAS_AGENTS = False + +needs_agents = pytest.mark.skipif(not _HAS_AGENTS, + reason="openai-agents not installed") + + +def _msg_item(text): + from openai.types.responses import (ResponseOutputMessage, + ResponseOutputText) + return ResponseOutputMessage( + id="msg_1", type="message", role="assistant", status="completed", + content=[ResponseOutputText(type="output_text", text=text, + annotations=[])]) + + +def _call_item(name, arguments, call_id="call_1"): + from openai.types.responses import ResponseFunctionToolCall + return ResponseFunctionToolCall( + id="fc_1", type="function_call", call_id=call_id, name=name, + arguments=json.dumps(arguments), status="completed") + + +def _usage(): + from agents.usage import Usage + return Usage(requests=1, input_tokens=10, output_tokens=5, + total_tokens=15) + + +if _HAS_AGENTS: + from agents.models.interface import Model # noqa: E402 +else: # pragma: no cover - placeholder so the class statement parses + Model = object + + +class FakeModel(Model): + """Scripted backend: one list of output items per model turn.""" + + def __init__(self, turns): + self.turns = list(turns) + self.inputs = [] + self.instructions = [] + self.deltas_emitted = 0 + + def _record(self, system_instructions, input): + self.instructions.append(system_instructions) + items = input if isinstance(input, list) else [input] + self.inputs.append( + [dict(item) if isinstance(item, dict) else item + for item in items]) + + async def get_response(self, system_instructions, input, model_settings, + tools, output_schema, handoffs, tracing, + **kwargs): + from agents.items import ModelResponse + self._record(system_instructions, input) + # Mimic the real model's transport hop when a test attaches one, so + # the transport-level status recorder sees each turn. + transport = getattr(getattr(self, "_client", None), "responses", None) + if transport is not None: + await transport.create() + return ModelResponse(output=self.turns.pop(0), usage=_usage(), + response_id=None) + + async def stream_response(self, system_instructions, input, + model_settings, tools, output_schema, handoffs, + tracing, **kwargs): + import asyncio as aio + from openai.types.responses import (Response, ResponseCompletedEvent, + ResponseTextDeltaEvent) + from openai.types.responses.response_usage import ( + InputTokensDetails, OutputTokensDetails, ResponseUsage) + block_from = getattr(self, "block_from", None) + if block_from is not None and len(self.inputs) + 1 >= block_from: + while True: # released only by task cancellation + await aio.sleep(0.01) + self._record(system_instructions, input) + output = self.turns.pop(0) + sequence = 0 + for item in output: + if item.type == "message": + pieces = getattr(self, "pieces", ("The ", "answer")) + for piece in pieces: + sequence += 1 + self.deltas_emitted += 1 + yield ResponseTextDeltaEvent( + type="response.output_text.delta", delta=piece, + content_index=0, item_id=item.id, output_index=0, + logprobs=[], sequence_number=sequence) + if getattr(self, "no_terminal", False): + return # backend died mid-stream: no terminal event + sequence += 1 + yield ResponseCompletedEvent( + type="response.completed", sequence_number=sequence, + response=Response( + id="resp_fake", created_at=0.0, model="fake", + object="response", output=output, parallel_tool_calls=False, + tool_choice="auto", tools=[], + usage=ResponseUsage( + input_tokens=10, output_tokens=5, total_tokens=15, + input_tokens_details=InputTokensDetails( + cached_tokens=0, cache_write_tokens=0), + output_tokens_details=OutputTokensDetails( + reasoning_tokens=0)))) + + +@pytest.fixture +def fake_model(monkeypatch): + state = {} + + def install(turns): + fake = FakeModel(turns) + state["protocols"] = [] + + def factory(protocol, model_name): + state["protocols"].append((protocol, model_name)) + return fake + + monkeypatch.setattr(local_chat, "_openai_model", factory) + return fake + + install.state = state + return install + + +# ── chat_completions ── + +@needs_agents +def test_chat_completions_end_to_end(client, store_path, fake_model): + seed_doc(store_path, "pi-a", "report.pdf") + fake = fake_model([ + [_call_item("get_document", {"doc_name": "report.pdf"})], + [_msg_item("The answer")], + ]) + result = client.chat_completions( + [{"role": "user", "content": "What status?"}]) + assert result["id"].startswith("chatcmpl-") + assert result["object"] == "chat.completion" + assert result["choices"][0]["message"] == {"role": "assistant", + "content": "The answer"} + assert result["choices"][0]["finish_reason"] == "stop" + assert result["usage"] == {"prompt_tokens": 20, "completion_tokens": 10, + "total_tokens": 30, + "prompt_tokens_details": {"cached_tokens": 0}, + "completion_tokens_details": + {"reasoning_tokens": 0}} + assert fake_model.state["protocols"][0][0] == "chat" + # The tool ran for real: turn 2's input carries its output. + turn2 = json.dumps(fake.inputs[1]) + assert "report.pdf" in turn2 and "completed" in turn2 + # Managed instructions: header + the local agent guidance. + assert fake.instructions[0].startswith(CHAT_HEADER) + assert "READING WORKFLOW" in fake.instructions[0] + + +@needs_agents +def test_chat_completions_system_and_doc_block(client, store_path, fake_model): + doc_id = seed_doc(store_path, "pi-a", "report.pdf") + fake = fake_model([[_msg_item("ok")]]) + client.chat_completions( + [{"role": "system", "content": "Answer in French."}, + {"role": "user", "content": "hi"}], + doc_id=doc_id) + assert fake.instructions[0].endswith("Answer in French.") + first_item = fake.inputs[0][0] + assert "The user has specified document: report.pdf" in first_item["content"] + + +@needs_agents +def test_chat_completions_accepts_query_string(client, store_path, fake_model): + seed_doc(store_path, "pi-a", "report.pdf") + fake = fake_model([[_msg_item("Answer")]]) + result = client.chat_completions("What status?") + assert result["choices"][0]["message"]["content"] == "Answer" + assert fake.inputs[0][-1] == {"role": "user", "content": "What status?"} + with pytest.raises(PageIndexAPIError, match="non-empty string"): + client.chat_completions(" ") + + +@needs_agents +def test_chat_completions_validation(client, store_path, fake_model): + fake_model([[_msg_item("ok")]]) + with pytest.raises(PageIndexAPIError, match="cloud-only"): + client.chat_completions([{"role": "user", "content": "x"}], + enable_citations=True) + with pytest.raises(PageIndexAPIError, match="responses\\(\\) or messages"): + client.chat_completions([{"role": "tool", "content": "x"}]) + with pytest.raises(PageIndexAPIError, match="must be a string"): + client.chat_completions([{"role": "user", "content": [1]}]) + with pytest.raises(PageIndexAPIError, match="non-empty"): + client.chat_completions([]) + with pytest.raises(PageIndexAPIError, + match="Documents not found or access denied: a, b"): + client.chat_completions([{"role": "user", "content": "x"}], + doc_id=["a", "b"]) + + +@needs_agents +def test_chat_completions_stream_modes(client, store_path, fake_model): + fake_model([[_msg_item("The answer")]]) + pieces = list(client.chat_completions( + [{"role": "user", "content": "q"}], stream=True)) + assert pieces == ["The ", "answer"] + + fake_model([[_msg_item("The answer")]]) + chunks = list(client.chat_completions( + [{"role": "user", "content": "q"}], stream=True, + stream_metadata=True)) + assert chunks[0]["choices"][0]["delta"] == {"role": "assistant", + "content": ""} + assert chunks[-2]["choices"][0]["finish_reason"] == "stop" + assert chunks[-1]["choices"] == [] + assert chunks[-1]["usage"]["total_tokens"] == 15 + assert all(c["object"] == "chat.completion.chunk" for c in chunks[:-1]) + + +def test_chat_completions_missing_framework(client, monkeypatch): + monkeypatch.setitem(sys.modules, "agents", None) + with pytest.raises(PageIndexAPIError, match="pageindex\\[openai\\]"): + client.chat_completions([{"role": "user", "content": "x"}]) + + +def test_cloud_guards(): + cloud = PageIndexCloudClient(api_key="pi-test-key") + with pytest.raises(PageIndexAPIError, match="local-mode parameters"): + cloud.chat_completions([{"role": "user", "content": "x"}], model="m") + with pytest.raises(PageIndexAPIError, match="not available on PageIndex " + "cloud yet"): + cloud.responses("x") + with pytest.raises(PageIndexAPIError, match="not available on PageIndex " + "cloud yet"): + cloud.messages([{"role": "user", "content": "x"}], model="m", + max_tokens=10) + + +# ── responses ── + +@needs_agents +def test_responses_end_to_end(client, store_path, fake_model): + seed_doc(store_path, "pi-a", "report.pdf") + fake_model([ + [_call_item("get_document", {"doc_name": "report.pdf"})], + [_msg_item("The answer")], + ]) + result = client.responses("What status?") + assert result["id"].startswith("resp_") + assert result["object"] == "response" + assert result["status"] == "completed" + assert result["usage"] == { + "input_tokens": 20, + "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0}, + "output_tokens": 10, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 30} + assert fake_model.state["protocols"][0][0] == "responses" + assert [item.get("type", "message") for item in result["output"]] == [ + "function_call", "message"] + assert [item.get("type", "message") for item in result["items"]] == [ + "function_call", "function_call_output", "message"] + # The final item is the assistant answer. + assert "The answer" in json.dumps(result["output"][-1]) + + +@needs_agents +def test_responses_round_trip_extends_prefix(client, store_path, fake_model): + """The cache contract: a round-tripped call's first model input must + extend the previous call's final model input item-for-item.""" + seed_doc(store_path, "pi-a", "report.pdf") + first = fake_model([ + [_call_item("get_document", {"doc_name": "report.pdf"})], + [_msg_item("The answer")], + ]) + result = client.responses("What status?") + + second = fake_model([[_msg_item("Done")]]) + follow_up = ([{"role": "user", "content": "What status?"}] + + result["items"] + + [{"role": "user", "content": "and now?"}]) + client.responses(follow_up) + previous_final = first.inputs[-1] + assert second.inputs[0][:len(previous_final)] == previous_final + + +@needs_agents +def test_responses_round_trip_prefix_with_doc_id(client, store_path, fake_model): + """Same contract with doc targeting: re-passing the same doc_id re-sets + an identical leading block, so the prefix still extends item-for-item.""" + seed_doc(store_path, "pi-a", "report.pdf") + first = fake_model([ + [_call_item("get_document", {"doc_name": "report.pdf"})], + [_msg_item("The answer")], + ]) + result = client.responses("What status?", doc_id="pi-a") + + second = fake_model([[_msg_item("Done")]]) + follow_up = ([{"role": "user", "content": "What status?"}] + + result["items"] + + [{"role": "user", "content": "and now?"}]) + client.responses(follow_up, doc_id="pi-a") + previous_final = first.inputs[-1] + assert second.inputs[0][:len(previous_final)] == previous_final + + +@needs_agents +def test_doc_id_conversations_get_distinct_cache_keys(client, store_path, + fake_model, + monkeypatch): + """The doc-targeting block is byte-identical for every conversation + about a document — seeding the cache key on items[0] pooled them all + under one prompt_cache_key.""" + seed_doc(store_path, "pi-a", "report.pdf") + keys = [] + real = local_chat._run_kwargs + + def spy(max_turns, group_id): + keys.append(group_id) + return real(max_turns, group_id) + + monkeypatch.setattr(local_chat, "_run_kwargs", spy) + + fake_model([[_msg_item("a")]]) + result = client.responses("What is the CAGR?", doc_id="pi-a") + fake_model([[_msg_item("b")]]) + client.responses("Summarize section 3.", doc_id="pi-a") + assert keys[0] != keys[1] # unrelated conversations never pool + + fake_model([[_msg_item("c")]]) + follow_up = ([{"role": "user", "content": "What is the CAGR?"}] + + result["items"] + + [{"role": "user", "content": "and now?"}]) + client.responses(follow_up, doc_id="pi-a") + assert keys[2] == keys[0] # a continuation keeps its conversation's key + + fake_model([[_msg_item("d")]]) + client.chat_completions("What is the CAGR?", doc_id="pi-a") + fake_model([[_msg_item("e")]]) + client.chat_completions("Summarize section 3.", doc_id="pi-a") + assert keys[3] != keys[4] # same property on the chat surface + + +@needs_agents +def test_responses_stream_passthrough(client, store_path, fake_model): + seed_doc(store_path, "pi-a", "report.pdf") + fake_model([ + [_call_item("get_document", {"doc_name": "report.pdf"})], + [_msg_item("The answer")], + ]) + events = list(client.responses("q", stream=True)) + types = [event.get("type") for event in events] + assert "response.output_text.delta" in types + assert not [event for event in events + if event.get("item", {}).get("type") == "function_call_output"] + assert types[-1] == "response.completed" + final = events[-1]["response"] + assert final["status"] == "completed" + assert final["usage"]["total_tokens"] == 30 + assert [item.get("type", "message") for item in final["output"]] == [ + "function_call", "message"] + assert [item.get("type", "message") for item in final["items"]] == [ + "function_call", "function_call_output", "message"] + # output_index addresses the logical response.output: turn 2's deltas + # are re-based past turn 1's item instead of restarting at 0. + last_delta = [event for event in events + if event.get("type") == "response.output_text.delta"][-1] + assert (final["output"][last_delta["output_index"]] + .get("type", "message") == "message") + + +@needs_agents +def test_responses_envelope_validates_as_official_response(client, store_path, + fake_model): + """The conformance contract: the envelope parses with the official + openai SDK types, and the transcript survives in the extension field.""" + from openai.types.responses import Response + seed_doc(store_path, "pi-a", "report.pdf") + fake_model([ + [_call_item("get_document", {"doc_name": "report.pdf"})], + [_msg_item("The answer")], + ]) + result = client.responses("What status?") + parsed = Response.model_validate(result) + assert [item.type for item in parsed.output] == ["function_call", + "message"] + assert parsed.model_dump()["items"] == result["items"] + + +@needs_agents +def test_responses_stream_events_validate_as_official_events( + client, store_path, fake_model): + """Every stream event, terminal envelope included, parses with the + official event union.""" + from pydantic import TypeAdapter + from openai.types.responses import ResponseStreamEvent + seed_doc(store_path, "pi-a", "report.pdf") + fake_model([ + [_call_item("get_document", {"doc_name": "report.pdf"})], + [_msg_item("The answer")], + ]) + adapter = TypeAdapter(ResponseStreamEvent) + events = list(client.responses("q", stream=True)) + assert events + for event in events: + adapter.validate_python(event) + + +# ── messages (Anthropic engine) ── + +try: + import anthropic + _HAS_ANTHROPIC = True +except ImportError: + _HAS_ANTHROPIC = False + +needs_anthropic = pytest.mark.skipif(not _HAS_ANTHROPIC, + reason="anthropic not installed") + + +def _anthropic_message(content, stop_reason): + return { + "id": "msg_fake", "type": "message", "role": "assistant", + "model": "claude-test", "content": content, + "stop_reason": stop_reason, "stop_sequence": None, + "usage": {"input_tokens": 10, "output_tokens": 5}, + } + + +@pytest.fixture +def fake_anthropic(monkeypatch): + state = {"calls": []} + + def install(responses): + state["calls"].clear() + + def handler(request): + state["calls"].append(json.loads(request.content)) + body = responses[len(state["calls"]) - 1] + if isinstance(body, str): # pre-rendered SSE + return httpx.Response( + 200, content=body.encode(), + headers={"content-type": "text/event-stream"}) + return httpx.Response(200, json=body) + + fake = anthropic.Anthropic( + api_key="test", + http_client=httpx.Client(transport=httpx.MockTransport(handler))) + monkeypatch.setattr(local_chat, "_anthropic_client", lambda: fake) + return state["calls"] + + return install + + +@needs_anthropic +def test_messages_end_to_end(client, store_path, fake_anthropic): + seed_doc(store_path, "pi-a", "report.pdf") + calls = fake_anthropic([ + _anthropic_message( + [{"type": "tool_use", "id": "tu_1", "name": "get_document", + "input": {"doc_name": "report.pdf"}}], "tool_use"), + _anthropic_message([{"type": "text", "text": "The answer"}], + "end_turn"), + ]) + result = client.messages([{"role": "user", "content": "What status?"}], + model="claude-test", max_tokens=100) + assert result["stop_reason"] == "end_turn" + assert result["content"][0]["text"] == "The answer" + assert result["usage"]["input_tokens"] == 20 + assert result["usage"]["output_tokens"] == 10 + # Full new-turn sequence, valid for verbatim history append. + roles = [message["role"] for message in result["messages"]] + assert roles == ["assistant", "user", "assistant"] + tool_result = json.dumps(result["messages"][1]) + assert "tool_result" in tool_result and "report.pdf" in tool_result + + request = calls[0] + assert request["system"][0]["text"].startswith(CHAT_HEADER) + assert request["system"][0]["cache_control"] == {"type": "ephemeral"} + browse = next(t for t in request["tools"] + if t["name"] == "browse_documents") + assert "folder_id" not in browse["input_schema"]["properties"] + # Native prefix continuation: request 2 extends request 1's messages. + assert calls[1]["messages"][:len(calls[0]["messages"])] \ + == calls[0]["messages"] + + +@needs_anthropic +def test_messages_doc_block_and_system(client, store_path, fake_anthropic): + doc_id = seed_doc(store_path, "pi-a", "report.pdf") + calls = fake_anthropic([ + _anthropic_message([{"type": "text", "text": "ok"}], "end_turn"), + ]) + client.messages([{"role": "user", "content": "hi"}], model="claude-test", + max_tokens=100, doc_id=doc_id, system="Answer in French.") + system = calls[0]["system"] + assert "The user has specified document: report.pdf" in system[1]["text"] + assert system[-1]["text"] == "Answer in French." + + +@needs_anthropic +def test_messages_stream_passthrough(client, store_path, fake_anthropic): + sse = "\n".join([ + 'event: message_start', + 'data: {"type":"message_start","message":{"id":"msg_1","type":"message","role":"assistant","model":"claude-test","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":1}}}', + "", + 'event: content_block_start', + 'data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}', + "", + 'event: content_block_delta', + 'data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"The answer"}}', + "", + 'event: content_block_stop', + 'data: {"type":"content_block_stop","index":0}', + "", + 'event: message_delta', + 'data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":5}}', + "", + 'event: message_stop', + 'data: {"type":"message_stop"}', + "", + "", + ]) + fake_anthropic([sse]) + events = list(client.messages([{"role": "user", "content": "q"}], + model="claude-test", max_tokens=100, + stream=True)) + types = [event.type for event in events] + assert "content_block_delta" in types and "message_stop" in types + + +@needs_anthropic +def test_messages_accepts_query_string(client, fake_anthropic): + calls = fake_anthropic([ + _anthropic_message([{"type": "text", "text": "ok"}], "end_turn"), + ]) + result = client.messages("What status?", model="claude-test") + assert result["content"][0]["text"] == "ok" + assert calls[0]["messages"] == [{"role": "user", + "content": "What status?"}] + # The wire-required budget is table-setting, not a user obligation. + assert calls[0]["max_tokens"] == 8192 + with pytest.raises(PageIndexAPIError, match="non-empty string"): + client.messages(" ", model="claude-test") + + +@needs_anthropic +def test_messages_validation(client, fake_anthropic): + fake_anthropic([]) + with pytest.raises(PageIndexAPIError, match="non-empty"): + client.messages([], model="claude-test", max_tokens=100) + with pytest.raises(PageIndexAPIError, + match="Documents not found or access denied"): + client.messages([{"role": "user", "content": "x"}], + model="claude-test", max_tokens=100, doc_id="ghost") + + +@needs_anthropic +def test_messages_raises_when_runner_params_unreadable(client, fake_anthropic, + monkeypatch): + """The conversation is read back through set_messages_params (a mutator + used as a reader); if a vendor change stops it delivering params, the + envelope silently lost every tool turn — it must raise instead.""" + from anthropic.lib.tools import BetaToolRunner + fake_anthropic([ + _anthropic_message([{"type": "text", "text": "ok"}], "end_turn"), + ]) + monkeypatch.setattr(BetaToolRunner, "set_messages_params", + lambda self, params: None) + with pytest.raises(PageIndexAPIError, match="anthropic version"): + client.messages([{"role": "user", "content": "hi"}], + model="claude-test", max_tokens=100) + + +def test_messages_missing_framework(client, monkeypatch): + monkeypatch.setitem(sys.modules, "anthropic", None) + with pytest.raises(PageIndexAPIError, match="pageindex\\[anthropic\\]"): + client.messages([{"role": "user", "content": "x"}], + model="claude-test", max_tokens=100) + + +# ── review-round regressions ── + +def _anthropic_tool_use(tool_use_id="tu_1"): + return {"type": "tool_use", "id": tool_use_id, "name": "get_document", + "input": {"doc_name": "report.pdf"}} + + +@needs_agents +@pytest.mark.parametrize("surface", ["chat_completions", "responses"]) +@pytest.mark.parametrize("streaming", [False, True]) +def test_max_turns_wrapped(client, store_path, fake_model, surface, streaming): + """MaxTurnsExceeded is an engine-internal type; callers get the SDK's + own error, with the engine exception kept as the cause — on every + surface and both the non-stream and stream paths.""" + seed_doc(store_path, "pi-a", "report.pdf") + fake_model([ + [_call_item("get_document", {"doc_name": "report.pdf"})], + [_call_item("get_document", {"doc_name": "report.pdf"}, "call_2")], + [_msg_item("never reached")], + ]) + with pytest.raises(PageIndexAPIError, match=r"max_turns \(1\)") as caught: + result = getattr(client, surface)("q", max_turns=1, stream=streaming) + if streaming: + list(result) + assert type(caught.value.__cause__).__name__ == "MaxTurnsExceeded" + + +@needs_agents +def test_max_turns_rejects_non_positive(client, store_path, fake_model): + seed_doc(store_path, "pi-a", "report.pdf") + with pytest.raises(PageIndexAPIError, match="positive integer"): + client.chat_completions([{"role": "user", "content": "q"}], + max_turns=0) + + +def test_enable_citations_rejected_before_framework_check(client, monkeypatch): + monkeypatch.setitem(sys.modules, "agents", None) + with pytest.raises(PageIndexAPIError, match="cloud-only"): + client.chat_completions([{"role": "user", "content": "x"}], + enable_citations=True) + + +@needs_agents +def test_chat_stream_role_chunk_even_with_empty_output(client, fake_model): + fake_model([[]]) + chunks = list(client.chat_completions([{"role": "user", "content": "q"}], + stream=True, stream_metadata=True)) + assert chunks[0]["choices"][0]["delta"] == {"role": "assistant", + "content": ""} + assert chunks[-2]["choices"][0]["finish_reason"] == "stop" + + +@needs_agents +def test_responses_stream_single_completed_monotonic_sequence( + client, store_path, fake_model): + """One logical response per call: per-turn backend lifecycle events are + collapsed and sequence numbers never go backwards.""" + seed_doc(store_path, "pi-a", "report.pdf") + fake_model([ + [_call_item("get_document", {"doc_name": "report.pdf"})], + [_msg_item("The answer")], + ]) + events = list(client.responses("q", stream=True)) + completed = [event for event in events + if event.get("type") == "response.completed"] + assert len(completed) == 1 and events[-1] is completed[0] + sequences = [event["sequence_number"] for event in events + if "sequence_number" in event] + assert sequences == sorted(sequences) + assert len(set(sequences)) == len(sequences) + + +@needs_agents +def test_responses_envelope_fields_and_cache_group(client, store_path, + fake_model): + seed_doc(store_path, "pi-a", "report.pdf") + fake_model([[_msg_item("ok")]]) + result = client.responses("q") + names = {tool["name"] for tool in result["tools"]} + assert names == {"browse_documents", "get_document", + "get_document_structure", "get_page_content"} + assert all(tool["type"] == "function" for tool in result["tools"]) + assert result["instructions"].startswith(CHAT_HEADER) + assert result["parallel_tool_calls"] is True + assert result["tool_choice"] == "auto" + + +def test_conversation_group_id_stable_per_conversation(): + """Cache-routing key: openai-agents hashes group_id into the OpenAI + prompt_cache_key. A conversation's continuations must share one key + (same model/instructions/first item), and unrelated conversations must + not pool under it.""" + turn1 = [{"role": "user", "content": "q"}] + continuation = turn1 + [{"role": "assistant", "content": "a"}, + {"role": "user", "content": "and?"}] + key = local_chat._conversation_group_id("m", "sys", turn1) + assert key == local_chat._conversation_group_id("m", "sys", continuation) + assert key != local_chat._conversation_group_id( + "m", "sys", [{"role": "user", "content": "other"}]) + assert key != local_chat._conversation_group_id("m2", "sys", turn1) + assert key != local_chat._conversation_group_id("m", "sys2", turn1) + + +@needs_agents +def test_run_kwargs_sets_conversation_group_id(): + key = "pageindex-test" + assert (local_chat._run_kwargs(None, key)["run_config"].group_id == key) + + +@needs_agents +def test_responses_input_validation(client, fake_model): + fake_model([]) + for bad in ("", " ", [], [1], None): + with pytest.raises(PageIndexAPIError, match="input must be"): + client.responses(bad) + + +@needs_agents +def test_doc_id_scopes_tools_to_targeted_documents(client, store_path, + fake_model): + """doc_id is enforcement, not just a prompt: name-addressed reads of + out-of-scope documents fail and browse lists only the targeted set.""" + seed_doc(store_path, "pi-a", "report.pdf") + seed_doc(store_path, "pi-b", "payroll.pdf") + fake = fake_model([ + [_call_item("get_page_content", + {"doc_name": "payroll.pdf", "pages": "1"})], + [_call_item("browse_documents", {}, "call_2")], + [_msg_item("done")], + ]) + client.chat_completions("q", doc_id="pi-a") + + def tool_outputs(items): + return [item["output"] for item in items + if item.get("type") == "function_call_output"] + + assert "NOT_FOUND" in tool_outputs(fake.inputs[1])[-1] + browse = json.loads(tool_outputs(fake.inputs[2])[-1]) + assert [doc["name"] for doc in browse["documents"]] == ["report.pdf"] + + +@needs_agents +def test_empty_doc_id_is_an_empty_allowlist(client, store_path, fake_model): + """doc_id=[] scopes the agent to nothing; `or None` used to wash it + into unscoped full-library access.""" + seed_doc(store_path, "pi-a", "report.pdf") + fake = fake_model([ + [_call_item("browse_documents", {})], + [_msg_item("done")], + ]) + client.chat_completions("q", doc_id=[]) + outputs = [item["output"] for item in fake.inputs[1] + if item.get("type") == "function_call_output"] + assert json.loads(outputs[-1])["documents"] == [] + + +@needs_agents +def test_openai_model_resolves_provider_prefixes(): + """retrieve_model arrives normalized (litellm//); the + OpenAI SDK must never see that prefix as a wire model name.""" + pytest.importorskip("litellm") + from agents.extensions.models.litellm_model import LitellmModel + from agents.models.openai_chatcompletions import OpenAIChatCompletionsModel + from agents.models.openai_responses import OpenAIResponsesModel + + model = local_chat._openai_model("chat", "litellm/anthropic/claude-x") + assert isinstance(model, LitellmModel) and model.model == "anthropic/claude-x" + model = local_chat._openai_model("chat", "anthropic/claude-x") + assert isinstance(model, LitellmModel) and model.model == "anthropic/claude-x" + model = local_chat._openai_model("chat", "openai/gpt-5.2") + assert isinstance(model, OpenAIChatCompletionsModel) + assert str(model.model) == "gpt-5.2" + model = local_chat._openai_model("responses", "gpt-5.2") + assert isinstance(model, OpenAIResponsesModel) + assert str(model.model) == "gpt-5.2" + + +@needs_agents +def test_chat_refuses_unknown_litellm_provider(): + """A HuggingFace-style id (vLLM serving Qwen/...) must fail at build + time with the openai/ escape, not inside LiteLLM at request time.""" + pytest.importorskip("litellm") + for name in ("Qwen/Qwen2.5-7B-Instruct", "litellm/Qwen/Qwen2.5-7B-Instruct"): + with pytest.raises(PageIndexAPIError, match="openai/Qwen"): + local_chat._openai_model("chat", name) + + +@needs_agents +def test_responses_refuses_litellm_routed_models(store_path): + """LiteLLM speaks chat.completions, not /responses — the responses + protocol must refuse the silent downgrade, at agent-build time and + before any backend call.""" + for name in ("anthropic/claude-x", "litellm/anthropic/claude-x"): + with pytest.raises(PageIndexAPIError, match="Responses API"): + local_chat._openai_model("responses", name) + client = PageIndexLocalClient(storage_path=store_path, + retrieve_model="anthropic/claude-x") + with pytest.raises(PageIndexAPIError, match="chat_completions"): + client.responses("q") + + +@needs_agents +def test_envelope_model_strips_litellm_routing_prefix(store_path, fake_model): + """litellm/ is the SDK's routing marker, not a model name — the + OpenAI-shaped envelopes must report the model the provider serves.""" + seed_doc(store_path, "pi-a", "report.pdf") + client = PageIndexLocalClient(storage_path=store_path, + retrieve_model="anthropic/claude-x") + assert client.retrieve_model == "litellm/anthropic/claude-x" + fake_model([[_msg_item("ok")]]) + result = client.chat_completions("q") + assert result["model"] == "anthropic/claude-x" + fake_model([[_msg_item("ok")]]) + chunks = list(client.chat_completions("q", stream=True, + stream_metadata=True)) + assert {c["model"] for c in chunks} == {"anthropic/claude-x"} + + +@needs_agents +def test_envelope_model_strips_openai_routing_prefix(store_path, fake_model): + """openai/ is the other routing marker — both OpenAI-shaped envelopes + must report the name the provider actually serves.""" + seed_doc(store_path, "pi-a", "report.pdf") + client = PageIndexLocalClient(storage_path=store_path, + retrieve_model="openai/gpt-5.2") + fake_model([[_msg_item("ok")]]) + result = client.chat_completions("q") + assert result["model"] == "gpt-5.2" + fake_model([[_msg_item("ok")]]) + result = client.responses("q") + assert result["model"] == "gpt-5.2" + + +@needs_agents +def test_chat_missing_openai_key_fails_loud(monkeypatch): + """A missing backend credential surfaces as the SDK's own error type, + like every other precondition on the chat surfaces.""" + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + with pytest.raises(PageIndexAPIError, match="OPENAI_API_KEY"): + local_chat._openai_model("chat", "gpt-4o") + + +@needs_agents +def test_record_response_status_captures_last_status(): + class _Dumpable: + def __init__(self, data): + self._data = data + + def model_dump(self, mode=None): + return dict(self._data) + + async def create(*args, **kwargs): + return types.SimpleNamespace( + status="incomplete", + incomplete_details=_Dumpable({"reason": "max_output_tokens"}), + error=None) + + agent = types.SimpleNamespace(model=types.SimpleNamespace( + _client=types.SimpleNamespace( + responses=types.SimpleNamespace(create=create)))) + recorded = {} + local_chat._record_response_status(agent, recorded) + asyncio.run(agent.model._client.responses.create()) + assert recorded == {"status": "incomplete", + "incomplete_details": {"reason": "max_output_tokens"}, + "error": None} + + +@needs_agents +def test_responses_envelope_reports_backend_truncation(client, store_path, + fake_model): + """A final turn the backend reports as status "incomplete" must not be + dressed up as a clean completion.""" + seed_doc(store_path, "pi-a", "report.pdf") + fake = fake_model([[_msg_item("cut off mid-answer")]]) + + async def create(*args, **kwargs): + return types.SimpleNamespace( + status="incomplete", + incomplete_details={"reason": "max_output_tokens"}, + error=None) + + fake._client = types.SimpleNamespace( + responses=types.SimpleNamespace(create=create)) + result = client.responses("q") + assert result["status"] == "incomplete" + assert result["incomplete_details"] == {"reason": "max_output_tokens"} + assert result["error"] is None + + +@needs_agents +def test_chat_completions_wraps_framework_errors(client, store_path, + fake_model, monkeypatch): + """Both chat_completions paths surface engine failures as the SDK's + own error type, like responses().""" + from agents.exceptions import ModelBehaviorError + seed_doc(store_path, "pi-a", "report.pdf") + fake = fake_model([[_msg_item("never terminal")]]) + fake.no_terminal = True + with pytest.raises(PageIndexAPIError, match="agent backend failed"): + list(client.chat_completions("q", stream=True)) + + fake = fake_model([[_msg_item("x")]]) + + async def boom(*args, **kwargs): + raise ModelBehaviorError("backend broke") + + monkeypatch.setattr(fake, "get_response", boom) + with pytest.raises(PageIndexAPIError, match="agent backend failed"): + client.chat_completions("q") + + +@needs_agents +def test_responses_stream_wraps_framework_errors(client, store_path, + fake_model): + """A backend stream that dies without a terminal event surfaces as the + SDK's own error type, not a raw openai-agents exception.""" + seed_doc(store_path, "pi-a", "report.pdf") + fake = fake_model([[_msg_item("never terminal")]]) + fake.no_terminal = True + with pytest.raises(PageIndexAPIError, match="agent backend failed"): + list(client.responses("q", stream=True)) + + +class _TerminalModel(FakeModel): + """Engine-faithful backend terminal: openai-agents yields the + response.failed/response.incomplete lifecycle event, then raises.""" + terminal = "incomplete" + + async def stream_response(self, system_instructions, input, + model_settings, tools, output_schema, + handoffs, tracing, **kwargs): + from agents.exceptions import ModelBehaviorError + from openai.types.responses import (Response, ResponseFailedEvent, + ResponseIncompleteEvent, + ResponseTextDeltaEvent) + from openai.types.responses.response import IncompleteDetails + from openai.types.responses.response_error import ResponseError + self._record(system_instructions, input) + yield ResponseTextDeltaEvent( + type="response.output_text.delta", delta="partial ", + content_index=0, item_id="item_x", output_index=0, + logprobs=[], sequence_number=1) + response = Response( + id="resp_fake", created_at=0.0, model="fake", object="response", + output=[], parallel_tool_calls=False, tool_choice="auto", + tools=[], status=self.terminal, + incomplete_details=(IncompleteDetails(reason="max_output_tokens") + if self.terminal == "incomplete" else None), + error=(ResponseError(code="server_error", message="boom") + if self.terminal == "failed" else None)) + event_type = (ResponseIncompleteEvent if self.terminal == "incomplete" + else ResponseFailedEvent) + yield event_type(type=f"response.{self.terminal}", response=response, + sequence_number=2) + raise ModelBehaviorError(f"terminal: {self.terminal}") + + +@needs_agents +@pytest.mark.parametrize("terminal", ["incomplete", "failed"]) +def test_responses_stream_backend_terminal_states_are_events( + client, store_path, monkeypatch, terminal): + """response.failed / response.incomplete are protocol terminal states, + not engine failures: the stream must end with the honest terminal + event carrying the backend's status, not raise away the run.""" + seed_doc(store_path, "pi-a", "report.pdf") + fake = _TerminalModel([[]]) + fake.terminal = terminal + monkeypatch.setattr(local_chat, "_openai_model", + lambda protocol, model_name: fake) + events = list(client.responses("q", stream=True)) + assert events[0]["type"] == "response.output_text.delta" + last = events[-1] + assert last["type"] == f"response.{terminal}" + assert last["response"]["status"] == terminal + if terminal == "incomplete": + assert (last["response"]["incomplete_details"] + == {"reason": "max_output_tokens"}) + else: + assert last["response"]["error"]["message"] == "boom" + numbers = [event["sequence_number"] for event in events] + assert numbers == sorted(numbers) and len(set(numbers)) == len(numbers) + + +@needs_agents +def test_provider_errors_wrap_as_sdk_errors(client, store_path, fake_model, + monkeypatch): + """Raw provider exceptions (network, auth, rate limit) surface as + PageIndexAPIError on every OpenAI-engine path, never as openai types.""" + import openai + seed_doc(store_path, "pi-a", "report.pdf") + request = httpx.Request("POST", "https://backend.test") + + async def conn_err(*args, **kwargs): + raise openai.APIConnectionError(request=request) + + async def conn_err_stream(*args, **kwargs): + raise openai.APIConnectionError(request=request) + yield # unreached: makes this an async generator + + fake = fake_model([[_msg_item("x")], [_msg_item("x")]]) + monkeypatch.setattr(fake, "get_response", conn_err) + with pytest.raises(PageIndexAPIError, match="model backend failed"): + client.chat_completions("q") + with pytest.raises(PageIndexAPIError, match="model backend failed"): + client.responses("q") + monkeypatch.setattr(fake, "stream_response", conn_err_stream) + with pytest.raises(PageIndexAPIError, match="model backend failed"): + list(client.chat_completions("q", stream=True)) + with pytest.raises(PageIndexAPIError, match="model backend failed"): + list(client.responses("q", stream=True)) + + +@needs_anthropic +def test_messages_provider_errors_wrap_as_sdk_errors(client, store_path, + monkeypatch): + """Anthropic transport errors surface as PageIndexAPIError on both + messages() paths, never as anthropic types.""" + seed_doc(store_path, "pi-a", "report.pdf") + + def handler(request): + return httpx.Response(429, json={ + "type": "error", + "error": {"type": "rate_limit_error", "message": "slow down"}}) + + fake = anthropic.Anthropic( + api_key="test", max_retries=0, + http_client=httpx.Client(transport=httpx.MockTransport(handler))) + monkeypatch.setattr(local_chat, "_anthropic_client", lambda: fake) + with pytest.raises(PageIndexAPIError, match="model backend failed"): + client.messages("q", model="claude-test") + with pytest.raises(PageIndexAPIError, match="model backend failed"): + list(client.messages("q", model="claude-test", stream=True)) + + +@needs_agents +def test_chat_stream_close_at_opening_chunk_cancels_run(client, store_path, + fake_model, + monkeypatch): + """GeneratorExit at the opening chunk must still cancel the agent task: + the first yield sits inside the generator's try/finally.""" + seed_doc(store_path, "pi-a", "report.pdf") + fake = fake_model([[_msg_item("never")]]) + fake.block_from = 1 # turn 1 hangs until cancelled + captured = {} + + def capture(agen_factory): + captured["factory"] = agen_factory + return iter(()) # drive the async generator by hand instead + + monkeypatch.setattr(local_chat, "_stream_sync", capture) + client.chat_completions("q", stream=True, stream_metadata=True) + + async def drive(): + agen = captured["factory"]() + first = await agen.__anext__() + assert first["choices"][0]["delta"] == {"role": "assistant", + "content": ""} + await agen.aclose() + deadline = asyncio.get_running_loop().time() + 2.0 + pending = [] + while asyncio.get_running_loop().time() < deadline: + pending = [task for task in asyncio.all_tasks() + if task is not asyncio.current_task() + and not task.done()] + if not pending: + break + await asyncio.sleep(0.01) + return pending + + assert asyncio.run(drive()) == [] + + +@needs_agents +def test_stream_abandonment_cancels_pending_turn(client, store_path, + fake_model): + """Closing the iterator cancels the run even while it is awaiting the + backend: the blocked turn is torn down (pump thread exits) instead of + running — and billing — to completion in the background.""" + import threading + import time as time_mod + seed_doc(store_path, "pi-a", "report.pdf") + baseline = threading.active_count() + fake = fake_model([ + [_call_item("get_document", {"doc_name": "report.pdf"})], + [_msg_item("The answer")], + ]) + fake.block_from = 2 # turn 2 hangs until cancelled + stream = client.chat_completions([{"role": "user", "content": "q"}], + stream=True, stream_metadata=True) + next(stream) # the opening role chunk + stream.close() + deadline = time_mod.monotonic() + 3.0 + while (threading.active_count() > baseline + and time_mod.monotonic() < deadline): + time_mod.sleep(0.05) + assert threading.active_count() <= baseline + assert fake.deltas_emitted == 0 # turn 2 never produced output + + +@needs_anthropic +def test_messages_max_tokens_default_resolves_per_model(client, fake_anthropic): + """The wire-required budget must not exceed the model's ceiling: the + claude-3 generation caps output at 4096.""" + calls = fake_anthropic([ + _anthropic_message([{"type": "text", "text": "ok"}], "end_turn")]) + client.messages("q", model="claude-3-opus-20240229") + assert calls[0]["max_tokens"] == 4096 + calls = fake_anthropic([ + _anthropic_message([{"type": "text", "text": "ok"}], "end_turn")]) + client.messages("q", model="claude-sonnet-4-5") + assert calls[0]["max_tokens"] == 8192 + calls = fake_anthropic([ + _anthropic_message([{"type": "text", "text": "ok"}], "end_turn")]) + client.messages("q", model="claude-3-opus-20240229", max_tokens=1234) + assert calls[0]["max_tokens"] == 1234 + + +@needs_anthropic +def test_messages_tool_error_flagged_and_scoped(client, store_path, + fake_anthropic): + """Through the real runner: a failed call reaches Claude as a + tool_result with is_error true, and doc_id scoping makes out-of-scope + documents unreachable by name.""" + seed_doc(store_path, "pi-a", "report.pdf") + seed_doc(store_path, "pi-b", "secret.pdf") + calls = fake_anthropic([ + _anthropic_message([{"type": "tool_use", "id": "tu_1", + "name": "get_document", + "input": {"doc_name": "secret.pdf"}}], + "tool_use"), + _anthropic_message([{"type": "text", "text": "ok"}], "end_turn"), + ]) + client.messages("q", model="claude-test", doc_id="pi-a") + tool_result = calls[1]["messages"][-1]["content"][0] + assert tool_result["type"] == "tool_result" + assert tool_result.get("is_error") is True + assert "NOT_FOUND" in json.dumps(tool_result["content"]) + + +@needs_anthropic +def test_messages_envelope_json_and_no_internal_fields(client, store_path, + fake_anthropic): + seed_doc(store_path, "pi-a", "report.pdf") + fake_anthropic([ + _anthropic_message([_anthropic_tool_use()], "tool_use"), + _anthropic_message([{"type": "text", "text": "The answer"}], + "end_turn"), + ]) + result = client.messages([{"role": "user", "content": "q"}], + model="claude-test", max_tokens=100) + dumped = json.dumps(result) # the whole envelope must serialize + assert "parsed_output" not in dumped + + +@needs_anthropic +def test_messages_max_turns_truncation_round_trippable(client, store_path, + fake_anthropic): + """On a max_turns cut the runner has already appended the final turn — + no duplicate append, and the history stays valid for continuation.""" + seed_doc(store_path, "pi-a", "report.pdf") + calls = fake_anthropic([ + _anthropic_message([_anthropic_tool_use()], "tool_use"), + _anthropic_message([_anthropic_tool_use("tu_2")], "tool_use"), + ]) + result = client.messages([{"role": "user", "content": "q"}], + model="claude-test", max_tokens=100, max_turns=1) + assert len(calls) == 1 + assert result["stop_reason"] == "tool_use" + roles = [message["role"] for message in result["messages"]] + assert roles == ["assistant", "user"] # tool_use, tool_result — no dup + assert json.dumps(result).count('"tu_1"') == \ + json.dumps(result["messages"][0]).count('"tu_1"') \ + + json.dumps(result["messages"][1]).count('"tu_1"') \ + + json.dumps(result["content"]).count('"tu_1"') + json.dumps(result) + + +@needs_anthropic +def test_messages_tool_use_cut_by_max_tokens_not_duplicated(client, + store_path, + fake_anthropic): + """A max_tokens turn with complete tool_use blocks still executes and + is appended by the runner — keying the re-append guard on stop_reason + duplicated the tool_use id and broke verbatim continuation.""" + seed_doc(store_path, "pi-a", "report.pdf") + calls = fake_anthropic([ + _anthropic_message([_anthropic_tool_use()], "max_tokens"), + _anthropic_message([_anthropic_tool_use("tu_2")], "tool_use"), + ]) + result = client.messages([{"role": "user", "content": "q"}], + model="claude-test", max_tokens=100, max_turns=1) + assert len(calls) == 1 + assert result["stop_reason"] == "max_tokens" + roles = [message["role"] for message in result["messages"]] + assert roles == ["assistant", "user"] # tool_use, tool_result — no dup + assert json.dumps(result["messages"]).count('"tu_1"') == 2 # use + result + + +@needs_anthropic +def test_messages_refusal_with_tool_use_stays_appendable(client, store_path, + fake_anthropic): + """A refusal turn is never executed by the runner; its tool_use blocks + have no tool_result and must not enter the appendable history.""" + seed_doc(store_path, "pi-a", "report.pdf") + fake_anthropic([ + _anthropic_message([{"type": "text", "text": "I can't help."}, + _anthropic_tool_use()], "refusal"), + ]) + result = client.messages([{"role": "user", "content": "q"}], + model="claude-test", max_tokens=100) + assert result["stop_reason"] == "refusal" + message, = result["messages"] + assert message["role"] == "assistant" + assert [block["type"] for block in message["content"]] == ["text"] + assert message["content"][0]["text"] == "I can't help." + # The envelope's own content still carries the full turn verbatim. + assert [block["type"] for block in result["content"]] \ + == ["text", "tool_use"] + + +@needs_anthropic +def test_messages_default_cap(client, store_path, fake_anthropic): + seed_doc(store_path, "pi-a", "report.pdf") + calls = fake_anthropic([ + _anthropic_message([_anthropic_tool_use(f"tu_{index}")], "tool_use") + for index in range(30) + ]) + result = client.messages([{"role": "user", "content": "q"}], + model="claude-test", max_tokens=100) + assert len(calls) == 10 # bounded like the OpenAI surfaces + assert result["stop_reason"] == "tool_use" + json.dumps(result) + + +@needs_anthropic +def test_messages_edge_validation(client, store_path, fake_anthropic): + calls = fake_anthropic([ + _anthropic_message([{"type": "text", "text": "ok"}], "end_turn"), + ]) + client.messages([{"role": "user", "content": "q"}], model="claude-test", + max_tokens=100, system=" ") + assert all(block["text"].strip() for block in calls[0]["system"]) + with pytest.raises(PageIndexAPIError, match="message dicts"): + client.messages(["not a dict"], model="claude-test", max_tokens=100) + with pytest.raises(PageIndexAPIError, match="doc_id"): + client.messages([{"role": "user", "content": "q"}], + model="claude-test", max_tokens=100, doc_id=123) diff --git a/tests/test_package_surface.py b/tests/test_package_surface.py index e10c3e8c5..6d985c818 100644 --- a/tests/test_package_surface.py +++ b/tests/test_package_surface.py @@ -60,3 +60,44 @@ def test_import_pageindex_is_lazy(): out = subprocess.run([sys.executable, "-c", probe], capture_output=True, text=True, check=True) assert out.stdout.split() == ["clean", "function"] + + +def test_sdk_submodules_reachable_and_dunder_probes_stay_lazy(): + """The 0.2.10 modules resolve as attributes, and underscore probes (the + frequent unknown names: copy/pickle/inspect dunders) raise without + dragging in the indexing stack. A non-underscore unknown name still + raises AttributeError — after the compat fallthrough's one classic + import, which is the pre-0.2.10 behavior.""" + probe = ( + "import sys, pageindex\n" + "pageindex.agent_tools; pageindex.local_chat\n" + "pageindex.mcp_bridge; pageindex.integrations\n" + "assert not hasattr(pageindex, '__wrapped__')\n" + "heavy = [m for m in ('pageindex.page_index_classic', " + "'pageindex.flash', 'pageindex.utils') if m in sys.modules]\n" + "print(','.join(heavy) or 'clean')\n" + "try:\n" + " pageindex.definitely_missing\n" + " raise SystemExit('no AttributeError')\n" + "except AttributeError:\n" + " pass\n" + ) + out = subprocess.run([sys.executable, "-c", probe], + capture_output=True, text=True, check=True) + assert out.stdout.strip() == "clean" + + +def test_classic_compat_surface_still_reachable(): + """The pre-0.2.10 catch-all made every classic/utils public name a + package attribute; dropping it broke `from pageindex import + ConfigLoader` on upgrade with no deprecation path.""" + probe = ( + "import pageindex\n" + "assert callable(pageindex.count_tokens)\n" + "assert isinstance(pageindex.ConfigLoader, type)\n" + "from pageindex import check_toc # noqa: F401\n" + "print('ok')\n" + ) + out = subprocess.run([sys.executable, "-c", probe], + capture_output=True, text=True, check=True) + assert out.stdout.strip() == "ok"