From 6b343f55eeb075082f09a7547843922e6de364e0 Mon Sep 17 00:00:00 2001 From: Ray Date: Sun, 9 Aug 2026 19:21:39 +0800 Subject: [PATCH 001/137] =?UTF-8?q?feat:=20agent=20tools=20=E2=80=94=20the?= =?UTF-8?q?=20cloud=20MCP=20tool=20contract=20on=20the=20client?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Four new client methods make PageIndex documents available to agent frameworks, in both modes, with the mode decided solely by the client constructor: - agent_tools(): plain functions (browse_documents, get_document, get_document_structure, get_page_content) matching the PageIndex cloud MCP server's tools/list โ€” same names, schemas, descriptions, and JSON response envelopes โ€” so agent prompts port unchanged between the cloud MCP connection and these in-process tools. Tools never raise; errors come back in the same envelope. remove_document ships behind include_management=False. - as_openai_tools(): the same tools wrapped for the OpenAI Agents SDK. - as_claude_mcp(): one mcp_servers entry for the Claude Agent SDK โ€” cloud clients get the remote MCP config (the framework connects to api.pageindex.ai/mcp and discovers the full cloud tool set), local clients get an in-process SDK MCP server. - agent_instructions(doc_id=None): orchestration guidance for the agent's system prompt; doc_id (same shape as chat_completions) appends the target documents. submit_document() gains wait=True: poll get_document status until completed, raise on failed or after 30 minutes โ€” the manual polling loop every cloud caller writes today spins forever on a failed document. Neither framework becomes a dependency: imports happen at call time with actionable errors, and pageindex[openai] / pageindex[claude] extras are floor-only pins. tests/data/cloud_mcp_contract.json freezes the tool contract; a parity test guards against drift. 36 new tests (95 total), plus a live OpenAI Agents SDK run over a seeded local store verifying the structure-first navigation flow end to end. --- README.md | 59 +- examples/agentic_vectorless_rag_demo.py | 58 +- pageindex/agent_tools.py | 1326 ++++++++++++++++++++ pageindex/client.py | 148 ++- pageindex/integrations/__init__.py | 5 + pageindex/integrations/claude_agent_sdk.py | 67 + pageindex/integrations/openai_agents.py | 36 + pageindex/mcp_bridge.py | 181 +++ pyproject.toml | 8 + tests/data/cloud_mcp_contract.json | 197 +++ tests/test_agent_tools.py | 879 +++++++++++++ 11 files changed, 2915 insertions(+), 49 deletions(-) create mode 100644 pageindex/agent_tools.py create mode 100644 pageindex/integrations/__init__.py create mode 100644 pageindex/integrations/claude_agent_sdk.py create mode 100644 pageindex/integrations/openai_agents.py create mode 100644 pageindex/mcp_bridge.py create mode 100644 tests/data/cloud_mcp_contract.json create mode 100644 tests/test_agent_tools.py diff --git a/README.md b/README.md index 5ce0ca5e6..5dbafc141 100644 --- a/README.md +++ b/README.md @@ -207,13 +207,68 @@ python3 run_pageindex.py --md_path /path/to/your/document.md > > Add `--optimize` to refine the tree structure for more efficient retrieval (with an LLM expansion pass). +## ๐Ÿ Python SDK: Cloud & Local + +The `pageindex` package on PyPI is the Python SDK for the [PageIndex API](https://docs.pageindex.ai) โ€” and the same client now also runs fully **locally**, powered by this repo's indexing pipeline (including Flash). + +```bash +pip3 install --upgrade pageindex # local mode ships in pageindex >= 0.2.9; earlier versions are cloud-only +``` + +```python +from pageindex import PageIndexClient + +client = PageIndexClient(api_key="YOUR_PAGEINDEX_API_KEY") # cloud: managed OCR, tree building, retrieval +client = PageIndexClient() # local: same methods on your machine, using your LLM key (e.g. OPENAI_API_KEY) + +doc_id = client.submit_document("doc.pdf")["doc_id"] # local mode blocks until indexing finishes +doc_id = client.submit_document("doc.pdf", mode="flash")["doc_id"] # local mode with PageIndex Flash + +tree = client.get_tree(doc_id, node_summary=True)["result"] + +answer = client.chat_completions( + messages=[{"role": "user", "content": "Summarize the key findings"}], + doc_id=doc_id, +)["choices"][0]["message"]["content"] +``` + +Local documents are stored as plain JSON under `./.pageindex` (configurable via `storage_path`). Local mode supports PDFs; folders, `beta_headers`, `enable_citations`, and the deprecated retrieval API (`submit_query`/`get_retrieval`) remain cloud-only โ€” each method's docstring spells out the differences. To pin the mode at construction instead of inferring it from `api_key`, use `PageIndexCloudClient` (fails without a real key) or `PageIndexLocalClient` (has no key parameter). + +### ๐Ÿค– Agent integration + +The client exposes its documents as **agent tools**, following one rule: **cloud clients always serve the live tool set of the [PageIndex MCP server](https://docs.pageindex.ai/mcp)** (search, folders, images โ€” as enabled for your key, discovered dynamically; management tools like delete/upload sit behind `include_management=True` or the framework's approval layer), while local clients serve the same contract's built-in navigation subset (`browse_documents`, `get_document`, `get_document_structure`, `get_page_content`). Tool names and schemas are shared, so agent prompts port unchanged, and switching local โ†” cloud is just the client constructor line: + +```python +client = PageIndexLocalClient() # or PageIndexCloudClient(api_key=...) +client.submit_document("doc.pdf", wait=True) # wait=True: return once the doc is ready (both modes) + +# OpenAI Agents SDK (pip install "pageindex[openai]") +agent = Agent( + name="PageIndex", + instructions=client.agent_instructions(), # retrieval playbook for the agent's system prompt + tools=client.as_openai_tools(), # local: in-process tools; cloud: the full cloud MCP tool set (any model backend) +) # cloud + OpenAI models: hosted=True runs tool calls server-side (fastest) + +# Claude Agent SDK (pip install "pageindex[claude]") +options = ClaudeAgentOptions( + system_prompt=client.agent_instructions(), + mcp_servers={"pageindex": client.as_claude_mcp()}, # local: in-process server; cloud: connects to api.pageindex.ai/mcp + allowed_tools=["mcp__pageindex__*"], +) + +# Any other framework: plain functions, wrap with your framework's one-liner +tools = client.agent_tools() # local: built-in tools; cloud: full live tool set over MCP + # e.g. [StructuredTool.from_function(f) for f in tools] +``` + +Neither framework is a required dependency โ€” each is imported only when its method is called. Claude Code / Cursor and other MCP hosts connect to cloud documents via the hosted MCP server directly (no SDK needed); see the [MCP docs](https://docs.pageindex.ai/mcp). ## ๐Ÿš€ 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 +# Install with the OpenAI Agents SDK extra +pip3 install "pageindex[openai]" # 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..e8ed4a50a 100644 --- a/examples/agentic_vectorless_rag_demo.py +++ b/examples/agentic_vectorless_rag_demo.py @@ -6,20 +6,21 @@ chunking, PageIndex builds a hierarchical tree index and uses agentic LLM reasoning for human-like, context-aware retrieval. -Agent tools: - - get_document() โ€” document metadata (status, page count, etc.) - - get_document_structure() โ€” tree structure index of a document - - get_page_content() โ€” retrieve text content of specific pages +The agent tools come straight from the SDK โ€” ``client.as_openai_tools()`` +exposes the PageIndex tool contract (browse_documents, get_document, +get_document_structure, get_page_content) and ``client.agent_instructions()`` +provides the retrieval playbook, so the whole agent is a few lines. Swap +``PageIndexLocalClient()`` for ``PageIndexCloudClient(api_key=...)`` and the +same code runs against the cloud. Steps: 1 โ€” Index a PDF locally and view its tree structure index 2 โ€” View document metadata 3 โ€” Ask a question (agent reasons over the index and auto-calls tools) -Requirements: pip install openai-agents; OPENAI_API_KEY in the environment. +Requirements: pip install "pageindex[openai]"; OPENAI_API_KEY in the environment. """ import sys -import json import asyncio import concurrent.futures from pathlib import Path @@ -27,12 +28,12 @@ sys.path.insert(0, str(Path(__file__).parent.parent)) -from agents import Agent, Runner, function_tool, set_tracing_disabled +from agents import Agent, Runner, set_tracing_disabled from agents.model_settings import ModelSettings 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,47 +42,18 @@ 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), + instructions=client.agent_instructions(doc_id=doc_id), + tools=client.as_openai_tools(), + model=client.retrieve_model, # model_settings=ModelSettings(reasoning={"effort": "low", "summary": "auto"}), # Uncomment to enable reasoning ) @@ -152,7 +124,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) @@ -166,7 +138,7 @@ async def _run(): 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/agent_tools.py b/pageindex/agent_tools.py new file mode 100644 index 000000000..10f0340dd --- /dev/null +++ b/pageindex/agent_tools.py @@ -0,0 +1,1326 @@ +"""Agent tools: the cloud MCP tool contract, executed against a PageIndexClient. + +Tool names, input schemas, and descriptions match the PageIndex cloud MCP +server, so agent prompts work unchanged 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). + +Tools never raise: every outcome, including errors, is returned as the same +JSON envelope the cloud emits ({"success": true, ...} / {"error": ...}). +""" +from __future__ import annotations + +import copy +import difflib +import json +import re +import time +from typing import Any, Callable, Optional + +from .errors import PageIndexAPIError + +TOOL_RESPONSE_CHAR_LIMIT = 100_000 +STRUCTURE_FIRST_PAGE_THRESHOLD = 20 + +_CHAR_BUDGET = int(TOOL_RESPONSE_CHAR_LIMIT * 0.95) +_PAGES_SPEC_RE = re.compile(r"^(\d+(-\d+)?)(,\s*\d+(-\d+)?)*$") +_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, + "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": { + "type": ["string", "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": { + "type": ["string", "null"], + "description": _FOLDER_ID_DISAMBIGUATOR_DESCRIPTION, + }, + "part": { + "type": "integer", + "minimum": 1, + "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": { + "type": ["string", "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": { + "type": ["string", "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, indent=2, ensure_ascii=False) + + +# โ”€โ”€ document listing / name resolution โ”€โ”€ + +def _all_documents(client) -> list[dict[str, Any]]: + """Every document the client can list, newest first (both modes list + newest-first; paging preserves that order).""" + documents: list[dict[str, Any]] = [] + offset = 0 + while True: + page = client.list_documents(limit=100, offset=offset) + batch = page.get("documents") or [] + documents.extend(batch) + offset += 100 + if not batch or offset >= page.get("total", 0): + 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 _resolve_document( + client, doc_name: str, +) -> "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).""" + documents = _all_documents(client) + matches = [doc for doc in documents if doc.get("name") == doc_name] + if matches: + return max(matches, key=lambda d: d.get("createdAt") or ""), None + names = [str(doc.get("name")) for doc in documents if doc.get("name")] + similar = difflib.get_close_matches(doc_name, names, n=_SIMILAR_NAMES_LIMIT, + cutoff=0.5) + message = ( + "Document not found. Did you mean: " + + ", ".join(f'"{name}"' for name in similar) + "?" + if similar else "Document not found or you do not have access to it" + ) + return None, _failure( + message, + {"doc_name": doc_name, "similar_files": similar}, + { + "summary": "The requested document does not exist or is not accessible", + "options": [ + "Verify the document name is correct", + "Use browse_documents() to see your recent documents", + "Check if the document was deleted", + ], + }, + "NOT_FOUND", + ) + + +def _refetch_entry(client, doc_id: str) -> Optional[dict[str, Any]]: + try: + return client.get_document(doc_id) + except PageIndexAPIError: + return None + + +def _await_completion(client, entry: dict[str, Any], wait: bool) -> dict[str, Any]: + """Re-poll a processing document for up to 3 minutes when wait is set.""" + doc_id = entry.get("id") + if not wait or not doc_id or entry.get("status") in ("completed", "failed"): + return entry + deadline = time.monotonic() + _TOOL_WAIT_TIMEOUT + current = entry + while time.monotonic() < deadline: + time.sleep(_TOOL_WAIT_INTERVAL) + refreshed = _refetch_entry(client, doc_id) + if refreshed is None: + return current + refreshed.setdefault("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 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 available here โ€” omit {param}.", + None, + { + "summary": "This library has no folders", + "options": ["Retry the call without a folder_id", + "Use browse_documents() to list the library root"], + }, + "INVALID_INPUT", + ) + + +# โ”€โ”€ page spec handling โ”€โ”€ + +def _parse_page_spec( + pages: str, doc_name: str, +) -> "tuple[Optional[list[int]], Optional[_ToolResult]]": + """Expand '1-3,7' into a sorted, deduplicated page list, or an error.""" + invalid = _failure( + "Invalid page specification format", + {"doc_name": doc_name}, + { + "summary": "Failed to parse the pages parameter", + "options": [ + 'Use valid formats: "5", "3,7,10", "5-10", or "1-3,7,9-12"', + "Ensure page numbers are positive integers", + ], + }, + "INVALID_INPUT", + ) + if not isinstance(pages, str) or not _PAGES_SPEC_RE.match(pages.strip()): + return None, invalid + expanded: set[int] = set() + for part in pages.split(","): + part = part.strip() + if "-" in part: + start, end = (int(x) for x in part.split("-", 1)) + if start > end: + return None, invalid + expanded.update(range(start, end + 1)) + else: + expanded.add(int(part)) + if any(page < 1 for page in expanded): + return None, _failure( + "Invalid page numbers. Page numbers must be positive integers", + {"doc_name": doc_name}, + { + "summary": "Invalid page numbers provided", + "options": [ + "Page numbers must be positive integers (>= 1)", + "Check the page specification format", + ], + }, + "INVALID_INPUT", + ) + return sorted(expanded), None + + +def _format_page_spec(pages: list[int]) -> str: + """Compress [1,2,3,5] into '1-3,5'.""" + if not pages: + return "" + ordered = sorted(set(pages)) + ranges = [] + start = prev = ordered[0] + for page in ordered[1:]: + if page == prev + 1: + prev = page + continue + ranges.append(f"{start}" if start == prev else f"{start}-{prev}") + start = prev = page + ranges.append(f"{start}" if start == prev else f"{start}-{prev}") + return ",".join(ranges) + + +# โ”€โ”€ structure formatting / splitting โ”€โ”€ + +_STRUCTURE_KEY_ORDER = ("title", "node_id", "start_index", "end_index", + "page_index", "prefix_summary", "summary", "nodes") + + +def _format_structure(node: Any) -> Any: + """Drop node text and normalize key order, recursively.""" + if isinstance(node, list): + return [_format_structure(item) for item in node] + if isinstance(node, dict): + stripped = {key: value for key, value in node.items() if key != "text"} + if "nodes" in stripped: + stripped["nodes"] = _format_structure(stripped["nodes"]) + ordered = {key: stripped[key] for key in _STRUCTURE_KEY_ORDER + if key in stripped} + ordered.update({key: value for key, value in stripped.items() + if key not in ordered}) + return ordered + return node + + +def _serialized_size(value: Any) -> int: + return len(json.dumps(value, ensure_ascii=False)) + + +def _split_structure(structure: Any, budget: int) -> list[Any]: + """Split a formatted structure into chunks of at most ~budget serialized + chars. The paginated response shape matches the cloud tool; chunk + boundaries are implementation-defined.""" + 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 if len(group) > 1 else group[0]) + group, group_size = [], 0 + chunks.extend(_split_oversized_node(node, budget)) + continue + if group and group_size + size > budget: + chunks.append(group if len(group) > 1 else group[0]) + group, group_size = [], 0 + group.append(node) + group_size += size + if group: + chunks.append(group if len(group) > 1 else group[0]) + 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): + 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) -> tuple[dict, bool]: + if folder_id != "root": + return _folder_unsupported("folder_id") + if sort not in ("time", "relevance"): + return _failure('sort must be "time" or "relevance"', None, + {"summary": "Invalid sort mode", + "options": ['Use sort="time" or sort="relevance"']}, + "INVALID_INPUT") + if sort == "relevance" and not query: + return _failure('query is required when sort is "relevance"', None, + {"summary": "Missing query for relevance ranking", + "options": ['Pass query alongside sort="relevance"']}, + "INVALID_INPUT") + if sort == "time" and query: + return _failure('query is only allowed when sort is "relevance"', None, + {"summary": "query does not apply to the time sort", + "options": ["Drop query, or set sort=\"relevance\""]}, + "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") + + documents = _all_documents(client) + if sort == "relevance": + tokens = [token for token in (query or "").lower().split() if token] + scored = [] + for doc in documents: + haystack = f"{doc.get('name') or ''} {doc.get('description') or ''}".lower() + score = sum(1 for token in tokens if token in haystack) + if score: + scored.append((score, doc)) + # Stable sort: equal scores keep the newest-first listing order. + scored.sort(key=lambda pair: pair[0], reverse=True) + documents = [doc for _, doc in scored] + + window = documents[offset:offset + limit] + has_more = offset + limit < len(documents) + next_offset = offset + limit 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": ( + ["No documents matched this query. Rephrase with synonyms or " + "alternative terms and retry browse_documents(sort=\"relevance\")."] + if sort == "relevance" + else ["Nothing here. Index documents with " + "PageIndexClient.submit_document() to get started."] + ), + "auto_retry": ( + "Rephrase the query and retry browse_documents(sort=\"relevance\")" + if sort == "relevance" + else "Index a document with 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, rephrase the query and " + "retry browse_documents(sort=\"relevance\"). 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) -> tuple[dict, bool]: + if folder_id not in (None, "root"): + return _folder_unsupported("folder_id") + entry, error = _resolve_document(client, doc_name) + 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'Start with first few pages: get_page_content(doc_name: "{name}", pages: "1-3")', + f'Or view structure first: get_document_structure(doc_name: "{name}")', + ]) + else: + suggestions.append("Document processing failed. Index the document " + "again with 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) -> tuple[dict, bool]: + if folder_id not in (None, "root"): + return _folder_unsupported("folder_id") + entry, error = _resolve_document(client, doc_name) + if error is not None: + return error + assert entry is not None + entry = _await_completion(client, entry, wait_for_completion) + if entry.get("status") != "completed": + return _not_ready_error(doc_name, entry.get("status"), + "structure retrieval", wait_for_completion) + + try: + # Prefer the raw stored tree: its nodes carry start_index/end_index + # like the cloud structure tool, where client.get_tree() drops + # end_index and renames fields. + store = getattr(getattr(client, "_api", None), "_store", None) + tree = store.get_tree(entry["id"]) if store is not None else None + if tree is None: + tree = client.get_tree(entry["id"], node_summary=True).get("result") + except PageIndexAPIError as exc: + return _failure( + f"Failed to retrieve document structure: {exc}", + {"doc_name": doc_name}, + { + "summary": "Failed to retrieve document structure due to an error", + "options": [ + "The document may not exist or is not accessible", + "Check if the document name is correct", + "Try again in a few moments", + ], + }, + "INTERNAL_ERROR", + ) + if tree is None: + return _failure( + "Structure not available for this document", + {"doc_name": doc_name}, + { + "summary": "Structure not available for this document", + "options": [ + "The document may not have been processed correctly or " + "structure extraction may have failed", + "Try processing the document again if possible", + ], + }, + "INTERNAL_ERROR", + ) + + formatted = _format_structure(copy.deepcopy(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) -> tuple[dict, bool]: + if folder_id not in (None, "root"): + return _folder_unsupported("folder_id") + entry, error = _resolve_document(client, doc_name) + if error is not None: + return error + assert entry is not None + 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", wait_for_completion) + + 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: {', '.join(map(str, out_of_range))}", + { + "doc_name": doc_name, + "max_pages": max_page, + "requested_pages": _format_page_spec(out_of_range), + }, + { + "summary": "All requested pages are out of range for this document", + "options": [ + f"Request pages between 1 and {max_page}", + "Use get_document() to check document page count", + ], + }, + "INVALID_INPUT", + ) + + content = [] + included: list[int] = [] + remaining: list[int] = [] + budget = _CHAR_BUDGET + for page in valid_pages: + item = by_index.get(page) + markdown = item.get("markdown") if item else None + text = (markdown if isinstance(markdown, str) + else f"Page {page} content not available") + if not included or budget - len(text) >= 0: + content.append({"page": page, "text": text}) + included.append(page) + budget -= len(text) + else: + remaining.append(page) + + options = [ + "Use get_document_structure() to understand document organization", + "Request additional pages as needed", + ] + if remaining: + options.insert(0, f"For remaining pages, request: {_format_page_spec(remaining)}") + if out_of_range: + options.insert(0, f"Document has {max_page} pages total - request " + f"pages 1-{max_page}") + summary = ( + f"Retrieved {len(included)} pages. Pages " + f"{', '.join(map(str, out_of_range))} were out of range." + if out_of_range + else f"Returned {len(included)} of {len(requested)} requested pages " + "due to response size limits." + if remaining + else 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) -> 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") + 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") + results = [] + for doc_name in doc_names: + entry, error = _resolve_document(client, doc_name) + 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 PageIndexAPIError as exc: + 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 call_tool(client, name: str, arguments: dict[str, Any]) -> tuple[str, bool]: + """Run one contract tool; returns (envelope_json, is_error). Never raises + for tool-level failures โ€” unexpected exceptions become error envelopes.""" + implementation = _IMPLEMENTATIONS[name] + try: + payload, is_error = implementation(client, **arguments) + 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", + ) + 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) + + +def _docstring(name: str) -> str: + contract = TOOL_CONTRACT[name] + return _tool_docstring(contract["description"], + contract["schema"]["properties"]) + + +_SCHEMA_TYPE_MAP = {"string": str, "integer": int, "number": float, + "boolean": bool, "array": list, "object": dict} + + +def _annotation_for(spec: dict) -> Any: + schema_type = spec.get("type") + if isinstance(schema_type, list): + bases = [t for t in schema_type if t != "null"] + base = _SCHEMA_TYPE_MAP.get(bases[0], Any) if bases else Any + return Optional[base] if "null" in schema_type else base + return _SCHEMA_TYPE_MAP.get(schema_type, Any) + + +def _make_bridge_function(bridge, meta: dict) -> Callable[..., str]: + """One plain function for a cloud tool: real signature and docstring from + the server's schema, invocation proxied over MCP, errors contained.""" + import keyword + + name = str(meta.get("name") or "") + schema = meta.get("inputSchema") or {} + properties: dict[str, Any] = schema.get("properties") or {} + required = set(schema.get("required") or []) + + def _invoke(arguments: dict[str, Any]) -> str: + # None โ‰ก omitted, matching the contract's "omit if ..." semantics. + arguments = {key: value for key, value in arguments.items() + if value is not None} + try: + return bridge.call_tool(name, arguments) + except Exception as exc: + payload, _ = _failure( + f"{name} failed: {exc}", None, + {"summary": "Unexpected error while running the tool", + "options": ["Try the request again"], + "auto_retry": "This is likely a temporary issue - you can " + "try the request again"}, + "INTERNAL_ERROR", + ) + return _dumps(payload) + + params_usable = all(param.isidentifier() and not keyword.iskeyword(param) + and param != "_invoke" + for param in properties) + if not params_usable: + def proxy(**kwargs: Any) -> str: + return _invoke(kwargs) + 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})", namespace) + proxy = namespace["_synthesized"] + annotations: dict[str, Any] = {} + for p in ordered: + annotation = _annotation_for(properties[p]) + if p not in required and "default" not in properties[p]: + # Absent-but-non-nullable params must admit None, or strict + # schemas force the model to always send a value. + annotation = Optional[annotation] + annotations[p] = annotation + annotations["return"] = str + proxy.__annotations__ = annotations + proxy.__name__ = proxy.__qualname__ = name or "tool" + proxy.__doc__ = _tool_docstring(meta.get("description", ""), properties) + return proxy + + +def _build_cloud_agent_tools(client, include_management: bool) -> list[Callable[..., str]]: + from .mcp_bridge import McpBridge + bridge = McpBridge( + f"{client.BASE_URL}/mcp", + {"Authorization": f"Bearer {client.api_key}"}, + ) + tools_meta = bridge.list_tools() + if not include_management: + # Plain functions have no framework permission layer, so the + # management gate lives here: only tools the server marks read-only. + 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." + ) + tools_meta = filtered + return [_make_bridge_function(bridge, meta) for meta in tools_meta] + + +def build_agent_tools(client, include_management: bool = False) -> list[Callable[..., str]]: + """Plain synchronous functions bound to `client`. + + Cloud: one function per tool of the live cloud MCP tool set, signatures + synthesized from the server's schemas, calls proxied over MCP. Local: + the built-in contract tools over the local store. Every function returns + the JSON envelope as a string and never raises. + """ + if getattr(client, "api_key", None): + return _build_cloud_agent_tools(client, include_management) + + def browse_documents(folder_id: str = "root", recursive: bool = False, + sort: str = "time", query: Optional[str] = None, + offset: int = 0, limit: int = 10) -> str: + return call_tool(client, "browse_documents", { + "folder_id": folder_id, "recursive": recursive, "sort": sort, + "query": query, "offset": offset, "limit": limit, + })[0] + + def get_document(doc_name: str, folder_id: Optional[str] = None, + wait_for_completion: bool = False) -> str: + return call_tool(client, "get_document", { + "doc_name": doc_name, "folder_id": folder_id, + "wait_for_completion": wait_for_completion, + })[0] + + def get_document_structure(doc_name: str, folder_id: Optional[str] = None, + part: int = 1, + wait_for_completion: bool = False) -> str: + return call_tool(client, "get_document_structure", { + "doc_name": doc_name, "folder_id": folder_id, "part": part, + "wait_for_completion": wait_for_completion, + })[0] + + def get_page_content(doc_name: str, pages: str, + folder_id: Optional[str] = None, + wait_for_completion: bool = False) -> str: + return call_tool(client, "get_page_content", { + "doc_name": doc_name, "pages": pages, "folder_id": folder_id, + "wait_for_completion": wait_for_completion, + })[0] + + def remove_document(doc_names: list[str], + folder_id: Optional[str] = None) -> str: + return call_tool(client, "remove_document", { + "doc_names": doc_names, "folder_id": folder_id, + })[0] + + functions = { + "browse_documents": browse_documents, + "get_document": get_document, + "get_document_structure": get_document_structure, + "get_page_content": get_page_content, + "remove_document": remove_document, + } + tools = [] + for name in tool_names(include_management): + function = functions[name] + function.__doc__ = _docstring(name) + tools.append(function) + return tools + + +# โ”€โ”€ agent instructions โ”€โ”€ + +_INSTRUCTIONS_HEADER = ( + "PageIndex by Vectify AI is a document platform for uploading and " + "managing long PDFs (research papers, financial reports, legal docs, " + "textbooks, etc.)." +) + +_READING_WORKFLOW = f"""\ +READING WORKFLOW: +- For documents over {STRUCTURE_FIRST_PAGE_THRESHOLD} pages: call get_document_structure() first to locate relevant sections, then get_page_content() with targeted page ranges. +- For small documents ({STRUCTURE_FIRST_PAGE_THRESHOLD} pages or fewer): call get_page_content() directly.""" + +_TOOL_USAGE_RULES = """\ +TOOL USAGE RULES: +- Invoke a tool only when all required parameters are present or clearly inferable. Never invent placeholder values. +- If a tool returns an error, present the provided next_steps/options to the user instead of retrying blindly.""" + +_DISCOVERY = """\ +DOCUMENT DISCOVERY: +- browse_documents() โ€” DEFAULT discovery tool, first choice for any document-related question. The bare call returns your documents. Use sort="relevance" + query for semantic ranking.""" + +_DECISION = """\ +DECISION: +- "What do I have / list / recent" โ†’ browse_documents (time) +- ANY question that needs a document to answer (including "find THE paper about Y") โ†’ browse_documents(sort="relevance", query=โ€ฆ)""" + +_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(sort="relevance", query=โ€ฆ) with the original intent +2. Rephrase the query with synonyms or alternative terms โ†’ browse_documents(sort="relevance") again +3. browse_documents(recursive=true) to flatten the library into one list โ€” MANDATORY, must be attempted at least once before concluding "not found" +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 build_agent_instructions(client, doc_id=None) -> str: + """Orchestration guidance for document QA agents; with doc_id, appends + the target documents and directs the agent to work within them.""" + if doc_id is None: + return AGENT_INSTRUCTIONS + doc_ids = [doc_id] if isinstance(doc_id, str) else list(doc_id) + if not doc_ids: + return AGENT_INSTRUCTIONS + details = [client.get_document(one_id) for one_id in doc_ids] + context = json.dumps(details, ensure_ascii=False) + if len(details) == 1: + block = ( + 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()." + ) + else: + names = ", ".join(str(item.get("name")) for item in details) + block = ( + 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()." + ) + return AGENT_INSTRUCTIONS + "\n\n" + block diff --git a/pageindex/client.py b/pageindex/client.py index 158c9b6f7..2a402a485 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -1,7 +1,8 @@ """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 time +from typing import Any, Callable, Iterator, Optional, Union from .errors import PageIndexAPIError @@ -126,12 +127,14 @@ 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. + 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 (it blocks while your LLM builds the tree โ€” minutes for a standard index of a long document), @@ -151,14 +154,53 @@ def submit_document( 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': ...} """ - 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, ) + if wait: + self._wait_until_ready(result["doc_id"]) + return result + + def _wait_until_ready(self, doc_id: str, timeout: float = 1800.0) -> None: + 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: + # Tolerate transient poll failures; a 30-minute wait should + # not die on one 502. + poll_failures += 1 + if poll_failures >= 3: + raise + 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 ---------- @@ -365,6 +407,104 @@ def list_documents( """ return self._api.list_documents(limit=limit, offset=offset, folder_id=folder_id) + # ---------- AGENT INTEGRATION ---------- + + def agent_tools(self, include_management: bool = False) -> list[Callable[..., str]]: + """ + Plain functions for any agent framework (LangChain, PydanticAI, ...). + For the OpenAI / Claude Agent SDKs, prefer ``as_openai_tools()`` / + ``as_claude_mcp()``. + + Cloud: the full cloud tool set, discovered live from the PageIndex + MCP server when this method is called โ€” one function per tool, + signature and docstring synthesized from the server's schemas, calls + executed from your process over MCP. Raises PageIndexAPIError if the + server cannot be reached. Local: the built-in tools over the local + store (``browse_documents``, ``get_document``, + ``get_document_structure``, ``get_page_content``). + + Each function takes JSON-serializable arguments, returns a JSON + string, and reports failures inside that JSON instead of raising. + + Args: + include_management (bool): Also expose tools that modify the + library. Local: adds ``remove_document``. Cloud: by default + only tools the server marks read-only are exposed; True + exposes the server's complete list (upload, delete, ...). + """ + from .agent_tools import build_agent_tools + return build_agent_tools(self, include_management) + + def as_openai_tools(self, include_management: bool = False, + hosted: bool = False) -> list: + """ + Tools for the OpenAI Agents SDK โ€” pass to ``Agent(tools=...)``. + + 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. 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). + + Local: the in-process tools, any model backend; ``hosted`` does + not apply. (The framework's own ``MCPServerStreamableHttp`` + against ``{BASE_URL}/mcp`` is the async-native alternative for + its ``mcp_servers=`` slot.) + + Requires ``openai-agents`` (``pip install 'pageindex[openai]'``), + imported only when this method is called. + + Args: + include_management (bool): Also expose tools that modify the + library (delete, upload). Default off: the cloud default + serves only server-annotated read-only tools, and + ``hosted=True`` routes non-read-only tools through the + Responses API approval flow instead. + hosted (bool): Cloud only โ€” hand the MCP connection to OpenAI + for server-side tool execution (OpenAI models only). + """ + from .integrations.openai_agents import build_openai_tools + return build_openai_tools(self, include_management, hosted) + + def as_claude_mcp(self, include_management: bool = False): + """ + ``mcp_servers`` entry for the Claude Agent SDK. + + Cloud: returns the remote PageIndex MCP config โ€” the framework + connects to api.pageindex.ai/mcp directly and discovers the full + cloud tool set. ``include_management`` has no effect there; gate + destructive tools with the framework's permission layer (e.g. list + read tools in ``allowed_tools`` instead of the ``*`` wildcard, or + add ``disallowed_tools=["mcp__pageindex__remove_document"]``). + Local: returns an in-process SDK MCP server exposing the agent + tools (requires ``claude-agent-sdk``; + ``pip install 'pageindex[claude]'``). + + Usage:: + + options = ClaudeAgentOptions( + mcp_servers={"pageindex": client.as_claude_mcp()}, + allowed_tools=["mcp__pageindex__*"], + ) + """ + from .integrations.claude_agent_sdk import build_claude_mcp + return build_claude_mcp(self, include_management) + + def agent_instructions(self, doc_id: Optional[Union[str, list[str]]] = None) -> str: + """ + Orchestration guidance for document QA agents โ€” pass as the agent's + system prompt (or append to your own). + + 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. + """ + from .agent_tools import build_agent_instructions + return build_agent_instructions(self, doc_id) + # ---------- FOLDER MANAGEMENT ---------- def create_folder( 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/claude_agent_sdk.py b/pageindex/integrations/claude_agent_sdk.py new file mode 100644 index 000000000..d0e2316ae --- /dev/null +++ b/pageindex/integrations/claude_agent_sdk.py @@ -0,0 +1,67 @@ +"""Claude Agent SDK adapter: one value for the mcp_servers slot. + +Cloud clients get the remote PageIndex MCP config (the framework connects +directly and discovers the full cloud tool set); local clients get an +in-process SDK MCP server over the same tool contract. +""" +from __future__ import annotations + +import asyncio +from typing import Any + +from ..errors import PageIndexAPIError + + +def _sdk_version() -> str: + try: + from importlib.metadata import version + return version("pageindex") + except Exception: + return "0.0.0" + + +def build_claude_mcp(client, include_management: bool = False): + if getattr(client, "api_key", None): + return { + "type": "http", + "url": f"{client.BASE_URL}/mcp", + "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, call_tool, tool_names + + def make_handler(name: str): + async def handler(arguments: dict[str, Any]) -> dict[str, Any]: + text, is_error = await asyncio.to_thread( + call_tool, client, name, arguments or {} + ) + 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, TOOL_CONTRACT[name]["description"], + TOOL_CONTRACT[name]["schema"], **tool_kwargs(name))(make_handler(name)) + for name in tool_names(include_management) + ] + return create_sdk_mcp_server(name="pageindex", version=_sdk_version(), + tools=tools) diff --git a/pageindex/integrations/openai_agents.py b/pageindex/integrations/openai_agents.py new file mode 100644 index 000000000..91ed6ebe3 --- /dev/null +++ b/pageindex/integrations/openai_agents.py @@ -0,0 +1,36 @@ +"""OpenAI Agents SDK adapter for the Agent(tools=...) slot. + +Cloud clients get one hosted MCP tool (the model connects to the PageIndex +cloud MCP server from OpenAI's side and discovers the full cloud tool set); +local clients get the in-process tools wrapped as FunctionTools. +""" +from __future__ import annotations + +from ..errors import PageIndexAPIError + + +def build_openai_tools(client, include_management: bool = False, + hosted: bool = False) -> list: + try: + from agents import HostedMCPTool, function_tool + except ImportError as exc: + raise PageIndexAPIError( + "as_openai_tools requires the OpenAI Agents SDK โ€” " + "pip install openai-agents (or pip install 'pageindex[openai]')." + ) from exc + if getattr(client, "api_key", None) and hosted: + # Same gate as the in-process path, enforced by OpenAI: tools the + # server annotates read-only run freely, everything else goes + # through the Responses API approval flow. + require_approval = ("never" if include_management + else {"never": {"read_only": True}}) + return [HostedMCPTool(tool_config={ + "type": "mcp", + "server_label": "pageindex", + "server_url": f"{client.BASE_URL}/mcp", + "headers": {"Authorization": f"Bearer {client.api_key}"}, + "require_approval": require_approval, + })] + from ..agent_tools import build_agent_tools + return [function_tool(tool) + for tool in build_agent_tools(client, include_management)] diff --git a/pageindex/mcp_bridge.py b/pageindex/mcp_bridge.py new file mode 100644 index 000000000..d144cee92 --- /dev/null +++ b/pageindex/mcp_bridge.py @@ -0,0 +1,181 @@ +"""Minimal MCP client (streamable HTTP) for the PageIndex cloud MCP server. + +Backs the cloud branch of ``client.agent_tools()``: ``tools/list`` discovers +the live tool set, ``tools/call`` executes a tool. Synchronous, requests-only. +Works against both stateful and stateless servers: a session id returned by +``initialize`` is echoed back, and a request rejected after session expiry +re-initializes once and retries. +""" +from __future__ import annotations + +import json +import threading +from typing import Any, Optional + +import requests + +from .errors import PageIndexAPIError + +_PROTOCOL_VERSION = "2025-06-18" +_TIMEOUT = (10, 240) # tools may wait server-side (wait_for_completion: 3 min) + + +def _sdk_version() -> str: + try: + from importlib.metadata import version + return version("pageindex") + except Exception: + return "0.0.0" + + +def _parse_sse(text: str) -> list[dict]: + """JSON-RPC messages out of a text/event-stream body.""" + messages = [] + text = text.replace("\r\n", "\n").replace("\r", "\n") + for block in text.split("\n\n"): + data_lines = [line[5:].removeprefix(" ") for line in block.splitlines() + if line.startswith("data:")] + if not data_lines: + continue + try: + messages.append(json.loads("\n".join(data_lines))) + except ValueError: + continue + return messages + + +class McpBridge: + def __init__(self, url: str, headers: dict[str, str]): + self._url = url + self._auth_headers = dict(headers) + self._session_id: Optional[str] = None + self._protocol_version: Optional[str] = None + self._initialized = False + self._lock = threading.Lock() + self._next_id = 0 + + # โ”€โ”€ JSON-RPC over streamable HTTP โ”€โ”€ + + def _post(self, payload: dict) -> requests.Response: + headers = { + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + **self._auth_headers, + } + if self._session_id: + headers["Mcp-Session-Id"] = self._session_id + if self._protocol_version: + headers["MCP-Protocol-Version"] = self._protocol_version + try: + return requests.post(self._url, json=payload, headers=headers, + timeout=_TIMEOUT) + except requests.RequestException as exc: + raise PageIndexAPIError( + f"Could not reach the PageIndex MCP server: {exc}" + ) from exc + + def _extract_result(self, response: requests.Response, request_id: int) -> Any: + content_type = response.headers.get("Content-Type", "") + if "text/event-stream" in content_type: + # SSE is UTF-8 by spec; requests guesses latin-1 for charset-less + # text/* and would mojibake every non-ASCII character. + messages = _parse_sse(response.content.decode("utf-8", + errors="replace")) + else: + try: + messages = [response.json()] + except ValueError as exc: + raise PageIndexAPIError( + f"MCP server returned a non-JSON response " + f"(HTTP {response.status_code})." + ) from exc + reply = next((m for m in messages if m.get("id") == request_id), + next((m for m in messages + if "result" in m or "error" in m), None)) + if reply is None: + raise PageIndexAPIError("MCP server response contained no reply.") + 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 + payload: dict[str, Any] = {"jsonrpc": "2.0", "id": request_id, + "method": method} + if params is not None: + payload["params"] = params + response = self._post(payload) + if response.status_code in (400, 404) and self._initialized and _retry: + # Session expired (stateful servers): start over, retry once. + with self._lock: + self._initialized = False + self._session_id = None + return self._request(method, params, _retry=False) + if response.status_code >= 400: + raise PageIndexAPIError( + f"MCP request failed: HTTP {response.status_code} " + f"({response.text[:200]})" + ) + return self._extract_result(response, request_id) + + def _ensure_initialized(self) -> None: + with self._lock: + if self._initialized: + return + self._next_id += 1 + request_id = self._next_id + response = self._post({ + "jsonrpc": "2.0", "id": request_id, "method": "initialize", + "params": { + "protocolVersion": _PROTOCOL_VERSION, + "capabilities": {}, + "clientInfo": {"name": "pageindex-python-sdk", + "version": _sdk_version()}, + }, + }) + if response.status_code >= 400: + raise PageIndexAPIError( + f"Could not connect to the PageIndex MCP server: HTTP " + f"{response.status_code} ({response.text[:200]}). Check " + "your API key." + ) + result = self._extract_result(response, request_id) or {} + self._session_id = response.headers.get("Mcp-Session-Id") + self._protocol_version = result.get("protocolVersion", + _PROTOCOL_VERSION) + self._initialized = True + try: + self._post({"jsonrpc": "2.0", + "method": "notifications/initialized"}) + except PageIndexAPIError: + pass # advisory; a server that required it fails the next request + + # โ”€โ”€ public surface โ”€โ”€ + + def list_tools(self) -> list[dict]: + tools: list[dict] = [] + cursor: Optional[str] = None + while True: + params = {"cursor": cursor} if cursor else {} + result = self._request("tools/list", params) or {} + tools.extend(result.get("tools") or []) + cursor = result.get("nextCursor") + if not cursor: + return tools + + def call_tool(self, name: str, arguments: dict[str, Any]) -> str: + result = self._request("tools/call", + {"name": name, "arguments": arguments}) or {} + blocks = result.get("content") or [] + texts = [block.get("text", "") for block in blocks + if isinstance(block, dict) and block.get("type") == "text"] + if len(texts) == len(blocks): + return "\n".join(texts) + return json.dumps(blocks, ensure_ascii=False) diff --git a/pyproject.toml b/pyproject.toml index deac66be3..65f68646b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,6 +38,14 @@ sortedcontainers = ">=2.4.0" regex = ">=2024.0.0" python-dotenv = ">=1.0.0" pyyaml = ">=6.0" +claude-agent-sdk = { version = ">=0.1.0", optional = true } +# 0.8.0 offloads sync tools to a thread; older versions run them inline and +# a blocking bridge call would freeze the agent event loop. +openai-agents = { version = ">=0.8.0", optional = true } + +[tool.poetry.extras] +claude = ["claude-agent-sdk"] +openai = ["openai-agents"] [tool.poetry.group.dev.dependencies] pytest = ">=7.0" diff --git a/tests/data/cloud_mcp_contract.json b/tests/data/cloud_mcp_contract.json new file mode 100644 index 000000000..71743aee2 --- /dev/null +++ b/tests/data/cloud_mcp_contract.json @@ -0,0 +1,197 @@ +{ + "_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, + "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": { + "type": [ + "string", + "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": { + "type": [ + "string", + "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, + "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": { + "type": [ + "string", + "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": { + "type": [ + "string", + "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..289c2dada --- /dev/null +++ b/tests/test_agent_tools.py @@ -0,0 +1,879 @@ +"""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 json +import os +import sys +from pathlib import Path + +import pytest + +import pageindex.client as client_module +from pageindex import PageIndexAPIError, PageIndexCloudClient, PageIndexLocalClient +from pageindex.agent_tools import ( + AGENT_INSTRUCTIONS, + TOOL_CONTRACT, + call_tool, + tool_names, +) +from pageindex.local_store import DocStore + +SNAPSHOT_PATH = Path(__file__).parent / "data" / "cloud_mcp_contract.json" + + +def seed_doc(storage_path, doc_id, name, *, created_at="2026-08-01T10:00:00.123000", + description="A test document", metadata=None, tree=None, pages=None, + page_num=None): + pages = pages if pages is not None else [ + {"page_index": 1, "markdown": "Page one text about apples"}, + {"page_index": 2, "markdown": "Page two text about bananas"}, + ] + tree = tree if tree is not None else [{ + "title": "Doc", "node_id": "0000", "start_index": 1, "end_index": 2, + "summary": "root summary", "text": "ROOT TEXT", + "nodes": [ + {"title": "Intro", "node_id": "0001", "start_index": 1, + "end_index": 1, "summary": "intro summary", "text": "INTRO TEXT"}, + {"title": "Body", "node_id": "0002", "start_index": 2, + "end_index": 2, "summary": "body summary", "text": "BODY TEXT"}, + ], + }] + meta = { + "id": doc_id, "name": name, "description": description, + "status": "completed", "createdAt": created_at, + "pageNum": page_num if page_num is not None else len(pages), + "folderId": None, "metadata": metadata, "mode": "standard", + } + DocStore(storage_path).save_document(doc_id, meta, tree, pages) + return doc_id + + +@pytest.fixture +def store_path(tmp_path): + return str(tmp_path / "store") + + +@pytest.fixture +def client(store_path): + return PageIndexLocalClient(storage_path=store_path) + + +def run(client, name, **arguments): + text, is_error = call_tool(client, name, arguments) + return json.loads(text), is_error + + +# โ”€โ”€ contract parity โ”€โ”€ + +def test_contract_matches_snapshot(): + snapshot = json.loads(SNAPSHOT_PATH.read_text(encoding="utf-8")) + assert snapshot["tools"] == TOOL_CONTRACT + + +def test_tool_surface_and_docstrings(client): + 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 tools: + contract = TOOL_CONTRACT[tool.__name__] + assert tool.__doc__.startswith(contract["description"]) + for param in contract["schema"]["properties"]: + assert param in tool.__doc__ + + +# โ”€โ”€ 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 + 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 + + +def test_browse_documents_relevance(client, store_path): + seed_doc(store_path, "pi-a", "annual-report.pdf", + description="Financial results for the year") + seed_doc(store_path, "pi-b", "attention.pdf", + description="Transformers and attention mechanisms") + payload, is_error = run(client, "browse_documents", sort="relevance", + query="attention transformers") + assert not is_error + assert [d["name"] for d in payload["documents"]] == ["attention.pdf"] + assert payload["sort"] == "relevance" + + missing_query, is_error = run(client, "browse_documents", sort="relevance") + assert is_error and missing_query["errorCode"] == "INVALID_INPUT" + stray_query, is_error = run(client, "browse_documents", query="x") + assert is_error and "relevance" in stray_query["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) + chunk = payload["structure"] + nodes = chunk if isinstance(chunk, list) else [chunk] + titles.extend(node["title"] for node in nodes) + 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"] + + +# โ”€โ”€ 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 + + +@pytest.mark.parametrize("bad_spec", ["abc", "5-3", "1,,2", "-3", ""]) +def test_page_content_invalid_spec(client, store_path, bad_spec): + seed_doc(store_path, "pi-a", "report.pdf") + payload, is_error = run(client, "get_page_content", doc_name="report.pdf", + pages=bad_spec) + assert is_error and payload["errorCode"] == "INVALID_INPUT" + + +def test_page_content_zero_page_rejected(client, store_path): + seed_doc(store_path, "pi-a", "report.pdf") + payload, is_error = run(client, "get_page_content", doc_name="report.pdf", + pages="0") + assert is_error + assert "positive integers" in payload["error"] + + +def test_page_content_preserves_blank_pages(client, store_path): + pages = [ + {"page_index": 1, "markdown": ""}, + {"page_index": 2, "markdown": "content"}, + ] + seed_doc(store_path, "pi-a", "blanks.pdf", pages=pages) + payload, is_error = run(client, "get_page_content", doc_name="blanks.pdf", + pages="1-2") + assert not is_error + assert payload["content"][0] == {"page": 1, "text": ""} + assert payload["content"][1] == {"page": 2, "text": "content"} + + +def test_created_at_accepts_z_suffixed_input(client, store_path): + seed_doc(store_path, "pi-a", "cloudlike.pdf", + created_at="2026-08-01T10:00:00.123Z") + payload, _ = run(client, "browse_documents") + assert payload["documents"][0]["created_at"] == "2026-08-01T10:00:00.123Z" + + +def test_page_content_char_budget(client, store_path): + pages = [ + {"page_index": 1, "markdown": "x" * 96_000}, + {"page_index": 2, "markdown": "short"}, + ] + seed_doc(store_path, "pi-a", "huge.pdf", pages=pages) + payload, is_error = run(client, "get_page_content", doc_name="huge.pdf", + pages="1-2") + assert not is_error + assert payload["returned_pages"] == "1" + assert any("For remaining pages, request: 2" in option + for option in payload["next_steps"]["options"]) + + +# โ”€โ”€ 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_management_tools_hidden_by_default(client): + assert "remove_document" not in [t.__name__ for t in client.agent_tools()] + + +# โ”€โ”€ 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" + + +# โ”€โ”€ 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" + 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_claude_mcp_cloud_needs_no_framework(monkeypatch): + monkeypatch.setitem(sys.modules, "claude_agent_sdk", None) + cloud = PageIndexCloudClient(api_key="pi-test-key") + config = cloud.as_claude_mcp() + assert config == { + "type": "http", + "url": "https://api.pageindex.ai/mcp", + "headers": {"Authorization": "Bearer pi-test-key"}, + } + + +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_agent_tools_work_without_frameworks(client, store_path, monkeypatch): + monkeypatch.setitem(sys.modules, "agents", None) + monkeypatch.setitem(sys.modules, "claude_agent_sdk", 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": {"type": ["string", "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 call_tool(self, name, arguments): + self.calls.append((name, arguments)) + return json.dumps({"success": True, "tool": name, "args": arguments}) + + +@pytest.fixture +def cloud_with_fake_bridge(monkeypatch): + import pageindex.mcp_bridge as mcp_bridge + created = {} + + def factory(url, headers): + created["bridge"] = _FakeBridge(url, headers) + return created["bridge"] + + monkeypatch.setattr(mcp_bridge, "McpBridge", factory) + return PageIndexCloudClient(api_key="pi-test-key"), created + + +def test_cloud_agent_tools_discover_live_tool_set(cloud_with_fake_bridge): + cloud, created = cloud_with_fake_bridge + tools = cloud.agent_tools() + bridge = created["bridge"] + assert bridge.url == "https://api.pageindex.ai/mcp" + assert bridge.headers == {"Authorization": "Bearer pi-test-key"} + # Default: only tools the server marks read-only; unannotated tools are + # treated as non-read-only. + assert [t.__name__ for t in tools] == ["search_documents", "get_document"] + assert "ESCALATION tool" in tools[0].__doc__ + + +def test_cloud_agent_tools_management_gate(cloud_with_fake_bridge): + cloud, _ = cloud_with_fake_bridge + names = [t.__name__ for t in cloud.agent_tools(include_management=True)] + assert names == ["search_documents", "get_document", "remove_document", + "unannotated_tool"] + + +def test_cloud_agent_tools_signatures_from_schema(cloud_with_fake_bridge): + import inspect + cloud, _ = cloud_with_fake_bridge + search, get_document = cloud.agent_tools() + params = inspect.signature(search).parameters + assert list(params) == ["query", "limit"] + assert params["query"].default is inspect.Parameter.empty + assert params["limit"].default == 10 + assert search.__annotations__["query"] is str + folder_param = inspect.signature(get_document).parameters["folder_id"] + assert folder_param.default is None + + +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_call_errors_contained(cloud_with_fake_bridge): + cloud, created = cloud_with_fake_bridge + search, _ = cloud.agent_tools() + created["bridge"].call_tool = lambda *a, **k: (_ for _ in ()).throw( + RuntimeError("network down")) + payload = json.loads(search(query="x")) + assert payload["errorCode"] == "INTERNAL_ERROR" + assert "network down" in payload["error"] + + +def test_cloud_agent_tools_list_failure_raises(monkeypatch): + import pageindex.mcp_bridge as mcp_bridge + + class _DeadBridge: + def __init__(self, url, headers): + pass + + def list_tools(self): + raise PageIndexAPIError("Could not connect") + + monkeypatch.setattr(mcp_bridge, "McpBridge", _DeadBridge) + cloud = PageIndexCloudClient(api_key="pi-test-key") + with pytest.raises(PageIndexAPIError, match="Could not connect"): + cloud.agent_tools() + + +def test_mcp_bridge_protocol(monkeypatch): + 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"}}, + {"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}") + + monkeypatch.setattr(mcp_bridge.requests, "post", fake_post) + bridge = McpBridge("https://api.pageindex.ai/mcp", + {"Authorization": "Bearer k"}) + + tools = bridge.list_tools() + assert tools == [{"name": "t1", "description": "reads โ€” never writes"}] + 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 = bridge.call_tool("t1", {"a": 1}) + assert text == "hello\nworld" + methods = [p["payload"]["method"] for p in posts] + assert methods.count("initialize") == 2 + + +# โ”€โ”€ review-round regressions โ”€โ”€ + +def test_synth_optional_no_default_param_is_nullable(): + """A non-required, no-default schema param must annotate Optional, or + strict schemas force the model to always send a value (browse.query).""" + from pageindex.agent_tools import _make_bridge_function, TOOL_CONTRACT + from typing import get_args + + class _Bridge: + def call_tool(self, name, args): + return json.dumps(args) + + meta = {"name": "browse_documents", + "description": "d", + "inputSchema": TOOL_CONTRACT["browse_documents"]["schema"]} + fn = _make_bridge_function(_Bridge(), meta) + assert type(None) in get_args(fn.__annotations__["query"]) + + +def test_synth_escape_hatches(): + from pageindex.agent_tools import _make_bridge_function + + calls = [] + + class _Bridge: + def call_tool(self, name, args): + calls.append((name, args)) + return "ok" + + # Tool named "_invoke" must not recurse into itself. + invoke_named = _make_bridge_function(_Bridge(), { + "name": "_invoke", "description": "d", + "inputSchema": {"type": "object", "properties": {"x": {"type": "string"}}, + "required": ["x"]}}) + assert invoke_named("v") == "ok" + assert calls[-1] == ("_invoke", {"x": "v"}) + + # Param named "dict" must not shadow the builtin. + dict_param = _make_bridge_function(_Bridge(), { + "name": "t", "description": "d", + "inputSchema": {"type": "object", "properties": {"dict": {"type": "string"}}, + "required": ["dict"]}}) + assert dict_param("v") == "ok" + assert calls[-1] == ("t", {"dict": "v"}) + + # Non-identifier tool name still gets a real signature. + import inspect + dashed = _make_bridge_function(_Bridge(), { + "name": "page-content.v2", "description": "d", + "inputSchema": {"type": "object", "properties": {"a": {"type": "string"}}, + "required": ["a"]}}) + assert dashed.__name__ == "page-content.v2" + assert list(inspect.signature(dashed).parameters) == ["a"] + assert dashed("v") == "ok" + + +def test_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_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, "post", dead_post) + bridge = McpBridge("https://api.pageindex.ai/mcp", {}) + with pytest.raises(PageIndexAPIError, match="Could not reach"): + bridge.list_tools() + + +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 + entry = {"id": "pi-a", "name": "broken.pdf", "status": "failed"} + 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_approval_gate(): + pytest.importorskip("agents") + cloud = PageIndexCloudClient(api_key="pi-test-key") + gated = cloud.as_openai_tools(hosted=True)[0].tool_config + assert gated["require_approval"] == {"never": {"read_only": True}} + open_config = cloud.as_openai_tools(hosted=True, + include_management=True)[0].tool_config + assert open_config["require_approval"] == "never" + + +def test_wait_tolerates_transient_poll_failures(fake_cloud_client, monkeypatch): + cloud = fake_cloud_client(["processing", "completed"]) + original = cloud._api.get_document + state = {"raised": False} + + def flaky(doc_id): + if not state["raised"]: + state["raised"] = True + raise PageIndexAPIError("502") + return original(doc_id) + + monkeypatch.setattr(cloud._api, "get_document", flaky) + assert cloud.submit_document("x.pdf", wait=True) == {"doc_id": "pi-fake"} + + +LIVE_KEY = os.getenv("PAGEINDEX_API_KEY") + + +@pytest.mark.skipif(not LIVE_KEY, reason="PAGEINDEX_API_KEY not set") +def test_live_cloud_contract_parity(): + """Real-drift detector: the frozen contract must match the live server + on every shared tool, including the annotations the gates rely on.""" + from pageindex.mcp_bridge import McpBridge + bridge = McpBridge("https://api.pageindex.ai/mcp", + {"Authorization": f"Bearer {LIVE_KEY}"}) + live = {t["name"]: t for t in bridge.list_tools()} + for name, ours in TOOL_CONTRACT.items(): + real = live.get(name) + assert real is not None, f"{name} missing from live tools/list" + assert real.get("description") == ours["description"], name + real_schema = real.get("inputSchema") or {} + real_props = real_schema.get("properties") or {} + assert set(real_props) == set(ours["schema"]["properties"]), name + for param, spec in ours["schema"]["properties"].items(): + assert (real_props[param].get("description") + == spec.get("description")), (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) + + +# โ”€โ”€ 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 + + +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") + + +# โ”€โ”€ 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, "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, "monotonic", fake_monotonic) + 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 From 873779990ad447b2a8f3fce8cf1cdad47381a487 Mon Sep 17 00:00:00 2001 From: Ray Date: Mon, 10 Aug 2026 17:30:12 +0800 Subject: [PATCH 002/137] =?UTF-8?q?fix:=20agent=20tools=20review=20?= =?UTF-8?q?=E2=80=94=20next=5Fsteps=20order,=20resolve=20caching,=20error?= =?UTF-8?q?=20semantics?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Large-doc next_steps now says structure-first, consistent with tool descriptions and agent instructions - _remove_document fetches document list once instead of per-name - call_tool returns error envelope for unknown names instead of raising - _not_ready_error timed_out flag reflects actual wait outcome - openai_agents.py docstring corrected to match default (FunctionTools) - Removed unused ModelSettings import from demo --- examples/agentic_vectorless_rag_demo.py | 3 +-- pageindex/agent_tools.py | 30 +++++++++++++++++++------ pageindex/integrations/openai_agents.py | 7 +++--- 3 files changed, 28 insertions(+), 12 deletions(-) diff --git a/examples/agentic_vectorless_rag_demo.py b/examples/agentic_vectorless_rag_demo.py index e8ed4a50a..0682bf2c7 100644 --- a/examples/agentic_vectorless_rag_demo.py +++ b/examples/agentic_vectorless_rag_demo.py @@ -29,7 +29,6 @@ sys.path.insert(0, str(Path(__file__).parent.parent)) from agents import Agent, Runner, set_tracing_disabled -from agents.model_settings import ModelSettings from agents.stream_events import RawResponsesStreamEvent, RunItemStreamEvent from openai.types.responses import ResponseTextDeltaEvent, ResponseReasoningSummaryTextDeltaEvent @@ -54,7 +53,7 @@ def query_agent(client: PageIndexLocalClient, doc_id: str, prompt: str, verbose: instructions=client.agent_instructions(doc_id=doc_id), tools=client.as_openai_tools(), model=client.retrieve_model, - # model_settings=ModelSettings(reasoning={"effort": "low", "summary": "auto"}), # Uncomment to enable reasoning + # model_settings=ModelSettings(reasoning={"effort": "low", "summary": "auto"}), # from agents.model_settings import ModelSettings ) async def _run(): diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index 10f0340dd..916536c7b 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -356,10 +356,12 @@ def _flat_metadata(value: Any) -> Optional[dict[str, Any]]: def _resolve_document( client, doc_name: str, + documents: Optional[list[dict[str, Any]]] = 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).""" - documents = _all_documents(client) + if documents is None: + documents = _all_documents(client) 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 @@ -756,8 +758,8 @@ def _get_document(client, doc_name: str, folder_id: Optional[str] = None, else: suggestions.extend([ f"This is a large document with {page_num} pages.", - f'Start with first few pages: get_page_content(doc_name: "{name}", pages: "1-3")', - f'Or view structure first: get_document_structure(doc_name: "{name}")', + 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 " @@ -794,10 +796,12 @@ def _get_document_structure(client, doc_name: str, 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", wait_for_completion) + "structure retrieval", + waited and entry.get("status") != "failed") try: # Prefer the raw stored tree: its nodes carry start_index/end_index @@ -897,10 +901,12 @@ def _get_page_content(client, doc_name: str, pages: str, 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", wait_for_completion) + "page content retrieval", + waited and entry.get("status") != "failed") requested, error = _parse_page_spec(pages, doc_name) if error is not None: @@ -1013,9 +1019,10 @@ def _remove_document(client, doc_names: list[str], {"summary": "Too many documents in one call", "options": ["Delete at most 10 documents per call"]}, "INVALID_INPUT") + documents = _all_documents(client) results = [] for doc_name in doc_names: - entry, error = _resolve_document(client, doc_name) + entry, error = _resolve_document(client, doc_name, documents=documents) if error is not None or entry is None: results.append({"doc_name": doc_name, "status": "not_found"}) continue @@ -1051,7 +1058,16 @@ def tool_names(include_management: bool = False) -> tuple[str, ...]: def call_tool(client, name: str, arguments: dict[str, Any]) -> tuple[str, bool]: """Run one contract tool; returns (envelope_json, is_error). Never raises for tool-level failures โ€” unexpected exceptions become error envelopes.""" - implementation = _IMPLEMENTATIONS[name] + 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 json.dumps(payload), True try: payload, is_error = implementation(client, **arguments) except TypeError as exc: diff --git a/pageindex/integrations/openai_agents.py b/pageindex/integrations/openai_agents.py index 91ed6ebe3..f33c4b587 100644 --- a/pageindex/integrations/openai_agents.py +++ b/pageindex/integrations/openai_agents.py @@ -1,8 +1,9 @@ """OpenAI Agents SDK adapter for the Agent(tools=...) slot. -Cloud clients get one hosted MCP tool (the model connects to the PageIndex -cloud MCP server from OpenAI's side and discovers the full cloud tool set); -local clients get the in-process tools wrapped as FunctionTools. +Cloud clients default to the full live 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). +Local clients get the in-process tools wrapped as FunctionTools. """ from __future__ import annotations From 2ba9035569f1f66e24d52aa7d9662fe363f8431f Mon Sep 17 00:00:00 2001 From: Ray Date: Mon, 10 Aug 2026 17:59:51 +0800 Subject: [PATCH 003/137] =?UTF-8?q?fix:=20agent=20tools=20review=202=20?= =?UTF-8?q?=E2=80=94=20bridge=20thread=20safety,=20browse=20paging,=20meta?= =?UTF-8?q?data=20merge?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - McpBridge reads session/protocol headers under the lock (now RLock: _ensure_initialized posts while holding it). openai-agents runs sync tools on threads and executes parallel tool calls concurrently, so bridge functions genuinely race; a torn read sent a new session id with a stale protocol header. Measured: one session expiry under 8 threads cost 4 initializations before, minimal 2 after. - Session-expiry retry also resets the negotiated protocol version, so the re-handshake carries no stale MCP-Protocol-Version header. - browse_documents time sort pages list_documents natively instead of fetching the whole library to slice one window (relevance still needs the full list for scoring). - _await_completion: a status refetch that nulls out metadata no longer clobbers the listing's copy (setdefault was a no-op on existing None). - Structure tool reads the raw stored tree via a named LocalAPI raw_tree() seam instead of reaching into _api._store internals; drop the redundant deepcopy before _format_structure (store re-reads from disk, formatting builds fresh containers). - Shared pageindex/_version.py replaces _sdk_version duplicated in mcp_bridge and the Claude integration. Left as-is after source verification against the cloud MCP: first-page budget bypass, pageNum falsy-zero, and the page-gap fallback text are letter-for-letter cloud behavior โ€” parity wins over local repair. --- pageindex/_version.py | 10 +++++ pageindex/agent_tools.py | 26 +++++++------ pageindex/integrations/claude_agent_sdk.py | 11 +----- pageindex/local_api.py | 5 +++ pageindex/mcp_bridge.py | 25 ++++++------- tests/test_agent_tools.py | 43 ++++++++++++++++++++++ 6 files changed, 86 insertions(+), 34 deletions(-) create mode 100644 pageindex/_version.py 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 index 916536c7b..866b130a0 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -10,7 +10,6 @@ """ from __future__ import annotations -import copy import difflib import json import re @@ -407,7 +406,10 @@ def _await_completion(client, entry: dict[str, Any], wait: bool) -> dict[str, An refreshed = _refetch_entry(client, doc_id) if refreshed is None: return current - refreshed.setdefault("metadata", current.get("metadata")) + 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 @@ -631,21 +633,23 @@ def _browse_documents(client, folder_id: str = "root", recursive: bool = False, "options": ["Pass integer offset and limit values"]}, "INVALID_INPUT") - documents = _all_documents(client) if sort == "relevance": tokens = [token for token in (query or "").lower().split() if token] scored = [] - for doc in documents: + for doc in _all_documents(client): haystack = f"{doc.get('name') or ''} {doc.get('description') or ''}".lower() score = sum(1 for token in tokens if token in haystack) if score: scored.append((score, doc)) # Stable sort: equal scores keep the newest-first listing order. scored.sort(key=lambda pair: pair[0], reverse=True) - documents = [doc for _, doc in scored] - - window = documents[offset:offset + limit] - has_more = offset + limit < len(documents) + ranked = [doc for _, doc in scored] + window = ranked[offset:offset + limit] + has_more = offset + limit < len(ranked) + else: + listing = client.list_documents(limit=limit, offset=offset) + window = listing.get("documents") or [] + has_more = offset + limit < listing.get("total", 0) next_offset = offset + limit if has_more else None page_has_processing = False @@ -807,8 +811,8 @@ def _get_document_structure(client, doc_name: str, # Prefer the raw stored tree: its nodes carry start_index/end_index # like the cloud structure tool, where client.get_tree() drops # end_index and renames fields. - store = getattr(getattr(client, "_api", None), "_store", None) - tree = store.get_tree(entry["id"]) if store is not None else None + raw_tree = getattr(getattr(client, "_api", None), "raw_tree", None) + tree = raw_tree(entry["id"]) if raw_tree is not None else None if tree is None: tree = client.get_tree(entry["id"], node_summary=True).get("result") except PageIndexAPIError as exc: @@ -840,7 +844,7 @@ def _get_document_structure(client, doc_name: str, "INTERNAL_ERROR", ) - formatted = _format_structure(copy.deepcopy(tree)) + formatted = _format_structure(tree) chunks = _split_structure(formatted, _CHAR_BUDGET) total_parts = max(1, len(chunks)) try: diff --git a/pageindex/integrations/claude_agent_sdk.py b/pageindex/integrations/claude_agent_sdk.py index d0e2316ae..0fb77d2de 100644 --- a/pageindex/integrations/claude_agent_sdk.py +++ b/pageindex/integrations/claude_agent_sdk.py @@ -9,17 +9,10 @@ import asyncio from typing import Any +from .._version import sdk_version from ..errors import PageIndexAPIError -def _sdk_version() -> str: - try: - from importlib.metadata import version - return version("pageindex") - except Exception: - return "0.0.0" - - def build_claude_mcp(client, include_management: bool = False): if getattr(client, "api_key", None): return { @@ -63,5 +56,5 @@ def tool_kwargs(name: str) -> dict: TOOL_CONTRACT[name]["schema"], **tool_kwargs(name))(make_handler(name)) for name in tool_names(include_management) ] - return create_sdk_mcp_server(name="pageindex", version=_sdk_version(), + return create_sdk_mcp_server(name="pageindex", version=sdk_version(), tools=tools) diff --git a/pageindex/local_api.py b/pageindex/local_api.py index 0e9f682c8..7b7351cca 100644 --- a/pageindex/local_api.py +++ b/pageindex/local_api.py @@ -190,6 +190,11 @@ def _load_tree_with_text(self, doc_id: str, error_prefix: str) -> list: add_node_text(structure, pdf_pages) return structure + def raw_tree(self, doc_id: str) -> list | None: + """Stored tree verbatim โ€” keeps start_index/end_index, which + get_tree's cloud wire shape renames and drops.""" + return self._store.get_tree(doc_id) + def get_tree(self, doc_id: str, node_summary: bool = False, include_text: bool = True) -> dict[str, Any]: meta = self._require_doc(doc_id, "Failed to get tree result") diff --git a/pageindex/mcp_bridge.py b/pageindex/mcp_bridge.py index d144cee92..fad0baaf8 100644 --- a/pageindex/mcp_bridge.py +++ b/pageindex/mcp_bridge.py @@ -14,20 +14,13 @@ 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 _sdk_version() -> str: - try: - from importlib.metadata import version - return version("pageindex") - except Exception: - return "0.0.0" - - def _parse_sse(text: str) -> list[dict]: """JSON-RPC messages out of a text/event-stream body.""" messages = [] @@ -51,21 +44,24 @@ def __init__(self, url: str, headers: dict[str, str]): self._session_id: Optional[str] = None self._protocol_version: Optional[str] = None self._initialized = False - self._lock = threading.Lock() + self._lock = threading.RLock() self._next_id = 0 # โ”€โ”€ JSON-RPC over streamable HTTP โ”€โ”€ def _post(self, payload: dict) -> requests.Response: + with self._lock: + session_id = self._session_id + protocol_version = self._protocol_version headers = { "Content-Type": "application/json", "Accept": "application/json, text/event-stream", **self._auth_headers, } - if self._session_id: - headers["Mcp-Session-Id"] = self._session_id - if self._protocol_version: - headers["MCP-Protocol-Version"] = self._protocol_version + if session_id: + headers["Mcp-Session-Id"] = session_id + if protocol_version: + headers["MCP-Protocol-Version"] = protocol_version try: return requests.post(self._url, json=payload, headers=headers, timeout=_TIMEOUT) @@ -117,6 +113,7 @@ def _request(self, method: str, params: Optional[dict] = None, with self._lock: 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( @@ -137,7 +134,7 @@ def _ensure_initialized(self) -> None: "protocolVersion": _PROTOCOL_VERSION, "capabilities": {}, "clientInfo": {"name": "pageindex-python-sdk", - "version": _sdk_version()}, + "version": sdk_version()}, }, }) if response.status_code >= 400: diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 289c2dada..9ab865655 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -630,6 +630,11 @@ def fake_post(url, json=None, headers=None, timeout=None): assert text == "hello\nworld" 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"] # โ”€โ”€ review-round regressions โ”€โ”€ @@ -730,6 +735,44 @@ def dead_post(*args, **kwargs): 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, "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_failed_document_status_message(client, store_path): seed_doc(store_path, "pi-a", "broken.pdf") import pageindex.agent_tools as agent_tools_mod From 3f131584e25ada6d97de07832d6c6265fcf62df1 Mon Sep 17 00:00:00 2001 From: Ray Date: Mon, 10 Aug 2026 18:49:27 +0800 Subject: [PATCH 004/137] =?UTF-8?q?fix:=20agent=20tools=20review=203=20?= =?UTF-8?q?=E2=80=94=20page-span=20cap,=20duplicate=20names,=20wait=20resi?= =?UTF-8?q?lience,=20contract=20drift?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - _parse_page_spec bounds the requested span arithmetically (10k pages) before materializing it; pages="1-1000000000" previously expanded to a billion integers inside the caller's process. - Local submit_document uniquifies document names the way the cloud upload does (taken name -> _1.._99, then reject with the cloud's own message). Same-name duplicates broke name-addressed tools: resolution always picks the newest, so older duplicates were unreachable. - agent_instructions(doc_id=...) now fails loud when the pinned doc's name is shadowed by a newer same-name document (legacy stores predate the rename) โ€” it previews resolution with the same _resolve_document the tools use, so the check cannot drift from actual behavior. - submit_document(wait=True) tolerates transient network errors, not just API errors; a dropped connection at minute 25 of a 30-minute wait no longer kills it. Third strike wraps into PageIndexAPIError per the documented contract. - The live contract-parity test compares full per-param schemas, not just names and descriptions. It immediately caught real drift the shallow check had been passing: the server now emits nullables as anyOf unions and stamps MAX_SAFE_INTEGER maxima on offset/part. Contract and snapshot updated to the served wire form; _annotation_for learned anyOf so bridge signatures stay Optional[str] instead of degrading to Any. Adjudicated, not changed: the allowed_tools wildcard example stays (docstring advice covers scoping; Ray's call), and raw-length response accounting stays (letter-for-letter cloud behavior, parity wins). --- pageindex/agent_tools.py | 54 +++++++++++++++++++++++---- pageindex/client.py | 14 +++++-- pageindex/local_api.py | 18 ++++++++- tests/data/cloud_mcp_contract.json | 42 +++++++++++++++------ tests/test_agent_tools.py | 60 ++++++++++++++++++++++++++++-- tests/test_client.py | 24 ++++++++++++ 6 files changed, 185 insertions(+), 27 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index 866b130a0..68af57eaf 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -23,6 +23,7 @@ _CHAR_BUDGET = int(TOOL_RESPONSE_CHAR_LIMIT * 0.95) _PAGES_SPEC_RE = re.compile(r"^(\d+(-\d+)?)(,\s*\d+(-\d+)?)*$") +_MAX_REQUESTED_PAGES = 10_000 _SIMILAR_NAMES_LIMIT = 3 _TOOL_WAIT_TIMEOUT = 180.0 # "up to 3 minutes", per the wait_for_completion schema _TOOL_WAIT_INTERVAL = 5.0 @@ -114,6 +115,7 @@ "offset": { "type": "integer", "minimum": 0, + "maximum": 9007199254740991, "default": 0, "description": ( "Zero-based pagination offset. Pass the value of " @@ -152,7 +154,7 @@ "description": _DOC_NAME_DESCRIPTION, }, "folder_id": { - "type": ["string", "null"], + "anyOf": [{"type": "string"}, {"type": "null"}], "description": _FOLDER_ID_DISAMBIGUATOR_DESCRIPTION, }, "wait_for_completion": { @@ -183,12 +185,13 @@ "description": _DOC_NAME_DESCRIPTION, }, "folder_id": { - "type": ["string", "null"], + "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 " @@ -224,7 +227,7 @@ "description": _DOC_NAME_DESCRIPTION, }, "folder_id": { - "type": ["string", "null"], + "anyOf": [{"type": "string"}, {"type": "null"}], "description": _FOLDER_ID_DISAMBIGUATOR_DESCRIPTION, }, "pages": { @@ -273,7 +276,7 @@ ), }, "folder_id": { - "type": ["string", "null"], + "anyOf": [{"type": "string"}, {"type": "null"}], "description": _FOLDER_ID_DISAMBIGUATOR_DESCRIPTION, }, }, @@ -492,15 +495,34 @@ def _parse_page_spec( if not isinstance(pages, str) or not _PAGES_SPEC_RE.match(pages.strip()): return None, invalid expanded: set[int] = set() + requested_total = 0 for part in pages.split(","): part = part.strip() if "-" in part: start, end = (int(x) for x in part.split("-", 1)) if start > end: return None, invalid - expanded.update(range(start, end + 1)) else: - expanded.add(int(part)) + start = end = int(part) + # Bound the span arithmetically before materializing it: a spec like + # "1-1000000000" would otherwise expand to billions of integers + # inside the caller's process. + requested_total += end - start + 1 + if requested_total > _MAX_REQUESTED_PAGES: + 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", + ) + expanded.update(range(start, end + 1)) if any(page < 1 for page in expanded): return None, _failure( "Invalid page numbers. Page numbers must be positive integers", @@ -1114,6 +1136,10 @@ def _docstring(name: str) -> str: 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}]. + schema_type = [option.get("type") for option in spec["anyOf"] + if isinstance(option, dict) and option.get("type")] if isinstance(schema_type, list): bases = [t for t in schema_type if t != "null"] base = _SCHEMA_TYPE_MAP.get(bases[0], Any) if bases else Any @@ -1320,13 +1346,27 @@ def remove_document(doc_names: list[str], def build_agent_instructions(client, doc_id=None) -> str: """Orchestration guidance for document QA agents; with doc_id, appends - the target documents and directs the agent to work within them.""" + the target documents and directs the agent to work within them. Raises + when a doc_id's name is shadowed by a newer same-name document โ€” the + name-addressed tools could not reach it.""" if doc_id is None: return AGENT_INSTRUCTIONS doc_ids = [doc_id] if isinstance(doc_id, str) else list(doc_id) if not doc_ids: return AGENT_INSTRUCTIONS details = [client.get_document(one_id) for one_id in doc_ids] + documents = _all_documents(client) + 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." + ) context = json.dumps(details, ensure_ascii=False) if len(details) == 1: block = ( diff --git a/pageindex/client.py b/pageindex/client.py index 2a402a485..54ae9c018 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -172,6 +172,7 @@ def submit_document( 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 @@ -179,12 +180,16 @@ def _wait_until_ready(self, doc_id: str, timeout: float = 1800.0) -> None: try: status = self.get_document(doc_id).get("status") poll_failures = 0 - except PageIndexAPIError: + except (PageIndexAPIError, requests.RequestException) as exc: # Tolerate transient poll failures; a 30-minute wait should - # not die on one 502. + # not die on one 502 or dropped connection. poll_failures += 1 if poll_failures >= 3: - raise + if isinstance(exc, PageIndexAPIError): + raise + raise PageIndexAPIError( + f"Could not poll document status: {exc}" + ) from exc status = None if status == "completed": return @@ -500,7 +505,8 @@ def agent_instructions(self, doc_id: Optional[Union[str, list[str]]] = None) -> 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. + not exist, or if its name is shadowed by a newer same-name document + (the name-addressed tools could not reach it). """ from .agent_tools import build_agent_instructions return build_agent_instructions(self, doc_id) diff --git a/pageindex/local_api.py b/pageindex/local_api.py index 7b7351cca..5820656ac 100644 --- a/pageindex/local_api.py +++ b/pageindex/local_api.py @@ -115,7 +115,7 @@ def submit_document( doc_id = "pi-" + uuid.uuid4().hex meta = { "id": doc_id, - "name": os.path.basename(file_path), + "name": self._unique_doc_name(os.path.basename(file_path)), "description": description, "status": "completed", "createdAt": _now_iso(), @@ -131,6 +131,22 @@ def submit_document( doc_id, meta, remove_fields(structure, fields=["text"]), pages) return {"doc_id": doc_id} + 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]: import PyPDF2 diff --git a/tests/data/cloud_mcp_contract.json b/tests/data/cloud_mcp_contract.json index 71743aee2..25711a6ba 100644 --- a/tests/data/cloud_mcp_contract.json +++ b/tests/data/cloud_mcp_contract.json @@ -36,6 +36,7 @@ "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." }, @@ -65,9 +66,13 @@ "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": { - "type": [ - "string", - "null" + "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." }, @@ -97,15 +102,20 @@ "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": { - "type": [ - "string", - "null" + "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." }, @@ -135,9 +145,13 @@ "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": { - "type": [ - "string", - "null" + "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." }, @@ -181,9 +195,13 @@ "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": { - "type": [ - "string", - "null" + "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." } diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 9ab865655..5039bd031 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -451,7 +451,8 @@ def __init__(self, url, headers): "type": "object", "properties": { "doc_name": {"type": "string"}, - "folder_id": {"type": ["string", "null"]}, + "folder_id": {"anyOf": [{"type": "string"}, + {"type": "null"}]}, }, "required": ["doc_name"], }, @@ -525,6 +526,10 @@ def test_cloud_agent_tools_signatures_from_schema(cloud_with_fake_bridge): 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): @@ -693,6 +698,17 @@ def call_tool(self, name, args): 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 @@ -773,6 +789,43 @@ def spy(**kwargs): assert payload["has_more"] is True and payload["next_offset"] == 2 +def test_page_spec_span_bomb_rejected(client, store_path): + """An absurd range must be rejected arithmetically, not expanded into + billions of integers in the caller's process.""" + seed_doc(store_path, "pi-a", "report.pdf") + payload, is_error = run(client, "get_page_content", doc_name="report.pdf", + pages="1-1000000000") + assert is_error and payload["errorCode"] == "INVALID_INPUT" + assert "Too many pages" in payload["error"] + + +def test_agent_instructions_shadowed_doc_id_raises(client, store_path): + seed_doc(store_path, "pi-old", "report.pdf", + created_at="2026-08-01T10:00:00.000000") + seed_doc(store_path, "pi-new", "report.pdf", + created_at="2026-08-02T10:00:00.000000") + with pytest.raises(PageIndexAPIError, match="shadowed"): + client.agent_instructions(doc_id="pi-old") + text = client.agent_instructions(doc_id="pi-new") + assert "report.pdf" in text + + +def test_wait_tolerates_transient_network_failures(fake_cloud_client, monkeypatch): + import requests as requests_mod + cloud = fake_cloud_client(["processing", "completed"]) + original = cloud._api.get_document + state = {"raised": False} + + def flaky(doc_id): + if not state["raised"]: + state["raised"] = True + raise requests_mod.ConnectionError("network blip") + return original(doc_id) + + monkeypatch.setattr(cloud._api, "get_document", flaky) + assert cloud.submit_document("x.pdf", wait=True) == {"doc_id": "pi-fake"} + + def test_failed_document_status_message(client, store_path): seed_doc(store_path, "pi-a", "broken.pdf") import pageindex.agent_tools as agent_tools_mod @@ -828,9 +881,10 @@ def test_live_cloud_contract_parity(): 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].get("description") - == spec.get("description")), (name, param) + 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(): diff --git a/tests/test_client.py b/tests/test_client.py index 50b7f5178..3c9385831 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -162,6 +162,30 @@ def fake_page_index_main(doc, opt=None, logger=None, page_list=None): 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): + 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)["doc_id"] + second = local_client.submit_document(sample_pdf)["doc_id"] + names = {d["id"]: d["name"] + for d in local_client.list_documents()["documents"]} + assert names[first] == "sample.pdf" + assert names[second] == "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_flash(local_client, sample_pdf, monkeypatch): calls = {} def fake_flash(pdf, summary=True, summary_model=None, **kwargs): From 40b1706e8c4c4a5608bf624c72d441d6cb0d8e68 Mon Sep 17 00:00:00 2001 From: Ray Date: Mon, 10 Aug 2026 23:51:55 +0800 Subject: [PATCH 005/137] feat: surface the stored document name from submit_document MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Compute PR #558 makes /doc/ return {"doc_id", "name"} carrying the post-dedup-rename name. Mirror it end to end: local submit returns the stored name, the client warns when it differs from the uploaded file name (read via .get so older cloud servers stay compatible), the local name-exhaustion check runs before indexing instead of after the LLM spend, and the demo caches doc_id in a file instead of name-matching โ€” a renamed document made the name lookup re-index on every run. --- examples/agentic_vectorless_rag_demo.py | 17 +++++++++----- pageindex/client.py | 15 +++++++++++-- pageindex/cloud_api.py | 4 +++- pageindex/local_api.py | 5 ++++- tests/test_agent_tools.py | 9 ++++++++ tests/test_client.py | 30 ++++++++++++++++++++----- 6 files changed, 65 insertions(+), 15 deletions(-) diff --git a/examples/agentic_vectorless_rag_demo.py b/examples/agentic_vectorless_rag_demo.py index 0682bf2c7..f35c3c2e7 100644 --- a/examples/agentic_vectorless_rag_demo.py +++ b/examples/agentic_vectorless_rag_demo.py @@ -32,13 +32,14 @@ from agents.stream_events import RawResponsesStreamEvent, RunItemStreamEvent from openai.types.responses import ResponseTextDeltaEvent, ResponseReasoningSummaryTextDeltaEvent -from pageindex import PageIndexLocalClient +from pageindex import PageIndexAPIError, PageIndexLocalClient import pageindex.utils as utils PDF_URL = "https://arxiv.org/pdf/2603.15031" _EXAMPLES_DIR = Path(__file__).parent PDF_PATH = _EXAMPLES_DIR / "documents" / "attention-residuals.pdf" +DOC_ID_PATH = _EXAMPLES_DIR / "documents" / "attention-residuals.doc_id" STORAGE_PATH = _EXAMPLES_DIR / ".pageindex" @@ -129,15 +130,19 @@ async def _run(): print("=" * 60) print("Step 1: Index PDF and view tree structure") print("=" * 60) - doc_id = next( - (doc["id"] for doc in client.list_documents(limit=100)["documents"] - if doc["name"] == PDF_PATH.name), - None, - ) + doc_id = None + if DOC_ID_PATH.exists(): + cached = DOC_ID_PATH.read_text().strip() + try: + client.get_document(cached) + doc_id = cached + except PageIndexAPIError: + DOC_ID_PATH.unlink() if doc_id: print(f"\nLoaded cached doc_id: {doc_id}") else: doc_id = client.submit_document(str(PDF_PATH), wait=True)["doc_id"] + DOC_ID_PATH.write_text(doc_id) print(f"\nIndexed. doc_id: {doc_id}") print("\nTree Structure (top-level sections):") structure = client.get_tree(doc_id, node_summary=True)["result"] diff --git a/pageindex/client.py b/pageindex/client.py index 54ae9c018..ff44aa376 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -1,7 +1,9 @@ """PageIndex SDK client: the 0.2.x cloud surface, now with a local mode.""" from __future__ import annotations +import os import time +import warnings from typing import Any, Callable, Iterator, Optional, Union from .errors import PageIndexAPIError @@ -130,7 +132,7 @@ def submit_document( wait: bool = False, ) -> dict[str, Any]: """ - Submit a PDF document for processing. Returns {'doc_id': ...}. + 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 @@ -161,12 +163,21 @@ def submit_document( 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'. """ 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 diff --git a/pageindex/cloud_api.py b/pageindex/cloud_api.py index b7597cc9a..ab7a9c885 100644 --- a/pageindex/cloud_api.py +++ b/pageindex/cloud_api.py @@ -55,7 +55,9 @@ def submit_document( returned in get_tree/get_ocr responses and list_documents entries. Defaults to None. Returns: - dict: {'doc_id': ...} + dict: {'doc_id': ...} โ€” plus 'name', the stored document name + (a taken name gains a numeric suffix), when the server + returns it. """ data = {'if_retrieval': True} if mode is not None: diff --git a/pageindex/local_api.py b/pageindex/local_api.py index 5820656ac..9ad909ad0 100644 --- a/pageindex/local_api.py +++ b/pageindex/local_api.py @@ -97,6 +97,9 @@ def submit_document( raise PageIndexAPIError( "Failed to submit document: PDF has no content. All pages are blank." ) + # Fail before paying for indexing when _1.._99 are all taken; the + # binding name resolution happens again at save. + self._unique_doc_name(os.path.basename(file_path)) try: if mode == "flash": @@ -129,7 +132,7 @@ def submit_document( from .utils import remove_fields self._store.save_document( doc_id, meta, remove_fields(structure, fields=["text"]), pages) - return {"doc_id": doc_id} + return {"doc_id": doc_id, "name": meta["name"]} def _unique_doc_name(self, name: str) -> str: """Mirror the cloud upload: a taken name gets _1.._99 appended, diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 5039bd031..a9cb4cd7c 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -974,3 +974,12 @@ 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" diff --git a/tests/test_client.py b/tests/test_client.py index 3c9385831..bdbcd713f 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -169,12 +169,15 @@ def fake_page_index_main(doc, opt=None, logger=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)["doc_id"] - second = local_client.submit_document(sample_pdf)["doc_id"] + first = local_client.submit_document(sample_pdf) + assert first["name"] == "sample.pdf" + with pytest.warns(UserWarning, match='stored as "sample_1.pdf"'): + second = local_client.submit_document(sample_pdf) + assert second["name"] == "sample_1.pdf" names = {d["id"]: d["name"] for d in local_client.list_documents()["documents"]} - assert names[first] == "sample.pdf" - assert names[second] == "sample_1.pdf" + assert names[first["doc_id"]] == "sample.pdf" + assert names[second["doc_id"]] == "sample_1.pdf" def test_submit_duplicate_name_exhaustion(local_client, monkeypatch): @@ -186,6 +189,22 @@ def test_submit_duplicate_name_exhaustion(local_client, monkeypatch): 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) + + def test_submit_flash(local_client, sample_pdf, monkeypatch): calls = {} def fake_flash(pdf, summary=True, summary_model=None, **kwargs): @@ -456,7 +475,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)["doc_id"] (tmp_path / "store" / "docs" / indexed_doc / "doc.json").write_text("{truncated") # manifest still holds a good copy of the meta โ€” served consistently From bf9e6dac5f5359febaf30b3f231ac63c2af6fcd9 Mon Sep 17 00:00:00 2001 From: Ray Date: Tue, 11 Aug 2026 22:03:05 +0800 Subject: [PATCH 006/137] fix: add missing page_list kwarg in duplicate-name test mock --- README.md | 1 + tests/test_client.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 5dbafc141..f4278a598 100644 --- a/README.md +++ b/README.md @@ -262,6 +262,7 @@ tools = client.agent_tools() # local: built-in tools; ``` Neither framework is a required dependency โ€” each is imported only when its method is called. Claude Code / Cursor and other MCP hosts connect to cloud documents via the hosted MCP server directly (no SDK needed); see the [MCP docs](https://docs.pageindex.ai/mcp). + ## ๐Ÿš€ 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). diff --git a/tests/test_client.py b/tests/test_client.py index bdbcd713f..f55d519e6 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -165,7 +165,7 @@ def fake_page_index_main(doc, opt=None, logger=None, page_list=None): 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): + 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) From dece6e61fc5a5bed9b9739dfbc7a2e7eef8b068d Mon Sep 17 00:00:00 2001 From: Ray Date: Tue, 11 Aug 2026 22:08:56 +0800 Subject: [PATCH 007/137] =?UTF-8?q?revert:=20keep=20README.md=20unchanged?= =?UTF-8?q?=20from=20main=20=E2=80=94=20SDK=20section=20deferred?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 60 ++----------------------------------------------------- 1 file changed, 2 insertions(+), 58 deletions(-) diff --git a/README.md b/README.md index f4278a598..5ce0ca5e6 100644 --- a/README.md +++ b/README.md @@ -207,69 +207,13 @@ python3 run_pageindex.py --md_path /path/to/your/document.md > > Add `--optimize` to refine the tree structure for more efficient retrieval (with an LLM expansion pass). -## ๐Ÿ Python SDK: Cloud & Local - -The `pageindex` package on PyPI is the Python SDK for the [PageIndex API](https://docs.pageindex.ai) โ€” and the same client now also runs fully **locally**, powered by this repo's indexing pipeline (including Flash). - -```bash -pip3 install --upgrade pageindex # local mode ships in pageindex >= 0.2.9; earlier versions are cloud-only -``` - -```python -from pageindex import PageIndexClient - -client = PageIndexClient(api_key="YOUR_PAGEINDEX_API_KEY") # cloud: managed OCR, tree building, retrieval -client = PageIndexClient() # local: same methods on your machine, using your LLM key (e.g. OPENAI_API_KEY) - -doc_id = client.submit_document("doc.pdf")["doc_id"] # local mode blocks until indexing finishes -doc_id = client.submit_document("doc.pdf", mode="flash")["doc_id"] # local mode with PageIndex Flash - -tree = client.get_tree(doc_id, node_summary=True)["result"] - -answer = client.chat_completions( - messages=[{"role": "user", "content": "Summarize the key findings"}], - doc_id=doc_id, -)["choices"][0]["message"]["content"] -``` - -Local documents are stored as plain JSON under `./.pageindex` (configurable via `storage_path`). Local mode supports PDFs; folders, `beta_headers`, `enable_citations`, and the deprecated retrieval API (`submit_query`/`get_retrieval`) remain cloud-only โ€” each method's docstring spells out the differences. To pin the mode at construction instead of inferring it from `api_key`, use `PageIndexCloudClient` (fails without a real key) or `PageIndexLocalClient` (has no key parameter). - -### ๐Ÿค– Agent integration - -The client exposes its documents as **agent tools**, following one rule: **cloud clients always serve the live tool set of the [PageIndex MCP server](https://docs.pageindex.ai/mcp)** (search, folders, images โ€” as enabled for your key, discovered dynamically; management tools like delete/upload sit behind `include_management=True` or the framework's approval layer), while local clients serve the same contract's built-in navigation subset (`browse_documents`, `get_document`, `get_document_structure`, `get_page_content`). Tool names and schemas are shared, so agent prompts port unchanged, and switching local โ†” cloud is just the client constructor line: - -```python -client = PageIndexLocalClient() # or PageIndexCloudClient(api_key=...) -client.submit_document("doc.pdf", wait=True) # wait=True: return once the doc is ready (both modes) - -# OpenAI Agents SDK (pip install "pageindex[openai]") -agent = Agent( - name="PageIndex", - instructions=client.agent_instructions(), # retrieval playbook for the agent's system prompt - tools=client.as_openai_tools(), # local: in-process tools; cloud: the full cloud MCP tool set (any model backend) -) # cloud + OpenAI models: hosted=True runs tool calls server-side (fastest) - -# Claude Agent SDK (pip install "pageindex[claude]") -options = ClaudeAgentOptions( - system_prompt=client.agent_instructions(), - mcp_servers={"pageindex": client.as_claude_mcp()}, # local: in-process server; cloud: connects to api.pageindex.ai/mcp - allowed_tools=["mcp__pageindex__*"], -) - -# Any other framework: plain functions, wrap with your framework's one-liner -tools = client.agent_tools() # local: built-in tools; cloud: full live tool set over MCP - # e.g. [StructuredTool.from_function(f) for f in tools] -``` - -Neither framework is a required dependency โ€” each is imported only when its method is called. Claude Code / Cursor and other MCP hosts connect to cloud documents via the hosted MCP server directly (no SDK needed); see the [MCP docs](https://docs.pageindex.ai/mcp). - ## ๐Ÿš€ 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 with the OpenAI Agents SDK extra -pip3 install "pageindex[openai]" +# Install optional dependency +pip3 install openai-agents # Run the demo python3 examples/agentic_vectorless_rag_demo.py From 3c37cdc51039b00e413cccf0e983eac0bac8280f Mon Sep 17 00:00:00 2001 From: Ray Date: Tue, 11 Aug 2026 22:43:43 +0800 Subject: [PATCH 008/137] feat: serve cloud agent instructions live from the MCP server MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The cloud MCP server publishes its agent instructions in the initialize result, adapted to each key's tool set. agent_instructions() previously returned the SDK's local-subset text in both modes โ€” a silently forked copy that lacks the guidance for cloud-only tools (search_documents escalation, folders, images) and drifts as the server's prompt evolves. Cloud clients now serve the server's live instructions, captured from the initialize handshake on a per-client bridge shared with agent_tools() (one session, no extra request). An empty server response raises instead of silently substituting the subset text โ€” same posture as the annotation-regression guard. The local constant stays as the honest subset for the in-process tools, with its provenance noted and a consistency test that every tool it names exists in the local registry. --- pageindex/agent_tools.py | 47 +++++++++++++++++++++++----- pageindex/client.py | 6 ++++ pageindex/mcp_bridge.py | 13 ++++++-- tests/test_agent_tools.py | 65 ++++++++++++++++++++++++++++++++++++++- 4 files changed, 120 insertions(+), 11 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index 68af57eaf..f1628d190 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -1207,12 +1207,22 @@ def proxy(**kwargs: Any) -> str: return proxy +def _cloud_bridge(client): + """One bridge per client instance: tool discovery and instructions share + a single MCP session.""" + bridge = getattr(client, "_mcp_bridge", None) + if bridge is None: + from .mcp_bridge import McpBridge + bridge = McpBridge( + f"{client.BASE_URL}/mcp", + {"Authorization": f"Bearer {client.api_key}"}, + ) + client._mcp_bridge = bridge + return bridge + + def _build_cloud_agent_tools(client, include_management: bool) -> list[Callable[..., str]]: - from .mcp_bridge import McpBridge - bridge = McpBridge( - f"{client.BASE_URL}/mcp", - {"Authorization": f"Bearer {client.api_key}"}, - ) + bridge = _cloud_bridge(client) tools_meta = bridge.list_tools() if not include_management: # Plain functions have no framework permission layer, so the @@ -1294,6 +1304,11 @@ def remove_document(doc_names: list[str], # โ”€โ”€ agent instructions โ”€โ”€ +# Local subset of the cloud MCP server's initialize instructions (its +# no-folders variant), trimmed to the tools that exist here: the +# search_documents escalation steps, get_document_image, and the shared +# read-only-folders block are removed. Cloud clients receive the server's +# live instructions instead โ€” see _base_instructions(). _INSTRUCTIONS_HEADER = ( "PageIndex by Vectify AI is a document platform for uploading and " @@ -1344,16 +1359,32 @@ def remove_document(doc_names: list[str], ]) +def _base_instructions(client) -> str: + """Cloud: the live instructions the MCP server serves for this key's + tool set. Local: the built-in subset instructions.""" + if not getattr(client, "api_key", None): + return AGENT_INSTRUCTIONS + instructions = _cloud_bridge(client).instructions() + if not instructions: + 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 build_agent_instructions(client, doc_id=None) -> str: """Orchestration guidance for document QA agents; with doc_id, appends the target documents and directs the agent to work within them. Raises when a doc_id's name is shadowed by a newer same-name document โ€” the name-addressed tools could not reach it.""" + base = _base_instructions(client) if doc_id is None: - return AGENT_INSTRUCTIONS + return base doc_ids = [doc_id] if isinstance(doc_id, str) else list(doc_id) if not doc_ids: - return AGENT_INSTRUCTIONS + return base details = [client.get_document(one_id) for one_id in doc_ids] documents = _all_documents(client) for one_id, detail in zip(doc_ids, details): @@ -1383,4 +1414,4 @@ def build_agent_instructions(client, doc_id=None) -> str: "Use these documents' names to retrieve their content with " "get_document_structure() and get_page_content()." ) - return AGENT_INSTRUCTIONS + "\n\n" + block + return base + "\n\n" + block diff --git a/pageindex/client.py b/pageindex/client.py index ff44aa376..70608a480 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -513,6 +513,12 @@ def agent_instructions(self, doc_id: Optional[Union[str, list[str]]] = None) -> 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 diff --git a/pageindex/mcp_bridge.py b/pageindex/mcp_bridge.py index fad0baaf8..95aba5c70 100644 --- a/pageindex/mcp_bridge.py +++ b/pageindex/mcp_bridge.py @@ -1,7 +1,9 @@ """Minimal MCP client (streamable HTTP) for the PageIndex cloud MCP server. -Backs the cloud branch of ``client.agent_tools()``: ``tools/list`` discovers -the live tool set, ``tools/call`` executes a tool. Synchronous, requests-only. +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 request rejected after session expiry re-initializes once and retries. @@ -43,6 +45,7 @@ def __init__(self, url: str, headers: dict[str, str]): self._auth_headers = dict(headers) self._session_id: Optional[str] = None self._protocol_version: Optional[str] = None + self._instructions: Optional[str] = None self._initialized = False self._lock = threading.RLock() self._next_id = 0 @@ -147,6 +150,7 @@ def _ensure_initialized(self) -> None: 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 try: self._post({"jsonrpc": "2.0", @@ -156,6 +160,11 @@ def _ensure_initialized(self) -> None: # โ”€โ”€ 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 diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index a9cb4cd7c..b7677f1d7 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -2,6 +2,7 @@ local store (no LLM calls; one live parity test gated on PAGEINDEX_API_KEY).""" import json import os +import re import sys from pathlib import Path @@ -594,7 +595,8 @@ def fake_post(url, json=None, headers=None, timeout=None): rid = json.get("id") if method == "initialize": return _Resp(200, {"jsonrpc": "2.0", "id": rid, - "result": {"protocolVersion": "2025-06-18"}}, + "result": {"protocolVersion": "2025-06-18", + "instructions": "SERVER GUIDANCE"}}, {"Content-Type": "application/json", "Mcp-Session-Id": "sess-1"}) if method == "notifications/initialized": @@ -625,6 +627,10 @@ def fake_post(url, json=None, headers=None, timeout=None): 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" @@ -891,6 +897,16 @@ def test_live_cloud_contract_parity(): 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_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): @@ -916,6 +932,53 @@ def test_agent_instructions_with_doc_id(client, store_path): 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_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: From 50fd61862e062a49fe2670b0d32671605cb73e18 Mon Sep 17 00:00:00 2001 From: Ray Date: Tue, 11 Aug 2026 23:01:04 +0800 Subject: [PATCH 009/137] fix: local relevance sort answers honestly instead of imitating sort="relevance" is cloud-side semantic ranking; the local substring imitation could satisfy the letter of the interface while silently missing semantically relevant documents. Per the honest-subset rule (same treatment as folders), local now returns the "not available here" envelope for sort="relevance" or a stray query, and the local instructions steer discovery through name/description matching plus full-library paging instead of prescribing a capability that does not exist here. The tool schema keeps the cloud contract verbatim, like folder_id: honesty lives in the runtime answer, not a forked contract. --- pageindex/agent_tools.py | 85 +++++++++++++++------------------------ tests/test_agent_tools.py | 20 ++++----- 2 files changed, 43 insertions(+), 62 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index f1628d190..a8c15d761 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -636,16 +636,19 @@ def _browse_documents(client, folder_id: str = "root", recursive: bool = False, {"summary": "Invalid sort mode", "options": ['Use sort="time" or sort="relevance"']}, "INVALID_INPUT") - if sort == "relevance" and not query: - return _failure('query is required when sort is "relevance"', None, - {"summary": "Missing query for relevance ranking", - "options": ['Pass query alongside sort="relevance"']}, - "INVALID_INPUT") - if sort == "time" and query: - return _failure('query is only allowed when sort is "relevance"', None, - {"summary": "query does not apply to the time sort", - "options": ["Drop query, or set sort=\"relevance\""]}, - "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 available here โ€” use the default " + "time sort.", None, + {"summary": "Semantic ranking is not available in this library", + "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`"]}, + "INVALID_INPUT", + ) try: offset = max(int(offset), 0) limit = min(max(int(limit), 1), 50) @@ -655,23 +658,9 @@ def _browse_documents(client, folder_id: str = "root", recursive: bool = False, "options": ["Pass integer offset and limit values"]}, "INVALID_INPUT") - if sort == "relevance": - tokens = [token for token in (query or "").lower().split() if token] - scored = [] - for doc in _all_documents(client): - haystack = f"{doc.get('name') or ''} {doc.get('description') or ''}".lower() - score = sum(1 for token in tokens if token in haystack) - if score: - scored.append((score, doc)) - # Stable sort: equal scores keep the newest-first listing order. - scored.sort(key=lambda pair: pair[0], reverse=True) - ranked = [doc for _, doc in scored] - window = ranked[offset:offset + limit] - has_more = offset + limit < len(ranked) - else: - listing = client.list_documents(limit=limit, offset=offset) - window = listing.get("documents") or [] - has_more = offset + limit < listing.get("total", 0) + listing = client.list_documents(limit=limit, offset=offset) + window = listing.get("documents") or [] + has_more = offset + limit < listing.get("total", 0) next_offset = offset + limit if has_more else None page_has_processing = False @@ -706,18 +695,9 @@ def _browse_documents(client, folder_id: str = "root", recursive: bool = False, if not items and offset == 0: next_steps = { "summary": "Nothing to show", - "options": ( - ["No documents matched this query. Rephrase with synonyms or " - "alternative terms and retry browse_documents(sort=\"relevance\")."] - if sort == "relevance" - else ["Nothing here. Index documents with " - "PageIndexClient.submit_document() to get started."] - ), - "auto_retry": ( - "Rephrase the query and retry browse_documents(sort=\"relevance\")" - if sort == "relevance" - else "Index a document with submit_document() to get started" - ), + "options": ["Nothing here. Index documents with " + "PageIndexClient.submit_document() to get started."], + "auto_retry": "Index a document with submit_document() to get started", } return _success(data, next_steps) @@ -727,9 +707,8 @@ def _browse_documents(client, folder_id: str = "root", recursive: bool = False, 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, rephrase the query and " - "retry browse_documents(sort=\"relevance\"). Do NOT use general " - "knowledge as a substitute." + "before proceeding. If they do not match, page through the rest " + "of the library. Do NOT use general knowledge as a substitute." ) if page_has_processing: options.append("Some documents on this page are still processing. " @@ -1305,10 +1284,12 @@ def remove_document(doc_names: list[str], # โ”€โ”€ agent instructions โ”€โ”€ # Local subset of the cloud MCP server's initialize instructions (its -# no-folders variant), trimmed to the tools that exist here: the -# search_documents escalation steps, get_document_image, and the shared -# read-only-folders block are removed. Cloud clients receive the server's -# live instructions instead โ€” see _base_instructions(). +# no-folders variant), trimmed to what exists here: the search_documents +# escalation steps, get_document_image, and the shared read-only-folders +# block are removed, and the sort="relevance" guidance is replaced with +# name/description matching (semantic ranking is cloud-side). Cloud +# clients receive the server's live instructions instead โ€” see +# _base_instructions(). _INSTRUCTIONS_HEADER = ( "PageIndex by Vectify AI is a document platform for uploading and " @@ -1328,12 +1309,12 @@ def remove_document(doc_names: list[str], _DISCOVERY = """\ DOCUMENT DISCOVERY: -- browse_documents() โ€” DEFAULT discovery tool, first choice for any document-related question. The bare call returns your documents. Use sort="relevance" + query for semantic ranking.""" +- 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 (time) -- ANY question that needs a document to answer (including "find THE paper about Y") โ†’ browse_documents(sort="relevance", query=โ€ฆ)""" +- "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"). @@ -1343,9 +1324,9 @@ def remove_document(doc_names: list[str], _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(sort="relevance", query=โ€ฆ) with the original intent -2. Rephrase the query with synonyms or alternative terms โ†’ browse_documents(sort="relevance") again -3. browse_documents(recursive=true) to flatten the library into one list โ€” MANDATORY, must be attempted at least once before concluding "not found" +1. browse_documents() and compare every returned name/description against the user's intent +2. Page through the ENTIRE library with `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([ diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index b7677f1d7..69c8e9fee 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -119,21 +119,20 @@ def test_browse_documents_pagination(client, store_path): assert second["has_more"] is False -def test_browse_documents_relevance(client, store_path): - seed_doc(store_path, "pi-a", "annual-report.pdf", - description="Financial results for the year") - seed_doc(store_path, "pi-b", "attention.pdf", +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 not is_error - assert [d["name"] for d in payload["documents"]] == ["attention.pdf"] - assert payload["sort"] == "relevance" + assert is_error and payload["errorCode"] == "INVALID_INPUT" + assert "not available" in payload["error"] - missing_query, is_error = run(client, "browse_documents", sort="relevance") - assert is_error and missing_query["errorCode"] == "INVALID_INPUT" stray_query, is_error = run(client, "browse_documents", query="x") - assert is_error and "relevance" in stray_query["error"] + assert is_error and "not available" in stray_query["error"] + bad_sort, is_error = run(client, "browse_documents", sort="banana") + assert is_error and bad_sort["errorCode"] == "INVALID_INPUT" def test_browse_documents_empty_and_folder_error(client): @@ -916,6 +915,7 @@ def test_agent_instructions_default(client): 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): From f1301e7241cc313585b5f8be0070f49948182fb7 Mon Sep 17 00:00:00 2001 From: Ray Date: Tue, 11 Aug 2026 23:07:38 +0800 Subject: [PATCH 010/137] docs: note the cloud+Claude instructions duplication trade-off in as_claude_mcp --- pageindex/client.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/pageindex/client.py b/pageindex/client.py index 70608a480..a4f475ed1 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -498,6 +498,12 @@ def as_claude_mcp(self, include_management: bool = False): tools (requires ``claude-agent-sdk``; ``pip install 'pageindex[claude]'``). + 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:: options = ClaudeAgentOptions( From 20ed81f6b737005bdc9b1674d8d7fee76d365e66 Mon Sep 17 00:00:00 2001 From: Ray Date: Tue, 11 Aug 2026 23:17:41 +0800 Subject: [PATCH 011/137] fix: unsupported-capability envelopes say local-mode-yet, point to cloud "Not available here" read as a broken feature; the honest framing is that folders and semantic ranking exist on PageIndex cloud and are not in local mode yet. Both envelopes now say so and name the cloud client in next_steps, so agents relay an accurate story to the user. --- pageindex/agent_tools.py | 18 +++++++++++------- tests/test_agent_tools.py | 4 ++-- 2 files changed, 13 insertions(+), 9 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index a8c15d761..e0c813576 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -463,12 +463,14 @@ def _not_ready_error(doc_name: str, status: Any, operation: str, def _folder_unsupported(param: str) -> tuple[dict, bool]: return _failure( - f"Folders are not available here โ€” omit {param}.", + f"Folders are not supported in local mode yet โ€” omit {param}.", None, { - "summary": "This library has no folders", + "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"], + "Use browse_documents() to list the library root", + "Folders are available on PageIndex cloud " + "(PageIndexCloudClient with an API key)"], }, "INVALID_INPUT", ) @@ -640,13 +642,15 @@ def _browse_documents(client, folder_id: str = "root", recursive: bool = False, # Semantic ranking is a cloud capability; like folders, it is not # imitated here. return _failure( - "Relevance ranking is not available here โ€” use the default " - "time sort.", None, - {"summary": "Semantic ranking is not available in this library", + "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`"]}, + "`offset: next_offset`", + "Semantic ranking is available on PageIndex cloud " + "(PageIndexCloudClient with an API key)"]}, "INVALID_INPUT", ) try: diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 69c8e9fee..3a44cfeae 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -127,10 +127,10 @@ def test_browse_documents_relevance_unsupported(client, store_path): payload, is_error = run(client, "browse_documents", sort="relevance", query="attention transformers") assert is_error and payload["errorCode"] == "INVALID_INPUT" - assert "not available" in payload["error"] + assert "not supported in local mode" in payload["error"] stray_query, is_error = run(client, "browse_documents", query="x") - assert is_error and "not available" in stray_query["error"] + 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" From 8dc929ff1f8896aef28f5aa962c77f07a63f00a9 Mon Sep 17 00:00:00 2001 From: Ray Date: Tue, 11 Aug 2026 23:25:54 +0800 Subject: [PATCH 012/137] fix: local tool descriptions pre-announce cloud-only capabilities MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The cloud-verbatim browse_documents description invites sort="relevance" and folder drilling, so a local agent's first semantic search attempt was a guaranteed dead end discovered only from the runtime error envelope. Local registration now appends a LOCAL MODE note to the description โ€” the agent learns what is cloud-only before calling; the runtime envelope stays as the backstop for prompts that ignore descriptions. The cloud-facing contract stays byte-verbatim. --- pageindex/agent_tools.py | 23 +++++++++++++++++++--- pageindex/integrations/claude_agent_sdk.py | 5 +++-- tests/test_agent_tools.py | 10 ++++++++++ 3 files changed, 33 insertions(+), 5 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index e0c813576..a9c89be8e 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -1107,10 +1107,27 @@ def _tool_docstring(description: str, properties: dict[str, Any]) -> str: return "\n".join(lines) +#: Appended to the cloud-verbatim description when a tool is served locally, +#: so the agent learns what is cloud-only before calling instead of from the +#: runtime error envelope. +_LOCAL_DESCRIPTION_NOTES = { + "browse_documents": ( + 'LOCAL MODE: folder_id and sort="relevance"/query are not supported ' + "yet (they work on PageIndex cloud) โ€” use the default time sort and " + "page with offset." + ), +} + + +def _local_description(name: str) -> str: + description = TOOL_CONTRACT[name]["description"] + note = _LOCAL_DESCRIPTION_NOTES.get(name) + return f"{description}\n\n{note}" if note else description + + def _docstring(name: str) -> str: - contract = TOOL_CONTRACT[name] - return _tool_docstring(contract["description"], - contract["schema"]["properties"]) + return _tool_docstring(_local_description(name), + TOOL_CONTRACT[name]["schema"]["properties"]) _SCHEMA_TYPE_MAP = {"string": str, "integer": int, "number": float, diff --git a/pageindex/integrations/claude_agent_sdk.py b/pageindex/integrations/claude_agent_sdk.py index 0fb77d2de..77cc3d7a3 100644 --- a/pageindex/integrations/claude_agent_sdk.py +++ b/pageindex/integrations/claude_agent_sdk.py @@ -28,7 +28,8 @@ def build_claude_mcp(client, include_management: bool = False): "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, call_tool, tool_names + from ..agent_tools import (TOOL_CONTRACT, _local_description, call_tool, + tool_names) def make_handler(name: str): async def handler(arguments: dict[str, Any]) -> dict[str, Any]: @@ -52,7 +53,7 @@ def tool_kwargs(name: str) -> dict: return {"annotations": ToolAnnotations(**annotations)} tools = [ - tool(name, TOOL_CONTRACT[name]["description"], + tool(name, _local_description(name), TOOL_CONTRACT[name]["schema"], **tool_kwargs(name))(make_handler(name)) for name in tool_names(include_management) ] diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 3a44cfeae..41f710e7d 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -119,6 +119,16 @@ def test_browse_documents_pagination(client, store_path): assert second["has_more"] is False +def test_local_docstrings_preannounce_cloud_only_capabilities(client): + """The cloud-verbatim description invites sort="relevance" and folder + drilling; the local registration appends a LOCAL MODE note so the agent + learns the dead ends before calling, not from the runtime error.""" + browse = client.agent_tools()[0] + assert browse.__doc__.startswith(TOOL_CONTRACT["browse_documents"]["description"]) + assert "LOCAL MODE" in browse.__doc__ + assert "not supported yet" in browse.__doc__ + + 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.""" From 2b929eedfb5c78a0fe810522275d688b6b02d25c Mon Sep 17 00:00:00 2001 From: Ray Date: Tue, 11 Aug 2026 23:33:52 +0800 Subject: [PATCH 013/137] refactor: localized tool guidance replaces the appended LOCAL MODE note MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Appending a retraction to the cloud-verbatim description left the model parsing an instruction and its negation โ€” and kept the cloud text recommending search_documents and get_folder_structure, tools that are not registered locally (get_page_content likewise pointed at get_document_image). Guidance now adapts to the local surface the way AGENT_INSTRUCTIONS already does: schema structure stays byte-identical to the contract (mechanically asserted by a strip-descriptions test), while local description strings teach only what works here and point to PageIndex cloud for the rest. A dead-reference test forbids local guidance from naming tools outside the local registry, so a contract refresh that reintroduces a cloud-only reference fails loudly. --- pageindex/agent_tools.py | 90 ++++++++++++++++++---- pageindex/integrations/claude_agent_sdk.py | 6 +- tests/test_agent_tools.py | 55 +++++++++---- 3 files changed, 120 insertions(+), 31 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index a9c89be8e..ed161bf50 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -1,15 +1,19 @@ """Agent tools: the cloud MCP tool contract, executed against a PageIndexClient. -Tool names, input schemas, and descriptions match the PageIndex cloud MCP -server, so agent prompts work unchanged 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). +Tool names and input-schema structure match the PageIndex cloud MCP server, +so agent prompts work unchanged 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": ...}). """ from __future__ import annotations +import copy import difflib import json import re @@ -1107,27 +1111,83 @@ def _tool_docstring(description: str, properties: dict[str, Any]) -> str: return "\n".join(lines) -#: Appended to the cloud-verbatim description when a tool is served locally, -#: so the agent learns what is cloud-only before calling instead of from the -#: runtime error envelope. -_LOCAL_DESCRIPTION_NOTES = { +# Local guidance layer: schema STRUCTURE stays byte-identical to the cloud +# contract, but description strings adapt to the local surface the same way +# AGENT_INSTRUCTIONS does โ€” guidance must not teach capabilities (folders, +# semantic ranking) or tools (search_documents, get_document_image) that do +# not exist here. Guard tests assert both properties; a contract refresh +# that reintroduces a cloud-only reference fails the dead-reference test. + +_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_FOLDER_ID_DESCRIPTION = ( + "Not needed in local mode: document names are unique and folders are " + 'not supported yet (they work on PageIndex cloud). Omit, or pass "root".' +) + +_LOCAL_DESCRIPTIONS: dict[str, str] = { "browse_documents": ( - 'LOCAL MODE: folder_id and sort="relevance"/query are not supported ' - "yet (they work on PageIndex cloud) โ€” use the default time sort and " - "page with offset." + "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` while `has_more` is true. " + 'Folder browsing and semantic ranking (sort="relevance") are not ' + "supported in local mode yet โ€” they work on PageIndex cloud." + ), + # The image sentence points at a tool that is not registered locally. + "get_page_content": TOOL_CONTRACT["get_page_content"]["description"] + .replace(" Embedded image paths in the response feed into " + "`get_document_image()`.", ""), +} + +_LOCAL_PARAM_DESCRIPTIONS: dict[tuple[str, str], str] = { + ("browse_documents", "folder_id"): _LOCAL_FOLDER_ID_DESCRIPTION, + ("browse_documents", "recursive"): ( + "Kept for cloud compatibility; a local library has no folders, so " + "recursive and non-recursive return the same documents." + ), + ("browse_documents", "sort"): ( + 'Only "time" (newest first) is supported in local mode; ' + '"relevance" is cloud-only for now.' + ), + ("browse_documents", "query"): ( + 'Cloud-only for now (semantic ranking with sort="relevance") โ€” ' + "omit in local mode." ), + ("get_document", "doc_name"): _LOCAL_DOC_NAME_DESCRIPTION, + ("get_document", "folder_id"): _LOCAL_FOLDER_ID_DESCRIPTION, + ("get_document_structure", "doc_name"): _LOCAL_DOC_NAME_DESCRIPTION, + ("get_document_structure", "folder_id"): _LOCAL_FOLDER_ID_DESCRIPTION, + ("get_page_content", "doc_name"): _LOCAL_DOC_NAME_DESCRIPTION, + ("get_page_content", "folder_id"): _LOCAL_FOLDER_ID_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.' + ), + ("remove_document", "folder_id"): _LOCAL_FOLDER_ID_DESCRIPTION, } def _local_description(name: str) -> str: - description = TOOL_CONTRACT[name]["description"] - note = _LOCAL_DESCRIPTION_NOTES.get(name) - return f"{description}\n\n{note}" if note else description + 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 (tool_name, param), text in _LOCAL_PARAM_DESCRIPTIONS.items(): + if tool_name == name and param in schema["properties"]: + schema["properties"][param]["description"] = text + return schema def _docstring(name: str) -> str: return _tool_docstring(_local_description(name), - TOOL_CONTRACT[name]["schema"]["properties"]) + _local_schema(name)["properties"]) _SCHEMA_TYPE_MAP = {"string": str, "integer": int, "number": float, diff --git a/pageindex/integrations/claude_agent_sdk.py b/pageindex/integrations/claude_agent_sdk.py index 77cc3d7a3..b5e599da0 100644 --- a/pageindex/integrations/claude_agent_sdk.py +++ b/pageindex/integrations/claude_agent_sdk.py @@ -28,8 +28,8 @@ def build_claude_mcp(client, include_management: bool = False): "as_claude_mcp in local mode requires the Claude Agent SDK โ€” " "pip install claude-agent-sdk (or pip install 'pageindex[claude]')." ) from exc - from ..agent_tools import (TOOL_CONTRACT, _local_description, call_tool, - tool_names) + from ..agent_tools import (TOOL_CONTRACT, _local_description, + _local_schema, call_tool, tool_names) def make_handler(name: str): async def handler(arguments: dict[str, Any]) -> dict[str, Any]: @@ -54,7 +54,7 @@ def tool_kwargs(name: str) -> dict: tools = [ tool(name, _local_description(name), - TOOL_CONTRACT[name]["schema"], **tool_kwargs(name))(make_handler(name)) + _local_schema(name), **tool_kwargs(name))(make_handler(name)) for name in tool_names(include_management) ] return create_sdk_mcp_server(name="pageindex", version=sdk_version(), diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 41f710e7d..327cfa574 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -76,10 +76,49 @@ def test_tool_surface_and_docstrings(client): with_management = client.agent_tools(include_management=True) assert [tool.__name__ for tool in with_management][-1] == "remove_document" for tool in tools: - contract = TOOL_CONTRACT[tool.__name__] - assert tool.__doc__.startswith(contract["description"]) - for param in contract["schema"]["properties"]: + for param in TOOL_CONTRACT[tool.__name__]["schema"]["properties"]: assert param in tool.__doc__ + 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 guidance layer may localize description strings only โ€” + names, types, defaults, bounds, and required stay byte-identical.""" + import copy + from pageindex.agent_tools import _local_schema + + def stripped(schema): + schema = copy.deepcopy(schema) + for spec in schema["properties"].values(): + spec.pop("description", None) + return schema + + for name, contract in TOOL_CONTRACT.items(): + assert stripped(_local_schema(name)) == stripped(contract["schema"]), 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): + browse = client.agent_tools()[0].__doc__ + assert "not supported in local mode yet" in browse + assert "PageIndex cloud" in browse + assert "search_documents" not in browse + assert "get_folder_structure" not in browse # โ”€โ”€ browse_documents โ”€โ”€ @@ -119,16 +158,6 @@ def test_browse_documents_pagination(client, store_path): assert second["has_more"] is False -def test_local_docstrings_preannounce_cloud_only_capabilities(client): - """The cloud-verbatim description invites sort="relevance" and folder - drilling; the local registration appends a LOCAL MODE note so the agent - learns the dead ends before calling, not from the runtime error.""" - browse = client.agent_tools()[0] - assert browse.__doc__.startswith(TOOL_CONTRACT["browse_documents"]["description"]) - assert "LOCAL MODE" in browse.__doc__ - assert "not supported yet" in browse.__doc__ - - 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.""" From e790c375fefb7837c8b6ac36393999667c98ce0b Mon Sep 17 00:00:00 2001 From: Ray Date: Tue, 11 Aug 2026 23:41:27 +0800 Subject: [PATCH 014/137] feat: hide cloud-only parameters from the local tool surface MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit folder_id, sort, query, and recursive were exposed locally with localized "cloud-only" descriptions, leaving the dead-end calls expressible and discovered at runtime. Schema constraints beat guidance: the local surface now serves the contract minus these parameters, so strict-schema frameworks make the calls inexpressible and a prompt that insists on sort="relevance" degrades to the bare call (the correct local behavior) instead of an error round-trip. The implementations still accept the hidden parameters and answer with the guided "works on PageIndex cloud" envelope โ€” the backstop for direct call_tool callers and hosts without schema enforcement. wait_for_completion stays: seeded or torn stores can hold documents that are genuinely not completed. The structural guard now asserts the local schema equals the contract minus the documented hidden set, descriptions aside. --- pageindex/agent_tools.py | 72 ++++++++++++++++----------------------- tests/test_agent_tools.py | 31 +++++++++++++---- 2 files changed, 54 insertions(+), 49 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index ed161bf50..b6ed15c6f 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -1112,21 +1112,31 @@ def _tool_docstring(description: str, properties: dict[str, Any]) -> str: # Local guidance layer: schema STRUCTURE stays byte-identical to the cloud -# contract, but description strings adapt to the local surface the same way -# AGENT_INSTRUCTIONS does โ€” guidance must not teach capabilities (folders, -# semantic ranking) or tools (search_documents, get_document_image) that do -# not exist here. Guard tests assert both properties; a contract refresh -# that reintroduces a cloud-only reference fails the dead-reference test. +# contract minus the hidden cloud-only parameters, and description strings +# adapt to the local surface the same way AGENT_INSTRUCTIONS does โ€” guidance +# must not teach capabilities (folders, semantic ranking) or tools +# (search_documents, get_document_image) that do not exist here. Guard tests +# assert both properties; a contract refresh that reintroduces a cloud-only +# reference fails the dead-reference test. + +#: Cloud-only parameters hidden from the local surface โ€” strict-schema +#: frameworks then make the dead-end calls inexpressible. The +#: implementations still accept them and answer with the guided error +#: envelope, for direct call_tool callers and hosts without schema +#: enforcement. +_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_FOLDER_ID_DESCRIPTION = ( - "Not needed in local mode: document names are unique and folders are " - 'not supported yet (they work on PageIndex cloud). Omit, or pass "root".' -) _LOCAL_DESCRIPTIONS: dict[str, str] = { "browse_documents": ( @@ -1144,32 +1154,15 @@ def _tool_docstring(description: str, properties: dict[str, Any]) -> str: } _LOCAL_PARAM_DESCRIPTIONS: dict[tuple[str, str], str] = { - ("browse_documents", "folder_id"): _LOCAL_FOLDER_ID_DESCRIPTION, - ("browse_documents", "recursive"): ( - "Kept for cloud compatibility; a local library has no folders, so " - "recursive and non-recursive return the same documents." - ), - ("browse_documents", "sort"): ( - 'Only "time" (newest first) is supported in local mode; ' - '"relevance" is cloud-only for now.' - ), - ("browse_documents", "query"): ( - 'Cloud-only for now (semantic ranking with sort="relevance") โ€” ' - "omit in local mode." - ), ("get_document", "doc_name"): _LOCAL_DOC_NAME_DESCRIPTION, - ("get_document", "folder_id"): _LOCAL_FOLDER_ID_DESCRIPTION, ("get_document_structure", "doc_name"): _LOCAL_DOC_NAME_DESCRIPTION, - ("get_document_structure", "folder_id"): _LOCAL_FOLDER_ID_DESCRIPTION, ("get_page_content", "doc_name"): _LOCAL_DOC_NAME_DESCRIPTION, - ("get_page_content", "folder_id"): _LOCAL_FOLDER_ID_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.' ), - ("remove_document", "folder_id"): _LOCAL_FOLDER_ID_DESCRIPTION, } @@ -1179,6 +1172,8 @@ def _local_description(name: str) -> str: 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 @@ -1311,41 +1306,34 @@ def build_agent_tools(client, include_management: bool = False) -> list[Callable if getattr(client, "api_key", None): return _build_cloud_agent_tools(client, include_management) - def browse_documents(folder_id: str = "root", recursive: bool = False, - sort: str = "time", query: Optional[str] = None, - offset: int = 0, limit: int = 10) -> str: + def browse_documents(offset: int = 0, limit: int = 10) -> str: return call_tool(client, "browse_documents", { - "folder_id": folder_id, "recursive": recursive, "sort": sort, - "query": query, "offset": offset, "limit": limit, + "offset": offset, "limit": limit, })[0] - def get_document(doc_name: str, folder_id: Optional[str] = None, - wait_for_completion: bool = False) -> str: + def get_document(doc_name: str, wait_for_completion: bool = False) -> str: return call_tool(client, "get_document", { - "doc_name": doc_name, "folder_id": folder_id, + "doc_name": doc_name, "wait_for_completion": wait_for_completion, })[0] - def get_document_structure(doc_name: str, folder_id: Optional[str] = None, - part: int = 1, + def get_document_structure(doc_name: str, part: int = 1, wait_for_completion: bool = False) -> str: return call_tool(client, "get_document_structure", { - "doc_name": doc_name, "folder_id": folder_id, "part": part, + "doc_name": doc_name, "part": part, "wait_for_completion": wait_for_completion, })[0] def get_page_content(doc_name: str, pages: str, - folder_id: Optional[str] = None, wait_for_completion: bool = False) -> str: return call_tool(client, "get_page_content", { - "doc_name": doc_name, "pages": pages, "folder_id": folder_id, + "doc_name": doc_name, "pages": pages, "wait_for_completion": wait_for_completion, })[0] - def remove_document(doc_names: list[str], - folder_id: Optional[str] = None) -> str: + def remove_document(doc_names: list[str]) -> str: return call_tool(client, "remove_document", { - "doc_names": doc_names, "folder_id": folder_id, + "doc_names": doc_names, })[0] functions = { diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 327cfa574..1d8b7e6d5 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -71,13 +71,23 @@ def test_contract_matches_snapshot(): 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 tools: - for param in TOOL_CONTRACT[tool.__name__]["schema"]["properties"]: + 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. @@ -88,19 +98,26 @@ def test_tool_surface_and_docstrings(client): def test_local_schema_structure_matches_contract(): - """The local guidance layer may localize description strings only โ€” - names, types, defaults, bounds, and required stay byte-identical.""" + """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_schema + from pageindex.agent_tools import _LOCAL_HIDDEN_PARAMS, _local_schema - def stripped(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(): - assert stripped(_local_schema(name)) == stripped(contract["schema"]), name + 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): From 1fa3eb7fb3d23763598abc238c5ec0e6583f2778 Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 12 Aug 2026 00:18:08 +0800 Subject: [PATCH 015/137] =?UTF-8?q?fix:=20incremental-review=20findings=20?= =?UTF-8?q?=E2=80=94=20bridge=20cache,=20guards,=20envelope=20drift?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three independent review passes over the agent-instructions increment surfaced six fixes: - The per-client bridge moved off the instance into a weak-keyed, lock-guarded module cache: cloud clients stay picklable (threading.RLock no longer rides on the client) and concurrent first calls can no longer construct duplicate bridges/sessions. - Blank or non-string initialize.instructions now hit the same honest error as a missing one โ€” a whitespace-only or structured value could previously become the system prompt (or crash the doc_id append with a raw TypeError). - The invalid-sort envelope no longer prescribes sort="relevance" โ€” the one error text that still taught the cloud-only value it would then reject. - "Page through the rest of the library" is emitted only when has_more is true; a fully-listed library no longer instructs a pointless call. - The mandatory full-library paging step now says limit: 50 โ€” 6 calls instead of 30 on a 300-document library. - Docstrings and comments rescoped to what is actually true: the never-raise contract covers invocations the signatures accept (unknown params fail at the Python boundary; call_tool answers them with the guided envelope), recursive is accepted as the identity rather than errored, lenient framework arg models drop hidden params pre-call, and the module header no longer claims full schema parity. The capability-phrase guard now covers every local docstring, not just browse_documents. --- examples/documents/attention-residuals.doc_id | 1 + pageindex/agent_tools.py | 102 +++++++++++------- tests/test_agent_tools.py | 71 +++++++++++- 3 files changed, 133 insertions(+), 41 deletions(-) create mode 100644 examples/documents/attention-residuals.doc_id diff --git a/examples/documents/attention-residuals.doc_id b/examples/documents/attention-residuals.doc_id new file mode 100644 index 000000000..19003f3ce --- /dev/null +++ b/examples/documents/attention-residuals.doc_id @@ -0,0 +1 @@ +pi-c4a794161a904216bde92b3d2c269a19 \ No newline at end of file diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index b6ed15c6f..2399ca0e7 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -1,15 +1,19 @@ """Agent tools: the cloud MCP tool contract, executed against a PageIndexClient. -Tool names and input-schema structure match the PageIndex cloud MCP server, -so agent prompts work unchanged 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": ...}). +Tool names and the surviving input-schema structure match the PageIndex +cloud MCP server โ€” the local surface hides the documented cloud-only +parameters โ€” so agent prompts port across the cloud MCP connection and +this in-process layer. Only the tools that exist in every mode are +registered (no folders, search_documents, or get_document_image), and the +guidance strings (tool descriptions) adapt to the local surface the same +way the agent instructions do โ€” they never teach capabilities that only +exist on the cloud. + +Tools never raise for any invocation their signatures accept: every +outcome, including errors, is returned as the same JSON envelope the cloud +emits ({"success": true, ...} / {"error": ...}). Arguments outside a pruned +local signature fail at the Python call boundary; the call_tool path +answers them with the guided error envelope instead. """ from __future__ import annotations @@ -17,7 +21,9 @@ import difflib import json import re +import threading import time +import weakref from typing import Any, Callable, Optional from .errors import PageIndexAPIError @@ -638,10 +644,15 @@ def _browse_documents(client, folder_id: str = "root", recursive: bool = False, if folder_id != "root": return _folder_unsupported("folder_id") if sort not in ("time", "relevance"): - return _failure('sort must be "time" or "relevance"', None, - {"summary": "Invalid sort mode", - "options": ['Use sort="time" or sort="relevance"']}, - "INVALID_INPUT") + 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. @@ -715,8 +726,10 @@ def _browse_documents(client, folder_id: str = "root", recursive: bool = False, 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. Do NOT use general knowledge as a substitute." + "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. " @@ -1115,15 +1128,19 @@ def _tool_docstring(description: str, properties: dict[str, Any]) -> str: # contract minus the hidden cloud-only parameters, and description strings # adapt to the local surface the same way AGENT_INSTRUCTIONS does โ€” guidance # must not teach capabilities (folders, semantic ranking) or tools -# (search_documents, get_document_image) that do not exist here. Guard tests -# assert both properties; a contract refresh that reintroduces a cloud-only -# reference fails the dead-reference test. +# (search_documents, get_document_image) that do not exist here. Guard +# tests pin structure (contract-minus-hidden equality), tool references +# (the dead-reference test), and capability phrases (the per-docstring +# phrase test) โ€” a contract refresh that reintroduces a cloud-only +# reference fails loudly. #: Cloud-only parameters hidden from the local surface โ€” strict-schema -#: frameworks then make the dead-end calls inexpressible. The -#: implementations still accept them and answer with the guided error -#: envelope, for direct call_tool callers and hosts without schema -#: enforcement. +#: frameworks make the dead-end calls inexpressible, and lenient framework +#: argument models drop them before the call (degrading to the bare call). +#: The call_tool path still answers folder_id/sort/query with the guided +#: error envelope; recursive is simply accepted (flattening a folderless +#: library is the identity). Plain functions reject unknown parameters at +#: the Python call boundary. _LOCAL_HIDDEN_PARAMS: dict[str, tuple[str, ...]] = { "browse_documents": ("folder_id", "recursive", "sort", "query"), "get_document": ("folder_id",), @@ -1143,7 +1160,8 @@ def _tool_docstring(description: str, properties: dict[str, Any]) -> str: "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` while `has_more` is true. " + "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." ), @@ -1262,18 +1280,24 @@ def proxy(**kwargs: Any) -> str: return proxy +_BRIDGES: "weakref.WeakKeyDictionary" = weakref.WeakKeyDictionary() +_BRIDGES_LOCK = threading.Lock() + + def _cloud_bridge(client): - """One bridge per client instance: tool discovery and instructions share - a single MCP session.""" - bridge = getattr(client, "_mcp_bridge", None) - if bridge is None: - from .mcp_bridge import McpBridge - bridge = McpBridge( - f"{client.BASE_URL}/mcp", - {"Authorization": f"Bearer {client.api_key}"}, - ) - client._mcp_bridge = bridge - return bridge + """One bridge per client: tool discovery and instructions share a single + MCP session. Weak-keyed off the instance so clients stay picklable; the + lock closes the check-then-set race under concurrent first calls.""" + with _BRIDGES_LOCK: + bridge = _BRIDGES.get(client) + if bridge is None: + from .mcp_bridge import McpBridge + bridge = McpBridge( + f"{client.BASE_URL}/mcp", + {"Authorization": f"Bearer {client.api_key}"}, + ) + _BRIDGES[client] = bridge + return bridge def _build_cloud_agent_tools(client, include_management: bool) -> list[Callable[..., str]]: @@ -1301,7 +1325,9 @@ def build_agent_tools(client, include_management: bool = False) -> list[Callable 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. + the JSON envelope as a string and never raises for arguments its + signature accepts (cloud-only parameters are absent from the local + signatures; the call_tool path answers them with the guided envelope). """ if getattr(client, "api_key", None): return _build_cloud_agent_tools(client, include_management) @@ -1394,7 +1420,7 @@ def remove_document(doc_names: list[str]) -> str: 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 `offset: next_offset` until has_more is false โ€” MANDATORY, must be completed before concluding "not found" +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.""" @@ -1415,7 +1441,7 @@ def _base_instructions(client) -> str: if not getattr(client, "api_key", None): return AGENT_INSTRUCTIONS instructions = _cloud_bridge(client).instructions() - if not 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 " diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 1d8b7e6d5..63189b3ba 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -131,11 +131,20 @@ def test_local_guidance_references_only_local_tools(client): def test_local_guidance_points_cloud_only_capabilities_at_cloud(client): - browse = client.agent_tools()[0].__doc__ + 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 - assert "search_documents" not in browse - assert "get_folder_structure" not 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 โ”€โ”€ @@ -170,9 +179,12 @@ def test_browse_documents_pagination(client, store_path): 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): @@ -189,6 +201,9 @@ def test_browse_documents_relevance_unsupported(client, store_path): 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): @@ -1017,6 +1032,56 @@ def instructions(self): assert len(created) == 1 +def test_cloud_bridge_cache_threadsafe_and_pickle_clean(monkeypatch): + """One bridge per client even under concurrent first calls, and the + bridge lives off the instance so cloud clients stay picklable.""" + import pickle + import threading + import time as time_mod + import pageindex.mcp_bridge as mcp_bridge + created = [] + + class _Bridge(_FakeBridge): + def __init__(self, url, headers): + time_mod.sleep(0.01) # widen the construction window + super().__init__(url, headers) + created.append(self) + + def instructions(self): + return "LIVE" + + monkeypatch.setattr(mcp_bridge, "McpBridge", _Bridge) + cloud = PageIndexCloudClient(api_key="pi-test-key") + workers = ([threading.Thread(target=cloud.agent_tools) for _ in range(4)] + + [threading.Thread(target=cloud.agent_instructions) + for _ in range(4)]) + for worker in workers: + worker.start() + for worker in workers: + worker.join() + assert len(created) == 1 + pickle.dumps(cloud) + + +def test_cloud_agent_instructions_blank_or_nonstring_raises(monkeypatch): + """Whitespace-only or non-string initialize.instructions must hit the + same honest error as a missing one โ€” never a blank system prompt.""" + import pageindex.mcp_bridge as mcp_bridge + + for bad in (" \n\t ", {"not": "a string"}): + class _SilentBridge: + def __init__(self, url, headers): + pass + + def instructions(self, _value=bad): + return _value + + monkeypatch.setattr(mcp_bridge, "McpBridge", _SilentBridge) + cloud = PageIndexCloudClient(api_key="pi-test-key") + with pytest.raises(PageIndexAPIError, match="no agent instructions"): + cloud.agent_instructions() + + def test_cloud_agent_instructions_empty_raises(monkeypatch): """An empty server response must raise, not silently substitute the subset guidance โ€” same posture as the annotation-regression guard.""" From 63b767f70889999e9be7b0a1d61272fe54ac5ab3 Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 12 Aug 2026 00:18:24 +0800 Subject: [PATCH 016/137] chore: keep the demo's doc_id cache file out of the repo --- .gitignore | 1 + examples/documents/attention-residuals.doc_id | 1 - 2 files changed, 1 insertion(+), 1 deletion(-) delete mode 100644 examples/documents/attention-residuals.doc_id diff --git a/.gitignore b/.gitignore index 5193735ca..b5c223b31 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,4 @@ __pycache__ logs/ .pageindex/ dist/ +*.doc_id diff --git a/examples/documents/attention-residuals.doc_id b/examples/documents/attention-residuals.doc_id deleted file mode 100644 index 19003f3ce..000000000 --- a/examples/documents/attention-residuals.doc_id +++ /dev/null @@ -1 +0,0 @@ -pi-c4a794161a904216bde92b3d2c269a19 \ No newline at end of file From 6c9fe2e544c400674daded224b2f399bcf0152c3 Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 12 Aug 2026 00:29:48 +0800 Subject: [PATCH 017/137] test: live envelope field-parity guard against cloud response drift The frozen contract guards tools/list, but the response envelopes the local tools emit were hand-built to mirror the cloud's and had no drift detector. A key-gated live test now asserts every field local emits exists in the live cloud response for the analogous call (top-level keys, next_steps, document entries, structure nodes, content entries). Guidance wording is deliberately localized and not compared. Verified green against the live server: local and cloud field structures currently match exactly. --- tests/test_agent_tools.py | 57 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 57 insertions(+) diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 63189b3ba..2a3cbc70b 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -967,6 +967,63 @@ def test_live_cloud_contract_parity(): 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})) + 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})), + "get_document_structure": json.loads(bridge.call_tool( + "get_document_structure", {"doc_name": doc_name})), + "get_page_content": json.loads(bridge.call_tool( + "get_page_content", {"doc_name": doc_name, "pages": "1"})), + } + + 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 From a45b55418d5a2de231683b6c8c6b6b73d01aeca8 Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 12 Aug 2026 01:44:05 +0800 Subject: [PATCH 018/137] =?UTF-8?q?feat:=20local=20chat=20=E2=80=94=20thre?= =?UTF-8?q?e=20protocol=20surfaces=20over=20the=20agent=20tools=20(v0.2.10?= =?UTF-8?q?)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Local mode gains managed document QA: an agent over the #393 local tool set, reachable through three wire protocols, each 1:1 with the backend and with no translation layer. - chat_completions(): standard chat.completions semantics on any OpenAI-compatible backend (openai-agents engine). Final answer only, cross-turn aggregated usage, streaming as text pieces or chunk dicts (the existing cloud signature, now implemented locally; model and max_turns are local-only additions). - responses(): the agentic surface โ€” OpenAI Responses format, the tool process is standard output items, streaming forwards native events (tool outputs emitted as response.output_item.done, the way the platform streams its own server-side tools). Round-tripping output into the next input keeps provider prompt-cache prefix continuity and the agent's memory โ€” live-verified: the follow-up call answered from round-tripped tool output with zero new tool calls. - messages(): Anthropic-native via the SDK's own tool runner (new pageindex[anthropic] extra, floor 0.68.0 verified for tool_runner/beta_tool(input_schema)). tool_use/tool_result round-trip is the format's native behavior; the envelope is the final message with aggregated usage plus the full new-turn sequence; the managed system blocks carry cache_control breakpoints. Shared skeleton: thin chat header + the local AGENT_INSTRUCTIONS (caller system content is appended, not rejected), the doc_id targeting block as a leading context item (factored out of build_agent_instructions), read-only toolset, structural-only validation (no arbitrary caps โ€” backend limits govern), sampling params passed through, per-run tracing disabled, enable_citations rejected as cloud-only. Design basis is industry-standard formats rather than the cloud chat endpoint; responses()/messages() raise on cloud clients until the cloud converges. Tests run the real engines against scripted backends (a Model fake for openai-agents, a mock HTTP transport under the real anthropic SDK) with real tool execution against a seeded store, including the round-trip prefix-extension assertions on both engines. --- .github/workflows/tests.yml | 2 +- pageindex/agent_tools.py | 40 ++-- pageindex/client.py | 145 ++++++++++- pageindex/local_chat.py | 464 ++++++++++++++++++++++++++++++++++++ pyproject.toml | 6 +- tests/test_client.py | 8 +- tests/test_local_chat.py | 426 +++++++++++++++++++++++++++++++++ 7 files changed, 1059 insertions(+), 32 deletions(-) create mode 100644 pageindex/local_chat.py create mode 100644 tests/test_local_chat.py diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index d8d9dbb38..6b349a09b 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -31,5 +31,5 @@ jobs: cache: pip - run: pip install -r requirements.txt pytest - if: matrix.agent-frameworks == 'with' - run: pip install openai-agents claude-agent-sdk + run: pip install openai-agents claude-agent-sdk anthropic - run: python -m pytest -q diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index 2399ca0e7..bf8627d91 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -1450,17 +1450,17 @@ def _base_instructions(client) -> str: return instructions -def build_agent_instructions(client, doc_id=None) -> str: - """Orchestration guidance for document QA agents; with doc_id, appends - the target documents and directs the agent to work within them. Raises - when a doc_id's name is shadowed by a newer same-name document โ€” the +def doc_targeting_block(client, doc_id) -> 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 (which place it as a leading conversation item). Raises when a + doc_id's name is shadowed by a newer same-name document โ€” the name-addressed tools could not reach it.""" - base = _base_instructions(client) if doc_id is None: - return base + return None doc_ids = [doc_id] if isinstance(doc_id, str) else list(doc_id) if not doc_ids: - return base + return None details = [client.get_document(one_id) for one_id in doc_ids] documents = _all_documents(client) for one_id, detail in zip(doc_ids, details): @@ -1476,18 +1476,24 @@ def build_agent_instructions(client, doc_id=None) -> str: ) context = json.dumps(details, ensure_ascii=False) if len(details) == 1: - block = ( + 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()." ) - else: - names = ", ".join(str(item.get("name")) for item in details) - block = ( - 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()." - ) - return base + "\n\n" + block + 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) -> str: + """Orchestration guidance for document QA agents; with doc_id, appends + the target documents and directs the agent to work within them.""" + base = _base_instructions(client) + block = doc_targeting_block(client, doc_id) + return base if block is None else base + "\n\n" + block diff --git a/pageindex/client.py b/pageindex/client.py index a4f475ed1..91e4f18a8 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -344,38 +344,161 @@ def chat_completions( temperature: Optional[float] = None, stream_metadata: bool = False, enable_citations: bool = False, + model: Optional[str] = None, + max_turns: Optional[int] = None, ) -> Union[dict[str, Any], Iterator[str], Iterator[dict[str, Any]]]: """ - PageIndex Chat Completions, scoped to specific PageIndex documents. + PageIndex Chat Completions: document QA in one call. + + Cloud: the hosted chat endpoint. Local: a managed document-QA agent + run over the local tools against your own LLM backend's + /chat/completions (requires ``pageindex[openai]``; the OpenAI SDK's + usual env config โ€” OPENAI_API_KEY, OPENAI_BASE_URL โ€” selects the + backend, so any OpenAI-compatible server works). The response + carries the final answer only; for the tool-use process and + prompt-cache round-trip use ``responses()`` or ``messages()``. Args: messages: Conversation messages with 'role' and 'content' keys. + Local also accepts system/developer messages โ€” their content + is appended to the managed system prompt. stream: Enable streaming responses. doc_id: Document ID or list of IDs to scope the conversation. - temperature: Sampling temperature (0.0-1.0). + temperature: Sampling temperature, passed through to the model. stream_metadata: With stream=True, yield chunk dicts instead of text pieces. - enable_citations: Enable citation instructions in responses. + enable_citations: Cloud-only โ€” local mode raises (citations need + block-level OCR data local mode does not store). + model: Local only โ€” backend model name (defaults to + ``retrieve_model``). The cloud endpoint selects its own. + max_turns: Local only โ€” cap on agent turns per call. Returns: - stream=False: complete response dict ({'id', 'object', 'created', 'choices', 'usage'}) - stream=True, stream_metadata=False: iterator of text chunks - stream=True, stream_metadata=True: iterator of chunk dicts - - Local: not yet supported โ€” raises PageIndexAPIError. Agent-based - local chat arrives in a later release. """ - return self._require_cloud( - "chat_completions is not yet supported in local mode โ€” it arrives " - "in a later release. Create the client with an api_key to use " - "cloud chat." - ).chat_completions( + from .cloud_api import CloudAPI + if not isinstance(self._api, CloudAPI): + from .local_chat import run_chat_completions + return run_chat_completions( + self, messages, stream=stream, doc_id=doc_id, + temperature=temperature, stream_metadata=stream_metadata, + enable_citations=enable_citations, model=model, + max_turns=max_turns, + ) + if model is not None or max_turns is not None: + raise PageIndexAPIError( + "model and max_turns are local-mode parameters โ€” the cloud " + "chat endpoint selects its own model." + ) + return self._api.chat_completions( messages=messages, stream=stream, doc_id=doc_id, temperature=temperature, stream_metadata=stream_metadata, enable_citations=enable_citations, ) + def responses( + self, + input: Union[str, list[dict[str, Any]]], + model: Optional[str] = None, + stream: bool = False, + doc_id: Optional[Union[str, list[str]]] = None, + instructions: Optional[str] = None, + temperature: Optional[float] = None, + top_p: Optional[float] = None, + max_turns: Optional[int] = None, + ) -> Union[dict[str, Any], Iterator[dict[str, Any]]]: + """ + Document QA over the OpenAI Responses protocol โ€” the agentic surface. + + Local only for now. Drives your backend's /responses end to end (no + translation layer), so the ``output`` carries the whole process as + standard items โ€” messages, function calls, and function outputs + (the SDK executes the tools). Append the returned ``output`` to your + next call's ``input`` verbatim to keep provider prompt-cache prefix + continuity and the agent's memory of what it already read. + + Requires ``pageindex[openai]`` and a backend that supports the + Responses API; backends that only speak chat.completions should use + ``chat_completions()``. + + Args: + input: A user message string, or a list of Responses input items + (round-trip prior ``output`` items here). + model: Backend model name (defaults to ``retrieve_model``). + stream: Yield native Responses stream events as dicts; tool + outputs are emitted as ``response.output_item.done`` events + and the final event is ``response.completed``. + doc_id: Document ID or list of IDs to scope the conversation. + instructions: Appended to the managed system prompt. + temperature / top_p: Passed through to the model. + max_turns: Cap on agent turns per call. + """ + from .cloud_api import CloudAPI + if isinstance(self._api, CloudAPI): + raise PageIndexAPIError( + "responses is not available on PageIndex cloud yet โ€” it is " + "a local-mode surface for now." + ) + from .local_chat import run_responses + return run_responses( + self, input, model=model, stream=stream, doc_id=doc_id, + instructions=instructions, temperature=temperature, top_p=top_p, + max_turns=max_turns, + ) + + def messages( + self, + messages: list[dict[str, Any]], + model: str, + max_tokens: int, + stream: bool = False, + doc_id: Optional[Union[str, list[str]]] = None, + system: Optional[Union[str, list[dict[str, Any]]]] = None, + temperature: Optional[float] = None, + top_p: Optional[float] = None, + top_k: Optional[int] = None, + stop_sequences: Optional[list[str]] = None, + max_turns: Optional[int] = None, + ) -> Union[dict[str, Any], Iterator[Any]]: + """ + Document QA over the Anthropic Messages protocol โ€” Claude-native. + + Local only for now. Drives Anthropic's /v1/messages via the + Anthropic SDK's own tool runner (requires ``pageindex[anthropic]``; + ANTHROPIC_API_KEY selects the backend). ``tool_use``/``tool_result`` + round-trip is the format's native behavior: the response is the + final message envelope with cross-turn aggregated ``usage`` plus a + ``messages`` field โ€” the full new turn sequence, valid for verbatim + append to your history. The managed system prompt and the doc + targeting block carry ``cache_control`` breakpoints. + + Args: + messages: Native Messages-format history (including prior + tool_use/tool_result blocks on round-trip). + model / max_tokens: Required by the Messages API; passed through. + stream: Yield the native event stream across turns, verbatim. + doc_id: Document ID or list of IDs to scope the conversation. + system: Appended after the managed system blocks. + temperature / top_p / top_k / stop_sequences: Passed through. + max_turns: Cap on agent turns per call. + """ + from .cloud_api import CloudAPI + if isinstance(self._api, CloudAPI): + raise PageIndexAPIError( + "messages is not available on PageIndex cloud yet โ€” it is " + "a local-mode surface for now." + ) + from .local_chat import run_messages + return run_messages( + self, messages, model=model, max_tokens=max_tokens, + stream=stream, doc_id=doc_id, system=system, + temperature=temperature, top_p=top_p, top_k=top_k, + stop_sequences=stop_sequences, max_turns=max_turns, + ) + # ---------- DOCUMENT MANAGEMENT ---------- def get_document(self, doc_id: str) -> dict[str, Any]: diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py new file mode 100644 index 000000000..3497f3970 --- /dev/null +++ b/pageindex/local_chat.py @@ -0,0 +1,464 @@ +"""Managed local chat: document-QA agents over the local tools. + +Three methods, three wire protocols, 1:1 with the backend and no translation +layer: ``chat_completions`` drives the backend's /chat/completions (any +OpenAI-compatible backend, final answer only), ``responses`` drives +/responses (process items are standard output; round-trip them for provider +prompt-cache continuation and agent memory), ``messages`` drives Anthropic's +/v1/messages via the SDK's own tool runner (tool_use/tool_result round-trip +is the format's native behavior). + +Content passes through untouched โ€” the caller's messages, the model's +answers, tool outputs, finish/stop reasons. The SDK owns only gatekeeping +(structural validation), table-setting (managed instructions, tools, doc +targeting), tool execution, and billing (usage aggregation, envelope ids). +""" +from __future__ import annotations + +import asyncio +import concurrent.futures +import queue +import threading +import time +import uuid +from typing import Any, Iterator, Optional, Union + +from .agent_tools import (AGENT_INSTRUCTIONS, _local_description, + _local_schema, call_tool, doc_targeting_block, + tool_names) +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 + doc_ids = [doc_id] if isinstance(doc_id, str) else list(doc_id) + missing = [] + for one_id in doc_ids: + try: + client.get_document(one_id) + except PageIndexAPIError: + missing.append(str(one_id)) + if missing: + raise PageIndexAPIError( + "Documents not found or access denied: " + ", ".join(missing) + ) + return doc_targeting_block(client, doc_id) + + +def _system_text(content: Any) -> str: + """Text of a system/developer message: a string, or text parts joined.""" + if isinstance(content, str): + return content + if isinstance(content, list): + texts = [part.get("text") for part in content + if isinstance(part, dict) and isinstance(part.get("text"), str)] + if texts: + return "\n".join(texts) + raise PageIndexAPIError( + "system message content must be a string or a list of text parts." + ) + + +def _split_chat_messages(messages) -> "tuple[list[str], list[dict]]": + """Validate the chat_completions surface's messages: system/developer + content joins the managed instructions; user/assistant history passes + through. Tool-history round-trips belong to responses()/messages().""" + if not isinstance(messages, list) or not messages: + raise PageIndexAPIError("messages must be a non-empty list.") + system_texts: list[str] = [] + history: list[dict] = [] + for message in messages: + if not isinstance(message, dict) or "role" not in message: + raise PageIndexAPIError( + "Each message must be a dict with 'role' and 'content'.") + role = message["role"] + if role in ("system", "developer"): + system_texts.append(_system_text(message.get("content"))) + elif role in ("user", "assistant"): + content = message.get("content") + if not isinstance(content, str): + raise PageIndexAPIError( + "chat_completions content must be a string; for " + "structured items use responses() or messages()." + ) + history.append({"role": role, "content": content}) + else: + raise PageIndexAPIError( + f"Unsupported role for chat_completions: {role!r}. Tool " + "history round-trips belong to responses() or messages()." + ) + if not history: + raise PageIndexAPIError("messages must contain a user or assistant " + "message.") + return system_texts, history + + +def _run_sync(coro): + try: + asyncio.get_running_loop() + except RuntimeError: + return asyncio.run(coro) + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: + return pool.submit(asyncio.run, coro).result() + + +_SENTINEL = object() + + +def _stream_sync(agen_factory) -> Iterator[Any]: + """Drive an async generator from a background thread; yield synchronously.""" + items: "queue.Queue[Any]" = queue.Queue() + + def pump(): + async def consume(): + async for item in agen_factory(): + items.put(item) + + try: + asyncio.run(consume()) + except BaseException as exc: # re-raised on the consumer thread + items.put(exc) + return + items.put(_SENTINEL) + + threading.Thread(target=pump, daemon=True).start() + while True: + item = items.get() + if item is _SENTINEL: + return + if isinstance(item, BaseException): + raise item + yield item + + +# โ”€โ”€ OpenAI engine (chat_completions / responses) โ”€โ”€ + +def _require_openai_agents(method: str) -> None: + try: + import agents # noqa: F401 + except ImportError as exc: + raise PageIndexAPIError( + f"{method} in local mode requires the OpenAI Agents SDK โ€” " + "pip install openai-agents (or pip install 'pageindex[openai]')." + ) from exc + + +def _openai_model(protocol: str, model_name: str): + """The backend protocol driver โ€” the seam tests replace with a fake.""" + from openai import AsyncOpenAI + if protocol == "chat": + from agents.models.openai_chatcompletions import ( + OpenAIChatCompletionsModel) + return OpenAIChatCompletionsModel(model_name, AsyncOpenAI()) + from agents.models.openai_responses import OpenAIResponsesModel + return OpenAIResponsesModel(model_name, openai_client=AsyncOpenAI()) + + +def _openai_agent(client, protocol: str, model_name: str, instructions: str, + temperature, top_p): + from agents import Agent, ModelSettings + from .integrations.openai_agents import build_openai_tools + return Agent( + name="PageIndex", + instructions=instructions, + tools=build_openai_tools(client), + model=_openai_model(protocol, model_name), + model_settings=ModelSettings(temperature=temperature, top_p=top_p), + ) + + +def _run_kwargs(max_turns) -> dict: + # Managed runs never export traces โ€” the caller opted into document 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 _openai_usage(raw_responses) -> dict: + prompt = sum(r.usage.input_tokens for r in raw_responses) + completion = sum(r.usage.output_tokens for r in raw_responses) + return {"prompt_tokens": prompt, "completion_tokens": completion, + "total_tokens": prompt + completion} + + +def run_chat_completions(client, messages, stream: bool = False, + doc_id=None, temperature: Optional[float] = None, + stream_metadata: bool = False, + enable_citations: bool = False, + model: Optional[str] = None, + max_turns: Optional[int] = None, + ) -> Union[dict, Iterator[str], Iterator[dict]]: + _require_openai_agents("chat_completions") + if enable_citations: + raise PageIndexAPIError( + "enable_citations is cloud-only โ€” citations need block-level OCR " + "data that local mode does not store." + ) + system_texts, history = _split_chat_messages(messages) + block = _doc_block(client, doc_id) + items = ([{"role": "user", "content": block}] if block else []) + history + model_name = model or client.retrieve_model + agent = _openai_agent(client, "chat", model_name, + _managed_instructions(system_texts), + temperature, None) + from agents import Runner + if not stream: + result = _run_sync( + Runner.run(agent, input=items, **_run_kwargs(max_turns))) + return { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": int(time.time()), + "model": model_name, + "choices": [{ + "index": 0, + "message": {"role": "assistant", + "content": result.final_output or ""}, + "finish_reason": "stop", + }], + "usage": _openai_usage(result.raw_responses), + } + + chat_id = f"chatcmpl-{uuid.uuid4().hex}" + created = int(time.time()) + + def chunk(delta: dict, finish=None) -> dict: + return { + "id": chat_id, "object": "chat.completion.chunk", + "created": created, "model": model_name, + "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(max_turns)) + first = True + async for event in streamed.stream_events(): + if (event.type == "raw_response_event" + and isinstance(event.data, ResponseTextDeltaEvent)): + if first: + yield chunk({"role": "assistant", "content": ""}) + first = False + yield chunk({"content": event.data.delta}) + yield chunk({}, finish="stop") + yield { + "id": chat_id, "object": "chat.completion.chunk", + "created": created, "model": model_name, "choices": [], + "usage": _openai_usage(streamed.raw_responses), + } + + if stream_metadata: + return _stream_sync(agen) + return (piece["choices"][0]["delta"]["content"] + for piece in _stream_sync(agen) + if piece.get("choices") + and "content" in piece["choices"][0]["delta"] + and piece["choices"][0]["delta"]["content"]) + + +def run_responses(client, input, model: Optional[str] = None, + stream: bool = False, doc_id=None, + instructions: Optional[str] = None, + temperature: Optional[float] = None, + top_p: Optional[float] = None, + max_turns: Optional[int] = None, + ) -> Union[dict, Iterator[dict]]: + _require_openai_agents("responses") + if isinstance(input, str): + 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) + if block: + items = [{"role": "user", "content": block}] + items + extra = [instructions] if instructions else [] + model_name = model or client.retrieve_model + agent = _openai_agent(client, "responses", model_name, + _managed_instructions(extra), temperature, top_p) + from agents import Runner + + def envelope(output: list, raw_responses) -> dict: + usage = _openai_usage(raw_responses) + return { + "id": f"resp_{uuid.uuid4().hex}", + "object": "response", + "created_at": int(time.time()), + "model": model_name, + "status": "completed", + "output": output, + "usage": {"input_tokens": usage["prompt_tokens"], + "output_tokens": usage["completion_tokens"], + "total_tokens": usage["total_tokens"]}, + } + + if not stream: + result = _run_sync( + Runner.run(agent, input=[dict(item) for item in items], + **_run_kwargs(max_turns))) + output = result.to_input_list()[len(items):] + return envelope(output, result.raw_responses) + + async def agen(): + streamed = Runner.run_streamed(agent, + input=[dict(item) for item in items], + **_run_kwargs(max_turns)) + async for event in streamed.stream_events(): + if event.type == "raw_response_event": + yield event.data.model_dump(exclude_unset=True) + elif (event.type == "run_item_stream_event" + and event.item.type == "tool_call_output_item"): + # We are the tool executor, so we emit the output item the + # way the platform streams its own server-side tools. + yield {"type": "response.output_item.done", + "item": dict(event.item.to_input_item())} + output = streamed.to_input_list()[len(items):] + yield {"type": "response.completed", + "response": envelope(output, 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 + + +def _anthropic_client(): + """The backend client โ€” the seam tests replace with a fake transport.""" + import anthropic + return anthropic.Anthropic() + + +def _runnable_tools(client) -> list: + from anthropic import beta_tool + + def make(name: str): + def _fn(**kwargs: Any) -> str: + return call_tool(client, name, kwargs)[0] + + _fn.__name__ = name + return beta_tool(_fn, name=name, description=_local_description(name), + input_schema=_local_schema(name)) + + return [make(name) for name in tool_names()] + + +def _anthropic_system(extra_system, block: Optional[str]) -> list[dict]: + """System blocks with cache_control on the stable managed prefix; 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, + "cache_control": {"type": "ephemeral"}}) + if extra_system is None: + return blocks + if isinstance(extra_system, str): + return blocks + [{"type": "text", "text": extra_system}] + if isinstance(extra_system, list): + return blocks + list(extra_system) + raise PageIndexAPIError("system must be a string or a list of blocks.") + + +def _anthropic_usage(turns) -> dict: + fields = ("input_tokens", "output_tokens", + "cache_creation_input_tokens", "cache_read_input_tokens") + totals = {field: 0 for field in fields} + for turn in turns: + for field in fields: + value = getattr(turn.usage, field, None) + if isinstance(value, int): + totals[field] += value + return totals + + +def run_messages(client, messages, model: str, max_tokens: int, + stream: bool = False, doc_id=None, system=None, + temperature: Optional[float] = None, + top_p: Optional[float] = None, + top_k: Optional[int] = None, + stop_sequences: Optional[list[str]] = None, + max_turns: Optional[int] = None, + ) -> Union[dict, Iterator[Any]]: + _require_anthropic() + if not isinstance(messages, list) or not messages: + raise PageIndexAPIError("messages must be a non-empty list.") + block = _doc_block(client, doc_id) + prepared = [dict(message) for message in messages] + passthrough = {key: value for key, value in { + "temperature": temperature, "top_p": top_p, "top_k": top_k, + "stop_sequences": stop_sequences, + }.items() if value is not None} + runner = _anthropic_client().beta.messages.tool_runner( + max_tokens=max_tokens, + messages=prepared, + model=model, + tools=_runnable_tools(client), + system=_anthropic_system(system, block), + stream=stream, + **({"max_iterations": max_turns} if max_turns is not None else {}), + **passthrough, + ) + + if stream: + def events() -> Iterator[Any]: + for turn_stream in runner: + for event in turn_stream: + yield event + return events() + + turns = [turn for turn in runner] + 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) + conversation = list(captured.get("messages") or []) + final = turns[-1] + envelope = final.model_dump(mode="json") + envelope["usage"] = _anthropic_usage(turns) + # The full turn sequence (assistant tool_use + user tool_result + final), + # valid for verbatim append to the caller's history. The runner appends + # intermediate turns to its params but not the final assistant message. + new_messages = conversation[len(prepared):] + if not new_messages or new_messages[-1].get("role") != "assistant": + new_messages = new_messages + [{ + "role": "assistant", + "content": [block.model_dump(mode="json") + for block in final.content], + }] + envelope["messages"] = new_messages + return envelope diff --git a/pyproject.toml b/pyproject.toml index 65f68646b..df947424b 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" @@ -42,10 +42,14 @@ claude-agent-sdk = { version = ">=0.1.0", optional = true } # 0.8.0 offloads sync tools to a thread; older versions run them inline and # a blocking bridge call would freeze the agent event loop. openai-agents = { version = ">=0.8.0", optional = true } +# messages() drives the SDK's beta tool runner; 0.68.0 is the first release +# with tool_runner(stream/system/max_iterations) and beta_tool(input_schema). +anthropic = { version = ">=0.68.0", optional = true } [tool.poetry.extras] claude = ["claude-agent-sdk"] openai = ["openai-agents"] +anthropic = ["anthropic"] [tool.poetry.group.dev.dependencies] pytest = ">=7.0" diff --git a/tests/test_client.py b/tests/test_client.py index f55d519e6..e02107432 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -643,8 +643,12 @@ def test_retrieval_endpoints_cloud_only(local_client): local_client.get_retrieval("any") -def test_chat_completions_cloud_only(local_client): - with pytest.raises(PageIndexAPIError, match="not yet supported in local mode"): +def test_chat_completions_local_needs_agents_extra(local_client, monkeypatch): + """Local chat is implemented (see test_local_chat.py); without the + openai-agents extra it raises the actionable install error.""" + import sys + monkeypatch.setitem(sys.modules, "agents", None) + with pytest.raises(PageIndexAPIError, match="pageindex\\[openai\\]"): local_client.chat_completions( messages=[{"role": "user", "content": "q"}]) diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py new file mode 100644 index 000000000..8ecc6177a --- /dev/null +++ b/tests/test_local_chat.py @@ -0,0 +1,426 @@ +"""Local chat surfaces: three protocols over fake backends โ€” no network, +no LLM keys. Tool execution runs for real against a seeded local store.""" +import json +import sys +from pathlib import Path + +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 + +sys.path.insert(0, str(Path(__file__).parent.parent)) + + +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) โ”€โ”€ + +agents = pytest.importorskip("agents") + + +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) + + +from agents.models.interface import Model # noqa: E402 + + +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 = [] + + 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) + 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): + from openai.types.responses import (Response, ResponseCompletedEvent, + ResponseTextDeltaEvent) + from openai.types.responses.response_usage import ( + InputTokensDetails, OutputTokensDetails, ResponseUsage) + self._record(system_instructions, input) + output = self.turns.pop(0) + sequence = 0 + for item in output: + if item.type == "message": + for piece in ("The ", "answer"): + sequence += 1 + yield ResponseTextDeltaEvent( + type="response.output_text.delta", delta=piece, + content_index=0, item_id=item.id, output_index=0, + logprobs=[], sequence_number=sequence) + sequence += 1 + yield ResponseCompletedEvent( + type="response.completed", sequence_number=sequence, + response=Response( + id="resp_fake", created_at=0.0, model="fake", + object="response", output=output, parallel_tool_calls=False, + tool_choice="auto", tools=[], + usage=ResponseUsage( + input_tokens=10, output_tokens=5, total_tokens=15, + input_tokens_details=InputTokensDetails( + cached_tokens=0, cache_write_tokens=0), + output_tokens_details=OutputTokensDetails( + reasoning_tokens=0)))) + + +@pytest.fixture +def fake_model(monkeypatch): + state = {} + + def install(turns): + fake = FakeModel(turns) + state["protocols"] = [] + + def factory(protocol, model_name): + state["protocols"].append((protocol, model_name)) + return fake + + monkeypatch.setattr(local_chat, "_openai_model", factory) + return fake + + install.state = state + return install + + +# โ”€โ”€ chat_completions โ”€โ”€ + +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} + 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] + + +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"] + + +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"]) + + +def test_chat_completions_stream_modes(client, store_path, fake_model): + fake_model([[_msg_item("The answer")]]) + pieces = list(client.chat_completions( + [{"role": "user", "content": "q"}], stream=True)) + assert pieces == ["The ", "answer"] + + fake_model([[_msg_item("The answer")]]) + chunks = list(client.chat_completions( + [{"role": "user", "content": "q"}], stream=True, + stream_metadata=True)) + assert chunks[0]["choices"][0]["delta"] == {"role": "assistant", + "content": ""} + assert chunks[-2]["choices"][0]["finish_reason"] == "stop" + assert chunks[-1]["choices"] == [] + assert chunks[-1]["usage"]["total_tokens"] == 15 + assert all(c["object"] == "chat.completion.chunk" for c in chunks[:-1]) + + +def test_chat_completions_missing_framework(client, monkeypatch): + monkeypatch.setitem(sys.modules, "agents", None) + with pytest.raises(PageIndexAPIError, match="pageindex\\[openai\\]"): + client.chat_completions([{"role": "user", "content": "x"}]) + + +def test_cloud_guards(): + cloud = PageIndexCloudClient(api_key="pi-test-key") + with pytest.raises(PageIndexAPIError, match="local-mode parameters"): + cloud.chat_completions([{"role": "user", "content": "x"}], model="m") + with pytest.raises(PageIndexAPIError, match="not available on PageIndex " + "cloud yet"): + cloud.responses("x") + with pytest.raises(PageIndexAPIError, match="not available on PageIndex " + "cloud yet"): + cloud.messages([{"role": "user", "content": "x"}], model="m", + max_tokens=10) + + +# โ”€โ”€ responses โ”€โ”€ + +def test_responses_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.responses("What status?") + assert result["id"].startswith("resp_") + assert result["object"] == "response" + assert result["status"] == "completed" + assert result["usage"] == {"input_tokens": 20, "output_tokens": 10, + "total_tokens": 30} + assert fake_model.state["protocols"][0][0] == "responses" + types = [item.get("type", "message") for item in result["output"]] + assert "function_call" in types and "function_call_output" in types + # The final item is the assistant answer. + assert "The answer" in json.dumps(result["output"][-1]) + + +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["output"] + + [{"role": "user", "content": "and now?"}]) + client.responses(follow_up) + previous_final = first.inputs[-1] + assert second.inputs[0][:len(previous_final)] == previous_final + + +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 + tool_events = [event for event in events + if event.get("type") == "response.output_item.done" + and event.get("item", {}).get("type") + == "function_call_output"] + assert tool_events, types + assert types[-1] == "response.completed" + final = events[-1]["response"] + assert final["status"] == "completed" + assert final["usage"]["total_tokens"] == 30 + + +# โ”€โ”€ messages (Anthropic engine) โ”€โ”€ + +anthropic = pytest.importorskip("anthropic") +import httpx # noqa: E402 (anthropic depends on httpx) + + +def _anthropic_message(content, stop_reason): + return { + "id": "msg_fake", "type": "message", "role": "assistant", + "model": "claude-test", "content": content, + "stop_reason": stop_reason, "stop_sequence": None, + "usage": {"input_tokens": 10, "output_tokens": 5}, + } + + +@pytest.fixture +def fake_anthropic(monkeypatch): + state = {"calls": []} + + def install(responses): + state["calls"].clear() + + def handler(request): + state["calls"].append(json.loads(request.content)) + body = responses[len(state["calls"]) - 1] + if isinstance(body, str): # pre-rendered SSE + return httpx.Response( + 200, content=body.encode(), + headers={"content-type": "text/event-stream"}) + return httpx.Response(200, json=body) + + fake = anthropic.Anthropic( + api_key="test", + http_client=httpx.Client(transport=httpx.MockTransport(handler))) + monkeypatch.setattr(local_chat, "_anthropic_client", lambda: fake) + return state["calls"] + + return install + + +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"] + + +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." + + +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 + + +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") + + +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) From daac9d2dc09c17b2ea81bea97dd00c7f034bf975 Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 12 Aug 2026 02:31:12 +0800 Subject: [PATCH 019/137] =?UTF-8?q?fix:=20local-chat=20review=20findings?= =?UTF-8?q?=20=E2=80=94=20truncation,=20serialization,=20streams?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three independent review passes (bug scan, claims-vs-code, adversarial runtime probes) over the local-chat increment; every fix below was reproduced before being fixed. messages(): - A max_turns cut no longer duplicates the final assistant turn: the runner has already appended it when iterations exhaust, so the round-trip history carried a duplicate tool_use id and ended on an unanswered tool_use โ€” a guaranteed 400 on continuation. The append now keys on stop_reason, and truncation reads natively as stop_reason: "tool_use" with a continuable history. - The envelope is JSON-serializable end to end: runner-stored turns carry pydantic content blocks; everything is dumped to plain dicts, excluding SDK-internal __api_exclude__ fields (parsed_output) that the API rejects on round-trip. - Bounded by default (max_iterations 10, like the OpenAI surfaces); usage aggregation now preserves the final turn's native fields and sums the token counters None-safely; empty caller system strings are skipped; non-dict message entries and bad doc_id types raise PageIndexAPIError; anthropic < 0.68 gets an actionable version error; the doc block no longer spends a cache_control breakpoint. chat_completions()/responses(): - MaxTurnsExceeded wraps into PageIndexAPIError on all four run paths. - responses(stream=True) is one logical response: per-turn backend lifecycle events are collapsed (a canonical consumer previously stopped at turn 1's response.completed and never saw the answer), sequence numbers are reassigned monotonically, and the synthesized tool-output event carries output_index/sequence_number. - The responses envelope carries the real request surface (instructions, the actual function tool definitions, tool_choice, parallel_tool_calls, error/incomplete_details). - RunConfig(group_id) pins a stable prompt_cache_key: openai-agents otherwise stamps each run with a fresh key, tagging round-tripped prefixes as different cache groups and defeating the feature the round-trip exists for. - Abandoning a stream now cancels the run: a watchdog task lets the cancellation land even while the pump awaits the backend, and the per-call AsyncOpenAI client is closed before its loop ends (fixes "Task exception was never retrieved" noise). The opening role chunk is emitted even for empty outputs; empty responses() input and enable_citations-before-extra ordering fixed. Docs rescoped to what is true: finish_reason/status reflect loop completion on the OpenAI surfaces (the engine does not surface per-turn backend reasons); chat streaming yields visible narration including pre-tool text; messages(stream=True) forwards the Anthropic SDK's native event objects (not wire-verbatim); the doc block is a leading conversation item on OpenAI surfaces and a system block on messages(). Tests: 25 in the file (11 new), with per-extra skip sections so a machine with only one framework still covers the other surface; without-frameworks matrix re-verified; live smoke re-run green with a clean exit. --- pageindex/agent_tools.py | 6 +- pageindex/client.py | 27 +++- pageindex/local_chat.py | 328 ++++++++++++++++++++++++++++++--------- tests/test_local_chat.py | 244 ++++++++++++++++++++++++++++- 4 files changed, 516 insertions(+), 89 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index bf8627d91..83c2d23cc 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -1453,9 +1453,9 @@ def _base_instructions(client) -> str: def doc_targeting_block(client, doc_id) -> 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 (which place it as a leading conversation item). Raises when a - doc_id's name is shadowed by a newer same-name document โ€” the - name-addressed tools could not reach it.""" + 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.""" if doc_id is None: return None doc_ids = [doc_id] if isinstance(doc_id, str) else list(doc_id) diff --git a/pageindex/client.py b/pageindex/client.py index 91e4f18a8..c3cab1278 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -354,9 +354,13 @@ def chat_completions( run over the local tools against your own LLM backend's /chat/completions (requires ``pageindex[openai]``; the OpenAI SDK's usual env config โ€” OPENAI_API_KEY, OPENAI_BASE_URL โ€” selects the - backend, so any OpenAI-compatible server works). The response - carries the final answer only; for the tool-use process and - prompt-cache round-trip use ``responses()`` or ``messages()``. + backend, so any OpenAI-compatible server works). 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. @@ -428,9 +432,11 @@ def responses( input: A user message string, or a list of Responses input items (round-trip prior ``output`` items here). model: Backend model name (defaults to ``retrieve_model``). - stream: Yield native Responses stream events as dicts; tool - outputs are emitted as ``response.output_item.done`` events - and the final event is ``response.completed``. + stream: Yield Responses stream events as dicts โ€” one logical + response per call: per-turn backend lifecycle events are + collapsed and sequence numbers reassigned monotonically; + tool outputs are emitted as ``response.output_item.done`` + events and the single final event is ``response.completed``. doc_id: Document ID or list of IDs to scope the conversation. instructions: Appended to the managed system prompt. temperature / top_p: Passed through to the model. @@ -479,11 +485,16 @@ def messages( messages: Native Messages-format history (including prior tool_use/tool_result blocks on round-trip). model / max_tokens: Required by the Messages API; passed through. - stream: Yield the native event stream across turns, verbatim. + 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. system: Appended after the managed system blocks. temperature / top_p / top_k / stop_sequences: Passed through. - max_turns: Cap on agent turns per call. + max_turns: Cap on agent turns per call (default 10, like the + OpenAI surfaces). A truncated run reports + ``stop_reason: "tool_use"`` and its ``messages`` remain + valid for continuation. """ from .cloud_api import CloudAPI if isinstance(self._api, CloudAPI): diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 3497f3970..d78207bd6 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -1,15 +1,18 @@ """Managed local chat: document-QA agents over the local tools. -Three methods, three wire protocols, 1:1 with the backend and no translation -layer: ``chat_completions`` drives the backend's /chat/completions (any -OpenAI-compatible backend, final answer only), ``responses`` drives -/responses (process items are standard output; round-trip them for provider -prompt-cache continuation and agent memory), ``messages`` drives Anthropic's -/v1/messages via the SDK's own tool runner (tool_use/tool_result round-trip -is the format's native behavior). +Three methods, three backend protocols, routed 1:1: ``chat_completions`` +drives the backend's /chat/completions (any OpenAI-compatible backend, +final answer only), ``responses`` drives /responses (process items are +standard output; round-trip them for provider prompt-cache continuation and +agent memory), ``messages`` drives Anthropic's /v1/messages via the SDK's +own tool runner (tool_use/tool_result round-trip is the format's native +behavior). Content passes through untouched โ€” the caller's messages, the model's -answers, tool outputs, finish/stop reasons. The SDK owns only gatekeeping +answers, tool outputs. Native stop reasons pass through on ``messages``; +the OpenAI engine's abstraction does not surface per-turn finish reasons, +so ``chat_completions`` reports loop completion as ``"stop"`` and +``responses`` as ``status: "completed"``. The SDK owns gatekeeping (structural validation), table-setting (managed instructions, tools, doc targeting), tool execution, and billing (usage aggregation, envelope ids). """ @@ -43,6 +46,9 @@ def _managed_instructions(extra_system: list[str]) -> str: def _doc_block(client, doc_id) -> Optional[str]: if doc_id is None: return None + if not isinstance(doc_id, (str, list)): + raise PageIndexAPIError("doc_id must be a string or a list of " + "strings.") doc_ids = [doc_id] if isinstance(doc_id, str) else list(doc_id) missing = [] for one_id in doc_ids: @@ -118,29 +124,69 @@ def _run_sync(coro): def _stream_sync(agen_factory) -> Iterator[Any]: - """Drive an async generator from a background thread; yield synchronously.""" - items: "queue.Queue[Any]" = queue.Queue() + """Drive an async generator from a background thread; yield synchronously. + + Closing (or abandoning) 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(): - async for item in agen_factory(): - items.put(item) + 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 - items.put(exc) + deliver(exc) return - items.put(_SENTINEL) + deliver(_SENTINEL) threading.Thread(target=pump, daemon=True).start() - while True: - item = items.get() - if item is _SENTINEL: - return - if isinstance(item, BaseException): - raise item - yield item + 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) โ”€โ”€ @@ -179,16 +225,54 @@ def _openai_agent(client, protocol: str, model_name: str, instructions: str, ) +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 _run_kwargs(max_turns) -> dict: # Managed runs never export traces โ€” the caller opted into document QA, - # not telemetry. + # not telemetry. The stable group_id keys OpenAI's prompt-cache routing: + # without it openai-agents stamps every run with a fresh + # prompt_cache_key, tagging a round-tripped prefix as a different cache + # group. from agents import RunConfig - kwargs: dict = {"run_config": RunConfig(tracing_disabled=True)} + kwargs: dict = {"run_config": RunConfig(tracing_disabled=True, + group_id="pageindex-local-chat")} if max_turns is not None: kwargs["max_turns"] = max_turns return kwargs +async def _aclose_backend(agent) -> None: + """Close the per-call AsyncOpenAI client before its event loop ends โ€” + otherwise httpx tears down pooled connections on a closed loop and + emits 'Task exception was never retrieved' noise.""" + backend = getattr(getattr(agent, "model", None), "_client", None) + close = getattr(backend, "close", None) + if close is not None: + try: + await close() + except Exception: + pass + + +async def _run_closing(agent, coro): + try: + return await coro + finally: + await _aclose_backend(agent) + + +def _wrap_max_turns(exc, 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 _openai_usage(raw_responses) -> dict: prompt = sum(r.usage.input_tokens for r in raw_responses) completion = sum(r.usage.output_tokens for r in raw_responses) @@ -203,12 +287,13 @@ def run_chat_completions(client, messages, stream: bool = False, model: Optional[str] = None, max_turns: Optional[int] = None, ) -> Union[dict, Iterator[str], Iterator[dict]]: - _require_openai_agents("chat_completions") 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 @@ -217,9 +302,13 @@ def run_chat_completions(client, messages, stream: bool = False, _managed_instructions(system_texts), temperature, None) from agents import Runner + from agents.exceptions import MaxTurnsExceeded if not stream: - result = _run_sync( - Runner.run(agent, input=items, **_run_kwargs(max_turns))) + try: + result = _run_sync(_run_closing(agent, + Runner.run(agent, input=items, **_run_kwargs(max_turns)))) + except MaxTurnsExceeded as exc: + raise _wrap_max_turns(exc, max_turns) from exc return { "id": f"chatcmpl-{uuid.uuid4().hex}", "object": "chat.completion", @@ -249,14 +338,20 @@ async def agen(): from openai.types.responses import ResponseTextDeltaEvent streamed = Runner.run_streamed(agent, input=items, **_run_kwargs(max_turns)) - first = True - async for event in streamed.stream_events(): - if (event.type == "raw_response_event" - and isinstance(event.data, ResponseTextDeltaEvent)): - if first: - yield chunk({"role": "assistant", "content": ""}) - first = False - yield chunk({"content": event.data.delta}) + yield chunk({"role": "assistant", "content": ""}) + completed = False + try: + 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(exc, max_turns) from exc + finally: + if not completed and hasattr(streamed, "cancel"): + streamed.cancel() # abandoned/failed: stop the agent task + await _aclose_backend(agent) yield chunk({}, finish="stop") yield { "id": chat_id, "object": "chat.completion.chunk", @@ -281,7 +376,8 @@ def run_responses(client, input, model: Optional[str] = None, max_turns: Optional[int] = None, ) -> Union[dict, Iterator[dict]]: _require_openai_agents("responses") - if isinstance(input, str): + _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)): @@ -294,9 +390,11 @@ def run_responses(client, input, model: Optional[str] = None, items = [{"role": "user", "content": block}] + items extra = [instructions] if instructions else [] model_name = model or client.retrieve_model - agent = _openai_agent(client, "responses", model_name, - _managed_instructions(extra), temperature, top_p) + managed = _managed_instructions(extra) + agent = _openai_agent(client, "responses", model_name, managed, + temperature, top_p) from agents import Runner + from agents.exceptions import MaxTurnsExceeded def envelope(output: list, raw_responses) -> dict: usage = _openai_usage(raw_responses) @@ -310,30 +408,74 @@ def envelope(output: list, raw_responses) -> dict: "usage": {"input_tokens": usage["prompt_tokens"], "output_tokens": usage["completion_tokens"], "total_tokens": usage["total_tokens"]}, + "instructions": managed, + "tools": [{"type": "function", "name": tool.name, + "description": tool.description, + "parameters": tool.params_json_schema, + "strict": getattr(tool, "strict_json_schema", True)} + for tool in agent.tools], + "tool_choice": "auto", + "parallel_tool_calls": True, + "temperature": temperature, + "top_p": top_p, + "max_output_tokens": None, + "error": None, + "incomplete_details": None, + "metadata": None, } if not stream: - result = _run_sync( - Runner.run(agent, input=[dict(item) for item in items], - **_run_kwargs(max_turns))) + try: + result = _run_sync(_run_closing(agent, + Runner.run(agent, input=[dict(item) for item in items], + **_run_kwargs(max_turns)))) + except MaxTurnsExceeded as exc: + raise _wrap_max_turns(exc, max_turns) from exc output = result.to_input_list()[len(items):] return envelope(output, result.raw_responses) + # One logical response per call: per-turn backend lifecycle events + # (created/completed/...) are collapsed โ€” forwarding them verbatim would + # end a canonical consumer at the first turn โ€” and sequence numbers are + # reassigned monotonically across the whole run. + 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(max_turns)) - async for event in streamed.stream_events(): - if event.type == "raw_response_event": - yield event.data.model_dump(exclude_unset=True) - elif (event.type == "run_item_stream_event" - and event.item.type == "tool_call_output_item"): - # We are the tool executor, so we emit the output item the - # way the platform streams its own server-side tools. - yield {"type": "response.output_item.done", - "item": dict(event.item.to_input_item())} + sequence = 0 + completed = False + try: + async for event in streamed.stream_events(): + if event.type == "raw_response_event": + data = event.data.model_dump(exclude_unset=True) + if data.get("type") in lifecycle: + continue + sequence += 1 + data["sequence_number"] = sequence + yield data + elif (event.type == "run_item_stream_event" + and event.item.type == "tool_call_output_item"): + # We are the tool executor, so we emit the output item + # the way the platform streams its own server-side tools. + sequence += 1 + yield {"type": "response.output_item.done", + "output_index": sequence, + "sequence_number": sequence, + "item": dict(event.item.to_input_item())} + completed = True + except MaxTurnsExceeded as exc: + raise _wrap_max_turns(exc, max_turns) from exc + finally: + if not completed and hasattr(streamed, "cancel"): + streamed.cancel() # abandoned/failed: stop the agent task + await _aclose_backend(agent) output = streamed.to_input_list()[len(items):] - yield {"type": "response.completed", + sequence += 1 + yield {"type": "response.completed", "sequence_number": sequence, "response": envelope(output, streamed.raw_responses)} return _stream_sync(agen) @@ -349,6 +491,13 @@ def _require_anthropic() -> None: "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 + except ImportError as exc: + raise PageIndexAPIError( + "messages in local mode requires anthropic >= 0.68.0 (the tool " + "runner) โ€” pip install -U anthropic." + ) from exc def _anthropic_client(): @@ -372,32 +521,54 @@ def _fn(**kwargs: Any) -> str: def _anthropic_system(extra_system, block: Optional[str]) -> list[dict]: - """System blocks with cache_control on the stable managed prefix; the - doc block and caller system content follow as their own blocks.""" + """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, - "cache_control": {"type": "ephemeral"}}) + blocks.append({"type": "text", "text": block}) if extra_system is None: return blocks if isinstance(extra_system, str): - return blocks + [{"type": "text", "text": extra_system}] + 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 _anthropic_usage(turns) -> dict: - fields = ("input_tokens", "output_tokens", - "cache_creation_input_tokens", "cache_read_input_tokens") - totals = {field: 0 for field in fields} - for turn in turns: - for field in fields: - value = getattr(turn.usage, field, None) - if isinstance(value, int): - totals[field] += value +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 @@ -410,8 +581,11 @@ def run_messages(client, messages, model: str, max_tokens: int, max_turns: Optional[int] = None, ) -> Union[dict, Iterator[Any]]: _require_anthropic() - if not isinstance(messages, list) or not messages: - raise PageIndexAPIError("messages must be a non-empty list.") + _validate_max_turns(max_turns) + 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 list of " + "message dicts.") block = _doc_block(client, doc_id) prepared = [dict(message) for message in messages] passthrough = {key: value for key, value in { @@ -425,7 +599,8 @@ def run_messages(client, messages, model: str, max_tokens: int, tools=_runnable_tools(client), system=_anthropic_system(system, block), stream=stream, - **({"max_iterations": max_turns} if max_turns is not None else {}), + # Bounded like the OpenAI surfaces (their framework default is 10). + max_iterations=max_turns if max_turns is not None else 10, **passthrough, ) @@ -449,16 +624,23 @@ def capture(params): conversation = list(captured.get("messages") or []) final = turns[-1] envelope = final.model_dump(mode="json") - envelope["usage"] = _anthropic_usage(turns) + 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 - # intermediate turns to its params but not the final assistant message. - new_messages = conversation[len(prepared):] - if not new_messages or new_messages[-1].get("role") != "assistant": + # a turn to its params only when it executed tools, so the final + # assistant message is missing exactly when the run ended naturally + # (stop_reason != "tool_use"); on a max_turns cut the last appended + # turn IS the final message and appending again would duplicate its + # tool_use ids. + new_messages = [_dump_message(message) + for message in conversation[len(prepared):]] + if (final.stop_reason != "tool_use" + and (not new_messages + or new_messages[-1].get("role") != "assistant")): new_messages = new_messages + [{ "role": "assistant", - "content": [block.model_dump(mode="json") - for block in final.content], + "content": [_dump_block(item) for item in final.content], }] envelope["messages"] = new_messages return envelope diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index 8ecc6177a..50cbc22eb 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -39,8 +39,18 @@ def client(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. -agents = pytest.importorskip("agents") +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") +pytestmark_openai = needs_agents def _msg_item(text): @@ -65,7 +75,10 @@ def _usage(): total_tokens=15) -from agents.models.interface import Model # noqa: E402 +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): @@ -75,6 +88,7 @@ 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) @@ -94,17 +108,24 @@ async def get_response(self, system_instructions, input, model_settings, async def stream_response(self, system_instructions, input, model_settings, tools, output_schema, handoffs, tracing, **kwargs): + import asyncio as aio from openai.types.responses import (Response, ResponseCompletedEvent, ResponseTextDeltaEvent) from openai.types.responses.response_usage import ( InputTokensDetails, OutputTokensDetails, ResponseUsage) + block_from = getattr(self, "block_from", None) + if block_from is not None and len(self.inputs) + 1 >= block_from: + while True: # released only by task cancellation + await aio.sleep(0.01) self._record(system_instructions, input) output = self.turns.pop(0) sequence = 0 for item in output: if item.type == "message": - for piece in ("The ", "answer"): + 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, @@ -145,6 +166,7 @@ def factory(protocol, model_name): # โ”€โ”€ 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([ @@ -169,6 +191,7 @@ def test_chat_completions_end_to_end(client, store_path, fake_model): 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")]]) @@ -181,6 +204,7 @@ def test_chat_completions_system_and_doc_block(client, store_path, fake_model): assert "The user has specified document: report.pdf" in first_item["content"] +@needs_agents def test_chat_completions_validation(client, store_path, fake_model): fake_model([[_msg_item("ok")]]) with pytest.raises(PageIndexAPIError, match="cloud-only"): @@ -198,6 +222,7 @@ def test_chat_completions_validation(client, store_path, fake_model): 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( @@ -237,6 +262,7 @@ def test_cloud_guards(): # โ”€โ”€ responses โ”€โ”€ +@needs_agents def test_responses_end_to_end(client, store_path, fake_model): seed_doc(store_path, "pi-a", "report.pdf") fake = fake_model([ @@ -256,6 +282,7 @@ def test_responses_end_to_end(client, store_path, fake_model): 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.""" @@ -275,6 +302,7 @@ def test_responses_round_trip_extends_prefix(client, store_path, fake_model): assert second.inputs[0][:len(previous_final)] == previous_final +@needs_agents def test_responses_stream_passthrough(client, store_path, fake_model): seed_doc(store_path, "pi-a", "report.pdf") fake_model([ @@ -297,8 +325,15 @@ def test_responses_stream_passthrough(client, store_path, fake_model): # โ”€โ”€ messages (Anthropic engine) โ”€โ”€ -anthropic = pytest.importorskip("anthropic") -import httpx # noqa: E402 (anthropic depends on httpx) +try: + import anthropic + import httpx + _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): @@ -335,6 +370,7 @@ def handler(request): 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([ @@ -367,6 +403,7 @@ def test_messages_end_to_end(client, store_path, fake_anthropic): == 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([ @@ -379,6 +416,7 @@ def test_messages_doc_block_and_system(client, store_path, fake_anthropic): 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', @@ -409,6 +447,7 @@ def test_messages_stream_passthrough(client, store_path, fake_anthropic): assert "content_block_delta" in types and "message_stop" in types +@needs_anthropic def test_messages_validation(client, fake_anthropic): fake_anthropic([]) with pytest.raises(PageIndexAPIError, match="non-empty"): @@ -424,3 +463,198 @@ def test_messages_missing_framework(client, monkeypatch): 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 +def test_chat_completions_max_turns_wrapped(client, store_path, fake_model): + """MaxTurnsExceeded is an engine-internal type; callers get the SDK's + own error โ€” on 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="max_turns"): + client.chat_completions([{"role": "user", "content": "q"}], + max_turns=1) + 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="max_turns"): + list(client.chat_completions([{"role": "user", "content": "q"}], + stream=True, max_turns=1)) + with pytest.raises(PageIndexAPIError, match="positive integer"): + client.chat_completions([{"role": "user", "content": "q"}], + max_turns=0) + + +def test_enable_citations_rejected_before_framework_check(client, monkeypatch): + monkeypatch.setitem(sys.modules, "agents", None) + with pytest.raises(PageIndexAPIError, match="cloud-only"): + client.chat_completions([{"role": "user", "content": "x"}], + enable_citations=True) + + +@needs_agents +def test_chat_stream_role_chunk_even_with_empty_output(client, fake_model): + fake_model([[]]) + chunks = list(client.chat_completions([{"role": "user", "content": "q"}], + stream=True, stream_metadata=True)) + assert chunks[0]["choices"][0]["delta"] == {"role": "assistant", + "content": ""} + assert chunks[-2]["choices"][0]["finish_reason"] == "stop" + + +@needs_agents +def test_responses_stream_single_completed_monotonic_sequence( + client, store_path, fake_model): + """One logical response per call: per-turn backend lifecycle events are + collapsed and sequence numbers never go backwards.""" + seed_doc(store_path, "pi-a", "report.pdf") + fake_model([ + [_call_item("get_document", {"doc_name": "report.pdf"})], + [_msg_item("The answer")], + ]) + events = list(client.responses("q", stream=True)) + completed = [event for event in events + if event.get("type") == "response.completed"] + assert len(completed) == 1 and events[-1] is completed[0] + sequences = [event["sequence_number"] for event in events + if "sequence_number" in event] + assert sequences == sorted(sequences) + assert len(set(sequences)) == len(sequences) + tool_done = next(event for event in events + if event.get("type") == "response.output_item.done" + and event["item"]["type"] == "function_call_output") + assert "sequence_number" in tool_done and "output_index" in tool_done + + +@needs_agents +def test_responses_envelope_fields_and_cache_group(client, store_path, + fake_model): + seed_doc(store_path, "pi-a", "report.pdf") + fake_model([[_msg_item("ok")]]) + result = client.responses("q") + names = {tool["name"] for tool in result["tools"]} + assert names == {"browse_documents", "get_document", + "get_document_structure", "get_page_content"} + assert all(tool["type"] == "function" for tool in result["tools"]) + assert result["instructions"].startswith(CHAT_HEADER) + assert result["parallel_tool_calls"] is True + assert result["tool_choice"] == "auto" + # Stable cache group: without it openai-agents stamps each run with a + # fresh prompt_cache_key, defeating round-trip cache routing. + assert (local_chat._run_kwargs(None)["run_config"].group_id + == "pageindex-local-chat") + + +@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_stream_abandonment_cancels_pending_turn(client, store_path, + fake_model): + """Closing the iterator cancels the run even while it is awaiting the + backend: the blocked turn is torn down (pump thread exits) instead of + running โ€” and billing โ€” to completion in the background.""" + import threading + import time as time_mod + seed_doc(store_path, "pi-a", "report.pdf") + baseline = threading.active_count() + fake = fake_model([ + [_call_item("get_document", {"doc_name": "report.pdf"})], + [_msg_item("The answer")], + ]) + fake.block_from = 2 # turn 2 hangs until cancelled + stream = client.chat_completions([{"role": "user", "content": "q"}], + stream=True, stream_metadata=True) + next(stream) # the opening role chunk + stream.close() + deadline = time_mod.monotonic() + 3.0 + while (threading.active_count() > baseline + and time_mod.monotonic() < deadline): + time_mod.sleep(0.05) + assert threading.active_count() <= baseline + assert fake.deltas_emitted == 0 # turn 2 never produced output + + +@needs_anthropic +def test_messages_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_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) From 4590dd855c4e16190f7f1f305071df9ed59892cc Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 12 Aug 2026 03:03:08 +0800 Subject: [PATCH 020/137] =?UTF-8?q?feat:=20as=5Fanthropic=5Ftools=20?= =?UTF-8?q?=E2=80=94=20Anthropic=20tool-runner=20export,=20both=20modes?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fills the last cell of the agent-connection matrix: users driving their own anthropic tool_runner loop get runnable tools directly. Cloud wraps the live MCP tool set with input schemas passing through verbatim (MCP inputSchema is the Messages API schema shape); local exposes the same set messages() runs internally. The beta_tool wrapping moves from local_chat into integrations/anthropic_sdk.py, parallel to openai_agents.py, and messages() now consumes the shared builder. agent_tools grows _bridge_invoker/_read_only_tools so the plain-function and beta_tool cloud paths share invocation containment and the read-only gate. --- pageindex/agent_tools.py | 57 +++++++++++++--------- pageindex/client.py | 27 +++++++++++ pageindex/integrations/anthropic_sdk.py | 60 +++++++++++++++++++++++ pageindex/local_chat.py | 22 ++------- pyproject.toml | 5 +- tests/test_agent_tools.py | 64 +++++++++++++++++++++++++ 6 files changed, 192 insertions(+), 43 deletions(-) create mode 100644 pageindex/integrations/anthropic_sdk.py diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index 83c2d23cc..e0069faeb 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -1220,18 +1220,11 @@ def _annotation_for(spec: dict) -> Any: return _SCHEMA_TYPE_MAP.get(schema_type, Any) -def _make_bridge_function(bridge, meta: dict) -> Callable[..., str]: - """One plain function for a cloud tool: real signature and docstring from - the server's schema, invocation proxied over MCP, errors contained.""" - import keyword - - name = str(meta.get("name") or "") - schema = meta.get("inputSchema") or {} - properties: dict[str, Any] = schema.get("properties") or {} - required = set(schema.get("required") or []) - +def _bridge_invoker(bridge, name: str) -> Callable[[dict], str]: + """One cloud tool call proxied over MCP: None-valued arguments are + dropped (None โ‰ก omitted, matching the contract's "omit if ..." + semantics) and failures are contained in the error envelope.""" def _invoke(arguments: dict[str, Any]) -> str: - # None โ‰ก omitted, matching the contract's "omit if ..." semantics. arguments = {key: value for key, value in arguments.items() if value is not None} try: @@ -1246,6 +1239,19 @@ def _invoke(arguments: dict[str, Any]) -> str: "INTERNAL_ERROR", ) return _dumps(payload) + return _invoke + + +def _make_bridge_function(bridge, meta: dict) -> Callable[..., str]: + """One plain function for a cloud tool: real signature and docstring from + the server's schema, invocation proxied over MCP, errors contained.""" + import keyword + + name = str(meta.get("name") or "") + schema = meta.get("inputSchema") or {} + properties: dict[str, Any] = schema.get("properties") or {} + required = set(schema.get("required") or []) + _invoke = _bridge_invoker(bridge, name) params_usable = all(param.isidentifier() and not keyword.iskeyword(param) and param != "_invoke" @@ -1300,22 +1306,27 @@ def _cloud_bridge(client): return bridge +def _read_only_tools(tools_meta: list[dict]) -> list[dict]: + """The management gate for consumers without a framework permission + layer: only tools the server marks read-only, guarded against a server + annotation regression silently disabling every tool.""" + filtered = [meta for meta in tools_meta + if (meta.get("annotations") or {}).get("readOnlyHint") is True] + if tools_meta and not filtered: + raise PageIndexAPIError( + "The MCP server returned tools but none are annotated " + "read-only โ€” a server annotation regression would otherwise " + "silently disable every tool. Pass include_management=True " + "to expose the unfiltered list." + ) + return filtered + + def _build_cloud_agent_tools(client, include_management: bool) -> list[Callable[..., str]]: bridge = _cloud_bridge(client) tools_meta = bridge.list_tools() if not include_management: - # Plain functions have no framework permission layer, so the - # management gate lives here: only tools the server marks read-only. - 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." - ) - tools_meta = filtered + tools_meta = _read_only_tools(tools_meta) return [_make_bridge_function(bridge, meta) for meta in tools_meta] diff --git a/pageindex/client.py b/pageindex/client.py index c3cab1278..b10f8c542 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -618,6 +618,33 @@ def as_openai_tools(self, include_management: bool = False, from .integrations.openai_agents import build_openai_tools return build_openai_tools(self, include_management, hosted) + def as_anthropic_tools(self, include_management: bool = False) -> list: + """ + Runnable tools for the Anthropic SDK's tool runner โ€” pass to + ``client.beta.messages.tool_runner(tools=...)``. + + 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). The Messages API's MCP connector (``mcp_servers=`` + pointing at ``{BASE_URL}/mcp``) is the server-side alternative + with no client-side tools involved. Local: the in-process tools โ€” + the same set ``messages()`` runs internally. + + Requires ``anthropic>=0.68.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, ...). + """ + from .integrations.anthropic_sdk import build_anthropic_tools + return build_anthropic_tools(self, include_management) + def as_claude_mcp(self, include_management: bool = False): """ ``mcp_servers`` entry for the Claude Agent SDK. diff --git a/pageindex/integrations/anthropic_sdk.py b/pageindex/integrations/anthropic_sdk.py new file mode 100644 index 000000000..36869f4e0 --- /dev/null +++ b/pageindex/integrations/anthropic_sdk.py @@ -0,0 +1,60 @@ +"""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. +""" +from __future__ import annotations + +from typing import Any + +from ..errors import PageIndexAPIError + + +def build_anthropic_tools(client, include_management: bool = False) -> list: + try: + from anthropic import beta_tool + except ImportError as exc: + raise PageIndexAPIError( + "as_anthropic_tools requires the Anthropic SDK tool runner " + "(anthropic>=0.68.0) โ€” pip install -U anthropic (or pip install " + "'pageindex[anthropic]')." + ) from exc + + if getattr(client, "api_key", None): + from ..agent_tools import (_bridge_invoker, _cloud_bridge, + _read_only_tools) + bridge = _cloud_bridge(client) + tools_meta = bridge.list_tools() + if not include_management: + tools_meta = _read_only_tools(tools_meta) + + def make_cloud(meta: dict): + name = str(meta.get("name") or "tool") + invoke = _bridge_invoker(bridge, name) + + def _fn(**kwargs: Any) -> str: + return invoke(kwargs) + + _fn.__name__ = name + return beta_tool( + _fn, name=name, description=meta.get("description", ""), + input_schema=meta.get("inputSchema") + or {"type": "object", "properties": {}}, + ) + + return [make_cloud(meta) for meta in tools_meta] + + from ..agent_tools import (_local_description, _local_schema, call_tool, + tool_names) + + def make_local(name: str): + def _fn(**kwargs: Any) -> str: + return call_tool(client, name, kwargs)[0] + + _fn.__name__ = name + return beta_tool(_fn, name=name, description=_local_description(name), + input_schema=_local_schema(name)) + + return [make_local(name) for name in tool_names(include_management)] diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index d78207bd6..192ef27e4 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -26,9 +26,7 @@ import uuid from typing import Any, Iterator, Optional, Union -from .agent_tools import (AGENT_INSTRUCTIONS, _local_description, - _local_schema, call_tool, doc_targeting_block, - tool_names) +from .agent_tools import AGENT_INSTRUCTIONS, doc_targeting_block from .errors import PageIndexAPIError CHAT_HEADER = ( @@ -506,20 +504,6 @@ def _anthropic_client(): return anthropic.Anthropic() -def _runnable_tools(client) -> list: - from anthropic import beta_tool - - def make(name: str): - def _fn(**kwargs: Any) -> str: - return call_tool(client, name, kwargs)[0] - - _fn.__name__ = name - return beta_tool(_fn, name=name, description=_local_description(name), - input_schema=_local_schema(name)) - - return [make(name) for name in tool_names()] - - 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 @@ -580,6 +564,8 @@ def run_messages(client, messages, model: str, max_tokens: int, stop_sequences: Optional[list[str]] = None, max_turns: Optional[int] = None, ) -> Union[dict, Iterator[Any]]: + from .integrations.anthropic_sdk import build_anthropic_tools + _require_anthropic() _validate_max_turns(max_turns) if (not isinstance(messages, list) or not messages @@ -596,7 +582,7 @@ def run_messages(client, messages, model: str, max_tokens: int, max_tokens=max_tokens, messages=prepared, model=model, - tools=_runnable_tools(client), + tools=build_anthropic_tools(client), system=_anthropic_system(system, block), stream=stream, # Bounded like the OpenAI surfaces (their framework default is 10). diff --git a/pyproject.toml b/pyproject.toml index df947424b..e55de0287 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,8 +42,9 @@ claude-agent-sdk = { version = ">=0.1.0", optional = true } # 0.8.0 offloads sync tools to a thread; older versions run them inline and # a blocking bridge call would freeze the agent event loop. openai-agents = { version = ">=0.8.0", optional = true } -# messages() drives the SDK's beta tool runner; 0.68.0 is the first release -# with tool_runner(stream/system/max_iterations) and beta_tool(input_schema). +# messages() and as_anthropic_tools() need the SDK's beta tool runner; +# 0.68.0 is the first release with tool_runner(stream/system/max_iterations) +# and beta_tool(input_schema). anthropic = { version = ">=0.68.0", optional = true } [tool.poetry.extras] diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 2a3cbc70b..afd40ccbc 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -484,9 +484,73 @@ def test_as_claude_mcp_local_when_installed(client): assert server.get("type") != "http" +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 pageindex.agent_tools import _local_description, _local_schema + tools = client.as_anthropic_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_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_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"] + 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_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): + pytest.importorskip("anthropic") + cloud, created = cloud_with_fake_bridge + tools = cloud.as_anthropic_tools() + + def boom(name, arguments): + raise RuntimeError("bridge down") + + created["bridge"].call_tool = boom + payload = json.loads(tools[0].call({"query": "q"})) + 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() From 02022df40fd6aaa4a4d2a7bd1286d8f6109dbf11 Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 12 Aug 2026 03:27:52 +0800 Subject: [PATCH 021/137] =?UTF-8?q?fix:=20as=5Fanthropic=5Ftools=20review?= =?UTF-8?q?=20findings=20=E2=80=94=20async=20flavor,=20schema=20isolation?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adversarial + best-practice review of 4590dd8 (three independent passes) surfaced two holes. The export was sync-only: AsyncAnthropic's runner accepts only BetaAsyncFunctionTool and splices anything else into the request body unserialized, so the first call died with an opaque TypeError โ€” asynchronous=True now builds beta_async_tool runnables (present since the 0.68.0 floor) that run the blocking bridge/store call in a worker thread, keeping I/O off the caller's event loop. And beta_tool stores input_schema by reference, so cloud tools aliased the bridge's cached metas while the local path deep-copied โ€” the builder now copies, and the passthrough test asserts equal-but-not-aliased so it can no longer compare an object with itself. Docstring fixes from the same round: the MCP-connector pointer now carries the full live-verified shape (authorization_token was missing โ€” following it literally gave a 401), and the manual messages.create loop's to_dict() serialization is documented. Tests pin the runnable flavor both ways (isinstance), which existing tests could not distinguish. --- pageindex/client.py | 25 +++++++--- pageindex/integrations/anthropic_sdk.py | 62 ++++++++++++++++--------- tests/test_agent_tools.py | 30 ++++++++++++ 3 files changed, 88 insertions(+), 29 deletions(-) diff --git a/pageindex/client.py b/pageindex/client.py index b10f8c542..8687695e4 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -618,19 +618,26 @@ def as_openai_tools(self, include_management: bool = False, from .integrations.openai_agents import build_openai_tools return build_openai_tools(self, include_management, hosted) - def as_anthropic_tools(self, include_management: bool = False) -> list: + def as_anthropic_tools(self, include_management: bool = False, + asynchronous: bool = False) -> list: """ Runnable tools for the Anthropic SDK's tool runner โ€” pass to - ``client.beta.messages.tool_runner(tools=...)``. + ``client.beta.messages.tool_runner(tools=...)``. 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). The Messages API's MCP connector (``mcp_servers=`` - pointing at ``{BASE_URL}/mcp``) is the server-side alternative - with no client-side tools involved. Local: the in-process tools โ€” - the same set ``messages()`` runs internally. + shape). The server-side alternative is the Messages API's beta + MCP connector โ€” ``mcp_servers=[{"type": "url", "name": + "pageindex", "url": f"{BASE_URL}/mcp", "authorization_token": + }]`` โ€” with no client-side tools + involved. Local: the in-process tools โ€” the same set + ``messages()`` runs internally. Requires ``anthropic>=0.68.0`` (``pip install 'pageindex[anthropic]'``), imported only when this @@ -641,9 +648,13 @@ def as_anthropic_tools(self, include_management: bool = False) -> list: 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. """ from .integrations.anthropic_sdk import build_anthropic_tools - return build_anthropic_tools(self, include_management) + return build_anthropic_tools(self, include_management, asynchronous) def as_claude_mcp(self, include_management: bool = False): """ diff --git a/pageindex/integrations/anthropic_sdk.py b/pageindex/integrations/anthropic_sdk.py index 36869f4e0..74d96e51a 100644 --- a/pageindex/integrations/anthropic_sdk.py +++ b/pageindex/integrations/anthropic_sdk.py @@ -7,14 +7,17 @@ """ from __future__ import annotations -from typing import Any +import asyncio +import copy +from typing import Any, Callable from ..errors import PageIndexAPIError -def build_anthropic_tools(client, include_management: bool = False) -> list: +def build_anthropic_tools(client, include_management: bool = False, + asynchronous: bool = False) -> list: try: - from anthropic import beta_tool + from anthropic import beta_async_tool, beta_tool except ImportError as exc: raise PageIndexAPIError( "as_anthropic_tools requires the Anthropic SDK tool runner " @@ -22,6 +25,27 @@ def build_anthropic_tools(client, include_management: bool = False) -> list: "'pageindex[anthropic]')." ) from exc + def wrap(name: str, description: str, schema: dict, + invoke: Callable[[dict], str]): + """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.""" + if asynchronous: + async def _afn(**kwargs: Any) -> str: + return await asyncio.to_thread(invoke, kwargs) + + _afn.__name__ = name + return beta_async_tool(_afn, name=name, description=description, + input_schema=schema) + + def _fn(**kwargs: Any) -> str: + return invoke(kwargs) + + _fn.__name__ = name + return beta_tool(_fn, name=name, description=description, + input_schema=schema) + if getattr(client, "api_key", None): from ..agent_tools import (_bridge_invoker, _cloud_bridge, _read_only_tools) @@ -32,29 +56,23 @@ def build_anthropic_tools(client, include_management: bool = False) -> list: def make_cloud(meta: dict): name = str(meta.get("name") or "tool") - invoke = _bridge_invoker(bridge, name) - - def _fn(**kwargs: Any) -> str: - return invoke(kwargs) - - _fn.__name__ = name - return beta_tool( - _fn, name=name, description=meta.get("description", ""), - input_schema=meta.get("inputSchema") - or {"type": "object", "properties": {}}, - ) + # beta_tool keeps the schema dict by reference โ€” hand out a copy, + # as _local_schema already does for the local contract. + schema = (copy.deepcopy(meta.get("inputSchema")) + or {"type": "object", "properties": {}}) + return wrap(name, meta.get("description", ""), schema, + _bridge_invoker(bridge, name)) return [make_cloud(meta) for meta in tools_meta] from ..agent_tools import (_local_description, _local_schema, call_tool, tool_names) - def make_local(name: str): - def _fn(**kwargs: Any) -> str: - return call_tool(client, name, kwargs)[0] - - _fn.__name__ = name - return beta_tool(_fn, name=name, description=_local_description(name), - input_schema=_local_schema(name)) + def local_invoke(name: str) -> Callable[[dict], str]: + def invoke(arguments: dict) -> str: + return call_tool(client, name, arguments)[0] + return invoke - return [make_local(name) for name in tool_names(include_management)] + return [wrap(name, _local_description(name), _local_schema(name), + local_invoke(name)) + for name in tool_names(include_management)] diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index afd40ccbc..39a73e3dc 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -1,5 +1,6 @@ """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 @@ -492,8 +493,12 @@ def test_as_anthropic_tools_missing_dependency(client, monkeypatch): 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") @@ -502,6 +507,17 @@ def test_as_anthropic_tools_local_in_process(client, store_path): 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 @@ -517,6 +533,9 @@ def test_as_anthropic_tools_cloud_schemas_pass_through(cloud_with_fake_bridge): 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}) @@ -524,6 +543,17 @@ def test_as_anthropic_tools_cloud_schemas_pass_through(cloud_with_fake_bridge): 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 From adb2f1fdddcb6720570629d1156f53dd814616f0 Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 12 Aug 2026 04:15:44 +0800 Subject: [PATCH 022/137] =?UTF-8?q?docs:=20doc=5Fid=20is=20per-call=20tabl?= =?UTF-8?q?e-setting=20=E2=80=94=20keep=20it=20identical=20across=20a=20co?= =?UTF-8?q?nversation?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The targeting block doc_id adds is re-set on every call and sits in the cached prompt prefix, so a round-trip that drops (or changes) doc_id silently diverges the prefix and loses the cache continuation. State the rule on all three chat surfaces' doc_id docs, and pin it with a prefix test that passes the same doc_id on both calls. --- pageindex/client.py | 9 +++++++++ tests/test_local_chat.py | 19 +++++++++++++++++++ 2 files changed, 28 insertions(+) diff --git a/pageindex/client.py b/pageindex/client.py index 8687695e4..5f3f18fee 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -368,6 +368,9 @@ def chat_completions( is appended to the managed system prompt. stream: Enable streaming responses. 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. temperature: Sampling temperature, passed through to the model. stream_metadata: With stream=True, yield chunk dicts instead of text pieces. @@ -438,6 +441,9 @@ def responses( tool outputs are emitted as ``response.output_item.done`` events and the single final event is ``response.completed``. 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. @@ -489,6 +495,9 @@ def messages( (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 and is part + of the cached prompt prefix. 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 diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index 50cbc22eb..95a1e9493 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -302,6 +302,25 @@ def test_responses_round_trip_extends_prefix(client, store_path, fake_model): assert second.inputs[0][:len(previous_final)] == previous_final +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["output"] + + [{"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_responses_stream_passthrough(client, store_path, fake_model): seed_doc(store_path, "pi-a", "report.pdf") From 3f1919b33dd611f5a654607463b4d97dcd3d3d59 Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 12 Aug 2026 04:28:13 +0800 Subject: [PATCH 023/137] feat: every chat surface takes a bare query string MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit query + doc_id is the minimal PageIndex contract, so it now works uniformly: chat_completions and messages accept a plain string (one user message), as responses always did per its wire format. The wrap is input sugar at the SDK surface, not a translation layer โ€” the outgoing wire is unchanged, and managed agent surfaces taking strings is the ecosystem convention (Runner.run, claude_agent_sdk.query). Cloud chat_completions gains the same acceptance; blank strings raise on every path. --- pageindex/client.py | 16 ++++++++++++---- pageindex/local_chat.py | 6 ++++-- tests/test_client.py | 10 ++++++++++ tests/test_local_chat.py | 24 ++++++++++++++++++++++++ 4 files changed, 50 insertions(+), 6 deletions(-) diff --git a/pageindex/client.py b/pageindex/client.py index 5f3f18fee..c7e4433a5 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -338,7 +338,7 @@ def get_retrieval(self, retrieval_id: str) -> dict[str, Any]: 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, @@ -363,7 +363,8 @@ def chat_completions( ``responses()`` or ``messages()``. Args: - messages: Conversation messages with 'role' and 'content' keys. + messages: Conversation messages with 'role' and 'content' keys, + or a bare query string (it becomes a single user message). Local also accepts system/developer messages โ€” their content is appended to the managed system prompt. stream: Enable streaming responses. @@ -386,6 +387,12 @@ def chat_completions( - stream=True, stream_metadata=False: iterator of text chunks - stream=True, stream_metadata=True: iterator of chunk dicts """ + 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 @@ -463,7 +470,7 @@ def responses( def messages( self, - messages: list[dict[str, Any]], + messages: Union[str, list[dict[str, Any]]], model: str, max_tokens: int, stream: bool = False, @@ -489,7 +496,8 @@ def messages( Args: messages: Native Messages-format history (including prior - tool_use/tool_result blocks on round-trip). + tool_use/tool_result blocks on round-trip), or a bare query + string (it becomes a single user message). model / max_tokens: Required by the Messages API; passed through. stream: Yield the Anthropic SDK's event stream across turns (its native event objects, including SDK-synthesized diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 192ef27e4..ea7f3bc35 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -568,10 +568,12 @@ def run_messages(client, messages, model: str, max_tokens: int, _require_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 list of " - "message dicts.") + 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 { diff --git a/tests/test_client.py b/tests/test_client.py index e02107432..006d015f6 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -758,3 +758,13 @@ 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(" ") diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index 95a1e9493..39f426d8a 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -205,6 +205,16 @@ def test_chat_completions_system_and_doc_block(client, store_path, fake_model): @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(" ") + + def test_chat_completions_validation(client, store_path, fake_model): fake_model([[_msg_item("ok")]]) with pytest.raises(PageIndexAPIError, match="cloud-only"): @@ -466,6 +476,20 @@ def test_messages_stream_passthrough(client, store_path, fake_anthropic): 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", + max_tokens=100) + assert result["content"][0]["text"] == "ok" + assert calls[0]["messages"] == [{"role": "user", + "content": "What status?"}] + with pytest.raises(PageIndexAPIError, match="non-empty string"): + client.messages(" ", model="claude-test", max_tokens=100) + + @needs_anthropic def test_messages_validation(client, fake_anthropic): fake_anthropic([]) From d28a8af84691ff455d745be8ad53804a31d52ac2 Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 12 Aug 2026 04:50:13 +0800 Subject: [PATCH 024/137] feat: messages() defaults max_tokens to 4096 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Messages API requires a per-turn output budget on the wire, but that is table-setting, not a PageIndex-layer user obligation โ€” the simple call is now a question + model + doc_id. The knob stays overridable (passthrough intact); model stays required because no cross-vendor default is honest to guess. --- pageindex/client.py | 7 +++++-- pageindex/local_chat.py | 2 +- tests/test_local_chat.py | 7 ++++--- 3 files changed, 10 insertions(+), 6 deletions(-) diff --git a/pageindex/client.py b/pageindex/client.py index c7e4433a5..3869da409 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -472,7 +472,7 @@ def messages( self, messages: Union[str, list[dict[str, Any]]], model: str, - max_tokens: int, + max_tokens: int = 4096, stream: bool = False, doc_id: Optional[Union[str, list[str]]] = None, system: Optional[Union[str, list[dict[str, Any]]]] = None, @@ -498,7 +498,10 @@ def messages( 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 / max_tokens: Required by the Messages API; passed through. + model: Required โ€” there is no cross-vendor default to guess. + max_tokens: Per-turn output budget the Messages API requires on + the wire; defaults to 4096 so the simple call needs only a + question. Passed through. stream: Yield the Anthropic SDK's event stream across turns (its native event objects, including SDK-synthesized convenience events), one message sequence per turn. diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index ea7f3bc35..2589767be 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -556,7 +556,7 @@ def _anthropic_usage(turns, final_usage: dict) -> dict: return totals -def run_messages(client, messages, model: str, max_tokens: int, +def run_messages(client, messages, model: str, max_tokens: int = 4096, stream: bool = False, doc_id=None, system=None, temperature: Optional[float] = None, top_p: Optional[float] = None, diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index 39f426d8a..10043f9b4 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -481,13 +481,14 @@ 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", - max_tokens=100) + 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"] == 4096 with pytest.raises(PageIndexAPIError, match="non-empty string"): - client.messages(" ", model="claude-test", max_tokens=100) + client.messages(" ", model="claude-test") @needs_anthropic From 2de3b18e134edca23e1c50e137f7f7b712cf7089 Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 12 Aug 2026 04:51:25 +0800 Subject: [PATCH 025/137] fix: raise messages() max_tokens default to 8192 max_tokens is a cap, not consumption, so the default should be the highest universally safe value: 4096 could truncate long-form answers (whole-document summaries), while 8192 is the output ceiling every non-EOL Claude model accepts and stays under the SDK's non-streaming long-request threshold. --- pageindex/client.py | 6 +++--- pageindex/local_chat.py | 2 +- tests/test_local_chat.py | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/pageindex/client.py b/pageindex/client.py index 3869da409..a3cc91fd1 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -472,7 +472,7 @@ def messages( self, messages: Union[str, list[dict[str, Any]]], model: str, - max_tokens: int = 4096, + max_tokens: int = 8192, stream: bool = False, doc_id: Optional[Union[str, list[str]]] = None, system: Optional[Union[str, list[dict[str, Any]]]] = None, @@ -500,8 +500,8 @@ def messages( 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; defaults to 4096 so the simple call needs only a - question. Passed through. + the wire; defaults to 8192 โ€” the ceiling every current model + accepts โ€” so the simple call needs only a question. Passed through. stream: Yield the Anthropic SDK's event stream across turns (its native event objects, including SDK-synthesized convenience events), one message sequence per turn. diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 2589767be..8742d0f6f 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -556,7 +556,7 @@ def _anthropic_usage(turns, final_usage: dict) -> dict: return totals -def run_messages(client, messages, model: str, max_tokens: int = 4096, +def run_messages(client, messages, model: str, max_tokens: int = 8192, stream: bool = False, doc_id=None, system=None, temperature: Optional[float] = None, top_p: Optional[float] = None, diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index 10043f9b4..c82711b45 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -486,7 +486,7 @@ def test_messages_accepts_query_string(client, fake_anthropic): 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"] == 4096 + assert calls[0]["max_tokens"] == 8192 with pytest.raises(PageIndexAPIError, match="non-empty string"): client.messages(" ", model="claude-test") From 2607a86b070ef1d3e37822f1148bb5c86e0f346c Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 12 Aug 2026 05:01:18 +0800 Subject: [PATCH 026/137] fix: restore per-extra skip markers the string-input tests displaced Inserting tests above decorated ones absorbed their @needs_agents markers, so two tests ran (and failed) in the without-frameworks CI job. Both simulated-bare and full runs are green again. --- tests/test_local_chat.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index c82711b45..98a2c0c9e 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -215,6 +215,7 @@ def test_chat_completions_accepts_query_string(client, store_path, fake_model): 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"): @@ -312,6 +313,7 @@ def test_responses_round_trip_extends_prefix(client, store_path, fake_model): 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.""" From d87fa8933e51eab625bcf801dd244b6848216c52 Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 12 Aug 2026 14:48:11 +0800 Subject: [PATCH 027/137] fix: close 17 findings from the v0.2.10 max review MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Tool layer: - anthropic adapter: failed tool calls raise ToolError so the runner emits tool_result is_error:true; McpBridge.call_tool returns (text, is_error) and surfaces the server's MCP isError marking - as_openai_tools builds FunctionTool with the contract/server schema verbatim (strict off) โ€” function_tool() regenerated schemas from signatures, dropping items/enum/pattern/bounds and aborting the whole list on object-typed params; shared _tool_specs() feeds both adapters - remove_document validates every name before deleting anything; call_tool classifies only bind-time TypeErrors as INVALID_INPUT - unknown-tool envelope formatted with _dumps like every other envelope Local chat: - doc_id is enforced at the tool layer (allowlist threaded through call_tool and the adapters), not just prompted; the shadow check runs inside the scope - _openai_model routes litellm/ and provider/ paths via LitellmModel and strips openai/ โ€” the normalized retrieve_model 404'd as a raw wire name - responses() reports the backend's real terminal status (recorded at the transport client; the framework discards Response.status) and wraps framework exceptions in PageIndexAPIError - chat_completions streaming yields its opening chunk inside try, so an abandoned iterator still cancels the run and closes the backend - prompt-cache group_id is per-conversation (model+instructions+first item) instead of one global constant pooling every user - messages() max_tokens default resolves per model (claude-3 caps at 4096) Packaging / surface: - __init__ registers the 0.2.10 modules in _SUBMODULES; unknown names raise AttributeError instead of eagerly importing page_index_classic - anthropic floor 0.84.0: first release with ToolError whose runner also executes the final turn's tools on a max_iterations cut - client docstrings caught up with local chat landing Claude Agent SDK gate: - claude_allowed_tools(mcp_servers) derives mcp____ entries from the caller's own registration map (live server annotations on cloud, the contract locally) โ€” no name is ever spelled twice - claude_agent_config() bundles the three slots as one-call sugar over the explicit form Examples: - demo runs against cloud again (getattr for local-only attrs) and finds an existing indexed copy by name before re-indexing Tests: monkeypatches replace the consuming module's binding instead of mutating the shared time/requests modules; 185 -> 211. --- examples/agentic_vectorless_rag_demo.py | 12 +- pageindex/__init__.py | 21 +- pageindex/agent_tools.py | 133 +++++++-- pageindex/client.py | 112 ++++++-- pageindex/integrations/anthropic_sdk.py | 59 ++-- pageindex/integrations/claude_agent_sdk.py | 34 +++ pageindex/integrations/openai_agents.py | 35 ++- pageindex/local_chat.py | 158 ++++++++--- pageindex/mcp_bridge.py | 10 +- pyproject.toml | 8 +- tests/test_agent_tools.py | 300 +++++++++++++++++++-- tests/test_local_chat.py | 216 ++++++++++++++- tests/test_package_surface.py | 21 ++ 13 files changed, 951 insertions(+), 168 deletions(-) diff --git a/examples/agentic_vectorless_rag_demo.py b/examples/agentic_vectorless_rag_demo.py index f35c3c2e7..7a83776ff 100644 --- a/examples/agentic_vectorless_rag_demo.py +++ b/examples/agentic_vectorless_rag_demo.py @@ -53,7 +53,9 @@ def query_agent(client: PageIndexLocalClient, doc_id: str, prompt: str, verbose: name="PageIndex", instructions=client.agent_instructions(doc_id=doc_id), tools=client.as_openai_tools(), - model=client.retrieve_model, + # retrieve_model is a local-mode attribute; cloud clients fall back + # to the framework's default model. + model=getattr(client, "retrieve_model", None), # model_settings=ModelSettings(reasoning={"effort": "low", "summary": "auto"}), # from agents.model_settings import ModelSettings ) @@ -138,7 +140,15 @@ async def _run(): doc_id = cached except PageIndexAPIError: DOC_ID_PATH.unlink() + if doc_id is None: + # The .doc_id cache is gitignored โ€” on a fresh clone with an + # existing store, find the already-indexed copy by name instead of + # re-indexing it. + doc_id = next( + (doc["id"] for doc in client.list_documents(limit=100)["documents"] + if doc["name"] == PDF_PATH.name), None) if doc_id: + DOC_ID_PATH.write_text(doc_id) print(f"\nLoaded cached doc_id: {doc_id}") else: doc_id = client.submit_document(str(PDF_PATH), wait=True)["doc_id"] diff --git a/pageindex/__init__.py b/pageindex/__init__.py index 3513668a2..5a19bd895 100644 --- a/pageindex/__init__.py +++ b/pageindex/__init__.py @@ -18,26 +18,27 @@ ] _LAZY = { + "page_index": ".page_index_classic", + "page_index_main": ".page_index_classic", "page_index_flash": ".flash", "optimize_tree": ".tree_optimize", "md_to_tree": ".page_index_md", } -_SUBMODULES = {"client", "cloud_api", "errors", "flash", "local_api", - "local_store", "page_index_classic", "page_index_md", "tree_optimize", - "utils"} +_SUBMODULES = {"agent_tools", "client", "cloud_api", "errors", "flash", + "integrations", "local_api", "local_chat", "local_store", + "mcp_bridge", "page_index_classic", "page_index_md", + "tree_optimize", "utils"} def __getattr__(name): - if name.startswith("_"): - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") import importlib if name in _SUBMODULES: return importlib.import_module(f".{name}", __name__) - module = importlib.import_module(_LAZY.get(name, ".page_index_classic"), __name__) - try: - value = getattr(module, name) - except AttributeError: - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None + if name not in _LAZY: + # Unknown names must not fall through to an eager import of the + # heavy indexing stack. + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + value = getattr(importlib.import_module(_LAZY[name], __name__), name) globals()[name] = value return value diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index e0069faeb..02c006036 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -19,6 +19,7 @@ import copy import difflib +import inspect import json import re import threading @@ -366,14 +367,23 @@ def _flat_metadata(value: Any) -> Optional[dict[str, Any]]: return flat or None +def _scope_documents(documents: list[dict[str, Any]], + allowed_ids: Optional[frozenset]) -> list[dict[str, Any]]: + if allowed_ids is None: + return documents + return [doc for doc in documents if doc.get("id") in allowed_ids] + + def _resolve_document( client, doc_name: str, documents: Optional[list[dict[str, Any]]] = None, + allowed_ids: Optional[frozenset] = None, ) -> "tuple[Optional[dict[str, Any]], Optional[_ToolResult]]": """Resolve doc_name to a list entry. Same-name duplicates resolve to the newest match. Returns (entry, None) or (None, error_payload_pair).""" if documents is None: documents = _all_documents(client) + documents = _scope_documents(documents, allowed_ids) matches = [doc for doc in documents if doc.get("name") == doc_name] if matches: return max(matches, key=lambda d: d.get("createdAt") or ""), None @@ -640,7 +650,8 @@ def _split_oversized_node(node: Any, budget: int) -> list[Any]: def _browse_documents(client, folder_id: str = "root", recursive: bool = False, sort: str = "time", query: Optional[str] = None, - offset: int = 0, limit: int = 10) -> tuple[dict, bool]: + 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"): @@ -677,9 +688,14 @@ def _browse_documents(client, folder_id: str = "root", recursive: bool = False, "options": ["Pass integer offset and limit values"]}, "INVALID_INPUT") - listing = client.list_documents(limit=limit, offset=offset) - window = listing.get("documents") or [] - has_more = offset + limit < listing.get("total", 0) + if _allowed_ids is None: + listing = client.list_documents(limit=limit, offset=offset) + window = listing.get("documents") or [] + total = listing.get("total", 0) + else: + scoped = _scope_documents(_all_documents(client), _allowed_ids) + window, total = scoped[offset:offset + limit], len(scoped) + has_more = offset + limit < total next_offset = offset + limit if has_more else None page_has_processing = False @@ -747,10 +763,11 @@ def _browse_documents(client, folder_id: str = "root", recursive: bool = False, def _get_document(client, doc_name: str, folder_id: Optional[str] = None, - wait_for_completion: bool = False) -> tuple[dict, bool]: + 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) + entry, error = _resolve_document(client, doc_name, allowed_ids=_allowed_ids) if error is not None: return error assert entry is not None @@ -815,10 +832,11 @@ def _get_document(client, doc_name: str, folder_id: Optional[str] = None, def _get_document_structure(client, doc_name: str, folder_id: Optional[str] = None, part: int = 1, - wait_for_completion: bool = False) -> tuple[dict, bool]: + 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) + entry, error = _resolve_document(client, doc_name, allowed_ids=_allowed_ids) if error is not None: return error assert entry is not None @@ -920,10 +938,11 @@ def _get_document_structure(client, doc_name: str, def _get_page_content(client, doc_name: str, pages: str, folder_id: Optional[str] = None, - wait_for_completion: bool = False) -> tuple[dict, bool]: + 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) + entry, error = _resolve_document(client, doc_name, allowed_ids=_allowed_ids) if error is not None: return error assert entry is not None @@ -1032,7 +1051,8 @@ def _get_page_content(client, doc_name: str, pages: str, def _remove_document(client, doc_names: list[str], - folder_id: Optional[str] = None) -> tuple[dict, bool]: + 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: @@ -1040,6 +1060,16 @@ def _remove_document(client, doc_names: list[str], {"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") if len(doc_names) > 10: return _failure("Maximum 10 documents can be deleted at once", None, {"summary": "Too many documents in one call", @@ -1048,7 +1078,8 @@ def _remove_document(client, doc_names: list[str], documents = _all_documents(client) results = [] for doc_name in doc_names: - entry, error = _resolve_document(client, doc_name, documents=documents) + 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 @@ -1081,9 +1112,12 @@ def tool_names(include_management: bool = False) -> tuple[str, ...]: return _READ_TOOLS + (_MANAGEMENT_TOOLS if include_management else ()) -def call_tool(client, name: str, arguments: dict[str, Any]) -> tuple[str, bool]: +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.""" + 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( @@ -1093,9 +1127,16 @@ def call_tool(client, name: str, arguments: dict[str, Any]) -> tuple[str, bool]: "options": [f"Available tools: {', '.join(_IMPLEMENTATIONS)}"]}, "INVALID_INPUT", ) - return json.dumps(payload), True + return _dumps(payload), True + # Underscore-prefixed keys are the SDK's private channel (the scope + # below), never model arguments. + kwargs = {key: value for key, value in arguments.items() + if not key.startswith("_")} + 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) try: - payload, is_error = implementation(client, **arguments) + bound = inspect.signature(implementation).bind(client, **kwargs) except TypeError as exc: payload, is_error = _failure( f"Invalid arguments for {name}: {exc}", None, @@ -1103,6 +1144,9 @@ def call_tool(client, name: str, arguments: dict[str, Any]) -> tuple[str, bool]: "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, @@ -1220,11 +1264,12 @@ def _annotation_for(spec: dict) -> Any: return _SCHEMA_TYPE_MAP.get(schema_type, Any) -def _bridge_invoker(bridge, name: str) -> Callable[[dict], str]: +def _bridge_invoker(bridge, name: str) -> "Callable[[dict], tuple[str, bool]]": """One cloud tool call proxied over MCP: None-valued arguments are dropped (None โ‰ก omitted, matching the contract's "omit if ..." - semantics) and failures are contained in the error envelope.""" - def _invoke(arguments: dict[str, Any]) -> str: + semantics) and failures are contained in the error envelope. Returns + (envelope_text, is_error), like call_tool.""" + def _invoke(arguments: dict[str, Any]) -> tuple[str, bool]: arguments = {key: value for key, value in arguments.items() if value is not None} try: @@ -1238,7 +1283,7 @@ def _invoke(arguments: dict[str, Any]) -> str: "try the request again"}, "INTERNAL_ERROR", ) - return _dumps(payload) + return _dumps(payload), True return _invoke @@ -1258,7 +1303,7 @@ def _make_bridge_function(bridge, meta: dict) -> Callable[..., str]: for param in properties) if not params_usable: def proxy(**kwargs: Any) -> str: - return _invoke(kwargs) + 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]) @@ -1269,7 +1314,7 @@ def proxy(**kwargs: Any) -> str: 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})", namespace) + f" return _invoke({args_literal})[0]", namespace) proxy = namespace["_synthesized"] annotations: dict[str, Any] = {} for p in ordered: @@ -1330,6 +1375,39 @@ def _build_cloud_agent_tools(client, include_management: bool) -> list[Callable[ return [_make_bridge_function(bridge, meta) for meta in tools_meta] +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.""" + if getattr(client, "api_key", None): + if doc_ids is not None: + raise PageIndexAPIError( + "doc_ids scoping applies to local tools only โ€” cloud calls " + "are scoped server-side." + ) + bridge = _cloud_bridge(client) + tools_meta = bridge.list_tools() + if not include_management: + tools_meta = _read_only_tools(tools_meta) + return [(str(meta.get("name") or "tool"), + meta.get("description") or "", + copy.deepcopy(meta.get("inputSchema")) + or {"type": "object", "properties": {}}, + _bridge_invoker(bridge, str(meta.get("name") or "tool"))) + for meta in tools_meta] + + def local_invoke(name: str) -> "Callable[[dict], tuple[str, bool]]": + def invoke(arguments: dict) -> tuple[str, bool]: + return call_tool(client, name, arguments, doc_ids=doc_ids) + return invoke + + return [(name, _local_description(name), _local_schema(name), + local_invoke(name)) + for name in tool_names(include_management)] + + def build_agent_tools(client, include_management: bool = False) -> list[Callable[..., str]]: """Plain synchronous functions bound to `client`. @@ -1461,19 +1539,24 @@ def _base_instructions(client) -> str: return instructions -def doc_targeting_block(client, doc_id) -> Optional[str]: +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.""" + same-name document โ€” the name-addressed tools could not reach it. With + ``scoped`` (the chat surfaces, whose tools resolve names inside the + doc_id allowlist) only a same-name duplicate within the targeted set + shadows.""" if doc_id is None: return None doc_ids = [doc_id] if isinstance(doc_id, str) else list(doc_id) if not doc_ids: return None details = [client.get_document(one_id) for one_id in doc_ids] - documents = _all_documents(client) + documents = ([{**detail, "id": one_id} + for one_id, detail in zip(doc_ids, details)] + if scoped else _all_documents(client)) for one_id, detail in zip(doc_ids, details): entry, _ = _resolve_document(client, str(detail.get("name")), documents=documents) diff --git a/pageindex/client.py b/pageindex/client.py index a3cc91fd1..868e37e2f 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -50,10 +50,9 @@ class PageIndexClient: trees. Defaults to the packaged config (see pageindex/config.yaml). summary_model (str, optional): Local mode only โ€” LLM used for node summaries and document descriptions. - retrieve_model (str, optional): Local mode only โ€” exposed as - ``client.retrieve_model`` (the agent demo reads it); the SDK - itself consumes it once agent-based local chat lands in a - later release. + retrieve_model (str, optional): Local mode only โ€” the model the + local chat surfaces (``chat_completions``, ``responses``) + default to, exposed as ``client.retrieve_model``. storage_path (str, optional): Local mode only โ€” directory where indexed documents are stored. Defaults to ``./.pageindex``. @@ -65,10 +64,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" @@ -314,11 +312,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]: @@ -327,11 +325,11 @@ 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 ---------- @@ -472,7 +470,7 @@ def messages( self, messages: Union[str, list[dict[str, Any]]], model: str, - max_tokens: int = 8192, + 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, @@ -500,8 +498,9 @@ def messages( 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; defaults to 8192 โ€” the ceiling every current model - accepts โ€” so the simple call needs only a question. Passed through. + the wire; the default is resolved per model (8192, or 4096 + for the claude-3 generation whose ceiling is lower) so the + simple call needs only a question. Passed through. stream: Yield the Anthropic SDK's event stream across turns (its native event objects, including SDK-synthesized convenience events), one message sequence per turn. @@ -659,7 +658,7 @@ def as_anthropic_tools(self, include_management: bool = False, involved. Local: the in-process tools โ€” the same set ``messages()`` runs internally. - Requires ``anthropic>=0.68.0`` + Requires ``anthropic>=0.84.0`` (``pip install 'pageindex[anthropic]'``), imported only when this method is called. @@ -682,12 +681,11 @@ def as_claude_mcp(self, include_management: bool = False): Cloud: returns the remote PageIndex MCP config โ€” the framework connects to api.pageindex.ai/mcp directly and discovers the full - cloud tool set. ``include_management`` has no effect there; gate - destructive tools with the framework's permission layer (e.g. list - read tools in ``allowed_tools`` instead of the ``*`` wildcard, or - add ``disallowed_tools=["mcp__pageindex__remove_document"]``). - Local: returns an in-process SDK MCP server exposing the agent - tools (requires ``claude-agent-sdk``; + cloud tool set. A remote server cannot be filtered client-side, + so ``include_management`` has no effect there โ€” the gate is + ``allowed_tools``, built from your registration map by + ``claude_allowed_tools()``. Local: returns an in-process SDK MCP + server exposing the agent tools (requires ``claude-agent-sdk``; ``pip install 'pageindex[claude]'``). Cloud hosts that surface MCP server instructions receive the same @@ -696,16 +694,80 @@ def as_claude_mcp(self, include_management: bool = False): recommended channel: it is guaranteed delivery, carries ``doc_id`` targeting, and is the only channel local mode has. - Usage:: + Usage (or ``claude_agent_config()`` for all three slots in one + call):: + servers = {"pageindex": client.as_claude_mcp()} options = ClaudeAgentOptions( - mcp_servers={"pageindex": client.as_claude_mcp()}, - allowed_tools=["mcp__pageindex__*"], + system_prompt=client.agent_instructions(), + mcp_servers=servers, + allowed_tools=client.claude_allowed_tools(servers), ) """ from .integrations.claude_agent_sdk import build_claude_mcp return build_claude_mcp(self, include_management) + def claude_allowed_tools(self, mcp_servers: dict[str, Any], + include_management: bool = False) -> list[str]: + """ + ``allowed_tools`` entries for the PageIndex servers in your + ``mcp_servers`` map โ€” pass the same dict you hand to + ``ClaudeAgentOptions``. The framework bakes the registration key + into every tool id (``mcp____``), so the keys are read + from the map rather than spelled a second time, and the tool + names are the read-only gate every other adapter applies โ€” live + server annotations on cloud, the tool contract locally. Nothing + is hand-maintained, and no framework install is needed. + + Raises PageIndexAPIError when the map holds no PageIndex entry โ€” + a gate list that silently matched nothing would disable every + tool. + + Args: + mcp_servers: The registration map; non-PageIndex entries are + ignored. + include_management (bool): Also allow tools that modify the + library (``remove_document``, and on cloud the server's + full management list). + """ + from .integrations.claude_agent_sdk import build_claude_allowed_tools + return build_claude_allowed_tools(self, mcp_servers, + include_management) + + 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``), the server entry (``as_claude_mcp``), + and the matching ``allowed_tools`` gate + (``claude_allowed_tools``), with one ``include_management`` and + ``server_name`` applied everywhere. To customize (your own + system prompt, extra servers), switch to those three methods + directly. + + Args: + doc_id: Document ID or list of IDs to target, as in + ``agent_instructions``. + include_management (bool): Also allow tools that modify the + library. + server_name (str): Key the server is registered under. + """ + servers = {server_name: self.as_claude_mcp(include_management)} + return { + "system_prompt": self.agent_instructions(doc_id=doc_id), + "mcp_servers": servers, + "allowed_tools": self.claude_allowed_tools(servers, + include_management), + } + def agent_instructions(self, doc_id: Optional[Union[str, list[str]]] = None) -> str: """ Orchestration guidance for document QA agents โ€” pass as the agent's diff --git a/pageindex/integrations/anthropic_sdk.py b/pageindex/integrations/anthropic_sdk.py index 74d96e51a..4405df5a1 100644 --- a/pageindex/integrations/anthropic_sdk.py +++ b/pageindex/integrations/anthropic_sdk.py @@ -3,76 +3,57 @@ 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. +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 -import copy -from typing import Any, Callable +from typing import Any from ..errors import PageIndexAPIError def build_anthropic_tools(client, include_management: bool = False, - asynchronous: bool = False) -> list: + 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.68.0) โ€” pip install -U anthropic (or pip install " + "(anthropic>=0.84.0) โ€” pip install -U anthropic (or pip install " "'pageindex[anthropic]')." ) from exc + from ..agent_tools import _tool_specs - def wrap(name: str, description: str, schema: dict, - invoke: Callable[[dict], str]): + 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(invoke, kwargs) + 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 invoke(kwargs) + return run(kwargs) _fn.__name__ = name return beta_tool(_fn, name=name, description=description, input_schema=schema) - if getattr(client, "api_key", None): - from ..agent_tools import (_bridge_invoker, _cloud_bridge, - _read_only_tools) - bridge = _cloud_bridge(client) - tools_meta = bridge.list_tools() - if not include_management: - tools_meta = _read_only_tools(tools_meta) - - def make_cloud(meta: dict): - name = str(meta.get("name") or "tool") - # beta_tool keeps the schema dict by reference โ€” hand out a copy, - # as _local_schema already does for the local contract. - schema = (copy.deepcopy(meta.get("inputSchema")) - or {"type": "object", "properties": {}}) - return wrap(name, meta.get("description", ""), schema, - _bridge_invoker(bridge, name)) - - return [make_cloud(meta) for meta in tools_meta] - - from ..agent_tools import (_local_description, _local_schema, call_tool, - tool_names) - - def local_invoke(name: str) -> Callable[[dict], str]: - def invoke(arguments: dict) -> str: - return call_tool(client, name, arguments)[0] - return invoke - - return [wrap(name, _local_description(name), _local_schema(name), - local_invoke(name)) - for name in tool_names(include_management)] + 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 index b5e599da0..a28cd155f 100644 --- a/pageindex/integrations/claude_agent_sdk.py +++ b/pageindex/integrations/claude_agent_sdk.py @@ -13,6 +13,40 @@ from ..errors import PageIndexAPIError +def build_claude_allowed_tools(client, mcp_servers, + include_management: bool = False) -> list[str]: + """``allowed_tools`` entries for the PageIndex entries of an + mcp_servers map. The framework scopes every tool id by the map key + (``mcp____``), so the keys are read from the map instead of + being spelled a second time; tool names are the gated set โ€” live + server annotations on cloud, the contract locally. Needs no framework + import.""" + from ..agent_tools import _tool_specs + if not isinstance(mcp_servers, dict) or not mcp_servers: + raise PageIndexAPIError( + "claude_allowed_tools takes the mcp_servers dict you register " + "with the framework (the {name: server} map)." + ) + + def is_pageindex(value) -> bool: + get = (value.get if isinstance(value, dict) + else lambda key, default=None: getattr(value, key, default)) + url = get("url") + if isinstance(url, str): + return url.startswith(f"{client.BASE_URL}/mcp") + return get("type") == "sdk" and get("name") == "pageindex" + + keys = [key for key, value in mcp_servers.items() if is_pageindex(value)] + if not keys: + raise PageIndexAPIError( + "No PageIndex server found in mcp_servers โ€” register " + "client.as_claude_mcp() under a key first (an allowed_tools " + "list built from this map would match nothing)." + ) + names = [spec[0] for spec in _tool_specs(client, include_management)] + return [f"mcp__{key}__{name}" for key in keys for name in names] + + def build_claude_mcp(client, include_management: bool = False): if getattr(client, "api_key", None): return { diff --git a/pageindex/integrations/openai_agents.py b/pageindex/integrations/openai_agents.py index f33c4b587..d6cf52948 100644 --- a/pageindex/integrations/openai_agents.py +++ b/pageindex/integrations/openai_agents.py @@ -3,17 +3,25 @@ Cloud clients default to the full live 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). -Local clients get the in-process tools wrapped as FunctionTools. +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) -> list: + hosted: bool = False, doc_ids=None) -> list: try: - from agents import HostedMCPTool, function_tool + from agents import FunctionTool, HostedMCPTool except ImportError as exc: raise PageIndexAPIError( "as_openai_tools requires the OpenAI Agents SDK โ€” " @@ -32,6 +40,21 @@ def build_openai_tools(client, include_management: bool = False, "headers": {"Authorization": f"Bearer {client.api_key}"}, "require_approval": require_approval, })] - from ..agent_tools import build_agent_tools - return [function_tool(tool) - for tool in build_agent_tools(client, include_management)] + from ..agent_tools import _tool_specs + + def wrap(name, description, schema, invoke): + async def on_invoke_tool(ctx: Any, args_json: str) -> str: + arguments = {key: value for key, value + in (json.loads(args_json) if args_json else {}).items() + if value is not None} + text, _ = await asyncio.to_thread(invoke, arguments) + return text + + return FunctionTool(name=name, description=description, + params_json_schema=schema, + on_invoke_tool=on_invoke_tool, + strict_json_schema=False) + + return [wrap(*spec) + for spec in _tool_specs(client, include_management, + doc_ids=doc_ids)] diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 8742d0f6f..8dd01f64c 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -11,15 +11,19 @@ Content passes through untouched โ€” the caller's messages, the model's answers, tool outputs. Native stop reasons pass through on ``messages``; the OpenAI engine's abstraction does not surface per-turn finish reasons, -so ``chat_completions`` reports loop completion as ``"stop"`` and -``responses`` as ``status: "completed"``. The SDK owns gatekeeping -(structural validation), table-setting (managed instructions, tools, doc -targeting), tool execution, and billing (usage aggregation, envelope ids). +so ``chat_completions`` reports loop completion as ``"stop"``, while +``responses`` reports the backend's terminal ``status`` where the wire +surfaces one (recorded at the transport layer โ€” the framework discards +it). The SDK owns gatekeeping (structural validation), table-setting +(managed instructions, tools, doc targeting), tool execution, and billing +(usage aggregation, envelope ids). """ from __future__ import annotations import asyncio import concurrent.futures +import hashlib +import json import queue import threading import time @@ -58,7 +62,10 @@ def _doc_block(client, doc_id) -> Optional[str]: raise PageIndexAPIError( "Documents not found or access denied: " + ", ".join(missing) ) - return doc_targeting_block(client, doc_id) + # 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: @@ -200,8 +207,17 @@ def _require_openai_agents(method: str) -> None: def _openai_model(protocol: str, model_name: str): - """The backend protocol driver โ€” the seam tests replace with a fake.""" + """The backend protocol driver โ€” the seam tests replace with a fake. + + ``litellm//`` (the client's normalized retrieve_model + form) and bare ``/`` paths drive the provider through + LiteLLM; an ``openai/`` prefix strips to the OpenAI SDK; bare names go + to the OpenAI SDK as-is.""" + if "/" in model_name and not model_name.startswith("openai/"): + from agents.extensions.models.litellm_model import LitellmModel + return LitellmModel(model_name.removeprefix("litellm/")) from openai import AsyncOpenAI + model_name = model_name.removeprefix("openai/") if protocol == "chat": from agents.models.openai_chatcompletions import ( OpenAIChatCompletionsModel) @@ -211,13 +227,13 @@ def _openai_model(protocol: str, model_name: str): def _openai_agent(client, protocol: str, model_name: str, instructions: str, - temperature, top_p): + temperature, top_p, doc_ids=None): from agents import Agent, ModelSettings from .integrations.openai_agents import build_openai_tools return Agent( name="PageIndex", instructions=instructions, - tools=build_openai_tools(client), + tools=build_openai_tools(client, doc_ids=doc_ids), model=_openai_model(protocol, model_name), model_settings=ModelSettings(temperature=temperature, top_p=top_p), ) @@ -229,20 +245,55 @@ def _validate_max_turns(max_turns) -> None: raise PageIndexAPIError("max_turns must be a positive integer.") -def _run_kwargs(max_turns) -> dict: +def _conversation_group_id(model_name: str, instructions: str, items) -> str: + """Stable per-conversation cache-routing key: openai-agents hashes + RunConfig.group_id into the OpenAI prompt_cache_key, and without one it + stamps every run with a fresh key, tagging a round-tripped prefix as a + different cache group. Keyed on the prefix identity โ€” model, + instructions, first input item โ€” so a conversation's continuations + share one route without pooling unrelated conversations.""" + seed = json.dumps([model_name, instructions, + items[0] if items else None], + sort_keys=True, default=str) + return "pageindex-" + hashlib.sha256(seed.encode()).hexdigest()[:16] + + +def _run_kwargs(max_turns, group_id: str) -> dict: # Managed runs never export traces โ€” the caller opted into document QA, - # not telemetry. The stable group_id keys OpenAI's prompt-cache routing: - # without it openai-agents stamps every run with a fresh - # prompt_cache_key, tagging a round-tripped prefix as a different cache - # group. + # not telemetry. from agents import RunConfig kwargs: dict = {"run_config": RunConfig(tracing_disabled=True, - group_id="pageindex-local-chat")} + group_id=group_id)} if max_turns is not None: kwargs["max_turns"] = max_turns return kwargs +def _record_response_status(agent, recorded: dict) -> None: + """Capture each turn's terminal Response status at the transport client: + openai-agents' non-streaming path discards Response.status, so a final + turn truncated at the output cap would otherwise report as a clean + completion. No-op for backends without an OpenAI responses resource + (the streaming path records from lifecycle events instead).""" + responses = getattr(getattr(getattr(agent, "model", None), "_client", None), + "responses", None) + create = getattr(responses, "create", None) + if create is None: + return + + async def recording_create(*args, **kwargs): + response = await create(*args, **kwargs) + if getattr(response, "status", None): + recorded["status"] = response.status + for field in ("incomplete_details", "error"): + value = getattr(response, field, None) + recorded[field] = (value.model_dump(mode="json") + if hasattr(value, "model_dump") else value) + return response + + responses.create = recording_create + + async def _aclose_backend(agent) -> None: """Close the per-call AsyncOpenAI client before its event loop ends โ€” otherwise httpx tears down pooled connections on a closed loop and @@ -296,15 +347,17 @@ def run_chat_completions(client, messages, stream: bool = False, block = _doc_block(client, doc_id) items = ([{"role": "user", "content": block}] if block else []) + history model_name = model or client.retrieve_model - agent = _openai_agent(client, "chat", model_name, - _managed_instructions(system_texts), - temperature, None) + managed = _managed_instructions(system_texts) + agent = _openai_agent(client, "chat", model_name, managed, + temperature, None, doc_ids=doc_id or None) + run_kwargs = _run_kwargs(max_turns, + _conversation_group_id(model_name, managed, items)) from agents import Runner from agents.exceptions import MaxTurnsExceeded if not stream: try: result = _run_sync(_run_closing(agent, - Runner.run(agent, input=items, **_run_kwargs(max_turns)))) + Runner.run(agent, input=items, **run_kwargs))) except MaxTurnsExceeded as exc: raise _wrap_max_turns(exc, max_turns) from exc return { @@ -334,11 +387,12 @@ def chunk(delta: dict, finish=None) -> dict: async def agen(): from openai.types.responses import ResponseTextDeltaEvent - streamed = Runner.run_streamed(agent, input=items, - **_run_kwargs(max_turns)) - yield chunk({"role": "assistant", "content": ""}) + 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)): @@ -390,9 +444,12 @@ def run_responses(client, input, model: Optional[str] = None, model_name = model or client.retrieve_model managed = _managed_instructions(extra) agent = _openai_agent(client, "responses", model_name, managed, - temperature, top_p) + temperature, top_p, doc_ids=doc_id or None) + run_kwargs = _run_kwargs(max_turns, + _conversation_group_id(model_name, managed, items)) + recorded: dict = {} from agents import Runner - from agents.exceptions import MaxTurnsExceeded + from agents.exceptions import AgentsException, MaxTurnsExceeded def envelope(output: list, raw_responses) -> dict: usage = _openai_usage(raw_responses) @@ -401,7 +458,7 @@ def envelope(output: list, raw_responses) -> dict: "object": "response", "created_at": int(time.time()), "model": model_name, - "status": "completed", + "status": recorded.get("status") or "completed", "output": output, "usage": {"input_tokens": usage["prompt_tokens"], "output_tokens": usage["completion_tokens"], @@ -417,18 +474,22 @@ def envelope(output: list, raw_responses) -> dict: "temperature": temperature, "top_p": top_p, "max_output_tokens": None, - "error": None, - "incomplete_details": None, + "error": recorded.get("error"), + "incomplete_details": recorded.get("incomplete_details"), "metadata": None, } if not stream: + _record_response_status(agent, recorded) try: result = _run_sync(_run_closing(agent, Runner.run(agent, input=[dict(item) for item in items], - **_run_kwargs(max_turns)))) + **run_kwargs))) except MaxTurnsExceeded as exc: raise _wrap_max_turns(exc, max_turns) from exc + except AgentsException as exc: + raise PageIndexAPIError( + f"The agent backend failed: {exc}") from exc output = result.to_input_list()[len(items):] return envelope(output, result.raw_responses) @@ -443,7 +504,7 @@ def envelope(output: list, raw_responses) -> dict: async def agen(): streamed = Runner.run_streamed(agent, input=[dict(item) for item in items], - **_run_kwargs(max_turns)) + **run_kwargs) sequence = 0 completed = False try: @@ -451,6 +512,15 @@ async def agen(): if event.type == "raw_response_event": data = event.data.model_dump(exclude_unset=True) if data.get("type") in lifecycle: + if data["type"] in ("response.completed", + "response.incomplete", + "response.failed"): + # Per-turn terminal state; the last turn's wins + # and feeds the final envelope below. + state = data.get("response") or {} + for field in ("status", "incomplete_details", + "error"): + recorded[field] = state.get(field) continue sequence += 1 data["sequence_number"] = sequence @@ -467,13 +537,20 @@ async def agen(): completed = True except MaxTurnsExceeded as exc: raise _wrap_max_turns(exc, max_turns) from exc + except AgentsException as exc: + raise PageIndexAPIError( + f"The agent 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) output = streamed.to_input_list()[len(items):] sequence += 1 - yield {"type": "response.completed", "sequence_number": sequence, + 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(output, streamed.raw_responses)} return _stream_sync(agen) @@ -491,10 +568,11 @@ def _require_anthropic() -> None: ) 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.68.0 (the tool " - "runner) โ€” pip install -U anthropic." + "messages in local mode requires anthropic >= 0.84.0 (the tool " + "runner with ToolError) โ€” pip install -U anthropic." ) from exc @@ -556,7 +634,18 @@ def _anthropic_usage(turns, final_usage: dict) -> dict: return totals -def run_messages(client, messages, model: str, max_tokens: int = 8192, +_CLAUDE_4096_MODELS = ("claude-3-opus", "claude-3-sonnet", "claude-3-haiku", + "claude-3-5-sonnet-20240620") + + +def _default_max_tokens(model: str) -> int: + """The wire-required per-turn budget when the caller sets none: 8192, + except the claude-3 generation whose output ceiling is 4096.""" + return 4096 if model.startswith(_CLAUDE_4096_MODELS) else 8192 + + +def run_messages(client, messages, model: str, + max_tokens: Optional[int] = None, stream: bool = False, doc_id=None, system=None, temperature: Optional[float] = None, top_p: Optional[float] = None, @@ -581,10 +670,11 @@ def run_messages(client, messages, model: str, max_tokens: int = 8192, "stop_sequences": stop_sequences, }.items() if value is not None} runner = _anthropic_client().beta.messages.tool_runner( - max_tokens=max_tokens, + max_tokens=(max_tokens if max_tokens is not None + else _default_max_tokens(model)), messages=prepared, model=model, - tools=build_anthropic_tools(client), + tools=build_anthropic_tools(client, doc_ids=doc_id or None), system=_anthropic_system(system, block), stream=stream, # Bounded like the OpenAI surfaces (their framework default is 10). diff --git a/pageindex/mcp_bridge.py b/pageindex/mcp_bridge.py index 95aba5c70..f6ba5809f 100644 --- a/pageindex/mcp_bridge.py +++ b/pageindex/mcp_bridge.py @@ -176,12 +176,16 @@ def list_tools(self) -> list[dict]: if not cursor: return tools - def call_tool(self, name: str, arguments: dict[str, Any]) -> str: + 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")) blocks = result.get("content") or [] texts = [block.get("text", "") for block in blocks if isinstance(block, dict) and block.get("type") == "text"] if len(texts) == len(blocks): - return "\n".join(texts) - return json.dumps(blocks, ensure_ascii=False) + return "\n".join(texts), is_error + return json.dumps(blocks, ensure_ascii=False), is_error diff --git a/pyproject.toml b/pyproject.toml index e55de0287..a84b083e3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,9 +43,11 @@ claude-agent-sdk = { version = ">=0.1.0", optional = true } # a blocking bridge call would freeze the agent event loop. openai-agents = { version = ">=0.8.0", optional = true } # messages() and as_anthropic_tools() need the SDK's beta tool runner; -# 0.68.0 is the first release with tool_runner(stream/system/max_iterations) -# and beta_tool(input_schema). -anthropic = { version = ">=0.68.0", optional = true } +# 0.84.0 is the first release with ToolError (failed tool calls flagged +# is_error) whose runner also executes the final turn's tools on a +# max_iterations cut (0.75.0 ordering) โ€” older runners return truncated +# histories with no tool_result. +anthropic = { version = ">=0.84.0", optional = true } [tool.poetry.extras] claude = ["claude-agent-sdk"] diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 39a73e3dc..221f26c88 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -5,6 +5,8 @@ import os import re import sys +import time +import types from pathlib import Path import pytest @@ -392,10 +394,55 @@ def test_remove_document(client, store_path): 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_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 + + +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): @@ -415,6 +462,26 @@ def test_unknown_argument_becomes_error_envelope(client, store_path): 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), indent=2, ensure_ascii=False) + + # โ”€โ”€ framework adapters โ”€โ”€ def test_as_openai_tools_missing_dependency(client, monkeypatch): @@ -460,6 +527,67 @@ def test_as_openai_tools_local_ignores_hosted(client): == 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_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") @@ -485,6 +613,66 @@ def test_as_claude_mcp_local_when_installed(client): assert server.get("type") != "http" +def test_claude_allowed_tools_reads_keys_from_registration(client, monkeypatch): + """The registration key is data in the caller's mcp_servers map โ€” the + gate entries derive from it, it is never spelled a second time. Needs + no framework installed.""" + monkeypatch.setitem(sys.modules, "claude_agent_sdk", None) + servers = { + "docs": {"type": "sdk", "name": "pageindex", "instance": object()}, + "other": {"type": "http", "url": "https://example.com/mcp"}, + } + assert (client.claude_allowed_tools(servers) + == [f"mcp__docs__{name}" for name in tool_names()]) + managed = client.claude_allowed_tools(servers, include_management=True) + assert "mcp__docs__remove_document" in managed + with pytest.raises(PageIndexAPIError, match="No PageIndex server"): + client.claude_allowed_tools( + {"other": {"type": "http", "url": "https://example.com/mcp"}}) + with pytest.raises(PageIndexAPIError, match="mcp_servers dict"): + client.claude_allowed_tools("path/to/.mcp.json") + + +def test_claude_allowed_tools_recognizes_real_local_server(client): + pytest.importorskip("claude_agent_sdk") + servers = {"pi": client.as_claude_mcp()} + assert (client.claude_allowed_tools(servers) + == [f"mcp__pi__{name}" for name in tool_names()]) + + +def test_claude_allowed_tools_cloud_from_registration(cloud_with_fake_bridge): + cloud, _ = cloud_with_fake_bridge + servers = {"docs": cloud.as_claude_mcp(), + "other": {"type": "http", "url": "https://example.com/mcp"}} + assert cloud.claude_allowed_tools(servers) == [ + "mcp__docs__search_documents", "mcp__docs__get_document"] + assert ("mcp__docs__remove_document" + in cloud.claude_allowed_tools(servers, include_management=True)) + + +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" + assert config["mcp_servers"]["pageindex"]["type"] == "http" + assert config["allowed_tools"] == cloud.claude_allowed_tools( + config["mcp_servers"]) + renamed = cloud.claude_agent_config(server_name="docs", + include_management=True) + assert set(renamed["mcp_servers"]) == {"docs"} + assert "mcp__docs__remove_document" in renamed["allowed_tools"] + + +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"] + == [f"mcp__pageindex__{name}" for name in tool_names()]) + + def test_as_anthropic_tools_missing_dependency(client, monkeypatch): monkeypatch.setitem(sys.modules, "anthropic", None) with pytest.raises(PageIndexAPIError, match="anthropic"): @@ -526,6 +714,34 @@ def test_as_anthropic_tools_local_management_opt_in(client): 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 @@ -564,7 +780,11 @@ def test_as_anthropic_tools_cloud_management_opt_in(cloud_with_fake_bridge): 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() @@ -572,7 +792,9 @@ def boom(name, arguments): raise RuntimeError("bridge down") created["bridge"].call_tool = boom - payload = json.loads(tools[0].call({"query": "q"})) + 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"] @@ -643,9 +865,13 @@ def __init__(self, url, headers): 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}) + return json.dumps({"success": True, "tool": name, + "args": arguments}), False @pytest.fixture @@ -733,6 +959,7 @@ def list_tools(self): def test_mcp_bridge_protocol(monkeypatch): + import requests as requests_mod from pageindex.mcp_bridge import McpBridge import pageindex.mcp_bridge as mcp_bridge @@ -785,7 +1012,10 @@ def fake_post(url, json=None, headers=None, timeout=None): {"type": "text", "text": "world"}]}}) raise AssertionError(f"unexpected method {method}") - monkeypatch.setattr(mcp_bridge.requests, "post", fake_post) + # Replace the module's own `requests` binding โ€” patching the shared + # requests module would leak the fake process-wide. + monkeypatch.setattr(mcp_bridge, "requests", types.SimpleNamespace( + post=fake_post, RequestException=requests_mod.RequestException)) bridge = McpBridge("https://api.pageindex.ai/mcp", {"Authorization": "Bearer k"}) @@ -801,8 +1031,8 @@ def fake_post(url, json=None, headers=None, timeout=None): assert list_headers["Authorization"] == "Bearer k" # First tools/call 404s (expired session) โ†’ re-initialize โ†’ retry succeeds. - text = bridge.call_tool("t1", {"a": 1}) - assert text == "hello\nworld" + 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 @@ -822,7 +1052,7 @@ def test_synth_optional_no_default_param_is_nullable(): class _Bridge: def call_tool(self, name, args): - return json.dumps(args) + return json.dumps(args), False meta = {"name": "browse_documents", "description": "d", @@ -839,7 +1069,7 @@ def test_synth_escape_hatches(): class _Bridge: def call_tool(self, name, args): calls.append((name, args)) - return "ok" + return "ok", False # Tool named "_invoke" must not recurse into itself. invoke_named = _make_bridge_function(_Bridge(), { @@ -898,6 +1128,41 @@ def list_tools(self): assert len(cloud.agent_tools(include_management=True)) == 1 +def test_bridge_call_tool_surfaces_iserror(monkeypatch): + import requests as requests_mod + import pageindex.mcp_bridge as mcp_bridge + from pageindex.mcp_bridge import McpBridge + + class _Resp: + def __init__(self, status, body=None): + self.status_code = status + self._body = body + self.headers = {"Content-Type": "application/json"} + self.text = json.dumps(body) if body else "" + self.content = self.text.encode("utf-8") + + def json(self): + if self._body is None: + raise ValueError("no body") + return self._body + + def fake_post(url, json=None, headers=None, timeout=None): + method = json.get("method") + rid = json.get("id") + if method == "initialize": + return _Resp(200, {"jsonrpc": "2.0", "id": rid, "result": {}}) + if method == "notifications/initialized": + return _Resp(202) + return _Resp(200, {"jsonrpc": "2.0", "id": rid, "result": { + "isError": True, + "content": [{"type": "text", "text": '{"error": "denied"}'}]}}) + + monkeypatch.setattr(mcp_bridge, "requests", types.SimpleNamespace( + post=fake_post, RequestException=requests_mod.RequestException)) + bridge = McpBridge("https://api.pageindex.ai/mcp", {}) + assert bridge.call_tool("t", {}) == ('{"error": "denied"}', True) + + def test_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' @@ -915,7 +1180,8 @@ def test_bridge_transport_error_is_pageindex_error(monkeypatch): def dead_post(*args, **kwargs): raise requests_mod.ConnectionError("dns down") - monkeypatch.setattr(mcp_bridge.requests, "post", dead_post) + monkeypatch.setattr(mcp_bridge, "requests", types.SimpleNamespace( + post=dead_post, RequestException=requests_mod.RequestException)) bridge = McpBridge("https://api.pageindex.ai/mcp", {}) with pytest.raises(PageIndexAPIError, match="Could not reach"): bridge.list_tools() @@ -925,7 +1191,8 @@ 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, "sleep", lambda seconds: None) + monkeypatch.setattr(agent_tools_mod, "time", types.SimpleNamespace( + monotonic=time.monotonic, sleep=lambda seconds: None)) class _Client: def get_document(self, doc_id): @@ -1070,17 +1337,18 @@ def test_live_cloud_envelope_field_parity(tmp_path): 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})) + 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})), + "get_document", {"doc_name": doc_name})[0]), "get_document_structure": json.loads(bridge.call_tool( - "get_document_structure", {"doc_name": doc_name})), + "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"})), + "get_page_content", {"doc_name": doc_name, "pages": "1"})[0]), } store = str(tmp_path / "store") @@ -1270,7 +1538,8 @@ def get_document(self, doc_id): @pytest.fixture def fake_cloud_client(tmp_path, monkeypatch): - monkeypatch.setattr(client_module.time, "sleep", lambda seconds: None) + 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")) @@ -1299,7 +1568,8 @@ def fake_monotonic(): clock["now"] += 700.0 return clock["now"] - monkeypatch.setattr(client_module.time, "monotonic", fake_monotonic) + 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) diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index 98a2c0c9e..1517e93a8 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -1,8 +1,9 @@ """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 -from pathlib import Path +import types import pytest @@ -12,8 +13,6 @@ from pageindex.local_chat import CHAT_HEADER from pageindex.local_store import DocStore -sys.path.insert(0, str(Path(__file__).parent.parent)) - def seed_doc(storage_path, doc_id, name): pages = [{"page_index": 1, "markdown": "Page one text about apples"}] @@ -102,6 +101,11 @@ async def get_response(self, system_instructions, input, model_settings, **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) @@ -130,6 +134,8 @@ async def stream_response(self, system_instructions, input, 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, @@ -598,10 +604,23 @@ def test_responses_envelope_fields_and_cache_group(client, store_path, assert result["instructions"].startswith(CHAT_HEADER) assert result["parallel_tool_calls"] is True assert result["tool_choice"] == "auto" - # Stable cache group: without it openai-agents stamps each run with a - # fresh prompt_cache_key, defeating round-trip cache routing. - assert (local_chat._run_kwargs(None)["run_config"].group_id - == "pageindex-local-chat") + + +def test_conversation_group_id_stable_per_conversation(): + """Cache-routing key: openai-agents hashes group_id into the OpenAI + prompt_cache_key. A conversation's continuations must share one key + (same model/instructions/first item), and unrelated conversations must + not pool under it.""" + turn1 = [{"role": "user", "content": "q"}] + continuation = turn1 + [{"role": "assistant", "content": "a"}, + {"role": "user", "content": "and?"}] + key = local_chat._conversation_group_id("m", "sys", turn1) + assert key == local_chat._conversation_group_id("m", "sys", continuation) + assert key != local_chat._conversation_group_id( + "m", "sys", [{"role": "user", "content": "other"}]) + assert key != local_chat._conversation_group_id("m2", "sys", turn1) + assert key != local_chat._conversation_group_id("m", "sys2", turn1) + assert (local_chat._run_kwargs(None, key)["run_config"].group_id == key) @needs_agents @@ -612,6 +631,149 @@ def test_responses_input_validation(client, fake_model): 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_openai_model_resolves_provider_prefixes(): + """retrieve_model arrives normalized (litellm//); the + OpenAI SDK must never see that prefix as a wire model name.""" + pytest.importorskip("litellm") + from agents.extensions.models.litellm_model import LitellmModel + from agents.models.openai_chatcompletions import OpenAIChatCompletionsModel + from agents.models.openai_responses import OpenAIResponsesModel + + model = local_chat._openai_model("chat", "litellm/anthropic/claude-x") + assert isinstance(model, LitellmModel) and model.model == "anthropic/claude-x" + model = local_chat._openai_model("responses", "anthropic/claude-x") + assert isinstance(model, LitellmModel) and model.model == "anthropic/claude-x" + model = local_chat._openai_model("chat", "openai/gpt-5.2") + assert isinstance(model, OpenAIChatCompletionsModel) + assert str(model.model) == "gpt-5.2" + model = local_chat._openai_model("responses", "gpt-5.2") + assert isinstance(model, OpenAIResponsesModel) + assert str(model.model) == "gpt-5.2" + + +@needs_agents +def test_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_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)) + + +@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): @@ -639,6 +801,46 @@ def test_stream_abandonment_cancels_pending_turn(client, store_path, assert fake.deltas_emitted == 0 # turn 2 never produced output +@needs_anthropic +def test_messages_max_tokens_default_resolves_per_model(client, fake_anthropic): + """The wire-required budget must not exceed the model's ceiling: the + claude-3 generation caps output at 4096.""" + calls = fake_anthropic([ + _anthropic_message([{"type": "text", "text": "ok"}], "end_turn")]) + client.messages("q", model="claude-3-opus-20240229") + assert calls[0]["max_tokens"] == 4096 + calls = fake_anthropic([ + _anthropic_message([{"type": "text", "text": "ok"}], "end_turn")]) + client.messages("q", model="claude-sonnet-4-5") + assert calls[0]["max_tokens"] == 8192 + calls = fake_anthropic([ + _anthropic_message([{"type": "text", "text": "ok"}], "end_turn")]) + client.messages("q", model="claude-3-opus-20240229", max_tokens=1234) + assert calls[0]["max_tokens"] == 1234 + + +@needs_anthropic +def test_messages_tool_error_flagged_and_scoped(client, store_path, + fake_anthropic): + """Through the real runner: a failed call reaches Claude as a + tool_result with is_error true, and doc_id scoping makes out-of-scope + documents unreachable by name.""" + seed_doc(store_path, "pi-a", "report.pdf") + seed_doc(store_path, "pi-b", "secret.pdf") + calls = fake_anthropic([ + _anthropic_message([{"type": "tool_use", "id": "tu_1", + "name": "get_document", + "input": {"doc_name": "secret.pdf"}}], + "tool_use"), + _anthropic_message([{"type": "text", "text": "ok"}], "end_turn"), + ]) + client.messages("q", model="claude-test", doc_id="pi-a") + tool_result = calls[1]["messages"][-1]["content"][0] + assert tool_result["type"] == "tool_result" + assert tool_result.get("is_error") is True + assert "NOT_FOUND" in json.dumps(tool_result["content"]) + + @needs_anthropic def test_messages_envelope_json_and_no_internal_fields(client, store_path, fake_anthropic): diff --git a/tests/test_package_surface.py b/tests/test_package_surface.py index e10c3e8c5..d93e0de8c 100644 --- a/tests/test_package_surface.py +++ b/tests/test_package_surface.py @@ -60,3 +60,24 @@ 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_unknown_names_stay_lazy(): + """The 0.2.10 modules resolve as attributes, and an unknown name raises + AttributeError without dragging in the indexing stack.""" + probe = ( + "import sys, pageindex\n" + "pageindex.agent_tools; pageindex.local_chat\n" + "pageindex.mcp_bridge; pageindex.integrations\n" + "try:\n" + " pageindex.definitely_missing\n" + " raise SystemExit('no AttributeError')\n" + "except AttributeError:\n" + " pass\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" + ) + out = subprocess.run([sys.executable, "-c", probe], + capture_output=True, text=True, check=True) + assert out.stdout.strip() == "clean" From dbe585707383a037928b55266b98080b38bd5834 Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 12 Aug 2026 15:03:08 +0800 Subject: [PATCH 028/137] feat: one-call config bundles for every bring-your-own-framework surface MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit claude_agent_config() gets two symmetric siblings, so each framework's front door is a single splat over the same explicit primitives: - openai_agent_config(): Agent(**...) kwargs โ€” instructions, tools, and the local retrieve_model (cloud omits model for the framework default) - anthropic_runner_config(): tool_runner(**...) kwargs โ€” system, tools, and the messages() defaults (per-model max_tokens, 10-iteration bound); only the user's messages remain Bundles stay pure sugar: doc_id rides agent_instructions, no extra semantics over the explicit form, docstrings point both ways. The demo agent shrinks to Agent(**client.openai_agent_config(doc_id=...)). Construction is pinned against the real frameworks in tests (Agent and tool_runner both built offline), so an upstream kwargs rename fails loudly; 211 -> 215 tests. --- examples/agentic_vectorless_rag_demo.py | 7 +- pageindex/client.py | 94 ++++++++++++++++++++++++- tests/test_agent_tools.py | 60 ++++++++++++++++ 3 files changed, 152 insertions(+), 9 deletions(-) diff --git a/examples/agentic_vectorless_rag_demo.py b/examples/agentic_vectorless_rag_demo.py index 7a83776ff..ac0db360a 100644 --- a/examples/agentic_vectorless_rag_demo.py +++ b/examples/agentic_vectorless_rag_demo.py @@ -50,12 +50,7 @@ def query_agent(client: PageIndexLocalClient, doc_id: str, prompt: str, verbose: Tool calls are always printed; verbose=True also prints arguments and output previews. """ agent = Agent( - name="PageIndex", - instructions=client.agent_instructions(doc_id=doc_id), - tools=client.as_openai_tools(), - # retrieve_model is a local-mode attribute; cloud clients fall back - # to the framework's default model. - model=getattr(client, "retrieve_model", None), + **client.openai_agent_config(doc_id=doc_id), # model_settings=ModelSettings(reasoning={"effort": "low", "summary": "auto"}), # from agents.model_settings import ModelSettings ) diff --git a/pageindex/client.py b/pageindex/client.py index 868e37e2f..4fbb86c3e 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -607,7 +607,9 @@ def agent_tools(self, include_management: bool = False) -> list[Callable[..., st def as_openai_tools(self, include_management: bool = False, hosted: bool = False) -> list: """ - Tools for the OpenAI Agents SDK โ€” pass to ``Agent(tools=...)``. + 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, @@ -637,12 +639,49 @@ def as_openai_tools(self, include_management: bool = False, from .integrations.openai_agents import build_openai_tools return build_openai_tools(self, include_management, hosted) + def openai_agent_config( + self, + doc_id: Optional[Union[str, list[str]]] = None, + include_management: bool = False, + model: Optional[str] = None, + ) -> dict[str, Any]: + """ + Document QA ``Agent`` kwargs for the OpenAI Agents SDK in one + call:: + + agent = Agent(**client.openai_agent_config()) + + Sugar over the explicit form โ€” ``agent_instructions`` (with + ``doc_id`` targeting) as the instructions and + ``as_openai_tools`` as the tools; local clients also carry their + configured ``retrieve_model`` (cloud omits ``model`` so the + framework default applies). To customize further, switch to + those methods directly. + + Args: + doc_id: Document ID or list of IDs to target, as in + ``agent_instructions``. + include_management (bool): Also expose tools that modify the + library. + model: Backend model name; overrides the local default. + """ + config: dict[str, Any] = { + "name": "PageIndex", + "instructions": self.agent_instructions(doc_id=doc_id), + "tools": self.as_openai_tools(include_management), + } + model = model or getattr(self, "retrieve_model", None) + if model: + config["model"] = model + return config + def as_anthropic_tools(self, include_management: bool = False, asynchronous: bool = False) -> list: """ Runnable tools for the Anthropic SDK's tool runner โ€” pass to - ``client.beta.messages.tool_runner(tools=...)``. The default - flavor is for the sync ``Anthropic`` client; pass + ``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 ...]``. @@ -675,6 +714,55 @@ def as_anthropic_tools(self, include_management: bool = False, from .integrations.anthropic_sdk import build_anthropic_tools return build_anthropic_tools(self, include_management, asynchronous) + def anthropic_runner_config( + self, + model: str, + doc_id: Optional[Union[str, list[str]]] = None, + include_management: bool = False, + asynchronous: bool = False, + max_tokens: Optional[int] = None, + max_turns: Optional[int] = None, + ) -> dict[str, Any]: + """ + Document QA ``tool_runner`` kwargs for the Anthropic SDK in one + call โ€” only your ``messages`` remain:: + + runner = anthropic_client.beta.messages.tool_runner( + **client.anthropic_runner_config(model="claude-sonnet-4-5"), + messages=[{"role": "user", "content": "..."}], + ) + + Sugar over the explicit form โ€” ``agent_instructions`` (with + ``doc_id`` targeting) as the system prompt and + ``as_anthropic_tools`` as the tools โ€” plus the same defaults + ``messages()`` applies: a per-model ``max_tokens`` and a + ``max_iterations`` bound of 10. To customize further, switch to + those methods directly. + + Args: + model: Backend model name (also resolves the ``max_tokens`` + default). + doc_id: Document ID or list of IDs to target, as in + ``agent_instructions``. + include_management (bool): Also expose tools that modify the + library. + asynchronous (bool): Build async runnables for + ``AsyncAnthropic``. + max_tokens: Per-turn output budget; default resolved per + model. + max_turns: Agent-loop bound; default 10. + """ + from .local_chat import _default_max_tokens + return { + "model": model, + "max_tokens": (max_tokens if max_tokens is not None + else _default_max_tokens(model)), + "system": self.agent_instructions(doc_id=doc_id), + "tools": self.as_anthropic_tools(include_management, + asynchronous), + "max_iterations": max_turns if max_turns is not None else 10, + } + def as_claude_mcp(self, include_management: bool = False): """ ``mcp_servers`` entry for the Claude Agent SDK. diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 221f26c88..64e0994e4 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -673,6 +673,66 @@ def test_claude_agent_config_local(client, store_path): == [f"mcp__pageindex__{name}" for name in tool_names()]) +def test_openai_agent_config_local(client, store_path): + pytest.importorskip("agents") + from agents import Agent + seed_doc(store_path, "pi-a", "report.pdf") + config = client.openai_agent_config(doc_id="pi-a") + assert config["name"] == "PageIndex" + assert "report.pdf" in config["instructions"] + assert [tool.name for tool in config["tools"]] == list(tool_names()) + assert config["model"] == client.retrieve_model + assert client.openai_agent_config(model="gpt-x")["model"] == "gpt-x" + assert Agent(**client.openai_agent_config()).name == "PageIndex" + + +def test_openai_agent_config_cloud_omits_model(cloud_with_fake_bridge): + pytest.importorskip("agents") + cloud, _ = cloud_with_fake_bridge + config = cloud.openai_agent_config() + assert "model" not in config + assert config["instructions"] == "SERVER GUIDANCE" + assert [tool.name for tool in config["tools"]] == ["search_documents", + "get_document"] + + +def test_anthropic_runner_config_shapes(client, store_path): + pytest.importorskip("anthropic") + import anthropic + from anthropic.lib.tools import BetaAsyncFunctionTool + seed_doc(store_path, "pi-a", "report.pdf") + config = client.anthropic_runner_config(model="claude-3-opus-20240229", + doc_id="pi-a") + assert config["max_tokens"] == 4096 + assert config["max_iterations"] == 10 + assert "report.pdf" in config["system"] + assert [tool.name for tool in config["tools"]] == list(tool_names()) + assert (client.anthropic_runner_config(model="claude-sonnet-4-5") + ["max_tokens"] == 8192) + override = client.anthropic_runner_config(model="claude-sonnet-4-5", + max_tokens=99, max_turns=3) + assert override["max_tokens"] == 99 and override["max_iterations"] == 3 + async_tools = client.anthropic_runner_config( + model="claude-sonnet-4-5", asynchronous=True)["tools"] + assert all(isinstance(tool, BetaAsyncFunctionTool) + for tool in async_tools) + # The kwargs must construct a real runner (construction is offline โ€” + # requests start on iteration), pinning tool_runner's parameter names. + runner = anthropic.Anthropic(api_key="test").beta.messages.tool_runner( + **client.anthropic_runner_config(model="claude-sonnet-4-5"), + messages=[{"role": "user", "content": "q"}]) + assert runner is not None + + +def test_anthropic_runner_config_cloud(cloud_with_fake_bridge): + pytest.importorskip("anthropic") + cloud, _ = cloud_with_fake_bridge + config = cloud.anthropic_runner_config(model="claude-sonnet-4-5") + assert config["system"] == "SERVER GUIDANCE" + assert [tool.name for tool in config["tools"]] == ["search_documents", + "get_document"] + + def test_as_anthropic_tools_missing_dependency(client, monkeypatch): monkeypatch.setitem(sys.modules, "anthropic", None) with pytest.raises(PageIndexAPIError, match="anthropic"): From b135711918a4d571d6883eae83a08378c65242dc Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 12 Aug 2026 15:12:55 +0800 Subject: [PATCH 029/137] =?UTF-8?q?fix:=20three=20more=20review=20findings?= =?UTF-8?q?=20=E2=80=94=20partial-read=20reporting,=20reply=20correlation,?= =?UTF-8?q?=20output=5Findex=20axis?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - get_page_content: the summary is additive, not either/or โ€” a call that both truncates for size and has out-of-range pages reported only the latter, telling the agent every in-range page was returned (#2) - McpBridge._extract_result: strict request-id correlation only; the eager fallback could hand back a stale or mis-correlated JSON-RPC message as this call's reply (#16) - responses() streaming: output_index now addresses the logical response.output โ€” backend per-turn indexes are re-based past prior turns' items and the SDK-injected tool outputs take the next slot on that axis, instead of reusing the event-sequence counter (#15) 215 -> 217 tests. --- pageindex/agent_tools.py | 25 ++++++++++------- pageindex/client.py | 8 +++--- pageindex/local_chat.py | 13 ++++++++- pageindex/mcp_bridge.py | 10 ++++--- tests/test_agent_tools.py | 57 +++++++++++++++++++++++++++++++++++++++ tests/test_local_chat.py | 9 +++++++ 6 files changed, 104 insertions(+), 18 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index 02c006036..ff58d9913 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -1028,16 +1028,21 @@ def _get_page_content(client, doc_name: str, pages: str, if out_of_range: options.insert(0, f"Document has {max_page} pages total - request " f"pages 1-{max_page}") - summary = ( - f"Retrieved {len(included)} pages. Pages " - f"{', '.join(map(str, out_of_range))} were out of range." - if out_of_range - else f"Returned {len(included)} of {len(requested)} requested pages " - "due to response size limits." - if remaining - else f"Successfully retrieved content for {len(content)} " - f"page{'' if len(content) == 1 else 's'}." - ) + # Additive, not either/or: a call can both truncate for size and have + # out-of-range pages โ€” hiding either would misreport what was returned. + 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 {', '.join(map(str, 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, diff --git a/pageindex/client.py b/pageindex/client.py index 4fbb86c3e..89bdb0da4 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -442,9 +442,11 @@ def responses( model: Backend model name (defaults to ``retrieve_model``). stream: Yield Responses stream events as dicts โ€” one logical response per call: per-turn backend lifecycle events are - collapsed and sequence numbers reassigned monotonically; - tool outputs are emitted as ``response.output_item.done`` - events and the single final event is ``response.completed``. + collapsed, sequence numbers are reassigned monotonically, + and ``output_index`` is re-based onto the single logical + ``output``; tool outputs are emitted as + ``response.output_item.done`` events and the single final + event is the terminal ``response.*`` for the run's status. 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 diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 8dd01f64c..59cf1431f 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -506,6 +506,13 @@ async def agen(): 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 + # โ€” and the tool outputs the SDK injects between turns take the + # next slot on that same axis. + output_offset = 0 completed = False try: async for event in streamed.stream_events(): @@ -521,7 +528,10 @@ async def agen(): for field in ("status", "incomplete_details", "error"): recorded[field] = state.get(field) + output_offset += len(state.get("output") or []) continue + if isinstance(data.get("output_index"), int): + data["output_index"] += output_offset sequence += 1 data["sequence_number"] = sequence yield data @@ -531,9 +541,10 @@ async def agen(): # the way the platform streams its own server-side tools. sequence += 1 yield {"type": "response.output_item.done", - "output_index": sequence, + "output_index": output_offset, "sequence_number": sequence, "item": dict(event.item.to_input_item())} + output_offset += 1 completed = True except MaxTurnsExceeded as exc: raise _wrap_max_turns(exc, max_turns) from exc diff --git a/pageindex/mcp_bridge.py b/pageindex/mcp_bridge.py index f6ba5809f..f23575e20 100644 --- a/pageindex/mcp_bridge.py +++ b/pageindex/mcp_bridge.py @@ -88,11 +88,13 @@ def _extract_result(self, response: requests.Response, request_id: int) -> Any: f"MCP server returned a non-JSON response " f"(HTTP {response.status_code})." ) from exc - reply = next((m for m in messages if m.get("id") == request_id), - next((m for m in messages - if "result" in m or "error" in m), None)) + # 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.") + raise PageIndexAPIError( + "MCP server response contained no reply matching the request." + ) if "error" in reply: error = reply["error"] or {} raise PageIndexAPIError( diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 64e0994e4..fb7aa593a 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -376,10 +376,29 @@ def test_page_content_char_budget(client, store_path): pages="1-2") assert not is_error assert payload["returned_pages"] == "1" + assert "size limits" in payload["next_steps"]["summary"] assert any("For remaining pages, request: 2" in option for option in payload["next_steps"]["options"]) +def test_page_content_reports_truncation_and_out_of_range_together( + client, store_path): + """Size truncation must not hide behind the out-of-range report (or + vice versa) โ€” the agent otherwise believes it holds every in-range + page.""" + pages = [ + {"page_index": 1, "markdown": "x" * 96_000}, + {"page_index": 2, "markdown": "short"}, + ] + seed_doc(store_path, "pi-a", "huge.pdf", pages=pages) + payload, is_error = run(client, "get_page_content", doc_name="huge.pdf", + pages="1-2,99") + assert not is_error + assert payload["returned_pages"] == "1" + summary = payload["next_steps"]["summary"] + assert "size limits" in summary and "out of range" in summary + + # โ”€โ”€ remove_document (management-gated) โ”€โ”€ def test_remove_document(client, store_path): @@ -1223,6 +1242,44 @@ def fake_post(url, json=None, headers=None, timeout=None): assert bridge.call_tool("t", {}) == ('{"error": "denied"}', True) +def test_bridge_rejects_mismatched_reply_id(monkeypatch): + """A result-bearing message with the wrong id must not be returned as + this call's reply.""" + import requests as requests_mod + import pageindex.mcp_bridge as mcp_bridge + from pageindex.mcp_bridge import McpBridge + + class _Resp: + def __init__(self, status, body=None): + self.status_code = status + self._body = body + self.headers = {"Content-Type": "application/json"} + self.text = json.dumps(body) if body else "" + self.content = self.text.encode("utf-8") + + def json(self): + if self._body is None: + raise ValueError("no body") + return self._body + + def fake_post(url, json=None, headers=None, timeout=None): + method = json.get("method") + rid = json.get("id") + if method == "initialize": + return _Resp(200, {"jsonrpc": "2.0", "id": rid, "result": {}}) + if method == "notifications/initialized": + return _Resp(202) + return _Resp(200, {"jsonrpc": "2.0", "id": rid - 1, # stale reply + "result": {"content": [{"type": "text", + "text": "old"}]}}) + + monkeypatch.setattr(mcp_bridge, "requests", types.SimpleNamespace( + post=fake_post, RequestException=requests_mod.RequestException)) + bridge = McpBridge("https://api.pageindex.ai/mcp", {}) + with pytest.raises(PageIndexAPIError, match="no reply matching"): + bridge.call_tool("t", {}) + + def test_sse_crlf_multi_message(): from pageindex.mcp_bridge import _parse_sse body = ('event: message\r\ndata: {"jsonrpc":"2.0","method":"notifications/progress"}\r\n\r\n' diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index 1517e93a8..adccc62b6 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -358,6 +358,15 @@ def test_responses_stream_passthrough(client, store_path, fake_model): final = events[-1]["response"] assert final["status"] == "completed" assert final["usage"]["total_tokens"] == 30 + # output_index addresses the logical response.output: the tool output + # slots in after turn 1's item, and turn 2's deltas are re-based past + # both instead of restarting at 0. + assert (final["output"][tool_events[0]["output_index"]]["type"] + == "function_call_output") + 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") # โ”€โ”€ messages (Anthropic engine) โ”€โ”€ From eb1a2301a5606d7f70b5dda7b7be58d295d434bb Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 12 Aug 2026 15:52:50 +0800 Subject: [PATCH 030/137] feat: gate the config-handoff surfaces by the read-only MCP endpoint MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit pageindex-chat#448 adds /mcp?tools=read โ€” the server registers only readOnlyHint-annotated tools โ€” so the URL itself becomes the gate for every surface that hands a config to a third party: - as_claude_mcp: include_management now picks the endpoint on cloud; the parameter is real in both modes - as_openai_tools(hosted=True): OpenAI connects to the read-only endpoint by default and require_approval simplifies to "never" โ€” the approval-flow middle ground becomes hard absence, matching every other surface's default - claude_allowed_tools() retired before ever shipping: with the server gated, allowed_tools degenerates to whole-server pre-approval, which claude_agent_config emits as the constant ["mcp__"] โ€” no setup-time bridge round-trip remains - in-process surfaces (agent_tools, as_openai_tools, as_anthropic_tools over the bridge) keep bare /mcp + client-side annotation filtering: they materialize tools locally and hand no URL to anyone Release ordering: 0.2.10 must ship after pageindex-chat#448 deploys โ€” an older server ignores unknown query params and would silently serve the full set behind a URL that promises read-only. --- pageindex/client.py | 81 ++++++++-------------- pageindex/integrations/claude_agent_sdk.py | 46 +++--------- pageindex/integrations/openai_agents.py | 27 ++++---- tests/test_agent_tools.py | 70 ++++++------------- 4 files changed, 69 insertions(+), 155 deletions(-) diff --git a/pageindex/client.py b/pageindex/client.py index 89bdb0da4..3bf7c2e3e 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -631,10 +631,10 @@ def as_openai_tools(self, include_management: bool = False, Args: include_management (bool): Also expose tools that modify the - library (delete, upload). Default off: the cloud default - serves only server-annotated read-only tools, and - ``hosted=True`` routes non-read-only tools through the - Responses API approval flow instead. + 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). """ @@ -694,9 +694,10 @@ def as_anthropic_tools(self, include_management: bool = False, through verbatim (MCP and the Messages API share the schema shape). The server-side alternative is the Messages API's beta MCP connector โ€” ``mcp_servers=[{"type": "url", "name": - "pageindex", "url": f"{BASE_URL}/mcp", "authorization_token": - }]`` โ€” with no client-side tools - involved. Local: the in-process tools โ€” the same set + "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.84.0`` @@ -769,13 +770,13 @@ def as_claude_mcp(self, include_management: bool = False): """ ``mcp_servers`` entry for the Claude Agent SDK. - Cloud: returns the remote PageIndex MCP config โ€” the framework - connects to api.pageindex.ai/mcp directly and discovers the full - cloud tool set. A remote server cannot be filtered client-side, - so ``include_management`` has no effect there โ€” the gate is - ``allowed_tools``, built from your registration map by - ``claude_allowed_tools()``. Local: returns an in-process SDK MCP - server exposing the agent tools (requires ``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]'``). Cloud hosts that surface MCP server instructions receive the same @@ -787,43 +788,16 @@ def as_claude_mcp(self, include_management: bool = False): Usage (or ``claude_agent_config()`` for all three slots in one call):: - servers = {"pageindex": client.as_claude_mcp()} options = ClaudeAgentOptions( system_prompt=client.agent_instructions(), - mcp_servers=servers, - allowed_tools=client.claude_allowed_tools(servers), + 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) - def claude_allowed_tools(self, mcp_servers: dict[str, Any], - include_management: bool = False) -> list[str]: - """ - ``allowed_tools`` entries for the PageIndex servers in your - ``mcp_servers`` map โ€” pass the same dict you hand to - ``ClaudeAgentOptions``. The framework bakes the registration key - into every tool id (``mcp____``), so the keys are read - from the map rather than spelled a second time, and the tool - names are the read-only gate every other adapter applies โ€” live - server annotations on cloud, the tool contract locally. Nothing - is hand-maintained, and no framework install is needed. - - Raises PageIndexAPIError when the map holds no PageIndex entry โ€” - a gate list that silently matched nothing would disable every - tool. - - Args: - mcp_servers: The registration map; non-PageIndex entries are - ignored. - include_management (bool): Also allow tools that modify the - library (``remove_document``, and on cloud the server's - full management list). - """ - from .integrations.claude_agent_sdk import build_claude_allowed_tools - return build_claude_allowed_tools(self, mcp_servers, - include_management) - def claude_agent_config( self, doc_id: Optional[Union[str, list[str]]] = None, @@ -836,12 +810,11 @@ def claude_agent_config( options = ClaudeAgentOptions(**client.claude_agent_config()) Sugar over the explicit form โ€” the managed system prompt - (``agent_instructions``), the server entry (``as_claude_mcp``), - and the matching ``allowed_tools`` gate - (``claude_allowed_tools``), with one ``include_management`` and - ``server_name`` applied everywhere. To customize (your own - system prompt, extra servers), switch to those three methods - directly. + (``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 @@ -850,12 +823,12 @@ def claude_agent_config( library. server_name (str): Key the server is registered under. """ - servers = {server_name: self.as_claude_mcp(include_management)} return { "system_prompt": self.agent_instructions(doc_id=doc_id), - "mcp_servers": servers, - "allowed_tools": self.claude_allowed_tools(servers, - include_management), + "mcp_servers": {server_name: self.as_claude_mcp(include_management)}, + # Pre-approval only โ€” the server itself is already gated (the + # read-only endpoint on cloud, the registered set locally). + "allowed_tools": [f"mcp__{server_name}"], } def agent_instructions(self, doc_id: Optional[Union[str, list[str]]] = None) -> str: diff --git a/pageindex/integrations/claude_agent_sdk.py b/pageindex/integrations/claude_agent_sdk.py index a28cd155f..58b9da4e6 100644 --- a/pageindex/integrations/claude_agent_sdk.py +++ b/pageindex/integrations/claude_agent_sdk.py @@ -1,8 +1,9 @@ """Claude Agent SDK adapter: one value for the mcp_servers slot. -Cloud clients get the remote PageIndex MCP config (the framework connects -directly and discovers the full cloud tool set); local clients get an -in-process SDK MCP server over the same tool contract. +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 @@ -13,45 +14,14 @@ from ..errors import PageIndexAPIError -def build_claude_allowed_tools(client, mcp_servers, - include_management: bool = False) -> list[str]: - """``allowed_tools`` entries for the PageIndex entries of an - mcp_servers map. The framework scopes every tool id by the map key - (``mcp____``), so the keys are read from the map instead of - being spelled a second time; tool names are the gated set โ€” live - server annotations on cloud, the contract locally. Needs no framework - import.""" - from ..agent_tools import _tool_specs - if not isinstance(mcp_servers, dict) or not mcp_servers: - raise PageIndexAPIError( - "claude_allowed_tools takes the mcp_servers dict you register " - "with the framework (the {name: server} map)." - ) - - def is_pageindex(value) -> bool: - get = (value.get if isinstance(value, dict) - else lambda key, default=None: getattr(value, key, default)) - url = get("url") - if isinstance(url, str): - return url.startswith(f"{client.BASE_URL}/mcp") - return get("type") == "sdk" and get("name") == "pageindex" - - keys = [key for key, value in mcp_servers.items() if is_pageindex(value)] - if not keys: - raise PageIndexAPIError( - "No PageIndex server found in mcp_servers โ€” register " - "client.as_claude_mcp() under a key first (an allowed_tools " - "list built from this map would match nothing)." - ) - names = [spec[0] for spec in _tool_specs(client, include_management)] - return [f"mcp__{key}__{name}" for key in keys for name in names] - - def build_claude_mcp(client, include_management: bool = False): 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", + "url": f"{client.BASE_URL}/mcp{suffix}", "headers": {"Authorization": f"Bearer {client.api_key}"}, } diff --git a/pageindex/integrations/openai_agents.py b/pageindex/integrations/openai_agents.py index d6cf52948..9f5df5064 100644 --- a/pageindex/integrations/openai_agents.py +++ b/pageindex/integrations/openai_agents.py @@ -1,13 +1,13 @@ """OpenAI Agents SDK adapter for the Agent(tools=...) slot. -Cloud clients default to the full live tool set as plain FunctionTools via +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). -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. +(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 @@ -28,17 +28,16 @@ def build_openai_tools(client, include_management: bool = False, "pip install openai-agents (or pip install 'pageindex[openai]')." ) from exc if getattr(client, "api_key", None) and hosted: - # Same gate as the in-process path, enforced by OpenAI: tools the - # server annotates read-only run freely, everything else goes - # through the Responses API approval flow. - require_approval = ("never" if include_management - else {"never": {"read_only": True}}) + # 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", + "server_url": f"{client.BASE_URL}/mcp{suffix}", "headers": {"Authorization": f"Bearer {client.api_key}"}, - "require_approval": require_approval, + "require_approval": "never", })] from ..agent_tools import _tool_specs diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index fb7aa593a..f9e8d1c70 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -534,7 +534,7 @@ def test_as_openai_tools_cloud_hosted_opt_in(): assert len(tools) == 1 assert isinstance(tools[0], HostedMCPTool) config = tools[0].tool_config - assert config["server_url"] == "https://api.pageindex.ai/mcp" + assert config["server_url"] == "https://api.pageindex.ai/mcp?tools=read" assert config["headers"] == {"Authorization": "Bearer pi-test-key"} assert config["server_label"] == "pageindex" @@ -610,12 +610,15 @@ def call_tool(self, name, arguments): 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") - config = cloud.as_claude_mcp() - assert config == { + # 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", + "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): @@ -632,55 +635,21 @@ def test_as_claude_mcp_local_when_installed(client): assert server.get("type") != "http" -def test_claude_allowed_tools_reads_keys_from_registration(client, monkeypatch): - """The registration key is data in the caller's mcp_servers map โ€” the - gate entries derive from it, it is never spelled a second time. Needs - no framework installed.""" - monkeypatch.setitem(sys.modules, "claude_agent_sdk", None) - servers = { - "docs": {"type": "sdk", "name": "pageindex", "instance": object()}, - "other": {"type": "http", "url": "https://example.com/mcp"}, - } - assert (client.claude_allowed_tools(servers) - == [f"mcp__docs__{name}" for name in tool_names()]) - managed = client.claude_allowed_tools(servers, include_management=True) - assert "mcp__docs__remove_document" in managed - with pytest.raises(PageIndexAPIError, match="No PageIndex server"): - client.claude_allowed_tools( - {"other": {"type": "http", "url": "https://example.com/mcp"}}) - with pytest.raises(PageIndexAPIError, match="mcp_servers dict"): - client.claude_allowed_tools("path/to/.mcp.json") - - -def test_claude_allowed_tools_recognizes_real_local_server(client): - pytest.importorskip("claude_agent_sdk") - servers = {"pi": client.as_claude_mcp()} - assert (client.claude_allowed_tools(servers) - == [f"mcp__pi__{name}" for name in tool_names()]) - - -def test_claude_allowed_tools_cloud_from_registration(cloud_with_fake_bridge): - cloud, _ = cloud_with_fake_bridge - servers = {"docs": cloud.as_claude_mcp(), - "other": {"type": "http", "url": "https://example.com/mcp"}} - assert cloud.claude_allowed_tools(servers) == [ - "mcp__docs__search_documents", "mcp__docs__get_document"] - assert ("mcp__docs__remove_document" - in cloud.claude_allowed_tools(servers, include_management=True)) - - 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" - assert config["mcp_servers"]["pageindex"]["type"] == "http" - assert config["allowed_tools"] == cloud.claude_allowed_tools( - config["mcp_servers"]) + 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 "mcp__docs__remove_document" in renamed["allowed_tools"] + 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): @@ -688,8 +657,7 @@ def test_claude_agent_config_local(client, store_path): 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"] - == [f"mcp__pageindex__{name}" for name in tool_names()]) + assert config["allowed_tools"] == ["mcp__pageindex"] def test_openai_agent_config_local(client, store_path): @@ -1392,13 +1360,17 @@ def test_failed_document_status_message(client, store_path): for option in payload["next_steps"]["options"]) -def test_hosted_approval_gate(): +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["require_approval"] == {"never": {"read_only": True}} + 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" From 5f35a33eb94084434bf9071b1a7c565865aa9042 Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 12 Aug 2026 16:30:40 +0800 Subject: [PATCH 031/137] fix: chat_completions wraps framework exceptions like responses() The AgentsException -> PageIndexAPIError wrap from the responses() fix covered only that surface; a backend stream dying without a terminal event (or any engine failure) still escaped chat_completions as a raw openai-agents exception type on both its paths. --- pageindex/local_chat.py | 8 +++++++- tests/test_local_chat.py | 22 ++++++++++++++++++++++ 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 59cf1431f..787208e3e 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -353,13 +353,16 @@ def run_chat_completions(client, messages, stream: bool = False, run_kwargs = _run_kwargs(max_turns, _conversation_group_id(model_name, managed, items)) from agents import Runner - from agents.exceptions import MaxTurnsExceeded + 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(exc, max_turns) from exc + except AgentsException as exc: + raise PageIndexAPIError( + f"The agent backend failed: {exc}") from exc return { "id": f"chatcmpl-{uuid.uuid4().hex}", "object": "chat.completion", @@ -400,6 +403,9 @@ async def agen(): completed = True except MaxTurnsExceeded as exc: raise _wrap_max_turns(exc, max_turns) from exc + except AgentsException as exc: + raise PageIndexAPIError( + f"The agent backend failed: {exc}") from exc finally: if not completed and hasattr(streamed, "cancel"): streamed.cancel() # abandoned/failed: stop the agent task diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index adccc62b6..ee1630650 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -733,6 +733,28 @@ async def create(*args, **kwargs): assert result["error"] is None +@needs_agents +def test_chat_completions_wraps_framework_errors(client, store_path, + fake_model, monkeypatch): + """Both chat_completions paths surface engine failures as the SDK's + own error type, like responses().""" + from agents.exceptions import ModelBehaviorError + seed_doc(store_path, "pi-a", "report.pdf") + fake = fake_model([[_msg_item("never terminal")]]) + fake.no_terminal = True + with pytest.raises(PageIndexAPIError, match="agent backend failed"): + list(client.chat_completions("q", stream=True)) + + fake = fake_model([[_msg_item("x")]]) + + async def boom(*args, **kwargs): + raise ModelBehaviorError("backend broke") + + monkeypatch.setattr(fake, "get_response", boom) + with pytest.raises(PageIndexAPIError, match="agent backend failed"): + client.chat_completions("q") + + @needs_agents def test_responses_stream_wraps_framework_errors(client, store_path, fake_model): From 9d16dbfebdd965fab2faa00cf576d719d2d24026 Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 12 Aug 2026 16:36:31 +0800 Subject: [PATCH 032/137] fix: same-name documents in different folders no longer refuse doc_id targeting MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The shadow check in doc_targeting_block compared names across the whole library, so agent_instructions(doc_id=...) hard-raised for a legal cloud layout โ€” one file name in two folders โ€” with advice (rename/remove) that contradicts the contract, whose folder_id parameter exists precisely to disambiguate this case. Shadowing is now judged per folder: only a newer same-name document in the SAME folder makes the name unreachable and raises. A same-name document in another folder serves the call, and the targeting block adds a directive to pass folder_id on every tool call โ€” dropping the raise alone would have traded a loud refusal for the agent silently reading the newer document. Local mode (folderId always None) and the scoped chat path (allowlist resolution, fixed with the doc_id enforcement) are behaviorally unchanged. --- pageindex/agent_tools.py | 50 +++++++++++++++++++++++++++------------ tests/test_agent_tools.py | 19 +++++++++++++-- 2 files changed, 52 insertions(+), 17 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index ff58d9913..737f20cda 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -1549,10 +1549,13 @@ def doc_targeting_block(client, doc_id, scoped: bool = False) -> Optional[str]: 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`` (the chat surfaces, whose tools resolve names inside the - doc_id allowlist) only a same-name duplicate within the targeted set - shadows.""" + same-name document in the same folder โ€” the name-addressed tools could + not reach it. A same-name document in ANOTHER folder is the cloud + contract's supported case: no refusal, the block instead directs the + agent to pass folder_id (the documented disambiguator) on every call. + With ``scoped`` (the chat 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) @@ -1562,32 +1565,49 @@ def doc_targeting_block(client, doc_id, scoped: bool = False) -> Optional[str]: documents = ([{**detail, "id": one_id} for one_id, detail in zip(doc_ids, details)] if scoped else _all_documents(client)) + folder_notes: list[str] = [] for one_id, detail in zip(doc_ids, details): - entry, _ = _resolve_document(client, str(detail.get("name")), - documents=documents) + name = str(detail.get("name")) + pool = (documents if scoped else + [doc for doc in documents + if doc.get("folderId") == detail.get("folderId")]) + entry, _ = _resolve_document(client, name, documents=pool) if entry is not None and entry.get("id") != one_id: raise PageIndexAPIError( - f'Document "{detail.get("name")}" (doc_id: {one_id}) is ' + f'Document "{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." ) + if not scoped and any(doc.get("name") == name + and doc.get("id") != one_id + for doc in documents): + folder = detail.get("folderId") or "root" + folder_notes.append( + f'A document named "{name}" also exists in another folder ' + f'โ€” pass folder_id "{folder}" together with doc_name in ' + "every tool call to address the targeted one." + ) context = json.dumps(details, ensure_ascii=False) if len(details) == 1: - return ( + block = ( 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()." - ) + else: + names = ", ".join(str(item.get("name")) for item in details) + block = ( + 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()." + ) + if folder_notes: + block += "\n" + "\n".join(folder_notes) + return block def build_agent_instructions(client, doc_id=None) -> str: diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index f9e8d1c70..321d4cd20 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -26,7 +26,7 @@ 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): + page_num=None, folder_id=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"}, @@ -45,7 +45,7 @@ def seed_doc(storage_path, doc_id, name, *, created_at="2026-08-01T10:00:00.1230 "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", + "folderId": folder_id, "metadata": metadata, "mode": "standard", } DocStore(storage_path).save_document(doc_id, meta, tree, pages) return doc_id @@ -1332,6 +1332,21 @@ def test_agent_instructions_shadowed_doc_id_raises(client, store_path): assert "report.pdf" in text +def test_same_name_in_another_folder_directs_instead_of_refusing(client, + store_path): + """A same-name document in a different folder is the cloud contract's + supported case: the targeting block serves the call and directs the + agent to disambiguate with folder_id, instead of raising.""" + seed_doc(store_path, "pi-old", "report.pdf", folder_id="f-reports", + created_at="2026-08-01T10:00:00.000000") + seed_doc(store_path, "pi-new", "report.pdf", folder_id="f-archive", + created_at="2026-08-02T10:00:00.000000") + text = client.agent_instructions(doc_id="pi-old") + assert 'pass folder_id "f-reports"' in text + newer = client.agent_instructions(doc_id="pi-new") + assert 'pass folder_id "f-archive"' in newer + + def test_wait_tolerates_transient_network_failures(fake_cloud_client, monkeypatch): import requests as requests_mod cloud = fake_cloud_client(["processing", "completed"]) From ec58189a3609dba215e86fef3fbf99328aae154f Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 12 Aug 2026 17:31:47 +0800 Subject: [PATCH 033/137] Revert "fix: same-name documents in different folders no longer refuse doc_id targeting" MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This reverts commit 9d16dbf6, whose premise collapsed on verification against the cloud upload paths. Review finding #10 inferred from the folder_id tool description that one file name in two folders is a legal cloud layout; both upload paths actually dedup names per USER SPACE with no folder dimension โ€” chat's getSignedUploadUrl queries fileName + sourceName + mode + owner (file-access.service.ts), and compute's get_upload_url probes the S3 key (user, source, file_name) โ€” so own documents cannot share a name across any folders. The only legitimate same-name source is shared mounts (shared-with-me/following), which the api-proxy surface this SDK talks to never carries. Cloud and local therefore share one invariant โ€” names unique per space, server-enforced โ€” and the original global shadow check was the right shape: a duplicate is an anomaly worth refusing loudly, not a layout to accommodate with per-folder adjudication and conditional prompt notes. The invariant is now stated in doc_targeting_block's docstring so the finding does not get re-raised. --- pageindex/agent_tools.py | 52 +++++++++++++-------------------------- tests/test_agent_tools.py | 19 ++------------ 2 files changed, 19 insertions(+), 52 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index 737f20cda..747d03255 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -1549,13 +1549,12 @@ def doc_targeting_block(client, doc_id, scoped: bool = False) -> Optional[str]: 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 in the same folder โ€” the name-addressed tools could - not reach it. A same-name document in ANOTHER folder is the cloud - contract's supported case: no refusal, the block instead directs the - agent to pass folder_id (the documented disambiguator) on every call. - With ``scoped`` (the chat surfaces, whose tools resolve names inside - the doc_id allowlist) only a same-name duplicate within the targeted - set shadows.""" + same-name document โ€” the name-addressed tools could not reach it. Names + are unique per user space by upload-time dedup on both cloud surfaces + and locally, so a duplicate is an anomaly worth refusing loudly, not a + layout to accommodate. With ``scoped`` (the chat 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) @@ -1565,49 +1564,32 @@ def doc_targeting_block(client, doc_id, scoped: bool = False) -> Optional[str]: documents = ([{**detail, "id": one_id} for one_id, detail in zip(doc_ids, details)] if scoped else _all_documents(client)) - folder_notes: list[str] = [] for one_id, detail in zip(doc_ids, details): - name = str(detail.get("name")) - pool = (documents if scoped else - [doc for doc in documents - if doc.get("folderId") == detail.get("folderId")]) - entry, _ = _resolve_document(client, name, documents=pool) + 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 "{name}" (doc_id: {one_id}) is ' + 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." ) - if not scoped and any(doc.get("name") == name - and doc.get("id") != one_id - for doc in documents): - folder = detail.get("folderId") or "root" - folder_notes.append( - f'A document named "{name}" also exists in another folder ' - f'โ€” pass folder_id "{folder}" together with doc_name in ' - "every tool call to address the targeted one." - ) context = json.dumps(details, ensure_ascii=False) if len(details) == 1: - block = ( + 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()." ) - else: - names = ", ".join(str(item.get("name")) for item in details) - block = ( - 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()." - ) - if folder_notes: - block += "\n" + "\n".join(folder_notes) - return block + 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) -> str: diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 321d4cd20..f9e8d1c70 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -26,7 +26,7 @@ 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, folder_id=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"}, @@ -45,7 +45,7 @@ def seed_doc(storage_path, doc_id, name, *, created_at="2026-08-01T10:00:00.1230 "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": folder_id, "metadata": metadata, "mode": "standard", + "folderId": None, "metadata": metadata, "mode": "standard", } DocStore(storage_path).save_document(doc_id, meta, tree, pages) return doc_id @@ -1332,21 +1332,6 @@ def test_agent_instructions_shadowed_doc_id_raises(client, store_path): assert "report.pdf" in text -def test_same_name_in_another_folder_directs_instead_of_refusing(client, - store_path): - """A same-name document in a different folder is the cloud contract's - supported case: the targeting block serves the call and directs the - agent to disambiguate with folder_id, instead of raising.""" - seed_doc(store_path, "pi-old", "report.pdf", folder_id="f-reports", - created_at="2026-08-01T10:00:00.000000") - seed_doc(store_path, "pi-new", "report.pdf", folder_id="f-archive", - created_at="2026-08-02T10:00:00.000000") - text = client.agent_instructions(doc_id="pi-old") - assert 'pass folder_id "f-reports"' in text - newer = client.agent_instructions(doc_id="pi-new") - assert 'pass folder_id "f-archive"' in newer - - def test_wait_tolerates_transient_network_failures(fake_cloud_client, monkeypatch): import requests as requests_mod cloud = fake_cloud_client(["processing", "completed"]) From 001493c6d107f7983f65d2bd7264633eb177c290 Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 12 Aug 2026 17:38:25 +0800 Subject: [PATCH 034/137] docs: state the name-uniqueness invariant in library terms --- pageindex/agent_tools.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index 747d03255..8e2fdb757 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -1550,11 +1550,11 @@ def doc_targeting_block(client, doc_id, scoped: bool = False) -> Optional[str]: 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. Names - are unique per user space by upload-time dedup on both cloud surfaces - and locally, so a duplicate is an anomaly worth refusing loudly, not a - layout to accommodate. With ``scoped`` (the chat surfaces, whose tools - resolve names inside the doc_id allowlist) only a same-name duplicate - within the targeted set shadows.""" + are unique per library in both modes (uploads deduplicate a taken name + with _1.._99 suffixes), so a duplicate is an anomaly worth refusing + loudly, not a layout to accommodate. With ``scoped`` (the chat + 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) From ec1fd84fad6bfdba1a13ccdd4c5322ddedd20ee1 Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 12 Aug 2026 17:39:59 +0800 Subject: [PATCH 035/137] docs: trim doc_targeting_block docstring to the contract --- pageindex/agent_tools.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index 8e2fdb757..ff58d9913 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -1549,12 +1549,10 @@ def doc_targeting_block(client, doc_id, scoped: bool = False) -> Optional[str]: 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. Names - are unique per library in both modes (uploads deduplicate a taken name - with _1.._99 suffixes), so a duplicate is an anomaly worth refusing - loudly, not a layout to accommodate. With ``scoped`` (the chat - surfaces, whose tools resolve names inside the doc_id allowlist) only - a same-name duplicate within the targeted set shadows.""" + same-name document โ€” the name-addressed tools could not reach it. With + ``scoped`` (the chat 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) From b14af3386d05f5803aea36dcf4956f065fa752f4 Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 12 Aug 2026 18:59:33 +0800 Subject: [PATCH 036/137] fix: compress out-of-range page lists in get_page_content MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two message strings enumerated every out-of-range page number one by one while the payload fields beside them already used _format_page_spec. Against a 2-page document, pages="3-10000" produced a 59,310-character error whose own requested_pages field expressed the identical set as "3-10000"; the mixed case pages="1-10000" produced 59,479. Both now render through the helper: 431 and 600 characters. This was inherited behaviour, not a local slip โ€” the cloud MCP server enumerated at the same two sites, so local reproduced it verbatim. The cloud fixed it first (pageindex-chat #449), and this follows to keep the strings byte-identical; the error message now matches the served one character for character. A differential run of the two compressors over 94 inputs (empty, single, unsorted, duplicated, 10k spans, 80 random) agrees on every one, separator included. The new test pins all three shapes, including the non-contiguous case ("5,9" must not collapse into a range) that the compressor had no direct coverage for. --- pageindex/agent_tools.py | 4 ++-- tests/test_agent_tools.py | 20 ++++++++++++++++++++ 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index ff58d9913..376537473 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -987,7 +987,7 @@ def _get_page_content(client, doc_name: str, pages: str, 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: {', '.join(map(str, out_of_range))}", + f"pages, but you requested pages: {_format_page_spec(out_of_range)}", { "doc_name": doc_name, "max_pages": max_page, @@ -1037,7 +1037,7 @@ def _get_page_content(client, doc_name: str, pages: str, parts.append(f"Pages {_format_page_spec(remaining)} were " "omitted due to response size limits.") if out_of_range: - parts.append(f"Pages {', '.join(map(str, out_of_range))} " + parts.append(f"Pages {_format_page_spec(out_of_range)} " "were out of range.") summary = " ".join(parts) else: diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index f9e8d1c70..90e0a3f07 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -330,6 +330,26 @@ def test_page_content_out_of_range(client, store_path): 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") From f523485fa7e2f4e32d71415e9f5dc7eb56fe7abb Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 12 Aug 2026 18:59:33 +0800 Subject: [PATCH 037/137] docs: messages() marks only the managed prefix with cache_control MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The method docstring claimed the doc targeting block carries a cache_control breakpoint too, and the doc_id note called that block part of the cached prompt prefix. _anthropic_system deliberately marks only the stable managed prefix โ€” the API allows four breakpoints and the varying doc block must not consume one โ€” and the block is appended after the sole breakpoint, so it is never cached. a45b554 added the block with cache_control, making the claim true when written; daac9d2 removed it without touching the docstring, and adb2f1f then added the "cached prompt prefix" sentence after the fact. The same phrase at the chat_completions and responses docstrings is correct โ€” there the block is a leading conversation item inside the auto-cached prefix โ€” so only the Messages surface is reworded. The doc_id advice itself stands: the block is per-call table-setting and should stay identical across a conversation. Only the caching rationale was wrong. --- pageindex/client.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/pageindex/client.py b/pageindex/client.py index 3bf7c2e3e..371f9acb2 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -491,8 +491,8 @@ def messages( 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 and the doc - targeting block carry ``cache_control`` breakpoints. + append to your history. The managed system prompt carries a + ``cache_control`` breakpoint. Args: messages: Native Messages-format history (including prior @@ -508,8 +508,7 @@ def messages( 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 and is part - of the cached prompt prefix. + 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 From 166c60c17fc0ab86d69f9e3ad81fa5441ad689a0 Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 12 Aug 2026 18:59:33 +0800 Subject: [PATCH 038/137] docs: _stream_sync cancels on close, not on abandonment The docstring promised that closing "or abandoning" the iterator cancels the run. Abandoning only works when refcounting collects the generator: a caller that breaks out of the loop while keeping the reference never runs the finally that sets the cancel event, so the pump thread stays parked on the full queue and the backend client is never released. Closing is correct and is what the dedicated test exercises. Narrowing the promise to the behaviour the code actually provides is the honest fix; a watchdog or finalizer would be machinery bought for a shape the sync surface is not meant to serve, and the async client planned for 0.2.11 gets native task cancellation instead. --- pageindex/local_chat.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 787208e3e..bf8d63226 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -131,10 +131,10 @@ def _run_sync(coro): def _stream_sync(agen_factory) -> Iterator[Any]: """Drive an async generator from a background thread; yield synchronously. - Closing (or abandoning) 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. + 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() From 18171d63586900e2f8321babfeb5b1a76822d59e Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 12 Aug 2026 18:59:33 +0800 Subject: [PATCH 039/137] fix: raise the openai-agents floor to 0.14.0 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _conversation_group_id feeds RunConfig.group_id into OpenAI's prompt_cache_key so a round-tripped prefix stays in one cache group. That wiring first appears in openai-agents 0.14.0: 0.8.0 through 0.13.x have no prompt_cache_key at all, and group_id there is a tracing group id only โ€” inert, since tracing is disabled on the line above. An install resolving to the declared floor lost the cache continuity that the responses() docstring sells, silently and with no test able to catch it. The old floor's rationale (0.8.0 offloads sync tools to a thread) is subsumed by the new one. Every symbol the package imports predates 0.14.0, so nothing else constrains the bound. --- pyproject.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index a84b083e3..530c963d2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,9 +39,9 @@ regex = ">=2024.0.0" python-dotenv = ">=1.0.0" pyyaml = ">=6.0" claude-agent-sdk = { version = ">=0.1.0", optional = true } -# 0.8.0 offloads sync tools to a thread; older versions run them inline and -# a blocking bridge call would freeze the agent event loop. -openai-agents = { version = ">=0.8.0", optional = true } +# 0.14.0 is the first release that feeds RunConfig.group_id into the OpenAI +# prompt_cache_key; below it the conversation cache group is inert. +openai-agents = { version = ">=0.14.0", optional = true } # messages() and as_anthropic_tools() need the SDK's beta tool runner; # 0.84.0 is the first release with ToolError (failed tool calls flagged # is_error) whose runner also executes the final turn's tools on a From 9f67fdd17bfc68cd74d6e21220af7f35a887da22 Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 12 Aug 2026 21:10:56 +0800 Subject: [PATCH 040/137] fix: enforce doc_id at the tool layer in the framework config helpers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit openai_agent_config / anthropic_runner_config / claude_agent_config accepted doc_id but built unscoped tools, so the parameter that is a structural allowlist on chat_completions() was prompt-only advice here โ€” the agent could read every document in the store regardless. - as_openai_tools / as_anthropic_tools / as_claude_mcp take a doc_id tail parameter and thread it to the existing _allowed_ids channel; the config helpers pass it through in local mode - cloud config helpers keep prompt-level targeting (tool scoping is server-side there, documented); explicit as_*(doc_id=...) raises on cloud instead of silently dropping the allowlist โ€” including the hosted branch, which returned before _tool_specs' existing guard - _require_local_scope consolidates the cloud rejection that was inlined in _tool_specs - doc_id=[] is an empty allowlist, not "unscoped": dropped the `or None` at the three local chat surfaces --- pageindex/agent_tools.py | 16 +++-- pageindex/client.py | 58 ++++++++++++---- pageindex/integrations/claude_agent_sdk.py | 8 ++- pageindex/integrations/openai_agents.py | 5 +- pageindex/local_chat.py | 6 +- tests/test_agent_tools.py | 80 ++++++++++++++++++++++ tests/test_local_chat.py | 20 ++++++ 7 files changed, 168 insertions(+), 25 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index 376537473..7e31f3acc 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -1380,18 +1380,24 @@ def _build_cloud_agent_tools(client, include_management: bool) -> list[Callable[ return [_make_bridge_function(bridge, meta) for meta in tools_meta] +def _require_local_scope(client, doc_ids) -> None: + """The allowlist is enforced in-process; cloud lookups run server-side, + so accepting doc_ids there would be advisory-only โ€” refuse loudly.""" + if doc_ids is not None and getattr(client, "api_key", None): + raise PageIndexAPIError( + "doc_ids scoping applies to local tools only โ€” cloud calls " + "are scoped server-side." + ) + + def _tool_specs(client, include_management: bool = False, doc_ids=None, ) -> "list[tuple[str, str, dict, Callable[[dict], tuple[str, bool]]]]": """(name, description, schema, invoke) per tool, for adapters that take the wire schema verbatim. ``invoke`` returns (envelope_text, is_error). Schemas are copies (frameworks keep the dict by reference). ``doc_ids`` is the local chat scope; cloud scoping is server-side.""" + _require_local_scope(client, doc_ids) if getattr(client, "api_key", None): - if doc_ids is not None: - raise PageIndexAPIError( - "doc_ids scoping applies to local tools only โ€” cloud calls " - "are scoped server-side." - ) bridge = _cloud_bridge(client) tools_meta = bridge.list_tools() if not include_management: diff --git a/pageindex/client.py b/pageindex/client.py index 371f9acb2..9ab0e936a 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -606,7 +606,8 @@ def agent_tools(self, include_management: bool = False) -> list[Callable[..., st return build_agent_tools(self, include_management) def as_openai_tools(self, include_management: bool = False, - hosted: bool = False) -> list: + 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 @@ -636,9 +637,20 @@ def as_openai_tools(self, include_management: bool = False, 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) + 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.""" + return None if getattr(self, "api_key", None) else doc_id def openai_agent_config( self, @@ -661,7 +673,9 @@ def openai_agent_config( Args: doc_id: Document ID or list of IDs to target, as in - ``agent_instructions``. + ``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. @@ -669,7 +683,8 @@ def openai_agent_config( config: dict[str, Any] = { "name": "PageIndex", "instructions": self.agent_instructions(doc_id=doc_id), - "tools": self.as_openai_tools(include_management), + "tools": self.as_openai_tools(include_management, + doc_id=self._local_doc_scope(doc_id)), } model = model or getattr(self, "retrieve_model", None) if model: @@ -677,7 +692,9 @@ def openai_agent_config( return config def as_anthropic_tools(self, include_management: bool = False, - asynchronous: bool = False) -> list: + 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 @@ -712,9 +729,14 @@ def as_anthropic_tools(self, include_management: bool = False, ``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) + return build_anthropic_tools(self, include_management, asynchronous, + doc_ids=doc_id) def anthropic_runner_config( self, @@ -745,7 +767,9 @@ def anthropic_runner_config( model: Backend model name (also resolves the ``max_tokens`` default). doc_id: Document ID or list of IDs to target, as in - ``agent_instructions``. + ``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 @@ -760,12 +784,13 @@ def anthropic_runner_config( "max_tokens": (max_tokens if max_tokens is not None else _default_max_tokens(model)), "system": self.agent_instructions(doc_id=doc_id), - "tools": self.as_anthropic_tools(include_management, - asynchronous), + "tools": self.as_anthropic_tools(include_management, asynchronous, + doc_id=self._local_doc_scope(doc_id)), "max_iterations": max_turns if max_turns is not None else 10, } - def as_claude_mcp(self, include_management: bool = False): + 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. @@ -776,7 +801,9 @@ def as_claude_mcp(self, include_management: bool = False): ``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]'``). + ``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 @@ -795,7 +822,7 @@ def as_claude_mcp(self, include_management: bool = False): ) """ from .integrations.claude_agent_sdk import build_claude_mcp - return build_claude_mcp(self, include_management) + return build_claude_mcp(self, include_management, doc_ids=doc_id) def claude_agent_config( self, @@ -817,14 +844,17 @@ def claude_agent_config( Args: doc_id: Document ID or list of IDs to target, as in - ``agent_instructions``. + ``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. """ return { "system_prompt": self.agent_instructions(doc_id=doc_id), - "mcp_servers": {server_name: self.as_claude_mcp(include_management)}, + "mcp_servers": {server_name: self.as_claude_mcp( + include_management, doc_id=self._local_doc_scope(doc_id))}, # 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}"], diff --git a/pageindex/integrations/claude_agent_sdk.py b/pageindex/integrations/claude_agent_sdk.py index 58b9da4e6..f3ab19c9c 100644 --- a/pageindex/integrations/claude_agent_sdk.py +++ b/pageindex/integrations/claude_agent_sdk.py @@ -14,7 +14,11 @@ from ..errors import PageIndexAPIError -def build_claude_mcp(client, include_management: bool = False): +def build_claude_mcp(client, include_management: bool = False, doc_ids=None): + from ..agent_tools import _require_local_scope + # The cloud branch returns a URL config โ€” reject cloud doc_ids so they + # are never silently dropped. + _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). @@ -38,7 +42,7 @@ def build_claude_mcp(client, include_management: bool = False): def make_handler(name: str): async def handler(arguments: dict[str, Any]) -> dict[str, Any]: text, is_error = await asyncio.to_thread( - call_tool, client, name, arguments or {} + call_tool, client, name, arguments or {}, doc_ids ) result: dict[str, Any] = {"content": [{"type": "text", "text": text}]} if is_error: diff --git a/pageindex/integrations/openai_agents.py b/pageindex/integrations/openai_agents.py index 9f5df5064..1e68f0a0f 100644 --- a/pageindex/integrations/openai_agents.py +++ b/pageindex/integrations/openai_agents.py @@ -27,6 +27,10 @@ def build_openai_tools(client, include_management: bool = False, "as_openai_tools requires the OpenAI Agents SDK โ€” " "pip install openai-agents (or pip install 'pageindex[openai]')." ) from exc + from ..agent_tools import _require_local_scope, _tool_specs + # The hosted branch returns before _tool_specs โ€” reject cloud doc_ids + # here so they are never silently dropped. + _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 @@ -39,7 +43,6 @@ def build_openai_tools(client, include_management: bool = False, "headers": {"Authorization": f"Bearer {client.api_key}"}, "require_approval": "never", })] - from ..agent_tools import _tool_specs def wrap(name, description, schema, invoke): async def on_invoke_tool(ctx: Any, args_json: str) -> str: diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index bf8d63226..c22697d3b 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -349,7 +349,7 @@ def run_chat_completions(client, messages, stream: bool = False, model_name = model or client.retrieve_model managed = _managed_instructions(system_texts) agent = _openai_agent(client, "chat", model_name, managed, - temperature, None, doc_ids=doc_id or None) + temperature, None, doc_ids=doc_id) run_kwargs = _run_kwargs(max_turns, _conversation_group_id(model_name, managed, items)) from agents import Runner @@ -450,7 +450,7 @@ def run_responses(client, input, model: Optional[str] = None, model_name = model or client.retrieve_model managed = _managed_instructions(extra) agent = _openai_agent(client, "responses", model_name, managed, - temperature, top_p, doc_ids=doc_id or None) + temperature, top_p, doc_ids=doc_id) run_kwargs = _run_kwargs(max_turns, _conversation_group_id(model_name, managed, items)) recorded: dict = {} @@ -691,7 +691,7 @@ def run_messages(client, messages, model: str, else _default_max_tokens(model)), messages=prepared, model=model, - tools=build_anthropic_tools(client, doc_ids=doc_id or None), + tools=build_anthropic_tools(client, doc_ids=doc_id), system=_anthropic_system(system, block), stream=stream, # Bounded like the OpenAI surfaces (their framework default is 10). diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 90e0a3f07..40c0c863f 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -471,6 +471,10 @@ def test_call_tool_doc_scope_limits_every_lookup(client, store_path): {"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 @@ -740,6 +744,82 @@ def test_anthropic_runner_config_cloud(cloud_with_fake_bridge): "get_document"] +# โ”€โ”€ config helpers: doc_id is structural in the tools, not just prompted โ”€โ”€ + +def test_openai_agent_config_doc_scope_enforced_in_tools(client, store_path): + pytest.importorskip("agents") + seed_doc(store_path, "pi-a", "report.pdf") + seed_doc(store_path, "pi-b", "payroll.pdf", + created_at="2026-08-02T10:00:00.123000") + tools = {tool.name: tool + for tool in client.openai_agent_config(doc_id="pi-a")["tools"]} + out = asyncio.run(tools["get_page_content"].on_invoke_tool( + None, json.dumps({"doc_name": "payroll.pdf", "pages": "1"}))) + assert json.loads(out)["errorCode"] == "NOT_FOUND" + out = asyncio.run(tools["browse_documents"].on_invoke_tool(None, "{}")) + assert [doc["name"] + for doc in json.loads(out)["documents"]] == ["report.pdf"] + + +def test_anthropic_runner_config_doc_scope_enforced_in_tools(client, + store_path): + pytest.importorskip("anthropic") + from anthropic.lib.tools import ToolError + seed_doc(store_path, "pi-a", "report.pdf") + seed_doc(store_path, "pi-b", "payroll.pdf", + created_at="2026-08-02T10:00:00.123000") + config = client.anthropic_runner_config(model="claude-sonnet-4-5", + doc_id="pi-a") + tools = {tool.name: tool for tool in config["tools"]} + with pytest.raises(ToolError, match="NOT_FOUND"): + tools["get_page_content"].call({"doc_name": "payroll.pdf", + "pages": "1"}) + browse = json.loads(tools["browse_documents"].call({})) + assert [doc["name"] for doc in browse["documents"]] == ["report.pdf"] + + +def test_claude_agent_config_doc_scope_enforced_in_tools(client, store_path): + pytest.importorskip("claude_agent_sdk") + from mcp.types import CallToolRequest, CallToolRequestParams + seed_doc(store_path, "pi-a", "report.pdf") + seed_doc(store_path, "pi-b", "payroll.pdf", + created_at="2026-08-02T10:00:00.123000") + config = client.claude_agent_config(doc_id="pi-a") + server = config["mcp_servers"]["pageindex"] + handler = server["instance"].request_handlers[CallToolRequest] + result = asyncio.run(handler(CallToolRequest( + method="tools/call", + params=CallToolRequestParams( + name="get_page_content", + arguments={"doc_name": "payroll.pdf", "pages": "1"})))) + payload = json.loads(result.root.content[0].text) + assert payload["errorCode"] == "NOT_FOUND" + + +def test_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"): diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index ee1630650..5331e490f 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -629,6 +629,11 @@ def test_conversation_group_id_stable_per_conversation(): "m", "sys", [{"role": "user", "content": "other"}]) assert key != local_chat._conversation_group_id("m2", "sys", turn1) assert key != local_chat._conversation_group_id("m", "sys2", turn1) + + +@needs_agents +def test_run_kwargs_sets_conversation_group_id(): + key = "pageindex-test" assert (local_chat._run_kwargs(None, key)["run_config"].group_id == key) @@ -664,6 +669,21 @@ def tool_outputs(items): assert [doc["name"] for doc in browse["documents"]] == ["report.pdf"] +@needs_agents +def test_empty_doc_id_is_an_empty_allowlist(client, store_path, fake_model): + """doc_id=[] scopes the agent to nothing; `or None` used to wash it + into unscoped full-library access.""" + seed_doc(store_path, "pi-a", "report.pdf") + fake = fake_model([ + [_call_item("browse_documents", {})], + [_msg_item("done")], + ]) + client.chat_completions("q", doc_id=[]) + outputs = [item["output"] for item in fake.inputs[1] + if item.get("type") == "function_call_output"] + assert json.loads(outputs[-1])["documents"] == [] + + @needs_agents def test_openai_model_resolves_provider_prefixes(): """retrieve_model arrives normalized (litellm//); the From 66912c6670eda0380b09be88834ad7be14afc7a2 Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 12 Aug 2026 22:19:41 +0800 Subject: [PATCH 041/137] =?UTF-8?q?fix:=20two=20chat=20findings=20?= =?UTF-8?q?=E2=80=94=20final-turn=20append=20and=20cache-key=20seeding?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit run_messages keyed its re-append guard on stop_reason, but the anthropic runner executes tools whenever the turn's content carries tool_use blocks (refusal excepted) โ€” a max_tokens turn with complete tool_use blocks was already appended by the runner, so the guard re-appended it, duplicating tool_use ids and 400ing the documented verbatim continuation. The guard now checks whether final's tool_use ids already sit in the appended history; unexecuted tool_use blocks (refusal turns) are stripped from the appendable history, as the SDK itself does when rebuilding params around an unresulted turn. _conversation_group_id seeded on items[0], which is the doc-targeting block whenever doc_id is set โ€” byte-identical across every conversation about a document, so all of them pooled under one prompt_cache_key and evicted each other's prefixes. Seed on the conversation's own first item instead: continuations keep their key, unrelated conversations never share one. Also drop the dead pytestmark_openai assignment (pytest's magic name is pytestmark; the section gate it implied never existed). --- pageindex/local_chat.py | 50 +++++++++++++++++-------- tests/test_local_chat.py | 81 +++++++++++++++++++++++++++++++++++++++- 2 files changed, 114 insertions(+), 17 deletions(-) diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index c22697d3b..efac51651 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -250,8 +250,11 @@ def _conversation_group_id(model_name: str, instructions: str, items) -> str: RunConfig.group_id into the OpenAI prompt_cache_key, and without one it stamps every run with a fresh key, tagging a round-tripped prefix as a different cache group. Keyed on the prefix identity โ€” model, - instructions, first input item โ€” so a conversation's continuations - share one route without pooling unrelated conversations.""" + instructions, first conversation item โ€” so a conversation's + continuations share one route without pooling unrelated conversations. + Callers pass the conversation's own items, never the SDK-prepended + doc-targeting block: that block is byte-identical for every + conversation about a document and would pool them all under one key.""" seed = json.dumps([model_name, instructions, items[0] if items else None], sort_keys=True, default=str) @@ -351,7 +354,8 @@ def run_chat_completions(client, messages, stream: bool = False, agent = _openai_agent(client, "chat", model_name, managed, temperature, None, doc_ids=doc_id) run_kwargs = _run_kwargs(max_turns, - _conversation_group_id(model_name, managed, items)) + _conversation_group_id(model_name, managed, + history)) from agents import Runner from agents.exceptions import AgentsException, MaxTurnsExceeded if not stream: @@ -444,6 +448,7 @@ def run_responses(client, input, model: Optional[str] = None, 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 [] @@ -452,7 +457,8 @@ def run_responses(client, input, model: Optional[str] = None, agent = _openai_agent(client, "responses", model_name, managed, temperature, top_p, doc_ids=doc_id) run_kwargs = _run_kwargs(max_turns, - _conversation_group_id(model_name, managed, items)) + _conversation_group_id(model_name, managed, + conversation)) recorded: dict = {} from agents import Runner from agents.exceptions import AgentsException, MaxTurnsExceeded @@ -723,19 +729,31 @@ def capture(params): 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, so the final - # assistant message is missing exactly when the run ended naturally - # (stop_reason != "tool_use"); on a max_turns cut the last appended - # turn IS the final message and appending again would duplicate its - # tool_use ids. + # 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):]] - if (final.stop_reason != "tool_use" - and (not new_messages - or new_messages[-1].get("role") != "assistant")): - new_messages = new_messages + [{ - "role": "assistant", - "content": [_dump_block(item) for item in final.content], - }] + final_blocks = [_dump_block(item) for item in final.content] + final_ids = {block["id"] for block in final_blocks + if block.get("type") == "tool_use"} + history_ids = {block.get("id") + for message in new_messages + if (message.get("role") == "assistant" + and isinstance(message.get("content"), list)) + for block in message["content"] + if (isinstance(block, dict) + and block.get("type") == "tool_use")} + if not final_ids or not final_ids <= history_ids: + # Unexecuted tool_use blocks (refusal turns) have no tool_result, + # so they cannot enter an appendable history โ€” strip them, as the + # SDK itself does when it rebuilds params around such a turn. + appendable = [block for block in final_blocks + if block.get("type") != "tool_use"] + if appendable: + new_messages = new_messages + [ + {"role": "assistant", "content": appendable}] envelope["messages"] = new_messages return envelope diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index 5331e490f..701710145 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -49,7 +49,6 @@ def client(store_path): needs_agents = pytest.mark.skipif(not _HAS_AGENTS, reason="openai-agents not installed") -pytestmark_openai = needs_agents def _msg_item(text): @@ -339,6 +338,43 @@ def test_responses_round_trip_prefix_with_doc_id(client, store_path, fake_model) assert second.inputs[0][:len(previous_final)] == previous_final +@needs_agents +def test_doc_id_conversations_get_distinct_cache_keys(client, store_path, + fake_model, + monkeypatch): + """The doc-targeting block is byte-identical for every conversation + about a document โ€” seeding the cache key on items[0] pooled them all + under one prompt_cache_key.""" + seed_doc(store_path, "pi-a", "report.pdf") + keys = [] + real = local_chat._run_kwargs + + def spy(max_turns, group_id): + keys.append(group_id) + return real(max_turns, group_id) + + monkeypatch.setattr(local_chat, "_run_kwargs", spy) + + fake_model([[_msg_item("a")]]) + result = client.responses("What is the CAGR?", doc_id="pi-a") + fake_model([[_msg_item("b")]]) + client.responses("Summarize section 3.", doc_id="pi-a") + assert keys[0] != keys[1] # unrelated conversations never pool + + fake_model([[_msg_item("c")]]) + follow_up = ([{"role": "user", "content": "What is the CAGR?"}] + + result["output"] + + [{"role": "user", "content": "and now?"}]) + client.responses(follow_up, doc_id="pi-a") + assert keys[2] == keys[0] # a continuation keeps its conversation's key + + fake_model([[_msg_item("d")]]) + client.chat_completions("What is the CAGR?", doc_id="pi-a") + fake_model([[_msg_item("e")]]) + client.chat_completions("Summarize section 3.", doc_id="pi-a") + assert keys[3] != keys[4] # same property on the chat surface + + @needs_agents def test_responses_stream_passthrough(client, store_path, fake_model): seed_doc(store_path, "pi-a", "report.pdf") @@ -930,6 +966,49 @@ def test_messages_max_turns_truncation_round_trippable(client, store_path, 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") From 31c9150930f20554c5dc53e924f265f0b4146f05 Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 12 Aug 2026 23:07:12 +0800 Subject: [PATCH 042/137] =?UTF-8?q?fix:=20six=20review=20findings=20?= =?UTF-8?q?=E2=80=94=20pagination,=20compat,=20and=20containment?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - _all_documents advances by what actually arrived and treats `total` as an optimization: absent/null totals and short pages silently truncated the library behind every name resolution - _make_bridge_function survives description: null (the parallel _tool_specs path already did) - as_openai_tools answers a malformed argument string with the guided error envelope instead of raising through the caller's whole run - the pre-0.2.10 package attributes (ConfigLoader, count_tokens, ...) resolve again: main's underscore-guarded fallthrough is restored โ€” dunder probes stay lazy, a non-underscore typo pays one classic import before its AttributeError - _split_structure chunks are always lists: the structure field no longer changes JSON type between parts of one paginated response - the bridge replays only session-carrying 404s (the spec's expiry status); 400 raises instead of re-running side effects, and the reset double-checks under the lock so concurrent retries cannot clobber a freshly re-initialized session --- pageindex/__init__.py | 21 ++-- pageindex/agent_tools.py | 29 +++-- pageindex/integrations/openai_agents.py | 24 ++++- pageindex/mcp_bridge.py | 35 +++--- tests/test_agent_tools.py | 135 +++++++++++++++++++++++- tests/test_package_surface.py | 32 ++++-- 6 files changed, 236 insertions(+), 40 deletions(-) diff --git a/pageindex/__init__.py b/pageindex/__init__.py index 5a19bd895..8a2383014 100644 --- a/pageindex/__init__.py +++ b/pageindex/__init__.py @@ -29,16 +29,25 @@ "mcp_bridge", "page_index_classic", "page_index_md", "tree_optimize", "utils"} - def __getattr__(name): + if name.startswith("_"): + # Dunder probes (copy, pickle, inspect) are the frequent unknown + # names โ€” they must not trigger the classic import below. + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") import importlib if name in _SUBMODULES: return importlib.import_module(f".{name}", __name__) - if name not in _LAZY: - # Unknown names must not fall through to an eager import of the - # heavy indexing stack. - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") - value = getattr(importlib.import_module(_LAZY[name], __name__), name) + # Pre-0.2.10 compat: unknown names fall through to the classic module, + # whose public surface (ConfigLoader, count_tokens, ...) resolved as + # package attributes. A non-underscore typo pays one classic import + # before its AttributeError โ€” not worth an allowlist. + module = importlib.import_module(_LAZY.get(name, ".page_index_classic"), + __name__) + try: + value = getattr(module, name) + except AttributeError: + raise AttributeError( + f"module {__name__!r} has no attribute {name!r}") from None globals()[name] = value return value diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index 7e31f3acc..859d61647 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -337,8 +337,13 @@ def _all_documents(client) -> list[dict[str, Any]]: page = client.list_documents(limit=100, offset=offset) batch = page.get("documents") or [] documents.extend(batch) - offset += 100 - if not batch or offset >= page.get("total", 0): + # 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 @@ -606,8 +611,11 @@ def _serialized_size(value: Any) -> int: 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; chunk - boundaries are implementation-defined.""" + 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] @@ -618,17 +626,18 @@ def _split_structure(structure: Any, budget: int) -> list[Any]: size = _serialized_size(node) if size > budget: if group: - chunks.append(group if len(group) > 1 else group[0]) + chunks.append(group) group, group_size = [], 0 - chunks.extend(_split_oversized_node(node, budget)) + chunks.extend([part] + for part in _split_oversized_node(node, budget)) continue if group and group_size + size > budget: - chunks.append(group if len(group) > 1 else group[0]) + chunks.append(group) group, group_size = [], 0 group.append(node) group_size += size if group: - chunks.append(group if len(group) > 1 else group[0]) + chunks.append(group) return chunks or [structure] @@ -641,6 +650,8 @@ def _split_oversized_node(node: Any, budget: int) -> list[Any]: 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 @@ -1332,7 +1343,7 @@ def proxy(**kwargs: Any) -> str: annotations["return"] = str proxy.__annotations__ = annotations proxy.__name__ = proxy.__qualname__ = name or "tool" - proxy.__doc__ = _tool_docstring(meta.get("description", ""), properties) + proxy.__doc__ = _tool_docstring(meta.get("description") or "", properties) return proxy diff --git a/pageindex/integrations/openai_agents.py b/pageindex/integrations/openai_agents.py index 1e68f0a0f..266f71618 100644 --- a/pageindex/integrations/openai_agents.py +++ b/pageindex/integrations/openai_agents.py @@ -27,7 +27,8 @@ def build_openai_tools(client, include_management: bool = False, "as_openai_tools requires the OpenAI Agents SDK โ€” " "pip install openai-agents (or pip install 'pageindex[openai]')." ) from exc - from ..agent_tools import _require_local_scope, _tool_specs + from ..agent_tools import (_dumps, _failure, _require_local_scope, + _tool_specs) # The hosted branch returns before _tool_specs โ€” reject cloud doc_ids # here so they are never silently dropped. _require_local_scope(client, doc_ids) @@ -46,8 +47,25 @@ def build_openai_tools(client, include_management: bool = False, def wrap(name, description, schema, invoke): async def on_invoke_tool(ctx: Any, args_json: str) -> str: - arguments = {key: value for key, value - in (json.loads(args_json) if args_json else {}).items() + # strict_json_schema is off, so the provider never validates the + # payload; a malformed or non-object argument string must come + # back as the guided error envelope โ€” raising here aborts the + # caller's whole run (hand-built FunctionTools have no + # failure_error_function to hand the error back to the model). + try: + parsed = json.loads(args_json) if args_json else {} + except ValueError: + parsed = None + if not isinstance(parsed, dict): + payload, _ = _failure( + f"Invalid arguments for {name}: expected a JSON object, " + f"got: {(args_json or '')[:200]!r}", None, + {"summary": "Malformed tool arguments", + "options": [f"Re-send the {name} call with a JSON " + "object of its parameters"]}, + "INVALID_INPUT") + return _dumps(payload) + arguments = {key: value for key, value in parsed.items() if value is not None} text, _ = await asyncio.to_thread(invoke, arguments) return text diff --git a/pageindex/mcp_bridge.py b/pageindex/mcp_bridge.py index f23575e20..6b7714550 100644 --- a/pageindex/mcp_bridge.py +++ b/pageindex/mcp_bridge.py @@ -5,8 +5,9 @@ ``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 request rejected after session expiry -re-initializes once and retries. +``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 @@ -52,10 +53,8 @@ def __init__(self, url: str, headers: dict[str, str]): # โ”€โ”€ JSON-RPC over streamable HTTP โ”€โ”€ - def _post(self, payload: dict) -> requests.Response: - with self._lock: - session_id = self._session_id - protocol_version = self._protocol_version + 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", @@ -108,17 +107,24 @@ def _request(self, method: str, params: Optional[dict] = None, 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) - if response.status_code in (400, 404) and self._initialized and _retry: - # Session expired (stateful servers): start over, retry once. + 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: - self._initialized = False - self._session_id = None - self._protocol_version = None + 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( @@ -154,9 +160,12 @@ def _ensure_initialized(self) -> None: _PROTOCOL_VERSION) self._instructions = result.get("instructions") self._initialized = True + session_id = self._session_id + protocol_version = self._protocol_version try: self._post({"jsonrpc": "2.0", - "method": "notifications/initialized"}) + "method": "notifications/initialized"}, + session_id, protocol_version) except PageIndexAPIError: pass # advisory; a server that required it fails the next request diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 40c0c863f..349c60370 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -289,9 +289,10 @@ def test_structure_multipart_pagination(client, store_path): for part in range(1, first["total_parts"] + 1): payload, _ = run(client, "get_document_structure", doc_name="big.pdf", part=part) - chunk = payload["structure"] - nodes = chunk if isinstance(chunk, list) else [chunk] - titles.extend(node["title"] for node in nodes) + # 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)] @@ -300,6 +301,22 @@ def test_structure_multipart_pagination(client, store_path): 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): @@ -594,6 +611,20 @@ def test_as_openai_tools_invocation_runs_call_tool(client, store_path): 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 {}.""" @@ -1079,6 +1110,27 @@ def test_cloud_agent_tools_proxy_and_drop_none(cloud_with_fake_bridge): assert created["bridge"].calls == [("get_document", {"doc_name": "report.pdf"})] +def test_cloud_agent_tools_null_description_survives(): + """A server may send description: null โ€” .get(key, default) does not + apply the default to it, and agent_tools() died with a TypeError while + the _tool_specs path handled the same payload fine.""" + from pageindex.agent_tools import _make_bridge_function + + class _Bridge: + def call_tool(self, name, arguments): + return json.dumps({"success": True}), False + + tool = _make_bridge_function(_Bridge(), { + "name": "search_documents", + "description": None, + "inputSchema": {"type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"]}, + }) + assert tool.__name__ == "search_documents" + assert json.loads(tool(query="x"))["success"] is True + + def test_cloud_agent_tools_call_errors_contained(cloud_with_fake_bridge): cloud, created = cloud_with_fake_bridge search, _ = cloud.agent_tools() @@ -1189,6 +1241,54 @@ def fake_post(url, json=None, headers=None, timeout=None): assert "Mcp-Session-Id" not in reinit["headers"] +def test_mcp_bridge_400_is_an_error_not_session_expiry(monkeypatch): + """The spec's expired-session status is 404; a 400 is an ordinary bad + request โ€” treating it as expiry replayed the rejected call (running a + management tool's side effect twice) behind a spurious re-initialize.""" + import requests as requests_mod + import pageindex.mcp_bridge as mcp_bridge + from pageindex.mcp_bridge import McpBridge + + posts = [] + + class _Resp: + def __init__(self, status, body=None, headers=None, text=""): + self.status_code = status + self._body = body + self.headers = headers or {"Content-Type": "application/json"} + self.text = text or (json.dumps(body) if body else "") + self.content = self.text.encode("utf-8") + + def json(self): + if self._body is None: + raise ValueError("no body") + return self._body + + def fake_post(url, json=None, headers=None, timeout=None): + posts.append(json.get("method")) + rid = json.get("id") + if json.get("method") == "initialize": + return _Resp(200, {"jsonrpc": "2.0", "id": rid, + "result": {"protocolVersion": "2025-06-18"}}, + {"Content-Type": "application/json", + "Mcp-Session-Id": "sess-1"}) + if json.get("method") == "notifications/initialized": + return _Resp(202) + return _Resp(400, text="unknown tool") + + monkeypatch.setattr(mcp_bridge, "requests", types.SimpleNamespace( + post=fake_post, RequestException=requests_mod.RequestException)) + bridge = McpBridge("https://api.pageindex.ai/mcp", + {"Authorization": "Bearer k"}) + with pytest.raises(PageIndexAPIError, match="HTTP 400"): + bridge.call_tool("nope", {}) + # Exactly one call attempt, no replay, no re-initialize; the live + # session survives for the next request. + assert posts.count("tools/call") == 1 + assert posts.count("initialize") == 1 + assert bridge._session_id == "sess-1" + + # โ”€โ”€ review-round regressions โ”€โ”€ def test_synth_optional_no_default_param_is_nullable(): @@ -1411,6 +1511,35 @@ def spy(**kwargs): assert payload["has_more"] is True and payload["next_offset"] == 2 +def test_all_documents_survives_short_pages_and_missing_total(): + """The full-library walk behind every name resolution must trust what + actually arrives: a server capping page size, omitting `total`, or + sending total: null silently truncated the library (or raised).""" + from pageindex.agent_tools import _all_documents + + docs = [{"id": f"pi-{index}"} for index in range(120)] + + def make_client(total_field, page_cap): + class _Client: + calls = 0 + + def list_documents(self, limit, offset): + type(self).calls += 1 + page = {"documents": docs[offset:offset + min(limit, + page_cap)]} + if total_field != "omit": + page["total"] = total_field + return page + return _Client() + + assert _all_documents(make_client(120, 50)) == docs # short pages + assert _all_documents(make_client("omit", 100)) == docs # no total + assert _all_documents(make_client(None, 100)) == docs # total: null + exact = make_client(120, 100) # well-behaved server: + assert _all_documents(exact) == docs + assert type(exact).calls == 2 # ...total still saves the empty page + + def test_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.""" diff --git a/tests/test_package_surface.py b/tests/test_package_surface.py index d93e0de8c..6d985c818 100644 --- a/tests/test_package_surface.py +++ b/tests/test_package_surface.py @@ -62,22 +62,42 @@ def test_import_pageindex_is_lazy(): assert out.stdout.split() == ["clean", "function"] -def test_sdk_submodules_reachable_and_unknown_names_stay_lazy(): - """The 0.2.10 modules resolve as attributes, and an unknown name raises - AttributeError without dragging in the indexing stack.""" +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" - "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" ) 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" From 32c4940f0b734b54aac0665df5cb254e439904e2 Mon Sep 17 00:00:00 2001 From: Ray Date: Thu, 13 Aug 2026 00:07:46 +0800 Subject: [PATCH 043/137] =?UTF-8?q?fix:=20five=20secondary=20review=20find?= =?UTF-8?q?ings=20=E2=80=94=20containment=20and=20guards?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - the bridge maps content blocks individually: base64 payloads (image/audio) become metadata stubs instead of handing the model the raw blob, text blocks pass verbatim, anything else keeps the JSON dump (revisit if tool results become real multimodal input) - cloud proxy annotations keep array item types (list[str], not bare list) so strict function calling accepts the round-trip; a type-array in items degrades to bare list instead of crashing the build - run_messages raises when set_messages_params stops delivering params instead of silently dropping every tool turn from the envelope - call_tool drops None-valued arguments (None โ‰ก omitted, the contract's semantics) โ€” adapters that forward the model's nulls verbatim no longer trip parameter validation - client._parse_pages bounds the span arithmetically before materializing it, like the tool layer: "1-999999999" raises instead of allocating a billion integers --- pageindex/agent_tools.py | 30 ++++++++++++++----- pageindex/client.py | 11 +++++-- pageindex/local_chat.py | 11 ++++++- pageindex/mcp_bridge.py | 20 +++++++++---- tests/test_agent_tools.py | 61 +++++++++++++++++++++++++++++++++++++++ tests/test_client.py | 10 +++++++ tests/test_local_chat.py | 17 +++++++++++ 7 files changed, 144 insertions(+), 16 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index 859d61647..b93d9f043 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -1145,9 +1145,10 @@ def call_tool(client, name: str, arguments: dict[str, Any], ) return _dumps(payload), True # Underscore-prefixed keys are the SDK's private channel (the scope - # below), never model arguments. + # 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.items() - if not key.startswith("_")} + if not key.startswith("_") and value is not None} 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) @@ -1271,13 +1272,28 @@ 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}]. - schema_type = [option.get("type") for option in spec["anyOf"] - if isinstance(option, dict) and option.get("type")] + 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"] - base = _SCHEMA_TYPE_MAP.get(bases[0], Any) if bases else Any - return Optional[base] if "null" in schema_type else base - return _SCHEMA_TYPE_MAP.get(schema_type, Any) + schema_type = bases[0] if bases else None + base = _SCHEMA_TYPE_MAP.get(schema_type or "", Any) + if base is list: + # Strict function calling rejects arrays whose item type was lost + # in the annotation round-trip; parameterize when it is known. + item_type = (spec["items"].get("type") + if isinstance(spec.get("items"), dict) else None) + element = (_SCHEMA_TYPE_MAP.get(item_type) + if isinstance(item_type, str) else None) + if element is not None: + base = list[element] + return Optional[base] if nullable else base def _bridge_invoker(bridge, name: str) -> "Callable[[dict], tuple[str, bool]]": diff --git a/pageindex/client.py b/pageindex/client.py index 9ab0e936a..bb23150c9 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -11,15 +11,22 @@ def _parse_pages(pages: str) -> list[int]: result = [] + total = 0 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)) + start = end = int(part) + # Bound the span arithmetically before materializing it โ€” a spec + # like "1-999999999" would otherwise expand to a billion integers. + total += end - start + 1 + if total > 10_000: + raise ValueError(f"Page specification '{pages}' spans more than " + "10000 pages; request a narrower range") + result.extend(range(start, end + 1)) return sorted(set(result)) diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index efac51651..a9ec88f3b 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -722,7 +722,16 @@ def capture(params): return params runner.set_messages_params(capture) - conversation = list(captured.get("messages") or []) + 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] diff --git a/pageindex/mcp_bridge.py b/pageindex/mcp_bridge.py index 6b7714550..0bcbc01aa 100644 --- a/pageindex/mcp_bridge.py +++ b/pageindex/mcp_bridge.py @@ -194,9 +194,17 @@ def call_tool(self, name: str, arguments: dict[str, Any]) -> "tuple[str, bool]": result = self._request("tools/call", {"name": name, "arguments": arguments}) or {} is_error = bool(result.get("isError")) - blocks = result.get("content") or [] - texts = [block.get("text", "") for block in blocks - if isinstance(block, dict) and block.get("type") == "text"] - if len(texts) == len(blocks): - return "\n".join(texts), is_error - return json.dumps(blocks, ensure_ascii=False), is_error + texts = [] + for block in result.get("content") or []: + if isinstance(block, dict) and block.get("type") == "text": + texts.append(block.get("text", "")) + elif isinstance(block, dict) and isinstance(block.get("data"), str): + # Base64 payloads (image/audio) become a metadata stub โ€” + # dumped verbatim they hand the model the raw blob. Revisit + # if tool results ever pass through as real multimodal input. + kind = block.get("mimeType") or block.get("type") or "binary" + size_kb = max(1, len(block["data"]) * 3 // 4096) + texts.append(f"[{kind} content omitted: ~{size_kb} KB]") + else: + texts.append(json.dumps(block, ensure_ascii=False)) + return "\n".join(texts), is_error diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 349c60370..1de495264 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -1289,6 +1289,24 @@ def fake_post(url, json=None, headers=None, timeout=None): assert bridge._session_id == "sess-1" +def test_mcp_bridge_blob_blocks_become_stubs(): + """Non-text content used to be json.dumps'd wholesale, handing the + model the raw base64 payload of an image tool's response.""" + from pageindex.mcp_bridge import McpBridge + + bridge = McpBridge("https://api.pageindex.ai/mcp", {}) + blob = "A" * 8192 # ~6 KB decoded + bridge._request = lambda method, params: {"content": [ + {"type": "text", "text": "Page 3 of report.pdf"}, + {"type": "image", "mimeType": "image/png", "data": blob}, + ]} + text, is_error = bridge.call_tool("get_document_image", {}) + assert not is_error + assert "Page 3 of report.pdf" in text + assert "AAAA" not in text + assert "[image/png content omitted: ~6 KB]" in text + + # โ”€โ”€ review-round regressions โ”€โ”€ def test_synth_optional_no_default_param_is_nullable(): @@ -1308,6 +1326,38 @@ def call_tool(self, name, args): assert type(None) in get_args(fn.__annotations__["query"]) +def test_synth_array_params_keep_their_item_type(): + """The schemaโ†’annotation round-trip flattened arrays to bare `list`; + function_tool then emits {"type": "array", "items": {}}, which strict + function calling rejects.""" + from typing import Optional + from pageindex.agent_tools import _make_bridge_function + + class _Bridge: + def call_tool(self, name, args): + return json.dumps(args), False + + meta = {"name": "remove_documents", "description": "d", + "inputSchema": { + "type": "object", + "properties": { + "doc_ids": {"type": "array", "items": {"type": "string"}}, + "tags": {"anyOf": [{"type": "array", + "items": {"type": "integer"}}, + {"type": "null"}]}, + "mixed": {"type": "array", + "items": {"type": ["string", "null"]}}, + }, + "required": ["doc_ids", "mixed"], + }} + fn = _make_bridge_function(_Bridge(), meta) + assert fn.__annotations__["doc_ids"] == list[str] + assert fn.__annotations__["tags"] == Optional[list[int]] + # A type-array in items (nullable elements) degrades to bare list โ€” + # it must not crash the build on an unhashable dict key. + assert fn.__annotations__["mixed"] == list + + def test_synth_escape_hatches(): from pageindex.agent_tools import _make_bridge_function @@ -1540,6 +1590,17 @@ def list_documents(self, limit, offset): assert type(exact).calls == 2 # ...total still saves the empty page +def test_null_arguments_mean_omitted(client, store_path): + """Adapters that forward the model's null values verbatim (the Claude + MCP handler) used to trip parameter validation โ€” None โ‰ก omitted is + enforced once, in call_tool.""" + seed_doc(store_path, "pi-a", "report.pdf") + payload, is_error = run(client, "browse_documents", folder_id=None, + sort=None, query=None) + assert not is_error + assert [doc["name"] for doc in payload["documents"]] == ["report.pdf"] + + def test_page_spec_span_bomb_rejected(client, store_path): """An absurd range must be rejected arithmetically, not expanded into billions of integers in the caller's process.""" diff --git a/tests/test_client.py b/tests/test_client.py index 006d015f6..59c90145d 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -151,6 +151,16 @@ def test_get_page_content(local_client, indexed_doc): local_client.get_page_content(indexed_doc, "abc") +def test_get_page_content_span_bomb_rejected(local_client, indexed_doc): + """An absurd range must be rejected arithmetically, not expanded into + a billion integers in the caller's process (the tool layer already + refused; the public client method did not).""" + with pytest.raises(ValueError, match="spans more than 10000"): + local_client.get_page_content(indexed_doc, "1-1000001") + # At the bound itself the spec still parses. + assert local_client.get_page_content(indexed_doc, "5-10004") == [] + + def test_submit_does_not_create_cwd_logs(local_client, sample_pdf, tmp_path, monkeypatch): monkeypatch.chdir(tmp_path) def fake_page_index_main(doc, opt=None, logger=None, page_list=None): diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index 701710145..dbed2b1e4 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -555,6 +555,23 @@ def test_messages_validation(client, fake_anthropic): 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\\]"): From ce1deafdff11b08947a46dfceb4a30f5b94a019e Mon Sep 17 00:00:00 2001 From: Ray Date: Thu, 13 Aug 2026 01:16:02 +0800 Subject: [PATCH 044/137] =?UTF-8?q?fix:=20three=20review=20findings=20?= =?UTF-8?q?=E2=80=94=20protocol=20honesty,=20model=20echo,=20containment?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - responses() promised the Responses protocol ("no translation layer") but _openai_model ignored protocol on the LiteLLM branch: provider- prefixed models silently ran chat.completions under a responses-shaped envelope, and with no transport hook to record status (LitellmModel has no _client.responses) a turn truncated at the output cap reported status "completed". The branch now raises for protocol == "responses" โ€” at agent-build time, before any backend call โ€” naming the routes out: chat_completions(), messages() for Anthropic models, or OPENAI_BASE_URL + a bare/openai/-prefixed name for backends that genuinely speak /responses. Refusal, not emulation: most providers have no /responses endpoint to drive. - chat_completions envelopes echoed retrieve_model verbatim, which carries the SDK's litellm/ routing marker after normalization โ€” a name no provider catalog contains, and a different string than the same model passed per-call. The envelope and every streaming chunk now report the name the provider actually serves; routing and the prompt-cache group key keep the prefixed form. responses() needs no change (post-refusal the prefix cannot reach its envelope), and the user-typed openai/ prefix stays echoed as typed. - _remove_document caught only PageIndexAPIError around the per-doc delete, so a bare OSError (local_store re-raises them) or a transport error (cloud delete_document wraps nothing) escaped mid-batch, discarded the entries for documents already irreversibly deleted, and surfaced as a generic INTERNAL_ERROR envelope inviting a retry โ€” which then reports the destroyed document as not_found. The loop now catches Exception, keeping the per-document results the contract promises. --- pageindex/agent_tools.py | 4 +++- pageindex/client.py | 5 ++++- pageindex/local_chat.py | 24 +++++++++++++++++++----- tests/test_agent_tools.py | 25 +++++++++++++++++++++++++ tests/test_local_chat.py | 33 ++++++++++++++++++++++++++++++++- 5 files changed, 83 insertions(+), 8 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index b93d9f043..3279b3252 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -1102,7 +1102,9 @@ def _remove_document(client, doc_names: list[str], try: client.delete_document(entry["id"]) results.append({"doc_name": doc_name, "status": "deleted"}) - except PageIndexAPIError as exc: + 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") diff --git a/pageindex/client.py b/pageindex/client.py index bb23150c9..5e55665f9 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -441,7 +441,10 @@ def responses( Requires ``pageindex[openai]`` and a backend that supports the Responses API; backends that only speak chat.completions should use - ``chat_completions()``. + ``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 diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index a9ec88f3b..195f7de88 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -211,9 +211,20 @@ def _openai_model(protocol: str, model_name: str): ``litellm//`` (the client's normalized retrieve_model form) and bare ``/`` paths drive the provider through - LiteLLM; an ``openai/`` prefix strips to the OpenAI SDK; bare names go - to the OpenAI SDK as-is.""" + LiteLLM โ€” chat.completions only, so the responses protocol refuses them + instead of silently downgrading; an ``openai/`` prefix strips to the + OpenAI SDK; bare names go to the OpenAI SDK as-is.""" if "/" in model_name and not model_name.startswith("openai/"): + if protocol == "responses": + raise PageIndexAPIError( + f"responses() cannot drive " + f"'{model_name.removeprefix('litellm/')}': provider-prefixed " + "models route through LiteLLM, which speaks chat.completions, " + "not the Responses API. Use chat_completions() (or messages() " + "for Anthropic models), or point OPENAI_BASE_URL at a " + "Responses-capable backend and use a bare or " + "'openai/'-prefixed model name." + ) from agents.extensions.models.litellm_model import LitellmModel return LitellmModel(model_name.removeprefix("litellm/")) from openai import AsyncOpenAI @@ -350,6 +361,9 @@ def run_chat_completions(client, messages, stream: bool = False, block = _doc_block(client, doc_id) items = ([{"role": "user", "content": block}] if block else []) + history model_name = model or client.retrieve_model + # litellm/ is the SDK's routing marker, not a model name โ€” report the + # name the provider actually serves. + reported_model = model_name.removeprefix("litellm/") managed = _managed_instructions(system_texts) agent = _openai_agent(client, "chat", model_name, managed, temperature, None, doc_ids=doc_id) @@ -371,7 +385,7 @@ def run_chat_completions(client, messages, stream: bool = False, "id": f"chatcmpl-{uuid.uuid4().hex}", "object": "chat.completion", "created": int(time.time()), - "model": model_name, + "model": reported_model, "choices": [{ "index": 0, "message": {"role": "assistant", @@ -387,7 +401,7 @@ def run_chat_completions(client, messages, stream: bool = False, def chunk(delta: dict, finish=None) -> dict: return { "id": chat_id, "object": "chat.completion.chunk", - "created": created, "model": model_name, + "created": created, "model": reported_model, "choices": [{"index": 0, "delta": delta, "finish_reason": finish}], } @@ -417,7 +431,7 @@ async def agen(): yield chunk({}, finish="stop") yield { "id": chat_id, "object": "chat.completion.chunk", - "created": created, "model": model_name, "choices": [], + "created": created, "model": reported_model, "choices": [], "usage": _openai_usage(streamed.raw_responses), } diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 1de495264..42498c775 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -461,6 +461,31 @@ def test_remove_document_rejects_non_string_names_before_deleting(client, 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()] diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index dbed2b1e4..33ec4b8bb 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -748,7 +748,7 @@ def test_openai_model_resolves_provider_prefixes(): 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("responses", "anthropic/claude-x") + model = local_chat._openai_model("chat", "anthropic/claude-x") assert isinstance(model, LitellmModel) and model.model == "anthropic/claude-x" model = local_chat._openai_model("chat", "openai/gpt-5.2") assert isinstance(model, OpenAIChatCompletionsModel) @@ -758,6 +758,37 @@ def test_openai_model_resolves_provider_prefixes(): assert str(model.model) == "gpt-5.2" +@needs_agents +def test_responses_refuses_litellm_routed_models(store_path): + """LiteLLM speaks chat.completions, not /responses โ€” the responses + protocol must refuse the silent downgrade, at agent-build time and + before any backend call.""" + for name in ("anthropic/claude-x", "litellm/anthropic/claude-x"): + with pytest.raises(PageIndexAPIError, match="Responses API"): + local_chat._openai_model("responses", name) + client = PageIndexLocalClient(storage_path=store_path, + retrieve_model="anthropic/claude-x") + with pytest.raises(PageIndexAPIError, match="chat_completions"): + client.responses("q") + + +@needs_agents +def test_envelope_model_strips_litellm_routing_prefix(store_path, fake_model): + """litellm/ is the SDK's routing marker, not a model name โ€” the + OpenAI-shaped envelopes must report the model the provider serves.""" + seed_doc(store_path, "pi-a", "report.pdf") + client = PageIndexLocalClient(storage_path=store_path, + retrieve_model="anthropic/claude-x") + assert client.retrieve_model == "litellm/anthropic/claude-x" + fake_model([[_msg_item("ok")]]) + result = client.chat_completions("q") + assert result["model"] == "anthropic/claude-x" + fake_model([[_msg_item("ok")]]) + chunks = list(client.chat_completions("q", stream=True, + stream_metadata=True)) + assert {c["model"] for c in chunks} == {"anthropic/claude-x"} + + @needs_agents def test_record_response_status_captures_last_status(): class _Dumpable: From adeb2095f0e6f8d8b2f3d9afddf5142b3dfacebd Mon Sep 17 00:00:00 2001 From: Ray Date: Thu, 13 Aug 2026 02:55:11 +0800 Subject: [PATCH 045/137] fix: config bundles use the scoped shadow check their tools earned MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 9f67fdd made the three config helpers enforce doc_id at the tool layer but left their instructions on doc_targeting_block's unscoped default, so a bundle refused any doc_id whose name a newer library-wide duplicate shadows โ€” a raise whose message ("the tools address documents by name and would read the newer one") had just become false: the bundle's own tools resolve names inside the allowlist and read the targeted document correctly. chat_completions() accepted the same doc_id via _doc_block's scoped=True. Each helper now computes scope = _local_doc_scope(doc_id) once and derives both slots from it โ€” scoped=scope is not None for the instructions, doc_id=scope for the tools โ€” so the check mode and the tool allowlist come from one fact and cannot drift apart again. build_agent_instructions grows a scoped passthrough; cloud stays on the whole-library check (scope is None there and the tools are genuinely unscoped), and the public agent_instructions() keeps its unscoped default for the same reason. An in-set duplicate still raises โ€” and in that case the message is true on every surface that emits it. --- pageindex/agent_tools.py | 8 ++++---- pageindex/client.py | 22 +++++++++++++++------- tests/test_agent_tools.py | 33 +++++++++++++++++++++++++++++++++ 3 files changed, 52 insertions(+), 11 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index 3279b3252..69bcb25a7 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -1585,8 +1585,8 @@ def doc_targeting_block(client, doc_id, scoped: bool = False) -> Optional[str]: 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`` (the chat surfaces, whose tools resolve names inside the - doc_id allowlist) only a same-name duplicate within the targeted set + ``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 @@ -1625,9 +1625,9 @@ def doc_targeting_block(client, doc_id, scoped: bool = False) -> Optional[str]: ) -def build_agent_instructions(client, doc_id=None) -> str: +def build_agent_instructions(client, doc_id=None, scoped: bool = False) -> str: """Orchestration guidance for document QA agents; with doc_id, appends the target documents and directs the agent to work within them.""" base = _base_instructions(client) - block = doc_targeting_block(client, doc_id) + 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 5e55665f9..b97934e4e 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -690,11 +690,13 @@ def openai_agent_config( library. model: Backend model name; overrides the local default. """ + from .agent_tools import build_agent_instructions + scope = self._local_doc_scope(doc_id) config: dict[str, Any] = { "name": "PageIndex", - "instructions": self.agent_instructions(doc_id=doc_id), - "tools": self.as_openai_tools(include_management, - doc_id=self._local_doc_scope(doc_id)), + "instructions": build_agent_instructions(self, doc_id, + scoped=scope is not None), + "tools": self.as_openai_tools(include_management, doc_id=scope), } model = model or getattr(self, "retrieve_model", None) if model: @@ -788,14 +790,17 @@ def anthropic_runner_config( model. max_turns: Agent-loop bound; default 10. """ + from .agent_tools import build_agent_instructions from .local_chat import _default_max_tokens + scope = self._local_doc_scope(doc_id) return { "model": model, "max_tokens": (max_tokens if max_tokens is not None else _default_max_tokens(model)), - "system": self.agent_instructions(doc_id=doc_id), + "system": build_agent_instructions(self, doc_id, + scoped=scope is not None), "tools": self.as_anthropic_tools(include_management, asynchronous, - doc_id=self._local_doc_scope(doc_id)), + doc_id=scope), "max_iterations": max_turns if max_turns is not None else 10, } @@ -861,10 +866,13 @@ def claude_agent_config( 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": self.agent_instructions(doc_id=doc_id), + "system_prompt": build_agent_instructions(self, doc_id, + scoped=scope is not None), "mcp_servers": {server_name: self.as_claude_mcp( - include_management, doc_id=self._local_doc_scope(doc_id))}, + 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}"], diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 42498c775..36a20186f 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -852,6 +852,39 @@ def test_claude_agent_config_doc_scope_enforced_in_tools(client, store_path): assert payload["errorCode"] == "NOT_FOUND" +def test_openai_agent_config_scoped_shadow_check(client, store_path): + """The bundles' tools resolve names inside the allowlist, so a same-name + document outside the target set must not block โ€” only an in-set + duplicate shadows.""" + pytest.importorskip("agents") + seed_doc(store_path, "pi-old", "report.pdf") + seed_doc(store_path, "pi-new", "report.pdf", + created_at="2026-08-02T10:00:00.123000") + config = client.openai_agent_config(doc_id="pi-old") + assert "report.pdf" in config["instructions"] + with pytest.raises(PageIndexAPIError, match="shadowed"): + client.openai_agent_config(doc_id=["pi-old", "pi-new"]) + + +def test_anthropic_runner_config_scoped_shadow_check(client, store_path): + pytest.importorskip("anthropic") + seed_doc(store_path, "pi-old", "report.pdf") + seed_doc(store_path, "pi-new", "report.pdf", + created_at="2026-08-02T10:00:00.123000") + config = client.anthropic_runner_config(model="claude-sonnet-4-5", + doc_id="pi-old") + assert "report.pdf" in config["system"] + + +def test_claude_agent_config_scoped_shadow_check(client, store_path): + pytest.importorskip("claude_agent_sdk") + seed_doc(store_path, "pi-old", "report.pdf") + seed_doc(store_path, "pi-new", "report.pdf", + created_at="2026-08-02T10:00:00.123000") + config = client.claude_agent_config(doc_id="pi-old") + assert "report.pdf" in config["system_prompt"] + + def test_doc_scope_rejected_on_cloud_openai(): pytest.importorskip("agents") cloud = PageIndexCloudClient(api_key="pi-test-key") From 2b62194170a75503a4e2db268adb8b225f8147c1 Mon Sep 17 00:00:00 2001 From: Ray Date: Thu, 13 Aug 2026 03:04:06 +0800 Subject: [PATCH 046/137] docs: as_openai_tools' remote-MCP note moves to the Cloud paragraph MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The MCPServerStreamableHttp alternative sat in the Local: paragraph pointing at bare {BASE_URL}/mcp โ€” a cloud-only route (BASE_URL is the hosted API; local has no HTTP MCP server) that as written would connect unauthenticated to the full tool set. Now stated where it applies, in the as_anthropic_tools connector-note form: Cloud paragraph, Bearer auth spelled out, ?tools=read default with the drop-it escape. --- pageindex/client.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/pageindex/client.py b/pageindex/client.py index b97934e4e..381f75d6f 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -629,12 +629,15 @@ def as_openai_tools(self, include_management: bool = False, process โ€” works with any model backend. 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). + 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. (The framework's own ``MCPServerStreamableHttp`` - against ``{BASE_URL}/mcp`` is the async-native alternative for - its ``mcp_servers=`` slot.) + not apply. Requires ``openai-agents`` (``pip install 'pageindex[openai]'``), imported only when this method is called. From b997e3b56bdd26789295b5d36943da88ce28fa0c Mon Sep 17 00:00:00 2001 From: Ray Date: Thu, 13 Aug 2026 03:33:24 +0800 Subject: [PATCH 047/137] =?UTF-8?q?fix:=20nine=20review=20findings=20?= =?UTF-8?q?=E2=80=94=20argument=20coercion,=20scope,=20and=20honest=20enve?= =?UTF-8?q?lopes?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - call_tool coerces string booleans per the TOOL_CONTRACT schema ("false"/"no"/"0" read as False, not a truthy 3-minute wait) and survives arguments: null (json.loads("null") reaches the seam as None) - _local_doc_scope raises on an explicitly empty doc_id on cloud: with no tool-layer allowlist there, dropping it silently widened an empty scope to the whole library - both page-spec caps count distinct pages instead of summing parts, so overlapping ranges (a parent section plus its children) within the 10k union pass again as they did in 0.2.9; the per-part arithmetic bound still rejects billion-page specs before materializing anything - _remove_document deduplicates doc_names: a repeated name is one deletion, not a second "failed" row with an internal error string - doc_targeting_block merges the user's metadata tags from the listing (local get_document keeps the 7-key cloud detail wire shape, which carries none) so the block delivers the metadata it promises - _wait_until_ready folds its two raise branches into one that carries the doc_id: a poll that dies no longer discards the handle to an uploaded, billed document - _reported_model strips both routing prefixes (litellm/ and openai/) and responses() now reports it too, instead of echoing a model id the provider never served - _openai_model wraps AsyncOpenAI() construction so a missing backend credential surfaces as PageIndexAPIError like every other gate on the chat surfaces (and builds the client once for both protocols) - _browse_documents advances its cursor by the rows that actually arrived and guards a null/absent total โ€” the same hazards _all_documents already guards โ€” and an empty window ends pagination instead of freezing the cursor --- pageindex/agent_tools.py | 78 +++++++++++++++++++++--------- pageindex/client.py | 37 ++++++++++----- pageindex/local_chat.py | 24 +++++++--- tests/test_agent_tools.py | 99 +++++++++++++++++++++++++++++++++++++++ tests/test_client.py | 8 ++++ tests/test_local_chat.py | 24 ++++++++++ 6 files changed, 227 insertions(+), 43 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index 69bcb25a7..63648d540 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -521,8 +521,20 @@ def _parse_page_spec( ) if not isinstance(pages, str) or not _PAGES_SPEC_RE.match(pages.strip()): return None, invalid + too_many = _failure( + f"Too many pages requested (over {_MAX_REQUESTED_PAGES})", + {"doc_name": doc_name}, + { + "summary": "The page specification spans too many pages", + "options": [ + "Request a narrower page range", + "The response holds only a few pages per call - page " + "through with several smaller requests", + ], + }, + "INVALID_INPUT", + ) expanded: set[int] = set() - requested_total = 0 for part in pages.split(","): part = part.strip() if "-" in part: @@ -531,25 +543,16 @@ def _parse_page_spec( return None, invalid else: start = end = int(part) - # Bound the span arithmetically before materializing it: a spec like + # Bound each part arithmetically before materializing it: a spec like # "1-1000000000" would otherwise expand to billions of integers - # inside the caller's process. - requested_total += end - start + 1 - if requested_total > _MAX_REQUESTED_PAGES: - 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", - ) + # inside the caller's process. The cap is on distinct pages, so + # overlapping parts (a parent section plus its children) don't + # double-count. + if end - start + 1 > _MAX_REQUESTED_PAGES: + return None, too_many expanded.update(range(start, end + 1)) + if len(expanded) > _MAX_REQUESTED_PAGES: + return None, too_many if any(page < 1 for page in expanded): return None, _failure( "Invalid page numbers. Page numbers must be positive integers", @@ -702,12 +705,17 @@ def _browse_documents(client, folder_id: str = "root", recursive: bool = False, if _allowed_ids is None: listing = client.list_documents(limit=limit, offset=offset) window = listing.get("documents") or [] - total = listing.get("total", 0) + total = listing.get("total") else: scoped = _scope_documents(_all_documents(client), _allowed_ids) window, total = scoped[offset:offset + limit], len(scoped) - has_more = offset + limit < total - next_offset = offset + limit if has_more else None + # Advance by what actually arrived โ€” a server may cap its page size โ€” + # and treat an absent/None total like _all_documents does: a full + # window means there may be more. + 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 @@ -1086,6 +1094,8 @@ def _remove_document(client, doc_names: list[str], "options": ["Copy each name verbatim from a browse_documents() " "response"]}, "INVALID_INPUT") + # A repeated name is one deletion, not a second "failed" row. + 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", @@ -1130,6 +1140,17 @@ def tool_names(include_management: bool = False) -> tuple[str, ...]: return _READ_TOOLS + (_MANAGEMENT_TOOLS if include_management else ()) +def _coerce_bool_args(name: str, kwargs: dict[str, Any]) -> None: + """Models routinely send booleans as JSON strings ("false"); the bare + truthiness tests downstream would read those as True.""" + properties = TOOL_CONTRACT.get(name, {}).get("schema", {}).get( + "properties", {}) + for key, spec in properties.items(): + value = kwargs.get(key) + if spec.get("type") == "boolean" and isinstance(value, str): + kwargs[key] = value.strip().lower() not in ("false", "no", "0", "") + + def call_tool(client, name: str, arguments: dict[str, Any], doc_ids=None) -> tuple[str, bool]: """Run one contract tool; returns (envelope_json, is_error). Never raises @@ -1149,8 +1170,9 @@ def call_tool(client, name: str, arguments: dict[str, Any], # 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.items() + kwargs = {key: value for key, value in (arguments or {}).items() if not key.startswith("_") and value is not None} + _coerce_bool_args(name, kwargs) 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) @@ -1594,9 +1616,10 @@ def doc_targeting_block(client, doc_id, scoped: bool = False) -> Optional[str]: if not doc_ids: return None details = [client.get_document(one_id) for one_id in doc_ids] + listing = _all_documents(client) documents = ([{**detail, "id": one_id} for one_id, detail in zip(doc_ids, details)] - if scoped else _all_documents(client)) + if scoped else listing) for one_id, detail in zip(doc_ids, details): entry, _ = _resolve_document(client, str(detail.get("name")), documents=documents) @@ -1608,6 +1631,15 @@ def doc_targeting_block(client, doc_id, scoped: bool = False) -> Optional[str]: "and would read the newer one. Rename or remove the " "duplicate, or pass the newer doc_id." ) + # get_document keeps the cloud detail wire shape, which local mode + # serves without the user's metadata tags; the listing carries them + # in both modes. + 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 ( diff --git a/pageindex/client.py b/pageindex/client.py index 381f75d6f..3221c6f4f 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -10,8 +10,9 @@ def _parse_pages(pages: str) -> list[int]: - result = [] - total = 0 + result: set[int] = set() + too_many = (f"Page specification '{pages}' spans more than " + "10000 pages; request a narrower range") for part in pages.split(","): part = part.strip() if "-" in part: @@ -20,14 +21,16 @@ def _parse_pages(pages: str) -> list[int]: raise ValueError(f"Invalid range '{part}': start must be <= end") else: start = end = int(part) - # Bound the span arithmetically before materializing it โ€” a spec + # Bound each part arithmetically before materializing it โ€” a spec # like "1-999999999" would otherwise expand to a billion integers. - total += end - start + 1 - if total > 10_000: - raise ValueError(f"Page specification '{pages}' spans more than " - "10000 pages; request a narrower range") - result.extend(range(start, end + 1)) - return sorted(set(result)) + # The cap is on distinct pages, so overlapping parts (a parent + # section plus its children) don't double-count. + if end - start + 1 > 10_000: + raise ValueError(too_many) + result.update(range(start, end + 1)) + if len(result) > 10_000: + raise ValueError(too_many) + return sorted(result) def _normalize_retrieve_model(model: str) -> str: @@ -201,10 +204,10 @@ def _wait_until_ready(self, doc_id: str, timeout: float = 1800.0) -> None: # not die on one 502 or dropped connection. poll_failures += 1 if poll_failures >= 3: - if isinstance(exc, PageIndexAPIError): - raise raise PageIndexAPIError( - f"Could not poll document status: {exc}" + 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": @@ -663,7 +666,15 @@ 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.""" - return None if getattr(self, "api_key", None) else doc_id + if not getattr(self, "api_key", None): + return doc_id + if doc_id is not None and not doc_id: + # Cloud has no tool-layer allowlist to make an empty scope mean + # "nothing"; dropping it would silently mean "everything". + raise PageIndexAPIError( + "doc_id is empty. Pass one or more document IDs, or omit " + "doc_id to give the agent the whole library.") + return None def openai_agent_config( self, diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 195f7de88..57e10c742 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -227,14 +227,24 @@ def _openai_model(protocol: str, model_name: str): ) from agents.extensions.models.litellm_model import LitellmModel return LitellmModel(model_name.removeprefix("litellm/")) - from openai import AsyncOpenAI + import openai model_name = model_name.removeprefix("openai/") + try: + backend = openai.AsyncOpenAI() + except openai.OpenAIError as exc: + raise PageIndexAPIError( + f"The OpenAI backend is not configured: {exc}") from exc if protocol == "chat": from agents.models.openai_chatcompletions import ( OpenAIChatCompletionsModel) - return OpenAIChatCompletionsModel(model_name, AsyncOpenAI()) + return OpenAIChatCompletionsModel(model_name, backend) from agents.models.openai_responses import OpenAIResponsesModel - return OpenAIResponsesModel(model_name, openai_client=AsyncOpenAI()) + return OpenAIResponsesModel(model_name, openai_client=backend) + + +def _reported_model(model_name: str) -> str: + """The name the provider actually serves โ€” routing prefixes stripped.""" + return model_name.removeprefix("litellm/").removeprefix("openai/") def _openai_agent(client, protocol: str, model_name: str, instructions: str, @@ -361,9 +371,9 @@ def run_chat_completions(client, messages, stream: bool = False, block = _doc_block(client, doc_id) items = ([{"role": "user", "content": block}] if block else []) + history model_name = model or client.retrieve_model - # litellm/ is the SDK's routing marker, not a model name โ€” report the - # name the provider actually serves. - reported_model = model_name.removeprefix("litellm/") + # litellm/ and openai/ are the SDK's routing markers, not model names โ€” + # report the name the provider actually serves. + reported_model = _reported_model(model_name) managed = _managed_instructions(system_texts) agent = _openai_agent(client, "chat", model_name, managed, temperature, None, doc_ids=doc_id) @@ -483,7 +493,7 @@ def envelope(output: list, raw_responses) -> dict: "id": f"resp_{uuid.uuid4().hex}", "object": "response", "created_at": int(time.time()), - "model": model_name, + "model": _reported_model(model_name), "status": recorded.get("status") or "completed", "output": output, "usage": {"input_tokens": usage["prompt_tokens"], diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 36a20186f..ef43525d9 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -11,6 +11,7 @@ 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 ( @@ -2025,3 +2026,101 @@ def test_submit_warns_when_stored_name_differs(fake_cloud_client): with pytest.warns(UserWarning, match='stored as "whatever_1.pdf"'): result = cloud.submit_document("docs/whatever.pdf") assert result["name"] == "whatever_1.pdf" + + +def test_submit_wait_poll_error_carries_doc_id(fake_cloud_client, monkeypatch): + """A poll that dies on transient errors must keep the uploaded doc_id + recoverable, like the timeout and failed branches do.""" + cloud = fake_cloud_client(["processing"]) + + def boom(doc_id): + raise PageIndexAPIError("Failed to get document metadata: 502") + + monkeypatch.setattr(cloud, "get_document", boom) + with pytest.raises(PageIndexAPIError, match="pi-fake"): + cloud.submit_document("whatever.pdf", wait=True) + + +def test_config_helpers_reject_empty_doc_id_on_cloud(): + """An explicitly empty scope must not silently widen to the whole + library โ€” cloud has no tool-layer allowlist to enforce it.""" + cloud = PageIndexCloudClient(api_key="pi-test-key") + with pytest.raises(PageIndexAPIError, match="doc_id is empty"): + cloud.openai_agent_config(doc_id=[]) + with pytest.raises(PageIndexAPIError, match="doc_id is empty"): + cloud.anthropic_runner_config(model="claude-sonnet-4-5", doc_id=[]) + with pytest.raises(PageIndexAPIError, match="doc_id is empty"): + cloud.claude_agent_config(doc_id=[]) + + +def test_call_tool_coerces_string_booleans(client, store_path, monkeypatch): + """Models routinely send booleans as JSON strings โ€” "false" must not + read as True (a full wait_for_completion stall).""" + seed_doc(store_path, "pi-1", "a.pdf") + seen = {} + real = agent_tools_module._await_completion + + def spy(spy_client, entry, wait): + seen["wait"] = wait + return real(spy_client, entry, wait) + + monkeypatch.setattr(agent_tools_module, "_await_completion", spy) + run(client, "get_document", doc_name="a.pdf", wait_for_completion="false") + assert seen["wait"] is False + run(client, "get_document", doc_name="a.pdf", wait_for_completion="true") + assert seen["wait"] is True + + +def test_remove_document_repeated_name_deletes_once(client, store_path): + seed_doc(store_path, "pi-1", "a.pdf") + payload, is_error = run(client, "remove_document", + doc_names=["a.pdf", "a.pdf"]) + assert not is_error + assert payload["results"] == [{"doc_name": "a.pdf", "status": "deleted"}] + assert "1 of 1" in payload["next_steps"]["summary"] + + +def test_page_spec_cap_counts_distinct_pages(): + """Overlapping parts are normal tree output (a parent section plus its + children) โ€” the cap is on the union, not the sum.""" + pages, error = agent_tools_module._parse_page_spec("1-5000,2000-9000", + "a.pdf") + assert error is None and pages is not None and len(pages) == 9000 + pages, error = agent_tools_module._parse_page_spec("1-10001", "a.pdf") + assert pages is None and "Too many pages" in error[0]["error"] + + +def test_browse_documents_pages_by_rows_returned(): + """A backend that caps its page size must not make the cursor skip + documents, and a null total must not crash (same guards as + _all_documents).""" + class _Capping: + def list_documents(self, limit, offset): + docs = [{"id": f"pi-{i}", "name": f"d{i}.pdf", + "status": "completed"} + for i in range(offset, min(offset + 5, 30))] + return {"documents": docs, "total": 30} + + payload, is_error = agent_tools_module._browse_documents(_Capping(), + limit=10) + assert not is_error + assert payload["has_more"] is True and payload["next_offset"] == 5 + + class _NullTotal: + def list_documents(self, limit, offset): + return {"documents": [{"name": "d.pdf", "status": "completed"}], + "total": None} + + payload, is_error = agent_tools_module._browse_documents(_NullTotal(), + limit=10) + assert not is_error + assert payload["has_more"] is False and payload["next_offset"] is None + + +def test_agent_instructions_carry_user_metadata(client, store_path): + """The targeting block promises names and metadata; local get_document + keeps the 7-key detail wire shape, so the tags come from the listing.""" + seed_doc(store_path, "pi-1", "report.pdf", + metadata={"quarter": "Q3", "year": 2025}) + text = client.agent_instructions(doc_id="pi-1") + assert '"quarter": "Q3"' in text and '"year": 2025' in text diff --git a/tests/test_client.py b/tests/test_client.py index 59c90145d..60375e0df 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -778,3 +778,11 @@ def test_cloud_chat_accepts_query_string(cloud): {"role": "user", "content": "What status?"}] with pytest.raises(PageIndexAPIError, match="non-empty string"): client.chat_completions(" ") + + +def test_parse_pages_overlap_counts_union(): + from pageindex.client import _parse_pages + pages = _parse_pages("1-5000,2000-9000") + assert len(pages) == 9000 and pages[0] == 1 and pages[-1] == 9000 + with pytest.raises(ValueError, match="spans more than"): + _parse_pages("1-10001") diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index 33ec4b8bb..023b9fa10 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -789,6 +789,30 @@ def test_envelope_model_strips_litellm_routing_prefix(store_path, fake_model): assert {c["model"] for c in chunks} == {"anthropic/claude-x"} +@needs_agents +def test_envelope_model_strips_openai_routing_prefix(store_path, fake_model): + """openai/ is the other routing marker โ€” both OpenAI-shaped envelopes + must report the name the provider actually serves.""" + seed_doc(store_path, "pi-a", "report.pdf") + client = PageIndexLocalClient(storage_path=store_path, + retrieve_model="openai/gpt-5.2") + fake_model([[_msg_item("ok")]]) + result = client.chat_completions("q") + assert result["model"] == "gpt-5.2" + fake_model([[_msg_item("ok")]]) + result = client.responses("q") + assert result["model"] == "gpt-5.2" + + +@needs_agents +def test_chat_missing_openai_key_fails_loud(monkeypatch): + """A missing backend credential surfaces as the SDK's own error type, + like every other precondition on the chat surfaces.""" + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + with pytest.raises(PageIndexAPIError, match="OPENAI_API_KEY"): + local_chat._openai_model("chat", "gpt-4o") + + @needs_agents def test_record_response_status_captures_last_status(): class _Dumpable: From bcd0bf4bb84dee33536a48a01c2af608604e7a95 Mon Sep 17 00:00:00 2001 From: Ray Date: Thu, 13 Aug 2026 03:58:09 +0800 Subject: [PATCH 048/137] =?UTF-8?q?fix:=20two=20chat=20findings=20?= =?UTF-8?q?=E2=80=94=20protocol=20terminal=20states,=20provider=20error=20?= =?UTF-8?q?types?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - responses(stream=True) raised PageIndexAPIError when the backend ended the response with response.failed / response.incomplete: openai-agents yields the terminal lifecycle event, then re-raises it as ModelBehaviorError, so the generic AgentsException wrap short-circuited the emit the agen's tail was built for โ€” its failed/incomplete terminal mapping was dead code against the real engine, and the caller lost both the partial output and the real status. The wrap now steps aside when the recorded terminal state is failed/incomplete, and the stream ends with the honest terminal event (committed output, real status, error/incomplete_details) โ€” the backend's terminal state is a protocol event, not an engine failure. Non-stream was already honest for incomplete via the transport recorder; a failed response arrives there as an HTTP error, covered below. Known ceiling: the truncated final turn's partial text was already streamed as deltas but is not reconstructed into the terminal event's output (the engine commits items only on turn completion). - Provider exceptions (network, auth, rate limit) leaked as raw openai/anthropic types through every chat surface, against the layer's own "never raw engine types" contract. Every engine boundary now wraps its vendor's base exception into PageIndexAPIError (chained): the four OpenAI-engine sites catch openai.OpenAIError โ€” LiteLLM's exception types subclass openai's, so one handler covers both routing paths โ€” and messages() catches anthropic.AnthropicError around the batch drive and the stream generator. --- pageindex/local_chat.py | 38 ++++++++++++-- tests/test_local_chat.py | 111 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 144 insertions(+), 5 deletions(-) diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 57e10c742..124b57b3b 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -380,6 +380,7 @@ def run_chat_completions(client, messages, stream: bool = False, run_kwargs = _run_kwargs(max_turns, _conversation_group_id(model_name, managed, history)) + import openai from agents import Runner from agents.exceptions import AgentsException, MaxTurnsExceeded if not stream: @@ -391,6 +392,9 @@ def run_chat_completions(client, messages, stream: bool = False, except AgentsException as exc: raise PageIndexAPIError( f"The agent backend failed: {exc}") from exc + except openai.OpenAIError as exc: + raise PageIndexAPIError( + f"The model backend failed: {exc}") from exc return { "id": f"chatcmpl-{uuid.uuid4().hex}", "object": "chat.completion", @@ -434,6 +438,9 @@ async def agen(): except AgentsException as exc: raise PageIndexAPIError( f"The agent backend failed: {exc}") from exc + except openai.OpenAIError as exc: + raise PageIndexAPIError( + f"The model backend failed: {exc}") from exc finally: if not completed and hasattr(streamed, "cancel"): streamed.cancel() # abandoned/failed: stop the agent task @@ -484,6 +491,7 @@ def run_responses(client, input, model: Optional[str] = None, _conversation_group_id(model_name, managed, conversation)) recorded: dict = {} + import openai from agents import Runner from agents.exceptions import AgentsException, MaxTurnsExceeded @@ -526,6 +534,9 @@ def envelope(output: list, raw_responses) -> dict: 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 output = result.to_input_list()[len(items):] return envelope(output, result.raw_responses) @@ -585,8 +596,16 @@ async def agen(): except MaxTurnsExceeded as exc: raise _wrap_max_turns(exc, 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 + # response.failed / response.incomplete: the engine re-raises + # the backend's terminal state as an exception โ€” it is a + # protocol event, emitted as the terminal event below. + completed = True + except openai.OpenAIError as exc: raise PageIndexAPIError( - f"The agent backend failed: {exc}") from exc + f"The model backend failed: {exc}") from exc finally: if not completed and hasattr(streamed, "cancel"): streamed.cancel() # abandoned/failed: stop the agent task @@ -703,6 +722,7 @@ def run_messages(client, messages, model: str, 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}] @@ -731,12 +751,20 @@ def run_messages(client, messages, model: str, if stream: def events() -> Iterator[Any]: - for turn_stream in runner: - for event in turn_stream: - yield event + try: + for turn_stream in runner: + for event in turn_stream: + yield event + except anthropic.AnthropicError as exc: + raise PageIndexAPIError( + f"The model backend failed: {exc}") from exc return events() - turns = [turn for turn in runner] + try: + turns = [turn for turn in runner] + except anthropic.AnthropicError as exc: + raise PageIndexAPIError( + f"The model backend failed: {exc}") from exc if not turns: raise PageIndexAPIError("The model returned no response.") captured: dict = {} diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index 023b9fa10..e04a9dcc6 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -895,6 +895,117 @@ def test_responses_stream_wraps_framework_errors(client, store_path, list(client.responses("q", stream=True)) +class _TerminalModel(FakeModel): + """Engine-faithful backend terminal: openai-agents yields the + response.failed/response.incomplete lifecycle event, then raises.""" + terminal = "incomplete" + + async def stream_response(self, system_instructions, input, + model_settings, tools, output_schema, + handoffs, tracing, **kwargs): + from agents.exceptions import ModelBehaviorError + from openai.types.responses import (Response, ResponseFailedEvent, + ResponseIncompleteEvent, + ResponseTextDeltaEvent) + from openai.types.responses.response import IncompleteDetails + from openai.types.responses.response_error import ResponseError + self._record(system_instructions, input) + yield ResponseTextDeltaEvent( + type="response.output_text.delta", delta="partial ", + content_index=0, item_id="item_x", output_index=0, + logprobs=[], sequence_number=1) + response = Response( + id="resp_fake", created_at=0.0, model="fake", object="response", + output=[], parallel_tool_calls=False, tool_choice="auto", + tools=[], status=self.terminal, + incomplete_details=(IncompleteDetails(reason="max_output_tokens") + if self.terminal == "incomplete" else None), + error=(ResponseError(code="server_error", message="boom") + if self.terminal == "failed" else None)) + event_type = (ResponseIncompleteEvent if self.terminal == "incomplete" + else ResponseFailedEvent) + yield event_type(type=f"response.{self.terminal}", response=response, + sequence_number=2) + raise ModelBehaviorError(f"terminal: {self.terminal}") + + +@needs_agents +@pytest.mark.parametrize("terminal", ["incomplete", "failed"]) +def test_responses_stream_backend_terminal_states_are_events( + client, store_path, monkeypatch, terminal): + """response.failed / response.incomplete are protocol terminal states, + not engine failures: the stream must end with the honest terminal + event carrying the backend's status, not raise away the run.""" + seed_doc(store_path, "pi-a", "report.pdf") + fake = _TerminalModel([[]]) + fake.terminal = terminal + monkeypatch.setattr(local_chat, "_openai_model", + lambda protocol, model_name: fake) + events = list(client.responses("q", stream=True)) + assert events[0]["type"] == "response.output_text.delta" + last = events[-1] + assert last["type"] == f"response.{terminal}" + assert last["response"]["status"] == terminal + if terminal == "incomplete": + assert (last["response"]["incomplete_details"] + == {"reason": "max_output_tokens"}) + else: + assert last["response"]["error"]["message"] == "boom" + numbers = [event["sequence_number"] for event in events] + assert numbers == sorted(numbers) and len(set(numbers)) == len(numbers) + + +@needs_agents +def test_provider_errors_wrap_as_sdk_errors(client, store_path, fake_model, + monkeypatch): + """Raw provider exceptions (network, auth, rate limit) surface as + PageIndexAPIError on every OpenAI-engine path, never as openai types.""" + import openai + seed_doc(store_path, "pi-a", "report.pdf") + request = httpx.Request("POST", "https://backend.test") + + async def conn_err(*args, **kwargs): + raise openai.APIConnectionError(request=request) + + async def conn_err_stream(*args, **kwargs): + raise openai.APIConnectionError(request=request) + yield # unreached: makes this an async generator + + fake = fake_model([[_msg_item("x")], [_msg_item("x")]]) + monkeypatch.setattr(fake, "get_response", conn_err) + with pytest.raises(PageIndexAPIError, match="model backend failed"): + client.chat_completions("q") + with pytest.raises(PageIndexAPIError, match="model backend failed"): + client.responses("q") + monkeypatch.setattr(fake, "stream_response", conn_err_stream) + with pytest.raises(PageIndexAPIError, match="model backend failed"): + list(client.chat_completions("q", stream=True)) + with pytest.raises(PageIndexAPIError, match="model backend failed"): + list(client.responses("q", stream=True)) + + +@needs_anthropic +def test_messages_provider_errors_wrap_as_sdk_errors(client, store_path, + monkeypatch): + """Anthropic transport errors surface as PageIndexAPIError on both + messages() paths, never as anthropic types.""" + seed_doc(store_path, "pi-a", "report.pdf") + + def handler(request): + return httpx.Response(429, json={ + "type": "error", + "error": {"type": "rate_limit_error", "message": "slow down"}}) + + fake = anthropic.Anthropic( + api_key="test", max_retries=0, + http_client=httpx.Client(transport=httpx.MockTransport(handler))) + monkeypatch.setattr(local_chat, "_anthropic_client", lambda: fake) + with pytest.raises(PageIndexAPIError, match="model backend failed"): + client.messages("q", model="claude-test") + with pytest.raises(PageIndexAPIError, match="model backend failed"): + list(client.messages("q", model="claude-test", stream=True)) + + @needs_agents def test_chat_stream_close_at_opening_chunk_cancels_run(client, store_path, fake_model, From bc25f72b787f26405e034e62e5c1ecd6e170504f Mon Sep 17 00:00:00 2001 From: Ray Date: Thu, 13 Aug 2026 04:04:07 +0800 Subject: [PATCH 049/137] fix: guided failure for unknown LiteLLM providers, non-object call_tool args MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - _openai_model pre-checks the first path segment against litellm.provider_list (fail-open if the attribute ever disappears): a HuggingFace repo id like Qwen/Qwen2.5-7B-Instruct on an OpenAI-compatible server now fails at build time with the escape spelled out โ€” 'openai/' plus OPENAI_BASE_URL โ€” instead of at request time inside LiteLLM with "LLM Provider NOT provided". The slash-means-provider routing convention itself is unchanged; the retrieve_model and chat_completions docstrings now document it where they promise "any OpenAI-compatible server works" - call_tool answers a non-dict arguments value (a JSON array or scalar from a misbehaving caller) with the guided INVALID_INPUT envelope instead of raising AttributeError through the agent loop, matching the openai adapter's own non-object guard --- pageindex/agent_tools.py | 10 ++++++++++ pageindex/client.py | 8 +++++++- pageindex/local_chat.py | 20 +++++++++++++++++--- tests/test_agent_tools.py | 9 +++++++++ tests/test_local_chat.py | 10 ++++++++++ 5 files changed, 53 insertions(+), 4 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index 63648d540..717de2f99 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -1167,6 +1167,16 @@ def call_tool(client, name: str, arguments: dict[str, Any], "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). diff --git a/pageindex/client.py b/pageindex/client.py index 3221c6f4f..e105dddce 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -63,6 +63,9 @@ class PageIndexClient: retrieve_model (str, optional): Local mode only โ€” the model the local chat surfaces (``chat_completions``, ``responses``) default to, exposed as ``client.retrieve_model``. + ``provider/model`` names route through LiteLLM; for an + OpenAI-compatible server that itself serves slashed model ids + (vLLM, TGI), prefix ``openai/`` (e.g. ``openai/Qwen/...``). storage_path (str, optional): Local mode only โ€” directory where indexed documents are stored. Defaults to ``./.pageindex``. @@ -362,7 +365,10 @@ def chat_completions( run over the local tools against your own LLM backend's /chat/completions (requires ``pageindex[openai]``; the OpenAI SDK's usual env config โ€” OPENAI_API_KEY, OPENAI_BASE_URL โ€” selects the - backend, so any OpenAI-compatible server works). The non-stream + backend, so any OpenAI-compatible server works; a ``/`` in the + model name means LiteLLM provider routing, so prefix ``openai/`` + when the backend itself serves slashed ids, e.g. + ``openai/Qwen/...`` on vLLM). The non-stream response carries the final answer only; streaming yields the agent's visible text as it is produced, including narration before tool calls. ``finish_reason`` reports loop completion ("stop") โ€” diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 124b57b3b..b4fa46b86 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -212,8 +212,11 @@ def _openai_model(protocol: str, model_name: str): ``litellm//`` (the client's normalized retrieve_model form) and bare ``/`` paths drive the provider through LiteLLM โ€” chat.completions only, so the responses protocol refuses them - instead of silently downgrading; an ``openai/`` prefix strips to the - OpenAI SDK; bare names go to the OpenAI SDK as-is.""" + instead of silently downgrading; a first segment LiteLLM does not know + (a HuggingFace repo id like ``Qwen/...``) is refused with the + ``openai/`` escape instead of failing inside LiteLLM at request time; + an ``openai/`` prefix strips to the OpenAI SDK; bare names go to the + OpenAI SDK as-is.""" if "/" in model_name and not model_name.startswith("openai/"): if protocol == "responses": raise PageIndexAPIError( @@ -226,7 +229,18 @@ def _openai_model(protocol: str, model_name: str): "'openai/'-prefixed model name." ) from agents.extensions.models.litellm_model import LitellmModel - return LitellmModel(model_name.removeprefix("litellm/")) + import litellm + wire = model_name.removeprefix("litellm/") + providers = getattr(litellm, "provider_list", None) + if providers and wire.split("/", 1)[0] not in providers: + raise PageIndexAPIError( + f"'{wire}' routes through LiteLLM, but " + f"'{wire.split('/', 1)[0]}' is not a LiteLLM provider. For an " + "OpenAI-compatible server (vLLM, TGI, Ollama) serving this " + f"model id, use 'openai/{wire}' and point OPENAI_BASE_URL " + "at the server." + ) + return LitellmModel(wire) import openai model_name = model_name.removeprefix("openai/") try: diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index ef43525d9..3b2c2b9ab 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -2071,6 +2071,15 @@ def spy(spy_client, entry, wait): 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", diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index e04a9dcc6..a370deffa 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -758,6 +758,16 @@ def test_openai_model_resolves_provider_prefixes(): 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 From 031d411541c461e9c7d617ac67be0bbd600b6867 Mon Sep 17 00:00:00 2001 From: Ray Date: Thu, 13 Aug 2026 04:42:48 +0800 Subject: [PATCH 050/137] fix: wrap litellm import in PageIndexAPIError when not installed --- pageindex/local_chat.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index b4fa46b86..9c0511626 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -228,8 +228,14 @@ def _openai_model(protocol: str, model_name: str): "Responses-capable backend and use a bare or " "'openai/'-prefixed model name." ) - from agents.extensions.models.litellm_model import LitellmModel - import litellm + try: + from agents.extensions.models.litellm_model import LitellmModel + import litellm + except ImportError: + raise PageIndexAPIError( + f"'{model_name}' routes through LiteLLM, but litellm is not " + "installed. Run: pip install 'litellm>=1.30'" + ) wire = model_name.removeprefix("litellm/") providers = getattr(litellm, "provider_list", None) if providers and wire.split("/", 1)[0] not in providers: From 1c863faea12dfb57ab8b3e4030bd11a0ee57247e Mon Sep 17 00:00:00 2001 From: Ray Date: Thu, 13 Aug 2026 04:49:46 +0800 Subject: [PATCH 051/137] =?UTF-8?q?fix:=20silence=20CodeQL=20findings=20?= =?UTF-8?q?=E2=80=94=20merge=20implicit=20string=20concat,=20drop=20unused?= =?UTF-8?q?=20vars?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pageindex/agent_tools.py | 21 +++++++-------------- tests/test_agent_tools.py | 1 - tests/test_local_chat.py | 2 +- 3 files changed, 8 insertions(+), 16 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index 717de2f99..b8628be45 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -494,8 +494,7 @@ def _folder_unsupported(param: str) -> tuple[dict, bool]: "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)"], + "Folders are available on PageIndex cloud (PageIndexCloudClient with an API key)"], }, "INVALID_INPUT", ) @@ -528,8 +527,7 @@ def _parse_page_spec( "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", + "The response holds only a few pages per call - page through with several smaller requests", ], }, "INVALID_INPUT", @@ -674,8 +672,7 @@ def _browse_documents(client, folder_id: str = "root", recursive: bool = False, "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)"]}, + "Semantic ranking is available on PageIndex cloud (PageIndexCloudClient with an API key)"]}, "INVALID_INPUT", ) if sort == "relevance" or query: @@ -685,12 +682,9 @@ def _browse_documents(client, folder_id: str = "root", recursive: bool = False, "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)"]}, + "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: @@ -895,8 +889,7 @@ def _get_document_structure(client, doc_name: str, { "summary": "Structure not available for this document", "options": [ - "The document may not have been processed correctly or " - "structure extraction may have failed", + "The document may not have been processed correctly or structure extraction may have failed", "Try processing the document again if possible", ], }, diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 3b2c2b9ab..de300a62b 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -1700,7 +1700,6 @@ def flaky(doc_id): 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 - entry = {"id": "pi-a", "name": "broken.pdf", "status": "failed"} payload, is_error = agent_tools_mod._not_ready_error( "broken.pdf", "failed", "structure retrieval", timed_out=False) assert is_error diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index a370deffa..bb162441f 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -281,7 +281,7 @@ def test_cloud_guards(): @needs_agents def test_responses_end_to_end(client, store_path, fake_model): seed_doc(store_path, "pi-a", "report.pdf") - fake = fake_model([ + fake_model([ [_call_item("get_document", {"doc_name": "report.pdf"})], [_msg_item("The answer")], ]) From 15eecee435308dd045f2346250ae9517debb55d3 Mon Sep 17 00:00:00 2001 From: Ray Date: Thu, 13 Aug 2026 05:35:45 +0800 Subject: [PATCH 052/137] =?UTF-8?q?fix:=20two=20external=20review=20findin?= =?UTF-8?q?gs=20=E2=80=94=20init-notification=20race,=20SDK=20floor?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit notifications/initialized moves inside the bridge lock: a concurrent first use could send tools/list between the handshake and the notification, which strict MCP servers reject with a 400 the bridge never replays. Regression test races two threads through a stalled notification window. claude-agent-sdk floor rises to 0.1.53 โ€” below it, string prompts with SDK MCP servers (the documented local-mode flow) hit invisible registration (#597) and a deadlock (#780). --- pageindex/mcp_bridge.py | 16 ++++---- pyproject.toml | 5 ++- tests/test_agent_tools.py | 78 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 90 insertions(+), 9 deletions(-) diff --git a/pageindex/mcp_bridge.py b/pageindex/mcp_bridge.py index 0bcbc01aa..7d8d153c2 100644 --- a/pageindex/mcp_bridge.py +++ b/pageindex/mcp_bridge.py @@ -160,14 +160,14 @@ def _ensure_initialized(self) -> None: _PROTOCOL_VERSION) self._instructions = result.get("instructions") self._initialized = True - session_id = self._session_id - protocol_version = self._protocol_version - try: - self._post({"jsonrpc": "2.0", - "method": "notifications/initialized"}, - session_id, protocol_version) - except PageIndexAPIError: - pass # advisory; a server that required it fails the next request + # 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 โ”€โ”€ diff --git a/pyproject.toml b/pyproject.toml index 530c963d2..e37a54e59 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,7 +38,10 @@ sortedcontainers = ">=2.4.0" regex = ">=2024.0.0" python-dotenv = ">=1.0.0" pyyaml = ">=6.0" -claude-agent-sdk = { version = ">=0.1.0", optional = true } +# 0.1.53 is the first release where string prompts work with SDK MCP +# servers (invisible registration #597, deadlock #780) โ€” the documented +# local-mode flow. +claude-agent-sdk = { version = ">=0.1.53", optional = true } # 0.14.0 is the first release that feeds RunConfig.group_id into the OpenAI # prompt_cache_key; below it the conversation cache group is inert. openai-agents = { version = ">=0.14.0", optional = true } diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index de300a62b..6d54bff0b 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -1348,6 +1348,84 @@ def fake_post(url, json=None, headers=None, timeout=None): assert bridge._session_id == "sess-1" +def test_mcp_bridge_init_notification_bars_concurrent_requests(monkeypatch): + """No thread may send a request between the initialize handshake and + notifications/initialized โ€” strict servers reject such requests with + HTTP 400, which the bridge never replays. The notification's fake + transport stalls to hold that window open; a racing thread would post + its tools/list inside it.""" + import threading + import requests as requests_mod + import pageindex.mcp_bridge as mcp_bridge + from pageindex.mcp_bridge import McpBridge + + events = [] + events_lock = threading.Lock() + in_notification = threading.Event() + + class _Resp: + def __init__(self, status, body=None): + self.status_code = status + self._body = body + self.headers = {"Content-Type": "application/json"} + self.text = json.dumps(body) if body else "" + self.content = self.text.encode("utf-8") + + def json(self): + if self._body is None: + raise ValueError("no body") + return self._body + + def fake_post(url, json=None, headers=None, timeout=None): + method = json.get("method") + with events_lock: + events.append(("start", method)) + if method == "notifications/initialized": + in_notification.set() + time.sleep(0.2) + rid = json.get("id") + if method == "initialize": + resp = _Resp(200, {"jsonrpc": "2.0", "id": rid, + "result": {"protocolVersion": "2025-06-18"}}) + elif method == "notifications/initialized": + resp = _Resp(202) + else: + resp = _Resp(200, {"jsonrpc": "2.0", "id": rid, + "result": {"tools": [], "nextCursor": None}}) + with events_lock: + events.append(("end", method)) + return resp + + monkeypatch.setattr(mcp_bridge, "requests", types.SimpleNamespace( + post=fake_post, RequestException=requests_mod.RequestException)) + bridge = McpBridge("https://api.pageindex.ai/mcp", + {"Authorization": "Bearer k"}) + + errors = [] + + def list_tools(): + try: + bridge.list_tools() + except BaseException as exc: + errors.append(exc) + + first = threading.Thread(target=list_tools) + first.start() + assert in_notification.wait(5), "handshake never reached the notification" + second = threading.Thread(target=list_tools) + second.start() + first.join(5) + second.join(5) + assert not first.is_alive() and not second.is_alive() + assert not errors + + notified = events.index(("end", "notifications/initialized")) + first_list = events.index(("start", "tools/list")) + assert notified < first_list, ( + f"tools/list overtook notifications/initialized: {events}") + assert events.count(("start", "initialize")) == 1 + + def test_mcp_bridge_blob_blocks_become_stubs(): """Non-text content used to be json.dumps'd wholesale, handing the model the raw base64 payload of an image tool's response.""" From ce0dbf04ca5b2425f83d043e907d9471e4825734 Mon Sep 17 00:00:00 2001 From: Ray Date: Thu, 13 Aug 2026 06:09:09 +0800 Subject: [PATCH 053/137] refactor: drop the unused exc parameter from _wrap_max_turns MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The parameter was dead from the moment it was introduced (daac9d2): the body reads only max_turns, and every call site already carries the cause via `raise ... from exc`. The signature implied the helper inspected the engine exception, which it never did. No behavior change โ€” message text and __cause__ chaining verified identical across all four call sites (chat_completions and responses, stream and non-stream). --- pageindex/local_chat.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 9c0511626..ca94a7d1f 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -358,7 +358,7 @@ async def _run_closing(agent, coro): await _aclose_backend(agent) -def _wrap_max_turns(exc, max_turns) -> PageIndexAPIError: +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 " @@ -408,7 +408,7 @@ def run_chat_completions(client, messages, stream: bool = False, result = _run_sync(_run_closing(agent, Runner.run(agent, input=items, **run_kwargs))) except MaxTurnsExceeded as exc: - raise _wrap_max_turns(exc, max_turns) from exc + raise _wrap_max_turns(max_turns) from exc except AgentsException as exc: raise PageIndexAPIError( f"The agent backend failed: {exc}") from exc @@ -454,7 +454,7 @@ async def agen(): yield chunk({"content": event.data.delta}) completed = True except MaxTurnsExceeded as exc: - raise _wrap_max_turns(exc, max_turns) from exc + raise _wrap_max_turns(max_turns) from exc except AgentsException as exc: raise PageIndexAPIError( f"The agent backend failed: {exc}") from exc @@ -550,7 +550,7 @@ def envelope(output: list, raw_responses) -> dict: Runner.run(agent, input=[dict(item) for item in items], **run_kwargs))) except MaxTurnsExceeded as exc: - raise _wrap_max_turns(exc, max_turns) from exc + raise _wrap_max_turns(max_turns) from exc except AgentsException as exc: raise PageIndexAPIError( f"The agent backend failed: {exc}") from exc @@ -614,7 +614,7 @@ async def agen(): output_offset += 1 completed = True except MaxTurnsExceeded as exc: - raise _wrap_max_turns(exc, max_turns) from exc + raise _wrap_max_turns(max_turns) from exc except AgentsException as exc: if recorded.get("status") not in ("failed", "incomplete"): raise PageIndexAPIError( From f58cca181e05784dc59c7949a02f3194b528c009 Mon Sep 17 00:00:00 2001 From: Ray Date: Thu, 13 Aug 2026 06:21:44 +0800 Subject: [PATCH 054/137] fix: raise the anthropic and openai-agents floors past broken releases MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Both declared floors named a version that cannot work, and CI never caught either because it installs the latest. anthropic >=0.84.0 -> >=0.108.0. Probed against a mock transport: on a turn with stop_reason="refusal" carrying a tool_use block, 0.84.0, 0.92.0 and 0.100.0 all execute the tool and post the tool_result back; 0.108.0 and later stop at the refusal. test_messages_refusal_with_ tool_use_stays_appendable asserts the latter, so that test was false at the floor. messages() is unaffected in practice (it never passes include_management, so remove_document is not registered), but as_anthropic_tools(include_management=True) hands it to a caller's own runner. openai-agents >=0.14.0 -> >=0.18.1. 0.14.0 and 0.16.0 raise pydantic ValidationError on InputTokensDetails.cache_write_tokens before any request reaches the transport when paired with openai 2.54.0 โ€” and they declare openai <3,>=2.26.0, so pip resolves exactly that pair. 0.18.1 is clean. The 0.14.0 rationale (RunConfig.group_id -> prompt_cache_key) still holds above the new floor. The three extras' floor comments are cut to the binding constraint; the reasoning lives here. --- pyproject.toml | 17 +++++------------ 1 file changed, 5 insertions(+), 12 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index e37a54e59..9ea877290 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,19 +38,12 @@ sortedcontainers = ">=2.4.0" regex = ">=2024.0.0" python-dotenv = ">=1.0.0" pyyaml = ">=6.0" -# 0.1.53 is the first release where string prompts work with SDK MCP -# servers (invisible registration #597, deadlock #780) โ€” the documented -# local-mode flow. +# Older releases break string prompts with SDK MCP servers (#597, #780). claude-agent-sdk = { version = ">=0.1.53", optional = true } -# 0.14.0 is the first release that feeds RunConfig.group_id into the OpenAI -# prompt_cache_key; below it the conversation cache group is inert. -openai-agents = { version = ">=0.14.0", optional = true } -# messages() and as_anthropic_tools() need the SDK's beta tool runner; -# 0.84.0 is the first release with ToolError (failed tool calls flagged -# is_error) whose runner also executes the final turn's tools on a -# max_iterations cut (0.75.0 ordering) โ€” older runners return truncated -# histories with no tool_result. -anthropic = { version = ">=0.84.0", optional = true } +# Older releases crash on current openai before the request is sent. +openai-agents = { version = ">=0.18.1", optional = true } +# Older releases execute a refusal turn's tool_use blocks. +anthropic = { version = ">=0.108.0", optional = true } [tool.poetry.extras] claude = ["claude-agent-sdk"] From c88f9d0f63e5c1724a843c72ace1cb0c180cbb72 Mon Sep 17 00:00:00 2001 From: Ray Date: Thu, 13 Aug 2026 06:27:15 +0800 Subject: [PATCH 055/137] test: cover max_turns wrapping on every chat surface test_chat_completions_max_turns_wrapped only drove chat_completions, so the two responses() call sites had no coverage, and no test asserted that the engine exception survives as __cause__. Parametrized over both surfaces and both stream modes; the non-positive max_turns rejection splits out, since it is input validation rather than wrapping. --- tests/test_local_chat.py | 28 +++++++++++++++------------- 1 file changed, 15 insertions(+), 13 deletions(-) diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index bb162441f..950122e8c 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -587,26 +587,28 @@ def _anthropic_tool_use(tool_use_id="tu_1"): @needs_agents -def test_chat_completions_max_turns_wrapped(client, store_path, fake_model): +@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 โ€” on both the non-stream and stream paths.""" + 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="max_turns"): - client.chat_completions([{"role": "user", "content": "q"}], - max_turns=1) - 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="max_turns"): - list(client.chat_completions([{"role": "user", "content": "q"}], - stream=True, max_turns=1)) + 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) From eebed64d888c04456569c4bfabe350aca4fe3908 Mon Sep 17 00:00:00 2001 From: Ray Date: Thu, 13 Aug 2026 06:39:48 +0800 Subject: [PATCH 056/137] =?UTF-8?q?fix:=20four=20review=20findings=20?= =?UTF-8?q?=E2=80=94=20envelope=20size=20honesty,=20contained=20tool=20err?= =?UTF-8?q?ors?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - _dumps drops indent=2: emission now matches _serialized_size's compact accounting, so the pagination budget bounds what is actually sent (indented parts measured under 95k but emitted ~1.8x the 100k cap) - call_tool builds the _allowed_ids frozenset inside the guarded block: a non-iterable doc_id returns the INVALID_INPUT envelope instead of raising into the agent loop; same move for _bridge_invoker's arguments normalization - next_steps strings qualify submit_document() as PageIndexClient.submit_document() (three sites), matching the one already-qualified site โ€” it is a client method, not a registered tool - tests: import httpx at module scope (guaranteed via the hard openai dependency) so agents-gated tests survive an install without the anthropic extra; formatting assertion follows the compact envelope --- pageindex/agent_tools.py | 22 +++++++++++++--------- tests/test_agent_tools.py | 2 +- tests/test_local_chat.py | 2 +- 3 files changed, 15 insertions(+), 11 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index b8628be45..98fc5d206 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -323,7 +323,9 @@ def _failure(error: str, details: Optional[dict[str, Any]], def _dumps(payload: dict[str, Any]) -> str: - return json.dumps(payload, indent=2, ensure_ascii=False) + # Compact, matching _serialized_size โ€” so the size budget measures + # what is actually emitted. + return json.dumps(payload, ensure_ascii=False) # โ”€โ”€ document listing / name resolution โ”€โ”€ @@ -453,7 +455,8 @@ def _not_ready_error(doc_name: str, status: Any, operation: str, { "summary": "Document processing has failed", "options": [ - "Index the document again with submit_document()", + "Index the document again with " + "PageIndexClient.submit_document()", "Use browse_documents() to work with other documents", ], }, @@ -745,7 +748,8 @@ def _browse_documents(client, folder_id: str = "root", recursive: bool = False, "summary": "Nothing to show", "options": ["Nothing here. Index documents with " "PageIndexClient.submit_document() to get started."], - "auto_retry": "Index a document with submit_document() to get started", + "auto_retry": "Index a document with " + "PageIndexClient.submit_document() to get started", } return _success(data, next_steps) @@ -819,7 +823,7 @@ def _get_document(client, doc_name: str, folder_id: Optional[str] = None, ]) else: suggestions.append("Document processing failed. Index the document " - "again with submit_document().") + "again with PageIndexClient.submit_document().") data: dict[str, Any] = { "name": name, @@ -1176,10 +1180,10 @@ def call_tool(client, name: str, arguments: dict[str, Any], kwargs = {key: value for key, value in (arguments or {}).items() if not key.startswith("_") and value is not None} _coerce_bool_args(name, kwargs) - 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) 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( @@ -1329,9 +1333,9 @@ def _bridge_invoker(bridge, name: str) -> "Callable[[dict], tuple[str, bool]]": semantics) and failures are contained in the error envelope. Returns (envelope_text, is_error), like call_tool.""" def _invoke(arguments: dict[str, Any]) -> tuple[str, bool]: - arguments = {key: value for key, value in arguments.items() - if value is not None} try: + arguments = {key: value for key, value in arguments.items() + if value is not None} return bridge.call_tool(name, arguments) except Exception as exc: payload, _ = _failure( diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 6d54bff0b..5c7f7d368 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -565,7 +565,7 @@ def test_execution_type_error_is_internal_not_invalid_input(client, store_path, 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), indent=2, ensure_ascii=False) + assert text == json.dumps(json.loads(text), ensure_ascii=False) # โ”€โ”€ framework adapters โ”€โ”€ diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index 950122e8c..a215baf71 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -5,6 +5,7 @@ import sys import types +import httpx # via the hard `openai` dependency import pytest import pageindex.local_chat as local_chat @@ -409,7 +410,6 @@ def test_responses_stream_passthrough(client, store_path, fake_model): try: import anthropic - import httpx _HAS_ANTHROPIC = True except ImportError: _HAS_ANTHROPIC = False From 9001ef95d9f5cbd156e4bf46f099855ea12274c8 Mon Sep 17 00:00:00 2001 From: Ray Date: Thu, 13 Aug 2026 13:53:00 +0800 Subject: [PATCH 057/137] =?UTF-8?q?fix:=20conformant=20responses()=20envel?= =?UTF-8?q?ope=20=E2=80=94=20official=20output,=20transcript=20in=20items,?= =?UTF-8?q?=20full=20usage=20details?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - output now carries only model-produced items, so the envelope parses with the official openai SDK types (function_call_output is input vocabulary โ€” the real API never returns it in output) - the full process transcript moves to the new items field; round-trip appends items instead of output (same bytes, so the provider prompt-cache prefix contract is unchanged) - usage aggregates token details across turns (cached_tokens, cache_write_tokens, reasoning_tokens) on both OpenAI surfaces โ€” cache hits are now observable instead of discarded - streaming stops synthesizing the nonstandard tool-output event; every stream event now validates against the official event union, tool results arrive in the terminal envelope's items - tests: two conformance tests pin the contract (non-stream model_validate + per-event stream validation); round-trip prefix tests append items Verified: 267 tests green; live A/B against the real OpenAI API โ€” field-identical to the official hand-rolled flow, round-trip accepted with zero repeat tool calls. --- pageindex/client.py | 20 ++++----- pageindex/local_chat.py | 80 ++++++++++++++++++++++-------------- tests/test_local_chat.py | 87 ++++++++++++++++++++++++++++++---------- 3 files changed, 126 insertions(+), 61 deletions(-) diff --git a/pageindex/client.py b/pageindex/client.py index e105dddce..d6351ad7e 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -442,11 +442,13 @@ def responses( 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), so the ``output`` carries the whole process as - standard items โ€” messages, function calls, and function outputs - (the SDK executes the tools). Append the returned ``output`` to your - next call's ``input`` verbatim to keep provider prompt-cache prefix - continuity and the agent's memory of what it already read. + translation layer). The envelope is official Responses shape โ€” + ``output`` carries the model-produced items and parses with the + openai SDK types โ€” and the whole process transcript (including the + tool outputs the SDK executed) rides in the extra ``items`` field. + Append the returned ``items`` to your next call's ``input`` verbatim + to keep provider prompt-cache prefix continuity and the agent's + memory of what it already read. Requires ``pageindex[openai]`` and a backend that supports the Responses API; backends that only speak chat.completions should use @@ -457,15 +459,15 @@ def responses( Args: input: A user message string, or a list of Responses input items - (round-trip prior ``output`` items here). + (round-trip prior ``items`` here). model: Backend model name (defaults to ``retrieve_model``). stream: Yield Responses stream events as dicts โ€” one logical response per call: per-turn backend lifecycle events are collapsed, sequence numbers are reassigned monotonically, and ``output_index`` is re-based onto the single logical - ``output``; tool outputs are emitted as - ``response.output_item.done`` events and the single final - event is the terminal ``response.*`` for the run's status. + ``output``. The single final event is the terminal + ``response.*`` for the run's status; its ``response`` + carries the tool outputs in ``items``. doc_id: Document ID or list of IDs to scope the conversation. Keep it identical across a conversation's calls โ€” the targeting block it adds is re-set each call and is part diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index ca94a7d1f..082752861 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -2,11 +2,11 @@ Three methods, three backend protocols, routed 1:1: ``chat_completions`` drives the backend's /chat/completions (any OpenAI-compatible backend, -final answer only), ``responses`` drives /responses (process items are -standard output; round-trip them for provider prompt-cache continuation and -agent memory), ``messages`` drives Anthropic's /v1/messages via the SDK's -own tool runner (tool_use/tool_result round-trip is the format's native -behavior). +final answer only), ``responses`` drives /responses (official-shape +envelope; the full process transcript rides in ``items`` โ€” round-trip it +for provider prompt-cache continuation and agent memory), ``messages`` +drives Anthropic's /v1/messages via the SDK's own tool runner +(tool_use/tool_result round-trip is the format's native behavior). Content passes through untouched โ€” the caller's messages, the model's answers, tool outputs. Native stop reasons pass through on ``messages``; @@ -366,10 +366,37 @@ def _wrap_max_turns(max_turns) -> PageIndexAPIError: ) +def _usage_sums(raw_responses) -> "tuple[int, int, int, int, int]": + prompt = completion = cached = cache_write = reasoning = 0 + for r in raw_responses: + prompt += r.usage.input_tokens + completion += r.usage.output_tokens + details = getattr(r.usage, "input_tokens_details", None) + cached += getattr(details, "cached_tokens", 0) or 0 + cache_write += getattr(details, "cache_write_tokens", 0) or 0 + details = getattr(r.usage, "output_tokens_details", None) + reasoning += getattr(details, "reasoning_tokens", 0) or 0 + return prompt, completion, cached, cache_write, reasoning + + def _openai_usage(raw_responses) -> dict: - prompt = sum(r.usage.input_tokens for r in raw_responses) - completion = sum(r.usage.output_tokens for r in raw_responses) + """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} @@ -515,18 +542,21 @@ def run_responses(client, input, model: Optional[str] = None, from agents import Runner from agents.exceptions import AgentsException, MaxTurnsExceeded - def envelope(output: list, raw_responses) -> dict: - usage = _openai_usage(raw_responses) + def envelope(transcript: list, raw_responses) -> dict: + # function_call_output is input vocabulary โ€” the official response + # shape does not admit it in ``output``. The conformant ``output`` + # keeps the model-produced items; the full transcript (the + # round-trip payload) rides in ``items``. return { "id": f"resp_{uuid.uuid4().hex}", "object": "response", "created_at": int(time.time()), "model": _reported_model(model_name), "status": recorded.get("status") or "completed", - "output": output, - "usage": {"input_tokens": usage["prompt_tokens"], - "output_tokens": usage["completion_tokens"], - "total_tokens": usage["total_tokens"]}, + "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, @@ -557,8 +587,8 @@ def envelope(output: list, raw_responses) -> dict: except openai.OpenAIError as exc: raise PageIndexAPIError( f"The model backend failed: {exc}") from exc - output = result.to_input_list()[len(items):] - return envelope(output, result.raw_responses) + transcript = result.to_input_list()[len(items):] + return envelope(transcript, result.raw_responses) # One logical response per call: per-turn backend lifecycle events # (created/completed/...) are collapsed โ€” forwarding them verbatim would @@ -576,9 +606,9 @@ async def agen(): # 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 - # โ€” and the tool outputs the SDK injects between turns take the - # next slot on that same axis. + # re-based by the count of items already committed by prior turns. + # Tool outputs are not output items โ€” they ride only in the + # envelope's ``items``. output_offset = 0 completed = False try: @@ -602,16 +632,6 @@ async def agen(): sequence += 1 data["sequence_number"] = sequence yield data - elif (event.type == "run_item_stream_event" - and event.item.type == "tool_call_output_item"): - # We are the tool executor, so we emit the output item - # the way the platform streams its own server-side tools. - sequence += 1 - yield {"type": "response.output_item.done", - "output_index": output_offset, - "sequence_number": sequence, - "item": dict(event.item.to_input_item())} - output_offset += 1 completed = True except MaxTurnsExceeded as exc: raise _wrap_max_turns(max_turns) from exc @@ -630,14 +650,14 @@ async def agen(): if not completed and hasattr(streamed, "cancel"): streamed.cancel() # abandoned/failed: stop the agent task await _aclose_backend(agent) - output = streamed.to_input_list()[len(items):] + 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(output, streamed.raw_responses)} + "response": envelope(transcript, streamed.raw_responses)} return _stream_sync(agen) diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index a215baf71..b05a7475c 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -187,7 +187,10 @@ def test_chat_completions_end_to_end(client, store_path, fake_model): "content": "The answer"} assert result["choices"][0]["finish_reason"] == "stop" assert result["usage"] == {"prompt_tokens": 20, "completion_tokens": 10, - "total_tokens": 30} + "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]) @@ -290,11 +293,18 @@ def test_responses_end_to_end(client, store_path, fake_model): assert result["id"].startswith("resp_") assert result["object"] == "response" assert result["status"] == "completed" - assert result["usage"] == {"input_tokens": 20, "output_tokens": 10, - "total_tokens": 30} + 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" - types = [item.get("type", "message") for item in result["output"]] - assert "function_call" in types and "function_call_output" in types + # Conformant output (model items only); the full transcript in items. + 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]) @@ -312,7 +322,7 @@ def test_responses_round_trip_extends_prefix(client, store_path, fake_model): second = fake_model([[_msg_item("Done")]]) follow_up = ([{"role": "user", "content": "What status?"}] - + result["output"] + + result["items"] + [{"role": "user", "content": "and now?"}]) client.responses(follow_up) previous_final = first.inputs[-1] @@ -332,7 +342,7 @@ def test_responses_round_trip_prefix_with_doc_id(client, store_path, fake_model) second = fake_model([[_msg_item("Done")]]) follow_up = ([{"role": "user", "content": "What status?"}] - + result["output"] + + result["items"] + [{"role": "user", "content": "and now?"}]) client.responses(follow_up, doc_id="pi-a") previous_final = first.inputs[-1] @@ -364,7 +374,7 @@ def spy(max_turns, group_id): fake_model([[_msg_item("c")]]) follow_up = ([{"role": "user", "content": "What is the CAGR?"}] - + result["output"] + + 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 @@ -386,26 +396,63 @@ def test_responses_stream_passthrough(client, store_path, fake_model): events = list(client.responses("q", stream=True)) types = [event.get("type") for event in events] assert "response.output_text.delta" in types - tool_events = [event for event in events - if event.get("type") == "response.output_item.done" - and event.get("item", {}).get("type") - == "function_call_output"] - assert tool_events, types + # Tool outputs are not stream events (official vocabulary only) โ€” they + # arrive in the terminal envelope's items. + 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 - # output_index addresses the logical response.output: the tool output - # slots in after turn 1's item, and turn 2's deltas are re-based past - # both instead of restarting at 0. - assert (final["output"][tool_events[0]["output_index"]]["type"] - == "function_call_output") + assert [item.get("type", "message") for item in final["output"]] == [ + "function_call", "message"] + assert [item.get("type", "message") for item in final["items"]] == [ + "function_call", "function_call_output", "message"] + # output_index addresses the logical response.output: turn 2's deltas + # are re-based past turn 1's item instead of restarting at 0. last_delta = [event for event in events if event.get("type") == "response.output_text.delta"][-1] assert (final["output"][last_delta["output_index"]] .get("type", "message") == "message") +@needs_agents +def test_responses_envelope_validates_as_official_response(client, store_path, + fake_model): + """The conformance contract: the envelope parses with the official + openai SDK types, and the transcript survives in the extension field.""" + from openai.types.responses import Response + seed_doc(store_path, "pi-a", "report.pdf") + fake_model([ + [_call_item("get_document", {"doc_name": "report.pdf"})], + [_msg_item("The answer")], + ]) + result = client.responses("What status?") + parsed = Response.model_validate(result) + assert [item.type for item in parsed.output] == ["function_call", + "message"] + assert parsed.model_dump()["items"] == result["items"] + + +@needs_agents +def test_responses_stream_events_validate_as_official_events( + client, store_path, fake_model): + """Every stream event, terminal envelope included, parses with the + official event union.""" + from pydantic import TypeAdapter + from openai.types.responses import ResponseStreamEvent + seed_doc(store_path, "pi-a", "report.pdf") + fake_model([ + [_call_item("get_document", {"doc_name": "report.pdf"})], + [_msg_item("The answer")], + ]) + adapter = TypeAdapter(ResponseStreamEvent) + events = list(client.responses("q", stream=True)) + assert events + for event in events: + adapter.validate_python(event) + + # โ”€โ”€ messages (Anthropic engine) โ”€โ”€ try: @@ -649,10 +696,6 @@ def test_responses_stream_single_completed_monotonic_sequence( if "sequence_number" in event] assert sequences == sorted(sequences) assert len(set(sequences)) == len(sequences) - tool_done = next(event for event in events - if event.get("type") == "response.output_item.done" - and event["item"]["type"] == "function_call_output") - assert "sequence_number" in tool_done and "output_index" in tool_done @needs_agents From 48889862257ade1f64306296e952efec3fc008cc Mon Sep 17 00:00:00 2001 From: Ray Date: Thu, 13 Aug 2026 13:57:45 +0800 Subject: [PATCH 058/137] =?UTF-8?q?fix:=20stale=20anthropic>=3D0.84.0=20hi?= =?UTF-8?q?nts=20=E2=80=94=20the=20supported=20floor=20is=200.108.0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit pyproject raised the floor in f58cca1 (0.84-0.107 execute a refusal turn's tool_use blocks); the three user-facing strings still pointed hand-installers at the broken range. --- pageindex/client.py | 2 +- pageindex/integrations/anthropic_sdk.py | 2 +- pageindex/local_chat.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pageindex/client.py b/pageindex/client.py index d6351ad7e..4fbc9ad8a 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -750,7 +750,7 @@ def as_anthropic_tools(self, include_management: bool = False, tools involved. Local: the in-process tools โ€” the same set ``messages()`` runs internally. - Requires ``anthropic>=0.84.0`` + Requires ``anthropic>=0.108.0`` (``pip install 'pageindex[anthropic]'``), imported only when this method is called. diff --git a/pageindex/integrations/anthropic_sdk.py b/pageindex/integrations/anthropic_sdk.py index 4405df5a1..089b0809f 100644 --- a/pageindex/integrations/anthropic_sdk.py +++ b/pageindex/integrations/anthropic_sdk.py @@ -23,7 +23,7 @@ def build_anthropic_tools(client, include_management: bool = False, except ImportError as exc: raise PageIndexAPIError( "as_anthropic_tools requires the Anthropic SDK tool runner " - "(anthropic>=0.84.0) โ€” pip install -U anthropic (or pip install " + "(anthropic>=0.108.0) โ€” pip install -U anthropic (or pip install " "'pageindex[anthropic]')." ) from exc from ..agent_tools import _tool_specs diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 082752861..91fba3ae3 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -677,7 +677,7 @@ def _require_anthropic() -> None: from anthropic.lib.tools import ToolError # noqa: F401 except ImportError as exc: raise PageIndexAPIError( - "messages in local mode requires anthropic >= 0.84.0 (the tool " + "messages in local mode requires anthropic >= 0.108.0 (the tool " "runner with ToolError) โ€” pip install -U anthropic." ) from exc From f521fe7446de0e6839b592c04e9755bdbb8eebda Mon Sep 17 00:00:00 2001 From: Ray Date: Thu, 13 Aug 2026 14:00:04 +0800 Subject: [PATCH 059/137] docs: disclose the bridge's binary-stub behavior on the two image-advertising tool surfaces as_openai_tools / as_anthropic_tools cloud docstrings advertised the image tool without mentioning that the in-process bridge replaces base64 payloads with text placeholder stubs (mcp_bridge call_tool). --- pageindex/client.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/pageindex/client.py b/pageindex/client.py index 4fbc9ad8a..eebd4e7bc 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -637,7 +637,9 @@ def as_openai_tools(self, include_management: bool = False, 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. Pass ``hosted=True`` to + 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 @@ -742,7 +744,9 @@ def as_anthropic_tools(self, include_management: bool = False, 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). The server-side alternative is the Messages API's beta + 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 From 56c28b75ba3bb7fffd286b6993cdf35064ee2442 Mon Sep 17 00:00:00 2001 From: Ray Date: Thu, 13 Aug 2026 14:01:16 +0800 Subject: [PATCH 060/137] refactor: trim the envelope-change comments to the essential constraint --- pageindex/local_chat.py | 8 ++------ tests/test_local_chat.py | 4 +--- 2 files changed, 3 insertions(+), 9 deletions(-) diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 91fba3ae3..4ccb379b0 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -543,10 +543,8 @@ def run_responses(client, input, model: Optional[str] = None, from agents.exceptions import AgentsException, MaxTurnsExceeded def envelope(transcript: list, raw_responses) -> dict: - # function_call_output is input vocabulary โ€” the official response - # shape does not admit it in ``output``. The conformant ``output`` - # keeps the model-produced items; the full transcript (the - # round-trip payload) rides in ``items``. + # The official output shape admits no function_call_output; the + # round-trip transcript rides in items. return { "id": f"resp_{uuid.uuid4().hex}", "object": "response", @@ -607,8 +605,6 @@ async def agen(): # 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. - # Tool outputs are not output items โ€” they ride only in the - # envelope's ``items``. output_offset = 0 completed = False try: diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index b05a7475c..1022f83e2 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -300,7 +300,6 @@ def test_responses_end_to_end(client, store_path, fake_model): "output_tokens_details": {"reasoning_tokens": 0}, "total_tokens": 30} assert fake_model.state["protocols"][0][0] == "responses" - # Conformant output (model items only); the full transcript in items. assert [item.get("type", "message") for item in result["output"]] == [ "function_call", "message"] assert [item.get("type", "message") for item in result["items"]] == [ @@ -396,8 +395,7 @@ def test_responses_stream_passthrough(client, store_path, fake_model): events = list(client.responses("q", stream=True)) types = [event.get("type") for event in events] assert "response.output_text.delta" in types - # Tool outputs are not stream events (official vocabulary only) โ€” they - # arrive in the terminal envelope's items. + # Tool outputs arrive only in the terminal envelope's items. assert not [event for event in events if event.get("item", {}).get("type") == "function_call_output"] assert types[-1] == "response.completed" From fee6890474e5c79445f87daff065645a15004705 Mon Sep 17 00:00:00 2001 From: Ray Date: Thu, 13 Aug 2026 14:09:47 +0800 Subject: [PATCH 061/137] =?UTF-8?q?fix:=20declare=20the=20real=20python=20?= =?UTF-8?q?floor=20=E2=80=94=20>=3D3.10?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit litellm's stable channel (every release satisfying our >=1.84.0 floor) and both agent extras require 3.10; on 3.9 pip resolution fails on the hard deps (verified in a clean venv โ€” zero packages install). A clean 3.10 venv with all three extras runs the full suite green. CI already tests 3.10/3.13 only. The >=3.7 claim was inherited from the two-dep 0.2.8 client and was already unsatisfiable then (openai>=1.70 needs 3.8). Closes recurring review finding #10. --- pyproject.toml | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 9ea877290..e65a989a6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,9 +10,6 @@ classifiers = [ "Intended Audience :: Developers", "License :: OSI Approved :: MIT License", "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.7", - "Programming Language :: Python :: 3.8", - "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", @@ -28,7 +25,8 @@ include = [ exclude = ["pageindex/flash/assets"] [tool.poetry.dependencies] -python = ">=3.7" +# litellm's stable channel and both agent extras require 3.10. +python = ">=3.10" requests = ">=2.28.0" openai = ">=1.70.0" litellm = ">=1.84.0" From 525caa45b750e580075f41f94d8bebdc52ae11dd Mon Sep 17 00:00:00 2001 From: Ray Date: Thu, 13 Aug 2026 14:21:48 +0800 Subject: [PATCH 062/137] chore: trim rationale comments from this session's commits --- pageindex/local_chat.py | 8 ++------ pyproject.toml | 1 - tests/test_local_chat.py | 1 - 3 files changed, 2 insertions(+), 8 deletions(-) diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 4ccb379b0..434bdf5ae 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -303,8 +303,7 @@ def _conversation_group_id(model_name: str, instructions: str, items) -> str: def _run_kwargs(max_turns, group_id: str) -> dict: - # Managed runs never export traces โ€” the caller opted into document QA, - # not telemetry. + # No traces โ€” the caller opted into QA, not telemetry. from agents import RunConfig kwargs: dict = {"run_config": RunConfig(tracing_disabled=True, group_id=group_id)} @@ -418,8 +417,7 @@ def run_chat_completions(client, messages, stream: bool = False, block = _doc_block(client, doc_id) items = ([{"role": "user", "content": block}] if block else []) + history model_name = model or client.retrieve_model - # litellm/ and openai/ are the SDK's routing markers, not model names โ€” - # report the name the provider actually serves. + # Strip routing prefixes โ€” report the name the provider serves. reported_model = _reported_model(model_name) managed = _managed_instructions(system_texts) agent = _openai_agent(client, "chat", model_name, managed, @@ -543,8 +541,6 @@ def run_responses(client, input, model: Optional[str] = None, from agents.exceptions import AgentsException, MaxTurnsExceeded def envelope(transcript: list, raw_responses) -> dict: - # The official output shape admits no function_call_output; the - # round-trip transcript rides in items. return { "id": f"resp_{uuid.uuid4().hex}", "object": "response", diff --git a/pyproject.toml b/pyproject.toml index e65a989a6..c316451b6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,6 @@ include = [ exclude = ["pageindex/flash/assets"] [tool.poetry.dependencies] -# litellm's stable channel and both agent extras require 3.10. python = ">=3.10" requests = ">=2.28.0" openai = ">=1.70.0" diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index 1022f83e2..37553b27f 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -395,7 +395,6 @@ def test_responses_stream_passthrough(client, store_path, fake_model): events = list(client.responses("q", stream=True)) types = [event.get("type") for event in events] assert "response.output_text.delta" in types - # Tool outputs arrive only in the terminal envelope's items. assert not [event for event in events if event.get("item", {}).get("type") == "function_call_output"] assert types[-1] == "response.completed" From db78209512258772d687efec6d04747aa1f7d89d Mon Sep 17 00:00:00 2001 From: Ray Date: Thu, 13 Aug 2026 14:34:45 +0800 Subject: [PATCH 063/137] chore: trim non-essential comments across the PR 53 comment lines removed: rationale that belongs in commit messages, descriptions restating what adjacent code or function names already show, and cloud-implementation provenance notes. Section headers and constraint comments (protocol invariants, safety guards) kept. --- pageindex/agent_tools.py | 39 ---------------------- pageindex/integrations/claude_agent_sdk.py | 2 -- pageindex/integrations/openai_agents.py | 2 -- pageindex/local_api.py | 2 -- pageindex/local_chat.py | 8 ----- 5 files changed, 53 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index 98fc5d206..e00eb3ca9 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -323,8 +323,6 @@ def _failure(error: str, details: Optional[dict[str, Any]], def _dumps(payload: dict[str, Any]) -> str: - # Compact, matching _serialized_size โ€” so the size budget measures - # what is actually emitted. return json.dumps(payload, ensure_ascii=False) @@ -706,9 +704,6 @@ def _browse_documents(client, folder_id: str = "root", recursive: bool = False, else: scoped = _scope_documents(_all_documents(client), _allowed_ids) window, total = scoped[offset:offset + limit], len(scoped) - # Advance by what actually arrived โ€” a server may cap its page size โ€” - # and treat an absent/None total like _all_documents does: a full - # window means there may be more. window_end = offset + len(window) has_more = bool(window) and (window_end < total if isinstance(total, int) else len(window) == limit) @@ -865,9 +860,6 @@ def _get_document_structure(client, doc_name: str, waited and entry.get("status") != "failed") try: - # Prefer the raw stored tree: its nodes carry start_index/end_index - # like the cloud structure tool, where client.get_tree() drops - # end_index and renames fields. 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: @@ -1044,8 +1036,6 @@ def _get_page_content(client, doc_name: str, pages: str, if out_of_range: options.insert(0, f"Document has {max_page} pages total - request " f"pages 1-{max_page}") - # Additive, not either/or: a call can both truncate for size and have - # out-of-range pages โ€” hiding either would misreport what was returned. if remaining or out_of_range: parts = [f"Retrieved {len(included)} of {len(requested)} " "requested pages."] @@ -1091,7 +1081,6 @@ def _remove_document(client, doc_names: list[str], "options": ["Copy each name verbatim from a browse_documents() " "response"]}, "INVALID_INPUT") - # A repeated name is one deletion, not a second "failed" row. doc_names = list(dict.fromkeys(doc_names)) if len(doc_names) > 10: return _failure("Maximum 10 documents can be deleted at once", None, @@ -1216,23 +1205,6 @@ def _tool_docstring(description: str, properties: dict[str, Any]) -> str: return "\n".join(lines) -# Local guidance layer: schema STRUCTURE stays byte-identical to the cloud -# contract minus the hidden cloud-only parameters, and description strings -# adapt to the local surface the same way AGENT_INSTRUCTIONS does โ€” guidance -# must not teach capabilities (folders, semantic ranking) or tools -# (search_documents, get_document_image) that do not exist here. Guard -# tests pin structure (contract-minus-hidden equality), tool references -# (the dead-reference test), and capability phrases (the per-docstring -# phrase test) โ€” a contract refresh that reintroduces a cloud-only -# reference fails loudly. - -#: Cloud-only parameters hidden from the local surface โ€” strict-schema -#: frameworks make the dead-end calls inexpressible, and lenient framework -#: argument models drop them before the call (degrading to the bare call). -#: The call_tool path still answers folder_id/sort/query with the guided -#: error envelope; recursive is simply accepted (flattening a folderless -#: library is the identity). Plain functions reject unknown parameters at -#: the Python call boundary. _LOCAL_HIDDEN_PARAMS: dict[str, tuple[str, ...]] = { "browse_documents": ("folder_id", "recursive", "sort", "query"), "get_document": ("folder_id",), @@ -1257,7 +1229,6 @@ def _tool_docstring(description: str, properties: dict[str, Any]) -> str: 'Folder browsing and semantic ranking (sort="relevance") are not ' "supported in local mode yet โ€” they work on PageIndex cloud." ), - # The image sentence points at a tool that is not registered locally. "get_page_content": TOOL_CONTRACT["get_page_content"]["description"] .replace(" Embedded image paths in the response feed into " "`get_document_image()`.", ""), @@ -1536,13 +1507,6 @@ def remove_document(doc_names: list[str]) -> str: # โ”€โ”€ agent instructions โ”€โ”€ -# Local subset of the cloud MCP server's initialize instructions (its -# no-folders variant), trimmed to what exists here: the search_documents -# escalation steps, get_document_image, and the shared read-only-folders -# block are removed, and the sort="relevance" guidance is replaced with -# name/description matching (semantic ranking is cloud-side). Cloud -# clients receive the server's live instructions instead โ€” see -# _base_instructions(). _INSTRUCTIONS_HEADER = ( "PageIndex by Vectify AI is a document platform for uploading and " @@ -1638,9 +1602,6 @@ def doc_targeting_block(client, doc_id, scoped: bool = False) -> Optional[str]: "and would read the newer one. Rename or remove the " "duplicate, or pass the newer doc_id." ) - # get_document keeps the cloud detail wire shape, which local mode - # serves without the user's metadata tags; the listing carries them - # in both modes. 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: diff --git a/pageindex/integrations/claude_agent_sdk.py b/pageindex/integrations/claude_agent_sdk.py index f3ab19c9c..8c76cb434 100644 --- a/pageindex/integrations/claude_agent_sdk.py +++ b/pageindex/integrations/claude_agent_sdk.py @@ -16,8 +16,6 @@ def build_claude_mcp(client, include_management: bool = False, doc_ids=None): from ..agent_tools import _require_local_scope - # The cloud branch returns a URL config โ€” reject cloud doc_ids so they - # are never silently dropped. _require_local_scope(client, doc_ids) if getattr(client, "api_key", None): # include_management picks the endpoint โ€” the URL itself is the diff --git a/pageindex/integrations/openai_agents.py b/pageindex/integrations/openai_agents.py index 266f71618..36c062d2f 100644 --- a/pageindex/integrations/openai_agents.py +++ b/pageindex/integrations/openai_agents.py @@ -29,8 +29,6 @@ def build_openai_tools(client, include_management: bool = False, ) from exc from ..agent_tools import (_dumps, _failure, _require_local_scope, _tool_specs) - # The hosted branch returns before _tool_specs โ€” reject cloud doc_ids - # here so they are never silently dropped. _require_local_scope(client, doc_ids) if getattr(client, "api_key", None) and hosted: # include_management picks the endpoint โ€” the URL itself is the diff --git a/pageindex/local_api.py b/pageindex/local_api.py index 9ad909ad0..8b1e6f184 100644 --- a/pageindex/local_api.py +++ b/pageindex/local_api.py @@ -97,8 +97,6 @@ def submit_document( raise PageIndexAPIError( "Failed to submit document: PDF has no content. All pages are blank." ) - # Fail before paying for indexing when _1.._99 are all taken; the - # binding name resolution happens again at save. self._unique_doc_name(os.path.basename(file_path)) try: diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 434bdf5ae..3dddcf18c 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -417,7 +417,6 @@ def run_chat_completions(client, messages, stream: bool = False, block = _doc_block(client, doc_id) items = ([{"role": "user", "content": block}] if block else []) + history model_name = model or client.retrieve_model - # Strip routing prefixes โ€” report the name the provider serves. reported_model = _reported_model(model_name) managed = _managed_instructions(system_texts) agent = _openai_agent(client, "chat", model_name, managed, @@ -584,10 +583,6 @@ def envelope(transcript: list, raw_responses) -> dict: transcript = result.to_input_list()[len(items):] return envelope(transcript, result.raw_responses) - # One logical response per call: per-turn backend lifecycle events - # (created/completed/...) are collapsed โ€” forwarding them verbatim would - # end a canonical consumer at the first turn โ€” and sequence numbers are - # reassigned monotonically across the whole run. lifecycle = {"response.created", "response.in_progress", "response.completed", "response.failed", "response.incomplete", "response.queued"} @@ -631,9 +626,6 @@ async def agen(): if recorded.get("status") not in ("failed", "incomplete"): raise PageIndexAPIError( f"The agent backend failed: {exc}") from exc - # response.failed / response.incomplete: the engine re-raises - # the backend's terminal state as an exception โ€” it is a - # protocol event, emitted as the terminal event below. completed = True except openai.OpenAIError as exc: raise PageIndexAPIError( From 1a9721880ad6f16f36e9fd585f6515bc76986980 Mon Sep 17 00:00:00 2001 From: Ray Date: Thu, 13 Aug 2026 15:17:41 +0800 Subject: [PATCH 064/137] fix: break the phantom exception chain in _run_sync Move asyncio.run(coro) out of the except RuntimeError block so real errors no longer carry a bogus "no running event loop" context in their traceback. --- pageindex/local_chat.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 3dddcf18c..df890ebb6 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -120,6 +120,10 @@ def _run_sync(coro): try: asyncio.get_running_loop() except RuntimeError: + has_loop = False + else: + has_loop = True + if not has_loop: return asyncio.run(coro) with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: return pool.submit(asyncio.run, coro).result() From d3880c6db9cc7f3471f837b111731b1ad81edc23 Mon Sep 17 00:00:00 2001 From: Ray Date: Thu, 13 Aug 2026 18:12:56 +0800 Subject: [PATCH 065/137] feat: Flash with full optimization becomes the default local indexing mode MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every entrance now defaults to Flash with the full optimize pass (deterministic merge, then LLM expand), replacing the standard LLM-built tree as the default: - submit_document(): mode=None now means "flash"; pass mode="standard" for the LLM-built tree. _index_flash runs optimize="full" with the expand model = summary_model, and fails fast with the missing key name(s) via litellm.validate_environment before any work. - page_index_flash(): optimize takes "full" (default) / "merge" / False; True is accepted as "full" for compatibility, unknown values raise instead of silently degrading to merge-only. optimize_expand stays honored for legacy callers. - CLI: --mode {flash,standard} replaces --flash (kept as a hidden compatibility alias that forces flash). --optimize defaults to full in flash mode with an `off` choice; explicitly passing it outside flash still errors. Standard-only tuning flags (--toc-check-pages, --max-*-per-node, --if-add-*) now error in flash mode instead of being silently ignored, mirroring the existing flash-only flag errors. The key pre-check runs only when an LLM will actually be called, so --no-summary --optimize off|merge works keyless. Output drops the _structure_flash suffix โ€” always _structure.json. On the Disney earnings PDF the optimized default is also faster than unoptimized flash (fewer nodes to summarize) and fixes hierarchy mistakes; both modes emit identical schemas end to end. Docs updated to match (mode flag, defaults, LLM usage honesty); tests pin the new defaults: stored mode == "flash", optimize passthrough, and the unknown-optimize rejection. --- README.md | 9 ++--- pageindex/client.py | 18 +++++----- pageindex/flash/README.md | 13 ++++---- pageindex/flash/api.py | 15 +++++++-- pageindex/local_api.py | 16 +++++++-- run_pageindex.py | 70 ++++++++++++++++++++++++--------------- tests/test_client.py | 39 +++++++++++++++++----- 7 files changed, 120 insertions(+), 60 deletions(-) diff --git a/README.md b/README.md index 5ce0ca5e6..27e3084a4 100644 --- a/README.md +++ b/README.md @@ -173,9 +173,10 @@ python3 run_pageindex.py --pdf_path /path/to/your/document.pdf
Optional parameters
-You can customize the processing with additional optional arguments: +You can customize the processing with additional optional arguments (the structure-tuning flags below require --mode standard): ``` +--mode Processing mode: flash (default) or standard --model LLM model to use (default: gpt-4o-2024-11-20) --toc-check-pages Pages to check for table of contents (default: 20) --max-pages-per-node Max pages per node (default: 10) @@ -199,13 +200,13 @@ python3 run_pageindex.py --md_path /path/to/your/document.md
> ### โšก PageIndex Flash *(preview)* -> **PageIndex Flash** ([`pageindex/flash`](pageindex/flash)) generates tree structures from PDFs in seconds. Structure extraction is purely heuristic-based, no LLM needed. LLM is only used to generate node summaries. +> **PageIndex Flash** ([`pageindex/flash`](pageindex/flash)) generates tree structures from PDFs in seconds. Structure extraction is purely heuristic-based, no LLM needed. An LLM is used only for node summaries and the optimization's expansion pass. > > ```bash -> python3 run_pageindex.py --flash --pdf_path /path/to/your/document.pdf +> python3 run_pageindex.py --mode flash --pdf_path /path/to/your/document.pdf > ``` > -> Add `--optimize` to refine the tree structure for more efficient retrieval (with an LLM expansion pass). +> Tree optimization for retrieval (a deterministic merge, then an LLM expansion pass) is on by default; pass `--optimize off` to disable. ## ๐Ÿš€ Agentic Vectorless RAG: An Example diff --git a/pageindex/client.py b/pageindex/client.py index eebd4e7bc..009b7cf41 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -149,19 +149,19 @@ def submit_document( ``wait=True`` to block until the document is ready, or poll ``get_document(doc_id)['status']`` yourself. - 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 + 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 diff --git a/pageindex/flash/README.md b/pageindex/flash/README.md index 99d236181..0d23a3c35 100644 --- a/pageindex/flash/README.md +++ b/pageindex/flash/README.md @@ -11,9 +11,9 @@ an LLM. ```python from pageindex.flash import page_index_flash -tree = page_index_flash("paper.pdf") -tree = page_index_flash("paper.pdf", summary=False) # tree structure only, no LLM -tree = page_index_flash("paper.pdf", optimize=True) # refined tree for retrieval +tree = page_index_flash("paper.pdf") # optimized tree + summaries +tree = page_index_flash("paper.pdf", summary=False, optimize=False) # raw tree only, no LLM +tree = page_index_flash("paper.pdf", optimize="merge") # deterministic merge, no LLM expand ``` Takes a file path or an `io.BytesIO` stream and returns the tree as a dict. @@ -22,12 +22,11 @@ Summaries are on by default and need an LLM API key. ### Command line ```bash -python3 run_pageindex.py --pdf_path document.pdf --flash -python3 run_pageindex.py --pdf_path document.pdf --flash --no-summary # tree structure only, no LLM -python3 run_pageindex.py --pdf_path document.pdf --flash --optimize # refined tree for retrieval +python3 run_pageindex.py --mode flash --pdf_path document.pdf # optimized tree + summaries +python3 run_pageindex.py --mode flash --pdf_path document.pdf --no-summary --optimize off # raw tree only, no LLM ``` -Writes the tree to `results/_structure_flash.json`. +Writes the tree to `results/_structure.json`. ## Output diff --git a/pageindex/flash/api.py b/pageindex/flash/api.py index 74656d162..bf62d9657 100644 --- a/pageindex/flash/api.py +++ b/pageindex/flash/api.py @@ -96,15 +96,24 @@ def _optimize(structure, page_texts, do_expand, model): def page_index_flash(pdf, summary=True, summary_model=None, - optimize=False, optimize_expand=True, + optimize: str | bool = "full", optimize_expand=None, optimize_model=None, summary_concurrency=None, use_embedded_toc=True) -> dict: - """Build a PageIndex tree structure from a PDF using layout statistics, without an LLM. Args: pdf: path to a PDF file (``str`` or ``pathlib.Path``) or an in-memory binary stream (``io.BytesIO``). summary: if True, generate LLM summaries for each node (requires ``summary_model``). summary_model: the LLM model identifier to use for summary generation. optimize: if True, refine the tree for search cost before summaries: a deterministic merge collapses subtrees whose structure does not beat a linear scan, keeping the removed titles on the parent as ``key_items``, then an LLM pass expands oversized sections. Without it the extracted tree is returned unchanged. optimize_expand: if False, run the merge but skip the LLM expansion. optimize_model: the LLM model for expand (defaults to the summary model). summary_concurrency: maximum simultaneous summary model calls; None uses the library default. use_embedded_toc: if True, consume the PDF's embedded bookmarks when trustworthy: deep bookmarks become the frame and the detected sections they lack are grafted back in after noise filtering, coarse ones become the chapter frame with detected nodes re-hung under them (deeper sparse entries are filled in when the page text confirms them, and garbled extracted titles are repaired from the bookmark strings), garbage ones are ignored; adds a ``toc_source`` key to the result. On by default; pass False for the pure detected structure. Returns: dict with keys ``doc_name``, ``doc_title``, ``structure`` (a list of nested ``{"title", "start_index", "end_index", "nodes"}`` dicts; page indexes are 1-based) and ``has_abstract_or_references_section`` (True when a top-level entry is an abstract or references heading). With ``optimize`` an ``optimize`` key reports merge/expand counts and before/after search-cost metrics. """ + """Build a PageIndex tree structure from a PDF using layout statistics, without an LLM. Args: pdf: path to a PDF file (``str`` or ``pathlib.Path``) or an in-memory binary stream (``io.BytesIO``). summary: if True, generate LLM summaries for each node (requires ``summary_model``). summary_model: the LLM model identifier to use for summary generation. optimize: ``"full"`` for merge + LLM expand, ``"merge"`` for deterministic merge only, ``False`` to disable. ``True`` is accepted as ``"full"`` for backward compatibility. optimize_model: the LLM model for expand (defaults to the summary model). summary_concurrency: maximum simultaneous summary model calls; None uses the library default. use_embedded_toc: if True, consume the PDF's embedded bookmarks when trustworthy: deep bookmarks become the frame and the detected sections they lack are grafted back in after noise filtering, coarse ones become the chapter frame with detected nodes re-hung under them (deeper sparse entries are filled in when the page text confirms them, and garbled extracted titles are repaired from the bookmark strings), garbage ones are ignored; adds a ``toc_source`` key to the result. On by default; pass False for the pure detected structure. Returns: dict with keys ``doc_name``, ``doc_title``, ``structure`` (a list of nested ``{"title", "start_index", "end_index", "nodes"}`` dicts; page indexes are 1-based) and ``has_abstract_or_references_section`` (True when a top-level entry is an abstract or references heading). With ``optimize`` an ``optimize`` key reports merge/expand counts and before/after search-cost metrics. """ + if optimize is True: + optimize = "full" + if not optimize: + optimize = False + elif optimize not in ("full", "merge"): + raise ValueError( + f"optimize must be 'full', 'merge', or False, got {optimize!r}") + if optimize_expand is not None and optimize: + optimize = "full" if optimize_expand else "merge" result = extract_toc(_validate_pdf(pdf), use_embedded_toc=use_embedded_toc) structure = result.get("structure", []) if optimize and structure: result["optimize"] = _optimize(structure, result.get("page_texts") or [], - optimize_expand, + optimize == "full", optimize_model or summary_model) if summary and structure: import asyncio diff --git a/pageindex/local_api.py b/pageindex/local_api.py index 8b1e6f184..9a82c48f9 100644 --- a/pageindex/local_api.py +++ b/pageindex/local_api.py @@ -75,8 +75,10 @@ def submit_document( if mode not in (None, "standard", "flash"): raise PageIndexAPIError( f"Failed to submit document: unknown local processing mode {mode!r}. " - "Supported: None or 'standard' for standard indexing, or 'flash'." + "Supported: 'flash' (default) or 'standard'." ) + if mode is None: + mode = "flash" file_path = os.path.abspath(os.path.expanduser(str(file_path))) if not os.path.isfile(file_path): raise FileNotFoundError(f"No such file: {file_path}") @@ -123,7 +125,7 @@ def submit_document( "pageNum": len(page_texts), "folderId": None, "metadata": metadata, - "mode": mode or "standard", + "mode": mode, } pages = [{"page_index": i + 1, "markdown": text} for i, text in enumerate(page_texts)] @@ -180,8 +182,16 @@ def _index_flash(self, file_path: str, page_texts: list[str]) -> tuple[list, str from .flash import page_index_flash from .utils import (add_node_text, create_clean_structure_for_description, generate_doc_description, write_node_id) + import litellm + env = litellm.validate_environment(self._summary_model) + if not env["keys_in_environment"]: + raise PageIndexAPIError( + f"Failed to submit document: missing API key for " + f"{self._summary_model}: {', '.join(env['missing_keys'])}") result = page_index_flash(file_path, summary=True, - summary_model=self._summary_model) + summary_model=self._summary_model, + optimize="full", + optimize_model=self._summary_model) structure = result.get("structure", []) if not structure: raise PageIndexAPIError( diff --git a/run_pageindex.py b/run_pageindex.py index 452f08174..80c01f16f 100644 --- a/run_pageindex.py +++ b/run_pageindex.py @@ -10,15 +10,18 @@ parser = argparse.ArgumentParser(description='Process PDF or Markdown document and generate structure') parser.add_argument('--pdf_path', type=str, help='Path to the PDF file') parser.add_argument('--md_path', type=str, help='Path to the Markdown file') - parser.add_argument('--flash', action='store_true', help='Use PageIndex Flash (with --pdf_path)') + parser.add_argument('--mode', choices=['flash', 'standard'], default='flash', + help='Processing mode (default: flash)') + parser.add_argument('--flash', action='store_true', default=False, + help=argparse.SUPPRESS) parser.add_argument('--embedded-toc', action=argparse.BooleanOptionalAction, default=None, - help='Use the PDF\'s embedded bookmarks when trustworthy (default: on with --flash)') + help='Use the PDF\'s embedded bookmarks when trustworthy (default: on in flash mode)') parser.add_argument('--summary', action=argparse.BooleanOptionalAction, default=None, - help='Generate node summaries with an LLM (default: on with --flash)') - parser.add_argument('--optimize', nargs='?', const='full', choices=['full', 'merge'], + help='Generate node summaries with an LLM (default: on in flash mode)') + parser.add_argument('--optimize', nargs='?', const='full', choices=['full', 'merge', 'off'], default=None, - help='Refine the tree for search cost: a deterministic merge, then an ' - 'LLM expansion pass; pass `merge` to run the merge alone (PDF only)') + help='Refine the tree for search cost (default: full in flash mode). ' + '`merge` for deterministic merge only; `off` to disable') parser.add_argument('--model', type=str, default=None, help='Model to use (overrides config.yaml)') parser.add_argument('--summary-model', type=str, default=None, @@ -48,18 +51,32 @@ parser.add_argument('--summary-token-threshold', type=int, default=200, help='Token threshold for generating summaries (markdown only)') args = parser.parse_args() - + if args.flash: + args.mode = 'flash' + # Validate that exactly one file type is specified if not args.pdf_path and not args.md_path: raise ValueError("Either --pdf_path or --md_path must be specified") if args.pdf_path and args.md_path: raise ValueError("Only one of --pdf_path or --md_path can be specified") - if args.optimize and not (args.pdf_path and args.flash): - raise ValueError("--optimize requires --flash with --pdf_path") - if args.embedded_toc is not None and not (args.pdf_path and args.flash): - raise ValueError("--embedded-toc requires --flash with --pdf_path") - if args.summary is not None and not (args.pdf_path and args.flash): - raise ValueError("--summary requires --flash with --pdf_path") + if args.optimize in ('full', 'merge') and not (args.pdf_path and args.mode == 'flash'): + raise ValueError("--optimize requires Flash mode with --pdf_path") + if args.optimize is None: + args.optimize = 'full' if args.mode == 'flash' else 'off' + if args.embedded_toc is not None and not (args.pdf_path and args.mode == 'flash'): + raise ValueError("--embedded-toc requires Flash mode with --pdf_path") + if args.summary is not None and not (args.pdf_path and args.mode == 'flash'): + raise ValueError("--summary requires Flash mode with --pdf_path") + if args.pdf_path and args.mode == 'flash': + for flag, value in (('--toc-check-pages', args.toc_check_pages), + ('--max-pages-per-node', args.max_pages_per_node), + ('--max-tokens-per-node', args.max_tokens_per_node), + ('--if-add-node-id', args.if_add_node_id), + ('--if-add-node-summary', args.if_add_node_summary), + ('--if-add-doc-description', args.if_add_doc_description), + ('--if-add-node-text', args.if_add_node_text)): + if value is not None: + raise ValueError(f"{flag} is not supported in flash mode; use --mode standard") if args.pdf_path: # Validate PDF file @@ -68,22 +85,23 @@ if not os.path.isfile(args.pdf_path): raise ValueError(f"PDF file not found: {args.pdf_path}") - if args.flash: + if args.mode == 'flash': from pageindex.flash import page_index_flash - if args.optimize == 'full': - from pageindex.tree_optimize import default_model - from pageindex.utils import _is_openai_model - expand_model = args.model or default_model() - if _is_openai_model(expand_model) and not os.getenv("OPENAI_API_KEY"): - raise SystemExit(f"OPENAI_API_KEY is not set (expand model: {expand_model}).") + summary_model = args.summary_model or args.model + will_summarize = args.summary if args.summary is not None else True + if summary_model and (will_summarize or args.optimize == 'full'): + import litellm + env = litellm.validate_environment(summary_model) + if not env["keys_in_environment"]: + raise SystemExit( + f"Missing API key for {summary_model}: {', '.join(env['missing_keys'])}") toc_with_page_number = page_index_flash( args.pdf_path, - optimize=args.optimize is not None, - optimize_expand=args.optimize == 'full', - optimize_model=args.model, - summary_model=args.summary_model or args.model, + optimize=args.optimize if args.optimize != 'off' else False, + optimize_model=summary_model, + summary_model=summary_model, use_embedded_toc=args.embedded_toc if args.embedded_toc is not None else True, - summary=args.summary if args.summary is not None else True, + summary=will_summarize, ) if 'optimize' in toc_with_page_number: o = toc_with_page_number['optimize'] @@ -110,7 +128,7 @@ # Save results pdf_name = os.path.splitext(os.path.basename(args.pdf_path))[0] - suffix = '_structure_flash' if args.flash else '_structure' + suffix = '_structure' output_dir = './results' output_file = f'{output_dir}/{pdf_name}{suffix}.json' os.makedirs(output_dir, exist_ok=True) diff --git a/tests/test_client.py b/tests/test_client.py index 60375e0df..aad4e3979 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -47,7 +47,7 @@ def fake_page_index_main(doc, opt=None, logger=None, page_list=None): "doc_description": "A test document.", "structure": json.loads(json.dumps(STRUCTURE))} monkeypatch.setattr(page_index_module, "page_index_main", fake_page_index_main) - return local_client.submit_document(sample_pdf)["doc_id"] + return local_client.submit_document(sample_pdf, mode="standard")["doc_id"] # โ”€โ”€ constructor โ”€โ”€ @@ -168,7 +168,7 @@ 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() @@ -179,10 +179,10 @@ 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) + 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) + 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"]} @@ -212,7 +212,7 @@ def test_submit_name_exhaustion_rejects_before_indexing( "indexer ran despite name exhaustion"), ) with pytest.raises(PageIndexAPIError, match="Too many files"): - local_client.submit_document(sample_pdf) + local_client.submit_document(sample_pdf, mode="standard") def test_submit_flash(local_client, sample_pdf, monkeypatch): @@ -220,6 +220,8 @@ def test_submit_flash(local_client, sample_pdf, monkeypatch): 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": []}]} @@ -227,13 +229,34 @@ def fake_flash(pdf, summary=True, summary_model=None, **kwargs): monkeypatch.setattr(pageindex.utils, "llm_completion", lambda model, prompt, **kw: "Flash description.") doc_id = local_client.submit_document(sample_pdf, mode="flash")["doc_id"] - assert calls == {"summary": True, "summary_model": local_client.summary_model} + assert calls == {"summary": True, "summary_model": local_client.summary_model, + "optimize": "full", + "optimize_model": local_client.summary_model} root = local_client.get_tree(doc_id)["result"][0] assert root["node_id"] == "0000" assert "Hello page one" in root["text"] assert local_client.get_document(doc_id)["description"] == "Flash description." +def test_submit_defaults_to_flash(local_client, sample_pdf, monkeypatch): + monkeypatch.setattr( + pageindex.flash, "page_index_flash", + lambda pdf, **kwargs: { + "doc_name": "sample.pdf", + "structure": [{"title": "Flash Root", "start_index": 1, + "end_index": 2, "summary": "s", "nodes": []}]}) + monkeypatch.setattr(pageindex.utils, "llm_completion", + lambda model, prompt, **kw: "Flash description.") + doc_id = local_client.submit_document(sample_pdf)["doc_id"] + assert local_client._api._store.get_meta(doc_id)["mode"] == "flash" + + +def test_page_index_flash_rejects_unknown_optimize(): + from pageindex.flash import page_index_flash + with pytest.raises(ValueError, match="optimize must be"): + page_index_flash("never-opened.pdf", optimize="off") + + def test_llm_completion_missing_key_raises_immediately(monkeypatch): import openai monkeypatch.delenv("OPENAI_API_KEY", raising=False) @@ -335,7 +358,7 @@ def test_submit_with_metadata(local_client, sample_pdf, monkeypatch): "doc_name": "sample.pdf", "doc_description": None, "structure": json.loads(json.dumps(STRUCTURE))}) tags = {"project": "alpha", "year": 2026} - doc_id = local_client.submit_document(sample_pdf, metadata=tags)["doc_id"] + doc_id = local_client.submit_document(sample_pdf, mode="standard", metadata=tags)["doc_id"] assert local_client.get_tree(doc_id)["metadata"] == tags assert local_client.get_ocr(doc_id)["metadata"] == tags assert local_client.list_documents()["documents"][0]["metadata"] == tags @@ -486,7 +509,7 @@ 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): with pytest.warns(UserWarning): # same-name resubmit โ†’ stored as sample_1.pdf - second = local_client.submit_document(sample_pdf)["doc_id"] + 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 From 97d4c41e312752680750c096a2b7208968a420bb Mon Sep 17 00:00:00 2001 From: Ray Date: Fri, 14 Aug 2026 18:27:17 +0800 Subject: [PATCH 066/137] =?UTF-8?q?feat:=20chat()=20=E2=80=94=20the=20answ?= =?UTF-8?q?er-out=20front=20door=20over=20chat=5Fcompletions?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pageindex/client.py | 41 +++++++++++++++++++++++++++++++++++-- tests/test_local_chat.py | 44 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 83 insertions(+), 2 deletions(-) diff --git a/pageindex/client.py b/pageindex/client.py index 009b7cf41..394b2c95b 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -4,7 +4,7 @@ import os import time import warnings -from typing import Any, Callable, Iterator, Optional, Union +from typing import Any, Callable, Iterator, Optional, Union, cast from .errors import PageIndexAPIError @@ -345,7 +345,44 @@ def get_retrieval(self, retrieval_id: str) -> dict[str, Any]: "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, + ) -> 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 + ``retrieve_model``). + + 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) + if stream: + return cast(Iterator[str], result) + envelope = cast(dict[str, Any], result) + return envelope["choices"][0]["message"]["content"] or "" def chat_completions( self, diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index 37553b27f..53287f54b 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -280,6 +280,50 @@ def test_cloud_guards(): max_tokens=10) +# โ”€โ”€ 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 + + +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 From 73c6a00d833a33c26878e98d2067027bd16d8a0b Mon Sep 17 00:00:00 2001 From: Ray Date: Fri, 14 Aug 2026 19:35:34 +0800 Subject: [PATCH 067/137] feat: cache-mark the managed prefix on anthropic-routed LiteLLM models --- pageindex/client.py | 3 ++- pageindex/local_chat.py | 16 +++++++++++++++- tests/test_local_chat.py | 15 +++++++++++++++ 3 files changed, 32 insertions(+), 2 deletions(-) diff --git a/pageindex/client.py b/pageindex/client.py index 394b2c95b..760f2f4c5 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -405,7 +405,8 @@ def chat_completions( backend, so any OpenAI-compatible server works; a ``/`` in the model name means LiteLLM provider routing, so prefix ``openai/`` when the backend itself serves slashed ids, e.g. - ``openai/Qwen/...`` on vLLM). The non-stream + ``openai/Qwen/...`` on vLLM; Anthropic-routed 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") โ€” diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index df890ebb6..537507750 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -271,6 +271,19 @@ def _reported_model(model_name: str) -> str: return model_name.removeprefix("litellm/").removeprefix("openai/") +def _cache_extra_args(model_name: str) -> Optional[dict]: + """Anthropic's prompt caching is opt-in per request: on + anthropic-routed LiteLLM models, mark the managed system prefix via + LiteLLM's injection param so the loop's later turns and a + conversation's next calls read it instead of repaying full price.""" + wire = model_name.removeprefix("litellm/") + if ("/" in model_name and not model_name.startswith("openai/") + and wire.split("/", 1)[0] == "anthropic"): + return {"cache_control_injection_points": [ + {"location": "message", "role": "system"}]} + return None + + def _openai_agent(client, protocol: str, model_name: str, instructions: str, temperature, top_p, doc_ids=None): from agents import Agent, ModelSettings @@ -280,7 +293,8 @@ def _openai_agent(client, protocol: str, model_name: str, instructions: str, instructions=instructions, tools=build_openai_tools(client, doc_ids=doc_ids), model=_openai_model(protocol, model_name), - model_settings=ModelSettings(temperature=temperature, top_p=top_p), + model_settings=ModelSettings(temperature=temperature, top_p=top_p, + extra_args=_cache_extra_args(model_name)), ) diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index 53287f54b..f15068498 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -280,6 +280,21 @@ def test_cloud_guards(): 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"}]} + for name in ("anthropic/claude-x", "litellm/anthropic/claude-x"): + 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"): + agent = _openai_agent(client, "chat", name, "sys", None, None) + assert agent.model_settings.extra_args is None + + # โ”€โ”€ chat (front door) โ”€โ”€ @needs_agents From 9f14287e782763bfcfb59db1a21c23c92bc4c442 Mon Sep 17 00:00:00 2001 From: Ray Date: Fri, 14 Aug 2026 20:27:25 +0800 Subject: [PATCH 068/137] refactor: cache predicate asks litellm's own provider resolution --- pageindex/local_chat.py | 16 ++++++++++++---- tests/test_local_chat.py | 4 +++- 2 files changed, 15 insertions(+), 5 deletions(-) diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 537507750..e52f07fcf 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -275,10 +275,18 @@ def _cache_extra_args(model_name: str) -> Optional[dict]: """Anthropic's prompt caching is opt-in per request: on anthropic-routed LiteLLM models, mark the managed system prefix via LiteLLM's injection param so the loop's later turns and a - conversation's next calls read it instead of repaying full price.""" - wire = model_name.removeprefix("litellm/") - if ("/" in model_name and not model_name.startswith("openai/") - and wire.split("/", 1)[0] == "anthropic"): + conversation's next calls read it 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 + _, provider, _, _ = get_llm_provider( + model=model_name.removeprefix("litellm/")) + except Exception: + return None + if provider == "anthropic": return {"cache_control_injection_points": [ {"location": "message", "role": "system"}]} return None diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index f15068498..fed4c1a74 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -290,7 +290,9 @@ def test_anthropic_routed_models_mark_managed_prefix_for_cache( for name in ("anthropic/claude-x", "litellm/anthropic/claude-x"): 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/vertex Claude stay unmarked until verified live + for name in ("gpt-5", "openai/Qwen/x", "litellm/groq/x", + "bedrock/anthropic.claude-v1", "vertex_ai/claude-x"): agent = _openai_agent(client, "chat", name, "sys", None, None) assert agent.model_settings.extra_args is None From 97cfc065cc9a146eb5d3b8f4a86ef536de71d12f Mon Sep 17 00:00:00 2001 From: Ray Date: Fri, 14 Aug 2026 20:50:56 +0800 Subject: [PATCH 069/137] =?UTF-8?q?feat:=20extend=20cache=20marking=20to?= =?UTF-8?q?=20Claude=20on=20Bedrock=20and=20Vertex=20=E2=80=94=20both=20li?= =?UTF-8?q?ve-verified?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pageindex/client.py | 5 +++-- pageindex/local_chat.py | 18 ++++++++++-------- tests/test_local_chat.py | 8 +++++--- 3 files changed, 18 insertions(+), 13 deletions(-) diff --git a/pageindex/client.py b/pageindex/client.py index 760f2f4c5..f79d48445 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -405,8 +405,9 @@ def chat_completions( backend, so any OpenAI-compatible server works; a ``/`` in the model name means LiteLLM provider routing, so prefix ``openai/`` when the backend itself serves slashed ids, e.g. - ``openai/Qwen/...`` on vLLM; Anthropic-routed models get the - managed prompt prefix cache-marked automatically). The non-stream + ``openai/Qwen/...`` on vLLM; LiteLLM-routed Claude models โ€” + Anthropic direct, Bedrock, Vertex โ€” 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") โ€” diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index e52f07fcf..dbcb9b478 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -272,21 +272,23 @@ def _reported_model(model_name: str) -> str: def _cache_extra_args(model_name: str) -> Optional[dict]: - """Anthropic's prompt caching is opt-in per request: on - anthropic-routed LiteLLM models, mark the managed system prefix via - LiteLLM's injection param so the loop's later turns and a - conversation's next calls read it instead of repaying full price. - Provider resolution is LiteLLM's own, so this predicate can never - disagree with where the request actually routes.""" + """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 via LiteLLM's + injection param so the loop's later turns and a conversation's next + calls read it 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 - _, provider, _, _ = get_llm_provider( + model, provider, _, _ = get_llm_provider( model=model_name.removeprefix("litellm/")) except Exception: return None - if provider == "anthropic": + if provider == "anthropic" or (provider in ("bedrock", "vertex_ai") + and "claude" in model.lower()): return {"cache_control_injection_points": [ {"location": "message", "role": "system"}]} return None diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index fed4c1a74..5c9781395 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -287,12 +287,14 @@ def test_anthropic_routed_models_mark_managed_prefix_for_cache( from pageindex.local_chat import _openai_agent marked = {"cache_control_injection_points": [ {"location": "message", "role": "system"}]} - for name in ("anthropic/claude-x", "litellm/anthropic/claude-x"): + 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 - # bedrock/vertex Claude stay unmarked until verified live for name in ("gpt-5", "openai/Qwen/x", "litellm/groq/x", - "bedrock/anthropic.claude-v1", "vertex_ai/claude-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 From 797d54540f8192a519fd148257d9f650ed2e2a9d Mon Sep 17 00:00:00 2001 From: Ray Date: Fri, 14 Aug 2026 22:02:22 +0800 Subject: [PATCH 070/137] fix: point anthropic-extra users at messages(); guard the two silent vendor chains --- pageindex/local_chat.py | 3 ++- tests/test_local_chat.py | 58 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 60 insertions(+), 1 deletion(-) diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index dbcb9b478..60d06acc5 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -206,7 +206,8 @@ def _require_openai_agents(method: str) -> None: except ImportError as exc: raise PageIndexAPIError( f"{method} in local mode requires the OpenAI Agents SDK โ€” " - "pip install openai-agents (or pip install 'pageindex[openai]')." + "pip install openai-agents (or pip install 'pageindex[openai]'). " + "messages() runs on the anthropic extra instead." ) from exc diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index 5c9781395..18a79cbc0 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -299,6 +299,64 @@ def test_anthropic_routed_models_mark_managed_prefix_for_cache( 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"]) + assert result["choices"][0]["message"]["content"] == "ok" + + # โ”€โ”€ chat (front door) โ”€โ”€ @needs_agents From a098dd73e1416d3e60e22e0dc1d35cf367b58dc0 Mon Sep 17 00:00:00 2001 From: Ray Date: Fri, 14 Aug 2026 23:11:08 +0800 Subject: [PATCH 071/137] =?UTF-8?q?docs:=20the=20max-tokens=20table=20is?= =?UTF-8?q?=20a=20closed=20set=20=E2=80=94=20litellm's=20map=20prunes=20EO?= =?UTF-8?q?L'd=20entries,=20so=20it=20cannot=20replace=20it?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pageindex/local_chat.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 60d06acc5..c953df8ec 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -759,7 +759,9 @@ def _anthropic_usage(turns, final_usage: dict) -> dict: def _default_max_tokens(model: str) -> int: """The wire-required per-turn budget when the caller sets none: 8192, - except the claude-3 generation whose output ceiling is 4096.""" + except the claude-3 generation whose output ceiling is 4096 โ€” a + closed historical set (every later model supports >=8192), so the + table needs no new entries and no live capability source.""" return 4096 if model.startswith(_CLAUDE_4096_MODELS) else 8192 From 152ddbe70efb072a84f1540438fd1ab0720416ba Mon Sep 17 00:00:00 2001 From: Ray Date: Fri, 14 Aug 2026 23:16:47 +0800 Subject: [PATCH 072/137] chore: trim the max-tokens docstring to the contract --- pageindex/local_chat.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index c953df8ec..60d06acc5 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -759,9 +759,7 @@ def _anthropic_usage(turns, final_usage: dict) -> dict: def _default_max_tokens(model: str) -> int: """The wire-required per-turn budget when the caller sets none: 8192, - except the claude-3 generation whose output ceiling is 4096 โ€” a - closed historical set (every later model supports >=8192), so the - table needs no new entries and no live capability source.""" + except the claude-3 generation whose output ceiling is 4096.""" return 4096 if model.startswith(_CLAUDE_4096_MODELS) else 8192 From f0598534ff8f1540dfa55e934e41ffe24986597f Mon Sep 17 00:00:00 2001 From: Ray Date: Fri, 14 Aug 2026 23:25:39 +0800 Subject: [PATCH 073/137] =?UTF-8?q?docs:=20state=20the=20local=20text-only?= =?UTF-8?q?=20history=20contract=20=E2=80=94=20cloud=20forwards=20tool=20t?= =?UTF-8?q?urns,=20local=20rejects;=20extra=20fields=20drop?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pageindex/client.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/pageindex/client.py b/pageindex/client.py index f79d48445..5329cfd73 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -419,7 +419,10 @@ def chat_completions( 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. + 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. Keep it identical across a conversation's calls โ€” the From c78c01fc33153a8394ad656dd53ca8fe23a110c0 Mon Sep 17 00:00:00 2001 From: Ray Date: Sat, 15 Aug 2026 02:16:27 +0800 Subject: [PATCH 074/137] =?UTF-8?q?feat:=20openai-agents=20becomes=20a=20b?= =?UTF-8?q?ase=20dependency=20=E2=80=94=20the=20chat=20engine=20ships=20wi?= =?UTF-8?q?th=20the=20SDK?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit chat() is the SDK's front door, and its engine lived behind a vendor-named extra: pip install pageindex could index a document but failed on the first chat call, and chatting with Claude required installing '[openai]'. Measured before moving: the base tree already carries litellm (75 MB) + openai (13 MB), openai-agents adds ~15 MB (agents 8.1 + mcp 1.7 + griffe 1.4 + small pure-python deps), and current litellm's openai range (>=2.20,<3) intersects cleanly with openai-agents' (>=2.45,<3). The [openai] extra stays declared but empty, so existing pip install 'pageindex[openai]' commands keep resolving. Error messages and docstrings drop the extra; requirements.txt gains the dependency, so CI now runs the openai-agents test lane instead of skipping it. Extras now mean exactly one thing: a vendor's own SDK surface ([anthropic] for messages()/tool runner, [claude] for the Claude Agent SDK). --- examples/agentic_vectorless_rag_demo.py | 2 +- pageindex/client.py | 7 +++---- pageindex/integrations/openai_agents.py | 2 +- pageindex/local_chat.py | 2 +- pyproject.toml | 7 ++++--- requirements.txt | 2 +- tests/test_client.py | 8 ++++---- tests/test_local_chat.py | 2 +- 8 files changed, 16 insertions(+), 16 deletions(-) diff --git a/examples/agentic_vectorless_rag_demo.py b/examples/agentic_vectorless_rag_demo.py index ac0db360a..5b25c7638 100644 --- a/examples/agentic_vectorless_rag_demo.py +++ b/examples/agentic_vectorless_rag_demo.py @@ -18,7 +18,7 @@ 2 โ€” View document metadata 3 โ€” Ask a question (agent reasons over the index and auto-calls tools) -Requirements: pip install "pageindex[openai]"; OPENAI_API_KEY in the environment. +Requirements: pip install pageindex; OPENAI_API_KEY in the environment. """ import sys import asyncio diff --git a/pageindex/client.py b/pageindex/client.py index 5329cfd73..7eb58af68 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -400,7 +400,7 @@ def chat_completions( Cloud: the hosted chat endpoint. Local: a managed document-QA agent run over the local tools against your own LLM backend's - /chat/completions (requires ``pageindex[openai]``; the OpenAI SDK's + /chat/completions (the OpenAI SDK's usual env config โ€” OPENAI_API_KEY, OPENAI_BASE_URL โ€” selects the backend, so any OpenAI-compatible server works; a ``/`` in the model name means LiteLLM provider routing, so prefix ``openai/`` @@ -492,7 +492,7 @@ def responses( to keep provider prompt-cache prefix continuity and the agent's memory of what it already read. - Requires ``pageindex[openai]`` and a backend that supports the + 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 @@ -694,8 +694,7 @@ def as_openai_tools(self, include_management: bool = False, Local: the in-process tools, any model backend; ``hosted`` does not apply. - Requires ``openai-agents`` (``pip install 'pageindex[openai]'``), - imported only when this method is called. + ``openai-agents`` is imported only when this method is called. Args: include_management (bool): Also expose tools that modify the diff --git a/pageindex/integrations/openai_agents.py b/pageindex/integrations/openai_agents.py index 36c062d2f..4c8ca8ddc 100644 --- a/pageindex/integrations/openai_agents.py +++ b/pageindex/integrations/openai_agents.py @@ -25,7 +25,7 @@ def build_openai_tools(client, include_management: bool = False, except ImportError as exc: raise PageIndexAPIError( "as_openai_tools requires the OpenAI Agents SDK โ€” " - "pip install openai-agents (or pip install 'pageindex[openai]')." + "pip install openai-agents." ) from exc from ..agent_tools import (_dumps, _failure, _require_local_scope, _tool_specs) diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 60d06acc5..e0f3ab58e 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -206,7 +206,7 @@ def _require_openai_agents(method: str) -> None: except ImportError as exc: raise PageIndexAPIError( f"{method} in local mode requires the OpenAI Agents SDK โ€” " - "pip install openai-agents (or pip install 'pageindex[openai]'). " + "pip install openai-agents. " "messages() runs on the anthropic extra instead." ) from exc diff --git a/pyproject.toml b/pyproject.toml index c316451b6..c9aa729f2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,6 +28,8 @@ exclude = ["pageindex/flash/assets"] python = ">=3.10" requests = ">=2.28.0" openai = ">=1.70.0" +# Older releases crash on current openai before the request is sent. +openai-agents = ">=0.18.1" litellm = ">=1.84.0" PyPDF2 = ">=3.0.0" pypdfium2 = ">=4.30.0" @@ -37,14 +39,13 @@ python-dotenv = ">=1.0.0" pyyaml = ">=6.0" # Older releases break string prompts with SDK MCP servers (#597, #780). claude-agent-sdk = { version = ">=0.1.53", optional = true } -# Older releases crash on current openai before the request is sent. -openai-agents = { version = ">=0.18.1", optional = true } # Older releases execute a refusal turn's tool_use blocks. anthropic = { version = ">=0.108.0", optional = true } [tool.poetry.extras] claude = ["claude-agent-sdk"] -openai = ["openai-agents"] +# Empty on purpose: keeps pip install "pageindex[openai]" valid. +openai = [] anthropic = ["anthropic"] [tool.poetry.group.dev.dependencies] diff --git a/requirements.txt b/requirements.txt index 5fd4f2e4e..1516a3979 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,7 @@ litellm==1.84.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 diff --git a/tests/test_client.py b/tests/test_client.py index aad4e3979..7bd505342 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -676,12 +676,12 @@ def test_retrieval_endpoints_cloud_only(local_client): local_client.get_retrieval("any") -def test_chat_completions_local_needs_agents_extra(local_client, monkeypatch): - """Local chat is implemented (see test_local_chat.py); without the - openai-agents extra it raises the actionable install error.""" +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="pageindex\\[openai\\]"): + with pytest.raises(PageIndexAPIError, match="pip install openai-agents"): local_client.chat_completions( messages=[{"role": "user", "content": "q"}]) diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index 18a79cbc0..9c1c911d5 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -263,7 +263,7 @@ def test_chat_completions_stream_modes(client, store_path, fake_model): def test_chat_completions_missing_framework(client, monkeypatch): monkeypatch.setitem(sys.modules, "agents", None) - with pytest.raises(PageIndexAPIError, match="pageindex\\[openai\\]"): + with pytest.raises(PageIndexAPIError, match="pip install openai-agents"): client.chat_completions([{"role": "user", "content": "x"}]) From 75a59c80ffc6cd23da1da3938eb4e6bd39a03d01 Mon Sep 17 00:00:00 2001 From: Ray Date: Sat, 15 Aug 2026 02:33:24 +0800 Subject: [PATCH 075/137] =?UTF-8?q?docs:=20drop=20the=20demo's=20install?= =?UTF-8?q?=20step=20=E2=80=94=20openai-agents=20ships=20with=20the=20SDK?= =?UTF-8?q?=20now?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 4 ---- 1 file changed, 4 deletions(-) diff --git a/README.md b/README.md index 27e3084a4..3f0a1b972 100644 --- a/README.md +++ b/README.md @@ -213,10 +213,6 @@ python3 run_pageindex.py --md_path /path/to/your/document.md 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 ``` From ad239a991538b598a93bc17fede184f11bb11dfe Mon Sep 17 00:00:00 2001 From: Ray Date: Sun, 16 Aug 2026 20:16:31 +0800 Subject: [PATCH 076/137] =?UTF-8?q?feat:=20the=20chat=20lane=20routes=20ev?= =?UTF-8?q?ery=20model=20through=20LiteLLM=20=E2=80=94=20bare=20names=20in?= =?UTF-8?q?cluded?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The direct-OpenAI special case existed to dodge LiteLLM's import cost, and it made OpenAI's own Responses-first models fail on the front door: gpt-5.6-sol 400s on chatcmpl+tools while reasoning is on (server-side policy โ€” wire-captured with no reasoning_effort in our request). LiteLLM 1.97 translates such calls onto /v1/responses; 1.84 does not, so the sol-class 400 now carries its two exits (upgrade litellm / responses()). Routing after the flip: chat protocol โ€” bare names are OpenAI-compatible shorthand (wire form openai/; OPENAI_API_KEY / OPENAI_BASE_URL still select the backend, and the missing key stays a build-time failure), litellm/ strips, openai/ opts out to the OpenAI SDK directly; responses protocol unchanged (OpenAI-SDK native, LiteLLM refused). The import cost is handled instead of dodged: local clients preload litellm on a background thread (first call then perceives 0.0s), and pageindex sets LITELLM_LOCAL_MODEL_COST_MAP=True via setdefault โ€” LiteLLM's import otherwise blocks on a network fetch of its price map (fresh venv: 5.6s -> 1.3s; offline it hangs to the timeout). Also restores prompt_cache_key delivery, found dead during the flip's gating verification: openai-agents 0.20 no longer derives it from RunConfig.group_id, so both lanes sent nothing. ModelSettings.extra_body is the one channel all three model classes put on the wire (the bare kwarg is dropped by LiteLLM; extra_args[extra_body] collides with the responses model's own parameter โ€” both wire-verified), and it is scoped to OpenAI destinations: LiteLLM plants extra_body as a literal field in other providers' bodies, and Anthropic rejects unknown fields โ€” the anthropic wire test now pins the absence. Verified before landing: mock-server matrix (OPENAI_BASE_URL + bare name works through LiteLLM; gpt-named self-hosted models are NOT bridged off a custom base_url; prompt_cache_key on the wire in every OpenAI lane with distinct per-conversation keys; anthropic body clean) and live (sol answers through chat(), gpt-5.4 unchanged, responses() bare unchanged with the key on its wire). --- pageindex/__init__.py | 5 ++ pageindex/client.py | 39 +++++++++----- pageindex/local_chat.py | 107 ++++++++++++++++++++++++++------------- tests/test_local_chat.py | 79 +++++++++++++++++++++-------- 4 files changed, 161 insertions(+), 69 deletions(-) diff --git a/pageindex/__init__.py b/pageindex/__init__.py index 8a2383014..88ca32ff8 100644 --- a/pageindex/__init__.py +++ b/pageindex/__init__.py @@ -1,6 +1,11 @@ """PageIndex SDK.""" +import os as _os from typing import TYPE_CHECKING as _TYPE_CHECKING +# LiteLLM's import otherwise fetches its model map over the network โ€” seconds +# of blocking (or a hang offline). setdefault, so an explicit user choice wins. +_os.environ.setdefault("LITELLM_LOCAL_MODEL_COST_MAP", "True") + from .client import PageIndexClient, PageIndexCloudClient, PageIndexLocalClient from .errors import PageIndexAPIError diff --git a/pageindex/client.py b/pageindex/client.py index 7eb58af68..5b1a03d42 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -2,6 +2,7 @@ from __future__ import annotations import os +import threading import time import warnings from typing import Any, Callable, Iterator, Optional, Union, cast @@ -9,6 +10,13 @@ from .errors import PageIndexAPIError +def _preload_litellm() -> None: + try: + import litellm # noqa: F401 + except Exception: + pass + + def _parse_pages(pages: str) -> list[int]: result: set[int] = set() too_many = (f"Page specification '{pages}' spans more than " @@ -62,10 +70,12 @@ class PageIndexClient: summaries and document descriptions. retrieve_model (str, optional): Local mode only โ€” the model the local chat surfaces (``chat_completions``, ``responses``) - default to, exposed as ``client.retrieve_model``. - ``provider/model`` names route through LiteLLM; for an - OpenAI-compatible server that itself serves slashed model ids - (vLLM, TGI), prefix ``openai/`` (e.g. ``openai/Qwen/...``). + default to, exposed as ``client.retrieve_model``. On the chat + lane every name routes through LiteLLM (bare names are + OpenAI-compatible shorthand); prefix ``openai/`` to drive the + OpenAI SDK directly โ€” also the form for an OpenAI-compatible + server that itself serves slashed model ids (vLLM, TGI, + e.g. ``openai/Qwen/...``). storage_path (str, optional): Local mode only โ€” directory where indexed documents are stored. Defaults to ``./.pageindex``. @@ -130,6 +140,9 @@ def __init__( summary_model=self.summary_model, retrieve_model=self.retrieve_model, ) + # LiteLLM's multi-second import would otherwise land on the + # first chat call; failures resurface there with real context. + threading.Thread(target=_preload_litellm, daemon=True).start() # ---------- DOCUMENT SUBMISSION ---------- @@ -399,15 +412,15 @@ def chat_completions( PageIndex Chat Completions: document QA in one call. Cloud: the hosted chat endpoint. Local: a managed document-QA agent - run over the local tools against your own LLM backend's - /chat/completions (the OpenAI SDK's - usual env config โ€” OPENAI_API_KEY, OPENAI_BASE_URL โ€” selects the - backend, so any OpenAI-compatible server works; a ``/`` in the - model name means LiteLLM provider routing, so prefix ``openai/`` - when the backend itself serves slashed ids, e.g. - ``openai/Qwen/...`` on vLLM; LiteLLM-routed Claude models โ€” - Anthropic direct, Bedrock, Vertex โ€” get the managed prompt - prefix cache-marked automatically). The non-stream + run over the local tools against your own LLM backend, routed + through LiteLLM (bare names are OpenAI-compatible shorthand: the + OpenAI SDK's usual env config โ€” OPENAI_API_KEY, OPENAI_BASE_URL โ€” + still selects the backend, so any OpenAI-compatible server works, + and prefixing ``openai/`` opts out of LiteLLM to the OpenAI SDK + directly, e.g. ``openai/Qwen/...`` on vLLM; 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") โ€” diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index e0f3ab58e..8c55fb874 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -24,6 +24,7 @@ import concurrent.futures import hashlib import json +import os import queue import threading import time @@ -214,14 +215,18 @@ def _require_openai_agents(method: str) -> None: def _openai_model(protocol: str, model_name: str): """The backend protocol driver โ€” the seam tests replace with a fake. - ``litellm//`` (the client's normalized retrieve_model - form) and bare ``/`` paths drive the provider through - LiteLLM โ€” chat.completions only, so the responses protocol refuses them - instead of silently downgrading; a first segment LiteLLM does not know - (a HuggingFace repo id like ``Qwen/...``) is refused with the - ``openai/`` escape instead of failing inside LiteLLM at request time; - an ``openai/`` prefix strips to the OpenAI SDK; bare names go to the - OpenAI SDK as-is.""" + chat protocol: every model routes through LiteLLM. Bare names are + OpenAI-compatible shorthand (wire form ``openai/``, so + OPENAI_API_KEY / OPENAI_BASE_URL keep selecting the backend), a + ``litellm/`` prefix strips, a first segment LiteLLM does not know (a + HuggingFace repo id like ``Qwen/...``) is refused with the ``openai/`` + escape instead of failing inside LiteLLM at request time, and an + ``openai/`` prefix opts out to the OpenAI SDK directly. + + responses protocol: the Responses API is OpenAI-SDK native โ€” LiteLLM's + completion surface speaks the chat.completions format, so + provider-prefixed models are refused instead of silently downgrading; + bare and ``openai/`` names drive the OpenAI SDK.""" if "/" in model_name and not model_name.startswith("openai/"): if protocol == "responses": raise PageIndexAPIError( @@ -233,6 +238,7 @@ def _openai_model(protocol: str, model_name: str): "Responses-capable backend and use a bare or " "'openai/'-prefixed model name." ) + if protocol == "chat" and not model_name.startswith("openai/"): try: from agents.extensions.models.litellm_model import LitellmModel import litellm @@ -242,6 +248,14 @@ def _openai_model(protocol: str, model_name: str): "installed. Run: pip install 'litellm>=1.30'" ) wire = model_name.removeprefix("litellm/") + if "/" not in wire: + if not os.environ.get("OPENAI_API_KEY"): + raise PageIndexAPIError( + "The OpenAI backend is not configured: set the " + "OPENAI_API_KEY environment variable (any value works " + "for keyless OPENAI_BASE_URL servers)." + ) + wire = f"openai/{wire}" providers = getattr(litellm, "provider_list", None) if providers and wire.split("/", 1)[0] not in providers: raise PageIndexAPIError( @@ -296,16 +310,26 @@ def _cache_extra_args(model_name: str) -> Optional[dict]: def _openai_agent(client, protocol: str, model_name: str, instructions: str, - temperature, top_p, doc_ids=None): + temperature, top_p, doc_ids=None, cache_key=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. + wire = model_name.removeprefix("litellm/") + openai_backend = "/" not in wire or wire.startswith("openai/") return Agent( name="PageIndex", instructions=instructions, tools=build_openai_tools(client, doc_ids=doc_ids), model=_openai_model(protocol, model_name), - model_settings=ModelSettings(temperature=temperature, top_p=top_p, - extra_args=_cache_extra_args(model_name)), + model_settings=ModelSettings( + temperature=temperature, top_p=top_p, + extra_body=({"prompt_cache_key": cache_key} + if cache_key and openai_backend else None), + extra_args=_cache_extra_args(model_name)), ) @@ -315,27 +339,40 @@ def _validate_max_turns(max_turns) -> None: raise PageIndexAPIError("max_turns must be a positive integer.") -def _conversation_group_id(model_name: str, instructions: str, items) -> str: - """Stable per-conversation cache-routing key: openai-agents hashes - RunConfig.group_id into the OpenAI prompt_cache_key, and without one it - stamps every run with a fresh key, tagging a round-tripped prefix as a - different cache group. Keyed on the prefix identity โ€” model, - instructions, first conversation item โ€” so a conversation's - continuations share one route without pooling unrelated conversations. - Callers pass the conversation's own items, never the SDK-prepended - doc-targeting block: that block is byte-identical for every - conversation about a document and would pool them all under one key.""" +def _conversation_cache_key(model_name: str, instructions: str, items) -> str: + """Stable per-conversation cache-routing key, sent as the OpenAI + ``prompt_cache_key`` through ModelSettings.extra_args (openai-agents + 0.20 no longer derives it from RunConfig.group_id โ€” verified against a + captured wire). Keyed on the prefix identity โ€” model, instructions, + first conversation item โ€” so a conversation's continuations share one + route without pooling unrelated conversations. Callers pass the + conversation's own items, never the SDK-prepended doc-targeting block: + that block is byte-identical for every conversation about a document + and would pool them all under one key.""" seed = json.dumps([model_name, instructions, items[0] if items else None], sort_keys=True, default=str) return "pageindex-" + hashlib.sha256(seed.encode()).hexdigest()[:16] -def _run_kwargs(max_turns, group_id: str) -> dict: +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) 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, - group_id=group_id)} + kwargs: dict = {"run_config": RunConfig(tracing_disabled=True)} if max_turns is not None: kwargs["max_turns"] = max_turns return kwargs @@ -449,10 +486,10 @@ def run_chat_completions(client, messages, stream: bool = False, reported_model = _reported_model(model_name) managed = _managed_instructions(system_texts) agent = _openai_agent(client, "chat", model_name, managed, - temperature, None, doc_ids=doc_id) - run_kwargs = _run_kwargs(max_turns, - _conversation_group_id(model_name, managed, - history)) + temperature, None, doc_ids=doc_id, + cache_key=_conversation_cache_key(model_name, + managed, history)) + run_kwargs = _run_kwargs(max_turns) import openai from agents import Runner from agents.exceptions import AgentsException, MaxTurnsExceeded @@ -466,8 +503,7 @@ def run_chat_completions(client, messages, stream: bool = False, 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 + raise _model_backend_error(exc) from exc return { "id": f"chatcmpl-{uuid.uuid4().hex}", "object": "chat.completion", @@ -512,8 +548,7 @@ async def agen(): 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 + raise _model_backend_error(exc) from exc finally: if not completed and hasattr(streamed, "cancel"): streamed.cancel() # abandoned/failed: stop the agent task @@ -559,10 +594,10 @@ def run_responses(client, input, model: Optional[str] = None, model_name = model or client.retrieve_model managed = _managed_instructions(extra) agent = _openai_agent(client, "responses", model_name, managed, - temperature, top_p, doc_ids=doc_id) - run_kwargs = _run_kwargs(max_turns, - _conversation_group_id(model_name, managed, - conversation)) + temperature, top_p, doc_ids=doc_id, + cache_key=_conversation_cache_key(model_name, managed, + conversation)) + run_kwargs = _run_kwargs(max_turns) recorded: dict = {} import openai from agents import Runner diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index 9c1c911d5..27333a0e5 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -354,6 +354,9 @@ def fake_post(self, url=None, *args, **kwargs): "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" @@ -478,13 +481,14 @@ def test_doc_id_conversations_get_distinct_cache_keys(client, store_path, under one prompt_cache_key.""" seed_doc(store_path, "pi-a", "report.pdf") keys = [] - real = local_chat._run_kwargs + real = local_chat._conversation_cache_key - def spy(max_turns, group_id): - keys.append(group_id) - return real(max_turns, group_id) + def spy(model_name, instructions, items): + key = real(model_name, instructions, items) + keys.append(key) + return key - monkeypatch.setattr(local_chat, "_run_kwargs", spy) + monkeypatch.setattr(local_chat, "_conversation_cache_key", spy) fake_model([[_msg_item("a")]]) result = client.responses("What is the CAGR?", doc_id="pi-a") @@ -831,26 +835,58 @@ def test_responses_envelope_fields_and_cache_group(client, store_path, assert result["tool_choice"] == "auto" -def test_conversation_group_id_stable_per_conversation(): - """Cache-routing key: openai-agents hashes group_id into the OpenAI - prompt_cache_key. A conversation's continuations must share one key - (same model/instructions/first item), and unrelated conversations must - not pool under it.""" +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 both exits.""" + 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) + 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 / first item), and unrelated conversations must not pool + under it.""" turn1 = [{"role": "user", "content": "q"}] continuation = turn1 + [{"role": "assistant", "content": "a"}, {"role": "user", "content": "and?"}] - key = local_chat._conversation_group_id("m", "sys", turn1) - assert key == local_chat._conversation_group_id("m", "sys", continuation) - assert key != local_chat._conversation_group_id( + key = local_chat._conversation_cache_key("m", "sys", turn1) + assert key == local_chat._conversation_cache_key("m", "sys", continuation) + assert key != local_chat._conversation_cache_key( "m", "sys", [{"role": "user", "content": "other"}]) - assert key != local_chat._conversation_group_id("m2", "sys", turn1) - assert key != local_chat._conversation_group_id("m", "sys2", turn1) + assert key != local_chat._conversation_cache_key("m2", "sys", turn1) + assert key != local_chat._conversation_cache_key("m", "sys2", turn1) @needs_agents -def test_run_kwargs_sets_conversation_group_id(): - key = "pageindex-test" - assert (local_chat._run_kwargs(None, key)["run_config"].group_id == key) +def test_agent_carries_prompt_cache_key_in_extra_args(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 + 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 @@ -902,8 +938,9 @@ def test_empty_doc_id_is_an_empty_allowlist(client, store_path, fake_model): @needs_agents def test_openai_model_resolves_provider_prefixes(): - """retrieve_model arrives normalized (litellm//); the - OpenAI SDK must never see that prefix as a wire model name.""" + """The chat lane routes everything through LiteLLM โ€” bare names as the + openai/ shorthand, routing prefixes never leak as wire model names. + openai/ opts out to the OpenAI SDK; responses stays OpenAI-SDK native.""" pytest.importorskip("litellm") from agents.extensions.models.litellm_model import LitellmModel from agents.models.openai_chatcompletions import OpenAIChatCompletionsModel @@ -913,6 +950,8 @@ def test_openai_model_resolves_provider_prefixes(): 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, OpenAIChatCompletionsModel) assert str(model.model) == "gpt-5.2" From 82e33548823ac37503bdd6dbceaa180dd56dbf9b Mon Sep 17 00:00:00 2001 From: Ray Date: Sun, 16 Aug 2026 20:31:20 +0800 Subject: [PATCH 077/137] =?UTF-8?q?refactor:=20no=20prefix-triggered=20dir?= =?UTF-8?q?ect=20lane=20=E2=80=94=20chat=20model=20names=20are=20LiteLLM's?= =?UTF-8?q?,=20verbatim?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ray's ruling on the flip's remaining carve-out: a routing decision must never hide in a model-name prefix. openai/ now means what LiteLLM says it means (its openai provider), like every other name on the chat lane โ€” the grammar is LiteLLM's with zero exceptions. The two defenses for keeping a direct carve-out had no concrete victim: debugging isolation (litellm is unavoidable in indexing anyway, and responses() IS the OpenAI-SDK-native door), and endpoint determinism (litellm sends chatcmpl for openai-provider models except the gpt-5 bridge, which never fires against a custom base_url โ€” wire-verified). If a direct escape is ever needed, it will be a declared parameter, never name grammar. openai/-prefixed names keep the build-time OPENAI_API_KEY check for parity with bare names; responses() is untouched (bare and openai/ still drive the OpenAI SDK โ€” LiteLLM cannot speak that protocol). --- pageindex/client.py | 27 ++++++------- pageindex/local_chat.py | 85 +++++++++++++++++++--------------------- tests/test_local_chat.py | 19 +++++---- 3 files changed, 64 insertions(+), 67 deletions(-) diff --git a/pageindex/client.py b/pageindex/client.py index 5b1a03d42..3aeba52e3 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -70,12 +70,11 @@ class PageIndexClient: summaries and document descriptions. retrieve_model (str, optional): Local mode only โ€” the model the local chat surfaces (``chat_completions``, ``responses``) - default to, exposed as ``client.retrieve_model``. On the chat - lane every name routes through LiteLLM (bare names are - OpenAI-compatible shorthand); prefix ``openai/`` to drive the - OpenAI SDK directly โ€” also the form for an OpenAI-compatible - server that itself serves slashed model ids (vLLM, TGI, - e.g. ``openai/Qwen/...``). + default to, exposed as ``client.retrieve_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). storage_path (str, optional): Local mode only โ€” directory where indexed documents are stored. Defaults to ``./.pageindex``. @@ -413,14 +412,14 @@ def chat_completions( Cloud: the hosted chat endpoint. Local: a managed document-QA agent run over the local tools against your own LLM backend, routed - through LiteLLM (bare names are OpenAI-compatible shorthand: the - OpenAI SDK's usual env config โ€” OPENAI_API_KEY, OPENAI_BASE_URL โ€” - still selects the backend, so any OpenAI-compatible server works, - and prefixing ``openai/`` opts out of LiteLLM to the OpenAI SDK - directly, e.g. ``openai/Qwen/...`` on vLLM; provider-prefixed - names โ€” ``anthropic/โ€ฆ``, ``bedrock/โ€ฆ`` โ€” reach that provider, and - LiteLLM-routed Claude models get the managed prompt prefix - cache-marked automatically). The non-stream + 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") โ€” diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 8c55fb874..4271056a3 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -215,20 +215,19 @@ def _require_openai_agents(method: str) -> None: def _openai_model(protocol: str, model_name: str): """The backend protocol driver โ€” the seam tests replace with a fake. - chat protocol: every model routes through LiteLLM. Bare names are - OpenAI-compatible shorthand (wire form ``openai/``, so - OPENAI_API_KEY / OPENAI_BASE_URL keep selecting the backend), a - ``litellm/`` prefix strips, a first segment LiteLLM does not know (a - HuggingFace repo id like ``Qwen/...``) is refused with the ``openai/`` - escape instead of failing inside LiteLLM at request time, and an - ``openai/`` prefix opts out to the OpenAI SDK directly. + chat protocol: LiteLLM, full stop โ€” model names mean what LiteLLM says + they mean. Bare names are OpenAI-compatible shorthand (wire form + ``openai/``, so OPENAI_API_KEY / OPENAI_BASE_URL keep selecting + the backend), a ``litellm/`` prefix strips, and a first segment LiteLLM + does not know (a HuggingFace repo id like ``Qwen/...``) is refused with + the ``openai/`` form instead of failing inside LiteLLM at request time. responses protocol: the Responses API is OpenAI-SDK native โ€” LiteLLM's completion surface speaks the chat.completions format, so provider-prefixed models are refused instead of silently downgrading; bare and ``openai/`` names drive the OpenAI SDK.""" - if "/" in model_name and not model_name.startswith("openai/"): - if protocol == "responses": + if protocol == "responses": + if "/" in model_name and not model_name.startswith("openai/"): raise PageIndexAPIError( f"responses() cannot drive " f"'{model_name.removeprefix('litellm/')}': provider-prefixed " @@ -238,47 +237,43 @@ def _openai_model(protocol: str, model_name: str): "Responses-capable backend and use a bare or " "'openai/'-prefixed model name." ) - if protocol == "chat" and not model_name.startswith("openai/"): + import openai + model_name = model_name.removeprefix("openai/") try: - from agents.extensions.models.litellm_model import LitellmModel - import litellm - except ImportError: + backend = openai.AsyncOpenAI() + except openai.OpenAIError as exc: raise PageIndexAPIError( - f"'{model_name}' routes through LiteLLM, but litellm is not " - "installed. Run: pip install 'litellm>=1.30'" - ) - wire = model_name.removeprefix("litellm/") - if "/" not in wire: - if not os.environ.get("OPENAI_API_KEY"): - raise PageIndexAPIError( - "The OpenAI backend is not configured: set the " - "OPENAI_API_KEY environment variable (any value works " - "for keyless OPENAI_BASE_URL servers)." - ) - wire = f"openai/{wire}" - providers = getattr(litellm, "provider_list", None) - if providers and wire.split("/", 1)[0] not in providers: + f"The OpenAI backend is not configured: {exc}") from exc + from agents.models.openai_responses import OpenAIResponsesModel + return OpenAIResponsesModel(model_name, openai_client=backend) + try: + from agents.extensions.models.litellm_model import LitellmModel + import litellm + except ImportError: + raise PageIndexAPIError( + f"'{model_name}' routes through LiteLLM, but litellm is not " + "installed. Run: pip install 'litellm>=1.30'" + ) + wire = model_name.removeprefix("litellm/") + if "/" not in wire or wire.startswith("openai/"): + if not os.environ.get("OPENAI_API_KEY"): raise PageIndexAPIError( - f"'{wire}' routes through LiteLLM, but " - f"'{wire.split('/', 1)[0]}' is not a LiteLLM provider. For an " - "OpenAI-compatible server (vLLM, TGI, Ollama) serving this " - f"model id, use 'openai/{wire}' and point OPENAI_BASE_URL " - "at the server." + "The OpenAI backend is not configured: set the " + "OPENAI_API_KEY environment variable (any value works " + "for keyless OPENAI_BASE_URL servers)." ) - return LitellmModel(wire) - import openai - model_name = model_name.removeprefix("openai/") - try: - backend = openai.AsyncOpenAI() - except openai.OpenAIError as exc: + if "/" not in wire: + wire = f"openai/{wire}" + providers = getattr(litellm, "provider_list", None) + if providers and wire.split("/", 1)[0] not in providers: raise PageIndexAPIError( - f"The OpenAI backend is not configured: {exc}") from exc - if protocol == "chat": - from agents.models.openai_chatcompletions import ( - OpenAIChatCompletionsModel) - return OpenAIChatCompletionsModel(model_name, backend) - from agents.models.openai_responses import OpenAIResponsesModel - return OpenAIResponsesModel(model_name, openai_client=backend) + f"'{wire}' routes through LiteLLM, but " + f"'{wire.split('/', 1)[0]}' is not a LiteLLM provider. For an " + "OpenAI-compatible server (vLLM, TGI, Ollama) serving this " + f"model id, use 'openai/{wire}' and point OPENAI_BASE_URL " + "at the server." + ) + return LitellmModel(wire) def _reported_model(model_name: str) -> str: diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index 27333a0e5..c721ee857 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -938,12 +938,12 @@ def test_empty_doc_id_is_an_empty_allowlist(client, store_path, fake_model): @needs_agents def test_openai_model_resolves_provider_prefixes(): - """The chat lane routes everything through LiteLLM โ€” bare names as the - openai/ shorthand, routing prefixes never leak as wire model names. - openai/ opts out to the OpenAI SDK; responses stays OpenAI-SDK native.""" + """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_chatcompletions import OpenAIChatCompletionsModel from agents.models.openai_responses import OpenAIResponsesModel model = local_chat._openai_model("chat", "litellm/anthropic/claude-x") @@ -953,11 +953,13 @@ def test_openai_model_resolves_provider_prefixes(): 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, OpenAIChatCompletionsModel) - assert str(model.model) == "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" @needs_agents @@ -1021,8 +1023,9 @@ def test_chat_missing_openai_key_fails_loud(monkeypatch): """A missing backend credential surfaces as the SDK's own error type, like every other precondition on the chat surfaces.""" monkeypatch.delenv("OPENAI_API_KEY", raising=False) - with pytest.raises(PageIndexAPIError, match="OPENAI_API_KEY"): - local_chat._openai_model("chat", "gpt-4o") + for name in ("gpt-4o", "openai/gpt-4o"): + with pytest.raises(PageIndexAPIError, match="OPENAI_API_KEY"): + local_chat._openai_model("chat", name) @needs_agents From 1ecc322d085ea7ed12b35fe0bdaa4ae86ec33839 Mon Sep 17 00:00:00 2001 From: Ray Date: Sun, 16 Aug 2026 20:51:55 +0800 Subject: [PATCH 078/137] =?UTF-8?q?fix:=20third-audit=20findings=20?= =?UTF-8?q?=E2=80=94=20extra=5Fbody=20naming=20drift,=20litellm=20floor=20?= =?UTF-8?q?hint?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The prompt_cache_key delivery channel went through three iterations and settled on ModelSettings.extra_body; the _conversation_cache_key docstring still named extra_args from the middle iteration. The litellm install hint said >=1.30, below both our own pyproject floor (>=1.84.0) and the floor openai-agents' litellm extra declares (>=1.83). A user in a broken environment following it would land on a version the package itself rules out. The hint now matches the declared floor. --- pageindex/local_chat.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 4271056a3..564068cf5 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -252,7 +252,7 @@ def _openai_model(protocol: str, model_name: str): except ImportError: raise PageIndexAPIError( f"'{model_name}' routes through LiteLLM, but litellm is not " - "installed. Run: pip install 'litellm>=1.30'" + "installed. Run: pip install 'litellm>=1.84'" ) wire = model_name.removeprefix("litellm/") if "/" not in wire or wire.startswith("openai/"): @@ -336,7 +336,7 @@ def _validate_max_turns(max_turns) -> None: def _conversation_cache_key(model_name: str, instructions: str, items) -> str: """Stable per-conversation cache-routing key, sent as the OpenAI - ``prompt_cache_key`` through ModelSettings.extra_args (openai-agents + ``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, first conversation item โ€” so a conversation's continuations share one From 0307a2c27392d5c69017c19e7fe0cc14ba7ec7ab Mon Sep 17 00:00:00 2001 From: Ray Date: Sun, 16 Aug 2026 20:52:15 +0800 Subject: [PATCH 079/137] test: pin the bundle door's Agents-SDK model grammar; rename the cache-key test to extra_body MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The config bundle hands its model string to the Agents SDK's own MultiProvider grammar, which refuses unknown prefixes (probe on 0.20: 'anthropic/x' -> UserError: Unknown prefix). _normalize_retrieve_model's litellm/ spelling is what keeps that door working โ€” a link the existing self-referential assert (config["model"] == client.retrieve_model) could not catch. Pinned with a provider-slashed name. Also renames the cache-key delivery test to its real channel, extra_body โ€” the extra_args name survived from the superseded delivery attempt. --- tests/test_agent_tools.py | 10 ++++++++++ tests/test_local_chat.py | 2 +- 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 5c7f7d368..d08ab547c 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -754,6 +754,16 @@ def test_openai_agent_config_local(client, store_path): 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.""" + client = PageIndexLocalClient(storage_path=str(tmp_path / "s"), + retrieve_model="anthropic/claude-x") + assert (client.openai_agent_config()["model"] + == "litellm/anthropic/claude-x") + + def test_openai_agent_config_cloud_omits_model(cloud_with_fake_bridge): pytest.importorskip("agents") cloud, _ = cloud_with_fake_bridge diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index c721ee857..f0b5978ee 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -863,7 +863,7 @@ def test_conversation_cache_key_stable_per_conversation(): @needs_agents -def test_agent_carries_prompt_cache_key_in_extra_args(monkeypatch): +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 From d411e5a63ea707dc50e71d4ebbb8f03b00111c3a Mon Sep 17 00:00:00 2001 From: Ray Date: Sun, 16 Aug 2026 21:12:03 +0800 Subject: [PATCH 080/137] =?UTF-8?q?refactor:=20name=20the=20retrieve=5Fmod?= =?UTF-8?q?el=20helper=20for=20its=20reason=20=E2=80=94=20the=20Agents=20S?= =?UTF-8?q?DK's=20grammar?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _normalize_retrieve_model said what it does, not why. The litellm/ spelling exists because the Agents SDK resolves raw model strings with its own prefix grammar and refuses unknown prefixes โ€” the name now points at that constraint. --- pageindex/client.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pageindex/client.py b/pageindex/client.py index 3aeba52e3..2cda81616 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -41,7 +41,7 @@ def _parse_pages(pages: str) -> list[int]: return sorted(result) -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: @@ -129,7 +129,7 @@ def __init__( 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( + self.retrieve_model = _agents_sdk_model_name( getattr(opt, "retrieve_model", None) or opt.model) self.storage_path = storage_path or ".pageindex" from .local_api import LocalAPI From 4ee6a5e31fb10e22b16c6c8be255df686435afc6 Mon Sep 17 00:00:00 2001 From: Ray Date: Sun, 16 Aug 2026 21:23:04 +0800 Subject: [PATCH 081/137] test: the bundle-grammar test skips without openai-agents, like its file's siblings Every agents-dependent test in this file importorskips; without the guard this one errors where the others skip. --- tests/test_agent_tools.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index d08ab547c..3873cd36e 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -758,6 +758,7 @@ 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"] From 7244ee4954f4c79fc4f2c83dfd48955dcc2ea168 Mon Sep 17 00:00:00 2001 From: Ray Date: Sun, 16 Aug 2026 22:09:07 +0800 Subject: [PATCH 082/137] =?UTF-8?q?feat:=20index=5Fmodel=20+=20chat=5Fmode?= =?UTF-8?q?l=20=E2=80=94=20two-knob=20model=20surface=20with=20full=20lega?= =?UTF-8?q?cy=20fallback?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The documented surface becomes two role knobs: index_model builds the index, chat_model answers on the chat surfaces. model turns into the set-both umbrella (its 0.2.8 indexing semantics are a strict subset, so old configs run unchanged); summary_model and retrieve_model stay accepted as legacy role names. Resolution lives in ConfigLoader.load(), the one seam every consumer already passes through (client, CLI standard/md paths, flash's summary fallback, tree_optimize's default_model): new names win over old, specific over general, model sets every role, and code constants close each chain. The packaged yaml no longer ships model keys โ€” key presence is what separates a user's explicit choice from a built-in default, and _validate_keys accepts the five model names explicitly. Consequences: with no config at all, classic-mode structure extraction now uses DEFAULT_INDEX_MODEL (gpt-5.6-luna) instead of the yaml's old gpt-4o-2024-11-20 line (ratified; flash-default users see no change). client.retrieve_model becomes a read-only alias for client.chat_model. The resolution matrix test pins one row per released generation: 0.2.8 (model), 0.3.0.dev (model+retrieve_model), 0.2.10.dev (all three legacy names), the new pair, umbrella-only, and mixed. --- pageindex/client.py | 64 +++++++++++++++++++++++++++-------------- pageindex/config.yaml | 14 +++++---- pageindex/local_chat.py | 4 +-- pageindex/utils.py | 29 ++++++++++++++++++- tests/test_client.py | 26 +++++++++++++++++ 5 files changed, 106 insertions(+), 31 deletions(-) diff --git a/pageindex/client.py b/pageindex/client.py index 2cda81616..6b652cad0 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -64,17 +64,27 @@ 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 โ€” the model the - local chat surfaces (``chat_completions``, ``responses``) - default to, exposed as ``client.retrieve_model``. Chat names + 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``. @@ -97,6 +107,8 @@ 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, @@ -107,9 +119,11 @@ def __init__( "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) passed = [name for name, value in local_only.items() if value is not None] if passed: raise PageIndexAPIError( @@ -122,27 +136,30 @@ 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 = _agents_sdk_model_name( - getattr(opt, "retrieve_model", None) or opt.model) + self.index_model = opt.index_model + self.summary_model = opt.summary_model + self.chat_model = _agents_sdk_model_name(opt.chat_model) 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, + retrieve_model=self.chat_model, ) # LiteLLM's multi-second import would otherwise land on the # first chat call; failures resurface there with real context. threading.Thread(target=_preload_litellm, daemon=True).start() + @property + def retrieve_model(self): + """Legacy name for ``chat_model``.""" + return self.chat_model + # ---------- DOCUMENT SUBMISSION ---------- def submit_document( @@ -383,7 +400,7 @@ def chat( 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 - ``retrieve_model``). + ``chat_model``). Returns: - stream=False: the answer string @@ -446,7 +463,7 @@ def chat_completions( enable_citations: Cloud-only โ€” local mode raises (citations need block-level OCR data local mode does not store). model: Local only โ€” backend model name (defaults to - ``retrieve_model``). The cloud endpoint selects its own. + ``chat_model``). The cloud endpoint selects its own. max_turns: Local only โ€” cap on agent turns per call. Returns: @@ -514,7 +531,7 @@ def responses( Args: input: A user message string, or a list of Responses input items (round-trip prior ``items`` here). - model: Backend model name (defaults to ``retrieve_model``). + 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, sequence numbers are reassigned monotonically, @@ -754,7 +771,7 @@ def openai_agent_config( Sugar over the explicit form โ€” ``agent_instructions`` (with ``doc_id`` targeting) as the instructions and ``as_openai_tools`` as the tools; local clients also carry their - configured ``retrieve_model`` (cloud omits ``model`` so the + configured ``chat_model`` (cloud omits ``model`` so the framework default applies). To customize further, switch to those methods directly. @@ -775,7 +792,7 @@ def openai_agent_config( scoped=scope is not None), "tools": self.as_openai_tools(include_management, doc_id=scope), } - model = model or getattr(self, "retrieve_model", None) + model = model or getattr(self, "chat_model", None) if model: config["model"] = model return config @@ -1032,10 +1049,13 @@ 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, ): - super().__init__(None, model=model, summary_model=summary_model, + 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) diff --git a/pageindex/config.yaml b/pageindex/config.yaml index 73a512c7a..a592cfb97 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. +# Indexing names without a provider prefix use the OpenAI SDK directly; +# for other providers, use "provider/model" (e.g. "anthropic/claude-sonnet-4-6"). +# index_model: "gpt-5.6-luna" +# chat_model: "gpt-5.4" toc_check_page_num: 20 max_page_num_each_node: 10 max_token_num_each_node: 20000 diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 564068cf5..0bec5539f 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -477,7 +477,7 @@ def run_chat_completions(client, messages, stream: bool = False, system_texts, history = _split_chat_messages(messages) block = _doc_block(client, doc_id) items = ([{"role": "user", "content": block}] if block else []) + history - model_name = model or client.retrieve_model + 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, @@ -586,7 +586,7 @@ def run_responses(client, input, model: Optional[str] = None, if block: items = [{"role": "user", "content": block}] + items extra = [instructions] if instructions else [] - model_name = model or client.retrieve_model + 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, diff --git a/pageindex/utils.py b/pageindex/utils.py index 97f60a942..52433c3bd 100644 --- a/pageindex/utils.py +++ b/pageindex/utils.py @@ -954,6 +954,31 @@ def thin(nodes, total_nodes): return structure +DEFAULT_INDEX_MODEL = "gpt-5.6-luna" +DEFAULT_CHAT_MODEL = "gpt-5.4" + +# Every model name any released generation shipped: 0.2.8 (model), +# 0.3.0.dev (model, retrieve_model), 0.2.10.dev (model, summary_model, +# retrieve_model), plus the current pair (index_model, chat_model). +_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 +991,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 +1011,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/tests/test_client.py b/tests/test_client.py index 7bd505342..2feb9a89a 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -78,6 +78,32 @@ def resolved(retrieve_model): assert resolved(already_routable) == already_routable +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): from pageindex import PageIndexCloudClient, PageIndexLocalClient From d0ca51d4464abc9363a412184277d7752a6aca7d Mon Sep 17 00:00:00 2001 From: Ray Date: Sun, 16 Aug 2026 22:10:19 +0800 Subject: [PATCH 083/137] feat: --index-model on the CLI; README flag docs follow MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The CLI leads with --index-model; --model stays as its legacy synonym (the CLI only indexes, so the umbrella and the index role coincide). The flash branch's summary fallback gains the index position, and the standard branch now forwards --summary-model, which it had silently ignored โ€” the flag's help always claimed it worked there. The md branch's unfiltered model=None no longer clobbers the default: the resolver treats None as unset. --- README.md | 2 +- run_pageindex.py | 12 +++++++++--- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 3f0a1b972..dca5dbc60 100644 --- a/README.md +++ b/README.md @@ -177,7 +177,7 @@ You can customize the processing with additional optional arguments (the structu ``` --mode Processing mode: flash (default) or standard ---model LLM model to use (default: gpt-4o-2024-11-20) +--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) diff --git a/run_pageindex.py b/run_pageindex.py index 80c01f16f..1cbea0c72 100644 --- a/run_pageindex.py +++ b/run_pageindex.py @@ -23,9 +23,12 @@ 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)') @@ -87,7 +90,7 @@ if args.mode == 'flash': from pageindex.flash import page_index_flash - summary_model = args.summary_model or args.model + summary_model = args.summary_model or args.index_model or args.model will_summarize = args.summary if args.summary is not None else True if summary_model and (will_summarize or args.optimize == 'full'): import litellm @@ -112,7 +115,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, @@ -157,6 +162,7 @@ # 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, From 96c5330d7c4ca684389495220dd1b2fa9bdd1dfb Mon Sep 17 00:00:00 2001 From: Ray Date: Sun, 16 Aug 2026 22:24:16 +0800 Subject: [PATCH 084/137] feat: default chat model becomes gpt-5.6-sol MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ray's pick for the out-of-box QA default; indexing stays on luna. sol runs tools on the Responses lane โ€” current litellm bridges chat() there automatically; older litellm gets the guided 400 naming both exits. --- pageindex/config.yaml | 2 +- pageindex/utils.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pageindex/config.yaml b/pageindex/config.yaml index a592cfb97..4fde6f60e 100644 --- a/pageindex/config.yaml +++ b/pageindex/config.yaml @@ -5,7 +5,7 @@ # Indexing names without a provider prefix use the OpenAI SDK directly; # for other providers, use "provider/model" (e.g. "anthropic/claude-sonnet-4-6"). # index_model: "gpt-5.6-luna" -# chat_model: "gpt-5.4" +# 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/utils.py b/pageindex/utils.py index 52433c3bd..41fa2f04c 100644 --- a/pageindex/utils.py +++ b/pageindex/utils.py @@ -955,7 +955,7 @@ def thin(nodes, total_nodes): DEFAULT_INDEX_MODEL = "gpt-5.6-luna" -DEFAULT_CHAT_MODEL = "gpt-5.4" +DEFAULT_CHAT_MODEL = "gpt-5.6-sol" # Every model name any released generation shipped: 0.2.8 (model), # 0.3.0.dev (model, retrieve_model), 0.2.10.dev (model, summary_model, From a9f5cb024a371b2c60f13e6732ce6a6c1e59b709 Mon Sep 17 00:00:00 2001 From: Ray Date: Sun, 16 Aug 2026 22:25:43 +0800 Subject: [PATCH 085/137] chore: trim the model-keys comment to the constraint The per-generation history lives in 7244ee4's message. --- pageindex/utils.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/pageindex/utils.py b/pageindex/utils.py index 41fa2f04c..9e6243c96 100644 --- a/pageindex/utils.py +++ b/pageindex/utils.py @@ -957,9 +957,7 @@ def thin(nodes, total_nodes): DEFAULT_INDEX_MODEL = "gpt-5.6-luna" DEFAULT_CHAT_MODEL = "gpt-5.6-sol" -# Every model name any released generation shipped: 0.2.8 (model), -# 0.3.0.dev (model, retrieve_model), 0.2.10.dev (model, summary_model, -# retrieve_model), plus the current pair (index_model, chat_model). +# Each of the five names has shipped in a release; all stay accepted. _MODEL_KEYS = ("model", "summary_model", "retrieve_model", "index_model", "chat_model") From c7c3f6c569807d4717a82832b240489b4aca806e Mon Sep 17 00:00:00 2001 From: Ray Date: Sun, 16 Aug 2026 23:42:02 +0800 Subject: [PATCH 086/137] =?UTF-8?q?feat:=20per-door=20reasoning=20passthro?= =?UTF-8?q?ugh=20=E2=80=94=20reasoning=5Feffort=20/=20reasoning=20/=20thin?= =?UTF-8?q?king?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Each chat door gains its own protocol's native thinking control, forwarded verbatim with no invented vocabulary and no default of ours: chat_completions(reasoning_effort=...), responses(reasoning={...}), messages(thinking={...}). Unset sends nothing, so backend defaults (sol: medium, adaptive) are untouched. chat() stays answer-only. Delivery channels, each verified: the chat door rides extra_args["reasoning_effort"] โ€” LiteLLM's own top-level kwarg on every supported openai-agents version, admitting non-enum values ("none"); newer openai-agents promotes it to the top-level argument and pops the duplicate. Wire-captured on a mock backend (/v1/chat/completions body carries it) and coexists with the Claude cache marker in one dict. The responses door rides ModelSettings.reasoning โ€” coerced to the typed openai Reasoning object and forwarded verbatim by the Responses model; the envelope echoes the caller's dict. The messages door joins the existing anthropic passthrough dict, asserted through the real tool runner. LiteLLM semantics observed and accepted as-is: unknown models refuse the param loudly with LiteLLM's own remedies, and gpt-5.4+ names with an explicit effort route to /v1/responses even against a custom api_base (its documented pre-existing arm). The sol-class 400 guidance now names the third exit โ€” an explicit effort routes on older litellm releases too. Cloud chat_completions rejects the new parameter like model/max_turns; responses()/messages() are local-only already. --- pageindex/client.py | 27 +++++++++++++--- pageindex/local_chat.py | 28 +++++++++++++---- tests/test_local_chat.py | 66 +++++++++++++++++++++++++++++++++++++++- 3 files changed, 109 insertions(+), 12 deletions(-) diff --git a/pageindex/client.py b/pageindex/client.py index 6b652cad0..4e72afcce 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -423,6 +423,7 @@ def chat_completions( enable_citations: bool = False, model: Optional[str] = None, max_turns: Optional[int] = None, + reasoning_effort: Optional[str] = None, ) -> Union[dict[str, Any], Iterator[str], Iterator[dict[str, Any]]]: """ PageIndex Chat Completions: document QA in one call. @@ -465,6 +466,11 @@ def chat_completions( 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. + 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). Returns: - stream=False: complete response dict ({'id', 'object', 'created', @@ -485,12 +491,13 @@ def 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, + max_turns=max_turns, reasoning_effort=reasoning_effort, ) - if model is not None or max_turns is not None: + if (model is not None or max_turns is not None + or reasoning_effort is not None): raise PageIndexAPIError( - "model and max_turns are local-mode parameters โ€” the cloud " - "chat endpoint selects its own model." + "model, max_turns and reasoning_effort 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, @@ -508,6 +515,7 @@ def responses( temperature: Optional[float] = None, top_p: Optional[float] = None, max_turns: Optional[int] = None, + reasoning: Optional[dict[str, Any]] = None, ) -> Union[dict[str, Any], Iterator[dict[str, Any]]]: """ Document QA over the OpenAI Responses protocol โ€” the agentic surface. @@ -546,6 +554,10 @@ def responses( instructions: Appended to the managed system prompt. temperature / top_p: Passed through to the model. max_turns: Cap on agent turns per call. + 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). """ from .cloud_api import CloudAPI if isinstance(self._api, CloudAPI): @@ -557,7 +569,7 @@ def 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_turns=max_turns, reasoning=reasoning, ) def messages( @@ -573,6 +585,7 @@ def messages( top_k: Optional[int] = None, stop_sequences: Optional[list[str]] = None, max_turns: Optional[int] = None, + thinking: Optional[dict[str, Any]] = None, ) -> Union[dict[str, Any], Iterator[Any]]: """ Document QA over the Anthropic Messages protocol โ€” Claude-native. @@ -607,6 +620,9 @@ def messages( 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. """ from .cloud_api import CloudAPI if isinstance(self._api, CloudAPI): @@ -620,6 +636,7 @@ def messages( 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, ) # ---------- DOCUMENT MANAGEMENT ---------- diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 0bec5539f..54828f6cc 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -305,7 +305,8 @@ def _cache_extra_args(model_name: str) -> Optional[dict]: def _openai_agent(client, protocol: str, model_name: str, instructions: str, - temperature, top_p, doc_ids=None, cache_key=None): + temperature, top_p, doc_ids=None, cache_key=None, + reasoning=None, reasoning_effort=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 @@ -315,6 +316,13 @@ def _openai_agent(client, protocol: str, model_name: str, instructions: str, # other providers' request bodies, which Anthropic rejects as unknown. wire = model_name.removeprefix("litellm/") openai_backend = "/" not in wire or wire.startswith("openai/") + # 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} return Agent( name="PageIndex", instructions=instructions, @@ -322,9 +330,10 @@ def _openai_agent(client, protocol: str, model_name: str, instructions: str, model=_openai_model(protocol, model_name), model_settings=ModelSettings( temperature=temperature, top_p=top_p, + reasoning=reasoning, extra_body=({"prompt_cache_key": cache_key} if cache_key and openai_backend else None), - extra_args=_cache_extra_args(model_name)), + extra_args=extra_args), ) @@ -358,7 +367,8 @@ def _model_backend_error(exc) -> PageIndexAPIError: 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) or " + "litellm (newer releases route it there automatically), pass " + "reasoning_effort (older litellm routes explicit efforts), or " "call responses() instead." ) return PageIndexAPIError(message) @@ -466,6 +476,7 @@ def run_chat_completions(client, messages, stream: bool = False, enable_citations: bool = False, model: Optional[str] = None, max_turns: Optional[int] = None, + reasoning_effort: Optional[str] = None, ) -> Union[dict, Iterator[str], Iterator[dict]]: if enable_citations: raise PageIndexAPIError( @@ -483,7 +494,8 @@ def run_chat_completions(client, messages, stream: bool = False, agent = _openai_agent(client, "chat", model_name, managed, temperature, None, doc_ids=doc_id, cache_key=_conversation_cache_key(model_name, - managed, history)) + managed, history), + reasoning_effort=reasoning_effort) run_kwargs = _run_kwargs(max_turns) import openai from agents import Runner @@ -570,6 +582,7 @@ def run_responses(client, input, model: Optional[str] = None, temperature: Optional[float] = None, top_p: Optional[float] = None, max_turns: Optional[int] = None, + reasoning: Optional[dict] = None, ) -> Union[dict, Iterator[dict]]: _require_openai_agents("responses") _validate_max_turns(max_turns) @@ -591,7 +604,8 @@ def run_responses(client, input, model: Optional[str] = None, agent = _openai_agent(client, "responses", model_name, managed, temperature, top_p, doc_ids=doc_id, cache_key=_conversation_cache_key(model_name, managed, - conversation)) + conversation), + reasoning=reasoning) run_kwargs = _run_kwargs(max_turns) recorded: dict = {} import openai @@ -619,6 +633,7 @@ def envelope(transcript: list, raw_responses) -> dict: "parallel_tool_calls": True, "temperature": temperature, "top_p": top_p, + "reasoning": reasoning, "max_output_tokens": None, "error": recorded.get("error"), "incomplete_details": recorded.get("incomplete_details"), @@ -801,6 +816,7 @@ def run_messages(client, messages, model: str, top_k: Optional[int] = None, stop_sequences: Optional[list[str]] = None, max_turns: Optional[int] = None, + thinking: Optional[dict] = None, ) -> Union[dict, Iterator[Any]]: from .integrations.anthropic_sdk import build_anthropic_tools @@ -817,7 +833,7 @@ def run_messages(client, messages, model: str, 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, + "stop_sequences": stop_sequences, "thinking": thinking, }.items() if value is not None} runner = _anthropic_client().beta.messages.tool_runner( max_tokens=(max_tokens if max_tokens is not None diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index f0b5978ee..55958d654 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -271,6 +271,9 @@ 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="not available on PageIndex " "cloud yet"): cloud.responses("x") @@ -837,11 +840,12 @@ def test_responses_envelope_fields_and_cache_group(client, store_path, 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 both exits.""" + 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) @@ -889,6 +893,51 @@ def test_agent_carries_prompt_cache_key_in_extra_body(monkeypatch): 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_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([]) @@ -1304,6 +1353,21 @@ def test_messages_max_tokens_default_resolves_per_model(client, fake_anthropic): 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_tool_error_flagged_and_scoped(client, store_path, fake_anthropic): From 7ee64aa5bea1a419d11d6195614fad37efb06bdd Mon Sep 17 00:00:00 2001 From: Ray Date: Mon, 17 Aug 2026 00:12:34 +0800 Subject: [PATCH 087/137] feat: extra_body escape hatch on the three protocol doors The industry-standard per-request extension channel (openai/anthropic SDK trio): a dict merged verbatim into the backend request, last, so caller keys win over SDK-set ones. Routing per door: OpenAI-compatible destinations get a true body merge (ModelSettings.extra_body / the anthropic SDK's native extra_body); LiteLLM-routed providers take the keys as LiteLLM's own top-level kwargs instead, since LiteLLM plants extra_body as literal fields other providers reject. Cloud mode rejects it like the other local-only knobs. chat() stays answer-only. --- pageindex/client.py | 25 ++++++++++++++++----- pageindex/local_chat.py | 23 ++++++++++++++----- tests/test_local_chat.py | 48 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 86 insertions(+), 10 deletions(-) diff --git a/pageindex/client.py b/pageindex/client.py index 4e72afcce..ddf663908 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -424,6 +424,7 @@ def chat_completions( model: Optional[str] = None, max_turns: Optional[int] = None, reasoning_effort: Optional[str] = None, + extra_body: Optional[dict[str, Any]] = None, ) -> Union[dict[str, Any], Iterator[str], Iterator[dict[str, Any]]]: """ PageIndex Chat Completions: document QA in one call. @@ -471,6 +472,11 @@ def chat_completions( 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). Returns: - stream=False: complete response dict ({'id', 'object', 'created', @@ -492,12 +498,14 @@ def chat_completions( temperature=temperature, stream_metadata=stream_metadata, enable_citations=enable_citations, model=model, max_turns=max_turns, reasoning_effort=reasoning_effort, + extra_body=extra_body, ) if (model is not None or max_turns is not None - or reasoning_effort is not None): + or reasoning_effort is not None or extra_body is not None): raise PageIndexAPIError( - "model, max_turns and reasoning_effort are local-mode " - "parameters โ€” the cloud chat endpoint selects its own model." + "model, max_turns, reasoning_effort and extra_body 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, @@ -516,6 +524,7 @@ def responses( top_p: Optional[float] = None, max_turns: Optional[int] = None, reasoning: Optional[dict[str, Any]] = None, + extra_body: Optional[dict[str, Any]] = None, ) -> Union[dict[str, Any], Iterator[dict[str, Any]]]: """ Document QA over the OpenAI Responses protocol โ€” the agentic surface. @@ -558,6 +567,9 @@ def responses( (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). """ from .cloud_api import CloudAPI if isinstance(self._api, CloudAPI): @@ -569,7 +581,7 @@ def 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, reasoning=reasoning, + max_turns=max_turns, reasoning=reasoning, extra_body=extra_body, ) def messages( @@ -586,6 +598,7 @@ def messages( stop_sequences: Optional[list[str]] = None, max_turns: Optional[int] = None, thinking: Optional[dict[str, Any]] = None, + extra_body: Optional[dict[str, Any]] = None, ) -> Union[dict[str, Any], Iterator[Any]]: """ Document QA over the Anthropic Messages protocol โ€” Claude-native. @@ -623,6 +636,8 @@ def messages( 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. """ from .cloud_api import CloudAPI if isinstance(self._api, CloudAPI): @@ -636,7 +651,7 @@ def messages( 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, + thinking=thinking, extra_body=extra_body, ) # ---------- DOCUMENT MANAGEMENT ---------- diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 54828f6cc..a9d355123 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -306,7 +306,7 @@ def _cache_extra_args(model_name: str) -> Optional[dict]: 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): + reasoning=None, reasoning_effort=None, extra_body=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 @@ -323,6 +323,15 @@ def _openai_agent(client, protocol: str, model_name: str, instructions: str, if reasoning_effort is not None: extra_args = {**(extra_args or {}), "reasoning_effort": reasoning_effort} + 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, @@ -331,8 +340,7 @@ def _openai_agent(client, protocol: str, model_name: str, instructions: str, model_settings=ModelSettings( temperature=temperature, top_p=top_p, reasoning=reasoning, - extra_body=({"prompt_cache_key": cache_key} - if cache_key and openai_backend else None), + extra_body=body, extra_args=extra_args), ) @@ -477,6 +485,7 @@ def run_chat_completions(client, messages, stream: bool = False, model: Optional[str] = None, max_turns: Optional[int] = None, reasoning_effort: Optional[str] = None, + extra_body: Optional[dict] = None, ) -> Union[dict, Iterator[str], Iterator[dict]]: if enable_citations: raise PageIndexAPIError( @@ -495,7 +504,8 @@ def run_chat_completions(client, messages, stream: bool = False, temperature, None, doc_ids=doc_id, cache_key=_conversation_cache_key(model_name, managed, history), - reasoning_effort=reasoning_effort) + reasoning_effort=reasoning_effort, + extra_body=extra_body) run_kwargs = _run_kwargs(max_turns) import openai from agents import Runner @@ -583,6 +593,7 @@ def run_responses(client, input, model: Optional[str] = None, top_p: Optional[float] = None, max_turns: Optional[int] = None, reasoning: Optional[dict] = None, + extra_body: Optional[dict] = None, ) -> Union[dict, Iterator[dict]]: _require_openai_agents("responses") _validate_max_turns(max_turns) @@ -605,7 +616,7 @@ def run_responses(client, input, model: Optional[str] = None, temperature, top_p, doc_ids=doc_id, cache_key=_conversation_cache_key(model_name, managed, conversation), - reasoning=reasoning) + reasoning=reasoning, extra_body=extra_body) run_kwargs = _run_kwargs(max_turns) recorded: dict = {} import openai @@ -817,6 +828,7 @@ def run_messages(client, messages, model: str, stop_sequences: Optional[list[str]] = None, max_turns: Optional[int] = None, thinking: Optional[dict] = None, + extra_body: Optional[dict] = None, ) -> Union[dict, Iterator[Any]]: from .integrations.anthropic_sdk import build_anthropic_tools @@ -834,6 +846,7 @@ def run_messages(client, messages, model: str, 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, }.items() if value is not None} runner = _anthropic_client().beta.messages.tool_runner( max_tokens=(max_tokens if max_tokens is not None diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index 55958d654..639feb515 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -274,6 +274,9 @@ def test_cloud_guards(): 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="not available on PageIndex " "cloud yet"): cloud.responses("x") @@ -928,6 +931,36 @@ def test_reasoning_passthrough_reaches_each_engine(monkeypatch): 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_responses_envelope_echoes_reasoning(client, store_path, fake_model): seed_doc(store_path, "pi-a", "report.pdf") @@ -1368,6 +1401,21 @@ def test_messages_thinking_passes_through(client, fake_anthropic): 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): From c5e8326ddcb454ae757331084fb83723fef02af9 Mon Sep 17 00:00:00 2001 From: Ray Date: Mon, 17 Aug 2026 00:15:51 +0800 Subject: [PATCH 088/137] chore: litellm floor 1.84 -> 1.97.0 1.97.0 is where the unset-effort chatcmpl->responses bridge landed (responses_api_bridge_check's on_constraint_enforcing_endpoint arm, A/B-verified against 1.96.2), so sol-class models work through the chat lane out of the box instead of 400ing until a manual upgrade. Three spots move together: the pyproject floor, the requirements.txt CI pin, and the install hint. --- pageindex/local_chat.py | 2 +- pyproject.toml | 2 +- requirements.txt | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index a9d355123..68dee5fc6 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -252,7 +252,7 @@ def _openai_model(protocol: str, model_name: str): except ImportError: raise PageIndexAPIError( f"'{model_name}' routes through LiteLLM, but litellm is not " - "installed. Run: pip install 'litellm>=1.84'" + "installed. Run: pip install 'litellm>=1.97'" ) wire = model_name.removeprefix("litellm/") if "/" not in wire or wire.startswith("openai/"): diff --git a/pyproject.toml b/pyproject.toml index c9aa729f2..499291d3d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,7 +30,7 @@ requests = ">=2.28.0" openai = ">=1.70.0" # Older releases crash on current openai before the request is sent. openai-agents = ">=0.18.1" -litellm = ">=1.84.0" +litellm = ">=1.97.0" PyPDF2 = ">=3.0.0" pypdfium2 = ">=4.30.0" sortedcontainers = ">=2.4.0" diff --git a/requirements.txt b/requirements.txt index 1516a3979..3202169ce 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ -litellm==1.84.0 +litellm==1.97.0 openai>=1.70.0 requests>=2.28.0 openai-agents>=0.18.1 From 347fdd10a930822da924efc8f298be45c52478cb Mon Sep 17 00:00:00 2001 From: Ray Date: Mon, 17 Aug 2026 00:30:39 +0800 Subject: [PATCH 089/137] feat: named top_p/max_tokens on the chat door, max_output_tokens on responses MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit extra_body could not carry these: openai-agents' LitellmModel passes every ModelSettings sampling field as an explicit keyword and unpacks extra_args into the same call, so the common knobs collided with a bare TypeError on the LiteLLM lane (reproduced against a stub โ€” Python call semantics, callee-independent). Named params ride ModelSettings fields, the one channel clean on every lane; responses() uses the protocol's own name (openai_responses maps ModelSettings.max_tokens to max_output_tokens on the wire) and the envelope now echoes the real value instead of a constant None. Both caps bound each backend call in the agent loop, not the whole run โ€” documented. Cloud rejects them like the other local-only knobs; the long tail (frequency_penalty etc.) stays extra_body-blocked-loudly on that lane by choice. --- pageindex/client.py | 30 ++++++++++++++++++++++-------- pageindex/local_chat.py | 17 +++++++++++------ tests/test_local_chat.py | 34 ++++++++++++++++++++++++++++++++++ 3 files changed, 67 insertions(+), 14 deletions(-) diff --git a/pageindex/client.py b/pageindex/client.py index ddf663908..9a12434da 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -423,6 +423,8 @@ def chat_completions( 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, ) -> Union[dict[str, Any], Iterator[str], Iterator[dict[str, Any]]]: @@ -467,6 +469,11 @@ def chat_completions( 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 @@ -497,15 +504,16 @@ def 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, reasoning_effort=reasoning_effort, - extra_body=extra_body, + max_turns=max_turns, top_p=top_p, max_tokens=max_tokens, + reasoning_effort=reasoning_effort, extra_body=extra_body, ) - if (model is not None or max_turns is not None - or reasoning_effort is not None or extra_body is not None): + 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): raise PageIndexAPIError( - "model, max_turns, reasoning_effort and extra_body are " - "local-mode parameters โ€” the cloud chat endpoint selects " - "its own model." + "model, max_turns, top_p, max_tokens, reasoning_effort and " + "extra_body 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, @@ -523,6 +531,7 @@ def responses( 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, ) -> Union[dict[str, Any], Iterator[dict[str, Any]]]: @@ -563,6 +572,10 @@ def responses( 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 @@ -581,7 +594,8 @@ def 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, reasoning=reasoning, extra_body=extra_body, + max_turns=max_turns, max_output_tokens=max_output_tokens, + reasoning=reasoning, extra_body=extra_body, ) def messages( diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 68dee5fc6..23971a243 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -306,7 +306,8 @@ def _cache_extra_args(model_name: str) -> Optional[dict]: 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): + reasoning=None, reasoning_effort=None, extra_body=None, + max_tokens=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 @@ -338,7 +339,7 @@ def _openai_agent(client, protocol: str, model_name: str, instructions: str, tools=build_openai_tools(client, doc_ids=doc_ids), model=_openai_model(protocol, model_name), model_settings=ModelSettings( - temperature=temperature, top_p=top_p, + temperature=temperature, top_p=top_p, max_tokens=max_tokens, reasoning=reasoning, extra_body=body, extra_args=extra_args), @@ -484,6 +485,8 @@ def run_chat_completions(client, messages, stream: 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, ) -> Union[dict, Iterator[str], Iterator[dict]]: @@ -501,11 +504,11 @@ def run_chat_completions(client, messages, stream: bool = False, reported_model = _reported_model(model_name) managed = _managed_instructions(system_texts) agent = _openai_agent(client, "chat", model_name, managed, - temperature, None, doc_ids=doc_id, + temperature, top_p, doc_ids=doc_id, cache_key=_conversation_cache_key(model_name, managed, history), reasoning_effort=reasoning_effort, - extra_body=extra_body) + extra_body=extra_body, max_tokens=max_tokens) run_kwargs = _run_kwargs(max_turns) import openai from agents import Runner @@ -592,6 +595,7 @@ def run_responses(client, input, model: 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, ) -> Union[dict, Iterator[dict]]: @@ -616,7 +620,8 @@ def run_responses(client, input, model: Optional[str] = None, temperature, top_p, doc_ids=doc_id, cache_key=_conversation_cache_key(model_name, managed, conversation), - reasoning=reasoning, extra_body=extra_body) + reasoning=reasoning, extra_body=extra_body, + max_tokens=max_output_tokens) run_kwargs = _run_kwargs(max_turns) recorded: dict = {} import openai @@ -645,7 +650,7 @@ def envelope(transcript: list, raw_responses) -> dict: "temperature": temperature, "top_p": top_p, "reasoning": reasoning, - "max_output_tokens": None, + "max_output_tokens": max_output_tokens, "error": recorded.get("error"), "incomplete_details": recorded.get("incomplete_details"), "metadata": None, diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index 639feb515..aeb0d2875 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -277,6 +277,11 @@ def test_cloud_guards(): 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="not available on PageIndex " "cloud yet"): cloud.responses("x") @@ -961,6 +966,35 @@ def test_extra_body_passthrough_reaches_each_engine(monkeypatch): 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") From 874210190eeba50ccced4ed556e5bc3fea40f691 Mon Sep 17 00:00:00 2001 From: Ray Date: Mon, 17 Aug 2026 00:33:05 +0800 Subject: [PATCH 090/137] chore: the missing-key error names chat_model as the other exit The default chat_model is what put keyless users on the OpenAI lane, so the error now points at the knob that picks a different backend. --- pageindex/local_chat.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 23971a243..c17619b20 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -260,7 +260,8 @@ def _openai_model(protocol: str, model_name: str): raise PageIndexAPIError( "The OpenAI backend is not configured: set the " "OPENAI_API_KEY environment variable (any value works " - "for keyless OPENAI_BASE_URL servers)." + "for keyless OPENAI_BASE_URL servers), or point " + "chat_model at another provider (e.g. 'anthropic/...')." ) if "/" not in wire: wire = f"openai/{wire}" From 6e7d826f0cb074ad834dc66a6d37be1e7b97fa5e Mon Sep 17 00:00:00 2001 From: Ray Date: Mon, 17 Aug 2026 01:32:19 +0800 Subject: [PATCH 091/137] fix: rebuild litellm's Message/Delta types on Python 3.10 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit litellm 1.97.0 ships Message and Delta annotations whose nested forward refs (ChatCompletionReasoningSummaryTextBlock et al) do not resolve on 3.10, so every completion() dies constructing its response object โ€” non-stream and stream alike (upstream BerriAI/litellm#36384, open, no patch release; 1.96.2 is clean, so the floor raise surfaced it, and pydantic 2.12/2.13 both reproduce). The repair rebuilds the two models once with their defining modules' namespaces at our three completion gateways; version-gated to <3.11 and best-effort, so it is a no-op on healthy interpreters and future fixed litellm releases. Verified on a 3.10 venv: the previously failing anthropic wire test and the full suite pass (250 green, matching CI's matrix leg). --- pageindex/local_chat.py | 2 ++ pageindex/utils.py | 20 ++++++++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index c17619b20..5cf7be538 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -254,6 +254,8 @@ def _openai_model(protocol: str, model_name: str): f"'{model_name}' routes through LiteLLM, but litellm is not " "installed. Run: pip install 'litellm>=1.97'" ) + from .utils import _repair_litellm_types + _repair_litellm_types() wire = model_name.removeprefix("litellm/") if "/" not in wire or wire.startswith("openai/"): if not os.environ.get("OPENAI_API_KEY"): diff --git a/pageindex/utils.py b/pageindex/utils.py index 9e6243c96..d12b32894 100644 --- a/pageindex/utils.py +++ b/pageindex/utils.py @@ -1,5 +1,6 @@ import logging import os +import sys import textwrap from datetime import datetime import time @@ -19,6 +20,23 @@ # litellm is imported inside the functions that use it; eager import is slow # and fetches a remote model-cost map. + +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"): os.environ["OPENAI_API_KEY"] = os.getenv("CHATGPT_API_KEY") @@ -81,6 +99,7 @@ def llm_completion(model, prompt, chat_history=None, return_finish_reason=False) ) else: import litellm + _repair_litellm_types() response = litellm.completion( model=model, messages=messages, @@ -127,6 +146,7 @@ async def llm_acompletion(model, prompt): ) else: import litellm + _repair_litellm_types() response = await litellm.acompletion( model=model, messages=messages, From 4d345bdef44303c3b9a1fdf6570422cca8c372d6 Mon Sep 17 00:00:00 2001 From: Ray Date: Mon, 17 Aug 2026 01:33:40 +0800 Subject: [PATCH 092/137] =?UTF-8?q?feat:=20chat()=20takes=20reasoning=5Fef?= =?UTF-8?q?fort=20=E2=80=94=20the=20front=20door's=20one=20thinking=20knob?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ruled in as a business-level control alongside model: who answers, and how hard it thinks. Same name, values, and verbatim semantics as chat_completions underneath (LiteLLM's cross-provider tier string); unset sends nothing so each backend's own default behavior applies. Sampling and wire-level knobs deliberately stay off the front door. --- pageindex/client.py | 8 +++++++- tests/test_local_chat.py | 24 ++++++++++++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/pageindex/client.py b/pageindex/client.py index 9a12434da..b69168267 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -382,6 +382,7 @@ def chat( 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. @@ -401,13 +402,18 @@ def chat( 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) + doc_id=doc_id, model=model, + reasoning_effort=reasoning_effort) if stream: return cast(Iterator[str], result) envelope = cast(dict[str, Any], result) diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index aeb0d2875..c7c707b38 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -282,6 +282,8 @@ def test_cloud_guards(): 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="not available on PageIndex " "cloud yet"): cloud.responses("x") @@ -403,6 +405,28 @@ def test_chat_multi_turn_history(client, store_path, fake_model): 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") From 7f9b1a485e08d3e485aeac8475e8869c6029b1aa Mon Sep 17 00:00:00 2001 From: Ray Date: Mon, 17 Aug 2026 16:32:53 +0800 Subject: [PATCH 093/137] feat: backend connection overrides; extra_headers on every local door MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two clients, two configs โ€” the gap this closes. index_backend / chat_backend on the constructor (and per-call backend on the chat doors, mirroring per-call model) carry connection params in each lane's own vocabulary, verbatim: the indexing gateways take the dict as LiteLLM call kwargs (a contextvar scopes it per operation, and the env-var key pre-check yields to it), the chat lane lifts api_key / base_url into LitellmModel's two pinned constructor slots and rides the rest as call kwargs, responses() and messages() hand the dict to their SDK client constructors. Per-call keys win over the client's; the openai-SDK fast path normalizes LiteLLM's api_base spelling. Config bundles deliberately don't carry it โ€” you run those in your own environment (docstring says so). extra_headers lands on all three protocol doors, each engine merging caller headers verbatim (anthropic-beta wire-proven on messages). Wire-probed exception, documented on the chat door: LiteLLM's anthropic adapter owns the anthropic-beta header and drops the caller's value โ€” Anthropic beta flags belong on messages(). Cloud rejects the new knobs like the other local-only params, and the docstrings now say credentials belong in backend, never extra_body (on the bare lane they would leak into the JSON body without touching auth โ€” wire-probed). --- pageindex/client.py | 65 ++++++++++++++++++++--- pageindex/local_api.py | 29 +++++++--- pageindex/local_chat.py | 60 ++++++++++++++++----- pageindex/utils.py | 49 +++++++++++++---- tests/test_client.py | 69 ++++++++++++++++++++++++ tests/test_local_chat.py | 111 +++++++++++++++++++++++++++++++++++++-- 6 files changed, 339 insertions(+), 44 deletions(-) diff --git a/pageindex/client.py b/pageindex/client.py index b69168267..cfa68537d 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -113,6 +113,8 @@ def __init__( 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( @@ -123,7 +125,9 @@ def __init__( "model": model, "summary_model": summary_model, "retrieve_model": retrieve_model} if api_key is not None: - local_only = dict(model_args, 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( @@ -143,6 +147,7 @@ def __init__( self.index_model = opt.index_model self.summary_model = opt.summary_model self.chat_model = _agents_sdk_model_name(opt.chat_model) + self.chat_backend = chat_backend self.storage_path = storage_path or ".pageindex" from .local_api import LocalAPI self._api = LocalAPI( @@ -150,6 +155,7 @@ def __init__( model=self.model, summary_model=self.summary_model, retrieve_model=self.chat_model, + index_backend=index_backend, ) # LiteLLM's multi-second import would otherwise land on the # first chat call; failures resurface there with real context. @@ -433,6 +439,8 @@ def chat_completions( 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: document QA in one call. @@ -490,6 +498,17 @@ def chat_completions( 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', @@ -512,14 +531,16 @@ def chat_completions( 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_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 and " - "extra_body are local-mode parameters โ€” the cloud chat " - "endpoint selects its own model." + "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, @@ -540,6 +561,8 @@ def responses( 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. @@ -588,7 +611,15 @@ def responses( 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). + 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): @@ -602,6 +633,7 @@ def responses( 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( @@ -619,6 +651,8 @@ def messages( 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. @@ -658,6 +692,15 @@ def messages( 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): @@ -672,6 +715,7 @@ def messages( 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 ---------- @@ -825,7 +869,9 @@ def openai_agent_config( ``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. + 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. Args: doc_id: Document ID or list of IDs to target, as in @@ -1107,7 +1153,10 @@ def __init__( 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, index_model=index_model, chat_model=chat_model, model=model, summary_model=summary_model, - retrieve_model=retrieve_model, storage_path=storage_path) + retrieve_model=retrieve_model, storage_path=storage_path, + index_backend=index_backend, chat_backend=chat_backend) diff --git a/pageindex/local_api.py b/pageindex/local_api.py index 9a82c48f9..88ad02ed0 100644 --- a/pageindex/local_api.py +++ b/pageindex/local_api.py @@ -35,14 +35,25 @@ 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): + 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_indexer 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( @@ -104,11 +115,12 @@ def submit_document( try: if mode == "flash": structure, description = _run_indexer( - self._index_flash, file_path, page_texts + self._with_backend, self._index_flash, file_path, page_texts ) else: structure, description = _run_indexer( - self._index_standard, file_path, page_texts + self._with_backend, self._index_standard, file_path, + page_texts ) except PageIndexAPIError: raise @@ -183,11 +195,12 @@ def _index_flash(self, file_path: str, page_texts: list[str]) -> tuple[list, str from .utils import (add_node_text, create_clean_structure_for_description, generate_doc_description, write_node_id) import litellm - env = litellm.validate_environment(self._summary_model) - if not env["keys_in_environment"]: - raise PageIndexAPIError( - f"Failed to submit document: missing API key for " - f"{self._summary_model}: {', '.join(env['missing_keys'])}") + if not self._index_backend: + env = litellm.validate_environment(self._summary_model) + if not env["keys_in_environment"]: + raise PageIndexAPIError( + f"Failed to submit document: missing API key for " + f"{self._summary_model}: {', '.join(env['missing_keys'])}") result = page_index_flash(file_path, summary=True, summary_model=self._summary_model, optimize="full", diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 5cf7be538..ed66081a9 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -212,7 +212,7 @@ def _require_openai_agents(method: str) -> None: ) from exc -def _openai_model(protocol: str, model_name: str): +def _openai_model(protocol: str, model_name: str, backend=None): """The backend protocol driver โ€” the seam tests replace with a fake. chat protocol: LiteLLM, full stop โ€” model names mean what LiteLLM says @@ -240,12 +240,12 @@ def _openai_model(protocol: str, model_name: str): import openai model_name = model_name.removeprefix("openai/") try: - backend = openai.AsyncOpenAI() - except openai.OpenAIError as exc: + sdk_client = openai.AsyncOpenAI(**(backend or {})) + except (openai.OpenAIError, TypeError) as exc: raise PageIndexAPIError( f"The OpenAI backend is not configured: {exc}") from exc from agents.models.openai_responses import OpenAIResponsesModel - return OpenAIResponsesModel(model_name, openai_client=backend) + return OpenAIResponsesModel(model_name, openai_client=sdk_client) try: from agents.extensions.models.litellm_model import LitellmModel import litellm @@ -258,7 +258,8 @@ def _openai_model(protocol: str, model_name: str): _repair_litellm_types() wire = model_name.removeprefix("litellm/") if "/" not in wire or wire.startswith("openai/"): - if not os.environ.get("OPENAI_API_KEY"): + if (not os.environ.get("OPENAI_API_KEY") + and not (backend or {}).get("api_key")): raise PageIndexAPIError( "The OpenAI backend is not configured: set the " "OPENAI_API_KEY environment variable (any value works " @@ -276,7 +277,8 @@ def _openai_model(protocol: str, model_name: str): f"model id, use 'openai/{wire}' and point OPENAI_BASE_URL " "at the server." ) - return LitellmModel(wire) + return LitellmModel(wire, api_key=(backend or {}).get("api_key"), + base_url=(backend or {}).get("base_url")) def _reported_model(model_name: str) -> str: @@ -307,10 +309,18 @@ def _cache_extra_args(model_name: str) -> Optional[dict]: return None +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): + 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 @@ -327,6 +337,16 @@ def _openai_agent(client, protocol: str, model_name: str, instructions: str, if reasoning_effort is not None: extra_args = {**(extra_args or {}), "reasoning_effort": reasoning_effort} + conn = dict(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", conn.pop("api_base", 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 @@ -340,11 +360,12 @@ def _openai_agent(client, protocol: str, model_name: str, instructions: str, name="PageIndex", instructions=instructions, tools=build_openai_tools(client, doc_ids=doc_ids), - model=_openai_model(protocol, model_name), + model=_openai_model(protocol, model_name, conn or None), model_settings=ModelSettings( temperature=temperature, top_p=top_p, max_tokens=max_tokens, reasoning=reasoning, extra_body=body, + extra_headers=extra_headers, extra_args=extra_args), ) @@ -492,6 +513,8 @@ def run_chat_completions(client, messages, stream: bool = False, 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( @@ -511,7 +534,9 @@ def run_chat_completions(client, messages, stream: bool = False, cache_key=_conversation_cache_key(model_name, managed, history), reasoning_effort=reasoning_effort, - extra_body=extra_body, max_tokens=max_tokens) + extra_body=extra_body, max_tokens=max_tokens, + backend=_merged_backend(client, backend), + extra_headers=extra_headers) run_kwargs = _run_kwargs(max_turns) import openai from agents import Runner @@ -601,6 +626,8 @@ def run_responses(client, input, model: Optional[str] = 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) @@ -624,7 +651,9 @@ def run_responses(client, input, model: Optional[str] = None, cache_key=_conversation_cache_key(model_name, managed, conversation), reasoning=reasoning, extra_body=extra_body, - max_tokens=max_output_tokens) + max_tokens=max_output_tokens, + backend=_merged_backend(client, backend), + extra_headers=extra_headers) run_kwargs = _run_kwargs(max_turns) recorded: dict = {} import openai @@ -759,10 +788,10 @@ def _require_anthropic() -> None: ) from exc -def _anthropic_client(): +def _anthropic_client(backend=None): """The backend client โ€” the seam tests replace with a fake transport.""" import anthropic - return anthropic.Anthropic() + return anthropic.Anthropic(**(backend or {})) def _anthropic_system(extra_system, block: Optional[str]) -> list[dict]: @@ -837,6 +866,8 @@ def run_messages(client, messages, model: str, 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 @@ -854,9 +885,10 @@ def run_messages(client, messages, model: str, 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_body": extra_body, "extra_headers": extra_headers, }.items() if value is not None} - runner = _anthropic_client().beta.messages.tool_runner( + runner = _anthropic_client(_merged_backend(client, backend)) \ + .beta.messages.tool_runner( max_tokens=(max_tokens if max_tokens is not None else _default_max_tokens(model)), messages=prepared, diff --git a/pageindex/utils.py b/pageindex/utils.py index d12b32894..114f04668 100644 --- a/pageindex/utils.py +++ b/pageindex/utils.py @@ -1,3 +1,4 @@ +import contextvars import logging import os import sys @@ -21,6 +22,20 @@ # 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 _openai_sdk_kwargs(backend: dict) -> dict: + """The same backend dict works on both gateway paths: LiteLLM accepts + either endpoint spelling, the openai SDK only ``base_url``.""" + return {("base_url" if key == "api_base" else key): value + for key, value in backend.items()} + + 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 @@ -85,15 +100,21 @@ def llm_completion(model, prompt, chat_history=None, return_finish_reason=False) model = _strip_prefix(model, "openai/") max_retries = 10 messages = list(chat_history) + [{"role": "user", "content": prompt}] if chat_history else [{"role": "user", "content": prompt}] + backend = _llm_backend.get() if use_openai_sdk: - global _openai_sync_client - if _openai_sync_client is None: - import openai - _openai_sync_client = openai.OpenAI(max_retries=0) + import openai + if backend: + oai_client = openai.OpenAI(**{"max_retries": 0, + **_openai_sdk_kwargs(backend)}) + else: + global _openai_sync_client + if _openai_sync_client is None: + _openai_sync_client = openai.OpenAI(max_retries=0) + oai_client = _openai_sync_client for i in range(max_retries): try: if use_openai_sdk: - response = _openai_sync_client.chat.completions.create( + response = oai_client.chat.completions.create( model=model, messages=messages, ) @@ -105,6 +126,7 @@ def llm_completion(model, prompt, chat_history=None, return_finish_reason=False) messages=messages, temperature=0, drop_params=True, + **(backend or {}), ) content = response.choices[0].message.content if return_finish_reason: @@ -132,15 +154,21 @@ async def llm_acompletion(model, prompt): model = _strip_prefix(model, "openai/") max_retries = 10 messages = [{"role": "user", "content": prompt}] + backend = _llm_backend.get() if use_openai_sdk: - global _openai_async_client - if _openai_async_client is None: - import openai - _openai_async_client = openai.AsyncOpenAI(max_retries=0) + import openai + if backend: + oai_client = openai.AsyncOpenAI(**{"max_retries": 0, + **_openai_sdk_kwargs(backend)}) + else: + global _openai_async_client + if _openai_async_client is None: + _openai_async_client = openai.AsyncOpenAI(max_retries=0) + oai_client = _openai_async_client for i in range(max_retries): try: if use_openai_sdk: - response = await _openai_async_client.chat.completions.create( + response = await oai_client.chat.completions.create( model=model, messages=messages, ) @@ -152,6 +180,7 @@ async def llm_acompletion(model, prompt): messages=messages, temperature=0, drop_params=True, + **(backend or {}), ) return response.choices[0].message.content except Exception as e: diff --git a/tests/test_client.py b/tests/test_client.py index 2feb9a89a..d3c17ef52 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -835,3 +835,72 @@ def test_parse_pages_overlap_counts_union(): assert len(pages) == 9000 and pages[0] == 1 and pages[-1] == 9000 with pytest.raises(ValueError, match="spans more than"): _parse_pages("1-10001") + + +# โ”€โ”€ backend: the indexing lane โ”€โ”€ + +def test_backend_scopes_the_index_lane(tmp_path, monkeypatch): + """index_backend reaches both gateway paths โ€” LiteLLM's call kwargs + and a fresh openai-SDK client โ€” scoped to the operation, with the + endpoint spelling normalized for the openai SDK.""" + pytest.importorskip("litellm") + import litellm + import openai + 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", retrieve_model="r", + 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]) + api._with_backend(lambda: llm_completion("anthropic/claude-x", "p")) + assert captured["api_key"] == "ik" + assert captured["api_base"] == "http://b" + + seen = {} + + class _FakeOpenAI: + def __init__(self, **kw): + seen.update(kw) + self.chat = SimpleNamespace(completions=SimpleNamespace( + create=lambda **_: reply)) + + monkeypatch.setattr(openai, "OpenAI", _FakeOpenAI) + api._with_backend(lambda: llm_completion("gpt-4o", "p")) + assert seen["api_key"] == "ik" + assert seen["base_url"] == "http://b" + + +def test_index_backend_skips_the_env_precheck(tmp_path, monkeypatch): + """The missing-key pre-check reads environment variables only, so a + backend-supplied key must bypass it instead of being refused.""" + pytest.importorskip("litellm") + import litellm + from pageindex.local_api import LocalAPI + + api = LocalAPI(storage_path=str(tmp_path / "s"), model="m", + summary_model="s", retrieve_model="r", + index_backend={"api_key": "ik"}) + monkeypatch.setattr(litellm, "validate_environment", + lambda *a, **k: pytest.fail("env pre-check ran")) + monkeypatch.setattr(pageindex.flash, "page_index_flash", + lambda *a, **k: (_ for _ in ()).throw( + RuntimeError("reached-flash"))) + with pytest.raises(RuntimeError, match="reached-flash"): + api._index_flash("f.pdf", ["text"]) + + +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"}) diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index c7c707b38..ea9eb7a09 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -159,8 +159,9 @@ def install(turns): fake = FakeModel(turns) state["protocols"] = [] - def factory(protocol, model_name): + 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) @@ -284,6 +285,12 @@ def test_cloud_guards(): 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") @@ -650,7 +657,8 @@ def handler(request): fake = anthropic.Anthropic( api_key="test", http_client=httpx.Client(transport=httpx.MockTransport(handler))) - monkeypatch.setattr(local_chat, "_anthropic_client", lambda: fake) + monkeypatch.setattr(local_chat, "_anthropic_client", + lambda backend=None: fake) return state["calls"] return install @@ -1295,7 +1303,7 @@ def test_responses_stream_backend_terminal_states_are_events( fake = _TerminalModel([[]]) fake.terminal = terminal monkeypatch.setattr(local_chat, "_openai_model", - lambda protocol, model_name: fake) + 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] @@ -1354,7 +1362,8 @@ def handler(request): fake = anthropic.Anthropic( api_key="test", max_retries=0, http_client=httpx.Client(transport=httpx.MockTransport(handler))) - monkeypatch.setattr(local_chat, "_anthropic_client", lambda: fake) + 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"): @@ -1604,3 +1613,97 @@ def test_messages_edge_validation(client, store_path, fake_anthropic): 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" + + +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" + + 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_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" From e13332d8915e823e840f42eee830eb098923e19a Mon Sep 17 00:00:00 2001 From: Ray Date: Mon, 17 Aug 2026 17:02:42 +0800 Subject: [PATCH 094/137] docs: constructor Args list gains index_backend / chat_backend The class docstring documents every constructor argument; the backend wave added two without entries. Same phrasing as the per-call docs: index lane is LiteLLM vocabulary verbatim, chat_backend reaches whichever door runs (api_key/base_url portable across all three). --- pageindex/client.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/pageindex/client.py b/pageindex/client.py index cfa68537d..8125aa525 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -87,6 +87,15 @@ class PageIndexClient: ``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 From 0553ac20aefd27d0c93e051eea9e1316c2924261 Mon Sep 17 00:00:00 2001 From: Ray Date: Mon, 17 Aug 2026 17:13:36 +0800 Subject: [PATCH 095/137] chore: three audit leftovers The classic pipeline's ThreadPoolExecutor import (unused since the 0.2.9 merge) goes. The bare-model missing-key error now also names the backend={'api_key': ...} route, which satisfies the same check. messages() joins the uniform: a bad backend dict wraps as PageIndexAPIError like the responses door, instead of leaking the SDK's raw constructor error. --- pageindex/local_chat.py | 13 +++++++++---- pageindex/page_index_classic.py | 1 - tests/test_local_chat.py | 7 +++++++ 3 files changed, 16 insertions(+), 5 deletions(-) diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index ed66081a9..4f9edb543 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -262,9 +262,10 @@ def _openai_model(protocol: str, model_name: str, backend=None): and not (backend or {}).get("api_key")): raise PageIndexAPIError( "The OpenAI backend is not configured: set the " - "OPENAI_API_KEY environment variable (any value works " - "for keyless OPENAI_BASE_URL servers), or point " - "chat_model at another provider (e.g. 'anthropic/...')." + "OPENAI_API_KEY environment variable or pass " + "backend={'api_key': ...} (any value works for keyless " + "OPENAI_BASE_URL servers), or point chat_model at " + "another provider (e.g. 'anthropic/...')." ) if "/" not in wire: wire = f"openai/{wire}" @@ -791,7 +792,11 @@ def _require_anthropic() -> None: def _anthropic_client(backend=None): """The backend client โ€” the seam tests replace with a fake transport.""" import anthropic - return anthropic.Anthropic(**(backend or {})) + try: + return anthropic.Anthropic(**(backend or {})) + except (anthropic.AnthropicError, 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]: 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/tests/test_local_chat.py b/tests/test_local_chat.py index ea9eb7a09..d34ff9a47 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -1669,6 +1669,13 @@ def test_messages_backend_merges_and_reaches_the_client(client, fake_anthropic, 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 From 4430eba8fc2818185b01c2f88483e4fc668aa205 Mon Sep 17 00:00:00 2001 From: Ray Date: Mon, 17 Aug 2026 17:18:38 +0800 Subject: [PATCH 096/137] chore: narrow the anthropic wrap to TypeError; key advice names chat_backend MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Anthropic's constructor raises only TypeError in our supported range (unknown kwargs, conflicting credentials) โ€” a missing key defers to request time, so catching AnthropicError there guarded an impossible case. And chat() has no backend parameter, so the missing-key advice now names chat_backend alongside the per-call route. --- pageindex/local_chat.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 4f9edb543..5b9bfb406 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -262,8 +262,8 @@ def _openai_model(protocol: str, model_name: str, backend=None): and not (backend or {}).get("api_key")): raise PageIndexAPIError( "The OpenAI backend is not configured: set the " - "OPENAI_API_KEY environment variable or pass " - "backend={'api_key': ...} (any value works for keyless " + "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/...')." ) @@ -794,7 +794,7 @@ def _anthropic_client(backend=None): import anthropic try: return anthropic.Anthropic(**(backend or {})) - except (anthropic.AnthropicError, TypeError) as exc: + except TypeError as exc: raise PageIndexAPIError( f"The Anthropic backend is not configured: {exc}") from exc From 7cbc3be4b7706e2f6d86d45fa3e9664e4719cca2 Mon Sep 17 00:00:00 2001 From: Ray Date: Mon, 17 Aug 2026 18:21:55 +0800 Subject: [PATCH 097/137] fix: the bundle's per-call model override speaks chat_model's grammar MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit openai_agent_config(model=) went into the config raw while the constructor default was normalized for the Agents SDK's prefix resolver โ€” the same slashed name worked from one knob and raised UserError('Unknown prefix') from the other. The normalizer is idempotent and prefix-preserving, so every input that worked before resolves identically; only guaranteed-error inputs become working. Found by an external agent review; the grammar test now covers the override path it missed. --- pageindex/client.py | 6 ++++-- tests/test_agent_tools.py | 5 +++++ 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/pageindex/client.py b/pageindex/client.py index 8125aa525..bf3b09a24 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -889,7 +889,9 @@ def openai_agent_config( (tool scoping is server-side). include_management (bool): Also expose tools that modify the library. - model: Backend model name; overrides the local default. + 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) @@ -901,7 +903,7 @@ def openai_agent_config( } model = model or getattr(self, "chat_model", None) if model: - config["model"] = model + config["model"] = _agents_sdk_model_name(model) return config def as_anthropic_tools(self, include_management: bool = False, diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 3873cd36e..f92a477ee 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -763,6 +763,11 @@ def test_openai_agent_config_model_speaks_the_agents_sdk_grammar(tmp_path): 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_openai_agent_config_cloud_omits_model(cloud_with_fake_bridge): From 57a77f63f876bdcf418fe4d5d01242b8c4666164 Mon Sep 17 00:00:00 2001 From: Ray Date: Mon, 17 Aug 2026 19:02:48 +0800 Subject: [PATCH 098/137] fix: support pypdfium2 5.x (GetFontName removed, bookmark API changed) pypdfium2 5.x removed FPDFFont_GetFontName (now GetBaseFontName) and changed PdfBookmark from attribute access to method calls. The existing >=4.30.0 lower bound already resolved to 5.x, breaking fresh installs. --- pageindex/flash/embedded_toc.py | 12 +++++++++--- pageindex/flash/parser_pdfium_charlevel/geometry.py | 4 +++- requirements.txt | 2 +- 3 files changed, 13 insertions(+), 5 deletions(-) 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/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/requirements.txt b/requirements.txt index 3202169ce..2406eef11 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,7 +4,7 @@ requests>=2.28.0 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 From 6a9104a6a374f41247a3f0b8d05648f38ec8339f Mon Sep 17 00:00:00 2001 From: Ray Date: Mon, 17 Aug 2026 19:17:32 +0800 Subject: [PATCH 099/137] fix: one openai client per backend, not one per call MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit index_backend opted straight out of the module-level client cache 548d3201 added: every llm_completion / llm_acompletion built a fresh openai.OpenAI, and the indexing lane calls those once per node. A 200-node PDF paid ~46 ms of SSL-context construction each time (~9 s of it blocking inside the summarize coroutines) and reused no connection, re-handshaking TLS per summary. Output was always correct and the clients close on refcount, so this was waste, not a leak. The two bare globals become one dict keyed by (is_async, backend), so the default path keeps the singleton it had and every distinct backend gets its own reused client. The key is repr(sorted(items)) rather than a JSON dump because a backend may legitimately carry a non-serializable value (http_client=httpx.Client(...)), and sorting item tuples only ever compares the unique string keys, never the values. Capped at 8 entries and cleared wholesale past that: a backend can carry per-tenant credentials, and thrashing degrades to today's build-every-call rather than growing without bound. is_async is keyword-only so neither call site reads as a bare boolean, and both gateway branches collapse to one line each โ€” product code is net shorter. Tests swap the whole dict via monkeypatch rather than clearing it, so the fakes they cache are restored away with the attribute the way the old globals were; clearing left them in the module, including under the no-backend key every default call reads. test_backend_scopes_the_index_lane also relied on construction happening per call to capture the kwargs, so it swaps too and its docstring drops "fresh". --- pageindex/utils.py | 38 ++++++++++++++++++-------------------- tests/test_client.py | 35 +++++++++++++++++++++++++++++++---- 2 files changed, 49 insertions(+), 24 deletions(-) diff --git a/pageindex/utils.py b/pageindex/utils.py index 114f04668..74cfad1cb 100644 --- a/pageindex/utils.py +++ b/pageindex/utils.py @@ -77,8 +77,22 @@ def _is_openai_model(model): return '/' not in model or model.startswith('openai/') -_openai_sync_client = None -_openai_async_client = None +_openai_clients: dict = {} + + +def _openai_client(backend, *, is_async: bool = False): + """One client per distinct backend, reused: constructing one rebuilds the + SSL context and opens a fresh connection pool, and the indexing lane calls + this once per node. Capped, since a backend can carry per-tenant keys.""" + import openai + key = (is_async, repr(sorted((backend or {}).items()))) + if key not in _openai_clients: + if len(_openai_clients) >= 8: + _openai_clients.clear() + cls = openai.AsyncOpenAI if is_async else openai.OpenAI + _openai_clients[key] = cls(**{"max_retries": 0, + **_openai_sdk_kwargs(backend or {})}) + return _openai_clients[key] # Misconfiguration: no retry can fix a rejected key or a model that does not @@ -102,15 +116,7 @@ def llm_completion(model, prompt, chat_history=None, return_finish_reason=False) messages = list(chat_history) + [{"role": "user", "content": prompt}] if chat_history else [{"role": "user", "content": prompt}] backend = _llm_backend.get() if use_openai_sdk: - import openai - if backend: - oai_client = openai.OpenAI(**{"max_retries": 0, - **_openai_sdk_kwargs(backend)}) - else: - global _openai_sync_client - if _openai_sync_client is None: - _openai_sync_client = openai.OpenAI(max_retries=0) - oai_client = _openai_sync_client + oai_client = _openai_client(backend) for i in range(max_retries): try: if use_openai_sdk: @@ -156,15 +162,7 @@ async def llm_acompletion(model, prompt): messages = [{"role": "user", "content": prompt}] backend = _llm_backend.get() if use_openai_sdk: - import openai - if backend: - oai_client = openai.AsyncOpenAI(**{"max_retries": 0, - **_openai_sdk_kwargs(backend)}) - else: - global _openai_async_client - if _openai_async_client is None: - _openai_async_client = openai.AsyncOpenAI(max_retries=0) - oai_client = _openai_async_client + oai_client = _openai_client(backend, is_async=True) for i in range(max_retries): try: if use_openai_sdk: diff --git a/tests/test_client.py b/tests/test_client.py index d3c17ef52..4ea9a12fc 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -286,8 +286,7 @@ def test_page_index_flash_rejects_unknown_optimize(): def test_llm_completion_missing_key_raises_immediately(monkeypatch): import openai monkeypatch.delenv("OPENAI_API_KEY", raising=False) - monkeypatch.setattr(pageindex.utils, "_openai_sync_client", None) - monkeypatch.setattr(pageindex.utils, "_openai_async_client", None) + monkeypatch.setattr(pageindex.utils, "_openai_clients", {}) with pytest.raises(openai.OpenAIError): pageindex.utils.llm_completion("gpt-4o", "probe") with pytest.raises(openai.OpenAIError): @@ -296,7 +295,7 @@ def test_llm_completion_missing_key_raises_immediately(monkeypatch): def test_submit_missing_llm_key_fails_loud(local_client, sample_pdf, monkeypatch): monkeypatch.delenv("OPENAI_API_KEY", raising=False) - monkeypatch.setattr(pageindex.utils, "_openai_sync_client", None) + monkeypatch.setattr(pageindex.utils, "_openai_clients", {}) 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) @@ -841,7 +840,7 @@ def test_parse_pages_overlap_counts_union(): def test_backend_scopes_the_index_lane(tmp_path, monkeypatch): """index_backend reaches both gateway paths โ€” LiteLLM's call kwargs - and a fresh openai-SDK client โ€” scoped to the operation, with the + and the openai-SDK client โ€” scoped to the operation, with the endpoint spelling normalized for the openai SDK.""" pytest.importorskip("litellm") import litellm @@ -875,11 +874,39 @@ def __init__(self, **kw): create=lambda **_: reply)) monkeypatch.setattr(openai, "OpenAI", _FakeOpenAI) + monkeypatch.setattr(pageindex.utils, "_openai_clients", {}) api._with_backend(lambda: llm_completion("gpt-4o", "p")) assert seen["api_key"] == "ik" assert seen["base_url"] == "http://b" +def test_openai_client_is_reused_per_backend(monkeypatch): + """One client per distinct backend, built once โ€” the indexing lane calls + the gateway once per node, and rebuilding costs an SSL context each time.""" + import openai + from pageindex.utils import _openai_client + + built = [] + + class _FakeOpenAI: + def __init__(self, **kw): + built.append(kw) + + monkeypatch.setattr(openai, "OpenAI", _FakeOpenAI) + monkeypatch.setattr(openai, "AsyncOpenAI", _FakeOpenAI) + monkeypatch.setattr(pageindex.utils, "_openai_clients", {}) + + one = _openai_client({"api_key": "k", "api_base": "http://b"}) + assert _openai_client({"api_base": "http://b", "api_key": "k"}) is one + assert len(built) == 1 and built[0]["base_url"] == "http://b" + + assert _openai_client({"api_key": "other"}) is not one + assert _openai_client({"api_key": "k", "api_base": "http://b"}, + is_async=True) is not one + assert _openai_client(None) is not one + assert len(built) == 4 + + def test_index_backend_skips_the_env_precheck(tmp_path, monkeypatch): """The missing-key pre-check reads environment variables only, so a backend-supplied key must bypass it instead of being refused.""" From c4ffa3142bcd10d080529c180cdfc42039dd9bcc Mon Sep 17 00:00:00 2001 From: Ray Date: Mon, 17 Aug 2026 20:21:40 +0800 Subject: [PATCH 100/137] fix: reassemble UTF-16 surrogate pairs from PDFium's FPDFText_GetUnicode MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PDFium returns astral characters (U+10000+, e.g. mathematical italic ๐‘ž) as two consecutive UTF-16 surrogate code units instead of one UTF-32 codepoint. The lone surrogates survive through the pipeline and cause UnicodeEncodeError when the OpenAI SDK serializes the summary prompt (ensure_ascii=False + .encode('utf-8')), silently producing empty summaries for every node whose pages contain these characters. Detect high surrogates at the FPDFText_GetUnicode read site, peek the next textpage slot for the low surrogate, and combine into the full codepoint before chr(). Downstream processing (glyph width, span merge, heading detection) now sees one correct character instead of two invalid ones. --- .../flash/parser_pdfium_charlevel/char_extract.py | 11 +++++++++++ 1 file changed, 11 insertions(+) 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 From 8e059aa908ef5da48edb2d2b6f484c306416e208 Mon Sep 17 00:00:00 2001 From: Ray Date: Mon, 17 Aug 2026 20:23:19 +0800 Subject: [PATCH 101/137] fix: CLI flash precheck resolves the default model; summarize_tree stops masking fatal errors MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two fixes, same root: a missing or invalid API key on the default CLI command could silently produce a tree with empty summaries. run_pageindex.py โ€” the missing-key pre-check never fired on `python3 run_pageindex.py --pdf_path doc.pdf` because summary_model was None (no CLI flag) and `if summary_model and (...)` short-circuited. Now resolved via ConfigLoader โ€” the same source page_index_flash uses internally โ€” so the pre-check and execution always agree on the model. Also fixes a pre-existing issue where optimize_model=None was passed to flash's expand stage. pageindex/utils.py โ€” summarize_tree's visit() caught every exception (including 401 auth errors) and set summary="". A single leaf under 200 tokens โ€” whose raw text is used as the summary without an LLM call โ€” then satisfied _any_summary, so the function returned "successfully" with one real summary and the rest empty. Now visit() re-raises unrecoverable errors (401/403/404), and the gather results are checked before _any_summary. Transient errors (429, 500, timeout) still degrade gracefully. --- pageindex/utils.py | 11 ++++++++--- run_pageindex.py | 6 ++++-- 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/pageindex/utils.py b/pageindex/utils.py index 74cfad1cb..2fd2a5197 100644 --- a/pageindex/utils.py +++ b/pageindex/utils.py @@ -891,11 +891,16 @@ async def visit(node): 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 []) diff --git a/run_pageindex.py b/run_pageindex.py index 1cbea0c72..f91f2b751 100644 --- a/run_pageindex.py +++ b/run_pageindex.py @@ -90,9 +90,11 @@ if args.mode == 'flash': from pageindex.flash import page_index_flash - summary_model = args.summary_model or args.index_model or args.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 summary_model and (will_summarize or args.optimize == 'full'): + if will_summarize or args.optimize == 'full': import litellm env = litellm.validate_environment(summary_model) if not env["keys_in_environment"]: From 78d44b4b18064b76b8558f883ce3498c177dd940 Mon Sep 17 00:00:00 2001 From: Ray Date: Mon, 17 Aug 2026 21:50:22 +0800 Subject: [PATCH 102/137] =?UTF-8?q?fix:=20the=20indexing=20lane=20routes?= =?UTF-8?q?=20every=20model=20through=20LiteLLM=20=E2=80=94=20direct=20lan?= =?UTF-8?q?e=20removed?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit index_backend documented LiteLLM's connection params (api_key, api_base, api_version, aws_*, verbatim), but bare model names โ€” including the default โ€” took an openai-SDK direct lane whose constructor rejects everything beyond api_key/base_url: TypeError on the first summary call. Root fix per Ray's chat-lane ruling (ad239a9, 82e3354): a routing decision must never hide in a model-name prefix. The direct lane (548d320) is gone; indexing names now share the chat lane's grammar โ€” litellm/ strips, bare names are OpenAI-compatible shorthand (wire form openai/), and a first segment LiteLLM does not know is refused with the openai/ escape instead of burning the retry loop on 400s. _litellm_model also fail-fasts on a missing provider key. litellm reports one as a retryable 500, which would spin the 10-attempt loop and then be absorbed per-node by the summary and optimize passes; the guard's 401 (and the provider guard's 404) are unrecoverable to those layers, so misconfiguration fails loud. Sitting in the shared lane, it replaces the flash-only preflight in local_api and now also covers _index_standard, the classic CLI path, and tree_optimize โ€” whose own precheck broadens from OPENAI_API_KEY-only to validate_environment. temperature=0 does not return with the LiteLLM call: the direct lane existed precisely because gpt-5.x rejects it and litellm 1.97.0's map does not know gpt-5.6-luna, so drop_params cannot save it. Provider- prefixed indexing models therefore move from temperature=0 to the provider default. Supersedes the 6a9104a client cache โ€” litellm manages its own connection pools. Tests pin the wire-name matrix (bare, openai/, litellm/, provider- prefixed, unknown-bare), the provider refusal, and backend bypassing the env pre-check. The missing-key tests import litellm before delenv: its first import may load a .env and would hand the deleted key back. --- pageindex/config.yaml | 4 +- pageindex/local_api.py | 7 --- pageindex/tree_optimize.py | 11 ++-- pageindex/utils.py | 125 +++++++++++++++---------------------- tests/test_client.py | 98 +++++++++-------------------- 5 files changed, 88 insertions(+), 157 deletions(-) diff --git a/pageindex/config.yaml b/pageindex/config.yaml index 4fde6f60e..d3786ca12 100644 --- a/pageindex/config.yaml +++ b/pageindex/config.yaml @@ -2,8 +2,8 @@ # 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. -# Indexing names without a provider prefix use the OpenAI SDK directly; -# for other providers, use "provider/model" (e.g. "anthropic/claude-sonnet-4-6"). +# 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 diff --git a/pageindex/local_api.py b/pageindex/local_api.py index 88ad02ed0..180b13516 100644 --- a/pageindex/local_api.py +++ b/pageindex/local_api.py @@ -194,13 +194,6 @@ def _index_flash(self, file_path: str, page_texts: list[str]) -> tuple[list, str from .flash import page_index_flash from .utils import (add_node_text, create_clean_structure_for_description, generate_doc_description, write_node_id) - import litellm - if not self._index_backend: - env = litellm.validate_environment(self._summary_model) - if not env["keys_in_environment"]: - raise PageIndexAPIError( - f"Failed to submit document: missing API key for " - f"{self._summary_model}: {', '.join(env['missing_keys'])}") result = page_index_flash(file_path, summary=True, summary_model=self._summary_model, optimize="full", diff --git a/pageindex/tree_optimize.py b/pageindex/tree_optimize.py index 04719ccb2..026f69dde 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, llm_acompletion, strip_internal_keys) TRIGGER_PAGES = 5 # only look ahead on nodes larger than this @@ -872,9 +872,12 @@ 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: + import litellm + env = litellm.validate_environment(model) + if not env["keys_in_environment"]: + sys.exit(f"{', '.join(env['missing_keys'])} 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 2fd2a5197..ee5f549e8 100644 --- a/pageindex/utils.py +++ b/pageindex/utils.py @@ -29,13 +29,6 @@ "pageindex_llm_backend", default=None) -def _openai_sdk_kwargs(backend: dict) -> dict: - """The same backend dict works on both gateway paths: LiteLLM accepts - either endpoint spelling, the openai SDK only ``base_url``.""" - return {("base_url" if key == "api_base" else key): value - for key, value in backend.items()} - - 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 @@ -69,30 +62,38 @@ 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_clients: dict = {} - - -def _openai_client(backend, *, is_async: bool = False): - """One client per distinct backend, reused: constructing one rebuilds the - SSL context and opens a fresh connection pool, and the indexing lane calls - this once per node. Capped, since a backend can carry per-tenant keys.""" - import openai - key = (is_async, repr(sorted((backend or {}).items()))) - if key not in _openai_clients: - if len(_openai_clients) >= 8: - _openai_clients.clear() - cls = openai.AsyncOpenAI if is_async else openai.OpenAI - _openai_clients[key] = cls(**{"max_retries": 0, - **_openai_sdk_kwargs(backend or {})}) - return _openai_clients[key] +def _litellm_model(model, backend): + """Normalize to LiteLLM's grammar โ€” same as the chat lane: bare names + are OpenAI-compatible shorthand (wire form ``openai/``), a + ``litellm/`` prefix strips โ€” and fail fast on misconfiguration: + litellm reports a missing key as a retryable 500 and an unknown + provider as a 400, either of which would burn the whole retry loop. + A backend override carries its own credentials, so it skips the key + check; the 401/404 status codes make both errors unrecoverable to + the summary and optimize passes instead of silently absorbed.""" + if not model: + return 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) + if providers and provider not in providers: + 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: + env = litellm.validate_environment(model) + if not env["keys_in_environment"]: + raise litellm.AuthenticationError( + f"missing API key for {model}: " + f"{', '.join(env['missing_keys'])}", + llm_provider=None, model=model) + return model # Misconfiguration: no retry can fix a rejected key or a model that does not @@ -107,33 +108,20 @@ def _is_unrecoverable(exc: Exception) -> bool: 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}] backend = _llm_backend.get() - if use_openai_sdk: - oai_client = _openai_client(backend) + model = _litellm_model(model, backend) + _repair_litellm_types() for i in range(max_retries): try: - if use_openai_sdk: - response = oai_client.chat.completions.create( - model=model, - messages=messages, - ) - else: - import litellm - _repair_litellm_types() - response = litellm.completion( - model=model, - messages=messages, - temperature=0, - drop_params=True, - **(backend or {}), - ) + response = litellm.completion( + model=model, + messages=messages, + drop_params=True, + **(backend or {}), + ) content = response.choices[0].message.content if return_finish_reason: finish_reason = "max_output_reached" if response.choices[0].finish_reason == "length" else "finished" @@ -153,33 +141,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}] backend = _llm_backend.get() - if use_openai_sdk: - oai_client = _openai_client(backend, is_async=True) + model = _litellm_model(model, backend) + _repair_litellm_types() for i in range(max_retries): try: - if use_openai_sdk: - response = await oai_client.chat.completions.create( - model=model, - messages=messages, - ) - else: - import litellm - _repair_litellm_types() - response = await litellm.acompletion( - model=model, - messages=messages, - temperature=0, - drop_params=True, - **(backend or {}), - ) + response = await litellm.acompletion( + model=model, + messages=messages, + drop_params=True, + **(backend or {}), + ) return response.choices[0].message.content except Exception as e: if _is_unrecoverable(e): diff --git a/tests/test_client.py b/tests/test_client.py index 4ea9a12fc..fac192989 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -285,17 +285,31 @@ def test_page_index_flash_rejects_unknown_optimize(): 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_clients", {}) - 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") + + +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_clients", {}) 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) @@ -839,12 +853,11 @@ def test_parse_pages_overlap_counts_union(): # โ”€โ”€ backend: the indexing lane โ”€โ”€ def test_backend_scopes_the_index_lane(tmp_path, monkeypatch): - """index_backend reaches both gateway paths โ€” LiteLLM's call kwargs - and the openai-SDK client โ€” scoped to the operation, with the - endpoint spelling normalized for the openai SDK.""" + """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 - import openai from types import SimpleNamespace from pageindex.local_api import LocalAPI from pageindex.utils import _llm_backend, llm_completion @@ -861,69 +874,16 @@ def test_backend_scopes_the_index_lane(tmp_path, monkeypatch): captured = {} monkeypatch.setattr(litellm, "completion", lambda **kw: (captured.update(kw), reply)[1]) - api._with_backend(lambda: llm_completion("anthropic/claude-x", "p")) - assert captured["api_key"] == "ik" - assert captured["api_base"] == "http://b" - - seen = {} - - class _FakeOpenAI: - def __init__(self, **kw): - seen.update(kw) - self.chat = SimpleNamespace(completions=SimpleNamespace( - create=lambda **_: reply)) - - monkeypatch.setattr(openai, "OpenAI", _FakeOpenAI) - monkeypatch.setattr(pageindex.utils, "_openai_clients", {}) - api._with_backend(lambda: llm_completion("gpt-4o", "p")) - assert seen["api_key"] == "ik" - assert seen["base_url"] == "http://b" - - -def test_openai_client_is_reused_per_backend(monkeypatch): - """One client per distinct backend, built once โ€” the indexing lane calls - the gateway once per node, and rebuilding costs an SSL context each time.""" - import openai - from pageindex.utils import _openai_client - - built = [] - - class _FakeOpenAI: - def __init__(self, **kw): - built.append(kw) - - monkeypatch.setattr(openai, "OpenAI", _FakeOpenAI) - monkeypatch.setattr(openai, "AsyncOpenAI", _FakeOpenAI) - monkeypatch.setattr(pageindex.utils, "_openai_clients", {}) - - one = _openai_client({"api_key": "k", "api_base": "http://b"}) - assert _openai_client({"api_base": "http://b", "api_key": "k"}) is one - assert len(built) == 1 and built[0]["base_url"] == "http://b" - - assert _openai_client({"api_key": "other"}) is not one - assert _openai_client({"api_key": "k", "api_base": "http://b"}, - is_async=True) is not one - assert _openai_client(None) is not one - assert len(built) == 4 - - -def test_index_backend_skips_the_env_precheck(tmp_path, monkeypatch): - """The missing-key pre-check reads environment variables only, so a - backend-supplied key must bypass it instead of being refused.""" - pytest.importorskip("litellm") - import litellm - from pageindex.local_api import LocalAPI - - api = LocalAPI(storage_path=str(tmp_path / "s"), model="m", - summary_model="s", retrieve_model="r", - index_backend={"api_key": "ik"}) monkeypatch.setattr(litellm, "validate_environment", lambda *a, **k: pytest.fail("env pre-check ran")) - monkeypatch.setattr(pageindex.flash, "page_index_flash", - lambda *a, **k: (_ for _ in ()).throw( - RuntimeError("reached-flash"))) - with pytest.raises(RuntimeError, match="reached-flash"): - api._index_flash("f.pdf", ["text"]) + 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_backend_args_are_local_only(): From a8f6ffdf061e09ceabf72949135201972272f424 Mon Sep 17 00:00:00 2001 From: Ray Date: Mon, 17 Aug 2026 23:39:01 +0800 Subject: [PATCH 103/137] =?UTF-8?q?fix:=20four=20review=20findings=20?= =?UTF-8?q?=E2=80=94=20empty=20doc=5Fid=20refused,=20backend=20spelling,?= =?UTF-8?q?=20chat=20unwrap,=20MCP=20pagination?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit doc_id=[] (an empty selection โ€” the natural `[d.id for d in selected]` idiom) built an agent scoped to nothing: frozenset() allowlist, no targeting block, no error โ€” browse returned an empty library and the agent confidently reported the documents don't exist. Cloud already refused it in _local_doc_scope. Now every entry family refuses with that same message: doc_targeting_block (chat surfaces, agent_instructions), _require_local_scope (the three tool builders), and _local_doc_scope with its existing raise hoisted above the local branch (config helpers). 9f67fdd's property โ€” [] must never wash into full-library access โ€” is preserved and upgraded from silent-nothing to loud. chat_backend spoke three vocabularies: chat_completions() lifted api_base to base_url, while responses() and messages() splatted the dict into SDK constructors that reject it. _sdk_backend now does the lift at both construction sites, so the index_backend docs' spelling works on every door. chat() unwrapped the cloud envelope with choices[0] bare โ€” a filtered or malformed reply leaked IndexError/KeyError, the only raw builtin on the cloud lane. The unwrap now wraps shape failures in PageIndexAPIError carrying the envelope's head. McpBridge.list_tools() was the file's one unbounded loop: termination solely on a falsy nextCursor, one 240s-read-timeout POST per pass. A server echoing its cursor now terminates via the no-progress guard; a cycling one hits the 50-page cap and errors instead of hanging agent_tools() forever. --- pageindex/agent_tools.py | 14 ++++++++++++-- pageindex/client.py | 17 ++++++++++++----- pageindex/local_chat.py | 11 +++++++++-- pageindex/mcp_bridge.py | 11 ++++++++--- tests/test_agent_tools.py | 19 +++++++++++++++++++ tests/test_client.py | 12 ++++++++++++ tests/test_local_chat.py | 31 ++++++++++++++++++++----------- 7 files changed, 92 insertions(+), 23 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index e00eb3ca9..dca6e8401 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -1411,7 +1411,13 @@ def _build_cloud_agent_tools(client, include_management: bool) -> list[Callable[ 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.""" + 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 " @@ -1585,7 +1591,11 @@ def doc_targeting_block(client, doc_id, scoped: bool = False) -> Optional[str]: return None doc_ids = [doc_id] if isinstance(doc_id, str) else list(doc_id) if not doc_ids: - return None + # 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 = [client.get_document(one_id) for one_id in doc_ids] listing = _all_documents(client) documents = ([{**detail, "id": one_id} diff --git a/pageindex/client.py b/pageindex/client.py index bf3b09a24..374f6ded8 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -432,7 +432,14 @@ def chat( if stream: return cast(Iterator[str], result) envelope = cast(dict[str, Any], result) - return envelope["choices"][0]["message"]["content"] or "" + try: + return envelope["choices"][0]["message"]["content"] or "" + except (KeyError, IndexError, TypeError) as exc: + # A cloud reply without answer choices (filtered / malformed) + # must surface as the SDK's error, not a bare builtin. + raise PageIndexAPIError( + "The chat response carries no answer: " + f"{str(envelope)[:200]}") from exc def chat_completions( self, @@ -851,14 +858,14 @@ def _local_doc_scope(self, doc_id): """doc_id for the tool layer: passed through locally (structural allowlist), dropped on cloud where scoping is server-side and the config helpers keep prompt-level targeting.""" - if not getattr(self, "api_key", None): - return doc_id if doc_id is not None and not doc_id: - # Cloud has no tool-layer allowlist to make an empty scope mean - # "nothing"; dropping it would silently mean "everything". + # 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( diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 5b9bfb406..f963c9b6a 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -212,6 +212,13 @@ def _require_openai_agents(method: str) -> None: ) 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. @@ -240,7 +247,7 @@ def _openai_model(protocol: str, model_name: str, backend=None): import openai model_name = model_name.removeprefix("openai/") try: - sdk_client = openai.AsyncOpenAI(**(backend or {})) + 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 @@ -793,7 +800,7 @@ def _anthropic_client(backend=None): """The backend client โ€” the seam tests replace with a fake transport.""" import anthropic try: - return anthropic.Anthropic(**(backend or {})) + return anthropic.Anthropic(**_sdk_backend(backend)) except TypeError as exc: raise PageIndexAPIError( f"The Anthropic backend is not configured: {exc}") from exc diff --git a/pageindex/mcp_bridge.py b/pageindex/mcp_bridge.py index 7d8d153c2..f428be2a7 100644 --- a/pageindex/mcp_bridge.py +++ b/pageindex/mcp_bridge.py @@ -179,13 +179,18 @@ def instructions(self) -> Optional[str]: def list_tools(self) -> list[dict]: tools: list[dict] = [] cursor: Optional[str] = None - while True: + # 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 []) - cursor = result.get("nextCursor") - if not cursor: + 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 diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index f92a477ee..b117cb8ba 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -1232,6 +1232,25 @@ def list_tools(self): cloud.agent_tools() +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 diff --git a/tests/test_client.py b/tests/test_client.py index fac192989..0971202f4 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -891,3 +891,15 @@ def test_backend_args_are_local_only(): 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") diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index d34ff9a47..309f16d11 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -1070,18 +1070,18 @@ def tool_outputs(items): @needs_agents -def test_empty_doc_id_is_an_empty_allowlist(client, store_path, fake_model): - """doc_id=[] scopes the agent to nothing; `or None` used to wash it - into unscoped full-library access.""" +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") - fake = fake_model([ - [_call_item("browse_documents", {})], - [_msg_item("done")], - ]) - client.chat_completions("q", doc_id=[]) - outputs = [item["output"] for item in fake.inputs[1] - if item.get("type") == "function_call_output"] - assert json.loads(outputs[-1])["documents"] == [] + 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 @@ -1638,6 +1638,12 @@ def test_backend_connection_reaches_each_engine(monkeypatch): 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" def test_merged_backend_precedence(): @@ -1655,6 +1661,9 @@ def test_messages_backend_merges_and_reaches_the_client(client, fake_anthropic, "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")]) From 959be4a3f51437cf8f855f5f0a5a0eb335924a92 Mon Sep 17 00:00:00 2001 From: Ray Date: Tue, 18 Aug 2026 00:00:41 +0800 Subject: [PATCH 104/137] fix: seven review findings across the bridge, envelope, stream, store and knobs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _cloud_bridge cached its McpBridge with a snapshotted Authorization header while CloudAPI re-reads client.api_key per request โ€” after a key rotation the same client's REST half worked and its MCP half 401ed forever, with no reset handle. The cache entry now carries its (BASE_URL, api_key) identity and a mismatch rebuilds the bridge. The responses() envelope reported tool_choice:"auto" / parallel_tool_calls:True as literals, but ModelSettings sets neither, so the request sends neither โ€” the envelope stated parameters that never hit the wire, and the fixture pinned the fabricated value against the backend's real one. Both capture points (the transport status recorder and the terminal stream event) now record the backend's echo, the envelope reports it, and the literals remain only as fallbacks for echo-less fakes. responses() refused litellm/gpt-4o: the provider-prefix guard ran on the unstripped name, then stripped the prefix inside the error โ€” calling a bare name provider-prefixed and recommending exactly what the user passed. The litellm/ prefix now strips before the guard, matching the grammar everywhere else. responses(stream=True) dropped every lifecycle event including response.created, so the stream opened mid-response with nothing to key state off. One opening now passes through (N per-turn openings collapse to one, not zero), carrying the same pre-generated id the terminal envelope reports. Document-name uniquing was advisory: computed before indexing, discarded, recomputed minutes later, with no store-level exclusion โ€” concurrent submitters of one filename both stored it, permanently shadowing the older doc_id. The check-then-write now runs under a store-level fcntl lock (absent on Windows, where the best-effort behavior stays). Cloud get_tree fallback pulled full document text on a 30s budget only for _format_structure to strip it โ€” include_text=False. retrieve_model regained its setter (0.2.9 allowed assignment), aliasing chat_model like the getter. The stream-abandonment test replaced its process-global threading.active_count() invariant โ€” flaky against litellm's background threads โ€” with tracking the pump thread itself. --- pageindex/agent_tools.py | 19 ++++++--- pageindex/client.py | 16 +++++--- pageindex/local_api.py | 30 ++++++++------ pageindex/local_chat.py | 41 ++++++++++++++++--- pageindex/local_store.py | 19 +++++++++ tests/test_agent_tools.py | 25 ++++++++++++ tests/test_client.py | 36 +++++++++++++++++ tests/test_local_chat.py | 84 ++++++++++++++++++++++++++++++++++----- 8 files changed, 231 insertions(+), 39 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index dca6e8401..93e2272c8 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -863,7 +863,9 @@ def _get_document_structure(client, doc_name: str, raw_tree = getattr(getattr(client, "_api", None), "raw_tree", None) tree = raw_tree(entry["id"]) if raw_tree is not None else None if tree is None: - tree = client.get_tree(entry["id"], node_summary=True).get("result") + # _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}", @@ -1374,14 +1376,19 @@ def _cloud_bridge(client): MCP session. Weak-keyed off the instance so clients stay picklable; the lock closes the check-then-set race under concurrent first calls.""" with _BRIDGES_LOCK: - bridge = _BRIDGES.get(client) - if bridge is None: + # Keyed by the connection identity too: CloudAPI re-reads + # client.api_key on every REST call, so a rotated key (or moved + # BASE_URL) must rebuild the bridge instead of serving the stale + # session's snapshot. + auth = (client.BASE_URL, client.api_key) + bridge, seen = _BRIDGES.get(client) or (None, None) + if bridge is None or seen != auth: from .mcp_bridge import McpBridge bridge = McpBridge( - f"{client.BASE_URL}/mcp", - {"Authorization": f"Bearer {client.api_key}"}, + f"{auth[0]}/mcp", + {"Authorization": f"Bearer {auth[1]}"}, ) - _BRIDGES[client] = bridge + _BRIDGES[client] = (bridge, auth) return bridge diff --git a/pageindex/client.py b/pageindex/client.py index 374f6ded8..7eca7e97a 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -175,6 +175,11 @@ 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 ---------- def submit_document( @@ -605,11 +610,12 @@ def responses( 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, sequence numbers are reassigned monotonically, - and ``output_index`` is re-based onto the single logical - ``output``. The single final event is the terminal - ``response.*`` for the run's status; its ``response`` - carries the tool outputs in ``items``. + 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 diff --git a/pageindex/local_api.py b/pageindex/local_api.py index 180b13516..8049b6da9 100644 --- a/pageindex/local_api.py +++ b/pageindex/local_api.py @@ -128,22 +128,26 @@ def submit_document( raise PageIndexAPIError(f"Failed to submit document: {e}") from e doc_id = "pi-" + uuid.uuid4().hex - 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, - } 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) + # Uniquing must observe concurrent submitters' saves, so the + # check-then-write runs under the store lock โ€” the early pre-check + # above is advisory fail-fast 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: diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index f963c9b6a..8d870c6c8 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -232,12 +232,13 @@ def _openai_model(protocol: str, model_name: str, backend=None): responses protocol: the Responses API is OpenAI-SDK native โ€” LiteLLM's completion surface speaks the chat.completions format, so provider-prefixed models are refused instead of silently downgrading; - bare and ``openai/`` names drive the OpenAI SDK.""" + a ``litellm/`` prefix strips first (it is routing grammar, not a + provider), then bare and ``openai/`` names drive the OpenAI SDK.""" 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 " - f"'{model_name.removeprefix('litellm/')}': provider-prefixed " + 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 " @@ -444,6 +445,13 @@ async def recording_create(*args, **kwargs): 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 โ€” the envelope + # must report these, not assumed values (the request sends neither). + 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 @@ -668,9 +676,11 @@ def run_responses(client, input, model: Optional[str] = None, from agents import Runner from agents.exceptions import AgentsException, MaxTurnsExceeded + response_id = f"resp_{uuid.uuid4().hex}" + def envelope(transcript: list, raw_responses) -> dict: return { - "id": f"resp_{uuid.uuid4().hex}", + "id": response_id, "object": "response", "created_at": int(time.time()), "model": _reported_model(model_name), @@ -685,8 +695,11 @@ def envelope(transcript: list, raw_responses) -> dict: "parameters": tool.params_json_schema, "strict": getattr(tool, "strict_json_schema", True)} for tool in agent.tools], - "tool_choice": "auto", - "parallel_tool_calls": True, + # The backend's own echo when captured (transport wrapper / + # terminal stream event); the request sends neither param, so + # without an echo the OpenAI server defaults apply. + "tool_choice": recorded.get("tool_choice", "auto"), + "parallel_tool_calls": recorded.get("parallel_tool_calls", True), "temperature": temperature, "top_p": top_p, "reasoning": reasoning, @@ -728,11 +741,23 @@ async def agen(): # 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, not zero: + # consumers key state off response.created, so + # the logical stream must open with it, carrying + # the same id the terminal event will report. + opened = True + (data.get("response") or {})["id"] = response_id + sequence += 1 + data["sequence_number"] = sequence + yield data + continue if data["type"] in ("response.completed", "response.incomplete", "response.failed"): @@ -742,6 +767,10 @@ async def agen(): 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): 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/tests/test_agent_tools.py b/tests/test_agent_tools.py index b117cb8ba..fc9515e56 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -2032,6 +2032,31 @@ def instructions(self): 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.""" diff --git a/tests/test_client.py b/tests/test_client.py index 0971202f4..7ec259ad3 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -903,3 +903,39 @@ def test_chat_wraps_answerless_cloud_reply(monkeypatch): 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"} diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index 309f16d11..94b798533 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -124,6 +124,15 @@ async def stream_response(self, system_instructions, input, 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")) @@ -580,6 +589,28 @@ def test_responses_stream_passthrough(client, store_path, fake_model): .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 terminal["response"]["parallel_tool_calls"] is False # echo + + @needs_agents def test_responses_envelope_validates_as_official_response(client, store_path, fake_model): @@ -874,6 +905,9 @@ def test_responses_envelope_fields_and_cache_group(client, store_path, "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 in this fixture, so these are the + # documented fallbacks (the OpenAI server defaults), not assertions + # about the request. assert result["parallel_tool_calls"] is True assert result["tool_choice"] == "auto" @@ -1108,6 +1142,11 @@ def test_openai_model_resolves_provider_prefixes(): 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 @@ -1224,6 +1263,25 @@ async def create(*args, **kwargs): 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): @@ -1410,14 +1468,24 @@ async def drive(): @needs_agents def test_stream_abandonment_cancels_pending_turn(client, store_path, - fake_model): + 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.""" + 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 - import time as time_mod seed_doc(store_path, "pi-a", "report.pdf") - baseline = threading.active_count() + 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")], @@ -1427,11 +1495,9 @@ def test_stream_abandonment_cancels_pending_turn(client, store_path, stream=True, stream_metadata=True) next(stream) # the opening role chunk stream.close() - deadline = time_mod.monotonic() + 3.0 - while (threading.active_count() > baseline - and time_mod.monotonic() < deadline): - time_mod.sleep(0.05) - assert threading.active_count() <= baseline + assert 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 From 4cdc90ed79c4b2120687cb06a2696879efdeafb7 Mon Sep 17 00:00:00 2001 From: Ray Date: Tue, 18 Aug 2026 00:14:59 +0800 Subject: [PATCH 105/137] =?UTF-8?q?chore:=20trim=20rationale=20comments=20?= =?UTF-8?q?to=20the=20essentials=20=E2=80=94=20the=20why=20lives=20in=20th?= =?UTF-8?q?e=20commit=20messages?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pageindex/agent_tools.py | 5 +---- pageindex/client.py | 2 -- pageindex/local_api.py | 5 ++--- pageindex/local_chat.py | 16 ++++++---------- pageindex/utils.py | 12 ++++-------- tests/test_local_chat.py | 4 +--- 6 files changed, 14 insertions(+), 30 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index 93e2272c8..94d997894 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -1376,10 +1376,7 @@ def _cloud_bridge(client): MCP session. Weak-keyed off the instance so clients stay picklable; the lock closes the check-then-set race under concurrent first calls.""" with _BRIDGES_LOCK: - # Keyed by the connection identity too: CloudAPI re-reads - # client.api_key on every REST call, so a rotated key (or moved - # BASE_URL) must rebuild the bridge instead of serving the stale - # session's snapshot. + # A rotated api_key or moved BASE_URL rebuilds the bridge. auth = (client.BASE_URL, client.api_key) bridge, seen = _BRIDGES.get(client) or (None, None) if bridge is None or seen != auth: diff --git a/pageindex/client.py b/pageindex/client.py index 7eca7e97a..7db9cf695 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -440,8 +440,6 @@ def chat( try: return envelope["choices"][0]["message"]["content"] or "" except (KeyError, IndexError, TypeError) as exc: - # A cloud reply without answer choices (filtered / malformed) - # must surface as the SDK's error, not a bare builtin. raise PageIndexAPIError( "The chat response carries no answer: " f"{str(envelope)[:200]}") from exc diff --git a/pageindex/local_api.py b/pageindex/local_api.py index 8049b6da9..88bae0007 100644 --- a/pageindex/local_api.py +++ b/pageindex/local_api.py @@ -131,9 +131,8 @@ def submit_document( pages = [{"page_index": i + 1, "markdown": text} for i, text in enumerate(page_texts)] from .utils import remove_fields - # Uniquing must observe concurrent submitters' saves, so the - # check-then-write runs under the store lock โ€” the early pre-check - # above is advisory fail-fast only. + # Check-then-write under the store lock; the early pre-check above + # is advisory only. with self._store.lock(): meta = { "id": doc_id, diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 8d870c6c8..9578cba00 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -445,8 +445,7 @@ async def recording_create(*args, **kwargs): 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 โ€” the envelope - # must report these, not assumed values (the request sends neither). + # 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: @@ -695,9 +694,7 @@ def envelope(transcript: list, raw_responses) -> dict: "parameters": tool.params_json_schema, "strict": getattr(tool, "strict_json_schema", True)} for tool in agent.tools], - # The backend's own echo when captured (transport wrapper / - # terminal stream event); the request sends neither param, so - # without an echo the OpenAI server defaults apply. + # 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, @@ -748,12 +745,11 @@ async def agen(): 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, not zero: - # consumers key state off response.created, so - # the logical stream must open with it, carrying - # the same id the terminal event will report. + # N per-turn openings collapse to one, carrying + # the id the terminal event will report. opened = True - (data.get("response") or {})["id"] = response_id + if data.get("response"): + data["response"]["id"] = response_id sequence += 1 data["sequence_number"] = sequence yield data diff --git a/pageindex/utils.py b/pageindex/utils.py index ee5f549e8..1b4efbd44 100644 --- a/pageindex/utils.py +++ b/pageindex/utils.py @@ -63,14 +63,10 @@ def _strip_prefix(s, prefix): def _litellm_model(model, backend): - """Normalize to LiteLLM's grammar โ€” same as the chat lane: bare names - are OpenAI-compatible shorthand (wire form ``openai/``), a - ``litellm/`` prefix strips โ€” and fail fast on misconfiguration: - litellm reports a missing key as a retryable 500 and an unknown - provider as a 400, either of which would burn the whole retry loop. - A backend override carries its own credentials, so it skips the key - check; the 401/404 status codes make both errors unrecoverable to - the summary and optimize passes instead of silently absorbed.""" + """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 model = _strip_prefix(model, "litellm/") diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index 94b798533..232a3f2b5 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -905,9 +905,7 @@ def test_responses_envelope_fields_and_cache_group(client, store_path, "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 in this fixture, so these are the - # documented fallbacks (the OpenAI server defaults), not assertions - # about the request. + # No transport echo attached here, so these are the fallbacks. assert result["parallel_tool_calls"] is True assert result["tool_choice"] == "auto" From d429fe414657134ea26c804efe94ec4bc33e903d Mon Sep 17 00:00:00 2001 From: Ray Date: Tue, 18 Aug 2026 00:36:27 +0800 Subject: [PATCH 106/137] chore: deprecation notice when the CHATGPT_API_KEY alias fires MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit FutureWarning, not DeprecationWarning โ€” the audience is app users who set an env var, and Python hides DeprecationWarning from them by default. The alias itself stays: it is documented first-commit README behavior, and env aliasing for third-party consumers (litellm, the openai SDK read the variable themselves) has no cleaner implementation than the one-time import shim. --- pageindex/utils.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pageindex/utils.py b/pageindex/utils.py index 1b4efbd44..fd3a0abba 100644 --- a/pageindex/utils.py +++ b/pageindex/utils.py @@ -47,6 +47,9 @@ def _repair_litellm_types() -> None: # 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): From 77205580473a58bfa2ec264583e1a78fdde5690c Mon Sep 17 00:00:00 2001 From: Ray Date: Tue, 18 Aug 2026 00:43:12 +0800 Subject: [PATCH 107/137] fix: connection reuse, one preload thread, runner max_turns gate, description-edit pin MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit McpBridge gets a requests.Session โ€” agent tool calls come in bursts, and the bare requests.post paid a fresh TCP+TLS handshake on every one. The litellm preload thread starts once per process instead of once per client: the import it warms is process-global, and per-request clients were spawning a throwaway thread each. anthropic_runner_config was the one max_turns door without _validate_max_turns โ€” a bad value sailed into the runner while every sibling door refused it at the edge. The local get_page_content description drops the image-tool sentence from the contract text by exact .replace โ€” a wording change upstream would silently ship the sentence back. A test now pins that the edit actually removed something. The repr()'d schema defaults finding resolves as not-a-bug: the only repr site renders Python literals into an exec'd def, where that is the required dialect; frameworks read real default objects off the signature, and docstrings carry no defaults. --- pageindex/client.py | 26 ++++++++++++++++++++------ pageindex/mcp_bridge.py | 5 +++-- tests/test_agent_tools.py | 28 ++++++++++++++++++++++------ tests/test_local_chat.py | 4 ++++ 4 files changed, 49 insertions(+), 14 deletions(-) diff --git a/pageindex/client.py b/pageindex/client.py index 7db9cf695..1d40823c3 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -10,11 +10,24 @@ from .errors import PageIndexAPIError +_litellm_preload_started = False + + def _preload_litellm() -> None: - try: - import litellm # noqa: F401 - except Exception: - pass + """Start litellm's multi-second import in the background, once per + process โ€” a per-client thread would churn under per-request clients.""" + 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]: @@ -168,7 +181,7 @@ def __init__( ) # LiteLLM's multi-second import would otherwise land on the # first chat call; failures resurface there with real context. - threading.Thread(target=_preload_litellm, daemon=True).start() + _preload_litellm() @property def retrieve_model(self): @@ -1007,7 +1020,8 @@ def anthropic_runner_config( max_turns: Agent-loop bound; default 10. """ from .agent_tools import build_agent_instructions - from .local_chat import _default_max_tokens + 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, diff --git a/pageindex/mcp_bridge.py b/pageindex/mcp_bridge.py index f428be2a7..203226ced 100644 --- a/pageindex/mcp_bridge.py +++ b/pageindex/mcp_bridge.py @@ -44,6 +44,7 @@ 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 @@ -65,8 +66,8 @@ def _post(self, payload: dict, session_id: Optional[str] = None, if protocol_version: headers["MCP-Protocol-Version"] = protocol_version try: - return requests.post(self._url, json=payload, headers=headers, - timeout=_TIMEOUT) + 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}" diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index fc9515e56..db04fb91d 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -1232,6 +1232,16 @@ def list_tools(self): 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.""" @@ -1308,7 +1318,8 @@ def fake_post(url, json=None, headers=None, timeout=None): # Replace the module's own `requests` binding โ€” patching the shared # requests module would leak the fake process-wide. monkeypatch.setattr(mcp_bridge, "requests", types.SimpleNamespace( - post=fake_post, RequestException=requests_mod.RequestException)) + Session=lambda: types.SimpleNamespace(post=fake_post), + RequestException=requests_mod.RequestException)) bridge = McpBridge("https://api.pageindex.ai/mcp", {"Authorization": "Bearer k"}) @@ -1371,7 +1382,8 @@ def fake_post(url, json=None, headers=None, timeout=None): return _Resp(400, text="unknown tool") monkeypatch.setattr(mcp_bridge, "requests", types.SimpleNamespace( - post=fake_post, RequestException=requests_mod.RequestException)) + 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"): @@ -1432,7 +1444,8 @@ def fake_post(url, json=None, headers=None, timeout=None): return resp monkeypatch.setattr(mcp_bridge, "requests", types.SimpleNamespace( - post=fake_post, RequestException=requests_mod.RequestException)) + Session=lambda: types.SimpleNamespace(post=fake_post), + RequestException=requests_mod.RequestException)) bridge = McpBridge("https://api.pageindex.ai/mcp", {"Authorization": "Bearer k"}) @@ -1627,7 +1640,8 @@ def fake_post(url, json=None, headers=None, timeout=None): "content": [{"type": "text", "text": '{"error": "denied"}'}]}}) monkeypatch.setattr(mcp_bridge, "requests", types.SimpleNamespace( - post=fake_post, RequestException=requests_mod.RequestException)) + 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) @@ -1664,7 +1678,8 @@ def fake_post(url, json=None, headers=None, timeout=None): "text": "old"}]}}) monkeypatch.setattr(mcp_bridge, "requests", types.SimpleNamespace( - post=fake_post, RequestException=requests_mod.RequestException)) + 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", {}) @@ -1688,7 +1703,8 @@ def dead_post(*args, **kwargs): raise requests_mod.ConnectionError("dns down") monkeypatch.setattr(mcp_bridge, "requests", types.SimpleNamespace( - post=dead_post, RequestException=requests_mod.RequestException)) + 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() diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index 232a3f2b5..edde90b58 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -855,6 +855,10 @@ def test_max_turns_rejects_non_positive(client, store_path, fake_model): 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): From d0007c2b054743c1077d68b14c4de197bba27261 Mon Sep 17 00:00:00 2001 From: Ray Date: Tue, 18 Aug 2026 00:49:30 +0800 Subject: [PATCH 108/137] refactor: the image-sentence edit anchors on the tool name, not the wording A contract rewording that keeps naming get_document_image() now strips cleanly instead of requiring a needle update; the refresh-pin test stays as the backstop for sentence splits or a renamed tool. --- pageindex/agent_tools.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index 94d997894..437622129 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -1231,9 +1231,11 @@ def _tool_docstring(description: str, properties: dict[str, Any]) -> str: 'Folder browsing and semantic ranking (sort="relevance") are not ' "supported in local mode yet โ€” they work on PageIndex cloud." ), - "get_page_content": TOOL_CONTRACT["get_page_content"]["description"] - .replace(" Embedded image paths in the response feed into " - "`get_document_image()`.", ""), + # 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] = { From 535d75bdf5701406b6647aae88c0b8b144874aeb Mon Sep 17 00:00:00 2001 From: Ray Date: Tue, 18 Aug 2026 01:07:24 +0800 Subject: [PATCH 109/137] =?UTF-8?q?refactor:=20five=20cleanup=20findings?= =?UTF-8?q?=20=E2=80=94=20one=20synthesizer,=20one=20parser,=20one=20loop?= =?UTF-8?q?=20bridge?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit agent_tools() local mode carried five hand-written wrappers whose signatures had already drifted from the contract (limit: int vs the schema's number). _make_bridge_function generalizes to _make_tool_function(name, description, schema, invoke), and build_agent_tools now renders _tool_specs for both modes โ€” the same pruned local schemas and descriptions the other doors already serve, so the local surface stays the subset by construction. build_claude_mcp drops its parallel hand-rolled iteration for the same _tool_specs. The page-spec parser existed twice โ€” client._parse_pages (raising) and agent_tools._parse_page_spec (envelopes) โ€” and disagreed on page 0. One core (_expand_pages, _PageSpecError with a rendering code) now backs both skins; the client surface gains the tool layer's page >= 1 rule, so get_ocr-by-pages rejects 0 instead of passing it downstream. The run-a-blocking-call-off-the-loop bridge existed twice (local_api._run_indexer, local_chat._run_sync); both now ride utils.run_off_loop. test_contract_matches_snapshot renamed to what it can honestly claim: both copies live in this repo, so it cannot detect server drift โ€” it makes TOOL_CONTRACT edits deliberate two-file changes. --- pageindex/agent_tools.py | 196 +++++++++------------ pageindex/client.py | 23 +-- pageindex/integrations/claude_agent_sdk.py | 16 +- pageindex/local_api.py | 16 +- pageindex/local_chat.py | 13 +- pageindex/utils.py | 12 ++ tests/test_agent_tools.py | 32 ++-- tests/test_client.py | 3 + 8 files changed, 138 insertions(+), 173 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index 437622129..655224ad5 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -33,7 +33,6 @@ STRUCTURE_FIRST_PAGE_THRESHOLD = 20 _CHAR_BUDGET = int(TOOL_RESPONSE_CHAR_LIMIT * 0.95) -_PAGES_SPEC_RE = re.compile(r"^(\d+(-\d+)?)(,\s*\d+(-\d+)?)*$") _MAX_REQUESTED_PAGES = 10_000 _SIMILAR_NAMES_LIMIT = 3 _TOOL_WAIT_TIMEOUT = 180.0 # "up to 3 minutes", per the wait_for_completion schema @@ -503,69 +502,102 @@ def _folder_unsupported(param: str) -> tuple[dict, bool]: # โ”€โ”€ page spec handling โ”€โ”€ -def _parse_page_spec( - pages: str, doc_name: str, -) -> "tuple[Optional[list[int]], Optional[_ToolResult]]": - """Expand '1-3,7' into a sorted, deduplicated page list, or an error.""" - invalid = _failure( - "Invalid page specification format", - {"doc_name": doc_name}, - { - "summary": "Failed to parse the pages parameter", - "options": [ - 'Use valid formats: "5", "3,7,10", "5-10", or "1-3,7,9-12"', - "Ensure page numbers are positive integers", - ], - }, - "INVALID_INPUT", - ) - if not isinstance(pages, str) or not _PAGES_SPEC_RE.match(pages.strip()): - return None, invalid - too_many = _failure( - f"Too many pages requested (over {_MAX_REQUESTED_PAGES})", - {"doc_name": doc_name}, - { - "summary": "The page specification spans too many pages", - "options": [ - "Request a narrower page range", - "The response holds only a few pages per call - page through with several smaller requests", - ], - }, - "INVALID_INPUT", - ) +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() - if "-" in part: - start, end = (int(x) for x in part.split("-", 1)) - if start > end: - return None, invalid - else: - start = end = int(part) + 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: - return None, too_many + raise _PageSpecError("too_many", too_many) expanded.update(range(start, end + 1)) if len(expanded) > _MAX_REQUESTED_PAGES: - return None, too_many - if any(page < 1 for page in expanded): + 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 numbers. Page numbers must be positive integers", + "Invalid page specification format", {"doc_name": doc_name}, { - "summary": "Invalid page numbers provided", + "summary": "Failed to parse the pages parameter", "options": [ - "Page numbers must be positive integers (>= 1)", - "Check the page specification format", + 'Use valid formats: "5", "3,7,10", "5-10", or "1-3,7,9-12"', + "Ensure page numbers are positive integers", ], }, "INVALID_INPUT", ) - return sorted(expanded), None def _format_page_spec(pages: list[int]) -> str: @@ -1265,9 +1297,6 @@ def _local_schema(name: str) -> dict[str, Any]: return schema -def _docstring(name: str) -> str: - return _tool_docstring(_local_description(name), - _local_schema(name)["properties"]) _SCHEMA_TYPE_MAP = {"string": str, "integer": int, "number": float, @@ -1325,16 +1354,16 @@ def _invoke(arguments: dict[str, Any]) -> tuple[str, bool]: return _invoke -def _make_bridge_function(bridge, meta: dict) -> Callable[..., str]: - """One plain function for a cloud tool: real signature and docstring from - the server's schema, invocation proxied over MCP, errors contained.""" +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.""" import keyword - name = str(meta.get("name") or "") - schema = meta.get("inputSchema") or {} properties: dict[str, Any] = schema.get("properties") or {} required = set(schema.get("required") or []) - _invoke = _bridge_invoker(bridge, name) + _invoke = invoke params_usable = all(param.isidentifier() and not keyword.iskeyword(param) and param != "_invoke" @@ -1365,7 +1394,7 @@ def proxy(**kwargs: Any) -> str: annotations["return"] = str proxy.__annotations__ = annotations proxy.__name__ = proxy.__qualname__ = name or "tool" - proxy.__doc__ = _tool_docstring(meta.get("description") or "", properties) + proxy.__doc__ = _tool_docstring(description or "", properties) return proxy @@ -1407,14 +1436,6 @@ def _read_only_tools(tools_meta: list[dict]) -> list[dict]: return filtered -def _build_cloud_agent_tools(client, include_management: bool) -> list[Callable[..., str]]: - bridge = _cloud_bridge(client) - tools_meta = bridge.list_tools() - if not include_management: - tools_meta = _read_only_tools(tools_meta) - return [_make_bridge_function(bridge, meta) for meta in tools_meta] - - def _require_local_scope(client, doc_ids) -> None: """The allowlist is enforced in-process; cloud lookups run server-side, so accepting doc_ids there would be advisory-only โ€” refuse loudly. @@ -1470,52 +1491,9 @@ def build_agent_tools(client, include_management: bool = False) -> list[Callable signature accepts (cloud-only parameters are absent from the local signatures; the call_tool path answers them with the guided envelope). """ - if getattr(client, "api_key", None): - return _build_cloud_agent_tools(client, include_management) - - def browse_documents(offset: int = 0, limit: int = 10) -> str: - return call_tool(client, "browse_documents", { - "offset": offset, "limit": limit, - })[0] - - def get_document(doc_name: str, wait_for_completion: bool = False) -> str: - return call_tool(client, "get_document", { - "doc_name": doc_name, - "wait_for_completion": wait_for_completion, - })[0] - - def get_document_structure(doc_name: str, part: int = 1, - wait_for_completion: bool = False) -> str: - return call_tool(client, "get_document_structure", { - "doc_name": doc_name, "part": part, - "wait_for_completion": wait_for_completion, - })[0] - - def get_page_content(doc_name: str, pages: str, - wait_for_completion: bool = False) -> str: - return call_tool(client, "get_page_content", { - "doc_name": doc_name, "pages": pages, - "wait_for_completion": wait_for_completion, - })[0] - - def remove_document(doc_names: list[str]) -> str: - return call_tool(client, "remove_document", { - "doc_names": doc_names, - })[0] - - functions = { - "browse_documents": browse_documents, - "get_document": get_document, - "get_document_structure": get_document_structure, - "get_page_content": get_page_content, - "remove_document": remove_document, - } - tools = [] - for name in tool_names(include_management): - function = functions[name] - function.__doc__ = _docstring(name) - tools.append(function) - return tools + return [_make_tool_function(name, description, schema, invoke) + for name, description, schema, invoke + in _tool_specs(client, include_management)] # โ”€โ”€ agent instructions โ”€โ”€ diff --git a/pageindex/client.py b/pageindex/client.py index 1d40823c3..5d440cd61 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -31,27 +31,8 @@ def _import() -> None: def _parse_pages(pages: str) -> list[int]: - result: set[int] = set() - too_many = (f"Page specification '{pages}' spans more than " - "10000 pages; request a narrower range") - for part in pages.split(","): - part = part.strip() - if "-" in part: - start, end = (int(x) for x in part.split("-", 1)) - if start > end: - raise ValueError(f"Invalid range '{part}': start must be <= end") - else: - start = end = int(part) - # Bound each part arithmetically before materializing it โ€” a spec - # like "1-999999999" would otherwise expand to a billion integers. - # The cap is on distinct pages, so overlapping parts (a parent - # section plus its children) don't double-count. - if end - start + 1 > 10_000: - raise ValueError(too_many) - result.update(range(start, end + 1)) - if len(result) > 10_000: - raise ValueError(too_many) - return sorted(result) + from .agent_tools import _expand_pages + return _expand_pages(pages) def _agents_sdk_model_name(model: str) -> str: diff --git a/pageindex/integrations/claude_agent_sdk.py b/pageindex/integrations/claude_agent_sdk.py index 8c76cb434..b3e67947d 100644 --- a/pageindex/integrations/claude_agent_sdk.py +++ b/pageindex/integrations/claude_agent_sdk.py @@ -34,14 +34,11 @@ def build_claude_mcp(client, include_management: bool = False, doc_ids=None): "as_claude_mcp in local mode requires the Claude Agent SDK โ€” " "pip install claude-agent-sdk (or pip install 'pageindex[claude]')." ) from exc - from ..agent_tools import (TOOL_CONTRACT, _local_description, - _local_schema, call_tool, tool_names) + from ..agent_tools import TOOL_CONTRACT, _tool_specs - def make_handler(name: str): + def make_handler(invoke): async def handler(arguments: dict[str, Any]) -> dict[str, Any]: - text, is_error = await asyncio.to_thread( - call_tool, client, name, arguments or {}, doc_ids - ) + 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 @@ -59,9 +56,10 @@ def tool_kwargs(name: str) -> dict: return {"annotations": ToolAnnotations(**annotations)} tools = [ - tool(name, _local_description(name), - _local_schema(name), **tool_kwargs(name))(make_handler(name)) - for name in tool_names(include_management) + 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/local_api.py b/pageindex/local_api.py index 88bae0007..83e846a63 100644 --- a/pageindex/local_api.py +++ b/pageindex/local_api.py @@ -1,17 +1,16 @@ """Local implementation of the PageIndex SDK surface.""" from __future__ import annotations -import asyncio import json import logging 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,13 +21,6 @@ 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: @@ -46,7 +38,7 @@ def __init__(self, storage_path: str, model: str, summary_model: str, def _with_backend(self, func, *args): """Scope the indexing lane's connection overrides around one - operation โ€” runs inside whatever thread _run_indexer picked.""" + operation โ€” runs inside whatever thread run_off_loop picked.""" from .utils import _llm_backend token = _llm_backend.set(self._index_backend) try: @@ -114,11 +106,11 @@ def submit_document( try: if mode == "flash": - structure, description = _run_indexer( + structure, description = run_off_loop( self._with_backend, self._index_flash, file_path, page_texts ) else: - structure, description = _run_indexer( + structure, description = run_off_loop( self._with_backend, self._index_standard, file_path, page_texts ) diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 9578cba00..0fe4fe787 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -21,7 +21,6 @@ from __future__ import annotations import asyncio -import concurrent.futures import hashlib import json import os @@ -118,16 +117,8 @@ def _split_chat_messages(messages) -> "tuple[list[str], list[dict]]": def _run_sync(coro): - try: - asyncio.get_running_loop() - except RuntimeError: - has_loop = False - else: - has_loop = True - if not has_loop: - return asyncio.run(coro) - with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: - return pool.submit(asyncio.run, coro).result() + from .utils import run_off_loop + return run_off_loop(asyncio.run, coro) _SENTINEL = object() diff --git a/pageindex/utils.py b/pageindex/utils.py index fd3a0abba..b6377009a 100644 --- a/pageindex/utils.py +++ b/pageindex/utils.py @@ -65,6 +65,18 @@ def _strip_prefix(s, prefix): return s +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 _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 diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index db04fb91d..f260e0dd8 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -69,7 +69,10 @@ def run(client, name, **arguments): # โ”€โ”€ contract parity โ”€โ”€ -def test_contract_matches_snapshot(): +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 @@ -1189,13 +1192,12 @@ def test_cloud_agent_tools_null_description_survives(): """A server may send description: null โ€” .get(key, default) does not apply the default to it, and agent_tools() died with a TypeError while the _tool_specs path handled the same payload fine.""" - from pageindex.agent_tools import _make_bridge_function class _Bridge: def call_tool(self, name, arguments): return json.dumps({"success": True}), False - tool = _make_bridge_function(_Bridge(), { + tool = _synth(_Bridge(), { "name": "search_documents", "description": None, "inputSchema": {"type": "object", @@ -1494,10 +1496,20 @@ def test_mcp_bridge_blob_blocks_become_stubs(): # โ”€โ”€ 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)) + + def test_synth_optional_no_default_param_is_nullable(): """A non-required, no-default schema param must annotate Optional, or strict schemas force the model to always send a value (browse.query).""" - from pageindex.agent_tools import _make_bridge_function, TOOL_CONTRACT + from pageindex.agent_tools import TOOL_CONTRACT from typing import get_args class _Bridge: @@ -1507,7 +1519,7 @@ def call_tool(self, name, args): meta = {"name": "browse_documents", "description": "d", "inputSchema": TOOL_CONTRACT["browse_documents"]["schema"]} - fn = _make_bridge_function(_Bridge(), meta) + fn = _synth(_Bridge(), meta) assert type(None) in get_args(fn.__annotations__["query"]) @@ -1516,7 +1528,6 @@ def test_synth_array_params_keep_their_item_type(): function_tool then emits {"type": "array", "items": {}}, which strict function calling rejects.""" from typing import Optional - from pageindex.agent_tools import _make_bridge_function class _Bridge: def call_tool(self, name, args): @@ -1535,7 +1546,7 @@ def call_tool(self, name, args): }, "required": ["doc_ids", "mixed"], }} - fn = _make_bridge_function(_Bridge(), meta) + 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 โ€” @@ -1544,7 +1555,6 @@ def call_tool(self, name, args): def test_synth_escape_hatches(): - from pageindex.agent_tools import _make_bridge_function calls = [] @@ -1554,7 +1564,7 @@ def call_tool(self, name, args): return "ok", False # Tool named "_invoke" must not recurse into itself. - invoke_named = _make_bridge_function(_Bridge(), { + invoke_named = _synth(_Bridge(), { "name": "_invoke", "description": "d", "inputSchema": {"type": "object", "properties": {"x": {"type": "string"}}, "required": ["x"]}}) @@ -1562,7 +1572,7 @@ def call_tool(self, name, args): assert calls[-1] == ("_invoke", {"x": "v"}) # Param named "dict" must not shadow the builtin. - dict_param = _make_bridge_function(_Bridge(), { + dict_param = _synth(_Bridge(), { "name": "t", "description": "d", "inputSchema": {"type": "object", "properties": {"dict": {"type": "string"}}, "required": ["dict"]}}) @@ -1571,7 +1581,7 @@ def call_tool(self, name, args): # Non-identifier tool name still gets a real signature. import inspect - dashed = _make_bridge_function(_Bridge(), { + dashed = _synth(_Bridge(), { "name": "page-content.v2", "description": "d", "inputSchema": {"type": "object", "properties": {"a": {"type": "string"}}, "required": ["a"]}}) diff --git a/tests/test_client.py b/tests/test_client.py index 7ec259ad3..6e983c477 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -848,6 +848,9 @@ def test_parse_pages_overlap_counts_union(): assert len(pages) == 9000 and pages[0] == 1 and pages[-1] == 9000 with pytest.raises(ValueError, match="spans more than"): _parse_pages("1-10001") + # one parser with the tool layer now: page 0 is rejected, not passed on + with pytest.raises(ValueError, match="positive"): + _parse_pages("0-3") # โ”€โ”€ backend: the indexing lane โ”€โ”€ From 28eaab8acccfae1b54177ace4dd55611f86af846 Mon Sep 17 00:00:00 2001 From: Ray Date: Tue, 18 Aug 2026 01:07:24 +0800 Subject: [PATCH 110/137] chore: dotenv-proof the chat missing-key test; page_num falsy-zero re-checked MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit test_chat_missing_openai_key_fails_loud imports litellm before delenv โ€” its first import may load a .env and hand the deleted key back mid-test. The other delenv site (backend-connection test) is ordering-safe: its chat-agent build imports litellm before the delenv runs. The old page_num falsy-zero note closes as benign: both sites (get_document guidance, page_count) conflate 0 with missing only where both mean "omit the hint", and local submits guarantee pageNum >= 1. --- tests/test_local_chat.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index edde90b58..1cd8cfd86 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -1211,6 +1211,7 @@ def test_envelope_model_strips_openai_routing_prefix(store_path, fake_model): 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"): From 228426d02494c5ace8aa76faf8063c9f479504ca Mon Sep 17 00:00:00 2001 From: Ray Date: Tue, 18 Aug 2026 16:51:05 +0800 Subject: [PATCH 111/137] =?UTF-8?q?fix:=20close=20the=20pypdfium2=20dual-s?= =?UTF-8?q?emantics=20window=20=E2=80=94=20floor=20>=3D5,=20sentinel=20tes?= =?UTF-8?q?t?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 57a77f6 treated FPDFFont_GetBaseFontName as a rename of GetFontName; pdfium actually split that call (GetFamilyName carries the old family semantics), so font_name moved from family to per-face granularity. A 9-PDF / 1574-page A/B against 4.30: zero characters lost, tree skeletons stable, word segmentation ~6:1 net better (1036 spurious math-notation spaces removed vs 178 added, 173 of those in four-lectures.pdf pseudocode โ€” accepted), 124 pages reorder within the page (table linearization, no ground truth either way), one figure axis label recovered. Keeping the per-face semantics. But pyproject's >=4.30.0 floor left both semantics simultaneously installable โ€” clone/CI resolved 5.13 while pip users on satisfied envs kept 4.30: same SDK version, different extracted text, no signal. Floor now >=5; the getattr fallback stays so un-upgraded envs still run until pip next resolves. The sentinel test pins the 5.x extraction behavior (skips below 5.x). --- pyproject.toml | 2 +- tests/test_flash_extraction.py | 23 +++++++++++++++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) create mode 100644 tests/test_flash_extraction.py diff --git a/pyproject.toml b/pyproject.toml index 499291d3d..f5bec8c9a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,7 +32,7 @@ openai = ">=1.70.0" 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" diff --git a/tests/test_flash_extraction.py b/tests/test_flash_extraction.py new file mode 100644 index 000000000..735d31ce4 --- /dev/null +++ b/tests/test_flash_extraction.py @@ -0,0 +1,23 @@ +"""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 From 54341111fc4c22414f71fb70fe89117095cf7ffb Mon Sep 17 00:00:00 2001 From: Ray Date: Tue, 18 Aug 2026 17:03:32 +0800 Subject: [PATCH 112/137] fix: the missing-key pre-check covers only OpenAI-shaped models MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The blanket litellm.validate_environment gate (all three indexing entry points) only sees environment variables, so it hard-failed providers whose credentials resolve at call time โ€” Ollama (litellm's own localhost:11434 default), Bedrock IAM chains on bare EC2, Vertex ADC โ€” before any request was sent, on the default flash path. The chat lane already draws the line at OpenAI-shaped names (bare or openai/) and lets every other provider resolve natively; the indexing lane now uses one shared helper (_openai_missing_keys) with the same rule at _litellm_model, the tree_optimize CLI, and the run_pageindex flash pre-check. Unknown-provider grammar validation is unchanged. --- pageindex/tree_optimize.py | 9 ++++----- pageindex/utils.py | 20 ++++++++++++++++---- run_pageindex.py | 11 +++++------ tests/test_client.py | 23 +++++++++++++++++++++++ 4 files changed, 48 insertions(+), 15 deletions(-) diff --git a/pageindex/tree_optimize.py b/pageindex/tree_optimize.py index 026f69dde..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_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 @@ -873,10 +873,9 @@ async def main(): model = args.model or default_model() if args.expand and not args.plan: - import litellm - env = litellm.validate_environment(model) - if not env["keys_in_environment"]: - sys.exit(f"{', '.join(env['missing_keys'])} is not set " + 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)) diff --git a/pageindex/utils.py b/pageindex/utils.py index b6377009a..75531cc5d 100644 --- a/pageindex/utils.py +++ b/pageindex/utils.py @@ -77,6 +77,19 @@ def run_off_loop(func, *args): 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.""" + import litellm + wire = _strip_prefix(model, "litellm/") + if "/" in wire and not wire.startswith("openai/"): + return [] + env = litellm.validate_environment(wire if "/" in wire else f"openai/{wire}") + return [] if env["keys_in_environment"] else env["missing_keys"] + + 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 @@ -98,11 +111,10 @@ def _litellm_model(model, backend): f"OPENAI_BASE_URL at the server.", llm_provider=None, model=model) if not backend: - env = litellm.validate_environment(model) - if not env["keys_in_environment"]: + missing = _openai_missing_keys(model) + if missing: raise litellm.AuthenticationError( - f"missing API key for {model}: " - f"{', '.join(env['missing_keys'])}", + f"missing API key for {model}: {', '.join(missing)}", llm_provider=None, model=model) return model diff --git a/run_pageindex.py b/run_pageindex.py index f91f2b751..054ebf429 100644 --- a/run_pageindex.py +++ b/run_pageindex.py @@ -3,7 +3,7 @@ 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 if __name__ == "__main__": # Set up argument parser @@ -95,11 +95,10 @@ or ConfigLoader().load().summary_model) will_summarize = args.summary if args.summary is not None else True if will_summarize or args.optimize == 'full': - import litellm - env = litellm.validate_environment(summary_model) - if not env["keys_in_environment"]: + missing = _openai_missing_keys(summary_model) + if missing: raise SystemExit( - f"Missing API key for {summary_model}: {', '.join(env['missing_keys'])}") + f"Missing API key for {summary_model}: {', '.join(missing)}") toc_with_page_number = page_index_flash( args.pdf_path, optimize=args.optimize if args.optimize != 'off' else False, @@ -159,7 +158,7 @@ import asyncio # Use ConfigLoader to get consistent defaults (matching PDF behavior) - from pageindex.utils import ConfigLoader + from pageindex.utils import ConfigLoader, _openai_missing_keys config_loader = ConfigLoader() # Create options dict with user args diff --git a/tests/test_client.py b/tests/test_client.py index 6e983c477..63c16a598 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -889,6 +889,29 @@ def test_backend_scopes_the_index_lane(tmp_path, monkeypatch): 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_backend_args_are_local_only(): with pytest.raises(PageIndexAPIError, match="chat_backend"): PageIndexClient(api_key="pi-k", chat_backend={"api_key": "x"}) From 86d6833c69c7fcae071175db5aa349b2e07e61d0 Mon Sep 17 00:00:00 2001 From: Ray Date: Tue, 18 Aug 2026 17:03:39 +0800 Subject: [PATCH 113/137] docs: page_index_flash's opening claim scoped to what is actually LLM-free Since optimize defaults to "full", the bare call makes LLM calls, so "without an LLM" was false for the default run. The tree extraction itself still uses none; the docstring now says exactly that and names the fully LLM-free form (summary=False, optimize=False), matching the flash README. --- pageindex/flash/api.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pageindex/flash/api.py b/pageindex/flash/api.py index bf62d9657..26018a6c0 100644 --- a/pageindex/flash/api.py +++ b/pageindex/flash/api.py @@ -99,7 +99,7 @@ def page_index_flash(pdf, summary=True, summary_model=None, optimize: str | bool = "full", optimize_expand=None, optimize_model=None, summary_concurrency=None, use_embedded_toc=True) -> dict: - """Build a PageIndex tree structure from a PDF using layout statistics, without an LLM. Args: pdf: path to a PDF file (``str`` or ``pathlib.Path``) or an in-memory binary stream (``io.BytesIO``). summary: if True, generate LLM summaries for each node (requires ``summary_model``). summary_model: the LLM model identifier to use for summary generation. optimize: ``"full"`` for merge + LLM expand, ``"merge"`` for deterministic merge only, ``False`` to disable. ``True`` is accepted as ``"full"`` for backward compatibility. optimize_model: the LLM model for expand (defaults to the summary model). summary_concurrency: maximum simultaneous summary model calls; None uses the library default. use_embedded_toc: if True, consume the PDF's embedded bookmarks when trustworthy: deep bookmarks become the frame and the detected sections they lack are grafted back in after noise filtering, coarse ones become the chapter frame with detected nodes re-hung under them (deeper sparse entries are filled in when the page text confirms them, and garbled extracted titles are repaired from the bookmark strings), garbage ones are ignored; adds a ``toc_source`` key to the result. On by default; pass False for the pure detected structure. Returns: dict with keys ``doc_name``, ``doc_title``, ``structure`` (a list of nested ``{"title", "start_index", "end_index", "nodes"}`` dicts; page indexes are 1-based) and ``has_abstract_or_references_section`` (True when a top-level entry is an abstract or references heading). With ``optimize`` an ``optimize`` key reports merge/expand counts and before/after search-cost metrics. """ + """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, ``"merge"`` for deterministic merge only, ``False`` to disable. ``True`` is accepted as ``"full"`` for backward compatibility. optimize_model: the LLM model for expand (defaults to the summary model). summary_concurrency: maximum simultaneous summary model calls; None uses the library default. use_embedded_toc: if True, consume the PDF's embedded bookmarks when trustworthy: deep bookmarks become the frame and the detected sections they lack are grafted back in after noise filtering, coarse ones become the chapter frame with detected nodes re-hung under them (deeper sparse entries are filled in when the page text confirms them, and garbled extracted titles are repaired from the bookmark strings), garbage ones are ignored; adds a ``toc_source`` key to the result. On by default; pass False for the pure detected structure. Returns: dict with keys ``doc_name``, ``doc_title``, ``structure`` (a list of nested ``{"title", "start_index", "end_index", "nodes"}`` dicts; page indexes are 1-based) and ``has_abstract_or_references_section`` (True when a top-level entry is an abstract or references heading). With ``optimize`` an ``optimize`` key reports merge/expand counts and before/after search-cost metrics. """ if optimize is True: optimize = "full" if not optimize: From 2570cfa55e23b60f427bd3257efce3b17f5d57b6 Mon Sep 17 00:00:00 2001 From: Ray Date: Tue, 18 Aug 2026 17:10:44 +0800 Subject: [PATCH 114/137] chore: __init__.py diff carries only the lines the feature needed MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reverts the ride-along edits: comments explaining pre-existing behavior (the underscore guard, the classic fallthrough), two re-wrapped statements, and two _LAZY entries mapping to the value _LAZY.get() already defaults to โ€” provably inert. What stays is what 0.2.10 required: the litellm model-map setdefault and the new module names in _SUBMODULES. --- pageindex/__init__.py | 15 +++------------ 1 file changed, 3 insertions(+), 12 deletions(-) diff --git a/pageindex/__init__.py b/pageindex/__init__.py index 88ca32ff8..e85278704 100644 --- a/pageindex/__init__.py +++ b/pageindex/__init__.py @@ -23,8 +23,6 @@ ] _LAZY = { - "page_index": ".page_index_classic", - "page_index_main": ".page_index_classic", "page_index_flash": ".flash", "optimize_tree": ".tree_optimize", "md_to_tree": ".page_index_md", @@ -34,25 +32,18 @@ "mcp_bridge", "page_index_classic", "page_index_md", "tree_optimize", "utils"} + def __getattr__(name): if name.startswith("_"): - # Dunder probes (copy, pickle, inspect) are the frequent unknown - # names โ€” they must not trigger the classic import below. raise AttributeError(f"module {__name__!r} has no attribute {name!r}") import importlib if name in _SUBMODULES: return importlib.import_module(f".{name}", __name__) - # Pre-0.2.10 compat: unknown names fall through to the classic module, - # whose public surface (ConfigLoader, count_tokens, ...) resolved as - # package attributes. A non-underscore typo pays one classic import - # before its AttributeError โ€” not worth an allowlist. - module = importlib.import_module(_LAZY.get(name, ".page_index_classic"), - __name__) + module = importlib.import_module(_LAZY.get(name, ".page_index_classic"), __name__) try: value = getattr(module, name) except AttributeError: - raise AttributeError( - f"module {__name__!r} has no attribute {name!r}") from None + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None globals()[name] = value return value From 03ffab38886ebb86ad7e7075258f3feb7e740172 Mon Sep 17 00:00:00 2001 From: Ray Date: Tue, 18 Aug 2026 19:40:06 +0800 Subject: [PATCH 115/137] =?UTF-8?q?fix:=20ten=20review=20findings=20?= =?UTF-8?q?=E2=80=94=20swallowed=20401s,=20key=20gate,=20tool=20envelopes,?= =?UTF-8?q?=20model=20grammar,=20prompt=20caching?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit summarize_tree's fail-fast only saw root-node failures: the recursive child gather ran with return_exceptions=True and discarded the results, so a per-leaf 401 stored a blank subtree as "completed". The child results are now checked like the root's, and the flat generate_summaries_for_structure sibling re-raises unrecoverable failures the same way instead of blanking the node. The missing-key pre-check trusted litellm.validate_environment, which reports a blank or whitespace OPENAI_API_KEY as present โ€” the fail-fast never fired and the failure resurfaced as a 10-retry stall per node with a misleading message. The gate is a truthiness check now, matching the chat lane's, and covers the CLI and tree_optimize copies through the shared helper. The exec-synthesized plain functions were the one tool surface that raised instead of answering: a cloud-only parameter pruned from the local signature TypeErrored out of the call and aborted the framework loop. A wrapper turns binding rejections into the same INVALID_INPUT envelope call_tool returns, with the schema-bearing signature preserved via __signature__. A non-string doc_name likewise crashed get_close_matches, converting a clean NOT_FOUND (with similar names) into an INTERNAL_ERROR whose next_steps invited retrying the identical bad call. _openai_agent's nested pop evaluated its default eagerly, so a per-call api_base was popped and discarded whenever the client carried base_url โ€” and the chat lane resolved a different host than the SDK constructors for byte-identical config. The backend now normalizes through _sdk_backend at entry, where insertion order makes the per-call spelling win. chat_model stored the litellm/ routing prefix, so the documented model=client.chat_model hand-off sent a non-existent id to the Anthropic SDK. The attribute keeps the caller's spelling; the only consumer that needs the prefix (openai_agent_config) applies it at the config door โ€” which now also runs the py3.10 litellm type repair, since the BYO runner resolves that model through LiteLLM outside our completion helpers. Prompt caching gets its second breakpoint on every Claude lane: the chat-lane injection points add the newest-message mark (the pair LiteLLM itself seeds for Anthropic/Bedrock, extended to Vertex), messages() and anthropic_runner_config set the top-level cache_control so each loop turn re-reads the growing conversation (live-verified: delta-writes work, and a fifth breakpoint 400s โ€” messages() counts the caller's marks and stands down at four). The indexing lane is the inverse case โ€” single-shot unique prompts โ€” so its calls carry a no-match injection point that stops litellm 1.97 from seeding cache writes nothing ever reads back. --- pageindex/agent_tools.py | 32 ++++++++++++--- pageindex/client.py | 35 ++++++++++++++--- pageindex/local_chat.py | 42 +++++++++++++++----- pageindex/utils.py | 34 ++++++++++++---- tests/test_agent_tools.py | 44 +++++++++++++++++++++ tests/test_client.py | 82 +++++++++++++++++++++++++++++++++++++-- tests/test_local_chat.py | 34 +++++++++++++++- 7 files changed, 268 insertions(+), 35 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index 655224ad5..13922c96d 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -392,8 +392,8 @@ def _resolve_document( if matches: return max(matches, key=lambda d: d.get("createdAt") or ""), None names = [str(doc.get("name")) for doc in documents if doc.get("name")] - similar = difflib.get_close_matches(doc_name, names, n=_SIMILAR_NAMES_LIMIT, - cutoff=0.5) + 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) + "?" @@ -1358,7 +1358,8 @@ 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.""" + 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 {} @@ -1369,7 +1370,7 @@ def _make_tool_function(name: str, description: str, schema: dict, and param != "_invoke" for param in properties) if not params_usable: - def proxy(**kwargs: Any) -> str: + def inner(**kwargs: Any) -> str: return _invoke(kwargs)[0] else: ordered = ([p for p in properties if p in required] @@ -1382,7 +1383,7 @@ def proxy(**kwargs: Any) -> str: namespace: dict[str, Any] = {"_invoke": _invoke} exec(f"def _synthesized({rendered}):\n" f" return _invoke({args_literal})[0]", namespace) - proxy = namespace["_synthesized"] + inner = namespace["_synthesized"] annotations: dict[str, Any] = {} for p in ordered: annotation = _annotation_for(properties[p]) @@ -1392,7 +1393,26 @@ def proxy(**kwargs: Any) -> str: annotation = Optional[annotation] annotations[p] = annotation annotations["return"] = str - proxy.__annotations__ = annotations + inner.__annotations__ = annotations + + def proxy(*args: Any, **kwargs: Any) -> str: + # The invoker never raises, so a TypeError here is the binding + # rejecting the arguments (e.g. cloud-only parameters pruned + # from the local signature) โ€” answer with the same guided + # envelope call_tool returns for them. + 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 diff --git a/pageindex/client.py b/pageindex/client.py index 5d440cd61..3ea54ef43 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -149,7 +149,7 @@ def __init__( self.model = opt.model self.index_model = opt.index_model self.summary_model = opt.summary_model - self.chat_model = _agents_sdk_model_name(opt.chat_model) + self.chat_model = opt.chat_model self.chat_backend = chat_backend self.storage_path = storage_path or ".pageindex" from .local_api import LocalAPI @@ -678,7 +678,10 @@ def messages( 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. + ``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 all four + breakpoints. Args: messages: Native Messages-format history (including prior @@ -887,6 +890,16 @@ def openai_agent_config( 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 @@ -909,6 +922,11 @@ def openai_agent_config( 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, @@ -980,10 +998,14 @@ def anthropic_runner_config( Sugar over the explicit form โ€” ``agent_instructions`` (with ``doc_id`` targeting) as the system prompt and - ``as_anthropic_tools`` as the tools โ€” plus the same defaults - ``messages()`` applies: a per-model ``max_tokens`` and a - ``max_iterations`` bound of 10. To customize further, switch to - those methods directly. + ``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`` @@ -1013,6 +1035,7 @@ def anthropic_runner_config( "tools": self.as_anthropic_tools(include_management, asynchronous, doc_id=scope), "max_iterations": max_turns if max_turns is not None else 10, + "cache_control": {"type": "ephemeral"}, } def as_claude_mcp(self, include_management: bool = False, diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 0fe4fe787..cd10207c4 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -289,11 +289,11 @@ def _reported_model(model_name: str) -> str: 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 via LiteLLM's - injection param so the loop's later turns and a conversation's next - calls read it instead of repaying full price. Provider resolution is - LiteLLM's own, so this predicate can never disagree with where the - request actually routes.""" + 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: @@ -304,8 +304,12 @@ def _cache_extra_args(model_name: str) -> Optional[dict]: 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", "role": "system"}, + {"location": "message", "index": -1}]} return None @@ -337,12 +341,12 @@ def _openai_agent(client, protocol: str, model_name: str, instructions: str, if reasoning_effort is not None: extra_args = {**(extra_args or {}), "reasoning_effort": reasoning_effort} - conn = dict(backend) if backend else {} + 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", conn.pop("api_base", 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() @@ -843,6 +847,18 @@ def _anthropic_system(extra_system, block: Optional[str]) -> list[dict]: 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).""" @@ -915,6 +931,13 @@ def run_messages(client, messages, model: str, "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 {}) runner = _anthropic_client(_merged_backend(client, backend)) \ .beta.messages.tool_runner( max_tokens=(max_tokens if max_tokens is not None @@ -922,11 +945,12 @@ def run_messages(client, messages, model: str, messages=prepared, model=model, tools=build_anthropic_tools(client, doc_ids=doc_id), - system=_anthropic_system(system, block), + 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: diff --git a/pageindex/utils.py b/pageindex/utils.py index 75531cc5d..cde6b4d11 100644 --- a/pageindex/utils.py +++ b/pageindex/utils.py @@ -81,13 +81,14 @@ 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.""" - import litellm + invisible to env inspection โ€” the chat lane draws the same line. + Truthiness, not litellm's validate_environment, which reports a blank + exported key as present.""" wire = _strip_prefix(model, "litellm/") if "/" in wire and not wire.startswith("openai/"): return [] - env = litellm.validate_environment(wire if "/" in wire else f"openai/{wire}") - return [] if env["keys_in_environment"] else env["missing_keys"] + return ([] if (os.getenv("OPENAI_API_KEY") or "").strip() + else ["OPENAI_API_KEY"]) def _litellm_model(model, backend): @@ -130,6 +131,18 @@ 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): import litellm max_retries = 10 @@ -143,7 +156,7 @@ def llm_completion(model, prompt, chat_history=None, return_finish_reason=False) model=model, messages=messages, drop_params=True, - **(backend or {}), + **_no_cache_seeding_kwargs(backend), ) content = response.choices[0].message.content if return_finish_reason: @@ -176,7 +189,7 @@ async def llm_acompletion(model, prompt): model=model, messages=messages, drop_params=True, - **(backend or {}), + **_no_cache_seeding_kwargs(backend), ) return response.choices[0].message.content except Exception as e: @@ -710,6 +723,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( @@ -883,8 +898,11 @@ 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: diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index f260e0dd8..29592a7f8 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -773,6 +773,49 @@ def test_openai_agent_config_model_speaks_the_agents_sdk_grammar(tmp_path): == "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 @@ -792,6 +835,7 @@ def test_anthropic_runner_config_shapes(client, store_path): 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") diff --git a/tests/test_client.py b/tests/test_client.py index 63c16a598..8522d40ac 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -68,14 +68,19 @@ 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): @@ -294,6 +299,11 @@ def test_llm_completion_missing_key_raises_immediately(monkeypatch): # 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): @@ -665,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) diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index 1cd8cfd86..d1d16e653 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -315,7 +315,8 @@ def test_anthropic_routed_models_mark_managed_prefix_for_cache( fake_model([[_msg_item("ok")]]) from pageindex.local_chat import _openai_agent marked = {"cache_control_injection_points": [ - {"location": "message", "role": "system"}]} + {"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"): @@ -960,6 +961,9 @@ def test_agent_carries_prompt_cache_key_in_extra_body(monkeypatch): 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, @@ -1182,7 +1186,7 @@ def test_envelope_model_strips_litellm_routing_prefix(store_path, fake_model): seed_doc(store_path, "pi-a", "report.pdf") client = PageIndexLocalClient(storage_path=store_path, retrieve_model="anthropic/claude-x") - assert client.retrieve_model == "litellm/anthropic/claude-x" + assert client.retrieve_model == "anthropic/claude-x" fake_model([[_msg_item("ok")]]) result = client.chat_completions("q") assert result["model"] == "anthropic/claude-x" @@ -1713,6 +1717,32 @@ def test_backend_connection_reaches_each_engine(monkeypatch): 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 199584273ea775cd1aa7f867c3d6b1574be51d29 Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 19 Aug 2026 01:35:41 +0800 Subject: [PATCH 116/137] chore: drop LocalAPI's dead retrieve_model plumbing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit d375c00's tree-search retrieval consumed it; when retrieval became agent-driven the consumer went away and the constructor parameter stayed โ€” required, stored, never read. The client-level retrieve_model alias (property, setter, ConfigLoader chain) is unrelated and stays. --- pageindex/client.py | 1 - pageindex/local_api.py | 3 +-- tests/test_client.py | 2 +- 3 files changed, 2 insertions(+), 4 deletions(-) diff --git a/pageindex/client.py b/pageindex/client.py index 3ea54ef43..bde11627e 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -157,7 +157,6 @@ def __init__( storage_path=self.storage_path, model=self.model, summary_model=self.summary_model, - retrieve_model=self.chat_model, index_backend=index_backend, ) # LiteLLM's multi-second import would otherwise land on the diff --git a/pageindex/local_api.py b/pageindex/local_api.py index 83e846a63..59d12bccb 100644 --- a/pageindex/local_api.py +++ b/pageindex/local_api.py @@ -27,11 +27,10 @@ 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): + 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() diff --git a/tests/test_client.py b/tests/test_client.py index 8522d40ac..e74af3236 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -940,7 +940,7 @@ def test_backend_scopes_the_index_lane(tmp_path, monkeypatch): from pageindex.utils import _llm_backend, llm_completion api = LocalAPI(storage_path=str(tmp_path / "s"), model="m", - summary_model="s", retrieve_model="r", + 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"} From db3a0071872a7307a5d6db1e2a2c32cbbebe5b03 Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 19 Aug 2026 01:46:01 +0800 Subject: [PATCH 117/137] fix: optimize='full' without a key fails fast with a guided error MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The default flash call reached the missing-key failure only after the full PDF layout pass, surfacing as a raw litellm.AuthenticationError naming a model the caller never chose. The check now runs at entry โ€” before any PDF work โ€” and raises PageIndexAPIError naming both exits: configure a key, or take the LLM-free tree via optimize='merge'/False. Skipped when the indexing-lane backend contextvar carries connection overrides, and drawn on the same env-inspectable line as the shared pre-check (provider-prefixed models keep resolving credentials at call time). Ruled over the degrade-to-merge alternative: the tree stays identical across environments โ€” loud when impossible, never quietly less. --- pageindex/flash/api.py | 12 +++++++++++- tests/test_flash_extraction.py | 30 ++++++++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 1 deletion(-) diff --git a/pageindex/flash/api.py b/pageindex/flash/api.py index 26018a6c0..f6ba8d99c 100644 --- a/pageindex/flash/api.py +++ b/pageindex/flash/api.py @@ -99,7 +99,7 @@ def page_index_flash(pdf, summary=True, summary_model=None, optimize: str | bool = "full", optimize_expand=None, optimize_model=None, summary_concurrency=None, use_embedded_toc=True) -> dict: - """Build a PageIndex tree structure from a PDF using layout statistics. 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, ``"merge"`` for deterministic merge only, ``False`` to disable. ``True`` is accepted as ``"full"`` for backward compatibility. optimize_model: the LLM model for expand (defaults to the summary model). summary_concurrency: maximum simultaneous summary model calls; None uses the library default. use_embedded_toc: if True, consume the PDF's embedded bookmarks when trustworthy: deep bookmarks become the frame and the detected sections they lack are grafted back in after noise filtering, coarse ones become the chapter frame with detected nodes re-hung under them (deeper sparse entries are filled in when the page text confirms them, and garbled extracted titles are repaired from the bookmark strings), garbage ones are ignored; adds a ``toc_source`` key to the result. On by default; pass False for the pure detected structure. Returns: dict with keys ``doc_name``, ``doc_title``, ``structure`` (a list of nested ``{"title", "start_index", "end_index", "nodes"}`` dicts; page indexes are 1-based) and ``has_abstract_or_references_section`` (True when a top-level entry is an abstract or references heading). With ``optimize`` an ``optimize`` key reports merge/expand counts and before/after search-cost metrics. """ + """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. optimize_model: the LLM model for expand (defaults to the summary model). summary_concurrency: maximum simultaneous summary model calls; None uses the library default. use_embedded_toc: if True, consume the PDF's embedded bookmarks when trustworthy: deep bookmarks become the frame and the detected sections they lack are grafted back in after noise filtering, coarse ones become the chapter frame with detected nodes re-hung under them (deeper sparse entries are filled in when the page text confirms them, and garbled extracted titles are repaired from the bookmark strings), garbage ones are ignored; adds a ``toc_source`` key to the result. On by default; pass False for the pure detected structure. Returns: dict with keys ``doc_name``, ``doc_title``, ``structure`` (a list of nested ``{"title", "start_index", "end_index", "nodes"}`` dicts; page indexes are 1-based) and ``has_abstract_or_references_section`` (True when a top-level entry is an abstract or references heading). With ``optimize`` an ``optimize`` key reports merge/expand counts and before/after search-cost metrics. """ if optimize is True: optimize = "full" if not optimize: @@ -109,6 +109,16 @@ def page_index_flash(pdf, summary=True, summary_model=None, f"optimize must be 'full', 'merge', or False, got {optimize!r}") if optimize_expand is not None and optimize: optimize = "full" if optimize_expand else "merge" + 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: diff --git a/tests/test_flash_extraction.py b/tests/test_flash_extraction.py index 735d31ce4..4e9645f9f 100644 --- a/tests/test_flash_extraction.py +++ b/tests/test_flash_extraction.py @@ -21,3 +21,33 @@ def test_page_text_pins_pdfium5_semantics(): 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_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 + 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) From 7abf88a93b5a7acce45a20543942a89836cee1b7 Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 19 Aug 2026 02:20:18 +0800 Subject: [PATCH 118/137] =?UTF-8?q?fix:=20five=20review=20findings=20?= =?UTF-8?q?=E2=80=94=20one=20doc=20fetch,=20safe=20id=20reads,=20docstring?= =?UTF-8?q?s=20trimmed?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - doc_targeting_block collects missing IDs in its own fetch loop; the _doc_block pre-check that fetched every document a second time is deleted. Chat-surface error text unchanged (test-pinned); agent_instructions now reports the same batched message instead of get_document's raw error. - final_ids reads block.get("id"), matching the history_ids comprehension one line below. - _usage_sums skips a None usage. Both agents model classes guarantee a Usage instance today (verified in 0.20.0 and the 0.18.1 floor wheel); the guard is insurance against a future version dropping that. - Module and _openai_model docstrings cut to one line; the routing rationale lives in the history (28eaab8, d0007c2 era) and here: chat = LiteLLM (bare names get openai/ shorthand), responses = OpenAI-SDK native so provider-prefixed names are refused. --- pageindex/agent_tools.py | 11 ++++++++- pageindex/local_chat.py | 51 ++++------------------------------------ 2 files changed, 15 insertions(+), 47 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index 13922c96d..e7e6b6425 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -1600,7 +1600,16 @@ def doc_targeting_block(client, doc_id, scoped: bool = False) -> Optional[str]: raise PageIndexAPIError( "doc_id is empty. Pass one or more document IDs, or omit " "doc_id to give the agent the whole library.") - details = [client.get_document(one_id) for one_id in doc_ids] + details = [] + missing = [] + for one_id in doc_ids: + try: + details.append(client.get_document(one_id)) + except PageIndexAPIError: + missing.append(str(one_id)) + if missing: + raise PageIndexAPIError( + "Documents not found or access denied: " + ", ".join(missing)) listing = _all_documents(client) documents = ([{**detail, "id": one_id} for one_id, detail in zip(doc_ids, details)] diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index cd10207c4..2f0b2eba1 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -1,23 +1,4 @@ -"""Managed local chat: document-QA agents over the local tools. - -Three methods, three backend protocols, routed 1:1: ``chat_completions`` -drives the backend's /chat/completions (any OpenAI-compatible backend, -final answer only), ``responses`` drives /responses (official-shape -envelope; the full process transcript rides in ``items`` โ€” round-trip it -for provider prompt-cache continuation and agent memory), ``messages`` -drives Anthropic's /v1/messages via the SDK's own tool runner -(tool_use/tool_result round-trip is the format's native behavior). - -Content passes through untouched โ€” the caller's messages, the model's -answers, tool outputs. Native stop reasons pass through on ``messages``; -the OpenAI engine's abstraction does not surface per-turn finish reasons, -so ``chat_completions`` reports loop completion as ``"stop"``, while -``responses`` reports the backend's terminal ``status`` where the wire -surfaces one (recorded at the transport layer โ€” the framework discards -it). The SDK owns gatekeeping (structural validation), table-setting -(managed instructions, tools, doc targeting), tool execution, and billing -(usage aggregation, envelope ids). -""" +"""Managed local chat: document-QA agents over the local tools.""" from __future__ import annotations import asyncio @@ -51,17 +32,6 @@ def _doc_block(client, doc_id) -> Optional[str]: if not isinstance(doc_id, (str, list)): raise PageIndexAPIError("doc_id must be a string or a list of " "strings.") - doc_ids = [doc_id] if isinstance(doc_id, str) else list(doc_id) - missing = [] - for one_id in doc_ids: - try: - client.get_document(one_id) - except PageIndexAPIError: - missing.append(str(one_id)) - if missing: - raise PageIndexAPIError( - "Documents not found or access denied: " + ", ".join(missing) - ) # scoped: the chat surfaces also pass doc_id into the tool layer, so # name resolution happens inside the allowlist โ€” only a duplicate name # within the targeted set shadows. @@ -211,20 +181,7 @@ def _sdk_backend(backend) -> dict: def _openai_model(protocol: str, model_name: str, backend=None): - """The backend protocol driver โ€” the seam tests replace with a fake. - - chat protocol: LiteLLM, full stop โ€” model names mean what LiteLLM says - they mean. Bare names are OpenAI-compatible shorthand (wire form - ``openai/``, so OPENAI_API_KEY / OPENAI_BASE_URL keep selecting - the backend), a ``litellm/`` prefix strips, and a first segment LiteLLM - does not know (a HuggingFace repo id like ``Qwen/...``) is refused with - the ``openai/`` form instead of failing inside LiteLLM at request time. - - responses protocol: the Responses API is OpenAI-SDK native โ€” LiteLLM's - completion surface speaks the chat.completions format, so - provider-prefixed models are refused instead of silently downgrading; - a ``litellm/`` prefix strips first (it is routing grammar, not a - provider), then bare and ``openai/`` names drive the OpenAI SDK.""" + """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/"): @@ -482,6 +439,8 @@ def _wrap_max_turns(max_turns) -> PageIndexAPIError: 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 completion += r.usage.output_tokens details = getattr(r.usage, "input_tokens_details", None) @@ -1002,7 +961,7 @@ def capture(params): new_messages = [_dump_message(message) for message in conversation[len(prepared):]] final_blocks = [_dump_block(item) for item in final.content] - final_ids = {block["id"] for block in final_blocks + 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 From c129c2ebcef04f1c64161c98d19fc059345052db Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 19 Aug 2026 03:10:54 +0800 Subject: [PATCH 119/137] fix: guard _usage_sums against None token counts from non-standard backends --- pageindex/local_chat.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 2f0b2eba1..8268eb91b 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -441,8 +441,8 @@ def _usage_sums(raw_responses) -> "tuple[int, int, int, int, int]": for r in raw_responses: if r.usage is None: continue - prompt += r.usage.input_tokens - completion += r.usage.output_tokens + 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 From 86c257df9f2cf6c3a8621e562b4247c5a32088bc Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 19 Aug 2026 03:11:35 +0800 Subject: [PATCH 120/137] fix: raise non-stream chat_completions timeout to 600s 300s was below the industry standard (OpenAI/Anthropic SDKs both default to 600s) and could clip complex multi-document agent runs the cloud endpoint processes synchronously. --- pageindex/cloud_api.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pageindex/cloud_api.py b/pageindex/cloud_api.py index ab7a9c885..e42363522 100644 --- a/pageindex/cloud_api.py +++ b/pageindex/cloud_api.py @@ -225,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: From fc2dd3cb94061d7acac5365ae6a0141ffc4ff5f7 Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 19 Aug 2026 03:34:51 +0800 Subject: [PATCH 121/137] =?UTF-8?q?fix:=20three=20review=20findings=20?= =?UTF-8?q?=E2=80=94=20page=20budget=20in=20emitted=20units,=20demo=20cach?= =?UTF-8?q?e=20cut,=20breakpoint=20docs?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit get_page_content sized its response budget in raw characters while the envelope is emitted as JSON, where every quote, backslash, newline and tab costs an escape character โ€” quote-dense tables measured 108,777 and code-dense pages 116,506 against the 100,000 limit while the accounting said 94,400. The same sizing-units-vs-emitted-units defect was raised on the structure path in #393 and fixed there via _serialized_size; the sibling never got the fix. The budget now measures the serialized page entry plus json's ", " item separator, closing the two residual gaps an adversarial sweep still breached (per-entry shell: 300 short pages; separators: 4000 tiny pages, 102,274 emitted). All shapes now land under the limit except the deliberate first-page exception. The budget test's pages are now escape-dense โ€” raw length under the budget, serialized length over โ€” so it fails on the old accounting, and it asserts the real property: the emitted envelope stays within TOOL_RESPONSE_CHAR_LIMIT. The demo's .doc_id sidecar cache is deleted: the name lookup kept right below it resolves every re-run (submit's uniquing suffixes later duplicates only, the first copy keeps its plain name), so the cache bought one list_documents call on a local store ahead of a multi-minute index step, at the cost of a read/validate/unlink block, two write-backs, a disk artifact, and a .gitignore entry. Step 1 returns to its 0.2.9 shape, wait=True kept. Verified end to end against the seeded store. messages() docstring promised the top-level cache_control is skipped "when your own blocks already use all four breakpoints" โ€” the managed prefix always holds one, so a caller placing four puts five marks on the wire and the API rejects the request. Now states the caller budget of three, matching the _cache_marks < 4 gate. Non-changes: _split_structure has the same separator undercount but its node counts are too small to breach (no measured overshoot); the first-page overshoot exception stays (one page must always return); the _split_oversized_node budget floor stays (trigger is ~16x beyond real summary sizes). --- .gitignore | 1 - examples/agentic_vectorless_rag_demo.py | 23 ++++------------------- pageindex/agent_tools.py | 8 +++++--- pageindex/client.py | 4 ++-- tests/test_agent_tools.py | 8 ++++++-- 5 files changed, 17 insertions(+), 27 deletions(-) diff --git a/.gitignore b/.gitignore index b5c223b31..5193735ca 100644 --- a/.gitignore +++ b/.gitignore @@ -6,4 +6,3 @@ __pycache__ logs/ .pageindex/ dist/ -*.doc_id diff --git a/examples/agentic_vectorless_rag_demo.py b/examples/agentic_vectorless_rag_demo.py index 5b25c7638..93a00735c 100644 --- a/examples/agentic_vectorless_rag_demo.py +++ b/examples/agentic_vectorless_rag_demo.py @@ -32,14 +32,13 @@ from agents.stream_events import RawResponsesStreamEvent, RunItemStreamEvent from openai.types.responses import ResponseTextDeltaEvent, ResponseReasoningSummaryTextDeltaEvent -from pageindex import PageIndexAPIError, PageIndexLocalClient +from pageindex import PageIndexLocalClient import pageindex.utils as utils PDF_URL = "https://arxiv.org/pdf/2603.15031" _EXAMPLES_DIR = Path(__file__).parent PDF_PATH = _EXAMPLES_DIR / "documents" / "attention-residuals.pdf" -DOC_ID_PATH = _EXAMPLES_DIR / "documents" / "attention-residuals.doc_id" STORAGE_PATH = _EXAMPLES_DIR / ".pageindex" @@ -127,27 +126,13 @@ async def _run(): print("=" * 60) print("Step 1: Index PDF and view tree structure") print("=" * 60) - doc_id = None - if DOC_ID_PATH.exists(): - cached = DOC_ID_PATH.read_text().strip() - try: - client.get_document(cached) - doc_id = cached - except PageIndexAPIError: - DOC_ID_PATH.unlink() - if doc_id is None: - # The .doc_id cache is gitignored โ€” on a fresh clone with an - # existing store, find the already-indexed copy by name instead of - # re-indexing it. - doc_id = next( - (doc["id"] for doc in client.list_documents(limit=100)["documents"] - if doc["name"] == PDF_PATH.name), None) + doc_id = next( + (doc["id"] for doc in client.list_documents(limit=100)["documents"] + if doc["name"] == PDF_PATH.name), None) if doc_id: - DOC_ID_PATH.write_text(doc_id) print(f"\nLoaded cached doc_id: {doc_id}") else: doc_id = client.submit_document(str(PDF_PATH), wait=True)["doc_id"] - DOC_ID_PATH.write_text(doc_id) print(f"\nIndexed. doc_id: {doc_id}") print("\nTree Structure (top-level sections):") structure = client.get_tree(doc_id, node_summary=True)["result"] diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index e7e6b6425..feb0e5e07 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -1054,10 +1054,12 @@ def _get_page_content(client, doc_name: str, pages: str, markdown = item.get("markdown") if item else None text = (markdown if isinstance(markdown, str) else f"Page {page} content not available") - if not included or budget - len(text) >= 0: - content.append({"page": page, "text": text}) + 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 -= len(text) + budget -= size else: remaining.append(page) diff --git a/pageindex/client.py b/pageindex/client.py index bde11627e..53c5502e1 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -679,8 +679,8 @@ def messages( 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 all four - breakpoints. + 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 diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 29592a7f8..e56a026f2 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -408,9 +408,11 @@ def test_created_at_accepts_z_suffixed_input(client, store_path): 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": "x" * 96_000}, - {"page_index": 2, "markdown": "short"}, + {"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", @@ -420,6 +422,8 @@ def test_page_content_char_budget(client, store_path): 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( From b988fc688c1ee919b0dec4500bfcb01ed113a3d4 Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 19 Aug 2026 03:49:25 +0800 Subject: [PATCH 122/137] =?UTF-8?q?docs:=20two=20review=20wording=20fixes?= =?UTF-8?q?=20=E2=80=94=20agent=5Ftools()=20read-only=20default,=20README?= =?UTF-8?q?=20flag=20scoping?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit agent_tools()'s lead sentence promised "the full cloud tool set" while the default exposes only tools the server marks read-only โ€” the Args block and both sibling exports already say "the full live read tool set"; the lead now matches. The README's optional-arguments note scoped the whole block to --mode standard; --index-model works in the default flash mode, so the note now names the flags it covers. --- README.md | 2 +- pageindex/client.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index dca5dbc60..4a086ed6f 100644 --- a/README.md +++ b/README.md @@ -173,7 +173,7 @@ python3 run_pageindex.py --pdf_path /path/to/your/document.pdf
Optional parameters
-You can customize the processing with additional optional arguments (the structure-tuning flags below require --mode standard): +You can customize the processing with additional optional arguments (the structure-tuning flags from --toc-check-pages down require --mode standard): ``` --mode Processing mode: flash (default) or standard diff --git a/pageindex/client.py b/pageindex/client.py index 53c5502e1..9fdae9356 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -789,7 +789,7 @@ def agent_tools(self, include_management: bool = False) -> list[Callable[..., st For the OpenAI / Claude Agent SDKs, prefer ``as_openai_tools()`` / ``as_claude_mcp()``. - Cloud: the full cloud tool set, discovered live from the PageIndex + 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 From 39b13b529b7956c2ce5ddf9ad9d4313693d8834b Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 19 Aug 2026 03:59:48 +0800 Subject: [PATCH 123/137] fix: annotate mixed-type dict literals for Pyright --- pageindex/cloud_api.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pageindex/cloud_api.py b/pageindex/cloud_api.py index e42363522..f930c1805 100644 --- a/pageindex/cloud_api.py +++ b/pageindex/cloud_api.py @@ -59,7 +59,7 @@ def submit_document( (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: @@ -356,7 +356,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 From 3b71c70bd3189f1bfa6dda2acafc6df39d1d1eb7 Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 19 Aug 2026 04:11:49 +0800 Subject: [PATCH 124/137] =?UTF-8?q?fix:=20three=20review=20findings=20?= =?UTF-8?q?=E2=80=94=20surrogate=20census,=20bridge=20bool=20coercion,=20l?= =?UTF-8?q?isting=20early-exit?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - The page-mode unicode walk consumes char_extract's raw_chars (surrogate pairs already merged) instead of re-reading the textpage, which split astral chars back into two lone-surrogate slots, desynced the walk against their one-char cmap targets, and silently dropped the whole page's patch. One census for both modes makes the desync structurally impossible; negative-codepoint chars now follow that census too (absent, conservative rollback) instead of walking as "\x00". A/B over the nine example PDFs: structures byte-identical. - _coerce_bool_args is schema-driven and runs on the cloud bridge invoker as well; the coercion previously landed only on call_tool, so a model's "false" reached the wire verbatim โ€” schema-validating servers rejected the call, lenient ones read it truthy (wait_for_completion blocking up to 3 minutes). Server-supplied schemas cover cloud-only tools the local contract never named. - _all_documents stops paging once every wanted id has been seen (both modes list newest-first): doc-scoped chat turns, scoped resolution, browse and remove no longer sweep the whole library to use a handful of entries; an id absent from the listing still costs the full sweep. _index_flash also drops its add_node_text call โ€” the text was built, excluded from the description prompt, and stripped before save. --- pageindex/agent_tools.py | 50 ++++++++++++------- .../flash/parser_pdfium_charlevel/pipeline.py | 2 +- .../parser_pdfium_charlevel/unicode_apply.py | 16 +++--- pageindex/local_api.py | 7 ++- tests/test_agent_tools.py | 33 +++++++++++- tests/test_flash_extraction.py | 22 ++++++++ 6 files changed, 97 insertions(+), 33 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index feb0e5e07..65c11ee50 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -327,15 +327,23 @@ def _dumps(payload: dict[str, Any]) -> str: # โ”€โ”€ document listing / name resolution โ”€โ”€ -def _all_documents(client) -> list[dict[str, Any]]: +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).""" + 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) @@ -386,7 +394,7 @@ def _resolve_document( """Resolve doc_name to a list entry. Same-name duplicates resolve to the newest match. Returns (entry, None) or (None, error_payload_pair).""" if documents is None: - documents = _all_documents(client) + documents = _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: @@ -734,7 +742,8 @@ def _browse_documents(client, folder_id: str = "root", recursive: bool = False, window = listing.get("documents") or [] total = listing.get("total") else: - scoped = _scope_documents(_all_documents(client), _allowed_ids) + 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) @@ -1123,7 +1132,7 @@ def _remove_document(client, doc_names: list[str], {"summary": "Too many documents in one call", "options": ["Delete at most 10 documents per call"]}, "INVALID_INPUT") - documents = _all_documents(client) + documents = _all_documents(client, stop_ids=_allowed_ids) results = [] for doc_name in doc_names: entry, error = _resolve_document(client, doc_name, documents=documents, @@ -1162,11 +1171,11 @@ def tool_names(include_management: bool = False) -> tuple[str, ...]: return _READ_TOOLS + (_MANAGEMENT_TOOLS if include_management else ()) -def _coerce_bool_args(name: str, kwargs: dict[str, Any]) -> None: +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.""" - properties = TOOL_CONTRACT.get(name, {}).get("schema", {}).get( - "properties", {}) + 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): @@ -1204,7 +1213,7 @@ def call_tool(client, name: str, arguments: dict[str, Any], # "omit if ..." semantics, same as the cloud bridge invoker). kwargs = {key: value for key, value in (arguments or {}).items() if not key.startswith("_") and value is not None} - _coerce_bool_args(name, kwargs) + _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 @@ -1333,15 +1342,18 @@ def _annotation_for(spec: dict) -> Any: return Optional[base] if nullable else base -def _bridge_invoker(bridge, name: str) -> "Callable[[dict], tuple[str, bool]]": - """One cloud tool call proxied over MCP: None-valued arguments are - dropped (None โ‰ก omitted, matching the contract's "omit if ..." - semantics) and failures are contained in the error envelope. Returns - (envelope_text, is_error), like call_tool.""" +def _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. 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: payload, _ = _failure( @@ -1490,7 +1502,8 @@ def _tool_specs(client, include_management: bool = False, doc_ids=None, meta.get("description") or "", copy.deepcopy(meta.get("inputSchema")) or {"type": "object", "properties": {}}, - _bridge_invoker(bridge, str(meta.get("name") or "tool"))) + _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]]": @@ -1612,7 +1625,10 @@ def doc_targeting_block(client, doc_id, scoped: bool = False) -> Optional[str]: if missing: raise PageIndexAPIError( "Documents not found or access denied: " + ", ".join(missing)) - listing = _all_documents(client) + # 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) 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/local_api.py b/pageindex/local_api.py index 59d12bccb..4fe9c617c 100644 --- a/pageindex/local_api.py +++ b/pageindex/local_api.py @@ -106,7 +106,7 @@ def submit_document( try: if mode == "flash": structure, description = run_off_loop( - self._with_backend, self._index_flash, file_path, page_texts + self._with_backend, self._index_flash, file_path ) else: structure, description = run_off_loop( @@ -184,9 +184,9 @@ 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, @@ -199,7 +199,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, diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index e56a026f2..05fd21deb 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -1551,7 +1551,27 @@ def _synth(bridge, meta): name = meta["name"] return _make_tool_function(name, meta.get("description"), meta["inputSchema"], - _bridge_invoker(bridge, name)) + _bridge_invoker(bridge, name, + meta["inputSchema"])) + + +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(): @@ -1835,6 +1855,17 @@ def list_documents(self, limit, offset): 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 diff --git a/tests/test_flash_extraction.py b/tests/test_flash_extraction.py index 4e9645f9f..9f1786670 100644 --- a/tests/test_flash_extraction.py +++ b/tests/test_flash_extraction.py @@ -23,6 +23,28 @@ def test_page_text_pins_pdfium5_semantics(): 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 From 4563223c7472229b17bbbe5afa1c9583a775e4df Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 19 Aug 2026 05:32:53 +0800 Subject: [PATCH 125/137] fix: nine approved review cleanups + arm the live drift tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - llm_completion/llm_acompletion pass max_retries=0 via a dict-merge a backend override still wins: the 10-try loop is the retry policy, and litellm's client default (2) silently multiplied it 3x after the SDK migration dropped the original openai.OpenAI(max_retries=0). 429/5xx keep flowing to the loop (_UNRECOVERABLE_STATUS is 401/403/404 only); anthropic's payload builder never sees the param (source-verified). - exec-synthesized tools set __name__/__qualname__ to the tool name โ€” binding TypeErrors quote __qualname__, so the model saw "_synthesized() got an unexpected keyword argument" with no tool name to correct against. Pinned by test. - McpBridge stubs EmbeddedResource blobs (resource.blob one level down); only top-level data payloads were stubbed, so a nested base64 image was json.dumps'd into the model's context. Pinned by test. - run_messages closes its anthropic client on both paths: try/finally after the non-stream iteration (the params read-back is offline), and a finally in events() that runs on exhaustion and abandonment alike. Each call leaked an httpx pool until GC. - run_pageindex: the md branch filters None user options before ConfigLoader.load โ€” merged {**defaults, **user} let unset CLI flags clobber yaml defaults, silently dropping node ids and summaries from md trees; --optimize rejects every value outside flash+pdf (is not None, matching the sibling guards) instead of letting 'off' slip; dead _openai_missing_keys import dropped. - tests.yml uninstalls openai-agents on the without-frameworks leg: requirements.txt now carries it, so the leg had become a duplicate of the with-leg and the 41 importorskip guards never skipped. - test_optimize_full_fails_fast_without_a_key imports litellm before delenv (its import may load a .env on a daemon thread after the client constructor returns; the sibling tests already guard this). - test_agent_tools imports pageindex.utils before reading PAGEINDEX_API_KEY: the lazy package import never triggered utils' load_dotenv, so the module-scope getenv always ran before the .env load and the three live drift tests had never executed even with a key present. First live run: all green, including frozen-contract parity against the real tools/list. --- .github/workflows/tests.yml | 3 +++ pageindex/agent_tools.py | 2 ++ pageindex/local_chat.py | 10 ++++++++-- pageindex/mcp_bridge.py | 8 ++++++++ pageindex/utils.py | 5 +++-- run_pageindex.py | 6 +++--- tests/test_agent_tools.py | 22 ++++++++++++++++++++++ tests/test_flash_extraction.py | 1 + 8 files changed, 50 insertions(+), 7 deletions(-) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 6b349a09b..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 anthropic - run: python -m pytest -q diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index 65c11ee50..9f942d9ef 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -1398,6 +1398,8 @@ def inner(**kwargs: Any) -> str: 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]) diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 8268eb91b..5121d2a5c 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -897,8 +897,8 @@ def run_messages(client, messages, model: str, cached: dict[str, Any] = ( {"cache_control": {"type": "ephemeral"}} if _cache_marks(system_blocks, prepared) < 4 else {}) - runner = _anthropic_client(_merged_backend(client, backend)) \ - .beta.messages.tool_runner( + backend_client = _anthropic_client(_merged_backend(client, backend)) + runner = backend_client.beta.messages.tool_runner( max_tokens=(max_tokens if max_tokens is not None else _default_max_tokens(model)), messages=prepared, @@ -921,6 +921,9 @@ def events() -> Iterator[Any]: except anthropic.AnthropicError as exc: raise PageIndexAPIError( f"The model backend failed: {exc}") from exc + finally: + # runs on exhaustion and abandonment (GeneratorExit) alike + backend_client.close() return events() try: @@ -928,6 +931,9 @@ def events() -> Iterator[Any]: 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 + backend_client.close() if not turns: raise PageIndexAPIError("The model returned no response.") captured: dict = {} diff --git a/pageindex/mcp_bridge.py b/pageindex/mcp_bridge.py index 203226ced..3ac9adbde 100644 --- a/pageindex/mcp_bridge.py +++ b/pageindex/mcp_bridge.py @@ -211,6 +211,14 @@ def call_tool(self, name: str, arguments: dict[str, Any]) -> "tuple[str, bool]": 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/utils.py b/pageindex/utils.py index cde6b4d11..842a3f268 100644 --- a/pageindex/utils.py +++ b/pageindex/utils.py @@ -156,7 +156,8 @@ def llm_completion(model, prompt, chat_history=None, return_finish_reason=False) model=model, messages=messages, drop_params=True, - **_no_cache_seeding_kwargs(backend), + # 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: @@ -189,7 +190,7 @@ async def llm_acompletion(model, prompt): model=model, messages=messages, drop_params=True, - **_no_cache_seeding_kwargs(backend), + **{"max_retries": 0, **_no_cache_seeding_kwargs(backend)}, ) return response.choices[0].message.content except Exception as e: diff --git a/run_pageindex.py b/run_pageindex.py index 054ebf429..91ddf0763 100644 --- a/run_pageindex.py +++ b/run_pageindex.py @@ -62,7 +62,7 @@ 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 in ('full', 'merge') and not (args.pdf_path and args.mode == 'flash'): + 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' @@ -158,7 +158,7 @@ import asyncio # Use ConfigLoader to get consistent defaults (matching PDF behavior) - from pageindex.utils import ConfigLoader, _openai_missing_keys + from pageindex.utils import ConfigLoader config_loader = ConfigLoader() # Create options dict with user args @@ -172,7 +172,7 @@ } # 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}) toc_with_page_number = asyncio.run(md_to_tree( md_path=args.md_path, diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 05fd21deb..8630e6c23 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -1534,12 +1534,15 @@ def test_mcp_bridge_blob_blocks_become_stubs(): 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 โ”€โ”€ @@ -1555,6 +1558,24 @@ def _synth(bridge, meta): 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, @@ -1955,6 +1976,7 @@ def flaky(doc_id): 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") diff --git a/tests/test_flash_extraction.py b/tests/test_flash_extraction.py index 9f1786670..f16c58495 100644 --- a/tests/test_flash_extraction.py +++ b/tests/test_flash_extraction.py @@ -54,6 +54,7 @@ def test_optimize_full_fails_fast_without_a_key(tmp_path, monkeypatch): 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) From 50d86bff7f3de051b1940ea1cf8566554eb21dfd Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 19 Aug 2026 05:52:59 +0800 Subject: [PATCH 126/137] fix: drop mcp internal API usage that broke on mcp 2.0 CI pulled mcp 2.0.0 (via openai-agents), which removed Server.request_handlers. Use _tool_specs invoke callables instead. --- tests/test_agent_tools.py | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 8630e6c23..877875e1e 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -903,21 +903,22 @@ def test_anthropic_runner_config_doc_scope_enforced_in_tools(client, def test_claude_agent_config_doc_scope_enforced_in_tools(client, store_path): + """claude_agent_config(doc_id=...) must wire scope all the way into the + tool invoke callables โ€” an out-of-scope document returns NOT_FOUND.""" pytest.importorskip("claude_agent_sdk") - from mcp.types import CallToolRequest, CallToolRequestParams + from pageindex.agent_tools import _tool_specs seed_doc(store_path, "pi-a", "report.pdf") seed_doc(store_path, "pi-b", "payroll.pdf", created_at="2026-08-02T10:00:00.123000") config = client.claude_agent_config(doc_id="pi-a") - server = config["mcp_servers"]["pageindex"] - handler = server["instance"].request_handlers[CallToolRequest] - result = asyncio.run(handler(CallToolRequest( - method="tools/call", - params=CallToolRequestParams( - name="get_page_content", - arguments={"doc_name": "payroll.pdf", "pages": "1"})))) - payload = json.loads(result.root.content[0].text) - assert payload["errorCode"] == "NOT_FOUND" + assert "report.pdf" in config["system_prompt"] + specs = dict((name, invoke) + for name, _desc, _schema, invoke + in _tool_specs(client, doc_ids=["pi-a"])) + text, is_error = specs["get_page_content"]( + {"doc_name": "payroll.pdf", "pages": "1"}) + assert is_error + assert json.loads(text)["errorCode"] == "NOT_FOUND" def test_openai_agent_config_scoped_shadow_check(client, store_path): From 16bd56f63e81ca04c879e6b1b78e5f2352c6750a Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 19 Aug 2026 06:39:08 +0800 Subject: [PATCH 127/137] =?UTF-8?q?fix:=20three=20review=20findings=20?= =?UTF-8?q?=E2=80=94=20transport-error=20fidelity,=20cache=20key=20doc=20s?= =?UTF-8?q?cope,=20one=20created=5Fat?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - doc_targeting_block batches a doc_id into "Documents not found or access denied" only on a definite 403/404 (or a local raise, which carries no status); a cloud 429/5xx in the fetch loop now propagates with its original text instead of masquerading as a missing document. PageIndexAPIError grows an optional status_code, set at the cloud get_document raise โ€” main let these errors through raw; the batching added in 7abf88a had widened to every failure. - _conversation_cache_key seeds doc_id alongside the first item: the same opening question against different documents is different conversations with different prefixes, yet pooled under one prompt_cache_key before this. A str doc_id and its one-item list form hash identically (same targeting, same prefix). - run_responses mints created_at once beside response_id: the streamed response.created event and the terminal envelope now carry the same timestamp for the same response id (the created event previously kept the backend's per-turn timestamp with only the id patched). --- pageindex/agent_tools.py | 6 +++++- pageindex/cloud_api.py | 4 +++- pageindex/errors.py | 7 ++++++- pageindex/local_chat.py | 33 ++++++++++++++++++++------------- tests/test_agent_tools.py | 22 ++++++++++++++++++++++ tests/test_local_chat.py | 30 +++++++++++++++++++++--------- 6 files changed, 77 insertions(+), 25 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index 9f942d9ef..826c3d6ea 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -1622,7 +1622,11 @@ def doc_targeting_block(client, doc_id, scoped: bool = False) -> Optional[str]: for one_id in doc_ids: try: details.append(client.get_document(one_id)) - except PageIndexAPIError: + 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( diff --git a/pageindex/cloud_api.py b/pageindex/cloud_api.py index f930c1805..f3a7740b6 100644 --- a/pageindex/cloud_api.py +++ b/pageindex/cloud_api.py @@ -312,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]: 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/local_chat.py b/pageindex/local_chat.py index 5121d2a5c..2c5dd49f8 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -337,17 +337,21 @@ def _validate_max_turns(max_turns) -> None: raise PageIndexAPIError("max_turns must be a positive integer.") -def _conversation_cache_key(model_name: str, instructions: str, items) -> str: +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, - first conversation item โ€” so a conversation's continuations share one - route without pooling unrelated conversations. Callers pass the - conversation's own items, never the SDK-prepended doc-targeting block: - that block is byte-identical for every conversation about a document - and would pool them all under one key.""" - seed = json.dumps([model_name, instructions, + 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] @@ -500,8 +504,8 @@ def run_chat_completions(client, messages, stream: bool = False, 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, history), + 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), @@ -617,8 +621,8 @@ def run_responses(client, input, model: Optional[str] = None, 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, - conversation), + 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), @@ -630,12 +634,13 @@ def run_responses(client, input, model: Optional[str] = None, 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": int(time.time()), + "created_at": created_at, "model": _reported_model(model_name), "status": recorded.get("status") or "completed", "output": [item for item in transcript @@ -700,10 +705,12 @@ async def agen(): if data.get("type") in lifecycle: if data["type"] == "response.created" and not opened: # N per-turn openings collapse to one, carrying - # the id the terminal event will report. + # 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 diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 877875e1e..43cd1e412 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -2318,6 +2318,28 @@ def test_config_helpers_reject_empty_doc_id_on_cloud(): 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).""" diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index d1d16e653..dc2c758eb 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -535,8 +535,8 @@ def test_doc_id_conversations_get_distinct_cache_keys(client, store_path, keys = [] real = local_chat._conversation_cache_key - def spy(model_name, instructions, items): - key = real(model_name, instructions, items) + def spy(model_name, instructions, doc_id, items): + key = real(model_name, instructions, doc_id, items) keys.append(key) return key @@ -561,6 +561,11 @@ def spy(model_name, instructions, items): 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): @@ -609,6 +614,8 @@ def test_responses_stream_opens_with_created(client, store_path, fake_model): 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 @@ -930,17 +937,22 @@ def test_sol_class_refusal_names_its_exits(): 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 / first item), and unrelated conversations must not pool - under it.""" + 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", turn1) - assert key == local_chat._conversation_cache_key("m", "sys", continuation) + 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", [{"role": "user", "content": "other"}]) - assert key != local_chat._conversation_cache_key("m2", "sys", turn1) - assert key != local_chat._conversation_cache_key("m", "sys2", turn1) + "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 From c57da99f313b6ed02f98e161220f6274c6367760 Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 19 Aug 2026 06:54:31 +0800 Subject: [PATCH 128/137] fix: re-arm the claude_agent_config doc-scope guard, correct a stale docstring - 50d86bf's mcp-2.0 migration rewrote test_claude_agent_config_doc_scope_enforced_in_tools to build its own callables via _tool_specs, so nothing drove the handlers build_claude_mcp actually registers. Dropping doc_ids at any hop of claude_agent_config -> as_claude_mcp -> build_claude_mcp left the whole suite green, and the test's own docstring ("wire scope all the way into the tool invoke callables") was no longer true. The assertion now spies create_sdk_mcp_server and calls the registered SdkMcpTool.handler -- public API, so it survives mcp 2.0 without reaching for the Server.request_handlers the migration had to drop. A browse assertion catches over-tight scope the way the openai/anthropic siblings do. All four mutations (three hops plus over-scoping) now fail. - The module docstring still said arguments outside a pruned local signature "fail at the Python call boundary". 03ffab3 made proxy() catch that binding TypeError and return the same guided envelope call_tool returns, as _make_tool_function's docstring and test_plain_functions_answer_bad_arguments_with_the_envelope already state; the header was the last copy of the old behavior. --- pageindex/agent_tools.py | 9 ++++----- tests/test_agent_tools.py | 34 +++++++++++++++++++++++----------- 2 files changed, 27 insertions(+), 16 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index 826c3d6ea..24af5fb53 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -9,11 +9,10 @@ way the agent instructions do โ€” they never teach capabilities that only exist on the cloud. -Tools never raise for any invocation their signatures accept: every -outcome, including errors, is returned as the same JSON envelope the cloud -emits ({"success": true, ...} / {"error": ...}). Arguments outside a pruned -local signature fail at the Python call boundary; the call_tool path -answers them with the guided error envelope instead. +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 diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 43cd1e412..1ac63dc63 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -902,23 +902,35 @@ def test_anthropic_runner_config_doc_scope_enforced_in_tools(client, assert [doc["name"] for doc in browse["documents"]] == ["report.pdf"] -def test_claude_agent_config_doc_scope_enforced_in_tools(client, store_path): +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 - tool invoke callables โ€” an out-of-scope document returns NOT_FOUND.""" - pytest.importorskip("claude_agent_sdk") - from pageindex.agent_tools import _tool_specs + 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"] - specs = dict((name, invoke) - for name, _desc, _schema, invoke - in _tool_specs(client, doc_ids=["pi-a"])) - text, is_error = specs["get_page_content"]( - {"doc_name": "payroll.pdf", "pages": "1"}) - assert is_error - assert json.loads(text)["errorCode"] == "NOT_FOUND" + 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): From 594c03dcc4020b775c8304dedd8df0247bf010f3 Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 19 Aug 2026 07:57:32 +0800 Subject: [PATCH 129/137] fix: stop spawn workers from re-running unguarded caller scripts Flash's parallel parser spawns workers for PDFs >= 64 pages; Python's spawn bootstrap re-imports the caller's __main__, so a script without the if __name__ == '__main__' guard re-ran wholesale in every worker. The parser's broad sequential fallback then swallowed the RuntimeError Python raises for exactly this case, so each worker completed a full duplicate index: N documents in the store and N times the LLM spend, with zero warnings. Default-on since flash became the default mode (v0.2.10.dev1). Three independent layers, in order of engagement: - _anonymous_main hides __main__'s __file__/__spec__ while workers spawn. Workers import everything by module name and never need the caller's script, so spawn simply skips the re-import: unguarded scripts now get full parallel speed (measured at parity with guarded ones; the old behavior was 3x slower wall-clock on top of the 8x duplication). - The fallback re-raises instead of running sequentially when the process is a spawn child mid-bootstrap (_inheriting), so if the hiding ever stops working the duplicate run dies loudly with Python's canonical guard message instead of silently completing. - submit_document refuses outright during a foreign spawn bootstrap, which also covers a user's own unguarded pool re-running top-level submits of small documents that never reach the parser. _inheriting is private CPython API; it is read via getattr with a False default and only powers the insurance layers, so its removal degrades to the old fallback rather than breaking anything. --- pageindex/flash/parser_pdfium_parallel.py | 32 ++++++++++- pageindex/local_api.py | 7 +++ tests/test_flash_extraction.py | 68 +++++++++++++++++++++++ 3 files changed, 106 insertions(+), 1 deletion(-) 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/local_api.py b/pageindex/local_api.py index 4fe9c617c..9b29ac6e9 100644 --- a/pageindex/local_api.py +++ b/pageindex/local_api.py @@ -3,6 +3,7 @@ import json import logging +import multiprocessing import os import uuid from datetime import datetime, timezone @@ -55,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." diff --git a/tests/test_flash_extraction.py b/tests/test_flash_extraction.py index f16c58495..d82a1e4ff 100644 --- a/tests/test_flash_extraction.py +++ b/tests/test_flash_extraction.py @@ -74,3 +74,71 @@ def test_optimize_full_fails_fast_without_a_key(tmp_path, monkeypatch): 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 From 914dc431a76c42bc263abf5552370a91f516efd6 Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 19 Aug 2026 16:45:55 +0800 Subject: [PATCH 130/137] =?UTF-8?q?fix:=20ten=20review=20findings=20?= =?UTF-8?q?=E2=80=94=20doc=5Fid=20scoping,=20honest=20chat=20envelopes,=20?= =?UTF-8?q?thinking-safe=20defaults?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ten approved findings from the Aug 19 max review of PR #400: - agent_tools() gains doc_id โ€” the one BYO tool surface without it: passed through to _tool_specs' enforced allowlist, refused loudly on cloud like its siblings. - Chat streams request usage (ModelSettings.include_usage=True): the stream_metadata terminal chunk carried all zeros. Live-verified real counts on both OpenAI and Anthropic backends; agents forwards the flag as stream_options only on streaming calls, so non-stream runs are untouched, and litellm consumes it itself for providers without native stream_options support. - chat_completions() reports the backend's native finish_reason on the non-stream envelope and the terminal chunk: _record_chat_finish captures it from the raw LiteLLM response at the model's fetch seam (openai-agents' ModelResponse discards it; this door has no transport client to hook, cf. _record_response_status). Live-verified: truncation reports "length", clean turns "stop". Degrades to the old "stop" literal if the private seam moves. - get_page_content()'s page-spec rejections surface as the documented PageIndexAPIError instead of the tool layer's private ValueError subclass; validation semantics unchanged. - The markdown CLI passes if_add_* verbatim again: filtering None out of user_opt let config.yaml's PDF-lane summary default switch on a per-node LLM pass the md CLI never ran โ€” unkeyed runs died with a raw traceback, keyed runs billed silently. - messages() resolves the default max_tokens to budget_tokens + 8192 when thinking is enabled: the wire requires max_tokens above the budget, so the flat 8192 default made every thinking call with a budget >= 8192 a hard 400. Explicit values pass through untouched. - The litellm/ routing prefix skips the env-only OPENAI_API_KEY pre-check (litellm resolves litellm.api_key and keyless OPENAI_BASE_URL servers itself), and both lanes' provider allowlists consult litellm.custom_provider_map, whose providers join provider_list only at completion time. The pre-checks stay: litellm wraps a typo'd provider as a 400 the retry loop treats as recoverable, so deleting them would burn ten retries per call. - agent_instructions(doc_id=...) passes scoped like the config sugar, so a newer same-name document elsewhere in the library no longer shadows an id the scoped tools reach by allowlist; duplicates within the targeted set still raise. - The LITELLM_LOCAL_MODEL_COST_MAP stamp moves from package import to _preload_litellm and the CLI head: importing pageindex no longer switches the host process's own litellm to the frozen bundled cost map. Library-direct callers without a client pay litellm's one-time fetch fallback instead โ€” accepted. - _await_completion polls through a transient refetch failure to the deadline instead of returning early, which the caller reported as the full three-minute timeout after five seconds. One superseded test removed (the agent_instructions shadow raise, replaced by the scoped-contract test), three stale pins updated to the documented error type. Suite: 335 passed. --- pageindex/__init__.py | 5 -- pageindex/agent_tools.py | 8 ++- pageindex/client.py | 38 ++++++++--- pageindex/local_chat.py | 78 +++++++++++++++++++-- pageindex/utils.py | 16 +++-- run_pageindex.py | 19 +++--- tests/test_agent_tools.py | 65 +++++++++++++++--- tests/test_client.py | 37 ++++++++-- tests/test_local_chat.py | 124 ++++++++++++++++++++++++++++++++++ tests/test_package_surface.py | 19 ++++++ tests/test_page_index_md.py | 28 ++++++++ 11 files changed, 385 insertions(+), 52 deletions(-) diff --git a/pageindex/__init__.py b/pageindex/__init__.py index e85278704..669f13629 100644 --- a/pageindex/__init__.py +++ b/pageindex/__init__.py @@ -1,11 +1,6 @@ """PageIndex SDK.""" -import os as _os from typing import TYPE_CHECKING as _TYPE_CHECKING -# LiteLLM's import otherwise fetches its model map over the network โ€” seconds -# of blocking (or a hang offline). setdefault, so an explicit user choice wins. -_os.environ.setdefault("LITELLM_LOCAL_MODEL_COST_MAP", "True") - from .client import PageIndexClient, PageIndexCloudClient, PageIndexLocalClient from .errors import PageIndexAPIError diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index 24af5fb53..e79e8dbfc 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -439,7 +439,7 @@ def _await_completion(client, entry: dict[str, Any], wait: bool) -> dict[str, An time.sleep(_TOOL_WAIT_INTERVAL) refreshed = _refetch_entry(client, doc_id) if refreshed is None: - return current + 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. @@ -1517,7 +1517,8 @@ def invoke(arguments: dict) -> tuple[str, bool]: for name in tool_names(include_management)] -def build_agent_tools(client, include_management: bool = False) -> list[Callable[..., str]]: +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 @@ -1526,10 +1527,11 @@ def build_agent_tools(client, include_management: bool = False) -> list[Callable 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)] + in _tool_specs(client, include_management, doc_ids)] # โ”€โ”€ agent instructions โ”€โ”€ diff --git a/pageindex/client.py b/pageindex/client.py index 9fdae9356..267ae8855 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -16,6 +16,11 @@ 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 @@ -31,8 +36,11 @@ def _import() -> None: def _parse_pages(pages: str) -> list[int]: - from .agent_tools import _expand_pages - return _expand_pages(pages) + from .agent_tools import _PageSpecError, _expand_pages + try: + return _expand_pages(pages) + except _PageSpecError as exc: + raise PageIndexAPIError(str(exc)) from exc def _agents_sdk_model_name(model: str) -> str: @@ -690,7 +698,9 @@ def messages( max_tokens: Per-turn output budget the Messages API requires on the wire; the default is resolved per model (8192, or 4096 for the claude-3 generation whose ceiling is lower) so the - simple call needs only a question. Passed through. + 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. @@ -783,7 +793,10 @@ def list_documents( # ---------- AGENT INTEGRATION ---------- - def agent_tools(self, include_management: bool = False) -> list[Callable[..., str]]: + 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()`` / @@ -805,9 +818,13 @@ def agent_tools(self, include_management: bool = False) -> list[Callable[..., st 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) + return build_agent_tools(self, include_management, doc_ids=doc_id) def as_openai_tools(self, include_management: bool = False, hosted: bool = False, @@ -1124,12 +1141,15 @@ def agent_instructions(self, doc_id: Optional[Union[str, list[str]]] = None) -> 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). + agent to work within them. Raises PageIndexAPIError if a doc_id + does not exist. Cloud also raises if its name is shadowed by a + newer same-name document (the name-addressed tools could not reach + it); local tools carry the doc_id scope, so only a duplicate name + within the targeted set shadows. """ from .agent_tools import build_agent_instructions - return build_agent_instructions(self, doc_id) + scope = self._local_doc_scope(doc_id) + return build_agent_instructions(self, doc_id, scoped=scope is not None) # ---------- FOLDER MANAGEMENT ---------- diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 2c5dd49f8..7ca2fbfb8 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -213,7 +213,10 @@ def _openai_model(protocol: str, model_name: str, backend=None): from .utils import _repair_litellm_types _repair_litellm_types() wire = model_name.removeprefix("litellm/") - if "/" not in wire or wire.startswith("openai/"): + # A litellm/ prefix is an explicit routing choice: litellm resolves + # credentials beyond the environment, so the key pre-check stands aside. + if not model_name.startswith("litellm/") and ( + "/" not in wire or wire.startswith("openai/")): if (not os.environ.get("OPENAI_API_KEY") and not (backend or {}).get("api_key")): raise PageIndexAPIError( @@ -226,10 +229,14 @@ def _openai_model(protocol: str, model_name: str, backend=None): if "/" not in wire: wire = f"openai/{wire}" providers = getattr(litellm, "provider_list", None) - if providers and wire.split("/", 1)[0] not in providers: + # 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 []} + provider = wire.split("/", 1)[0] + if providers and provider not in providers and provider not in custom: raise PageIndexAPIError( f"'{wire}' routes through LiteLLM, but " - f"'{wire.split('/', 1)[0]}' is not a LiteLLM provider. For an " + f"'{provider}' is not a LiteLLM provider. For an " "OpenAI-compatible server (vLLM, TGI, Ollama) serving this " f"model id, use 'openai/{wire}' and point OPENAI_BASE_URL " "at the server." @@ -325,6 +332,9 @@ def _openai_agent(client, protocol: str, model_name: str, instructions: str, 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), @@ -412,6 +422,53 @@ async def recording_create(*args, **kwargs): 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 @@ -510,6 +567,8 @@ def run_chat_completions(client, messages, stream: bool = False, 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 @@ -534,7 +593,7 @@ def run_chat_completions(client, messages, stream: bool = False, "index": 0, "message": {"role": "assistant", "content": result.final_output or ""}, - "finish_reason": "stop", + "finish_reason": recorded.get("finish_reason") or "stop", }], "usage": _openai_usage(result.raw_responses), } @@ -574,7 +633,7 @@ async def agen(): if not completed and hasattr(streamed, "cancel"): streamed.cancel() # abandoned/failed: stop the agent task await _aclose_backend(agent) - yield chunk({}, finish="stop") + yield chunk({}, finish=recorded.get("finish_reason") or "stop") yield { "id": chat_id, "object": "chat.completion.chunk", "created": created, "model": reported_model, "choices": [], @@ -905,9 +964,14 @@ def run_messages(client, messages, model: str, {"cache_control": {"type": "ephemeral"}} if _cache_marks(system_blocks, prepared) < 4 else {}) backend_client = _anthropic_client(_merged_backend(client, backend)) + if max_tokens is None: + budget = thinking.get("budget_tokens") if isinstance(thinking, dict) else None + # The wire requires max_tokens > thinking.budget_tokens; the flat + # default would 400 every thinking call with a budget >= 8192. + max_tokens = (budget + 8192 if isinstance(budget, int) + else _default_max_tokens(model)) runner = backend_client.beta.messages.tool_runner( - max_tokens=(max_tokens if max_tokens is not None - else _default_max_tokens(model)), + max_tokens=max_tokens, messages=prepared, model=model, tools=build_anthropic_tools(client, doc_ids=doc_id), diff --git a/pageindex/utils.py b/pageindex/utils.py index 842a3f268..dbeddb56d 100644 --- a/pageindex/utils.py +++ b/pageindex/utils.py @@ -82,10 +82,14 @@ def _openai_missing_keys(model): 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.""" - wire = _strip_prefix(model, "litellm/") - if "/" in wire and not wire.startswith("openai/"): + 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"]) @@ -98,13 +102,17 @@ def _litellm_model(model, backend): 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) - if providers and provider not in providers: + # 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 " @@ -112,7 +120,7 @@ def _litellm_model(model, backend): f"OPENAI_BASE_URL at the server.", llm_provider=None, model=model) if not backend: - missing = _openai_missing_keys(model) + missing = _openai_missing_keys(raw) if missing: raise litellm.AuthenticationError( f"missing API key for {model}: {', '.join(missing)}", diff --git a/run_pageindex.py b/run_pageindex.py index 91ddf0763..f2642b8a6 100644 --- a/run_pageindex.py +++ b/run_pageindex.py @@ -5,6 +5,10 @@ from pageindex.page_index_md import md_to_tree 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') @@ -165,25 +169,24 @@ 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({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/test_agent_tools.py b/tests/test_agent_tools.py index 1ac63dc63..637895d0e 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -1922,17 +1922,6 @@ def test_page_spec_span_bomb_rejected(client, store_path): assert "Too many pages" in payload["error"] -def test_agent_instructions_shadowed_doc_id_raises(client, store_path): - seed_doc(store_path, "pi-old", "report.pdf", - created_at="2026-08-01T10:00:00.000000") - seed_doc(store_path, "pi-new", "report.pdf", - created_at="2026-08-02T10:00:00.000000") - with pytest.raises(PageIndexAPIError, match="shadowed"): - client.agent_instructions(doc_id="pi-old") - text = client.agent_instructions(doc_id="pi-new") - assert "report.pdf" in text - - def test_wait_tolerates_transient_network_failures(fake_cloud_client, monkeypatch): import requests as requests_mod cloud = fake_cloud_client(["processing", "completed"]) @@ -2432,3 +2421,57 @@ def test_agent_instructions_carry_user_metadata(client, store_path): 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_ignores_out_of_scope_duplicates(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") + # The scoped tools reach pi-old by id, so a newer same-name document + # elsewhere in the library no longer shadows it... + text = client.agent_instructions(doc_id="pi-old") + assert "report.pdf" in text + # ...but a duplicate name inside the targeted set still does. + with pytest.raises(PageIndexAPIError, match="shadowed"): + client.agent_instructions(doc_id=["pi-old", "pi-new"]) + + +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 e74af3236..dd30ba27b 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -178,7 +178,7 @@ 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") @@ -186,7 +186,7 @@ def test_get_page_content_span_bomb_rejected(local_client, indexed_doc): """An absurd range must be rejected arithmetically, not expanded into a billion integers in the caller's process (the tool layer already refused; the public client method did not).""" - with pytest.raises(ValueError, match="spans more than 10000"): + 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") == [] @@ -920,10 +920,11 @@ def test_parse_pages_overlap_counts_union(): from pageindex.client import _parse_pages pages = _parse_pages("1-5000,2000-9000") assert len(pages) == 9000 and pages[0] == 1 and pages[-1] == 9000 - with pytest.raises(ValueError, match="spans more than"): + 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 - with pytest.raises(ValueError, match="positive"): + # 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") @@ -986,6 +987,31 @@ def test_index_precheck_covers_only_openai_shaped(monkeypatch): 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"}) @@ -1039,3 +1065,4 @@ def slow_flash(pdf, **kwargs): for worker in workers: worker.join() assert {r["name"] for r in results} == {"sample.pdf", "sample_1.pdf"} + diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index dc2c758eb..9f6ebb0b9 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -1834,3 +1834,127 @@ def handler(request): 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") diff --git a/tests/test_package_surface.py b/tests/test_package_surface.py index 6d985c818..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 @@ -101,3 +102,21 @@ def test_classic_compat_surface_still_reachable(): 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..f0a41ac83 100644 --- a/tests/test_page_index_md.py +++ b/tests/test_page_index_md.py @@ -21,3 +21,31 @@ def test_skips_bold_heading_with_only_whitespace(self): if __name__ == "__main__": unittest.main() + + +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()) From 49a24e1155e9ba2e3cf7febbe9ab597b104afcc8 Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 19 Aug 2026 21:35:59 +0800 Subject: [PATCH 131/137] =?UTF-8?q?fix:=20twelve=20max-review=20findings?= =?UTF-8?q?=20=E2=80=94=20shadow=20guard=20restored,=20caller=20transports?= =?UTF-8?q?=20survive,=20gated=20cloud=20endpoints,=20optimize=20precedenc?= =?UTF-8?q?e?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Correctness (silent behavior): - agent_instructions(doc_id=) returns to the strict shadow check: built alone it cannot know whether the caller's tools carry the same doc_id scope, so a newer same-name document blocks loudly again. The *_agent_config bundles keep the relaxed in-set check โ€” they build both sides. Undoes 914dc43's blanket scoped=True. - messages()/responses() no longer close a caller-owned http_client passed through chat_backend/backend ("passed verbatim"): messages() gates both closes on the merged dict, the responses lane marks the SDK client at build time and _aclose_backend leaves it open. - get_tree()/get_document_structure() keep key_items: the merge optimization (the local flash default) folds collapsed subsection titles into it, and the formatter silently dropped them while the agent tools still showed them. - _openai_agent classifies OpenAI-protocol destinations through litellm's own routing (get_llm_provider + openai_compatible_providers + azure/openrouter) instead of a name-prefix test: azure/openrouter/ deepseek/groq/xai destinations get prompt_cache_key again and their extra_body stays in the request body. - page_index_flash: an explicit optimize= now wins over the deprecated optimize_expand modifier (precedence was inverted); the modifier applies only to the legacy spellings (optimize absent or True), warns DeprecationWarning, and is back in the docstring. optimize=None is the unset sentinel and resolves to "full". Error contract: - messages() pre-checks Anthropic credentials (env, or any non-empty backend dict โ€” the same wide rule as the OpenAI lanes) instead of leaking the SDK's request-time bare TypeError on the most common misconfiguration. - anthropic_runner_config(thinking=) resolves the max_tokens default through the same thinking-aware helper messages() uses (budget+8192) and includes thinking in the returned kwargs, keeping the docstring's "the default messages() uses" claim true. - _bridge_invoker re-raises 401/403 instead of wrapping permanent auth failures in a "temporary, retry" envelope that burns the agent loop; bridge HTTP failures now carry status_code per the errors.py contract. - Cloud discovery and instructions ride the endpoint matching the tool gate (?tools=read unless include_management), threaded through the config bundles and agent_instructions. Live-verified: the gated endpoint serves the 7 read-only tools with annotations; instructions are byte-identical on both endpoints today, so this is protocol fidelity plus future-proofing, not a behavior change. - OpenAI Agents adapter: documented that is_error has no per-result non-aborting channel on hand-built FunctionTools at the 0.18.1 floor (raising aborts the run via UserError), so the envelope text is the whole signal. Comment only. Reuse: - _openai_model's chat branch delegates to utils._litellm_model โ€” the two hand-copies had already diverged on the keyless-backend rule; the wide rule (any non-empty backend dict stands aside) now covers both lanes, so one keyless OpenAI-compatible server works end to end. Tests: 12 new โ€” 11 fail on the pre-fix tree (stash A/B verified), plus a seam pin asserting LitellmModel._fetch_response exists so an upstream rename cannot silently report finish_reason "stop" for truncated turns. conftest now sets ANTHROPIC_API_KEY too: the suite no longer depends on the repo .env (verified green with .env removed, 343 passed + 3 live skipped). test_page_index_md's __main__ block moved to EOF so direct execution reaches both classes. Ruled not-fixed: one-shot doc_id iterables (off-contract input), a 5th caller cache_control breakpoint (loud 400, caller workaround), and malformed JSON-RPC defense (our own server). Attempted and withdrawn: stop_ids on the unscoped shadow sweep (trades an order-robust guard for a server-sort assumption), a bigger local listing page (LocalAPI mirrors cloud's 1-100 limit by design), and final_blocks reuse (the "duplicate" pass is load-bearing copy semantics: the appendable branch embeds those dicts into envelope["messages"] on every plain text answer, and callers currently get independent copies). --- pageindex/agent_tools.py | 50 ++++++----- pageindex/client.py | 44 ++++++---- pageindex/flash/api.py | 16 ++-- pageindex/integrations/openai_agents.py | 2 + pageindex/local_api.py | 2 + pageindex/local_chat.py | 105 ++++++++++++++---------- pageindex/mcp_bridge.py | 9 +- tests/conftest.py | 1 + tests/test_agent_tools.py | 69 ++++++++++++++-- tests/test_client.py | 13 +++ tests/test_flash_extraction.py | 26 ++++++ tests/test_local_chat.py | 80 ++++++++++++++++++ tests/test_page_index_md.py | 8 +- 13 files changed, 327 insertions(+), 98 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index e79e8dbfc..000ce0e6a 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -1346,8 +1346,8 @@ def _bridge_invoker(bridge, name: str, schema: dict, """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. Returns (envelope_text, is_error), - like call_tool.""" + 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() @@ -1355,6 +1355,9 @@ def _invoke(arguments: dict[str, Any]) -> tuple[str, bool]: _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", @@ -1411,10 +1414,8 @@ def inner(**kwargs: Any) -> str: inner.__annotations__ = annotations def proxy(*args: Any, **kwargs: Any) -> str: - # The invoker never raises, so a TypeError here is the binding - # rejecting the arguments (e.g. cloud-only parameters pruned - # from the local signature) โ€” answer with the same guided - # envelope call_tool returns for them. + # 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: @@ -1437,21 +1438,26 @@ def proxy(*args: Any, **kwargs: Any) -> str: _BRIDGES_LOCK = threading.Lock() -def _cloud_bridge(client): - """One bridge per client: tool discovery and instructions share a single - MCP session. Weak-keyed off the instance so clients stay picklable; the +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 bridge. + # A rotated api_key or moved BASE_URL rebuilds the bridges. auth = (client.BASE_URL, client.api_key) - bridge, seen = _BRIDGES.get(client) or (None, None) - if bridge is None or seen != auth: + 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", + f"{auth[0]}/mcp" + ("?tools=read" if gated else ""), {"Authorization": f"Bearer {auth[1]}"}, ) - _BRIDGES[client] = (bridge, auth) + bridges[gated] = bridge + _BRIDGES[client] = (bridges, auth) return bridge @@ -1495,7 +1501,7 @@ def _tool_specs(client, include_management: bool = False, doc_ids=None, 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) + bridge = _cloud_bridge(client, gated=not include_management) tools_meta = bridge.list_tools() if not include_management: tools_meta = _read_only_tools(tools_meta) @@ -1585,12 +1591,13 @@ def build_agent_tools(client, include_management: bool = False, ]) -def _base_instructions(client) -> str: - """Cloud: the live instructions the MCP server serves for this key's - tool set. Local: the built-in subset instructions.""" +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).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 " @@ -1673,9 +1680,10 @@ def doc_targeting_block(client, doc_id, scoped: bool = False) -> Optional[str]: ) -def build_agent_instructions(client, doc_id=None, scoped: bool = False) -> str: +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) + 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 267ae8855..c4d71eb15 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -931,8 +931,9 @@ def openai_agent_config( 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), + "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) @@ -1002,6 +1003,7 @@ def anthropic_runner_config( 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 @@ -1037,6 +1039,10 @@ def anthropic_runner_config( 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 @@ -1045,12 +1051,14 @@ def anthropic_runner_config( return { "model": model, "max_tokens": (max_tokens if max_tokens is not None - else _default_max_tokens(model)), - "system": build_agent_instructions(self, doc_id, - scoped=scope is not None), + 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"}, } @@ -1119,8 +1127,9 @@ def claude_agent_config( 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), + "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 @@ -1128,7 +1137,10 @@ def claude_agent_config( "allowed_tools": [f"mcp__{server_name}"], } - def agent_instructions(self, doc_id: Optional[Union[str, list[str]]] = None) -> str: + 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). @@ -1142,14 +1154,18 @@ def agent_instructions(self, doc_id: Optional[Union[str, list[str]]] = None) -> 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. Cloud also raises if its name is shadowed by a - newer same-name document (the name-addressed tools could not reach - it); local tools carry the doc_id scope, so only a duplicate name - within the targeted set shadows. + 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 - scope = self._local_doc_scope(doc_id) - return build_agent_instructions(self, doc_id, scoped=scope is not None) + return build_agent_instructions( + self, doc_id, include_management=include_management) # ---------- FOLDER MANAGEMENT ---------- diff --git a/pageindex/flash/api.py b/pageindex/flash/api.py index f6ba8d99c..72438325f 100644 --- a/pageindex/flash/api.py +++ b/pageindex/flash/api.py @@ -96,19 +96,23 @@ def _optimize(structure, page_texts, do_expand, model): def page_index_flash(pdf, summary=True, summary_model=None, - optimize: str | bool = "full", optimize_expand=None, + 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. 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. optimize_model: the LLM model for expand (defaults to the summary model). summary_concurrency: maximum simultaneous summary model calls; None uses the library default. use_embedded_toc: if True, consume the PDF's embedded bookmarks when trustworthy: deep bookmarks become the frame and the detected sections they lack are grafted back in after noise filtering, coarse ones become the chapter frame with detected nodes re-hung under them (deeper sparse entries are filled in when the page text confirms them, and garbled extracted titles are repaired from the bookmark strings), garbage ones are ignored; adds a ``toc_source`` key to the result. On by default; pass False for the pure detected structure. Returns: dict with keys ``doc_name``, ``doc_title``, ``structure`` (a list of nested ``{"title", "start_index", "end_index", "nodes"}`` dicts; page indexes are 1-based) and ``has_abstract_or_references_section`` (True when a top-level entry is an abstract or references heading). With ``optimize`` an ``optimize`` key reports merge/expand counts and before/after search-cost metrics. """ - if optimize is True: - optimize = "full" + """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_expand is not None and optimize: - optimize = "full" if optimize_expand else "merge" if optimize == "full": from ..errors import PageIndexAPIError from ..utils import ConfigLoader, _llm_backend, _openai_missing_keys diff --git a/pageindex/integrations/openai_agents.py b/pageindex/integrations/openai_agents.py index 4c8ca8ddc..862a35f0b 100644 --- a/pageindex/integrations/openai_agents.py +++ b/pageindex/integrations/openai_agents.py @@ -65,6 +65,8 @@ async def on_invoke_tool(ctx: Any, args_json: str) -> str: 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 diff --git a/pageindex/local_api.py b/pageindex/local_api.py index 9b29ac6e9..f3c6afa65 100644 --- a/pageindex/local_api.py +++ b/pageindex/local_api.py @@ -349,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 index 7ca2fbfb8..e4dcadef3 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -200,6 +200,8 @@ def _openai_model(protocol: str, model_name: str, backend=None): 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: @@ -210,37 +212,20 @@ def _openai_model(protocol: str, model_name: str, backend=None): f"'{model_name}' routes through LiteLLM, but litellm is not " "installed. Run: pip install 'litellm>=1.97'" ) - from .utils import _repair_litellm_types + from .utils import _litellm_model, _repair_litellm_types _repair_litellm_types() - wire = model_name.removeprefix("litellm/") - # A litellm/ prefix is an explicit routing choice: litellm resolves - # credentials beyond the environment, so the key pre-check stands aside. - if not model_name.startswith("litellm/") and ( - "/" not in wire or wire.startswith("openai/")): - if (not os.environ.get("OPENAI_API_KEY") - and not (backend or {}).get("api_key")): - 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/...')." - ) - if "/" not in wire: - wire = f"openai/{wire}" - 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 []} - provider = wire.split("/", 1)[0] - if providers and provider not in providers and provider not in custom: + try: + wire = _litellm_model(model_name, backend) + except litellm.AuthenticationError as exc: raise PageIndexAPIError( - f"'{wire}' routes through LiteLLM, but " - f"'{provider}' is not a LiteLLM provider. For an " - "OpenAI-compatible server (vLLM, TGI, Ollama) serving this " - f"model id, use 'openai/{wire}' and point OPENAI_BASE_URL " - "at the server." - ) + "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")) @@ -277,6 +262,25 @@ def _cache_extra_args(model_name: str) -> Optional[dict]: 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.""" @@ -296,8 +300,7 @@ def _openai_agent(client, protocol: str, model_name: str, instructions: str, # 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. - wire = model_name.removeprefix("litellm/") - openai_backend = "/" not in wire or wire.startswith("openai/") + 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"). @@ -472,8 +475,11 @@ async def recording_fetch(*args, **kwargs): 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.""" + 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: @@ -919,9 +925,15 @@ def _anthropic_usage(turns, final_usage: dict) -> dict: "claude-3-5-sonnet-20240620") -def _default_max_tokens(model: str) -> int: +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.""" + 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 @@ -963,13 +975,20 @@ def run_messages(client, messages, model: str, cached: dict[str, Any] = ( {"cache_control": {"type": "ephemeral"}} if _cache_marks(system_blocks, prepared) < 4 else {}) - backend_client = _anthropic_client(_merged_backend(client, backend)) + 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: - budget = thinking.get("budget_tokens") if isinstance(thinking, dict) else None - # The wire requires max_tokens > thinking.budget_tokens; the flat - # default would 400 every thinking call with a budget >= 8192. - max_tokens = (budget + 8192 if isinstance(budget, int) - else _default_max_tokens(model)) + max_tokens = _default_max_tokens(model, thinking) runner = backend_client.beta.messages.tool_runner( max_tokens=max_tokens, messages=prepared, @@ -994,7 +1013,8 @@ def events() -> Iterator[Any]: f"The model backend failed: {exc}") from exc finally: # runs on exhaustion and abandonment (GeneratorExit) alike - backend_client.close() + if owns_transport: + backend_client.close() return events() try: @@ -1004,7 +1024,8 @@ def events() -> Iterator[Any]: f"The model backend failed: {exc}") from exc finally: # safe here: the params read-back below does no HTTP - backend_client.close() + if owns_transport: + backend_client.close() if not turns: raise PageIndexAPIError("The model returned no response.") captured: dict = {} diff --git a/pageindex/mcp_bridge.py b/pageindex/mcp_bridge.py index 3ac9adbde..32ec09884 100644 --- a/pageindex/mcp_bridge.py +++ b/pageindex/mcp_bridge.py @@ -86,7 +86,8 @@ def _extract_result(self, response: requests.Response, request_id: int) -> Any: except ValueError as exc: raise PageIndexAPIError( f"MCP server returned a non-JSON response " - f"(HTTP {response.status_code})." + 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. @@ -130,7 +131,8 @@ def _request(self, method: str, params: Optional[dict] = None, if response.status_code >= 400: raise PageIndexAPIError( f"MCP request failed: HTTP {response.status_code} " - f"({response.text[:200]})" + f"({response.text[:200]})", + status_code=response.status_code, ) return self._extract_result(response, request_id) @@ -153,7 +155,8 @@ def _ensure_initialized(self) -> None: raise PageIndexAPIError( f"Could not connect to the PageIndex MCP server: HTTP " f"{response.status_code} ({response.text[:200]}). Check " - "your API key." + "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") 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/test_agent_tools.py b/tests/test_agent_tools.py index 637895d0e..28324d23d 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -966,6 +966,60 @@ def test_claude_agent_config_scoped_shadow_check(client, store_path): 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") @@ -1208,7 +1262,9 @@ def test_cloud_agent_tools_discover_live_tool_set(cloud_with_fake_bridge): cloud, created = cloud_with_fake_bridge tools = cloud.agent_tools() bridge = created["bridge"] - assert bridge.url == "https://api.pageindex.ai/mcp" + # 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. @@ -2444,18 +2500,15 @@ def get_document(self, doc_id): assert calls["n"] == 2 -def test_agent_instructions_doc_id_ignores_out_of_scope_duplicates(client, store_path): +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") - # The scoped tools reach pi-old by id, so a newer same-name document - # elsewhere in the library no longer shadows it... - text = client.agent_instructions(doc_id="pi-old") - assert "report.pdf" in text - # ...but a duplicate name inside the targeted set still does. + # 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", "pi-new"]) + client.agent_instructions(doc_id="pi-old") def test_agent_tools_doc_id_scopes_the_functions(client, store_path): diff --git a/tests/test_client.py b/tests/test_client.py index dd30ba27b..2ace6ba2e 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -1066,3 +1066,16 @@ def slow_flash(pdf, **kwargs): 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 index d82a1e4ff..1e9dcb87d 100644 --- a/tests/test_flash_extraction.py +++ b/tests/test_flash_extraction.py @@ -142,3 +142,29 @@ def test_unguarded_script_parses_parallel_without_reexecution(tmp_path): 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 index 9f6ebb0b9..a2a3d591e 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -1958,3 +1958,83 @@ def test_chat_gate_honors_litellm_routing_and_custom_providers(monkeypatch): 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_page_index_md.py b/tests/test_page_index_md.py index f0a41ac83..37b8c2a88 100644 --- a/tests/test_page_index_md.py +++ b/tests/test_page_index_md.py @@ -19,10 +19,6 @@ def test_skips_bold_heading_with_only_whitespace(self): ) -if __name__ == "__main__": - unittest.main() - - 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 @@ -49,3 +45,7 @@ def test_md_cli_runs_without_llm_or_key(self): out = Path(tmp) / "results" / "notes_structure.json" self.assertTrue(out.exists(), res.stdout.decode()) json.loads(out.read_text()) + + +if __name__ == "__main__": + unittest.main() From 0667e3b8a48b3364d97051a192e21556531b9a47 Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 19 Aug 2026 22:14:50 +0800 Subject: [PATCH 132/137] =?UTF-8?q?ci:=20dev=20tags=20publish=20to=20PyPI?= =?UTF-8?q?=20only=20=E2=80=94=20skip=20the=20GitHub=20Release?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Dev builds need an explicit ==pin to install, so their GitHub Releases carry no install value and double the feed next to the same-day stable (v0.2.10 + v0.2.10.dev6 showed as two near-identical cards). The dev6 Release was deleted by hand; this keeps future dev tags PyPI-only. --- .github/workflows/publish.yml | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) 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 }} From c0c715c56840cb7efa5605b4c944cefbb1a19c1a Mon Sep 17 00:00:00 2001 From: Ray Date: Thu, 20 Aug 2026 15:11:58 +0800 Subject: [PATCH 133/137] =?UTF-8?q?fix:=20six=20review=20findings=20?= =?UTF-8?q?=E2=80=94=20poison=20docs=20rejected,=20lone=20surrogates=20scr?= =?UTF-8?q?ubbed,=20honest=20poll=20errors?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Silent breakage: - submit_document restores the index-time page-bounds gate the PR lost with add_node_text: a tree whose span escapes 1..page_count (pdfium and PyPDF2 disagreeing about the page tree) now fails the submit instead of storing a completed document whose every read IndexErrors. - Lone UTF-16 surrogates can no longer reach stored text, where they survive in memory but kill every utf-8 JSON save after the LLM spend. All three producers are guarded: PDFium textpage output (unpaired halves the pair reassembly cannot compose), font-map walk targets (uniD83D glyph names, surrogate-band CIDs via the chr fallbacks โ€” scrubbed at targets_for, the map's single consumer, so walk patches and glyph synthesis cannot write one back), and PyPDF2 page texts (its _cmap decodes broken ToUnicode with surrogatepass). Each becomes U+FFFD in place, so census slots and the unicode walk stay synced. - _anonymous_main is depth-counted under a lock: overlapping windows from concurrent submits restore the true __main__ identity instead of a mid-window snapshot that left the host's __spec__ None forever. - page_index_flash skips the LLM expand pass when the extraction carries no page texts (bookmark-only documents): expand burned its retries on IndexErrors per node and expanded nothing; merge still runs. Error contracts: - messages() pre-checks credentials whenever the merged backend lacks api_key/auth_token/default_headers, not only when it is empty โ€” a backend={"timeout": 30} or client-level chat_backend no longer lets the SDK's request-time bare TypeError escape the PageIndexAPIError contract. - _wait_until_ready re-raises definite poll answers (401/403/404) immediately instead of retrying and advising to keep polling a document that will never turn up. --- pageindex/client.py | 2 + pageindex/flash/api.py | 6 +- .../parser_pdfium_charlevel/char_extract.py | 2 + .../parser_pdfium_charlevel/unicode_apply.py | 10 +- pageindex/flash/parser_pdfium_parallel.py | 26 ++++- pageindex/local_api.py | 26 ++++- pageindex/local_chat.py | 7 +- tests/test_agent_tools.py | 18 ++++ tests/test_client.py | 35 +++++++ tests/test_flash_extraction.py | 96 +++++++++++++++++++ tests/test_local_chat.py | 11 +++ 11 files changed, 227 insertions(+), 12 deletions(-) diff --git a/pageindex/client.py b/pageindex/client.py index c4d71eb15..73902ee13 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -253,6 +253,8 @@ def _wait_until_ready(self, doc_id: str, timeout: float = 1800.0) -> None: status = self.get_document(doc_id).get("status") poll_failures = 0 except (PageIndexAPIError, requests.RequestException) as exc: + if getattr(exc, "status_code", None) in (401, 403, 404): + raise # a definite answer, not a poll failure # Tolerate transient poll failures; a 30-minute wait should # not die on one 502 or dropped connection. poll_failures += 1 diff --git a/pageindex/flash/api.py b/pageindex/flash/api.py index 72438325f..2a2182ae2 100644 --- a/pageindex/flash/api.py +++ b/pageindex/flash/api.py @@ -126,8 +126,10 @@ def page_index_flash(pdf, summary=True, summary_model=None, 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 == "full", + # bookmark-only extractions carry no page_texts; expand needs them + pages = result.get("page_texts") or [] + result["optimize"] = _optimize(structure, pages, + optimize == "full" and bool(pages), optimize_model or summary_model) if summary and structure: import asyncio diff --git a/pageindex/flash/parser_pdfium_charlevel/char_extract.py b/pageindex/flash/parser_pdfium_charlevel/char_extract.py index 56d6288f7..e143fe23c 100644 --- a/pageindex/flash/parser_pdfium_charlevel/char_extract.py +++ b/pageindex/flash/parser_pdfium_charlevel/char_extract.py @@ -71,6 +71,8 @@ def _extract_raw_chars(page, text_page) -> tuple[list[dict], list[dict]]: if 0xDC00 <= low <= 0xDFFF: codepoint = ((codepoint & 0x3FF) << 10) + (low & 0x3FF) + 0x10000 skip_next = True + if 0xD800 <= codepoint <= 0xDFFF: + codepoint = 0xFFFD # unpaired surrogate: not utf-8 encodable # 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/unicode_apply.py b/pageindex/flash/parser_pdfium_charlevel/unicode_apply.py index 0ae314236..0053da52c 100644 --- a/pageindex/flash/parser_pdfium_charlevel/unicode_apply.py +++ b/pageindex/flash/parser_pdfium_charlevel/unicode_apply.py @@ -4,6 +4,7 @@ import bisect import difflib +import re from collections import Counter from .text_normalize import _is_whitespace @@ -13,6 +14,8 @@ _walk_codes, ) +_SURROGATES = re.compile("[\ud800-\udfff]") + def _apply_font_unicode( raw_chars: list[dict], @@ -37,9 +40,12 @@ def targets_for(font_xref: int | None, other_numbers: tuple[int, ...]) -> list[s if entry is None: return None next_block, measure_item = entry + # Broken font data (uniD83D glyph names, surrogate-band CIDs) yields + # lone-surrogate targets; patched into chars they crash utf-8 saves. if next_block == 1: - return [measure_item.get(code) or chr(code) for code in other_numbers] - return [measure_item.get((other_numbers[key_value] << 8) | other_numbers[key_value + 1]) or chr((other_numbers[key_value] << 8) | other_numbers[key_value + 1]) + return [_SURROGATES.sub("\ufffd", measure_item.get(code) or chr(code)) + for code in other_numbers] + return [_SURROGATES.sub("\ufffd", measure_item.get((other_numbers[key_value] << 8) | other_numbers[key_value + 1]) or chr((other_numbers[key_value] << 8) | other_numbers[key_value + 1])) for key_value in range(0, len(other_numbers) - 1, 2)] def apply(patches: list[tuple[int, str]], drops: list[int], diff --git a/pageindex/flash/parser_pdfium_parallel.py b/pageindex/flash/parser_pdfium_parallel.py index ec4c8dbc9..d70c1c270 100644 --- a/pageindex/flash/parser_pdfium_parallel.py +++ b/pageindex/flash/parser_pdfium_parallel.py @@ -24,6 +24,7 @@ import multiprocessing import os import sys +import threading from concurrent.futures import ProcessPoolExecutor from contextlib import contextmanager from io import BytesIO @@ -56,27 +57,42 @@ class _Type3Detected(Exception): _worker_font_maps: dict = {} +_window_lock = threading.Lock() +_window_depth = 0 +_window_saved: 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__. + everything by module name and never need __main__. Depth-counted so + overlapping windows restore the true originals, not a mid-window snapshot. ponytail: window covers the whole map; a concurrent pool spawned from another thread whose tasks live in __main__ would break during it.""" + global _window_depth, _window_saved 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 + with _window_lock: + _window_depth += 1 + if _window_depth == 1: + _window_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) + with _window_lock: + _window_depth -= 1 + if _window_depth == 0: + d.pop("__spec__", None) + d.update(_window_saved) + _window_saved = {} def _init_worker(kind: str, payload) -> None: diff --git a/pageindex/local_api.py b/pageindex/local_api.py index f3c6afa65..9603f386f 100644 --- a/pageindex/local_api.py +++ b/pageindex/local_api.py @@ -5,6 +5,7 @@ import logging import multiprocessing import os +import re import uuid from datetime import datetime, timezone from typing import Any @@ -15,6 +16,8 @@ logger = logging.getLogger(__name__) +_SURROGATES = re.compile("[\ud800-\udfff]") + def _now_iso() -> str: """Naive UTC, millisecond precision.""" @@ -124,6 +127,7 @@ def submit_document( raise except Exception as e: raise PageIndexAPIError(f"Failed to submit document: {e}") from e + self._check_page_bounds(structure, len(page_texts)) doc_id = "pi-" + uuid.uuid4().hex pages = [{"page_index": i + 1, "markdown": text} @@ -163,12 +167,32 @@ def _unique_doc_name(self, name: str) -> str: "Please use a different file name." ) + @staticmethod + def _check_page_bounds(structure: list, page_count: int) -> None: + """The tree (pdfium) and stored pages (PyPDF2) come from different + parsers; a span outside 1..page_count IndexErrors every later read.""" + stack = list(structure) + while stack: + node = stack.pop() + start, end = node.get("start_index"), node.get("end_index") + if (start is not None and end is not None + and not (1 <= start and end <= page_count)): + raise PageIndexAPIError( + f"Failed to submit document: the extracted structure " + f"references pages {start}-{end} outside the PDF's " + f"{page_count} readable pages." + ) + stack.extend(node.get("nodes") or []) + @staticmethod def _extract_page_texts(file_path: str) -> list[str]: import PyPDF2 with open(file_path, "rb") as f: reader = PyPDF2.PdfReader(f) - return [page.extract_text() or "" for page in reader.pages] + # PyPDF2 decodes broken ToUnicode maps with surrogatepass; lone + # surrogates would crash every utf-8 JSON save downstream. + return [_SURROGATES.sub("\ufffd", page.extract_text() or "") + for page in reader.pages] def _index_standard(self, file_path: str, page_texts: list[str]) -> tuple[list, str | None]: from .page_index_classic import page_index_main diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index e4dcadef3..38bb446ce 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -978,8 +978,11 @@ def run_messages(client, messages, model: str, 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")): + # default_headers counts: it can carry auth (or the SDK's Omit escape). + if (not any(key in (merged or {}) + for key in ("api_key", "auth_token", "default_headers")) + 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 " diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 28324d23d..ba593e884 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -2363,6 +2363,24 @@ def boom(doc_id): cloud.submit_document("whatever.pdf", wait=True) +def test_submit_wait_reraises_definite_poll_answers(fake_cloud_client, + monkeypatch): + """A 401/403/404 poll answer is final: re-raised untouched, no retries, no keep-polling advice.""" + cloud = fake_cloud_client(["processing"]) + polls = {"n": 0} + + def denied(doc_id): + polls["n"] += 1 + raise PageIndexAPIError("Failed to get document metadata: 401", + status_code=401) + + monkeypatch.setattr(cloud, "get_document", denied) + with pytest.raises(PageIndexAPIError, match="401") as err: + cloud.submit_document("whatever.pdf", wait=True) + assert polls["n"] == 1 + assert "Processing continues" not in str(err.value) + + 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.""" diff --git a/tests/test_client.py b/tests/test_client.py index 2ace6ba2e..c3bf83fa9 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -282,6 +282,41 @@ def test_submit_defaults_to_flash(local_client, sample_pdf, monkeypatch): assert local_client._api._store.get_meta(doc_id)["mode"] == "flash" +def test_submit_rejects_structure_beyond_stored_pages(local_client, sample_pdf, + monkeypatch): + """A tree spanning pages the store lacks fails submit instead of saving a doc whose reads IndexError.""" + monkeypatch.setattr( + pageindex.flash, "page_index_flash", + lambda pdf, **kwargs: { + "doc_name": "sample.pdf", + "structure": [{"title": "Root", "start_index": 1, + "end_index": 3, "summary": "s", "nodes": []}]}) + monkeypatch.setattr(pageindex.utils, "llm_completion", + lambda model, prompt, **kw: "d.") + with pytest.raises(PageIndexAPIError, match="pages 1-3 outside"): + local_client.submit_document(sample_pdf) + assert local_client._api._store.list_metas() == [] + + +def test_submit_survives_pypdf2_lone_surrogates(local_client, sample_pdf, + monkeypatch): + """PyPDF2 decodes broken ToUnicode with surrogatepass; the store gets U+FFFD, not a utf-8-fatal str.""" + import PyPDF2 + monkeypatch.setattr(PyPDF2.PageObject, "extract_text", + lambda self: "\ud83dello broken") + monkeypatch.setattr( + pageindex.flash, "page_index_flash", + lambda p, **kwargs: { + "structure": [{"title": "T", "start_index": 1, + "end_index": 1, "summary": "s", "nodes": []}]}) + monkeypatch.setattr(pageindex.utils, "llm_completion", + lambda model, prompt, **kw: "d.") + doc_id = local_client.submit_document(sample_pdf)["doc_id"] + markdown = local_client.get_ocr(doc_id)["result"][0]["markdown"] + assert "\ud83d" not in markdown + assert markdown.startswith("๏ฟฝello") + + def test_page_index_flash_rejects_unknown_optimize(): from pageindex.flash import page_index_flash with pytest.raises(ValueError, match="optimize must be"): diff --git a/tests/test_flash_extraction.py b/tests/test_flash_extraction.py index 1e9dcb87d..9ec0c6439 100644 --- a/tests/test_flash_extraction.py +++ b/tests/test_flash_extraction.py @@ -168,3 +168,99 @@ def test_optimize_wins_over_deprecated_optimize_expand(tmp_path, monkeypatch): with pytest.raises(PageIndexAPIError, match="optimize='merge'"): page_index_flash(str(tmp_path / "missing.pdf"), summary=False, optimize=None) + + +def test_lone_surrogate_from_broken_tounicode_is_replaced(monkeypatch): + """An unpaired UTF-16 surrogate leaves as U+FFFD, not a str that crashes utf-8 save.""" + import json + from io import BytesIO + + import pypdfium2 as pdfium + import pypdfium2.raw as pdfium_c + from conftest import build_pdf + from pageindex.flash.parser_pdfium_charlevel.char_extract import ( + _extract_raw_chars) + + orig = pdfium_c.FPDFText_GetUnicode + monkeypatch.setattr(pdfium_c, "FPDFText_GetUnicode", + lambda tp, i: 0xD83D if i == 0 else orig(tp, i)) + pdf = pdfium.PdfDocument(BytesIO(build_pdf(["Hello broken cmap"]))) + page = pdf[0] + raw_chars, _objects = _extract_raw_chars(page, page.get_textpage().raw) + text = "".join(char["ch"] for char in raw_chars) + assert "\ud83d" not in text + assert text.startswith("๏ฟฝello") + json.dumps(text) # the save-time crash this guards against + + +def test_lone_surrogate_targets_never_patched_into_chars(): + """A surrogate-band code with no cmap entry (chr fallback) must not patch a lone surrogate back in.""" + from pageindex.flash.parser_pdfium_charlevel.unicode_apply import ( + _apply_font_unicode) + + char = {"i": 0, "ch": "X", "is_gen": False} + show_codes = [(7, (0xD8, 0x3D), 100.0)] + map_cache = {7: (2, {})} # Identity map, no ToUnicode: target = chr(0xD83D) + + _apply_font_unicode([char], [], show_codes, None, map_cache) + + assert char["ch"] == "๏ฟฝ" + + +def test_anonymous_main_overlapping_windows_restore(monkeypatch): + """The last window out must restore the true originals, not a mid-window snapshot.""" + import sys + import threading + + from pageindex.flash.parser_pdfium_parallel import _anonymous_main + + main = sys.modules["__main__"] + spec = object() + monkeypatch.setattr(main, "__file__", "sentinel-file", raising=False) + monkeypatch.setattr(main, "__spec__", spec, raising=False) + a_in, b_in, a_out = (threading.Event() for _ in range(3)) + + def first(): + with _anonymous_main(): + a_in.set() + assert b_in.wait(5) + a_out.set() + + def second(): + assert a_in.wait(5) + with _anonymous_main(): + b_in.set() + assert a_out.wait(5) + + threads = [threading.Thread(target=first), threading.Thread(target=second)] + for thread in threads: + thread.start() + for thread in threads: + thread.join(10) + assert main.__spec__ is spec + assert main.__file__ == "sentinel-file" + + +def test_optimize_full_skips_expand_without_page_texts(tmp_path, monkeypatch): + """A bookmark-only extraction (no page_texts) skips expand; merge still runs.""" + from conftest import build_pdf + from pageindex.flash import api as flash_api + + monkeypatch.setenv("OPENAI_API_KEY", "k") + calls = {} + + def fake_optimize(structure, pages, do_expand, model): + calls["pages"] = pages + calls["do_expand"] = do_expand + return {"merges": 0} + + monkeypatch.setattr(flash_api, "_optimize", fake_optimize) + monkeypatch.setattr(flash_api, "extract_toc", + lambda pdf, use_embedded_toc=True: { + "structure": [{"title": "T", "start_index": 1, + "end_index": 1, "nodes": []}]}) + pdf = tmp_path / "doc.pdf" + pdf.write_bytes(build_pdf(["x"])) + result = flash_api.page_index_flash(str(pdf), summary=False) + assert calls == {"pages": [], "do_expand": False} + assert result["optimize"] == {"merges": 0} diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index a2a3d591e..56973e8ff 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -2038,3 +2038,14 @@ def test_messages_without_credentials_raises_contract_error(client, with pytest.raises(PageIndexAPIError, match="Anthropic backend is not configured"): client.messages("q", model="claude-test") + + +@needs_anthropic +def test_messages_keyless_backend_still_raises_contract_error(client, + monkeypatch): + """A backend dict without credentials must not disarm the pre-check.""" + 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", backend={"timeout": 30}) From 3504664d033c10a47928f8b4360488725265470d Mon Sep 17 00:00:00 2001 From: Ray Date: Thu, 20 Aug 2026 15:23:09 +0800 Subject: [PATCH 134/137] fix: drop the inverted cache-seeding kwarg, true up five doc contracts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _no_cache_seeding_kwargs rested on a false premise โ€” litellm 1.97 seeds no cache marks unprompted (its hook fires only on explicit injection points) โ€” and with a system message present it caused the paid cache write it claimed to prevent. Deleted; backend keys still merge through, and the rewritten test pins the honest shape: no cache params sent, backend keys win the merge. Docs that contradicted shipped behavior: the three "tools never raise" sites now admit the deliberate 401/403 re-raise (49a24e1 updated only the private ones); chat_completions' finish_reason line caught up with 914dc43's native surfacing; errors.py stops implying status_code=None means local; char_extract's IsGenerated comment no longer points at a deleted second read site; the parallel parser's parity contract admits the spawn-child re-raise. --- pageindex/agent_tools.py | 7 +++++-- pageindex/client.py | 8 +++++--- pageindex/errors.py | 4 ++-- .../parser_pdfium_charlevel/char_extract.py | 6 +++--- pageindex/flash/parser_pdfium_parallel.py | 4 +++- pageindex/local_chat.py | 7 ++++--- pageindex/utils.py | 16 ++-------------- tests/test_client.py | 15 +++++++-------- 8 files changed, 31 insertions(+), 36 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index 000ce0e6a..093d41316 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -12,7 +12,9 @@ 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. +that envelope too, on the direct and the call_tool path alike. One +exception: a cloud 401/403 re-raises PageIndexAPIError โ€” a dead key is +for the caller to fix, not for the model to retry. """ from __future__ import annotations @@ -1531,7 +1533,8 @@ def build_agent_tools(client, include_management: bool = False, 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 + signature accepts โ€” except a cloud 401/403, which re-raises + PageIndexAPIError (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``. """ diff --git a/pageindex/client.py b/pageindex/client.py index 73902ee13..e3feb88ce 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -479,8 +479,9 @@ def chat_completions( 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 + tool calls. ``finish_reason`` carries the final turn's native + finish reason โ€” "stop", or the backend's "length" / + "content_filter" when the last turn was cut short. For the tool-use process and prompt-cache round-trip use ``responses()`` or ``messages()``. @@ -813,7 +814,8 @@ def agent_tools( ``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. + string, and reports failures inside that JSON instead of raising โ€” + except a cloud 401/403, which raises PageIndexAPIError. Args: include_management (bool): Also expose tools that modify the diff --git a/pageindex/errors.py b/pageindex/errors.py index 608ba6e4f..fc6b09eb9 100644 --- a/pageindex/errors.py +++ b/pageindex/errors.py @@ -1,6 +1,6 @@ class PageIndexAPIError(Exception): - """status_code carries the HTTP status when the failure came from a - non-200 cloud response; None otherwise (local mode, client-side).""" + """status_code carries the HTTP status when the raising site passes it + (some cloud paths raise bare), so None does not imply local/client-side.""" def __init__(self, *args: object, status_code: int | None = None) -> None: super().__init__(*args) diff --git a/pageindex/flash/parser_pdfium_charlevel/char_extract.py b/pageindex/flash/parser_pdfium_charlevel/char_extract.py index e143fe23c..2bef05fe2 100644 --- a/pageindex/flash/parser_pdfium_charlevel/char_extract.py +++ b/pageindex/flash/parser_pdfium_charlevel/char_extract.py @@ -81,9 +81,9 @@ def _extract_raw_chars(page, text_page) -> tuple[list[dict], list[dict]]: ch_str = chr(codepoint) is_ws = js_is_ws(codepoint) # FPDFText_IsGenerated returns a c_int: 1 generated, 0 real, -1 error. - # Only a POSITIVE 1 may mark a char generated -- the -1 has to read the - # same way here as it does in the page-mode unicode walk, or the two - # char sets disagree and that walk desyncs. + # Only a POSITIVE 1 may mark a char generated. This is the package's + # only read: the page-mode unicode walk consumes this flag rather than + # re-reading PDFium (a second read is how astral chars desynced it). is_gen = is_generated(text_page, index_value) == 1 # PDFium inserts is_generated chars as layout placeholders for # Td/Tm jumps with no literal content-stream char (typically diff --git a/pageindex/flash/parser_pdfium_parallel.py b/pageindex/flash/parser_pdfium_parallel.py index d70c1c270..2c45da9e9 100644 --- a/pageindex/flash/parser_pdfium_parallel.py +++ b/pageindex/flash/parser_pdfium_parallel.py @@ -12,7 +12,9 @@ poison the run the moment any page accumulates an extent; the driver then discards the parallel attempt and reruns the document on the sequential path, which is the source of truth. Any other worker failure falls back the -same way, so this entry can only ever return sequential-identical output. +same way, so this entry returns sequential-identical output โ€” except in a +spawn child re-importing an unguarded __main__, where it re-raises instead +of silently duplicating the caller's whole run per worker. Worker startup pays the full package import chain plus its own document open; ``min_pages`` routes documents too small to amortize that to the diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 38bb446ce..523879ba8 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -253,9 +253,10 @@ def _cache_extra_args(model_name: str) -> Optional[dict]: 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. + # The stable prefix plus the newest message, so each turn re-reads + # the turns before it. LiteLLM seeds nothing on its own (its hook + # fires only when injection points are passed), so this pair is the + # sole source of the marks on all three channels. return {"cache_control_injection_points": [ {"location": "message", "role": "system"}, {"location": "message", "index": -1}]} diff --git a/pageindex/utils.py b/pageindex/utils.py index dbeddb56d..8075b0a88 100644 --- a/pageindex/utils.py +++ b/pageindex/utils.py @@ -139,18 +139,6 @@ 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): import litellm max_retries = 10 @@ -165,7 +153,7 @@ def llm_completion(model, prompt, chat_history=None, return_finish_reason=False) 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)}, + **{"max_retries": 0, **(backend or {})}, ) content = response.choices[0].message.content if return_finish_reason: @@ -198,7 +186,7 @@ async def llm_acompletion(model, prompt): model=model, messages=messages, drop_params=True, - **{"max_retries": 0, **_no_cache_seeding_kwargs(backend)}, + **{"max_retries": 0, **(backend or {})}, ) return response.choices[0].message.content except Exception as e: diff --git a/tests/test_client.py b/tests/test_client.py index c3bf83fa9..72b88e729 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -747,10 +747,9 @@ async def deny_alpha(model, prompt): 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.""" +def test_llm_completion_backend_reaches_litellm(monkeypatch): + """Indexing sends no cache params of its own (litellm seeds nothing + unprompted); backend keys pass through and win the merge.""" import litellm captured = {} @@ -763,15 +762,15 @@ def fake_completion(**kwargs): 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"}] + assert "cache_control_injection_points" not in captured token = pageindex.utils._llm_backend.set( - {"api_key": "x", "cache_control_injection_points": []}) + {"api_key": "x", "max_retries": 3}) try: pageindex.utils.llm_completion("gpt-4o", "probe") finally: pageindex.utils._llm_backend.reset(token) - assert captured["cache_control_injection_points"] == [] + assert captured["api_key"] == "x" + assert captured["max_retries"] == 3 def test_delete_survives_marker_tamper(local_client, tmp_path): From 5c48ad81aa292015e0372f7290cb5f5edd042a4b Mon Sep 17 00:00:00 2001 From: Ray Date: Thu, 20 Aug 2026 15:25:18 +0800 Subject: [PATCH 135/137] chore: trim the new comments to essentials --- pageindex/agent_tools.py | 3 +-- pageindex/errors.py | 4 ++-- pageindex/flash/parser_pdfium_charlevel/char_extract.py | 4 ++-- pageindex/flash/parser_pdfium_parallel.py | 3 +-- pageindex/local_chat.py | 5 ++--- tests/test_client.py | 3 +-- 6 files changed, 9 insertions(+), 13 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index 093d41316..30fcb3dbe 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -13,8 +13,7 @@ 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. One -exception: a cloud 401/403 re-raises PageIndexAPIError โ€” a dead key is -for the caller to fix, not for the model to retry. +exception: a cloud 401/403 re-raises PageIndexAPIError. """ from __future__ import annotations diff --git a/pageindex/errors.py b/pageindex/errors.py index fc6b09eb9..b9ccf7a05 100644 --- a/pageindex/errors.py +++ b/pageindex/errors.py @@ -1,6 +1,6 @@ class PageIndexAPIError(Exception): - """status_code carries the HTTP status when the raising site passes it - (some cloud paths raise bare), so None does not imply local/client-side.""" + """status_code carries the HTTP status when the raising site passes it; + None does not imply local/client-side.""" def __init__(self, *args: object, status_code: int | None = None) -> None: super().__init__(*args) diff --git a/pageindex/flash/parser_pdfium_charlevel/char_extract.py b/pageindex/flash/parser_pdfium_charlevel/char_extract.py index 2bef05fe2..ffdbc7266 100644 --- a/pageindex/flash/parser_pdfium_charlevel/char_extract.py +++ b/pageindex/flash/parser_pdfium_charlevel/char_extract.py @@ -82,8 +82,8 @@ def _extract_raw_chars(page, text_page) -> tuple[list[dict], list[dict]]: is_ws = js_is_ws(codepoint) # FPDFText_IsGenerated returns a c_int: 1 generated, 0 real, -1 error. # Only a POSITIVE 1 may mark a char generated. This is the package's - # only read: the page-mode unicode walk consumes this flag rather than - # re-reading PDFium (a second read is how astral chars desynced it). + # only read: the page-mode unicode walk consumes this flag rather + # than re-reading PDFium. is_gen = is_generated(text_page, index_value) == 1 # PDFium inserts is_generated chars as layout placeholders for # Td/Tm jumps with no literal content-stream char (typically diff --git a/pageindex/flash/parser_pdfium_parallel.py b/pageindex/flash/parser_pdfium_parallel.py index 2c45da9e9..0ee715c4a 100644 --- a/pageindex/flash/parser_pdfium_parallel.py +++ b/pageindex/flash/parser_pdfium_parallel.py @@ -13,8 +13,7 @@ discards the parallel attempt and reruns the document on the sequential path, which is the source of truth. Any other worker failure falls back the same way, so this entry returns sequential-identical output โ€” except in a -spawn child re-importing an unguarded __main__, where it re-raises instead -of silently duplicating the caller's whole run per worker. +spawn child re-importing an unguarded __main__, where it re-raises. Worker startup pays the full package import chain plus its own document open; ``min_pages`` routes documents too small to amortize that to the diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 523879ba8..f77a8699d 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -254,9 +254,8 @@ def _cache_extra_args(model_name: str) -> Optional[dict]: if provider == "anthropic" or (provider in ("bedrock", "vertex_ai") and "claude" in model.lower()): # The stable prefix plus the newest message, so each turn re-reads - # the turns before it. LiteLLM seeds nothing on its own (its hook - # fires only when injection points are passed), so this pair is the - # sole source of the marks on all three channels. + # the turns before it. LiteLLM seeds nothing unprompted, so this + # pair is the marks' sole source. return {"cache_control_injection_points": [ {"location": "message", "role": "system"}, {"location": "message", "index": -1}]} diff --git a/tests/test_client.py b/tests/test_client.py index 72b88e729..1cf544793 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -748,8 +748,7 @@ async def deny_alpha(model, prompt): def test_llm_completion_backend_reaches_litellm(monkeypatch): - """Indexing sends no cache params of its own (litellm seeds nothing - unprompted); backend keys pass through and win the merge.""" + """No cache params of our own; backend keys reach litellm and win the merge.""" import litellm captured = {} From f4fc566805ad3940a398c58cfe8e69a04510d87d Mon Sep 17 00:00:00 2001 From: Ray Date: Thu, 20 Aug 2026 16:35:11 +0800 Subject: [PATCH 136/137] fix: openai_agent_config bundles the Claude cache marks its doc promised MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The docstring claimed LiteLLM seeds cache_control for Anthropic- and Bedrock-hosted Claude on its own, with Vertex the lone manual exception. litellm 1.97 seeds nothing unprompted โ€” its hook fires only on explicit cache_control_injection_points, and the auto-seed path is gated behind LITELLM_ENABLE_ANTHROPIC_PROMPT_CACHING, default off โ€” so every LiteLLM-routed Claude agent built from the bundle ran uncached, paying full prompt price per turn, while chat_completions() marked the very same model through _cache_extra_args. 3504664 trued up the sibling doc sites and deleted the false-premise kwarg but missed this paragraph. Attach the same helper's marks as model_settings when non-None: bare/gpt names and cloud-without-model are unchanged (key absent), unknown providers and litellm failures degrade to None. Verified on all three channels through the real Agent(**config) -> Runner -> LitellmModel path: anthropic direct real-key E2E (cache_creation 2131 -> cache_read 2131, 3 uncached tokens), bedrock converse cachePoint and vertex rawPredict cache_control both captured on the outbound body. Per-run RunConfig overrides keep the marks โ€” ModelSettings.resolve() dict-merges extra_args โ€” as the rewritten docstring now states. Claude-Session: https://claude.ai/code/session_01YXqF4bsMfWwt1G3nTXJX6b --- pageindex/client.py | 19 ++++++++++--------- tests/test_agent_tools.py | 18 ++++++++++++++++++ 2 files changed, 28 insertions(+), 9 deletions(-) diff --git a/pageindex/client.py b/pageindex/client.py index e3feb88ce..301523e0c 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -910,15 +910,11 @@ def openai_agent_config( 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}]})) + Prompt caching: OpenAI models cache server-side on their own; + LiteLLM-routed Claude (Anthropic, Bedrock, Vertex) gets its + cache marks from the bundled ``model_settings``. Replace that + key wholesale and the marks go with it โ€” per-run overrides via + ``RunConfig(model_settings=...)`` merge instead. Args: doc_id: Document ID or list of IDs to target, as in @@ -948,6 +944,11 @@ def openai_agent_config( # caller's process, outside our completion helpers. from .utils import _repair_litellm_types _repair_litellm_types() + from .local_chat import _cache_extra_args + extra_args = _cache_extra_args(model) + if extra_args: + from agents import ModelSettings + config["model_settings"] = ModelSettings(extra_args=extra_args) return config def as_anthropic_tools(self, include_management: bool = False, diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index ba593e884..98ffc7411 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -777,6 +777,24 @@ def test_openai_agent_config_model_speaks_the_agents_sdk_grammar(tmp_path): == "litellm/groq/llama-x") +def test_openai_agent_config_carries_cache_marks_for_litellm_claude(tmp_path): + """LiteLLM-routed Claude gets the same cache marks the engine + attaches in chat_completions(); OpenAI-bound models stay unmarked + (their caching is server-side, and LiteLLM seeds nothing on its + own).""" + pytest.importorskip("agents") + from pageindex.local_chat import _cache_extra_args + client = PageIndexLocalClient(storage_path=str(tmp_path / "s"), + chat_model="anthropic/claude-x") + settings = client.openai_agent_config()["model_settings"] + assert settings.extra_args == _cache_extra_args("anthropic/claude-x") + assert "cache_control_injection_points" in settings.extra_args + assert "model_settings" not in client.openai_agent_config(model="gpt-x") + # The per-call override is marked by its own routing, not the default's. + marked = client.openai_agent_config(model="bedrock/claude-y") + assert "cache_control_injection_points" in marked["model_settings"].extra_args + + def test_plain_functions_answer_bad_arguments_with_the_envelope(client, store_path): """agent_tools() functions must not raise into a framework loop: From a15e9159356bc019d67a68226e902976196859c6 Mon Sep 17 00:00:00 2001 From: Ray Date: Thu, 20 Aug 2026 18:06:32 +0800 Subject: [PATCH 137/137] feat: model_settings and name become openai_agent_config parameters MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit model_settings merges on top of the bundled cache marks via the SDK's own ModelSettings.resolve() โ€” caller fields win, extra_args dict-merge โ€” so customizing one knob no longer silently drops the marks the way wholesale replacement of the returned key does; with no marks in play the caller's object rides through verbatim (cloud included). The alongside-the-splat spelling Agent(**cfg, model_settings=...) cannot be rescued at the class level โ€” Python raises the duplicate-keyword TypeError while assembling the call, before any __init__ runs โ€” so the merge lives where the two settings meet as data, mirroring the SDK's own RunConfig merge point. name joins as a pure-forward parameter by explicit ruling: it is cosmetic on this bundle (unlike claude_agent_config's server_name, which must stay in sync with the mcp__ allowed_tools pre-approval string), but in composition it seeds the SDK-derived handoff and as_tool names โ€” live-checked: handoff(agent).tool_name follows it. The demo's commented model_settings template moves inside the config call, so uncommenting it is safe under any model. Claude-Session: https://claude.ai/code/session_01YXqF4bsMfWwt1G3nTXJX6b --- examples/agentic_vectorless_rag_demo.py | 6 ++++-- pageindex/client.py | 20 ++++++++++++++++---- tests/test_agent_tools.py | 18 ++++++++++++++++++ 3 files changed, 38 insertions(+), 6 deletions(-) diff --git a/examples/agentic_vectorless_rag_demo.py b/examples/agentic_vectorless_rag_demo.py index 93a00735c..fef4f794e 100644 --- a/examples/agentic_vectorless_rag_demo.py +++ b/examples/agentic_vectorless_rag_demo.py @@ -49,8 +49,10 @@ def query_agent(client: PageIndexLocalClient, doc_id: str, prompt: str, verbose: Tool calls are always printed; verbose=True also prints arguments and output previews. """ agent = Agent( - **client.openai_agent_config(doc_id=doc_id), - # model_settings=ModelSettings(reasoning={"effort": "low", "summary": "auto"}), # from agents.model_settings import ModelSettings + **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(): diff --git a/pageindex/client.py b/pageindex/client.py index 301523e0c..d292c5f9b 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -894,6 +894,8 @@ def openai_agent_config( doc_id: Optional[Union[str, list[str]]] = None, include_management: bool = False, model: Optional[str] = None, + model_settings: Optional[Any] = None, + name: str = "PageIndex", ) -> dict[str, Any]: """ Document QA ``Agent`` kwargs for the OpenAI Agents SDK in one @@ -912,9 +914,10 @@ def openai_agent_config( Prompt caching: OpenAI models cache server-side on their own; LiteLLM-routed Claude (Anthropic, Bedrock, Vertex) gets its - cache marks from the bundled ``model_settings``. Replace that - key wholesale and the marks go with it โ€” per-run overrides via - ``RunConfig(model_settings=...)`` merge instead. + cache marks from the bundled ``model_settings``. Pass + ``model_settings`` here to layer your own on top โ€” your fields + win and ``extra_args`` merge. Replacing the returned key + wholesale drops the marks instead. Args: doc_id: Document ID or list of IDs to target, as in @@ -926,11 +929,16 @@ def openai_agent_config( model: Backend model name; overrides the local default. Same grammar as ``chat_model`` (LiteLLM names; bare names are OpenAI-compatible shorthand). + model_settings: Your own ``ModelSettings``, merged on top of + the bundled cache marks; included verbatim when no marks + apply. + name (str): Agent display name; in composition it also seeds + the SDK-derived handoff and ``as_tool`` names. """ from .agent_tools import build_agent_instructions scope = self._local_doc_scope(doc_id) config: dict[str, Any] = { - "name": "PageIndex", + "name": name, "instructions": build_agent_instructions( self, doc_id, scoped=scope is not None, include_management=include_management), @@ -949,6 +957,10 @@ def openai_agent_config( if extra_args: from agents import ModelSettings config["model_settings"] = ModelSettings(extra_args=extra_args) + if model_settings is not None: + marks = config.get("model_settings") + config["model_settings"] = (marks.resolve(model_settings) + if marks else model_settings) return config def as_anthropic_tools(self, include_management: bool = False, diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 98ffc7411..c6c008f70 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -759,6 +759,7 @@ def test_openai_agent_config_local(client, store_path): 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" + assert client.openai_agent_config(name="Researcher")["name"] == "Researcher" def test_openai_agent_config_model_speaks_the_agents_sdk_grammar(tmp_path): @@ -795,6 +796,23 @@ def test_openai_agent_config_carries_cache_marks_for_litellm_claude(tmp_path): assert "cache_control_injection_points" in marked["model_settings"].extra_args +def test_openai_agent_config_merges_caller_model_settings(tmp_path): + """Caller model_settings merge on top of the bundled cache marks + (caller fields win, extra_args dict-merge); with no marks the + caller's object rides through verbatim.""" + pytest.importorskip("agents") + from agents import ModelSettings + client = PageIndexLocalClient(storage_path=str(tmp_path / "s"), + chat_model="anthropic/claude-x") + mine = ModelSettings(temperature=0.2, extra_args={"top_k": 5}) + merged = client.openai_agent_config(model_settings=mine)["model_settings"] + assert merged.temperature == 0.2 + assert merged.extra_args["top_k"] == 5 + assert "cache_control_injection_points" in merged.extra_args + verbatim = client.openai_agent_config(model="gpt-x", model_settings=mine) + assert verbatim["model_settings"] is mine + + def test_plain_functions_answer_bad_arguments_with_the_envelope(client, store_path): """agent_tools() functions must not raise into a framework loop: