From 6b343f55eeb075082f09a7547843922e6de364e0 Mon Sep 17 00:00:00 2001 From: Ray Date: Sun, 9 Aug 2026 19:21:39 +0800 Subject: [PATCH 01/65] =?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 02/65] =?UTF-8?q?fix:=20agent=20tools=20review=20=E2=80=94?= =?UTF-8?q?=20next=5Fsteps=20order,=20resolve=20caching,=20error=20semanti?= =?UTF-8?q?cs?= 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 03/65] =?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 04/65] =?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 05/65] 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 06/65] 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 07/65] =?UTF-8?q?revert:=20keep=20README.md=20unchanged=20?= =?UTF-8?q?from=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 08/65] 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 09/65] 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 10/65] 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 11/65] 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 12/65] 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 13/65] 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 14/65] 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 15/65] =?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 16/65] 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 17/65] 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 18/65] =?UTF-8?q?feat:=20local=20chat=20=E2=80=94=20three?= =?UTF-8?q?=20protocol=20surfaces=20over=20the=20agent=20tools=20(v0.2.10)?= 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 19/65] =?UTF-8?q?fix:=20local-chat=20review=20findings=20?= =?UTF-8?q?=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 20/65] =?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 21/65] =?UTF-8?q?fix:=20as=5Fanthropic=5Ftools=20review=20?= =?UTF-8?q?findings=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 22/65] =?UTF-8?q?docs:=20doc=5Fid=20is=20per-call=20table-?= =?UTF-8?q?setting=20=E2=80=94=20keep=20it=20identical=20across=20a=20conv?= =?UTF-8?q?ersation?= 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 23/65] 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 24/65] 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 25/65] 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 26/65] 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 27/65] 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 28/65] 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 29/65] =?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 30/65] 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 31/65] 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 32/65] 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 33/65] 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 34/65] 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 35/65] 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 36/65] 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 37/65] 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 38/65] 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 39/65] 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 40/65] 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 41/65] =?UTF-8?q?fix:=20two=20chat=20findings=20=E2=80=94?= =?UTF-8?q?=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 42/65] =?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 43/65] =?UTF-8?q?fix:=20five=20secondary=20review=20findin?= =?UTF-8?q?gs=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 44/65] =?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 45/65] 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 46/65] 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 47/65] =?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 48/65] =?UTF-8?q?fix:=20two=20chat=20findings=20=E2=80=94?= =?UTF-8?q?=20protocol=20terminal=20states,=20provider=20error=20types?= 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 49/65] 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 50/65] 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 51/65] =?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 52/65] =?UTF-8?q?fix:=20two=20external=20review=20findings?= =?UTF-8?q?=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 53/65] 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 54/65] 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 55/65] 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 56/65] =?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 57/65] =?UTF-8?q?fix:=20conformant=20responses()=20envelop?= =?UTF-8?q?e=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 58/65] =?UTF-8?q?fix:=20stale=20anthropic>=3D0.84.0=20hint?= =?UTF-8?q?s=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 59/65] 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 60/65] 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 61/65] =?UTF-8?q?fix:=20declare=20the=20real=20python=20fl?= =?UTF-8?q?oor=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 62/65] 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 63/65] 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 64/65] 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 65/65] 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