diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 6b349a09b..b37db059a 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -30,6 +30,9 @@ jobs: python-version: ${{ matrix.python-version }} cache: pip - run: pip install -r requirements.txt pytest + - if: matrix.agent-frameworks == 'without' + # requirements.txt carries it; this leg tests the no-framework paths + run: pip uninstall -y openai-agents - if: matrix.agent-frameworks == 'with' run: pip install openai-agents claude-agent-sdk anthropic - run: python -m pytest -q diff --git a/.gitignore b/.gitignore index b5c223b31..5193735ca 100644 --- a/.gitignore +++ b/.gitignore @@ -6,4 +6,3 @@ __pycache__ logs/ .pageindex/ dist/ -*.doc_id diff --git a/README.md b/README.md index dca5dbc60..4a086ed6f 100644 --- a/README.md +++ b/README.md @@ -173,7 +173,7 @@ python3 run_pageindex.py --pdf_path /path/to/your/document.pdf
Optional parameters
-You can customize the processing with additional optional arguments (the structure-tuning flags below require --mode standard): +You can customize the processing with additional optional arguments (the structure-tuning flags from --toc-check-pages down require --mode standard): ``` --mode Processing mode: flash (default) or standard diff --git a/examples/agentic_vectorless_rag_demo.py b/examples/agentic_vectorless_rag_demo.py index 5b25c7638..93a00735c 100644 --- a/examples/agentic_vectorless_rag_demo.py +++ b/examples/agentic_vectorless_rag_demo.py @@ -32,14 +32,13 @@ from agents.stream_events import RawResponsesStreamEvent, RunItemStreamEvent from openai.types.responses import ResponseTextDeltaEvent, ResponseReasoningSummaryTextDeltaEvent -from pageindex import PageIndexAPIError, PageIndexLocalClient +from pageindex import PageIndexLocalClient import pageindex.utils as utils PDF_URL = "https://arxiv.org/pdf/2603.15031" _EXAMPLES_DIR = Path(__file__).parent PDF_PATH = _EXAMPLES_DIR / "documents" / "attention-residuals.pdf" -DOC_ID_PATH = _EXAMPLES_DIR / "documents" / "attention-residuals.doc_id" STORAGE_PATH = _EXAMPLES_DIR / ".pageindex" @@ -127,27 +126,13 @@ async def _run(): print("=" * 60) print("Step 1: Index PDF and view tree structure") print("=" * 60) - doc_id = None - if DOC_ID_PATH.exists(): - cached = DOC_ID_PATH.read_text().strip() - try: - client.get_document(cached) - doc_id = cached - except PageIndexAPIError: - DOC_ID_PATH.unlink() - if doc_id is None: - # The .doc_id cache is gitignored — on a fresh clone with an - # existing store, find the already-indexed copy by name instead of - # re-indexing it. - doc_id = next( - (doc["id"] for doc in client.list_documents(limit=100)["documents"] - if doc["name"] == PDF_PATH.name), None) + doc_id = next( + (doc["id"] for doc in client.list_documents(limit=100)["documents"] + if doc["name"] == PDF_PATH.name), None) if doc_id: - DOC_ID_PATH.write_text(doc_id) print(f"\nLoaded cached doc_id: {doc_id}") else: doc_id = client.submit_document(str(PDF_PATH), wait=True)["doc_id"] - DOC_ID_PATH.write_text(doc_id) print(f"\nIndexed. doc_id: {doc_id}") print("\nTree Structure (top-level sections):") structure = client.get_tree(doc_id, node_summary=True)["result"] diff --git a/pageindex/__init__.py b/pageindex/__init__.py index 88ca32ff8..669f13629 100644 --- a/pageindex/__init__.py +++ b/pageindex/__init__.py @@ -1,11 +1,6 @@ """PageIndex SDK.""" -import os as _os from typing import TYPE_CHECKING as _TYPE_CHECKING -# LiteLLM's import otherwise fetches its model map over the network — seconds -# of blocking (or a hang offline). setdefault, so an explicit user choice wins. -_os.environ.setdefault("LITELLM_LOCAL_MODEL_COST_MAP", "True") - from .client import PageIndexClient, PageIndexCloudClient, PageIndexLocalClient from .errors import PageIndexAPIError @@ -23,8 +18,6 @@ ] _LAZY = { - "page_index": ".page_index_classic", - "page_index_main": ".page_index_classic", "page_index_flash": ".flash", "optimize_tree": ".tree_optimize", "md_to_tree": ".page_index_md", @@ -34,25 +27,18 @@ "mcp_bridge", "page_index_classic", "page_index_md", "tree_optimize", "utils"} + def __getattr__(name): if name.startswith("_"): - # Dunder probes (copy, pickle, inspect) are the frequent unknown - # names — they must not trigger the classic import below. raise AttributeError(f"module {__name__!r} has no attribute {name!r}") import importlib if name in _SUBMODULES: return importlib.import_module(f".{name}", __name__) - # Pre-0.2.10 compat: unknown names fall through to the classic module, - # whose public surface (ConfigLoader, count_tokens, ...) resolved as - # package attributes. A non-underscore typo pays one classic import - # before its AttributeError — not worth an allowlist. - module = importlib.import_module(_LAZY.get(name, ".page_index_classic"), - __name__) + module = importlib.import_module(_LAZY.get(name, ".page_index_classic"), __name__) try: value = getattr(module, name) except AttributeError: - raise AttributeError( - f"module {__name__!r} has no attribute {name!r}") from None + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None globals()[name] = value return value diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index e00eb3ca9..000ce0e6a 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -9,11 +9,10 @@ way the agent instructions do — they never teach capabilities that only exist on the cloud. -Tools never raise for any invocation their signatures accept: every -outcome, including errors, is returned as the same JSON envelope the cloud -emits ({"success": true, ...} / {"error": ...}). Arguments outside a pruned -local signature fail at the Python call boundary; the call_tool path -answers them with the guided error envelope instead. +Tools never raise: every outcome, including errors, is returned as the +same JSON envelope the cloud emits ({"success": true, ...} / +{"error": ...}) — arguments outside a pruned local signature come back as +that envelope too, on the direct and the call_tool path alike. """ from __future__ import annotations @@ -33,7 +32,6 @@ STRUCTURE_FIRST_PAGE_THRESHOLD = 20 _CHAR_BUDGET = int(TOOL_RESPONSE_CHAR_LIMIT * 0.95) -_PAGES_SPEC_RE = re.compile(r"^(\d+(-\d+)?)(,\s*\d+(-\d+)?)*$") _MAX_REQUESTED_PAGES = 10_000 _SIMILAR_NAMES_LIMIT = 3 _TOOL_WAIT_TIMEOUT = 180.0 # "up to 3 minutes", per the wait_for_completion schema @@ -328,15 +326,23 @@ def _dumps(payload: dict[str, Any]) -> str: # ── document listing / name resolution ── -def _all_documents(client) -> list[dict[str, Any]]: +def _all_documents(client, stop_ids=None) -> list[dict[str, Any]]: """Every document the client can list, newest first (both modes list - newest-first; paging preserves that order).""" + newest-first; paging preserves that order). With ``stop_ids``, paging + stops early once every one of those ids has been seen — for callers + that only need those entries; an id absent from the listing still + costs a full sweep.""" documents: list[dict[str, Any]] = [] offset = 0 + remaining = {str(one_id) for one_id in stop_ids} if stop_ids else None while True: page = client.list_documents(limit=100, offset=offset) batch = page.get("documents") or [] documents.extend(batch) + if remaining is not None: + remaining.difference_update(str(doc.get("id")) for doc in batch) + if not remaining: + return documents # Advance by what actually arrived — stepping by the requested # limit skips documents whenever a server caps its page size. offset += len(batch) @@ -387,14 +393,14 @@ def _resolve_document( """Resolve doc_name to a list entry. Same-name duplicates resolve to the newest match. Returns (entry, None) or (None, error_payload_pair).""" if documents is None: - documents = _all_documents(client) + documents = _all_documents(client, stop_ids=allowed_ids) documents = _scope_documents(documents, allowed_ids) matches = [doc for doc in documents if doc.get("name") == doc_name] if matches: return max(matches, key=lambda d: d.get("createdAt") or ""), None names = [str(doc.get("name")) for doc in documents if doc.get("name")] - similar = difflib.get_close_matches(doc_name, names, n=_SIMILAR_NAMES_LIMIT, - cutoff=0.5) + similar = difflib.get_close_matches(str(doc_name), names, + n=_SIMILAR_NAMES_LIMIT, cutoff=0.5) message = ( "Document not found. Did you mean: " + ", ".join(f'"{name}"' for name in similar) + "?" @@ -433,7 +439,7 @@ def _await_completion(client, entry: dict[str, Any], wait: bool) -> dict[str, An time.sleep(_TOOL_WAIT_INTERVAL) refreshed = _refetch_entry(client, doc_id) if refreshed is None: - return current + continue # transient refetch failure: poll on to the deadline if refreshed.get("metadata") is None: # Status refetches omit (or null out) custom metadata; keep the # listing's copy. @@ -503,69 +509,102 @@ def _folder_unsupported(param: str) -> tuple[dict, bool]: # ── page spec handling ── -def _parse_page_spec( - pages: str, doc_name: str, -) -> "tuple[Optional[list[int]], Optional[_ToolResult]]": - """Expand '1-3,7' into a sorted, deduplicated page list, or an error.""" - invalid = _failure( - "Invalid page specification format", - {"doc_name": doc_name}, - { - "summary": "Failed to parse the pages parameter", - "options": [ - 'Use valid formats: "5", "3,7,10", "5-10", or "1-3,7,9-12"', - "Ensure page numbers are positive integers", - ], - }, - "INVALID_INPUT", - ) - if not isinstance(pages, str) or not _PAGES_SPEC_RE.match(pages.strip()): - return None, invalid - too_many = _failure( - f"Too many pages requested (over {_MAX_REQUESTED_PAGES})", - {"doc_name": doc_name}, - { - "summary": "The page specification spans too many pages", - "options": [ - "Request a narrower page range", - "The response holds only a few pages per call - page through with several smaller requests", - ], - }, - "INVALID_INPUT", - ) +class _PageSpecError(ValueError): + """Shared page-spec rejection; ``code`` picks the caller's rendering.""" + + def __init__(self, code: str, message: str): + super().__init__(message) + self.code = code + + +def _expand_pages(pages) -> list[int]: + """Expand '1-3,7' into sorted distinct pages — the one parser for the + SDK surface and the tool layer. Raises _PageSpecError with code + 'invalid', 'too_many', or 'nonpositive'.""" + if not isinstance(pages, str): + raise _PageSpecError("invalid", + f"Invalid page specification: {pages!r}") + too_many = (f"Page specification '{pages}' spans more than " + f"{_MAX_REQUESTED_PAGES} pages; request a narrower range") expanded: set[int] = set() for part in pages.split(","): part = part.strip() - if "-" in part: - start, end = (int(x) for x in part.split("-", 1)) - if start > end: - return None, invalid - else: - start = end = int(part) + try: + if "-" in part: + start, end = (int(x) for x in part.split("-", 1)) + if start > end: + raise _PageSpecError( + "invalid", + f"Invalid range '{part}': start must be <= end") + else: + start = end = int(part) + except _PageSpecError: + raise + except ValueError as exc: + raise _PageSpecError( + "invalid", f"Invalid page specification '{pages}'") from exc # Bound each part arithmetically before materializing it: a spec like # "1-1000000000" would otherwise expand to billions of integers # inside the caller's process. The cap is on distinct pages, so # overlapping parts (a parent section plus its children) don't # double-count. if end - start + 1 > _MAX_REQUESTED_PAGES: - return None, too_many + raise _PageSpecError("too_many", too_many) expanded.update(range(start, end + 1)) if len(expanded) > _MAX_REQUESTED_PAGES: - return None, too_many - if any(page < 1 for page in expanded): + raise _PageSpecError("too_many", too_many) + if min(expanded) < 1: + raise _PageSpecError( + "nonpositive", + "Invalid page numbers. Page numbers must be positive integers") + return sorted(expanded) + + +def _parse_page_spec( + pages: str, doc_name: str, +) -> "tuple[Optional[list[int]], Optional[_ToolResult]]": + """Expand '1-3,7' into a sorted, deduplicated page list, or an error.""" + try: + return _expand_pages(pages), None + except _PageSpecError as exc: + if exc.code == "too_many": + return None, _failure( + f"Too many pages requested (over {_MAX_REQUESTED_PAGES})", + {"doc_name": doc_name}, + { + "summary": "The page specification spans too many pages", + "options": [ + "Request a narrower page range", + "The response holds only a few pages per call - page through with several smaller requests", + ], + }, + "INVALID_INPUT", + ) + if exc.code == "nonpositive": + return None, _failure( + "Invalid page numbers. Page numbers must be positive integers", + {"doc_name": doc_name}, + { + "summary": "Invalid page numbers provided", + "options": [ + "Page numbers must be positive integers (>= 1)", + "Check the page specification format", + ], + }, + "INVALID_INPUT", + ) return None, _failure( - "Invalid page numbers. Page numbers must be positive integers", + "Invalid page specification format", {"doc_name": doc_name}, { - "summary": "Invalid page numbers provided", + "summary": "Failed to parse the pages parameter", "options": [ - "Page numbers must be positive integers (>= 1)", - "Check the page specification format", + 'Use valid formats: "5", "3,7,10", "5-10", or "1-3,7,9-12"', + "Ensure page numbers are positive integers", ], }, "INVALID_INPUT", ) - return sorted(expanded), None def _format_page_spec(pages: list[int]) -> str: @@ -702,7 +741,8 @@ def _browse_documents(client, folder_id: str = "root", recursive: bool = False, window = listing.get("documents") or [] total = listing.get("total") else: - scoped = _scope_documents(_all_documents(client), _allowed_ids) + scoped = _scope_documents(_all_documents(client, stop_ids=_allowed_ids), + _allowed_ids) window, total = scoped[offset:offset + limit], len(scoped) window_end = offset + len(window) has_more = bool(window) and (window_end < total if isinstance(total, int) @@ -863,7 +903,9 @@ def _get_document_structure(client, doc_name: str, raw_tree = getattr(getattr(client, "_api", None), "raw_tree", None) tree = raw_tree(entry["id"]) if raw_tree is not None else None if tree is None: - tree = client.get_tree(entry["id"], node_summary=True).get("result") + # _format_structure strips text anyway — don't download it. + tree = client.get_tree(entry["id"], node_summary=True, + include_text=False).get("result") except PageIndexAPIError as exc: return _failure( f"Failed to retrieve document structure: {exc}", @@ -1020,10 +1062,12 @@ def _get_page_content(client, doc_name: str, pages: str, markdown = item.get("markdown") if item else None text = (markdown if isinstance(markdown, str) else f"Page {page} content not available") - if not included or budget - len(text) >= 0: - content.append({"page": page, "text": text}) + entry = {"page": page, "text": text} + size = _serialized_size(entry) + 2 # +2: json ", " item separator + if not included or budget - size >= 0: + content.append(entry) included.append(page) - budget -= len(text) + budget -= size else: remaining.append(page) @@ -1087,7 +1131,7 @@ def _remove_document(client, doc_names: list[str], {"summary": "Too many documents in one call", "options": ["Delete at most 10 documents per call"]}, "INVALID_INPUT") - documents = _all_documents(client) + documents = _all_documents(client, stop_ids=_allowed_ids) results = [] for doc_name in doc_names: entry, error = _resolve_document(client, doc_name, documents=documents, @@ -1126,11 +1170,11 @@ def tool_names(include_management: bool = False) -> tuple[str, ...]: return _READ_TOOLS + (_MANAGEMENT_TOOLS if include_management else ()) -def _coerce_bool_args(name: str, kwargs: dict[str, Any]) -> None: +def _coerce_bool_args(schema: dict, kwargs: dict[str, Any]) -> None: """Models routinely send booleans as JSON strings ("false"); the bare - truthiness tests downstream would read those as True.""" - properties = TOOL_CONTRACT.get(name, {}).get("schema", {}).get( - "properties", {}) + truthiness tests downstream would read those as True. Runs on both + dispatch paths — call_tool and the cloud bridge invoker.""" + properties = (schema or {}).get("properties", {}) for key, spec in properties.items(): value = kwargs.get(key) if spec.get("type") == "boolean" and isinstance(value, str): @@ -1168,7 +1212,7 @@ def call_tool(client, name: str, arguments: dict[str, Any], # "omit if ..." semantics, same as the cloud bridge invoker). kwargs = {key: value for key, value in (arguments or {}).items() if not key.startswith("_") and value is not None} - _coerce_bool_args(name, kwargs) + _coerce_bool_args(TOOL_CONTRACT.get(name, {}).get("schema", {}), kwargs) try: if doc_ids is not None: ids = [doc_ids] if isinstance(doc_ids, str) else doc_ids @@ -1229,9 +1273,11 @@ def _tool_docstring(description: str, properties: dict[str, Any]) -> str: 'Folder browsing and semantic ranking (sort="relevance") are not ' "supported in local mode yet — they work on PageIndex cloud." ), - "get_page_content": TOOL_CONTRACT["get_page_content"]["description"] - .replace(" Embedded image paths in the response feed into " - "`get_document_image()`.", ""), + # Drop the sentence naming the cloud-only image tool, whatever its + # wording; the contract-refresh test pins that something was removed. + "get_page_content": re.sub( + r"\s*[^.]*`get_document_image\(\)`[^.]*\.", "", + TOOL_CONTRACT["get_page_content"]["description"]), } _LOCAL_PARAM_DESCRIPTIONS: dict[tuple[str, str], str] = { @@ -1261,9 +1307,6 @@ def _local_schema(name: str) -> dict[str, Any]: return schema -def _docstring(name: str) -> str: - return _tool_docstring(_local_description(name), - _local_schema(name)["properties"]) _SCHEMA_TYPE_MAP = {"string": str, "integer": int, "number": float, @@ -1298,17 +1341,23 @@ def _annotation_for(spec: dict) -> Any: return Optional[base] if nullable else base -def _bridge_invoker(bridge, name: str) -> "Callable[[dict], tuple[str, bool]]": - """One cloud tool call proxied over MCP: None-valued arguments are - dropped (None ≡ omitted, matching the contract's "omit if ..." - semantics) and failures are contained in the error envelope. Returns - (envelope_text, is_error), like call_tool.""" +def _bridge_invoker(bridge, name: str, schema: dict, + ) -> "Callable[[dict], tuple[str, bool]]": + """One cloud tool call proxied over MCP: string booleans are coerced + (same as call_tool), None-valued arguments are dropped (None ≡ omitted, + matching the contract's "omit if ..." semantics) and failures are + contained in the error envelope — except 401/403, which re-raise. + Returns (envelope_text, is_error), like call_tool.""" def _invoke(arguments: dict[str, Any]) -> tuple[str, bool]: try: arguments = {key: value for key, value in arguments.items() if value is not None} + _coerce_bool_args(schema, arguments) return bridge.call_tool(name, arguments) except Exception as exc: + if (isinstance(exc, PageIndexAPIError) + and exc.status_code in (401, 403)): + raise payload, _ = _failure( f"{name} failed: {exc}", None, {"summary": "Unexpected error while running the tool", @@ -1321,22 +1370,23 @@ def _invoke(arguments: dict[str, Any]) -> tuple[str, bool]: return _invoke -def _make_bridge_function(bridge, meta: dict) -> Callable[..., str]: - """One plain function for a cloud tool: real signature and docstring from - the server's schema, invocation proxied over MCP, errors contained.""" +def _make_tool_function(name: str, description: str, schema: dict, + invoke: "Callable[[dict], tuple[str, bool]]", + ) -> Callable[..., str]: + """One plain function for a tool: real signature and docstring from the + schema, errors contained by the invoker; arguments the signature + rejects come back as the guided envelope instead of raising.""" import keyword - name = str(meta.get("name") or "") - schema = meta.get("inputSchema") or {} properties: dict[str, Any] = schema.get("properties") or {} required = set(schema.get("required") or []) - _invoke = _bridge_invoker(bridge, name) + _invoke = invoke params_usable = all(param.isidentifier() and not keyword.iskeyword(param) and param != "_invoke" for param in properties) if not params_usable: - def proxy(**kwargs: Any) -> str: + def inner(**kwargs: Any) -> str: return _invoke(kwargs)[0] else: ordered = ([p for p in properties if p in required] @@ -1349,7 +1399,9 @@ def proxy(**kwargs: Any) -> str: namespace: dict[str, Any] = {"_invoke": _invoke} exec(f"def _synthesized({rendered}):\n" f" return _invoke({args_literal})[0]", namespace) - proxy = namespace["_synthesized"] + inner = namespace["_synthesized"] + # binding TypeErrors quote __qualname__, not __name__ + inner.__name__ = inner.__qualname__ = name or "tool" annotations: dict[str, Any] = {} for p in ordered: annotation = _annotation_for(properties[p]) @@ -1359,9 +1411,26 @@ def proxy(**kwargs: Any) -> str: annotation = Optional[annotation] annotations[p] = annotation annotations["return"] = str - proxy.__annotations__ = annotations + inner.__annotations__ = annotations + + def proxy(*args: Any, **kwargs: Any) -> str: + # The invoker lets only 401/403 auth failures through, so a + # TypeError here is the binding rejecting the arguments. + try: + return inner(*args, **kwargs) + except TypeError as exc: + payload, _ = _failure( + f"Invalid arguments for {name}: {exc}", None, + {"summary": "Invalid tool arguments", + "options": [f"Check the {name}() parameter names " + "and types"]}, + "INVALID_INPUT", + ) + return _dumps(payload) + proxy.__signature__ = inspect.signature(inner) # type: ignore[attr-defined] + proxy.__annotations__ = dict(inner.__annotations__) proxy.__name__ = proxy.__qualname__ = name or "tool" - proxy.__doc__ = _tool_docstring(meta.get("description") or "", properties) + proxy.__doc__ = _tool_docstring(description or "", properties) return proxy @@ -1369,19 +1438,26 @@ def proxy(**kwargs: Any) -> str: _BRIDGES_LOCK = threading.Lock() -def _cloud_bridge(client): - """One bridge per client: tool discovery and instructions share a single - MCP session. Weak-keyed off the instance so clients stay picklable; the +def _cloud_bridge(client, gated: bool = False): + """One bridge per client and endpoint gate (``gated`` = the read-only + ?tools=read endpoint); tool discovery and instructions share a session + per gate. Weak-keyed off the instance so clients stay picklable; the lock closes the check-then-set race under concurrent first calls.""" with _BRIDGES_LOCK: - bridge = _BRIDGES.get(client) + # A rotated api_key or moved BASE_URL rebuilds the bridges. + auth = (client.BASE_URL, client.api_key) + bridges, seen = _BRIDGES.get(client) or ({}, None) + if seen != auth: + bridges = {} + bridge = bridges.get(gated) if bridge is None: from .mcp_bridge import McpBridge bridge = McpBridge( - f"{client.BASE_URL}/mcp", - {"Authorization": f"Bearer {client.api_key}"}, + f"{auth[0]}/mcp" + ("?tools=read" if gated else ""), + {"Authorization": f"Bearer {auth[1]}"}, ) - _BRIDGES[client] = bridge + bridges[gated] = bridge + _BRIDGES[client] = (bridges, auth) return bridge @@ -1401,17 +1477,15 @@ def _read_only_tools(tools_meta: list[dict]) -> list[dict]: return filtered -def _build_cloud_agent_tools(client, include_management: bool) -> list[Callable[..., str]]: - bridge = _cloud_bridge(client) - tools_meta = bridge.list_tools() - if not include_management: - tools_meta = _read_only_tools(tools_meta) - return [_make_bridge_function(bridge, meta) for meta in tools_meta] - - def _require_local_scope(client, doc_ids) -> None: """The allowlist is enforced in-process; cloud lookups run server-side, - so accepting doc_ids there would be advisory-only — refuse loudly.""" + so accepting doc_ids there would be advisory-only — refuse loudly. + An empty allowlist is refused too: it would scope the agent to + nothing, with no signal to the caller.""" + if doc_ids is not None and not doc_ids: + raise PageIndexAPIError( + "doc_id is empty. Pass one or more document IDs, or omit " + "doc_id to give the agent the whole library.") if doc_ids is not None and getattr(client, "api_key", None): raise PageIndexAPIError( "doc_ids scoping applies to local tools only — cloud calls " @@ -1427,7 +1501,7 @@ def _tool_specs(client, include_management: bool = False, doc_ids=None, is the local chat scope; cloud scoping is server-side.""" _require_local_scope(client, doc_ids) if getattr(client, "api_key", None): - bridge = _cloud_bridge(client) + bridge = _cloud_bridge(client, gated=not include_management) tools_meta = bridge.list_tools() if not include_management: tools_meta = _read_only_tools(tools_meta) @@ -1435,7 +1509,8 @@ def _tool_specs(client, include_management: bool = False, doc_ids=None, meta.get("description") or "", copy.deepcopy(meta.get("inputSchema")) or {"type": "object", "properties": {}}, - _bridge_invoker(bridge, str(meta.get("name") or "tool"))) + _bridge_invoker(bridge, str(meta.get("name") or "tool"), + meta.get("inputSchema") or {})) for meta in tools_meta] def local_invoke(name: str) -> "Callable[[dict], tuple[str, bool]]": @@ -1448,7 +1523,8 @@ def invoke(arguments: dict) -> tuple[str, bool]: for name in tool_names(include_management)] -def build_agent_tools(client, include_management: bool = False) -> list[Callable[..., str]]: +def build_agent_tools(client, include_management: bool = False, + doc_ids=None) -> list[Callable[..., str]]: """Plain synchronous functions bound to `client`. Cloud: one function per tool of the live cloud MCP tool set, signatures @@ -1457,53 +1533,11 @@ def build_agent_tools(client, include_management: bool = False) -> list[Callable the JSON envelope as a string and never raises for arguments its signature accepts (cloud-only parameters are absent from the local signatures; the call_tool path answers them with the guided envelope). + ``doc_ids`` is the local allowlist, as in ``_tool_specs``. """ - if getattr(client, "api_key", None): - return _build_cloud_agent_tools(client, include_management) - - def browse_documents(offset: int = 0, limit: int = 10) -> str: - return call_tool(client, "browse_documents", { - "offset": offset, "limit": limit, - })[0] - - def get_document(doc_name: str, wait_for_completion: bool = False) -> str: - return call_tool(client, "get_document", { - "doc_name": doc_name, - "wait_for_completion": wait_for_completion, - })[0] - - def get_document_structure(doc_name: str, part: int = 1, - wait_for_completion: bool = False) -> str: - return call_tool(client, "get_document_structure", { - "doc_name": doc_name, "part": part, - "wait_for_completion": wait_for_completion, - })[0] - - def get_page_content(doc_name: str, pages: str, - wait_for_completion: bool = False) -> str: - return call_tool(client, "get_page_content", { - "doc_name": doc_name, "pages": pages, - "wait_for_completion": wait_for_completion, - })[0] - - def remove_document(doc_names: list[str]) -> str: - return call_tool(client, "remove_document", { - "doc_names": doc_names, - })[0] - - functions = { - "browse_documents": browse_documents, - "get_document": get_document, - "get_document_structure": get_document_structure, - "get_page_content": get_page_content, - "remove_document": remove_document, - } - tools = [] - for name in tool_names(include_management): - function = functions[name] - function.__doc__ = _docstring(name) - tools.append(function) - return tools + return [_make_tool_function(name, description, schema, invoke) + for name, description, schema, invoke + in _tool_specs(client, include_management, doc_ids)] # ── agent instructions ── @@ -1557,12 +1591,13 @@ def remove_document(doc_names: list[str]) -> 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.""" +def _base_instructions(client, include_management: bool = False) -> str: + """Cloud: the live instructions the MCP server serves for the tool set + actually shipped. Local: the built-in subset instructions.""" if not getattr(client, "api_key", None): return AGENT_INSTRUCTIONS - instructions = _cloud_bridge(client).instructions() + instructions = _cloud_bridge( + client, gated=not include_management).instructions() if not isinstance(instructions, str) or not instructions.strip(): raise PageIndexAPIError( "The MCP server returned no agent instructions — refusing to " @@ -1585,9 +1620,29 @@ def doc_targeting_block(client, doc_id, scoped: bool = False) -> Optional[str]: return None doc_ids = [doc_id] if isinstance(doc_id, str) else list(doc_id) if not doc_ids: - return None - details = [client.get_document(one_id) for one_id in doc_ids] - listing = _all_documents(client) + # An empty selection must fail loud: washing it to None would mean + # "everything", and the tool-layer allowlist would mean "nothing". + raise PageIndexAPIError( + "doc_id is empty. Pass one or more document IDs, or omit " + "doc_id to give the agent the whole library.") + details = [] + missing = [] + for one_id in doc_ids: + try: + details.append(client.get_document(one_id)) + except PageIndexAPIError as exc: + # Batch only a definite not-found/denied (local raises carry no + # status); a cloud transport failure (429/5xx) propagates raw. + if exc.status_code not in (None, 403, 404): + raise + missing.append(str(one_id)) + if missing: + raise PageIndexAPIError( + "Documents not found or access denied: " + ", ".join(missing)) + # Scoped: the listing only backfills the target docs' metadata (list + # entries carry it, get_document does not — cloud parity), so paging + # can stop at those ids. Unscoped needs it all for the shadow check. + listing = _all_documents(client, stop_ids=doc_ids if scoped else None) documents = ([{**detail, "id": one_id} for one_id, detail in zip(doc_ids, details)] if scoped else listing) @@ -1625,9 +1680,10 @@ def doc_targeting_block(client, doc_id, scoped: bool = False) -> Optional[str]: ) -def build_agent_instructions(client, doc_id=None, scoped: bool = False) -> str: +def build_agent_instructions(client, doc_id=None, scoped: bool = False, + include_management: bool = False) -> str: """Orchestration guidance for document QA agents; with doc_id, appends the target documents and directs the agent to work within them.""" - base = _base_instructions(client) + base = _base_instructions(client, include_management) block = doc_targeting_block(client, doc_id, scoped=scoped) return base if block is None else base + "\n\n" + block diff --git a/pageindex/client.py b/pageindex/client.py index 8125aa525..c4d71eb15 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -10,35 +10,37 @@ from .errors import PageIndexAPIError +_litellm_preload_started = False + + def _preload_litellm() -> None: - try: - import litellm # noqa: F401 - except Exception: - pass + """Start litellm's multi-second import in the background, once per + process — a per-client thread would churn under per-request clients.""" + # LiteLLM's import otherwise fetches its model map over the network — + # seconds of blocking (or a hang offline). Stamped here, not at package + # import, so merely importing pageindex leaves the host process's own + # litellm untouched; setdefault, so an explicit user choice wins. + os.environ.setdefault("LITELLM_LOCAL_MODEL_COST_MAP", "True") + global _litellm_preload_started + if _litellm_preload_started: + return + _litellm_preload_started = True + + def _import() -> None: + try: + import litellm # noqa: F401 + except Exception: + pass + + threading.Thread(target=_import, daemon=True).start() def _parse_pages(pages: str) -> list[int]: - result: set[int] = set() - too_many = (f"Page specification '{pages}' spans more than " - "10000 pages; request a narrower range") - for part in pages.split(","): - part = part.strip() - if "-" in part: - start, end = (int(x) for x in part.split("-", 1)) - if start > end: - raise ValueError(f"Invalid range '{part}': start must be <= end") - else: - start = end = int(part) - # Bound each part arithmetically before materializing it — a spec - # like "1-999999999" would otherwise expand to a billion integers. - # The cap is on distinct pages, so overlapping parts (a parent - # section plus its children) don't double-count. - if end - start + 1 > 10_000: - raise ValueError(too_many) - result.update(range(start, end + 1)) - if len(result) > 10_000: - raise ValueError(too_many) - return sorted(result) + from .agent_tools import _PageSpecError, _expand_pages + try: + return _expand_pages(pages) + except _PageSpecError as exc: + raise PageIndexAPIError(str(exc)) from exc def _agents_sdk_model_name(model: str) -> str: @@ -155,7 +157,7 @@ def __init__( self.model = opt.model self.index_model = opt.index_model self.summary_model = opt.summary_model - self.chat_model = _agents_sdk_model_name(opt.chat_model) + self.chat_model = opt.chat_model self.chat_backend = chat_backend self.storage_path = storage_path or ".pageindex" from .local_api import LocalAPI @@ -163,18 +165,22 @@ def __init__( storage_path=self.storage_path, model=self.model, summary_model=self.summary_model, - retrieve_model=self.chat_model, index_backend=index_backend, ) # LiteLLM's multi-second import would otherwise land on the # first chat call; failures resurface there with real context. - threading.Thread(target=_preload_litellm, daemon=True).start() + _preload_litellm() @property def retrieve_model(self): """Legacy name for ``chat_model``.""" return self.chat_model + @retrieve_model.setter + def retrieve_model(self, value): + # 0.2.9 allowed assignment; keep the write path working too. + self.chat_model = value + # ---------- DOCUMENT SUBMISSION ---------- def submit_document( @@ -432,7 +438,12 @@ def chat( if stream: return cast(Iterator[str], result) envelope = cast(dict[str, Any], result) - return envelope["choices"][0]["message"]["content"] or "" + try: + return envelope["choices"][0]["message"]["content"] or "" + except (KeyError, IndexError, TypeError) as exc: + raise PageIndexAPIError( + "The chat response carries no answer: " + f"{str(envelope)[:200]}") from exc def chat_completions( self, @@ -598,11 +609,12 @@ def responses( model: Backend model name (defaults to ``chat_model``). stream: Yield Responses stream events as dicts — one logical response per call: per-turn backend lifecycle events are - collapsed, sequence numbers are reassigned monotonically, - and ``output_index`` is re-based onto the single logical - ``output``. The single final event is the terminal - ``response.*`` for the run's status; its ``response`` - carries the tool outputs in ``items``. + collapsed to one opening ``response.created`` and one + final terminal event, sequence numbers are reassigned + monotonically, and ``output_index`` is re-based onto the + single logical ``output``. The final event is the + terminal ``response.*`` for the run's status; its + ``response`` carries the tool outputs in ``items``. doc_id: Document ID or list of IDs to scope the conversation. Keep it identical across a conversation's calls — the targeting block it adds is re-set each call and is part @@ -673,7 +685,10 @@ def messages( final message envelope with cross-turn aggregated ``usage`` plus a ``messages`` field — the full new turn sequence, valid for verbatim append to your history. The managed system prompt carries a - ``cache_control`` breakpoint. + ``cache_control`` breakpoint, and the request sets the top-level + ``cache_control`` so each turn re-reads the growing conversation + from cache — skipped when your own blocks already use the three + remaining breakpoints (the managed prompt holds the fourth). Args: messages: Native Messages-format history (including prior @@ -683,7 +698,9 @@ def messages( max_tokens: Per-turn output budget the Messages API requires on the wire; the default is resolved per model (8192, or 4096 for the claude-3 generation whose ceiling is lower) so the - simple call needs only a question. Passed through. + simple call needs only a question, and rises to + budget_tokens + 8192 when ``thinking`` is enabled (the wire + requires max_tokens above the budget). Passed through. stream: Yield the Anthropic SDK's event stream across turns (its native event objects, including SDK-synthesized convenience events), one message sequence per turn. @@ -776,13 +793,16 @@ def list_documents( # ---------- AGENT INTEGRATION ---------- - def agent_tools(self, include_management: bool = False) -> list[Callable[..., str]]: + def agent_tools( + self, include_management: bool = False, + doc_id: Optional[Union[str, list[str]]] = None, + ) -> list[Callable[..., str]]: """ Plain functions for any agent framework (LangChain, PydanticAI, ...). For the OpenAI / Claude Agent SDKs, prefer ``as_openai_tools()`` / ``as_claude_mcp()``. - Cloud: the full cloud tool set, discovered live from the PageIndex + Cloud: the full live read tool set, discovered from the PageIndex MCP server when this method is called — one function per tool, signature and docstring synthesized from the server's schemas, calls executed from your process over MCP. Raises PageIndexAPIError if the @@ -798,9 +818,13 @@ def agent_tools(self, include_management: bool = False) -> list[Callable[..., st library. Local: adds ``remove_document``. Cloud: by default only tools the server marks read-only are exposed; True exposes the server's complete list (upload, delete, ...). + doc_id: Local only — restrict the tools to this document ID + (or list of IDs), enforced at the tool layer: out-of-scope + lookups return NOT_FOUND. Raises on cloud, where scoping + is server-side. """ from .agent_tools import build_agent_tools - return build_agent_tools(self, include_management) + return build_agent_tools(self, include_management, doc_ids=doc_id) def as_openai_tools(self, include_management: bool = False, hosted: bool = False, @@ -851,14 +875,14 @@ def _local_doc_scope(self, doc_id): """doc_id for the tool layer: passed through locally (structural allowlist), dropped on cloud where scoping is server-side and the config helpers keep prompt-level targeting.""" - if not getattr(self, "api_key", None): - return doc_id if doc_id is not None and not doc_id: - # Cloud has no tool-layer allowlist to make an empty scope mean - # "nothing"; dropping it would silently mean "everything". + # An empty scope means "nothing" locally (empty allowlist) and + # cannot be represented on cloud; both refuse it loudly. raise PageIndexAPIError( "doc_id is empty. Pass one or more document IDs, or omit " "doc_id to give the agent the whole library.") + if not getattr(self, "api_key", None): + return doc_id return None def openai_agent_config( @@ -882,6 +906,16 @@ def openai_agent_config( environment, so its model auth comes from there — ``chat_backend`` does not travel with it. + Prompt caching configures itself for most destinations (OpenAI + server-side; Anthropic- and Bedrock-hosted Claude via LiteLLM's + defaults). Vertex-hosted Claude is the exception — pass the + injection points yourself:: + + Agent(**config, model_settings=ModelSettings(extra_args={ + "cache_control_injection_points": [ + {"location": "message", "role": "system"}, + {"location": "message", "index": -1}]})) + Args: doc_id: Document ID or list of IDs to target, as in ``agent_instructions``. Local: also enforced at the tool @@ -889,19 +923,27 @@ def openai_agent_config( (tool scoping is server-side). include_management (bool): Also expose tools that modify the library. - model: Backend model name; overrides the local default. + model: Backend model name; overrides the local default. Same + grammar as ``chat_model`` (LiteLLM names; bare names are + OpenAI-compatible shorthand). """ from .agent_tools import build_agent_instructions scope = self._local_doc_scope(doc_id) config: dict[str, Any] = { "name": "PageIndex", - "instructions": build_agent_instructions(self, doc_id, - scoped=scope is not None), + "instructions": build_agent_instructions( + self, doc_id, scoped=scope is not None, + include_management=include_management), "tools": self.as_openai_tools(include_management, doc_id=scope), } model = model or getattr(self, "chat_model", None) if model: - config["model"] = model + config["model"] = _agents_sdk_model_name(model) + if config["model"].startswith("litellm/"): + # The runner resolves this model through LiteLLM in the + # caller's process, outside our completion helpers. + from .utils import _repair_litellm_types + _repair_litellm_types() return config def as_anthropic_tools(self, include_management: bool = False, @@ -961,6 +1003,7 @@ def anthropic_runner_config( asynchronous: bool = False, max_tokens: Optional[int] = None, max_turns: Optional[int] = None, + thinking: Optional[dict] = None, ) -> dict[str, Any]: """ Document QA ``tool_runner`` kwargs for the Anthropic SDK in one @@ -973,10 +1016,14 @@ def anthropic_runner_config( Sugar over the explicit form — ``agent_instructions`` (with ``doc_id`` targeting) as the system prompt and - ``as_anthropic_tools`` as the tools — plus the same defaults - ``messages()`` applies: a per-model ``max_tokens`` and a - ``max_iterations`` bound of 10. To customize further, switch to - those methods directly. + ``as_anthropic_tools`` as the tools — plus the ``max_tokens`` + default and 10-turn ``max_iterations`` bound ``messages()`` uses, + and a top-level ``cache_control`` so each loop turn re-reads the + growing prompt from cache (pop the key if you place your own + breakpoints — the API allows four). Unlike ``messages()``, + ``system`` here is the bare instructions string, without the chat + header or its block-level breakpoint. To customize further, + switch to those methods directly. Args: model: Backend model name (also resolves the ``max_tokens`` @@ -992,19 +1039,27 @@ def anthropic_runner_config( max_tokens: Per-turn output budget; default resolved per model. max_turns: Agent-loop bound; default 10. + thinking: Anthropic ``thinking`` config, included in the + kwargs; an enabled budget also lifts the ``max_tokens`` + default above it. Pass it here, not alongside the + unpacked config, so the default stays valid. """ from .agent_tools import build_agent_instructions - from .local_chat import _default_max_tokens + from .local_chat import _default_max_tokens, _validate_max_turns + _validate_max_turns(max_turns) scope = self._local_doc_scope(doc_id) return { "model": model, "max_tokens": (max_tokens if max_tokens is not None - else _default_max_tokens(model)), - "system": build_agent_instructions(self, doc_id, - scoped=scope is not None), + else _default_max_tokens(model, thinking)), + "system": build_agent_instructions( + self, doc_id, scoped=scope is not None, + include_management=include_management), "tools": self.as_anthropic_tools(include_management, asynchronous, doc_id=scope), "max_iterations": max_turns if max_turns is not None else 10, + **({"thinking": thinking} if thinking is not None else {}), + "cache_control": {"type": "ephemeral"}, } def as_claude_mcp(self, include_management: bool = False, @@ -1072,8 +1127,9 @@ def claude_agent_config( from .agent_tools import build_agent_instructions scope = self._local_doc_scope(doc_id) return { - "system_prompt": build_agent_instructions(self, doc_id, - scoped=scope is not None), + "system_prompt": build_agent_instructions( + self, doc_id, scoped=scope is not None, + include_management=include_management), "mcp_servers": {server_name: self.as_claude_mcp( include_management, doc_id=scope)}, # Pre-approval only — the server itself is already gated (the @@ -1081,7 +1137,10 @@ def claude_agent_config( "allowed_tools": [f"mcp__{server_name}"], } - def agent_instructions(self, doc_id: Optional[Union[str, list[str]]] = None) -> str: + def agent_instructions( + self, doc_id: Optional[Union[str, list[str]]] = None, + include_management: bool = False, + ) -> str: """ Orchestration guidance for document QA agents — pass as the agent's system prompt (or append to your own). @@ -1094,12 +1153,19 @@ def agent_instructions(self, doc_id: Optional[Union[str, list[str]]] = None) -> With ``doc_id`` (str or list, same shape as ``chat_completions``), appends the target documents' names and metadata and directs the - agent to work within them. Raises PageIndexAPIError if a doc_id does - not exist, or if its name is shadowed by a newer same-name document - (the name-addressed tools could not reach it). + agent to work within them. Raises PageIndexAPIError if a doc_id + does not exist, or if its name is shadowed by a newer same-name + document — the name-addressed tools could not reach it (the + ``*_agent_config`` bundles, whose tools carry the doc_id scope, + relax this to duplicates within the targeted set). + + ``include_management``: fetch the guidance for the full tool set, + matching tools built with ``include_management=True`` (cloud; + local guidance is a single set). """ from .agent_tools import build_agent_instructions - return build_agent_instructions(self, doc_id) + return build_agent_instructions( + self, doc_id, include_management=include_management) # ---------- FOLDER MANAGEMENT ---------- diff --git a/pageindex/cloud_api.py b/pageindex/cloud_api.py index ab7a9c885..f3a7740b6 100644 --- a/pageindex/cloud_api.py +++ b/pageindex/cloud_api.py @@ -59,7 +59,7 @@ def submit_document( (a taken name gains a numeric suffix), when the server returns it. """ - data = {'if_retrieval': True} + data: Dict[str, Any] = {'if_retrieval': True} if mode is not None: data['mode'] = mode if beta_headers is not None: @@ -225,7 +225,7 @@ def chat_completions( headers=self._headers(), json=payload, stream=stream, - timeout=120 if stream else 300 + timeout=120 if stream else 600 ) if response.status_code != 200: @@ -312,7 +312,9 @@ def get_document(self, doc_id: str) -> Dict[str, Any]: timeout=30 ) if response.status_code != 200: - raise PageIndexAPIError(f"Failed to get document metadata: {response.text}") + raise PageIndexAPIError( + f"Failed to get document metadata: {response.text}", + status_code=response.status_code) return response.json() def delete_document(self, doc_id: str) -> Dict[str, Any]: @@ -356,7 +358,7 @@ def list_documents(self, limit: int = 50, offset: int = 0, folder_id: Optional[s if offset < 0: raise ValueError("offset must be non-negative") - params = {"limit": limit, "offset": offset} + params: Dict[str, Any] = {"limit": limit, "offset": offset} if folder_id is not None: params["folder_id"] = folder_id diff --git a/pageindex/config.yaml b/pageindex/config.yaml index 4fde6f60e..d3786ca12 100644 --- a/pageindex/config.yaml +++ b/pageindex/config.yaml @@ -2,8 +2,8 @@ # chat_model answers questions on the chat surfaces; set model to use one # for both. Unset keys use the SDK defaults shown below. Legacy keys # (model, summary_model, retrieve_model) keep working. -# Indexing names without a provider prefix use the OpenAI SDK directly; -# for other providers, use "provider/model" (e.g. "anthropic/claude-sonnet-4-6"). +# Model names are LiteLLM's: bare names are OpenAI models; for other +# providers, use "provider/model" (e.g. "anthropic/claude-sonnet-4-6"). # index_model: "gpt-5.6-luna" # chat_model: "gpt-5.6-sol" toc_check_page_num: 20 diff --git a/pageindex/errors.py b/pageindex/errors.py index e460a956b..608ba6e4f 100644 --- a/pageindex/errors.py +++ b/pageindex/errors.py @@ -1,2 +1,7 @@ class PageIndexAPIError(Exception): - pass + """status_code carries the HTTP status when the failure came from a + non-200 cloud response; None otherwise (local mode, client-side).""" + + def __init__(self, *args: object, status_code: int | None = None) -> None: + super().__init__(*args) + self.status_code = status_code diff --git a/pageindex/flash/api.py b/pageindex/flash/api.py index bf62d9657..72438325f 100644 --- a/pageindex/flash/api.py +++ b/pageindex/flash/api.py @@ -96,19 +96,33 @@ def _optimize(structure, page_texts, do_expand, model): def page_index_flash(pdf, summary=True, summary_model=None, - optimize: str | bool = "full", optimize_expand=None, + optimize: str | bool | None = None, optimize_expand=None, optimize_model=None, summary_concurrency=None, use_embedded_toc=True) -> dict: - """Build a PageIndex tree structure from a PDF using layout statistics, 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" + """Build a PageIndex tree structure from a PDF using layout statistics. The tree extraction itself uses no LLM; by default an LLM writes node summaries and expands the tree (``summary=False, optimize=False`` runs fully LLM-free). Args: pdf: path to a PDF file (``str`` or ``pathlib.Path``) or an in-memory binary stream (``io.BytesIO``). summary: if True, generate LLM summaries for each node (requires ``summary_model``). summary_model: the LLM model identifier to use for summary generation. optimize: ``"full"`` for merge + LLM expand (fails fast with ``PageIndexAPIError`` when no LLM key is configured), ``"merge"`` for deterministic merge only, ``False`` to disable. ``True`` is accepted as ``"full"`` for backward compatibility; defaults to ``"full"``. optimize_expand: deprecated — use ``optimize``. Honored only when ``optimize`` is not passed (or is the legacy ``True``): ``False`` maps to ``"merge"``, ``True`` to ``"full"``. optimize_model: the LLM model for expand (defaults to the summary model). summary_concurrency: maximum simultaneous summary model calls; None uses the library default. use_embedded_toc: if True, consume the PDF's embedded bookmarks when trustworthy: deep bookmarks become the frame and the detected sections they lack are grafted back in after noise filtering, coarse ones become the chapter frame with detected nodes re-hung under them (deeper sparse entries are filled in when the page text confirms them, and garbled extracted titles are repaired from the bookmark strings), garbage ones are ignored; adds a ``toc_source`` key to the result. On by default; pass False for the pure detected structure. Returns: dict with keys ``doc_name``, ``doc_title``, ``structure`` (a list of nested ``{"title", "start_index", "end_index", "nodes"}`` dicts; page indexes are 1-based) and ``has_abstract_or_references_section`` (True when a top-level entry is an abstract or references heading). With ``optimize`` an ``optimize`` key reports merge/expand counts and before/after search-cost metrics. """ + if optimize_expand is not None: + import warnings + warnings.warn( + "optimize_expand is deprecated: pass optimize='full', 'merge', " + "or False.", DeprecationWarning, stacklevel=2) + if optimize is None or optimize is True: + # legacy spellings only — an explicit 'full'/'merge' wins + optimize = "merge" if optimize_expand is False else "full" if not optimize: optimize = False elif optimize not in ("full", "merge"): raise ValueError( f"optimize must be 'full', 'merge', or False, got {optimize!r}") - if optimize_expand is not None and optimize: - optimize = "full" if optimize_expand else "merge" + if optimize == "full": + from ..errors import PageIndexAPIError + from ..utils import ConfigLoader, _llm_backend, _openai_missing_keys + model = (optimize_model or summary_model + or ConfigLoader().load().summary_model) + if not _llm_backend.get() and _openai_missing_keys(model): + raise PageIndexAPIError( + "optimize='full' runs LLM expand and no LLM key is " + "configured — set OPENAI_API_KEY, or pass optimize='merge' " + "or optimize=False for the LLM-free tree.") result = extract_toc(_validate_pdf(pdf), use_embedded_toc=use_embedded_toc) structure = result.get("structure", []) if optimize and structure: diff --git a/pageindex/flash/embedded_toc.py b/pageindex/flash/embedded_toc.py index 6bd38704d..285d7dcbb 100644 --- a/pageindex/flash/embedded_toc.py +++ b/pageindex/flash/embedded_toc.py @@ -123,13 +123,19 @@ def read_bookmarks(doc_handle: Union[str, Path, BytesIO]) -> list[dict]: doc = pdfium.PdfDocument(handle) entries = [] for item in doc.get_toc(): - if item.page_index is None: + if hasattr(item, "get_dest"): + dest = item.get_dest() + page_idx = dest.get_index() if dest else None + title = item.get_title() + else: + page_idx, title = item.page_index, item.title + if page_idx is None: continue - title = (item.title or "").strip() + title = (title or "").strip() if not title: continue entries.append( - {"title": title, "level": item.level + 1, "page": item.page_index + 1} + {"title": title, "level": item.level + 1, "page": page_idx + 1} ) return entries except (pdfium.PdfiumError, OSError, ValueError, TypeError): diff --git a/pageindex/flash/parser_pdfium_charlevel/char_extract.py b/pageindex/flash/parser_pdfium_charlevel/char_extract.py index d55b4620c..56d6288f7 100644 --- a/pageindex/flash/parser_pdfium_charlevel/char_extract.py +++ b/pageindex/flash/parser_pdfium_charlevel/char_extract.py @@ -56,10 +56,21 @@ def _extract_raw_chars(page, text_page) -> tuple[list[dict], list[dict]]: name_cache: dict[bytes, str] = {} raw_chars: list[dict] = [] last_obj: dict | None = None + skip_next = False for index_value in range(count_item): + if skip_next: + skip_next = False + continue codepoint = get_unicode(text_page, index_value) if codepoint < 0: continue + # PDFium returns astral characters (U+10000+) as two UTF-16 surrogate + # code units in consecutive textpage slots. Reassemble before chr(). + if 0xD800 <= codepoint <= 0xDBFF and index_value + 1 < count_item: + low = get_unicode(text_page, index_value + 1) + if 0xDC00 <= low <= 0xDFFF: + codepoint = ((codepoint & 0x3FF) << 10) + (low & 0x3FF) + 0x10000 + skip_next = True # u == 0 (PDFium found no unicode for the glyph) is KEPT as '\x00': # text extraction emits the raw charcode for unmapped codes, so its items # really contain chr(0) for extension-font pieces at code 0, and the diff --git a/pageindex/flash/parser_pdfium_charlevel/geometry.py b/pageindex/flash/parser_pdfium_charlevel/geometry.py index e3a4b072d..436d7dde4 100644 --- a/pageindex/flash/parser_pdfium_charlevel/geometry.py +++ b/pageindex/flash/parser_pdfium_charlevel/geometry.py @@ -6,6 +6,8 @@ import math import pypdfium2.raw as pdfium_c +_get_font_name = getattr(pdfium_c, "FPDFFont_GetBaseFontName", None) or pdfium_c.FPDFFont_GetFontName + def _obj_rotation(value: float, other_item: float, candidate_item: float, reference_item: float) -> int: """Classify a text-object matrix as upright, cardinal rotation, or oblique. Near-cardinal matrices snap to the cardinal bucket; genuinely oblique matrices use the baseline remerge path.""" @@ -128,7 +130,7 @@ def iter_text_objs(parent, anc_mtx, depth): # Snap back to the shortest decimal so knife-edge font-size comparisons # match the content-stream value. fs_eff = float(f"{fs_eff:.6g}") - name = pdfium_c.FPDFFont_GetFontName(font, font_name_buffer, 256) + name = _get_font_name(font, font_name_buffer, 256) font_name = ( bytes(font_name_buffer[:name]).decode("latin-1", errors="replace").rstrip("\x00") if name > 1 else "" diff --git a/pageindex/flash/parser_pdfium_charlevel/pipeline.py b/pageindex/flash/parser_pdfium_charlevel/pipeline.py index 5a4caf9f0..460ac703f 100644 --- a/pageindex/flash/parser_pdfium_charlevel/pipeline.py +++ b/pageindex/flash/parser_pdfium_charlevel/pipeline.py @@ -95,7 +95,7 @@ def _page_pass1(pdf, pdf_doc, page_idx: int, type3_ext: dict, font_map_cache: di # PDFium's output). if raw_chars: _apply_font_unicode( - text_page.raw, raw_chars, objects, show_codes, pdf_doc, + raw_chars, objects, show_codes, pdf_doc, font_map_cache) except Exception: pass diff --git a/pageindex/flash/parser_pdfium_charlevel/unicode_apply.py b/pageindex/flash/parser_pdfium_charlevel/unicode_apply.py index a7490bd20..0ae314236 100644 --- a/pageindex/flash/parser_pdfium_charlevel/unicode_apply.py +++ b/pageindex/flash/parser_pdfium_charlevel/unicode_apply.py @@ -5,7 +5,6 @@ import bisect import difflib from collections import Counter -import pypdfium2.raw as pdfium_c from .text_normalize import _is_whitespace from .font_unicode import _font_unicode_map @@ -16,7 +15,6 @@ def _apply_font_unicode( - text_page, raw_chars: list[dict], objects: list[dict], show_codes: list[tuple[int | None, tuple[int, ...], float]], @@ -287,14 +285,12 @@ def _run_window(window: list[int]) -> None: _synthesize_dropped_glyphs(kept, raw_chars, chars_by_index) return - # Page mode. - seq: list[tuple[int, str]] = [] - char_count = pdfium_c.FPDFText_CountChars(text_page) - for char_index in range(char_count): - if pdfium_c.FPDFText_IsGenerated(text_page, char_index) == 1: - continue - codepoint = pdfium_c.FPDFText_GetUnicode(text_page, char_index) - seq.append((char_index, chr(codepoint) if codepoint > 0 else "\x00")) + # Page mode. Walk the char census char_extract built (surrogate pairs + # already merged there): re-reading the textpage would split astral + # chars back into two lone-surrogate slots, desync the walk against + # their one-char cmap targets, and drop the whole page's patch. + seq = [(raw_char["i"], raw_char["ch"]) + for raw_char in raw_chars if not raw_char["is_gen"]] targets: list[str] = [] for font_index, encoded_text, _tz in show_codes: if not encoded_text: diff --git a/pageindex/flash/parser_pdfium_parallel.py b/pageindex/flash/parser_pdfium_parallel.py index f0d87b287..ec4c8dbc9 100644 --- a/pageindex/flash/parser_pdfium_parallel.py +++ b/pageindex/flash/parser_pdfium_parallel.py @@ -23,7 +23,9 @@ import multiprocessing import os +import sys from concurrent.futures import ProcessPoolExecutor +from contextlib import contextmanager from io import BytesIO from pathlib import Path from typing import Union @@ -54,6 +56,29 @@ class _Type3Detected(Exception): _worker_font_maps: dict = {} +@contextmanager +def _anonymous_main(): + """Hide __main__'s import identity while workers spawn: spawn re-executes + the caller's script in every worker otherwise, which for an unguarded + script means one duplicate full run per worker. Our workers import + everything by module name and never need __main__. + + ponytail: window covers the whole map; a concurrent pool spawned from + another thread whose tasks live in __main__ would break during it.""" + main = sys.modules.get("__main__") + if main is None: + yield + return + d = main.__dict__ + saved = {k: d.pop(k) for k in ("__file__", "__spec__") if k in d} + d["__spec__"] = None # get_preparation_data reads it via attribute access + try: + yield + finally: + d.pop("__spec__", None) + d.update(saved) + + def _init_worker(kind: str, payload) -> None: global _worker_pdf, _worker_pdf_doc, _worker_font_maps # Open the document exactly as parse_charlevel_meta does, including @@ -122,12 +147,17 @@ def parse_charlevel_meta_parallel( initargs=src, ) try: - results = list(executor.map(_run_page, range(n_pages))) + with _anonymous_main(): + results = list(executor.map(_run_page, range(n_pages))) except Exception: # _Type3Detected or any worker/pool failure. Cancel what is queued # and rerun sequentially; in-flight pages finish in their workers # and are discarded (separate processes, no shared PDFium state). executor.shutdown(wait=False, cancel_futures=True) + if getattr(multiprocessing.current_process(), "_inheriting", False): + # Spawn child re-importing an unguarded __main__; a sequential rerun + # here would silently duplicate the caller's whole run per worker. + raise return parse_charlevel_meta(doc_handle) executor.shutdown() diff --git a/pageindex/integrations/claude_agent_sdk.py b/pageindex/integrations/claude_agent_sdk.py index 8c76cb434..b3e67947d 100644 --- a/pageindex/integrations/claude_agent_sdk.py +++ b/pageindex/integrations/claude_agent_sdk.py @@ -34,14 +34,11 @@ def build_claude_mcp(client, include_management: bool = False, doc_ids=None): "as_claude_mcp in local mode requires the Claude Agent SDK — " "pip install claude-agent-sdk (or pip install 'pageindex[claude]')." ) from exc - from ..agent_tools import (TOOL_CONTRACT, _local_description, - _local_schema, call_tool, tool_names) + from ..agent_tools import TOOL_CONTRACT, _tool_specs - def make_handler(name: str): + def make_handler(invoke): async def handler(arguments: dict[str, Any]) -> dict[str, Any]: - text, is_error = await asyncio.to_thread( - call_tool, client, name, arguments or {}, doc_ids - ) + text, is_error = await asyncio.to_thread(invoke, arguments or {}) result: dict[str, Any] = {"content": [{"type": "text", "text": text}]} if is_error: result["is_error"] = True @@ -59,9 +56,10 @@ def tool_kwargs(name: str) -> dict: return {"annotations": ToolAnnotations(**annotations)} tools = [ - tool(name, _local_description(name), - _local_schema(name), **tool_kwargs(name))(make_handler(name)) - for name in tool_names(include_management) + tool(name, description, schema, + **tool_kwargs(name))(make_handler(invoke)) + for name, description, schema, invoke + in _tool_specs(client, include_management, doc_ids) ] return create_sdk_mcp_server(name="pageindex", version=sdk_version(), tools=tools) diff --git a/pageindex/integrations/openai_agents.py b/pageindex/integrations/openai_agents.py index 4c8ca8ddc..862a35f0b 100644 --- a/pageindex/integrations/openai_agents.py +++ b/pageindex/integrations/openai_agents.py @@ -65,6 +65,8 @@ async def on_invoke_tool(ctx: Any, args_json: str) -> str: return _dumps(payload) arguments = {key: value for key, value in parsed.items() if value is not None} + # is_error has no per-result channel on hand-built FunctionTools + # (raising aborts the run — see above); the text is the signal. text, _ = await asyncio.to_thread(invoke, arguments) return text diff --git a/pageindex/local_api.py b/pageindex/local_api.py index 88ad02ed0..f3c6afa65 100644 --- a/pageindex/local_api.py +++ b/pageindex/local_api.py @@ -1,17 +1,17 @@ """Local implementation of the PageIndex SDK surface.""" from __future__ import annotations -import asyncio import json import logging +import multiprocessing import os import uuid -from concurrent.futures import ThreadPoolExecutor from datetime import datetime, timezone from typing import Any from .errors import PageIndexAPIError from .local_store import DocStore +from .utils import run_off_loop logger = logging.getLogger(__name__) @@ -22,31 +22,23 @@ def _now_iso() -> str: return now.replace(microsecond=now.microsecond // 1000 * 1000).isoformat() -def _run_indexer(func, *args, **kwargs): - try: - asyncio.get_running_loop() - except RuntimeError: - return func(*args, **kwargs) - with ThreadPoolExecutor(max_workers=1) as pool: - return pool.submit(func, *args, **kwargs).result() class LocalAPI: """Backs PageIndexClient's local mode. One instance per client.""" def __init__(self, storage_path: str, model: str, summary_model: str, - retrieve_model: str, index_backend: dict | None = None): + index_backend: dict | None = None): self._store = DocStore(storage_path) self._model = model self._summary_model = summary_model - self._retrieve_model = retrieve_model self._index_backend = index_backend from .utils import ConfigLoader self._config_loader = ConfigLoader() def _with_backend(self, func, *args): """Scope the indexing lane's connection overrides around one - operation — runs inside whatever thread _run_indexer picked.""" + operation — runs inside whatever thread run_off_loop picked.""" from .utils import _llm_backend token = _llm_backend.set(self._index_backend) try: @@ -64,6 +56,12 @@ def submit_document( folder_id: str | None = None, metadata: dict | None = None, ) -> dict[str, Any]: + if getattr(multiprocessing.current_process(), "_inheriting", False): + raise PageIndexAPIError( + "Failed to submit document: called again while a spawned worker " + "process was importing your script. Put your top-level code under " + "if __name__ == '__main__': so worker processes do not re-run it." + ) if beta_headers is not None: raise PageIndexAPIError( "Failed to submit document: beta_headers is not supported in local mode." @@ -114,11 +112,11 @@ def submit_document( try: if mode == "flash": - structure, description = _run_indexer( - self._with_backend, self._index_flash, file_path, page_texts + structure, description = run_off_loop( + self._with_backend, self._index_flash, file_path ) else: - structure, description = _run_indexer( + structure, description = run_off_loop( self._with_backend, self._index_standard, file_path, page_texts ) @@ -128,22 +126,25 @@ def submit_document( raise PageIndexAPIError(f"Failed to submit document: {e}") from e doc_id = "pi-" + uuid.uuid4().hex - meta = { - "id": doc_id, - "name": self._unique_doc_name(os.path.basename(file_path)), - "description": description, - "status": "completed", - "createdAt": _now_iso(), - "pageNum": len(page_texts), - "folderId": None, - "metadata": metadata, - "mode": mode, - } pages = [{"page_index": i + 1, "markdown": text} for i, text in enumerate(page_texts)] from .utils import remove_fields - self._store.save_document( - doc_id, meta, remove_fields(structure, fields=["text"]), pages) + # Check-then-write under the store lock; the early pre-check above + # is advisory only. + with self._store.lock(): + meta = { + "id": doc_id, + "name": self._unique_doc_name(os.path.basename(file_path)), + "description": description, + "status": "completed", + "createdAt": _now_iso(), + "pageNum": len(page_texts), + "folderId": None, + "metadata": metadata, + "mode": mode, + } + self._store.save_document( + doc_id, meta, remove_fields(structure, fields=["text"]), pages) return {"doc_id": doc_id, "name": meta["name"]} def _unique_doc_name(self, name: str) -> str: @@ -190,17 +191,10 @@ def _index_standard(self, file_path: str, page_texts: list[str]) -> tuple[list, ) return structure, result.get("doc_description") - def _index_flash(self, file_path: str, page_texts: list[str]) -> tuple[list, str | None]: + def _index_flash(self, file_path: str) -> tuple[list, str | None]: from .flash import page_index_flash - from .utils import (add_node_text, create_clean_structure_for_description, + from .utils import (create_clean_structure_for_description, generate_doc_description, write_node_id) - import litellm - if not self._index_backend: - env = litellm.validate_environment(self._summary_model) - if not env["keys_in_environment"]: - raise PageIndexAPIError( - f"Failed to submit document: missing API key for " - f"{self._summary_model}: {', '.join(env['missing_keys'])}") result = page_index_flash(file_path, summary=True, summary_model=self._summary_model, optimize="full", @@ -212,7 +206,6 @@ def _index_flash(self, file_path: str, page_texts: list[str]) -> tuple[list, str "a structure from this PDF." ) write_node_id(structure) - add_node_text(structure, [(text, 0) for text in page_texts]) description = generate_doc_description( create_clean_structure_for_description(structure), model=self._summary_model, @@ -356,6 +349,8 @@ def _format_tree_node(node: dict, node_summary: bool) -> dict: "node_id": node.get("node_id"), "page_index": node.get("start_index"), } + if node.get("key_items"): + out["key_items"] = node["key_items"] if node_summary: summary = node.get("summary") if summary is not None: diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index 5b9bfb406..e4dcadef3 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -1,27 +1,7 @@ -"""Managed local chat: document-QA agents over the local tools. - -Three methods, three backend protocols, routed 1:1: ``chat_completions`` -drives the backend's /chat/completions (any OpenAI-compatible backend, -final answer only), ``responses`` drives /responses (official-shape -envelope; the full process transcript rides in ``items`` — round-trip it -for provider prompt-cache continuation and agent memory), ``messages`` -drives Anthropic's /v1/messages via the SDK's own tool runner -(tool_use/tool_result round-trip is the format's native behavior). - -Content passes through untouched — the caller's messages, the model's -answers, tool outputs. Native stop reasons pass through on ``messages``; -the OpenAI engine's abstraction does not surface per-turn finish reasons, -so ``chat_completions`` reports loop completion as ``"stop"``, while -``responses`` reports the backend's terminal ``status`` where the wire -surfaces one (recorded at the transport layer — the framework discards -it). The SDK owns gatekeeping (structural validation), table-setting -(managed instructions, tools, doc targeting), tool execution, and billing -(usage aggregation, envelope ids). -""" +"""Managed local chat: document-QA agents over the local tools.""" from __future__ import annotations import asyncio -import concurrent.futures import hashlib import json import os @@ -52,17 +32,6 @@ def _doc_block(client, doc_id) -> Optional[str]: if not isinstance(doc_id, (str, list)): raise PageIndexAPIError("doc_id must be a string or a list of " "strings.") - doc_ids = [doc_id] if isinstance(doc_id, str) else list(doc_id) - missing = [] - for one_id in doc_ids: - try: - client.get_document(one_id) - except PageIndexAPIError: - missing.append(str(one_id)) - if missing: - raise PageIndexAPIError( - "Documents not found or access denied: " + ", ".join(missing) - ) # scoped: the chat surfaces also pass doc_id into the tool layer, so # name resolution happens inside the allowlist — only a duplicate name # within the targeted set shadows. @@ -118,16 +87,8 @@ def _split_chat_messages(messages) -> "tuple[list[str], list[dict]]": def _run_sync(coro): - try: - asyncio.get_running_loop() - except RuntimeError: - has_loop = False - else: - has_loop = True - if not has_loop: - return asyncio.run(coro) - with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: - return pool.submit(asyncio.run, coro).result() + from .utils import run_off_loop + return run_off_loop(asyncio.run, coro) _SENTINEL = object() @@ -212,25 +173,20 @@ def _require_openai_agents(method: str) -> None: ) from exc +def _sdk_backend(backend) -> dict: + """chat_backend for an SDK constructor: LiteLLM takes either endpoint + spelling, the openai and anthropic SDKs only ``base_url``.""" + return {("base_url" if key == "api_base" else key): value + for key, value in (backend or {}).items()} + + def _openai_model(protocol: str, model_name: str, backend=None): - """The backend protocol driver — the seam tests replace with a fake. - - chat protocol: LiteLLM, full stop — model names mean what LiteLLM says - they mean. Bare names are OpenAI-compatible shorthand (wire form - ``openai/``, so OPENAI_API_KEY / OPENAI_BASE_URL keep selecting - the backend), a ``litellm/`` prefix strips, and a first segment LiteLLM - does not know (a HuggingFace repo id like ``Qwen/...``) is refused with - the ``openai/`` form instead of failing inside LiteLLM at request time. - - responses protocol: the Responses API is OpenAI-SDK native — LiteLLM's - completion surface speaks the chat.completions format, so - provider-prefixed models are refused instead of silently downgrading; - bare and ``openai/`` names drive the OpenAI SDK.""" + """The backend protocol driver — the seam tests replace with a fake.""" if protocol == "responses": + model_name = model_name.removeprefix("litellm/") if "/" in model_name and not model_name.startswith("openai/"): raise PageIndexAPIError( - f"responses() cannot drive " - f"'{model_name.removeprefix('litellm/')}': provider-prefixed " + f"responses() cannot drive '{model_name}': provider-prefixed " "models route through LiteLLM, which speaks chat.completions, " "not the Responses API. Use chat_completions() (or messages() " "for Anthropic models), or point OPENAI_BASE_URL at a " @@ -240,10 +196,12 @@ def _openai_model(protocol: str, model_name: str, backend=None): import openai model_name = model_name.removeprefix("openai/") try: - sdk_client = openai.AsyncOpenAI(**(backend or {})) + sdk_client = openai.AsyncOpenAI(**_sdk_backend(backend)) except (openai.OpenAIError, TypeError) as exc: raise PageIndexAPIError( f"The OpenAI backend is not configured: {exc}") from exc + # A caller-owned transport must survive the per-call close. + sdk_client._pageindex_caller_http = "http_client" in (backend or {}) from agents.models.openai_responses import OpenAIResponsesModel return OpenAIResponsesModel(model_name, openai_client=sdk_client) try: @@ -254,30 +212,20 @@ def _openai_model(protocol: str, model_name: str, backend=None): f"'{model_name}' routes through LiteLLM, but litellm is not " "installed. Run: pip install 'litellm>=1.97'" ) - from .utils import _repair_litellm_types + from .utils import _litellm_model, _repair_litellm_types _repair_litellm_types() - wire = model_name.removeprefix("litellm/") - if "/" not in wire or wire.startswith("openai/"): - if (not os.environ.get("OPENAI_API_KEY") - and not (backend or {}).get("api_key")): - raise PageIndexAPIError( - "The OpenAI backend is not configured: set the " - "OPENAI_API_KEY environment variable, pass an api_key " - "in chat_backend / backend (any value works for keyless " - "OPENAI_BASE_URL servers), or point chat_model at " - "another provider (e.g. 'anthropic/...')." - ) - if "/" not in wire: - wire = f"openai/{wire}" - providers = getattr(litellm, "provider_list", None) - if providers and wire.split("/", 1)[0] not in providers: + try: + wire = _litellm_model(model_name, backend) + except litellm.AuthenticationError as exc: raise PageIndexAPIError( - f"'{wire}' routes through LiteLLM, but " - f"'{wire.split('/', 1)[0]}' is not a LiteLLM provider. For an " - "OpenAI-compatible server (vLLM, TGI, Ollama) serving this " - f"model id, use 'openai/{wire}' and point OPENAI_BASE_URL " - "at the server." - ) + "The OpenAI backend is not configured: set the " + "OPENAI_API_KEY environment variable, pass an api_key " + "in chat_backend / backend (any value works for keyless " + "OPENAI_BASE_URL servers), or point chat_model at " + "another provider (e.g. 'anthropic/...')." + ) from exc + except litellm.NotFoundError as exc: + raise PageIndexAPIError(str(exc)) from exc return LitellmModel(wire, api_key=(backend or {}).get("api_key"), base_url=(backend or {}).get("base_url")) @@ -290,11 +238,11 @@ def _reported_model(model_name: str) -> str: def _cache_extra_args(model_name: str) -> Optional[dict]: """Claude's prompt caching is opt-in per request: on Claude models routed through LiteLLM (Anthropic direct, Bedrock, Vertex — each - channel live-verified), mark the managed system prefix via LiteLLM's - injection param so the loop's later turns and a conversation's next - calls read it instead of repaying full price. Provider resolution is - LiteLLM's own, so this predicate can never disagree with where the - request actually routes.""" + channel live-verified), mark the managed system prefix and the newest + message via LiteLLM's injection param so the loop's later turns and a + conversation's next calls read them instead of repaying full price. + Provider resolution is LiteLLM's own, so this predicate can never + disagree with where the request actually routes.""" if "/" not in model_name or model_name.startswith("openai/"): return None try: @@ -305,11 +253,34 @@ def _cache_extra_args(model_name: str) -> Optional[dict]: return None if provider == "anthropic" or (provider in ("bedrock", "vertex_ai") and "claude" in model.lower()): + # The pair LiteLLM itself seeds for Anthropic and Bedrock: the + # stable prefix plus the newest message, so each turn re-reads + # the turns before it. Passing it explicitly extends it to Vertex. return {"cache_control_injection_points": [ - {"location": "message", "role": "system"}]} + {"location": "message", "role": "system"}, + {"location": "message", "index": -1}]} return None +def _openai_protocol(model_name: str) -> bool: + """Destinations that speak the OpenAI protocol on the wire, where + prompt_cache_key means something and extra_body lands in the request + body. Resolution is LiteLLM's own (same as _cache_extra_args), so the + answer follows actual routing; azure and openrouter ride the OpenAI + protocol without appearing in openai_compatible_providers.""" + wire = model_name.removeprefix("litellm/") + if "/" not in wire or wire.startswith("openai/"): + return True + try: + import litellm + _, provider, _, _ = litellm.get_llm_provider(model=wire) + except Exception: + return False + return (provider in ("openai", "azure", "openrouter") + or provider in getattr(litellm, "openai_compatible_providers", + ())) + + def _merged_backend(client, backend): """This call's connection overrides: the client's ``chat_backend`` under the per-call dict, per-call keys winning.""" @@ -329,8 +300,7 @@ def _openai_agent(client, protocol: str, model_name: str, instructions: str, # and both OpenAI model classes pass extra_body through verbatim. OpenAI # destinations only — LiteLLM plants extra_body as a literal field in # other providers' request bodies, which Anthropic rejects as unknown. - wire = model_name.removeprefix("litellm/") - openai_backend = "/" not in wire or wire.startswith("openai/") + openai_backend = _openai_protocol(model_name) # Chat-lane effort rides extra_args: LiteLLM takes it as its own # top-level kwarg on every supported openai-agents version, and the # channel admits values outside the OpenAI enum ("none"). @@ -338,12 +308,12 @@ def _openai_agent(client, protocol: str, model_name: str, instructions: str, if reasoning_effort is not None: extra_args = {**(extra_args or {}), "reasoning_effort": reasoning_effort} - conn = dict(backend) if backend else {} + conn = _sdk_backend(backend) if backend else {} if conn and protocol == "chat": # LiteLLM takes connection params per call, except the two names # LitellmModel pins as its own keywords — those ride its constructor. lifted = {"api_key": conn.pop("api_key", None), - "base_url": conn.pop("base_url", conn.pop("api_base", None))} + "base_url": conn.pop("base_url", None)} if conn: extra_args = {**(extra_args or {}), **conn} conn = {key: value for key, value in lifted.items() @@ -365,6 +335,9 @@ def _openai_agent(client, protocol: str, model_name: str, instructions: str, model_settings=ModelSettings( temperature=temperature, top_p=top_p, max_tokens=max_tokens, reasoning=reasoning, + # Streamed runs otherwise carry no usage at all (agents forwards + # this as stream_options only on streaming calls). + include_usage=True, extra_body=body, extra_headers=extra_headers, extra_args=extra_args), @@ -377,17 +350,21 @@ def _validate_max_turns(max_turns) -> None: raise PageIndexAPIError("max_turns must be a positive integer.") -def _conversation_cache_key(model_name: str, instructions: str, items) -> str: +def _conversation_cache_key(model_name: str, instructions: str, doc_id, + items) -> str: """Stable per-conversation cache-routing key, sent as the OpenAI ``prompt_cache_key`` through ModelSettings.extra_body (openai-agents 0.20 no longer derives it from RunConfig.group_id — verified against a captured wire). Keyed on the prefix identity — model, instructions, - first conversation item — so a conversation's continuations share one - route without pooling unrelated conversations. Callers pass the - conversation's own items, never the SDK-prepended doc-targeting block: - that block is byte-identical for every conversation about a document - and would pool them all under one key.""" - seed = json.dumps([model_name, instructions, + doc targeting, first conversation item — so a conversation's + continuations share one route without pooling unrelated conversations. + Callers pass the conversation's own items, never the SDK-prepended + doc-targeting block: that block is byte-identical for every + conversation about a document and would pool them all under one key. + doc_id carries the targeting identity instead — the same opening + question against different documents is different conversations.""" + scope = [doc_id] if isinstance(doc_id, str) else doc_id + seed = json.dumps([model_name, instructions, scope, items[0] if items else None], sort_keys=True, default=str) return "pageindex-" + hashlib.sha256(seed.encode()).hexdigest()[:16] @@ -437,16 +414,72 @@ async def recording_create(*args, **kwargs): value = getattr(response, field, None) recorded[field] = (value.model_dump(mode="json") if hasattr(value, "model_dump") else value) + # The backend's echo of what it actually ran with. + for field in ("tool_choice", "parallel_tool_calls"): + value = getattr(response, field, None) + if value is not None: + recorded[field] = (value.model_dump(mode="json") + if hasattr(value, "model_dump") else value) return response responses.create = recording_create +def _record_chat_finish(agent, recorded: dict) -> None: + """Capture each turn's native finish_reason from the raw LiteLLM + response: openai-agents' ModelResponse drops it, so a truncated or + content-filtered final turn would otherwise report as a clean "stop". + The chat protocol has no transport client to hook (cf. + _record_response_status), so this wraps the model's response fetch; + no-op if that private seam moves.""" + model = getattr(agent, "model", None) + fetch = getattr(model, "_fetch_response", None) + if fetch is None: + return + + def note(item) -> None: + choices = getattr(item, "choices", None) + finish = getattr(choices[0], "finish_reason", None) if choices else None + if finish: + recorded["finish_reason"] = finish + + class _Tee: + """Iteration passthrough that notes each chunk; everything else + (aclose/close/...) delegates to the provider stream itself.""" + + def __init__(self, inner): + self._inner = inner + + def __aiter__(self): + return self + + async def __anext__(self): + chunk = await self._inner.__anext__() + note(chunk) + return chunk + + def __getattr__(self, name): + return getattr(self._inner, name) + + async def recording_fetch(*args, **kwargs): + result = await fetch(*args, **kwargs) + if isinstance(result, tuple): + response, stream = result + return response, _Tee(stream) + note(result) + return result + + model._fetch_response = recording_fetch + + async def _aclose_backend(agent) -> None: """Close the per-call AsyncOpenAI client before its event loop ends — otherwise httpx tears down pooled connections on a closed loop and - emits 'Task exception was never retrieved' noise.""" + emits 'Task exception was never retrieved' noise. A client built on a + caller-owned http_client stays open.""" backend = getattr(getattr(agent, "model", None), "_client", None) + if getattr(backend, "_pageindex_caller_http", False): + return close = getattr(backend, "close", None) if close is not None: try: @@ -473,8 +506,10 @@ 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 + if r.usage is None: + continue + prompt += r.usage.input_tokens or 0 + completion += r.usage.output_tokens or 0 details = getattr(r.usage, "input_tokens_details", None) cached += getattr(details, "cached_tokens", 0) or 0 cache_write += getattr(details, "cache_write_tokens", 0) or 0 @@ -532,12 +567,14 @@ def run_chat_completions(client, messages, stream: bool = False, managed = _managed_instructions(system_texts) agent = _openai_agent(client, "chat", model_name, managed, temperature, top_p, doc_ids=doc_id, - cache_key=_conversation_cache_key(model_name, - managed, history), + cache_key=_conversation_cache_key( + model_name, managed, doc_id, history), reasoning_effort=reasoning_effort, extra_body=extra_body, max_tokens=max_tokens, backend=_merged_backend(client, backend), extra_headers=extra_headers) + recorded: dict = {} + _record_chat_finish(agent, recorded) run_kwargs = _run_kwargs(max_turns) import openai from agents import Runner @@ -562,7 +599,7 @@ def run_chat_completions(client, messages, stream: bool = False, "index": 0, "message": {"role": "assistant", "content": result.final_output or ""}, - "finish_reason": "stop", + "finish_reason": recorded.get("finish_reason") or "stop", }], "usage": _openai_usage(result.raw_responses), } @@ -602,7 +639,7 @@ async def agen(): if not completed and hasattr(streamed, "cancel"): streamed.cancel() # abandoned/failed: stop the agent task await _aclose_backend(agent) - yield chunk({}, finish="stop") + yield chunk({}, finish=recorded.get("finish_reason") or "stop") yield { "id": chat_id, "object": "chat.completion.chunk", "created": created, "model": reported_model, "choices": [], @@ -649,8 +686,8 @@ def run_responses(client, input, model: Optional[str] = None, managed = _managed_instructions(extra) agent = _openai_agent(client, "responses", model_name, managed, temperature, top_p, doc_ids=doc_id, - cache_key=_conversation_cache_key(model_name, managed, - conversation), + cache_key=_conversation_cache_key( + model_name, managed, doc_id, conversation), reasoning=reasoning, extra_body=extra_body, max_tokens=max_output_tokens, backend=_merged_backend(client, backend), @@ -661,11 +698,14 @@ def run_responses(client, input, model: Optional[str] = None, from agents import Runner from agents.exceptions import AgentsException, MaxTurnsExceeded + response_id = f"resp_{uuid.uuid4().hex}" + created_at = int(time.time()) + def envelope(transcript: list, raw_responses) -> dict: return { - "id": f"resp_{uuid.uuid4().hex}", + "id": response_id, "object": "response", - "created_at": int(time.time()), + "created_at": created_at, "model": _reported_model(model_name), "status": recorded.get("status") or "completed", "output": [item for item in transcript @@ -678,8 +718,9 @@ def envelope(transcript: list, raw_responses) -> dict: "parameters": tool.params_json_schema, "strict": getattr(tool, "strict_json_schema", True)} for tool in agent.tools], - "tool_choice": "auto", - "parallel_tool_calls": True, + # Backend echo when captured; the request sends neither param. + "tool_choice": recorded.get("tool_choice", "auto"), + "parallel_tool_calls": recorded.get("parallel_tool_calls", True), "temperature": temperature, "top_p": top_p, "reasoning": reasoning, @@ -721,11 +762,24 @@ async def agen(): # re-based by the count of items already committed by prior turns. output_offset = 0 completed = False + opened = False try: async for event in streamed.stream_events(): if event.type == "raw_response_event": data = event.data.model_dump(exclude_unset=True) if data.get("type") in lifecycle: + if data["type"] == "response.created" and not opened: + # N per-turn openings collapse to one, carrying + # the id and created_at the terminal event will + # report. + opened = True + if data.get("response"): + data["response"]["id"] = response_id + data["response"]["created_at"] = created_at + sequence += 1 + data["sequence_number"] = sequence + yield data + continue if data["type"] in ("response.completed", "response.incomplete", "response.failed"): @@ -735,6 +789,10 @@ async def agen(): for field in ("status", "incomplete_details", "error"): recorded[field] = state.get(field) + for field in ("tool_choice", + "parallel_tool_calls"): + if state.get(field) is not None: + recorded[field] = state[field] output_offset += len(state.get("output") or []) continue if isinstance(data.get("output_index"), int): @@ -793,7 +851,7 @@ def _anthropic_client(backend=None): """The backend client — the seam tests replace with a fake transport.""" import anthropic try: - return anthropic.Anthropic(**(backend or {})) + return anthropic.Anthropic(**_sdk_backend(backend)) except TypeError as exc: raise PageIndexAPIError( f"The Anthropic backend is not configured: {exc}") from exc @@ -820,6 +878,18 @@ def _anthropic_system(extra_system, block: Optional[str]) -> list[dict]: raise PageIndexAPIError("system must be a string or a list of blocks.") +def _cache_marks(system_blocks, messages) -> int: + """Breakpoints already on the request. The API allows 4 total; the + top-level moving breakpoint is only added when it fits.""" + blocks = list(system_blocks) + for message in messages: + content = message.get("content") + if isinstance(content, list): + blocks += [b for b in content if isinstance(b, dict)] + return sum(1 for b in blocks + if isinstance(b, dict) and b.get("cache_control")) + + def _dump_block(block) -> Any: """A content block as a plain JSON dict, minus SDK-internal fields the API rejects (ParsedBetaTextBlock.__api_exclude__, e.g. parsed_output).""" @@ -855,9 +925,15 @@ def _anthropic_usage(turns, final_usage: dict) -> dict: "claude-3-5-sonnet-20240620") -def _default_max_tokens(model: str) -> int: +def _default_max_tokens(model: str, thinking=None) -> int: """The wire-required per-turn budget when the caller sets none: 8192, - except the claude-3 generation whose output ceiling is 4096.""" + except the claude-3 generation whose output ceiling is 4096. The wire + also requires max_tokens > thinking.budget_tokens, so an enabled + budget lifts the default above itself.""" + budget = (thinking.get("budget_tokens") + if isinstance(thinking, dict) else None) + if isinstance(budget, int): + return budget + 8192 return 4096 if model.startswith(_CLAUDE_4096_MODELS) else 8192 @@ -892,18 +968,38 @@ def run_messages(client, messages, model: str, "stop_sequences": stop_sequences, "thinking": thinking, "extra_body": extra_body, "extra_headers": extra_headers, }.items() if value is not None} - runner = _anthropic_client(_merged_backend(client, backend)) \ - .beta.messages.tool_runner( - max_tokens=(max_tokens if max_tokens is not None - else _default_max_tokens(model)), + system_blocks = _anthropic_system(system, block) + # Top-level cache_control: the server re-marks the newest block each + # turn, so the loop re-reads the growing conversation from cache. + # Counts toward the 4-breakpoint limit (live-verified 400 past it). + cached: dict[str, Any] = ( + {"cache_control": {"type": "ephemeral"}} + if _cache_marks(system_blocks, prepared) < 4 else {}) + merged = _merged_backend(client, backend) + # The SDK defers credential resolution to request time and raises a + # bare TypeError there — pre-check for the contract's PageIndexAPIError. + if not merged and not (os.environ.get("ANTHROPIC_API_KEY") + or os.environ.get("ANTHROPIC_AUTH_TOKEN")): + raise PageIndexAPIError( + "The Anthropic backend is not configured: set the " + "ANTHROPIC_API_KEY environment variable, or pass an api_key " + "in chat_backend / backend.") + backend_client = _anthropic_client(merged) + # A caller-owned http_client must survive the per-call closes below. + owns_transport = "http_client" not in (merged or {}) + if max_tokens is None: + max_tokens = _default_max_tokens(model, thinking) + runner = backend_client.beta.messages.tool_runner( + max_tokens=max_tokens, messages=prepared, model=model, tools=build_anthropic_tools(client, doc_ids=doc_id), - system=_anthropic_system(system, block), + system=system_blocks, stream=stream, # Bounded like the OpenAI surfaces (their framework default is 10). max_iterations=max_turns if max_turns is not None else 10, **passthrough, + **cached, ) if stream: @@ -915,6 +1011,10 @@ def events() -> Iterator[Any]: except anthropic.AnthropicError as exc: raise PageIndexAPIError( f"The model backend failed: {exc}") from exc + finally: + # runs on exhaustion and abandonment (GeneratorExit) alike + if owns_transport: + backend_client.close() return events() try: @@ -922,6 +1022,10 @@ def events() -> Iterator[Any]: except anthropic.AnthropicError as exc: raise PageIndexAPIError( f"The model backend failed: {exc}") from exc + finally: + # safe here: the params read-back below does no HTTP + if owns_transport: + backend_client.close() if not turns: raise PageIndexAPIError("The model returned no response.") captured: dict = {} @@ -955,7 +1059,7 @@ def capture(params): new_messages = [_dump_message(message) for message in conversation[len(prepared):]] final_blocks = [_dump_block(item) for item in final.content] - final_ids = {block["id"] for block in final_blocks + final_ids = {block.get("id") for block in final_blocks if block.get("type") == "tool_use"} history_ids = {block.get("id") for message in new_messages diff --git a/pageindex/local_store.py b/pageindex/local_store.py index 37108f93c..cdba19b71 100644 --- a/pageindex/local_store.py +++ b/pageindex/local_store.py @@ -6,6 +6,7 @@ import os import shutil import uuid +from contextlib import contextmanager from pathlib import Path logger = logging.getLogger(__name__) @@ -89,6 +90,24 @@ def _write_manifest(self, docs: dict) -> None: except OSError: pass + @contextmanager + def lock(self): + """Cross-process mutex for check-then-write sequences (name + uniquing before save). fcntl is absent on Windows, where the + pre-existing best-effort behavior stays.""" + try: + import fcntl + except ImportError: + yield + return + self._root.mkdir(parents=True, exist_ok=True) + with open(self._root / ".lock", "w") as handle: + fcntl.flock(handle, fcntl.LOCK_EX) + try: + yield + finally: + fcntl.flock(handle, fcntl.LOCK_UN) + # ── documents ── def save_document(self, doc_id: str, meta: dict, tree: list, pages: list) -> None: doc_dir = self._doc_dir(doc_id) diff --git a/pageindex/mcp_bridge.py b/pageindex/mcp_bridge.py index 7d8d153c2..32ec09884 100644 --- a/pageindex/mcp_bridge.py +++ b/pageindex/mcp_bridge.py @@ -44,6 +44,7 @@ class McpBridge: def __init__(self, url: str, headers: dict[str, str]): self._url = url self._auth_headers = dict(headers) + self._session = requests.Session() # agent tool calls come in bursts self._session_id: Optional[str] = None self._protocol_version: Optional[str] = None self._instructions: Optional[str] = None @@ -65,8 +66,8 @@ def _post(self, payload: dict, session_id: Optional[str] = None, if protocol_version: headers["MCP-Protocol-Version"] = protocol_version try: - return requests.post(self._url, json=payload, headers=headers, - timeout=_TIMEOUT) + return self._session.post(self._url, json=payload, + headers=headers, timeout=_TIMEOUT) except requests.RequestException as exc: raise PageIndexAPIError( f"Could not reach the PageIndex MCP server: {exc}" @@ -85,7 +86,8 @@ def _extract_result(self, response: requests.Response, request_id: int) -> Any: except ValueError as exc: raise PageIndexAPIError( f"MCP server returned a non-JSON response " - f"(HTTP {response.status_code})." + f"(HTTP {response.status_code}).", + status_code=response.status_code, ) from exc # Strict id correlation only — accepting any result-bearing message # would return a stale or mis-correlated reply as this call's. @@ -129,7 +131,8 @@ def _request(self, method: str, params: Optional[dict] = None, if response.status_code >= 400: raise PageIndexAPIError( f"MCP request failed: HTTP {response.status_code} " - f"({response.text[:200]})" + f"({response.text[:200]})", + status_code=response.status_code, ) return self._extract_result(response, request_id) @@ -152,7 +155,8 @@ def _ensure_initialized(self) -> None: raise PageIndexAPIError( f"Could not connect to the PageIndex MCP server: HTTP " f"{response.status_code} ({response.text[:200]}). Check " - "your API key." + "your API key.", + status_code=response.status_code, ) result = self._extract_result(response, request_id) or {} self._session_id = response.headers.get("Mcp-Session-Id") @@ -179,13 +183,18 @@ def instructions(self) -> Optional[str]: def list_tools(self) -> list[dict]: tools: list[dict] = [] cursor: Optional[str] = None - while True: + # A server echoing its cursor (or cycling) must not hang the client: + # no-progress terminates, the page cap turns a cycle into an error. + for _ in range(50): params = {"cursor": cursor} if cursor else {} result = self._request("tools/list", params) or {} tools.extend(result.get("tools") or []) - cursor = result.get("nextCursor") - if not cursor: + next_cursor = result.get("nextCursor") + if not next_cursor or next_cursor == cursor: return tools + cursor = next_cursor + raise PageIndexAPIError( + "MCP tools/list pagination did not terminate within 50 pages.") def call_tool(self, name: str, arguments: dict[str, Any]) -> "tuple[str, bool]": """Returns (text, is_error) — is_error is the server's MCP isError @@ -205,6 +214,14 @@ def call_tool(self, name: str, arguments: dict[str, Any]) -> "tuple[str, bool]": kind = block.get("mimeType") or block.get("type") or "binary" size_kb = max(1, len(block["data"]) * 3 // 4096) texts.append(f"[{kind} content omitted: ~{size_kb} KB]") + elif (isinstance(block, dict) + and isinstance(block.get("resource"), dict) + and isinstance(block["resource"].get("blob"), str)): + # EmbeddedResource nests its base64 one level down. + resource = block["resource"] + kind = resource.get("mimeType") or "binary" + size_kb = max(1, len(resource["blob"]) * 3 // 4096) + texts.append(f"[{kind} content omitted: ~{size_kb} KB]") else: texts.append(json.dumps(block, ensure_ascii=False)) return "\n".join(texts), is_error diff --git a/pageindex/tree_optimize.py b/pageindex/tree_optimize.py index 04719ccb2..67d2e4728 100644 --- a/pageindex/tree_optimize.py +++ b/pageindex/tree_optimize.py @@ -61,7 +61,7 @@ import sys from types import SimpleNamespace -from .utils import (ConfigLoader, _is_openai_model, _is_unrecoverable, +from .utils import (ConfigLoader, _is_unrecoverable, _openai_missing_keys, llm_acompletion, strip_internal_keys) TRIGGER_PAGES = 5 # only look ahead on nodes larger than this @@ -872,9 +872,11 @@ async def main(): args = parser.parse_args() model = args.model or default_model() - if args.expand and not args.plan and _is_openai_model(model) \ - and not os.getenv("OPENAI_API_KEY"): - sys.exit(f"OPENAI_API_KEY is not set (expand model: {model}).") + if args.expand and not args.plan: + missing = _openai_missing_keys(model) + if missing: + sys.exit(f"{', '.join(missing)} is not set " + f"(expand model: {model}).") original = json.load(open(args.structure)) structure = copy.deepcopy(original["structure"]) diff --git a/pageindex/utils.py b/pageindex/utils.py index 114f04668..dbeddb56d 100644 --- a/pageindex/utils.py +++ b/pageindex/utils.py @@ -29,13 +29,6 @@ "pageindex_llm_backend", default=None) -def _openai_sdk_kwargs(backend: dict) -> dict: - """The same backend dict works on both gateway paths: LiteLLM accepts - either endpoint spelling, the openai SDK only ``base_url``.""" - return {("base_url" if key == "api_base" else key): value - for key, value in backend.items()} - - def _repair_litellm_types() -> None: """litellm 1.97.0's Message/Delta annotations carry nested forward refs Python 3.10 cannot resolve (BerriAI/litellm#36384), so every completion @@ -54,6 +47,9 @@ def _repair_litellm_types() -> None: # Backward compatibility: support CHATGPT_API_KEY as alias for OPENAI_API_KEY if not os.getenv("OPENAI_API_KEY") and os.getenv("CHATGPT_API_KEY"): + import warnings + warnings.warn("CHATGPT_API_KEY is deprecated — set OPENAI_API_KEY " + "instead.", FutureWarning) os.environ["OPENAI_API_KEY"] = os.getenv("CHATGPT_API_KEY") def count_tokens(text, model=None): @@ -69,16 +65,67 @@ def _strip_prefix(s, prefix): return s -def _is_openai_model(model): - """Models without a provider prefix (no '/') use the openai SDK directly. - For other providers, use 'provider/model' format (e.g. 'anthropic/claude-sonnet-4-6').""" - if not model or model.startswith('litellm/'): - return False - return '/' not in model or model.startswith('openai/') - - -_openai_sync_client = None -_openai_async_client = None +def run_off_loop(func, *args): + """Run func now, or on a worker thread when this thread already runs an + asyncio loop (func may itself call asyncio.run).""" + try: + asyncio.get_running_loop() + except RuntimeError: + return func(*args) + from concurrent.futures import ThreadPoolExecutor + with ThreadPoolExecutor(max_workers=1) as pool: + return pool.submit(func, *args).result() + + +def _openai_missing_keys(model): + """Missing env keys for the pre-check, which covers only OpenAI-shaped + names (bare or ``openai/``): other providers resolve credentials their + own way at call time (IAM chains, ADC, Ollama's localhost default), + invisible to env inspection — the chat lane draws the same line. + ``litellm/``-prefixed names are exempt: the prefix is an explicit + routing choice, and litellm resolves credentials beyond the + environment (litellm.api_key, a keyless OPENAI_BASE_URL server). + Truthiness, not litellm's validate_environment, which reports a blank + exported key as present.""" + if model.startswith("litellm/"): + return [] + if "/" in model and not model.startswith("openai/"): + return [] + return ([] if (os.getenv("OPENAI_API_KEY") or "").strip() + else ["OPENAI_API_KEY"]) + + +def _litellm_model(model, backend): + """Normalize to LiteLLM's grammar (``litellm/`` strips, bare names get + the ``openai/`` wire form — same as the chat lane) and fail fast on a + missing key or unknown provider, with status codes the retry loop and + the summary/optimize passes treat as unrecoverable.""" + if not model: + return model + raw = model + model = _strip_prefix(model, "litellm/") + if "/" not in model: + model = f"openai/{model}" + import litellm + provider = model.split("/", 1)[0] + providers = getattr(litellm, "provider_list", None) + # custom_provider_map providers join provider_list only at call time. + custom = {entry.get("provider") for entry + in getattr(litellm, "custom_provider_map", None) or []} + if providers and provider not in providers and provider not in custom: + raise litellm.NotFoundError( + f"'{model}' routes through LiteLLM, but '{provider}' is not a " + f"LiteLLM provider. For an OpenAI-compatible server serving " + f"this model id, use 'openai/{model}' and point " + f"OPENAI_BASE_URL at the server.", + llm_provider=None, model=model) + if not backend: + missing = _openai_missing_keys(raw) + if missing: + raise litellm.AuthenticationError( + f"missing API key for {model}: {', '.join(missing)}", + llm_provider=None, model=model) + return model # Misconfiguration: no retry can fix a rejected key or a model that does not @@ -92,42 +139,34 @@ def _is_unrecoverable(exc: Exception) -> bool: return getattr(exc, "status_code", None) in _UNRECOVERABLE_STATUS +def _no_cache_seeding_kwargs(backend): + """litellm 1.97 auto-marks Claude requests for prompt caching (system + + last message); indexing prompts are single-shot and unique, so every call + would pay the cache-write premium with nothing ever read back. A + system-role-only injection point matches no indexing message, and its + presence stops litellm seeding its own defaults; backend keys still + win.""" + return {"cache_control_injection_points": + [{"location": "message", "role": "system"}], + **(backend or {})} + + def llm_completion(model, prompt, chat_history=None, return_finish_reason=False): - use_openai_sdk = _is_openai_model(model) - if model: - model = _strip_prefix(model, "litellm/") - if use_openai_sdk: - model = _strip_prefix(model, "openai/") + import litellm max_retries = 10 messages = list(chat_history) + [{"role": "user", "content": prompt}] if chat_history else [{"role": "user", "content": prompt}] backend = _llm_backend.get() - if use_openai_sdk: - import openai - if backend: - oai_client = openai.OpenAI(**{"max_retries": 0, - **_openai_sdk_kwargs(backend)}) - else: - global _openai_sync_client - if _openai_sync_client is None: - _openai_sync_client = openai.OpenAI(max_retries=0) - oai_client = _openai_sync_client + model = _litellm_model(model, backend) + _repair_litellm_types() for i in range(max_retries): try: - if use_openai_sdk: - response = oai_client.chat.completions.create( - model=model, - messages=messages, - ) - else: - import litellm - _repair_litellm_types() - response = litellm.completion( - model=model, - messages=messages, - temperature=0, - drop_params=True, - **(backend or {}), - ) + response = litellm.completion( + model=model, + messages=messages, + drop_params=True, + # the loop is the retry policy; the merge lets a backend override win + **{"max_retries": 0, **_no_cache_seeding_kwargs(backend)}, + ) content = response.choices[0].message.content if return_finish_reason: finish_reason = "max_output_reached" if response.choices[0].finish_reason == "length" else "finished" @@ -147,41 +186,20 @@ def llm_completion(model, prompt, chat_history=None, return_finish_reason=False) async def llm_acompletion(model, prompt): - use_openai_sdk = _is_openai_model(model) - if model: - model = _strip_prefix(model, "litellm/") - if use_openai_sdk: - model = _strip_prefix(model, "openai/") + import litellm max_retries = 10 messages = [{"role": "user", "content": prompt}] backend = _llm_backend.get() - if use_openai_sdk: - import openai - if backend: - oai_client = openai.AsyncOpenAI(**{"max_retries": 0, - **_openai_sdk_kwargs(backend)}) - else: - global _openai_async_client - if _openai_async_client is None: - _openai_async_client = openai.AsyncOpenAI(max_retries=0) - oai_client = _openai_async_client + model = _litellm_model(model, backend) + _repair_litellm_types() for i in range(max_retries): try: - if use_openai_sdk: - response = await oai_client.chat.completions.create( - model=model, - messages=messages, - ) - else: - import litellm - _repair_litellm_types() - response = await litellm.acompletion( - model=model, - messages=messages, - temperature=0, - drop_params=True, - **(backend or {}), - ) + response = await litellm.acompletion( + model=model, + messages=messages, + drop_params=True, + **{"max_retries": 0, **_no_cache_seeding_kwargs(backend)}, + ) return response.choices[0].message.content except Exception as e: if _is_unrecoverable(e): @@ -714,6 +732,8 @@ async def generate_summaries_for_structure(structure, model=None): summaries = await asyncio.gather(*tasks, return_exceptions=True) for node, summary in zip(nodes, summaries): + if isinstance(summary, Exception) and _is_unrecoverable(summary): + raise summary node['summary'] = "" if isinstance(summary, BaseException) else summary if nodes and not any(node['summary'] for node in nodes): raise RuntimeError( @@ -887,17 +907,25 @@ async def parent_summary(node): async def visit(node): children = node.get('nodes') or [] if children: - await asyncio.gather(*(visit(child) for child in children), - return_exceptions=True) + done = await asyncio.gather(*(visit(child) for child in children), + return_exceptions=True) + for result in done: + if isinstance(result, Exception) and _is_unrecoverable(result): + raise result if node.get('summary'): return try: node['summary'] = await (parent_summary(node) if children else leaf_summary(node)) - except Exception: + except Exception as e: node['summary'] = "" + if _is_unrecoverable(e): + raise - await asyncio.gather(*(visit(root) for root in structure), - return_exceptions=True) + results = await asyncio.gather(*(visit(root) for root in structure), + return_exceptions=True) + for r in results: + if isinstance(r, Exception) and _is_unrecoverable(r): + raise r def _any_summary(nodes): return any(n.get('summary') or _any_summary(n.get('nodes') or []) diff --git a/pyproject.toml b/pyproject.toml index 499291d3d..f5bec8c9a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,7 +32,7 @@ openai = ">=1.70.0" openai-agents = ">=0.18.1" litellm = ">=1.97.0" PyPDF2 = ">=3.0.0" -pypdfium2 = ">=4.30.0" +pypdfium2 = ">=5" sortedcontainers = ">=2.4.0" regex = ">=2024.0.0" python-dotenv = ">=1.0.0" diff --git a/requirements.txt b/requirements.txt index 3202169ce..2406eef11 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,7 +4,7 @@ requests>=2.28.0 openai-agents>=0.18.1 # pymupdf # optional PyPDF2==3.0.1 -pypdfium2==4.30.0 +pypdfium2==5.13.0 python-dotenv==1.2.2 pyyaml==6.0.2 regex>=2024.0.0 diff --git a/run_pageindex.py b/run_pageindex.py index 1cbea0c72..f2642b8a6 100644 --- a/run_pageindex.py +++ b/run_pageindex.py @@ -3,7 +3,11 @@ import json from pageindex import * from pageindex.page_index_md import md_to_tree -from pageindex.utils import ConfigLoader +from pageindex.utils import ConfigLoader, _openai_missing_keys + +# Keep LiteLLM's import off the network (frozen bundled model-cost map); +# an explicit user setting wins. +os.environ.setdefault("LITELLM_LOCAL_MODEL_COST_MAP", "True") if __name__ == "__main__": # Set up argument parser @@ -62,7 +66,7 @@ raise ValueError("Either --pdf_path or --md_path must be specified") if args.pdf_path and args.md_path: raise ValueError("Only one of --pdf_path or --md_path can be specified") - if args.optimize in ('full', 'merge') and not (args.pdf_path and args.mode == 'flash'): + if args.optimize is not None and not (args.pdf_path and args.mode == 'flash'): raise ValueError("--optimize requires Flash mode with --pdf_path") if args.optimize is None: args.optimize = 'full' if args.mode == 'flash' else 'off' @@ -90,14 +94,15 @@ if args.mode == 'flash': from pageindex.flash import page_index_flash - summary_model = args.summary_model or args.index_model or args.model + summary_model = (args.summary_model or args.index_model + or args.model + or ConfigLoader().load().summary_model) will_summarize = args.summary if args.summary is not None else True - if summary_model and (will_summarize or args.optimize == 'full'): - import litellm - env = litellm.validate_environment(summary_model) - if not env["keys_in_environment"]: + if will_summarize or args.optimize == 'full': + missing = _openai_missing_keys(summary_model) + if missing: raise SystemExit( - f"Missing API key for {summary_model}: {', '.join(env['missing_keys'])}") + f"Missing API key for {summary_model}: {', '.join(missing)}") toc_with_page_number = page_index_flash( args.pdf_path, optimize=args.optimize if args.optimize != 'off' else False, @@ -164,25 +169,24 @@ user_opt = { 'index_model': args.index_model, 'model': args.model, - 'if_add_node_summary': args.if_add_node_summary, - 'if_add_doc_description': args.if_add_doc_description, - 'if_add_node_text': args.if_add_node_text, - 'if_add_node_id': args.if_add_node_id } # Load config with defaults from config.yaml - opt = config_loader.load(user_opt) + opt = config_loader.load({k: v for k, v in user_opt.items() if v is not None}) + # if_add_* pass through as given (absent = off, as before this CLI + # used config.yaml): the PDF defaults there must not switch on LLM + # passes the markdown CLI never ran. toc_with_page_number = asyncio.run(md_to_tree( md_path=args.md_path, if_thinning=args.if_thinning.lower() == 'yes', min_token_threshold=args.thinning_threshold, - if_add_node_summary=opt.if_add_node_summary, + if_add_node_summary=args.if_add_node_summary, summary_token_threshold=args.summary_token_threshold, model=opt.model, - if_add_doc_description=opt.if_add_doc_description, - if_add_node_text=opt.if_add_node_text, - if_add_node_id=opt.if_add_node_id + if_add_doc_description=args.if_add_doc_description, + if_add_node_text=args.if_add_node_text, + if_add_node_id=args.if_add_node_id )) print('Parsing done, saving to file...') diff --git a/tests/conftest.py b/tests/conftest.py index 40a4a9199..04af63745 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -5,6 +5,7 @@ def _llm_key(monkeypatch): """Deterministic key presence for every test; missing-key tests delenv.""" monkeypatch.setenv("OPENAI_API_KEY", "test-key") + monkeypatch.setenv("ANTHROPIC_API_KEY", "test-key") def build_pdf(page_texts): diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 3873cd36e..28324d23d 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -69,7 +69,10 @@ def run(client, name, **arguments): # ── contract parity ── -def test_contract_matches_snapshot(): +def test_contract_edits_are_deliberate(): + """The committed snapshot cannot detect drift from the live cloud + server — both copies live in this repo. It exists so a TOOL_CONTRACT + edit must touch two files in one change, never land by accident.""" snapshot = json.loads(SNAPSHOT_PATH.read_text(encoding="utf-8")) assert snapshot["tools"] == TOOL_CONTRACT @@ -405,9 +408,11 @@ def test_created_at_accepts_z_suffixed_input(client, store_path): def test_page_content_char_budget(client, store_path): + # Escape-dense pages: raw length fits the budget, JSON-serialized + # length does not — the budget must count emitted characters. pages = [ - {"page_index": 1, "markdown": "x" * 96_000}, - {"page_index": 2, "markdown": "short"}, + {"page_index": 1, "markdown": '"' * 30_000}, + {"page_index": 2, "markdown": '"' * 20_000}, ] seed_doc(store_path, "pi-a", "huge.pdf", pages=pages) payload, is_error = run(client, "get_page_content", doc_name="huge.pdf", @@ -417,6 +422,8 @@ def test_page_content_char_budget(client, store_path): assert "size limits" in payload["next_steps"]["summary"] assert any("For remaining pages, request: 2" in option for option in payload["next_steps"]["options"]) + emitted = json.dumps(payload, ensure_ascii=False) + assert len(emitted) <= agent_tools_module.TOOL_RESPONSE_CHAR_LIMIT def test_page_content_reports_truncation_and_out_of_range_together( @@ -763,6 +770,54 @@ def test_openai_agent_config_model_speaks_the_agents_sdk_grammar(tmp_path): retrieve_model="anthropic/claude-x") assert (client.openai_agent_config()["model"] == "litellm/anthropic/claude-x") + # The per-call override speaks the same grammar as chat_model. + assert (client.openai_agent_config(model="anthropic/claude-y")["model"] + == "litellm/anthropic/claude-y") + assert (client.openai_agent_config(model="litellm/groq/llama-x")["model"] + == "litellm/groq/llama-x") + + +def test_plain_functions_answer_bad_arguments_with_the_envelope(client, + store_path): + """agent_tools() functions must not raise into a framework loop: + cloud-only parameters pruned from the local signature come back as + the guided envelope, and the schema-bearing signature survives.""" + import inspect + seed_doc(store_path, "pi-a", "report.pdf") + tools = {f.__name__: f for f in client.agent_tools()} + fn = tools["get_document"] + assert "doc_name" in inspect.signature(fn).parameters + payload = json.loads(fn(doc_name="report.pdf", folder_id="root")) + assert payload["errorCode"] == "INVALID_INPUT" + ok = json.loads(fn(doc_name="report.pdf")) + assert not ok.get("errorCode") + + +def test_non_string_doc_name_stays_not_found(client, store_path): + """A type-loose model argument must not turn NOT_FOUND into the + retry-inviting INTERNAL_ERROR (strict_json_schema is off on the + OpenAI adapter, so nothing upstream validates the type).""" + seed_doc(store_path, "pi-a", "report.pdf") + payload, is_error = run(client, "get_document", doc_name=5) + assert is_error and payload["errorCode"] == "NOT_FOUND" + + +def test_openai_agent_config_repairs_litellm_types(tmp_path, monkeypatch): + """The BYO path resolves its model through LiteLLM in the caller's + process, outside our completion helpers — the py3.10 type repair must + run at config time, and only for LiteLLM-routed models.""" + pytest.importorskip("agents") + import pageindex.utils + calls = [] + monkeypatch.setattr(pageindex.utils, "_repair_litellm_types", + lambda: calls.append(True)) + client = PageIndexLocalClient(storage_path=str(tmp_path / "s"), + chat_model="anthropic/claude-x") + client.openai_agent_config() + assert calls + calls.clear() + client.openai_agent_config(model="gpt-plain") + assert not calls def test_openai_agent_config_cloud_omits_model(cloud_with_fake_bridge): @@ -784,6 +839,7 @@ def test_anthropic_runner_config_shapes(client, store_path): doc_id="pi-a") assert config["max_tokens"] == 4096 assert config["max_iterations"] == 10 + assert config["cache_control"] == {"type": "ephemeral"} assert "report.pdf" in config["system"] assert [tool.name for tool in config["tools"]] == list(tool_names()) assert (client.anthropic_runner_config(model="claude-sonnet-4-5") @@ -846,22 +902,35 @@ def test_anthropic_runner_config_doc_scope_enforced_in_tools(client, assert [doc["name"] for doc in browse["documents"]] == ["report.pdf"] -def test_claude_agent_config_doc_scope_enforced_in_tools(client, store_path): - pytest.importorskip("claude_agent_sdk") - from mcp.types import CallToolRequest, CallToolRequestParams +def test_claude_agent_config_doc_scope_enforced_in_tools(client, store_path, + monkeypatch): + """claude_agent_config(doc_id=...) must wire scope all the way into the + handlers it registers — an out-of-scope document returns NOT_FOUND. The + assertion has to drive those handlers, or build_claude_mcp's doc_ids + pass-through goes unguarded.""" + claude_agent_sdk = pytest.importorskip("claude_agent_sdk") seed_doc(store_path, "pi-a", "report.pdf") seed_doc(store_path, "pi-b", "payroll.pdf", created_at="2026-08-02T10:00:00.123000") + + registered = {} + create_server = claude_agent_sdk.create_sdk_mcp_server + + def capture(**kwargs): + registered.update(kwargs) + return create_server(**kwargs) + + monkeypatch.setattr(claude_agent_sdk, "create_sdk_mcp_server", capture) config = client.claude_agent_config(doc_id="pi-a") - server = config["mcp_servers"]["pageindex"] - handler = server["instance"].request_handlers[CallToolRequest] - result = asyncio.run(handler(CallToolRequest( - method="tools/call", - params=CallToolRequestParams( - name="get_page_content", - arguments={"doc_name": "payroll.pdf", "pages": "1"})))) - payload = json.loads(result.root.content[0].text) - assert payload["errorCode"] == "NOT_FOUND" + assert "report.pdf" in config["system_prompt"] + handlers = {spec.name: spec.handler for spec in registered["tools"]} + result = asyncio.run(handlers["get_page_content"]( + {"doc_name": "payroll.pdf", "pages": "1"})) + assert result.get("is_error") + assert json.loads(result["content"][0]["text"])["errorCode"] == "NOT_FOUND" + browse = asyncio.run(handlers["browse_documents"]({})) + listed = json.loads(browse["content"][0]["text"])["documents"] + assert [doc["name"] for doc in listed] == ["report.pdf"] def test_openai_agent_config_scoped_shadow_check(client, store_path): @@ -897,6 +966,60 @@ def test_claude_agent_config_scoped_shadow_check(client, store_path): assert "report.pdf" in config["system_prompt"] +def test_anthropic_runner_config_thinking_lifts_max_tokens(client): + pytest.importorskip("anthropic") + config = client.anthropic_runner_config( + model="claude-sonnet-4-5", + thinking={"type": "enabled", "budget_tokens": 10000}) + assert config["max_tokens"] == 10000 + 8192 + assert config["thinking"] == {"type": "enabled", "budget_tokens": 10000} + assert "thinking" not in client.anthropic_runner_config( + model="claude-sonnet-4-5") + + +def test_bridge_invoker_reraises_auth_failures(): + class Revoked: + def call_tool(self, name, arguments): + raise PageIndexAPIError("HTTP 401", status_code=401) + + invoke = agent_tools_module._bridge_invoker(Revoked(), "get_document", {}) + with pytest.raises(PageIndexAPIError, match="401"): + invoke({}) + # Transport blips stay contained in the retryable envelope. + + class Down: + def call_tool(self, name, arguments): + raise PageIndexAPIError("HTTP 503", status_code=503) + + text, is_error = agent_tools_module._bridge_invoker( + Down(), "get_document", {})({}) + assert is_error + assert json.loads(text)["errorCode"] == "INTERNAL_ERROR" + + +def test_cloud_bridge_gates_the_endpoint(monkeypatch): + """Instructions come from the same endpoint the tools register: the + read-gated URL by default, the full one with include_management.""" + created = [] + + class FakeBridge: + def __init__(self, url, headers): + created.append(url) + + def instructions(self): + return "SERVED" + + monkeypatch.setattr("pageindex.mcp_bridge.McpBridge", FakeBridge) + cloud = PageIndexCloudClient(api_key="pi-k") + assert cloud.agent_instructions() == "SERVED" + assert created == [f"{cloud.BASE_URL}/mcp?tools=read"] + cloud.agent_instructions(include_management=True) + assert created[1:] == [f"{cloud.BASE_URL}/mcp"] + cloud.agent_instructions() + cloud.agent_instructions(include_management=True) + assert len(created) == 2 # cached per gate + + def test_doc_scope_rejected_on_cloud_openai(): pytest.importorskip("agents") cloud = PageIndexCloudClient(api_key="pi-test-key") @@ -1139,7 +1262,9 @@ def test_cloud_agent_tools_discover_live_tool_set(cloud_with_fake_bridge): cloud, created = cloud_with_fake_bridge tools = cloud.agent_tools() bridge = created["bridge"] - assert bridge.url == "https://api.pageindex.ai/mcp" + # Default discovery rides the read-gated endpoint, matching the + # instructions fetch and the hosted/MCP registrations. + assert bridge.url == "https://api.pageindex.ai/mcp?tools=read" assert bridge.headers == {"Authorization": "Bearer pi-test-key"} # Default: only tools the server marks read-only; unannotated tools are # treated as non-read-only. @@ -1184,13 +1309,12 @@ def test_cloud_agent_tools_null_description_survives(): """A server may send description: null — .get(key, default) does not apply the default to it, and agent_tools() died with a TypeError while the _tool_specs path handled the same payload fine.""" - from pageindex.agent_tools import _make_bridge_function class _Bridge: def call_tool(self, name, arguments): return json.dumps({"success": True}), False - tool = _make_bridge_function(_Bridge(), { + tool = _synth(_Bridge(), { "name": "search_documents", "description": None, "inputSchema": {"type": "object", @@ -1227,6 +1351,35 @@ def list_tools(self): cloud.agent_tools() +def test_local_description_edit_actually_removes_the_image_sentence(): + """The local get_page_content description edits the contract text by + exact string replace — a contract wording change must fail here, not + silently ship the image-tool sentence to local models.""" + from pageindex.agent_tools import TOOL_CONTRACT, _LOCAL_DESCRIPTIONS + local = _LOCAL_DESCRIPTIONS["get_page_content"] + assert "get_document_image" not in local + assert len(local) < len(TOOL_CONTRACT["get_page_content"]["description"]) + + +def test_list_tools_pagination_is_bounded(): + """A server echoing its nextCursor terminates (no-progress guard); + a cycling one hits the page cap instead of hanging forever.""" + from pageindex.mcp_bridge import McpBridge + + bridge = McpBridge("http://x/mcp", {}) + pages = {None: {"tools": [{"name": "a"}], "nextCursor": "c1"}, + "c1": {"tools": [{"name": "b"}], "nextCursor": "c1"}} + bridge._request = lambda method, params=None: pages[ + (params or {}).get("cursor")] + assert [t["name"] for t in bridge.list_tools()] == ["a", "b"] + + bridge._request = lambda method, params=None: { + "tools": [], + "nextCursor": {"c1": "c2"}.get((params or {}).get("cursor"), "c1")} + with pytest.raises(PageIndexAPIError, match="did not terminate"): + bridge.list_tools() + + def test_mcp_bridge_protocol(monkeypatch): import requests as requests_mod from pageindex.mcp_bridge import McpBridge @@ -1284,7 +1437,8 @@ def fake_post(url, json=None, headers=None, timeout=None): # Replace the module's own `requests` binding — patching the shared # requests module would leak the fake process-wide. monkeypatch.setattr(mcp_bridge, "requests", types.SimpleNamespace( - post=fake_post, RequestException=requests_mod.RequestException)) + Session=lambda: types.SimpleNamespace(post=fake_post), + RequestException=requests_mod.RequestException)) bridge = McpBridge("https://api.pageindex.ai/mcp", {"Authorization": "Bearer k"}) @@ -1347,7 +1501,8 @@ def fake_post(url, json=None, headers=None, timeout=None): return _Resp(400, text="unknown tool") monkeypatch.setattr(mcp_bridge, "requests", types.SimpleNamespace( - post=fake_post, RequestException=requests_mod.RequestException)) + Session=lambda: types.SimpleNamespace(post=fake_post), + RequestException=requests_mod.RequestException)) bridge = McpBridge("https://api.pageindex.ai/mcp", {"Authorization": "Bearer k"}) with pytest.raises(PageIndexAPIError, match="HTTP 400"): @@ -1408,7 +1563,8 @@ def fake_post(url, json=None, headers=None, timeout=None): return resp monkeypatch.setattr(mcp_bridge, "requests", types.SimpleNamespace( - post=fake_post, RequestException=requests_mod.RequestException)) + Session=lambda: types.SimpleNamespace(post=fake_post), + RequestException=requests_mod.RequestException)) bridge = McpBridge("https://api.pageindex.ai/mcp", {"Authorization": "Bearer k"}) @@ -1447,20 +1603,71 @@ def test_mcp_bridge_blob_blocks_become_stubs(): bridge._request = lambda method, params: {"content": [ {"type": "text", "text": "Page 3 of report.pdf"}, {"type": "image", "mimeType": "image/png", "data": blob}, + {"type": "resource", + "resource": {"mimeType": "image/jpeg", "blob": blob}}, ]} text, is_error = bridge.call_tool("get_document_image", {}) assert not is_error assert "Page 3 of report.pdf" in text assert "AAAA" not in text assert "[image/png content omitted: ~6 KB]" in text + assert "[image/jpeg content omitted: ~6 KB]" in text # ── review-round regressions ── +def _synth(bridge, meta): + """Build a tool function the way the cloud lane does: signature + synthesis over a bridge invoker.""" + from pageindex.agent_tools import _bridge_invoker, _make_tool_function + name = meta["name"] + return _make_tool_function(name, meta.get("description"), + meta["inputSchema"], + _bridge_invoker(bridge, name, + meta["inputSchema"])) + + +def test_synth_binding_error_names_the_tool(): + """Binding TypeErrors quote the function's __qualname__; the model used + to see \"_synthesized() got an unexpected keyword argument\" and had no + tool name to correct against.""" + from pageindex.agent_tools import TOOL_CONTRACT + + class _Bridge: + def call_tool(self, name, args): + return json.dumps(args), False + + meta = {"name": "browse_documents", "description": "d", + "inputSchema": TOOL_CONTRACT["browse_documents"]["schema"]} + fn = _synth(_Bridge(), meta) + payload = json.loads(fn(bogus_param=1)) + assert "browse_documents" in payload["error"] + assert "_synthesized" not in payload["error"] + + +def test_bridge_invoker_coerces_string_booleans(): + """Identical model output must behave the same on both dispatch paths: + call_tool coerced "false" but the cloud bridge forwarded it verbatim, + turning "don't wait" into a 3-minute wait on lenient servers.""" + from pageindex.agent_tools import TOOL_CONTRACT, _bridge_invoker + + seen = {} + + class _Bridge: + def call_tool(self, name, args): + seen.update(args) + return "{}", False + + invoke = _bridge_invoker(_Bridge(), "get_document", + TOOL_CONTRACT["get_document"]["schema"]) + invoke({"doc_name": "q.pdf", "wait_for_completion": "false"}) + assert seen["wait_for_completion"] is False + + def test_synth_optional_no_default_param_is_nullable(): """A non-required, no-default schema param must annotate Optional, or strict schemas force the model to always send a value (browse.query).""" - from pageindex.agent_tools import _make_bridge_function, TOOL_CONTRACT + from pageindex.agent_tools import TOOL_CONTRACT from typing import get_args class _Bridge: @@ -1470,7 +1677,7 @@ def call_tool(self, name, args): meta = {"name": "browse_documents", "description": "d", "inputSchema": TOOL_CONTRACT["browse_documents"]["schema"]} - fn = _make_bridge_function(_Bridge(), meta) + fn = _synth(_Bridge(), meta) assert type(None) in get_args(fn.__annotations__["query"]) @@ -1479,7 +1686,6 @@ def test_synth_array_params_keep_their_item_type(): function_tool then emits {"type": "array", "items": {}}, which strict function calling rejects.""" from typing import Optional - from pageindex.agent_tools import _make_bridge_function class _Bridge: def call_tool(self, name, args): @@ -1498,7 +1704,7 @@ def call_tool(self, name, args): }, "required": ["doc_ids", "mixed"], }} - fn = _make_bridge_function(_Bridge(), meta) + fn = _synth(_Bridge(), meta) assert fn.__annotations__["doc_ids"] == list[str] assert fn.__annotations__["tags"] == Optional[list[int]] # A type-array in items (nullable elements) degrades to bare list — @@ -1507,7 +1713,6 @@ def call_tool(self, name, args): def test_synth_escape_hatches(): - from pageindex.agent_tools import _make_bridge_function calls = [] @@ -1517,7 +1722,7 @@ def call_tool(self, name, args): return "ok", False # Tool named "_invoke" must not recurse into itself. - invoke_named = _make_bridge_function(_Bridge(), { + invoke_named = _synth(_Bridge(), { "name": "_invoke", "description": "d", "inputSchema": {"type": "object", "properties": {"x": {"type": "string"}}, "required": ["x"]}}) @@ -1525,7 +1730,7 @@ def call_tool(self, name, args): assert calls[-1] == ("_invoke", {"x": "v"}) # Param named "dict" must not shadow the builtin. - dict_param = _make_bridge_function(_Bridge(), { + dict_param = _synth(_Bridge(), { "name": "t", "description": "d", "inputSchema": {"type": "object", "properties": {"dict": {"type": "string"}}, "required": ["dict"]}}) @@ -1534,7 +1739,7 @@ def call_tool(self, name, args): # Non-identifier tool name still gets a real signature. import inspect - dashed = _make_bridge_function(_Bridge(), { + dashed = _synth(_Bridge(), { "name": "page-content.v2", "description": "d", "inputSchema": {"type": "object", "properties": {"a": {"type": "string"}}, "required": ["a"]}}) @@ -1603,7 +1808,8 @@ def fake_post(url, json=None, headers=None, timeout=None): "content": [{"type": "text", "text": '{"error": "denied"}'}]}}) monkeypatch.setattr(mcp_bridge, "requests", types.SimpleNamespace( - post=fake_post, RequestException=requests_mod.RequestException)) + Session=lambda: types.SimpleNamespace(post=fake_post), + RequestException=requests_mod.RequestException)) bridge = McpBridge("https://api.pageindex.ai/mcp", {}) assert bridge.call_tool("t", {}) == ('{"error": "denied"}', True) @@ -1640,7 +1846,8 @@ def fake_post(url, json=None, headers=None, timeout=None): "text": "old"}]}}) monkeypatch.setattr(mcp_bridge, "requests", types.SimpleNamespace( - post=fake_post, RequestException=requests_mod.RequestException)) + Session=lambda: types.SimpleNamespace(post=fake_post), + RequestException=requests_mod.RequestException)) bridge = McpBridge("https://api.pageindex.ai/mcp", {}) with pytest.raises(PageIndexAPIError, match="no reply matching"): bridge.call_tool("t", {}) @@ -1664,7 +1871,8 @@ def dead_post(*args, **kwargs): raise requests_mod.ConnectionError("dns down") monkeypatch.setattr(mcp_bridge, "requests", types.SimpleNamespace( - post=dead_post, RequestException=requests_mod.RequestException)) + Session=lambda: types.SimpleNamespace(post=dead_post), + RequestException=requests_mod.RequestException)) bridge = McpBridge("https://api.pageindex.ai/mcp", {}) with pytest.raises(PageIndexAPIError, match="Could not reach"): bridge.list_tools() @@ -1737,6 +1945,17 @@ def list_documents(self, limit, offset): assert _all_documents(exact) == docs assert type(exact).calls == 2 # ...total still saves the empty page + # stop_ids ends the walk once every wanted id has been seen — a + # doc-scoped chat turn must not page the whole library... + early = make_client(120, 100) + listed = _all_documents(early, stop_ids=frozenset({"pi-3"})) + assert type(early).calls == 1 + assert any(doc["id"] == "pi-3" for doc in listed) + # ...while an id the listing lacks still costs the full sweep. + full = make_client(120, 100) + assert _all_documents(full, stop_ids=frozenset({"pi-missing"})) == docs + assert type(full).calls == 2 + def test_null_arguments_mean_omitted(client, store_path): """Adapters that forward the model's null values verbatim (the Claude @@ -1759,17 +1978,6 @@ def test_page_spec_span_bomb_rejected(client, store_path): assert "Too many pages" in payload["error"] -def test_agent_instructions_shadowed_doc_id_raises(client, store_path): - seed_doc(store_path, "pi-old", "report.pdf", - created_at="2026-08-01T10:00:00.000000") - seed_doc(store_path, "pi-new", "report.pdf", - created_at="2026-08-02T10:00:00.000000") - with pytest.raises(PageIndexAPIError, match="shadowed"): - client.agent_instructions(doc_id="pi-old") - text = client.agent_instructions(doc_id="pi-new") - assert "report.pdf" in text - - def test_wait_tolerates_transient_network_failures(fake_cloud_client, monkeypatch): import requests as requests_mod cloud = fake_cloud_client(["processing", "completed"]) @@ -1826,6 +2034,7 @@ def flaky(doc_id): assert cloud.submit_document("x.pdf", wait=True) == {"doc_id": "pi-fake"} +import pageindex.utils # noqa: F401 — its import loads .env LIVE_KEY = os.getenv("PAGEINDEX_API_KEY") @@ -2008,6 +2217,31 @@ def instructions(self): pickle.dumps(cloud) +def test_cloud_bridge_rebuilds_on_credential_change(monkeypatch): + """CloudAPI re-reads client.api_key on every REST call; the MCP half + must not keep authenticating with a rotation-stale snapshot.""" + import pageindex.mcp_bridge as mcp_bridge + from pageindex.agent_tools import _cloud_bridge + built = [] + + class _Bridge(_FakeBridge): + def __init__(self, url, headers): + super().__init__(url, headers) + built.append((url, dict(headers))) + + monkeypatch.setattr(mcp_bridge, "McpBridge", _Bridge) + cloud = PageIndexCloudClient(api_key="pi-old") + first = _cloud_bridge(cloud) + assert _cloud_bridge(cloud) is first # unchanged credentials: cached + cloud.api_key = "pi-new" + second = _cloud_bridge(cloud) + assert second is not first + assert built[-1][1]["Authorization"] == "Bearer pi-new" + cloud.BASE_URL = "https://alt.example" + assert _cloud_bridge(cloud) is not second + assert built[-1][0] == "https://alt.example/mcp" + + def test_cloud_agent_instructions_blank_or_nonstring_raises(monkeypatch): """Whitespace-only or non-string initialize.instructions must hit the same honest error as a missing one — never a blank system prompt.""" @@ -2141,6 +2375,28 @@ def test_config_helpers_reject_empty_doc_id_on_cloud(): cloud.claude_agent_config(doc_id=[]) +def test_doc_targeting_keeps_transport_errors_out_of_not_found(): + """A cloud 429/5xx during the doc_id fetch is an outage, not a missing + document — only a definite not-found/denied batches into the + "Documents not found" message; anything else propagates raw.""" + class Stub: + def __init__(self, status): + self.status = status + + def get_document(self, doc_id): + raise PageIndexAPIError( + f"Failed to get document metadata: {self.status}", + status_code=self.status) + + for status in (429, 500): + with pytest.raises(PageIndexAPIError, match=f"metadata: {status}"): + agent_tools_module.doc_targeting_block(Stub(status), "pi-a") + for status in (403, 404): + with pytest.raises(PageIndexAPIError, + match="Documents not found or access denied: pi-a"): + agent_tools_module.doc_targeting_block(Stub(status), "pi-a") + + def test_call_tool_coerces_string_booleans(client, store_path, monkeypatch): """Models routinely send booleans as JSON strings — "false" must not read as True (a full wait_for_completion stall).""" @@ -2221,3 +2477,54 @@ def test_agent_instructions_carry_user_metadata(client, store_path): metadata={"quarter": "Q3", "year": 2025}) text = client.agent_instructions(doc_id="pi-1") assert '"quarter": "Q3"' in text and '"year": 2025' in text + + +# ── wait-poll resilience, instruction scoping, agent_tools doc_id ── + +def test_await_completion_polls_through_transient_refetch_failure(monkeypatch): + """A refetch that fails once must not end the wait early — the caller + would report that 5-second exit as the full 3-minute timeout.""" + monkeypatch.setattr(agent_tools_module, "_TOOL_WAIT_INTERVAL", 0.0) + calls = {"n": 0} + + class Flaky: + def get_document(self, doc_id): + calls["n"] += 1 + if calls["n"] == 1: + raise PageIndexAPIError("transient listing failure") + return {"id": doc_id, "status": "completed"} + + result = agent_tools_module._await_completion( + Flaky(), {"id": "pi-x", "status": "processing"}, wait=True) + assert result["status"] == "completed" + assert calls["n"] == 2 + + +def test_agent_instructions_doc_id_shadow_check(client, store_path): + seed_doc(store_path, "pi-old", "report.pdf", + created_at="2026-08-01T10:00:00.123000") + seed_doc(store_path, "pi-new", "report.pdf", + created_at="2026-08-05T10:00:00.123000") + # standalone instructions get the strict check; only the *_agent_config + # bundles (which build the tools too) relax it + with pytest.raises(PageIndexAPIError, match="shadowed"): + client.agent_instructions(doc_id="pi-old") + + +def test_agent_tools_doc_id_scopes_the_functions(client, store_path): + seed_doc(store_path, "pi-a", "alpha.pdf") + seed_doc(store_path, "pi-b", "secret.pdf") + funcs = {fn.__name__: fn for fn in client.agent_tools(doc_id="pi-a")} + blocked = json.loads(funcs["get_page_content"](doc_name="secret.pdf", + pages="1")) + assert "success" not in blocked + assert blocked["errorCode"] == "NOT_FOUND" + allowed = json.loads(funcs["get_page_content"](doc_name="alpha.pdf", + pages="1")) + assert allowed["success"] is True + + +def test_agent_tools_doc_id_refused_on_cloud(): + cloud = PageIndexCloudClient(api_key="pi-test-key") + with pytest.raises(PageIndexAPIError, match="local tools only"): + cloud.agent_tools(doc_id="pi-a") diff --git a/tests/test_client.py b/tests/test_client.py index d3c17ef52..2ace6ba2e 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -68,14 +68,19 @@ def test_local_client_does_not_touch_disk(tmp_path): assert not storage.exists() -def test_retrieve_model_carries_agents_sdk_prefix(tmp_path): +def test_retrieve_model_stays_as_configured(tmp_path): + """The public attribute keeps the caller's spelling — ``litellm/`` is + Agents SDK routing grammar, applied at the config door + (openai_agent_config), never baked into ``chat_model``: handed to the + Anthropic SDK or a raw request, the prefixed form is a 404.""" def resolved(retrieve_model): return PageIndexClient(retrieve_model=retrieve_model, storage_path=str(tmp_path / "s")).retrieve_model - assert resolved("anthropic/claude-sonnet-4-6") == "litellm/anthropic/claude-sonnet-4-6" - for already_routable in ("gpt-4o", "openai/gpt-4o", "litellm/anthropic/claude-sonnet-4-6"): - assert resolved(already_routable) == already_routable + for as_configured in ("anthropic/claude-sonnet-4-6", "gpt-4o", + "openai/gpt-4o", + "litellm/anthropic/claude-sonnet-4-6"): + assert resolved(as_configured) == as_configured def test_model_resolution_covers_every_generation(tmp_path): @@ -173,7 +178,7 @@ def test_get_page_content(local_client, indexed_doc): assert local_client.get_page_content(indexed_doc, "99") == [] - with pytest.raises(ValueError): + with pytest.raises(PageIndexAPIError): local_client.get_page_content(indexed_doc, "abc") @@ -181,7 +186,7 @@ def test_get_page_content_span_bomb_rejected(local_client, indexed_doc): """An absurd range must be rejected arithmetically, not expanded into a billion integers in the caller's process (the tool layer already refused; the public client method did not).""" - with pytest.raises(ValueError, match="spans more than 10000"): + with pytest.raises(PageIndexAPIError, match="spans more than 10000"): local_client.get_page_content(indexed_doc, "1-1000001") # At the bound itself the spec still parses. assert local_client.get_page_content(indexed_doc, "5-10004") == [] @@ -285,18 +290,36 @@ def test_page_index_flash_rejects_unknown_optimize(): def test_llm_completion_missing_key_raises_immediately(monkeypatch): import openai + import litellm # first import may load a .env; delenv after it monkeypatch.delenv("OPENAI_API_KEY", raising=False) - monkeypatch.setattr(pageindex.utils, "_openai_sync_client", None) - monkeypatch.setattr(pageindex.utils, "_openai_async_client", None) - with pytest.raises(openai.OpenAIError): + with pytest.raises(openai.OpenAIError, match="OPENAI_API_KEY"): pageindex.utils.llm_completion("gpt-4o", "probe") - with pytest.raises(openai.OpenAIError): + with pytest.raises(openai.OpenAIError, match="OPENAI_API_KEY"): asyncio.run(pageindex.utils.llm_acompletion("gpt-4o", "probe")) + # unknown bare names are OpenAI shorthand, so the same check applies + with pytest.raises(openai.OpenAIError, match="OPENAI_API_KEY"): + pageindex.utils.llm_completion("my-finetune-v2", "probe") + # a blank exported key is as missing as no key (litellm's + # validate_environment reports it present) + monkeypatch.setenv("OPENAI_API_KEY", " ") + with pytest.raises(openai.OpenAIError, match="OPENAI_API_KEY"): + pageindex.utils.llm_completion("gpt-4o", "probe") + + +def test_llm_completion_refuses_unknown_provider(monkeypatch): + """A first segment LiteLLM does not know (a HuggingFace repo id like + Qwen/...) is refused with the openai/ escape before the retry loop, + instead of burning it on per-call 400s.""" + import litellm + monkeypatch.setattr(litellm, "completion", + lambda **kw: pytest.fail("reached the wire")) + with pytest.raises(Exception, match="not a LiteLLM provider"): + pageindex.utils.llm_completion("Qwen/my-model", "probe") def test_submit_missing_llm_key_fails_loud(local_client, sample_pdf, monkeypatch): + import litellm # first import may load a .env; delenv after it monkeypatch.delenv("OPENAI_API_KEY", raising=False) - monkeypatch.setattr(pageindex.utils, "_openai_sync_client", None) def first_llm_call(*args, **kwargs): return pageindex.utils.llm_completion("gpt-4o", "probe") monkeypatch.setattr(page_index_module, "page_index_main", first_llm_call) @@ -652,6 +675,70 @@ async def flaky(model, prompt): assert summaries == {"A": "", "B": "ok"} +def test_generate_summaries_unrecoverable_raises(monkeypatch): + """A per-node 401 must abort, not store a blank node as completed.""" + class Denied(Exception): + status_code = 401 + + async def deny_t1(model, prompt): + if "t1" in prompt: + raise Denied("key rejected") + return "ok" + monkeypatch.setattr(pageindex.utils, "llm_acompletion", deny_t1) + structure = [{"title": "A", "text": "t1", + "nodes": [{"title": "B", "text": "t2"}]}] + with pytest.raises(Denied): + asyncio.run(pageindex.utils.generate_summaries_for_structure(structure)) + + +def test_summarize_tree_child_unrecoverable_raises(monkeypatch): + """A 401 on a leaf must abort the run, not store a blank subtree as + completed: the child gather's exceptions are checked, not discarded.""" + class Denied(Exception): + status_code = 401 + + async def deny_alpha(model, prompt): + if "alpha" in prompt: + raise Denied("key rejected") + return '{"points": [], "summary": "ok"}' + monkeypatch.setattr(pageindex.utils, "llm_acompletion", deny_alpha) + pdf_pages = [("alpha " * 5, 5), ("beta " * 5, 5)] + structure = [{"title": "R", "start_index": 1, "end_index": 2, + "nodes": [ + {"title": "A", "start_index": 1, "end_index": 1}, + {"title": "B", "start_index": 2, "end_index": 2}]}] + with pytest.raises(Denied): + asyncio.run(pageindex.utils.summarize_tree( + structure, pdf_pages, small_node_tokens=0)) + + +def test_llm_completion_suppresses_litellm_cache_seeding(monkeypatch): + """Indexing prompts are single-shot: without an explicit injection + point litellm 1.97 seeds its own cache marks and every call pays the + write premium for nothing. Backend keys still override ours.""" + import litellm + captured = {} + + def fake_completion(**kwargs): + captured.clear() + captured.update(kwargs) + message = types.SimpleNamespace(content="ok") + choice = types.SimpleNamespace(message=message, finish_reason="stop") + return types.SimpleNamespace(choices=[choice]) + monkeypatch.setattr(litellm, "completion", fake_completion) + monkeypatch.setenv("OPENAI_API_KEY", "k") + assert pageindex.utils.llm_completion("gpt-4o", "probe") == "ok" + assert captured["cache_control_injection_points"] == [ + {"location": "message", "role": "system"}] + token = pageindex.utils._llm_backend.set( + {"api_key": "x", "cache_control_injection_points": []}) + try: + pageindex.utils.llm_completion("gpt-4o", "probe") + finally: + pageindex.utils._llm_backend.reset(token) + assert captured["cache_control_injection_points"] == [] + + def test_delete_survives_marker_tamper(local_client, tmp_path): tampered = tmp_path / "store" / "docs" / "tampered" / "doc.json" tampered.mkdir(parents=True) @@ -833,25 +920,28 @@ def test_parse_pages_overlap_counts_union(): from pageindex.client import _parse_pages pages = _parse_pages("1-5000,2000-9000") assert len(pages) == 9000 and pages[0] == 1 and pages[-1] == 9000 - with pytest.raises(ValueError, match="spans more than"): + with pytest.raises(PageIndexAPIError, match="spans more than"): _parse_pages("1-10001") + # one parser with the tool layer now: page 0 is rejected, not passed + # on — surfaced as the documented SDK error type + with pytest.raises(PageIndexAPIError, match="positive"): + _parse_pages("0-3") # ── backend: the indexing lane ── def test_backend_scopes_the_index_lane(tmp_path, monkeypatch): - """index_backend reaches both gateway paths — LiteLLM's call kwargs - and a fresh openai-SDK client — scoped to the operation, with the - endpoint spelling normalized for the openai SDK.""" + """index_backend reaches the indexing lane's LiteLLM call kwargs + verbatim — bare and provider-prefixed models alike — scoped to the + operation, bypassing the env pre-check.""" pytest.importorskip("litellm") import litellm - import openai from types import SimpleNamespace from pageindex.local_api import LocalAPI from pageindex.utils import _llm_backend, llm_completion api = LocalAPI(storage_path=str(tmp_path / "s"), model="m", - summary_model="s", retrieve_model="r", + summary_model="s", index_backend={"api_key": "ik", "api_base": "http://b"}) assert api._with_backend(_llm_backend.get) == {"api_key": "ik", "api_base": "http://b"} @@ -862,41 +952,64 @@ def test_backend_scopes_the_index_lane(tmp_path, monkeypatch): captured = {} monkeypatch.setattr(litellm, "completion", lambda **kw: (captured.update(kw), reply)[1]) - api._with_backend(lambda: llm_completion("anthropic/claude-x", "p")) - assert captured["api_key"] == "ik" - assert captured["api_base"] == "http://b" + monkeypatch.setattr(litellm, "validate_environment", + lambda *a, **k: pytest.fail("env pre-check ran")) + for model, wire in (("anthropic/claude-x", "anthropic/claude-x"), + ("gpt-4o", "openai/gpt-4o"), + ("my-finetune-v2", "openai/my-finetune-v2")): + captured.clear() + api._with_backend(lambda: llm_completion(model, "p")) + assert captured["model"] == wire + assert captured["api_key"] == "ik" + assert captured["api_base"] == "http://b" + + +def test_index_precheck_covers_only_openai_shaped(monkeypatch): + """The missing-key pre-check fires only for OpenAI-shaped names — other + providers resolve credentials at call time (IAM chains, ADC, Ollama's + localhost default), invisible to env inspection, so the lane must not + block them up front.""" + pytest.importorskip("litellm") + import litellm + from types import SimpleNamespace + from pageindex.utils import llm_completion + + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + with pytest.raises(litellm.AuthenticationError, match="missing API key"): + llm_completion("my-finetune-v2", "p") + + reply = SimpleNamespace(choices=[SimpleNamespace( + message=SimpleNamespace(content="ok"), finish_reason="stop")]) + monkeypatch.setattr(litellm, "completion", lambda **kw: reply) + monkeypatch.setattr(litellm, "validate_environment", + lambda *a, **k: pytest.fail("env pre-check ran")) + assert llm_completion("ollama/llama3", "p") == "ok" + assert llm_completion("bedrock/anthropic.claude-sonnet", "p") == "ok" - seen = {} - class _FakeOpenAI: - def __init__(self, **kw): - seen.update(kw) - self.chat = SimpleNamespace(completions=SimpleNamespace( - create=lambda **_: reply)) +def test_litellm_routing_prefix_skips_key_precheck(monkeypatch): + """litellm/-prefixed names are an explicit routing choice: litellm + resolves credentials beyond the environment (litellm.api_key, a + keyless OPENAI_BASE_URL server), so the env pre-check stands aside.""" + pytest.importorskip("litellm") + import litellm # noqa: F401 — first import may load a .env; delenv after + from pageindex.utils import _litellm_model, _openai_missing_keys - monkeypatch.setattr(openai, "OpenAI", _FakeOpenAI) - api._with_backend(lambda: llm_completion("gpt-4o", "p")) - assert seen["api_key"] == "ik" - assert seen["base_url"] == "http://b" + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + assert _openai_missing_keys("litellm/gpt-4o") == [] + assert _litellm_model("litellm/gpt-4o", None) == "openai/gpt-4o" -def test_index_backend_skips_the_env_precheck(tmp_path, monkeypatch): - """The missing-key pre-check reads environment variables only, so a - backend-supplied key must bypass it instead of being refused.""" +def test_custom_provider_map_passes_provider_precheck(monkeypatch): + """litellm appends custom_provider_map providers to provider_list only + at completion time, so the pre-check must consult the map itself.""" pytest.importorskip("litellm") import litellm - from pageindex.local_api import LocalAPI + from pageindex.utils import _litellm_model - api = LocalAPI(storage_path=str(tmp_path / "s"), model="m", - summary_model="s", retrieve_model="r", - index_backend={"api_key": "ik"}) - monkeypatch.setattr(litellm, "validate_environment", - lambda *a, **k: pytest.fail("env pre-check ran")) - monkeypatch.setattr(pageindex.flash, "page_index_flash", - lambda *a, **k: (_ for _ in ()).throw( - RuntimeError("reached-flash"))) - with pytest.raises(RuntimeError, match="reached-flash"): - api._index_flash("f.pdf", ["text"]) + monkeypatch.setattr(litellm, "custom_provider_map", + [{"provider": "my-llm", "custom_handler": object()}]) + assert _litellm_model("my-llm/model-a", None) == "my-llm/model-a" def test_backend_args_are_local_only(): @@ -904,3 +1017,65 @@ def test_backend_args_are_local_only(): PageIndexClient(api_key="pi-k", chat_backend={"api_key": "x"}) with pytest.raises(PageIndexAPIError, match="index_backend"): PageIndexClient(api_key="pi-k", index_backend={"api_key": "x"}) + + +def test_chat_wraps_answerless_cloud_reply(monkeypatch): + """A cloud reply without choices (filtered / malformed) surfaces as + the SDK's error, not a bare IndexError/KeyError.""" + client = PageIndexClient(api_key="pi-k") + for reply in ({"id": "x", "object": "chat.completion", "choices": []}, + {"id": "x"}): + monkeypatch.setattr(client, "chat_completions", + lambda *a, _r=reply, **k: _r) + with pytest.raises(PageIndexAPIError, match="carries no answer"): + client.chat("hi") + + +def test_retrieve_model_assignment_still_works(local_client): + """0.2.9 allowed `client.retrieve_model = ...`; the legacy property + keeps the write path as an alias for chat_model.""" + local_client.retrieve_model = "gpt-x" + assert local_client.chat_model == "gpt-x" + assert local_client.retrieve_model == "gpt-x" + + +def test_concurrent_same_name_submits_store_unique_names(local_client, + sample_pdf, + monkeypatch): + """Name uniquing runs under the store lock at save time, so two + clients indexing the same filename concurrently cannot both store it + — a stored duplicate would shadow the older doc_id forever.""" + import threading + import time as time_mod + + def slow_flash(pdf, **kwargs): + time_mod.sleep(0.1) # both threads index before either saves + return {"structure": [{"title": "T", "start_index": 1, + "end_index": 2, "summary": "s", "nodes": []}]} + + monkeypatch.setattr(pageindex.flash, "page_index_flash", slow_flash) + monkeypatch.setattr(pageindex.utils, "llm_completion", + lambda *a, **k: "d") + results = [] + workers = [threading.Thread( + target=lambda: results.append(local_client.submit_document(sample_pdf))) + for _ in range(2)] + for worker in workers: + worker.start() + for worker in workers: + worker.join() + assert {r["name"] for r in results} == {"sample.pdf", "sample_1.pdf"} + + + +def test_format_tree_node_keeps_key_items(): + """key_items from the merge optimization survive the get_tree formatter.""" + from pageindex.local_api import _format_tree_node + + node = {"title": "Chapter 1", "node_id": "0000", "start_index": 1, + "summary": "s", + "key_items": ["1.1 Alpha", "1.2 Beta", "1.3 Gamma"]} + out = _format_tree_node(node, node_summary=True) + assert out["key_items"] == ["1.1 Alpha", "1.2 Beta", "1.3 Gamma"] + assert "key_items" not in _format_tree_node( + {"title": "t", "node_id": "0001", "start_index": 1}, False) diff --git a/tests/test_flash_extraction.py b/tests/test_flash_extraction.py new file mode 100644 index 000000000..1e9dcb87d --- /dev/null +++ b/tests/test_flash_extraction.py @@ -0,0 +1,170 @@ +"""Pins the pdfium 5.x text-extraction semantics the parser is calibrated to. + +pdfium split FPDFFont_GetFontName into GetBaseFontName (/BaseFont, per-face) +and GetFamilyName (old family semantics); the parser uses the per-face names, +which changes word joining and figure-label pickup. These sentinels come from +a 4.30-vs-5.13 corpus A/B and fail if the semantics move again. +""" +from importlib.metadata import version +from pathlib import Path + +import pytest + +PDF = Path(__file__).parent.parent / "examples" / "documents" / "earthmover.pdf" + + +@pytest.mark.skipif(int(version("pypdfium2").split(".")[0]) < 5, + reason="extraction is pinned to pdfium 5.x font-name semantics") +def test_page_text_pins_pdfium5_semantics(): + from pageindex.flash.main import extract_toc + + page7 = extract_toc(str(PDF))["page_texts"][6] + assert "p5\nEMD\n1.0" in page7 # figure axis label pdfium 4.x dropped + assert "break loop\n5: if lbp" in page7 # pseudocode lines no longer glued + + +def test_page_mode_walk_uses_merged_surrogate_census(): + """The page-mode unicode walk must consume the char census char_extract + built (astral chars merged to one entry at the high-surrogate slot). + Re-reading the textpage split them back into two lone-surrogate slots, + desynced the walk against their one-char cmap targets, and silently + dropped every patch on any page containing an astral char.""" + from pageindex.flash.parser_pdfium_charlevel.unicode_apply import ( + _apply_font_unicode) + + astral = {"i": 0, "ch": "\U0001d44e", "is_gen": False} # slots 0-1 merged + unmapped = {"i": 2, "ch": "\x00", "is_gen": False} # PDFium found no unicode + raw_chars = [astral, unmapped] + show_codes = [(7, (5, 6), 100.0)] + map_cache = {7: (1, {5: "\U0001d44e", 6: "β"})} + + # objects vs show ops count differs -> page mode. + _apply_font_unicode(raw_chars, [], show_codes, None, map_cache) + + assert astral["ch"] == "\U0001d44e" + assert unmapped["ch"] == "β" + + +def test_optimize_full_fails_fast_without_a_key(tmp_path, monkeypatch): + """optimize='full' runs LLM expand: with no key configured it must be + an instant, guided PageIndexAPIError — raised before any PDF work (a + bogus path proves the ordering) — while the LLM-free spellings and a + backend-carrying indexing scope stay untouched.""" + from conftest import build_pdf + from pageindex import PageIndexAPIError + from pageindex.flash import page_index_flash + from pageindex.utils import _llm_backend + import litellm # noqa: F401 — first import may load a .env; delenv after it + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + monkeypatch.delenv("CHATGPT_API_KEY", raising=False) + + with pytest.raises(PageIndexAPIError, match="optimize='merge'"): + page_index_flash(str(tmp_path / "missing.pdf"), summary=False) + + pdf = tmp_path / "doc.pdf" + pdf.write_bytes(build_pdf(["1 Introduction", "Body text"])) + result = page_index_flash(str(pdf), summary=False, optimize="merge") + assert "structure" in result + result = page_index_flash(str(pdf), summary=False, optimize=False) + assert "structure" in result + + token = _llm_backend.set({"api_key": "k"}) + try: + with pytest.raises(FileNotFoundError): + page_index_flash(str(tmp_path / "missing.pdf"), summary=False) + finally: + _llm_backend.reset(token) + + +def test_bootstrap_reimport_is_not_swallowed(monkeypatch): + # An unguarded caller script re-imported by a spawn worker must die loudly, + # not fall back to a silent full sequential rerun in every worker. + import multiprocessing + import sys + + from pageindex.flash import parser_pdfium_parallel as mod + + class BoomExecutor: + def __init__(self, *a, **k): + pass + + def map(self, *a, **k): + raise RuntimeError("start a new process before bootstrapping") + + def shutdown(self, *a, **k): + pass + + monkeypatch.setattr(mod, "ProcessPoolExecutor", BoomExecutor) + cur = multiprocessing.current_process() + + monkeypatch.setattr(cur, "_inheriting", True, raising=False) + with pytest.raises(RuntimeError): + mod.parse_charlevel_meta_parallel(str(PDF), workers=2, min_pages=1) + assert hasattr(sys.modules["__main__"], "__file__") # window restored on error + + monkeypatch.delattr(cur, "_inheriting") + out, meta = mod.parse_charlevel_meta_parallel(str(PDF), workers=2, min_pages=1) + assert len(out) == len(meta) > 0 # normal failures still fall back sequentially + + +def test_submit_document_refuses_during_bootstrap(tmp_path, monkeypatch): + import multiprocessing + + from pageindex import PageIndexAPIError, PageIndexLocalClient + + c = PageIndexLocalClient(storage_path=str(tmp_path)) + monkeypatch.setattr( + multiprocessing.current_process(), "_inheriting", True, raising=False + ) + with pytest.raises(PageIndexAPIError, match="__main__"): + c.submit_document("whatever.pdf") + + +def test_unguarded_script_parses_parallel_without_reexecution(tmp_path): + # spawn workers must not re-run an unguarded caller script: one completion, + # no dead-worker noise (dying workers would trip the sequential fallback). + import os + import subprocess + import sys + + marker = tmp_path / "runs.txt" + script = tmp_path / "unguarded.py" + script.write_text( + "from pageindex.flash.parser_pdfium_parallel import parse_charlevel_meta_parallel\n" + f"out, meta = parse_charlevel_meta_parallel({str(PDF)!r}, workers=2, min_pages=1)\n" + "assert len(out) == len(meta) > 0\n" + f"open({str(marker)!r}, 'a').write('ran\\n')\n" + ) + env = {**os.environ, "PYTHONPATH": str(Path(__file__).parent.parent)} + res = subprocess.run( + [sys.executable, str(script)], capture_output=True, env=env, timeout=120 + ) + assert res.returncode == 0, res.stderr.decode() + assert marker.read_text() == "ran\n" + assert b"Traceback" not in res.stderr + + +def test_optimize_wins_over_deprecated_optimize_expand(tmp_path, monkeypatch): + """Explicit optimize= beats optimize_expand; legacy True still honors it.""" + from conftest import build_pdf + from pageindex.flash import page_index_flash + import litellm # noqa: F401 — first import may load a .env; delenv after it + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + monkeypatch.delenv("CHATGPT_API_KEY", raising=False) + + pdf = tmp_path / "doc.pdf" + pdf.write_bytes(build_pdf(["1 Introduction", "Body text"])) + # resolved to "full" before the precedence fix (keyless → fail-fast) + with pytest.warns(DeprecationWarning): + result = page_index_flash(str(pdf), summary=False, + optimize="merge", optimize_expand=True) + assert "structure" in result + with pytest.warns(DeprecationWarning): + result = page_index_flash(str(pdf), summary=False, + optimize=True, optimize_expand=False) + assert "structure" in result + # optimize=None means unset ("full"), not off + from pageindex import PageIndexAPIError + with pytest.raises(PageIndexAPIError, match="optimize='merge'"): + page_index_flash(str(tmp_path / "missing.pdf"), summary=False, + optimize=None) diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index d34ff9a47..a2a3d591e 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -124,6 +124,15 @@ async def stream_response(self, system_instructions, input, self._record(system_instructions, input) output = self.turns.pop(0) sequence = 0 + if getattr(self, "emit_created", False): + from openai.types.responses import ResponseCreatedEvent + sequence += 1 + yield ResponseCreatedEvent( + type="response.created", sequence_number=sequence, + response=Response( + id="resp_backend_turn", created_at=0.0, model="fake", + object="response", output=[], parallel_tool_calls=False, + tool_choice="auto", tools=[], status="in_progress")) for item in output: if item.type == "message": pieces = getattr(self, "pieces", ("The ", "answer")) @@ -306,7 +315,8 @@ def test_anthropic_routed_models_mark_managed_prefix_for_cache( fake_model([[_msg_item("ok")]]) from pageindex.local_chat import _openai_agent marked = {"cache_control_injection_points": [ - {"location": "message", "role": "system"}]} + {"location": "message", "role": "system"}, + {"location": "message", "index": -1}]} for name in ("anthropic/claude-x", "litellm/anthropic/claude-x", "bedrock/us.anthropic.claude-sonnet-5", "vertex_ai/claude-sonnet-4-5"): @@ -525,8 +535,8 @@ def test_doc_id_conversations_get_distinct_cache_keys(client, store_path, keys = [] real = local_chat._conversation_cache_key - def spy(model_name, instructions, items): - key = real(model_name, instructions, items) + def spy(model_name, instructions, doc_id, items): + key = real(model_name, instructions, doc_id, items) keys.append(key) return key @@ -551,6 +561,11 @@ def spy(model_name, instructions, items): client.chat_completions("Summarize section 3.", doc_id="pi-a") assert keys[3] != keys[4] # same property on the chat surface + seed_doc(store_path, "pi-b", "contract.pdf") + fake_model([[_msg_item("f")]]) + client.responses("What is the CAGR?", doc_id="pi-b") + assert keys[5] != keys[0] # same opener, different doc: no pooling + @needs_agents def test_responses_stream_passthrough(client, store_path, fake_model): @@ -580,6 +595,30 @@ def test_responses_stream_passthrough(client, store_path, fake_model): .get("type", "message") == "message") +@needs_agents +def test_responses_stream_opens_with_created(client, store_path, fake_model): + """N per-turn openings collapse to one response.created, not zero — + the logical stream must open with a response object carrying the same + id the terminal event reports, and the terminal envelope reports the + backend's tool-param echo, not assumed values.""" + seed_doc(store_path, "pi-a", "report.pdf") + fake = fake_model([ + [_call_item("get_document", {"doc_name": "report.pdf"})], + [_msg_item("The answer")], + ]) + fake.emit_created = True # two turns emit two; one must pass through + events = list(client.responses("q", stream=True)) + created = [e for e in events if e["type"] == "response.created"] + assert len(created) == 1 and events[0] is created[0] + assert events[0]["sequence_number"] == 1 + terminal = events[-1] + assert terminal["type"] == "response.completed" + assert created[0]["response"]["id"] == terminal["response"]["id"] + assert (created[0]["response"]["created_at"] + == terminal["response"]["created_at"]) # one timestamp, not two + assert terminal["response"]["parallel_tool_calls"] is False # echo + + @needs_agents def test_responses_envelope_validates_as_official_response(client, store_path, fake_model): @@ -824,6 +863,10 @@ def test_max_turns_rejects_non_positive(client, store_path, fake_model): with pytest.raises(PageIndexAPIError, match="positive integer"): client.chat_completions([{"role": "user", "content": "q"}], max_turns=0) + # every door that takes max_turns validates it, the runner config too + with pytest.raises(PageIndexAPIError, match="positive integer"): + client.anthropic_runner_config(model="claude-sonnet-4-5", + max_turns=-1) def test_enable_citations_rejected_before_framework_check(client, monkeypatch): @@ -874,6 +917,7 @@ def test_responses_envelope_fields_and_cache_group(client, store_path, "get_document_structure", "get_page_content"} assert all(tool["type"] == "function" for tool in result["tools"]) assert result["instructions"].startswith(CHAT_HEADER) + # No transport echo attached here, so these are the fallbacks. assert result["parallel_tool_calls"] is True assert result["tool_choice"] == "auto" @@ -893,17 +937,22 @@ def test_sol_class_refusal_names_its_exits(): def test_conversation_cache_key_stable_per_conversation(): """Cache-routing key, sent as the OpenAI prompt_cache_key. A conversation's continuations must share one key (same model / - instructions / first item), and unrelated conversations must not pool - under it.""" + instructions / doc targeting / first item), and unrelated + conversations must not pool under it.""" turn1 = [{"role": "user", "content": "q"}] continuation = turn1 + [{"role": "assistant", "content": "a"}, {"role": "user", "content": "and?"}] - key = local_chat._conversation_cache_key("m", "sys", turn1) - assert key == local_chat._conversation_cache_key("m", "sys", continuation) + key = local_chat._conversation_cache_key("m", "sys", "d1", turn1) + assert key == local_chat._conversation_cache_key( + "m", "sys", "d1", continuation) + assert key == local_chat._conversation_cache_key( + "m", "sys", ["d1"], turn1) # str and one-item list: same targeting assert key != local_chat._conversation_cache_key( - "m", "sys", [{"role": "user", "content": "other"}]) - assert key != local_chat._conversation_cache_key("m2", "sys", turn1) - assert key != local_chat._conversation_cache_key("m", "sys2", turn1) + "m", "sys", "d1", [{"role": "user", "content": "other"}]) + assert key != local_chat._conversation_cache_key("m2", "sys", "d1", turn1) + assert key != local_chat._conversation_cache_key("m", "sys2", "d1", turn1) + assert key != local_chat._conversation_cache_key("m", "sys", "d2", turn1) + assert key != local_chat._conversation_cache_key("m", "sys", None, turn1) @needs_agents @@ -924,6 +973,9 @@ def test_agent_carries_prompt_cache_key_in_extra_body(monkeypatch): settings = agent.model_settings assert settings.extra_body is None assert "cache_control_injection_points" in settings.extra_args + assert settings.extra_args["cache_control_injection_points"] == [ + {"location": "message", "role": "system"}, + {"location": "message", "index": -1}] for name in ("gpt-test", "openai/gpt-test", "litellm/openai/gpt-test"): agent = local_chat._openai_agent( None, "chat", name, "sys", None, None, @@ -1070,18 +1122,18 @@ def tool_outputs(items): @needs_agents -def test_empty_doc_id_is_an_empty_allowlist(client, store_path, fake_model): - """doc_id=[] scopes the agent to nothing; `or None` used to wash it - into unscoped full-library access.""" +def test_empty_doc_id_is_refused(client, store_path): + """doc_id=[] fails loud on every local surface, like cloud already + did: washing it to None would mean "everything", and the empty + allowlist meant "nothing" — an agent confidently reporting the + documents don't exist, with no signal the scope was empty.""" seed_doc(store_path, "pi-a", "report.pdf") - fake = fake_model([ - [_call_item("browse_documents", {})], - [_msg_item("done")], - ]) - client.chat_completions("q", doc_id=[]) - outputs = [item["output"] for item in fake.inputs[1] - if item.get("type") == "function_call_output"] - assert json.loads(outputs[-1])["documents"] == [] + with pytest.raises(PageIndexAPIError, match="doc_id is empty"): + client.chat_completions("q", doc_id=[]) + with pytest.raises(PageIndexAPIError, match="doc_id is empty"): + client.as_openai_tools(doc_id=[]) + with pytest.raises(PageIndexAPIError, match="doc_id is empty"): + client.agent_instructions(doc_id=[]) @needs_agents @@ -1108,6 +1160,11 @@ def test_openai_model_resolves_provider_prefixes(): model = local_chat._openai_model("responses", "openai/gpt-5.2") assert isinstance(model, OpenAIResponsesModel) assert str(model.model) == "gpt-5.2" + # litellm/ is routing grammar, not a provider: it strips before the + # provider-prefix guard, so an OpenAI model stays reachable. + model = local_chat._openai_model("responses", "litellm/gpt-5.2") + assert isinstance(model, OpenAIResponsesModel) + assert str(model.model) == "gpt-5.2" @needs_agents @@ -1141,7 +1198,7 @@ def test_envelope_model_strips_litellm_routing_prefix(store_path, fake_model): seed_doc(store_path, "pi-a", "report.pdf") client = PageIndexLocalClient(storage_path=store_path, retrieve_model="anthropic/claude-x") - assert client.retrieve_model == "litellm/anthropic/claude-x" + assert client.retrieve_model == "anthropic/claude-x" fake_model([[_msg_item("ok")]]) result = client.chat_completions("q") assert result["model"] == "anthropic/claude-x" @@ -1170,6 +1227,7 @@ def test_envelope_model_strips_openai_routing_prefix(store_path, fake_model): def test_chat_missing_openai_key_fails_loud(monkeypatch): """A missing backend credential surfaces as the SDK's own error type, like every other precondition on the chat surfaces.""" + pytest.importorskip("litellm") # first import may load a .env; delenv after monkeypatch.delenv("OPENAI_API_KEY", raising=False) for name in ("gpt-4o", "openai/gpt-4o"): with pytest.raises(PageIndexAPIError, match="OPENAI_API_KEY"): @@ -1224,6 +1282,25 @@ async def create(*args, **kwargs): assert result["error"] is None +@needs_agents +def test_responses_envelope_reports_backend_tool_params(client, store_path, + fake_model): + """tool_choice / parallel_tool_calls come from the backend's echo — + the request sends neither, so the envelope must not assume values.""" + seed_doc(store_path, "pi-a", "report.pdf") + fake = fake_model([[_msg_item("ok")]]) + + async def create(*args, **kwargs): + return types.SimpleNamespace(status=None, tool_choice="none", + parallel_tool_calls=False) + + fake._client = types.SimpleNamespace( + responses=types.SimpleNamespace(create=create)) + result = client.responses("q") + assert result["tool_choice"] == "none" + assert result["parallel_tool_calls"] is False + + @needs_agents def test_chat_completions_wraps_framework_errors(client, store_path, fake_model, monkeypatch): @@ -1410,14 +1487,24 @@ async def drive(): @needs_agents def test_stream_abandonment_cancels_pending_turn(client, store_path, - fake_model): + fake_model, monkeypatch): """Closing the iterator cancels the run even while it is awaiting the backend: the blocked turn is torn down (pump thread exits) instead of - running — and billing — to completion in the background.""" + running — and billing — to completion in the background. The pump + thread is tracked directly — a process-global thread count would be + flaky against litellm's background threads.""" import threading - import time as time_mod seed_doc(store_path, "pi-a", "report.pdf") - baseline = threading.active_count() + pumps = [] + real_thread = threading.Thread + + class _Tracking(real_thread): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + if getattr(kwargs.get("target"), "__name__", "") == "pump": + pumps.append(self) + + monkeypatch.setattr(threading, "Thread", _Tracking) fake = fake_model([ [_call_item("get_document", {"doc_name": "report.pdf"})], [_msg_item("The answer")], @@ -1427,11 +1514,9 @@ def test_stream_abandonment_cancels_pending_turn(client, store_path, stream=True, stream_metadata=True) next(stream) # the opening role chunk stream.close() - deadline = time_mod.monotonic() + 3.0 - while (threading.active_count() > baseline - and time_mod.monotonic() < deadline): - time_mod.sleep(0.05) - assert threading.active_count() <= baseline + assert len(pumps) == 1 + pumps[0].join(timeout=3.0) + assert not pumps[0].is_alive() assert fake.deltas_emitted == 0 # turn 2 never produced output @@ -1638,6 +1723,38 @@ def test_backend_connection_reaches_each_engine(monkeypatch): agent = local_chat._openai_agent(None, "responses", "gpt-test", "sys", None, None, backend={"api_key": "k2"}) assert agent.model._client.api_key == "k2" + # the LiteLLM endpoint spelling works on the SDK-constructed door too + agent = local_chat._openai_agent(None, "responses", "gpt-test", "sys", + None, None, + backend={"api_key": "k3", + "api_base": "http://rb"}) + assert str(agent.model._client.base_url).rstrip("/") == "http://rb" + # both endpoint spellings on one merged dict: normalization keeps the + # later (per-call) key instead of the eager nested pop discarding it + agent = local_chat._openai_agent( + None, "chat", "anthropic/claude-x", "sys", None, None, + backend=local_chat._merged_backend( + types.SimpleNamespace(chat_backend={"base_url": "http://client"}), + {"api_base": "http://call"})) + assert agent.model.base_url == "http://call" + + +@needs_anthropic +def test_messages_top_level_cache_control(client, store_path, fake_anthropic): + """The moving breakpoint rides every request so each turn re-reads the + growing conversation; it stands down when the caller's own marks fill + the four-breakpoint budget (a fifth is a live-verified 400).""" + seed_doc(store_path, "pi-a", "report.pdf") + calls = fake_anthropic([ + _anthropic_message([{"type": "text", "text": "ok"}], "end_turn")]) + client.messages("q", model="claude-test", max_tokens=50) + assert calls[0]["cache_control"] == {"type": "ephemeral"} + marked = [{"type": "text", "text": f"b{i}", + "cache_control": {"type": "ephemeral"}} for i in range(3)] + calls = fake_anthropic([ + _anthropic_message([{"type": "text", "text": "ok"}], "end_turn")]) + client.messages("q", model="claude-test", max_tokens=50, system=marked) + assert "cache_control" not in calls[0] def test_merged_backend_precedence(): @@ -1655,6 +1772,9 @@ def test_messages_backend_merges_and_reaches_the_client(client, fake_anthropic, "base_url": "http://x"}) assert real.api_key == "kk" assert str(real.base_url).rstrip("/") == "http://x" + real = local_chat._anthropic_client({"api_key": "kk", + "api_base": "http://y"}) + assert str(real.base_url).rstrip("/") == "http://y" calls = fake_anthropic([ _anthropic_message([{"type": "text", "text": "ok"}], "end_turn")]) @@ -1714,3 +1834,207 @@ def handler(request): client.messages("q", model="claude-sonnet-4-5", extra_headers={"anthropic-beta": "context-1m-2025"}) assert seen["beta"] == "context-1m-2025" + + +@needs_agents +def test_chat_model_settings_request_stream_usage(monkeypatch): + """Without include_usage the streamed run carries no usage at all and + the terminal chunk reports zeros (agents forwards it as + stream_options only on streaming calls).""" + monkeypatch.setattr(local_chat, "_openai_model", lambda *a: None) + agent = local_chat._openai_agent(None, "chat", "gpt-test", "sys", + None, None) + assert agent.model_settings.include_usage is True + + +@needs_anthropic +def test_messages_default_max_tokens_clears_thinking_budget(client, fake_anthropic): + calls = fake_anthropic([ + _anthropic_message([{"type": "text", "text": "a"}], "end_turn"), + ]) + client.messages("q", model="claude-test", + thinking={"type": "enabled", "budget_tokens": 10000}) + assert calls[0]["max_tokens"] == 10000 + 8192 + assert calls[0]["thinking"] == {"type": "enabled", + "budget_tokens": 10000} + calls = fake_anthropic([ # fresh fake: each run closes its client + _anthropic_message([{"type": "text", "text": "b"}], "end_turn"), + ]) + client.messages("q", model="claude-test", max_tokens=11000, + thinking={"type": "enabled", "budget_tokens": 10000}) + assert calls[0]["max_tokens"] == 11000 # explicit value passes through + + +def test_record_chat_finish_records_and_delegates(): + recorded = {} + closed = {"n": 0} + + class Stream: + def __init__(self): + self.chunks = [ + types.SimpleNamespace(choices=[]), + types.SimpleNamespace(choices=[types.SimpleNamespace( + finish_reason="content_filter")]), + ] + + def __aiter__(self): + return self + + async def __anext__(self): + if not self.chunks: + raise StopAsyncIteration + return self.chunks.pop(0) + + async def aclose(self): + closed["n"] += 1 + + async def fetch(*args, **kwargs): + return "shell", Stream() + + model = types.SimpleNamespace(_fetch_response=fetch) + local_chat._record_chat_finish(types.SimpleNamespace(model=model), + recorded) + + async def drive(): + shell, tee = await model._fetch_response() + assert shell == "shell" + async for _chunk in tee: + pass + await tee.aclose() + + asyncio.run(drive()) + assert recorded == {"finish_reason": "content_filter"} + assert closed["n"] == 1 + # A model without the seam: silently a no-op. + local_chat._record_chat_finish( + types.SimpleNamespace(model=types.SimpleNamespace()), {}) + + +@needs_agents +def test_chat_completions_reports_native_finish_reason(client, store_path, + monkeypatch): + seed_doc(store_path, "pi-a", "report.pdf") + + class TruncatingModel(FakeModel): + async def _fetch_response(self, *args, **kwargs): + return types.SimpleNamespace(choices=[ + types.SimpleNamespace(finish_reason="length")]) + + async def get_response(self, *args, **kwargs): + await self._fetch_response() + return await super().get_response(*args, **kwargs) + + async def stream_response(self, *args, **kwargs): + await self._fetch_response() + async for event in super().stream_response(*args, **kwargs): + yield event + + fake = TruncatingModel([[_msg_item("cut ")], [_msg_item("cut ")]]) + monkeypatch.setattr(local_chat, "_openai_model", lambda *a: fake) + result = client.chat_completions("q") + assert result["choices"][0]["finish_reason"] == "length" + chunks = list(client.chat_completions("q", stream=True, + stream_metadata=True)) + assert chunks[-2]["choices"][0]["finish_reason"] == "length" + + +@needs_agents +def test_chat_gate_honors_litellm_routing_and_custom_providers(monkeypatch): + """Mirrors the indexing lane: an explicit litellm/ prefix skips the env + key pre-check, and custom_provider_map providers pass the allowlist; + a name LiteLLM cannot route is still refused up front.""" + pytest.importorskip("litellm") + import litellm # first import may load a .env; delenv after it + from agents.extensions.models.litellm_model import LitellmModel + + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + model = local_chat._openai_model("chat", "litellm/gpt-4o") + assert isinstance(model, LitellmModel) and model.model == "openai/gpt-4o" + + monkeypatch.setattr(litellm, "custom_provider_map", + [{"provider": "my-llm", "custom_handler": object()}]) + model = local_chat._openai_model("chat", "my-llm/model-a") + assert isinstance(model, LitellmModel) and model.model == "my-llm/model-a" + + with pytest.raises(PageIndexAPIError, match="not a LiteLLM provider"): + local_chat._openai_model("chat", "Qwen/my-model") + + +@needs_agents +def test_litellm_model_still_has_the_fetch_response_seam(): + # Guards the private seam _record_chat_finish rides (LitellmModel + # ._fetch_response): a vendor rename turns the recorder into a silent + # no-op and every truncated turn reports finish_reason "stop". + pytest.importorskip("litellm") + from agents.extensions.models.litellm_model import LitellmModel + + assert hasattr(LitellmModel, "_fetch_response") + + +def test_openai_protocol_predicate_follows_litellm_routing(): + pytest.importorskip("litellm") + for name in ("gpt-5", "openai/gpt-4o", "litellm/gpt-4o", + "azure/gpt-4o", "openrouter/openai/gpt-4o", + "deepseek/deepseek-chat", "groq/llama-3.3-70b-versatile", + "xai/grok-3"): + assert local_chat._openai_protocol(name), name + for name in ("anthropic/claude-sonnet-4-5", "gemini/gemini-2.5-pro", + "bedrock/us.anthropic.claude-sonnet-5", + "vertex_ai/claude-sonnet-4-5"): + assert not local_chat._openai_protocol(name), name + + +@needs_agents +def test_chat_backend_without_key_stands_aside_like_index_lane(monkeypatch): + """Any non-empty backend dict suppresses the key pre-check (utils rule).""" + pytest.importorskip("litellm") + import litellm # noqa: F401 — first import may load a .env; delenv after it + from agents.extensions.models.litellm_model import LitellmModel + + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + model = local_chat._openai_model( + "chat", "gpt-test", {"base_url": "http://localhost:9"}) + assert isinstance(model, LitellmModel) + + +@needs_agents +def test_responses_model_marks_caller_owned_transport(monkeypatch): + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + shared = httpx.AsyncClient() + caller = local_chat._openai_model("responses", "gpt-test", + {"http_client": shared}) + assert caller._client._pageindex_caller_http is True + owned = local_chat._openai_model("responses", "gpt-test") + assert owned._client._pageindex_caller_http is False + + async def run(): + await local_chat._aclose_backend(types.SimpleNamespace(model=caller)) + assert not shared.is_closed # caller-owned transport survives + await local_chat._aclose_backend(types.SimpleNamespace(model=owned)) + await shared.aclose() + + asyncio.run(run()) + + +@needs_anthropic +def test_messages_keeps_caller_owned_http_client_open(client): + body = _anthropic_message([{"type": "text", "text": "a"}], "end_turn") + shared = httpx.Client(transport=httpx.MockTransport( + lambda request: httpx.Response(200, json=body))) + out = client.messages("q", model="claude-test", + backend={"api_key": "t", "http_client": shared}) + assert out["content"][0]["text"] == "a" + assert not shared.is_closed + client.messages("q", model="claude-test", + backend={"api_key": "t", "http_client": shared}) + shared.close() + + +@needs_anthropic +def test_messages_without_credentials_raises_contract_error(client, + monkeypatch): + monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False) + monkeypatch.delenv("ANTHROPIC_AUTH_TOKEN", raising=False) + with pytest.raises(PageIndexAPIError, + match="Anthropic backend is not configured"): + client.messages("q", model="claude-test") diff --git a/tests/test_package_surface.py b/tests/test_package_surface.py index 6d985c818..8f68cea70 100644 --- a/tests/test_package_surface.py +++ b/tests/test_package_surface.py @@ -1,4 +1,5 @@ """What `pip install pageindex` exposes: 0.2.8 helper compat and import cost.""" +import os import subprocess import sys @@ -101,3 +102,21 @@ def test_classic_compat_surface_still_reachable(): out = subprocess.run([sys.executable, "-c", probe], capture_output=True, text=True, check=True) assert out.stdout.strip() == "ok" + + +def test_import_leaves_litellm_env_untouched(tmp_path): + """Importing the package must not configure litellm for the host + process; constructing a local client (which will use litellm) does.""" + env = {k: v for k, v in os.environ.items() + if k != "LITELLM_LOCAL_MODEL_COST_MAP"} + probe = ( + "import os, pageindex\n" + "assert 'LITELLM_LOCAL_MODEL_COST_MAP' not in os.environ, " + "'stamped at import'\n" + f"pageindex.PageIndexLocalClient(storage_path={str(tmp_path / 's')!r})\n" + "assert os.environ['LITELLM_LOCAL_MODEL_COST_MAP'] == 'True'\n" + "print('ok')\n" + ) + out = subprocess.run([sys.executable, "-c", probe], env=env, + capture_output=True, text=True, check=True) + assert out.stdout.strip() == "ok" diff --git a/tests/test_page_index_md.py b/tests/test_page_index_md.py index 12aa90890..37b8c2a88 100644 --- a/tests/test_page_index_md.py +++ b/tests/test_page_index_md.py @@ -19,5 +19,33 @@ def test_skips_bold_heading_with_only_whitespace(self): ) +class MarkdownCliTest(unittest.TestCase): + def test_md_cli_runs_without_llm_or_key(self): + """--md_path with no flags makes zero LLM calls: config.yaml's PDF + summary default must not leak in, so the run completes without any + provider key and writes the structure file.""" + import json + import os + import subprocess + import sys + import tempfile + from pathlib import Path + + script = Path(__file__).resolve().parent.parent / "run_pageindex.py" + with tempfile.TemporaryDirectory() as tmp: + md = Path(tmp) / "notes.md" + md.write_text("# Title\n\nIntro.\n\n## Section\n\nBody.\n") + env = {k: v for k, v in os.environ.items() + if k not in ("OPENAI_API_KEY", "CHATGPT_API_KEY")} + env["PYTHONPATH"] = str(script.parent) + res = subprocess.run( + [sys.executable, str(script), "--md_path", str(md)], + capture_output=True, cwd=tmp, env=env, timeout=180) + self.assertEqual(res.returncode, 0, res.stderr.decode()) + out = Path(tmp) / "results" / "notes_structure.json" + self.assertTrue(out.exists(), res.stdout.decode()) + json.loads(out.read_text()) + + if __name__ == "__main__": unittest.main()