diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml
index e96dc4cca..6d2cbd8a5 100644
--- a/.github/workflows/publish.yml
+++ b/.github/workflows/publish.yml
@@ -5,7 +5,8 @@ name: Publish to PyPI
# 2. git push origin v0.2.9
# 3. This workflow derives the version from the tag, injects it into
# pyproject.toml, builds, publishes to PyPI via OIDC trusted publishing
-# (no stored secret), and creates a GitHub Release with generated notes.
+# (no stored secret), and creates a GitHub Release with generated notes
+# (dev tags publish to PyPI only — no GitHub Release).
#
# The tag must be a PEP 440 version with a leading `v`:
# v0.2.9 v0.2.9rc1 v0.2.9.dev1
@@ -76,6 +77,9 @@ jobs:
uses: pypa/gh-action-pypi-publish@cef221092ed1bacb1cc03d23a2d87d1d172e277b # release/v1.14.0
- name: Create GitHub Release
+ # dev builds need an explicit pin to install; a Release for them
+ # only doubles the feed next to the stable that follows.
+ if: ${{ !contains(github.ref_name, 'dev') }}
uses: softprops/action-gh-release@718ea10b132b3b2eba29c1007bb80653f286566b # v3.0.1
with:
tag_name: ${{ github.ref_name }}
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index d8d9dbb38..b37db059a 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -30,6 +30,9 @@ jobs:
python-version: ${{ matrix.python-version }}
cache: pip
- run: pip install -r requirements.txt pytest
+ - if: matrix.agent-frameworks == 'without'
+ # requirements.txt carries it; this leg tests the no-framework paths
+ run: pip uninstall -y openai-agents
- 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/README.md b/README.md
index 5ce0ca5e6..4a086ed6f 100644
--- a/README.md
+++ b/README.md
@@ -173,10 +173,11 @@ 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 from --toc-check-pages down require --mode standard):
```
---model LLM model to use (default: gpt-4o-2024-11-20)
+--mode Processing mode: flash (default) or standard
+--index-model LLM model used to index the document (default: gpt-5.6-luna)
--toc-check-pages Pages to check for table of contents (default: 20)
--max-pages-per-node Max pages per node (default: 10)
--max-tokens-per-node Max tokens per node (default: 20000)
@@ -199,23 +200,19 @@ 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
For a simple, end-to-end **agentic vectorless RAG** example using **self-hosted PageIndex** (with OpenAI Agents SDK), see [`examples/agentic_vectorless_rag_demo.py`](examples/agentic_vectorless_rag_demo.py).
```bash
-# Install optional dependency
-pip3 install openai-agents
-
-# Run the demo
python3 examples/agentic_vectorless_rag_demo.py
```
diff --git a/examples/agentic_vectorless_rag_demo.py b/examples/agentic_vectorless_rag_demo.py
index 4fe5f179f..93a00735c 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_API_KEY in the environment.
"""
import sys
-import json
import asyncio
import concurrent.futures
from pathlib import Path
@@ -27,12 +28,11 @@
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 PageIndexLocalClient
import pageindex.utils as utils
PDF_URL = "https://arxiv.org/pdf/2603.15031"
@@ -41,48 +41,16 @@
PDF_PATH = _EXAMPLES_DIR / "documents" / "attention-residuals.pdf"
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,7 +120,7 @@ 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)
@@ -160,13 +128,11 @@ async def _run():
print("=" * 60)
doc_id = next(
(doc["id"] for doc in client.list_documents(limit=100)["documents"]
- if doc["name"] == PDF_PATH.name),
- None,
- )
+ if doc["name"] == PDF_PATH.name), None)
if 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"]
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..669f13629 100644
--- a/pageindex/__init__.py
+++ b/pageindex/__init__.py
@@ -22,9 +22,10 @@
"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):
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..000ce0e6a
--- /dev/null
+++ b/pageindex/agent_tools.py
@@ -0,0 +1,1689 @@
+"""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: every outcome, including errors, is returned as the
+same JSON envelope the cloud emits ({"success": true, ...} /
+{"error": ...}) — arguments outside a pruned local signature come back as
+that envelope too, on the direct and the call_tool path alike.
+"""
+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)
+_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, stop_ids=None) -> list[dict[str, Any]]:
+ """Every document the client can list, newest first (both modes list
+ newest-first; paging preserves that order). With ``stop_ids``, paging
+ stops early once every one of those ids has been seen — for callers
+ that only need those entries; an id absent from the listing still
+ costs a full sweep."""
+ documents: list[dict[str, Any]] = []
+ offset = 0
+ remaining = {str(one_id) for one_id in stop_ids} if stop_ids else None
+ while True:
+ page = client.list_documents(limit=100, offset=offset)
+ batch = page.get("documents") or []
+ documents.extend(batch)
+ if remaining is not None:
+ remaining.difference_update(str(doc.get("id")) for doc in batch)
+ if not remaining:
+ return documents
+ # 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, stop_ids=allowed_ids)
+ 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(str(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:
+ continue # transient refetch failure: poll on to the deadline
+ 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 ──
+
+class _PageSpecError(ValueError):
+ """Shared page-spec rejection; ``code`` picks the caller's rendering."""
+
+ def __init__(self, code: str, message: str):
+ super().__init__(message)
+ self.code = code
+
+
+def _expand_pages(pages) -> list[int]:
+ """Expand '1-3,7' into sorted distinct pages — the one parser for the
+ SDK surface and the tool layer. Raises _PageSpecError with code
+ 'invalid', 'too_many', or 'nonpositive'."""
+ if not isinstance(pages, str):
+ raise _PageSpecError("invalid",
+ f"Invalid page specification: {pages!r}")
+ too_many = (f"Page specification '{pages}' spans more than "
+ f"{_MAX_REQUESTED_PAGES} pages; request a narrower range")
+ expanded: set[int] = set()
+ for part in pages.split(","):
+ part = part.strip()
+ try:
+ if "-" in part:
+ start, end = (int(x) for x in part.split("-", 1))
+ if start > end:
+ raise _PageSpecError(
+ "invalid",
+ f"Invalid range '{part}': start must be <= end")
+ else:
+ start = end = int(part)
+ except _PageSpecError:
+ raise
+ except ValueError as exc:
+ raise _PageSpecError(
+ "invalid", f"Invalid page specification '{pages}'") from exc
+ # 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:
+ raise _PageSpecError("too_many", too_many)
+ expanded.update(range(start, end + 1))
+ if len(expanded) > _MAX_REQUESTED_PAGES:
+ raise _PageSpecError("too_many", too_many)
+ if min(expanded) < 1:
+ raise _PageSpecError(
+ "nonpositive",
+ "Invalid page numbers. Page numbers must be positive integers")
+ return sorted(expanded)
+
+
+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."""
+ try:
+ return _expand_pages(pages), None
+ except _PageSpecError as exc:
+ if exc.code == "too_many":
+ return None, _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",
+ )
+ if exc.code == "nonpositive":
+ 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 None, _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",
+ )
+
+
+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, stop_ids=_allowed_ids),
+ _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:
+ # _format_structure strips text anyway — don't download it.
+ tree = client.get_tree(entry["id"], node_summary=True,
+ include_text=False).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")
+ entry = {"page": page, "text": text}
+ size = _serialized_size(entry) + 2 # +2: json ", " item separator
+ if not included or budget - size >= 0:
+ content.append(entry)
+ included.append(page)
+ budget -= size
+ 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, stop_ids=_allowed_ids)
+ 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(schema: dict, kwargs: dict[str, Any]) -> None:
+ """Models routinely send booleans as JSON strings ("false"); the bare
+ truthiness tests downstream would read those as True. Runs on both
+ dispatch paths — call_tool and the cloud bridge invoker."""
+ properties = (schema or {}).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(TOOL_CONTRACT.get(name, {}).get("schema", {}), 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."
+ ),
+ # Drop the sentence naming the cloud-only image tool, whatever its
+ # wording; the contract-refresh test pins that something was removed.
+ "get_page_content": re.sub(
+ r"\s*[^.]*`get_document_image\(\)`[^.]*\.", "",
+ TOOL_CONTRACT["get_page_content"]["description"]),
+}
+
+_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
+
+
+
+
+_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, schema: dict,
+ ) -> "Callable[[dict], tuple[str, bool]]":
+ """One cloud tool call proxied over MCP: string booleans are coerced
+ (same as call_tool), None-valued arguments are dropped (None ≡ omitted,
+ matching the contract's "omit if ..." semantics) and failures are
+ contained in the error envelope — except 401/403, which re-raise.
+ 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}
+ _coerce_bool_args(schema, arguments)
+ return bridge.call_tool(name, arguments)
+ except Exception as exc:
+ if (isinstance(exc, PageIndexAPIError)
+ and exc.status_code in (401, 403)):
+ raise
+ 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_tool_function(name: str, description: str, schema: dict,
+ invoke: "Callable[[dict], tuple[str, bool]]",
+ ) -> Callable[..., str]:
+ """One plain function for a tool: real signature and docstring from the
+ schema, errors contained by the invoker; arguments the signature
+ rejects come back as the guided envelope instead of raising."""
+ import keyword
+
+ properties: dict[str, Any] = schema.get("properties") or {}
+ required = set(schema.get("required") or [])
+ _invoke = invoke
+
+ params_usable = all(param.isidentifier() and not keyword.iskeyword(param)
+ and param != "_invoke"
+ for param in properties)
+ if not params_usable:
+ def inner(**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)
+ inner = namespace["_synthesized"]
+ # binding TypeErrors quote __qualname__, not __name__
+ inner.__name__ = inner.__qualname__ = name or "tool"
+ 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
+ inner.__annotations__ = annotations
+
+ def proxy(*args: Any, **kwargs: Any) -> str:
+ # The invoker lets only 401/403 auth failures through, so a
+ # TypeError here is the binding rejecting the arguments.
+ try:
+ return inner(*args, **kwargs)
+ except TypeError as exc:
+ payload, _ = _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)
+ proxy.__signature__ = inspect.signature(inner) # type: ignore[attr-defined]
+ proxy.__annotations__ = dict(inner.__annotations__)
+ proxy.__name__ = proxy.__qualname__ = name or "tool"
+ proxy.__doc__ = _tool_docstring(description or "", properties)
+ return proxy
+
+
+_BRIDGES: "weakref.WeakKeyDictionary" = weakref.WeakKeyDictionary()
+_BRIDGES_LOCK = threading.Lock()
+
+
+def _cloud_bridge(client, gated: bool = False):
+ """One bridge per client and endpoint gate (``gated`` = the read-only
+ ?tools=read endpoint); tool discovery and instructions share a session
+ per gate. Weak-keyed off the instance so clients stay picklable; the
+ lock closes the check-then-set race under concurrent first calls."""
+ with _BRIDGES_LOCK:
+ # A rotated api_key or moved BASE_URL rebuilds the bridges.
+ auth = (client.BASE_URL, client.api_key)
+ bridges, seen = _BRIDGES.get(client) or ({}, None)
+ if seen != auth:
+ bridges = {}
+ bridge = bridges.get(gated)
+ if bridge is None:
+ from .mcp_bridge import McpBridge
+ bridge = McpBridge(
+ f"{auth[0]}/mcp" + ("?tools=read" if gated else ""),
+ {"Authorization": f"Bearer {auth[1]}"},
+ )
+ bridges[gated] = bridge
+ _BRIDGES[client] = (bridges, auth)
+ 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 _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.
+ An empty allowlist is refused too: it would scope the agent to
+ nothing, with no signal to the caller."""
+ if doc_ids is not None and not doc_ids:
+ raise PageIndexAPIError(
+ "doc_id is empty. Pass one or more document IDs, or omit "
+ "doc_id to give the agent the whole library.")
+ 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, gated=not include_management)
+ 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"),
+ meta.get("inputSchema") or {}))
+ 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,
+ doc_ids=None) -> 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).
+ ``doc_ids`` is the local allowlist, as in ``_tool_specs``.
+ """
+ return [_make_tool_function(name, description, schema, invoke)
+ for name, description, schema, invoke
+ in _tool_specs(client, include_management, doc_ids)]
+
+
+# ── 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, include_management: bool = False) -> str:
+ """Cloud: the live instructions the MCP server serves for the tool set
+ actually shipped. Local: the built-in subset instructions."""
+ if not getattr(client, "api_key", None):
+ return AGENT_INSTRUCTIONS
+ instructions = _cloud_bridge(
+ client, gated=not include_management).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:
+ # An empty selection must fail loud: washing it to None would mean
+ # "everything", and the tool-layer allowlist would mean "nothing".
+ raise PageIndexAPIError(
+ "doc_id is empty. Pass one or more document IDs, or omit "
+ "doc_id to give the agent the whole library.")
+ details = []
+ missing = []
+ for one_id in doc_ids:
+ try:
+ details.append(client.get_document(one_id))
+ except PageIndexAPIError as exc:
+ # Batch only a definite not-found/denied (local raises carry no
+ # status); a cloud transport failure (429/5xx) propagates raw.
+ if exc.status_code not in (None, 403, 404):
+ raise
+ missing.append(str(one_id))
+ if missing:
+ raise PageIndexAPIError(
+ "Documents not found or access denied: " + ", ".join(missing))
+ # Scoped: the listing only backfills the target docs' metadata (list
+ # entries carry it, get_document does not — cloud parity), so paging
+ # can stop at those ids. Unscoped needs it all for the shadow check.
+ listing = _all_documents(client, stop_ids=doc_ids if scoped else None)
+ 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,
+ include_management: 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, include_management)
+ 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..c4d71eb15 100644
--- a/pageindex/client.py
+++ b/pageindex/client.py
@@ -1,26 +1,49 @@
"""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 threading
+import time
+import warnings
+from typing import Any, Callable, Iterator, Optional, Union, cast
from .errors import PageIndexAPIError
+_litellm_preload_started = False
+
+
+def _preload_litellm() -> None:
+ """Start litellm's multi-second import in the background, once per
+ process — a per-client thread would churn under per-request clients."""
+ # LiteLLM's import otherwise fetches its model map over the network —
+ # seconds of blocking (or a hang offline). Stamped here, not at package
+ # import, so merely importing pageindex leaves the host process's own
+ # litellm untouched; setdefault, so an explicit user choice wins.
+ os.environ.setdefault("LITELLM_LOCAL_MODEL_COST_MAP", "True")
+ global _litellm_preload_started
+ if _litellm_preload_started:
+ return
+ _litellm_preload_started = True
+
+ def _import() -> None:
+ try:
+ import litellm # noqa: F401
+ except Exception:
+ pass
+
+ threading.Thread(target=_import, daemon=True).start()
+
+
def _parse_pages(pages: str) -> list[int]:
- result = []
- 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))
+ from .agent_tools import _PageSpecError, _expand_pages
+ try:
+ return _expand_pages(pages)
+ except _PageSpecError as exc:
+ raise PageIndexAPIError(str(exc)) from exc
-def _normalize_retrieve_model(model: str) -> str:
+def _agents_sdk_model_name(model: str) -> str:
"""Preserve supported Agents SDK prefixes and route other provider paths via LiteLLM."""
passthrough_prefixes = ("litellm/", "openai/")
if not model or "/" not in model:
@@ -43,16 +66,38 @@ class PageIndexClient:
Args:
api_key (str, optional): PageIndex cloud API key
(https://dash.pageindex.ai/api-keys). Omit for local mode.
- model (str, optional): Local mode only — LLM used to build document
- 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.
+ index_model (str, optional): Local mode only — LLM used to index
+ documents (structure and summaries). Defaults to the SDK
+ default (fast and cheap).
+ chat_model (str, optional): Local mode only — the model the chat
+ surfaces (``chat``, ``chat_completions``, ``responses``)
+ default to, exposed as ``client.chat_model``. Chat names
+ route through LiteLLM and mean what LiteLLM says they mean;
+ bare names are OpenAI-compatible shorthand, and
+ ``openai/Qwen/...`` is the form for an OpenAI-compatible
+ server that itself serves slashed model ids (vLLM, TGI).
+ Defaults to the SDK default (strong).
+ model (str, optional): Local mode only — one model for both roles:
+ sets the default for ``index_model`` and ``chat_model`` at
+ once. The role-specific arguments win over it. (Also the
+ 0.2.8-era name for the indexing model — old configs keep
+ working unchanged.)
+ summary_model (str, optional): Local mode only — legacy: overrides
+ the model used for node summaries and document descriptions;
+ ``index_model`` covers this.
+ retrieve_model (str, optional): Local mode only — legacy name for
+ ``chat_model``.
storage_path (str, optional): Local mode only — directory where
indexed documents are stored. Defaults to ``./.pageindex``.
+ index_backend (dict, optional): Local mode only — connection
+ overrides for the indexing lane's LLM calls. Keys are
+ LiteLLM's own connection params — ``api_key``, ``api_base``,
+ ``api_version``, ``aws_*``, … — passed through verbatim.
+ chat_backend (dict, optional): Local mode only — default
+ connection overrides for the chat surfaces; a call's own
+ ``backend`` keys win over it. The dict reaches whichever
+ door runs, in that door's vocabulary (see each method) —
+ ``api_key`` / ``base_url`` mean the same thing on all three.
Usage:
client = PageIndexClient(api_key="...") # cloud
@@ -62,10 +107,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"
@@ -74,19 +118,27 @@ def __init__(
self,
api_key: Optional[str] = None,
*,
+ index_model: Optional[str] = None,
+ chat_model: Optional[str] = None,
model: Optional[str] = None,
summary_model: Optional[str] = None,
retrieve_model: Optional[str] = None,
storage_path: Optional[str] = None,
+ index_backend: Optional[dict[str, Any]] = None,
+ chat_backend: Optional[dict[str, Any]] = None,
):
if api_key == "":
raise PageIndexAPIError(
"api_key is an empty string. Pass a real PageIndex API key for "
"cloud mode, or omit api_key entirely for local mode."
)
+ model_args = {"index_model": index_model, "chat_model": chat_model,
+ "model": model, "summary_model": summary_model,
+ "retrieve_model": retrieve_model}
if api_key is not None:
- local_only = {"model": model, "summary_model": summary_model,
- "retrieve_model": retrieve_model, "storage_path": storage_path}
+ local_only = dict(model_args, storage_path=storage_path,
+ index_backend=index_backend,
+ chat_backend=chat_backend)
passed = [name for name, value in local_only.items() if value is not None]
if passed:
raise PageIndexAPIError(
@@ -99,23 +151,35 @@ def __init__(
self._api = CloudAPI(self)
else:
from .utils import ConfigLoader
- overrides = {key: value for key, value in
- {"model": model, "summary_model": summary_model,
- "retrieve_model": retrieve_model}.items()
+ overrides = {key: value for key, value in model_args.items()
if value}
opt = ConfigLoader().load(overrides or None)
self.model = opt.model
- self.summary_model = getattr(opt, "summary_model", None) or opt.model
- self.retrieve_model = _normalize_retrieve_model(
- getattr(opt, "retrieve_model", None) or opt.model)
+ self.index_model = opt.index_model
+ self.summary_model = opt.summary_model
+ self.chat_model = opt.chat_model
+ self.chat_backend = chat_backend
self.storage_path = storage_path or ".pageindex"
from .local_api import LocalAPI
self._api = LocalAPI(
storage_path=self.storage_path,
model=self.model,
summary_model=self.summary_model,
- retrieve_model=self.retrieve_model,
+ index_backend=index_backend,
)
+ # LiteLLM's multi-second import would otherwise land on the
+ # first chat call; failures resurface there with real context.
+ _preload_litellm()
+
+ @property
+ def retrieve_model(self):
+ """Legacy name for ``chat_model``."""
+ return self.chat_model
+
+ @retrieve_model.setter
+ def retrieve_model(self, value):
+ # 0.2.9 allowed assignment; keep the write path working too.
+ self.chat_model = value
# ---------- DOCUMENT SUBMISSION ----------
@@ -126,39 +190,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 +375,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 +388,362 @@ 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 ----------
+ # ---------- CHAT ----------
+
+ def chat(
+ self,
+ messages: Union[str, list[dict[str, str]]],
+ doc_id: Optional[Union[str, list[str]]] = None,
+ stream: bool = False,
+ model: Optional[str] = None,
+ reasoning_effort: Optional[str] = None,
+ ) -> Union[str, Iterator[str]]:
+ """
+ Ask a question about your documents, get the answer.
+
+ Thin sugar over ``chat_completions()`` in both modes — same
+ engine, same wire, minus the envelope. Multi-turn: keep your own
+ role/content list of the visible conversation (append each answer
+ as an assistant message) and pass it back. For usage accounting,
+ streaming metadata, or the tool-use process, use the protocol
+ surfaces: ``chat_completions()``, ``responses()``, ``messages()``.
+
+ Args:
+ messages: A question string, or role/content conversation
+ history.
+ doc_id: Document ID or list of IDs to scope the conversation.
+ Keep it identical across a conversation's calls.
+ stream: Yield the answer as text chunks as it is produced.
+ model: Local only — backend model name (defaults to
+ ``chat_model``).
+ reasoning_effort: Local only — how hard the model thinks
+ (``"low"`` / ``"medium"`` / ``"high"``; what a backend
+ accepts is its own). Unset sends nothing — the model's
+ default behavior applies.
+
+ Returns:
+ - stream=False: the answer string
+ - stream=True: iterator of text chunks
+ """
+ result = self.chat_completions(messages, stream=stream,
+ doc_id=doc_id, model=model,
+ reasoning_effort=reasoning_effort)
+ if stream:
+ return cast(Iterator[str], result)
+ envelope = cast(dict[str, Any], result)
+ try:
+ return envelope["choices"][0]["message"]["content"] or ""
+ except (KeyError, IndexError, TypeError) as exc:
+ raise PageIndexAPIError(
+ "The chat response carries no answer: "
+ f"{str(envelope)[:200]}") from exc
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,
+ top_p: Optional[float] = None,
+ max_tokens: Optional[int] = None,
+ reasoning_effort: Optional[str] = None,
+ extra_body: Optional[dict[str, Any]] = None,
+ extra_headers: Optional[dict[str, str]] = None,
+ backend: Optional[dict[str, Any]] = 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, routed
+ through LiteLLM — model names mean what LiteLLM says they mean.
+ Bare names are OpenAI-compatible shorthand (the OpenAI SDK's usual
+ env config — OPENAI_API_KEY, OPENAI_BASE_URL — selects the
+ backend, so any OpenAI-compatible server works; write
+ ``openai/Qwen/...`` when the server itself serves slashed ids),
+ provider-prefixed names — ``anthropic/…``, ``bedrock/…`` — reach
+ that provider, and LiteLLM-routed Claude models get the managed
+ prompt prefix cache-marked automatically. 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. Local takes text
+ history only: tool-role turns are rejected (the cloud
+ endpoint forwards them verbatim), and message fields beyond
+ role/content are dropped.
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
+ ``chat_model``). The cloud endpoint selects its own.
+ max_turns: Local only — cap on agent turns per call.
+ top_p: Local only — nucleus sampling, passed through to the
+ model.
+ max_tokens: Local only — per-call output cap, passed through;
+ it bounds each backend call in the agent loop (the way
+ max_turns bounds the loop), not the whole run.
+ reasoning_effort: Local only — passed through verbatim as
+ LiteLLM's ``reasoning_effort``; each provider maps it to
+ its own thinking control, and the values mean what the
+ backend says they mean. Unset sends nothing (the
+ backend's default applies).
+ extra_body: Local only — extra request fields beyond this
+ method's parameters, merged last so they win.
+ OpenAI-compatible backends take them verbatim in the
+ request body; LiteLLM-routed providers take them as
+ LiteLLM's own params (mapped or refused per provider).
+ Credentials belong in ``backend``, never here.
+ extra_headers: Local only — extra HTTP headers merged into
+ each backend request; caller headers win. One exception:
+ LiteLLM's anthropic adapter owns the ``anthropic-beta``
+ header (your value is dropped there) — use ``messages()``
+ for Anthropic beta flags.
+ backend: Local only — connection overrides for this call's
+ backend, merged over the client's ``chat_backend``
+ (per-call keys win). Keys are LiteLLM's own connection
+ params — ``api_key``, ``base_url``, ``api_version``,
+ ``aws_*``, … — passed through verbatim.
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, top_p=top_p, max_tokens=max_tokens,
+ reasoning_effort=reasoning_effort, extra_body=extra_body,
+ extra_headers=extra_headers, backend=backend,
+ )
+ if (model is not None or max_turns is not None or top_p is not None
+ or max_tokens is not None or reasoning_effort is not None
+ or extra_body is not None or extra_headers is not None
+ or backend is not None):
+ raise PageIndexAPIError(
+ "model, max_turns, top_p, max_tokens, reasoning_effort, "
+ "extra_body, extra_headers and backend 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,
+ max_output_tokens: Optional[int] = None,
+ reasoning: Optional[dict[str, Any]] = None,
+ extra_body: Optional[dict[str, Any]] = None,
+ extra_headers: Optional[dict[str, str]] = None,
+ backend: Optional[dict[str, Any]] = 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 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 ``chat_model``).
+ stream: Yield Responses stream events as dicts — one logical
+ response per call: per-turn backend lifecycle events are
+ collapsed to one opening ``response.created`` and one
+ final terminal event, sequence numbers are reassigned
+ monotonically, and ``output_index`` is re-based onto the
+ single logical ``output``. The 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.
+ max_output_tokens: Per-call output cap, passed through; it
+ bounds each backend call in the agent loop (the way
+ max_turns bounds the loop), not the whole run. Echoed in
+ the envelope.
+ reasoning: Responses reasoning options, forwarded verbatim
+ (e.g. ``{"effort": "low", "summary": "auto"}``) — the
+ values mean what the backend says they mean. Unset sends
+ nothing (the backend's default applies).
+ extra_body: Extra request fields beyond this method's
+ parameters, merged verbatim into the request body (last,
+ so they win). Credentials belong in ``backend``, never
+ here.
+ extra_headers: Extra HTTP headers merged into each request;
+ caller headers win over defaults.
+ backend: Connection overrides for this call's backend client,
+ merged over the client's ``chat_backend`` (per-call keys
+ win). Keys are the openai SDK's client params —
+ ``api_key``, ``base_url``, ``organization``, … — passed
+ verbatim; unknown keys raise.
+ """
+ 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, max_output_tokens=max_output_tokens,
+ reasoning=reasoning, extra_body=extra_body,
+ extra_headers=extra_headers, backend=backend,
+ )
+
+ 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,
+ thinking: Optional[dict[str, Any]] = None,
+ extra_body: Optional[dict[str, Any]] = None,
+ extra_headers: Optional[dict[str, str]] = None,
+ backend: Optional[dict[str, Any]] = 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, and the request sets the top-level
+ ``cache_control`` so each turn re-reads the growing conversation
+ from cache — skipped when your own blocks already use the three
+ remaining breakpoints (the managed prompt holds the fourth).
+
+ 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, and rises to
+ budget_tokens + 8192 when ``thinking`` is enabled (the wire
+ requires max_tokens above the budget). 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.
+ thinking: Anthropic thinking configuration, forwarded verbatim
+ (e.g. ``{"type": "adaptive"}``) — the values and their
+ constraints are the backend's. Unset sends nothing.
+ extra_body: Extra request fields beyond this method's
+ parameters, merged verbatim into each request body.
+ Credentials belong in ``backend``, never here.
+ extra_headers: Extra HTTP headers merged into each request
+ (e.g. ``anthropic-beta`` feature flags); caller headers
+ win over defaults.
+ backend: Connection overrides for this call's backend client,
+ merged over the client's ``chat_backend`` (per-call keys
+ win). Keys are the anthropic SDK's client params —
+ ``api_key``, ``base_url``, ``auth_token``, … — passed
+ verbatim; unknown keys raise.
+ """
+ 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,
+ thinking=thinking, extra_body=extra_body,
+ extra_headers=extra_headers, backend=backend,
+ )
+
# ---------- DOCUMENT MANAGEMENT ----------
def get_document(self, doc_id: str) -> dict[str, Any]:
@@ -365,6 +791,382 @@ 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,
+ doc_id: Optional[Union[str, list[str]]] = None,
+ ) -> 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 live read tool set, discovered 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, ...).
+ 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 .agent_tools import build_agent_tools
+ return build_agent_tools(self, include_management, doc_ids=doc_id)
+
+ 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.
+
+ ``openai-agents`` is 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 doc_id is not None and not doc_id:
+ # An empty scope means "nothing" locally (empty allowlist) and
+ # cannot be represented on cloud; both refuse it loudly.
+ raise PageIndexAPIError(
+ "doc_id is empty. Pass one or more document IDs, or omit "
+ "doc_id to give the agent the whole library.")
+ if not getattr(self, "api_key", None):
+ return doc_id
+ 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 ``chat_model`` (cloud omits ``model`` so the
+ framework default applies). To customize further, switch to
+ those methods directly. You run this config in your own
+ environment, so its model auth comes from there —
+ ``chat_backend`` does not travel with it.
+
+ Prompt caching configures itself for most destinations (OpenAI
+ server-side; Anthropic- and Bedrock-hosted Claude via LiteLLM's
+ defaults). Vertex-hosted Claude is the exception — pass the
+ injection points yourself::
+
+ Agent(**config, model_settings=ModelSettings(extra_args={
+ "cache_control_injection_points": [
+ {"location": "message", "role": "system"},
+ {"location": "message", "index": -1}]}))
+
+ 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. Same
+ grammar as ``chat_model`` (LiteLLM names; bare names are
+ OpenAI-compatible shorthand).
+ """
+ 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,
+ include_management=include_management),
+ "tools": self.as_openai_tools(include_management, doc_id=scope),
+ }
+ model = model or getattr(self, "chat_model", None)
+ if model:
+ config["model"] = _agents_sdk_model_name(model)
+ if config["model"].startswith("litellm/"):
+ # The runner resolves this model through LiteLLM in the
+ # caller's process, outside our completion helpers.
+ from .utils import _repair_litellm_types
+ _repair_litellm_types()
+ 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,
+ thinking: Optional[dict] = 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 ``max_tokens``
+ default and 10-turn ``max_iterations`` bound ``messages()`` uses,
+ and a top-level ``cache_control`` so each loop turn re-reads the
+ growing prompt from cache (pop the key if you place your own
+ breakpoints — the API allows four). Unlike ``messages()``,
+ ``system`` here is the bare instructions string, without the chat
+ header or its block-level breakpoint. 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.
+ thinking: Anthropic ``thinking`` config, included in the
+ kwargs; an enabled budget also lifts the ``max_tokens``
+ default above it. Pass it here, not alongside the
+ unpacked config, so the default stays valid.
+ """
+ from .agent_tools import build_agent_instructions
+ from .local_chat import _default_max_tokens, _validate_max_turns
+ _validate_max_turns(max_turns)
+ 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, thinking)),
+ "system": build_agent_instructions(
+ self, doc_id, scoped=scope is not None,
+ include_management=include_management),
+ "tools": self.as_anthropic_tools(include_management, asynchronous,
+ doc_id=scope),
+ "max_iterations": max_turns if max_turns is not None else 10,
+ **({"thinking": thinking} if thinking is not None else {}),
+ "cache_control": {"type": "ephemeral"},
+ }
+
+ 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,
+ include_management=include_management),
+ "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,
+ include_management: bool = False,
+ ) -> 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 (the
+ ``*_agent_config`` bundles, whose tools carry the doc_id scope,
+ relax this to duplicates within the targeted set).
+
+ ``include_management``: fetch the guidance for the full tool set,
+ matching tools built with ``include_management=True`` (cloud;
+ local guidance is a single set).
+ """
+ from .agent_tools import build_agent_instructions
+ return build_agent_instructions(
+ self, doc_id, include_management=include_management)
+
# ---------- FOLDER MANAGEMENT ----------
def create_folder(
@@ -420,10 +1222,16 @@ class PageIndexLocalClient(PageIndexClient):
def __init__(
self,
*,
+ index_model: Optional[str] = None,
+ chat_model: Optional[str] = None,
model: Optional[str] = None,
summary_model: Optional[str] = None,
retrieve_model: Optional[str] = None,
storage_path: Optional[str] = None,
+ index_backend: Optional[dict[str, Any]] = None,
+ chat_backend: Optional[dict[str, Any]] = None,
):
- super().__init__(None, model=model, summary_model=summary_model,
- retrieve_model=retrieve_model, storage_path=storage_path)
+ super().__init__(None, index_model=index_model, chat_model=chat_model,
+ model=model, summary_model=summary_model,
+ retrieve_model=retrieve_model, storage_path=storage_path,
+ index_backend=index_backend, chat_backend=chat_backend)
diff --git a/pageindex/cloud_api.py b/pageindex/cloud_api.py
index b7597cc9a..f3a7740b6 100644
--- a/pageindex/cloud_api.py
+++ b/pageindex/cloud_api.py
@@ -55,9 +55,11 @@ 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}
+ data: Dict[str, Any] = {'if_retrieval': True}
if mode is not None:
data['mode'] = mode
if beta_headers is not None:
@@ -223,7 +225,7 @@ def chat_completions(
headers=self._headers(),
json=payload,
stream=stream,
- timeout=120 if stream else 300
+ timeout=120 if stream else 600
)
if response.status_code != 200:
@@ -310,7 +312,9 @@ def get_document(self, doc_id: str) -> Dict[str, Any]:
timeout=30
)
if response.status_code != 200:
- raise PageIndexAPIError(f"Failed to get document metadata: {response.text}")
+ raise PageIndexAPIError(
+ f"Failed to get document metadata: {response.text}",
+ status_code=response.status_code)
return response.json()
def delete_document(self, doc_id: str) -> Dict[str, Any]:
@@ -354,7 +358,7 @@ def list_documents(self, limit: int = 50, offset: int = 0, folder_id: Optional[s
if offset < 0:
raise ValueError("offset must be non-negative")
- params = {"limit": limit, "offset": offset}
+ params: Dict[str, Any] = {"limit": limit, "offset": offset}
if folder_id is not None:
params["folder_id"] = folder_id
diff --git a/pageindex/config.yaml b/pageindex/config.yaml
index 73a512c7a..d3786ca12 100644
--- a/pageindex/config.yaml
+++ b/pageindex/config.yaml
@@ -1,9 +1,11 @@
-# Models without a provider prefix use the OpenAI SDK directly.
-# For other providers, use "provider/model" format (e.g. "anthropic/claude-sonnet-4-6").
-model: "gpt-4o-2024-11-20"
-# model: "anthropic/claude-sonnet-4-6"
-summary_model: "gpt-5.6-luna"
-retrieve_model: "gpt-5.4" # defaults to `model` if not set
+# Models — index_model indexes documents (structure and summaries),
+# chat_model answers questions on the chat surfaces; set model to use one
+# for both. Unset keys use the SDK defaults shown below. Legacy keys
+# (model, summary_model, retrieve_model) keep working.
+# Model names are LiteLLM's: bare names are OpenAI models; for other
+# providers, use "provider/model" (e.g. "anthropic/claude-sonnet-4-6").
+# index_model: "gpt-5.6-luna"
+# chat_model: "gpt-5.6-sol"
toc_check_page_num: 20
max_page_num_each_node: 10
max_token_num_each_node: 20000
diff --git a/pageindex/errors.py b/pageindex/errors.py
index e460a956b..608ba6e4f 100644
--- a/pageindex/errors.py
+++ b/pageindex/errors.py
@@ -1,2 +1,7 @@
class PageIndexAPIError(Exception):
- pass
+ """status_code carries the HTTP status when the failure came from a
+ non-200 cloud response; None otherwise (local mode, client-side)."""
+
+ def __init__(self, *args: object, status_code: int | None = None) -> None:
+ super().__init__(*args)
+ self.status_code = status_code
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..72438325f 100644
--- a/pageindex/flash/api.py
+++ b/pageindex/flash/api.py
@@ -96,15 +96,38 @@ 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 | None = None, 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. The tree extraction itself uses no LLM; by default an LLM writes node summaries and expands the tree (``summary=False, optimize=False`` runs fully LLM-free). 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 (fails fast with ``PageIndexAPIError`` when no LLM key is configured), ``"merge"`` for deterministic merge only, ``False`` to disable. ``True`` is accepted as ``"full"`` for backward compatibility; defaults to ``"full"``. optimize_expand: deprecated — use ``optimize``. Honored only when ``optimize`` is not passed (or is the legacy ``True``): ``False`` maps to ``"merge"``, ``True`` to ``"full"``. 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_expand is not None:
+ import warnings
+ warnings.warn(
+ "optimize_expand is deprecated: pass optimize='full', 'merge', "
+ "or False.", DeprecationWarning, stacklevel=2)
+ if optimize is None or optimize is True:
+ # legacy spellings only — an explicit 'full'/'merge' wins
+ optimize = "merge" if optimize_expand is False else "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 == "full":
+ from ..errors import PageIndexAPIError
+ from ..utils import ConfigLoader, _llm_backend, _openai_missing_keys
+ model = (optimize_model or summary_model
+ or ConfigLoader().load().summary_model)
+ if not _llm_backend.get() and _openai_missing_keys(model):
+ raise PageIndexAPIError(
+ "optimize='full' runs LLM expand and no LLM key is "
+ "configured — set OPENAI_API_KEY, or pass optimize='merge' "
+ "or optimize=False for the LLM-free tree.")
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/flash/embedded_toc.py b/pageindex/flash/embedded_toc.py
index 6bd38704d..285d7dcbb 100644
--- a/pageindex/flash/embedded_toc.py
+++ b/pageindex/flash/embedded_toc.py
@@ -123,13 +123,19 @@ def read_bookmarks(doc_handle: Union[str, Path, BytesIO]) -> list[dict]:
doc = pdfium.PdfDocument(handle)
entries = []
for item in doc.get_toc():
- if item.page_index is None:
+ if hasattr(item, "get_dest"):
+ dest = item.get_dest()
+ page_idx = dest.get_index() if dest else None
+ title = item.get_title()
+ else:
+ page_idx, title = item.page_index, item.title
+ if page_idx is None:
continue
- title = (item.title or "").strip()
+ title = (title or "").strip()
if not title:
continue
entries.append(
- {"title": title, "level": item.level + 1, "page": item.page_index + 1}
+ {"title": title, "level": item.level + 1, "page": page_idx + 1}
)
return entries
except (pdfium.PdfiumError, OSError, ValueError, TypeError):
diff --git a/pageindex/flash/parser_pdfium_charlevel/char_extract.py b/pageindex/flash/parser_pdfium_charlevel/char_extract.py
index d55b4620c..56d6288f7 100644
--- a/pageindex/flash/parser_pdfium_charlevel/char_extract.py
+++ b/pageindex/flash/parser_pdfium_charlevel/char_extract.py
@@ -56,10 +56,21 @@ def _extract_raw_chars(page, text_page) -> tuple[list[dict], list[dict]]:
name_cache: dict[bytes, str] = {}
raw_chars: list[dict] = []
last_obj: dict | None = None
+ skip_next = False
for index_value in range(count_item):
+ if skip_next:
+ skip_next = False
+ continue
codepoint = get_unicode(text_page, index_value)
if codepoint < 0:
continue
+ # PDFium returns astral characters (U+10000+) as two UTF-16 surrogate
+ # code units in consecutive textpage slots. Reassemble before chr().
+ if 0xD800 <= codepoint <= 0xDBFF and index_value + 1 < count_item:
+ low = get_unicode(text_page, index_value + 1)
+ if 0xDC00 <= low <= 0xDFFF:
+ codepoint = ((codepoint & 0x3FF) << 10) + (low & 0x3FF) + 0x10000
+ skip_next = True
# u == 0 (PDFium found no unicode for the glyph) is KEPT as '\x00':
# text extraction emits the raw charcode for unmapped codes, so its items
# really contain chr(0) for extension-font pieces at code 0, and the
diff --git a/pageindex/flash/parser_pdfium_charlevel/geometry.py b/pageindex/flash/parser_pdfium_charlevel/geometry.py
index e3a4b072d..436d7dde4 100644
--- a/pageindex/flash/parser_pdfium_charlevel/geometry.py
+++ b/pageindex/flash/parser_pdfium_charlevel/geometry.py
@@ -6,6 +6,8 @@
import math
import pypdfium2.raw as pdfium_c
+_get_font_name = getattr(pdfium_c, "FPDFFont_GetBaseFontName", None) or pdfium_c.FPDFFont_GetFontName
+
def _obj_rotation(value: float, other_item: float, candidate_item: float, reference_item: float) -> int:
"""Classify a text-object matrix as upright, cardinal rotation, or oblique. Near-cardinal matrices snap to the cardinal bucket; genuinely oblique matrices use the baseline remerge path."""
@@ -128,7 +130,7 @@ def iter_text_objs(parent, anc_mtx, depth):
# Snap back to the shortest decimal so knife-edge font-size comparisons
# match the content-stream value.
fs_eff = float(f"{fs_eff:.6g}")
- name = pdfium_c.FPDFFont_GetFontName(font, font_name_buffer, 256)
+ name = _get_font_name(font, font_name_buffer, 256)
font_name = (
bytes(font_name_buffer[:name]).decode("latin-1", errors="replace").rstrip("\x00")
if name > 1 else ""
diff --git a/pageindex/flash/parser_pdfium_charlevel/pipeline.py b/pageindex/flash/parser_pdfium_charlevel/pipeline.py
index 5a4caf9f0..460ac703f 100644
--- a/pageindex/flash/parser_pdfium_charlevel/pipeline.py
+++ b/pageindex/flash/parser_pdfium_charlevel/pipeline.py
@@ -95,7 +95,7 @@ def _page_pass1(pdf, pdf_doc, page_idx: int, type3_ext: dict, font_map_cache: di
# PDFium's output).
if raw_chars:
_apply_font_unicode(
- text_page.raw, raw_chars, objects, show_codes, pdf_doc,
+ raw_chars, objects, show_codes, pdf_doc,
font_map_cache)
except Exception:
pass
diff --git a/pageindex/flash/parser_pdfium_charlevel/unicode_apply.py b/pageindex/flash/parser_pdfium_charlevel/unicode_apply.py
index a7490bd20..0ae314236 100644
--- a/pageindex/flash/parser_pdfium_charlevel/unicode_apply.py
+++ b/pageindex/flash/parser_pdfium_charlevel/unicode_apply.py
@@ -5,7 +5,6 @@
import bisect
import difflib
from collections import Counter
-import pypdfium2.raw as pdfium_c
from .text_normalize import _is_whitespace
from .font_unicode import _font_unicode_map
@@ -16,7 +15,6 @@
def _apply_font_unicode(
- text_page,
raw_chars: list[dict],
objects: list[dict],
show_codes: list[tuple[int | None, tuple[int, ...], float]],
@@ -287,14 +285,12 @@ def _run_window(window: list[int]) -> None:
_synthesize_dropped_glyphs(kept, raw_chars, chars_by_index)
return
- # Page mode.
- seq: list[tuple[int, str]] = []
- char_count = pdfium_c.FPDFText_CountChars(text_page)
- for char_index in range(char_count):
- if pdfium_c.FPDFText_IsGenerated(text_page, char_index) == 1:
- continue
- codepoint = pdfium_c.FPDFText_GetUnicode(text_page, char_index)
- seq.append((char_index, chr(codepoint) if codepoint > 0 else "\x00"))
+ # Page mode. Walk the char census char_extract built (surrogate pairs
+ # already merged there): re-reading the textpage would split astral
+ # chars back into two lone-surrogate slots, desync the walk against
+ # their one-char cmap targets, and drop the whole page's patch.
+ seq = [(raw_char["i"], raw_char["ch"])
+ for raw_char in raw_chars if not raw_char["is_gen"]]
targets: list[str] = []
for font_index, encoded_text, _tz in show_codes:
if not encoded_text:
diff --git a/pageindex/flash/parser_pdfium_parallel.py b/pageindex/flash/parser_pdfium_parallel.py
index f0d87b287..ec4c8dbc9 100644
--- a/pageindex/flash/parser_pdfium_parallel.py
+++ b/pageindex/flash/parser_pdfium_parallel.py
@@ -23,7 +23,9 @@
import multiprocessing
import os
+import sys
from concurrent.futures import ProcessPoolExecutor
+from contextlib import contextmanager
from io import BytesIO
from pathlib import Path
from typing import Union
@@ -54,6 +56,29 @@ class _Type3Detected(Exception):
_worker_font_maps: dict = {}
+@contextmanager
+def _anonymous_main():
+ """Hide __main__'s import identity while workers spawn: spawn re-executes
+ the caller's script in every worker otherwise, which for an unguarded
+ script means one duplicate full run per worker. Our workers import
+ everything by module name and never need __main__.
+
+ ponytail: window covers the whole map; a concurrent pool spawned from
+ another thread whose tasks live in __main__ would break during it."""
+ main = sys.modules.get("__main__")
+ if main is None:
+ yield
+ return
+ d = main.__dict__
+ saved = {k: d.pop(k) for k in ("__file__", "__spec__") if k in d}
+ d["__spec__"] = None # get_preparation_data reads it via attribute access
+ try:
+ yield
+ finally:
+ d.pop("__spec__", None)
+ d.update(saved)
+
+
def _init_worker(kind: str, payload) -> None:
global _worker_pdf, _worker_pdf_doc, _worker_font_maps
# Open the document exactly as parse_charlevel_meta does, including
@@ -122,12 +147,17 @@ def parse_charlevel_meta_parallel(
initargs=src,
)
try:
- results = list(executor.map(_run_page, range(n_pages)))
+ with _anonymous_main():
+ results = list(executor.map(_run_page, range(n_pages)))
except Exception:
# _Type3Detected or any worker/pool failure. Cancel what is queued
# and rerun sequentially; in-flight pages finish in their workers
# and are discarded (separate processes, no shared PDFium state).
executor.shutdown(wait=False, cancel_futures=True)
+ if getattr(multiprocessing.current_process(), "_inheriting", False):
+ # Spawn child re-importing an unguarded __main__; a sequential rerun
+ # here would silently duplicate the caller's whole run per worker.
+ raise
return parse_charlevel_meta(doc_handle)
executor.shutdown()
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..b3e67947d
--- /dev/null
+++ b/pageindex/integrations/claude_agent_sdk.py
@@ -0,0 +1,65 @@
+"""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, _tool_specs
+
+ def make_handler(invoke):
+ async def handler(arguments: dict[str, Any]) -> dict[str, Any]:
+ text, is_error = await asyncio.to_thread(invoke, arguments or {})
+ 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, description, schema,
+ **tool_kwargs(name))(make_handler(invoke))
+ for name, description, schema, invoke
+ in _tool_specs(client, include_management, doc_ids)
+ ]
+ 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..862a35f0b
--- /dev/null
+++ b/pageindex/integrations/openai_agents.py
@@ -0,0 +1,80 @@
+"""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."
+ ) 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}
+ # is_error has no per-result channel on hand-built FunctionTools
+ # (raising aborts the run — see above); the text is the signal.
+ 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..f3c6afa65 100644
--- a/pageindex/local_api.py
+++ b/pageindex/local_api.py
@@ -1,17 +1,17 @@
"""Local implementation of the PageIndex SDK surface."""
from __future__ import annotations
-import asyncio
import json
import logging
+import multiprocessing
import os
import uuid
-from concurrent.futures import ThreadPoolExecutor
from datetime import datetime, timezone
from typing import Any
from .errors import PageIndexAPIError
from .local_store import DocStore
+from .utils import run_off_loop
logger = logging.getLogger(__name__)
@@ -22,27 +22,30 @@ def _now_iso() -> str:
return now.replace(microsecond=now.microsecond // 1000 * 1000).isoformat()
-def _run_indexer(func, *args, **kwargs):
- try:
- asyncio.get_running_loop()
- except RuntimeError:
- return func(*args, **kwargs)
- with ThreadPoolExecutor(max_workers=1) as pool:
- return pool.submit(func, *args, **kwargs).result()
class LocalAPI:
"""Backs PageIndexClient's local mode. One instance per client."""
def __init__(self, storage_path: str, model: str, summary_model: str,
- retrieve_model: str):
+ index_backend: dict | None = None):
self._store = DocStore(storage_path)
self._model = model
self._summary_model = summary_model
- self._retrieve_model = retrieve_model
+ self._index_backend = index_backend
from .utils import ConfigLoader
self._config_loader = ConfigLoader()
+ def _with_backend(self, func, *args):
+ """Scope the indexing lane's connection overrides around one
+ operation — runs inside whatever thread run_off_loop picked."""
+ from .utils import _llm_backend
+ token = _llm_backend.set(self._index_backend)
+ try:
+ return func(*args)
+ finally:
+ _llm_backend.reset(token)
+
# ── indexing ──
def submit_document(
@@ -53,6 +56,12 @@ def submit_document(
folder_id: str | None = None,
metadata: dict | None = None,
) -> dict[str, Any]:
+ if getattr(multiprocessing.current_process(), "_inheriting", False):
+ raise PageIndexAPIError(
+ "Failed to submit document: called again while a spawned worker "
+ "process was importing your script. Put your top-level code under "
+ "if __name__ == '__main__': so worker processes do not re-run it."
+ )
if beta_headers is not None:
raise PageIndexAPIError(
"Failed to submit document: beta_headers is not supported in local mode."
@@ -75,8 +84,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,15 +108,17 @@ 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":
- structure, description = _run_indexer(
- self._index_flash, file_path, page_texts
+ structure, description = run_off_loop(
+ self._with_backend, self._index_flash, file_path
)
else:
- structure, description = _run_indexer(
- self._index_standard, file_path, page_texts
+ structure, description = run_off_loop(
+ self._with_backend, self._index_standard, file_path,
+ page_texts
)
except PageIndexAPIError:
raise
@@ -113,23 +126,42 @@ def submit_document(
raise PageIndexAPIError(f"Failed to submit document: {e}") from e
doc_id = "pi-" + uuid.uuid4().hex
- meta = {
- "id": doc_id,
- "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",
- }
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}
+ # Check-then-write under the store lock; the early pre-check above
+ # is advisory only.
+ with self._store.lock():
+ meta = {
+ "id": doc_id,
+ "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,
+ }
+ self._store.save_document(
+ doc_id, meta, remove_fields(structure, fields=["text"]), pages)
+ 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]:
@@ -159,12 +191,14 @@ def _index_standard(self, file_path: str, page_texts: list[str]) -> tuple[list,
)
return structure, result.get("doc_description")
- def _index_flash(self, file_path: str, page_texts: list[str]) -> tuple[list, str | None]:
+ def _index_flash(self, file_path: str) -> tuple[list, str | None]:
from .flash import page_index_flash
- from .utils import (add_node_text, create_clean_structure_for_description,
+ from .utils import (create_clean_structure_for_description,
generate_doc_description, write_node_id)
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(
@@ -172,7 +206,6 @@ def _index_flash(self, file_path: str, page_texts: list[str]) -> tuple[list, str
"a structure from this PDF."
)
write_node_id(structure)
- add_node_text(structure, [(text, 0) for text in page_texts])
description = generate_doc_description(
create_clean_structure_for_description(structure),
model=self._summary_model,
@@ -190,6 +223,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")
@@ -311,6 +349,8 @@ def _format_tree_node(node: dict, node_summary: bool) -> dict:
"node_id": node.get("node_id"),
"page_index": node.get("start_index"),
}
+ if node.get("key_items"):
+ out["key_items"] = node["key_items"]
if node_summary:
summary = node.get("summary")
if summary is not None:
diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py
new file mode 100644
index 000000000..e4dcadef3
--- /dev/null
+++ b/pageindex/local_chat.py
@@ -0,0 +1,1081 @@
+"""Managed local chat: document-QA agents over the local tools."""
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import json
+import os
+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.")
+ # 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):
+ from .utils import run_off_loop
+ return run_off_loop(asyncio.run, coro)
+
+
+_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. "
+ "messages() runs on the anthropic extra instead."
+ ) from exc
+
+
+def _sdk_backend(backend) -> dict:
+ """chat_backend for an SDK constructor: LiteLLM takes either endpoint
+ spelling, the openai and anthropic SDKs only ``base_url``."""
+ return {("base_url" if key == "api_base" else key): value
+ for key, value in (backend or {}).items()}
+
+
+def _openai_model(protocol: str, model_name: str, backend=None):
+ """The backend protocol driver — the seam tests replace with a fake."""
+ if protocol == "responses":
+ model_name = model_name.removeprefix("litellm/")
+ if "/" in model_name and not model_name.startswith("openai/"):
+ raise PageIndexAPIError(
+ f"responses() cannot drive '{model_name}': 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."
+ )
+ import openai
+ model_name = model_name.removeprefix("openai/")
+ try:
+ sdk_client = openai.AsyncOpenAI(**_sdk_backend(backend))
+ except (openai.OpenAIError, TypeError) as exc:
+ raise PageIndexAPIError(
+ f"The OpenAI backend is not configured: {exc}") from exc
+ # A caller-owned transport must survive the per-call close.
+ sdk_client._pageindex_caller_http = "http_client" in (backend or {})
+ from agents.models.openai_responses import OpenAIResponsesModel
+ return OpenAIResponsesModel(model_name, openai_client=sdk_client)
+ 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.97'"
+ )
+ from .utils import _litellm_model, _repair_litellm_types
+ _repair_litellm_types()
+ try:
+ wire = _litellm_model(model_name, backend)
+ except litellm.AuthenticationError as exc:
+ raise PageIndexAPIError(
+ "The OpenAI backend is not configured: set the "
+ "OPENAI_API_KEY environment variable, pass an api_key "
+ "in chat_backend / backend (any value works for keyless "
+ "OPENAI_BASE_URL servers), or point chat_model at "
+ "another provider (e.g. 'anthropic/...')."
+ ) from exc
+ except litellm.NotFoundError as exc:
+ raise PageIndexAPIError(str(exc)) from exc
+ return LitellmModel(wire, api_key=(backend or {}).get("api_key"),
+ base_url=(backend or {}).get("base_url"))
+
+
+def _reported_model(model_name: str) -> str:
+ """The name the provider actually serves — routing prefixes stripped."""
+ return model_name.removeprefix("litellm/").removeprefix("openai/")
+
+
+def _cache_extra_args(model_name: str) -> Optional[dict]:
+ """Claude's prompt caching is opt-in per request: on Claude models
+ routed through LiteLLM (Anthropic direct, Bedrock, Vertex — each
+ channel live-verified), mark the managed system prefix and the newest
+ message via LiteLLM's injection param so the loop's later turns and a
+ conversation's next calls read them instead of repaying full price.
+ Provider resolution is LiteLLM's own, so this predicate can never
+ disagree with where the request actually routes."""
+ if "/" not in model_name or model_name.startswith("openai/"):
+ return None
+ try:
+ from litellm import get_llm_provider
+ model, provider, _, _ = get_llm_provider(
+ model=model_name.removeprefix("litellm/"))
+ except Exception:
+ return None
+ if provider == "anthropic" or (provider in ("bedrock", "vertex_ai")
+ and "claude" in model.lower()):
+ # The pair LiteLLM itself seeds for Anthropic and Bedrock: the
+ # stable prefix plus the newest message, so each turn re-reads
+ # the turns before it. Passing it explicitly extends it to Vertex.
+ return {"cache_control_injection_points": [
+ {"location": "message", "role": "system"},
+ {"location": "message", "index": -1}]}
+ return None
+
+
+def _openai_protocol(model_name: str) -> bool:
+ """Destinations that speak the OpenAI protocol on the wire, where
+ prompt_cache_key means something and extra_body lands in the request
+ body. Resolution is LiteLLM's own (same as _cache_extra_args), so the
+ answer follows actual routing; azure and openrouter ride the OpenAI
+ protocol without appearing in openai_compatible_providers."""
+ wire = model_name.removeprefix("litellm/")
+ if "/" not in wire or wire.startswith("openai/"):
+ return True
+ try:
+ import litellm
+ _, provider, _, _ = litellm.get_llm_provider(model=wire)
+ except Exception:
+ return False
+ return (provider in ("openai", "azure", "openrouter")
+ or provider in getattr(litellm, "openai_compatible_providers",
+ ()))
+
+
+def _merged_backend(client, backend):
+ """This call's connection overrides: the client's ``chat_backend``
+ under the per-call dict, per-call keys winning."""
+ merged = {**(getattr(client, "chat_backend", None) or {}),
+ **(backend or {})}
+ return merged or None
+
+
+def _openai_agent(client, protocol: str, model_name: str, instructions: str,
+ temperature, top_p, doc_ids=None, cache_key=None,
+ reasoning=None, reasoning_effort=None, extra_body=None,
+ max_tokens=None, backend=None, extra_headers=None):
+ from agents import Agent, ModelSettings
+ from .integrations.openai_agents import build_openai_tools
+ # ModelSettings.extra_body is the one channel all three engines put on
+ # the wire: LiteLLM drops the bare prompt_cache_key kwarg (wire-verified),
+ # and both OpenAI model classes pass extra_body through verbatim. OpenAI
+ # destinations only — LiteLLM plants extra_body as a literal field in
+ # other providers' request bodies, which Anthropic rejects as unknown.
+ openai_backend = _openai_protocol(model_name)
+ # Chat-lane effort rides extra_args: LiteLLM takes it as its own
+ # top-level kwarg on every supported openai-agents version, and the
+ # channel admits values outside the OpenAI enum ("none").
+ extra_args = _cache_extra_args(model_name)
+ if reasoning_effort is not None:
+ extra_args = {**(extra_args or {}),
+ "reasoning_effort": reasoning_effort}
+ conn = _sdk_backend(backend) if backend else {}
+ if conn and protocol == "chat":
+ # LiteLLM takes connection params per call, except the two names
+ # LitellmModel pins as its own keywords — those ride its constructor.
+ lifted = {"api_key": conn.pop("api_key", None),
+ "base_url": conn.pop("base_url", None)}
+ if conn:
+ extra_args = {**(extra_args or {}), **conn}
+ conn = {key: value for key, value in lifted.items()
+ if value is not None}
+ body = ({"prompt_cache_key": cache_key}
+ if cache_key and openai_backend else None)
+ # Caller extras merge last, so they win over ours; non-OpenAI
+ # destinations take them as LiteLLM kwargs instead (see note above).
+ if extra_body:
+ if openai_backend:
+ body = {**(body or {}), **extra_body}
+ else:
+ extra_args = {**(extra_args or {}), **extra_body}
+ return Agent(
+ name="PageIndex",
+ instructions=instructions,
+ tools=build_openai_tools(client, doc_ids=doc_ids),
+ model=_openai_model(protocol, model_name, conn or None),
+ model_settings=ModelSettings(
+ temperature=temperature, top_p=top_p, max_tokens=max_tokens,
+ reasoning=reasoning,
+ # Streamed runs otherwise carry no usage at all (agents forwards
+ # this as stream_options only on streaming calls).
+ include_usage=True,
+ extra_body=body,
+ extra_headers=extra_headers,
+ extra_args=extra_args),
+ )
+
+
+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_cache_key(model_name: str, instructions: str, doc_id,
+ items) -> str:
+ """Stable per-conversation cache-routing key, sent as the OpenAI
+ ``prompt_cache_key`` through ModelSettings.extra_body (openai-agents
+ 0.20 no longer derives it from RunConfig.group_id — verified against a
+ captured wire). Keyed on the prefix identity — model, instructions,
+ doc targeting, 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.
+ doc_id carries the targeting identity instead — the same opening
+ question against different documents is different conversations."""
+ scope = [doc_id] if isinstance(doc_id, str) else doc_id
+ seed = json.dumps([model_name, instructions, scope,
+ items[0] if items else None],
+ sort_keys=True, default=str)
+ return "pageindex-" + hashlib.sha256(seed.encode()).hexdigest()[:16]
+
+
+def _model_backend_error(exc) -> PageIndexAPIError:
+ """Wrap a provider failure; the sol-class refusal (chatcmpl rejects
+ function tools while reasoning is on) gets its two documented exits
+ appended, since the fix is a different lane, not a retry."""
+ message = f"The model backend failed: {exc}"
+ if "Function tools with reasoning_effort" in str(exc):
+ message += (
+ " — this model runs tools on the Responses lane: upgrade "
+ "litellm (newer releases route it there automatically), pass "
+ "reasoning_effort (older litellm routes explicit efforts), or "
+ "call responses() instead."
+ )
+ return PageIndexAPIError(message)
+
+
+def _run_kwargs(max_turns) -> dict:
+ # No traces — the caller opted into QA, not telemetry.
+ from agents import RunConfig
+ kwargs: dict = {"run_config": RunConfig(tracing_disabled=True)}
+ 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)
+ # The backend's echo of what it actually ran with.
+ for field in ("tool_choice", "parallel_tool_calls"):
+ value = getattr(response, field, None)
+ if value is not None:
+ recorded[field] = (value.model_dump(mode="json")
+ if hasattr(value, "model_dump") else value)
+ return response
+
+ responses.create = recording_create
+
+
+def _record_chat_finish(agent, recorded: dict) -> None:
+ """Capture each turn's native finish_reason from the raw LiteLLM
+ response: openai-agents' ModelResponse drops it, so a truncated or
+ content-filtered final turn would otherwise report as a clean "stop".
+ The chat protocol has no transport client to hook (cf.
+ _record_response_status), so this wraps the model's response fetch;
+ no-op if that private seam moves."""
+ model = getattr(agent, "model", None)
+ fetch = getattr(model, "_fetch_response", None)
+ if fetch is None:
+ return
+
+ def note(item) -> None:
+ choices = getattr(item, "choices", None)
+ finish = getattr(choices[0], "finish_reason", None) if choices else None
+ if finish:
+ recorded["finish_reason"] = finish
+
+ class _Tee:
+ """Iteration passthrough that notes each chunk; everything else
+ (aclose/close/...) delegates to the provider stream itself."""
+
+ def __init__(self, inner):
+ self._inner = inner
+
+ def __aiter__(self):
+ return self
+
+ async def __anext__(self):
+ chunk = await self._inner.__anext__()
+ note(chunk)
+ return chunk
+
+ def __getattr__(self, name):
+ return getattr(self._inner, name)
+
+ async def recording_fetch(*args, **kwargs):
+ result = await fetch(*args, **kwargs)
+ if isinstance(result, tuple):
+ response, stream = result
+ return response, _Tee(stream)
+ note(result)
+ return result
+
+ model._fetch_response = recording_fetch
+
+
+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. A client built on a
+ caller-owned http_client stays open."""
+ backend = getattr(getattr(agent, "model", None), "_client", None)
+ if getattr(backend, "_pageindex_caller_http", False):
+ return
+ 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:
+ if r.usage is None:
+ continue
+ prompt += r.usage.input_tokens or 0
+ completion += r.usage.output_tokens or 0
+ 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,
+ top_p: Optional[float] = None,
+ max_tokens: Optional[int] = None,
+ reasoning_effort: Optional[str] = None,
+ extra_body: Optional[dict] = None,
+ extra_headers: Optional[dict] = None,
+ backend: Optional[dict] = 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.chat_model
+ reported_model = _reported_model(model_name)
+ managed = _managed_instructions(system_texts)
+ agent = _openai_agent(client, "chat", model_name, managed,
+ temperature, top_p, doc_ids=doc_id,
+ cache_key=_conversation_cache_key(
+ model_name, managed, doc_id, history),
+ reasoning_effort=reasoning_effort,
+ extra_body=extra_body, max_tokens=max_tokens,
+ backend=_merged_backend(client, backend),
+ extra_headers=extra_headers)
+ recorded: dict = {}
+ _record_chat_finish(agent, recorded)
+ run_kwargs = _run_kwargs(max_turns)
+ 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 _model_backend_error(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": recorded.get("finish_reason") or "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 _model_backend_error(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=recorded.get("finish_reason") or "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,
+ max_output_tokens: Optional[int] = None,
+ reasoning: Optional[dict] = None,
+ extra_body: Optional[dict] = None,
+ extra_headers: Optional[dict] = None,
+ backend: Optional[dict] = 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.chat_model
+ managed = _managed_instructions(extra)
+ agent = _openai_agent(client, "responses", model_name, managed,
+ temperature, top_p, doc_ids=doc_id,
+ cache_key=_conversation_cache_key(
+ model_name, managed, doc_id, conversation),
+ reasoning=reasoning, extra_body=extra_body,
+ max_tokens=max_output_tokens,
+ backend=_merged_backend(client, backend),
+ extra_headers=extra_headers)
+ run_kwargs = _run_kwargs(max_turns)
+ recorded: dict = {}
+ import openai
+ from agents import Runner
+ from agents.exceptions import AgentsException, MaxTurnsExceeded
+
+ response_id = f"resp_{uuid.uuid4().hex}"
+ created_at = int(time.time())
+
+ def envelope(transcript: list, raw_responses) -> dict:
+ return {
+ "id": response_id,
+ "object": "response",
+ "created_at": created_at,
+ "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],
+ # Backend echo when captured; the request sends neither param.
+ "tool_choice": recorded.get("tool_choice", "auto"),
+ "parallel_tool_calls": recorded.get("parallel_tool_calls", True),
+ "temperature": temperature,
+ "top_p": top_p,
+ "reasoning": reasoning,
+ "max_output_tokens": max_output_tokens,
+ "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
+ opened = 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"] == "response.created" and not opened:
+ # N per-turn openings collapse to one, carrying
+ # the id and created_at the terminal event will
+ # report.
+ opened = True
+ if data.get("response"):
+ data["response"]["id"] = response_id
+ data["response"]["created_at"] = created_at
+ sequence += 1
+ data["sequence_number"] = sequence
+ yield data
+ continue
+ 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)
+ for field in ("tool_choice",
+ "parallel_tool_calls"):
+ if state.get(field) is not None:
+ recorded[field] = state[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(backend=None):
+ """The backend client — the seam tests replace with a fake transport."""
+ import anthropic
+ try:
+ return anthropic.Anthropic(**_sdk_backend(backend))
+ except TypeError as exc:
+ raise PageIndexAPIError(
+ f"The Anthropic backend is not configured: {exc}") from exc
+
+
+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 _cache_marks(system_blocks, messages) -> int:
+ """Breakpoints already on the request. The API allows 4 total; the
+ top-level moving breakpoint is only added when it fits."""
+ blocks = list(system_blocks)
+ for message in messages:
+ content = message.get("content")
+ if isinstance(content, list):
+ blocks += [b for b in content if isinstance(b, dict)]
+ return sum(1 for b in blocks
+ if isinstance(b, dict) and b.get("cache_control"))
+
+
+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, thinking=None) -> int:
+ """The wire-required per-turn budget when the caller sets none: 8192,
+ except the claude-3 generation whose output ceiling is 4096. The wire
+ also requires max_tokens > thinking.budget_tokens, so an enabled
+ budget lifts the default above itself."""
+ budget = (thinking.get("budget_tokens")
+ if isinstance(thinking, dict) else None)
+ if isinstance(budget, int):
+ return budget + 8192
+ 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,
+ thinking: Optional[dict] = None,
+ extra_body: Optional[dict] = None,
+ extra_headers: Optional[dict] = None,
+ backend: Optional[dict] = 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, "thinking": thinking,
+ "extra_body": extra_body, "extra_headers": extra_headers,
+ }.items() if value is not None}
+ system_blocks = _anthropic_system(system, block)
+ # Top-level cache_control: the server re-marks the newest block each
+ # turn, so the loop re-reads the growing conversation from cache.
+ # Counts toward the 4-breakpoint limit (live-verified 400 past it).
+ cached: dict[str, Any] = (
+ {"cache_control": {"type": "ephemeral"}}
+ if _cache_marks(system_blocks, prepared) < 4 else {})
+ merged = _merged_backend(client, backend)
+ # The SDK defers credential resolution to request time and raises a
+ # bare TypeError there — pre-check for the contract's PageIndexAPIError.
+ if not merged and not (os.environ.get("ANTHROPIC_API_KEY")
+ or os.environ.get("ANTHROPIC_AUTH_TOKEN")):
+ raise PageIndexAPIError(
+ "The Anthropic backend is not configured: set the "
+ "ANTHROPIC_API_KEY environment variable, or pass an api_key "
+ "in chat_backend / backend.")
+ backend_client = _anthropic_client(merged)
+ # A caller-owned http_client must survive the per-call closes below.
+ owns_transport = "http_client" not in (merged or {})
+ if max_tokens is None:
+ max_tokens = _default_max_tokens(model, thinking)
+ runner = backend_client.beta.messages.tool_runner(
+ max_tokens=max_tokens,
+ messages=prepared,
+ model=model,
+ tools=build_anthropic_tools(client, doc_ids=doc_id),
+ system=system_blocks,
+ 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,
+ **cached,
+ )
+
+ 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
+ finally:
+ # runs on exhaustion and abandonment (GeneratorExit) alike
+ if owns_transport:
+ backend_client.close()
+ return events()
+
+ try:
+ turns = [turn for turn in runner]
+ except anthropic.AnthropicError as exc:
+ raise PageIndexAPIError(
+ f"The model backend failed: {exc}") from exc
+ finally:
+ # safe here: the params read-back below does no HTTP
+ if owns_transport:
+ backend_client.close()
+ 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.get("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/local_store.py b/pageindex/local_store.py
index 37108f93c..cdba19b71 100644
--- a/pageindex/local_store.py
+++ b/pageindex/local_store.py
@@ -6,6 +6,7 @@
import os
import shutil
import uuid
+from contextlib import contextmanager
from pathlib import Path
logger = logging.getLogger(__name__)
@@ -89,6 +90,24 @@ def _write_manifest(self, docs: dict) -> None:
except OSError:
pass
+ @contextmanager
+ def lock(self):
+ """Cross-process mutex for check-then-write sequences (name
+ uniquing before save). fcntl is absent on Windows, where the
+ pre-existing best-effort behavior stays."""
+ try:
+ import fcntl
+ except ImportError:
+ yield
+ return
+ self._root.mkdir(parents=True, exist_ok=True)
+ with open(self._root / ".lock", "w") as handle:
+ fcntl.flock(handle, fcntl.LOCK_EX)
+ try:
+ yield
+ finally:
+ fcntl.flock(handle, fcntl.LOCK_UN)
+
# ── documents ──
def save_document(self, doc_id: str, meta: dict, tree: list, pages: list) -> None:
doc_dir = self._doc_dir(doc_id)
diff --git a/pageindex/mcp_bridge.py b/pageindex/mcp_bridge.py
new file mode 100644
index 000000000..32ec09884
--- /dev/null
+++ b/pageindex/mcp_bridge.py
@@ -0,0 +1,227 @@
+"""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 = requests.Session() # agent tool calls come in bursts
+ 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 self._session.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}).",
+ status_code=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]})",
+ status_code=response.status_code,
+ )
+ 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.",
+ status_code=response.status_code,
+ )
+ 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
+ # A server echoing its cursor (or cycling) must not hang the client:
+ # no-progress terminates, the page cap turns a cycle into an error.
+ for _ in range(50):
+ params = {"cursor": cursor} if cursor else {}
+ result = self._request("tools/list", params) or {}
+ tools.extend(result.get("tools") or [])
+ next_cursor = result.get("nextCursor")
+ if not next_cursor or next_cursor == cursor:
+ return tools
+ cursor = next_cursor
+ raise PageIndexAPIError(
+ "MCP tools/list pagination did not terminate within 50 pages.")
+
+ 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]")
+ elif (isinstance(block, dict)
+ and isinstance(block.get("resource"), dict)
+ and isinstance(block["resource"].get("blob"), str)):
+ # EmbeddedResource nests its base64 one level down.
+ resource = block["resource"]
+ kind = resource.get("mimeType") or "binary"
+ size_kb = max(1, len(resource["blob"]) * 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/pageindex/page_index_classic.py b/pageindex/page_index_classic.py
index 5b846e6d0..446e8893b 100644
--- a/pageindex/page_index_classic.py
+++ b/pageindex/page_index_classic.py
@@ -7,7 +7,6 @@
from .utils import *
from .tree_optimize import merge_tree
import os
-from concurrent.futures import ThreadPoolExecutor, as_completed
######################### Hardening for prompt injection patterns ####################################################
_INJECTION_PATTERNS = re.compile(
diff --git a/pageindex/tree_optimize.py b/pageindex/tree_optimize.py
index 04719ccb2..67d2e4728 100644
--- a/pageindex/tree_optimize.py
+++ b/pageindex/tree_optimize.py
@@ -61,7 +61,7 @@
import sys
from types import SimpleNamespace
-from .utils import (ConfigLoader, _is_openai_model, _is_unrecoverable,
+from .utils import (ConfigLoader, _is_unrecoverable, _openai_missing_keys,
llm_acompletion, strip_internal_keys)
TRIGGER_PAGES = 5 # only look ahead on nodes larger than this
@@ -872,9 +872,11 @@ async def main():
args = parser.parse_args()
model = args.model or default_model()
- if args.expand and not args.plan and _is_openai_model(model) \
- and not os.getenv("OPENAI_API_KEY"):
- sys.exit(f"OPENAI_API_KEY is not set (expand model: {model}).")
+ if args.expand and not args.plan:
+ missing = _openai_missing_keys(model)
+ if missing:
+ sys.exit(f"{', '.join(missing)} is not set "
+ f"(expand model: {model}).")
original = json.load(open(args.structure))
structure = copy.deepcopy(original["structure"])
diff --git a/pageindex/utils.py b/pageindex/utils.py
index 97f60a942..dbeddb56d 100644
--- a/pageindex/utils.py
+++ b/pageindex/utils.py
@@ -1,5 +1,7 @@
+import contextvars
import logging
import os
+import sys
import textwrap
from datetime import datetime
import time
@@ -19,8 +21,35 @@
# litellm is imported inside the functions that use it; eager import is slow
# and fetches a remote model-cost map.
+
+# The indexing lane's connection overrides, scoped by LocalAPI around each
+# indexing operation — a contextvar, so the value reaches this module's
+# helpers and their asyncio tasks without threading it through every call.
+_llm_backend: contextvars.ContextVar = contextvars.ContextVar(
+ "pageindex_llm_backend", default=None)
+
+
+def _repair_litellm_types() -> None:
+ """litellm 1.97.0's Message/Delta annotations carry nested forward refs
+ Python 3.10 cannot resolve (BerriAI/litellm#36384), so every completion
+ dies constructing its response. Rebuild them once with the defining
+ modules' names; no-op on 3.11+ and on fixed litellm releases."""
+ if sys.version_info >= (3, 11):
+ return
+ try:
+ import litellm.types.llms.openai as openai_types
+ import litellm.types.utils as litellm_types
+ namespace = {**vars(openai_types), **vars(litellm_types)}
+ litellm_types.Message.model_rebuild(_types_namespace=namespace)
+ litellm_types.Delta.model_rebuild(_types_namespace=namespace)
+ except Exception:
+ pass # best-effort: a failed repair leaves litellm's own error
+
# Backward compatibility: support CHATGPT_API_KEY as alias for OPENAI_API_KEY
if not os.getenv("OPENAI_API_KEY") and os.getenv("CHATGPT_API_KEY"):
+ import warnings
+ warnings.warn("CHATGPT_API_KEY is deprecated — set OPENAI_API_KEY "
+ "instead.", FutureWarning)
os.environ["OPENAI_API_KEY"] = os.getenv("CHATGPT_API_KEY")
def count_tokens(text, model=None):
@@ -36,16 +65,67 @@ def _strip_prefix(s, prefix):
return s
-def _is_openai_model(model):
- """Models without a provider prefix (no '/') use the openai SDK directly.
- For other providers, use 'provider/model' format (e.g. 'anthropic/claude-sonnet-4-6')."""
- if not model or model.startswith('litellm/'):
- return False
- return '/' not in model or model.startswith('openai/')
-
-
-_openai_sync_client = None
-_openai_async_client = None
+def run_off_loop(func, *args):
+ """Run func now, or on a worker thread when this thread already runs an
+ asyncio loop (func may itself call asyncio.run)."""
+ try:
+ asyncio.get_running_loop()
+ except RuntimeError:
+ return func(*args)
+ from concurrent.futures import ThreadPoolExecutor
+ with ThreadPoolExecutor(max_workers=1) as pool:
+ return pool.submit(func, *args).result()
+
+
+def _openai_missing_keys(model):
+ """Missing env keys for the pre-check, which covers only OpenAI-shaped
+ names (bare or ``openai/``): other providers resolve credentials their
+ own way at call time (IAM chains, ADC, Ollama's localhost default),
+ invisible to env inspection — the chat lane draws the same line.
+ ``litellm/``-prefixed names are exempt: the prefix is an explicit
+ routing choice, and litellm resolves credentials beyond the
+ environment (litellm.api_key, a keyless OPENAI_BASE_URL server).
+ Truthiness, not litellm's validate_environment, which reports a blank
+ exported key as present."""
+ if model.startswith("litellm/"):
+ return []
+ if "/" in model and not model.startswith("openai/"):
+ return []
+ return ([] if (os.getenv("OPENAI_API_KEY") or "").strip()
+ else ["OPENAI_API_KEY"])
+
+
+def _litellm_model(model, backend):
+ """Normalize to LiteLLM's grammar (``litellm/`` strips, bare names get
+ the ``openai/`` wire form — same as the chat lane) and fail fast on a
+ missing key or unknown provider, with status codes the retry loop and
+ the summary/optimize passes treat as unrecoverable."""
+ if not model:
+ return model
+ raw = model
+ model = _strip_prefix(model, "litellm/")
+ if "/" not in model:
+ model = f"openai/{model}"
+ import litellm
+ provider = model.split("/", 1)[0]
+ providers = getattr(litellm, "provider_list", None)
+ # custom_provider_map providers join provider_list only at call time.
+ custom = {entry.get("provider") for entry
+ in getattr(litellm, "custom_provider_map", None) or []}
+ if providers and provider not in providers and provider not in custom:
+ raise litellm.NotFoundError(
+ f"'{model}' routes through LiteLLM, but '{provider}' is not a "
+ f"LiteLLM provider. For an OpenAI-compatible server serving "
+ f"this model id, use 'openai/{model}' and point "
+ f"OPENAI_BASE_URL at the server.",
+ llm_provider=None, model=model)
+ if not backend:
+ missing = _openai_missing_keys(raw)
+ if missing:
+ raise litellm.AuthenticationError(
+ f"missing API key for {model}: {', '.join(missing)}",
+ llm_provider=None, model=model)
+ return model
# Misconfiguration: no retry can fix a rejected key or a model that does not
@@ -59,34 +139,34 @@ def _is_unrecoverable(exc: Exception) -> bool:
return getattr(exc, "status_code", None) in _UNRECOVERABLE_STATUS
+def _no_cache_seeding_kwargs(backend):
+ """litellm 1.97 auto-marks Claude requests for prompt caching (system +
+ last message); indexing prompts are single-shot and unique, so every call
+ would pay the cache-write premium with nothing ever read back. A
+ system-role-only injection point matches no indexing message, and its
+ presence stops litellm seeding its own defaults; backend keys still
+ win."""
+ return {"cache_control_injection_points":
+ [{"location": "message", "role": "system"}],
+ **(backend or {})}
+
+
def llm_completion(model, prompt, chat_history=None, return_finish_reason=False):
- use_openai_sdk = _is_openai_model(model)
- if model:
- model = _strip_prefix(model, "litellm/")
- if use_openai_sdk:
- model = _strip_prefix(model, "openai/")
+ import litellm
max_retries = 10
messages = list(chat_history) + [{"role": "user", "content": prompt}] if chat_history else [{"role": "user", "content": prompt}]
- if use_openai_sdk:
- global _openai_sync_client
- if _openai_sync_client is None:
- import openai
- _openai_sync_client = openai.OpenAI(max_retries=0)
+ backend = _llm_backend.get()
+ model = _litellm_model(model, backend)
+ _repair_litellm_types()
for i in range(max_retries):
try:
- if use_openai_sdk:
- response = _openai_sync_client.chat.completions.create(
- model=model,
- messages=messages,
- )
- else:
- import litellm
- response = litellm.completion(
- model=model,
- messages=messages,
- temperature=0,
- drop_params=True,
- )
+ response = litellm.completion(
+ model=model,
+ messages=messages,
+ drop_params=True,
+ # the loop is the retry policy; the merge lets a backend override win
+ **{"max_retries": 0, **_no_cache_seeding_kwargs(backend)},
+ )
content = response.choices[0].message.content
if return_finish_reason:
finish_reason = "max_output_reached" if response.choices[0].finish_reason == "length" else "finished"
@@ -106,33 +186,20 @@ def llm_completion(model, prompt, chat_history=None, return_finish_reason=False)
async def llm_acompletion(model, prompt):
- use_openai_sdk = _is_openai_model(model)
- if model:
- model = _strip_prefix(model, "litellm/")
- if use_openai_sdk:
- model = _strip_prefix(model, "openai/")
+ import litellm
max_retries = 10
messages = [{"role": "user", "content": prompt}]
- if use_openai_sdk:
- global _openai_async_client
- if _openai_async_client is None:
- import openai
- _openai_async_client = openai.AsyncOpenAI(max_retries=0)
+ backend = _llm_backend.get()
+ model = _litellm_model(model, backend)
+ _repair_litellm_types()
for i in range(max_retries):
try:
- if use_openai_sdk:
- response = await _openai_async_client.chat.completions.create(
- model=model,
- messages=messages,
- )
- else:
- import litellm
- response = await litellm.acompletion(
- model=model,
- messages=messages,
- temperature=0,
- drop_params=True,
- )
+ response = await litellm.acompletion(
+ model=model,
+ messages=messages,
+ drop_params=True,
+ **{"max_retries": 0, **_no_cache_seeding_kwargs(backend)},
+ )
return response.choices[0].message.content
except Exception as e:
if _is_unrecoverable(e):
@@ -665,6 +732,8 @@ async def generate_summaries_for_structure(structure, model=None):
summaries = await asyncio.gather(*tasks, return_exceptions=True)
for node, summary in zip(nodes, summaries):
+ if isinstance(summary, Exception) and _is_unrecoverable(summary):
+ raise summary
node['summary'] = "" if isinstance(summary, BaseException) else summary
if nodes and not any(node['summary'] for node in nodes):
raise RuntimeError(
@@ -838,17 +907,25 @@ async def parent_summary(node):
async def visit(node):
children = node.get('nodes') or []
if children:
- await asyncio.gather(*(visit(child) for child in children),
- return_exceptions=True)
+ done = await asyncio.gather(*(visit(child) for child in children),
+ return_exceptions=True)
+ for result in done:
+ if isinstance(result, Exception) and _is_unrecoverable(result):
+ raise result
if node.get('summary'):
return
try:
node['summary'] = await (parent_summary(node) if children else leaf_summary(node))
- except Exception:
+ except Exception as e:
node['summary'] = ""
+ if _is_unrecoverable(e):
+ raise
- await asyncio.gather(*(visit(root) for root in structure),
- return_exceptions=True)
+ results = await asyncio.gather(*(visit(root) for root in structure),
+ return_exceptions=True)
+ for r in results:
+ if isinstance(r, Exception) and _is_unrecoverable(r):
+ raise r
def _any_summary(nodes):
return any(n.get('summary') or _any_summary(n.get('nodes') or [])
@@ -954,6 +1031,29 @@ def thin(nodes, total_nodes):
return structure
+DEFAULT_INDEX_MODEL = "gpt-5.6-luna"
+DEFAULT_CHAT_MODEL = "gpt-5.6-sol"
+
+# Each of the five names has shipped in a release; all stay accepted.
+_MODEL_KEYS = ("model", "summary_model", "retrieve_model",
+ "index_model", "chat_model")
+
+
+def _resolve_models(merged: dict) -> None:
+ """Fill the model roles from whichever names were given: new names win
+ over old, specific over general, ``model`` sets every role, and the
+ built-in defaults close each chain. Idempotent, so already-resolved
+ config objects can round-trip through load()."""
+ given = {key: merged.get(key) for key in _MODEL_KEYS}
+ index = given["index_model"] or given["model"] or DEFAULT_INDEX_MODEL
+ summary = (given["summary_model"] or given["index_model"]
+ or given["model"] or DEFAULT_INDEX_MODEL)
+ chat = (given["chat_model"] or given["retrieve_model"]
+ or given["model"] or DEFAULT_CHAT_MODEL)
+ merged.update(model=index, index_model=index, summary_model=summary,
+ chat_model=chat, retrieve_model=chat)
+
+
class ConfigLoader:
def __init__(self, default_path: str = None):
if default_path is None:
@@ -966,7 +1066,8 @@ def _load_yaml(path):
return yaml.safe_load(f) or {}
def _validate_keys(self, user_dict):
- unknown_keys = set(user_dict) - set(self._default_dict)
+ unknown_keys = (set(user_dict) - set(self._default_dict)
+ - set(_MODEL_KEYS))
if unknown_keys:
raise ValueError(f"Unknown config keys: {unknown_keys}")
@@ -985,6 +1086,7 @@ def load(self, user_opt=None) -> config:
self._validate_keys(user_dict)
merged = {**self._default_dict, **user_dict}
+ _resolve_models(merged)
return config(**merged)
def create_node_mapping(tree, include_page_ranges=False, max_page=None):
diff --git a/pyproject.toml b/pyproject.toml
index deac66be3..f5bec8c9a 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,16 +25,28 @@ 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"
+# Older releases crash on current openai before the request is sent.
+openai-agents = ">=0.18.1"
+litellm = ">=1.97.0"
PyPDF2 = ">=3.0.0"
-pypdfium2 = ">=4.30.0"
+pypdfium2 = ">=5"
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 execute a refusal turn's tool_use blocks.
+anthropic = { version = ">=0.108.0", optional = true }
+
+[tool.poetry.extras]
+claude = ["claude-agent-sdk"]
+# Empty on purpose: keeps pip install "pageindex[openai]" valid.
+openai = []
+anthropic = ["anthropic"]
[tool.poetry.group.dev.dependencies]
pytest = ">=7.0"
diff --git a/requirements.txt b/requirements.txt
index 5fd4f2e4e..2406eef11 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,10 +1,10 @@
-litellm==1.84.0
+litellm==1.97.0
openai>=1.70.0
requests>=2.28.0
-# openai-agents # optional
+openai-agents>=0.18.1
# pymupdf # optional
PyPDF2==3.0.1
-pypdfium2==4.30.0
+pypdfium2==5.13.0
python-dotenv==1.2.2
pyyaml==6.0.2
regex>=2024.0.0
diff --git a/run_pageindex.py b/run_pageindex.py
index 452f08174..f2642b8a6 100644
--- a/run_pageindex.py
+++ b/run_pageindex.py
@@ -3,26 +3,36 @@
import json
from pageindex import *
from pageindex.page_index_md import md_to_tree
-from pageindex.utils import ConfigLoader
+from pageindex.utils import ConfigLoader, _openai_missing_keys
+
+# Keep LiteLLM's import off the network (frozen bundled model-cost map);
+# an explicit user setting wins.
+os.environ.setdefault("LITELLM_LOCAL_MODEL_COST_MAP", "True")
if __name__ == "__main__":
# Set up argument parser
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('--index-model', type=str, default=None,
+ help='Model used to index the document (overrides config.yaml)')
+ parser.add_argument('--model', type=str, default=None,
+ help='(legacy) Same as --index-model')
parser.add_argument('--summary-model', type=str, default=None,
- help='Model for node summaries (defaults to --model, then config.yaml)')
+ help='Model for node summaries (defaults to --index-model, then --model, then config.yaml)')
parser.add_argument('--toc-check-pages', type=int, default=None,
help='Number of pages to check for table of contents (PDF only)')
@@ -48,18 +58,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 is not None 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 +92,24 @@
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.index_model
+ or args.model
+ or ConfigLoader().load().summary_model)
+ will_summarize = args.summary if args.summary is not None else True
+ if will_summarize or args.optimize == 'full':
+ missing = _openai_missing_keys(summary_model)
+ if missing:
+ raise SystemExit(
+ f"Missing API key for {summary_model}: {', '.join(missing)}")
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']
@@ -94,7 +120,9 @@
else:
# Process PDF file
user_opt = {
+ 'index_model': args.index_model,
'model': args.model,
+ 'summary_model': args.summary_model,
'toc_check_page_num': args.toc_check_pages,
'max_page_num_each_node': args.max_pages_per_node,
'max_token_num_each_node': args.max_tokens_per_node,
@@ -110,7 +138,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)
@@ -139,26 +167,26 @@
# Create options dict with user args
user_opt = {
+ 'index_model': args.index_model,
'model': args.model,
- '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_add_node_id': args.if_add_node_id
}
# Load config with defaults from config.yaml
- opt = config_loader.load(user_opt)
+ opt = config_loader.load({k: v for k, v in user_opt.items() if v is not None})
+ # if_add_* pass through as given (absent = off, as before this CLI
+ # used config.yaml): the PDF defaults there must not switch on LLM
+ # passes the markdown CLI never ran.
toc_with_page_number = asyncio.run(md_to_tree(
md_path=args.md_path,
if_thinning=args.if_thinning.lower() == 'yes',
min_token_threshold=args.thinning_threshold,
- if_add_node_summary=opt.if_add_node_summary,
+ if_add_node_summary=args.if_add_node_summary,
summary_token_threshold=args.summary_token_threshold,
model=opt.model,
- if_add_doc_description=opt.if_add_doc_description,
- if_add_node_text=opt.if_add_node_text,
- if_add_node_id=opt.if_add_node_id
+ if_add_doc_description=args.if_add_doc_description,
+ if_add_node_text=args.if_add_node_text,
+ if_add_node_id=args.if_add_node_id
))
print('Parsing done, saving to file...')
diff --git a/tests/conftest.py b/tests/conftest.py
index 40a4a9199..04af63745 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -5,6 +5,7 @@
def _llm_key(monkeypatch):
"""Deterministic key presence for every test; missing-key tests delenv."""
monkeypatch.setenv("OPENAI_API_KEY", "test-key")
+ monkeypatch.setenv("ANTHROPIC_API_KEY", "test-key")
def build_pdf(page_texts):
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..28324d23d
--- /dev/null
+++ b/tests/test_agent_tools.py
@@ -0,0 +1,2530 @@
+"""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_edits_are_deliberate():
+ """The committed snapshot cannot detect drift from the live cloud
+ server — both copies live in this repo. It exists so a TOOL_CONTRACT
+ edit must touch two files in one change, never land by accident."""
+ 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):
+ # Escape-dense pages: raw length fits the budget, JSON-serialized
+ # length does not — the budget must count emitted characters.
+ pages = [
+ {"page_index": 1, "markdown": '"' * 30_000},
+ {"page_index": 2, "markdown": '"' * 20_000},
+ ]
+ 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"])
+ emitted = json.dumps(payload, ensure_ascii=False)
+ assert len(emitted) <= agent_tools_module.TOOL_RESPONSE_CHAR_LIMIT
+
+
+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_model_speaks_the_agents_sdk_grammar(tmp_path):
+ """The bundle's model string is resolved by the Agents SDK's own
+ prefix grammar, which refuses unknown prefixes — the constructor's
+ normalized litellm/ spelling is what must reach this door."""
+ pytest.importorskip("agents")
+ client = PageIndexLocalClient(storage_path=str(tmp_path / "s"),
+ retrieve_model="anthropic/claude-x")
+ assert (client.openai_agent_config()["model"]
+ == "litellm/anthropic/claude-x")
+ # The per-call override speaks the same grammar as chat_model.
+ assert (client.openai_agent_config(model="anthropic/claude-y")["model"]
+ == "litellm/anthropic/claude-y")
+ assert (client.openai_agent_config(model="litellm/groq/llama-x")["model"]
+ == "litellm/groq/llama-x")
+
+
+def test_plain_functions_answer_bad_arguments_with_the_envelope(client,
+ store_path):
+ """agent_tools() functions must not raise into a framework loop:
+ cloud-only parameters pruned from the local signature come back as
+ the guided envelope, and the schema-bearing signature survives."""
+ import inspect
+ seed_doc(store_path, "pi-a", "report.pdf")
+ tools = {f.__name__: f for f in client.agent_tools()}
+ fn = tools["get_document"]
+ assert "doc_name" in inspect.signature(fn).parameters
+ payload = json.loads(fn(doc_name="report.pdf", folder_id="root"))
+ assert payload["errorCode"] == "INVALID_INPUT"
+ ok = json.loads(fn(doc_name="report.pdf"))
+ assert not ok.get("errorCode")
+
+
+def test_non_string_doc_name_stays_not_found(client, store_path):
+ """A type-loose model argument must not turn NOT_FOUND into the
+ retry-inviting INTERNAL_ERROR (strict_json_schema is off on the
+ OpenAI adapter, so nothing upstream validates the type)."""
+ seed_doc(store_path, "pi-a", "report.pdf")
+ payload, is_error = run(client, "get_document", doc_name=5)
+ assert is_error and payload["errorCode"] == "NOT_FOUND"
+
+
+def test_openai_agent_config_repairs_litellm_types(tmp_path, monkeypatch):
+ """The BYO path resolves its model through LiteLLM in the caller's
+ process, outside our completion helpers — the py3.10 type repair must
+ run at config time, and only for LiteLLM-routed models."""
+ pytest.importorskip("agents")
+ import pageindex.utils
+ calls = []
+ monkeypatch.setattr(pageindex.utils, "_repair_litellm_types",
+ lambda: calls.append(True))
+ client = PageIndexLocalClient(storage_path=str(tmp_path / "s"),
+ chat_model="anthropic/claude-x")
+ client.openai_agent_config()
+ assert calls
+ calls.clear()
+ client.openai_agent_config(model="gpt-plain")
+ assert not calls
+
+
+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 config["cache_control"] == {"type": "ephemeral"}
+ 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,
+ monkeypatch):
+ """claude_agent_config(doc_id=...) must wire scope all the way into the
+ handlers it registers — an out-of-scope document returns NOT_FOUND. The
+ assertion has to drive those handlers, or build_claude_mcp's doc_ids
+ pass-through goes unguarded."""
+ claude_agent_sdk = pytest.importorskip("claude_agent_sdk")
+ seed_doc(store_path, "pi-a", "report.pdf")
+ seed_doc(store_path, "pi-b", "payroll.pdf",
+ created_at="2026-08-02T10:00:00.123000")
+
+ registered = {}
+ create_server = claude_agent_sdk.create_sdk_mcp_server
+
+ def capture(**kwargs):
+ registered.update(kwargs)
+ return create_server(**kwargs)
+
+ monkeypatch.setattr(claude_agent_sdk, "create_sdk_mcp_server", capture)
+ config = client.claude_agent_config(doc_id="pi-a")
+ assert "report.pdf" in config["system_prompt"]
+ handlers = {spec.name: spec.handler for spec in registered["tools"]}
+ result = asyncio.run(handlers["get_page_content"](
+ {"doc_name": "payroll.pdf", "pages": "1"}))
+ assert result.get("is_error")
+ assert json.loads(result["content"][0]["text"])["errorCode"] == "NOT_FOUND"
+ browse = asyncio.run(handlers["browse_documents"]({}))
+ listed = json.loads(browse["content"][0]["text"])["documents"]
+ assert [doc["name"] for doc in listed] == ["report.pdf"]
+
+
+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_anthropic_runner_config_thinking_lifts_max_tokens(client):
+ pytest.importorskip("anthropic")
+ config = client.anthropic_runner_config(
+ model="claude-sonnet-4-5",
+ thinking={"type": "enabled", "budget_tokens": 10000})
+ assert config["max_tokens"] == 10000 + 8192
+ assert config["thinking"] == {"type": "enabled", "budget_tokens": 10000}
+ assert "thinking" not in client.anthropic_runner_config(
+ model="claude-sonnet-4-5")
+
+
+def test_bridge_invoker_reraises_auth_failures():
+ class Revoked:
+ def call_tool(self, name, arguments):
+ raise PageIndexAPIError("HTTP 401", status_code=401)
+
+ invoke = agent_tools_module._bridge_invoker(Revoked(), "get_document", {})
+ with pytest.raises(PageIndexAPIError, match="401"):
+ invoke({})
+ # Transport blips stay contained in the retryable envelope.
+
+ class Down:
+ def call_tool(self, name, arguments):
+ raise PageIndexAPIError("HTTP 503", status_code=503)
+
+ text, is_error = agent_tools_module._bridge_invoker(
+ Down(), "get_document", {})({})
+ assert is_error
+ assert json.loads(text)["errorCode"] == "INTERNAL_ERROR"
+
+
+def test_cloud_bridge_gates_the_endpoint(monkeypatch):
+ """Instructions come from the same endpoint the tools register: the
+ read-gated URL by default, the full one with include_management."""
+ created = []
+
+ class FakeBridge:
+ def __init__(self, url, headers):
+ created.append(url)
+
+ def instructions(self):
+ return "SERVED"
+
+ monkeypatch.setattr("pageindex.mcp_bridge.McpBridge", FakeBridge)
+ cloud = PageIndexCloudClient(api_key="pi-k")
+ assert cloud.agent_instructions() == "SERVED"
+ assert created == [f"{cloud.BASE_URL}/mcp?tools=read"]
+ cloud.agent_instructions(include_management=True)
+ assert created[1:] == [f"{cloud.BASE_URL}/mcp"]
+ cloud.agent_instructions()
+ cloud.agent_instructions(include_management=True)
+ assert len(created) == 2 # cached per gate
+
+
+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"]
+ # Default discovery rides the read-gated endpoint, matching the
+ # instructions fetch and the hosted/MCP registrations.
+ assert bridge.url == "https://api.pageindex.ai/mcp?tools=read"
+ 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."""
+
+ class _Bridge:
+ def call_tool(self, name, arguments):
+ return json.dumps({"success": True}), False
+
+ tool = _synth(_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_local_description_edit_actually_removes_the_image_sentence():
+ """The local get_page_content description edits the contract text by
+ exact string replace — a contract wording change must fail here, not
+ silently ship the image-tool sentence to local models."""
+ from pageindex.agent_tools import TOOL_CONTRACT, _LOCAL_DESCRIPTIONS
+ local = _LOCAL_DESCRIPTIONS["get_page_content"]
+ assert "get_document_image" not in local
+ assert len(local) < len(TOOL_CONTRACT["get_page_content"]["description"])
+
+
+def test_list_tools_pagination_is_bounded():
+ """A server echoing its nextCursor terminates (no-progress guard);
+ a cycling one hits the page cap instead of hanging forever."""
+ from pageindex.mcp_bridge import McpBridge
+
+ bridge = McpBridge("http://x/mcp", {})
+ pages = {None: {"tools": [{"name": "a"}], "nextCursor": "c1"},
+ "c1": {"tools": [{"name": "b"}], "nextCursor": "c1"}}
+ bridge._request = lambda method, params=None: pages[
+ (params or {}).get("cursor")]
+ assert [t["name"] for t in bridge.list_tools()] == ["a", "b"]
+
+ bridge._request = lambda method, params=None: {
+ "tools": [],
+ "nextCursor": {"c1": "c2"}.get((params or {}).get("cursor"), "c1")}
+ with pytest.raises(PageIndexAPIError, match="did not terminate"):
+ bridge.list_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(
+ Session=lambda: 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(
+ Session=lambda: 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(
+ Session=lambda: 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},
+ {"type": "resource",
+ "resource": {"mimeType": "image/jpeg", "blob": 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
+ assert "[image/jpeg content omitted: ~6 KB]" in text
+
+
+# ── review-round regressions ──
+
+def _synth(bridge, meta):
+ """Build a tool function the way the cloud lane does: signature
+ synthesis over a bridge invoker."""
+ from pageindex.agent_tools import _bridge_invoker, _make_tool_function
+ name = meta["name"]
+ return _make_tool_function(name, meta.get("description"),
+ meta["inputSchema"],
+ _bridge_invoker(bridge, name,
+ meta["inputSchema"]))
+
+
+def test_synth_binding_error_names_the_tool():
+ """Binding TypeErrors quote the function's __qualname__; the model used
+ to see \"_synthesized() got an unexpected keyword argument\" and had no
+ tool name to correct against."""
+ from pageindex.agent_tools import TOOL_CONTRACT
+
+ 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 = _synth(_Bridge(), meta)
+ payload = json.loads(fn(bogus_param=1))
+ assert "browse_documents" in payload["error"]
+ assert "_synthesized" not in payload["error"]
+
+
+def test_bridge_invoker_coerces_string_booleans():
+ """Identical model output must behave the same on both dispatch paths:
+ call_tool coerced "false" but the cloud bridge forwarded it verbatim,
+ turning "don't wait" into a 3-minute wait on lenient servers."""
+ from pageindex.agent_tools import TOOL_CONTRACT, _bridge_invoker
+
+ seen = {}
+
+ class _Bridge:
+ def call_tool(self, name, args):
+ seen.update(args)
+ return "{}", False
+
+ invoke = _bridge_invoker(_Bridge(), "get_document",
+ TOOL_CONTRACT["get_document"]["schema"])
+ invoke({"doc_name": "q.pdf", "wait_for_completion": "false"})
+ assert seen["wait_for_completion"] is False
+
+
+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 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 = _synth(_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
+
+ 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 = _synth(_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():
+
+ 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 = _synth(_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 = _synth(_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 = _synth(_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(
+ Session=lambda: 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(
+ Session=lambda: 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(
+ Session=lambda: 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
+
+ # stop_ids ends the walk once every wanted id has been seen — a
+ # doc-scoped chat turn must not page the whole library...
+ early = make_client(120, 100)
+ listed = _all_documents(early, stop_ids=frozenset({"pi-3"}))
+ assert type(early).calls == 1
+ assert any(doc["id"] == "pi-3" for doc in listed)
+ # ...while an id the listing lacks still costs the full sweep.
+ full = make_client(120, 100)
+ assert _all_documents(full, stop_ids=frozenset({"pi-missing"})) == docs
+ assert type(full).calls == 2
+
+
+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_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"}
+
+
+import pageindex.utils # noqa: F401 — its import loads .env
+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_bridge_rebuilds_on_credential_change(monkeypatch):
+ """CloudAPI re-reads client.api_key on every REST call; the MCP half
+ must not keep authenticating with a rotation-stale snapshot."""
+ import pageindex.mcp_bridge as mcp_bridge
+ from pageindex.agent_tools import _cloud_bridge
+ built = []
+
+ class _Bridge(_FakeBridge):
+ def __init__(self, url, headers):
+ super().__init__(url, headers)
+ built.append((url, dict(headers)))
+
+ monkeypatch.setattr(mcp_bridge, "McpBridge", _Bridge)
+ cloud = PageIndexCloudClient(api_key="pi-old")
+ first = _cloud_bridge(cloud)
+ assert _cloud_bridge(cloud) is first # unchanged credentials: cached
+ cloud.api_key = "pi-new"
+ second = _cloud_bridge(cloud)
+ assert second is not first
+ assert built[-1][1]["Authorization"] == "Bearer pi-new"
+ cloud.BASE_URL = "https://alt.example"
+ assert _cloud_bridge(cloud) is not second
+ assert built[-1][0] == "https://alt.example/mcp"
+
+
+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_doc_targeting_keeps_transport_errors_out_of_not_found():
+ """A cloud 429/5xx during the doc_id fetch is an outage, not a missing
+ document — only a definite not-found/denied batches into the
+ "Documents not found" message; anything else propagates raw."""
+ class Stub:
+ def __init__(self, status):
+ self.status = status
+
+ def get_document(self, doc_id):
+ raise PageIndexAPIError(
+ f"Failed to get document metadata: {self.status}",
+ status_code=self.status)
+
+ for status in (429, 500):
+ with pytest.raises(PageIndexAPIError, match=f"metadata: {status}"):
+ agent_tools_module.doc_targeting_block(Stub(status), "pi-a")
+ for status in (403, 404):
+ with pytest.raises(PageIndexAPIError,
+ match="Documents not found or access denied: pi-a"):
+ agent_tools_module.doc_targeting_block(Stub(status), "pi-a")
+
+
+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
+
+
+# ── wait-poll resilience, instruction scoping, agent_tools doc_id ──
+
+def test_await_completion_polls_through_transient_refetch_failure(monkeypatch):
+ """A refetch that fails once must not end the wait early — the caller
+ would report that 5-second exit as the full 3-minute timeout."""
+ monkeypatch.setattr(agent_tools_module, "_TOOL_WAIT_INTERVAL", 0.0)
+ calls = {"n": 0}
+
+ class Flaky:
+ def get_document(self, doc_id):
+ calls["n"] += 1
+ if calls["n"] == 1:
+ raise PageIndexAPIError("transient listing failure")
+ return {"id": doc_id, "status": "completed"}
+
+ result = agent_tools_module._await_completion(
+ Flaky(), {"id": "pi-x", "status": "processing"}, wait=True)
+ assert result["status"] == "completed"
+ assert calls["n"] == 2
+
+
+def test_agent_instructions_doc_id_shadow_check(client, store_path):
+ seed_doc(store_path, "pi-old", "report.pdf",
+ created_at="2026-08-01T10:00:00.123000")
+ seed_doc(store_path, "pi-new", "report.pdf",
+ created_at="2026-08-05T10:00:00.123000")
+ # standalone instructions get the strict check; only the *_agent_config
+ # bundles (which build the tools too) relax it
+ with pytest.raises(PageIndexAPIError, match="shadowed"):
+ client.agent_instructions(doc_id="pi-old")
+
+
+def test_agent_tools_doc_id_scopes_the_functions(client, store_path):
+ seed_doc(store_path, "pi-a", "alpha.pdf")
+ seed_doc(store_path, "pi-b", "secret.pdf")
+ funcs = {fn.__name__: fn for fn in client.agent_tools(doc_id="pi-a")}
+ blocked = json.loads(funcs["get_page_content"](doc_name="secret.pdf",
+ pages="1"))
+ assert "success" not in blocked
+ assert blocked["errorCode"] == "NOT_FOUND"
+ allowed = json.loads(funcs["get_page_content"](doc_name="alpha.pdf",
+ pages="1"))
+ assert allowed["success"] is True
+
+
+def test_agent_tools_doc_id_refused_on_cloud():
+ cloud = PageIndexCloudClient(api_key="pi-test-key")
+ with pytest.raises(PageIndexAPIError, match="local tools only"):
+ cloud.agent_tools(doc_id="pi-a")
diff --git a/tests/test_client.py b/tests/test_client.py
index 50b7f5178..2ace6ba2e 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 ──
@@ -68,14 +68,45 @@ def test_local_client_does_not_touch_disk(tmp_path):
assert not storage.exists()
-def test_retrieve_model_carries_agents_sdk_prefix(tmp_path):
+def test_retrieve_model_stays_as_configured(tmp_path):
+ """The public attribute keeps the caller's spelling — ``litellm/`` is
+ Agents SDK routing grammar, applied at the config door
+ (openai_agent_config), never baked into ``chat_model``: handed to the
+ Anthropic SDK or a raw request, the prefixed form is a 404."""
def resolved(retrieve_model):
return PageIndexClient(retrieve_model=retrieve_model,
storage_path=str(tmp_path / "s")).retrieve_model
- assert resolved("anthropic/claude-sonnet-4-6") == "litellm/anthropic/claude-sonnet-4-6"
- for already_routable in ("gpt-4o", "openai/gpt-4o", "litellm/anthropic/claude-sonnet-4-6"):
- assert resolved(already_routable) == already_routable
+ for as_configured in ("anthropic/claude-sonnet-4-6", "gpt-4o",
+ "openai/gpt-4o",
+ "litellm/anthropic/claude-sonnet-4-6"):
+ assert resolved(as_configured) == as_configured
+
+
+def test_model_resolution_covers_every_generation(tmp_path):
+ """New names win over old, specific over general, ``model`` sets every
+ role, and the built-in defaults close each chain. One row per released
+ surface: 0.2.8 (model only), 0.3.0.dev (model + retrieve_model),
+ 0.2.10.dev (all three legacy names), the current pair, plus the
+ umbrella and mixed forms."""
+ from pageindex.utils import DEFAULT_CHAT_MODEL, DEFAULT_INDEX_MODEL
+ cases = [
+ ({}, DEFAULT_INDEX_MODEL, DEFAULT_INDEX_MODEL, DEFAULT_CHAT_MODEL),
+ ({"model": "m"}, "m", "m", "m"),
+ ({"model": "m", "retrieve_model": "r"}, "m", "m", "r"),
+ ({"model": "m", "summary_model": "s", "retrieve_model": "r"},
+ "m", "s", "r"),
+ ({"index_model": "i", "chat_model": "c"}, "i", "i", "c"),
+ ({"model": "m", "index_model": "i"}, "i", "i", "m"),
+ ({"summary_model": "s"},
+ DEFAULT_INDEX_MODEL, "s", DEFAULT_CHAT_MODEL),
+ ]
+ for kwargs, index, summary, chat in cases:
+ client = PageIndexClient(storage_path=str(tmp_path / "s"), **kwargs)
+ assert (client.index_model, client.model) == (index, index), kwargs
+ assert client.summary_model == summary, kwargs
+ assert client.chat_model == chat, kwargs
+ assert client.retrieve_model == client.chat_model, kwargs
def test_explicit_mode_clients(tmp_path):
@@ -147,10 +178,20 @@ def test_get_page_content(local_client, indexed_doc):
assert local_client.get_page_content(indexed_doc, "99") == []
- with pytest.raises(ValueError):
+ with pytest.raises(PageIndexAPIError):
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(PageIndexAPIError, 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 +199,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,27 +260,66 @@ 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
+ import litellm # first import may load a .env; delenv after it
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
- monkeypatch.setattr(pageindex.utils, "_openai_sync_client", None)
- monkeypatch.setattr(pageindex.utils, "_openai_async_client", None)
- with pytest.raises(openai.OpenAIError):
+ with pytest.raises(openai.OpenAIError, match="OPENAI_API_KEY"):
pageindex.utils.llm_completion("gpt-4o", "probe")
- with pytest.raises(openai.OpenAIError):
+ with pytest.raises(openai.OpenAIError, match="OPENAI_API_KEY"):
asyncio.run(pageindex.utils.llm_acompletion("gpt-4o", "probe"))
+ # unknown bare names are OpenAI shorthand, so the same check applies
+ with pytest.raises(openai.OpenAIError, match="OPENAI_API_KEY"):
+ pageindex.utils.llm_completion("my-finetune-v2", "probe")
+ # a blank exported key is as missing as no key (litellm's
+ # validate_environment reports it present)
+ monkeypatch.setenv("OPENAI_API_KEY", " ")
+ with pytest.raises(openai.OpenAIError, match="OPENAI_API_KEY"):
+ pageindex.utils.llm_completion("gpt-4o", "probe")
+
+
+def test_llm_completion_refuses_unknown_provider(monkeypatch):
+ """A first segment LiteLLM does not know (a HuggingFace repo id like
+ Qwen/...) is refused with the openai/ escape before the retry loop,
+ instead of burning it on per-call 400s."""
+ import litellm
+ monkeypatch.setattr(litellm, "completion",
+ lambda **kw: pytest.fail("reached the wire"))
+ with pytest.raises(Exception, match="not a LiteLLM provider"):
+ pageindex.utils.llm_completion("Qwen/my-model", "probe")
def test_submit_missing_llm_key_fails_loud(local_client, sample_pdf, monkeypatch):
+ import litellm # first import may load a .env; delenv after it
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
- monkeypatch.setattr(pageindex.utils, "_openai_sync_client", None)
def first_llm_call(*args, **kwargs):
return pageindex.utils.llm_completion("gpt-4o", "probe")
monkeypatch.setattr(page_index_module, "page_index_main", first_llm_call)
@@ -282,7 +407,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 +557,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
@@ -549,6 +675,70 @@ async def flaky(model, prompt):
assert summaries == {"A": "", "B": "ok"}
+def test_generate_summaries_unrecoverable_raises(monkeypatch):
+ """A per-node 401 must abort, not store a blank node as completed."""
+ class Denied(Exception):
+ status_code = 401
+
+ async def deny_t1(model, prompt):
+ if "t1" in prompt:
+ raise Denied("key rejected")
+ return "ok"
+ monkeypatch.setattr(pageindex.utils, "llm_acompletion", deny_t1)
+ structure = [{"title": "A", "text": "t1",
+ "nodes": [{"title": "B", "text": "t2"}]}]
+ with pytest.raises(Denied):
+ asyncio.run(pageindex.utils.generate_summaries_for_structure(structure))
+
+
+def test_summarize_tree_child_unrecoverable_raises(monkeypatch):
+ """A 401 on a leaf must abort the run, not store a blank subtree as
+ completed: the child gather's exceptions are checked, not discarded."""
+ class Denied(Exception):
+ status_code = 401
+
+ async def deny_alpha(model, prompt):
+ if "alpha" in prompt:
+ raise Denied("key rejected")
+ return '{"points": [], "summary": "ok"}'
+ monkeypatch.setattr(pageindex.utils, "llm_acompletion", deny_alpha)
+ pdf_pages = [("alpha " * 5, 5), ("beta " * 5, 5)]
+ structure = [{"title": "R", "start_index": 1, "end_index": 2,
+ "nodes": [
+ {"title": "A", "start_index": 1, "end_index": 1},
+ {"title": "B", "start_index": 2, "end_index": 2}]}]
+ with pytest.raises(Denied):
+ asyncio.run(pageindex.utils.summarize_tree(
+ structure, pdf_pages, small_node_tokens=0))
+
+
+def test_llm_completion_suppresses_litellm_cache_seeding(monkeypatch):
+ """Indexing prompts are single-shot: without an explicit injection
+ point litellm 1.97 seeds its own cache marks and every call pays the
+ write premium for nothing. Backend keys still override ours."""
+ import litellm
+ captured = {}
+
+ def fake_completion(**kwargs):
+ captured.clear()
+ captured.update(kwargs)
+ message = types.SimpleNamespace(content="ok")
+ choice = types.SimpleNamespace(message=message, finish_reason="stop")
+ return types.SimpleNamespace(choices=[choice])
+ monkeypatch.setattr(litellm, "completion", fake_completion)
+ monkeypatch.setenv("OPENAI_API_KEY", "k")
+ assert pageindex.utils.llm_completion("gpt-4o", "probe") == "ok"
+ assert captured["cache_control_injection_points"] == [
+ {"location": "message", "role": "system"}]
+ token = pageindex.utils._llm_backend.set(
+ {"api_key": "x", "cache_control_injection_points": []})
+ try:
+ pageindex.utils.llm_completion("gpt-4o", "probe")
+ finally:
+ pageindex.utils._llm_backend.reset(token)
+ assert captured["cache_control_injection_points"] == []
+
+
def test_delete_survives_marker_tamper(local_client, tmp_path):
tampered = tmp_path / "store" / "docs" / "tampered" / "doc.json"
tampered.mkdir(parents=True)
@@ -599,8 +789,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_openai_agents(local_client, monkeypatch):
+ """Local chat is implemented (see test_local_chat.py); without
+ openai-agents installed it raises the actionable install error."""
+ import sys
+ monkeypatch.setitem(sys.modules, "agents", None)
+ with pytest.raises(PageIndexAPIError, match="pip install openai-agents"):
local_client.chat_completions(
messages=[{"role": "user", "content": "q"}])
@@ -710,3 +904,178 @@ 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(PageIndexAPIError, match="spans more than"):
+ _parse_pages("1-10001")
+ # one parser with the tool layer now: page 0 is rejected, not passed
+ # on — surfaced as the documented SDK error type
+ with pytest.raises(PageIndexAPIError, match="positive"):
+ _parse_pages("0-3")
+
+
+# ── backend: the indexing lane ──
+
+def test_backend_scopes_the_index_lane(tmp_path, monkeypatch):
+ """index_backend reaches the indexing lane's LiteLLM call kwargs
+ verbatim — bare and provider-prefixed models alike — scoped to the
+ operation, bypassing the env pre-check."""
+ pytest.importorskip("litellm")
+ import litellm
+ from types import SimpleNamespace
+ from pageindex.local_api import LocalAPI
+ from pageindex.utils import _llm_backend, llm_completion
+
+ api = LocalAPI(storage_path=str(tmp_path / "s"), model="m",
+ summary_model="s",
+ index_backend={"api_key": "ik", "api_base": "http://b"})
+ assert api._with_backend(_llm_backend.get) == {"api_key": "ik",
+ "api_base": "http://b"}
+ assert _llm_backend.get() is None
+
+ reply = SimpleNamespace(choices=[SimpleNamespace(
+ message=SimpleNamespace(content="ok"), finish_reason="stop")])
+ captured = {}
+ monkeypatch.setattr(litellm, "completion",
+ lambda **kw: (captured.update(kw), reply)[1])
+ monkeypatch.setattr(litellm, "validate_environment",
+ lambda *a, **k: pytest.fail("env pre-check ran"))
+ for model, wire in (("anthropic/claude-x", "anthropic/claude-x"),
+ ("gpt-4o", "openai/gpt-4o"),
+ ("my-finetune-v2", "openai/my-finetune-v2")):
+ captured.clear()
+ api._with_backend(lambda: llm_completion(model, "p"))
+ assert captured["model"] == wire
+ assert captured["api_key"] == "ik"
+ assert captured["api_base"] == "http://b"
+
+
+def test_index_precheck_covers_only_openai_shaped(monkeypatch):
+ """The missing-key pre-check fires only for OpenAI-shaped names — other
+ providers resolve credentials at call time (IAM chains, ADC, Ollama's
+ localhost default), invisible to env inspection, so the lane must not
+ block them up front."""
+ pytest.importorskip("litellm")
+ import litellm
+ from types import SimpleNamespace
+ from pageindex.utils import llm_completion
+
+ monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+ with pytest.raises(litellm.AuthenticationError, match="missing API key"):
+ llm_completion("my-finetune-v2", "p")
+
+ reply = SimpleNamespace(choices=[SimpleNamespace(
+ message=SimpleNamespace(content="ok"), finish_reason="stop")])
+ monkeypatch.setattr(litellm, "completion", lambda **kw: reply)
+ monkeypatch.setattr(litellm, "validate_environment",
+ lambda *a, **k: pytest.fail("env pre-check ran"))
+ assert llm_completion("ollama/llama3", "p") == "ok"
+ assert llm_completion("bedrock/anthropic.claude-sonnet", "p") == "ok"
+
+
+def test_litellm_routing_prefix_skips_key_precheck(monkeypatch):
+ """litellm/-prefixed names are an explicit routing choice: litellm
+ resolves credentials beyond the environment (litellm.api_key, a
+ keyless OPENAI_BASE_URL server), so the env pre-check stands aside."""
+ pytest.importorskip("litellm")
+ import litellm # noqa: F401 — first import may load a .env; delenv after
+ from pageindex.utils import _litellm_model, _openai_missing_keys
+
+ monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+ assert _openai_missing_keys("litellm/gpt-4o") == []
+ assert _litellm_model("litellm/gpt-4o", None) == "openai/gpt-4o"
+
+
+def test_custom_provider_map_passes_provider_precheck(monkeypatch):
+ """litellm appends custom_provider_map providers to provider_list only
+ at completion time, so the pre-check must consult the map itself."""
+ pytest.importorskip("litellm")
+ import litellm
+ from pageindex.utils import _litellm_model
+
+ monkeypatch.setattr(litellm, "custom_provider_map",
+ [{"provider": "my-llm", "custom_handler": object()}])
+ assert _litellm_model("my-llm/model-a", None) == "my-llm/model-a"
+
+
+def test_backend_args_are_local_only():
+ with pytest.raises(PageIndexAPIError, match="chat_backend"):
+ PageIndexClient(api_key="pi-k", chat_backend={"api_key": "x"})
+ with pytest.raises(PageIndexAPIError, match="index_backend"):
+ PageIndexClient(api_key="pi-k", index_backend={"api_key": "x"})
+
+
+def test_chat_wraps_answerless_cloud_reply(monkeypatch):
+ """A cloud reply without choices (filtered / malformed) surfaces as
+ the SDK's error, not a bare IndexError/KeyError."""
+ client = PageIndexClient(api_key="pi-k")
+ for reply in ({"id": "x", "object": "chat.completion", "choices": []},
+ {"id": "x"}):
+ monkeypatch.setattr(client, "chat_completions",
+ lambda *a, _r=reply, **k: _r)
+ with pytest.raises(PageIndexAPIError, match="carries no answer"):
+ client.chat("hi")
+
+
+def test_retrieve_model_assignment_still_works(local_client):
+ """0.2.9 allowed `client.retrieve_model = ...`; the legacy property
+ keeps the write path as an alias for chat_model."""
+ local_client.retrieve_model = "gpt-x"
+ assert local_client.chat_model == "gpt-x"
+ assert local_client.retrieve_model == "gpt-x"
+
+
+def test_concurrent_same_name_submits_store_unique_names(local_client,
+ sample_pdf,
+ monkeypatch):
+ """Name uniquing runs under the store lock at save time, so two
+ clients indexing the same filename concurrently cannot both store it
+ — a stored duplicate would shadow the older doc_id forever."""
+ import threading
+ import time as time_mod
+
+ def slow_flash(pdf, **kwargs):
+ time_mod.sleep(0.1) # both threads index before either saves
+ return {"structure": [{"title": "T", "start_index": 1,
+ "end_index": 2, "summary": "s", "nodes": []}]}
+
+ monkeypatch.setattr(pageindex.flash, "page_index_flash", slow_flash)
+ monkeypatch.setattr(pageindex.utils, "llm_completion",
+ lambda *a, **k: "d")
+ results = []
+ workers = [threading.Thread(
+ target=lambda: results.append(local_client.submit_document(sample_pdf)))
+ for _ in range(2)]
+ for worker in workers:
+ worker.start()
+ for worker in workers:
+ worker.join()
+ assert {r["name"] for r in results} == {"sample.pdf", "sample_1.pdf"}
+
+
+
+def test_format_tree_node_keeps_key_items():
+ """key_items from the merge optimization survive the get_tree formatter."""
+ from pageindex.local_api import _format_tree_node
+
+ node = {"title": "Chapter 1", "node_id": "0000", "start_index": 1,
+ "summary": "s",
+ "key_items": ["1.1 Alpha", "1.2 Beta", "1.3 Gamma"]}
+ out = _format_tree_node(node, node_summary=True)
+ assert out["key_items"] == ["1.1 Alpha", "1.2 Beta", "1.3 Gamma"]
+ assert "key_items" not in _format_tree_node(
+ {"title": "t", "node_id": "0001", "start_index": 1}, False)
diff --git a/tests/test_flash_extraction.py b/tests/test_flash_extraction.py
new file mode 100644
index 000000000..1e9dcb87d
--- /dev/null
+++ b/tests/test_flash_extraction.py
@@ -0,0 +1,170 @@
+"""Pins the pdfium 5.x text-extraction semantics the parser is calibrated to.
+
+pdfium split FPDFFont_GetFontName into GetBaseFontName (/BaseFont, per-face)
+and GetFamilyName (old family semantics); the parser uses the per-face names,
+which changes word joining and figure-label pickup. These sentinels come from
+a 4.30-vs-5.13 corpus A/B and fail if the semantics move again.
+"""
+from importlib.metadata import version
+from pathlib import Path
+
+import pytest
+
+PDF = Path(__file__).parent.parent / "examples" / "documents" / "earthmover.pdf"
+
+
+@pytest.mark.skipif(int(version("pypdfium2").split(".")[0]) < 5,
+ reason="extraction is pinned to pdfium 5.x font-name semantics")
+def test_page_text_pins_pdfium5_semantics():
+ from pageindex.flash.main import extract_toc
+
+ page7 = extract_toc(str(PDF))["page_texts"][6]
+ assert "p5\nEMD\n1.0" in page7 # figure axis label pdfium 4.x dropped
+ assert "break loop\n5: if lbp" in page7 # pseudocode lines no longer glued
+
+
+def test_page_mode_walk_uses_merged_surrogate_census():
+ """The page-mode unicode walk must consume the char census char_extract
+ built (astral chars merged to one entry at the high-surrogate slot).
+ Re-reading the textpage split them back into two lone-surrogate slots,
+ desynced the walk against their one-char cmap targets, and silently
+ dropped every patch on any page containing an astral char."""
+ from pageindex.flash.parser_pdfium_charlevel.unicode_apply import (
+ _apply_font_unicode)
+
+ astral = {"i": 0, "ch": "\U0001d44e", "is_gen": False} # slots 0-1 merged
+ unmapped = {"i": 2, "ch": "\x00", "is_gen": False} # PDFium found no unicode
+ raw_chars = [astral, unmapped]
+ show_codes = [(7, (5, 6), 100.0)]
+ map_cache = {7: (1, {5: "\U0001d44e", 6: "β"})}
+
+ # objects vs show ops count differs -> page mode.
+ _apply_font_unicode(raw_chars, [], show_codes, None, map_cache)
+
+ assert astral["ch"] == "\U0001d44e"
+ assert unmapped["ch"] == "β"
+
+
+def test_optimize_full_fails_fast_without_a_key(tmp_path, monkeypatch):
+ """optimize='full' runs LLM expand: with no key configured it must be
+ an instant, guided PageIndexAPIError — raised before any PDF work (a
+ bogus path proves the ordering) — while the LLM-free spellings and a
+ backend-carrying indexing scope stay untouched."""
+ from conftest import build_pdf
+ from pageindex import PageIndexAPIError
+ from pageindex.flash import page_index_flash
+ from pageindex.utils import _llm_backend
+ import litellm # noqa: F401 — first import may load a .env; delenv after it
+ monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+ monkeypatch.delenv("CHATGPT_API_KEY", raising=False)
+
+ with pytest.raises(PageIndexAPIError, match="optimize='merge'"):
+ page_index_flash(str(tmp_path / "missing.pdf"), summary=False)
+
+ pdf = tmp_path / "doc.pdf"
+ pdf.write_bytes(build_pdf(["1 Introduction", "Body text"]))
+ result = page_index_flash(str(pdf), summary=False, optimize="merge")
+ assert "structure" in result
+ result = page_index_flash(str(pdf), summary=False, optimize=False)
+ assert "structure" in result
+
+ token = _llm_backend.set({"api_key": "k"})
+ try:
+ with pytest.raises(FileNotFoundError):
+ page_index_flash(str(tmp_path / "missing.pdf"), summary=False)
+ finally:
+ _llm_backend.reset(token)
+
+
+def test_bootstrap_reimport_is_not_swallowed(monkeypatch):
+ # An unguarded caller script re-imported by a spawn worker must die loudly,
+ # not fall back to a silent full sequential rerun in every worker.
+ import multiprocessing
+ import sys
+
+ from pageindex.flash import parser_pdfium_parallel as mod
+
+ class BoomExecutor:
+ def __init__(self, *a, **k):
+ pass
+
+ def map(self, *a, **k):
+ raise RuntimeError("start a new process before bootstrapping")
+
+ def shutdown(self, *a, **k):
+ pass
+
+ monkeypatch.setattr(mod, "ProcessPoolExecutor", BoomExecutor)
+ cur = multiprocessing.current_process()
+
+ monkeypatch.setattr(cur, "_inheriting", True, raising=False)
+ with pytest.raises(RuntimeError):
+ mod.parse_charlevel_meta_parallel(str(PDF), workers=2, min_pages=1)
+ assert hasattr(sys.modules["__main__"], "__file__") # window restored on error
+
+ monkeypatch.delattr(cur, "_inheriting")
+ out, meta = mod.parse_charlevel_meta_parallel(str(PDF), workers=2, min_pages=1)
+ assert len(out) == len(meta) > 0 # normal failures still fall back sequentially
+
+
+def test_submit_document_refuses_during_bootstrap(tmp_path, monkeypatch):
+ import multiprocessing
+
+ from pageindex import PageIndexAPIError, PageIndexLocalClient
+
+ c = PageIndexLocalClient(storage_path=str(tmp_path))
+ monkeypatch.setattr(
+ multiprocessing.current_process(), "_inheriting", True, raising=False
+ )
+ with pytest.raises(PageIndexAPIError, match="__main__"):
+ c.submit_document("whatever.pdf")
+
+
+def test_unguarded_script_parses_parallel_without_reexecution(tmp_path):
+ # spawn workers must not re-run an unguarded caller script: one completion,
+ # no dead-worker noise (dying workers would trip the sequential fallback).
+ import os
+ import subprocess
+ import sys
+
+ marker = tmp_path / "runs.txt"
+ script = tmp_path / "unguarded.py"
+ script.write_text(
+ "from pageindex.flash.parser_pdfium_parallel import parse_charlevel_meta_parallel\n"
+ f"out, meta = parse_charlevel_meta_parallel({str(PDF)!r}, workers=2, min_pages=1)\n"
+ "assert len(out) == len(meta) > 0\n"
+ f"open({str(marker)!r}, 'a').write('ran\\n')\n"
+ )
+ env = {**os.environ, "PYTHONPATH": str(Path(__file__).parent.parent)}
+ res = subprocess.run(
+ [sys.executable, str(script)], capture_output=True, env=env, timeout=120
+ )
+ assert res.returncode == 0, res.stderr.decode()
+ assert marker.read_text() == "ran\n"
+ assert b"Traceback" not in res.stderr
+
+
+def test_optimize_wins_over_deprecated_optimize_expand(tmp_path, monkeypatch):
+ """Explicit optimize= beats optimize_expand; legacy True still honors it."""
+ from conftest import build_pdf
+ from pageindex.flash import page_index_flash
+ import litellm # noqa: F401 — first import may load a .env; delenv after it
+ monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+ monkeypatch.delenv("CHATGPT_API_KEY", raising=False)
+
+ pdf = tmp_path / "doc.pdf"
+ pdf.write_bytes(build_pdf(["1 Introduction", "Body text"]))
+ # resolved to "full" before the precedence fix (keyless → fail-fast)
+ with pytest.warns(DeprecationWarning):
+ result = page_index_flash(str(pdf), summary=False,
+ optimize="merge", optimize_expand=True)
+ assert "structure" in result
+ with pytest.warns(DeprecationWarning):
+ result = page_index_flash(str(pdf), summary=False,
+ optimize=True, optimize_expand=False)
+ assert "structure" in result
+ # optimize=None means unset ("full"), not off
+ from pageindex import PageIndexAPIError
+ with pytest.raises(PageIndexAPIError, match="optimize='merge'"):
+ page_index_flash(str(tmp_path / "missing.pdf"), summary=False,
+ optimize=None)
diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py
new file mode 100644
index 000000000..a2a3d591e
--- /dev/null
+++ b/tests/test_local_chat.py
@@ -0,0 +1,2040 @@
+"""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
+ if getattr(self, "emit_created", False):
+ from openai.types.responses import ResponseCreatedEvent
+ sequence += 1
+ yield ResponseCreatedEvent(
+ type="response.created", sequence_number=sequence,
+ response=Response(
+ id="resp_backend_turn", created_at=0.0, model="fake",
+ object="response", output=[], parallel_tool_calls=False,
+ tool_choice="auto", tools=[], status="in_progress"))
+ 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, backend=None):
+ state["protocols"].append((protocol, model_name))
+ state["backends"] = state.get("backends", []) + [backend]
+ 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="pip install openai-agents"):
+ 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="local-mode"):
+ cloud.chat_completions([{"role": "user", "content": "x"}],
+ reasoning_effort="low")
+ with pytest.raises(PageIndexAPIError, match="local-mode"):
+ cloud.chat_completions([{"role": "user", "content": "x"}],
+ extra_body={"service_tier": "auto"})
+ with pytest.raises(PageIndexAPIError, match="local-mode"):
+ cloud.chat_completions([{"role": "user", "content": "x"}], top_p=0.9)
+ with pytest.raises(PageIndexAPIError, match="local-mode"):
+ cloud.chat_completions([{"role": "user", "content": "x"}],
+ max_tokens=256)
+ with pytest.raises(PageIndexAPIError, match="local-mode"):
+ cloud.chat("x", reasoning_effort="low")
+ with pytest.raises(PageIndexAPIError, match="local-mode"):
+ cloud.chat_completions([{"role": "user", "content": "x"}],
+ backend={"api_key": "k"})
+ with pytest.raises(PageIndexAPIError, match="local-mode"):
+ cloud.chat_completions([{"role": "user", "content": "x"}],
+ extra_headers={"x-beta": "1"})
+ 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)
+
+
+@needs_agents
+def test_anthropic_routed_models_mark_managed_prefix_for_cache(
+ client, store_path, fake_model):
+ fake_model([[_msg_item("ok")]])
+ from pageindex.local_chat import _openai_agent
+ marked = {"cache_control_injection_points": [
+ {"location": "message", "role": "system"},
+ {"location": "message", "index": -1}]}
+ for name in ("anthropic/claude-x", "litellm/anthropic/claude-x",
+ "bedrock/us.anthropic.claude-sonnet-5",
+ "vertex_ai/claude-sonnet-4-5"):
+ agent = _openai_agent(client, "chat", name, "sys", None, None)
+ assert agent.model_settings.extra_args == marked
+ for name in ("gpt-5", "openai/Qwen/x", "litellm/groq/x",
+ "bedrock/meta.llama3-70b-instruct-v1:0",
+ "vertex_ai/gemini-2.5-pro"):
+ agent = _openai_agent(client, "chat", name, "sys", None, None)
+ assert agent.model_settings.extra_args is None
+
+
+@needs_agents
+def test_status_recorder_attaches_to_the_real_responses_model(monkeypatch):
+ # Guards the private-attribute chain the recorder rides
+ # (agent.model._client.responses.create): a vendor rename turns the
+ # recorder into a silent no-op and truncation reports as completion.
+ monkeypatch.setenv("OPENAI_API_KEY", "sk-test")
+ import openai
+ from agents.models.openai_responses import OpenAIResponsesModel
+ backend = openai.AsyncOpenAI()
+ model = OpenAIResponsesModel("gpt-test", openai_client=backend)
+ original = backend.responses.create
+ local_chat._record_response_status(types.SimpleNamespace(model=model), {})
+ assert backend.responses.create is not original
+ asyncio.run(backend.close())
+
+
+@needs_agents
+def test_cache_marker_reaches_the_anthropic_wire(client, store_path,
+ monkeypatch):
+ # End-to-end guard for the injection flag: through the real
+ # LitellmModel and litellm's request build, the marker must appear in
+ # the HTTP body — a regression in either vendor hop silently reverts
+ # anthropic-routed calls to full price.
+ pytest.importorskip("litellm")
+ from litellm.llms.custom_httpx.http_handler import (AsyncHTTPHandler,
+ HTTPHandler)
+ monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-ant-test")
+ captured = {}
+ reply = {"id": "msg_01", "type": "message", "role": "assistant",
+ "model": "claude-3-5-sonnet-20240620",
+ "content": [{"type": "text", "text": "ok"}],
+ "stop_reason": "end_turn", "stop_sequence": None,
+ "usage": {"input_tokens": 10, "output_tokens": 2}}
+
+ def _capture(url, kwargs):
+ body = kwargs.get("json")
+ if body is None and kwargs.get("data") is not None:
+ body = json.loads(kwargs["data"])
+ captured["url"] = str(url)
+ captured["body"] = body
+ return httpx.Response(200, json=reply,
+ request=httpx.Request("POST", str(url)))
+
+ async def fake_apost(self, url=None, *args, **kwargs):
+ return _capture(url, kwargs)
+
+ def fake_post(self, url=None, *args, **kwargs):
+ return _capture(url, kwargs)
+
+ monkeypatch.setattr(AsyncHTTPHandler, "post", fake_apost)
+ monkeypatch.setattr(HTTPHandler, "post", fake_post)
+ result = client.chat_completions(
+ "hi", model="anthropic/claude-3-5-sonnet-20240620")
+ assert "/v1/messages" in captured["url"]
+ assert '"cache_control"' in json.dumps(captured["body"])
+ # The OpenAI cache-routing hint must not leak here: LiteLLM plants
+ # extra_body as a literal field, and Anthropic rejects unknown fields.
+ assert "extra_body" not in json.dumps(captured["body"])
+ assert result["choices"][0]["message"]["content"] == "ok"
+
+
+# ── chat (front door) ──
+
+@needs_agents
+def test_chat_returns_answer_string(client, store_path, fake_model):
+ doc_id = seed_doc(store_path, "pi-a", "report.pdf")
+ fake = fake_model([
+ [_call_item("get_document", {"doc_name": "report.pdf"})],
+ [_msg_item("The answer")],
+ ])
+ assert client.chat("What status?", doc_id=doc_id) == "The answer"
+ first_item = fake.inputs[0][0]
+ assert "The user has specified document: report.pdf" in first_item["content"]
+
+
+@needs_agents
+def test_chat_stream_yields_text_chunks(client, store_path, fake_model):
+ fake_model([[_msg_item("The answer")]])
+ assert list(client.chat("q", stream=True)) == ["The ", "answer"]
+
+
+@needs_agents
+def test_chat_multi_turn_history(client, store_path, fake_model):
+ fake = fake_model([[_msg_item("Chapter 4 covers pears")]])
+ history = [
+ {"role": "user", "content": "What about chapter 3?"},
+ {"role": "assistant", "content": "Chapter 3 covers apples"},
+ {"role": "user", "content": "And chapter 4?"},
+ ]
+ assert client.chat(history) == "Chapter 4 covers pears"
+ assert fake.inputs[0][-3:] == history
+
+
+@needs_agents
+def test_chat_reasoning_effort_reaches_the_engine(client, store_path,
+ fake_model, monkeypatch):
+ """The business door's one thinking knob rides chat_completions'
+ channel unchanged; unset sends nothing."""
+ seen = {}
+ real = local_chat._openai_agent
+
+ def spy(*args, **kwargs):
+ agent = real(*args, **kwargs)
+ seen["settings"] = agent.model_settings
+ return agent
+
+ monkeypatch.setattr(local_chat, "_openai_agent", spy)
+ fake_model([[_msg_item("ok")]])
+ client.chat("q", reasoning_effort="low")
+ assert seen["settings"].extra_args["reasoning_effort"] == "low"
+ fake_model([[_msg_item("ok")]])
+ client.chat("q")
+ assert seen["settings"].extra_args is None
+
+
+def test_chat_cloud_unwraps_envelope(monkeypatch):
+ cloud = PageIndexCloudClient(api_key="pi-test-key")
+
+ def fake_cc(**kwargs):
+ assert kwargs["messages"] == [{"role": "user", "content": "q"}]
+ return {"choices": [{"message": {"role": "assistant",
+ "content": "cloud answer"}}]}
+
+ monkeypatch.setattr(cloud._api, "chat_completions", fake_cc)
+ assert cloud.chat("q") == "cloud answer"
+
+
+# ── 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._conversation_cache_key
+
+ def spy(model_name, instructions, doc_id, items):
+ key = real(model_name, instructions, doc_id, items)
+ keys.append(key)
+ return key
+
+ monkeypatch.setattr(local_chat, "_conversation_cache_key", 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
+
+ seed_doc(store_path, "pi-b", "contract.pdf")
+ fake_model([[_msg_item("f")]])
+ client.responses("What is the CAGR?", doc_id="pi-b")
+ assert keys[5] != keys[0] # same opener, different doc: no pooling
+
+
+@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_stream_opens_with_created(client, store_path, fake_model):
+ """N per-turn openings collapse to one response.created, not zero —
+ the logical stream must open with a response object carrying the same
+ id the terminal event reports, and the terminal envelope reports the
+ backend's tool-param echo, not assumed values."""
+ seed_doc(store_path, "pi-a", "report.pdf")
+ fake = fake_model([
+ [_call_item("get_document", {"doc_name": "report.pdf"})],
+ [_msg_item("The answer")],
+ ])
+ fake.emit_created = True # two turns emit two; one must pass through
+ events = list(client.responses("q", stream=True))
+ created = [e for e in events if e["type"] == "response.created"]
+ assert len(created) == 1 and events[0] is created[0]
+ assert events[0]["sequence_number"] == 1
+ terminal = events[-1]
+ assert terminal["type"] == "response.completed"
+ assert created[0]["response"]["id"] == terminal["response"]["id"]
+ assert (created[0]["response"]["created_at"]
+ == terminal["response"]["created_at"]) # one timestamp, not two
+ assert terminal["response"]["parallel_tool_calls"] is False # echo
+
+
+@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 backend=None: 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)
+ # every door that takes max_turns validates it, the runner config too
+ with pytest.raises(PageIndexAPIError, match="positive integer"):
+ client.anthropic_runner_config(model="claude-sonnet-4-5",
+ max_turns=-1)
+
+
+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)
+ # No transport echo attached here, so these are the fallbacks.
+ assert result["parallel_tool_calls"] is True
+ assert result["tool_choice"] == "auto"
+
+
+def test_sol_class_refusal_names_its_exits():
+ """The chatcmpl+tools-while-reasoning 400 is a lane problem, not a
+ retry problem — the wrapped error must name every exit."""
+ err = local_chat._model_backend_error(Exception(
+ "Error code: 400 - Function tools with reasoning_effort are not "
+ "supported for gpt-5.6-sol in /v1/chat/completions."))
+ assert "responses()" in str(err) and "litellm" in str(err)
+ assert "pass reasoning_effort" in str(err)
+ plain = local_chat._model_backend_error(Exception("rate limited"))
+ assert "responses()" not in str(plain)
+
+
+def test_conversation_cache_key_stable_per_conversation():
+ """Cache-routing key, sent as the OpenAI prompt_cache_key. A
+ conversation's continuations must share one key (same model /
+ instructions / doc targeting / 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_cache_key("m", "sys", "d1", turn1)
+ assert key == local_chat._conversation_cache_key(
+ "m", "sys", "d1", continuation)
+ assert key == local_chat._conversation_cache_key(
+ "m", "sys", ["d1"], turn1) # str and one-item list: same targeting
+ assert key != local_chat._conversation_cache_key(
+ "m", "sys", "d1", [{"role": "user", "content": "other"}])
+ assert key != local_chat._conversation_cache_key("m2", "sys", "d1", turn1)
+ assert key != local_chat._conversation_cache_key("m", "sys2", "d1", turn1)
+ assert key != local_chat._conversation_cache_key("m", "sys", "d2", turn1)
+ assert key != local_chat._conversation_cache_key("m", "sys", None, turn1)
+
+
+@needs_agents
+def test_agent_carries_prompt_cache_key_in_extra_body(monkeypatch):
+ """The key must reach the wire: openai-agents 0.20 dropped the
+ RunConfig.group_id -> prompt_cache_key derivation, so the agent's
+ ModelSettings.extra_body is the delivery channel. OpenAI destinations
+ only — prompt_cache_key is OpenAI's routing hint, and LiteLLM plants
+ extra_body as a literal field in other providers' bodies (Anthropic
+ rejects unknown fields); Claude routes keep their cache_control marker
+ in extra_args instead."""
+ monkeypatch.setattr(local_chat, "_openai_model", lambda *a: None)
+ monkeypatch.setattr("pageindex.integrations.openai_agents.build_openai_tools",
+ lambda *a, **k: [])
+ agent = local_chat._openai_agent(
+ None, "chat", "anthropic/claude-x", "sys", None, None,
+ doc_ids=None, cache_key="pageindex-k1")
+ settings = agent.model_settings
+ assert settings.extra_body is None
+ assert "cache_control_injection_points" in settings.extra_args
+ assert settings.extra_args["cache_control_injection_points"] == [
+ {"location": "message", "role": "system"},
+ {"location": "message", "index": -1}]
+ for name in ("gpt-test", "openai/gpt-test", "litellm/openai/gpt-test"):
+ agent = local_chat._openai_agent(
+ None, "chat", name, "sys", None, None,
+ doc_ids=None, cache_key="pageindex-k2")
+ assert agent.model_settings.extra_body == {
+ "prompt_cache_key": "pageindex-k2"}, name
+ assert agent.model_settings.extra_args is None
+
+
+@needs_agents
+def test_reasoning_passthrough_reaches_each_engine(monkeypatch):
+ """Per-door native reasoning, forwarded verbatim. The chat door's
+ effort rides extra_args — LiteLLM's own top-level kwarg on every
+ supported openai-agents version, and the channel admits values outside
+ the OpenAI enum ("none") — coexisting with the Claude cache marker.
+ The responses door's object rides ModelSettings.reasoning, which the
+ Responses model forwards verbatim. Unset sends nothing."""
+ pytest.importorskip("litellm")
+ monkeypatch.setattr(local_chat, "_openai_model", lambda *a: None)
+ monkeypatch.setattr("pageindex.integrations.openai_agents.build_openai_tools",
+ lambda *a, **k: [])
+ agent = local_chat._openai_agent(None, "chat", "gpt-test", "sys",
+ None, None, reasoning_effort="low")
+ assert agent.model_settings.extra_args == {"reasoning_effort": "low"}
+ assert agent.model_settings.reasoning is None
+ agent = local_chat._openai_agent(None, "chat", "anthropic/claude-x",
+ "sys", None, None,
+ reasoning_effort="none")
+ assert agent.model_settings.extra_args["reasoning_effort"] == "none"
+ assert "cache_control_injection_points" in agent.model_settings.extra_args
+ agent = local_chat._openai_agent(None, "responses", "gpt-test", "sys",
+ None, None,
+ reasoning={"effort": "low",
+ "summary": "auto"})
+ # ModelSettings coerces the dict into the typed openai Reasoning object.
+ assert agent.model_settings.reasoning.effort == "low"
+ assert agent.model_settings.reasoning.summary == "auto"
+ assert agent.model_settings.extra_args is None
+ agent = local_chat._openai_agent(None, "chat", "gpt-test", "sys",
+ None, None)
+ assert agent.model_settings.reasoning is None
+ assert agent.model_settings.extra_args is None
+
+
+@needs_agents
+def test_extra_body_passthrough_reaches_each_engine(monkeypatch):
+ """Caller extras merge last — over the cache key on OpenAI
+ destinations — and ride LiteLLM's own kwargs elsewhere, where
+ extra_body would plant literal fields providers reject."""
+ pytest.importorskip("litellm")
+ monkeypatch.setattr(local_chat, "_openai_model", lambda *a: None)
+ monkeypatch.setattr("pageindex.integrations.openai_agents.build_openai_tools",
+ lambda *a, **k: [])
+ agent = local_chat._openai_agent(None, "chat", "gpt-test", "sys",
+ None, None, cache_key="pageindex-k",
+ extra_body={"logit_bias": {"1": 5},
+ "prompt_cache_key": "mine"})
+ assert agent.model_settings.extra_body == {
+ "prompt_cache_key": "mine", "logit_bias": {"1": 5}}
+ assert agent.model_settings.extra_args is None
+ agent = local_chat._openai_agent(None, "chat", "anthropic/claude-x",
+ "sys", None, None,
+ reasoning_effort="low",
+ extra_body={"top_k": 20})
+ assert agent.model_settings.extra_body is None
+ assert agent.model_settings.extra_args["top_k"] == 20
+ assert agent.model_settings.extra_args["reasoning_effort"] == "low"
+ assert "cache_control_injection_points" in agent.model_settings.extra_args
+ agent = local_chat._openai_agent(None, "responses", "gpt-test", "sys",
+ None, None,
+ extra_body={"service_tier": "flex"})
+ assert agent.model_settings.extra_body == {"service_tier": "flex"}
+
+
+@needs_agents
+def test_sampling_knobs_ride_model_settings(client, store_path, fake_model,
+ monkeypatch):
+ """top_p/max_tokens ride ModelSettings fields — the one channel clean
+ on every lane (extra_body collides with LitellmModel's explicit
+ kwargs). responses' max_output_tokens is the same field's wire name,
+ echoed in the envelope."""
+ seed_doc(store_path, "pi-a", "report.pdf")
+ seen = {}
+ real = local_chat._openai_agent
+
+ def spy(*args, **kwargs):
+ agent = real(*args, **kwargs)
+ seen[args[1]] = agent.model_settings
+ return agent
+
+ monkeypatch.setattr(local_chat, "_openai_agent", spy)
+ fake_model([[_msg_item("ok")]])
+ client.chat_completions("q", top_p=0.9, max_tokens=256)
+ assert seen["chat"].top_p == 0.9
+ assert seen["chat"].max_tokens == 256
+ fake_model([[_msg_item("ok")]])
+ result = client.responses("q", max_output_tokens=321)
+ assert seen["responses"].max_tokens == 321
+ assert result["max_output_tokens"] == 321
+ fake_model([[_msg_item("ok")]])
+ assert client.responses("q")["max_output_tokens"] is None
+
+
+@needs_agents
+def test_responses_envelope_echoes_reasoning(client, store_path, fake_model):
+ seed_doc(store_path, "pi-a", "report.pdf")
+ fake_model([[_msg_item("ok")]])
+ result = client.responses("q", reasoning={"effort": "low"})
+ assert result["reasoning"] == {"effort": "low"}
+ fake_model([[_msg_item("ok")]])
+ assert client.responses("q")["reasoning"] is None
+
+
+@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_refused(client, store_path):
+ """doc_id=[] fails loud on every local surface, like cloud already
+ did: washing it to None would mean "everything", and the empty
+ allowlist meant "nothing" — an agent confidently reporting the
+ documents don't exist, with no signal the scope was empty."""
+ seed_doc(store_path, "pi-a", "report.pdf")
+ with pytest.raises(PageIndexAPIError, match="doc_id is empty"):
+ client.chat_completions("q", doc_id=[])
+ with pytest.raises(PageIndexAPIError, match="doc_id is empty"):
+ client.as_openai_tools(doc_id=[])
+ with pytest.raises(PageIndexAPIError, match="doc_id is empty"):
+ client.agent_instructions(doc_id=[])
+
+
+@needs_agents
+def test_openai_model_resolves_provider_prefixes():
+ """The chat lane is LiteLLM, full stop — model names mean what LiteLLM
+ says they mean, bare names are the openai/ shorthand, and routing
+ prefixes never leak as wire model names. responses stays OpenAI-SDK
+ native."""
+ pytest.importorskip("litellm")
+ from agents.extensions.models.litellm_model import LitellmModel
+ 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", "gpt-5.2")
+ assert isinstance(model, LitellmModel) and model.model == "openai/gpt-5.2"
+ model = local_chat._openai_model("chat", "openai/gpt-5.2")
+ assert isinstance(model, LitellmModel) and model.model == "openai/gpt-5.2"
+ model = local_chat._openai_model("responses", "gpt-5.2")
+ assert isinstance(model, OpenAIResponsesModel)
+ assert str(model.model) == "gpt-5.2"
+ model = local_chat._openai_model("responses", "openai/gpt-5.2")
+ assert isinstance(model, OpenAIResponsesModel)
+ assert str(model.model) == "gpt-5.2"
+ # litellm/ is routing grammar, not a provider: it strips before the
+ # provider-prefix guard, so an OpenAI model stays reachable.
+ model = local_chat._openai_model("responses", "litellm/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 == "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."""
+ pytest.importorskip("litellm") # first import may load a .env; delenv after
+ monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+ for name in ("gpt-4o", "openai/gpt-4o"):
+ with pytest.raises(PageIndexAPIError, match="OPENAI_API_KEY"):
+ local_chat._openai_model("chat", name)
+
+
+@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_responses_envelope_reports_backend_tool_params(client, store_path,
+ fake_model):
+ """tool_choice / parallel_tool_calls come from the backend's echo —
+ the request sends neither, so the envelope must not assume values."""
+ seed_doc(store_path, "pi-a", "report.pdf")
+ fake = fake_model([[_msg_item("ok")]])
+
+ async def create(*args, **kwargs):
+ return types.SimpleNamespace(status=None, tool_choice="none",
+ parallel_tool_calls=False)
+
+ fake._client = types.SimpleNamespace(
+ responses=types.SimpleNamespace(create=create))
+ result = client.responses("q")
+ assert result["tool_choice"] == "none"
+ assert result["parallel_tool_calls"] is False
+
+
+@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, backend=None: 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 backend=None: 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, monkeypatch):
+ """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. The pump
+ thread is tracked directly — a process-global thread count would be
+ flaky against litellm's background threads."""
+ import threading
+ seed_doc(store_path, "pi-a", "report.pdf")
+ pumps = []
+ real_thread = threading.Thread
+
+ class _Tracking(real_thread):
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+ if getattr(kwargs.get("target"), "__name__", "") == "pump":
+ pumps.append(self)
+
+ monkeypatch.setattr(threading, "Thread", _Tracking)
+ 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()
+ assert len(pumps) == 1
+ pumps[0].join(timeout=3.0)
+ assert not pumps[0].is_alive()
+ 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_thinking_passes_through(client, fake_anthropic):
+ """Anthropic-native thinking config, forwarded verbatim; unset sends
+ nothing so the backend default applies."""
+ calls = fake_anthropic([
+ _anthropic_message([{"type": "text", "text": "ok"}], "end_turn")])
+ client.messages("q", model="claude-sonnet-4-5",
+ thinking={"type": "adaptive"})
+ assert calls[0]["thinking"] == {"type": "adaptive"}
+ calls = fake_anthropic([
+ _anthropic_message([{"type": "text", "text": "ok"}], "end_turn")])
+ client.messages("q", model="claude-sonnet-4-5")
+ assert "thinking" not in calls[0]
+
+
+@needs_anthropic
+def test_messages_extra_body_merges_into_the_wire_body(client, fake_anthropic):
+ """The anthropic SDK merges extra_body keys into the request JSON —
+ asserted on the captured wire body, not the SDK call."""
+ calls = fake_anthropic([
+ _anthropic_message([{"type": "text", "text": "ok"}], "end_turn")])
+ client.messages("q", model="claude-sonnet-4-5",
+ extra_body={"service_tier": "auto"})
+ assert calls[0]["service_tier"] == "auto"
+ calls = fake_anthropic([
+ _anthropic_message([{"type": "text", "text": "ok"}], "end_turn")])
+ client.messages("q", model="claude-sonnet-4-5")
+ assert "service_tier" not in calls[0]
+
+
+@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)
+
+
+# ── backend + extra_headers: the chat doors ──
+
+@needs_agents
+def test_backend_connection_reaches_each_engine(monkeypatch):
+ """api_key/base_url ride each engine's client construction; the
+ LiteLLM lane's remaining keys ride its call kwargs; a backend key
+ satisfies the responses lane's missing-key check."""
+ pytest.importorskip("litellm")
+ monkeypatch.setattr("pageindex.integrations.openai_agents.build_openai_tools",
+ lambda *a, **k: [])
+ agent = local_chat._openai_agent(None, "chat", "anthropic/claude-x",
+ "sys", None, None,
+ backend={"api_key": "k1",
+ "api_base": "http://lb",
+ "api_version": "v9"})
+ assert agent.model.api_key == "k1"
+ assert agent.model.base_url == "http://lb"
+ assert agent.model_settings.extra_args["api_version"] == "v9"
+ assert "api_key" not in agent.model_settings.extra_args
+ monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+ agent = local_chat._openai_agent(None, "responses", "gpt-test", "sys",
+ None, None, backend={"api_key": "k2"})
+ assert agent.model._client.api_key == "k2"
+ # the LiteLLM endpoint spelling works on the SDK-constructed door too
+ agent = local_chat._openai_agent(None, "responses", "gpt-test", "sys",
+ None, None,
+ backend={"api_key": "k3",
+ "api_base": "http://rb"})
+ assert str(agent.model._client.base_url).rstrip("/") == "http://rb"
+ # both endpoint spellings on one merged dict: normalization keeps the
+ # later (per-call) key instead of the eager nested pop discarding it
+ agent = local_chat._openai_agent(
+ None, "chat", "anthropic/claude-x", "sys", None, None,
+ backend=local_chat._merged_backend(
+ types.SimpleNamespace(chat_backend={"base_url": "http://client"}),
+ {"api_base": "http://call"}))
+ assert agent.model.base_url == "http://call"
+
+
+@needs_anthropic
+def test_messages_top_level_cache_control(client, store_path, fake_anthropic):
+ """The moving breakpoint rides every request so each turn re-reads the
+ growing conversation; it stands down when the caller's own marks fill
+ the four-breakpoint budget (a fifth is a live-verified 400)."""
+ seed_doc(store_path, "pi-a", "report.pdf")
+ calls = fake_anthropic([
+ _anthropic_message([{"type": "text", "text": "ok"}], "end_turn")])
+ client.messages("q", model="claude-test", max_tokens=50)
+ assert calls[0]["cache_control"] == {"type": "ephemeral"}
+ marked = [{"type": "text", "text": f"b{i}",
+ "cache_control": {"type": "ephemeral"}} for i in range(3)]
+ calls = fake_anthropic([
+ _anthropic_message([{"type": "text", "text": "ok"}], "end_turn")])
+ client.messages("q", model="claude-test", max_tokens=50, system=marked)
+ assert "cache_control" not in calls[0]
+
+
+def test_merged_backend_precedence():
+ from types import SimpleNamespace
+ stub = SimpleNamespace(chat_backend={"api_key": "a", "api_version": "v1"})
+ assert local_chat._merged_backend(stub, {"api_key": "b"}) == {
+ "api_key": "b", "api_version": "v1"}
+ assert local_chat._merged_backend(SimpleNamespace(), None) is None
+
+
+@needs_anthropic
+def test_messages_backend_merges_and_reaches_the_client(client, fake_anthropic,
+ monkeypatch):
+ real = local_chat._anthropic_client({"api_key": "kk",
+ "base_url": "http://x"})
+ assert real.api_key == "kk"
+ assert str(real.base_url).rstrip("/") == "http://x"
+ real = local_chat._anthropic_client({"api_key": "kk",
+ "api_base": "http://y"})
+ assert str(real.base_url).rstrip("/") == "http://y"
+
+ calls = fake_anthropic([
+ _anthropic_message([{"type": "text", "text": "ok"}], "end_turn")])
+ fixture_client = local_chat._anthropic_client
+ seen = {}
+ monkeypatch.setattr(
+ local_chat, "_anthropic_client",
+ lambda backend=None: (seen.setdefault("backend", backend),
+ fixture_client())[1])
+ client.chat_backend = {"base_url": "http://cb"}
+ client.messages("q", model="claude-sonnet-4-5", backend={"api_key": "z"})
+ assert seen["backend"] == {"base_url": "http://cb", "api_key": "z"}
+
+
+@needs_anthropic
+def test_messages_bad_backend_wraps_like_the_other_doors():
+ with pytest.raises(PageIndexAPIError,
+ match="Anthropic backend is not configured"):
+ local_chat._anthropic_client({"no_such_param": 1})
+
+
+@needs_agents
+def test_extra_headers_ride_model_settings(monkeypatch):
+ """Both openai-agents doors merge ModelSettings.extra_headers into
+ their requests (wire-probed: LiteLLM's chatcmpl adapters forward
+ custom headers; its anthropic adapter owns anthropic-beta only)."""
+ monkeypatch.setattr(local_chat, "_openai_model", lambda *a: None)
+ monkeypatch.setattr("pageindex.integrations.openai_agents.build_openai_tools",
+ lambda *a, **k: [])
+ agent = local_chat._openai_agent(None, "chat", "gpt-test", "sys",
+ None, None,
+ extra_headers={"x-beta": "1"})
+ assert agent.model_settings.extra_headers == {"x-beta": "1"}
+ agent = local_chat._openai_agent(None, "responses", "gpt-test", "sys",
+ None, None,
+ extra_headers={"x-beta": "2"})
+ assert agent.model_settings.extra_headers == {"x-beta": "2"}
+ agent = local_chat._openai_agent(None, "chat", "gpt-test", "sys",
+ None, None)
+ assert agent.model_settings.extra_headers is None
+
+
+@needs_anthropic
+def test_messages_extra_headers_reach_the_wire(client, monkeypatch):
+ import anthropic
+ seen = {}
+
+ def handler(request):
+ seen["beta"] = request.headers.get("anthropic-beta")
+ return httpx.Response(200, json=_anthropic_message(
+ [{"type": "text", "text": "ok"}], "end_turn"))
+
+ fake = anthropic.Anthropic(api_key="t", http_client=httpx.Client(
+ transport=httpx.MockTransport(handler)))
+ monkeypatch.setattr(local_chat, "_anthropic_client",
+ lambda backend=None: fake)
+ client.messages("q", model="claude-sonnet-4-5",
+ extra_headers={"anthropic-beta": "context-1m-2025"})
+ assert seen["beta"] == "context-1m-2025"
+
+
+@needs_agents
+def test_chat_model_settings_request_stream_usage(monkeypatch):
+ """Without include_usage the streamed run carries no usage at all and
+ the terminal chunk reports zeros (agents forwards it as
+ stream_options only on streaming calls)."""
+ monkeypatch.setattr(local_chat, "_openai_model", lambda *a: None)
+ agent = local_chat._openai_agent(None, "chat", "gpt-test", "sys",
+ None, None)
+ assert agent.model_settings.include_usage is True
+
+
+@needs_anthropic
+def test_messages_default_max_tokens_clears_thinking_budget(client, fake_anthropic):
+ calls = fake_anthropic([
+ _anthropic_message([{"type": "text", "text": "a"}], "end_turn"),
+ ])
+ client.messages("q", model="claude-test",
+ thinking={"type": "enabled", "budget_tokens": 10000})
+ assert calls[0]["max_tokens"] == 10000 + 8192
+ assert calls[0]["thinking"] == {"type": "enabled",
+ "budget_tokens": 10000}
+ calls = fake_anthropic([ # fresh fake: each run closes its client
+ _anthropic_message([{"type": "text", "text": "b"}], "end_turn"),
+ ])
+ client.messages("q", model="claude-test", max_tokens=11000,
+ thinking={"type": "enabled", "budget_tokens": 10000})
+ assert calls[0]["max_tokens"] == 11000 # explicit value passes through
+
+
+def test_record_chat_finish_records_and_delegates():
+ recorded = {}
+ closed = {"n": 0}
+
+ class Stream:
+ def __init__(self):
+ self.chunks = [
+ types.SimpleNamespace(choices=[]),
+ types.SimpleNamespace(choices=[types.SimpleNamespace(
+ finish_reason="content_filter")]),
+ ]
+
+ def __aiter__(self):
+ return self
+
+ async def __anext__(self):
+ if not self.chunks:
+ raise StopAsyncIteration
+ return self.chunks.pop(0)
+
+ async def aclose(self):
+ closed["n"] += 1
+
+ async def fetch(*args, **kwargs):
+ return "shell", Stream()
+
+ model = types.SimpleNamespace(_fetch_response=fetch)
+ local_chat._record_chat_finish(types.SimpleNamespace(model=model),
+ recorded)
+
+ async def drive():
+ shell, tee = await model._fetch_response()
+ assert shell == "shell"
+ async for _chunk in tee:
+ pass
+ await tee.aclose()
+
+ asyncio.run(drive())
+ assert recorded == {"finish_reason": "content_filter"}
+ assert closed["n"] == 1
+ # A model without the seam: silently a no-op.
+ local_chat._record_chat_finish(
+ types.SimpleNamespace(model=types.SimpleNamespace()), {})
+
+
+@needs_agents
+def test_chat_completions_reports_native_finish_reason(client, store_path,
+ monkeypatch):
+ seed_doc(store_path, "pi-a", "report.pdf")
+
+ class TruncatingModel(FakeModel):
+ async def _fetch_response(self, *args, **kwargs):
+ return types.SimpleNamespace(choices=[
+ types.SimpleNamespace(finish_reason="length")])
+
+ async def get_response(self, *args, **kwargs):
+ await self._fetch_response()
+ return await super().get_response(*args, **kwargs)
+
+ async def stream_response(self, *args, **kwargs):
+ await self._fetch_response()
+ async for event in super().stream_response(*args, **kwargs):
+ yield event
+
+ fake = TruncatingModel([[_msg_item("cut ")], [_msg_item("cut ")]])
+ monkeypatch.setattr(local_chat, "_openai_model", lambda *a: fake)
+ result = client.chat_completions("q")
+ assert result["choices"][0]["finish_reason"] == "length"
+ chunks = list(client.chat_completions("q", stream=True,
+ stream_metadata=True))
+ assert chunks[-2]["choices"][0]["finish_reason"] == "length"
+
+
+@needs_agents
+def test_chat_gate_honors_litellm_routing_and_custom_providers(monkeypatch):
+ """Mirrors the indexing lane: an explicit litellm/ prefix skips the env
+ key pre-check, and custom_provider_map providers pass the allowlist;
+ a name LiteLLM cannot route is still refused up front."""
+ pytest.importorskip("litellm")
+ import litellm # first import may load a .env; delenv after it
+ from agents.extensions.models.litellm_model import LitellmModel
+
+ monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+ model = local_chat._openai_model("chat", "litellm/gpt-4o")
+ assert isinstance(model, LitellmModel) and model.model == "openai/gpt-4o"
+
+ monkeypatch.setattr(litellm, "custom_provider_map",
+ [{"provider": "my-llm", "custom_handler": object()}])
+ model = local_chat._openai_model("chat", "my-llm/model-a")
+ assert isinstance(model, LitellmModel) and model.model == "my-llm/model-a"
+
+ with pytest.raises(PageIndexAPIError, match="not a LiteLLM provider"):
+ local_chat._openai_model("chat", "Qwen/my-model")
+
+
+@needs_agents
+def test_litellm_model_still_has_the_fetch_response_seam():
+ # Guards the private seam _record_chat_finish rides (LitellmModel
+ # ._fetch_response): a vendor rename turns the recorder into a silent
+ # no-op and every truncated turn reports finish_reason "stop".
+ pytest.importorskip("litellm")
+ from agents.extensions.models.litellm_model import LitellmModel
+
+ assert hasattr(LitellmModel, "_fetch_response")
+
+
+def test_openai_protocol_predicate_follows_litellm_routing():
+ pytest.importorskip("litellm")
+ for name in ("gpt-5", "openai/gpt-4o", "litellm/gpt-4o",
+ "azure/gpt-4o", "openrouter/openai/gpt-4o",
+ "deepseek/deepseek-chat", "groq/llama-3.3-70b-versatile",
+ "xai/grok-3"):
+ assert local_chat._openai_protocol(name), name
+ for name in ("anthropic/claude-sonnet-4-5", "gemini/gemini-2.5-pro",
+ "bedrock/us.anthropic.claude-sonnet-5",
+ "vertex_ai/claude-sonnet-4-5"):
+ assert not local_chat._openai_protocol(name), name
+
+
+@needs_agents
+def test_chat_backend_without_key_stands_aside_like_index_lane(monkeypatch):
+ """Any non-empty backend dict suppresses the key pre-check (utils rule)."""
+ pytest.importorskip("litellm")
+ import litellm # noqa: F401 — first import may load a .env; delenv after it
+ from agents.extensions.models.litellm_model import LitellmModel
+
+ monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+ model = local_chat._openai_model(
+ "chat", "gpt-test", {"base_url": "http://localhost:9"})
+ assert isinstance(model, LitellmModel)
+
+
+@needs_agents
+def test_responses_model_marks_caller_owned_transport(monkeypatch):
+ monkeypatch.setenv("OPENAI_API_KEY", "sk-test")
+ shared = httpx.AsyncClient()
+ caller = local_chat._openai_model("responses", "gpt-test",
+ {"http_client": shared})
+ assert caller._client._pageindex_caller_http is True
+ owned = local_chat._openai_model("responses", "gpt-test")
+ assert owned._client._pageindex_caller_http is False
+
+ async def run():
+ await local_chat._aclose_backend(types.SimpleNamespace(model=caller))
+ assert not shared.is_closed # caller-owned transport survives
+ await local_chat._aclose_backend(types.SimpleNamespace(model=owned))
+ await shared.aclose()
+
+ asyncio.run(run())
+
+
+@needs_anthropic
+def test_messages_keeps_caller_owned_http_client_open(client):
+ body = _anthropic_message([{"type": "text", "text": "a"}], "end_turn")
+ shared = httpx.Client(transport=httpx.MockTransport(
+ lambda request: httpx.Response(200, json=body)))
+ out = client.messages("q", model="claude-test",
+ backend={"api_key": "t", "http_client": shared})
+ assert out["content"][0]["text"] == "a"
+ assert not shared.is_closed
+ client.messages("q", model="claude-test",
+ backend={"api_key": "t", "http_client": shared})
+ shared.close()
+
+
+@needs_anthropic
+def test_messages_without_credentials_raises_contract_error(client,
+ monkeypatch):
+ monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False)
+ monkeypatch.delenv("ANTHROPIC_AUTH_TOKEN", raising=False)
+ with pytest.raises(PageIndexAPIError,
+ match="Anthropic backend is not configured"):
+ client.messages("q", model="claude-test")
diff --git a/tests/test_package_surface.py b/tests/test_package_surface.py
index e10c3e8c5..8f68cea70 100644
--- a/tests/test_package_surface.py
+++ b/tests/test_package_surface.py
@@ -1,4 +1,5 @@
"""What `pip install pageindex` exposes: 0.2.8 helper compat and import cost."""
+import os
import subprocess
import sys
@@ -60,3 +61,62 @@ 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"
+
+
+def test_import_leaves_litellm_env_untouched(tmp_path):
+ """Importing the package must not configure litellm for the host
+ process; constructing a local client (which will use litellm) does."""
+ env = {k: v for k, v in os.environ.items()
+ if k != "LITELLM_LOCAL_MODEL_COST_MAP"}
+ probe = (
+ "import os, pageindex\n"
+ "assert 'LITELLM_LOCAL_MODEL_COST_MAP' not in os.environ, "
+ "'stamped at import'\n"
+ f"pageindex.PageIndexLocalClient(storage_path={str(tmp_path / 's')!r})\n"
+ "assert os.environ['LITELLM_LOCAL_MODEL_COST_MAP'] == 'True'\n"
+ "print('ok')\n"
+ )
+ out = subprocess.run([sys.executable, "-c", probe], env=env,
+ capture_output=True, text=True, check=True)
+ assert out.stdout.strip() == "ok"
diff --git a/tests/test_page_index_md.py b/tests/test_page_index_md.py
index 12aa90890..37b8c2a88 100644
--- a/tests/test_page_index_md.py
+++ b/tests/test_page_index_md.py
@@ -19,5 +19,33 @@ def test_skips_bold_heading_with_only_whitespace(self):
)
+class MarkdownCliTest(unittest.TestCase):
+ def test_md_cli_runs_without_llm_or_key(self):
+ """--md_path with no flags makes zero LLM calls: config.yaml's PDF
+ summary default must not leak in, so the run completes without any
+ provider key and writes the structure file."""
+ import json
+ import os
+ import subprocess
+ import sys
+ import tempfile
+ from pathlib import Path
+
+ script = Path(__file__).resolve().parent.parent / "run_pageindex.py"
+ with tempfile.TemporaryDirectory() as tmp:
+ md = Path(tmp) / "notes.md"
+ md.write_text("# Title\n\nIntro.\n\n## Section\n\nBody.\n")
+ env = {k: v for k, v in os.environ.items()
+ if k not in ("OPENAI_API_KEY", "CHATGPT_API_KEY")}
+ env["PYTHONPATH"] = str(script.parent)
+ res = subprocess.run(
+ [sys.executable, str(script), "--md_path", str(md)],
+ capture_output=True, cwd=tmp, env=env, timeout=180)
+ self.assertEqual(res.returncode, 0, res.stderr.decode())
+ out = Path(tmp) / "results" / "notes_structure.json"
+ self.assertTrue(out.exists(), res.stdout.decode())
+ json.loads(out.read_text())
+
+
if __name__ == "__main__":
unittest.main()