From d375c00a5a6a45bb5caeca31790f777e222c154a Mon Sep 17 00:00:00 2001 From: Ray Date: Tue, 11 Aug 2026 21:52:03 +0800 Subject: [PATCH 01/18] feat: local mode for the PageIndex SDK (v0.2.9) (#389) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: add local mode to the PageIndex SDK client One PageIndexClient, two backends. With api_key: the 0.2.x cloud SDK, request for request (with the reviewed fixes: bounded timeouts on JSON endpoints, none on uploads, URL-encoded ids, 401 key hint, empty DELETE body tolerated). Without api_key: the same methods run locally — page_index builds the tree in submit_document (mode="flash" uses PageIndex Flash), documents are stored as plain JSON per doc under storage_path, submit_query is LLM tree search with retrieve_model, and chat_completions answers over the retrieved nodes with OpenAI-style responses and streaming. Local responses mirror the cloud wire shapes verified against the server source: tree nodes rename start_index to page_index and drop end_index, a non-leaf summary becomes prefix_summary, and the metadata/list/delete/ retrieval envelopes match key for key. Cloud-only features (folders, beta_headers, enable_citations) raise instead of pretending. Replaces the demo-only workspace client (index/get_document_structure/ get_page_content had no real users) and its retrieve.py helpers. page_index_main gains an optional logger param so the SDK can keep ./logs out of the caller's working directory; pymupdf import is now lazy (only the optional PyMuPDF parser path needs it). * chore: package pageindex 0.3.0.dev4 for PyPI Poetry packaging for the combined SDK + local pipeline: every production import is a declared dependency (openai and requests join requirements.txt for the same reason), config.yaml and the flash data tables ship in the wheel, the benchmark PNG does not. pymupdf drops to an optional note now that its import is lazy. dev4 follows the already-published 0.3.0.dev1-3; pip still resolves plain 'pip install pageindex' to 0.2.8 until a final 0.3.0 — install with --pre. * docs: add SDK section to README; move the agentic demo onto the SDK The demo keeps its flow and the post-cutoff demo paper, swapping the removed workspace client for PageIndexClient local mode (list_documents for the doc-id cache, get_tree/get_ocr behind the agent tools). The old examples/workspace JSONs demoed the removed format and go with it. * refactor: rebuild cloud_api on the 0.2.8 client text Ray's rule for the cloud half: his 0.2.8 code is the base; a Kylin-lineage change survives only when strictly better — invisible on healthy traffic while fixing a real failure mode. Kept under that bar: request timeouts (dead connections hung forever; uploads still pass none, exactly like 0.2.8), the upload handle closed via with (leaked on request errors), URL-encoded path ids (a crafted id could reroute the URL), the empty-DELETE-body guard, and stream hardening (choices guard, response.close in finally). Reverted as not strictly better: the lowercase summary param (the server accepts both spellings), the 401 message hint (visible text change; AUTH_HINT dropped from errors.py), and the _request/requests.request reorganization — every method body, docstring, and section comment is 0.2.8's text again. diff -w against ../pageindex_sdk/pageindex/client.py now reads as that surgical patch plus plumbing: CloudAPI reads BASE_URL/api_key through the owning client, PageIndexAPIError comes from errors.py, and is_retrieval_ready lives verbatim on PageIndexClient shared by both modes. A mocked-requests harness driving 0.2.8 and this file through 19 identical calls shows the only remaining request-level difference is timeout. * refactor: drop the local retrieval endpoints — cloud-only, deprecated Ray: the cloud already marks POST /retrieval/ and GET /retrieval/{id}/ deprecated in favor of chat completions, so local mode should not grow a fresh implementation of a retiring surface. submit_query/get_retrieval now raise in local mode with a pointer to chat_completions; cloud mode is untouched (the endpoint still works there and 0.2.8 code keeps running). The tree search that backed them stays as chat_completions' retrieval engine; the retrievals/ storage goes away. This also closes the one real cross-mode parity gap — the retrieved_nodes inner shape — by removing its local half. * feat: manifest.json — one-file document listings for the local store Ray wanted a metafile that shows every document in one place instead of per-directory reads. It is a cache, never a second source of truth: writers update it best-effort after save/delete (atomic replace, no locks), and list_metas trusts it only while its id set matches the docs/ directory names — documents are immutable, so matching names imply valid content. Any mismatch (lost concurrent update, crash, corrupt or deleted manifest) rebuilds it from the doc.json files, reading only the missing entries. Incomplete dirs (no doc.json) stay invisible and are never recorded, so a save that completes later is still picked up. 1000-doc listing: 37ms of per-dir reads -> 2.3ms warm (scandir names + one manifest read); one-time rebuild 160ms. * fix: align local doc_id prefix and createdAt format with the cloud Local doc ids now carry the cloud's pi- namespace prefix (random token stays uuid4 hex — nothing parses cuid internals, the prefix is the contract; chat ids already mirrored chatcmpl-). createdAt now matches the server byte for byte: the cloud emits the DB datetime's bare isoformat — naive UTC, second precision — while we emitted microseconds plus +00:00. * feat: optional metadata tags on submit_document (both modes) Ray's call after the alignment review: the metadata key in tree/OCR envelopes and list entries should carry real data, and the only honest way is exposing the field the server already accepts. Cloud mode forwards it as the existing metadata form field; local mode validates it early (a JSON-serializable dict, checked before any LLM spend), stores it in doc.json, and returns it from the same three places. get_document still omits it, mirroring the server, whose metadata-endpoint SQL never selects that column. Scope stays deliberately narrow: set at submit and read back — no metadata_filter, no update API. * docs: state that createdAt is UTC and show how to localize it The value is naive UTC in both modes (the cloud column is timestamp DEFAULT CURRENT_TIMESTAMP on a UTC server, emitted via bare isoformat). Wall-clock display is the consumer's layer: emitting local time under the same format would silently change meaning per machine, and adding an offset marker would break both the byte-format parity and string-order sorting. * fix: createdAt carries milliseconds, matching the cloud's datetime(3) The earlier second-precision alignment was reasoned from the postgres schema file, but production is MySQL (DATABASE_BACKEND defaults to mysql) and its FilePageIndex.createdAt is datetime(3) DEFAULT CURRENT_TIMESTAMP(3) — so the server isoformat()s a millisecond- precision naive-UTC datetime, emitting .XXX000 fractions (bare seconds only when the millisecond happens to be zero). Local now generates through the same mechanism. The docs' 2024-01-15T10:30:00.000Z sample is a JS-style string Python isoformat cannot produce — not evidence. * fix: createdAt at millisecond precision, matching the timestamp(3) column The cloud column is timestamp(3); its datetimes render through bare isoformat() as six fractional digits ending in 000 (or no fraction when the millisecond is exactly zero). Truncate to the millisecond and render the same way, replacing the second-precision guess from cd44ef0. Worth one confirmation against a live cloud response when a key is around. * chore: trim non-essential comments Alignment narration (mirrors the cloud, matches 0.2.8, trailing-period notes) moves out of the code — that rationale lives in the commit history. Kept only constraints the code can't show: the storage commit protocol and manifest trust rule, the id-escape guard, the silent logger's reason to exist, the client-reference indirection, and the tree-node reshape spec. * fix: close the local_store crash and corruption holes found in review Adversarial review reproduced three edge holes in the no-lock design: a torn delete (doc.json unlinked, rmtree unfinished) left a ghost the manifest kept listing forever; one truncated doc.json crashed every listing with a raw JSONDecodeError; and two concurrent deletes of the same id could crash the loser's rmtree. doc.json is now the existence marker in both directions — written last on save, unlinked first on delete — and list_metas serves an entry only after confirming it still exists, instead of trusting the manifest/dir name-set proxy. Atomic writes fsync before replace, closing the power-loss truncation window. An unreadable JSON file is logged and treated as absent; a document meta additionally falls back to the manifest copy (immutable docs make the cache a valid replica). rmtree runs with ignore_errors: the commit point has passed and cleanup is best-effort, which also absorbs the double-delete race. Also restores the reversed-page-range error in the demo's page parser (silently empty since the retrieve.py removal). Warm 1000-doc listing goes 2.3ms -> 8.3ms for the per-doc existence check; full parse remains 37ms. * docs: correct two docstring claims and the pymupdf note is_retrieval_ready reports only API errors as False — transport errors propagate; delete_document may return {} on an empty cloud body; pymupdf is also used by tree_optimize's page loading, not just get_page_tokens. * chore: target 0.2.9 for the local-mode release Ray's call: 0.2.x stays the no-collections line, so local mode ships as 0.2.9 and 0.3.0 stays reserved; plain pip installs never see the 0.3.0.devN pre-releases, and the cookbooks' existing 'pip install --upgrade pageindex' will deliver the new SDK without any --pre instructions. * ci: publish to PyPI on version tags Tag-driven releases: the pushed v-tag is the single source of truth for the version — validated as PEP 440, injected into pyproject, built, and published via OIDC trusted publishing with no stored credentials; a GitHub Release with the artifacts is created alongside. * chore: tighten the store docstring to essentials * fix: contain invalid-UTF-8 corruption; fail loud on unreadable data files The corruption guard caught JSONDecodeError but not the UnicodeDecodeError a torn multi-byte write produces — the exact scenario the guard targets whenever names or descriptions carry non-ASCII text — so that flavor crashed listings and gets raw, and regressed the old manifest read's broader ValueError guard. _read_json now catches ValueError, which covers both. Unreadable tree.json/pages.json under an intact doc.json previously served an empty tree with retrieval_ready true — a silent lie; those paths now raise 'stored document data is unreadable' and is_retrieval_ready honestly reports False. delete_document survives a doc.json tampered into a directory (cleans it, reports not-found) while real unlink failures such as permissions stay loud. * feat: LocalClient and CloudClient for explicit mode selection PageIndexClient(api_key=os.getenv(...)) with an unset variable gets None and silently falls into local mode — methods keep working against local storage on the caller's own LLM bill, the silent mode flip the empty-string guard can't see. The explicit classes close it at the type level: CloudClient raises on a missing key, LocalClient has no api_key parameter at all. Names per Ray. * feat: explicit-mode clients PageIndexCloudClient and PageIndexLocalClient PageIndexClient(api_key=os.getenv(...)) with an unset variable yields None and silently lands in local mode — with both modes fully working, that's a silent mode flip onto the user's own LLM bill. The explicit classes pin the mode at construction: the cloud one refuses a missing or empty key, the local one has no key parameter at all. Names follow the package's PageIndex- prefix convention. * fix local mode edge cases * fix: keep the litellm/ prefix normalization the demo depends on a108c02 added _normalize_retrieve_model because the agentic demo hands client.retrieve_model straight to the OpenAI Agents SDK, which routes a non-OpenAI provider only when the name carries a litellm/ prefix. The normalization sat in client.py, so the demo line never had to change -- and this rewrite dropped the helper while leaving that lone consumer untouched. A retrieve_model like anthropic/claude-sonnet-4-6, the form config.yaml documents, then died at Agent() with "Unknown prefix". Local mode is indifferent to which form it gets: llm_completion and _chat_llm both removeprefix("litellm/"), count_tokens returns the same count either way, and _is_openai_model classifies both as LiteLLM, so _require_llm_key still asks for no OPENAI_API_KEY. Models without a provider path (the packaged gpt-5.4 default) pass through untouched. * fix: restore the published 0.2.8 helper signatures the cookbooks call pyproject.toml makes this repo the source of the PyPI pageindex package, so this utils.py replaces the published 0.2.8 one -- whose helpers are the documented surface of the cookbook notebooks. Three had drifted: remove_fields lost max_len, create_node_mapping lost include_page_ranges/max_page, print_tree lost exclude_fields. Both README-linked notebooks open with `pip install --upgrade pageindex` and pass exactly those kwargs, so tagging v0.2.9 as-is would TypeError every Colab run of vision_RAG_pageindex.ipynb. remove_fields and create_node_mapping readopt the 0.2.8 bodies, strict supersets of the current ones (no in-repo caller passes the new params). print_tree keeps the outline view as its default and routes an explicit exclude_fields= to the 0.2.8 pprint view -- the two versions disagree on what the second positional means (indent vs exclude_fields), and the notebooks pass it by keyword. call_llm, the fifth published name, stays out: nothing imports it -- the one notebook using a call_llm defines its own, with a different signature. * perf: resolve the indexing stack lazily from pageindex/__init__ `import pageindex` eagerly pulled page_index, flash, and tree_optimize -- 0.73s warm, numpy and pypdfium2 in-process -- while the published 0.2.8 package imported in 0.08s on requests+openai alone. A cloud-only SDK user upgrading to 0.2.9 would pay for an indexing stack they never call, on every interpreter start. The eager surface shrinks to client and errors (2ms warm); everything else resolves on first attribute access via PEP 562 and is cached in the module namespace. `pageindex.page_index` stays the function, still shadowing its submodule as the old star-import had it. __all__ now names the public surface, so `from pageindex import *` binds the same working set as before instead of 124 names including stdlib modules. A TYPE_CHECKING block keeps real signatures visible to IDEs. The test suite's sys.modules lookup assumed the eager import chain; it now imports pageindex.page_index explicitly. * fix: wrap PDF read failures in PageIndexAPIError on local submit _extract_page_texts sat one line outside the try that wraps everything else in submit_document, so a corrupt PDF surfaced as a raw PyPDF2.errors.PdfReadError and a password-protected one as FileNotDecryptedError -- while a blank PDF, checked on the very next line, got a clean PageIndexAPIError. Callers handling the SDK error type crashed on exactly the malformed downloads and encrypted files an ingest loop sees most. The extraction gets its own wrap rather than joining the indexer try below, whose except would re-prefix the blank-PDF error into "Failed to submit document: Failed to submit document: ...". FileNotFoundError stays native, asserted by test_submit_rejections as cloud parity. * fix: address code review findings across local mode and publish workflow - _parse_json_reply: switch from extract_json to _reply_json (dead try/except, global None→null substitution, wrong error messages) - client.py: config override filter uses `is not None` instead of truthiness, so empty-string model args are no longer silently dropped - llm_completion/llm_acompletion: raise RuntimeError after retries exhausted instead of returning empty string - _require_llm_key: extend to anthropic/gemini/mistral providers - _chat_llm: reuse _openai_sync_client singleton from utils - _tree_search: drop redundant deepcopy before non-mutating remove_fields - _stream_chunks: move final chunk inside try so usage data is reachable; close stays in finally - _validate_chat_messages: accept system messages, merge them into the internal system prompt for cloud/local parity - _build_chat_context: accept pre-read metas and pass structure through to _tree_search, eliminating double get_meta and double tree.json reads - _index_standard: reuse ConfigLoader from construction; reject empty structure (parity with flash mode) - extract_json: bare except: → except Exception: (no longer swallows KeyboardInterrupt) - Replace all str.removeprefix() with _strip_prefix() helper to restore Python 3.7 compatibility; pyproject.toml back to python >= 3.7 - publish.yml: add test job (py3.10 + py3.13) gating the publish job - Remove unused bare `import pageindex` from tests * fix: tighten CI permissions, timestamp format, import style, and docstring - publish.yml: add permissions: {contents: read} to the test job - local_api: emit 3-digit ms timestamps via isoformat(timespec="milliseconds") - local_api: use relative import for pageindex.utils in _chat_llm - local_api: note the node_summary gate in _format_tree_node docstring - test_client: update createdAt regex to match the new 3-digit format * fix: accept GOOGLE_API_KEY for Gemini; mark pre-releases in GitHub - _require_llm_key: Gemini provider now accepts either GEMINI_API_KEY or GOOGLE_API_KEY (LiteLLM supports both) - publish.yml: set prerelease flag on rc/dev/alpha/beta tags so they are not shown as regular releases on GitHub * fix: preserve print_tree backward compat with 0.2.8 positional call Detect list passed as second positional arg (old 0.2.8 signature) and treat it as exclude_fields instead of indent. * refactor: print_tree param order — exclude_fields second for 0.2.8 compat Move exclude_fields back to the second position (matching 0.2.8) instead of detecting list-as-indent. Recursive call uses indent= keyword arg. * refactor: remove _require_llm_key pre-check entirely Let OpenAI SDK and LiteLLM report their own missing-key errors instead of maintaining a parallel provider-to-env-var map. The except Exception wrapper in chat_completions already converts these to PageIndexAPIError. * fix: let LLM provider errors propagate instead of wrapping them OpenAI/litellm auth, rate-limit, and other provider errors now reach the caller as their original type (e.g. openai.AuthenticationError) instead of being wrapped in PageIndexAPIError. PageIndexAPIError stays reserved for PageIndex's own errors (bad doc_id, invalid params, indexing failures). * refactor: catch only RuntimeError instead of isinstance check on openai Only our own _tree_search logic raises RuntimeError (bad JSON, missing node_list). Provider errors and unexpected bugs propagate naturally. * fix: restore createdAt to 6-digit .177000 format matching cloud DATETIME(3) Revert the timespec="milliseconds" that a linter introduced — it output .177 (3 digits) while the cloud server's isoformat() outputs .177000 (6 digits). Verified against pageindex-compute server/lib/db/planet.py: DATETIME(3) column + bare isoformat() = .177000. * fix: resolve 15 review findings from PR #389 - Rename page_index.py → page_index_classic.py to fix __getattr__ shadowing (function permanently replaced by submodule after import) - Add return_exceptions=True to verify_toc, generate_summaries, and summarize_tree gathers so one LLM failure doesn't abort the batch - Gracefully degrade generate_doc_description to "" on failure instead of discarding the entire completed index - Normalize file_path with str()/expanduser/abspath to support pathlib.Path and tilde paths - Strip text from tree.json at save time; reconstruct from pages.json on read via _load_tree_with_text - Filter empty-string model overrides in PageIndexClient constructor - Guard empty choices list before indexing response.choices[0] - Wrap streaming iteration errors as PageIndexAPIError inside the generator - Catch PermissionError in _read_json alongside FileNotFoundError - Set max_retries=3 for OpenAI and num_retries=3 for litellm in _chat_llm to match the retry behavior of the indexing path - Replace _SilentLogger with logging.getLogger(__name__) - Add .github/workflows/tests.yml for PR and push-to-main test runs * fix: close publish workflow injection and detect flash silent summary failure 1. publish.yml: pass VERSION via os.environ instead of shell interpolation into python -c, eliminating the command injection vector. 2. summarize_tree: raise RuntimeError when every node's summary generation fails (e.g. missing LLM credentials), instead of silently saving a document with all-empty summaries marked as completed. * refactor: make chat_completions cloud-only until agent-based local chat lands The local implementation was a tree-search RAG engine (retrieve prompt + context stuffing) that diverged from the cloud /chat/completions design, where an agent navigates documents through MCP tools. Rather than ship the divergent engine in 0.2.9, remove it; local chat returns as an agent loop built on the agent-tools layer in the 0.2.10 line. Also from the PR #389 review: - get_tree fails loud on unreadable pages.json instead of silently serving textless nodes - drop the wasted deepcopy in get_tree (_format_tree_node builds new dicts) - generate_doc_description catches only RuntimeError so provider errors (bad key, unknown model) propagate instead of storing "" - _read_json treats IsADirectoryError as unreadable - list_metas skips directory names that fail _is_safe_id - constructing with api_key plus local-only args raises PageIndexAPIError (was ValueError) to match the empty-api_key path * fix: verification follow-ups for the chat removal commit - retrieve_model docstring no longer promises the deleted tree-search machinery; the param is reserved for the coming agent-based local chat - empty pages.json ([]) fails loud like unreadable pages: a stored document can never legitimately have zero pages, and get_tree's "nodes always carry text" promise held only for the None case - README: chat example moved to its own cloud-only block so the local quickstart no longer ends in a raise - tests: pin IsADirectoryError loud-fail, list_metas unsafe-name skip, empty-pages loud-fail, and generate_doc_description's swallow-vs- propagate boundary * fix: detect classic-path silent summary failure like flash already does generate_summaries_for_structure absorbed every per-node error to summary: "" with no systemic check, so a bad key or model name produced an all-empty-summary tree with zero errors on the default CLI config (if_add_node_summary: yes, if_add_doc_description: no). Mirror flash's summarize_tree guard: partial failures still absorb, all-failed raises. * fix: accurate wording for retrieve_model docstring and empty-pages error - retrieve_model is not "unused": the agent demo drives its model from client.retrieve_model today; say so instead - an empty pages.json is invalid, not unreadable — dedicated _require_pages raises "stored document has no page content" so the operator reindexes instead of hunting for file corruption * chore: trim non-essential comments, fix three docstring issues Review fixes: - client.py: remove unverified "pending" from get_document status enum - cloud_api.py: update stale "kept line-for-line" module docstring - cloud_api.py: add node_summary to get_tree Args Comment trimming across __init__.py, client.py, cloud_api.py, errors.py, local_api.py, local_store.py — module docstrings shortened, explanatory inline comments removed, multi-line class docstrings collapsed to one line. * feat: add get_page_content and get_tree include_text parameter - client.get_page_content(doc_id, pages): convenience method wrapping get_ocr + page filtering; _parse_pages moved from the demo into client.py - client.get_tree(..., include_text=False): skips node text for structure-only views; local reads the stored tree directly (no text to begin with), cloud strips client-side via remove_fields - demo simplified: tools now call the new client methods directly * simplify: drop redundant try/except in demo get_page_content tool * feat: add get_document_structure convenience method * test: cover get_page_content, get_document_structure, include_text=False * fix: guard against four edge-case crashes found in PR #389 review - Demo: use getattr for retrieve_model so cloud clients don't AttributeError - utils: return empty string when start/end page index is None instead of TypeError - local_store: clean up temp file on _write_json_atomic failure - local_api: re-raise PageIndexAPIError before the catch-all Exception block * chore: trim verbose optional-dep comments in requirements.txt * revert: restore README.md to main — SDK section deferred to next version * fix: eliminate double PDF parse in local standard indexing _extract_page_texts already reads the PDF via PyPDF2; pass the pre-extracted texts as page_list to page_index_main so it skips its own get_page_tokens call. Token counts are computed once via litellm.token_counter in _index_standard. * test: assert page_list is passed and correctly shaped * fix: wire include_text to cloud API and guard get_page_content on processing docs cloud_api.py now sends include_text as a query param so the server can omit node text from tree responses, saving bandwidth. Backward-compatible: old servers ignore the param, client-side remove_fields still strips. get_page_content raises PageIndexAPIError instead of TypeError when the document is still processing (get_ocr returns result: null). --- .github/workflows/publish.yml | 85 +++ .github/workflows/tests.yml | 10 +- .gitignore | 2 + examples/agentic_vectorless_rag_demo.py | 38 +- .../12345678-abcd-4321-abcd-123456789abc.json | 274 ------- examples/workspace/_meta.json | 9 - pageindex/__init__.py | 51 +- pageindex/client.py | 617 +++++++++------ pageindex/cloud_api.py | 433 +++++++++++ pageindex/errors.py | 2 + pageindex/local_api.py | 325 ++++++++ pageindex/local_store.py | 164 ++++ .../{page_index.py => page_index_classic.py} | 23 +- pageindex/retrieve.py | 137 ---- pageindex/utils.py | 139 ++-- pyproject.toml | 53 ++ requirements.txt | 6 +- tests/conftest.py | 52 ++ tests/test_client.py | 712 ++++++++++++++++++ tests/test_issue_163.py | 48 +- tests/test_package_surface.py | 62 ++ tests/test_page_index.py | 30 +- 22 files changed, 2509 insertions(+), 763 deletions(-) create mode 100644 .github/workflows/publish.yml delete mode 100644 examples/workspace/12345678-abcd-4321-abcd-123456789abc.json delete mode 100644 examples/workspace/_meta.json create mode 100644 pageindex/cloud_api.py create mode 100644 pageindex/errors.py create mode 100644 pageindex/local_api.py create mode 100644 pageindex/local_store.py rename pageindex/{page_index.py => page_index_classic.py} (99%) delete mode 100644 pageindex/retrieve.py create mode 100644 pyproject.toml create mode 100644 tests/conftest.py create mode 100644 tests/test_client.py create mode 100644 tests/test_package_surface.py diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml new file mode 100644 index 000000000..e96dc4cca --- /dev/null +++ b/.github/workflows/publish.yml @@ -0,0 +1,85 @@ +name: Publish to PyPI + +# Release flow (the git tag IS the version — nothing to bump in the repo): +# 1. git tag -a v0.2.9 -m "Release 0.2.9" +# 2. git push origin v0.2.9 +# 3. This workflow derives the version from the tag, injects it into +# pyproject.toml, builds, publishes to PyPI via OIDC trusted publishing +# (no stored secret), and creates a GitHub Release with generated notes. +# +# The tag must be a PEP 440 version with a leading `v`: +# v0.2.9 v0.2.9rc1 v0.2.9.dev1 +# PyPI rejects duplicate version uploads, so each tag must be a new version. +# Plain `pip install pageindex` skips dev/rc pre-releases — install one +# explicitly with `pip install pageindex==0.2.9.dev1`. +# +# One-time setup this workflow depends on: +# - PyPI: add a Trusted Publisher on the `pageindex` project pointing at +# repo VectifyAI/PageIndex, workflow `publish.yml`, environment `pypi`. +# - GitHub: create an Environment named `pypi` (Settings -> Environments). + +on: + push: + tags: + - "v*" + +jobs: + test: + runs-on: ubuntu-latest + permissions: + contents: read + strategy: + fail-fast: false + matrix: + python-version: ["3.10", "3.13"] + name: test py${{ matrix.python-version }} + timeout-minutes: 15 + steps: + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + - uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0 + with: + python-version: ${{ matrix.python-version }} + cache: pip + - run: pip install -r requirements.txt pytest + - run: python -m pytest -q + + publish: + needs: test + runs-on: ubuntu-latest + environment: pypi + permissions: + id-token: write # OIDC trusted publishing to PyPI + contents: write # create the GitHub Release + steps: + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + + - uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0 + with: + python-version: "3.12" + + - name: Set version from tag and build + run: | + set -euo pipefail + python -m pip install --upgrade build packaging + VERSION="${GITHUB_REF_NAME#v}" + echo "Publishing version: $VERSION" + # Fail early on a malformed tag instead of publishing a junk version. + export VERSION + python -c "import os; from packaging.version import Version; Version(os.environ['VERSION'])" + # The git tag is the single source of truth; overwrite the static + # placeholder in [tool.poetry] so the built artifacts carry $VERSION. + sed -i "s/^version = .*/version = \"$VERSION\"/" pyproject.toml + grep '^version = ' pyproject.toml + python -m build + + - name: Publish to PyPI + uses: pypa/gh-action-pypi-publish@cef221092ed1bacb1cc03d23a2d87d1d172e277b # release/v1.14.0 + + - name: Create GitHub Release + uses: softprops/action-gh-release@718ea10b132b3b2eba29c1007bb80653f286566b # v3.0.1 + with: + tag_name: ${{ github.ref_name }} + name: ${{ github.ref_name }} + prerelease: ${{ contains(github.ref_name, 'rc') || contains(github.ref_name, 'dev') || contains(github.ref_name, 'alpha') || contains(github.ref_name, 'beta') }} + generate_release_notes: true + files: dist/* diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index fc25d6c45..d8d9dbb38 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -20,22 +20,16 @@ jobs: fail-fast: false matrix: python-version: ["3.10", "3.13"] - # "without" proves the package works with no agent framework - # installed; "with" covers the framework integration paths. agent-frameworks: [without, with] name: py${{ matrix.python-version }} (${{ matrix.agent-frameworks }} frameworks) timeout-minutes: 15 steps: - - uses: actions/checkout@v4 - - uses: actions/setup-python@v5 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + - uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0 with: python-version: ${{ matrix.python-version }} cache: pip - # requirements.txt works on every branch; the packaging metadata - # (pyproject) does not exist on all of them. - run: pip install -r requirements.txt pytest - if: matrix.agent-frameworks == 'with' run: pip install openai-agents claude-agent-sdk - # python -m pytest puts the repo root on sys.path, so the in-repo - # `pageindex` package is imported without an install step. - run: python -m pytest -q diff --git a/.gitignore b/.gitignore index 23d6b5655..5193735ca 100644 --- a/.gitignore +++ b/.gitignore @@ -4,3 +4,5 @@ __pycache__ .env* .venv/ logs/ +.pageindex/ +dist/ diff --git a/examples/agentic_vectorless_rag_demo.py b/examples/agentic_vectorless_rag_demo.py index b4ed9c2f8..4fe5f179f 100644 --- a/examples/agentic_vectorless_rag_demo.py +++ b/examples/agentic_vectorless_rag_demo.py @@ -1,10 +1,10 @@ """ Agentic Vectorless RAG with PageIndex - Demo -A simple example of building a document QA agent with self-hosted PageIndex -and the OpenAI Agents SDK. Instead of vector similarity search and chunking, -PageIndex builds a hierarchical tree index and uses agentic LLM reasoning for -human-like, context-aware retrieval. +A simple example of building a document QA agent with the PageIndex SDK in +local mode and the OpenAI Agents SDK. Instead of vector similarity search and +chunking, PageIndex builds a hierarchical tree index and uses agentic LLM +reasoning for human-like, context-aware retrieval. Agent tools: - get_document() — document metadata (status, page count, etc.) @@ -12,11 +12,11 @@ - get_page_content() — retrieve text content of specific pages Steps: - 1 — Index a PDF and view its tree structure index + 1 — Index a PDF locally and view its tree structure index 2 — View document metadata 3 — Ask a question (agent reasons over the index and auto-calls tools) -Requirements: pip install openai-agents +Requirements: pip install openai-agents; OPENAI_API_KEY in the environment. """ import sys import json @@ -39,12 +39,12 @@ _EXAMPLES_DIR = Path(__file__).parent PDF_PATH = _EXAMPLES_DIR / "documents" / "attention-residuals.pdf" -WORKSPACE = _EXAMPLES_DIR / "workspace" +STORAGE_PATH = _EXAMPLES_DIR / ".pageindex" AGENT_SYSTEM_PROMPT = """ You are PageIndex, a document QA assistant. TOOL USE: -- Call get_document() first to confirm status and page/line count. +- Call get_document() first to confirm status and page count. - Call get_document_structure() to identify relevant page ranges. - Call get_page_content(pages="5-7") with tight ranges; never fetch the whole document. - Before each tool call, output one short sentence explaining the reason. @@ -62,27 +62,26 @@ def query_agent(client: PageIndexClient, doc_id: str, prompt: str, verbose: bool @function_tool def get_document() -> str: """Get document metadata: status, page count, name, and description.""" - return client.get_document(doc_id) + return json.dumps(client.get_document(doc_id)) @function_tool def get_document_structure() -> str: """Get the document's full tree structure (without text) to find relevant sections.""" - return client.get_document_structure(doc_id) + return json.dumps(client.get_document_structure(doc_id), ensure_ascii=False) @function_tool def get_page_content(pages: str) -> str: """ - Get the text content of specific pages or line numbers. + Get the text content of specific pages. Use tight ranges: e.g. '5-7' for pages 5 to 7, '3,8' for pages 3 and 8, '12' for page 12. - For Markdown documents, use line numbers from the structure's line_num field. """ - return client.get_page_content(doc_id, pages) + return json.dumps(client.get_page_content(doc_id, pages), ensure_ascii=False) agent = Agent( name="PageIndex", instructions=AGENT_SYSTEM_PROMPT, tools=[get_document, get_document_structure, get_page_content], - model=client.retrieve_model, + model=getattr(client, "retrieve_model", None), # model_settings=ModelSettings(reasoning={"effort": "low", "summary": "auto"}), # Uncomment to enable reasoning ) @@ -152,24 +151,25 @@ async def _run(): f.write(chunk) print("Download complete.\n") - # Setup - client = PageIndexClient(workspace=WORKSPACE) + # Setup: local mode — no PageIndex API key needed, your LLM key does the work + client = PageIndexClient(storage_path=str(STORAGE_PATH)) # Step 1: Index PDF and view tree structure print("=" * 60) print("Step 1: Index PDF and view tree structure") print("=" * 60) doc_id = next( - (did for did, doc in client.documents.items() if doc.get('doc_name') == PDF_PATH.name), + (doc["id"] for doc in client.list_documents(limit=100)["documents"] + if doc["name"] == PDF_PATH.name), None, ) if doc_id: print(f"\nLoaded cached doc_id: {doc_id}") else: - doc_id = client.index(PDF_PATH) + doc_id = client.submit_document(str(PDF_PATH))["doc_id"] print(f"\nIndexed. doc_id: {doc_id}") print("\nTree Structure (top-level sections):") - structure = json.loads(client.get_document_structure(doc_id)) + structure = client.get_tree(doc_id, node_summary=True)["result"] utils.print_tree(structure) # Step 2: View document metadata diff --git a/examples/workspace/12345678-abcd-4321-abcd-123456789abc.json b/examples/workspace/12345678-abcd-4321-abcd-123456789abc.json deleted file mode 100644 index 23351d5c5..000000000 --- a/examples/workspace/12345678-abcd-4321-abcd-123456789abc.json +++ /dev/null @@ -1,274 +0,0 @@ -{ - "id": "12345678-abcd-4321-abcd-123456789abc", - "type": "pdf", - "path": "../documents/attention-residuals.pdf", - "doc_name": "attention-residuals.pdf", - "doc_description": "This document introduces \"Attention Residuals\" (AttnRes) and its scalable variant \"Block AttnRes,\" novel mechanisms for replacing fixed residual accumulation in neural networks with learned, input-dependent depth-wise attention, addressing limitations of standard residual connections while optimizing memory, computation, and scalability for large-scale training and inference.", - "page_count": 21, - "structure": [ - { - "title": "Preface", - "node_id": "0000", - "start_index": 1, - "end_index": 2, - "summary": "The partial document introduces \"Attention Residuals\" (AttnRes), a novel approach to replace fixed residual accumulation in large language models (LLMs) with learned, input-dependent softmax attention over preceding layer outputs. This method addresses issues like uncontrolled hidden-state growth and dilution of layer contributions caused by standard residual connections with PreNorm. To enhance scalability, the document proposes \"Block AttnRes,\" which partitions layers into blocks and applies attention at the block level, reducing memory and communication overhead while maintaining performance gains. The document highlights system optimizations, such as cross-stage caching and a two-phase computation strategy, to make Block AttnRes efficient for large-scale training. Experiments confirm consistent improvements across model sizes, with AttnRes mitigating PreNorm dilution, leading to more uniform output magnitudes, gradient distributions, and better downstream task performance. Key contributions include the introduction of AttnRes and Block AttnRes, scalable infrastructure optimizations, and comprehensive evaluations demonstrating their effectiveness." - }, - { - "title": "Introduction", - "node_id": "0001", - "start_index": 2, - "end_index": 3, - "summary": "The partial document introduces \"Attention Residuals\" (AttnRes), a novel mechanism that replaces fixed residual accumulation in deep networks with learned softmax attention over depth. It highlights the limitations of standard residual connections, such as uniform layer contributions, irreversible information loss, and output growth, and draws parallels between depth-wise accumulation and sequence modeling in RNNs. AttnRes enables selective, content-aware aggregation of information across layers using attention weights, addressing these limitations. The document also proposes a scalable variant, Block AttnRes, which reduces memory and communication overhead for large-scale training. Key contributions include the development of AttnRes and Block AttnRes, system optimizations for scalability, and comprehensive evaluations demonstrating improved training dynamics, bounded hidden-state magnitudes, and better gradient distribution. The approach is validated through scaling law experiments, ablations, and downstream benchmarks, showing consistent performance improvements over standard residual connections." - }, - { - "title": "Motivation", - "node_id": "0002", - "start_index": 3, - "end_index": 3, - "summary": "The partial document discusses the concept of Attention Residuals in the context of deep learning models, particularly Transformers. It begins by introducing the notation and structure of input sequences and layers in a Transformer model. The document then explains residual learning, highlighting its importance in training deep networks by enabling gradients to bypass transformations through identity mapping. It expands on the limitations of traditional residual connections and highway networks, such as lack of selective access to earlier layer outputs, irreversible information loss, and output growth issues that destabilize training. To address these limitations, the document proposes Attention Residuals (AttnRes), a mechanism inspired by the duality of time and depth in sequence modeling. This approach introduces layer-specific attention weights to selectively aggregate information from all preceding layers, offering a unified view of time and depth while maintaining computational feasibility." - }, - { - "title": "Attention Residuals: A Unified View of Time and Depth", - "node_id": "0003", - "start_index": 3, - "end_index": 4, - "summary": "The partial document discusses the concept of \"Attention Residuals\" as a mechanism to address limitations in training deep networks with residual connections. It begins by explaining residual learning, its benefits in gradient flow, and its limitations, such as lack of selective access, irreversible information loss, and output growth. The document introduces \"Attention Residuals\" (AttnRes), which generalizes residual connections by allowing layers to selectively aggregate information from all preceding layers using attention mechanisms. It describes \"Full Attention Residuals,\" which compute attention weights over depth with softmax normalization, and highlights their computational and memory overhead. To address scalability challenges, the document proposes \"Block Attention Residuals,\" which partition layers into blocks, reducing memory and communication overhead by applying attention at the block level. The text also outlines the intra-block accumulation process and its efficiency in distributed training setups.", - "nodes": [ - { - "title": "Full Attention Residuals", - "node_id": "0004", - "start_index": 4, - "end_index": 4, - "summary": "The partial document discusses \"Attention Residuals\" in neural networks, focusing on two main approaches: Full Attention Residuals and Block Attention Residuals. \n\n1. **Full Attention Residuals**: This method computes attention weights using a kernel function with RMS normalization to prevent large-magnitude outputs from dominating. It introduces no additional memory overhead during vanilla training but incurs communication and memory overhead in distributed training due to the need to retain and transmit layer outputs across stages. A blockwise optimization strategy is proposed to reduce memory I/O by batching attention computation within groups of layers.\n\n2. **Block Attention Residuals**: This approach partitions layers into blocks, reducing memory and communication overhead by summing layer outputs within each block and applying attention only to block-level representations. This reduces the complexity from O(Ld) to O(Nd), where N is the number of blocks. The method ensures normalization to avoid biases from magnitude differences between blocks.\n\nThe document highlights the trade-offs between memory, computation, and communication overheads in these methods and introduces strategies to optimize their efficiency in distributed training setups." - }, - { - "title": "Block Attention Residuals", - "node_id": "0005", - "start_index": 4, - "end_index": 5, - "summary": "The partial document discusses \"Attention Residuals,\" focusing on two main variants: Full Attention Residuals (Full AttnRes) and Block Attention Residuals (Block AttnRes). \n\n1. **Full Attention Residuals (Full AttnRes):**\n - Defines attention weights using a kernel function with RMS normalization to prevent large-magnitude outputs from dominating.\n - Requires O(L²d) arithmetic and O(Ld) memory, with no additional memory overhead during vanilla training.\n - Highlights challenges in large-scale training, such as memory and communication overhead under pipeline parallelism.\n - Introduces blockwise optimization to reduce memory I/O but notes that cross-stage communication remains a bottleneck.\n\n2. **Block Attention Residuals (Block AttnRes):**\n - Partitions layers into blocks, reducing memory and communication overhead from O(Ld) to O(Nd) by summing layer outputs within blocks and applying attention over block-level representations.\n - Provides PyTorch-style pseudocode for implementation, detailing intra-block accumulation and inter-block attention mechanisms.\n - Improves efficiency by reducing memory and computation requirements, with block count N interpolating between Full AttnRes (N=L) and standard residual connections (N=1).\n - Enhances inference latency and bounds KV cache size through blockwise optimization.\n\nThe document also addresses infrastructure challenges for large-scale training, emphasizing the need to manage communication overhead and optimize system design for block-based attention mechanisms." - } - ] - }, - { - "title": "Infrastructure Design", - "node_id": "0006", - "start_index": 5, - "end_index": 6, - "summary": "The partial document describes the concept and implementation of Block Attention Residuals (Block AttnRes), a mechanism designed to improve memory and computational efficiency in attention-based models. It introduces inter-block attention, where attention is computed over block representations and partial sums, reducing memory and computation from O(L) and O(L²) to O(N) and O(N²), respectively. The document provides PyTorch-style pseudocode for the implementation, detailing how block representations and partial sums are managed across layers. It highlights the efficiency benefits of using block representations instead of individual outputs, with empirical findings suggesting that a block count of N≈8 balances performance and resource usage. \n\nThe document also addresses infrastructure challenges in large-scale training and inference. It discusses pipeline communication optimizations, such as cross-stage caching, to reduce redundant data transmission and improve efficiency during distributed training. For inference, it proposes a two-phase computation strategy and memory-efficient prefilling to handle long-context scenarios. An example of cache-based pipeline communication is provided, illustrating how caching minimizes communication overhead in distributed systems.", - "nodes": [ - { - "title": "Training", - "node_id": "0007", - "start_index": 6, - "end_index": 7, - "summary": "The partial document discusses the optimization of Attention Residuals (AttnRes) in training and inference for large-scale distributed systems. It introduces cross-stage caching to address communication and memory overheads in pipeline parallelism, reducing redundant data transmission and improving efficiency. The document details a two-phase computation strategy for Block AttnRes, which includes parallel inter-block attention and sequential intra-block attention with online softmax merging. This approach minimizes memory access and I/O overhead while maintaining a low training overhead. Additionally, it highlights the memory-efficient prefilling scheme for long-context inputs and explains how Block AttnRes compresses representations to reduce storage requirements. The document also provides algorithmic details and performance improvements in both training and inference scenarios." - }, - { - "title": "Inference", - "node_id": "0008", - "start_index": 7, - "end_index": 8, - "summary": "The partial document describes the technical details and implementation of Attention Residuals (AttnRes) in neural network architectures. It introduces a two-phase computation strategy for block-based attention, optimizing memory and computational efficiency. Phase 1 handles parallel inter-block attention, while Phase 2 processes sequential intra-block attention with an online softmax merge. The document highlights memory overhead reduction through cross-stage caching, sequence-sharded prefilling, and kernel fusion, achieving minimal training and inference latency overhead. It compares memory access costs across different residual mechanisms and demonstrates the efficiency of AttnRes, particularly in Block AttnRes, which compresses block representations. Experimental results show that AttnRes improves scaling behavior and validation loss compared to baseline models, with negligible parameter overhead and consistent performance gains across compute ranges." - } - ] - }, - { - "title": "Experiments", - "node_id": "0009", - "start_index": 8, - "end_index": 8, - "summary": "The partial document discusses the technical details and performance of the Attention Residuals (AttnRes) mechanism in transformer architectures. It highlights the memory efficiency and reduced inference latency of AttnRes compared to prior residual mechanisms like mHC. The document provides a breakdown of memory access costs for different schemes, emphasizing the two-phase inference schedule of AttnRes and its memory-efficient prefilling strategy, which significantly reduces memory overhead through sharding and chunked prefill techniques. It also describes the integration of AttnRes into a Mixture-of-Experts (MoE) Transformer architecture, detailing its minimal parameter addition and initialization strategy to ensure stable training. Additionally, the document presents experimental results, including scaling laws and validation loss comparisons across model variants, demonstrating that AttnRes achieves consistently lower loss while maintaining similar scaling behavior to the baseline.", - "nodes": [ - { - "title": "Scaling Laws", - "node_id": "0010", - "start_index": 8, - "end_index": 9, - "summary": "The partial document discusses the implementation and evaluation of Attention Residuals (AttnRes) in transformer architectures. It highlights the memory efficiency and reduced inference latency of AttnRes compared to prior residual mechanisms like mHC. The document introduces a two-phase inference schedule for AttnRes, optimizing memory access costs and reducing per-device memory usage through sharding and chunked prefill techniques. It describes the integration of AttnRes into a Mixture-of-Experts (MoE) Transformer architecture, maintaining minimal parameter overhead and ensuring stable training through specific initialization strategies. Experiments compare scaling laws and validation loss across model sizes, showing that both Full and Block AttnRes outperform baselines and mHC in terms of loss and compute efficiency. The main results include training recipes for large-scale models, leveraging hybrid attention mechanisms and progressive sequence length extension without additional modifications." - }, - { - "title": "Main Results", - "node_id": "0011", - "start_index": 9, - "end_index": 11, - "summary": "The partial document discusses the concept of Attention Residuals (AttnRes) in transformer models, comparing its performance and efficiency against baseline models and other methods. Key points include:\n\n1. **Model Configurations and Validation Loss**: Comparison of Baseline, Block AttnRes, Full AttnRes, and mHC(-lite) models across various configurations, showing that AttnRes consistently achieves lower validation loss, with Block AttnRes closely tracking Full AttnRes.\n\n2. **Scaling Laws**: Analysis of scaling behavior, demonstrating that Block AttnRes achieves significant compute efficiency and narrows the performance gap with Full AttnRes at larger scales.\n\n3. **Training Recipe**: Description of the training process for large models, including pre-training and mid-training phases, use of hybrid attention mechanisms, and progressive sequence length extension.\n\n4. **Training Dynamics**: Examination of validation loss, output magnitude, and gradient magnitude during training, highlighting how Block AttnRes mitigates issues like PreNorm dilution and uneven gradient flow.\n\n5. **Downstream Performance**: Evaluation of AttnRes on various benchmarks for language understanding, reasoning, and code/math tasks, showing consistent improvements over the baseline, particularly in multi-step reasoning and compositional tasks.\n\n6. **Ablation Study**: Validation of key design choices in AttnRes, comparing it with prior methods like DenseFormer and mHC. Full AttnRes achieves the best performance, while Block AttnRes offers a memory-efficient trade-off with competitive results.\n\n7. **Cross-Layer Access**: Exploration of different granularities of cross-layer access, with Block AttnRes providing an effective balance between performance and memory efficiency, and Full AttnRes offering the best results at higher memory costs." - }, - { - "title": "Ablation Study", - "node_id": "0012", - "start_index": 11, - "end_index": 12, - "summary": "The partial document focuses on the development and evaluation of Attention Residuals (AttnRes), a novel mechanism for improving Transformer models. Key points include:\n\n1. **Ablation Studies**: The document evaluates the impact of various design choices in AttnRes, such as input-dependent queries, input-independent mixing, softmax vs. sigmoid, multihead attention, and RMSNorm. Results show that input-dependent queries and RMSNorm improve performance, while softmax outperforms sigmoid due to sharper selection.\n\n2. **Comparison with Prior Methods**: AttnRes is compared against baseline PreNorm, DenseFormer, and mHC. AttnRes achieves superior performance, with Full AttnRes and Block AttnRes showing significant improvements in validation loss.\n\n3. **Cross-Layer Access**: Different granularities of cross-layer access are analyzed. Full AttnRes achieves the best performance, while Block AttnRes offers a memory-efficient trade-off. Sliding-window aggregation (SWA) is less effective, highlighting the importance of selectively accessing distant layers.\n\n4. **Performance on Benchmarks**: AttnRes outperforms the baseline on various benchmarks, particularly in multi-step reasoning tasks, code generation, and knowledge-oriented tasks, demonstrating its effectiveness in compositional tasks.\n\n5. **Optimal Architecture Analysis**: The study explores how AttnRes reshapes architectural scaling under fixed compute and parameter budgets. AttnRes favors deeper models with a shift in the optimal depth–width–attention trade-off, achieving consistently lower losses across configurations compared to the baseline.\n\n6. **Validation Loss Trends**: The document provides detailed validation loss trends across different configurations and block sizes, showing graceful degradation with increasing block size and highlighting the efficiency of finer-grained configurations." - }, - { - "title": "Analysis", - "node_id": "0013", - "start_index": 12, - "end_index": 12, - "summary": "The partial document discusses the evaluation and analysis of Attention Residuals (AttnRes) in Transformer architectures. Key points include:\n\n1. **Architecture Sweep**: A study under fixed compute and parameter budgets to analyze validation loss across different configurations of model depth (dmodel/Lb) and attention heads (H/Lb). AttnRes consistently outperforms the baseline in all configurations, with a notable shift in optimal depth from dmodel/Lb ≈ 60 (baseline) to dmodel/Lb ≈ 45 (AttnRes).\n\n2. **Component Design Ablations**:\n - **Input-dependent query**: Improves performance but adds computational complexity.\n - **Input-independent mixing**: Degrades performance compared to learned queries.\n - **Softmax vs. Sigmoid**: Softmax performs better due to sharper selection among sources.\n - **Multihead Attention**: Depth aggregation across heads reduces performance, indicating uniform depth-wise mixtures are optimal.\n - **RMSNorm on Keys**: Removing RMSNorm negatively impacts performance, especially for block-level representations, by preventing bias in attention weights.\n\n3. **Optimal Architecture Analysis**: Investigates how AttnRes influences depth–width–attention trade-offs under fixed compute and parameter constraints. AttnRes favors deeper models and achieves lower loss compared to conventional Transformer designs.", - "nodes": [ - { - "title": "Optimal Architecture", - "node_id": "0014", - "start_index": 12, - "end_index": 13, - "summary": "The partial document discusses the concept of Attention Residuals (AttnRes) in Transformer architectures, focusing on their design, performance, and analysis. Key points include:\n\n1. **Component Design Ablations**: The document evaluates various modifications to the attention mechanism, such as input-dependent queries, input-independent mixing, softmax vs. sigmoid, multihead attention, and RMSNorm on keys. These experiments highlight the impact of each component on performance, with findings such as the importance of softmax for competitive normalization and RMSNorm for preventing bias in attention weights.\n\n2. **Optimal Architecture Analysis**: A controlled study under fixed compute and parameter budgets examines how AttnRes reshapes architectural scaling preferences. Results show that AttnRes favors deeper, narrower networks compared to baseline Transformers, achieving lower validation loss across configurations. The optimal configuration shifts to a lower dmodel/Lb ratio, indicating better exploitation of depth.\n\n3. **Learned AttnRes Patterns**: Visualization of learned attention weights reveals key insights:\n - Preserved locality with layers attending strongly to immediate predecessors while forming selective skip connections.\n - Layer specialization, with embeddings retaining weight and distinct patterns in pre-attention and pre-MLP layers.\n - Block AttnRes effectively preserves structural patterns while acting as implicit regularization.\n\n4. **Performance Trends**: AttnRes consistently outperforms the baseline across configurations, with lower validation loss and sharper, more decisive weight distributions in block attention settings." - }, - { - "title": "Analyzing Learned AttnRes Patterns", - "node_id": "0015", - "start_index": 13, - "end_index": 14, - "summary": "The partial document discusses Attention Residuals (AttnRes) in deep learning models, focusing on their structure, behavior, and benefits. Key points include:\n\n1. **Depth-wise Attention Weight Distributions**: Analysis of weight distributions in a 16-head model with full and block Attention Residuals, highlighting diagonal dominance (locality), learned skip connections, and sharper weight distributions in block settings.\n\n2. **Learned AttnRes Patterns**: Observations include preserved locality, layer specialization, and the ability of block AttnRes to maintain essential information pathways while acting as implicit regularization.\n\n3. **Comparison of Residual Update Mechanisms**: A detailed comparison of various residual connection methods, including their update rules, weight types (fixed, learned, or dynamic), and source access.\n\n4. **Sequence-Depth Duality**: Exploration of the analogy between residual connections and RNNs, emphasizing how AttnRes replaces depth-wise recurrence with direct cross-layer attention for improved information propagation.\n\n5. **Residual Connections as Structured Matrices**: Formalization of residual connections as depth mixing matrices, comparing different methods based on weight generation and structural constraints.\n\nThe document emphasizes the advantages of AttnRes in leveraging depth, preserving structure, and enabling efficient information flow across layers." - } - ] - } - ] - }, - { - "title": "Discussions", - "node_id": "0016", - "start_index": 14, - "end_index": 14, - "summary": "The partial document discusses various residual update mechanisms in neural network architectures, comparing their update rules, weight types (fixed, learned-static, or input-dependent), and sources of earlier representations. It categorizes methods into single-state recurrence, multi-state recurrence, and cross-layer access, providing examples like Residual, ReZero, LayerScale, Highway, DeepNorm, KEEL, DenseNet, DenseFormer, MRLA, and AttnRes. The document explores the sequence-depth duality, drawing parallels between residual connections and recurrent neural networks (RNNs), and highlights how AttnRes replaces depth-wise recurrence with direct cross-layer attention. Additionally, it formalizes residual connections as structured matrices, introducing a depth mixing matrix to analyze how different methods aggregate outputs from previous layers, and discusses their weight generation and rank constraints.", - "nodes": [ - { - "title": "Sequence-Depth Duality", - "node_id": "0017", - "start_index": 14, - "end_index": 14, - "summary": "The partial document discusses various residual update mechanisms in neural network architectures, comparing their update rules, weight types (fixed, learned-static, or input-dependent), and sources of earlier representations. It categorizes methods into single-state recurrence, multi-state recurrence, and cross-layer access, providing examples like Residual, ReZero, LayerScale, Highway, DeepNorm, KEEL, DenseNet, DenseFormer, MRLA, and AttnRes. The document explores the sequence-depth duality, drawing parallels between residual connections and recurrent neural networks (RNNs), and highlights how AttnRes replaces depth-wise recurrence with cross-layer attention. Additionally, it formalizes residual connections as structured matrices, introducing a depth mixing matrix to analyze how different methods aggregate outputs from previous layers, and discusses their weight generation and rank constraints." - }, - { - "title": "Residual Connections as Structured Matrices", - "node_id": "0018", - "start_index": 14, - "end_index": 16, - "summary": "The partial document discusses various residual update mechanisms in neural networks, comparing their weight types (fixed, learned, or input-dependent) and source accessibility. It introduces AttnRes, a novel approach that replaces fixed residual accumulation with learned, input-dependent depth-wise attention, inspired by the sequence-depth duality. The document explores structured matrix perspectives, showing how residual variants can be viewed as depth-wise linear attention. It highlights the limitations of existing methods like single-state recurrence and multi-state recurrence, and contrasts them with AttnRes, which provides selective access to earlier-layer outputs. The paper also introduces Block AttnRes, a scalable variant that partitions layers into blocks to reduce memory and computational overhead while retaining performance gains. Empirical results validate the effectiveness of AttnRes and Block AttnRes, with discussions on normalization, scaling, depth stability, and cross-layer connectivity. The document concludes by emphasizing the practicality and scalability of Block AttnRes for large-scale models." - }, - { - "title": "Prior Residuals as Depth-Wise Linear Attention", - "node_id": "0019", - "start_index": 16, - "end_index": 16, - "summary": "The partial document discusses the concept of Attention Residuals (AttnRes), which replaces traditional residual accumulation with learned, input-dependent depth-wise attention. It explores the structured-matrix perspective, sequence-depth duality, and the role of state expansion in depth-wise linear attention. The document addresses challenges in normalization, scaling, and depth stability, comparing PreNorm and PostNorm approaches and introducing AttnRes as a solution to avoid cumulative magnitude growth and gradient vanishing. It highlights multi-state recurrence methods, cross-layer connectivity strategies, and the advantages of AttnRes in selectively accessing earlier-layer outputs. The introduction of Block AttnRes is proposed to address memory constraints in large-scale models by partitioning layers into blocks, reducing computational overhead while maintaining performance. Empirical studies validate the effectiveness of AttnRes and Block AttnRes, with scalability and efficiency improvements highlighted as key contributions." - } - ] - }, - { - "title": "Related Work", - "node_id": "0020", - "start_index": 16, - "end_index": 16, - "summary": "The partial document discusses the concept of Attention Residuals (AttnRes) and its application as depth-wise attention in neural networks. It explores the structured-matrix perspective, highlighting how existing residual variants can be interpreted as linear attention mechanisms over the depth axis. The document addresses challenges in normalization, scaling, and depth stability, comparing PreNorm and PostNorm approaches and introducing AttnRes as a solution to avoid cumulative magnitude growth and gradient vanishing. It also examines multi-state recurrence methods, cross-layer connectivity strategies, and their limitations, proposing AttnRes as a method that selectively aggregates earlier-layer outputs with softmax-normalized, input-dependent weights. The introduction of Block AttnRes is detailed as a scalable alternative to Full AttnRes, reducing memory and computational overhead by partitioning layers into blocks while maintaining performance gains. Empirical validation and practical implementation strategies, such as cross-stage caching and two-phase computation, are also discussed." - }, - { - "title": "Conclusion", - "node_id": "0021", - "start_index": 16, - "end_index": 20, - "summary": "The partial document discusses the concept of Attention Residuals (AttnRes), introducing a novel approach to residual accumulation in neural networks by leveraging depth-wise attention mechanisms. It explores the sequence-depth duality, interpreting residual variants as linear attention over the depth axis. The document highlights the challenges of normalization placement and gradient propagation in standard residual updates, comparing PreNorm and PostNorm methods, and presents AttnRes as a solution that avoids cumulative magnitude growth and gradient vanishing. It also examines multi-state recurrence and cross-layer connectivity, contrasting AttnRes with existing methods like Hyper-Connections, DenseNet, and MUDDFormer, emphasizing its selective access to earlier-layer outputs and efficient scaling. The introduction of Block AttnRes addresses memory constraints by partitioning layers into blocks, reducing computational overhead while maintaining performance. Empirical studies validate the scalability and efficiency of AttnRes, with future directions focusing on finer-grained blocking and hardware advancements." - }, - { - "title": "Contributions", - "node_id": "0022", - "start_index": 20, - "end_index": 21, - "summary": "The partial document discusses the concept of \"Attention Residuals\" and provides a technical explanation of optimized inference input/output (I/O) for Full Attention Residuals. It highlights the inefficiencies of a naïve implementation, where memory traffic scales linearly with depth, and introduces a two-phase scheduling approach to reduce I/O costs. The document explains the partitioning of layers into blocks and details the two phases: Phase 1 (batched inter-block attention) and Phase 2 (sequential intra-block attention). It provides mathematical formulations for read and write costs during these phases and demonstrates how batching inter-block reads reduces per-layer I/O complexity from O(L) to O(S+N). The approach maintains the model architecture while optimizing inference efficiency. Additionally, the document lists the contributors to the work, with equal contributions noted for some authors." - }, - { - "title": "Optimized Inference I/O for Full Attention Residuals", - "node_id": "0023", - "start_index": 21, - "end_index": 21, - "summary": "The partial document discusses an optimized inference I/O strategy for Full Attention Residuals (Full AttnRes) to reduce memory traffic, which scales linearly with model depth in a naïve implementation. It introduces a two-phase scheduling approach for inference, dividing the model into blocks to batch inter-block and intra-block computations. Phase 1 handles batched inter-block attention, reducing redundant memory reads by reusing key-value pairs across layers within a block. Phase 2 processes sequential intra-block dependencies. The document provides detailed calculations for read and write costs during both phases, showing that the proposed method reduces per-layer I/O complexity from O(L) to O(S+N), where S is the block size and N is the number of blocks. The approach maintains the model architecture while optimizing memory efficiency during inference." - } - ], - "pages": [ - { - "page": 1, - "content": "ATTENTIONRESIDUALS\nTECHNICALREPORT OFATTENTIONRESIDUALS\nKimi Team\n/gtbhttps://github.com/MoonshotAI/Attention-Residuals\nABSTRACT\nResidual connections [12] with PreNorm [60] are standard in modern LLMs, yet they accumulate\nall layer outputs with fixed unit weights. This uniform aggregation causes uncontrolled hidden-state\ngrowth with depth, progressively diluting each layer’s contribution [27]. We proposeAttention\nResiduals (AttnRes), which replaces this fixed accumulation with softmax attention over preceding\nlayer outputs, allowing each layer to selectively aggregate earlier representations with learned, input-\ndependent weights. To address the memory and communication overhead of attending over all\npreceding layer outputs for large-scale model training, we introduceBlock AttnRes, which partitions\nlayers into blocks and attends over block-level representations, reducing the memory footprint while\npreserving most of the gains of full AttnRes. Combined with cache-based pipeline communication\nand a two-phase computation strategy, Block AttnRes becomes a practical drop-in replacement for\nstandard residual connections with minimal overhead.\nScaling law experiments confirm that the improvement is consistent across model sizes, and ablations\nvalidate the benefit of content-dependent depth-wise selection. We further integrate AttnRes into\nthe Kimi Linear architecture [69] (48B total / 3B activated parameters) and pre-train on 1.4T tokens,\nwhere AttnRes mitigates PreNorm dilution, yielding more uniform output magnitudes and gradient\ndistribution across depth, and improves downstream performance across all evaluated tasks.\nEmbedding...AttentionMoEAttentionMoEOutput\n(a) Standard ResidualsEmbedding...αAttentionαMoEαAttentionαMoE\nwwwwOutput\nαw\nααααα\n(b) Full Attention ResidualsEmbedding···Blockn-2Blockn-1AttentionMoEAttentionMoEOutput\nα\nαααα\nααααα\nwwwww\nAttnRes Op(α)wQKV\n(c) Block Attention Residuals\nFigure 1: Overview of Attention Residuals.(a)Standard Residuals: standard residual connections with uniform additive accumulation.\n(b)Full AttnRes: each layer selectively aggregates all previous layer outputs via learned attention weights.(c)Block AttnRes: layers\nare grouped into blocks, reducing memory fromO(Ld)toO(Nd).arXiv:2603.15031v1 [cs.CL] 16 Mar 2026" - }, - { - "page": 2, - "content": "Attention ResidualsTECHNICALREPORT\n1 Introduction\nStandard residual connections [12] are thede factobuilding block of modern LLMs [35, 51, 9]. The update hl=\nhl−1+fl−1(hl−1)is widely understood as agradient highwaythat lets gradients bypass transformations via identity\nmappings, enabling stable training at depth. Yet residuals also play a second role that has received less attention.\nUnrolling the recurrence shows that every layer receives the same uniformly-weighted sum of all prior layer outputs;\nresiduals define how information aggregates across depth. Unlike sequence mixing and expert routing, which now\nemploy learnable input-dependent weighting [53, 20, 9], this depth-wise aggregation remains governed by fixed unit\nweights, with no mechanism to selectively emphasize or suppress individual layer contributions.\nIn practice, PreNorm [60] has become the dominant paradigm, yet its unweighted accumulation causes hidden-state\nmagnitudes to grow as O(L) with depth, progressively diluting each layer’s relative contribution [27]. Early-layer\ninformation is buried and cannot be selectively retrieved; empirically, a significant fraction of layers can be pruned with\nminimal loss [11]. Recent efforts such as scaled residual paths [54] and multi-stream recurrences [72] remain bound to\nthe additive recurrence, while methods that do introduce cross-layer access [36, 56] are difficult to scale. The situation\nparallels the challenges that recurrent neural networks (RNNs) faced over the sequence dimension before attention\nmechanism provided an alternative.\nWe observe a formal duality between depth-wise accumulation and the sequential recurrence in RNNs. Building\non this duality, we proposeAttention Residuals (AttnRes), which replaces the fixed accumulation hl=P\nivi\nwithhl=P\niαi→l·vi, where αi→laresoftmax attention weights computed from a single learned pseudo-query\nwl∈Rdper layer. This lightweight mechanism enables selective, content-aware retrieval across depth with only one\nd-dimensional vector per layer. Indeed, standard residual connections and prior recurrence-based variants can all be\nshown to perform depth-wiselinearattention; AttnRes generalizes them to depth-wise softmax attention, completing\nfor depth the same linear-to-softmaxtransition that proved transformative over sequences (§6.2, §6.1).\nIn standard training, Full AttnRes adds negligible overhead, since the layer outputs it requires are already retained for\nbackpropagation. At scale, however, activation recomputation and pipeline parallelism are routinely employed, and these\nactivations must now be explicitly preserved and communicated across pipeline stages. We introduceBlock AttnResto\nmaintain efficiency in this regime: layers are partitioned into Nblocks, each reduced to a single representation via\nstandard residuals, with cross-block attention applied only over the Nblock-level summaries. This brings both memory\nand communication down to O(Nd) , and together with infrastructure optimizations (§4), Block AttnRes serves as a\ndrop-in replacement for standard residual connections with marginal training cost and negligible inference latency\noverhead.\nScaling law experiments confirm that AttnRes consistently outperforms the baseline across compute budgets, with\nBlock AttnRes matching the loss of a baseline trained with 1.25× more compute. We further integrate AttnRes into\nthe Kimi Linear architecture [69] (48B total / 3B activated parameters) and pre-train on 1.4T tokens. Analysis of\nthe resulting training dynamics reveals that AttnRes mitigates PreNorm dilution, with output magnitudes remaining\nbounded across depth and gradient norms distributing more uniformly across layers. On downstream benchmarks, our\nfinal model improves over the baseline across all evaluated tasks.\nContributions\n•Attention Residuals.We propose AttnRes, which replaces fixed residual accumulation with learned softmax\nattention over depth, and its scalable variant Block AttnRes that reduces memory and communication from O(Ld) to\nO(Nd) . Through a unified structured-matrix analysis, we show that standard residuals and prior recurrence-based\nvariants correspond to depth-wiselinearattention, while AttnRes performs depth-wisesoftmaxattention.\n•Infrastructure for scale.We develop system optimizations that make Block AttnRes practical and efficient at scale,\nincluding cross-stage caching that eliminates redundant transfers under pipeline parallelism and a two-phase inference\nstrategy that amortizes cross-block attention via online softmax [31]. The resulting training overhead is marginal,\nand the inference latency overhead is less than 2% on typical inference workloads.\n•Comprehensive evaluation and analysis.We validate AttnRes through scaling law experiments, component\nablations, and downstream benchmarks on a 48B-parameter model pre-trained on 1.4T tokens, demonstrating\nconsistent improvements over standard residual connections. Training dynamics analysis further reveals that AttnRes\nmitigates PreNorm dilution, yielding bounded hidden-state magnitudes and more uniform gradient distribution across\ndepth.\n2" - }, - { - "page": 3, - "content": "Attention ResidualsTECHNICALREPORT\n2 Motivation\nNotation.Consider a batch of input sequences with shape B×T×d , where Bis the batch size, Tis the sequence\nlength, and dis the hidden dimension. For clarity, we write formulas for a single token: hl∈Rddenotes the hidden state\nentering layer l, where l∈ {1, . . . , L} is the layer index and Lis the total number of layers. The token embedding is h1.\nThe function flrepresents the transformation applied by layer l. In Transformer models, we treat each self-attention or\nMLP as an individuallayer.\n2.1 Training Deep Networks via Residuals\nResidual Learning.Residual learning [12] proves to be a critical technique in training deep networks as it allows\ngradients to bypass transformations. Specifically, each layer updates the hidden state as:\nhl=hl−1+fl−1(hl−1)\nExpanding this recurrence, the hidden state at layer lis the sum of the embedding and all preceding layer outputs:\nhl=h 1+Pl−1\ni=1fi(hi). The key insight behind residual connections isidentity mapping: each layer preserves a direct\npath for both information and gradients to flow unchanged. During back-propagation, the gradient with respect to an\nintermediate hidden state is:\n∂L\n∂hl=∂L\n∂hL·L−1Y\nj=l\u0012\nI+∂fj\n∂hj\u0013\nExpanding this product yields Iplus higher-order terms involving the layer Jacobians ∂fj/∂hj. The identity term is\nalways preserved, providing a direct gradient path from the loss to any layer regardless of depth.\nGeneralizing Residuals.While effective, the fixed unit coefficients in the residual update treat every layer’s con-\ntribution uniformly, offering no mechanism to adapt the mixing across depth. Highway networks [45] relax this by\nintroducing learned element-wise gates:\nhl= (1−g l)⊙h l−1+gl⊙fl−1(hl−1)\nwhere gl∈[0,1]dinterpolates between the transformation and the identity path. More generally, both are instances\nof a weighted recurrence hl=α l·hl−1+βl·fl−1(hl−1), with residual setting αl=βl=1and Highway setting\nαl=1−g l, βl=gl.\nLimitations.Whether fixed or gated, both approaches share a fundamental constraint: each layer can only access\nits immediate input hl−1, a single compressed state that conflates all earlier layer outputs, rather than the individual\noutputs themselves. This entails several limitations: (1)no selective access: different layer types (e.g., attention vs.\nMLP) receive the same aggregated state, despite potentially benefiting from different weightings; (2)irreversible loss:\ninformation lost through aggregation cannot be selectively recovered in deeper layers; and (3)output growth: later\nlayers learn increasingly larger outputs to gain influence over the accumulated residual, which can destabilize training.\nThese limitations motivate a mechanism that lets each layer selectively aggregate information from all preceding layers.\n3 Attention Residuals: A Unified View of Time and Depth\nThe limitations discussed above are reminiscent of similar bottlenecks in sequence modeling, suggesting that we seek\nsimilar solutions for the depth dimension.\nThe Duality of Time and Depth.Like RNNs over time, residual connections compress all prior information into a\nsingle state hlover depth. For sequence modeling, the Transformer improved upon RNNs by replacing recurrence with\nattention [3, 52], allowing each position to selectively access all previous positions with data-dependent weights. We\npropose the same methodology for depth:\nhl=α 0→l·h1+l−1X\ni=1αi→l·fi(hi)(1)\nwhere αi→lare layer-specific attention weights satisfyingPl−1\ni=0αi→l= 1. Unlike sequence length (which can reach\nmillions of tokens), network depth is typically modest ( L <1000 ), making O(L2)attention over depth computationally\nfeasible. We call this approachAttention Residuals, abbreviated asAttnRes.\n3" - }, - { - "page": 4, - "content": "Attention ResidualsTECHNICALREPORT\n3.1 Full Attention Residuals\nThe attention weights can be written as αi→l=ϕ(q l,ki)for a kernel function ϕ:Rd×Rd→R≥0, where qland\nkiare query and key vectors [23, 70]. Different choices of ϕrecover different residual variants (§6.2); we adopt\nϕ(q,k) = exp\u0000\nq⊤RMSNorm(k)\u0001\n[66] with normalization, yieldingsoftmaxattention over depth:\nαi→l=ϕ(ql,ki)\nPl−1\nj=0ϕ(ql,kj)(2)\nFor each layerl, we define:\nql=w l,k i=vi=\u001ah1 i= 0\nfi(hi) 1≤i≤l−1(3)\nwhere the query ql=w lis a layer-specific learnable vector in Rd. The RMSNorm inside ϕprevents layers with\nlarge-magnitude outputs from dominating the attention weights. The input to layerlis then:\nhl=l−1X\ni=0αi→l·vi (4)\nWe call this formfull attention residuals. For each token, Full AttnRes requires O(L2d)arithmetic and O(Ld) memory\nto store layer outputs. Since depth is far smaller than sequence length, the arithmetic cost is modest.\nOverhead.The O(Ld) memory overlaps entirely with the activations already retained for backpropagation, so Full\nAttnRes introduces no additional memory overhead in vanilla training. At scale, however, activation recomputation and\npipeline parallelism are widely adopted: layer outputs that would otherwise be freed and recomputed must now be kept\nalive for all subsequent layers, and under pipeline parallelism each must further be transmitted across stage boundaries.\nBoth the memory and communication overhead then grow asO(Ld).\nBlockwise optimization.A deliberate design choice in Full AttnRes is that thepseudo-query wlis a learned parameter\ndecoupled from the layer’s forward computation. This independence means that attention weights for any group of\nlayers can be computed in parallel without waiting for their sequential outputs, and in particular permits grouping the L\nlayers into Nblocks of Slayers each and batching the attention computation within each block, reducing per-layer\nmemory I/O from O(Ld) toO((S+N)d) (we defer the detailed two-phase strategy to §4). Under current distributed\ntraining regimes, however, the dominant cost is not local memory bandwidth but cross-stage communication under\npipeline parallelism: every layer output must still be transmitted between stages, and this O(Ld) communication\noverhead cannot be alleviated by local batching. This motivates the Block AttnRes variant introduced below, which\nreduces the number of cross-stage representations from LtoN. We anticipate that future interconnect improvements\nwill make the fullO(Ld)communication practical, fully realizing the potential of Full AttnRes.\n3.2 Block Attention Residuals\nWe proposeBlock Attention Residuals, which partitions the Llayers into Nblocks: within each block, the layer outputs\nare reduced to a single representation via summation, and across blocks, we apply full attention over only Nblock-level\nrepresentations and the token embedding. This reduces both memory and communication overhead from O(Ld) to\nO(Nd).\nIntra-Block Accumulation.Specifically, we divide the Llayers into Nblocks of S=L/N layers each, assuming\nLis divisible by N; otherwise, the last block contains the remaining LmodN layers. Let Bndenote the set of layer\nindices in blockn(n= 1, . . . , N). To form a block, we sum all of its layer outputs:\nbn=X\nj∈Bnfj(hj)(5)\nWe further denote bi\nnas the partial sum over the first ilayers in Bn, so that bn=bS\nn. When Lis not divisible by N,\nthe final partial sum is taken as the last block’s representation. As in Full AttnRes, the RMSNorm inside ϕprevents\nmagnitude differences between complete blocks and partial sums from biasing the attention weights.\n4" - }, - { - "page": 5, - "content": "Attention ResidualsTECHNICALREPORT\n1 def block_attn_res(blocks: list[Tensor], partial_block: Tensor, proj: Linear, norm: RMSNorm) -> Tensor:\n2 \"\"\"\n3 Inter-block attention: attend over block reps + partial sum.\n4 blocks:\n5 N tensors of shape [B, T, D]: completed block representations for each previous block\n6 partial_block:\n7 [B, T, D]: intra-block partial sum (b_n^i)\n8 \"\"\"\n9 V = torch.stack(blocks + [partial_block]) # [N+1, B, T, D]\n10 K = norm(V)\n11 logits = torch.einsum('d, n b t d -> n b t', proj.weight.squeeze(), K)\n12 h = torch.einsum('n b t, n b t d -> b t d', logits.softmax(0), V)\n13 return h\n14\n15 def forward(self, blocks: list[Tensor], hidden_states: Tensor) -> tuple[list[Tensor], Tensor]:\n16 partial_block = hidden_states\n17 # apply block attnres before attn\n18 # blocks already include token embedding\n19 h = block_attn_res(blocks, partial_block, self.attn_res_proj, self.attn_res_norm)\n20\n21 # if reaches block boundary, start new block\n22 # block_size counts ATTN + MLP; each transformer layer has 2\n23 if self.layer_number % (self.block_size // 2) == 0:\n24 blocks.append(partial_block)\n25 partial_block = None\n26\n27 # self-attention layer\n28 attn_out = self.attn(self.attn_norm(h))\n29 partial_block = partial_block + attn_out if partial_block is not None else attn_out\n30\n31 # apply block attnres before MLP\n32 h = block_attn_res(blocks, partial_block, self.mlp_res_proj, self.mlp_res_norm)\n33\n34 # MLP layer\n35 mlp_out = self.mlp(self.mlp_norm(h))\n36 partial_block = partial_block + mlp_out\n37\n38 return blocks, partial_block\nFigure 2: PyTorch-style pseudo code for Block Attention Residuals. block_attn_res computes softmax attention over block\nrepresentations using a learned pseudo-query wl;forward is a single-layer pass that maintains partial_block (bi\nn, intra-block\nresidual) andblocks([b 0, . . . ,b n−1], inter-block history).\nInter-Block Attention.In Full AttnRes, the input to layer lis computed by attending over all outputs up to fl−1(hl−1).\nThe block-wise variant replaces these individual outputs with block representations, defining b0=h 1so that the token\nembedding is always included as a source. For thei-th layer in blockn, the value matrix is:\nV=\u001a[b0,b1, . . . ,b n−1]⊤ifi= 1(first layer of blockn)\n[b0,b1, . . . ,b n−1,bi−1\nn]⊤ifi≥2(subsequent layers)(6)\nKeys and attention weights follow Eq. 3 and Eq. 2. The input of the very first layer of the network is the token\nembeddings, i.e. b0=h 1. In each block, the first layer receives the previous block representations and the token\nembeddings, and the subsequent layers additionally attend to the partial sum bi−1\nn. The final output layer aggregates all\nNblock representations. Fig. 2 provides PyTorch-style pseudocode for Block AttnRes.\nEfficiency.Since each layer now attends over Nblock representations rather than Lindividual outputs, memory\nreduces from O(L) toO(N) and computation from O(L2)toO(N2). The block count Ninterpolates between two\nextremes: N=L recovers Full AttnRes, while N= 1 reduces to standard residual connections with the embedding\nisolated as b0. Empirically, we find that N≈8 recovers most of the benefit across model scales, requiring only eight\nstored hidden states per token (see § 5).\nBeyond memory and computation, the block structure also benefits inference latency: block boundaries define the\ndispatch granularity for the blockwise optimization described in §3, and the fixed block count Nbounds the KV cache\nsize. The parallel inter-block results are merged with the sequential intra-block partial sums via online softmax [31],\npreserving exact equivalence (§4).\n4 Infrastructure Design\nBlock AttnRes introduces additional system challenges compared to standard residual connections. For large-scale\nmodel training, block representations must be propagated across pipeline stages, causing heavy communication in a\n5" - }, - { - "page": 6, - "content": "Attention ResidualsTECHNICALREPORT\nRANK0\nRANK1\nRANK2\nRANK3[b0] [ ]\n[b0] [b1]\n[b0,b1] [ ]\n[b0,b1] [b2]+ [b 1,b2][ ]\n+ [b 1,b2][b3]\n+ [b 2,b3][ ]\n+ [b 2,b3][b4]VIRTUALSTAGE0 VIRTUALSTAGE1\n1 2\n1 2\n1 2\n1 21 2\n1 2\n1 2\n1 2\nFigure 3: Cache-based pipeline communication example with 4 physical ranks and 2 virtual stages per rank, where hatched boxes\ndenote end of AttnRes blocks. Numbers indicate micro-batch indices. Each rank caches previously received blocks; stage transitions\nonly transmit incremental blocks (+[b 1,b2]) instead of the full history.\nnaïve implementation. During inference, repeated access to accumulated block representations increases latency, while\nlong-context prefilling amplifies the memory cost of caching block representations. We address these challenges with\ncross-stage caching in training, and with a two-phase computation strategy together with a memory-efficient prefilling\nscheme in inference.\n4.1 Training\nFor small-scale training, AttnRes adds a tiny computation overhead and no extra memory usage, as the activations\nneed to be saved for backpropagation regardless. Under large-scale distributed training, pipeline parallelism poses the\nprimary infrastructure challenge for AttnRes. Full AttnRes requires all Llayer outputs to be transmitted across stages;\nBlock AttnRes reduces this to Nblock representations, and the optimizations below further minimize the remaining\noverhead.\nPipeline communication.With standard residual connections, pipeline parallelism [18] transfers a fixed-size hidden\nstate between adjacent stages, independent of pipeline depth. Block AttnRes requires all accumulated block representa-\ntions at each stage for inter-block attention, and naïvely transmitting the full history at every transition incurs redundant\ncommunication.\nConsider an interleaved pipeline schedule [33] with Pphysical stages and Vvirtual stages per physical stage. For\nsimplicity, assume each physical stage produces on average Npblock representations of dimension dper token.1With\nC=PV total chunks (each physical stage in each virtual stage), the j-th chunk accumulates jNpblocks. Naïvely\ntransmitting all accumulated blocks at every transition incurs per-token communication cost:\nComm naïve=C−1X\nj=1jNp·d=C(C−1)\n2Npd.(7)\nCross-stage caching.Since each physical stage processes multiple virtual stages in succession, we can eliminate\nthis redundancy by caching blocks locally: blocks received during earlier virtual stages remain in local memory and\nneed not be re-transmitted. The first virtual stage ( v= 1 ) has no cache and accumulates normally; for v≥2 , each\ntransition conveys only the ∼PN pincremental blocks accumulated since the receiver’s corresponding chunk in the\nprevious virtual stage. Total communication reduces to:\nComm cached =P(P−1)\n2Npd\n|{z}\nfirst virtual stage+ (V−1)P2Npd|{z }\nsubsequent virtual stages.(8)\nCaching reduces peak per-transition cost from O(C) toO(P) , aV× improvement that enables full overlap with\ncomputation during steady-state 1F1B. The backward pass benefits from the same scheme. Fig. 3 illustrates this\noptimization withP=4andV=2: for the second virtual stage, caching eliminates 6 redundant block transmissions.\n1In practice, block boundaries need not align with physical stage boundaries. For example, in Fig. 3, each block spans two\nphysical stages, so only every other transition involves a newly completed block.\n6" - }, - { - "page": 7, - "content": "Attention ResidualsTECHNICALREPORT\nAlgorithm 1:Two-phase computation for blockn\nInput:Pseudo queries{w l}l∈Bn, block representations{b 0, . . . ,b n−1}\n/* Phase 1: Parallel inter-block attention */\n1Q←[w l]l∈Bn //[S, d]\n2K,V←[b 0;. . .;b n−1]//[n, d]\n3{o(1)\nl, m(1)\nl, ℓ(1)\nl}l∈Bn←ATTNWITHSTATS(Q,K,V)// Return LSE\n4\n/* Phase 2: Sequential intra-block attention + Onlinesoftmaxmerge */\n5i←0\n6forl∈ B ndo\n7ifi= 0then\n8h l←o(1)\nl/ℓ(1)\nl// Inter-block only\n9else\n10o(2)\nl, m(2)\nl, ℓ(2)\nl←ATTNWITHSTATS(w l,bi\nn,bi\nn)// Intra-block\n11m l←max(m(1)\nl, m(2)\nl)\n12h l←em(1)\nl−mlo(1)\nl+em(2)\nl−mlo(2)\nl\nem(1)\nl−mlℓ(1)\nl+em(2)\nl−mlℓ(2)\nl// Online softmax merge\n13i←i+ 1\n14bi\nn←bi−1\nn+fl(hl)// Update partial sum;b0\nn:=0\n15return{h l}l∈Bn\nMemory overhead.With cross-stage caching, each block is stored exactly once across all Vvirtual stages, which\nbecomes negligible relative to standard per-layer activation cache. Crucially, the per-layer activation footprint remains\nidentical to standard architectures, as activation checkpointing eliminates all inter-block attention intermediates, and the\ncheckpointed inputp lmatches the memory size of the hidden stateh lit replaces.\nIn terms of wall-clock time, Block AttnRes adds negligible training overhead when pipeline parallelism is not enabled;\nunder pipeline parallelism, the measured end-to-end overhead is less than 4%.\n4.2 Inference\nThe two-phase computation strategy described below applies to both Full and Block AttnRes: in either case, layers are\ngrouped into blocks of size S, with Phase 1 batching the inter-block queries and Phase 2 handling sequential intra-block\nlookback. For Full AttnRes, this reduces per-layer I/O from O(Ld) toO((S+N)d) (detailed derivation shown in\nAppendix B); Block AttnRes further reduces the stored representations from LtoN, since each block is compressed\ninto a single vector. In what follows, we focus on Block AttnRes and detail the two-phase computation strategy together\nwith a sequence-sharded prefilling scheme for long-context inputs.\nTwo-phase computation strategy.The layer-wise attention computation of Block AttnRes resembles autoregressive\ndecoding, where block representations serve as a shared KV cache reused across layers. A naïve implementation\ncomputes the attention residual at every layer, each requiring a full pass over all preceding blocks, resulting in O(L·N)\nmemory accesses. Since the pseudo-query vectors are decoupled from the forward computation (§3), all S=L/N\nqueries within a block can be batched into a single matrix multiplication, amortizing memory access from Sreads to 1.\nAlgorithm 1 instantiates a two-phase computation strategy exploiting this property.\n•Phase 1computes inter-block attention for all Slayers simultaneously via a single batched query against the cached\nblock representations, returning both outputs and softmax statistics (max and log-sum-exp). This amortizes the\nmemory access cost, reducing reads fromStimes to just once per block.\n•Phase 2computes intra-block attention sequentially for each layer using the evolving partial sum, then merges with\nPhase 1 outputs through online softmax [31]. Because the online- softmax merge is elementwise, this phase naturally\nadmits kernel fusion with surrounding operations, further reducing I/O overhead.\nWith the two-phase design, Phase 2 preserves an I/O footprint similar to that of standard residual connections, whereas\nthe main additional cost arises from Phase 1 inter-block attention. Because these inter-block reads are amortized across\n7" - }, - { - "page": 8, - "content": "Attention ResidualsTECHNICALREPORT\nall layers in a block through batching, the total per-layer memory access cost remains only (N\nS+ 3)d reads and 2d\nwrites (Table 1). This is substantially lower than the residual-stream I/O of prior residual generalizations such as (m)HC\nunder typical settings. In practice, Phase 1 can also partially overlap with the computation of the first layer in the block,\nfurther reducing its wall-clock impact. As a result, the end-to-end inference latency overhead is less than 2% on typical\ninference workloads.\nTable 1: Memory access cost per token per layer incurred by the residual mechanism under each scheme. The internal I/O of the layer\nfunction flis excluded. For AttnRes, both Full and Block variants use the two-phase inference schedule described in Appendix B;\namortized costs are averaged overNlayers within a block. Typical values:L=128,N=8,S=L/N=16,m=4.\nOperation Read WriteTotal I/O\nSymbolic Typical\nStandard Residuals Residual Merge2d d3d3d\nmHC (mstreams)Computeα l,βl,Al md m2+2m\n(8m+2)d+2m2+4m 34dApplyα l md+m d\nApplyβ l d+m md\nApplyA l md+m2md\nResidual Merge2md md\nAttnResFullPhase 1 (amortized)(N−1)d d(S+N)d24dPhase 2(S−1)d d\nBlockPhase 1 (amortized)N\nSd d \u0000N\nS+5\u0001\nd 5.5dPhase 23d d\nMemory-efficient prefilling.Storing block representations during prefilling requires N·T·d elements, which incurs\n15 GB of memory for a 128K-token sequence with 8 blocks. We mitigate this by sharding these representations along\nthe sequence dimension across Ptensor-parallel devices, allowing Phase 1 to execute independently on local sequence\nshards. The Phase 2 online- softmax merge then integrates into the standard TP all-reduce communication path: the\noutput is reduce-scattered, merged locally, and reconstructed via all-gather, naturally admitting kernel fusion with\noperations like RMSNorm . This reduces the per-device memory footprint to N·(T/P)·d —lowering the 128K-context\nexample from 15 GB to roughly 1.9 GB per device. Combined with chunked prefill (e.g., 16K chunk size), the overhead\nfurther reduces to under 0.3 GB per device.\n5 Experiments\nArchitecture Details.Our architecture is identical to Kimi Linear [69], a Mixture-of-Experts (MoE) Transformer\nfollowing the Moonlight [28] / DeepSeek-V3 [9] design, which interleaves Kimi Delta Attention (KDA) and Multi-Head\nLatent Attention (MLA) layers in a 3:1 ratio, each followed by an MoE feed-forward layer. The only modification is the\naddition of AttnRes to the residual connections; all other components (model depth, hidden dimensions, expert routing,\nand MLP structure) remain unchanged. AttnRes introduces only one RMSNorm and one pseudo-query vector wl∈Rd\nper layer, amounting to a negligible fraction of the total parameter count. Crucially, all pseudo-query vectors must be\ninitialized to zero. This ensures that the initial attention weights αi→lare uniform across source layers, which reduces\nAttnRes to an equal-weight average at the start of training and prevents training volatility, as we validated empirically.\n5.1 Scaling Laws\nWe sweep five model sizes (Table 2) and train three variants per size: a PreNorm baseline, Full AttnRes, and Block\nAttnRes with ≈8blocks. They are trained with an 8192-token context window and a cosine learning rate schedule.\nWithin each scaling law size group, all variants share identical hyperparameters selected under the baseline to ensure\nfair comparison; this setup intentionally favors the baseline and thus makes the comparison conservative. Following\nstandard practice, we fit power-law curves of the form L=A×C−α[22, 15], where Lis validation loss and Cis\ncompute measured in PFLOP/s-days.\nScaling Behavior.Fig. 4 presents the fitted scaling curves. The Baseline follows L= 1.891×C−0.057, while Block\nAttnRes fits L= 1.870×C−0.058, and Full AttnRes fits L= 1.865×C−0.057. All three variants exhibit a similar\nslope, but AttnRes consistently achieves lower loss across the entire compute range. Based on the fitted curves, at 5.6\n8" - }, - { - "page": 9, - "content": "Attention ResidualsTECHNICALREPORT\nTable 2: Baseline vs Block AttnRes ( N= 8 ) vs Full AttnRes vs mHC(-lite) [64]: Model configurations, Hyperparameters, and\nValidation Loss.\n# Act.\nParams†TokensL bH d model dff lr batch size‡ Val. Loss\nBaseline Block AttnRes Full AttnRes mHC(-lite)\n194M 038.7B 12 12 0896 4002.99×10−3192 1.931 1.9091.8991.906\n241M 045.4B 13 13 0960 4322.80×10−3256 1.895 1.875 1.8741.869\n296M 062.1B 14 14 1024 4642.50×10−3320 1.829 1.8091.8041.807\n436M 087.9B 16 16 1168 5282.20×10−3384 1.766 1.7461.7371.747\n528M 119.0B 17 17 1264 5602.02×10−3432 1.719 1.6931.6921.694\n†Denotes the number of activated parameters in our MoE models, excluding embeddings.\n‡All models were trained with a context length of 8192.\n⋆Lb=L/2denotes the number of Transformer blocks.\n0.5 1 2 51.71.81.9\n1.25×\nPFLOP/s-daysLossBaseline:1.891×C−0.057\nFull AttnRes:1.865×C−0.057\nBlock AttnRes:1.870×C−0.058\nFigure 4: Scaling law curves for Attention Residuals. Both Full and Block AttnRes consistently outperform the baseline across all\nscales. Block AttnRes closely tracks Full AttnRes, recovering most of the gain at the largest scale.\nPFLOP/s-days, Block AttnRes reaches 1.692 versus the Baseline’s 1.714, equivalent to a 1.25× compute advantage.\nThe gap between Full and Block AttnRes narrows with scale, shrinking to just 0.001 at the largest size. We also list\nmHC(-lite) [64] in Table 2 for reference. Full AttnRes outperforms mHC, while Block AttnRes matches it at lower\nmemory I/O per layer:5.5dversus34dfor mHC withm=4streams (Table 1).\n5.2 Main Results\nTraining recipe.The largest models we study are based on the full Kimi Linear 48B configuration: 27 Transformer\nblocks (54 layers) with 8 out of 256 routed experts plus 1 shared expert, yielding 48B total and 3B activated parameters.\nThis model applies Block AttnRes with 6 layers per block, producing 9 blocks plus the token embedding for a total of\n10 depth-wise sources.\nWe follow the same data and training recipe as the Kimi Linear 1.4T-token runs [69]: all models are pre-trained with a\n4096-token context window, the Muon optimizer [28], and a WSD (Warmup–Stable–Decay) learning rate schedule [16],\nwith a global batch size of 8M tokens. Training of the final model proceeds in two stages: (i) a WSD pre-training phase\non 1T tokens, followed by (ii) a mid-training phase on ≈400B high-quality tokens, following the annealing recipe of\nMoonlight [28].\nAfter mid-training, we continue training with progressively longer sequence length of 32K tokens. Since our architecture\nuses hybrid KDA/MLA attention [69], where MLA operates without positional encodings (NoPE) [61], context extension\nrequires no modifications such as YaRN [37] or attention temperature rescaling.\n9" - }, - { - "page": 10, - "content": "Attention ResidualsTECHNICALREPORT\n20k 40k 60k 80k 100k1.21.31.41.5\nStep(a) Validation Loss\nBaseline\nBlock AttnRes\n0 10 20051015\nTransformer Block Index(b) Output Magnitude\n0 10 200123\nTransformer Block Index(c) Gradient Magnitude (×10−5)\nFigure 5: Training dynamics of Baseline and Block AttnRes.(a)Validation loss during training.(b)Each transformer block’s output\nmagnitude at the end of training.(c)Each transformer block’s gradient magnitude.\nTraining dynamics.We compare the training dynamics of our final Baseline and Block AttnRes models over 1T\ntokens in Fig. 5.\n•Validation loss:AttnRes achieves consistently lower validation loss throughout training, with the gap widening\nduring the decay phase and resulting in a notably lower final loss.\n•Output magnitude:The Baseline suffers from the PreNorm dilution problem [60, 27]: as hidden-state magnitudes\ngrow monotonically with depth, deeper layers are compelled to learn increasingly large outputs from fixed-scale\nnormalized inputs to remain influential. Block AttnRes confines this growth within each block, as selective aggregation\nat block boundaries resets the accumulation, yielding a bounded periodic pattern.\n•Gradient magnitude:With all residual weights fixed to 1, the Baseline provides no means of regulating gradient\nflow across depth, leading to disproportionately large gradients in the earliest layers. The learnable softmax weights\nin Block AttnRes (Fig. 8) introduce competition among sources for probability mass, resulting in a substantially more\nuniform gradient distribution.\nTable 3: Performance comparison of AttnRes with the baseline, both after the same pre-training recipe. Best per-row results are\nbolded.\nBaseline AttnRes\nGeneralMMLU 73.574.6\nMMLU-Pro52.2 52.2\nGPQA-Diamond 36.944.4\nBBH 76.378.0\nARC-Challenge 64.665.7\nHellaSwag 83.283.4\nTriviaQA 69.971.8\nMath & CodeGSM8K 81.782.4\nMGSM 64.966.1\nMath 53.557.1\nCMath 84.785.1\nHumanEval 59.162.2\nMBPP 72.073.9\nChineseCMMLU 82.082.9\nC-Eval 79.682.5\nDownstream performance.Following the evaluation protocol of Kimi Linear [69], we assess both models across\nthree areas (Table 3):\n10" - }, - { - "page": 11, - "content": "Attention ResidualsTECHNICALREPORT\nTable 4: Ablation on key components of AttnRes (16-layer\nmodel).\nVariant Loss\nBaseline (PreNorm) 1.766\nDenseFormer [36] 1.767\nmHC [59] 1.747\nAttnRes Full 1.737\nw/ input-dependent query1.731\nw/ input-independent mixing1.749\nw/sigmoid1.741\nw/oRMSNorm1.743\nSWA (W= 1 + 8) 1.764\nBlock (S= 4) 1.746\nw/ multihead (H= 16)1.752\nw/oRMSNorm1.75032 16 8 4 21.7351.7401.7451.7501.7551.7601.7651.770\n1.757\n1.753\n1.748\n1.746 1.746Baseline (1.766)\nFull AttnRes i.e. S=1 (1.737)\nBlock size (S)Validation lossBaseline\nFull AttnRes\nBlock AttnRes\nFigure 6: Effect of block size on validation loss (16-layer model).\n•Language understanding and reasoning: MMLU [13], MMLU-Pro Hard [55], GPQA-Diamond [41], BBH [48],\nARC-Challenge [6], HellaSwag [65], and TriviaQA [21].\n•Reasoning (Code and Math): GSM8K [7], MGSM [44], Math [25], CMath [14], HumanEval [5], and MBPP [1].\n•Chinese language understanding: CMMLU [26] and C-Eval [19].\nAs shown in Table 3, Block AttnRes matches or outperforms the baseline on all benchmarks. The improvements are\nparticularly pronounced on multi-step reasoning tasks such as GPQA-Diamond (+7.5) and Minerva Math (+3.6), as\nwell as code generation such as HumanEval (+3.1), while knowledge-oriented benchmarks such as MMLU (+1.1)\nand TriviaQA (+1.9) also show solid gains. This pattern is consistent with the hypothesis that improved depth-wise\ninformation flow benefits compositional tasks, where later layers can selectively retrieve and build upon earlier\nrepresentations.\n5.3 Ablation Study\nWe conduct ablation studies on the 16-head model from Table 2 to validate key design choices in AttnRes (Table 4). All\nmodels share identical hyperparameters and compute budget.\nComparison with prior methods.We compare AttnRes against the PreNorm baseline (loss 1.766) and two rep-\nresentative methods that generalize residual connections. DenseFormer [36] grants each layer access to all previous\noutputs but combines them with fixed, input-independent scalar coefficients; it shows no gain over the baseline (1.767),\nhighlighting the importance of input-dependent weighting. mHC [59] introduces input dependence through mparallel\nstreams with learned mixing matrices, improving to 1.747. AttnRes takes this further with explicit content-dependent\nselection via softmax attention: Full AttnRes achieves 1.737 and Block AttnRes 1.746, outperforming both methods\nwith only a single query vector per layer.\nCross-layer access.We compare three granularities of cross-layer access. Full AttnRes follows directly from the\ntime–depth duality (§ 3), applying attention over all previous layers, and achieves the lowest loss (1.737). A simple\nway to reduce its memory cost is sliding-window aggregation (SWA), which retains only the most recent W=8 layer\noutputs plus the token embedding; it improves over baseline (1.764) but falls well short of both Full and Block AttnRes,\nsuggesting that selectively accessing distant layers matters more than attending to many nearby ones.\nBlock AttnRes offers a better trade-off: with block size S=4 it reaches 1.746 while keeping memory overhead constant\nper layer. Fig. 6 sweeps Sacross the full spectrum from S=1 (i.e. Full AttnRes) to increasingly coarse groupings. Loss\ndegrades gracefully as Sgrows, with S=2,4,8 all landing near 1.746 while larger blocks ( S=16,32 ) move toward\nbaseline. In practice, we fix the number of blocks to ≈8for infrastructure efficiency (§ 4). As future hardware alleviates\nmemory capacity constraints, adopting finer-grained block sizes or Full AttnRes represents a natural pathway to further\nimprove performance.\n11" - }, - { - "page": 12, - "content": "Attention ResidualsTECHNICALREPORT\n15 30 45 60 750.30.40.50.60.7 2.017 1.909 1.875 1.851 1.858\n1.990 1.902 1.862 1.852 1.862\n1.973 1.883 1.859 1.849 1.854\n1.952 1.868 1.850 1.849 1.857\n1.926 1.857 1.851 1.858 1.847\ndmodel/LbH/L b\n(a) Baseline15 30 45 60 751.954 1.890 1.843 1.828 1.824\n1.931 1.863 1.830 1.817 1.818\n1.917 1.841 1.819 1.812 1.817\n1.893 1.823 1.815 1.813 1.813\n1.877 1.816 1.820 1.806 1.802\ndmodel/Lb\n1.841.881.921.962\n(b) Attention Residuals\nFigure 7: Architecture sweep under fixed compute ( ≈6.5×1019FLOPs, ≈2.3×108active parameters). Each cell reports\nvalidation loss for a (dmodel/Lb, H/L b)configuration, where Lb=L/2 is the number of Transformer blocks; the star marks the\noptimum.\nComponent design.We further ablate individual components of the attention mechanism:\n•Input-dependent query.A natural extension is to make the query input-dependent by projecting it from the current\nhidden state. This further lowers loss to 1.731, but introduces a d×d projection per layer and requires sequential\nmemory access during decoding, so we default to the learned query.\n•Input-independent mixing.We removed the query and key and replaced them with learnable, input-independent\nscalars to weigh previous layers, which hurts performance (1.749 vs. 1.737).\n•softmax vs.sigmoid .Replacing softmax withsigmoid degrades performance (1.741). We attribute this to softmax ’s\ncompetitive normalization, which forces sharper selection among sources.\n•Multihead attention.We test per-head depth aggregation ( H=16 ) on Block AttnRes, allowing different channel\ngroups to attend to different source layers. This hurts performance (1.752 vs. 1.746), indicating that the optimal\ndepth-wise mixture is largely uniform across channels: when a layer’s output is relevant, it is relevant as a whole.\n•RMSNorm on keys.Removing RMSNorm degrades both Full AttnRes (1.743) and Block AttnRes (1.750). For\nFull AttnRes, it prevents individual layers with naturally larger outputs from dominating the softmax . This becomes\neven more critical for Block AttnRes, as block-level representations accumulate over more layers and can develop\nlarge magnitude differences;RMSNormprevents these from biasing the attention weights.\n5.4 Analysis\n5.4.1 Optimal Architecture\nTo understand how AttnRes reshapes optimal architectural scaling, we perform a controlled capacity reallocation\nstudy under a fixed compute and parameter budget. Our central question is whether AttnRes alters the preferred\ndepth–width–attention trade-off, and in particular, given its potential strength on the depth dimension, whether it favors\ndeeper models compared to conventional Transformer design heuristics. To isolate structural factors directly coupled\nto depth, we fix the per-expert MLP expansion ratio based on internal empirical observations ( dff/dmodel≈0.45 ).\nWe further fix total training compute (FLOPs ≈6.5×1019) and active parameters ( ≈2.3×108), ensuring that any\nperformance variation arises purely from architectural reallocation rather than overall capacity differences. Under\nthis constrained budget, we enumerate 25 configurations on a 5×5 grid over dmodel/Lb∈ {15,30,45,60,75} and\nH/L b∈ {0.3,0.4,0.5,0.6,0.7} , where Lb=L/2 is the number of Transformer blocks and Hthe number of attention\nheads. The results are shown in Fig. 7.\nBoth heatmaps exhibit a shared pattern: loss decreases with growing dmodel/Lband shrinking H/L b, and both methods\nreach their optima at H/L b≈0.3 . Despite this shared trend, AttnRes achieves a lower loss than the baseline in each of\nthe 25 configurations, by 0.019 –0.063 . The most apparent difference lies in the location of the optimum: the baseline\nachieves its lowest loss at dmodel/Lb≈60 (1.847 ), whereas AttnRes shifts it to dmodel/Lb≈45 (1.802 ). Under a fixed\n12" - }, - { - "page": 13, - "content": "Attention ResidualsTECHNICALREPORT\n0 5 10 15 20 25 301\n2\n3\n4\n5\n6\n7\n8\n9\n10\n11\n12\n13\n14\n15\n16\nSource IndexLayerFull AttnRes, Pre-Attn\n0 5 10 15 20 25 30\nSource IndexFull AttnRes, Pre-MLP\n0 1 2 3 4 5 6 7 81\n2\n3\n4\n5\n6\n7\n8\n9\n10\n11\n12\n13\n14\n15\n16\nBlock IndexLayerBlock AttnRes, Pre-Attn\n0 1 2 3 4 5 6 7 8\nBlock IndexBlock AttnRes, Pre-MLP\n00.20.40.60.8Weight\nFigure 8: Depth-wise attention weight distributions for a 16-head model with full (top) and block (bottom) Attention Residuals,\naveraged over tokens. The model has 16 attention and 16 MLP layers. Each row shows how the lth attention (left) or MLP (right)\nlayer distributes weight over previous sources. Diagonal dominance indicates locality remains the primary information pathway,\nwhile persistent weights on source 0 (embedding) and occasional off-diagonal concentrations reveal learned skip connections. Block\nattention (N= 8) recovers the essential structure with sharper, more decisive weight distributions.\nparameter budget, a lower dmodel/Lbcorresponds to a deeper, narrower network, suggesting that AttnRes can exploit\nadditional depth more effectively. We note that this preference for depth does not directly translate to a deployment\nrecommendation, as deeper models generally incur higher inference latency due to their sequential computation [39].\nRather, this sweep serves as a diagnostic that reveals where AttnRes benefits most, and this depth preference can be\nfactored into the architecture selection alongside inference cost.\n5.4.2 Analyzing Learned AttnRes Patterns\nWe visualize the learned weights αi→lin Fig. 8 for the 16-head model (from Table 2) with both full and block ( N=8 )\nAttnRes. Each heatmap shows how the lth attention or MLP layer (rows) allocates its attention over previous sources\n(columns), with pre-attention and pre-MLP layers shown separately. We highlight three key observations:\n•Preserved locality.Each layer attends most strongly to its immediate predecessor, yet selective off-diagonal\nconcentrations emerge (e.g., layer 4 attending to early sources, layers 15–16 reaching back under the block setting),\nindicating learned skip connections beyond the standard residual path.\n•Layer specialization.The embedding h1retains non-trivial weight throughout, especially in pre-attention layers.\nPre-MLP inputs show sharper diagonal reliance on recent representations, while pre-attention inputs maintain broader\nreceptive fields, consistent with attention routing information across layers and MLPs operating locally.\n•Block AttnRes preserves structure.Diagonal dominance, embedding persistence, and layer specialization all\ntransfer from the full to the block variant, suggesting that block-wise compression acts as implicit regularization\nwhile preserving the essential information pathways.\n13" - }, - { - "page": 14, - "content": "Attention ResidualsTECHNICALREPORT\nTable 5: Comparison of residual update mechanisms.Weight: whether the mixing coefficients are architecture-fixed, learned-static\n(fixed after training), or input-dependent (dynamic).Source: which earlier representations layer lcan access. Normalization is\nomitted from most formulas for clarity.\nMethod Update rule Weight Source\nSingle-state recurrence: layerlreceives onlyh l−1\nResidual [12]h l=hl−1+fl−1(hl−1)Fixedh l−1\nReZero [2]h l=hl−1+αl·fl−1(hl−1)Statich l−1\nLayerScale [50]h l=hl−1+ diag(λ l)·fl−1(hl−1)Statich l−1\nHighway [45]h l= (1−g l)⊙h l−1+gl⊙fl−1(hl−1)Dynamich l−1\nDeepNorm [54]h l= Norm(αh l−1+fl−1(hl−1))Fixedh l−1\nKEEL [4]h l= Norm(αh l−1+fl−1(Norm(h l−1)))Fixedh l−1\nMulti-state recurrence: layerlreceivesmstreams\nSiameseNorm [27]h1\nl=Norm(h1\nl−1+yl−1);h2\nl=h2\nl−1+yl−1 Fixed 2 streams\nHC/mHC [72, 59]H l=H l−1Al+fl−1(Hl−1αl−1)β⊤\nl−1 Dynamicmstreams\nDDL [67]H l= (I−β lklk⊤\nl)Hl−1+βlklv⊤\nl Dynamicd vstreams\nCross-layer access: layerlcan access individual earlier-layer outputs\nDenseNet [17]h l= ConvPool([h 1;f1(h1);. . .;f l−1(hl−1)])Static[h 1, . . . ,h l−1]\nDenseFormer [36]h l=α 0→lh1+Pl−1\ni=1αi→lfi(hi)Static[h 1, . . . ,h l−1]\nMRLA [10]1hl=Pl−1\ni=1σ\u0000\nConvPool(f l−1(hl−1))\u0001⊤σ\u0000\nConvPool(f i(hi))\u0001\nConv(f i(hi))Dynamic[h 1, . . . ,h l−1]\nFull2hl∝Pl−1\ni=0ϕ(w l,ki)vi Dynamic [h1, . . . ,h l−1]AttnRes (ours)Block3hl∝Pn−1\ni=0ϕ(w l,ki)vi+ϕ(w l,kj\nn)vj\nn Dynamic [b0, . . . ,b n−1,bj\nn]\n1ConvPool: pooling operation followed by convolution (channel projection).\n2ϕ(q,k) = exp\u0000\nq⊤RMSNorm(k)\u0001\n;ki=vi;v0=h 1,vi≥1=fi(hi).softmaxjointly normalized over all sources.\n3Sameϕand normalization as Full;v i=bi,vj\nn=bj\nn.\n6 Discussions\n6.1 Sequence-Depth Duality\nResidual connections propagate information over depth via a fixed recurrence hl=hl−1+fl−1(hl−1), much as RNNs\npropagate information over time. Test-Time Training (TTT) [46] formalizes the sequence side of this analogy (cf. Fast\nWeight Programmers [43, 32]), casting each recurrent step as gradient descent on a self-supervised loss:\nWt=W t−1−η∇ℓ(W t−1;xt),(9)\nwhere a slow network parameterizes ℓand the state Wis updated once per token. When fis linear, this reduces to\nvanilla linear attention St=St−1+ktv⊤\nt. The standard residual exhibits the same additive form along depth, with hl\nserving as the state and each layerf lacting as one “gradient step.”\nAs noted by [4], this duality extends to richer variants (Table 5). Data-dependent gates on the sequence side [47, 63]\ncorrespond to Highway networks [45] on the depth side; the delta rule [42, 62, 69] corresponds to DDL [67]; and\nMRLA [10] mirrors GLA’s [63] gated linear attention. These methods all refine the recurrent update while remaining\nwithin the recurrence paradigm. AttnRes goes a step further and replaces depth-wise recurrence with direct cross-layer\nattention, just as Transformers replaced temporal recurrence with self-attention. Since the number of layers in current\narchitectures remains well within the practical regime of softmax attention, we adopt vanilla depth-wise attention.\nIncorporating more expressive yet memory-efficient (e.g. linear-complexity) alternatives is a natural direction for future\nwork.\n6.2 Residual Connections as Structured Matrices\nThe residual variants discussed above can all be viewed as weighted aggregations over previous layer outputs. We\nformalize this with adepth mixing matrix M∈RL×L, where Mi→lis the weight that layer lassigns to the output of\nlayer i. The variants differ in how these weights arise (fixed, learned, or input-dependent) and whether Mis constrained\nto low rank or allowed to be dense. The semiseparable rank ofM[8] offers a unified lens for comparing them.\nConcretely, the input to layer lishl=Pl−1\ni=0Mi→lvi, where v0=h 1(embedding) and vi=fi(hi)fori≥1 . Fig. 9\nvisualizesMfor representative methods; we derive each below.\n14" - }, - { - "page": 15, - "content": "Attention ResidualsTECHNICALREPORT\nHighway\n\n1\nγ×\n1→2g2\nγ×\n1→3g2γ×\n2→3g3\nγ×\n1→4g2γ×\n2→4g3γ×\n3→4g4\n(m)HC\n\nβ⊤\n0α1\nβ⊤\n0A×\n1→2α2 β⊤\n1α2\nβ⊤\n0A×\n1→3α3β⊤\n1A×\n2→3α3 β⊤\n2α3\nβ⊤\n0A×\n1→4α4β⊤\n1A×\n2→4α4β⊤\n2A×\n3→4α4 β⊤\n3α4\n\nFull AttnRes\n\nϕ(w 1,k0)\nϕ(w 2,k0) ϕ(w 2,k1)\nϕ(w 3,k0) ϕ(w 3,k1) ϕ(w 3,k2)\nϕ(w 4,k0) ϕ(w 4,k1) ϕ(w 4,k2) ϕ(w 4,k3)\nBlock AttnRes\n\nϕ(w 1,k0)\nϕ(w 2,k0) ϕ(w 2,k1)\nϕ(w 3,k0)\nϕ(w 4,k0) ϕ(w 4,k3)ϕ(w 3,k1+k 2)\nϕ(w 4,k1+k 2)\n\nFigure 9: Depth mixing matrices Mfor four residual variants ( L=4 ; Block AttnRes uses block size S=2 ). Highway is shown with\nscalar gates for clarity. AttnRes panels show unnormalized ϕscores; background colors group entries that share the same source\n(Full AttnRes) or the same source block (Block AttnRes).\n•Standard residual [12], hl=hl−1+fl−1(hl−1). Expanding gives hl=Pl−1\ni=0vi, soMi→l= 1for all i < l andM\nis an all-ones lower-triangular matrix:\n\nh1\nh2\n...\nhL\n=\n1\n1 1\n.........\n1 1···1\n\nv0\nv1\n...\nvL−1\n\n•Highway [45], hl= (1−g l)hl−1+glfl−1(hl−1)(written here with scalar gates for clarity; the element-wise\nextension is straightforward). Defining the carry product γ×\ni→l:=Ql\nj=i+1(1−g j), the weights are M0→l=γ×\n1→l\nfor the embedding and Mi→l=gi+1γ×\ni+1→lfori≥1 . Since the cumulative products factor through scalar gates, M\nis 1-semiseparable [8], the same rank as the standard residual but with input-dependent weights. The weights sum to\none by construction, making Highway a softmax-free depth-wise instance of stick-breaking attention [49].\n• (m)HC [72, 59] maintainmparallel streamsH l∈Rd×m, updated via\nHl=H l−1Al+fl−1(Hl−1αl−1)β⊤\nl−1,\nwhere Al∈Rm×mis a learned transition matrix, αl−1∈Rmmixes streams into a single input for fl−1, and\nβl−1∈Rmdistributes the output back across streams. Unrolling the recurrence gives the effective weight\nMi→l=β⊤\niA×\ni+1→lαl,(10)\nwhereA×\ni→j:=Qj\nk=i+1Ak. The m×m transitions render Mm -semiseparable [8]. mHC [59, 64] further constrains\neachA lto be doubly stochastic, stabilizing the cumulative products across depth.\n•Full AttnRes computes Mi→l=α i→lviaϕ(w l,ki) = exp\u0000\nw⊤\nlRMSNorm(k i)\u0001\nwith normalization, where\nki=viare input-dependent layer outputs, yielding a dense, rank-LM.\n•Block AttnRes partitions layers into Nblocks B1, . . . ,B N. For sources iin a completed earlier block Bn, all share\nthe block-level key/value bn, soMi→l=αn→lfor every i∈ B n. Within the current block, each layer additionally\nattends over the evolving partial sum bi−1\nn, introducing one extra distinct source per intra-block position. The effective\nrank of Mtherefore lies between NandN+S (where Sis the block size), interpolating between standard residual\n(N=1) and Full AttnRes (N=L).\nPracticality.The structured-matrix perspective serves two purposes. First, it enables analytical insights that are not\napparent from the recurrence form alone. The input-dependent Mof AttnRes, for instance, reveals depth-wise attention\nsinks (§5.4.2), where certain layers consistently attract high weight regardless of input, mirroring the same phenomenon\nin sequence-wise attention [57]. Second, it informs new designs by exposing which properties of the kernel ϕmatter. For\nexample, when ϕdecomposes as ϕ(q,k) =φ(q)⊤φ(k) for some feature map φ[23], depth-wise attention collapses\ninto a recurrence—precisely the structure underlying the MRLA–GLA and DDL–DeltaNet correspondences noted\nabove.\n15" - }, - { - "page": 16, - "content": "Attention ResidualsTECHNICALREPORT\nPrior Residuals as Depth-Wise Linear AttentionThe structured-matrix perspective further relates to the sequence-\ndepth duality by showing that existing residual variants are, in effect, instances oflinearattention over the depth axis.\nFor example, the unrolled (m)HC weight Mi→l=β⊤\niA×\ni+1→lαl(Eq. 10) admits a natural attention interpretation in\nwhich αlplays the role of a query issued by layer l,βiserves as a key summarizing the contribution of layer i, and\nthe cumulative transition A×\ni+1→lacts as a depth-relative positional operator [69] governing the query–key interaction\nacross intervening layers. Notably, themparallel streams correspond to state expansion [40, 29] along the depth axis,\nexpanding the recurrent state from dtod×m and thereby increasing the semiseparable rank of M. [58] show that\nreplacing A×\ni+1→lwith the identity matrix still yields competitive performance, highlighting the role of state expansion.\nThrough this lens, methods like (m)HC thus act as depth-wiselinearattention with matrix-valued states, while AttnRes\nacts as depth-wisesoftmaxattention.\n7 Related Work\nNormalization, Scaling, and Depth Stability.The standard residual update hl+1=h l+fl(hl)[12] presents a\nfundamental tension betweennormalization placementandgradient propagation. PostNorm [52] maintains bounded\nmagnitudes but distorts gradients, as repeated normalization on the residual path compounds into gradient vanishing at\ndepth [60]. PreNorm [34, 60] restores a clean identity path yet introduces unbounded magnitude growth: since ∥hl∥\ngrows as O(L) , each layer’s relative contribution shrinks, compelling deeper layers to produce ever-larger outputs\nand limiting effective depth [27]. Subsequent work reconciles both desiderata via scaled residual paths [54], hybrid\nnormalization [73], amplified skip connections [4], or learned element-wise gates [45] (see Table 5). AttnRes sidesteps\nthis tension by replacing the additive recurrence with selective aggregation over individual earlier-layer outputs, avoiding\nboth the cumulative magnitude growth of PreNorm and the repeated scale contraction of PostNorm.\nMulti-State Recurrence.All single-state methods above condition layer lonly on hl−1, from which individual\nearlier-layer contributions cannot be selectively retrieved. Several methods address this by widening the recurrence\nto multiple parallel streams: Hyper-Connections [72] and its stabilized variant mHC [59] maintain mstreams with\nlearned mixing matrices; DDL [67] maintains a matrix state updated via a delta-rule erase-and-write mechanism;\nSiameseNorm [27] maintains two parameter-shared streams—one PreNorm and one PostNorm—to preserve identity\ngradients and bounded representations. While these methods alleviate information compression, they still condition\non the immediate predecessor’s state; AttnRes is orthogonal, providing selective access to individual earlier-layer\noutputs while remaining compatible with any normalization or gating scheme. We discuss the formal connection to\nHyper-Connections in § 6.2.\nCross-Layer Connectivity.A separate line of work bypasses the single-state bottleneck by giving each layer direct\naccess to individual earlier-layer outputs. The simplest approach uses static weights: DenseNet [17] concatenates all\npreceding feature maps; ELMo [38] computes a softmax -weighted sum of layer representations with learned scalar\nweights; DenseFormer [36] and ANCRe [68] assign learned per-layer scalar coefficients fixed after training. For\ninput-dependent aggregation, MUDDFormer [56] generates position-dependent weights via a small MLP across four\ndecoupled streams; MRLA [10] applies element-wise sigmoid gating over all previous layers, though its separable\nquery–key product is closer to linear attention than softmax -based retrieval. Other methods trade full cross-layer access\nfor more targeted designs: Value Residual Learning [71] accesses only a single earlier layer; LAuReL [30] augments\nthe residual with low-rank projections over the previous kactivations; Dreamer [24] combines sequence attention with\ndepth attention and sparse experts. AttnRes combines softmax -normalized, input-dependent weights with selective\naccess to all preceding layers through a single d-dimensional pseudo-query per layer, and introduces a block structure\nreducing cost from O(L2)toO(LN) . Cache-based pipeline communication and a two-phase computation strategy\n(§ 4) make Block AttnRes practical at scale with negligible overhead.\nConclusion\nInspired by the duality between sequence and depth, we introduce AttnRes, which replaces fixed, uniform residual\naccumulation with learned, input-dependent depth-wise attention. We validate the method through ablation studies and\nscaling law experiments, showing that its gains persist across scales. Because Full AttnRes must access all preceding\nlayer outputs at every layer, the memory footprint of cross-layer aggregation grows as O(Ld) , which is prohibitive\nfor large-scale models on current hardware. We therefore introduce Block AttnRes, which partitions layers into N\nblocks and attends over block-level representations. Empirically, using about 8 blocks recovers most of the gains of Full\nAttnRes, while finer-grained blocking remains a promising direction as future hardware constraints relax. Together with\ncross-stage caching and a two-phase computation strategy, Block AttnRes is practical at scale, incurring only marginal\ntraining overhead and minimal inference overhead.\n16" - }, - { - "page": 17, - "content": "Attention ResidualsTECHNICALREPORT\nReferences\n[1] Jacob Austin et al.Program Synthesis with Large Language Models. 2021. arXiv: 2108.07732 [cs.PL] .URL:\nhttps://arxiv.org/abs/2108.07732.\n[2] Thomas Bachlechner et al.ReZero is All You Need: Fast Convergence at Large Depth. 2020. arXiv: 2003.04887\n[cs.LG].URL:https://arxiv.org/abs/2003.04887.\n[3] Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio.Neural Machine Translation by Jointly Learning to\nAlign and Translate. 2016. arXiv:1409.0473 [cs.CL].URL:https://arxiv.org/abs/1409.0473.\n[4] Chen Chen and Lai Wei.Post-LayerNorm Is Back: Stable, ExpressivE, and Deep. 2026. arXiv: 2601.19895\n[cs.LG].URL:https://arxiv.org/abs/2601.19895.\n[5] Mark Chen et al.Evaluating Large Language Models Trained on Code. 2021. arXiv: 2107.03374 [cs.LG] .\nURL:https://arxiv.org/abs/2107.03374.\n[6] Peter Clark et al. “Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge”. In:\narXiv:1803.05457v1(2018).\n[7] Karl Cobbe et al.Training Verifiers to Solve Math Word Problems. 2021. arXiv: 2110.14168 [cs.LG] .URL:\nhttps://arxiv.org/abs/2110.14168.\n[8] Tri Dao and Albert Gu. “Transformers are SSMs: Generalized Models and Efficient Algorithms Through\nStructured State Space Duality”. In:CoRRabs/2405.21060 (2024).DOI: 10.48550/ARXIV.2405.21060 . arXiv:\n2405.21060.URL:https://doi.org/10.48550/arXiv.2405.21060.\n[9] DeepSeek-AI et al.DeepSeek-V3 Technical Report. 2025. arXiv: 2412.19437 [cs.CL] .URL: https://arxiv.\norg/abs/2412.19437.\n[10] Yanwen Fang et al.Cross-Layer Retrospective Retrieving via Layer Attention. 2023. arXiv: 2302 . 03985\n[cs.CV].URL:https://arxiv.org/abs/2302.03985.\n[11] Andrey Gromov et al.The Unreasonable Ineffectiveness of the Deeper Layers. 2025. arXiv: 2403.17887\n[cs.CL].URL:https://arxiv.org/abs/2403.17887.\n[12] Kaiming He et al.Deep Residual Learning for Image Recognition. 2015. arXiv: 1512.03385 [cs.CV] .URL:\nhttps://arxiv.org/abs/1512.03385.\n[13] Dan Hendrycks et al.Measuring Massive Multitask Language Understanding. 2021. arXiv: 2009.03300\n[cs.CY].URL:https://arxiv.org/abs/2009.03300.\n[14] Dan Hendrycks et al.Measuring Mathematical Problem Solving With the MATH Dataset. 2021. arXiv: 2103.\n03874 [cs.LG].URL:https://arxiv.org/abs/2103.03874.\n[15] Jordan Hoffmann et al.Training Compute-Optimal Large Language Models. 2022. arXiv: 2203.15556 [cs.CL] .\nURL:https://arxiv.org/abs/2203.15556.\n[16] Shengding Hu et al.MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training\nStrategies. 2024. arXiv:2404.06395 [cs.CL].URL:https://arxiv.org/abs/2404.06395.\n[17] Gao Huang et al.Densely Connected Convolutional Networks. 2018. arXiv: 1608.06993 [cs.CV] .URL:\nhttps://arxiv.org/abs/1608.06993.\n[18] Yanping Huang et al. “GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism”. In:\nAdvances in NeurIPS. 2019.\n[19] Yuzhen Huang et al. “C-eval: A multi-level multi-discipline chinese evaluation suite for foundation models”. In:\nAdvances in NeurIPS36 (2023), pp. 62991–63010.\n[20] Robert A. Jacobs et al. “Adaptive Mixtures of Local Experts”. In:Neural Computation3.1 (1991), pp. 79–87.\nDOI:10.1162/neco.1991.3.1.79.\n[21] Mandar Joshi et al. “Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension”.\nIn:arXiv preprint arXiv:1705.03551(2017).\n[22] Jared Kaplan et al.Scaling Laws for Neural Language Models. 2020. arXiv: 2001.08361 [cs.LG] .URL:\nhttps://arxiv.org/abs/2001.08361.\n[23] Angelos Katharopoulos et al. “Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention”.\nIn:Proceedings of ICML. Ed. by Hal Daumé III and Aarti Singh. PMLR, 2020, pp. 5156–5165.URL: https:\n//proceedings.mlr.press/v119/katharopoulos20a.html.\n[24] Jonas Knupp et al.Depth-Recurrent Attention Mixtures: Giving Latent Reasoning the Attention it Deserves. 2026.\narXiv:2601.21582 [cs.AI].URL:https://arxiv.org/abs/2601.21582.\n[25] Aitor Lewkowycz et al.Solving Quantitative Reasoning Problems with Language Models. 2022. arXiv: 2206.\n14858 [cs.CL].URL:https://arxiv.org/abs/2206.14858.\n17" - }, - { - "page": 18, - "content": "Attention ResidualsTECHNICALREPORT\n[26] Haonan Li et al. “CMMLU: Measuring massive multitask language understanding in Chinese”. In:Findings\nof the Association for Computational Linguistics: ACL 2024. Ed. by Lun-Wei Ku, Andre Martins, and Vivek\nSrikumar. Bangkok, Thailand: Association for Computational Linguistics, Aug. 2024, pp. 11260–11285.DOI:\n10 . 18653 / v1 / 2024 . findings - acl . 671 .URL: https : / / aclanthology . org / 2024 . findings -\nacl.671/.\n[27] Tianyu Li et al.SiameseNorm: Breaking the Barrier to Reconciling Pre/Post-Norm. 2026. arXiv: 2602.08064\n[cs.LG].URL:https://arxiv.org/abs/2602.08064.\n[28] Jingyuan Liu et al.Muon is Scalable for LLM Training. 2025. arXiv: 2502.16982 [cs.LG] .URL: https:\n//arxiv.org/abs/2502.16982.\n[29] Brian Mak and Jeffrey Flanigan.Residual Matrix Transformers: Scaling the Size of the Residual Stream. 2025.\narXiv:2506.22696 [cs.LG].URL:https://arxiv.org/abs/2506.22696.\n[30] Gaurav Menghani, Ravi Kumar, and Sanjiv Kumar.LAuReL: Learned Augmented Residual Layer. 2025. arXiv:\n2411.07501 [cs.LG].URL:https://arxiv.org/abs/2411.07501.\n[31] Maxim Milakov and Natalia Gimelshein.Online normalizer calculation for softmax. 2018. arXiv: 1805.02867\n[cs.PF].URL:https://arxiv.org/abs/1805.02867.\n[32] Tsendsuren Munkhdalai et al. “Metalearned Neural Memory”. In:ArXivabs/1907.09720 (2019).URL: https:\n//api.semanticscholar.org/CorpusID:198179407.\n[33] Deepak Narayanan et al.Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM.\n2021. arXiv:2104.04473 [cs.CL].URL:https://arxiv.org/abs/2104.04473.\n[34] Toan Q. Nguyen and Julian Salazar. “Transformers without Tears: Improving the Normalization of Self-\nAttention”. In:Proceedings of IWSLT. Ed. by Jan Niehues et al. 2019.URL: https : / / aclanthology .\norg/2019.iwslt-1.17/.\n[35] OpenAI et al.GPT-4 Technical Report. 2024. arXiv: 2303.08774 [cs.CL] .URL: https://arxiv.org/abs/\n2303.08774.\n[36] Matteo Pagliardini et al.DenseFormer: Enhancing Information Flow in Transformers via Depth Weighted\nAveraging. 2024. arXiv:2402.02622 [cs.CL].URL:https://arxiv.org/abs/2402.02622.\n[37] Bowen Peng et al. “Yarn: Efficient context window extension of large language models”. In:arXiv preprint\narXiv:2309.00071(2023).\n[38] Matthew E. Peters et al. “Deep Contextualized Word Representations”. In:Proceedings of NAACL. 2018,\npp. 2227–2237.URL:https://aclanthology.org/N18-1202/.\n[39] Reiner Pope et al.Efficiently Scaling Transformer Inference. 2022. arXiv:2211.05102 [cs.LG].\n[40] Zhen Qin et al.HGRN2: Gated Linear RNNs with State Expansion. 2024. arXiv:2404.07904 [cs.CL].\n[41] David Rein et al. “Gpqa: A graduate-level google-proof q&a benchmark”. In:First Conference on Language\nModeling. 2024.\n[42] Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. “Linear Transformers Are Secretly Fast Weight Program-\nmers”. In:Proceedings of ICML. Ed. by Marina Meila and Tong Zhang. PMLR, 2021, pp. 9355–9366.URL:\nhttps://proceedings.mlr.press/v139/schlag21a.html.\n[43] Jürgen Schmidhuber. “Learning to control fast-weight memories: An alternative to dynamic recurrent networks”.\nIn:Neural Computation4.1 (1992), pp. 131–139.\n[44] Freda Shi et al.Language Models are Multilingual Chain-of-Thought Reasoners. 2022. arXiv: 2210.03057\n[cs.CL].URL:https://arxiv.org/abs/2210.03057.\n[45] Rupesh Kumar Srivastava, Klaus Greff, and Jürgen Schmidhuber.Highway Networks. 2015. arXiv: 1505.00387\n[cs.LG].URL:https://arxiv.org/abs/1505.00387.\n[46] Yu Sun et al. “Learning to (Learn at Test Time): RNNs with Expressive Hidden States”. In:ArXivabs/2407.04620\n(2024).URL:https://api.semanticscholar.org/CorpusID:271039606.\n[47] Yutao Sun et al.Retentive Network: A Successor to Transformer for Large Language Models. 2023. arXiv:\n2307.08621 [cs.CL].\n[48] Mirac Suzgun et al. “Challenging big-bench tasks and whether chain-of-thought can solve them”. In:arXiv\npreprint arXiv:2210.09261(2022).\n[49] Shawn Tan et al. “Scaling Stick-Breaking Attention: An Efficient Implementation and In-depth Study”. In:\nProceedings of ICLR. 2025.\n[50] Hugo Touvron et al.Going deeper with Image Transformers. 2021. arXiv: 2103.17239 [cs.CV] .URL: https:\n//arxiv.org/abs/2103.17239.\n[51] Hugo Touvron et al.LLaMA: Open and Efficient Foundation Language Models. 2023. arXiv: 2302.13971\n[cs.CL].\n18" - }, - { - "page": 19, - "content": "Attention ResidualsTECHNICALREPORT\n[52] Ashish Vaswani et al. “Attention is All you Need”. In:Advances in NeurIPS. Ed. by I. Guyon et al. Curran\nAssociates, Inc., 2017.URL: https://proceedings.neurips.cc/paper_files/paper/2017/file/\n3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf.\n[53] Ashish Vaswani et al. “Attention is All you Need”. In:Advances in NeurIPS. Ed. by I. Guyon et al. V ol. 30.\nCurran Associates, Inc., 2017.URL: https://proceedings.neurips.cc/paper_files/paper/2017/\nfile/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf.\n[54] Hongyu Wang et al.DeepNet: Scaling Transformers to 1,000 Layers. 2022. arXiv: 2203.00555 [cs.CL] .URL:\nhttps://arxiv.org/abs/2203.00555.\n[55] Yubo Wang et al. “Mmlu-pro: A more robust and challenging multi-task language understanding benchmark”. In:\nAdvances in NeurIPS37 (2024), pp. 95266–95290.\n[56] Da Xiao et al. “MUDDFormer: Breaking Residual Bottlenecks in Transformers via Multiway Dynamic Dense\nConnections”. In:Proceedings of ICML. 2025.\n[57] Guangxuan Xiao et al. “Efficient streaming language models with attention sinks”. In:arXiv preprint\narXiv:2309.17453(2023).\n[58] Tian Xie.Your DeepSeek mHC Might Not Need the “m”. Zhihu blog post. 2026.URL: https://zhuanlan.\nzhihu.com/p/2010852389670908320.\n[59] Zhenda Xie et al.mHC: Manifold-Constrained Hyper-Connections. 2026. arXiv: 2512.24880 [cs.CL] .URL:\nhttps://arxiv.org/abs/2512.24880.\n[60] Ruibin Xiong et al.On Layer Normalization in the Transformer Architecture. 2020. arXiv: 2002.04745 [cs.LG] .\nURL:https://arxiv.org/abs/2002.04745.\n[61] Bowen Yang et al.Rope to Nope and Back Again: A New Hybrid Attention Strategy. 2025. arXiv: 2501.18795\n[cs.CL].URL:https://arxiv.org/abs/2501.18795.\n[62] Songlin Yang, Jan Kautz, and Ali Hatamizadeh. “Gated Delta Networks: Improving Mamba2 with Delta Rule”.\nIn:Proceedings of ICLR. 2025.URL:https://openreview.net/forum?id=r8H7xhYPwz.\n[63] Songlin Yang et al. “Gated Linear Attention Transformers with Hardware-Efficient Training”. In:Proceedings of\nICML. PMLR, 2024.\n[64] Yongyi Yang and Jianyang Gao.mHC-lite: You Don’t Need 20 Sinkhorn-Knopp Iterations. 2026. arXiv: 2601.\n05732 [cs.LG].URL:https://arxiv.org/abs/2601.05732.\n[65] Rowan Zellers et al. “HellaSwag: Can a Machine Really Finish Your Sentence?” In:Proceedings of the 57th\nAnnual Meeting of the Association for Computational Linguistics. 2019.\n[66] Biao Zhang and Rico Sennrich. “Root mean square layer normalization”. In:Advances in NeurIPS32 (2019).\n[67] Yifan Zhang et al.Deep Delta Learning. 2026. arXiv: 2601.00417 [cs.LG] .URL: https://arxiv.org/\nabs/2601.00417.\n[68] Yilang Zhang et al.ANCRe: Adaptive Neural Connection Reassignment for Efficient Depth Scaling. 2026. arXiv:\n2602.09009 [cs.LG].URL:https://arxiv.org/abs/2602.09009.\n[69] Yu Zhang et al.Kimi Linear: An Expressive, Efficient Attention Architecture. 2025. arXiv: 2510.26692 [cs.CL] .\n[70] Shu Zhong et al.Understanding Transformer from the Perspective of Associative Memory. 2025. arXiv: 2505.\n19488 [cs.LG].URL:https://arxiv.org/abs/2505.19488.\n[71] Zhanchao Zhou et al. “Value Residual Learning”. In:Proceedings of ACL. Ed. by Wanxiang Che et al. Vienna,\nAustria, 2025, pp. 28341–28356.URL:https://aclanthology.org/2025.acl-long.1375/.\n[72] Defa Zhu et al.Hyper-Connections. 2025. arXiv: 2409.19606 [cs.LG] .URL: https://arxiv.org/abs/\n2409.19606.\n[73] Zhijian Zhuo et al.HybridNorm: Towards Stable and Efficient Transformer Training via Hybrid Normalization.\n2025. arXiv:2503.04598 [cs.CL].URL:https://arxiv.org/abs/2503.04598.\n19" - }, - { - "page": 20, - "content": "Attention ResidualsTECHNICALREPORT\nA Contributions\nThe authors are listed in order of the significance of their contributions, with those in project leadership roles appearing\nlast.\nGuangyu Chen∗\nYu Zhang∗\nJianlin Su∗\nWeixin Xu\nSiyuan Pan\nYaoyu Wang\nYucheng Wang\nGuanduo Chen\nBohong Yin\nYutian Chen\nJunjie Yan\nMing Wei\nY . Zhang\nFanqing Meng\nChao Hong\nXiaotong Xie\nShaowei Liu\nEnzhe Lu\nYunpeng TaiYanru Chen\nXin Men\nHaiqing Guo\nY . Charles\nHaoyu Lu\nLin Sui\nJinguo Zhu\nZaida Zhou\nWeiran He\nWeixiao Huang\nXinran Xu\nYuzhi Wang\nGuokun Lai\nYulun Du\nYuxin Wu\nZhilin Yang\nXinyu Zhou\n∗Equal contribution\n20" - }, - { - "page": 21, - "content": "Attention ResidualsTECHNICALREPORT\nB Optimized Inference I/O for Full Attention Residuals\nA naïve implementation of Full AttnRes scans all preceding layer outputs at every layer, so memory traffic scales\nlinearly with depth. As noted in §4.2, however, the pseudo-query wlis a learned parameter independent of both the\ninput and the hidden state. We can therefore batch inter-block accesses across layers in a two-phase schedule, bringing\ntotal I/O well below the naïve bound.\nNote that the block partition introduced below is purely an inference scheduling device. Unlike Block AttnRes, it leaves\nthe model architecture unchanged and does not replace per-layer sources with block summaries; it simply makes the\namortization argument concrete.\nSetupLet the model have Llayers and hidden dimension d, partitioned into Ncontiguous blocks of size S=L/N .\nInference proceeds one block at a time: Phase 1 jointly computes inter-block attention for all Slayers in the block\nagainst all preceding blocks, and Phase 2 walks through intra-block dependencies sequentially.\nPhase 1: Batched Inter-block Attention\nConsider block nwith its Slayers. The queries {wl}l∈Bnare all known before execution begins, so the (n−1)S\npreceding key–value pairs need only be read once from HBM and reused across all Squeries. The read cost for block n\nis therefore\nRead(n)\ninter= 2(n−1)Sd,(11)\nwhere the factor of2accounts for both keys and values. Summing over allNblocks and usingSN=L:\nRead inter=NX\nn=12(n−1)Sd= 2Sd·N(N−1)\n2=dL(N−1).(12)\nPhase 1 also writes oned-dimensional output per layer, givingWrite(n)\ninter=Sdper block and\nWrite inter=Ld(13)\nin total.\nPhase 2: Sequential Intra-block Attention\nPhase 1 covers all sources before the current block. Within the block, however, each layer depends on those before it,\nso these must be handled in order. Layer t(1≤t≤S ) reads t−1 intra-block key–value pairs at a cost of 2(t−1)d .\nSumming over one block:\nRead(n)\nintra=SX\nt=12(t−1)d=S(S−1)d.(14)\nPhase 2 also writes one output per layer, soWrite(n)\nintra=Sd.\nTotal Amortized I/O per Layer\nSumming both phases over allNblocks:\nRead total=dL(N−1) +N·S(S−1)d,Write total= 2Ld.(15)\nDividing byLand usingSN=L:\nRead per layer= (N−1)d+ (S−1)d= (S+N−2)d,Write per layer= 2d,(16)\nTotal I/O per layer= (S+N)d. (17)\nBatching inter-block reads thus brings per-layer I/O from O(L) down to O(S+N) . The schedule follows the same\ntwo-phase split as Block AttnRes: inter-block attention accounts for the bulk of the traffic, while sequential computation\nstays local within each block.\n21" - } - ] -} \ No newline at end of file diff --git a/examples/workspace/_meta.json b/examples/workspace/_meta.json deleted file mode 100644 index daf212c70..000000000 --- a/examples/workspace/_meta.json +++ /dev/null @@ -1,9 +0,0 @@ -{ - "12345678-abcd-4321-abcd-123456789abc": { - "type": "pdf", - "doc_name": "attention-residuals.pdf", - "doc_description": "This document introduces \"Attention Residuals\" (AttnRes) and its scalable variant \"Block AttnRes,\" novel mechanisms for replacing fixed residual accumulation in neural networks with learned, input-dependent depth-wise attention, addressing limitations of standard residual connections while optimizing memory, computation, and scalability for large-scale training and inference.", - "page_count": 21, - "path": "../documents/attention-residuals.pdf" - } -} \ No newline at end of file diff --git a/pageindex/__init__.py b/pageindex/__init__.py index 3f4df8d6d..3513668a2 100644 --- a/pageindex/__init__.py +++ b/pageindex/__init__.py @@ -1,5 +1,46 @@ -from .page_index import * -from .page_index_md import md_to_tree -from .retrieve import get_document, get_document_structure, get_page_content -from .client import PageIndexClient -from .tree_optimize import optimize_tree +"""PageIndex SDK.""" +from typing import TYPE_CHECKING as _TYPE_CHECKING + +from .client import PageIndexClient, PageIndexCloudClient, PageIndexLocalClient +from .errors import PageIndexAPIError + +if _TYPE_CHECKING: + from .flash import page_index_flash + from .page_index_classic import page_index, page_index_main + from .page_index_md import md_to_tree + from .tree_optimize import optimize_tree + +__all__ = [ + "PageIndexClient", "PageIndexCloudClient", "PageIndexLocalClient", + "PageIndexAPIError", + "page_index", "page_index_main", "page_index_flash", + "optimize_tree", "md_to_tree", +] + +_LAZY = { + "page_index_flash": ".flash", + "optimize_tree": ".tree_optimize", + "md_to_tree": ".page_index_md", +} +_SUBMODULES = {"client", "cloud_api", "errors", "flash", "local_api", + "local_store", "page_index_classic", "page_index_md", "tree_optimize", + "utils"} + + +def __getattr__(name): + if name.startswith("_"): + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + import importlib + if name in _SUBMODULES: + return importlib.import_module(f".{name}", __name__) + module = importlib.import_module(_LAZY.get(name, ".page_index_classic"), __name__) + try: + value = getattr(module, name) + except AttributeError: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None + globals()[name] = value + return value + + +def __dir__(): + return sorted(set(globals()) | set(__all__) | _SUBMODULES) diff --git a/pageindex/client.py b/pageindex/client.py index 894dab181..158c9b6f7 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -1,18 +1,23 @@ -import os -import uuid -import json -import asyncio -import concurrent.futures -from pathlib import Path +"""PageIndex SDK client: the 0.2.x cloud surface, now with a local mode.""" +from __future__ import annotations -import PyPDF2 +from typing import Any, Iterator, Optional, Union -from .page_index import page_index -from .page_index_md import md_to_tree -from .retrieve import get_document, get_document_structure, get_page_content -from .utils import ConfigLoader, remove_fields +from .errors import PageIndexAPIError -META_INDEX = "_meta.json" + +def _parse_pages(pages: str) -> list[int]: + result = [] + for part in pages.split(","): + part = part.strip() + if "-" in part: + start, end = (int(x) for x in part.split("-", 1)) + if start > end: + raise ValueError(f"Invalid range '{part}': start must be <= end") + result.extend(range(start, end + 1)) + else: + result.append(int(part)) + return sorted(set(result)) def _normalize_retrieve_model(model: str) -> str: @@ -27,208 +32,398 @@ def _normalize_retrieve_model(model: str) -> str: class PageIndexClient: """ - A client for indexing and retrieving document content. - Flow: index() -> get_document() / get_document_structure() / get_page_content() + Python SDK client for PageIndex. + + Cloud mode (an ``api_key`` is given) talks to the PageIndex API at + api.pageindex.ai, exactly like the 0.2.x SDK. Local mode (no ``api_key``) + runs the same operations on your machine: documents are indexed with the + open-source PageIndex pipeline using your own LLM provider key (e.g. + ``OPENAI_API_KEY`` in the environment) and stored under ``storage_path``. + + Args: + api_key (str, optional): PageIndex cloud API key + (https://dash.pageindex.ai/api-keys). Omit for local mode. + model (str, optional): Local mode only — LLM used to build document + trees. Defaults to the packaged config (see pageindex/config.yaml). + summary_model (str, optional): Local mode only — LLM used for node + summaries and document descriptions. + retrieve_model (str, optional): Local mode only — exposed as + ``client.retrieve_model`` (the agent demo reads it); the SDK + itself consumes it once agent-based local chat lands in a + later release. + storage_path (str, optional): Local mode only — directory where + indexed documents are stored. Defaults to ``./.pageindex``. - For agent-based QA, see examples/agentic_vectorless_rag_demo.py. + Usage: + client = PageIndexClient(api_key="...") # cloud + client = PageIndexClient() # local + + PageIndexCloudClient / PageIndexLocalClient pin the mode at construction + instead of inferring it from api_key. + + Local mode differences (all documented per method): indexing is + synchronous, only PDFs are supported, and ``chat_completions`` (until + agent-based local chat lands in a later release) / folders / + ``beta_headers`` / the deprecated retrieval API (``submit_query``, + ``get_retrieval``) are cloud-only. """ - def __init__(self, api_key: str = None, model: str = None, retrieve_model: str = None, workspace: str = None): - if api_key: - os.environ["OPENAI_API_KEY"] = api_key - elif not os.getenv("OPENAI_API_KEY") and os.getenv("CHATGPT_API_KEY"): - os.environ["OPENAI_API_KEY"] = os.getenv("CHATGPT_API_KEY") - self.workspace = Path(workspace).expanduser() if workspace else None - overrides = {} - if model: - overrides["model"] = model - if retrieve_model: - overrides["retrieve_model"] = retrieve_model - opt = ConfigLoader().load(overrides or None) - self.model = opt.model - self.retrieve_model = _normalize_retrieve_model(opt.retrieve_model or self.model) - if self.workspace: - self.workspace.mkdir(parents=True, exist_ok=True) - self.documents = {} - if self.workspace: - self._load_workspace() - - def index(self, file_path: str, mode: str = "auto") -> str: - """Index a document. Returns a document_id.""" - # Persist a canonical absolute path so workspace reloads do not - # reinterpret caller-relative paths against the workspace directory. - file_path = os.path.abspath(os.path.expanduser(file_path)) - if not os.path.exists(file_path): - raise FileNotFoundError(f"File not found: {file_path}") - - doc_id = str(uuid.uuid4()) - ext = os.path.splitext(file_path)[1].lower() - - is_pdf = ext == '.pdf' - is_md = ext in ['.md', '.markdown'] - - if mode == "pdf" or (mode == "auto" and is_pdf): - print(f"Indexing PDF: {file_path}") - result = page_index( - doc=file_path, - model=self.model, - if_add_node_summary='yes', - if_add_node_text='yes', - if_add_node_id='yes', - if_add_doc_description='yes' + + BASE_URL = "https://api.pageindex.ai" + + def __init__( + self, + api_key: Optional[str] = None, + *, + model: Optional[str] = None, + summary_model: Optional[str] = None, + retrieve_model: Optional[str] = None, + storage_path: Optional[str] = None, + ): + if api_key == "": + raise PageIndexAPIError( + "api_key is an empty string. Pass a real PageIndex API key for " + "cloud mode, or omit api_key entirely for local mode." ) - # Extract per-page text so queries don't need the original PDF - pages = [] - with open(file_path, 'rb') as f: - pdf_reader = PyPDF2.PdfReader(f) - for i, page in enumerate(pdf_reader.pages, 1): - pages.append({'page': i, 'content': page.extract_text() or ''}) - - self.documents[doc_id] = { - 'id': doc_id, - 'type': 'pdf', - 'path': file_path, - 'doc_name': result.get('doc_name', ''), - 'doc_description': result.get('doc_description', ''), - 'page_count': len(pages), - 'structure': result['structure'], - 'pages': pages, - } - - elif mode == "md" or (mode == "auto" and is_md): - print(f"Indexing Markdown: {file_path}") - coro = md_to_tree( - md_path=file_path, - if_thinning=False, - if_add_node_summary='yes', - summary_token_threshold=200, + if api_key is not None: + local_only = {"model": model, "summary_model": summary_model, + "retrieve_model": retrieve_model, "storage_path": storage_path} + passed = [name for name, value in local_only.items() if value is not None] + if passed: + raise PageIndexAPIError( + f"Local-mode arguments ({', '.join(passed)}) cannot be " + "combined with api_key — remove them, or omit api_key to " + "run locally." + ) + self.api_key = api_key + from .cloud_api import CloudAPI + self._api = CloudAPI(self) + else: + from .utils import ConfigLoader + overrides = {key: value for key, value in + {"model": model, "summary_model": summary_model, + "retrieve_model": retrieve_model}.items() + if value} + opt = ConfigLoader().load(overrides or None) + self.model = opt.model + self.summary_model = getattr(opt, "summary_model", None) or opt.model + self.retrieve_model = _normalize_retrieve_model( + getattr(opt, "retrieve_model", None) or opt.model) + self.storage_path = storage_path or ".pageindex" + from .local_api import LocalAPI + self._api = LocalAPI( + storage_path=self.storage_path, model=self.model, - if_add_doc_description='yes', - if_add_node_text='yes', - if_add_node_id='yes' + summary_model=self.summary_model, + retrieve_model=self.retrieve_model, ) - try: - asyncio.get_running_loop() - with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: - result = pool.submit(asyncio.run, coro).result() - except RuntimeError: - result = asyncio.run(coro) - self.documents[doc_id] = { - 'id': doc_id, - 'type': 'md', - 'path': file_path, - 'doc_name': result.get('doc_name', ''), - 'doc_description': result.get('doc_description', ''), - 'line_count': result.get('line_count', 0), - 'structure': result['structure'], - } - else: - raise ValueError(f"Unsupported file format for: {file_path}") - - print(f"Indexing complete. Document ID: {doc_id}") - if self.workspace: - self._save_doc(doc_id) - return doc_id - - @staticmethod - def _make_meta_entry(doc: dict) -> dict: - """Build a lightweight meta entry from a document dict.""" - entry = { - 'type': doc.get('type', ''), - 'doc_name': doc.get('doc_name', ''), - 'doc_description': doc.get('doc_description', ''), - 'path': doc.get('path', ''), - } - if doc.get('type') == 'pdf': - entry['page_count'] = doc.get('page_count') - elif doc.get('type') == 'md': - entry['line_count'] = doc.get('line_count') - return entry - - @staticmethod - def _read_json(path) -> dict | None: - """Read a JSON file, returning None on any error.""" + + # ---------- DOCUMENT SUBMISSION ---------- + + def submit_document( + self, + file_path: str, + mode: Optional[str] = None, + beta_headers: Optional[list[str]] = None, + folder_id: Optional[str] = None, + metadata: Optional[dict] = None, + ) -> dict[str, Any]: + """ + Submit a PDF document for processing. Returns {'doc_id': ...}. + + Cloud: uploads the file; processing is asynchronous — poll + ``is_retrieval_ready(doc_id)`` before retrieving. + + Local: indexes the document in this call (it blocks while your LLM + builds the tree — minutes for a standard index of a long document), + then stores it under ``storage_path``. Pass ``mode="flash"`` to build + the tree with PageIndex Flash (layout-based extraction, no LLM calls + for the structure; node summaries and the document description still + use ``summary_model``). ``beta_headers`` and ``folder_id`` are + cloud-only. + + Args: + file_path (str): Path to the PDF file. + mode (str, optional): Processing mode. Local mode supports + "standard" and "flash"; omit it for standard indexing. Cloud + modes are passed through (e.g. "mcp"). + beta_headers (list[str], optional): Cloud-only beta feature headers. + folder_id (str, optional): Cloud-only folder (workspace) ID. + metadata (dict, optional): Your own JSON-serializable tags for the + document; returned in get_tree/get_ocr responses and + list_documents entries (both modes). + + Returns: + dict: {'doc_id': ...} + """ + return self._api.submit_document( + file_path=file_path, mode=mode, + beta_headers=beta_headers, folder_id=folder_id, metadata=metadata, + ) + + # ---------- OCR FUNCTIONALITY ---------- + + def get_ocr(self, doc_id: str, format: str = "page") -> dict[str, Any]: + """ + Get OCR status and results. + + Args: + doc_id (str): Document ID. + format (str): 'page' for page-based results, 'node' for node-based + results, or 'raw' for concatenated markdown. + + Returns: + dict: {'doc_id', 'status', 'retrieval_ready', 'result', ...}. + With 'page', result entries are {'page_index', 'markdown', ...}. + + Local: the "OCR" result is the text extracted from the PDF while + indexing (no OCR model runs locally, so scanned/image-only PDFs have + no local text). + """ + return self._api.get_ocr(doc_id=doc_id, format=format) + + def get_page_content(self, doc_id: str, pages: str) -> list[dict[str, Any]]: + """ + Get text content of specific pages. + + Args: + doc_id (str): Document ID. + pages (str): Page specifier — '5-7', '3,8', or '12'. + + Returns: + list: Matching entries from get_ocr (format='page'). + """ + wanted = set(_parse_pages(pages)) + result = self.get_ocr(doc_id, format="page") + all_pages = result["result"] + if all_pages is None: + raise PageIndexAPIError( + f"Document '{doc_id}' is not ready " + f"(status: {result.get('status', 'unknown')})" + ) + return [p for p in all_pages if p["page_index"] in wanted] + + # ---------- TREE GENERATION ---------- + + def get_tree(self, doc_id: str, node_summary: bool = False, + include_text: bool = True) -> dict[str, Any]: + """ + Get tree generation status and results. + + Args: + doc_id (str): Document ID. + node_summary (bool): Include node summaries in the tree. + include_text (bool): Include node text (default True). + False is useful for structure-only views (saves tokens). + + Returns: + dict: {'doc_id', 'status', 'retrieval_ready', 'result', ...} where + result nodes are {'title', 'node_id', 'page_index', ('summary' / + 'prefix_summary',) ('text',) 'nodes'}. + """ + tree = self._api.get_tree(doc_id=doc_id, node_summary=node_summary, + include_text=include_text) + if not include_text and tree.get("result"): + from .utils import remove_fields + tree["result"] = remove_fields(tree["result"], fields=["text"]) + return tree + + def get_document_structure(self, doc_id: str) -> list[dict[str, Any]]: + """ + Get the document's tree structure without text — summaries included. + + Returns: + list: Tree nodes with titles, page ranges, and summaries. + """ + return self.get_tree(doc_id, node_summary=True, include_text=False)["result"] + + def is_retrieval_ready(self, doc_id: str) -> bool: + """ + Check if a document is ready for retrieval. API errors (including a + missing document) are reported as False; transport errors (connection + failures, timeouts) propagate. + """ try: - with open(path, "r", encoding="utf-8") as f: - return json.load(f) - except (json.JSONDecodeError, OSError) as e: - print(f"Warning: corrupt {Path(path).name}: {e}") - return None - - def _save_doc(self, doc_id: str): - doc = self.documents[doc_id].copy() - # Strip text from structure nodes — redundant with pages (PDF only) - if doc.get('structure') and doc.get('type') == 'pdf': - doc['structure'] = remove_fields(doc['structure'], fields=['text']) - path = self.workspace / f"{doc_id}.json" - with open(path, "w", encoding="utf-8") as f: - json.dump(doc, f, ensure_ascii=False, indent=2) - self._save_meta(doc_id, self._make_meta_entry(doc)) - # Drop heavy fields; will lazy-load on demand - self.documents[doc_id].pop('structure', None) - self.documents[doc_id].pop('pages', None) - - def _rebuild_meta(self) -> dict: - """Scan individual doc JSON files and return a meta dict.""" - meta = {} - for path in self.workspace.glob("*.json"): - if path.name == META_INDEX: - continue - doc = self._read_json(path) - if doc and isinstance(doc, dict): - meta[path.stem] = self._make_meta_entry(doc) - return meta - - def _read_meta(self) -> dict | None: - """Read and validate _meta.json, returning None on any corruption.""" - meta = self._read_json(self.workspace / META_INDEX) - if meta is not None and not isinstance(meta, dict): - print(f"Warning: {META_INDEX} is not a JSON object, ignoring") - return None - return meta - - def _save_meta(self, doc_id: str, entry: dict): - meta = self._read_meta() or self._rebuild_meta() - meta[doc_id] = entry - meta_path = self.workspace / META_INDEX - with open(meta_path, "w", encoding="utf-8") as f: - json.dump(meta, f, ensure_ascii=False, indent=2) - - def _load_workspace(self): - meta = self._read_meta() - if meta is None: - meta = self._rebuild_meta() - if meta: - print(f"Loaded {len(meta)} document(s) from workspace (legacy mode).") - for doc_id, entry in meta.items(): - doc = dict(entry, id=doc_id) - if doc.get('path') and not os.path.isabs(doc['path']): - doc['path'] = str((self.workspace / doc['path']).resolve()) - self.documents[doc_id] = doc - - def _ensure_doc_loaded(self, doc_id: str): - """Load full document JSON on demand (structure, pages, etc.).""" - doc = self.documents.get(doc_id) - if not doc or doc.get('structure') is not None: - return - full = self._read_json(self.workspace / f"{doc_id}.json") - if not full: - return - doc['structure'] = full.get('structure', []) - if full.get('pages'): - doc['pages'] = full['pages'] - - def get_document(self, doc_id: str) -> str: - """Return document metadata JSON.""" - return get_document(self.documents, doc_id) - - def get_document_structure(self, doc_id: str) -> str: - """Return document tree structure JSON (without text fields).""" - if self.workspace: - self._ensure_doc_loaded(doc_id) - return get_document_structure(self.documents, doc_id) - - def get_page_content(self, doc_id: str, pages: str) -> str: - """Return page content for the given pages string (e.g. '5-7', '3,8', '12').""" - if self.workspace: - self._ensure_doc_loaded(doc_id) - return get_page_content(self.documents, doc_id, pages) + result = self.get_tree(doc_id) + return result.get("retrieval_ready", False) + except PageIndexAPIError: + return False + + # ---------- RETRIEVAL (cloud-only, deprecated) ---------- + + def submit_query(self, doc_id: str, query: str, thinking: bool = False) -> dict[str, Any]: + """ + Submit a retrieval query for a document. Returns {'retrieval_id': ...}. + + Cloud-only: the cloud API marks this endpoint deprecated in favor of + chat completions, so local mode does not implement it — raises + PageIndexAPIError. Use ``chat_completions`` (cloud) instead. + """ + return self._require_cloud( + "submit_query is cloud-only — the retrieval API is deprecated in " + "favor of chat completions; use chat_completions in cloud mode." + ).submit_query(doc_id=doc_id, query=query, thinking=thinking) + + def get_retrieval(self, retrieval_id: str) -> dict[str, Any]: + """ + Get retrieval status and results for a submitted query. + + Cloud-only: the cloud API marks this endpoint deprecated in favor of + chat completions, so local mode does not implement it — raises + PageIndexAPIError. Use ``chat_completions`` (cloud) instead. + """ + return self._require_cloud( + "get_retrieval is cloud-only — the retrieval API is deprecated in " + "favor of chat completions; use chat_completions in cloud mode." + ).get_retrieval(retrieval_id=retrieval_id) + + # ---------- CHAT COMPLETIONS ---------- + + def chat_completions( + self, + messages: list[dict[str, str]], + stream: bool = False, + doc_id: Optional[Union[str, list[str]]] = None, + temperature: Optional[float] = None, + stream_metadata: bool = False, + enable_citations: bool = False, + ) -> Union[dict[str, Any], Iterator[str], Iterator[dict[str, Any]]]: + """ + PageIndex Chat Completions, scoped to specific PageIndex documents. + + Args: + messages: Conversation messages with 'role' and 'content' keys. + stream: Enable streaming responses. + doc_id: Document ID or list of IDs to scope the conversation. + temperature: Sampling temperature (0.0-1.0). + stream_metadata: With stream=True, yield chunk dicts instead of + text pieces. + enable_citations: Enable citation instructions in responses. + + Returns: + - stream=False: complete response dict ({'id', 'object', 'created', + 'choices', 'usage'}) + - stream=True, stream_metadata=False: iterator of text chunks + - stream=True, stream_metadata=True: iterator of chunk dicts + + Local: not yet supported — raises PageIndexAPIError. Agent-based + local chat arrives in a later release. + """ + return self._require_cloud( + "chat_completions is not yet supported in local mode — it arrives " + "in a later release. Create the client with an api_key to use " + "cloud chat." + ).chat_completions( + messages=messages, stream=stream, doc_id=doc_id, + temperature=temperature, stream_metadata=stream_metadata, + enable_citations=enable_citations, + ) + + # ---------- DOCUMENT MANAGEMENT ---------- + + def get_document(self, doc_id: str) -> dict[str, Any]: + """ + Get document metadata: {'id', 'name', 'description', 'status', + 'createdAt', 'pageNum', 'folderId'}. Status is one of "queued", + "processing", "completed", "failed" (local documents are + always "completed"; local 'folderId' is always None). + + 'createdAt' is UTC with no timezone marker, in both modes. To show + it in the user's timezone:: + + from datetime import datetime, timezone + datetime.fromisoformat(doc["createdAt"]).replace( + tzinfo=timezone.utc).astimezone() + """ + return self._api.get_document(doc_id=doc_id) + + def delete_document(self, doc_id: str) -> dict[str, Any]: + """ + Delete a PageIndex document and all its associated data. + + Returns: + dict: {'message': 'Document deleted successfully.'}, or an empty + dict if the cloud API responds with no body. + """ + return self._api.delete_document(doc_id=doc_id) + + def list_documents( + self, + limit: int = 50, + offset: int = 0, + folder_id: Optional[str] = None, + ) -> dict[str, Any]: + """ + List documents with pagination, newest first. + + Args: + limit (int): Maximum documents to return (1-100). + offset (int): Number of documents to skip. + folder_id (str, optional): Cloud-only folder filter. + + Returns: + dict: {'documents': [...], 'total', 'limit', 'offset'}. + """ + return self._api.list_documents(limit=limit, offset=offset, folder_id=folder_id) + + # ---------- FOLDER MANAGEMENT ---------- + + def create_folder( + self, + name: str, + description: Optional[str] = None, + parent_folder_id: Optional[str] = None, + ) -> dict[str, Any]: + """ + Create a folder (workspace). Cloud-only: local mode raises + PageIndexAPIError. + """ + return self._require_cloud( + "create_folder is cloud-only — folders are not supported in local " + "mode. Create the client with an api_key to use folders." + ).create_folder( + name=name, description=description, parent_folder_id=parent_folder_id, + ) + + def list_folders(self, parent_folder_id: Optional[str] = None) -> dict[str, Any]: + """ + List folders. Cloud-only: local mode raises PageIndexAPIError. + """ + return self._require_cloud( + "list_folders is cloud-only — folders are not supported in local " + "mode. Create the client with an api_key to use folders." + ).list_folders( + parent_folder_id=parent_folder_id, + ) + + def _require_cloud(self, message: str): + from .cloud_api import CloudAPI + if not isinstance(self._api, CloudAPI): + raise PageIndexAPIError(message) + return self._api + + +class PageIndexCloudClient(PageIndexClient): + """Cloud mode — requires a real API key at construction.""" + + def __init__(self, api_key: str): + if not api_key: + raise PageIndexAPIError( + "PageIndexCloudClient requires a PageIndex API key — get one " + "at https://dash.pageindex.ai/api-keys." + ) + super().__init__(api_key) + + +class PageIndexLocalClient(PageIndexClient): + """Local mode — no api_key parameter, no cloud access.""" + + def __init__( + self, + *, + model: Optional[str] = None, + summary_model: Optional[str] = None, + retrieve_model: Optional[str] = None, + storage_path: Optional[str] = None, + ): + super().__init__(None, model=model, summary_model=summary_model, + retrieve_model=retrieve_model, storage_path=storage_path) diff --git a/pageindex/cloud_api.py b/pageindex/cloud_api.py new file mode 100644 index 000000000..b7597cc9a --- /dev/null +++ b/pageindex/cloud_api.py @@ -0,0 +1,433 @@ +"""Cloud mode of the PageIndex SDK, based on the 0.2.8 client.""" +import requests +from typing import Optional, Dict, Any, List, Union, Iterator +import json +import urllib.parse + +from .errors import PageIndexAPIError + + +def _enc(value: str) -> str: + """URL-encode a path segment (ids may contain / ? # or spaces).""" + return urllib.parse.quote(str(value), safe="") + + +class CloudAPI: + """ + Python SDK client for the PageIndex API. + """ + + def __init__(self, client): + self._client = client + + @property + def BASE_URL(self) -> str: + return self._client.BASE_URL + + @property + def api_key(self) -> str: + return self._client.api_key + + def _headers(self) -> Dict[str, str]: + return {"api_key": self.api_key} + + # ---------- DOCUMENT SUBMISSION ---------- + + def submit_document( + self, + file_path: str, + mode: Optional[str] = None, + beta_headers: Optional[List[str]] = None, + folder_id: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """ + Upload a PDF document for processing. The system will automatically process both tree generation and OCR. + Immediately returns a document identifier (`doc_id`) for subsequent operations. + + Args: + file_path (str): Path to the PDF file. + mode (str, optional): Processing mode (e.g., "mcp"). Defaults to None. + beta_headers (List[str], optional): Beta feature headers (e.g., ["block_reference"] + to enable block-level content with bounding boxes). Defaults to None. + folder_id (str, optional): Folder (workspace) ID to assign the document to. Defaults to None. + metadata (dict, optional): Your own JSON-serializable tags for the document; + returned in get_tree/get_ocr responses and list_documents entries. Defaults to None. + + Returns: + dict: {'doc_id': ...} + """ + data = {'if_retrieval': True} + if mode is not None: + data['mode'] = mode + if beta_headers is not None: + data['beta_headers'] = json.dumps(beta_headers) + if folder_id is not None: + data['folder_id'] = folder_id + if metadata is not None: + data['metadata'] = json.dumps(metadata) + + with open(file_path, "rb") as f: + response = requests.post( + f"{self.BASE_URL}/doc/", + headers=self._headers(), + files={'file': f}, + data=data + ) + + if response.status_code != 200: + raise PageIndexAPIError(f"Failed to submit document: {response.text}") + return response.json() + + # ---------- OCR FUNCTIONALITY ---------- + + def get_ocr(self, doc_id: str, format: str = "page") -> Dict[str, Any]: + """ + Get OCR processing status and results. + + Args: + doc_id (str): Document ID. + format (str): Result format. Use 'page' for page-based results, 'node' for node-based results, or 'raw' for concatenated markdown. Defaults to 'page'. + + Returns: + dict: API response with status and, if ready, OCR results. + """ + if format not in ["page", "node", "raw"]: + raise ValueError("Format parameter must be 'page', 'node', or 'raw'") + + response = requests.get( + f"{self.BASE_URL}/doc/{_enc(doc_id)}/?type=ocr&format={format}", + headers=self._headers(), + timeout=30 + ) + if response.status_code != 200: + raise PageIndexAPIError(f"Failed to get OCR result: {response.text}") + return response.json() + + # ---------- TREE GENERATION ---------- + + def get_tree(self, doc_id: str, node_summary: bool = False, + include_text: bool = True) -> Dict[str, Any]: + """ + Get tree generation status and results. + + Args: + doc_id (str): Document ID. + node_summary (bool): Include node summaries (default False). + + Returns: + dict: API response with status and, if ready, tree structure. + """ + response = requests.get( + f"{self.BASE_URL}/doc/{_enc(doc_id)}/?type=tree&summary={node_summary}" + f"&include_text={str(include_text).lower()}", + headers=self._headers(), + timeout=30 + ) + if response.status_code != 200: + raise PageIndexAPIError(f"Failed to get tree result: {response.text}") + return response.json() + + # ---------- RETRIEVAL ---------- + + def submit_query(self, doc_id: str, query: str, thinking: bool = False) -> Dict[str, Any]: + """ + Submit a retrieval query for a specific PageIndex document. + + Args: + doc_id (str): Document ID. + query (str): User question or information need. + thinking (bool, optional): If true, enables deeper retrieval. Default is False. + + Returns: + dict: {'retrieval_id': ...} + """ + payload = { + "doc_id": doc_id, + "query": query, + "thinking": thinking + } + response = requests.post( + f"{self.BASE_URL}/retrieval/", + headers=self._headers(), + json=payload, + timeout=30 + ) + if response.status_code != 200: + raise PageIndexAPIError(f"Failed to submit retrieval: {response.text}") + return response.json() + + def get_retrieval(self, retrieval_id: str) -> Dict[str, Any]: + """ + Get retrieval status and results. + + Args: + retrieval_id (str): Retrieval ID. + + Returns: + dict: Retrieval status and results. + """ + response = requests.get( + f"{self.BASE_URL}/retrieval/{_enc(retrieval_id)}/", + headers=self._headers(), + timeout=30 + ) + if response.status_code != 200: + raise PageIndexAPIError(f"Failed to get retrieval result: {response.text}") + return response.json() + + # ---------- CHAT COMPLETIONS ---------- + + def chat_completions( + self, + messages: List[Dict[str, str]], + stream: bool = False, + doc_id: Optional[Union[str, List[str]]] = None, + temperature: Optional[float] = None, + stream_metadata: bool = False, + enable_citations: bool = False + ) -> Union[Dict[str, Any], Iterator[str], Iterator[Dict[str, Any]]]: + """ + PageIndex Chat Completions. Optionally scoped to specific PageIndex documents. + + Args: + messages (List[Dict[str, str]]): Conversation messages with 'role' and 'content' keys. + stream (bool, optional): Enable streaming responses. Default is False. + doc_id (Optional[Union[str, List[str]]], optional): Document ID(s) to scope the conversation. Can be a single ID or a list of IDs. + temperature (Optional[float], optional): Sampling temperature. Default is None (uses API default). + stream_metadata (bool, optional): If True and stream=True, return raw chunks with metadata instead of just text. Default is False. + enable_citations (bool, optional): Enable citation instructions in responses. Default is False. + + Returns: + Union[Dict[str, Any], Iterator[str], Iterator[Dict[str, Any]]]: + - If stream=False: Complete response dictionary + - If stream=True and stream_metadata=False: Iterator of text content chunks + - If stream=True and stream_metadata=True: Iterator of raw response chunks with metadata + """ + payload = { + "messages": messages, + "stream": stream + } + + if doc_id is not None: + payload["doc_id"] = doc_id + + if temperature is not None: + payload["temperature"] = temperature + + if enable_citations: + payload["enable_citations"] = enable_citations + + response = requests.post( + f"{self.BASE_URL}/chat/completions/", + headers=self._headers(), + json=payload, + stream=stream, + timeout=120 if stream else 300 + ) + + if response.status_code != 200: + raise PageIndexAPIError(f"Failed to get chat completion: {response.text}") + + if stream: + if stream_metadata: + return self._stream_chat_response_raw(response) + else: + return self._stream_chat_response(response) + else: + return response.json() + + def _stream_chat_response(self, response: requests.Response) -> Iterator[str]: + """ + Parse streaming chat completion response. + + Args: + response: Streaming HTTP response + + Yields: + str: Content chunks from the streaming response + """ + try: + for line in response.iter_lines(): + if line: + line = line.decode('utf-8') + if line.startswith('data: '): + data = line[6:] + if data == '[DONE]': + break + + try: + chunk = json.loads(data) + choices = chunk.get("choices") or [{}] + content = choices[0].get("delta", {}).get("content", "") + if content: + yield content + except json.JSONDecodeError: + continue + finally: + response.close() + + def _stream_chat_response_raw(self, response: requests.Response) -> Iterator[Dict[str, Any]]: + """Streaming chat completion with full metadata, including citation events.""" + try: + for line in response.iter_lines(): + if line: + line = line.decode('utf-8') + if line.startswith('data: '): + data = line[6:] + if data == '[DONE]': + break + + try: + chunk = json.loads(data) + yield chunk + except json.JSONDecodeError: + continue + finally: + response.close() + + # ---------- DOCUMENT MANAGEMENT ---------- + + def get_document(self, doc_id: str) -> Dict[str, Any]: + """ + Get document metadata including id, name, description, status, createdAt, and pageNum. + + Args: + doc_id (str): Document ID. + + Returns: + dict: Document metadata containing: + - id (str): Document ID + - name (str): Document name + - description (str): Document description + - status (str): Processing status (e.g., "queued", "processing", "completed", "failed") + - createdAt (str): Creation timestamp in ISO format + - pageNum (int): Number of pages in the document + """ + response = requests.get( + f"{self.BASE_URL}/doc/{_enc(doc_id)}/metadata/", + headers=self._headers(), + timeout=30 + ) + if response.status_code != 200: + raise PageIndexAPIError(f"Failed to get document metadata: {response.text}") + return response.json() + + def delete_document(self, doc_id: str) -> Dict[str, Any]: + """ + Delete a PageIndex document and all its associated data. + + Args: + doc_id (str): Document ID. + + Returns: + dict: API response. + """ + response = requests.delete( + f"{self.BASE_URL}/doc/{_enc(doc_id)}/", + headers=self._headers(), + timeout=30 + ) + if response.status_code != 200: + raise PageIndexAPIError(f"Failed to delete document: {response.text}") + return response.json() if response.content else {} + + def list_documents(self, limit: int = 50, offset: int = 0, folder_id: Optional[str] = None) -> Dict[str, Any]: + """ + List all documents for the authenticated user with pagination. + + Args: + limit (int, optional): Maximum number of documents to return (1-100). Defaults to 50. + offset (int, optional): Number of documents to skip. Defaults to 0. + folder_id (str, optional): Filter by folder (workspace) ID. If provided, only documents + in the specified folder are returned. Defaults to None (all documents). + + Returns: + dict: API response containing: + - documents (List[Dict]): List of document metadata objects (each includes folderId) + - total (int): Total number of documents + - limit (int): Applied limit + - offset (int): Applied offset + """ + if limit < 1 or limit > 100: + raise ValueError("limit must be between 1 and 100") + if offset < 0: + raise ValueError("offset must be non-negative") + + params = {"limit": limit, "offset": offset} + if folder_id is not None: + params["folder_id"] = folder_id + + response = requests.get( + f"{self.BASE_URL}/docs/", + headers=self._headers(), + params=params, + timeout=30 + ) + if response.status_code != 200: + raise PageIndexAPIError(f"Failed to list documents: {response.text}") + return response.json() + + # ---------- FOLDER MANAGEMENT ---------- + + def create_folder( + self, + name: str, + description: Optional[str] = None, + parent_folder_id: Optional[str] = None, + ) -> Dict[str, Any]: + """ + Create a new folder (workspace). + + Args: + name (str): Folder name. + description (str, optional): Folder description. Defaults to None. + parent_folder_id (str, optional): Parent folder ID for nesting. Defaults to None (root level). + + Returns: + dict: Created folder metadata containing: + - folder (dict): Folder info with id, name, description, parent_folder_id, + created_at, file_count, children_count + """ + payload = {"name": name} + if description is not None: + payload["description"] = description + if parent_folder_id is not None: + payload["parent_folder_id"] = parent_folder_id + + response = requests.post( + f"{self.BASE_URL}/folder/", + headers=self._headers(), + json=payload, + timeout=30 + ) + if response.status_code != 200: + raise PageIndexAPIError(f"Failed to create folder: {response.text}") + return response.json() + + def list_folders(self, parent_folder_id: Optional[str] = None) -> Dict[str, Any]: + """ + List folders. + + Args: + parent_folder_id (str, optional): Use "root" for root-level folders only, + a folder ID for subfolders, or omit for all folders. + + Returns: + dict: API response containing: + - folders (List[Dict]): List of folder metadata objects + - total (int): Total number of folders + """ + params = {} + if parent_folder_id is not None: + params["parent_folder_id"] = parent_folder_id + + response = requests.get( + f"{self.BASE_URL}/folders/", + headers=self._headers(), + params=params, + timeout=30 + ) + if response.status_code != 200: + raise PageIndexAPIError(f"Failed to list folders: {response.text}") + return response.json() diff --git a/pageindex/errors.py b/pageindex/errors.py new file mode 100644 index 000000000..e460a956b --- /dev/null +++ b/pageindex/errors.py @@ -0,0 +1,2 @@ +class PageIndexAPIError(Exception): + pass diff --git a/pageindex/local_api.py b/pageindex/local_api.py new file mode 100644 index 000000000..0e9f682c8 --- /dev/null +++ b/pageindex/local_api.py @@ -0,0 +1,325 @@ +"""Local implementation of the PageIndex SDK surface.""" +from __future__ import annotations + +import asyncio +import json +import logging +import os +import uuid +from concurrent.futures import ThreadPoolExecutor +from datetime import datetime, timezone +from typing import Any + +from .errors import PageIndexAPIError +from .local_store import DocStore + +logger = logging.getLogger(__name__) + + +def _now_iso() -> str: + """Naive UTC, millisecond precision.""" + now = datetime.now(timezone.utc).replace(tzinfo=None) + 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): + self._store = DocStore(storage_path) + self._model = model + self._summary_model = summary_model + self._retrieve_model = retrieve_model + from .utils import ConfigLoader + self._config_loader = ConfigLoader() + + # ── indexing ── + + def submit_document( + self, + file_path: str, + mode: str | None = None, + beta_headers: list[str] | None = None, + folder_id: str | None = None, + metadata: dict | None = None, + ) -> dict[str, Any]: + if beta_headers is not None: + raise PageIndexAPIError( + "Failed to submit document: beta_headers is not supported in local mode." + ) + if folder_id is not None: + raise PageIndexAPIError( + "Failed to submit document: folders are not supported in local mode." + ) + if metadata is not None: + if not isinstance(metadata, dict): + raise PageIndexAPIError( + "Failed to submit document: metadata must be a dict." + ) + try: + json.dumps(metadata) + except (TypeError, ValueError) as e: + raise PageIndexAPIError( + f"Failed to submit document: metadata must be valid JSON. {e}" + ) from e + if mode not in (None, "standard", "flash"): + raise PageIndexAPIError( + f"Failed to submit document: unknown local processing mode {mode!r}. " + "Supported: None or 'standard' for standard indexing, or 'flash'." + ) + file_path = os.path.abspath(os.path.expanduser(str(file_path))) + if not os.path.isfile(file_path): + raise FileNotFoundError(f"No such file: {file_path}") + if not file_path.lower().endswith(".pdf"): + raise PageIndexAPIError( + "Failed to submit document: only PDF files are supported in local mode." + ) + + try: + page_texts = self._extract_page_texts(file_path) + except PageIndexAPIError: + raise + except Exception as e: + raise PageIndexAPIError( + f"Failed to submit document: could not read PDF: {e}" + ) from e + if not any(text.strip() for text in page_texts): + raise PageIndexAPIError( + "Failed to submit document: PDF has no content. All pages are blank." + ) + + try: + if mode == "flash": + structure, description = _run_indexer( + self._index_flash, file_path, page_texts + ) + else: + structure, description = _run_indexer( + self._index_standard, file_path, page_texts + ) + except PageIndexAPIError: + raise + except Exception as e: + raise PageIndexAPIError(f"Failed to submit document: {e}") from e + + doc_id = "pi-" + uuid.uuid4().hex + meta = { + "id": doc_id, + "name": os.path.basename(file_path), + "description": description, + "status": "completed", + "createdAt": _now_iso(), + "pageNum": len(page_texts), + "folderId": None, + "metadata": metadata, + "mode": mode or "standard", + } + pages = [{"page_index": i + 1, "markdown": text} + for i, text in enumerate(page_texts)] + from .utils import remove_fields + self._store.save_document( + doc_id, meta, remove_fields(structure, fields=["text"]), pages) + return {"doc_id": doc_id} + + @staticmethod + def _extract_page_texts(file_path: str) -> list[str]: + import PyPDF2 + with open(file_path, "rb") as f: + reader = PyPDF2.PdfReader(f) + return [page.extract_text() or "" for page in reader.pages] + + def _index_standard(self, file_path: str, page_texts: list[str]) -> tuple[list, str | None]: + from .page_index_classic import page_index_main + import litellm + page_list = [(text, litellm.token_counter(model=self._model, text=text)) + for text in page_texts] + opt = self._config_loader.load({ + "model": self._model, + "summary_model": self._summary_model, + "if_add_node_id": "yes", + "if_add_node_summary": "yes", + "if_add_node_text": "yes", + "if_add_doc_description": "yes", + }) + result = page_index_main(file_path, opt, logger=logger, page_list=page_list) + structure = result.get("structure") or [] + if not structure: + raise PageIndexAPIError( + "Failed to submit document: standard indexing produced no structure." + ) + return structure, result.get("doc_description") + + def _index_flash(self, file_path: str, page_texts: list[str]) -> tuple[list, str | None]: + from .flash import page_index_flash + from .utils import (add_node_text, create_clean_structure_for_description, + generate_doc_description, write_node_id) + result = page_index_flash(file_path, summary=True, + summary_model=self._summary_model) + structure = result.get("structure", []) + if not structure: + raise PageIndexAPIError( + "Failed to submit document: PageIndex Flash could not extract " + "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, + ) + return structure, description + + # ── tree / ocr ── + + def _load_tree_with_text(self, doc_id: str, error_prefix: str) -> list: + from .utils import add_node_text + structure = self._require_data( + self._store.get_tree(doc_id), error_prefix) + pages = self._require_pages(doc_id, error_prefix) + pdf_pages = [(p.get("markdown", ""), 0) for p in pages] + add_node_text(structure, pdf_pages) + return structure + + def get_tree(self, doc_id: str, node_summary: bool = False, + include_text: bool = True) -> dict[str, Any]: + meta = self._require_doc(doc_id, "Failed to get tree result") + if include_text: + structure = self._load_tree_with_text(doc_id, "Failed to get tree result") + else: + structure = self._require_data( + self._store.get_tree(doc_id), "Failed to get tree result") + result = [_format_tree_node(node, node_summary) for node in structure] + return self._completed_envelope(doc_id, result, meta) + + def get_ocr(self, doc_id: str, format: str = "page") -> dict[str, Any]: + if format not in ["page", "node", "raw"]: + raise ValueError("Format parameter must be 'page', 'node', or 'raw'") + meta = self._require_doc(doc_id, "Failed to get OCR result") + if format == "node": + result: Any = [] + def _walk(nodes, level): + for node in nodes: + result.append({ + "title": node.get("title", ""), + "level": level, + "page_index": node.get("start_index"), + "text": node.get("text", ""), + }) + _walk(node.get("nodes") or [], level + 1) + _walk(self._load_tree_with_text(doc_id, + "Failed to get OCR result"), 1) + else: + pages = self._require_pages(doc_id, "Failed to get OCR result") + if format == "page": + result = pages + else: # raw + result = "\n\n".join(p.get("markdown", "") for p in pages) + return self._completed_envelope(doc_id, result, meta) + + @staticmethod + def _require_data(data, error_prefix: str): + if data is None: + raise PageIndexAPIError(f"{error_prefix}: stored document data is unreadable.") + return data + + def _require_pages(self, doc_id: str, error_prefix: str) -> list: + pages = self._require_data(self._store.get_pages(doc_id), error_prefix) + if not pages: + raise PageIndexAPIError( + f"{error_prefix}: stored document has no page content.") + return pages + + def _completed_envelope(self, doc_id: str, result, meta: dict) -> dict[str, Any]: + return { + "doc_id": doc_id, + "status": "completed", + "retrieval_ready": True, + "result": result, + "metadata": meta.get("metadata"), + "features": {}, + } + + # ── document management ── + + def _require_doc(self, doc_id: str, error_prefix: str) -> dict: + meta = self._store.get_meta(doc_id) + if meta is None: + raise PageIndexAPIError(f"{error_prefix}: Document not found.") + return meta + + def get_document(self, doc_id: str) -> dict[str, Any]: + meta = self._store.get_meta(doc_id) + if meta is None: + raise PageIndexAPIError("Failed to get document metadata: Document not found") + return {key: meta.get(key) for key in + ("id", "name", "description", "status", "createdAt", "pageNum", "folderId")} + + def delete_document(self, doc_id: str) -> dict[str, Any]: + if not self._store.delete_document(doc_id): + raise PageIndexAPIError("Failed to delete document: Document not found.") + return {"message": "Document deleted successfully."} + + def list_documents( + self, + limit: int = 50, + offset: int = 0, + folder_id: str | None = None, + ) -> dict[str, Any]: + if limit < 1 or limit > 100: + raise ValueError("limit must be between 1 and 100") + if offset < 0: + raise ValueError("offset must be non-negative") + if folder_id is not None: + raise PageIndexAPIError( + "Failed to list documents: folders are not supported in local mode." + ) + metas = sorted(self._store.list_metas(), key=lambda m: m.get("id") or "") + metas.sort(key=lambda m: m.get("createdAt") or "", reverse=True) + documents = [{ + "id": m.get("id"), + "name": m.get("name"), + "description": m.get("description"), + "status": m.get("status"), + "createdAt": m.get("createdAt"), + "pageNum": m.get("pageNum", 0), + "folderId": None, + "metadata": m.get("metadata"), + "features": {}, + } for m in metas[offset:offset + limit]] + return { + "documents": documents, + "total": len(metas), + "limit": limit, + "offset": offset, + } + + +def _format_tree_node(node: dict, node_summary: bool) -> dict: + children = node.get("nodes") or [] + out = { + "title": node.get("title", ""), + "node_id": node.get("node_id"), + "page_index": node.get("start_index"), + } + if node_summary: + summary = node.get("summary") + if summary is not None: + if children: + out["prefix_summary"] = summary + else: + out["summary"] = summary + if "text" in node: + out["text"] = node["text"] + if children: + out["nodes"] = [_format_tree_node(child, node_summary) for child in children] + return out diff --git a/pageindex/local_store.py b/pageindex/local_store.py new file mode 100644 index 000000000..37108f93c --- /dev/null +++ b/pageindex/local_store.py @@ -0,0 +1,164 @@ +"""On-disk document store behind PageIndexClient's local mode.""" +from __future__ import annotations + +import json +import logging +import os +import shutil +import uuid +from pathlib import Path + +logger = logging.getLogger(__name__) + + +def _write_json_atomic(path: Path, data) -> None: + tmp = path.with_name(path.name + f".{uuid.uuid4().hex}.tmp") + try: + with open(tmp, "w", encoding="utf-8") as f: + json.dump(data, f, ensure_ascii=False) + f.flush() + os.fsync(f.fileno()) + os.replace(tmp, path) + except BaseException: + tmp.unlink(missing_ok=True) + raise + + +def _read_json(path: Path): + try: + with open(path, "r", encoding="utf-8") as f: + return json.load(f) + except (FileNotFoundError, NotADirectoryError, IsADirectoryError, + PermissionError): + return None + except ValueError: + logger.warning("Unreadable JSON at %s; treating it as absent", path) + return None + + +def _is_safe_id(value: str) -> bool: + return ( + isinstance(value, str) + and value not in ("", ".", "..") + and os.path.basename(value) == value + and "\\" not in value + ) + + +def _is_valid_meta(meta, doc_id: str) -> bool: + if not isinstance(meta, dict) or meta.get("id") != doc_id: + return False + page_num = meta.get("pageNum") + return ( + isinstance(meta.get("name"), str) + and (meta.get("description") is None + or isinstance(meta.get("description"), str)) + and isinstance(meta.get("status"), str) + and isinstance(meta.get("createdAt"), str) + and isinstance(page_num, int) + and not isinstance(page_num, bool) + and page_num >= 0 + and (meta.get("folderId") is None + or isinstance(meta.get("folderId"), str)) + and (meta.get("metadata") is None + or isinstance(meta.get("metadata"), dict)) + and (meta.get("mode") is None or isinstance(meta.get("mode"), str)) + ) + + +class DocStore: + def __init__(self, storage_dir: str): + self._root = Path(storage_dir).expanduser() + self._docs = self._root / "docs" + self._manifest = self._root / "manifest.json" + + def _doc_dir(self, doc_id: str) -> Path | None: + if not _is_safe_id(doc_id): + return None + return self._docs / doc_id + + # ── manifest cache ── + def _read_manifest(self) -> dict: + data = _read_json(self._manifest) + docs = data.get("docs") if isinstance(data, dict) else None + return docs if isinstance(docs, dict) else {} + + def _write_manifest(self, docs: dict) -> None: + try: + _write_json_atomic(self._manifest, {"docs": docs}) + except OSError: + pass + + # ── documents ── + def save_document(self, doc_id: str, meta: dict, tree: list, pages: list) -> None: + doc_dir = self._doc_dir(doc_id) + if doc_dir is None: + raise ValueError(f"Invalid doc_id: {doc_id!r}") + doc_dir.mkdir(parents=True, exist_ok=True) + _write_json_atomic(doc_dir / "tree.json", tree) + _write_json_atomic(doc_dir / "pages.json", pages) + _write_json_atomic(doc_dir / "doc.json", meta) + manifest = self._read_manifest() + manifest[doc_id] = meta + self._write_manifest(manifest) + + def _read_doc_file(self, doc_id: str, name: str): + doc_dir = self._doc_dir(doc_id) + if doc_dir is None or not (doc_dir / "doc.json").is_file(): + return None + return _read_json(doc_dir / name) + + def get_meta(self, doc_id: str) -> dict | None: + doc_dir = self._doc_dir(doc_id) + if doc_dir is None or not (doc_dir / "doc.json").is_file(): + return None + meta = _read_json(doc_dir / "doc.json") + if not _is_valid_meta(meta, doc_id): + meta = self._read_manifest().get(doc_id) + return meta if _is_valid_meta(meta, doc_id) else None + + def get_tree(self, doc_id: str) -> list | None: + return self._read_doc_file(doc_id, "tree.json") + + def get_pages(self, doc_id: str) -> list | None: + return self._read_doc_file(doc_id, "pages.json") + + def list_metas(self) -> list[dict]: + if not self._docs.is_dir(): + return [] + with os.scandir(self._docs) as entries: + dir_names = {entry.name for entry in entries + if entry.is_dir() and _is_safe_id(entry.name)} + cached = self._read_manifest() + fresh = {} + for name in dir_names: + if not (self._docs / name / "doc.json").is_file(): + continue + meta = cached.get(name) + if not _is_valid_meta(meta, name): + meta = _read_json(self._docs / name / "doc.json") + if _is_valid_meta(meta, name): + fresh[name] = meta + if fresh != cached: + self._write_manifest(fresh) + return list(fresh.values()) + + def delete_document(self, doc_id: str) -> bool: + doc_dir = self._doc_dir(doc_id) + if doc_dir is None: + return False + try: + (doc_dir / "doc.json").unlink() + existed = True + except (FileNotFoundError, NotADirectoryError): + existed = False + except OSError: + if not (doc_dir / "doc.json").is_dir(): + raise + existed = False + if doc_dir.is_dir(): + shutil.rmtree(doc_dir, ignore_errors=True) + manifest = self._read_manifest() + if manifest.pop(doc_id, None) is not None: + self._write_manifest(manifest) + return existed diff --git a/pageindex/page_index.py b/pageindex/page_index_classic.py similarity index 99% rename from pageindex/page_index.py rename to pageindex/page_index_classic.py index c0b3ea935..5b846e6d0 100644 --- a/pageindex/page_index.py +++ b/pageindex/page_index_classic.py @@ -1100,19 +1100,21 @@ async def verify_toc(page_list, list_result, start_index=1, N=None, model=None): check_title_appearance(item, page_list, start_index, model) for item in indexed_sample_list ] - results = await asyncio.gather(*tasks) + results = await asyncio.gather(*tasks, return_exceptions=True) - # Process results + # Process results (skip exceptions from failed LLM calls) correct_count = 0 incorrect_results = [] for result in results: + if isinstance(result, BaseException): + continue if result['answer'] == 'yes': correct_count += 1 else: incorrect_results.append(result) - + # Calculate accuracy - checked_count = len(results) + checked_count = sum(1 for r in results if not isinstance(r, BaseException)) accuracy = correct_count / checked_count if checked_count > 0 else 0 print(f"accuracy: {accuracy*100:.2f}%") return accuracy, incorrect_results @@ -1229,18 +1231,19 @@ async def tree_parser(page_list, opt, doc=None, logger=None): return toc_tree -def page_index_main(doc, opt=None): - logger = JsonLogger(doc) - +def page_index_main(doc, opt=None, logger=None, page_list=None): + logger = logger or JsonLogger(doc) + is_valid_pdf = ( - (isinstance(doc, str) and os.path.isfile(doc) and doc.lower().endswith(".pdf")) or + (isinstance(doc, str) and os.path.isfile(doc) and doc.lower().endswith(".pdf")) or isinstance(doc, BytesIO) ) if not is_valid_pdf: raise ValueError("Unsupported input type. Expected a PDF file path or BytesIO object.") - print('Parsing PDF...') - page_list = get_page_tokens(doc, model=opt.model) + if page_list is None: + print('Parsing PDF...') + page_list = get_page_tokens(doc, model=opt.model) logger.info({'total_page_number': len(page_list)}) logger.info({'total_token': sum([page[1] for page in page_list])}) diff --git a/pageindex/retrieve.py b/pageindex/retrieve.py deleted file mode 100644 index 55c38509c..000000000 --- a/pageindex/retrieve.py +++ /dev/null @@ -1,137 +0,0 @@ -import json -import PyPDF2 - -try: - from .utils import get_number_of_pages, remove_fields -except ImportError: - from utils import get_number_of_pages, remove_fields - - -# ── Helpers ────────────────────────────────────────────────────────────────── - -def _parse_pages(pages: str) -> list[int]: - """Parse a pages string like '5-7', '3,8', or '12' into a sorted list of ints.""" - result = [] - for part in pages.split(','): - part = part.strip() - if '-' in part: - start, end = int(part.split('-', 1)[0].strip()), int(part.split('-', 1)[1].strip()) - if start > end: - raise ValueError(f"Invalid range '{part}': start must be <= end") - result.extend(range(start, end + 1)) - else: - result.append(int(part)) - return sorted(set(result)) - - -def _count_pages(doc_info: dict) -> int: - """Return total page count for a PDF document.""" - if doc_info.get('page_count'): - return doc_info['page_count'] - if doc_info.get('pages'): - return len(doc_info['pages']) - return get_number_of_pages(doc_info['path']) - - -def _get_pdf_page_content(doc_info: dict, page_nums: list[int]) -> list[dict]: - """Extract text for specific PDF pages (1-indexed). Prefer cached pages, fallback to PDF.""" - cached_pages = doc_info.get('pages') - if cached_pages: - page_map = {p['page']: p['content'] for p in cached_pages} - return [ - {'page': p, 'content': page_map[p]} - for p in page_nums if p in page_map - ] - path = doc_info['path'] - with open(path, 'rb') as f: - pdf_reader = PyPDF2.PdfReader(f) - total = len(pdf_reader.pages) - valid_pages = [p for p in page_nums if 1 <= p <= total] - return [ - {'page': p, 'content': pdf_reader.pages[p - 1].extract_text() or ''} - for p in valid_pages - ] - - -def _get_md_page_content(doc_info: dict, page_nums: list[int]) -> list[dict]: - """ - For Markdown documents, 'pages' are line numbers. - Find nodes whose line_num falls within [min(page_nums), max(page_nums)] and return their text. - """ - min_line, max_line = min(page_nums), max(page_nums) - results = [] - seen = set() - - def _traverse(nodes): - for node in nodes: - ln = node.get('line_num') - if ln and min_line <= ln <= max_line and ln not in seen: - seen.add(ln) - results.append({'page': ln, 'content': node.get('text', '')}) - if node.get('nodes'): - _traverse(node['nodes']) - - _traverse(doc_info.get('structure', [])) - results.sort(key=lambda x: x['page']) - return results - - -# ── Tool functions ──────────────────────────────────────────────────────────── - -def get_document(documents: dict, doc_id: str) -> str: - """Return JSON with document metadata: doc_id, doc_name, doc_description, type, status, page_count (PDF) or line_count (Markdown).""" - doc_info = documents.get(doc_id) - if not doc_info: - return json.dumps({'error': f'Document {doc_id} not found'}) - result = { - 'doc_id': doc_id, - 'doc_name': doc_info.get('doc_name', ''), - 'doc_description': doc_info.get('doc_description', ''), - 'type': doc_info.get('type', ''), - 'status': 'completed', - } - if doc_info.get('type') == 'pdf': - result['page_count'] = _count_pages(doc_info) - else: - result['line_count'] = doc_info.get('line_count', 0) - return json.dumps(result) - - -def get_document_structure(documents: dict, doc_id: str) -> str: - """Return tree structure JSON with text fields removed (saves tokens).""" - doc_info = documents.get(doc_id) - if not doc_info: - return json.dumps({'error': f'Document {doc_id} not found'}) - structure = doc_info.get('structure', []) - structure_no_text = remove_fields(structure, fields=['text']) - return json.dumps(structure_no_text, ensure_ascii=False) - - -def get_page_content(documents: dict, doc_id: str, pages: str) -> str: - """ - Retrieve page content for a document. - - pages format: '5-7', '3,8', or '12' - For PDF: pages are physical page numbers (1-indexed). - For Markdown: pages are line numbers corresponding to node headers. - - Returns JSON list of {'page': int, 'content': str}. - """ - doc_info = documents.get(doc_id) - if not doc_info: - return json.dumps({'error': f'Document {doc_id} not found'}) - - try: - page_nums = _parse_pages(pages) - except (ValueError, AttributeError) as e: - return json.dumps({'error': f'Invalid pages format: {pages!r}. Use "5-7", "3,8", or "12". Error: {e}'}) - - try: - if doc_info.get('type') == 'pdf': - content = _get_pdf_page_content(doc_info, page_nums) - else: - content = _get_md_page_content(doc_info, page_nums) - except Exception as e: - return json.dumps({'error': f'Failed to read page content: {e}'}) - - return json.dumps(content, ensure_ascii=False) diff --git a/pageindex/utils.py b/pageindex/utils.py index 92fc46d85..97f60a942 100644 --- a/pageindex/utils.py +++ b/pageindex/utils.py @@ -7,7 +7,6 @@ import PyPDF2 import copy import asyncio -import pymupdf from io import BytesIO from dotenv import load_dotenv load_dotenv() @@ -31,6 +30,12 @@ def count_tokens(text, model=None): return litellm.token_counter(model=model, text=text) +def _strip_prefix(s, prefix): + if s.startswith(prefix): + return s[len(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').""" @@ -57,18 +62,19 @@ def _is_unrecoverable(exc: Exception) -> bool: def llm_completion(model, prompt, chat_history=None, return_finish_reason=False): use_openai_sdk = _is_openai_model(model) if model: - model = model.removeprefix("litellm/") + model = _strip_prefix(model, "litellm/") if use_openai_sdk: - model = model.removeprefix("openai/") + model = _strip_prefix(model, "openai/") max_retries = 10 messages = list(chat_history) + [{"role": "user", "content": prompt}] if chat_history else [{"role": "user", "content": prompt}] + if use_openai_sdk: + global _openai_sync_client + if _openai_sync_client is None: + import openai + _openai_sync_client = openai.OpenAI(max_retries=0) for i in range(max_retries): try: if use_openai_sdk: - global _openai_sync_client - if _openai_sync_client is None: - import openai - _openai_sync_client = openai.OpenAI(max_retries=0) response = _openai_sync_client.chat.completions.create( model=model, messages=messages, @@ -94,27 +100,27 @@ def llm_completion(model, prompt, chat_history=None, return_finish_reason=False) if i < max_retries - 1: time.sleep(1) else: - logging.error('Max retries reached for prompt: ' + prompt) - if return_finish_reason: - return "", "error" - return "" + raise RuntimeError( + f"LLM completion failed after {max_retries} retries" + ) from e async def llm_acompletion(model, prompt): use_openai_sdk = _is_openai_model(model) if model: - model = model.removeprefix("litellm/") + model = _strip_prefix(model, "litellm/") if use_openai_sdk: - model = model.removeprefix("openai/") + model = _strip_prefix(model, "openai/") max_retries = 10 messages = [{"role": "user", "content": prompt}] + if use_openai_sdk: + global _openai_async_client + if _openai_async_client is None: + import openai + _openai_async_client = openai.AsyncOpenAI(max_retries=0) for i in range(max_retries): try: if use_openai_sdk: - global _openai_async_client - if _openai_async_client is None: - import openai - _openai_async_client = openai.AsyncOpenAI(max_retries=0) response = await _openai_async_client.chat.completions.create( model=model, messages=messages, @@ -136,10 +142,11 @@ async def llm_acompletion(model, prompt): if i < max_retries - 1: await asyncio.sleep(1) else: - logging.error('Max retries reached for prompt: ' + prompt) - return "" - - + raise RuntimeError( + f"LLM completion failed after {max_retries} retries" + ) from e + + def get_json_content(response): start_idx = response.find("```json") if start_idx != -1: @@ -180,7 +187,7 @@ def extract_json(content): # Remove any trailing commas before closing brackets/braces json_content = json_content.replace(',]', ']').replace(',}', '}') return json.loads(json_content) - except: + except Exception: logging.error("Failed to parse JSON even after cleanup") return {} except Exception as e: @@ -454,6 +461,7 @@ def get_page_tokens(pdf_path, model=None, pdf_parser="PyPDF2"): page_list.append((page_text, token_length)) return page_list elif pdf_parser == "PyMuPDF": + import pymupdf if isinstance(pdf_path, BytesIO): pdf_stream = pdf_path doc = pymupdf.open(stream=pdf_stream, filetype="pdf") @@ -471,12 +479,16 @@ def get_page_tokens(pdf_path, model=None, pdf_parser="PyPDF2"): def get_text_of_pdf_pages(pdf_pages, start_page, end_page): + if start_page is None or end_page is None: + return "" text = "" for page_num in range(start_page-1, end_page): text += pdf_pages[page_num][0] return text def get_text_of_pdf_pages_with_labels(pdf_pages, start_page, end_page): + if start_page is None or end_page is None: + return "" text = "" for page_num in range(start_page-1, end_page): text += f"\n{pdf_pages[page_num][0]}\n\n" @@ -522,12 +534,14 @@ def clean_structure_post(data): clean_structure_post(section) return data -def remove_fields(data, fields=['text']): +def remove_fields(data, fields=['text'], max_len=None): if isinstance(data, dict): - return {k: remove_fields(v, fields) + return {k: remove_fields(v, fields, max_len) for k, v in data.items() if k not in fields} elif isinstance(data, list): - return [remove_fields(item, fields) for item in data] + return [remove_fields(item, fields, max_len) for item in data] + elif isinstance(data, str): + return data[:max_len] + '...' if max_len is not None and len(data) > max_len else data return data def print_toc(tree, indent=0): @@ -648,10 +662,15 @@ async def generate_node_summary(node, model=None): async def generate_summaries_for_structure(structure, model=None): nodes = structure_to_list(structure) tasks = [generate_node_summary(node, model=model) for node in nodes] - summaries = await asyncio.gather(*tasks) + summaries = await asyncio.gather(*tasks, return_exceptions=True) for node, summary in zip(nodes, summaries): - node['summary'] = summary + node['summary'] = "" if isinstance(summary, BaseException) else summary + if nodes and not any(node['summary'] for node in nodes): + raise RuntimeError( + "Summary generation failed for all nodes " + "(check LLM credentials and model availability)" + ) return structure @@ -819,12 +838,27 @@ 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)) + await asyncio.gather(*(visit(child) for child in children), + return_exceptions=True) if node.get('summary'): return - node['summary'] = await (parent_summary(node) if children else leaf_summary(node)) + try: + node['summary'] = await (parent_summary(node) if children else leaf_summary(node)) + except Exception: + node['summary'] = "" + + await asyncio.gather(*(visit(root) for root in structure), + return_exceptions=True) + + def _any_summary(nodes): + return any(n.get('summary') or _any_summary(n.get('nodes') or []) + for n in nodes) + if not _any_summary(structure): + raise RuntimeError( + "Summary generation failed for all nodes " + "(check LLM credentials and model availability)" + ) - await asyncio.gather(*(visit(root) for root in structure)) strip_internal_keys(structure) return structure @@ -860,8 +894,10 @@ def generate_doc_description(structure, model=None): Directly return the description, do not include any other text. """ - response = llm_completion(model, prompt) - return response + try: + return llm_completion(model, prompt) + except RuntimeError: + return "" def reorder_dict(data, key_order): @@ -951,25 +987,42 @@ def load(self, user_opt=None) -> config: merged = {**self._default_dict, **user_dict} return config(**merged) -def create_node_mapping(tree): - """Create a flat dict mapping node_id to node for quick lookup.""" +def create_node_mapping(tree, include_page_ranges=False, max_page=None): + """Map node_id to node; with include_page_ranges, to {"node", "start_index", + "end_index"} (end = next node's page_index, or max_page for the last node).""" + def get_all_nodes(tree): + if isinstance(tree, dict): + return [tree] + [node for child in tree.get('nodes', []) for node in get_all_nodes(child)] + elif isinstance(tree, list): + return [node for item in tree for node in get_all_nodes(item)] + return [] + + all_nodes = get_all_nodes(tree) + if not include_page_ranges: + return {node["node_id"]: node for node in all_nodes if node.get("node_id")} mapping = {} - def _traverse(nodes): - for node in nodes: - if node.get('node_id'): - mapping[node['node_id']] = node - if node.get('nodes'): - _traverse(node['nodes']) - _traverse(tree) + for i, node in enumerate(all_nodes): + if node.get("node_id"): + end_page = all_nodes[i + 1].get("page_index") if i + 1 < len(all_nodes) else max_page + mapping[node["node_id"]] = { + "node": node, + "start_index": node["page_index"], + "end_index": end_page, + } return mapping -def print_tree(tree, indent=0): +def print_tree(tree, exclude_fields=None, indent=0): + """Outline view; passing exclude_fields gives the 0.2.8 pprint view.""" + if exclude_fields is not None: + from pprint import pprint + pprint(remove_fields(tree, exclude_fields, max_len=40), sort_dicts=False, width=100) + return for node in tree: summary = node.get('summary') or node.get('prefix_summary', '') summary_str = f" — {summary[:60]}..." if summary else "" print(' ' * indent + f"[{node.get('node_id', '?')}] {node.get('title', '')}{summary_str}") if node.get('nodes'): - print_tree(node['nodes'], indent + 1) + print_tree(node['nodes'], indent=indent + 1) def print_wrapped(text, width=100): for line in text.splitlines(): diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 000000000..deac66be3 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,53 @@ +[tool.poetry] +name = "pageindex" +version = "0.2.9" +description = "Python SDK for PageIndex — reasoning-based, vectorless document retrieval, cloud and local" +readme = "README.md" +license = "MIT" +authors = ["Ray "] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Developers", + "License :: OSI Approved :: MIT License", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.7", + "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Topic :: Scientific/Engineering :: Artificial Intelligence", +] +keywords = ["rag", "document", "retrieval", "llm", "pageindex", "vectorless"] +packages = [{ include = "pageindex" }] +include = [ + { path = "pageindex/config.yaml", format = ["sdist", "wheel"] }, + { path = "pageindex/flash/data/*.json", format = ["sdist", "wheel"] }, +] +exclude = ["pageindex/flash/assets"] + +[tool.poetry.dependencies] +python = ">=3.7" +requests = ">=2.28.0" +openai = ">=1.70.0" +litellm = ">=1.84.0" +PyPDF2 = ">=3.0.0" +pypdfium2 = ">=4.30.0" +sortedcontainers = ">=2.4.0" +regex = ">=2024.0.0" +python-dotenv = ">=1.0.0" +pyyaml = ">=6.0" + +[tool.poetry.group.dev.dependencies] +pytest = ">=7.0" + +[tool.poetry.urls] +Repository = "https://github.com/VectifyAI/PageIndex" +Homepage = "https://pageindex.ai" +Documentation = "https://docs.pageindex.ai" +Issues = "https://github.com/VectifyAI/PageIndex/issues" + +[build-system] +requires = ["poetry-core"] +build-backend = "poetry.core.masonry.api" diff --git a/requirements.txt b/requirements.txt index c0b76deb1..5fd4f2e4e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,8 @@ litellm==1.84.0 -# openai-agents # optional: required for examples/agentic_vectorless_rag_demo.py -pymupdf==1.26.4 +openai>=1.70.0 +requests>=2.28.0 +# openai-agents # optional +# pymupdf # optional PyPDF2==3.0.1 pypdfium2==4.30.0 python-dotenv==1.2.2 diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 000000000..40a4a9199 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,52 @@ +import pytest + + +@pytest.fixture(autouse=True) +def _llm_key(monkeypatch): + """Deterministic key presence for every test; missing-key tests delenv.""" + monkeypatch.setenv("OPENAI_API_KEY", "test-key") + + +def build_pdf(page_texts): + """Build a minimal, uncompressed PDF (one Helvetica line per page) whose + text PyPDF2 can extract. Returns the PDF file bytes.""" + n = len(page_texts) + objects = [] + kids = " ".join(f"{3 + i} 0 R" for i in range(n)) + objects.append(b"<< /Type /Catalog /Pages 2 0 R >>") + objects.append(f"<< /Type /Pages /Kids [{kids}] /Count {n} >>".encode()) + font_obj = 3 + 2 * n + for i in range(n): + objects.append( + f"<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] " + f"/Resources << /Font << /F1 {font_obj} 0 R >> >> " + f"/Contents {3 + n + i} 0 R >>".encode() + ) + for text in page_texts: + safe = text.replace("\\", r"\\").replace("(", r"\(").replace(")", r"\)") + stream = f"BT /F1 12 Tf 72 720 Td ({safe}) Tj ET".encode() + objects.append(b"<< /Length %d >>\nstream\n%s\nendstream" % (len(stream), stream)) + objects.append(b"<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>") + + out = bytearray(b"%PDF-1.4\n") + offsets = [] + for num, body in enumerate(objects, start=1): + offsets.append(len(out)) + out += b"%d 0 obj\n" % num + body + b"\nendobj\n" + xref_pos = len(out) + out += b"xref\n0 %d\n" % (len(objects) + 1) + out += b"0000000000 65535 f \n" + for offset in offsets: + out += b"%010d 00000 n \n" % offset + out += (b"trailer\n<< /Size %d /Root 1 0 R >>\nstartxref\n%d\n%%%%EOF\n" + % (len(objects) + 1, xref_pos)) + return bytes(out) + + +@pytest.fixture +def sample_pdf(tmp_path): + """A 2-page PDF with known, extractable text.""" + path = tmp_path / "sample.pdf" + path.write_bytes(build_pdf(["Hello page one about apples", + "Second page about bananas"])) + return str(path) diff --git a/tests/test_client.py b/tests/test_client.py new file mode 100644 index 000000000..50b7f5178 --- /dev/null +++ b/tests/test_client.py @@ -0,0 +1,712 @@ +"""SDK surface tests: PageIndexClient in local and cloud mode.""" +import asyncio +import importlib +import json +import re +import shutil +import types + +import pytest + +import pageindex.flash +import pageindex.utils +from pageindex import PageIndexClient, PageIndexAPIError + +page_index_module = importlib.import_module("pageindex.page_index_classic") + + +STRUCTURE = [ + { + "title": "Root Section", "node_id": "0000", + "start_index": 1, "end_index": 2, + "summary": "root summary", "text": "root text", + "nodes": [ + {"title": "Child Section", "node_id": "0001", + "start_index": 2, "end_index": 2, + "summary": "child summary", "text": "child text"}, + ], + }, +] + + +@pytest.fixture +def local_client(tmp_path): + return PageIndexClient(storage_path=str(tmp_path / "store")) + + +@pytest.fixture +def indexed_doc(local_client, sample_pdf, monkeypatch): + """A document indexed through a stubbed standard pipeline.""" + def fake_page_index_main(doc, opt=None, logger=None, page_list=None): + assert opt.if_add_node_summary == "yes" + assert opt.if_add_node_text == "yes" + assert logger is not None + assert page_list is not None + assert all(isinstance(t, tuple) and len(t) == 2 for t in page_list) + return {"doc_name": "sample.pdf", + "doc_description": "A test document.", + "structure": json.loads(json.dumps(STRUCTURE))} + monkeypatch.setattr(page_index_module, "page_index_main", fake_page_index_main) + return local_client.submit_document(sample_pdf)["doc_id"] + + +# ── constructor ── + +def test_empty_api_key_raises(): + with pytest.raises(PageIndexAPIError, match="empty string"): + PageIndexClient(api_key="") + + +def test_cloud_rejects_local_args(): + with pytest.raises(PageIndexAPIError, match="model, storage_path"): + PageIndexClient(api_key="k", model="m", storage_path="/tmp/x") + + +def test_local_client_does_not_touch_disk(tmp_path): + storage = tmp_path / "store" + PageIndexClient(storage_path=str(storage)) + assert not storage.exists() + + +def test_retrieve_model_carries_agents_sdk_prefix(tmp_path): + 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 + + +def test_explicit_mode_clients(tmp_path): + from pageindex import PageIndexCloudClient, PageIndexLocalClient + + for bad_key in (None, ""): + with pytest.raises(PageIndexAPIError, match="requires a PageIndex API key"): + PageIndexCloudClient(bad_key) + cloud = PageIndexCloudClient("k") + assert cloud.api_key == "k" and isinstance(cloud, PageIndexClient) + + local = PageIndexLocalClient(model="m", storage_path=str(tmp_path / "s")) + assert local.model == "m" and isinstance(local, PageIndexClient) + with pytest.raises(TypeError): + PageIndexLocalClient("k") + + +# ── local: indexing and reading ── + +def test_submit_and_get_tree(local_client, indexed_doc, tmp_path, monkeypatch): + tree = local_client.get_tree(indexed_doc, node_summary=True) + assert tree["status"] == "completed" + assert tree["retrieval_ready"] is True + root = tree["result"][0] + assert root["page_index"] == 1 + assert "start_index" not in root and "end_index" not in root + assert root["prefix_summary"] == "root summary" + assert "summary" not in root + child = root["nodes"][0] + assert child["summary"] == "child summary" + assert child["text"] == "Second page about bananas" + + no_summary = local_client.get_tree(indexed_doc)["result"][0] + assert "summary" not in no_summary and "prefix_summary" not in no_summary + + +def test_get_tree_include_text_false(local_client, indexed_doc): + tree = local_client.get_tree(indexed_doc, include_text=False) + root = tree["result"][0] + assert "text" not in root + assert "text" not in root["nodes"][0] + assert root["page_index"] == 1 + + with_text = local_client.get_tree(indexed_doc)["result"][0] + assert "text" in with_text + + +def test_get_document_structure(local_client, indexed_doc): + result = local_client.get_document_structure(indexed_doc) + assert isinstance(result, list) + root = result[0] + assert "text" not in root + assert "text" not in root["nodes"][0] + assert "prefix_summary" in root + assert root["nodes"][0]["summary"] == "child summary" + + +def test_get_page_content(local_client, indexed_doc): + pages = local_client.get_page_content(indexed_doc, "1") + assert len(pages) == 1 + assert pages[0]["page_index"] == 1 + assert "Hello page one" in pages[0]["markdown"] + + pages = local_client.get_page_content(indexed_doc, "1-2") + assert len(pages) == 2 + + pages = local_client.get_page_content(indexed_doc, "2,1") + assert [p["page_index"] for p in pages] == [1, 2] + + assert local_client.get_page_content(indexed_doc, "99") == [] + + with pytest.raises(ValueError): + local_client.get_page_content(indexed_doc, "abc") + + +def test_submit_does_not_create_cwd_logs(local_client, sample_pdf, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + def fake_page_index_main(doc, opt=None, logger=None, page_list=None): + logger.info({"probe": True}) + return {"doc_name": "sample.pdf", "doc_description": None, + "structure": json.loads(json.dumps(STRUCTURE))} + monkeypatch.setattr(page_index_module, "page_index_main", fake_page_index_main) + local_client.submit_document(sample_pdf) + assert not (tmp_path / "logs").exists() + + +def test_submit_flash(local_client, sample_pdf, monkeypatch): + calls = {} + def fake_flash(pdf, summary=True, summary_model=None, **kwargs): + calls["summary"] = summary + calls["summary_model"] = summary_model + return {"doc_name": "sample.pdf", + "structure": [{"title": "Flash Root", "start_index": 1, + "end_index": 2, "summary": "s", "nodes": []}]} + monkeypatch.setattr(pageindex.flash, "page_index_flash", fake_flash) + monkeypatch.setattr(pageindex.utils, "llm_completion", + lambda model, prompt, **kw: "Flash description.") + doc_id = local_client.submit_document(sample_pdf, mode="flash")["doc_id"] + assert calls == {"summary": True, "summary_model": local_client.summary_model} + root = local_client.get_tree(doc_id)["result"][0] + assert root["node_id"] == "0000" + assert "Hello page one" in root["text"] + assert local_client.get_document(doc_id)["description"] == "Flash description." + + +def test_llm_completion_missing_key_raises_immediately(monkeypatch): + import openai + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + monkeypatch.setattr(pageindex.utils, "_openai_sync_client", None) + monkeypatch.setattr(pageindex.utils, "_openai_async_client", None) + with pytest.raises(openai.OpenAIError): + pageindex.utils.llm_completion("gpt-4o", "probe") + with pytest.raises(openai.OpenAIError): + asyncio.run(pageindex.utils.llm_acompletion("gpt-4o", "probe")) + + +def test_submit_missing_llm_key_fails_loud(local_client, sample_pdf, monkeypatch): + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + monkeypatch.setattr(pageindex.utils, "_openai_sync_client", None) + 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) + monkeypatch.setattr(pageindex.flash, "page_index_flash", first_llm_call) + for kwargs in ({}, {"mode": "flash"}): + with pytest.raises(PageIndexAPIError, match="OPENAI_API_KEY"): + local_client.submit_document(sample_pdf, **kwargs) + assert local_client.list_documents()["total"] == 0 + + +def test_submit_rejections(local_client, sample_pdf, tmp_path): + with pytest.raises(FileNotFoundError): + local_client.submit_document(str(tmp_path / "missing.pdf")) + (tmp_path / "notes.txt").write_text("hi") + with pytest.raises(PageIndexAPIError, match="only PDF"): + local_client.submit_document(str(tmp_path / "notes.txt")) + with pytest.raises(PageIndexAPIError, match="unknown local processing mode"): + local_client.submit_document(sample_pdf, mode="mcp") + with pytest.raises(PageIndexAPIError, match="folders"): + local_client.submit_document(sample_pdf, folder_id="f1") + with pytest.raises(PageIndexAPIError, match="beta_headers"): + local_client.submit_document(sample_pdf, beta_headers=["block_reference"]) + + +def test_corrupt_pdf_raises_api_error(local_client, tmp_path): + bad = tmp_path / "bad.pdf" + bad.write_bytes(b"%PDF-1.4 garbage with no xref or trailer") + with pytest.raises(PageIndexAPIError, match="could not read PDF"): + local_client.submit_document(str(bad)) + + +def test_encrypted_pdf_raises_api_error(local_client, sample_pdf, tmp_path): + from PyPDF2 import PdfReader, PdfWriter + + writer = PdfWriter() + for page in PdfReader(sample_pdf).pages: + writer.add_page(page) + writer.encrypt("secret") + enc = tmp_path / "enc.pdf" + with open(enc, "wb") as f: + writer.write(f) + with pytest.raises(PageIndexAPIError, match="could not read PDF"): + local_client.submit_document(str(enc)) + + +def test_submit_explicit_standard_mode(local_client, sample_pdf, monkeypatch): + calls = [] + + def fake_page_index_main(doc, opt=None, logger=None, page_list=None): + calls.append(doc) + return { + "doc_name": "sample.pdf", + "doc_description": "A test document.", + "structure": json.loads(json.dumps(STRUCTURE)), + } + + monkeypatch.setattr(page_index_module, "page_index_main", fake_page_index_main) + doc_id = local_client.submit_document(sample_pdf, mode="standard")["doc_id"] + + assert calls == [sample_pdf] + assert local_client._api._store.get_meta(doc_id)["mode"] == "standard" + + +@pytest.mark.parametrize("mode", ["standard", "flash"]) +def test_submit_from_running_event_loop( + local_client, sample_pdf, monkeypatch, mode +): + def fake_index(*args): + asyncio.run(asyncio.sleep(0)) + return json.loads(json.dumps(STRUCTURE)), "A test document." + + monkeypatch.setattr(local_client._api, f"_index_{mode}", fake_index) + + async def submit(): + return local_client.submit_document(sample_pdf, mode=mode) + + doc_id = asyncio.run(submit())["doc_id"] + assert local_client.get_document(doc_id)["status"] == "completed" + + +def test_submit_with_metadata(local_client, sample_pdf, monkeypatch): + monkeypatch.setattr( + page_index_module, "page_index_main", + lambda doc, opt=None, logger=None, page_list=None: { + "doc_name": "sample.pdf", "doc_description": None, + "structure": json.loads(json.dumps(STRUCTURE))}) + tags = {"project": "alpha", "year": 2026} + doc_id = local_client.submit_document(sample_pdf, metadata=tags)["doc_id"] + assert local_client.get_tree(doc_id)["metadata"] == tags + assert local_client.get_ocr(doc_id)["metadata"] == tags + assert local_client.list_documents()["documents"][0]["metadata"] == tags + assert "metadata" not in local_client.get_document(doc_id) + + +def test_submit_metadata_validation(local_client, sample_pdf, monkeypatch): + indexed = [] + monkeypatch.setattr(page_index_module, "page_index_main", + lambda *args, **kwargs: indexed.append(1)) + with pytest.raises(PageIndexAPIError, match="metadata must be a dict"): + local_client.submit_document(sample_pdf, metadata=["not", "a", "dict"]) + with pytest.raises(PageIndexAPIError, match="valid JSON"): + local_client.submit_document(sample_pdf, metadata={"x": object()}) + assert indexed == [] + + +def test_blank_pdf_rejected(local_client, tmp_path): + from conftest import build_pdf + blank = tmp_path / "blank.pdf" + blank.write_bytes(build_pdf(["", ""])) + with pytest.raises(PageIndexAPIError, match="All pages are blank"): + local_client.submit_document(str(blank)) + + +def test_get_ocr(local_client, indexed_doc): + page = local_client.get_ocr(indexed_doc) + assert page["result"][0] == {"page_index": 1, + "markdown": "Hello page one about apples"} + raw = local_client.get_ocr(indexed_doc, format="raw") + assert raw["result"] == ("Hello page one about apples\n\n" + "Second page about bananas") + node = local_client.get_ocr(indexed_doc, format="node") + assert node["result"] == [ + {"title": "Root Section", "level": 1, "page_index": 1, + "text": "Hello page one about applesSecond page about bananas"}, + {"title": "Child Section", "level": 2, "page_index": 2, + "text": "Second page about bananas"}, + ] + with pytest.raises(ValueError): + local_client.get_ocr(indexed_doc, format="bogus") + + +def test_document_management(local_client, indexed_doc): + assert indexed_doc.startswith("pi-") + doc = local_client.get_document(indexed_doc) + assert doc["id"] == indexed_doc + assert re.fullmatch(r"\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(\.\d{3}000)?", + doc["createdAt"]) + assert doc["name"] == "sample.pdf" + assert doc["description"] == "A test document." + assert doc["status"] == "completed" + assert doc["pageNum"] == 2 + assert doc["folderId"] is None + + listing = local_client.list_documents() + assert listing["total"] == 1 + assert listing["limit"] == 50 and listing["offset"] == 0 + assert listing["documents"][0]["id"] == indexed_doc + + assert local_client.is_retrieval_ready(indexed_doc) is True + + assert local_client.delete_document(indexed_doc) == { + "message": "Document deleted successfully."} + with pytest.raises(PageIndexAPIError, match="Document not found"): + local_client.delete_document(indexed_doc) + assert local_client.is_retrieval_ready(indexed_doc) is False + + +def test_manifest_write_through_and_self_heal(local_client, indexed_doc, tmp_path): + manifest_path = tmp_path / "store" / "manifest.json" + manifest = json.loads(manifest_path.read_text()) + assert manifest["docs"][indexed_doc]["name"] == "sample.pdf" + + # corrupt cache → listings rebuild it from the doc.json files + manifest_path.write_text("{broken") + assert local_client.list_documents()["total"] == 1 + assert indexed_doc in json.loads(manifest_path.read_text())["docs"] + + # missing cache → same + manifest_path.unlink() + assert local_client.list_documents()["total"] == 1 + + # doc dir removed behind the store's back → healed, not served stale + shutil.rmtree(tmp_path / "store" / "docs" / indexed_doc) + assert local_client.list_documents()["total"] == 0 + + local_client_meta = json.loads(manifest_path.read_text()) + assert local_client_meta == {"docs": {}} + + +@pytest.mark.parametrize("bad_entry", ["corrupt-entry", {"id": "wrong"}]) +def test_manifest_invalid_entry_self_heals( + local_client, indexed_doc, tmp_path, bad_entry +): + manifest_path = tmp_path / "store" / "manifest.json" + manifest_path.write_text(json.dumps({"docs": {indexed_doc: bad_entry}})) + + listing = local_client.list_documents() + + assert listing["total"] == 1 + assert listing["documents"][0]["id"] == indexed_doc + healed = json.loads(manifest_path.read_text()) + assert healed["docs"][indexed_doc]["name"] == "sample.pdf" + + +def test_manifest_updated_on_delete(local_client, indexed_doc, tmp_path): + local_client.delete_document(indexed_doc) + manifest = json.loads((tmp_path / "store" / "manifest.json").read_text()) + assert manifest == {"docs": {}} + + +def test_manifest_picks_up_external_doc(local_client, indexed_doc, tmp_path): + # a doc whose manifest update was lost (e.g. concurrent writer) still lists + docs_dir = tmp_path / "store" / "docs" + external_id = "11111111-1111-4111-8111-111111111111" + shutil.copytree(docs_dir / indexed_doc, docs_dir / external_id) + meta_path = docs_dir / external_id / "doc.json" + meta = json.loads(meta_path.read_text()) + meta["id"] = external_id + meta_path.write_text(json.dumps(meta)) + + ids = {d["id"] for d in local_client.list_documents()["documents"]} + assert ids == {indexed_doc, external_id} + + +def test_manifest_ignores_incomplete_dir(local_client, indexed_doc, tmp_path): + (tmp_path / "store" / "docs" / "crashed-save").mkdir() + listing = local_client.list_documents() + assert listing["total"] == 1 + assert listing["documents"][0]["id"] == indexed_doc + + +def test_torn_delete_never_lists_ghost(local_client, indexed_doc, tmp_path): + # crash mid-delete: doc.json gone, dir and manifest entry remain + doc_dir = tmp_path / "store" / "docs" / indexed_doc + (doc_dir / "doc.json").unlink() + + assert local_client.list_documents()["total"] == 0 + manifest = json.loads((tmp_path / "store" / "manifest.json").read_text()) + assert manifest == {"docs": {}} + with pytest.raises(PageIndexAPIError): + local_client.get_document(indexed_doc) + with pytest.raises(PageIndexAPIError, match="Document not found"): + local_client.delete_document(indexed_doc) + assert not doc_dir.exists() + + +def test_corrupt_doc_json_is_contained(local_client, indexed_doc, sample_pdf, tmp_path): + second = local_client.submit_document(sample_pdf)["doc_id"] + (tmp_path / "store" / "docs" / indexed_doc / "doc.json").write_text("{truncated") + + # manifest still holds a good copy of the meta — served consistently + assert local_client.get_document(indexed_doc)["id"] == indexed_doc + assert local_client.list_documents()["total"] == 2 + + # without the manifest copy, the doc is treated as absent, not a crash + (tmp_path / "store" / "manifest.json").unlink() + listing = local_client.list_documents() + assert listing["total"] == 1 + assert listing["documents"][0]["id"] == second + with pytest.raises(PageIndexAPIError): + local_client.get_document(indexed_doc) + assert local_client.is_retrieval_ready(indexed_doc) is False + + +def test_invalid_utf8_is_contained(local_client, indexed_doc, tmp_path): + doc_json = tmp_path / "store" / "docs" / indexed_doc / "doc.json" + doc_json.write_bytes(b'{"id": "\xff\xfe broken') + + # the manifest copy keeps serving, consistently across list and get + assert local_client.get_document(indexed_doc)["id"] == indexed_doc + assert local_client.list_documents()["total"] == 1 + + # even with the manifest corrupted the same way: no crash, self-heals + (tmp_path / "store" / "manifest.json").write_bytes(b"\xff\xfe") + assert local_client.list_documents()["total"] == 0 + with pytest.raises(PageIndexAPIError): + local_client.get_document(indexed_doc) + + +def test_corrupt_data_files_fail_loud(local_client, indexed_doc, tmp_path): + doc_dir = tmp_path / "store" / "docs" / indexed_doc + (doc_dir / "tree.json").write_bytes(b"\xff\xfe") + with pytest.raises(PageIndexAPIError, match="unreadable"): + local_client.get_tree(indexed_doc) + assert local_client.is_retrieval_ready(indexed_doc) is False + + (doc_dir / "pages.json").write_text("{broken") + with pytest.raises(PageIndexAPIError, match="unreadable"): + local_client.get_ocr(indexed_doc) + + # the metadata itself is intact, so listings stay honest + assert local_client.list_documents()["total"] == 1 + + +def test_get_tree_fails_loud_on_broken_pages(local_client, indexed_doc, tmp_path): + doc_dir = tmp_path / "store" / "docs" / indexed_doc + (doc_dir / "pages.json").write_text("{broken") + with pytest.raises(PageIndexAPIError, match="unreadable"): + local_client.get_tree(indexed_doc) + + +def test_get_tree_fails_loud_on_empty_pages(local_client, indexed_doc, tmp_path): + doc_dir = tmp_path / "store" / "docs" / indexed_doc + (doc_dir / "pages.json").write_text("[]") + with pytest.raises(PageIndexAPIError, match="no page content"): + local_client.get_tree(indexed_doc) + + +def test_data_file_as_directory_fails_loud(local_client, indexed_doc, tmp_path): + tree_path = tmp_path / "store" / "docs" / indexed_doc / "tree.json" + tree_path.unlink() + tree_path.mkdir() + with pytest.raises(PageIndexAPIError, match="unreadable"): + local_client.get_tree(indexed_doc) + + +def test_list_documents_skips_unsafe_directory_names( + local_client, indexed_doc, tmp_path +): + bad_dir = tmp_path / "store" / "docs" / "bad\\name" + bad_dir.mkdir() + (bad_dir / "doc.json").write_text("{}") + listing = local_client.list_documents() + assert [d["id"] for d in listing["documents"]] == [indexed_doc] + + +def test_generate_doc_description_error_boundary(monkeypatch): + def raiser(exc): + def _f(*args, **kwargs): + raise exc + return _f + monkeypatch.setattr(pageindex.utils, "llm_completion", + raiser(RuntimeError("retries exhausted"))) + assert pageindex.utils.generate_doc_description([]) == "" + monkeypatch.setattr(pageindex.utils, "llm_completion", + raiser(ValueError("provider rejected the model"))) + with pytest.raises(ValueError): + pageindex.utils.generate_doc_description([]) + + +def test_generate_summaries_all_failed_raises(monkeypatch): + async def boom(model, prompt): + raise ValueError("bad key") + monkeypatch.setattr(pageindex.utils, "llm_acompletion", boom) + structure = [{"title": "A", "text": "t1", + "nodes": [{"title": "B", "text": "t2"}]}] + with pytest.raises(RuntimeError, match="all nodes"): + asyncio.run(pageindex.utils.generate_summaries_for_structure(structure)) + + +def test_generate_summaries_partial_failure_absorbed(monkeypatch): + async def flaky(model, prompt): + if "t1" in prompt: + raise ValueError("transient") + return "ok" + monkeypatch.setattr(pageindex.utils, "llm_acompletion", flaky) + structure = [{"title": "A", "text": "t1", + "nodes": [{"title": "B", "text": "t2"}]}] + result = asyncio.run(pageindex.utils.generate_summaries_for_structure(structure)) + summaries = {n["title"]: n["summary"] + for n in pageindex.utils.structure_to_list(result)} + assert summaries == {"A": "", "B": "ok"} + + +def test_delete_survives_marker_tamper(local_client, tmp_path): + tampered = tmp_path / "store" / "docs" / "tampered" / "doc.json" + tampered.mkdir(parents=True) + with pytest.raises(PageIndexAPIError, match="Document not found"): + local_client.delete_document("tampered") + assert not tampered.parent.exists() + + +def test_list_documents_validation(local_client): + with pytest.raises(ValueError): + local_client.list_documents(limit=0) + with pytest.raises(ValueError): + local_client.list_documents(offset=-1) + with pytest.raises(PageIndexAPIError, match="folders"): + local_client.list_documents(folder_id="f1") + + +def test_missing_document_errors(local_client): + with pytest.raises(PageIndexAPIError): + local_client.get_tree("nope") + with pytest.raises(PageIndexAPIError): + local_client.get_document("nope") + assert local_client.is_retrieval_ready("nope") is False + + +def test_traversal_ids_are_contained(local_client, indexed_doc, tmp_path): + store_root = tmp_path / "store" + with pytest.raises(PageIndexAPIError): + local_client.get_document("../../etc") + with pytest.raises(PageIndexAPIError): + local_client.delete_document("..") + assert (store_root / "docs").exists() + + +def test_folders_are_cloud_only(local_client): + with pytest.raises(PageIndexAPIError, match="cloud-only"): + local_client.create_folder("team") + with pytest.raises(PageIndexAPIError, match="cloud-only"): + local_client.list_folders() + + +# ── local: retrieval endpoints are cloud-only ── + +def test_retrieval_endpoints_cloud_only(local_client): + with pytest.raises(PageIndexAPIError, match="use chat_completions"): + local_client.submit_query("any", "q") + with pytest.raises(PageIndexAPIError, match="use chat_completions"): + local_client.get_retrieval("any") + + +def test_chat_completions_cloud_only(local_client): + with pytest.raises(PageIndexAPIError, match="not yet supported in local mode"): + local_client.chat_completions( + messages=[{"role": "user", "content": "q"}]) + + +# ── cloud mode: request wiring ── + +class FakeResponse: + def __init__(self, payload=None, status_code=200, text="", content=b"{}", + lines=None): + self._payload = payload if payload is not None else {} + self.status_code = status_code + self.text = text + self.content = content + self._lines = lines or [] + + def json(self): + return self._payload + + def iter_lines(self): + return iter(self._lines) + + def close(self): + pass + + +def _patch_requests(monkeypatch, handler): + """Replace cloud_api's requests module with per-verb fakes.""" + fake = types.SimpleNamespace( + post=lambda url, **kw: handler("POST", url, kw), + get=lambda url, **kw: handler("GET", url, kw), + delete=lambda url, **kw: handler("DELETE", url, kw), + Response=FakeResponse, + ) + monkeypatch.setattr("pageindex.cloud_api.requests", fake) + + +@pytest.fixture +def cloud(monkeypatch): + client = PageIndexClient(api_key="secret") + calls = [] + class Fake: + payload = {} + def handler(method, url, kw): + calls.append({"method": method, "url": url, **kw}) + return FakeResponse(Fake.payload) + _patch_requests(monkeypatch, handler) + return client, calls, Fake + + +def test_cloud_request_wiring(cloud, sample_pdf): + client, calls, fake = cloud + + fake.payload = {"doc_id": "pi-1"} + assert client.submit_document(sample_pdf) == {"doc_id": "pi-1"} + assert calls[-1]["url"] == "https://api.pageindex.ai/doc/" + assert calls[-1]["headers"] == {"api_key": "secret"} + assert calls[-1]["data"] == {"if_retrieval": True} + assert "timeout" not in calls[-1] + + client.submit_document(sample_pdf, metadata={"project": "alpha"}) + assert calls[-1]["data"]["metadata"] == json.dumps({"project": "alpha"}) + + fake.payload = {"status": "processing", "retrieval_ready": False} + client.get_tree("pi-1", node_summary=True) + assert calls[-1]["url"].endswith("/doc/pi-1/?type=tree&summary=True&include_text=true") + assert calls[-1]["timeout"] == 30 + assert client.is_retrieval_ready("pi-1") is False + + client.get_ocr("pi/../1") + assert "/doc/pi%2F..%2F1/" in calls[-1]["url"] + + client.BASE_URL = "https://staging.example" + client.api_key = "other" + client.get_document("pi-1") + assert calls[-1]["url"] == "https://staging.example/doc/pi-1/metadata/" + assert calls[-1]["headers"] == {"api_key": "other"} + + +def test_cloud_error_and_empty_delete(cloud, monkeypatch): + client, calls, fake = cloud + _patch_requests(monkeypatch, + lambda m, url, kw: FakeResponse(status_code=401, text="denied")) + with pytest.raises(PageIndexAPIError, + match="Failed to get document metadata: denied"): + client.get_document("pi-1") + + _patch_requests(monkeypatch, lambda m, url, kw: FakeResponse(content=b"")) + assert client.delete_document("pi-1") == {} + + +def test_cloud_chat_stream_parsing(cloud, monkeypatch): + client, calls, fake = cloud + lines = [ + b'data: {"choices": [{"delta": {"role": "assistant", "content": ""}}]}', + b'data: {"choices": [{"delta": {"content": "Hi"}}]}', + b"", + b'data: {"object": "chat.completion.citations", "citations": []}', + b'data: {"choices": [{"delta": {"content": " there"}}]}', + b"data: [DONE]", + ] + _patch_requests(monkeypatch, lambda m, url, kw: FakeResponse(lines=lines)) + pieces = list(client.chat_completions( + messages=[{"role": "user", "content": "q"}], stream=True)) + assert pieces == ["Hi", " there"] + + chunks = list(client.chat_completions( + messages=[{"role": "user", "content": "q"}], stream=True, + stream_metadata=True)) + assert {"object": "chat.completion.citations", "citations": []} in chunks diff --git a/tests/test_issue_163.py b/tests/test_issue_163.py index 7281cd14d..1e2999fe3 100644 --- a/tests/test_issue_163.py +++ b/tests/test_issue_163.py @@ -1,17 +1,11 @@ import pytest import sys import os -from importlib import import_module from unittest.mock import patch, MagicMock sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) -# Patch attributes on the module object: the package exports a *function* -# named page_index that shadows the submodule, and Python 3.10's mock -# resolves the string "pageindex.page_index" to that function. -page_index_module = import_module("pageindex.page_index") - -from pageindex.page_index import ( +from pageindex.page_index_classic import ( check_if_toc_extraction_is_complete, check_if_toc_transformation_is_complete, toc_detector_single_page, @@ -22,50 +16,50 @@ class TestRobustKeyAccess: - @patch.object(page_index_module, "llm_completion", return_value="") + @patch("pageindex.page_index_classic.llm_completion", return_value="") def test_toc_detector_empty_response(self, mock_llm): result = toc_detector_single_page("some content", model="test") assert result == "no" - @patch.object(page_index_module, "llm_completion", return_value='{"toc_detected": "yes"}') + @patch("pageindex.page_index_classic.llm_completion", return_value='{"toc_detected": "yes"}') def test_toc_detector_valid_response(self, mock_llm): result = toc_detector_single_page("some content", model="test") assert result == "yes" - @patch.object(page_index_module, "llm_completion", return_value="not json at all") + @patch("pageindex.page_index_classic.llm_completion", return_value="not json at all") def test_toc_detector_malformed_response(self, mock_llm): result = toc_detector_single_page("some content", model="test") assert result == "no" - @patch.object(page_index_module, "llm_completion", return_value="") + @patch("pageindex.page_index_classic.llm_completion", return_value="") def test_extraction_complete_empty_response(self, mock_llm): result = check_if_toc_extraction_is_complete("doc", "toc", model="test") assert result == "no" - @patch.object(page_index_module, "llm_completion", return_value='{"completed": "yes"}') + @patch("pageindex.page_index_classic.llm_completion", return_value='{"completed": "yes"}') def test_extraction_complete_valid_response(self, mock_llm): result = check_if_toc_extraction_is_complete("doc", "toc", model="test") assert result == "yes" - @patch.object(page_index_module, "llm_completion", return_value="") + @patch("pageindex.page_index_classic.llm_completion", return_value="") def test_transformation_complete_empty_response(self, mock_llm): result = check_if_toc_transformation_is_complete("raw", "cleaned", model="test") assert result == "no" - @patch.object(page_index_module, "llm_completion", return_value='{"thinking": "looks fine", "completed": "yes"}') + @patch("pageindex.page_index_classic.llm_completion", return_value='{"thinking": "looks fine", "completed": "yes"}') def test_transformation_complete_valid_response(self, mock_llm): result = check_if_toc_transformation_is_complete("raw", "cleaned", model="test") assert result == "yes" - @patch.object(page_index_module, "llm_completion", return_value="") + @patch("pageindex.page_index_classic.llm_completion", return_value="") def test_detect_page_index_empty_response(self, mock_llm): result = detect_page_index("toc text", model="test") assert result == "no" class TestExtractTocContentRetryLoop: - @patch.object(page_index_module, "check_if_toc_transformation_is_complete") - @patch.object(page_index_module, "llm_completion") + @patch("pageindex.page_index_classic.check_if_toc_transformation_is_complete") + @patch("pageindex.page_index_classic.llm_completion") def test_completes_on_first_try(self, mock_llm, mock_check): mock_llm.return_value = ("full toc content", "finished") mock_check.return_value = "yes" @@ -73,8 +67,8 @@ def test_completes_on_first_try(self, mock_llm, mock_check): assert result == "full toc content" assert mock_llm.call_count == 1 - @patch.object(page_index_module, "check_if_toc_transformation_is_complete") - @patch.object(page_index_module, "llm_completion") + @patch("pageindex.page_index_classic.check_if_toc_transformation_is_complete") + @patch("pageindex.page_index_classic.llm_completion") def test_continues_on_incomplete(self, mock_llm, mock_check): mock_llm.side_effect = [ ("partial toc", "max_output_reached"), @@ -85,8 +79,8 @@ def test_continues_on_incomplete(self, mock_llm, mock_check): assert result == "partial toc continued toc" assert mock_llm.call_count == 2 - @patch.object(page_index_module, "check_if_toc_transformation_is_complete") - @patch.object(page_index_module, "llm_completion") + @patch("pageindex.page_index_classic.check_if_toc_transformation_is_complete") + @patch("pageindex.page_index_classic.llm_completion") def test_max_retries_raises_exception(self, mock_llm, mock_check): mock_llm.return_value = ("chunk", "max_output_reached") mock_check.return_value = "no" @@ -94,8 +88,8 @@ def test_max_retries_raises_exception(self, mock_llm, mock_check): extract_toc_content("raw content", model="test") assert mock_llm.call_count == 6 - @patch.object(page_index_module, "check_if_toc_transformation_is_complete") - @patch.object(page_index_module, "llm_completion") + @patch("pageindex.page_index_classic.check_if_toc_transformation_is_complete") + @patch("pageindex.page_index_classic.llm_completion") def test_chat_history_grows_incrementally(self, mock_llm, mock_check): call_count = [0] @@ -120,8 +114,8 @@ def side_effect(*args, **kwargs): class TestTocTransformerRetryLoop: - @patch.object(page_index_module, "check_if_toc_transformation_is_complete") - @patch.object(page_index_module, "llm_completion") + @patch("pageindex.page_index_classic.check_if_toc_transformation_is_complete") + @patch("pageindex.page_index_classic.llm_completion") def test_completes_on_first_try(self, mock_llm, mock_check): mock_llm.return_value = ( '{"table_of_contents": [{"structure": "1", "title": "Intro", "page": 1}]}', @@ -132,8 +126,8 @@ def test_completes_on_first_try(self, mock_llm, mock_check): assert len(result) == 1 assert result[0]["title"] == "Intro" - @patch.object(page_index_module, "check_if_toc_transformation_is_complete") - @patch.object(page_index_module, "llm_completion") + @patch("pageindex.page_index_classic.check_if_toc_transformation_is_complete") + @patch("pageindex.page_index_classic.llm_completion") def test_handles_missing_table_of_contents_key(self, mock_llm, mock_check): mock_llm.return_value = ('{"other_key": "value"}', "finished") mock_check.return_value = "yes" diff --git a/tests/test_package_surface.py b/tests/test_package_surface.py new file mode 100644 index 000000000..e10c3e8c5 --- /dev/null +++ b/tests/test_package_surface.py @@ -0,0 +1,62 @@ +"""What `pip install pageindex` exposes: 0.2.8 helper compat and import cost.""" +import subprocess +import sys + +from pageindex.utils import create_node_mapping, print_tree, remove_fields + +TREE = [ + {"title": "Root", "node_id": "0000", "page_index": 1, + "text": "root text", + "nodes": [ + {"title": "Child", "node_id": "0001", "page_index": 3, + "text": "child text"}, + ]}, + {"title": "Tail", "node_id": "0002", "page_index": 5, "text": "tail text"}, +] + + +# ── the published 0.2.8 pageindex.utils surface, as the cookbooks call it ── + +def test_remove_fields_max_len(): + out = remove_fields({"keep": "x" * 50, "text": "gone"}, max_len=10) + assert out == {"keep": "x" * 10 + "..."} + assert remove_fields({"keep": "short"}, max_len=10) == {"keep": "short"} + + +def test_create_node_mapping_flat(): + mapping = create_node_mapping(TREE) + assert set(mapping) == {"0000", "0001", "0002"} + assert mapping["0001"]["title"] == "Child" + + +def test_create_node_mapping_page_ranges(): + mapping = create_node_mapping(TREE, include_page_ranges=True, max_page=9) + assert mapping["0000"] == {"node": TREE[0], "start_index": 1, "end_index": 3} + assert mapping["0001"]["start_index"] == 3 + assert mapping["0001"]["end_index"] == 5 + assert mapping["0002"] == {"node": TREE[1], "start_index": 5, "end_index": 9} + + +def test_print_tree_exclude_fields(capsys): + print_tree(TREE, exclude_fields=["text"]) + out = capsys.readouterr().out + assert "Root" in out and "'text'" not in out + + print_tree(TREE) + assert "[0000] Root" in capsys.readouterr().out + + +# ── import cost: the SDK must not pay for the indexing stack ── + +def test_import_pageindex_is_lazy(): + probe = ( + "import sys; import pageindex; " + "heavy = [m for m in ('pageindex.page_index_classic', 'pageindex.flash', " + "'pageindex.utils', 'pageindex.tree_optimize', 'numpy', 'PyPDF2') " + "if m in sys.modules]; " + "print(','.join(heavy) or 'clean'); " + "print(type(pageindex.page_index_main).__name__)" + ) + out = subprocess.run([sys.executable, "-c", probe], + capture_output=True, text=True, check=True) + assert out.stdout.split() == ["clean", "function"] diff --git a/tests/test_page_index.py b/tests/test_page_index.py index 1ffd51581..da69a9a79 100644 --- a/tests/test_page_index.py +++ b/tests/test_page_index.py @@ -1,13 +1,7 @@ import unittest -from importlib import import_module from unittest.mock import Mock, patch -# Patch attributes on the module object: the package exports a *function* -# named page_index that shadows the submodule, and Python 3.10's mock -# resolves the string "pageindex.page_index" to that function. -page_index_module = import_module("pageindex.page_index") - -from pageindex.page_index import ( +from pageindex.page_index_classic import ( _secure_doc_text, process_no_toc, process_toc_no_page_numbers, @@ -25,10 +19,10 @@ def test_rejects_same_length_reordered_llm_toc(self): {"structure": "1", "title": "First", "physical_index": ""}, ] - with patch.object(page_index_module, "toc_transformer", return_value=toc), \ - patch.object(page_index_module, "count_tokens", return_value=1), \ - patch.object(page_index_module, "page_list_to_group_text", return_value=[" "]), \ - patch.object(page_index_module, "add_page_number_to_toc", return_value=reordered): + with patch("pageindex.page_index_classic.toc_transformer", return_value=toc), \ + patch("pageindex.page_index_classic.count_tokens", return_value=1), \ + patch("pageindex.page_index_classic.page_list_to_group_text", return_value=[" "]), \ + patch("pageindex.page_index_classic.add_page_number_to_toc", return_value=reordered): with self.assertRaises(ValueError): process_toc_no_page_numbers( "toc", @@ -38,17 +32,17 @@ def test_rejects_same_length_reordered_llm_toc(self): ) def test_process_no_toc_validates_continuation_chunks(self): - with patch.object(page_index_module, "count_tokens", return_value=1), \ - patch.object( - page_index_module, "page_list_to_group_text", + with patch("pageindex.page_index_classic.count_tokens", return_value=1), \ + patch( + "pageindex.page_index_classic.page_list_to_group_text", return_value=["", ""], ), \ - patch.object( - page_index_module, "generate_toc_init", + patch( + "pageindex.page_index_classic.generate_toc_init", return_value=[{"title": "First", "physical_index": ""}], ), \ - patch.object( - page_index_module, "generate_toc_continue", + patch( + "pageindex.page_index_classic.generate_toc_continue", return_value=[{"title": "Second", "physical_index": ""}], ): result = process_no_toc( From 6b343f55eeb075082f09a7547843922e6de364e0 Mon Sep 17 00:00:00 2001 From: Ray Date: Sun, 9 Aug 2026 19:21:39 +0800 Subject: [PATCH 02/18] =?UTF-8?q?feat:=20agent=20tools=20=E2=80=94=20the?= =?UTF-8?q?=20cloud=20MCP=20tool=20contract=20on=20the=20client?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Four new client methods make PageIndex documents available to agent frameworks, in both modes, with the mode decided solely by the client constructor: - agent_tools(): plain functions (browse_documents, get_document, get_document_structure, get_page_content) matching the PageIndex cloud MCP server's tools/list — same names, schemas, descriptions, and JSON response envelopes — so agent prompts port unchanged between the cloud MCP connection and these in-process tools. Tools never raise; errors come back in the same envelope. remove_document ships behind include_management=False. - as_openai_tools(): the same tools wrapped for the OpenAI Agents SDK. - as_claude_mcp(): one mcp_servers entry for the Claude Agent SDK — cloud clients get the remote MCP config (the framework connects to api.pageindex.ai/mcp and discovers the full cloud tool set), local clients get an in-process SDK MCP server. - agent_instructions(doc_id=None): orchestration guidance for the agent's system prompt; doc_id (same shape as chat_completions) appends the target documents. submit_document() gains wait=True: poll get_document status until completed, raise on failed or after 30 minutes — the manual polling loop every cloud caller writes today spins forever on a failed document. Neither framework becomes a dependency: imports happen at call time with actionable errors, and pageindex[openai] / pageindex[claude] extras are floor-only pins. tests/data/cloud_mcp_contract.json freezes the tool contract; a parity test guards against drift. 36 new tests (95 total), plus a live OpenAI Agents SDK run over a seeded local store verifying the structure-first navigation flow end to end. --- README.md | 59 +- examples/agentic_vectorless_rag_demo.py | 58 +- pageindex/agent_tools.py | 1326 ++++++++++++++++++++ pageindex/client.py | 148 ++- pageindex/integrations/__init__.py | 5 + pageindex/integrations/claude_agent_sdk.py | 67 + pageindex/integrations/openai_agents.py | 36 + pageindex/mcp_bridge.py | 181 +++ pyproject.toml | 8 + tests/data/cloud_mcp_contract.json | 197 +++ tests/test_agent_tools.py | 879 +++++++++++++ 11 files changed, 2915 insertions(+), 49 deletions(-) create mode 100644 pageindex/agent_tools.py create mode 100644 pageindex/integrations/__init__.py create mode 100644 pageindex/integrations/claude_agent_sdk.py create mode 100644 pageindex/integrations/openai_agents.py create mode 100644 pageindex/mcp_bridge.py create mode 100644 tests/data/cloud_mcp_contract.json create mode 100644 tests/test_agent_tools.py diff --git a/README.md b/README.md index 5ce0ca5e6..5dbafc141 100644 --- a/README.md +++ b/README.md @@ -207,13 +207,68 @@ python3 run_pageindex.py --md_path /path/to/your/document.md > > Add `--optimize` to refine the tree structure for more efficient retrieval (with an LLM expansion pass). +## 🐍 Python SDK: Cloud & Local + +The `pageindex` package on PyPI is the Python SDK for the [PageIndex API](https://docs.pageindex.ai) — and the same client now also runs fully **locally**, powered by this repo's indexing pipeline (including Flash). + +```bash +pip3 install --upgrade pageindex # local mode ships in pageindex >= 0.2.9; earlier versions are cloud-only +``` + +```python +from pageindex import PageIndexClient + +client = PageIndexClient(api_key="YOUR_PAGEINDEX_API_KEY") # cloud: managed OCR, tree building, retrieval +client = PageIndexClient() # local: same methods on your machine, using your LLM key (e.g. OPENAI_API_KEY) + +doc_id = client.submit_document("doc.pdf")["doc_id"] # local mode blocks until indexing finishes +doc_id = client.submit_document("doc.pdf", mode="flash")["doc_id"] # local mode with PageIndex Flash + +tree = client.get_tree(doc_id, node_summary=True)["result"] + +answer = client.chat_completions( + messages=[{"role": "user", "content": "Summarize the key findings"}], + doc_id=doc_id, +)["choices"][0]["message"]["content"] +``` + +Local documents are stored as plain JSON under `./.pageindex` (configurable via `storage_path`). Local mode supports PDFs; folders, `beta_headers`, `enable_citations`, and the deprecated retrieval API (`submit_query`/`get_retrieval`) remain cloud-only — each method's docstring spells out the differences. To pin the mode at construction instead of inferring it from `api_key`, use `PageIndexCloudClient` (fails without a real key) or `PageIndexLocalClient` (has no key parameter). + +### 🤖 Agent integration + +The client exposes its documents as **agent tools**, following one rule: **cloud clients always serve the live tool set of the [PageIndex MCP server](https://docs.pageindex.ai/mcp)** (search, folders, images — as enabled for your key, discovered dynamically; management tools like delete/upload sit behind `include_management=True` or the framework's approval layer), while local clients serve the same contract's built-in navigation subset (`browse_documents`, `get_document`, `get_document_structure`, `get_page_content`). Tool names and schemas are shared, so agent prompts port unchanged, and switching local ↔ cloud is just the client constructor line: + +```python +client = PageIndexLocalClient() # or PageIndexCloudClient(api_key=...) +client.submit_document("doc.pdf", wait=True) # wait=True: return once the doc is ready (both modes) + +# OpenAI Agents SDK (pip install "pageindex[openai]") +agent = Agent( + name="PageIndex", + instructions=client.agent_instructions(), # retrieval playbook for the agent's system prompt + tools=client.as_openai_tools(), # local: in-process tools; cloud: the full cloud MCP tool set (any model backend) +) # cloud + OpenAI models: hosted=True runs tool calls server-side (fastest) + +# Claude Agent SDK (pip install "pageindex[claude]") +options = ClaudeAgentOptions( + system_prompt=client.agent_instructions(), + mcp_servers={"pageindex": client.as_claude_mcp()}, # local: in-process server; cloud: connects to api.pageindex.ai/mcp + allowed_tools=["mcp__pageindex__*"], +) + +# Any other framework: plain functions, wrap with your framework's one-liner +tools = client.agent_tools() # local: built-in tools; cloud: full live tool set over MCP + # e.g. [StructuredTool.from_function(f) for f in tools] +``` + +Neither framework is a required dependency — each is imported only when its method is called. Claude Code / Cursor and other MCP hosts connect to cloud documents via the hosted MCP server directly (no SDK needed); see the [MCP docs](https://docs.pageindex.ai/mcp). ## 🚀 Agentic Vectorless RAG: An Example For a simple, end-to-end **agentic vectorless RAG** example using **self-hosted PageIndex** (with OpenAI Agents SDK), see [`examples/agentic_vectorless_rag_demo.py`](examples/agentic_vectorless_rag_demo.py). ```bash -# Install optional dependency -pip3 install openai-agents +# Install with the OpenAI Agents SDK extra +pip3 install "pageindex[openai]" # Run the demo python3 examples/agentic_vectorless_rag_demo.py diff --git a/examples/agentic_vectorless_rag_demo.py b/examples/agentic_vectorless_rag_demo.py index 4fe5f179f..e8ed4a50a 100644 --- a/examples/agentic_vectorless_rag_demo.py +++ b/examples/agentic_vectorless_rag_demo.py @@ -6,20 +6,21 @@ chunking, PageIndex builds a hierarchical tree index and uses agentic LLM reasoning for human-like, context-aware retrieval. -Agent tools: - - get_document() — document metadata (status, page count, etc.) - - get_document_structure() — tree structure index of a document - - get_page_content() — retrieve text content of specific pages +The agent tools come straight from the SDK — ``client.as_openai_tools()`` +exposes the PageIndex tool contract (browse_documents, get_document, +get_document_structure, get_page_content) and ``client.agent_instructions()`` +provides the retrieval playbook, so the whole agent is a few lines. Swap +``PageIndexLocalClient()`` for ``PageIndexCloudClient(api_key=...)`` and the +same code runs against the cloud. Steps: 1 — Index a PDF locally and view its tree structure index 2 — View document metadata 3 — Ask a question (agent reasons over the index and auto-calls tools) -Requirements: pip install openai-agents; OPENAI_API_KEY in the environment. +Requirements: pip install "pageindex[openai]"; OPENAI_API_KEY in the environment. """ import sys -import json import asyncio import concurrent.futures from pathlib import Path @@ -27,12 +28,12 @@ sys.path.insert(0, str(Path(__file__).parent.parent)) -from agents import Agent, Runner, function_tool, set_tracing_disabled +from agents import Agent, Runner, set_tracing_disabled from agents.model_settings import ModelSettings from agents.stream_events import RawResponsesStreamEvent, RunItemStreamEvent from openai.types.responses import ResponseTextDeltaEvent, ResponseReasoningSummaryTextDeltaEvent -from pageindex import PageIndexClient +from pageindex import PageIndexLocalClient import pageindex.utils as utils PDF_URL = "https://arxiv.org/pdf/2603.15031" @@ -41,47 +42,18 @@ PDF_PATH = _EXAMPLES_DIR / "documents" / "attention-residuals.pdf" STORAGE_PATH = _EXAMPLES_DIR / ".pageindex" -AGENT_SYSTEM_PROMPT = """ -You are PageIndex, a document QA assistant. -TOOL USE: -- Call get_document() first to confirm status and page count. -- Call get_document_structure() to identify relevant page ranges. -- Call get_page_content(pages="5-7") with tight ranges; never fetch the whole document. -- Before each tool call, output one short sentence explaining the reason. -Answer based only on tool output. Be concise. -""" - -def query_agent(client: PageIndexClient, doc_id: str, prompt: str, verbose: bool = False) -> str: +def query_agent(client: PageIndexLocalClient, doc_id: str, prompt: str, verbose: bool = False) -> str: """Run a document QA agent using the OpenAI Agents SDK. Streams text output token-by-token and returns the full answer string. Tool calls are always printed; verbose=True also prints arguments and output previews. """ - - @function_tool - def get_document() -> str: - """Get document metadata: status, page count, name, and description.""" - return json.dumps(client.get_document(doc_id)) - - @function_tool - def get_document_structure() -> str: - """Get the document's full tree structure (without text) to find relevant sections.""" - return json.dumps(client.get_document_structure(doc_id), ensure_ascii=False) - - @function_tool - def get_page_content(pages: str) -> str: - """ - Get the text content of specific pages. - Use tight ranges: e.g. '5-7' for pages 5 to 7, '3,8' for pages 3 and 8, '12' for page 12. - """ - return json.dumps(client.get_page_content(doc_id, pages), ensure_ascii=False) - agent = Agent( name="PageIndex", - instructions=AGENT_SYSTEM_PROMPT, - tools=[get_document, get_document_structure, get_page_content], - model=getattr(client, "retrieve_model", None), + instructions=client.agent_instructions(doc_id=doc_id), + tools=client.as_openai_tools(), + model=client.retrieve_model, # model_settings=ModelSettings(reasoning={"effort": "low", "summary": "auto"}), # Uncomment to enable reasoning ) @@ -152,7 +124,7 @@ async def _run(): print("Download complete.\n") # Setup: local mode — no PageIndex API key needed, your LLM key does the work - client = PageIndexClient(storage_path=str(STORAGE_PATH)) + client = PageIndexLocalClient(storage_path=str(STORAGE_PATH)) # Step 1: Index PDF and view tree structure print("=" * 60) @@ -166,7 +138,7 @@ async def _run(): if doc_id: print(f"\nLoaded cached doc_id: {doc_id}") else: - doc_id = client.submit_document(str(PDF_PATH))["doc_id"] + doc_id = client.submit_document(str(PDF_PATH), wait=True)["doc_id"] print(f"\nIndexed. doc_id: {doc_id}") print("\nTree Structure (top-level sections):") structure = client.get_tree(doc_id, node_summary=True)["result"] diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py new file mode 100644 index 000000000..10f0340dd --- /dev/null +++ b/pageindex/agent_tools.py @@ -0,0 +1,1326 @@ +"""Agent tools: the cloud MCP tool contract, executed against a PageIndexClient. + +Tool names, input schemas, and descriptions match the PageIndex cloud MCP +server, so agent prompts work unchanged across the cloud MCP connection and +this in-process layer. Only the tools that exist in every mode are registered +(no folders, search_documents, or get_document_image). + +Tools never raise: every outcome, including errors, is returned as the same +JSON envelope the cloud emits ({"success": true, ...} / {"error": ...}). +""" +from __future__ import annotations + +import copy +import difflib +import json +import re +import time +from typing import Any, Callable, Optional + +from .errors import PageIndexAPIError + +TOOL_RESPONSE_CHAR_LIMIT = 100_000 +STRUCTURE_FIRST_PAGE_THRESHOLD = 20 + +_CHAR_BUDGET = int(TOOL_RESPONSE_CHAR_LIMIT * 0.95) +_PAGES_SPEC_RE = re.compile(r"^(\d+(-\d+)?)(,\s*\d+(-\d+)?)*$") +_SIMILAR_NAMES_LIMIT = 3 +_TOOL_WAIT_TIMEOUT = 180.0 # "up to 3 minutes", per the wait_for_completion schema +_TOOL_WAIT_INTERVAL = 5.0 + +_DOC_NAME_DESCRIPTION = ( + 'Copy the `name` field verbatim from a browse_documents() or ' + 'search_documents() response (case-sensitive, include extension). ' + 'Example: "Q3 Report.pdf". If the response shows two documents with the ' + 'same name, pass `folder_id` alongside to disambiguate.' +) +_FOLDER_ID_DISAMBIGUATOR_DESCRIPTION = ( + 'Disambiguator for same-name documents. Copy the `folder_id` from the ' + 'intended browse/search result; use "root" for root-level documents, or ' + '"shared-with-me"/"following" for the read-only folders at the library ' + 'root; omit if `doc_name` is unique. Copy any folder_id verbatim from a ' + 'browse_documents()/get_folder_structure() response, never construct one.' +) +_WAIT_FOR_COMPLETION_DESCRIPTION = ( + "If true and document is processing, automatically wait up to 3 minutes " + "until completed. Reduces repeated tool calls." +) + +#: Tool names, descriptions, and parameter schemas, identical to the cloud +#: MCP server's tools/list. +TOOL_CONTRACT: dict[str, dict[str, Any]] = { + "browse_documents": { + "annotations": {"readOnlyHint": True, "openWorldHint": False}, + "description": ( + "Primary document retrieval tool. After orienting with " + "get_folder_structure() (when available), use this for all " + "document-related questions. The bare call returns root-level " + "sub-folders and documents; pass folder_id to drill into a " + 'sub-folder level by level. Use sort="relevance" + query for ' + "semantic ranking. Do NOT jump to search_documents() first — it " + "is an escalation path, only after " + 'browse_documents(sort="relevance") has failed.' + ), + "schema": { + "type": "object", + "properties": { + "folder_id": { + "type": "string", + "default": "root", + "description": ( + 'Folder scope (default "root"). Pass a specific folder ' + 'ID to scope into that folder, or "root" to reference ' + "the library root. The read-only \"shared-with-me\" and " + '"following" folders live at the library root — pass ' + "one of those ids to browse them. Copy any folder_id " + "verbatim from a browse/tree response, never construct " + "one. Combine with `recursive` to control breadth." + ), + }, + "recursive": { + "type": "boolean", + "default": False, + "description": ( + "Whether to include documents from descendant folders. " + "When false (default), returns the direct contents of " + "folder_id along with its sub-folders — prefer this for " + "level-by-level exploration so you retain folder " + "hierarchy context. When true, flattens all descendant " + "documents into one list and omits sub-folders — use " + "only when a non-recursive browse of the target folder " + "returned no relevant results and you need to widen the " + "scope, or the user explicitly requests a flat listing." + ), + }, + "sort": { + "type": "string", + "enum": ["time", "relevance"], + "default": "time", + "description": ( + 'Sort order. "time" (default) sorts by upload date ' + '(newest first); "relevance" orders documents by ' + "semantic relevance to `query`. Relevance also works " + "inside the read-only shared folders — pass their " + "folder_id — but at the library root it ranks only " + "your own documents." + ), + }, + "query": { + "type": "string", + "description": ( + "Search query for relevance ranking. Required when " + 'sort="relevance"; must be omitted when sort="time".' + ), + }, + "offset": { + "type": "integer", + "minimum": 0, + "default": 0, + "description": ( + "Zero-based pagination offset. Pass the value of " + "`next_offset` from the previous response to fetch the " + "next page." + ), + }, + "limit": { + "type": "number", + "minimum": 1, + "maximum": 50, + "default": 10, + "description": ( + "Number of documents to return per page (1-50, " + "default 10)" + ), + }, + }, + "required": [], + }, + }, + "get_document": { + "annotations": {"readOnlyHint": True, "openWorldHint": False}, + "description": ( + "Check a document's processing status and metadata. `status` is " + 'one of "pending", "queued", "processing", "completed", or ' + '"failed" — call this before `get_document_structure()` or ' + "`get_page_content()` to confirm the document is ready." + ), + "schema": { + "type": "object", + "properties": { + "doc_name": { + "type": "string", + "minLength": 1, + "description": _DOC_NAME_DESCRIPTION, + }, + "folder_id": { + "type": ["string", "null"], + "description": _FOLDER_ID_DISAMBIGUATOR_DESCRIPTION, + }, + "wait_for_completion": { + "type": "boolean", + "default": False, + "description": _WAIT_FOR_COMPLETION_DESCRIPTION, + }, + }, + "required": ["doc_name"], + }, + }, + "get_document_structure": { + "annotations": {"readOnlyHint": True, "openWorldHint": False}, + "description": ( + "Extract a document's hierarchical outline (headers, sections, " + f"page references). REQUIRED for documents over " + f"{STRUCTURE_FIRST_PAGE_THRESHOLD} pages — call this first to " + "locate relevant sections, then pass their page numbers to " + "`get_page_content()`. Use the `part` parameter to iterate large " + "outlines until `pagination.has_more` is false." + ), + "schema": { + "type": "object", + "properties": { + "doc_name": { + "type": "string", + "minLength": 1, + "description": _DOC_NAME_DESCRIPTION, + }, + "folder_id": { + "type": ["string", "null"], + "description": _FOLDER_ID_DISAMBIGUATOR_DESCRIPTION, + }, + "part": { + "type": "integer", + "minimum": 1, + "default": 1, + "description": ( + "Part number for pagination (1-based, default 1). For " + "large outlines, increment until the response's " + "`pagination.has_more` becomes false." + ), + }, + "wait_for_completion": { + "type": "boolean", + "default": False, + "description": _WAIT_FOR_COMPLETION_DESCRIPTION, + }, + }, + "required": ["doc_name"], + }, + }, + "get_page_content": { + "annotations": {"readOnlyHint": True, "openWorldHint": False}, + "description": ( + "Extract page content from a processed document. Use tight, " + "targeted page ranges — never the whole document at once. For " + f"documents over {STRUCTURE_FIRST_PAGE_THRESHOLD} pages, call " + "`get_document_structure()` first to pick relevant sections. " + "Embedded image paths in the response feed into " + "`get_document_image()`." + ), + "schema": { + "type": "object", + "properties": { + "doc_name": { + "type": "string", + "minLength": 1, + "description": _DOC_NAME_DESCRIPTION, + }, + "folder_id": { + "type": ["string", "null"], + "description": _FOLDER_ID_DISAMBIGUATOR_DESCRIPTION, + }, + "pages": { + "type": "string", + "minLength": 1, + "pattern": r"^(\d+(-\d+)?)(,\s*\d+(-\d+)?)*$", + "description": ( + 'Page specification: "5", "3,7,10", "5-10", or ' + '"1-3,7,9-12"' + ), + }, + "wait_for_completion": { + "type": "boolean", + "default": False, + "description": _WAIT_FOR_COMPLETION_DESCRIPTION, + }, + }, + "required": ["doc_name", "pages"], + }, + }, + "remove_document": { + "annotations": {"readOnlyHint": False, "destructiveHint": True, + "idempotentHint": True, "openWorldHint": False}, + "description": ( + "Permanently delete documents and all associated data. Only invoke " + "when the user explicitly names the documents AND confirms " + "deletion. Returns `results` — one entry per requested document: " + '`{ doc_name, status: "deleted" | "not_found" | "failed", ' + "error? }`. Inspect each entry for per-document failures. This " + "action is irreversible." + ), + "schema": { + "type": "object", + "properties": { + "doc_names": { + "type": "array", + "items": {"type": "string", "minLength": 1}, + "minItems": 1, + "maxItems": 10, + "description": ( + "Array of document names to delete. Each name must be " + "copied verbatim from the `name` field of a " + "browse_documents() or search_documents() response " + "(case-sensitive, include extension). Example: " + '["Q3 Report.pdf", "draft.pdf"]. Max 10 per call.' + ), + }, + "folder_id": { + "type": ["string", "null"], + "description": _FOLDER_ID_DISAMBIGUATOR_DESCRIPTION, + }, + }, + "required": ["doc_names"], + }, + }, +} + +_READ_TOOLS = ("browse_documents", "get_document", "get_document_structure", + "get_page_content") +_MANAGEMENT_TOOLS = ("remove_document",) + + +# ── response envelopes ── + +_ToolResult = tuple[dict, bool] + + +def _success(data: dict[str, Any], next_steps: dict[str, Any]) -> tuple[dict, bool]: + return {"success": True, **data, "next_steps": next_steps}, False + + +def _failure(error: str, details: Optional[dict[str, Any]], + next_steps: dict[str, Any], error_code: Optional[str] = None, + ) -> tuple[dict, bool]: + payload: dict[str, Any] = {"error": error} + if error_code: + payload["errorCode"] = error_code + if details: + payload.update(details) + payload["next_steps"] = next_steps + return payload, True + + +def _dumps(payload: dict[str, Any]) -> str: + return json.dumps(payload, indent=2, ensure_ascii=False) + + +# ── document listing / name resolution ── + +def _all_documents(client) -> list[dict[str, Any]]: + """Every document the client can list, newest first (both modes list + newest-first; paging preserves that order).""" + documents: list[dict[str, Any]] = [] + offset = 0 + while True: + page = client.list_documents(limit=100, offset=offset) + batch = page.get("documents") or [] + documents.extend(batch) + offset += 100 + if not batch or offset >= page.get("total", 0): + return documents + + +def _normalize_created_at(value: Any) -> str: + """Emit the cloud tool format (ISO-8601 UTC with 'Z', millisecond + precision) from either mode's createdAt string.""" + if not isinstance(value, str) or not value: + return "" + try: + from datetime import datetime, timezone + parsed = datetime.fromisoformat(value.replace("Z", "+00:00")) + if parsed.tzinfo is None: + parsed = parsed.replace(tzinfo=timezone.utc) + parsed = parsed.astimezone(timezone.utc) + return parsed.isoformat(timespec="milliseconds").replace("+00:00", "Z") + except ValueError: + return value + + +def _flat_metadata(value: Any) -> Optional[dict[str, Any]]: + """User-facing string|number|boolean metadata fields only, or None.""" + if not isinstance(value, dict): + return None + flat = {key: val for key, val in value.items() + if isinstance(val, (str, int, float, bool))} + return flat or None + + +def _resolve_document( + client, doc_name: str, +) -> "tuple[Optional[dict[str, Any]], Optional[_ToolResult]]": + """Resolve doc_name to a list entry. Same-name duplicates resolve to the + newest match. Returns (entry, None) or (None, error_payload_pair).""" + documents = _all_documents(client) + matches = [doc for doc in documents if doc.get("name") == doc_name] + if matches: + return max(matches, key=lambda d: d.get("createdAt") or ""), None + names = [str(doc.get("name")) for doc in documents if doc.get("name")] + similar = difflib.get_close_matches(doc_name, names, n=_SIMILAR_NAMES_LIMIT, + cutoff=0.5) + message = ( + "Document not found. Did you mean: " + + ", ".join(f'"{name}"' for name in similar) + "?" + if similar else "Document not found or you do not have access to it" + ) + return None, _failure( + message, + {"doc_name": doc_name, "similar_files": similar}, + { + "summary": "The requested document does not exist or is not accessible", + "options": [ + "Verify the document name is correct", + "Use browse_documents() to see your recent documents", + "Check if the document was deleted", + ], + }, + "NOT_FOUND", + ) + + +def _refetch_entry(client, doc_id: str) -> Optional[dict[str, Any]]: + try: + return client.get_document(doc_id) + except PageIndexAPIError: + return None + + +def _await_completion(client, entry: dict[str, Any], wait: bool) -> dict[str, Any]: + """Re-poll a processing document for up to 3 minutes when wait is set.""" + doc_id = entry.get("id") + if not wait or not doc_id or entry.get("status") in ("completed", "failed"): + return entry + deadline = time.monotonic() + _TOOL_WAIT_TIMEOUT + current = entry + while time.monotonic() < deadline: + time.sleep(_TOOL_WAIT_INTERVAL) + refreshed = _refetch_entry(client, doc_id) + if refreshed is None: + return current + refreshed.setdefault("metadata", current.get("metadata")) + current = {**current, **refreshed} + if current.get("status") in ("completed", "failed"): + return current + return current + + +def _not_ready_error(doc_name: str, status: Any, operation: str, + timed_out: bool) -> tuple[dict, bool]: + if status == "failed": + return _failure( + f"Document processing failed. Current status: {status}", + {"doc_name": doc_name}, + { + "summary": "Document processing has failed", + "options": [ + "Index the document again with submit_document()", + "Use browse_documents() to work with other documents", + ], + }, + "INVALID_INPUT", + ) + if timed_out: + return _failure( + f"Document is still processing. Current status: {status}", + {"doc_name": doc_name}, + { + "summary": "Document processing timeout", + "options": [ + "Try again later when processing is complete", + "Check status with get_document()", + ], + }, + "INVALID_INPUT", + ) + return _failure( + f"Document is not ready for {operation}. Current status: {status}", + {"doc_name": doc_name}, + { + "summary": "Document is still processing", + "options": [ + "Wait for document processing to complete", + "Check status with browse_documents() or get_document()", + ], + }, + "INVALID_INPUT", + ) + + +def _folder_unsupported(param: str) -> tuple[dict, bool]: + return _failure( + f"Folders are not available here — omit {param}.", + None, + { + "summary": "This library has no folders", + "options": ["Retry the call without a folder_id", + "Use browse_documents() to list the library root"], + }, + "INVALID_INPUT", + ) + + +# ── page spec handling ── + +def _parse_page_spec( + pages: str, doc_name: str, +) -> "tuple[Optional[list[int]], Optional[_ToolResult]]": + """Expand '1-3,7' into a sorted, deduplicated page list, or an error.""" + invalid = _failure( + "Invalid page specification format", + {"doc_name": doc_name}, + { + "summary": "Failed to parse the pages parameter", + "options": [ + 'Use valid formats: "5", "3,7,10", "5-10", or "1-3,7,9-12"', + "Ensure page numbers are positive integers", + ], + }, + "INVALID_INPUT", + ) + if not isinstance(pages, str) or not _PAGES_SPEC_RE.match(pages.strip()): + return None, invalid + expanded: set[int] = set() + for part in pages.split(","): + part = part.strip() + if "-" in part: + start, end = (int(x) for x in part.split("-", 1)) + if start > end: + return None, invalid + expanded.update(range(start, end + 1)) + else: + expanded.add(int(part)) + if any(page < 1 for page in expanded): + return None, _failure( + "Invalid page numbers. Page numbers must be positive integers", + {"doc_name": doc_name}, + { + "summary": "Invalid page numbers provided", + "options": [ + "Page numbers must be positive integers (>= 1)", + "Check the page specification format", + ], + }, + "INVALID_INPUT", + ) + return sorted(expanded), None + + +def _format_page_spec(pages: list[int]) -> str: + """Compress [1,2,3,5] into '1-3,5'.""" + if not pages: + return "" + ordered = sorted(set(pages)) + ranges = [] + start = prev = ordered[0] + for page in ordered[1:]: + if page == prev + 1: + prev = page + continue + ranges.append(f"{start}" if start == prev else f"{start}-{prev}") + start = prev = page + ranges.append(f"{start}" if start == prev else f"{start}-{prev}") + return ",".join(ranges) + + +# ── structure formatting / splitting ── + +_STRUCTURE_KEY_ORDER = ("title", "node_id", "start_index", "end_index", + "page_index", "prefix_summary", "summary", "nodes") + + +def _format_structure(node: Any) -> Any: + """Drop node text and normalize key order, recursively.""" + if isinstance(node, list): + return [_format_structure(item) for item in node] + if isinstance(node, dict): + stripped = {key: value for key, value in node.items() if key != "text"} + if "nodes" in stripped: + stripped["nodes"] = _format_structure(stripped["nodes"]) + ordered = {key: stripped[key] for key in _STRUCTURE_KEY_ORDER + if key in stripped} + ordered.update({key: value for key, value in stripped.items() + if key not in ordered}) + return ordered + return node + + +def _serialized_size(value: Any) -> int: + return len(json.dumps(value, ensure_ascii=False)) + + +def _split_structure(structure: Any, budget: int) -> list[Any]: + """Split a formatted structure into chunks of at most ~budget serialized + chars. The paginated response shape matches the cloud tool; chunk + boundaries are implementation-defined.""" + if _serialized_size(structure) <= budget: + return [structure] + nodes = structure if isinstance(structure, list) else [structure] + chunks: list[Any] = [] + group: list[Any] = [] + group_size = 0 + for node in nodes: + size = _serialized_size(node) + if size > budget: + if group: + chunks.append(group if len(group) > 1 else group[0]) + group, group_size = [], 0 + chunks.extend(_split_oversized_node(node, budget)) + continue + if group and group_size + size > budget: + chunks.append(group if len(group) > 1 else group[0]) + group, group_size = [], 0 + group.append(node) + group_size += size + if group: + chunks.append(group if len(group) > 1 else group[0]) + return chunks or [structure] + + +def _split_oversized_node(node: Any, budget: int) -> list[Any]: + children = node.get("nodes") if isinstance(node, dict) else None + if not children: + return [node] + shell = {key: value for key, value in node.items() if key != "nodes"} + shell_size = _serialized_size(shell) + child_budget = max(budget - shell_size, budget // 2) + parts = [] + for chunk in _split_structure(children, child_budget): + parts.append({**shell, + "nodes": chunk if isinstance(chunk, list) else [chunk]}) + return parts + + +# ── tool implementations (client-backed; mode-blind) ── + +def _browse_documents(client, folder_id: str = "root", recursive: bool = False, + sort: str = "time", query: Optional[str] = None, + offset: int = 0, limit: int = 10) -> tuple[dict, bool]: + if folder_id != "root": + return _folder_unsupported("folder_id") + if sort not in ("time", "relevance"): + return _failure('sort must be "time" or "relevance"', None, + {"summary": "Invalid sort mode", + "options": ['Use sort="time" or sort="relevance"']}, + "INVALID_INPUT") + if sort == "relevance" and not query: + return _failure('query is required when sort is "relevance"', None, + {"summary": "Missing query for relevance ranking", + "options": ['Pass query alongside sort="relevance"']}, + "INVALID_INPUT") + if sort == "time" and query: + return _failure('query is only allowed when sort is "relevance"', None, + {"summary": "query does not apply to the time sort", + "options": ["Drop query, or set sort=\"relevance\""]}, + "INVALID_INPUT") + try: + offset = max(int(offset), 0) + limit = min(max(int(limit), 1), 50) + except (TypeError, ValueError): + return _failure("offset and limit must be numbers", None, + {"summary": "Invalid pagination parameters", + "options": ["Pass integer offset and limit values"]}, + "INVALID_INPUT") + + documents = _all_documents(client) + if sort == "relevance": + tokens = [token for token in (query or "").lower().split() if token] + scored = [] + for doc in documents: + haystack = f"{doc.get('name') or ''} {doc.get('description') or ''}".lower() + score = sum(1 for token in tokens if token in haystack) + if score: + scored.append((score, doc)) + # Stable sort: equal scores keep the newest-first listing order. + scored.sort(key=lambda pair: pair[0], reverse=True) + documents = [doc for _, doc in scored] + + window = documents[offset:offset + limit] + has_more = offset + limit < len(documents) + next_offset = offset + limit if has_more else None + + page_has_processing = False + page_has_failed = False + items = [] + for doc in window: + status = doc.get("status") or "unknown" + if status == "failed": + page_has_failed = True + elif status != "completed": + page_has_processing = True + item = { + "name": doc.get("name") or "Unknown Document", + "description": doc.get("description") or "No description provided", + "status": status, + "created_at": _normalize_created_at(doc.get("createdAt")), + } + metadata = _flat_metadata(doc.get("metadata")) + if metadata is not None: + item["metadata"] = metadata + items.append(item) + + data: dict[str, Any] = { + "documents": items, + "sort": sort, + "next_offset": next_offset, + "has_more": has_more, + } + if not recursive: + data["folders"] = [] + + if not items and offset == 0: + next_steps = { + "summary": "Nothing to show", + "options": ( + ["No documents matched this query. Rephrase with synonyms or " + "alternative terms and retry browse_documents(sort=\"relevance\")."] + if sort == "relevance" + else ["Nothing here. Index documents with " + "PageIndexClient.submit_document() to get started."] + ), + "auto_retry": ( + "Rephrase the query and retry browse_documents(sort=\"relevance\")" + if sort == "relevance" + else "Index a document with submit_document() to get started" + ), + } + return _success(data, next_steps) + + options = [] + if items: + options.append("Use get_document() with a document name to view details") + options.append( + "Results returned ≠ correct results. Verify these documents match " + "the user's actual intent (topic, time period, document type) " + "before proceeding. If they do not match, rephrase the query and " + "retry browse_documents(sort=\"relevance\"). Do NOT use general " + "knowledge as a substitute." + ) + if page_has_processing: + options.append("Some documents on this page are still processing. " + "Use get_document() to check individual status.") + if page_has_failed: + options.append("Some documents on this page failed processing. " + "Use get_document() to see error details.") + if has_more: + options.append("Use browse_documents() with `offset: next_offset` to " + "load more documents") + summary = (f"Showing {len(items)} document(s)" + + (" (more available)" if has_more else "") + if items else "Nothing to show") + return _success(data, {"summary": summary, "options": options}) + + +def _get_document(client, doc_name: str, folder_id: Optional[str] = None, + wait_for_completion: bool = False) -> tuple[dict, bool]: + if folder_id not in (None, "root"): + return _folder_unsupported("folder_id") + entry, error = _resolve_document(client, doc_name) + if error is not None: + return error + assert entry is not None + entry = _await_completion(client, entry, wait_for_completion) + + status = entry.get("status") or "unknown" + is_processing = status not in ("completed", "failed") + is_ready = status == "completed" + page_num = entry.get("pageNum") or 0 + name = entry.get("name") or "Unknown Document" + + suggestions: list[str] = [] + if is_processing: + suggestions.append("Document is still processing. Processing status " + "can be checked later.") + elif is_ready: + suggestions.append("Document is ready for analysis.") + if page_num > 0: + if page_num <= 5: + suggestions.extend([ + f"This is a short document with {page_num} pages.", + f'First explore structure: get_document_structure(doc_name: "{name}")', + f'Then extract all content: get_page_content(doc_name: "{name}", pages: "1-{page_num}")', + ]) + elif page_num <= STRUCTURE_FIRST_PAGE_THRESHOLD: + suggestions.extend([ + f"This document has {page_num} pages.", + f'First explore structure: get_document_structure(doc_name: "{name}")', + f'Then extract key pages: get_page_content(doc_name: "{name}", pages: "1,5,10")', + ]) + else: + suggestions.extend([ + f"This is a large document with {page_num} pages.", + f'Start with first few pages: get_page_content(doc_name: "{name}", pages: "1-3")', + f'Or view structure first: get_document_structure(doc_name: "{name}")', + ]) + else: + suggestions.append("Document processing failed. Index the document " + "again with submit_document().") + + data: dict[str, Any] = { + "name": name, + "description": entry.get("description") or "No description provided", + "status": status, + "created_at": _normalize_created_at(entry.get("createdAt")), + "page_count": page_num or None, + "folder_id": entry.get("folderId"), + } + metadata = _flat_metadata(entry.get("metadata")) + if metadata is not None: + data["metadata"] = metadata + + return _success(data, { + "summary": ("Document is ready for analysis and querying." if is_ready + else "Document is still being processed." if is_processing + else "Document processing has failed."), + "options": suggestions, + **({"auto_retry": "Document processing status can be monitored periodically"} + if is_processing else {}), + }) + + +def _get_document_structure(client, doc_name: str, + folder_id: Optional[str] = None, part: int = 1, + wait_for_completion: bool = False) -> tuple[dict, bool]: + if folder_id not in (None, "root"): + return _folder_unsupported("folder_id") + entry, error = _resolve_document(client, doc_name) + if error is not None: + return error + assert entry is not None + entry = _await_completion(client, entry, wait_for_completion) + if entry.get("status") != "completed": + return _not_ready_error(doc_name, entry.get("status"), + "structure retrieval", wait_for_completion) + + try: + # Prefer the raw stored tree: its nodes carry start_index/end_index + # like the cloud structure tool, where client.get_tree() drops + # end_index and renames fields. + store = getattr(getattr(client, "_api", None), "_store", None) + tree = store.get_tree(entry["id"]) if store is not None else None + if tree is None: + tree = client.get_tree(entry["id"], node_summary=True).get("result") + except PageIndexAPIError as exc: + return _failure( + f"Failed to retrieve document structure: {exc}", + {"doc_name": doc_name}, + { + "summary": "Failed to retrieve document structure due to an error", + "options": [ + "The document may not exist or is not accessible", + "Check if the document name is correct", + "Try again in a few moments", + ], + }, + "INTERNAL_ERROR", + ) + if tree is None: + return _failure( + "Structure not available for this document", + {"doc_name": doc_name}, + { + "summary": "Structure not available for this document", + "options": [ + "The document may not have been processed correctly or " + "structure extraction may have failed", + "Try processing the document again if possible", + ], + }, + "INTERNAL_ERROR", + ) + + formatted = _format_structure(copy.deepcopy(tree)) + chunks = _split_structure(formatted, _CHAR_BUDGET) + total_parts = max(1, len(chunks)) + try: + requested_part = int(part) + except (TypeError, ValueError): + requested_part = 1 + current = min(max(requested_part, 1), total_parts) + + if total_parts == 1: + return _success( + {"doc_name": doc_name, "structure": chunks[0]}, + { + "summary": "Document structure retrieved successfully.", + "options": [ + "Use get_page_content() to extract specific content from pages", + ], + }, + ) + + next_steps = ( + { + "summary": f"Showing part {current} of {total_parts}.", + "options": [ + f"Request next part with part: {current + 1}", + f"Jump to last part with part: {total_parts}", + "Proceed to get_page_content() for specific sections", + ], + } + if current < total_parts else + { + "summary": "All parts retrieved for current pagination.", + "options": [ + "Use get_page_content() to extract specific content from pages", + ], + } + ) + return _success( + { + "doc_name": doc_name, + "total_parts": total_parts, + "structure": chunks[current - 1], + "pagination": { + "part": current, + "total_parts": total_parts, + "has_more": current < total_parts, + }, + }, + next_steps, + ) + + +def _get_page_content(client, doc_name: str, pages: str, + folder_id: Optional[str] = None, + wait_for_completion: bool = False) -> tuple[dict, bool]: + if folder_id not in (None, "root"): + return _folder_unsupported("folder_id") + entry, error = _resolve_document(client, doc_name) + if error is not None: + return error + assert entry is not None + entry = _await_completion(client, entry, wait_for_completion) + if entry.get("status") != "completed": + return _not_ready_error(doc_name, entry.get("status"), + "page content retrieval", wait_for_completion) + + requested, error = _parse_page_spec(pages, doc_name) + if error is not None: + return error + assert requested is not None + + try: + page_data = client.get_ocr(entry["id"], format="page").get("result") or [] + except PageIndexAPIError as exc: + return _failure( + f"Failed to retrieve page content: {exc}", + {"doc_name": doc_name}, + { + "summary": "Unable to retrieve page content due to a service issue.", + "options": [ + "Verify the document name is correct using browse_documents()", + "Check if the document processing is complete with get_document()", + "Ensure the requested page numbers are valid", + ], + "auto_retry": "This may be a temporary issue - you can try " + "the request again", + }, + "INTERNAL_ERROR", + ) + + by_index = {item["page_index"]: item for item in page_data + if isinstance(item, dict) + and isinstance(item.get("page_index"), int)} + max_page = max(by_index, default=0) + + out_of_range = [page for page in requested if page > max_page] + valid_pages = [page for page in requested if page <= max_page] + if out_of_range and not valid_pages: + return _failure( + f"All requested pages are out of range. Document has {max_page} " + f"pages, but you requested pages: {', '.join(map(str, out_of_range))}", + { + "doc_name": doc_name, + "max_pages": max_page, + "requested_pages": _format_page_spec(out_of_range), + }, + { + "summary": "All requested pages are out of range for this document", + "options": [ + f"Request pages between 1 and {max_page}", + "Use get_document() to check document page count", + ], + }, + "INVALID_INPUT", + ) + + content = [] + included: list[int] = [] + remaining: list[int] = [] + budget = _CHAR_BUDGET + for page in valid_pages: + item = by_index.get(page) + markdown = item.get("markdown") if item else None + text = (markdown if isinstance(markdown, str) + else f"Page {page} content not available") + if not included or budget - len(text) >= 0: + content.append({"page": page, "text": text}) + included.append(page) + budget -= len(text) + else: + remaining.append(page) + + options = [ + "Use get_document_structure() to understand document organization", + "Request additional pages as needed", + ] + if remaining: + options.insert(0, f"For remaining pages, request: {_format_page_spec(remaining)}") + if out_of_range: + options.insert(0, f"Document has {max_page} pages total - request " + f"pages 1-{max_page}") + summary = ( + f"Retrieved {len(included)} pages. Pages " + f"{', '.join(map(str, out_of_range))} were out of range." + if out_of_range + else f"Returned {len(included)} of {len(requested)} requested pages " + "due to response size limits." + if remaining + else f"Successfully retrieved content for {len(content)} " + f"page{'' if len(content) == 1 else 's'}." + ) + return _success( + { + "doc_name": doc_name, + "total_pages": max_page, + "requested_pages": _format_page_spec(requested), + "returned_pages": _format_page_spec(included), + "content": content, + }, + {"summary": summary, "options": options}, + ) + + +def _remove_document(client, doc_names: list[str], + folder_id: Optional[str] = None) -> tuple[dict, bool]: + if folder_id not in (None, "root"): + return _folder_unsupported("folder_id") + if not isinstance(doc_names, list) or not doc_names: + return _failure("At least one document name is required", None, + {"summary": "No document names provided", + "options": ["Pass doc_names as a non-empty array"]}, + "INVALID_INPUT") + if len(doc_names) > 10: + return _failure("Maximum 10 documents can be deleted at once", None, + {"summary": "Too many documents in one call", + "options": ["Delete at most 10 documents per call"]}, + "INVALID_INPUT") + results = [] + for doc_name in doc_names: + entry, error = _resolve_document(client, doc_name) + if error is not None or entry is None: + results.append({"doc_name": doc_name, "status": "not_found"}) + continue + try: + client.delete_document(entry["id"]) + results.append({"doc_name": doc_name, "status": "deleted"}) + except PageIndexAPIError as exc: + results.append({"doc_name": doc_name, "status": "failed", + "error": str(exc)}) + deleted = sum(1 for item in results if item["status"] == "deleted") + return _success( + {"results": results}, + { + "summary": f"Deleted {deleted} of {len(doc_names)} document(s).", + "options": ["Use browse_documents() to review the remaining library"], + }, + ) + + +_IMPLEMENTATIONS: dict[str, Callable[..., tuple[dict, bool]]] = { + "browse_documents": _browse_documents, + "get_document": _get_document, + "get_document_structure": _get_document_structure, + "get_page_content": _get_page_content, + "remove_document": _remove_document, +} + + +def tool_names(include_management: bool = False) -> tuple[str, ...]: + return _READ_TOOLS + (_MANAGEMENT_TOOLS if include_management else ()) + + +def call_tool(client, name: str, arguments: dict[str, Any]) -> tuple[str, bool]: + """Run one contract tool; returns (envelope_json, is_error). Never raises + for tool-level failures — unexpected exceptions become error envelopes.""" + implementation = _IMPLEMENTATIONS[name] + try: + payload, is_error = implementation(client, **arguments) + except TypeError as exc: + payload, is_error = _failure( + f"Invalid arguments for {name}: {exc}", None, + {"summary": "Invalid tool arguments", + "options": [f"Check the {name}() parameter names and types"]}, + "INVALID_INPUT", + ) + except Exception as exc: # tool calls must never raise into the agent loop + payload, is_error = _failure( + f"{name} failed: {exc}", None, + {"summary": "Unexpected error while running the tool", + "options": ["Try the request again"], + "auto_retry": "This is likely a temporary issue - you can try " + "the request again"}, + "INTERNAL_ERROR", + ) + return _dumps(payload), is_error + + +# ── plain-function materialization (the `client.agent_tools()` surface) ── + +def _tool_docstring(description: str, properties: dict[str, Any]) -> str: + lines = [description, "", "Args:"] + for param, spec in properties.items(): + lines.append(f" {param}: {spec.get('description', '')}") + return "\n".join(lines) + + +def _docstring(name: str) -> str: + contract = TOOL_CONTRACT[name] + return _tool_docstring(contract["description"], + contract["schema"]["properties"]) + + +_SCHEMA_TYPE_MAP = {"string": str, "integer": int, "number": float, + "boolean": bool, "array": list, "object": dict} + + +def _annotation_for(spec: dict) -> Any: + schema_type = spec.get("type") + if isinstance(schema_type, list): + bases = [t for t in schema_type if t != "null"] + base = _SCHEMA_TYPE_MAP.get(bases[0], Any) if bases else Any + return Optional[base] if "null" in schema_type else base + return _SCHEMA_TYPE_MAP.get(schema_type, Any) + + +def _make_bridge_function(bridge, meta: dict) -> Callable[..., str]: + """One plain function for a cloud tool: real signature and docstring from + the server's schema, invocation proxied over MCP, errors contained.""" + import keyword + + name = str(meta.get("name") or "") + schema = meta.get("inputSchema") or {} + properties: dict[str, Any] = schema.get("properties") or {} + required = set(schema.get("required") or []) + + def _invoke(arguments: dict[str, Any]) -> str: + # None ≡ omitted, matching the contract's "omit if ..." semantics. + arguments = {key: value for key, value in arguments.items() + if value is not None} + try: + return bridge.call_tool(name, arguments) + except Exception as exc: + payload, _ = _failure( + f"{name} failed: {exc}", None, + {"summary": "Unexpected error while running the tool", + "options": ["Try the request again"], + "auto_retry": "This is likely a temporary issue - you can " + "try the request again"}, + "INTERNAL_ERROR", + ) + return _dumps(payload) + + params_usable = all(param.isidentifier() and not keyword.iskeyword(param) + and param != "_invoke" + for param in properties) + if not params_usable: + def proxy(**kwargs: Any) -> str: + return _invoke(kwargs) + else: + ordered = ([p for p in properties if p in required] + + [p for p in properties if p not in required]) + rendered = ", ".join( + p if p in required else f"{p}={properties[p].get('default')!r}" + for p in ordered + ) + args_literal = "{" + ", ".join(f"'{p}': {p}" for p in ordered) + "}" + namespace: dict[str, Any] = {"_invoke": _invoke} + exec(f"def _synthesized({rendered}):\n" + f" return _invoke({args_literal})", namespace) + proxy = namespace["_synthesized"] + annotations: dict[str, Any] = {} + for p in ordered: + annotation = _annotation_for(properties[p]) + if p not in required and "default" not in properties[p]: + # Absent-but-non-nullable params must admit None, or strict + # schemas force the model to always send a value. + annotation = Optional[annotation] + annotations[p] = annotation + annotations["return"] = str + proxy.__annotations__ = annotations + proxy.__name__ = proxy.__qualname__ = name or "tool" + proxy.__doc__ = _tool_docstring(meta.get("description", ""), properties) + return proxy + + +def _build_cloud_agent_tools(client, include_management: bool) -> list[Callable[..., str]]: + from .mcp_bridge import McpBridge + bridge = McpBridge( + f"{client.BASE_URL}/mcp", + {"Authorization": f"Bearer {client.api_key}"}, + ) + tools_meta = bridge.list_tools() + if not include_management: + # Plain functions have no framework permission layer, so the + # management gate lives here: only tools the server marks read-only. + filtered = [meta for meta in tools_meta + if (meta.get("annotations") or {}).get("readOnlyHint") is True] + if tools_meta and not filtered: + raise PageIndexAPIError( + "The MCP server returned tools but none are annotated " + "read-only — a server annotation regression would otherwise " + "silently disable every tool. Pass include_management=True " + "to expose the unfiltered list." + ) + tools_meta = filtered + return [_make_bridge_function(bridge, meta) for meta in tools_meta] + + +def build_agent_tools(client, include_management: bool = False) -> list[Callable[..., str]]: + """Plain synchronous functions bound to `client`. + + Cloud: one function per tool of the live cloud MCP tool set, signatures + synthesized from the server's schemas, calls proxied over MCP. Local: + the built-in contract tools over the local store. Every function returns + the JSON envelope as a string and never raises. + """ + if getattr(client, "api_key", None): + return _build_cloud_agent_tools(client, include_management) + + def browse_documents(folder_id: str = "root", recursive: bool = False, + sort: str = "time", query: Optional[str] = None, + offset: int = 0, limit: int = 10) -> str: + return call_tool(client, "browse_documents", { + "folder_id": folder_id, "recursive": recursive, "sort": sort, + "query": query, "offset": offset, "limit": limit, + })[0] + + def get_document(doc_name: str, folder_id: Optional[str] = None, + wait_for_completion: bool = False) -> str: + return call_tool(client, "get_document", { + "doc_name": doc_name, "folder_id": folder_id, + "wait_for_completion": wait_for_completion, + })[0] + + def get_document_structure(doc_name: str, folder_id: Optional[str] = None, + part: int = 1, + wait_for_completion: bool = False) -> str: + return call_tool(client, "get_document_structure", { + "doc_name": doc_name, "folder_id": folder_id, "part": part, + "wait_for_completion": wait_for_completion, + })[0] + + def get_page_content(doc_name: str, pages: str, + folder_id: Optional[str] = None, + wait_for_completion: bool = False) -> str: + return call_tool(client, "get_page_content", { + "doc_name": doc_name, "pages": pages, "folder_id": folder_id, + "wait_for_completion": wait_for_completion, + })[0] + + def remove_document(doc_names: list[str], + folder_id: Optional[str] = None) -> str: + return call_tool(client, "remove_document", { + "doc_names": doc_names, "folder_id": folder_id, + })[0] + + functions = { + "browse_documents": browse_documents, + "get_document": get_document, + "get_document_structure": get_document_structure, + "get_page_content": get_page_content, + "remove_document": remove_document, + } + tools = [] + for name in tool_names(include_management): + function = functions[name] + function.__doc__ = _docstring(name) + tools.append(function) + return tools + + +# ── agent instructions ── + +_INSTRUCTIONS_HEADER = ( + "PageIndex by Vectify AI is a document platform for uploading and " + "managing long PDFs (research papers, financial reports, legal docs, " + "textbooks, etc.)." +) + +_READING_WORKFLOW = f"""\ +READING WORKFLOW: +- For documents over {STRUCTURE_FIRST_PAGE_THRESHOLD} pages: call get_document_structure() first to locate relevant sections, then get_page_content() with targeted page ranges. +- For small documents ({STRUCTURE_FIRST_PAGE_THRESHOLD} pages or fewer): call get_page_content() directly.""" + +_TOOL_USAGE_RULES = """\ +TOOL USAGE RULES: +- Invoke a tool only when all required parameters are present or clearly inferable. Never invent placeholder values. +- If a tool returns an error, present the provided next_steps/options to the user instead of retrying blindly.""" + +_DISCOVERY = """\ +DOCUMENT DISCOVERY: +- browse_documents() — DEFAULT discovery tool, first choice for any document-related question. The bare call returns your documents. Use sort="relevance" + query for semantic ranking.""" + +_DECISION = """\ +DECISION: +- "What do I have / list / recent" → browse_documents (time) +- ANY question that needs a document to answer (including "find THE paper about Y") → browse_documents(sort="relevance", query=…)""" + +_AFTER_DISCOVERY = """\ +- Skip discovery ONLY for questions with NO possible document connection (e.g., "capital of France"). +- After discovery: 1 match or 1 clearly best match → proceed to read and answer without asking. Multiple equally relevant → ask user to pick. +- Results returned ≠ correct results. If the returned documents do not clearly match the user's intent (e.g., wrong topic, wrong time period, wrong document type), treat it the same as "not found" and continue the PERSISTENCE protocol below.""" + +_PERSISTENCE = """\ +PERSISTENCE (before concluding the target document is not in the library): +This protocol applies both when results are empty AND when results are returned but none match the user's intent. Do NOT give up after a single discovery attempt. Follow these steps in order: +1. browse_documents(sort="relevance", query=…) with the original intent +2. Rephrase the query with synonyms or alternative terms → browse_documents(sort="relevance") again +3. browse_documents(recursive=true) to flatten the library into one list — MANDATORY, must be attempted at least once before concluding "not found" +Only after ALL three steps have been tried may you conclude the document is not in the library. Do NOT fall back to general knowledge — if the user's question references their own documents, exhaust every discovery path first.""" + +AGENT_INSTRUCTIONS = "\n\n".join([ + _INSTRUCTIONS_HEADER, + _READING_WORKFLOW, + _TOOL_USAGE_RULES, + _DISCOVERY, + _DECISION, + _AFTER_DISCOVERY, + _PERSISTENCE, +]) + + +def build_agent_instructions(client, doc_id=None) -> str: + """Orchestration guidance for document QA agents; with doc_id, appends + the target documents and directs the agent to work within them.""" + if doc_id is None: + return AGENT_INSTRUCTIONS + doc_ids = [doc_id] if isinstance(doc_id, str) else list(doc_id) + if not doc_ids: + return AGENT_INSTRUCTIONS + details = [client.get_document(one_id) for one_id in doc_ids] + context = json.dumps(details, ensure_ascii=False) + if len(details) == 1: + block = ( + f"The user has specified document: {details[0].get('name')}\n" + f"Document metadata: {context}\n" + "Use this document's name to retrieve its content with " + "get_document_structure() and get_page_content()." + ) + else: + names = ", ".join(str(item.get("name")) for item in details) + block = ( + f"The user has specified documents: {names}\n" + f"Documents metadata: {context}\n" + "Use these documents' names to retrieve their content with " + "get_document_structure() and get_page_content()." + ) + return AGENT_INSTRUCTIONS + "\n\n" + block diff --git a/pageindex/client.py b/pageindex/client.py index 158c9b6f7..2a402a485 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -1,7 +1,8 @@ """PageIndex SDK client: the 0.2.x cloud surface, now with a local mode.""" from __future__ import annotations -from typing import Any, Iterator, Optional, Union +import time +from typing import Any, Callable, Iterator, Optional, Union from .errors import PageIndexAPIError @@ -126,12 +127,14 @@ def submit_document( beta_headers: Optional[list[str]] = None, folder_id: Optional[str] = None, metadata: Optional[dict] = None, + wait: bool = False, ) -> dict[str, Any]: """ Submit a PDF document for processing. Returns {'doc_id': ...}. - Cloud: uploads the file; processing is asynchronous — poll - ``is_retrieval_ready(doc_id)`` before retrieving. + Cloud: uploads the file; processing is asynchronous. Pass + ``wait=True`` to block until the document is ready, or poll + ``get_document(doc_id)['status']`` yourself. Local: indexes the document in this call (it blocks while your LLM builds the tree — minutes for a standard index of a long document), @@ -151,14 +154,53 @@ def submit_document( metadata (dict, optional): Your own JSON-serializable tags for the document; returned in get_tree/get_ocr responses and list_documents entries (both modes). + wait (bool): Return only once the document is ready for use. + Cloud: polls status until "completed" (raises on "failed" or + after 30 minutes). Local: indexing is synchronous already, so + this changes nothing. Leave False to submit many documents + concurrently and poll afterwards. Returns: dict: {'doc_id': ...} """ - return self._api.submit_document( + result = self._api.submit_document( file_path=file_path, mode=mode, beta_headers=beta_headers, folder_id=folder_id, metadata=metadata, ) + if wait: + self._wait_until_ready(result["doc_id"]) + return result + + def _wait_until_ready(self, doc_id: str, timeout: float = 1800.0) -> None: + interval = 2.0 + deadline = time.monotonic() + timeout + poll_failures = 0 + while True: + try: + status = self.get_document(doc_id).get("status") + poll_failures = 0 + except PageIndexAPIError: + # Tolerate transient poll failures; a 30-minute wait should + # not die on one 502. + poll_failures += 1 + if poll_failures >= 3: + raise + status = None + if status == "completed": + return + if status == "failed": + raise PageIndexAPIError( + f"Document processing failed (doc_id: {doc_id})." + ) + if time.monotonic() >= deadline: + raise PageIndexAPIError( + f"Timed out after {int(timeout)}s waiting for document " + f"processing (doc_id: {doc_id}, last status: {status}). " + "Processing continues in the cloud — poll " + "get_document(doc_id) for status." + ) + time.sleep(interval) + interval = min(interval * 1.5, 15.0) # ---------- OCR FUNCTIONALITY ---------- @@ -365,6 +407,104 @@ def list_documents( """ return self._api.list_documents(limit=limit, offset=offset, folder_id=folder_id) + # ---------- AGENT INTEGRATION ---------- + + def agent_tools(self, include_management: bool = False) -> list[Callable[..., str]]: + """ + Plain functions for any agent framework (LangChain, PydanticAI, ...). + For the OpenAI / Claude Agent SDKs, prefer ``as_openai_tools()`` / + ``as_claude_mcp()``. + + Cloud: the full cloud tool set, discovered live from the PageIndex + MCP server when this method is called — one function per tool, + signature and docstring synthesized from the server's schemas, calls + executed from your process over MCP. Raises PageIndexAPIError if the + server cannot be reached. Local: the built-in tools over the local + store (``browse_documents``, ``get_document``, + ``get_document_structure``, ``get_page_content``). + + Each function takes JSON-serializable arguments, returns a JSON + string, and reports failures inside that JSON instead of raising. + + Args: + include_management (bool): Also expose tools that modify the + library. Local: adds ``remove_document``. Cloud: by default + only tools the server marks read-only are exposed; True + exposes the server's complete list (upload, delete, ...). + """ + from .agent_tools import build_agent_tools + return build_agent_tools(self, include_management) + + def as_openai_tools(self, include_management: bool = False, + hosted: bool = False) -> list: + """ + Tools for the OpenAI Agents SDK — pass to ``Agent(tools=...)``. + + Cloud (default): the full live read tool set (search, folders, + images — as enabled for your key) as plain function tools, + discovered from the PageIndex MCP server and executed from your + process — works with any model backend. Pass ``hosted=True`` to + hand the connection to OpenAI instead: one hosted MCP tool, tool + calls executed server-side (lowest latency; requires an + OpenAI-hosted model on the Responses API). + + Local: the in-process tools, any model backend; ``hosted`` does + not apply. (The framework's own ``MCPServerStreamableHttp`` + against ``{BASE_URL}/mcp`` is the async-native alternative for + its ``mcp_servers=`` slot.) + + Requires ``openai-agents`` (``pip install 'pageindex[openai]'``), + imported only when this method is called. + + Args: + include_management (bool): Also expose tools that modify the + library (delete, upload). Default off: the cloud default + serves only server-annotated read-only tools, and + ``hosted=True`` routes non-read-only tools through the + Responses API approval flow instead. + hosted (bool): Cloud only — hand the MCP connection to OpenAI + for server-side tool execution (OpenAI models only). + """ + from .integrations.openai_agents import build_openai_tools + return build_openai_tools(self, include_management, hosted) + + def as_claude_mcp(self, include_management: bool = False): + """ + ``mcp_servers`` entry for the Claude Agent SDK. + + Cloud: returns the remote PageIndex MCP config — the framework + connects to api.pageindex.ai/mcp directly and discovers the full + cloud tool set. ``include_management`` has no effect there; gate + destructive tools with the framework's permission layer (e.g. list + read tools in ``allowed_tools`` instead of the ``*`` wildcard, or + add ``disallowed_tools=["mcp__pageindex__remove_document"]``). + Local: returns an in-process SDK MCP server exposing the agent + tools (requires ``claude-agent-sdk``; + ``pip install 'pageindex[claude]'``). + + Usage:: + + options = ClaudeAgentOptions( + mcp_servers={"pageindex": client.as_claude_mcp()}, + allowed_tools=["mcp__pageindex__*"], + ) + """ + from .integrations.claude_agent_sdk import build_claude_mcp + return build_claude_mcp(self, include_management) + + def agent_instructions(self, doc_id: Optional[Union[str, list[str]]] = None) -> str: + """ + Orchestration guidance for document QA agents — pass as the agent's + system prompt (or append to your own). + + With ``doc_id`` (str or list, same shape as ``chat_completions``), + appends the target documents' names and metadata and directs the + agent to work within them. Raises PageIndexAPIError if a doc_id does + not exist. + """ + from .agent_tools import build_agent_instructions + return build_agent_instructions(self, doc_id) + # ---------- FOLDER MANAGEMENT ---------- def create_folder( diff --git a/pageindex/integrations/__init__.py b/pageindex/integrations/__init__.py new file mode 100644 index 000000000..e42ccf64c --- /dev/null +++ b/pageindex/integrations/__init__.py @@ -0,0 +1,5 @@ +"""Framework adapters for the agent tools layer. + +These modules import their target frameworks lazily, at call time — the +frameworks are never required to install or import pageindex. +""" diff --git a/pageindex/integrations/claude_agent_sdk.py b/pageindex/integrations/claude_agent_sdk.py new file mode 100644 index 000000000..d0e2316ae --- /dev/null +++ b/pageindex/integrations/claude_agent_sdk.py @@ -0,0 +1,67 @@ +"""Claude Agent SDK adapter: one value for the mcp_servers slot. + +Cloud clients get the remote PageIndex MCP config (the framework connects +directly and discovers the full cloud tool set); local clients get an +in-process SDK MCP server over the same tool contract. +""" +from __future__ import annotations + +import asyncio +from typing import Any + +from ..errors import PageIndexAPIError + + +def _sdk_version() -> str: + try: + from importlib.metadata import version + return version("pageindex") + except Exception: + return "0.0.0" + + +def build_claude_mcp(client, include_management: bool = False): + if getattr(client, "api_key", None): + return { + "type": "http", + "url": f"{client.BASE_URL}/mcp", + "headers": {"Authorization": f"Bearer {client.api_key}"}, + } + + try: + from claude_agent_sdk import create_sdk_mcp_server, tool + except ImportError as exc: + raise PageIndexAPIError( + "as_claude_mcp in local mode requires the Claude Agent SDK — " + "pip install claude-agent-sdk (or pip install 'pageindex[claude]')." + ) from exc + from ..agent_tools import TOOL_CONTRACT, call_tool, tool_names + + def make_handler(name: str): + async def handler(arguments: dict[str, Any]) -> dict[str, Any]: + text, is_error = await asyncio.to_thread( + call_tool, client, name, arguments or {} + ) + result: dict[str, Any] = {"content": [{"type": "text", "text": text}]} + if is_error: + result["is_error"] = True + return result + return handler + + def tool_kwargs(name: str) -> dict: + annotations = TOOL_CONTRACT[name].get("annotations") + if not annotations: + return {} + try: + from claude_agent_sdk import ToolAnnotations + except ImportError: + return {} + return {"annotations": ToolAnnotations(**annotations)} + + tools = [ + tool(name, TOOL_CONTRACT[name]["description"], + TOOL_CONTRACT[name]["schema"], **tool_kwargs(name))(make_handler(name)) + for name in tool_names(include_management) + ] + return create_sdk_mcp_server(name="pageindex", version=_sdk_version(), + tools=tools) diff --git a/pageindex/integrations/openai_agents.py b/pageindex/integrations/openai_agents.py new file mode 100644 index 000000000..91ed6ebe3 --- /dev/null +++ b/pageindex/integrations/openai_agents.py @@ -0,0 +1,36 @@ +"""OpenAI Agents SDK adapter for the Agent(tools=...) slot. + +Cloud clients get one hosted MCP tool (the model connects to the PageIndex +cloud MCP server from OpenAI's side and discovers the full cloud tool set); +local clients get the in-process tools wrapped as FunctionTools. +""" +from __future__ import annotations + +from ..errors import PageIndexAPIError + + +def build_openai_tools(client, include_management: bool = False, + hosted: bool = False) -> list: + try: + from agents import HostedMCPTool, function_tool + except ImportError as exc: + raise PageIndexAPIError( + "as_openai_tools requires the OpenAI Agents SDK — " + "pip install openai-agents (or pip install 'pageindex[openai]')." + ) from exc + if getattr(client, "api_key", None) and hosted: + # Same gate as the in-process path, enforced by OpenAI: tools the + # server annotates read-only run freely, everything else goes + # through the Responses API approval flow. + require_approval = ("never" if include_management + else {"never": {"read_only": True}}) + return [HostedMCPTool(tool_config={ + "type": "mcp", + "server_label": "pageindex", + "server_url": f"{client.BASE_URL}/mcp", + "headers": {"Authorization": f"Bearer {client.api_key}"}, + "require_approval": require_approval, + })] + from ..agent_tools import build_agent_tools + return [function_tool(tool) + for tool in build_agent_tools(client, include_management)] diff --git a/pageindex/mcp_bridge.py b/pageindex/mcp_bridge.py new file mode 100644 index 000000000..d144cee92 --- /dev/null +++ b/pageindex/mcp_bridge.py @@ -0,0 +1,181 @@ +"""Minimal MCP client (streamable HTTP) for the PageIndex cloud MCP server. + +Backs the cloud branch of ``client.agent_tools()``: ``tools/list`` discovers +the live tool set, ``tools/call`` executes a tool. Synchronous, requests-only. +Works against both stateful and stateless servers: a session id returned by +``initialize`` is echoed back, and a request rejected after session expiry +re-initializes once and retries. +""" +from __future__ import annotations + +import json +import threading +from typing import Any, Optional + +import requests + +from .errors import PageIndexAPIError + +_PROTOCOL_VERSION = "2025-06-18" +_TIMEOUT = (10, 240) # tools may wait server-side (wait_for_completion: 3 min) + + +def _sdk_version() -> str: + try: + from importlib.metadata import version + return version("pageindex") + except Exception: + return "0.0.0" + + +def _parse_sse(text: str) -> list[dict]: + """JSON-RPC messages out of a text/event-stream body.""" + messages = [] + text = text.replace("\r\n", "\n").replace("\r", "\n") + for block in text.split("\n\n"): + data_lines = [line[5:].removeprefix(" ") for line in block.splitlines() + if line.startswith("data:")] + if not data_lines: + continue + try: + messages.append(json.loads("\n".join(data_lines))) + except ValueError: + continue + return messages + + +class McpBridge: + def __init__(self, url: str, headers: dict[str, str]): + self._url = url + self._auth_headers = dict(headers) + self._session_id: Optional[str] = None + self._protocol_version: Optional[str] = None + self._initialized = False + self._lock = threading.Lock() + self._next_id = 0 + + # ── JSON-RPC over streamable HTTP ── + + def _post(self, payload: dict) -> requests.Response: + headers = { + "Content-Type": "application/json", + "Accept": "application/json, text/event-stream", + **self._auth_headers, + } + if self._session_id: + headers["Mcp-Session-Id"] = self._session_id + if self._protocol_version: + headers["MCP-Protocol-Version"] = self._protocol_version + try: + return requests.post(self._url, json=payload, headers=headers, + timeout=_TIMEOUT) + except requests.RequestException as exc: + raise PageIndexAPIError( + f"Could not reach the PageIndex MCP server: {exc}" + ) from exc + + def _extract_result(self, response: requests.Response, request_id: int) -> Any: + content_type = response.headers.get("Content-Type", "") + if "text/event-stream" in content_type: + # SSE is UTF-8 by spec; requests guesses latin-1 for charset-less + # text/* and would mojibake every non-ASCII character. + messages = _parse_sse(response.content.decode("utf-8", + errors="replace")) + else: + try: + messages = [response.json()] + except ValueError as exc: + raise PageIndexAPIError( + f"MCP server returned a non-JSON response " + f"(HTTP {response.status_code})." + ) from exc + reply = next((m for m in messages if m.get("id") == request_id), + next((m for m in messages + if "result" in m or "error" in m), None)) + if reply is None: + raise PageIndexAPIError("MCP server response contained no reply.") + if "error" in reply: + error = reply["error"] or {} + raise PageIndexAPIError( + f"MCP error {error.get('code')}: {error.get('message')}" + ) + return reply.get("result") + + def _request(self, method: str, params: Optional[dict] = None, + _retry: bool = True) -> Any: + self._ensure_initialized() + with self._lock: + self._next_id += 1 + request_id = self._next_id + payload: dict[str, Any] = {"jsonrpc": "2.0", "id": request_id, + "method": method} + if params is not None: + payload["params"] = params + response = self._post(payload) + if response.status_code in (400, 404) and self._initialized and _retry: + # Session expired (stateful servers): start over, retry once. + with self._lock: + self._initialized = False + self._session_id = None + return self._request(method, params, _retry=False) + if response.status_code >= 400: + raise PageIndexAPIError( + f"MCP request failed: HTTP {response.status_code} " + f"({response.text[:200]})" + ) + return self._extract_result(response, request_id) + + def _ensure_initialized(self) -> None: + with self._lock: + if self._initialized: + return + self._next_id += 1 + request_id = self._next_id + response = self._post({ + "jsonrpc": "2.0", "id": request_id, "method": "initialize", + "params": { + "protocolVersion": _PROTOCOL_VERSION, + "capabilities": {}, + "clientInfo": {"name": "pageindex-python-sdk", + "version": _sdk_version()}, + }, + }) + if response.status_code >= 400: + raise PageIndexAPIError( + f"Could not connect to the PageIndex MCP server: HTTP " + f"{response.status_code} ({response.text[:200]}). Check " + "your API key." + ) + result = self._extract_result(response, request_id) or {} + self._session_id = response.headers.get("Mcp-Session-Id") + self._protocol_version = result.get("protocolVersion", + _PROTOCOL_VERSION) + self._initialized = True + try: + self._post({"jsonrpc": "2.0", + "method": "notifications/initialized"}) + except PageIndexAPIError: + pass # advisory; a server that required it fails the next request + + # ── public surface ── + + def list_tools(self) -> list[dict]: + tools: list[dict] = [] + cursor: Optional[str] = None + while True: + params = {"cursor": cursor} if cursor else {} + result = self._request("tools/list", params) or {} + tools.extend(result.get("tools") or []) + cursor = result.get("nextCursor") + if not cursor: + return tools + + def call_tool(self, name: str, arguments: dict[str, Any]) -> str: + result = self._request("tools/call", + {"name": name, "arguments": arguments}) or {} + blocks = result.get("content") or [] + texts = [block.get("text", "") for block in blocks + if isinstance(block, dict) and block.get("type") == "text"] + if len(texts) == len(blocks): + return "\n".join(texts) + return json.dumps(blocks, ensure_ascii=False) diff --git a/pyproject.toml b/pyproject.toml index deac66be3..65f68646b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,6 +38,14 @@ sortedcontainers = ">=2.4.0" regex = ">=2024.0.0" python-dotenv = ">=1.0.0" pyyaml = ">=6.0" +claude-agent-sdk = { version = ">=0.1.0", optional = true } +# 0.8.0 offloads sync tools to a thread; older versions run them inline and +# a blocking bridge call would freeze the agent event loop. +openai-agents = { version = ">=0.8.0", optional = true } + +[tool.poetry.extras] +claude = ["claude-agent-sdk"] +openai = ["openai-agents"] [tool.poetry.group.dev.dependencies] pytest = ">=7.0" diff --git a/tests/data/cloud_mcp_contract.json b/tests/data/cloud_mcp_contract.json new file mode 100644 index 000000000..71743aee2 --- /dev/null +++ b/tests/data/cloud_mcp_contract.json @@ -0,0 +1,197 @@ +{ + "_provenance": "Frozen copy of the PageIndex cloud MCP server's tool contract (names, input schemas, descriptions, and annotations as served via tools/list). The parity test asserts pageindex.agent_tools.TOOL_CONTRACT matches this file; update both together only when the cloud contract changes.", + "tools": { + "browse_documents": { + "annotations": { + "readOnlyHint": true, + "openWorldHint": false + }, + "description": "Primary document retrieval tool. After orienting with get_folder_structure() (when available), use this for all document-related questions. The bare call returns root-level sub-folders and documents; pass folder_id to drill into a sub-folder level by level. Use sort=\"relevance\" + query for semantic ranking. Do NOT jump to search_documents() first — it is an escalation path, only after browse_documents(sort=\"relevance\") has failed.", + "schema": { + "type": "object", + "properties": { + "folder_id": { + "type": "string", + "default": "root", + "description": "Folder scope (default \"root\"). Pass a specific folder ID to scope into that folder, or \"root\" to reference the library root. The read-only \"shared-with-me\" and \"following\" folders live at the library root — pass one of those ids to browse them. Copy any folder_id verbatim from a browse/tree response, never construct one. Combine with `recursive` to control breadth." + }, + "recursive": { + "type": "boolean", + "default": false, + "description": "Whether to include documents from descendant folders. When false (default), returns the direct contents of folder_id along with its sub-folders — prefer this for level-by-level exploration so you retain folder hierarchy context. When true, flattens all descendant documents into one list and omits sub-folders — use only when a non-recursive browse of the target folder returned no relevant results and you need to widen the scope, or the user explicitly requests a flat listing." + }, + "sort": { + "type": "string", + "enum": [ + "time", + "relevance" + ], + "default": "time", + "description": "Sort order. \"time\" (default) sorts by upload date (newest first); \"relevance\" orders documents by semantic relevance to `query`. Relevance also works inside the read-only shared folders — pass their folder_id — but at the library root it ranks only your own documents." + }, + "query": { + "type": "string", + "description": "Search query for relevance ranking. Required when sort=\"relevance\"; must be omitted when sort=\"time\"." + }, + "offset": { + "type": "integer", + "minimum": 0, + "default": 0, + "description": "Zero-based pagination offset. Pass the value of `next_offset` from the previous response to fetch the next page." + }, + "limit": { + "type": "number", + "minimum": 1, + "maximum": 50, + "default": 10, + "description": "Number of documents to return per page (1-50, default 10)" + } + }, + "required": [] + } + }, + "get_document": { + "annotations": { + "readOnlyHint": true, + "openWorldHint": false + }, + "description": "Check a document's processing status and metadata. `status` is one of \"pending\", \"queued\", \"processing\", \"completed\", or \"failed\" — call this before `get_document_structure()` or `get_page_content()` to confirm the document is ready.", + "schema": { + "type": "object", + "properties": { + "doc_name": { + "type": "string", + "minLength": 1, + "description": "Copy the `name` field verbatim from a browse_documents() or search_documents() response (case-sensitive, include extension). Example: \"Q3 Report.pdf\". If the response shows two documents with the same name, pass `folder_id` alongside to disambiguate." + }, + "folder_id": { + "type": [ + "string", + "null" + ], + "description": "Disambiguator for same-name documents. Copy the `folder_id` from the intended browse/search result; use \"root\" for root-level documents, or \"shared-with-me\"/\"following\" for the read-only folders at the library root; omit if `doc_name` is unique. Copy any folder_id verbatim from a browse_documents()/get_folder_structure() response, never construct one." + }, + "wait_for_completion": { + "type": "boolean", + "default": false, + "description": "If true and document is processing, automatically wait up to 3 minutes until completed. Reduces repeated tool calls." + } + }, + "required": [ + "doc_name" + ] + } + }, + "get_document_structure": { + "annotations": { + "readOnlyHint": true, + "openWorldHint": false + }, + "description": "Extract a document's hierarchical outline (headers, sections, page references). REQUIRED for documents over 20 pages — call this first to locate relevant sections, then pass their page numbers to `get_page_content()`. Use the `part` parameter to iterate large outlines until `pagination.has_more` is false.", + "schema": { + "type": "object", + "properties": { + "doc_name": { + "type": "string", + "minLength": 1, + "description": "Copy the `name` field verbatim from a browse_documents() or search_documents() response (case-sensitive, include extension). Example: \"Q3 Report.pdf\". If the response shows two documents with the same name, pass `folder_id` alongside to disambiguate." + }, + "folder_id": { + "type": [ + "string", + "null" + ], + "description": "Disambiguator for same-name documents. Copy the `folder_id` from the intended browse/search result; use \"root\" for root-level documents, or \"shared-with-me\"/\"following\" for the read-only folders at the library root; omit if `doc_name` is unique. Copy any folder_id verbatim from a browse_documents()/get_folder_structure() response, never construct one." + }, + "part": { + "type": "integer", + "minimum": 1, + "default": 1, + "description": "Part number for pagination (1-based, default 1). For large outlines, increment until the response's `pagination.has_more` becomes false." + }, + "wait_for_completion": { + "type": "boolean", + "default": false, + "description": "If true and document is processing, automatically wait up to 3 minutes until completed. Reduces repeated tool calls." + } + }, + "required": [ + "doc_name" + ] + } + }, + "get_page_content": { + "annotations": { + "readOnlyHint": true, + "openWorldHint": false + }, + "description": "Extract page content from a processed document. Use tight, targeted page ranges — never the whole document at once. For documents over 20 pages, call `get_document_structure()` first to pick relevant sections. Embedded image paths in the response feed into `get_document_image()`.", + "schema": { + "type": "object", + "properties": { + "doc_name": { + "type": "string", + "minLength": 1, + "description": "Copy the `name` field verbatim from a browse_documents() or search_documents() response (case-sensitive, include extension). Example: \"Q3 Report.pdf\". If the response shows two documents with the same name, pass `folder_id` alongside to disambiguate." + }, + "folder_id": { + "type": [ + "string", + "null" + ], + "description": "Disambiguator for same-name documents. Copy the `folder_id` from the intended browse/search result; use \"root\" for root-level documents, or \"shared-with-me\"/\"following\" for the read-only folders at the library root; omit if `doc_name` is unique. Copy any folder_id verbatim from a browse_documents()/get_folder_structure() response, never construct one." + }, + "pages": { + "type": "string", + "minLength": 1, + "pattern": "^(\\d+(-\\d+)?)(,\\s*\\d+(-\\d+)?)*$", + "description": "Page specification: \"5\", \"3,7,10\", \"5-10\", or \"1-3,7,9-12\"" + }, + "wait_for_completion": { + "type": "boolean", + "default": false, + "description": "If true and document is processing, automatically wait up to 3 minutes until completed. Reduces repeated tool calls." + } + }, + "required": [ + "doc_name", + "pages" + ] + } + }, + "remove_document": { + "annotations": { + "readOnlyHint": false, + "destructiveHint": true, + "idempotentHint": true, + "openWorldHint": false + }, + "description": "Permanently delete documents and all associated data. Only invoke when the user explicitly names the documents AND confirms deletion. Returns `results` — one entry per requested document: `{ doc_name, status: \"deleted\" | \"not_found\" | \"failed\", error? }`. Inspect each entry for per-document failures. This action is irreversible.", + "schema": { + "type": "object", + "properties": { + "doc_names": { + "type": "array", + "items": { + "type": "string", + "minLength": 1 + }, + "minItems": 1, + "maxItems": 10, + "description": "Array of document names to delete. Each name must be copied verbatim from the `name` field of a browse_documents() or search_documents() response (case-sensitive, include extension). Example: [\"Q3 Report.pdf\", \"draft.pdf\"]. Max 10 per call." + }, + "folder_id": { + "type": [ + "string", + "null" + ], + "description": "Disambiguator for same-name documents. Copy the `folder_id` from the intended browse/search result; use \"root\" for root-level documents, or \"shared-with-me\"/\"following\" for the read-only folders at the library root; omit if `doc_name` is unique. Copy any folder_id verbatim from a browse_documents()/get_folder_structure() response, never construct one." + } + }, + "required": [ + "doc_names" + ] + } + } + } +} diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py new file mode 100644 index 000000000..289c2dada --- /dev/null +++ b/tests/test_agent_tools.py @@ -0,0 +1,879 @@ +"""Agent tools layer: cloud-contract parity and behavior against a seeded +local store (no LLM calls; one live parity test gated on PAGEINDEX_API_KEY).""" +import json +import os +import sys +from pathlib import Path + +import pytest + +import pageindex.client as client_module +from pageindex import PageIndexAPIError, PageIndexCloudClient, PageIndexLocalClient +from pageindex.agent_tools import ( + AGENT_INSTRUCTIONS, + TOOL_CONTRACT, + call_tool, + tool_names, +) +from pageindex.local_store import DocStore + +SNAPSHOT_PATH = Path(__file__).parent / "data" / "cloud_mcp_contract.json" + + +def seed_doc(storage_path, doc_id, name, *, created_at="2026-08-01T10:00:00.123000", + description="A test document", metadata=None, tree=None, pages=None, + page_num=None): + pages = pages if pages is not None else [ + {"page_index": 1, "markdown": "Page one text about apples"}, + {"page_index": 2, "markdown": "Page two text about bananas"}, + ] + tree = tree if tree is not None else [{ + "title": "Doc", "node_id": "0000", "start_index": 1, "end_index": 2, + "summary": "root summary", "text": "ROOT TEXT", + "nodes": [ + {"title": "Intro", "node_id": "0001", "start_index": 1, + "end_index": 1, "summary": "intro summary", "text": "INTRO TEXT"}, + {"title": "Body", "node_id": "0002", "start_index": 2, + "end_index": 2, "summary": "body summary", "text": "BODY TEXT"}, + ], + }] + meta = { + "id": doc_id, "name": name, "description": description, + "status": "completed", "createdAt": created_at, + "pageNum": page_num if page_num is not None else len(pages), + "folderId": None, "metadata": metadata, "mode": "standard", + } + DocStore(storage_path).save_document(doc_id, meta, tree, pages) + return doc_id + + +@pytest.fixture +def store_path(tmp_path): + return str(tmp_path / "store") + + +@pytest.fixture +def client(store_path): + return PageIndexLocalClient(storage_path=store_path) + + +def run(client, name, **arguments): + text, is_error = call_tool(client, name, arguments) + return json.loads(text), is_error + + +# ── contract parity ── + +def test_contract_matches_snapshot(): + snapshot = json.loads(SNAPSHOT_PATH.read_text(encoding="utf-8")) + assert snapshot["tools"] == TOOL_CONTRACT + + +def test_tool_surface_and_docstrings(client): + tools = client.agent_tools() + assert [tool.__name__ for tool in tools] == list(tool_names()) + with_management = client.agent_tools(include_management=True) + assert [tool.__name__ for tool in with_management][-1] == "remove_document" + for tool in tools: + contract = TOOL_CONTRACT[tool.__name__] + assert tool.__doc__.startswith(contract["description"]) + for param in contract["schema"]["properties"]: + assert param in tool.__doc__ + + +# ── browse_documents ── + +def test_browse_documents_shape(client, store_path): + seed_doc(store_path, "pi-a", "older.pdf", created_at="2026-08-01T10:00:00.123000") + seed_doc(store_path, "pi-b", "newer.pdf", created_at="2026-08-02T10:00:00.456000", + metadata={"team": "research", "year": 2026, "nested": {"x": 1}}) + payload, is_error = run(client, "browse_documents") + assert not is_error + assert payload["success"] is True + assert payload["folders"] == [] + assert payload["has_more"] is False + assert payload["next_offset"] is None + names = [doc["name"] for doc in payload["documents"]] + assert names == ["newer.pdf", "older.pdf"] + newer = payload["documents"][0] + assert newer["status"] == "completed" + assert newer["created_at"] == "2026-08-02T10:00:00.456Z" + assert newer["metadata"] == {"team": "research", "year": 2026} + assert "folder_id" not in newer + assert "next_steps" in payload + + flat, _ = run(client, "browse_documents", recursive=True) + assert "folders" not in flat + + +def test_browse_documents_pagination(client, store_path): + for index in range(3): + seed_doc(store_path, f"pi-{index}", f"doc{index}.pdf", + created_at=f"2026-08-0{index + 1}T10:00:00.000000") + first, _ = run(client, "browse_documents", limit=2) + assert [d["name"] for d in first["documents"]] == ["doc2.pdf", "doc1.pdf"] + assert first["has_more"] is True and first["next_offset"] == 2 + second, _ = run(client, "browse_documents", limit=2, offset=2) + assert [d["name"] for d in second["documents"]] == ["doc0.pdf"] + assert second["has_more"] is False + + +def test_browse_documents_relevance(client, store_path): + seed_doc(store_path, "pi-a", "annual-report.pdf", + description="Financial results for the year") + seed_doc(store_path, "pi-b", "attention.pdf", + description="Transformers and attention mechanisms") + payload, is_error = run(client, "browse_documents", sort="relevance", + query="attention transformers") + assert not is_error + assert [d["name"] for d in payload["documents"]] == ["attention.pdf"] + assert payload["sort"] == "relevance" + + missing_query, is_error = run(client, "browse_documents", sort="relevance") + assert is_error and missing_query["errorCode"] == "INVALID_INPUT" + stray_query, is_error = run(client, "browse_documents", query="x") + assert is_error and "relevance" in stray_query["error"] + + +def test_browse_documents_empty_and_folder_error(client): + payload, is_error = run(client, "browse_documents") + assert not is_error + assert payload["documents"] == [] + assert "submit_document" in json.dumps(payload) + + folder, is_error = run(client, "browse_documents", folder_id="folder-123") + assert is_error and folder["errorCode"] == "INVALID_INPUT" + + +# ── get_document ── + +def test_get_document(client, store_path): + seed_doc(store_path, "pi-a", "report.pdf", metadata={"team": "research"}) + payload, is_error = run(client, "get_document", doc_name="report.pdf") + assert not is_error + assert payload["name"] == "report.pdf" + assert payload["status"] == "completed" + assert payload["page_count"] == 2 + assert payload["folder_id"] is None + assert payload["created_at"].endswith("Z") + assert payload["metadata"] == {"team": "research"} + assert any("short document" in option + for option in payload["next_steps"]["options"]) + + +def test_get_document_not_found_suggests_similar(client, store_path): + seed_doc(store_path, "pi-a", "annual-report.pdf") + payload, is_error = run(client, "get_document", doc_name="anual-report.pdf") + assert is_error + assert payload["errorCode"] == "NOT_FOUND" + assert "annual-report.pdf" in payload["similar_files"] + assert "Did you mean" in payload["error"] + + +def test_get_document_duplicate_names_resolve_newest(client, store_path): + seed_doc(store_path, "pi-old", "same.pdf", description="old copy", + created_at="2026-08-01T10:00:00.000000") + seed_doc(store_path, "pi-new", "same.pdf", description="new copy", + created_at="2026-08-02T10:00:00.000000") + payload, _ = run(client, "get_document", doc_name="same.pdf") + assert payload["description"] == "new copy" + + +# ── get_document_structure ── + +def test_structure_strips_text_and_orders_keys(client, store_path): + seed_doc(store_path, "pi-a", "report.pdf") + payload, is_error = run(client, "get_document_structure", doc_name="report.pdf") + assert not is_error + assert payload["doc_name"] == "report.pdf" + assert "pagination" not in payload and "total_parts" not in payload + serialized = json.dumps(payload["structure"]) + assert "ROOT TEXT" not in serialized and "INTRO TEXT" not in serialized + # Cloud structure node shape: start_index/end_index/summary (live-verified). + root = payload["structure"][0] + assert list(root)[:4] == ["title", "node_id", "start_index", "end_index"] + assert root["summary"] == "root summary" + assert (root["start_index"], root["end_index"]) == (1, 2) + assert root["nodes"][0]["summary"] == "intro summary" + assert root["nodes"][0]["end_index"] == 1 + + +def test_structure_multipart_pagination(client, store_path): + big_tree = [{ + "title": f"Chapter {index}", "node_id": f"{index:04d}", + "start_index": index + 1, "end_index": index + 1, + "summary": "s" * 4000, "text": "T", + } for index in range(60)] + seed_doc(store_path, "pi-big", "big.pdf", tree=big_tree, + pages=[{"page_index": 1, "markdown": "x"}]) + first, _ = run(client, "get_document_structure", doc_name="big.pdf") + assert first["total_parts"] > 1 + assert first["pagination"] == { + "part": 1, "total_parts": first["total_parts"], "has_more": True, + } + titles = [] + for part in range(1, first["total_parts"] + 1): + payload, _ = run(client, "get_document_structure", doc_name="big.pdf", + part=part) + chunk = payload["structure"] + nodes = chunk if isinstance(chunk, list) else [chunk] + titles.extend(node["title"] for node in nodes) + assert payload["pagination"]["has_more"] == (part < first["total_parts"]) + assert titles == [f"Chapter {index}" for index in range(60)] + + clamped, _ = run(client, "get_document_structure", doc_name="big.pdf", + part=999) + assert clamped["pagination"]["part"] == first["total_parts"] + + +# ── get_page_content ── + +def test_page_content(client, store_path): + seed_doc(store_path, "pi-a", "report.pdf") + payload, is_error = run(client, "get_page_content", doc_name="report.pdf", + pages="1-2") + assert not is_error + assert payload["total_pages"] == 2 + assert payload["requested_pages"] == "1-2" + assert payload["returned_pages"] == "1-2" + assert payload["content"] == [ + {"page": 1, "text": "Page one text about apples"}, + {"page": 2, "text": "Page two text about bananas"}, + ] + + +def test_page_content_out_of_range(client, store_path): + seed_doc(store_path, "pi-a", "report.pdf") + mixed, is_error = run(client, "get_page_content", doc_name="report.pdf", + pages="1,99") + assert not is_error + assert mixed["returned_pages"] == "1" + assert "out of range" in mixed["next_steps"]["summary"] + + all_out, is_error = run(client, "get_page_content", doc_name="report.pdf", + pages="99") + assert is_error and all_out["errorCode"] == "INVALID_INPUT" + assert all_out["max_pages"] == 2 + + +@pytest.mark.parametrize("bad_spec", ["abc", "5-3", "1,,2", "-3", ""]) +def test_page_content_invalid_spec(client, store_path, bad_spec): + seed_doc(store_path, "pi-a", "report.pdf") + payload, is_error = run(client, "get_page_content", doc_name="report.pdf", + pages=bad_spec) + assert is_error and payload["errorCode"] == "INVALID_INPUT" + + +def test_page_content_zero_page_rejected(client, store_path): + seed_doc(store_path, "pi-a", "report.pdf") + payload, is_error = run(client, "get_page_content", doc_name="report.pdf", + pages="0") + assert is_error + assert "positive integers" in payload["error"] + + +def test_page_content_preserves_blank_pages(client, store_path): + pages = [ + {"page_index": 1, "markdown": ""}, + {"page_index": 2, "markdown": "content"}, + ] + seed_doc(store_path, "pi-a", "blanks.pdf", pages=pages) + payload, is_error = run(client, "get_page_content", doc_name="blanks.pdf", + pages="1-2") + assert not is_error + assert payload["content"][0] == {"page": 1, "text": ""} + assert payload["content"][1] == {"page": 2, "text": "content"} + + +def test_created_at_accepts_z_suffixed_input(client, store_path): + seed_doc(store_path, "pi-a", "cloudlike.pdf", + created_at="2026-08-01T10:00:00.123Z") + payload, _ = run(client, "browse_documents") + assert payload["documents"][0]["created_at"] == "2026-08-01T10:00:00.123Z" + + +def test_page_content_char_budget(client, store_path): + pages = [ + {"page_index": 1, "markdown": "x" * 96_000}, + {"page_index": 2, "markdown": "short"}, + ] + seed_doc(store_path, "pi-a", "huge.pdf", pages=pages) + payload, is_error = run(client, "get_page_content", doc_name="huge.pdf", + pages="1-2") + assert not is_error + assert payload["returned_pages"] == "1" + assert any("For remaining pages, request: 2" in option + for option in payload["next_steps"]["options"]) + + +# ── remove_document (management-gated) ── + +def test_remove_document(client, store_path): + seed_doc(store_path, "pi-a", "report.pdf") + payload, is_error = run(client, "remove_document", + doc_names=["report.pdf", "ghost.pdf"]) + assert not is_error + assert payload["results"] == [ + {"doc_name": "report.pdf", "status": "deleted"}, + {"doc_name": "ghost.pdf", "status": "not_found"}, + ] + assert client.list_documents()["total"] == 0 + + +def test_management_tools_hidden_by_default(client): + assert "remove_document" not in [t.__name__ for t in client.agent_tools()] + + +# ── error containment ── + +def test_tools_never_raise(client, store_path, monkeypatch): + seed_doc(store_path, "pi-a", "report.pdf") + monkeypatch.setattr(client._api._store, "get_tree", + lambda *a, **k: (_ for _ in ()).throw(RuntimeError("boom"))) + payload, is_error = run(client, "get_document_structure", + doc_name="report.pdf") + assert is_error + assert "boom" in payload["error"] + + +def test_unknown_argument_becomes_error_envelope(client, store_path): + seed_doc(store_path, "pi-a", "report.pdf") + payload, is_error = run(client, "get_document", doc_name="report.pdf", + bogus=True) + assert is_error and payload["errorCode"] == "INVALID_INPUT" + + +# ── framework adapters ── + +def test_as_openai_tools_missing_dependency(client, monkeypatch): + monkeypatch.setitem(sys.modules, "agents", None) + with pytest.raises(PageIndexAPIError, match="openai-agents"): + client.as_openai_tools() + + +def test_as_openai_tools_local_in_process(client): + pytest.importorskip("agents") + tools = client.as_openai_tools() + assert [tool.name for tool in tools] == list(tool_names()) + + +def test_as_openai_tools_cloud_default_uses_bridge(monkeypatch): + pytest.importorskip("agents") + from agents import FunctionTool + import pageindex.mcp_bridge as mcp_bridge + monkeypatch.setattr(mcp_bridge, "McpBridge", _FakeBridge) + cloud = PageIndexCloudClient(api_key="pi-test-key") + tools = cloud.as_openai_tools() + assert all(isinstance(tool, FunctionTool) for tool in tools) + assert [tool.name for tool in tools] == ["search_documents", "get_document"] + + +def test_as_openai_tools_cloud_hosted_opt_in(): + pytest.importorskip("agents") + from agents import HostedMCPTool + cloud = PageIndexCloudClient(api_key="pi-test-key") + tools = cloud.as_openai_tools(hosted=True) + assert len(tools) == 1 + assert isinstance(tools[0], HostedMCPTool) + config = tools[0].tool_config + assert config["server_url"] == "https://api.pageindex.ai/mcp" + assert config["headers"] == {"Authorization": "Bearer pi-test-key"} + assert config["server_label"] == "pageindex" + + +def test_as_openai_tools_local_ignores_hosted(client): + pytest.importorskip("agents") + assert ([tool.name for tool in client.as_openai_tools(hosted=True)] + == [tool.name for tool in client.as_openai_tools()] + == list(tool_names())) + + +def test_as_claude_mcp_cloud_needs_no_framework(monkeypatch): + monkeypatch.setitem(sys.modules, "claude_agent_sdk", None) + cloud = PageIndexCloudClient(api_key="pi-test-key") + config = cloud.as_claude_mcp() + assert config == { + "type": "http", + "url": "https://api.pageindex.ai/mcp", + "headers": {"Authorization": "Bearer pi-test-key"}, + } + + +def test_as_claude_mcp_local_missing_dependency(client, monkeypatch): + monkeypatch.setitem(sys.modules, "claude_agent_sdk", None) + with pytest.raises(PageIndexAPIError, match="claude-agent-sdk"): + client.as_claude_mcp() + + +def test_as_claude_mcp_local_when_installed(client): + pytest.importorskip("claude_agent_sdk") + server = client.as_claude_mcp() + assert server is not None + if isinstance(server, dict): + assert server.get("type") != "http" + + +def test_agent_tools_work_without_frameworks(client, store_path, monkeypatch): + monkeypatch.setitem(sys.modules, "agents", None) + monkeypatch.setitem(sys.modules, "claude_agent_sdk", None) + seed_doc(store_path, "pi-a", "report.pdf") + browse = client.agent_tools()[0] + assert "report.pdf" in browse() + + +# ── cloud agent_tools: MCP bridge ── + +class _FakeBridge: + def __init__(self, url, headers): + self.url = url + self.headers = headers + self.calls = [] + read_only = {"readOnlyHint": True, "openWorldHint": False} + self.tools = [ + { + "name": "search_documents", + "description": "ESCALATION tool — keyword search.", + "annotations": read_only, + "inputSchema": { + "type": "object", + "properties": { + "query": {"type": "string", "description": "Keyword query."}, + "limit": {"type": "number", "default": 10}, + }, + "required": ["query"], + }, + }, + { + "name": "get_document", + "description": "Check a document's status.", + "annotations": read_only, + "inputSchema": { + "type": "object", + "properties": { + "doc_name": {"type": "string"}, + "folder_id": {"type": ["string", "null"]}, + }, + "required": ["doc_name"], + }, + }, + { + "name": "remove_document", + "description": "Permanently delete documents.", + "annotations": {"readOnlyHint": False, "destructiveHint": True}, + "inputSchema": { + "type": "object", + "properties": {"doc_names": {"type": "array"}}, + "required": ["doc_names"], + }, + }, + { + "name": "unannotated_tool", + "description": "A tool the server sent without annotations.", + "inputSchema": {"type": "object", "properties": {}, + "required": []}, + }, + ] + + def list_tools(self): + return self.tools + + def call_tool(self, name, arguments): + self.calls.append((name, arguments)) + return json.dumps({"success": True, "tool": name, "args": arguments}) + + +@pytest.fixture +def cloud_with_fake_bridge(monkeypatch): + import pageindex.mcp_bridge as mcp_bridge + created = {} + + def factory(url, headers): + created["bridge"] = _FakeBridge(url, headers) + return created["bridge"] + + monkeypatch.setattr(mcp_bridge, "McpBridge", factory) + return PageIndexCloudClient(api_key="pi-test-key"), created + + +def test_cloud_agent_tools_discover_live_tool_set(cloud_with_fake_bridge): + cloud, created = cloud_with_fake_bridge + tools = cloud.agent_tools() + bridge = created["bridge"] + assert bridge.url == "https://api.pageindex.ai/mcp" + assert bridge.headers == {"Authorization": "Bearer pi-test-key"} + # Default: only tools the server marks read-only; unannotated tools are + # treated as non-read-only. + assert [t.__name__ for t in tools] == ["search_documents", "get_document"] + assert "ESCALATION tool" in tools[0].__doc__ + + +def test_cloud_agent_tools_management_gate(cloud_with_fake_bridge): + cloud, _ = cloud_with_fake_bridge + names = [t.__name__ for t in cloud.agent_tools(include_management=True)] + assert names == ["search_documents", "get_document", "remove_document", + "unannotated_tool"] + + +def test_cloud_agent_tools_signatures_from_schema(cloud_with_fake_bridge): + import inspect + cloud, _ = cloud_with_fake_bridge + search, get_document = cloud.agent_tools() + params = inspect.signature(search).parameters + assert list(params) == ["query", "limit"] + assert params["query"].default is inspect.Parameter.empty + assert params["limit"].default == 10 + assert search.__annotations__["query"] is str + folder_param = inspect.signature(get_document).parameters["folder_id"] + assert folder_param.default is None + + +def test_cloud_agent_tools_proxy_and_drop_none(cloud_with_fake_bridge): + cloud, created = cloud_with_fake_bridge + _, get_document = cloud.agent_tools() + result = json.loads(get_document("report.pdf")) + assert result["tool"] == "get_document" + assert result["args"] == {"doc_name": "report.pdf"} # folder_id=None dropped + assert created["bridge"].calls == [("get_document", {"doc_name": "report.pdf"})] + + +def test_cloud_agent_tools_call_errors_contained(cloud_with_fake_bridge): + cloud, created = cloud_with_fake_bridge + search, _ = cloud.agent_tools() + created["bridge"].call_tool = lambda *a, **k: (_ for _ in ()).throw( + RuntimeError("network down")) + payload = json.loads(search(query="x")) + assert payload["errorCode"] == "INTERNAL_ERROR" + assert "network down" in payload["error"] + + +def test_cloud_agent_tools_list_failure_raises(monkeypatch): + import pageindex.mcp_bridge as mcp_bridge + + class _DeadBridge: + def __init__(self, url, headers): + pass + + def list_tools(self): + raise PageIndexAPIError("Could not connect") + + monkeypatch.setattr(mcp_bridge, "McpBridge", _DeadBridge) + cloud = PageIndexCloudClient(api_key="pi-test-key") + with pytest.raises(PageIndexAPIError, match="Could not connect"): + cloud.agent_tools() + + +def test_mcp_bridge_protocol(monkeypatch): + from pageindex.mcp_bridge import McpBridge + import pageindex.mcp_bridge as mcp_bridge + + posts = [] + + class _Resp: + def __init__(self, status, body=None, headers=None, text=""): + self.status_code = status + self._body = body + self.headers = headers or {"Content-Type": "application/json"} + self.text = text or (json.dumps(body) if body else "") + self.content = self.text.encode("utf-8") + + def json(self): + if self._body is None: + raise ValueError("no body") + return self._body + + session_alive = {"first": True} + + def fake_post(url, json=None, headers=None, timeout=None): + posts.append({"payload": json, "headers": headers}) + method = json.get("method") + rid = json.get("id") + if method == "initialize": + return _Resp(200, {"jsonrpc": "2.0", "id": rid, + "result": {"protocolVersion": "2025-06-18"}}, + {"Content-Type": "application/json", + "Mcp-Session-Id": "sess-1"}) + if method == "notifications/initialized": + return _Resp(202) + if method == "tools/list": + # SSE-framed response exercises the event-stream parser; the + # em-dash guards UTF-8 decoding (SSE is UTF-8 by spec). + body = {"jsonrpc": "2.0", "id": rid, + "result": {"tools": [{"name": "t1", + "description": "reads — never writes"}], + "nextCursor": None}} + import json as json_mod + return _Resp(200, None, + {"Content-Type": "text/event-stream"}, + f"event: message\ndata: {json_mod.dumps(body)}\n\n") + if method == "tools/call": + if session_alive["first"]: + session_alive["first"] = False + return _Resp(404, text="session expired") + return _Resp(200, {"jsonrpc": "2.0", "id": rid, "result": { + "content": [{"type": "text", "text": "hello"}, + {"type": "text", "text": "world"}]}}) + raise AssertionError(f"unexpected method {method}") + + monkeypatch.setattr(mcp_bridge.requests, "post", fake_post) + bridge = McpBridge("https://api.pageindex.ai/mcp", + {"Authorization": "Bearer k"}) + + tools = bridge.list_tools() + assert tools == [{"name": "t1", "description": "reads — never writes"}] + list_headers = posts[-1]["headers"] + assert list_headers["Mcp-Session-Id"] == "sess-1" + assert list_headers["MCP-Protocol-Version"] == "2025-06-18" + assert list_headers["Authorization"] == "Bearer k" + + # First tools/call 404s (expired session) → re-initialize → retry succeeds. + text = bridge.call_tool("t1", {"a": 1}) + assert text == "hello\nworld" + methods = [p["payload"]["method"] for p in posts] + assert methods.count("initialize") == 2 + + +# ── review-round regressions ── + +def test_synth_optional_no_default_param_is_nullable(): + """A non-required, no-default schema param must annotate Optional, or + strict schemas force the model to always send a value (browse.query).""" + from pageindex.agent_tools import _make_bridge_function, TOOL_CONTRACT + from typing import get_args + + class _Bridge: + def call_tool(self, name, args): + return json.dumps(args) + + meta = {"name": "browse_documents", + "description": "d", + "inputSchema": TOOL_CONTRACT["browse_documents"]["schema"]} + fn = _make_bridge_function(_Bridge(), meta) + assert type(None) in get_args(fn.__annotations__["query"]) + + +def test_synth_escape_hatches(): + from pageindex.agent_tools import _make_bridge_function + + calls = [] + + class _Bridge: + def call_tool(self, name, args): + calls.append((name, args)) + return "ok" + + # Tool named "_invoke" must not recurse into itself. + invoke_named = _make_bridge_function(_Bridge(), { + "name": "_invoke", "description": "d", + "inputSchema": {"type": "object", "properties": {"x": {"type": "string"}}, + "required": ["x"]}}) + assert invoke_named("v") == "ok" + assert calls[-1] == ("_invoke", {"x": "v"}) + + # Param named "dict" must not shadow the builtin. + dict_param = _make_bridge_function(_Bridge(), { + "name": "t", "description": "d", + "inputSchema": {"type": "object", "properties": {"dict": {"type": "string"}}, + "required": ["dict"]}}) + assert dict_param("v") == "ok" + assert calls[-1] == ("t", {"dict": "v"}) + + # Non-identifier tool name still gets a real signature. + import inspect + dashed = _make_bridge_function(_Bridge(), { + "name": "page-content.v2", "description": "d", + "inputSchema": {"type": "object", "properties": {"a": {"type": "string"}}, + "required": ["a"]}}) + assert dashed.__name__ == "page-content.v2" + assert list(inspect.signature(dashed).parameters) == ["a"] + assert dashed("v") == "ok" + + +def test_cloud_agent_tools_empty_filter_raises(monkeypatch): + import pageindex.mcp_bridge as mcp_bridge + + class _AllWriteBridge: + def __init__(self, url, headers): + pass + + def list_tools(self): + return [{"name": "remove_document", + "annotations": {"readOnlyHint": False}, + "inputSchema": {"type": "object", "properties": {}}}] + + monkeypatch.setattr(mcp_bridge, "McpBridge", _AllWriteBridge) + cloud = PageIndexCloudClient(api_key="pi-test-key") + with pytest.raises(PageIndexAPIError, match="annotation"): + cloud.agent_tools() + assert len(cloud.agent_tools(include_management=True)) == 1 + + +def test_sse_crlf_multi_message(): + from pageindex.mcp_bridge import _parse_sse + body = ('event: message\r\ndata: {"jsonrpc":"2.0","method":"notifications/progress"}\r\n\r\n' + 'event: message\r\ndata: {"jsonrpc":"2.0","id":7,"result":{"ok":true}}\r\n\r\n') + messages = _parse_sse(body) + assert len(messages) == 2 + assert messages[1]["result"] == {"ok": True} + + +def test_bridge_transport_error_is_pageindex_error(monkeypatch): + import requests as requests_mod + import pageindex.mcp_bridge as mcp_bridge + from pageindex.mcp_bridge import McpBridge + + def dead_post(*args, **kwargs): + raise requests_mod.ConnectionError("dns down") + + monkeypatch.setattr(mcp_bridge.requests, "post", dead_post) + bridge = McpBridge("https://api.pageindex.ai/mcp", {}) + with pytest.raises(PageIndexAPIError, match="Could not reach"): + bridge.list_tools() + + +def test_failed_document_status_message(client, store_path): + seed_doc(store_path, "pi-a", "broken.pdf") + import pageindex.agent_tools as agent_tools_mod + entry = {"id": "pi-a", "name": "broken.pdf", "status": "failed"} + payload, is_error = agent_tools_mod._not_ready_error( + "broken.pdf", "failed", "structure retrieval", timed_out=False) + assert is_error + assert "failed" in payload["error"] + assert any("submit_document" in option + for option in payload["next_steps"]["options"]) + + +def test_hosted_approval_gate(): + pytest.importorskip("agents") + cloud = PageIndexCloudClient(api_key="pi-test-key") + gated = cloud.as_openai_tools(hosted=True)[0].tool_config + assert gated["require_approval"] == {"never": {"read_only": True}} + open_config = cloud.as_openai_tools(hosted=True, + include_management=True)[0].tool_config + assert open_config["require_approval"] == "never" + + +def test_wait_tolerates_transient_poll_failures(fake_cloud_client, monkeypatch): + cloud = fake_cloud_client(["processing", "completed"]) + original = cloud._api.get_document + state = {"raised": False} + + def flaky(doc_id): + if not state["raised"]: + state["raised"] = True + raise PageIndexAPIError("502") + return original(doc_id) + + monkeypatch.setattr(cloud._api, "get_document", flaky) + assert cloud.submit_document("x.pdf", wait=True) == {"doc_id": "pi-fake"} + + +LIVE_KEY = os.getenv("PAGEINDEX_API_KEY") + + +@pytest.mark.skipif(not LIVE_KEY, reason="PAGEINDEX_API_KEY not set") +def test_live_cloud_contract_parity(): + """Real-drift detector: the frozen contract must match the live server + on every shared tool, including the annotations the gates rely on.""" + from pageindex.mcp_bridge import McpBridge + bridge = McpBridge("https://api.pageindex.ai/mcp", + {"Authorization": f"Bearer {LIVE_KEY}"}) + live = {t["name"]: t for t in bridge.list_tools()} + for name, ours in TOOL_CONTRACT.items(): + real = live.get(name) + assert real is not None, f"{name} missing from live tools/list" + assert real.get("description") == ours["description"], name + real_schema = real.get("inputSchema") or {} + real_props = real_schema.get("properties") or {} + assert set(real_props) == set(ours["schema"]["properties"]), name + for param, spec in ours["schema"]["properties"].items(): + assert (real_props[param].get("description") + == spec.get("description")), (name, param) + assert (sorted(real_schema.get("required") or []) + == sorted(ours["schema"].get("required", []))), name + for key, value in (ours.get("annotations") or {}).items(): + assert (real.get("annotations") or {}).get(key) == value, (name, key) + + +# ── agent_instructions ── + +def test_agent_instructions_default(client): + text = client.agent_instructions() + assert text == AGENT_INSTRUCTIONS + assert "READING WORKFLOW" in text + assert "browse_documents" in text + assert "search_documents" not in text + assert "get_folder_structure" not in text + + +def test_agent_instructions_with_doc_id(client, store_path): + seed_doc(store_path, "pi-a", "report.pdf") + text = client.agent_instructions(doc_id="pi-a") + assert text.startswith(AGENT_INSTRUCTIONS) + assert "The user has specified document: report.pdf" in text + + seed_doc(store_path, "pi-b", "other.pdf") + multi = client.agent_instructions(doc_id=["pi-a", "pi-b"]) + assert "The user has specified documents: report.pdf, other.pdf" in multi + + with pytest.raises(PageIndexAPIError): + client.agent_instructions(doc_id="pi-missing") + + +# ── submit_document(wait=True) ── + +class _FakeCloudAPI: + def __init__(self, statuses): + self._statuses = list(statuses) + self.polls = 0 + + def submit_document(self, **kwargs): + return {"doc_id": "pi-fake"} + + def get_document(self, doc_id): + self.polls += 1 + status = (self._statuses.pop(0) if len(self._statuses) > 1 + else self._statuses[0]) + return {"id": doc_id, "status": status} + + +@pytest.fixture +def fake_cloud_client(tmp_path, monkeypatch): + monkeypatch.setattr(client_module.time, "sleep", lambda seconds: None) + + def build(statuses): + cloud = PageIndexLocalClient(storage_path=str(tmp_path / "unused")) + cloud._api = _FakeCloudAPI(statuses) + return cloud + return build + + +def test_submit_wait_polls_until_completed(fake_cloud_client): + cloud = fake_cloud_client(["processing", "processing", "completed"]) + result = cloud.submit_document("whatever.pdf", wait=True) + assert result == {"doc_id": "pi-fake"} + assert cloud._api.polls == 3 + + +def test_submit_wait_raises_on_failed(fake_cloud_client): + cloud = fake_cloud_client(["processing", "failed"]) + with pytest.raises(PageIndexAPIError, match="failed"): + cloud.submit_document("whatever.pdf", wait=True) + + +def test_submit_wait_times_out(fake_cloud_client, monkeypatch): + clock = {"now": 0.0} + + def fake_monotonic(): + clock["now"] += 700.0 + return clock["now"] + + monkeypatch.setattr(client_module.time, "monotonic", fake_monotonic) + cloud = fake_cloud_client(["processing"]) + with pytest.raises(PageIndexAPIError, match="Timed out"): + cloud.submit_document("whatever.pdf", wait=True) + + +def test_submit_without_wait_does_not_poll(fake_cloud_client): + cloud = fake_cloud_client(["processing"]) + assert cloud.submit_document("whatever.pdf") == {"doc_id": "pi-fake"} + assert cloud._api.polls == 0 From 873779990ad447b2a8f3fce8cf1cdad47381a487 Mon Sep 17 00:00:00 2001 From: Ray Date: Mon, 10 Aug 2026 17:30:12 +0800 Subject: [PATCH 03/18] =?UTF-8?q?fix:=20agent=20tools=20review=20=E2=80=94?= =?UTF-8?q?=20next=5Fsteps=20order,=20resolve=20caching,=20error=20semanti?= =?UTF-8?q?cs?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Large-doc next_steps now says structure-first, consistent with tool descriptions and agent instructions - _remove_document fetches document list once instead of per-name - call_tool returns error envelope for unknown names instead of raising - _not_ready_error timed_out flag reflects actual wait outcome - openai_agents.py docstring corrected to match default (FunctionTools) - Removed unused ModelSettings import from demo --- examples/agentic_vectorless_rag_demo.py | 3 +-- pageindex/agent_tools.py | 30 +++++++++++++++++++------ pageindex/integrations/openai_agents.py | 7 +++--- 3 files changed, 28 insertions(+), 12 deletions(-) diff --git a/examples/agentic_vectorless_rag_demo.py b/examples/agentic_vectorless_rag_demo.py index e8ed4a50a..0682bf2c7 100644 --- a/examples/agentic_vectorless_rag_demo.py +++ b/examples/agentic_vectorless_rag_demo.py @@ -29,7 +29,6 @@ sys.path.insert(0, str(Path(__file__).parent.parent)) from agents import Agent, Runner, set_tracing_disabled -from agents.model_settings import ModelSettings from agents.stream_events import RawResponsesStreamEvent, RunItemStreamEvent from openai.types.responses import ResponseTextDeltaEvent, ResponseReasoningSummaryTextDeltaEvent @@ -54,7 +53,7 @@ def query_agent(client: PageIndexLocalClient, doc_id: str, prompt: str, verbose: instructions=client.agent_instructions(doc_id=doc_id), tools=client.as_openai_tools(), model=client.retrieve_model, - # model_settings=ModelSettings(reasoning={"effort": "low", "summary": "auto"}), # Uncomment to enable reasoning + # model_settings=ModelSettings(reasoning={"effort": "low", "summary": "auto"}), # from agents.model_settings import ModelSettings ) async def _run(): diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index 10f0340dd..916536c7b 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -356,10 +356,12 @@ def _flat_metadata(value: Any) -> Optional[dict[str, Any]]: def _resolve_document( client, doc_name: str, + documents: Optional[list[dict[str, Any]]] = None, ) -> "tuple[Optional[dict[str, Any]], Optional[_ToolResult]]": """Resolve doc_name to a list entry. Same-name duplicates resolve to the newest match. Returns (entry, None) or (None, error_payload_pair).""" - documents = _all_documents(client) + if documents is None: + documents = _all_documents(client) matches = [doc for doc in documents if doc.get("name") == doc_name] if matches: return max(matches, key=lambda d: d.get("createdAt") or ""), None @@ -756,8 +758,8 @@ def _get_document(client, doc_name: str, folder_id: Optional[str] = None, else: suggestions.extend([ f"This is a large document with {page_num} pages.", - f'Start with first few pages: get_page_content(doc_name: "{name}", pages: "1-3")', - f'Or view structure first: get_document_structure(doc_name: "{name}")', + f'First explore structure: get_document_structure(doc_name: "{name}")', + f'Then target specific sections: get_page_content(doc_name: "{name}", pages: "1-3")', ]) else: suggestions.append("Document processing failed. Index the document " @@ -794,10 +796,12 @@ def _get_document_structure(client, doc_name: str, if error is not None: return error assert entry is not None + waited = wait_for_completion and entry.get("status") not in ("completed", "failed") entry = _await_completion(client, entry, wait_for_completion) if entry.get("status") != "completed": return _not_ready_error(doc_name, entry.get("status"), - "structure retrieval", wait_for_completion) + "structure retrieval", + waited and entry.get("status") != "failed") try: # Prefer the raw stored tree: its nodes carry start_index/end_index @@ -897,10 +901,12 @@ def _get_page_content(client, doc_name: str, pages: str, if error is not None: return error assert entry is not None + waited = wait_for_completion and entry.get("status") not in ("completed", "failed") entry = _await_completion(client, entry, wait_for_completion) if entry.get("status") != "completed": return _not_ready_error(doc_name, entry.get("status"), - "page content retrieval", wait_for_completion) + "page content retrieval", + waited and entry.get("status") != "failed") requested, error = _parse_page_spec(pages, doc_name) if error is not None: @@ -1013,9 +1019,10 @@ def _remove_document(client, doc_names: list[str], {"summary": "Too many documents in one call", "options": ["Delete at most 10 documents per call"]}, "INVALID_INPUT") + documents = _all_documents(client) results = [] for doc_name in doc_names: - entry, error = _resolve_document(client, doc_name) + entry, error = _resolve_document(client, doc_name, documents=documents) if error is not None or entry is None: results.append({"doc_name": doc_name, "status": "not_found"}) continue @@ -1051,7 +1058,16 @@ def tool_names(include_management: bool = False) -> tuple[str, ...]: def call_tool(client, name: str, arguments: dict[str, Any]) -> tuple[str, bool]: """Run one contract tool; returns (envelope_json, is_error). Never raises for tool-level failures — unexpected exceptions become error envelopes.""" - implementation = _IMPLEMENTATIONS[name] + implementation = _IMPLEMENTATIONS.get(name) + if implementation is None: + payload, _ = _failure( + f"Unknown tool: {name}", + {"tool_name": name, "available_tools": list(_IMPLEMENTATIONS)}, + {"summary": "Tool not found", + "options": [f"Available tools: {', '.join(_IMPLEMENTATIONS)}"]}, + "INVALID_INPUT", + ) + return json.dumps(payload), True try: payload, is_error = implementation(client, **arguments) except TypeError as exc: diff --git a/pageindex/integrations/openai_agents.py b/pageindex/integrations/openai_agents.py index 91ed6ebe3..f33c4b587 100644 --- a/pageindex/integrations/openai_agents.py +++ b/pageindex/integrations/openai_agents.py @@ -1,8 +1,9 @@ """OpenAI Agents SDK adapter for the Agent(tools=...) slot. -Cloud clients get one hosted MCP tool (the model connects to the PageIndex -cloud MCP server from OpenAI's side and discovers the full cloud tool set); -local clients get the in-process tools wrapped as FunctionTools. +Cloud clients default to the full live tool set as plain FunctionTools via +the MCP bridge; pass hosted=True to use a single HostedMCPTool instead +(the model connects to the PageIndex cloud MCP server from OpenAI's side). +Local clients get the in-process tools wrapped as FunctionTools. """ from __future__ import annotations From 2ba9035569f1f66e24d52aa7d9662fe363f8431f Mon Sep 17 00:00:00 2001 From: Ray Date: Mon, 10 Aug 2026 17:59:51 +0800 Subject: [PATCH 04/18] =?UTF-8?q?fix:=20agent=20tools=20review=202=20?= =?UTF-8?q?=E2=80=94=20bridge=20thread=20safety,=20browse=20paging,=20meta?= =?UTF-8?q?data=20merge?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - McpBridge reads session/protocol headers under the lock (now RLock: _ensure_initialized posts while holding it). openai-agents runs sync tools on threads and executes parallel tool calls concurrently, so bridge functions genuinely race; a torn read sent a new session id with a stale protocol header. Measured: one session expiry under 8 threads cost 4 initializations before, minimal 2 after. - Session-expiry retry also resets the negotiated protocol version, so the re-handshake carries no stale MCP-Protocol-Version header. - browse_documents time sort pages list_documents natively instead of fetching the whole library to slice one window (relevance still needs the full list for scoring). - _await_completion: a status refetch that nulls out metadata no longer clobbers the listing's copy (setdefault was a no-op on existing None). - Structure tool reads the raw stored tree via a named LocalAPI raw_tree() seam instead of reaching into _api._store internals; drop the redundant deepcopy before _format_structure (store re-reads from disk, formatting builds fresh containers). - Shared pageindex/_version.py replaces _sdk_version duplicated in mcp_bridge and the Claude integration. Left as-is after source verification against the cloud MCP: first-page budget bypass, pageNum falsy-zero, and the page-gap fallback text are letter-for-letter cloud behavior — parity wins over local repair. --- pageindex/_version.py | 10 +++++ pageindex/agent_tools.py | 26 +++++++------ pageindex/integrations/claude_agent_sdk.py | 11 +----- pageindex/local_api.py | 5 +++ pageindex/mcp_bridge.py | 25 ++++++------- tests/test_agent_tools.py | 43 ++++++++++++++++++++++ 6 files changed, 86 insertions(+), 34 deletions(-) create mode 100644 pageindex/_version.py diff --git a/pageindex/_version.py b/pageindex/_version.py new file mode 100644 index 000000000..da5c00c2c --- /dev/null +++ b/pageindex/_version.py @@ -0,0 +1,10 @@ +"""Installed-package version, shared by every surface that reports it upstream.""" +from __future__ import annotations + + +def sdk_version() -> str: + try: + from importlib.metadata import version + return version("pageindex") + except Exception: + return "0.0.0" diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index 916536c7b..866b130a0 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -10,7 +10,6 @@ """ from __future__ import annotations -import copy import difflib import json import re @@ -407,7 +406,10 @@ def _await_completion(client, entry: dict[str, Any], wait: bool) -> dict[str, An refreshed = _refetch_entry(client, doc_id) if refreshed is None: return current - refreshed.setdefault("metadata", current.get("metadata")) + if refreshed.get("metadata") is None: + # Status refetches omit (or null out) custom metadata; keep the + # listing's copy. + refreshed["metadata"] = current.get("metadata") current = {**current, **refreshed} if current.get("status") in ("completed", "failed"): return current @@ -631,21 +633,23 @@ def _browse_documents(client, folder_id: str = "root", recursive: bool = False, "options": ["Pass integer offset and limit values"]}, "INVALID_INPUT") - documents = _all_documents(client) if sort == "relevance": tokens = [token for token in (query or "").lower().split() if token] scored = [] - for doc in documents: + for doc in _all_documents(client): haystack = f"{doc.get('name') or ''} {doc.get('description') or ''}".lower() score = sum(1 for token in tokens if token in haystack) if score: scored.append((score, doc)) # Stable sort: equal scores keep the newest-first listing order. scored.sort(key=lambda pair: pair[0], reverse=True) - documents = [doc for _, doc in scored] - - window = documents[offset:offset + limit] - has_more = offset + limit < len(documents) + ranked = [doc for _, doc in scored] + window = ranked[offset:offset + limit] + has_more = offset + limit < len(ranked) + else: + listing = client.list_documents(limit=limit, offset=offset) + window = listing.get("documents") or [] + has_more = offset + limit < listing.get("total", 0) next_offset = offset + limit if has_more else None page_has_processing = False @@ -807,8 +811,8 @@ def _get_document_structure(client, doc_name: str, # Prefer the raw stored tree: its nodes carry start_index/end_index # like the cloud structure tool, where client.get_tree() drops # end_index and renames fields. - store = getattr(getattr(client, "_api", None), "_store", None) - tree = store.get_tree(entry["id"]) if store is not None else None + raw_tree = getattr(getattr(client, "_api", None), "raw_tree", None) + tree = raw_tree(entry["id"]) if raw_tree is not None else None if tree is None: tree = client.get_tree(entry["id"], node_summary=True).get("result") except PageIndexAPIError as exc: @@ -840,7 +844,7 @@ def _get_document_structure(client, doc_name: str, "INTERNAL_ERROR", ) - formatted = _format_structure(copy.deepcopy(tree)) + formatted = _format_structure(tree) chunks = _split_structure(formatted, _CHAR_BUDGET) total_parts = max(1, len(chunks)) try: diff --git a/pageindex/integrations/claude_agent_sdk.py b/pageindex/integrations/claude_agent_sdk.py index d0e2316ae..0fb77d2de 100644 --- a/pageindex/integrations/claude_agent_sdk.py +++ b/pageindex/integrations/claude_agent_sdk.py @@ -9,17 +9,10 @@ import asyncio from typing import Any +from .._version import sdk_version from ..errors import PageIndexAPIError -def _sdk_version() -> str: - try: - from importlib.metadata import version - return version("pageindex") - except Exception: - return "0.0.0" - - def build_claude_mcp(client, include_management: bool = False): if getattr(client, "api_key", None): return { @@ -63,5 +56,5 @@ def tool_kwargs(name: str) -> dict: TOOL_CONTRACT[name]["schema"], **tool_kwargs(name))(make_handler(name)) for name in tool_names(include_management) ] - return create_sdk_mcp_server(name="pageindex", version=_sdk_version(), + return create_sdk_mcp_server(name="pageindex", version=sdk_version(), tools=tools) diff --git a/pageindex/local_api.py b/pageindex/local_api.py index 0e9f682c8..7b7351cca 100644 --- a/pageindex/local_api.py +++ b/pageindex/local_api.py @@ -190,6 +190,11 @@ def _load_tree_with_text(self, doc_id: str, error_prefix: str) -> list: add_node_text(structure, pdf_pages) return structure + def raw_tree(self, doc_id: str) -> list | None: + """Stored tree verbatim — keeps start_index/end_index, which + get_tree's cloud wire shape renames and drops.""" + return self._store.get_tree(doc_id) + def get_tree(self, doc_id: str, node_summary: bool = False, include_text: bool = True) -> dict[str, Any]: meta = self._require_doc(doc_id, "Failed to get tree result") diff --git a/pageindex/mcp_bridge.py b/pageindex/mcp_bridge.py index d144cee92..fad0baaf8 100644 --- a/pageindex/mcp_bridge.py +++ b/pageindex/mcp_bridge.py @@ -14,20 +14,13 @@ import requests +from ._version import sdk_version from .errors import PageIndexAPIError _PROTOCOL_VERSION = "2025-06-18" _TIMEOUT = (10, 240) # tools may wait server-side (wait_for_completion: 3 min) -def _sdk_version() -> str: - try: - from importlib.metadata import version - return version("pageindex") - except Exception: - return "0.0.0" - - def _parse_sse(text: str) -> list[dict]: """JSON-RPC messages out of a text/event-stream body.""" messages = [] @@ -51,21 +44,24 @@ def __init__(self, url: str, headers: dict[str, str]): self._session_id: Optional[str] = None self._protocol_version: Optional[str] = None self._initialized = False - self._lock = threading.Lock() + self._lock = threading.RLock() self._next_id = 0 # ── JSON-RPC over streamable HTTP ── def _post(self, payload: dict) -> requests.Response: + with self._lock: + session_id = self._session_id + protocol_version = self._protocol_version headers = { "Content-Type": "application/json", "Accept": "application/json, text/event-stream", **self._auth_headers, } - if self._session_id: - headers["Mcp-Session-Id"] = self._session_id - if self._protocol_version: - headers["MCP-Protocol-Version"] = self._protocol_version + if session_id: + headers["Mcp-Session-Id"] = session_id + if protocol_version: + headers["MCP-Protocol-Version"] = protocol_version try: return requests.post(self._url, json=payload, headers=headers, timeout=_TIMEOUT) @@ -117,6 +113,7 @@ def _request(self, method: str, params: Optional[dict] = None, with self._lock: self._initialized = False self._session_id = None + self._protocol_version = None return self._request(method, params, _retry=False) if response.status_code >= 400: raise PageIndexAPIError( @@ -137,7 +134,7 @@ def _ensure_initialized(self) -> None: "protocolVersion": _PROTOCOL_VERSION, "capabilities": {}, "clientInfo": {"name": "pageindex-python-sdk", - "version": _sdk_version()}, + "version": sdk_version()}, }, }) if response.status_code >= 400: diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 289c2dada..9ab865655 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -630,6 +630,11 @@ def fake_post(url, json=None, headers=None, timeout=None): assert text == "hello\nworld" methods = [p["payload"]["method"] for p in posts] assert methods.count("initialize") == 2 + # The expired session's negotiated state must not leak into the new + # handshake. + reinit = [p for p in posts if p["payload"].get("method") == "initialize"][1] + assert "MCP-Protocol-Version" not in reinit["headers"] + assert "Mcp-Session-Id" not in reinit["headers"] # ── review-round regressions ── @@ -730,6 +735,44 @@ def dead_post(*args, **kwargs): bridge.list_tools() +def test_await_completion_preserves_metadata_over_null_refetch(monkeypatch): + """A status refetch that nulls out metadata must not clobber the + listing's copy (setdefault is a no-op on an existing None value).""" + import pageindex.agent_tools as agent_tools_mod + monkeypatch.setattr(agent_tools_mod.time, "sleep", lambda seconds: None) + + class _Client: + def get_document(self, doc_id): + return {"id": doc_id, "status": "completed", "metadata": None} + + entry = {"id": "pi-x", "status": "processing", + "metadata": {"team": "research"}} + merged = agent_tools_mod._await_completion(_Client(), entry, True) + assert merged["status"] == "completed" + assert merged["metadata"] == {"team": "research"} + + +def test_browse_time_sort_uses_native_pagination(client, store_path, monkeypatch): + """Time-sorted browsing must page through list_documents directly, not + fetch the whole library to slice one window.""" + for index in range(3): + seed_doc(store_path, f"pi-{index}", f"doc{index}.pdf", + created_at=f"2026-08-0{index + 1}T10:00:00.000000") + calls = [] + original = client.list_documents + + def spy(**kwargs): + calls.append(kwargs) + return original(**kwargs) + + monkeypatch.setattr(client, "list_documents", spy) + payload, is_error = run(client, "browse_documents", limit=2) + assert not is_error + assert calls == [{"limit": 2, "offset": 0}] + assert [d["name"] for d in payload["documents"]] == ["doc2.pdf", "doc1.pdf"] + assert payload["has_more"] is True and payload["next_offset"] == 2 + + def test_failed_document_status_message(client, store_path): seed_doc(store_path, "pi-a", "broken.pdf") import pageindex.agent_tools as agent_tools_mod From 3f131584e25ada6d97de07832d6c6265fcf62df1 Mon Sep 17 00:00:00 2001 From: Ray Date: Mon, 10 Aug 2026 18:49:27 +0800 Subject: [PATCH 05/18] =?UTF-8?q?fix:=20agent=20tools=20review=203=20?= =?UTF-8?q?=E2=80=94=20page-span=20cap,=20duplicate=20names,=20wait=20resi?= =?UTF-8?q?lience,=20contract=20drift?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - _parse_page_spec bounds the requested span arithmetically (10k pages) before materializing it; pages="1-1000000000" previously expanded to a billion integers inside the caller's process. - Local submit_document uniquifies document names the way the cloud upload does (taken name -> _1.._99, then reject with the cloud's own message). Same-name duplicates broke name-addressed tools: resolution always picks the newest, so older duplicates were unreachable. - agent_instructions(doc_id=...) now fails loud when the pinned doc's name is shadowed by a newer same-name document (legacy stores predate the rename) — it previews resolution with the same _resolve_document the tools use, so the check cannot drift from actual behavior. - submit_document(wait=True) tolerates transient network errors, not just API errors; a dropped connection at minute 25 of a 30-minute wait no longer kills it. Third strike wraps into PageIndexAPIError per the documented contract. - The live contract-parity test compares full per-param schemas, not just names and descriptions. It immediately caught real drift the shallow check had been passing: the server now emits nullables as anyOf unions and stamps MAX_SAFE_INTEGER maxima on offset/part. Contract and snapshot updated to the served wire form; _annotation_for learned anyOf so bridge signatures stay Optional[str] instead of degrading to Any. Adjudicated, not changed: the allowed_tools wildcard example stays (docstring advice covers scoping; Ray's call), and raw-length response accounting stays (letter-for-letter cloud behavior, parity wins). --- pageindex/agent_tools.py | 54 +++++++++++++++++++++++---- pageindex/client.py | 14 +++++-- pageindex/local_api.py | 18 ++++++++- tests/data/cloud_mcp_contract.json | 42 +++++++++++++++------ tests/test_agent_tools.py | 60 ++++++++++++++++++++++++++++-- tests/test_client.py | 24 ++++++++++++ 6 files changed, 185 insertions(+), 27 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index 866b130a0..68af57eaf 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -23,6 +23,7 @@ _CHAR_BUDGET = int(TOOL_RESPONSE_CHAR_LIMIT * 0.95) _PAGES_SPEC_RE = re.compile(r"^(\d+(-\d+)?)(,\s*\d+(-\d+)?)*$") +_MAX_REQUESTED_PAGES = 10_000 _SIMILAR_NAMES_LIMIT = 3 _TOOL_WAIT_TIMEOUT = 180.0 # "up to 3 minutes", per the wait_for_completion schema _TOOL_WAIT_INTERVAL = 5.0 @@ -114,6 +115,7 @@ "offset": { "type": "integer", "minimum": 0, + "maximum": 9007199254740991, "default": 0, "description": ( "Zero-based pagination offset. Pass the value of " @@ -152,7 +154,7 @@ "description": _DOC_NAME_DESCRIPTION, }, "folder_id": { - "type": ["string", "null"], + "anyOf": [{"type": "string"}, {"type": "null"}], "description": _FOLDER_ID_DISAMBIGUATOR_DESCRIPTION, }, "wait_for_completion": { @@ -183,12 +185,13 @@ "description": _DOC_NAME_DESCRIPTION, }, "folder_id": { - "type": ["string", "null"], + "anyOf": [{"type": "string"}, {"type": "null"}], "description": _FOLDER_ID_DISAMBIGUATOR_DESCRIPTION, }, "part": { "type": "integer", "minimum": 1, + "maximum": 9007199254740991, "default": 1, "description": ( "Part number for pagination (1-based, default 1). For " @@ -224,7 +227,7 @@ "description": _DOC_NAME_DESCRIPTION, }, "folder_id": { - "type": ["string", "null"], + "anyOf": [{"type": "string"}, {"type": "null"}], "description": _FOLDER_ID_DISAMBIGUATOR_DESCRIPTION, }, "pages": { @@ -273,7 +276,7 @@ ), }, "folder_id": { - "type": ["string", "null"], + "anyOf": [{"type": "string"}, {"type": "null"}], "description": _FOLDER_ID_DISAMBIGUATOR_DESCRIPTION, }, }, @@ -492,15 +495,34 @@ def _parse_page_spec( if not isinstance(pages, str) or not _PAGES_SPEC_RE.match(pages.strip()): return None, invalid expanded: set[int] = set() + requested_total = 0 for part in pages.split(","): part = part.strip() if "-" in part: start, end = (int(x) for x in part.split("-", 1)) if start > end: return None, invalid - expanded.update(range(start, end + 1)) else: - expanded.add(int(part)) + start = end = int(part) + # Bound the span arithmetically before materializing it: a spec like + # "1-1000000000" would otherwise expand to billions of integers + # inside the caller's process. + requested_total += end - start + 1 + if requested_total > _MAX_REQUESTED_PAGES: + return None, _failure( + f"Too many pages requested (over {_MAX_REQUESTED_PAGES})", + {"doc_name": doc_name}, + { + "summary": "The page specification spans too many pages", + "options": [ + "Request a narrower page range", + "The response holds only a few pages per call - page " + "through with several smaller requests", + ], + }, + "INVALID_INPUT", + ) + expanded.update(range(start, end + 1)) if any(page < 1 for page in expanded): return None, _failure( "Invalid page numbers. Page numbers must be positive integers", @@ -1114,6 +1136,10 @@ def _docstring(name: str) -> str: def _annotation_for(spec: dict) -> Any: schema_type = spec.get("type") + if schema_type is None and isinstance(spec.get("anyOf"), list): + # Nullable unions arrive as anyOf: [{type: string}, {type: null}]. + schema_type = [option.get("type") for option in spec["anyOf"] + if isinstance(option, dict) and option.get("type")] if isinstance(schema_type, list): bases = [t for t in schema_type if t != "null"] base = _SCHEMA_TYPE_MAP.get(bases[0], Any) if bases else Any @@ -1320,13 +1346,27 @@ def remove_document(doc_names: list[str], def build_agent_instructions(client, doc_id=None) -> str: """Orchestration guidance for document QA agents; with doc_id, appends - the target documents and directs the agent to work within them.""" + the target documents and directs the agent to work within them. Raises + when a doc_id's name is shadowed by a newer same-name document — the + name-addressed tools could not reach it.""" if doc_id is None: return AGENT_INSTRUCTIONS doc_ids = [doc_id] if isinstance(doc_id, str) else list(doc_id) if not doc_ids: return AGENT_INSTRUCTIONS details = [client.get_document(one_id) for one_id in doc_ids] + documents = _all_documents(client) + for one_id, detail in zip(doc_ids, details): + entry, _ = _resolve_document(client, str(detail.get("name")), + documents=documents) + if entry is not None and entry.get("id") != one_id: + raise PageIndexAPIError( + f'Document "{detail.get("name")}" (doc_id: {one_id}) is ' + "shadowed by a newer document with the same name (doc_id: " + f'{entry.get("id")}). The tools address documents by name ' + "and would read the newer one. Rename or remove the " + "duplicate, or pass the newer doc_id." + ) context = json.dumps(details, ensure_ascii=False) if len(details) == 1: block = ( diff --git a/pageindex/client.py b/pageindex/client.py index 2a402a485..54ae9c018 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -172,6 +172,7 @@ def submit_document( return result def _wait_until_ready(self, doc_id: str, timeout: float = 1800.0) -> None: + import requests interval = 2.0 deadline = time.monotonic() + timeout poll_failures = 0 @@ -179,12 +180,16 @@ def _wait_until_ready(self, doc_id: str, timeout: float = 1800.0) -> None: try: status = self.get_document(doc_id).get("status") poll_failures = 0 - except PageIndexAPIError: + except (PageIndexAPIError, requests.RequestException) as exc: # Tolerate transient poll failures; a 30-minute wait should - # not die on one 502. + # not die on one 502 or dropped connection. poll_failures += 1 if poll_failures >= 3: - raise + if isinstance(exc, PageIndexAPIError): + raise + raise PageIndexAPIError( + f"Could not poll document status: {exc}" + ) from exc status = None if status == "completed": return @@ -500,7 +505,8 @@ def agent_instructions(self, doc_id: Optional[Union[str, list[str]]] = None) -> With ``doc_id`` (str or list, same shape as ``chat_completions``), appends the target documents' names and metadata and directs the agent to work within them. Raises PageIndexAPIError if a doc_id does - not exist. + not exist, or if its name is shadowed by a newer same-name document + (the name-addressed tools could not reach it). """ from .agent_tools import build_agent_instructions return build_agent_instructions(self, doc_id) diff --git a/pageindex/local_api.py b/pageindex/local_api.py index 7b7351cca..5820656ac 100644 --- a/pageindex/local_api.py +++ b/pageindex/local_api.py @@ -115,7 +115,7 @@ def submit_document( doc_id = "pi-" + uuid.uuid4().hex meta = { "id": doc_id, - "name": os.path.basename(file_path), + "name": self._unique_doc_name(os.path.basename(file_path)), "description": description, "status": "completed", "createdAt": _now_iso(), @@ -131,6 +131,22 @@ def submit_document( doc_id, meta, remove_fields(structure, fields=["text"]), pages) return {"doc_id": doc_id} + def _unique_doc_name(self, name: str) -> str: + """Mirror the cloud upload: a taken name gets _1.._99 appended, + beyond that the submit is rejected.""" + taken = {meta.get("name") for meta in self._store.list_metas()} + if name not in taken: + return name + base, ext = os.path.splitext(name) + for num in range(1, 100): + candidate = f"{base}_{num}{ext}" + if candidate not in taken: + return candidate + raise PageIndexAPIError( + "Failed to submit document: Too many files with similar names. " + "Please use a different file name." + ) + @staticmethod def _extract_page_texts(file_path: str) -> list[str]: import PyPDF2 diff --git a/tests/data/cloud_mcp_contract.json b/tests/data/cloud_mcp_contract.json index 71743aee2..25711a6ba 100644 --- a/tests/data/cloud_mcp_contract.json +++ b/tests/data/cloud_mcp_contract.json @@ -36,6 +36,7 @@ "offset": { "type": "integer", "minimum": 0, + "maximum": 9007199254740991, "default": 0, "description": "Zero-based pagination offset. Pass the value of `next_offset` from the previous response to fetch the next page." }, @@ -65,9 +66,13 @@ "description": "Copy the `name` field verbatim from a browse_documents() or search_documents() response (case-sensitive, include extension). Example: \"Q3 Report.pdf\". If the response shows two documents with the same name, pass `folder_id` alongside to disambiguate." }, "folder_id": { - "type": [ - "string", - "null" + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } ], "description": "Disambiguator for same-name documents. Copy the `folder_id` from the intended browse/search result; use \"root\" for root-level documents, or \"shared-with-me\"/\"following\" for the read-only folders at the library root; omit if `doc_name` is unique. Copy any folder_id verbatim from a browse_documents()/get_folder_structure() response, never construct one." }, @@ -97,15 +102,20 @@ "description": "Copy the `name` field verbatim from a browse_documents() or search_documents() response (case-sensitive, include extension). Example: \"Q3 Report.pdf\". If the response shows two documents with the same name, pass `folder_id` alongside to disambiguate." }, "folder_id": { - "type": [ - "string", - "null" + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } ], "description": "Disambiguator for same-name documents. Copy the `folder_id` from the intended browse/search result; use \"root\" for root-level documents, or \"shared-with-me\"/\"following\" for the read-only folders at the library root; omit if `doc_name` is unique. Copy any folder_id verbatim from a browse_documents()/get_folder_structure() response, never construct one." }, "part": { "type": "integer", "minimum": 1, + "maximum": 9007199254740991, "default": 1, "description": "Part number for pagination (1-based, default 1). For large outlines, increment until the response's `pagination.has_more` becomes false." }, @@ -135,9 +145,13 @@ "description": "Copy the `name` field verbatim from a browse_documents() or search_documents() response (case-sensitive, include extension). Example: \"Q3 Report.pdf\". If the response shows two documents with the same name, pass `folder_id` alongside to disambiguate." }, "folder_id": { - "type": [ - "string", - "null" + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } ], "description": "Disambiguator for same-name documents. Copy the `folder_id` from the intended browse/search result; use \"root\" for root-level documents, or \"shared-with-me\"/\"following\" for the read-only folders at the library root; omit if `doc_name` is unique. Copy any folder_id verbatim from a browse_documents()/get_folder_structure() response, never construct one." }, @@ -181,9 +195,13 @@ "description": "Array of document names to delete. Each name must be copied verbatim from the `name` field of a browse_documents() or search_documents() response (case-sensitive, include extension). Example: [\"Q3 Report.pdf\", \"draft.pdf\"]. Max 10 per call." }, "folder_id": { - "type": [ - "string", - "null" + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } ], "description": "Disambiguator for same-name documents. Copy the `folder_id` from the intended browse/search result; use \"root\" for root-level documents, or \"shared-with-me\"/\"following\" for the read-only folders at the library root; omit if `doc_name` is unique. Copy any folder_id verbatim from a browse_documents()/get_folder_structure() response, never construct one." } diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 9ab865655..5039bd031 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -451,7 +451,8 @@ def __init__(self, url, headers): "type": "object", "properties": { "doc_name": {"type": "string"}, - "folder_id": {"type": ["string", "null"]}, + "folder_id": {"anyOf": [{"type": "string"}, + {"type": "null"}]}, }, "required": ["doc_name"], }, @@ -525,6 +526,10 @@ def test_cloud_agent_tools_signatures_from_schema(cloud_with_fake_bridge): assert search.__annotations__["query"] is str folder_param = inspect.signature(get_document).parameters["folder_id"] assert folder_param.default is None + # The live server encodes nullables as anyOf; the annotation must still + # come out Optional[str], not Any. + from typing import Optional + assert get_document.__annotations__["folder_id"] == Optional[str] def test_cloud_agent_tools_proxy_and_drop_none(cloud_with_fake_bridge): @@ -693,6 +698,17 @@ def call_tool(self, name, args): assert dashed("v") == "ok" +def test_annotation_for_both_nullable_encodings(): + """Servers have emitted nullables as type-arrays and as anyOf unions; + both must map to Optional, not degrade to Any.""" + from typing import Optional + from pageindex.agent_tools import _annotation_for + assert _annotation_for({"type": "string"}) is str + assert _annotation_for({"type": ["string", "null"]}) == Optional[str] + assert (_annotation_for({"anyOf": [{"type": "string"}, {"type": "null"}]}) + == Optional[str]) + + def test_cloud_agent_tools_empty_filter_raises(monkeypatch): import pageindex.mcp_bridge as mcp_bridge @@ -773,6 +789,43 @@ def spy(**kwargs): assert payload["has_more"] is True and payload["next_offset"] == 2 +def test_page_spec_span_bomb_rejected(client, store_path): + """An absurd range must be rejected arithmetically, not expanded into + billions of integers in the caller's process.""" + seed_doc(store_path, "pi-a", "report.pdf") + payload, is_error = run(client, "get_page_content", doc_name="report.pdf", + pages="1-1000000000") + assert is_error and payload["errorCode"] == "INVALID_INPUT" + assert "Too many pages" in payload["error"] + + +def test_agent_instructions_shadowed_doc_id_raises(client, store_path): + seed_doc(store_path, "pi-old", "report.pdf", + created_at="2026-08-01T10:00:00.000000") + seed_doc(store_path, "pi-new", "report.pdf", + created_at="2026-08-02T10:00:00.000000") + with pytest.raises(PageIndexAPIError, match="shadowed"): + client.agent_instructions(doc_id="pi-old") + text = client.agent_instructions(doc_id="pi-new") + assert "report.pdf" in text + + +def test_wait_tolerates_transient_network_failures(fake_cloud_client, monkeypatch): + import requests as requests_mod + cloud = fake_cloud_client(["processing", "completed"]) + original = cloud._api.get_document + state = {"raised": False} + + def flaky(doc_id): + if not state["raised"]: + state["raised"] = True + raise requests_mod.ConnectionError("network blip") + return original(doc_id) + + monkeypatch.setattr(cloud._api, "get_document", flaky) + assert cloud.submit_document("x.pdf", wait=True) == {"doc_id": "pi-fake"} + + def test_failed_document_status_message(client, store_path): seed_doc(store_path, "pi-a", "broken.pdf") import pageindex.agent_tools as agent_tools_mod @@ -828,9 +881,10 @@ def test_live_cloud_contract_parity(): real_schema = real.get("inputSchema") or {} real_props = real_schema.get("properties") or {} assert set(real_props) == set(ours["schema"]["properties"]), name + # Full per-param equality: a drifted type, default, enum, or bound + # breaks calls just as surely as a renamed parameter. for param, spec in ours["schema"]["properties"].items(): - assert (real_props[param].get("description") - == spec.get("description")), (name, param) + assert real_props[param] == spec, (name, param) assert (sorted(real_schema.get("required") or []) == sorted(ours["schema"].get("required", []))), name for key, value in (ours.get("annotations") or {}).items(): diff --git a/tests/test_client.py b/tests/test_client.py index 50b7f5178..3c9385831 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -162,6 +162,30 @@ def fake_page_index_main(doc, opt=None, logger=None, page_list=None): assert not (tmp_path / "logs").exists() +def test_submit_duplicate_name_gets_suffix(local_client, sample_pdf, monkeypatch): + """Mirror the cloud upload: a second submit of the same file name is + stored as name_1, not as a same-name duplicate.""" + def fake_page_index_main(doc, opt=None, logger=None): + return {"doc_name": "sample.pdf", "doc_description": "d", + "structure": json.loads(json.dumps(STRUCTURE))} + monkeypatch.setattr(page_index_module, "page_index_main", fake_page_index_main) + first = local_client.submit_document(sample_pdf)["doc_id"] + second = local_client.submit_document(sample_pdf)["doc_id"] + names = {d["id"]: d["name"] + for d in local_client.list_documents()["documents"]} + assert names[first] == "sample.pdf" + assert names[second] == "sample_1.pdf" + + +def test_submit_duplicate_name_exhaustion(local_client, monkeypatch): + api = local_client._api + metas = ([{"name": "x.pdf"}] + + [{"name": f"x_{num}.pdf"} for num in range(1, 100)]) + monkeypatch.setattr(api._store, "list_metas", lambda: metas) + with pytest.raises(PageIndexAPIError, match="Too many files"): + api._unique_doc_name("x.pdf") + + def test_submit_flash(local_client, sample_pdf, monkeypatch): calls = {} def fake_flash(pdf, summary=True, summary_model=None, **kwargs): From 40b1706e8c4c4a5608bf624c72d441d6cb0d8e68 Mon Sep 17 00:00:00 2001 From: Ray Date: Mon, 10 Aug 2026 23:51:55 +0800 Subject: [PATCH 06/18] feat: surface the stored document name from submit_document MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Compute PR #558 makes /doc/ return {"doc_id", "name"} carrying the post-dedup-rename name. Mirror it end to end: local submit returns the stored name, the client warns when it differs from the uploaded file name (read via .get so older cloud servers stay compatible), the local name-exhaustion check runs before indexing instead of after the LLM spend, and the demo caches doc_id in a file instead of name-matching — a renamed document made the name lookup re-index on every run. --- examples/agentic_vectorless_rag_demo.py | 17 +++++++++----- pageindex/client.py | 15 +++++++++++-- pageindex/cloud_api.py | 4 +++- pageindex/local_api.py | 5 ++++- tests/test_agent_tools.py | 9 ++++++++ tests/test_client.py | 30 ++++++++++++++++++++----- 6 files changed, 65 insertions(+), 15 deletions(-) diff --git a/examples/agentic_vectorless_rag_demo.py b/examples/agentic_vectorless_rag_demo.py index 0682bf2c7..f35c3c2e7 100644 --- a/examples/agentic_vectorless_rag_demo.py +++ b/examples/agentic_vectorless_rag_demo.py @@ -32,13 +32,14 @@ from agents.stream_events import RawResponsesStreamEvent, RunItemStreamEvent from openai.types.responses import ResponseTextDeltaEvent, ResponseReasoningSummaryTextDeltaEvent -from pageindex import PageIndexLocalClient +from pageindex import PageIndexAPIError, PageIndexLocalClient import pageindex.utils as utils PDF_URL = "https://arxiv.org/pdf/2603.15031" _EXAMPLES_DIR = Path(__file__).parent PDF_PATH = _EXAMPLES_DIR / "documents" / "attention-residuals.pdf" +DOC_ID_PATH = _EXAMPLES_DIR / "documents" / "attention-residuals.doc_id" STORAGE_PATH = _EXAMPLES_DIR / ".pageindex" @@ -129,15 +130,19 @@ async def _run(): print("=" * 60) print("Step 1: Index PDF and view tree structure") print("=" * 60) - doc_id = next( - (doc["id"] for doc in client.list_documents(limit=100)["documents"] - if doc["name"] == PDF_PATH.name), - None, - ) + doc_id = None + if DOC_ID_PATH.exists(): + cached = DOC_ID_PATH.read_text().strip() + try: + client.get_document(cached) + doc_id = cached + except PageIndexAPIError: + DOC_ID_PATH.unlink() if doc_id: print(f"\nLoaded cached doc_id: {doc_id}") else: doc_id = client.submit_document(str(PDF_PATH), wait=True)["doc_id"] + DOC_ID_PATH.write_text(doc_id) print(f"\nIndexed. doc_id: {doc_id}") print("\nTree Structure (top-level sections):") structure = client.get_tree(doc_id, node_summary=True)["result"] diff --git a/pageindex/client.py b/pageindex/client.py index 54ae9c018..ff44aa376 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -1,7 +1,9 @@ """PageIndex SDK client: the 0.2.x cloud surface, now with a local mode.""" from __future__ import annotations +import os import time +import warnings from typing import Any, Callable, Iterator, Optional, Union from .errors import PageIndexAPIError @@ -130,7 +132,7 @@ def submit_document( wait: bool = False, ) -> dict[str, Any]: """ - Submit a PDF document for processing. Returns {'doc_id': ...}. + Submit a PDF document for processing. Returns {'doc_id': ..., 'name': ...}. Cloud: uploads the file; processing is asynchronous. Pass ``wait=True`` to block until the document is ready, or poll @@ -161,12 +163,21 @@ def submit_document( concurrently and poll afterwards. Returns: - dict: {'doc_id': ...} + dict: {'doc_id': ..., 'name': ...}. 'name' is the stored document + name: a taken name gains a numeric suffix (name_1..name_99) + and a UserWarning is emitted. Older cloud servers omit 'name'. """ result = self._api.submit_document( file_path=file_path, mode=mode, beta_headers=beta_headers, folder_id=folder_id, metadata=metadata, ) + stored = result.get("name") + if stored and stored != os.path.basename(file_path): + warnings.warn( + f'Document "{os.path.basename(file_path)}" was stored as ' + f'"{stored}".', + stacklevel=2, + ) if wait: self._wait_until_ready(result["doc_id"]) return result diff --git a/pageindex/cloud_api.py b/pageindex/cloud_api.py index b7597cc9a..ab7a9c885 100644 --- a/pageindex/cloud_api.py +++ b/pageindex/cloud_api.py @@ -55,7 +55,9 @@ def submit_document( returned in get_tree/get_ocr responses and list_documents entries. Defaults to None. Returns: - dict: {'doc_id': ...} + dict: {'doc_id': ...} — plus 'name', the stored document name + (a taken name gains a numeric suffix), when the server + returns it. """ data = {'if_retrieval': True} if mode is not None: diff --git a/pageindex/local_api.py b/pageindex/local_api.py index 5820656ac..9ad909ad0 100644 --- a/pageindex/local_api.py +++ b/pageindex/local_api.py @@ -97,6 +97,9 @@ def submit_document( raise PageIndexAPIError( "Failed to submit document: PDF has no content. All pages are blank." ) + # Fail before paying for indexing when _1.._99 are all taken; the + # binding name resolution happens again at save. + self._unique_doc_name(os.path.basename(file_path)) try: if mode == "flash": @@ -129,7 +132,7 @@ def submit_document( from .utils import remove_fields self._store.save_document( doc_id, meta, remove_fields(structure, fields=["text"]), pages) - return {"doc_id": doc_id} + return {"doc_id": doc_id, "name": meta["name"]} def _unique_doc_name(self, name: str) -> str: """Mirror the cloud upload: a taken name gets _1.._99 appended, diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 5039bd031..a9cb4cd7c 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -974,3 +974,12 @@ def test_submit_without_wait_does_not_poll(fake_cloud_client): cloud = fake_cloud_client(["processing"]) assert cloud.submit_document("whatever.pdf") == {"doc_id": "pi-fake"} assert cloud._api.polls == 0 + + +def test_submit_warns_when_stored_name_differs(fake_cloud_client): + cloud = fake_cloud_client(["processing"]) + cloud._api.submit_document = lambda **kwargs: { + "doc_id": "pi-fake", "name": "whatever_1.pdf"} + with pytest.warns(UserWarning, match='stored as "whatever_1.pdf"'): + result = cloud.submit_document("docs/whatever.pdf") + assert result["name"] == "whatever_1.pdf" diff --git a/tests/test_client.py b/tests/test_client.py index 3c9385831..bdbcd713f 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -169,12 +169,15 @@ def fake_page_index_main(doc, opt=None, logger=None): return {"doc_name": "sample.pdf", "doc_description": "d", "structure": json.loads(json.dumps(STRUCTURE))} monkeypatch.setattr(page_index_module, "page_index_main", fake_page_index_main) - first = local_client.submit_document(sample_pdf)["doc_id"] - second = local_client.submit_document(sample_pdf)["doc_id"] + first = local_client.submit_document(sample_pdf) + assert first["name"] == "sample.pdf" + with pytest.warns(UserWarning, match='stored as "sample_1.pdf"'): + second = local_client.submit_document(sample_pdf) + assert second["name"] == "sample_1.pdf" names = {d["id"]: d["name"] for d in local_client.list_documents()["documents"]} - assert names[first] == "sample.pdf" - assert names[second] == "sample_1.pdf" + assert names[first["doc_id"]] == "sample.pdf" + assert names[second["doc_id"]] == "sample_1.pdf" def test_submit_duplicate_name_exhaustion(local_client, monkeypatch): @@ -186,6 +189,22 @@ def test_submit_duplicate_name_exhaustion(local_client, monkeypatch): api._unique_doc_name("x.pdf") +def test_submit_name_exhaustion_rejects_before_indexing( + local_client, sample_pdf, monkeypatch, +): + api = local_client._api + metas = ([{"name": "sample.pdf"}] + + [{"name": f"sample_{num}.pdf"} for num in range(1, 100)]) + monkeypatch.setattr(api._store, "list_metas", lambda: metas) + monkeypatch.setattr( + page_index_module, "page_index_main", + lambda *args, **kwargs: pytest.fail( + "indexer ran despite name exhaustion"), + ) + with pytest.raises(PageIndexAPIError, match="Too many files"): + local_client.submit_document(sample_pdf) + + def test_submit_flash(local_client, sample_pdf, monkeypatch): calls = {} def fake_flash(pdf, summary=True, summary_model=None, **kwargs): @@ -456,7 +475,8 @@ def test_torn_delete_never_lists_ghost(local_client, indexed_doc, tmp_path): def test_corrupt_doc_json_is_contained(local_client, indexed_doc, sample_pdf, tmp_path): - second = local_client.submit_document(sample_pdf)["doc_id"] + with pytest.warns(UserWarning): # same-name resubmit → stored as sample_1.pdf + second = local_client.submit_document(sample_pdf)["doc_id"] (tmp_path / "store" / "docs" / indexed_doc / "doc.json").write_text("{truncated") # manifest still holds a good copy of the meta — served consistently From bf9e6dac5f5359febaf30b3f231ac63c2af6fcd9 Mon Sep 17 00:00:00 2001 From: Ray Date: Tue, 11 Aug 2026 22:03:05 +0800 Subject: [PATCH 07/18] fix: add missing page_list kwarg in duplicate-name test mock --- README.md | 1 + tests/test_client.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 5dbafc141..f4278a598 100644 --- a/README.md +++ b/README.md @@ -262,6 +262,7 @@ tools = client.agent_tools() # local: built-in tools; ``` Neither framework is a required dependency — each is imported only when its method is called. Claude Code / Cursor and other MCP hosts connect to cloud documents via the hosted MCP server directly (no SDK needed); see the [MCP docs](https://docs.pageindex.ai/mcp). + ## 🚀 Agentic Vectorless RAG: An Example For a simple, end-to-end **agentic vectorless RAG** example using **self-hosted PageIndex** (with OpenAI Agents SDK), see [`examples/agentic_vectorless_rag_demo.py`](examples/agentic_vectorless_rag_demo.py). diff --git a/tests/test_client.py b/tests/test_client.py index bdbcd713f..f55d519e6 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -165,7 +165,7 @@ def fake_page_index_main(doc, opt=None, logger=None, page_list=None): def test_submit_duplicate_name_gets_suffix(local_client, sample_pdf, monkeypatch): """Mirror the cloud upload: a second submit of the same file name is stored as name_1, not as a same-name duplicate.""" - def fake_page_index_main(doc, opt=None, logger=None): + def fake_page_index_main(doc, opt=None, logger=None, page_list=None): return {"doc_name": "sample.pdf", "doc_description": "d", "structure": json.loads(json.dumps(STRUCTURE))} monkeypatch.setattr(page_index_module, "page_index_main", fake_page_index_main) From dece6e61fc5a5bed9b9739dfbc7a2e7eef8b068d Mon Sep 17 00:00:00 2001 From: Ray Date: Tue, 11 Aug 2026 22:08:56 +0800 Subject: [PATCH 08/18] =?UTF-8?q?revert:=20keep=20README.md=20unchanged=20?= =?UTF-8?q?from=20main=20=E2=80=94=20SDK=20section=20deferred?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 60 ++----------------------------------------------------- 1 file changed, 2 insertions(+), 58 deletions(-) diff --git a/README.md b/README.md index f4278a598..5ce0ca5e6 100644 --- a/README.md +++ b/README.md @@ -207,69 +207,13 @@ python3 run_pageindex.py --md_path /path/to/your/document.md > > Add `--optimize` to refine the tree structure for more efficient retrieval (with an LLM expansion pass). -## 🐍 Python SDK: Cloud & Local - -The `pageindex` package on PyPI is the Python SDK for the [PageIndex API](https://docs.pageindex.ai) — and the same client now also runs fully **locally**, powered by this repo's indexing pipeline (including Flash). - -```bash -pip3 install --upgrade pageindex # local mode ships in pageindex >= 0.2.9; earlier versions are cloud-only -``` - -```python -from pageindex import PageIndexClient - -client = PageIndexClient(api_key="YOUR_PAGEINDEX_API_KEY") # cloud: managed OCR, tree building, retrieval -client = PageIndexClient() # local: same methods on your machine, using your LLM key (e.g. OPENAI_API_KEY) - -doc_id = client.submit_document("doc.pdf")["doc_id"] # local mode blocks until indexing finishes -doc_id = client.submit_document("doc.pdf", mode="flash")["doc_id"] # local mode with PageIndex Flash - -tree = client.get_tree(doc_id, node_summary=True)["result"] - -answer = client.chat_completions( - messages=[{"role": "user", "content": "Summarize the key findings"}], - doc_id=doc_id, -)["choices"][0]["message"]["content"] -``` - -Local documents are stored as plain JSON under `./.pageindex` (configurable via `storage_path`). Local mode supports PDFs; folders, `beta_headers`, `enable_citations`, and the deprecated retrieval API (`submit_query`/`get_retrieval`) remain cloud-only — each method's docstring spells out the differences. To pin the mode at construction instead of inferring it from `api_key`, use `PageIndexCloudClient` (fails without a real key) or `PageIndexLocalClient` (has no key parameter). - -### 🤖 Agent integration - -The client exposes its documents as **agent tools**, following one rule: **cloud clients always serve the live tool set of the [PageIndex MCP server](https://docs.pageindex.ai/mcp)** (search, folders, images — as enabled for your key, discovered dynamically; management tools like delete/upload sit behind `include_management=True` or the framework's approval layer), while local clients serve the same contract's built-in navigation subset (`browse_documents`, `get_document`, `get_document_structure`, `get_page_content`). Tool names and schemas are shared, so agent prompts port unchanged, and switching local ↔ cloud is just the client constructor line: - -```python -client = PageIndexLocalClient() # or PageIndexCloudClient(api_key=...) -client.submit_document("doc.pdf", wait=True) # wait=True: return once the doc is ready (both modes) - -# OpenAI Agents SDK (pip install "pageindex[openai]") -agent = Agent( - name="PageIndex", - instructions=client.agent_instructions(), # retrieval playbook for the agent's system prompt - tools=client.as_openai_tools(), # local: in-process tools; cloud: the full cloud MCP tool set (any model backend) -) # cloud + OpenAI models: hosted=True runs tool calls server-side (fastest) - -# Claude Agent SDK (pip install "pageindex[claude]") -options = ClaudeAgentOptions( - system_prompt=client.agent_instructions(), - mcp_servers={"pageindex": client.as_claude_mcp()}, # local: in-process server; cloud: connects to api.pageindex.ai/mcp - allowed_tools=["mcp__pageindex__*"], -) - -# Any other framework: plain functions, wrap with your framework's one-liner -tools = client.agent_tools() # local: built-in tools; cloud: full live tool set over MCP - # e.g. [StructuredTool.from_function(f) for f in tools] -``` - -Neither framework is a required dependency — each is imported only when its method is called. Claude Code / Cursor and other MCP hosts connect to cloud documents via the hosted MCP server directly (no SDK needed); see the [MCP docs](https://docs.pageindex.ai/mcp). - ## 🚀 Agentic Vectorless RAG: An Example For a simple, end-to-end **agentic vectorless RAG** example using **self-hosted PageIndex** (with OpenAI Agents SDK), see [`examples/agentic_vectorless_rag_demo.py`](examples/agentic_vectorless_rag_demo.py). ```bash -# Install with the OpenAI Agents SDK extra -pip3 install "pageindex[openai]" +# Install optional dependency +pip3 install openai-agents # Run the demo python3 examples/agentic_vectorless_rag_demo.py From 3c37cdc51039b00e413cccf0e983eac0bac8280f Mon Sep 17 00:00:00 2001 From: Ray Date: Tue, 11 Aug 2026 22:43:43 +0800 Subject: [PATCH 09/18] feat: serve cloud agent instructions live from the MCP server MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The cloud MCP server publishes its agent instructions in the initialize result, adapted to each key's tool set. agent_instructions() previously returned the SDK's local-subset text in both modes — a silently forked copy that lacks the guidance for cloud-only tools (search_documents escalation, folders, images) and drifts as the server's prompt evolves. Cloud clients now serve the server's live instructions, captured from the initialize handshake on a per-client bridge shared with agent_tools() (one session, no extra request). An empty server response raises instead of silently substituting the subset text — same posture as the annotation-regression guard. The local constant stays as the honest subset for the in-process tools, with its provenance noted and a consistency test that every tool it names exists in the local registry. --- pageindex/agent_tools.py | 47 +++++++++++++++++++++++----- pageindex/client.py | 6 ++++ pageindex/mcp_bridge.py | 13 ++++++-- tests/test_agent_tools.py | 65 ++++++++++++++++++++++++++++++++++++++- 4 files changed, 120 insertions(+), 11 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index 68af57eaf..f1628d190 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -1207,12 +1207,22 @@ def proxy(**kwargs: Any) -> str: return proxy +def _cloud_bridge(client): + """One bridge per client instance: tool discovery and instructions share + a single MCP session.""" + bridge = getattr(client, "_mcp_bridge", None) + if bridge is None: + from .mcp_bridge import McpBridge + bridge = McpBridge( + f"{client.BASE_URL}/mcp", + {"Authorization": f"Bearer {client.api_key}"}, + ) + client._mcp_bridge = bridge + return bridge + + def _build_cloud_agent_tools(client, include_management: bool) -> list[Callable[..., str]]: - from .mcp_bridge import McpBridge - bridge = McpBridge( - f"{client.BASE_URL}/mcp", - {"Authorization": f"Bearer {client.api_key}"}, - ) + bridge = _cloud_bridge(client) tools_meta = bridge.list_tools() if not include_management: # Plain functions have no framework permission layer, so the @@ -1294,6 +1304,11 @@ def remove_document(doc_names: list[str], # ── agent instructions ── +# Local subset of the cloud MCP server's initialize instructions (its +# no-folders variant), trimmed to the tools that exist here: the +# search_documents escalation steps, get_document_image, and the shared +# read-only-folders block are removed. Cloud clients receive the server's +# live instructions instead — see _base_instructions(). _INSTRUCTIONS_HEADER = ( "PageIndex by Vectify AI is a document platform for uploading and " @@ -1344,16 +1359,32 @@ def remove_document(doc_names: list[str], ]) +def _base_instructions(client) -> str: + """Cloud: the live instructions the MCP server serves for this key's + tool set. Local: the built-in subset instructions.""" + if not getattr(client, "api_key", None): + return AGENT_INSTRUCTIONS + instructions = _cloud_bridge(client).instructions() + if not instructions: + raise PageIndexAPIError( + "The MCP server returned no agent instructions — refusing to " + "substitute the SDK's local-subset guidance, which does not " + "cover the cloud tool set." + ) + return instructions + + def build_agent_instructions(client, doc_id=None) -> str: """Orchestration guidance for document QA agents; with doc_id, appends the target documents and directs the agent to work within them. Raises when a doc_id's name is shadowed by a newer same-name document — the name-addressed tools could not reach it.""" + base = _base_instructions(client) if doc_id is None: - return AGENT_INSTRUCTIONS + return base doc_ids = [doc_id] if isinstance(doc_id, str) else list(doc_id) if not doc_ids: - return AGENT_INSTRUCTIONS + return base details = [client.get_document(one_id) for one_id in doc_ids] documents = _all_documents(client) for one_id, detail in zip(doc_ids, details): @@ -1383,4 +1414,4 @@ def build_agent_instructions(client, doc_id=None) -> str: "Use these documents' names to retrieve their content with " "get_document_structure() and get_page_content()." ) - return AGENT_INSTRUCTIONS + "\n\n" + block + return base + "\n\n" + block diff --git a/pageindex/client.py b/pageindex/client.py index ff44aa376..70608a480 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -513,6 +513,12 @@ def agent_instructions(self, doc_id: Optional[Union[str, list[str]]] = None) -> Orchestration guidance for document QA agents — pass as the agent's system prompt (or append to your own). + Cloud: the live instructions the PageIndex MCP server serves for + your key's tool set, fetched over the same session as + ``agent_tools()`` — server-side guidance updates arrive without an + SDK release. Raises PageIndexAPIError if the server cannot be + reached. Local: the built-in guidance for the in-process tools. + With ``doc_id`` (str or list, same shape as ``chat_completions``), appends the target documents' names and metadata and directs the agent to work within them. Raises PageIndexAPIError if a doc_id does diff --git a/pageindex/mcp_bridge.py b/pageindex/mcp_bridge.py index fad0baaf8..95aba5c70 100644 --- a/pageindex/mcp_bridge.py +++ b/pageindex/mcp_bridge.py @@ -1,7 +1,9 @@ """Minimal MCP client (streamable HTTP) for the PageIndex cloud MCP server. -Backs the cloud branch of ``client.agent_tools()``: ``tools/list`` discovers -the live tool set, ``tools/call`` executes a tool. Synchronous, requests-only. +Backs the cloud branches of ``client.agent_tools()`` and +``client.agent_instructions()``: ``tools/list`` discovers the live tool set, +``tools/call`` executes a tool, and the ``initialize`` handshake carries the +server's agent instructions. Synchronous, requests-only. Works against both stateful and stateless servers: a session id returned by ``initialize`` is echoed back, and a request rejected after session expiry re-initializes once and retries. @@ -43,6 +45,7 @@ def __init__(self, url: str, headers: dict[str, str]): self._auth_headers = dict(headers) self._session_id: Optional[str] = None self._protocol_version: Optional[str] = None + self._instructions: Optional[str] = None self._initialized = False self._lock = threading.RLock() self._next_id = 0 @@ -147,6 +150,7 @@ def _ensure_initialized(self) -> None: self._session_id = response.headers.get("Mcp-Session-Id") self._protocol_version = result.get("protocolVersion", _PROTOCOL_VERSION) + self._instructions = result.get("instructions") self._initialized = True try: self._post({"jsonrpc": "2.0", @@ -156,6 +160,11 @@ def _ensure_initialized(self) -> None: # ── public surface ── + def instructions(self) -> Optional[str]: + """The server's agent instructions from the initialize handshake.""" + self._ensure_initialized() + return self._instructions + def list_tools(self) -> list[dict]: tools: list[dict] = [] cursor: Optional[str] = None diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index a9cb4cd7c..b7677f1d7 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -2,6 +2,7 @@ local store (no LLM calls; one live parity test gated on PAGEINDEX_API_KEY).""" import json import os +import re import sys from pathlib import Path @@ -594,7 +595,8 @@ def fake_post(url, json=None, headers=None, timeout=None): rid = json.get("id") if method == "initialize": return _Resp(200, {"jsonrpc": "2.0", "id": rid, - "result": {"protocolVersion": "2025-06-18"}}, + "result": {"protocolVersion": "2025-06-18", + "instructions": "SERVER GUIDANCE"}}, {"Content-Type": "application/json", "Mcp-Session-Id": "sess-1"}) if method == "notifications/initialized": @@ -625,6 +627,10 @@ def fake_post(url, json=None, headers=None, timeout=None): tools = bridge.list_tools() assert tools == [{"name": "t1", "description": "reads — never writes"}] + # Captured during the handshake — serving it must not post again. + posts_before = len(posts) + assert bridge.instructions() == "SERVER GUIDANCE" + assert len(posts) == posts_before list_headers = posts[-1]["headers"] assert list_headers["Mcp-Session-Id"] == "sess-1" assert list_headers["MCP-Protocol-Version"] == "2025-06-18" @@ -891,6 +897,16 @@ def test_live_cloud_contract_parity(): assert (real.get("annotations") or {}).get(key) == value, (name, key) +@pytest.mark.skipif(not LIVE_KEY, reason="PAGEINDEX_API_KEY not set") +def test_live_cloud_instructions_nonempty(): + """The empty-instructions guard raises for cloud clients; the real + server must actually serve instructions in its initialize result.""" + from pageindex.mcp_bridge import McpBridge + bridge = McpBridge("https://api.pageindex.ai/mcp", + {"Authorization": f"Bearer {LIVE_KEY}"}) + assert bridge.instructions() + + # ── agent_instructions ── def test_agent_instructions_default(client): @@ -916,6 +932,53 @@ def test_agent_instructions_with_doc_id(client, store_path): client.agent_instructions(doc_id="pi-missing") +def test_local_instructions_name_only_local_tools(): + """The local instructions are trimmed from the cloud server's; every + tool they name must exist in the local registry, or the trim drifted.""" + named = set(re.findall(r"\b(\w+)\(", AGENT_INSTRUCTIONS)) + assert named + assert named <= set(tool_names(include_management=True)) + + +def test_cloud_agent_instructions_served_live(monkeypatch): + """Cloud clients serve the server's live instructions from the MCP + initialize handshake — over the same bridge session as agent_tools().""" + import pageindex.mcp_bridge as mcp_bridge + created = [] + + class _Bridge(_FakeBridge): + def __init__(self, url, headers): + super().__init__(url, headers) + created.append(self) + + def instructions(self): + return "LIVE CLOUD GUIDANCE" + + monkeypatch.setattr(mcp_bridge, "McpBridge", _Bridge) + cloud = PageIndexCloudClient(api_key="pi-test-key") + cloud.agent_tools() + assert cloud.agent_instructions() == "LIVE CLOUD GUIDANCE" + assert len(created) == 1 + + +def test_cloud_agent_instructions_empty_raises(monkeypatch): + """An empty server response must raise, not silently substitute the + subset guidance — same posture as the annotation-regression guard.""" + import pageindex.mcp_bridge as mcp_bridge + + class _SilentBridge: + def __init__(self, url, headers): + pass + + def instructions(self): + return None + + monkeypatch.setattr(mcp_bridge, "McpBridge", _SilentBridge) + cloud = PageIndexCloudClient(api_key="pi-test-key") + with pytest.raises(PageIndexAPIError, match="no agent instructions"): + cloud.agent_instructions() + + # ── submit_document(wait=True) ── class _FakeCloudAPI: From 50fd61862e062a49fe2670b0d32671605cb73e18 Mon Sep 17 00:00:00 2001 From: Ray Date: Tue, 11 Aug 2026 23:01:04 +0800 Subject: [PATCH 10/18] fix: local relevance sort answers honestly instead of imitating sort="relevance" is cloud-side semantic ranking; the local substring imitation could satisfy the letter of the interface while silently missing semantically relevant documents. Per the honest-subset rule (same treatment as folders), local now returns the "not available here" envelope for sort="relevance" or a stray query, and the local instructions steer discovery through name/description matching plus full-library paging instead of prescribing a capability that does not exist here. The tool schema keeps the cloud contract verbatim, like folder_id: honesty lives in the runtime answer, not a forked contract. --- pageindex/agent_tools.py | 85 +++++++++++++++------------------------ tests/test_agent_tools.py | 20 ++++----- 2 files changed, 43 insertions(+), 62 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index f1628d190..a8c15d761 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -636,16 +636,19 @@ def _browse_documents(client, folder_id: str = "root", recursive: bool = False, {"summary": "Invalid sort mode", "options": ['Use sort="time" or sort="relevance"']}, "INVALID_INPUT") - if sort == "relevance" and not query: - return _failure('query is required when sort is "relevance"', None, - {"summary": "Missing query for relevance ranking", - "options": ['Pass query alongside sort="relevance"']}, - "INVALID_INPUT") - if sort == "time" and query: - return _failure('query is only allowed when sort is "relevance"', None, - {"summary": "query does not apply to the time sort", - "options": ["Drop query, or set sort=\"relevance\""]}, - "INVALID_INPUT") + if sort == "relevance" or query: + # Semantic ranking is a cloud capability; like folders, it is not + # imitated here. + return _failure( + "Relevance ranking is not available here — use the default " + "time sort.", None, + {"summary": "Semantic ranking is not available in this library", + "options": ["Retry without sort/query and match the returned " + "names and descriptions against the intent yourself", + "Page through the full library with " + "`offset: next_offset`"]}, + "INVALID_INPUT", + ) try: offset = max(int(offset), 0) limit = min(max(int(limit), 1), 50) @@ -655,23 +658,9 @@ def _browse_documents(client, folder_id: str = "root", recursive: bool = False, "options": ["Pass integer offset and limit values"]}, "INVALID_INPUT") - if sort == "relevance": - tokens = [token for token in (query or "").lower().split() if token] - scored = [] - for doc in _all_documents(client): - haystack = f"{doc.get('name') or ''} {doc.get('description') or ''}".lower() - score = sum(1 for token in tokens if token in haystack) - if score: - scored.append((score, doc)) - # Stable sort: equal scores keep the newest-first listing order. - scored.sort(key=lambda pair: pair[0], reverse=True) - ranked = [doc for _, doc in scored] - window = ranked[offset:offset + limit] - has_more = offset + limit < len(ranked) - else: - listing = client.list_documents(limit=limit, offset=offset) - window = listing.get("documents") or [] - has_more = offset + limit < listing.get("total", 0) + listing = client.list_documents(limit=limit, offset=offset) + window = listing.get("documents") or [] + has_more = offset + limit < listing.get("total", 0) next_offset = offset + limit if has_more else None page_has_processing = False @@ -706,18 +695,9 @@ def _browse_documents(client, folder_id: str = "root", recursive: bool = False, if not items and offset == 0: next_steps = { "summary": "Nothing to show", - "options": ( - ["No documents matched this query. Rephrase with synonyms or " - "alternative terms and retry browse_documents(sort=\"relevance\")."] - if sort == "relevance" - else ["Nothing here. Index documents with " - "PageIndexClient.submit_document() to get started."] - ), - "auto_retry": ( - "Rephrase the query and retry browse_documents(sort=\"relevance\")" - if sort == "relevance" - else "Index a document with submit_document() to get started" - ), + "options": ["Nothing here. Index documents with " + "PageIndexClient.submit_document() to get started."], + "auto_retry": "Index a document with submit_document() to get started", } return _success(data, next_steps) @@ -727,9 +707,8 @@ def _browse_documents(client, folder_id: str = "root", recursive: bool = False, options.append( "Results returned ≠ correct results. Verify these documents match " "the user's actual intent (topic, time period, document type) " - "before proceeding. If they do not match, rephrase the query and " - "retry browse_documents(sort=\"relevance\"). Do NOT use general " - "knowledge as a substitute." + "before proceeding. If they do not match, page through the rest " + "of the library. Do NOT use general knowledge as a substitute." ) if page_has_processing: options.append("Some documents on this page are still processing. " @@ -1305,10 +1284,12 @@ def remove_document(doc_names: list[str], # ── agent instructions ── # Local subset of the cloud MCP server's initialize instructions (its -# no-folders variant), trimmed to the tools that exist here: the -# search_documents escalation steps, get_document_image, and the shared -# read-only-folders block are removed. Cloud clients receive the server's -# live instructions instead — see _base_instructions(). +# no-folders variant), trimmed to what exists here: the search_documents +# escalation steps, get_document_image, and the shared read-only-folders +# block are removed, and the sort="relevance" guidance is replaced with +# name/description matching (semantic ranking is cloud-side). Cloud +# clients receive the server's live instructions instead — see +# _base_instructions(). _INSTRUCTIONS_HEADER = ( "PageIndex by Vectify AI is a document platform for uploading and " @@ -1328,12 +1309,12 @@ def remove_document(doc_names: list[str], _DISCOVERY = """\ DOCUMENT DISCOVERY: -- browse_documents() — DEFAULT discovery tool, first choice for any document-related question. The bare call returns your documents. Use sort="relevance" + query for semantic ranking.""" +- browse_documents() — DEFAULT discovery tool, first choice for any document-related question. It lists your documents newest first with names and descriptions; match them against the user's intent, and page through with `offset: next_offset` while has_more is true.""" _DECISION = """\ DECISION: -- "What do I have / list / recent" → browse_documents (time) -- ANY question that needs a document to answer (including "find THE paper about Y") → browse_documents(sort="relevance", query=…)""" +- "What do I have / list / recent" → browse_documents() +- ANY question that needs a document to answer (including "find THE paper about Y") → browse_documents(), then pick the documents whose name/description matches the question""" _AFTER_DISCOVERY = """\ - Skip discovery ONLY for questions with NO possible document connection (e.g., "capital of France"). @@ -1343,9 +1324,9 @@ def remove_document(doc_names: list[str], _PERSISTENCE = """\ PERSISTENCE (before concluding the target document is not in the library): This protocol applies both when results are empty AND when results are returned but none match the user's intent. Do NOT give up after a single discovery attempt. Follow these steps in order: -1. browse_documents(sort="relevance", query=…) with the original intent -2. Rephrase the query with synonyms or alternative terms → browse_documents(sort="relevance") again -3. browse_documents(recursive=true) to flatten the library into one list — MANDATORY, must be attempted at least once before concluding "not found" +1. browse_documents() and compare every returned name/description against the user's intent +2. Page through the ENTIRE library with `offset: next_offset` until has_more is false — MANDATORY, must be completed before concluding "not found" +3. Re-scan for loose matches: synonyms, abbreviations, and partial titles in names/descriptions can identify the target Only after ALL three steps have been tried may you conclude the document is not in the library. Do NOT fall back to general knowledge — if the user's question references their own documents, exhaust every discovery path first.""" AGENT_INSTRUCTIONS = "\n\n".join([ diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index b7677f1d7..69c8e9fee 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -119,21 +119,20 @@ def test_browse_documents_pagination(client, store_path): assert second["has_more"] is False -def test_browse_documents_relevance(client, store_path): - seed_doc(store_path, "pi-a", "annual-report.pdf", - description="Financial results for the year") - seed_doc(store_path, "pi-b", "attention.pdf", +def test_browse_documents_relevance_unsupported(client, store_path): + """Semantic ranking is cloud-side; like folders, local answers with an + honest error instead of a keyword imitation.""" + seed_doc(store_path, "pi-a", "attention.pdf", description="Transformers and attention mechanisms") payload, is_error = run(client, "browse_documents", sort="relevance", query="attention transformers") - assert not is_error - assert [d["name"] for d in payload["documents"]] == ["attention.pdf"] - assert payload["sort"] == "relevance" + assert is_error and payload["errorCode"] == "INVALID_INPUT" + assert "not available" in payload["error"] - missing_query, is_error = run(client, "browse_documents", sort="relevance") - assert is_error and missing_query["errorCode"] == "INVALID_INPUT" stray_query, is_error = run(client, "browse_documents", query="x") - assert is_error and "relevance" in stray_query["error"] + assert is_error and "not available" in stray_query["error"] + bad_sort, is_error = run(client, "browse_documents", sort="banana") + assert is_error and bad_sort["errorCode"] == "INVALID_INPUT" def test_browse_documents_empty_and_folder_error(client): @@ -916,6 +915,7 @@ def test_agent_instructions_default(client): assert "browse_documents" in text assert "search_documents" not in text assert "get_folder_structure" not in text + assert 'sort="relevance"' not in text # cloud-side capability def test_agent_instructions_with_doc_id(client, store_path): From f1301e7241cc313585b5f8be0070f49948182fb7 Mon Sep 17 00:00:00 2001 From: Ray Date: Tue, 11 Aug 2026 23:07:38 +0800 Subject: [PATCH 11/18] docs: note the cloud+Claude instructions duplication trade-off in as_claude_mcp --- pageindex/client.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/pageindex/client.py b/pageindex/client.py index 70608a480..a4f475ed1 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -498,6 +498,12 @@ def as_claude_mcp(self, include_management: bool = False): tools (requires ``claude-agent-sdk``; ``pip install 'pageindex[claude]'``). + Cloud hosts that surface MCP server instructions receive the same + guidance ``agent_instructions()`` returns natively — passing both + duplicates the text (harmless). ``system_prompt`` stays the + recommended channel: it is guaranteed delivery, carries ``doc_id`` + targeting, and is the only channel local mode has. + Usage:: options = ClaudeAgentOptions( From 20ed81f6b737005bdc9b1674d8d7fee76d365e66 Mon Sep 17 00:00:00 2001 From: Ray Date: Tue, 11 Aug 2026 23:17:41 +0800 Subject: [PATCH 12/18] fix: unsupported-capability envelopes say local-mode-yet, point to cloud "Not available here" read as a broken feature; the honest framing is that folders and semantic ranking exist on PageIndex cloud and are not in local mode yet. Both envelopes now say so and name the cloud client in next_steps, so agents relay an accurate story to the user. --- pageindex/agent_tools.py | 18 +++++++++++------- tests/test_agent_tools.py | 4 ++-- 2 files changed, 13 insertions(+), 9 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index a8c15d761..e0c813576 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -463,12 +463,14 @@ def _not_ready_error(doc_name: str, status: Any, operation: str, def _folder_unsupported(param: str) -> tuple[dict, bool]: return _failure( - f"Folders are not available here — omit {param}.", + f"Folders are not supported in local mode yet — omit {param}.", None, { - "summary": "This library has no folders", + "summary": "This local library does not have folders yet", "options": ["Retry the call without a folder_id", - "Use browse_documents() to list the library root"], + "Use browse_documents() to list the library root", + "Folders are available on PageIndex cloud " + "(PageIndexCloudClient with an API key)"], }, "INVALID_INPUT", ) @@ -640,13 +642,15 @@ def _browse_documents(client, folder_id: str = "root", recursive: bool = False, # Semantic ranking is a cloud capability; like folders, it is not # imitated here. return _failure( - "Relevance ranking is not available here — use the default " - "time sort.", None, - {"summary": "Semantic ranking is not available in this library", + "Relevance ranking is not supported in local mode yet — use " + "the default time sort.", None, + {"summary": "This local library does not have semantic ranking yet", "options": ["Retry without sort/query and match the returned " "names and descriptions against the intent yourself", "Page through the full library with " - "`offset: next_offset`"]}, + "`offset: next_offset`", + "Semantic ranking is available on PageIndex cloud " + "(PageIndexCloudClient with an API key)"]}, "INVALID_INPUT", ) try: diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 69c8e9fee..3a44cfeae 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -127,10 +127,10 @@ def test_browse_documents_relevance_unsupported(client, store_path): payload, is_error = run(client, "browse_documents", sort="relevance", query="attention transformers") assert is_error and payload["errorCode"] == "INVALID_INPUT" - assert "not available" in payload["error"] + assert "not supported in local mode" in payload["error"] stray_query, is_error = run(client, "browse_documents", query="x") - assert is_error and "not available" in stray_query["error"] + assert is_error and "not supported in local mode" in stray_query["error"] bad_sort, is_error = run(client, "browse_documents", sort="banana") assert is_error and bad_sort["errorCode"] == "INVALID_INPUT" From 8dc929ff1f8896aef28f5aa962c77f07a63f00a9 Mon Sep 17 00:00:00 2001 From: Ray Date: Tue, 11 Aug 2026 23:25:54 +0800 Subject: [PATCH 13/18] fix: local tool descriptions pre-announce cloud-only capabilities MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The cloud-verbatim browse_documents description invites sort="relevance" and folder drilling, so a local agent's first semantic search attempt was a guaranteed dead end discovered only from the runtime error envelope. Local registration now appends a LOCAL MODE note to the description — the agent learns what is cloud-only before calling; the runtime envelope stays as the backstop for prompts that ignore descriptions. The cloud-facing contract stays byte-verbatim. --- pageindex/agent_tools.py | 23 +++++++++++++++++++--- pageindex/integrations/claude_agent_sdk.py | 5 +++-- tests/test_agent_tools.py | 10 ++++++++++ 3 files changed, 33 insertions(+), 5 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index e0c813576..a9c89be8e 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -1107,10 +1107,27 @@ def _tool_docstring(description: str, properties: dict[str, Any]) -> str: return "\n".join(lines) +#: Appended to the cloud-verbatim description when a tool is served locally, +#: so the agent learns what is cloud-only before calling instead of from the +#: runtime error envelope. +_LOCAL_DESCRIPTION_NOTES = { + "browse_documents": ( + 'LOCAL MODE: folder_id and sort="relevance"/query are not supported ' + "yet (they work on PageIndex cloud) — use the default time sort and " + "page with offset." + ), +} + + +def _local_description(name: str) -> str: + description = TOOL_CONTRACT[name]["description"] + note = _LOCAL_DESCRIPTION_NOTES.get(name) + return f"{description}\n\n{note}" if note else description + + def _docstring(name: str) -> str: - contract = TOOL_CONTRACT[name] - return _tool_docstring(contract["description"], - contract["schema"]["properties"]) + return _tool_docstring(_local_description(name), + TOOL_CONTRACT[name]["schema"]["properties"]) _SCHEMA_TYPE_MAP = {"string": str, "integer": int, "number": float, diff --git a/pageindex/integrations/claude_agent_sdk.py b/pageindex/integrations/claude_agent_sdk.py index 0fb77d2de..77cc3d7a3 100644 --- a/pageindex/integrations/claude_agent_sdk.py +++ b/pageindex/integrations/claude_agent_sdk.py @@ -28,7 +28,8 @@ def build_claude_mcp(client, include_management: bool = False): "as_claude_mcp in local mode requires the Claude Agent SDK — " "pip install claude-agent-sdk (or pip install 'pageindex[claude]')." ) from exc - from ..agent_tools import TOOL_CONTRACT, call_tool, tool_names + from ..agent_tools import (TOOL_CONTRACT, _local_description, call_tool, + tool_names) def make_handler(name: str): async def handler(arguments: dict[str, Any]) -> dict[str, Any]: @@ -52,7 +53,7 @@ def tool_kwargs(name: str) -> dict: return {"annotations": ToolAnnotations(**annotations)} tools = [ - tool(name, TOOL_CONTRACT[name]["description"], + tool(name, _local_description(name), TOOL_CONTRACT[name]["schema"], **tool_kwargs(name))(make_handler(name)) for name in tool_names(include_management) ] diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 3a44cfeae..41f710e7d 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -119,6 +119,16 @@ def test_browse_documents_pagination(client, store_path): assert second["has_more"] is False +def test_local_docstrings_preannounce_cloud_only_capabilities(client): + """The cloud-verbatim description invites sort="relevance" and folder + drilling; the local registration appends a LOCAL MODE note so the agent + learns the dead ends before calling, not from the runtime error.""" + browse = client.agent_tools()[0] + assert browse.__doc__.startswith(TOOL_CONTRACT["browse_documents"]["description"]) + assert "LOCAL MODE" in browse.__doc__ + assert "not supported yet" in browse.__doc__ + + def test_browse_documents_relevance_unsupported(client, store_path): """Semantic ranking is cloud-side; like folders, local answers with an honest error instead of a keyword imitation.""" From 2b929eedfb5c78a0fe810522275d688b6b02d25c Mon Sep 17 00:00:00 2001 From: Ray Date: Tue, 11 Aug 2026 23:33:52 +0800 Subject: [PATCH 14/18] refactor: localized tool guidance replaces the appended LOCAL MODE note MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Appending a retraction to the cloud-verbatim description left the model parsing an instruction and its negation — and kept the cloud text recommending search_documents and get_folder_structure, tools that are not registered locally (get_page_content likewise pointed at get_document_image). Guidance now adapts to the local surface the way AGENT_INSTRUCTIONS already does: schema structure stays byte-identical to the contract (mechanically asserted by a strip-descriptions test), while local description strings teach only what works here and point to PageIndex cloud for the rest. A dead-reference test forbids local guidance from naming tools outside the local registry, so a contract refresh that reintroduces a cloud-only reference fails loudly. --- pageindex/agent_tools.py | 90 ++++++++++++++++++---- pageindex/integrations/claude_agent_sdk.py | 6 +- tests/test_agent_tools.py | 55 +++++++++---- 3 files changed, 120 insertions(+), 31 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index a9c89be8e..ed161bf50 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -1,15 +1,19 @@ """Agent tools: the cloud MCP tool contract, executed against a PageIndexClient. -Tool names, input schemas, and descriptions match the PageIndex cloud MCP -server, so agent prompts work unchanged across the cloud MCP connection and -this in-process layer. Only the tools that exist in every mode are registered -(no folders, search_documents, or get_document_image). +Tool names and input-schema structure match the PageIndex cloud MCP server, +so agent prompts work unchanged across the cloud MCP connection and this +in-process layer. Only the tools that exist in every mode are registered +(no folders, search_documents, or get_document_image), and the guidance +strings (tool descriptions) adapt to the local surface the same way the +agent instructions do — they never teach capabilities that only exist on +the cloud. Tools never raise: every outcome, including errors, is returned as the same JSON envelope the cloud emits ({"success": true, ...} / {"error": ...}). """ from __future__ import annotations +import copy import difflib import json import re @@ -1107,27 +1111,83 @@ def _tool_docstring(description: str, properties: dict[str, Any]) -> str: return "\n".join(lines) -#: Appended to the cloud-verbatim description when a tool is served locally, -#: so the agent learns what is cloud-only before calling instead of from the -#: runtime error envelope. -_LOCAL_DESCRIPTION_NOTES = { +# Local guidance layer: schema STRUCTURE stays byte-identical to the cloud +# contract, but description strings adapt to the local surface the same way +# AGENT_INSTRUCTIONS does — guidance must not teach capabilities (folders, +# semantic ranking) or tools (search_documents, get_document_image) that do +# not exist here. Guard tests assert both properties; a contract refresh +# that reintroduces a cloud-only reference fails the dead-reference test. + +_LOCAL_DOC_NAME_DESCRIPTION = ( + 'Copy the `name` field verbatim from a browse_documents() response ' + '(case-sensitive, include extension). Example: "Q3 Report.pdf". ' + "Document names are unique in a local library." +) +_LOCAL_FOLDER_ID_DESCRIPTION = ( + "Not needed in local mode: document names are unique and folders are " + 'not supported yet (they work on PageIndex cloud). Omit, or pass "root".' +) + +_LOCAL_DESCRIPTIONS: dict[str, str] = { "browse_documents": ( - 'LOCAL MODE: folder_id and sort="relevance"/query are not supported ' - "yet (they work on PageIndex cloud) — use the default time sort and " - "page with offset." + "Primary document retrieval tool — first choice for any " + "document-related question. Lists your documents newest first with " + "names and descriptions; match them against the user's intent and " + "page through with `offset: next_offset` while `has_more` is true. " + 'Folder browsing and semantic ranking (sort="relevance") are not ' + "supported in local mode yet — they work on PageIndex cloud." + ), + # The image sentence points at a tool that is not registered locally. + "get_page_content": TOOL_CONTRACT["get_page_content"]["description"] + .replace(" Embedded image paths in the response feed into " + "`get_document_image()`.", ""), +} + +_LOCAL_PARAM_DESCRIPTIONS: dict[tuple[str, str], str] = { + ("browse_documents", "folder_id"): _LOCAL_FOLDER_ID_DESCRIPTION, + ("browse_documents", "recursive"): ( + "Kept for cloud compatibility; a local library has no folders, so " + "recursive and non-recursive return the same documents." + ), + ("browse_documents", "sort"): ( + 'Only "time" (newest first) is supported in local mode; ' + '"relevance" is cloud-only for now.' + ), + ("browse_documents", "query"): ( + 'Cloud-only for now (semantic ranking with sort="relevance") — ' + "omit in local mode." ), + ("get_document", "doc_name"): _LOCAL_DOC_NAME_DESCRIPTION, + ("get_document", "folder_id"): _LOCAL_FOLDER_ID_DESCRIPTION, + ("get_document_structure", "doc_name"): _LOCAL_DOC_NAME_DESCRIPTION, + ("get_document_structure", "folder_id"): _LOCAL_FOLDER_ID_DESCRIPTION, + ("get_page_content", "doc_name"): _LOCAL_DOC_NAME_DESCRIPTION, + ("get_page_content", "folder_id"): _LOCAL_FOLDER_ID_DESCRIPTION, + ("remove_document", "doc_names"): ( + "Array of document names to delete. Each name must be copied " + "verbatim from the `name` field of a browse_documents() response " + '(case-sensitive, include extension). Example: ["Q3 Report.pdf", ' + '"draft.pdf"]. Max 10 per call.' + ), + ("remove_document", "folder_id"): _LOCAL_FOLDER_ID_DESCRIPTION, } def _local_description(name: str) -> str: - description = TOOL_CONTRACT[name]["description"] - note = _LOCAL_DESCRIPTION_NOTES.get(name) - return f"{description}\n\n{note}" if note else description + return _LOCAL_DESCRIPTIONS.get(name) or TOOL_CONTRACT[name]["description"] + + +def _local_schema(name: str) -> dict[str, Any]: + schema = copy.deepcopy(TOOL_CONTRACT[name]["schema"]) + for (tool_name, param), text in _LOCAL_PARAM_DESCRIPTIONS.items(): + if tool_name == name and param in schema["properties"]: + schema["properties"][param]["description"] = text + return schema def _docstring(name: str) -> str: return _tool_docstring(_local_description(name), - TOOL_CONTRACT[name]["schema"]["properties"]) + _local_schema(name)["properties"]) _SCHEMA_TYPE_MAP = {"string": str, "integer": int, "number": float, diff --git a/pageindex/integrations/claude_agent_sdk.py b/pageindex/integrations/claude_agent_sdk.py index 77cc3d7a3..b5e599da0 100644 --- a/pageindex/integrations/claude_agent_sdk.py +++ b/pageindex/integrations/claude_agent_sdk.py @@ -28,8 +28,8 @@ def build_claude_mcp(client, include_management: bool = False): "as_claude_mcp in local mode requires the Claude Agent SDK — " "pip install claude-agent-sdk (or pip install 'pageindex[claude]')." ) from exc - from ..agent_tools import (TOOL_CONTRACT, _local_description, call_tool, - tool_names) + from ..agent_tools import (TOOL_CONTRACT, _local_description, + _local_schema, call_tool, tool_names) def make_handler(name: str): async def handler(arguments: dict[str, Any]) -> dict[str, Any]: @@ -54,7 +54,7 @@ def tool_kwargs(name: str) -> dict: tools = [ tool(name, _local_description(name), - TOOL_CONTRACT[name]["schema"], **tool_kwargs(name))(make_handler(name)) + _local_schema(name), **tool_kwargs(name))(make_handler(name)) for name in tool_names(include_management) ] return create_sdk_mcp_server(name="pageindex", version=sdk_version(), diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 41f710e7d..327cfa574 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -76,10 +76,49 @@ def test_tool_surface_and_docstrings(client): with_management = client.agent_tools(include_management=True) assert [tool.__name__ for tool in with_management][-1] == "remove_document" for tool in tools: - contract = TOOL_CONTRACT[tool.__name__] - assert tool.__doc__.startswith(contract["description"]) - for param in contract["schema"]["properties"]: + for param in TOOL_CONTRACT[tool.__name__]["schema"]["properties"]: assert param in tool.__doc__ + docs = {tool.__name__: tool.__doc__ for tool in tools} + # Tools whose cloud description has no cloud-only content keep it + # verbatim; browse_documents serves the localized guidance. + assert docs["get_document"].startswith( + TOOL_CONTRACT["get_document"]["description"]) + assert docs["browse_documents"].startswith( + "Primary document retrieval tool") + + +def test_local_schema_structure_matches_contract(): + """The local guidance layer may localize description strings only — + names, types, defaults, bounds, and required stay byte-identical.""" + import copy + from pageindex.agent_tools import _local_schema + + def stripped(schema): + schema = copy.deepcopy(schema) + for spec in schema["properties"].values(): + spec.pop("description", None) + return schema + + for name, contract in TOOL_CONTRACT.items(): + assert stripped(_local_schema(name)) == stripped(contract["schema"]), name + + +def test_local_guidance_references_only_local_tools(client): + """Local descriptions must not send the agent to tools that are not + registered here (the cloud text names search_documents, + get_folder_structure, and get_document_image).""" + registered = set(tool_names(include_management=True)) + for tool in client.agent_tools(include_management=True): + named = set(re.findall(r"\b(\w+)\(", tool.__doc__)) + assert named <= registered, (tool.__name__, named - registered) + + +def test_local_guidance_points_cloud_only_capabilities_at_cloud(client): + browse = client.agent_tools()[0].__doc__ + assert "not supported in local mode yet" in browse + assert "PageIndex cloud" in browse + assert "search_documents" not in browse + assert "get_folder_structure" not in browse # ── browse_documents ── @@ -119,16 +158,6 @@ def test_browse_documents_pagination(client, store_path): assert second["has_more"] is False -def test_local_docstrings_preannounce_cloud_only_capabilities(client): - """The cloud-verbatim description invites sort="relevance" and folder - drilling; the local registration appends a LOCAL MODE note so the agent - learns the dead ends before calling, not from the runtime error.""" - browse = client.agent_tools()[0] - assert browse.__doc__.startswith(TOOL_CONTRACT["browse_documents"]["description"]) - assert "LOCAL MODE" in browse.__doc__ - assert "not supported yet" in browse.__doc__ - - def test_browse_documents_relevance_unsupported(client, store_path): """Semantic ranking is cloud-side; like folders, local answers with an honest error instead of a keyword imitation.""" From e790c375fefb7837c8b6ac36393999667c98ce0b Mon Sep 17 00:00:00 2001 From: Ray Date: Tue, 11 Aug 2026 23:41:27 +0800 Subject: [PATCH 15/18] feat: hide cloud-only parameters from the local tool surface MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit folder_id, sort, query, and recursive were exposed locally with localized "cloud-only" descriptions, leaving the dead-end calls expressible and discovered at runtime. Schema constraints beat guidance: the local surface now serves the contract minus these parameters, so strict-schema frameworks make the calls inexpressible and a prompt that insists on sort="relevance" degrades to the bare call (the correct local behavior) instead of an error round-trip. The implementations still accept the hidden parameters and answer with the guided "works on PageIndex cloud" envelope — the backstop for direct call_tool callers and hosts without schema enforcement. wait_for_completion stays: seeded or torn stores can hold documents that are genuinely not completed. The structural guard now asserts the local schema equals the contract minus the documented hidden set, descriptions aside. --- pageindex/agent_tools.py | 72 ++++++++++++++++----------------------- tests/test_agent_tools.py | 31 +++++++++++++---- 2 files changed, 54 insertions(+), 49 deletions(-) diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index ed161bf50..b6ed15c6f 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -1112,21 +1112,31 @@ def _tool_docstring(description: str, properties: dict[str, Any]) -> str: # Local guidance layer: schema STRUCTURE stays byte-identical to the cloud -# contract, but description strings adapt to the local surface the same way -# AGENT_INSTRUCTIONS does — guidance must not teach capabilities (folders, -# semantic ranking) or tools (search_documents, get_document_image) that do -# not exist here. Guard tests assert both properties; a contract refresh -# that reintroduces a cloud-only reference fails the dead-reference test. +# contract minus the hidden cloud-only parameters, and description strings +# adapt to the local surface the same way AGENT_INSTRUCTIONS does — guidance +# must not teach capabilities (folders, semantic ranking) or tools +# (search_documents, get_document_image) that do not exist here. Guard tests +# assert both properties; a contract refresh that reintroduces a cloud-only +# reference fails the dead-reference test. + +#: Cloud-only parameters hidden from the local surface — strict-schema +#: frameworks then make the dead-end calls inexpressible. The +#: implementations still accept them and answer with the guided error +#: envelope, for direct call_tool callers and hosts without schema +#: enforcement. +_LOCAL_HIDDEN_PARAMS: dict[str, tuple[str, ...]] = { + "browse_documents": ("folder_id", "recursive", "sort", "query"), + "get_document": ("folder_id",), + "get_document_structure": ("folder_id",), + "get_page_content": ("folder_id",), + "remove_document": ("folder_id",), +} _LOCAL_DOC_NAME_DESCRIPTION = ( 'Copy the `name` field verbatim from a browse_documents() response ' '(case-sensitive, include extension). Example: "Q3 Report.pdf". ' "Document names are unique in a local library." ) -_LOCAL_FOLDER_ID_DESCRIPTION = ( - "Not needed in local mode: document names are unique and folders are " - 'not supported yet (they work on PageIndex cloud). Omit, or pass "root".' -) _LOCAL_DESCRIPTIONS: dict[str, str] = { "browse_documents": ( @@ -1144,32 +1154,15 @@ def _tool_docstring(description: str, properties: dict[str, Any]) -> str: } _LOCAL_PARAM_DESCRIPTIONS: dict[tuple[str, str], str] = { - ("browse_documents", "folder_id"): _LOCAL_FOLDER_ID_DESCRIPTION, - ("browse_documents", "recursive"): ( - "Kept for cloud compatibility; a local library has no folders, so " - "recursive and non-recursive return the same documents." - ), - ("browse_documents", "sort"): ( - 'Only "time" (newest first) is supported in local mode; ' - '"relevance" is cloud-only for now.' - ), - ("browse_documents", "query"): ( - 'Cloud-only for now (semantic ranking with sort="relevance") — ' - "omit in local mode." - ), ("get_document", "doc_name"): _LOCAL_DOC_NAME_DESCRIPTION, - ("get_document", "folder_id"): _LOCAL_FOLDER_ID_DESCRIPTION, ("get_document_structure", "doc_name"): _LOCAL_DOC_NAME_DESCRIPTION, - ("get_document_structure", "folder_id"): _LOCAL_FOLDER_ID_DESCRIPTION, ("get_page_content", "doc_name"): _LOCAL_DOC_NAME_DESCRIPTION, - ("get_page_content", "folder_id"): _LOCAL_FOLDER_ID_DESCRIPTION, ("remove_document", "doc_names"): ( "Array of document names to delete. Each name must be copied " "verbatim from the `name` field of a browse_documents() response " '(case-sensitive, include extension). Example: ["Q3 Report.pdf", ' '"draft.pdf"]. Max 10 per call.' ), - ("remove_document", "folder_id"): _LOCAL_FOLDER_ID_DESCRIPTION, } @@ -1179,6 +1172,8 @@ def _local_description(name: str) -> str: def _local_schema(name: str) -> dict[str, Any]: schema = copy.deepcopy(TOOL_CONTRACT[name]["schema"]) + for param in _LOCAL_HIDDEN_PARAMS.get(name, ()): + schema["properties"].pop(param, None) for (tool_name, param), text in _LOCAL_PARAM_DESCRIPTIONS.items(): if tool_name == name and param in schema["properties"]: schema["properties"][param]["description"] = text @@ -1311,41 +1306,34 @@ def build_agent_tools(client, include_management: bool = False) -> list[Callable if getattr(client, "api_key", None): return _build_cloud_agent_tools(client, include_management) - def browse_documents(folder_id: str = "root", recursive: bool = False, - sort: str = "time", query: Optional[str] = None, - offset: int = 0, limit: int = 10) -> str: + def browse_documents(offset: int = 0, limit: int = 10) -> str: return call_tool(client, "browse_documents", { - "folder_id": folder_id, "recursive": recursive, "sort": sort, - "query": query, "offset": offset, "limit": limit, + "offset": offset, "limit": limit, })[0] - def get_document(doc_name: str, folder_id: Optional[str] = None, - wait_for_completion: bool = False) -> str: + def get_document(doc_name: str, wait_for_completion: bool = False) -> str: return call_tool(client, "get_document", { - "doc_name": doc_name, "folder_id": folder_id, + "doc_name": doc_name, "wait_for_completion": wait_for_completion, })[0] - def get_document_structure(doc_name: str, folder_id: Optional[str] = None, - part: int = 1, + def get_document_structure(doc_name: str, part: int = 1, wait_for_completion: bool = False) -> str: return call_tool(client, "get_document_structure", { - "doc_name": doc_name, "folder_id": folder_id, "part": part, + "doc_name": doc_name, "part": part, "wait_for_completion": wait_for_completion, })[0] def get_page_content(doc_name: str, pages: str, - folder_id: Optional[str] = None, wait_for_completion: bool = False) -> str: return call_tool(client, "get_page_content", { - "doc_name": doc_name, "pages": pages, "folder_id": folder_id, + "doc_name": doc_name, "pages": pages, "wait_for_completion": wait_for_completion, })[0] - def remove_document(doc_names: list[str], - folder_id: Optional[str] = None) -> str: + def remove_document(doc_names: list[str]) -> str: return call_tool(client, "remove_document", { - "doc_names": doc_names, "folder_id": folder_id, + "doc_names": doc_names, })[0] functions = { diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 327cfa574..1d8b7e6d5 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -71,13 +71,23 @@ def test_contract_matches_snapshot(): def test_tool_surface_and_docstrings(client): + import inspect + from pageindex.agent_tools import _LOCAL_HIDDEN_PARAMS, _local_schema tools = client.agent_tools() assert [tool.__name__ for tool in tools] == list(tool_names()) with_management = client.agent_tools(include_management=True) assert [tool.__name__ for tool in with_management][-1] == "remove_document" - for tool in tools: - for param in TOOL_CONTRACT[tool.__name__]["schema"]["properties"]: + for tool in with_management: + exposed = list(_local_schema(tool.__name__)["properties"]) + assert list(inspect.signature(tool).parameters) == exposed + for param in exposed: assert param in tool.__doc__ + # Cloud-only params are hidden, not documented-then-retracted: + # strict-schema frameworks cannot express the dead-end calls at all. + # (The description may still mention them as cloud capabilities.) + args_section = tool.__doc__.split("Args:", 1)[1] + for hidden in _LOCAL_HIDDEN_PARAMS.get(tool.__name__, ()): + assert f"{hidden}:" not in args_section docs = {tool.__name__: tool.__doc__ for tool in tools} # Tools whose cloud description has no cloud-only content keep it # verbatim; browse_documents serves the localized guidance. @@ -88,19 +98,26 @@ def test_tool_surface_and_docstrings(client): def test_local_schema_structure_matches_contract(): - """The local guidance layer may localize description strings only — - names, types, defaults, bounds, and required stay byte-identical.""" + """The local surface is the contract minus the documented cloud-only + params; the surviving params' names, types, defaults, bounds, and + required stay byte-identical — localization may only touch description + strings.""" import copy - from pageindex.agent_tools import _local_schema + from pageindex.agent_tools import _LOCAL_HIDDEN_PARAMS, _local_schema - def stripped(schema): + def stripped(schema, drop=()): schema = copy.deepcopy(schema) + for param in drop: + schema["properties"].pop(param, None) for spec in schema["properties"].values(): spec.pop("description", None) return schema for name, contract in TOOL_CONTRACT.items(): - assert stripped(_local_schema(name)) == stripped(contract["schema"]), name + hidden = _LOCAL_HIDDEN_PARAMS.get(name, ()) + assert not (set(hidden) & set(contract["schema"].get("required", []))), name + assert stripped(_local_schema(name)) == stripped(contract["schema"], + drop=hidden), name def test_local_guidance_references_only_local_tools(client): From 1fa3eb7fb3d23763598abc238c5ec0e6583f2778 Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 12 Aug 2026 00:18:08 +0800 Subject: [PATCH 16/18] =?UTF-8?q?fix:=20incremental-review=20findings=20?= =?UTF-8?q?=E2=80=94=20bridge=20cache,=20guards,=20envelope=20drift?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three independent review passes over the agent-instructions increment surfaced six fixes: - The per-client bridge moved off the instance into a weak-keyed, lock-guarded module cache: cloud clients stay picklable (threading.RLock no longer rides on the client) and concurrent first calls can no longer construct duplicate bridges/sessions. - Blank or non-string initialize.instructions now hit the same honest error as a missing one — a whitespace-only or structured value could previously become the system prompt (or crash the doc_id append with a raw TypeError). - The invalid-sort envelope no longer prescribes sort="relevance" — the one error text that still taught the cloud-only value it would then reject. - "Page through the rest of the library" is emitted only when has_more is true; a fully-listed library no longer instructs a pointless call. - The mandatory full-library paging step now says limit: 50 — 6 calls instead of 30 on a 300-document library. - Docstrings and comments rescoped to what is actually true: the never-raise contract covers invocations the signatures accept (unknown params fail at the Python boundary; call_tool answers them with the guided envelope), recursive is accepted as the identity rather than errored, lenient framework arg models drop hidden params pre-call, and the module header no longer claims full schema parity. The capability-phrase guard now covers every local docstring, not just browse_documents. --- examples/documents/attention-residuals.doc_id | 1 + pageindex/agent_tools.py | 102 +++++++++++------- tests/test_agent_tools.py | 71 +++++++++++- 3 files changed, 133 insertions(+), 41 deletions(-) create mode 100644 examples/documents/attention-residuals.doc_id diff --git a/examples/documents/attention-residuals.doc_id b/examples/documents/attention-residuals.doc_id new file mode 100644 index 000000000..19003f3ce --- /dev/null +++ b/examples/documents/attention-residuals.doc_id @@ -0,0 +1 @@ +pi-c4a794161a904216bde92b3d2c269a19 \ No newline at end of file diff --git a/pageindex/agent_tools.py b/pageindex/agent_tools.py index b6ed15c6f..2399ca0e7 100644 --- a/pageindex/agent_tools.py +++ b/pageindex/agent_tools.py @@ -1,15 +1,19 @@ """Agent tools: the cloud MCP tool contract, executed against a PageIndexClient. -Tool names and input-schema structure match the PageIndex cloud MCP server, -so agent prompts work unchanged across the cloud MCP connection and this -in-process layer. Only the tools that exist in every mode are registered -(no folders, search_documents, or get_document_image), and the guidance -strings (tool descriptions) adapt to the local surface the same way the -agent instructions do — they never teach capabilities that only exist on -the cloud. - -Tools never raise: every outcome, including errors, is returned as the same -JSON envelope the cloud emits ({"success": true, ...} / {"error": ...}). +Tool names and the surviving input-schema structure match the PageIndex +cloud MCP server — the local surface hides the documented cloud-only +parameters — so agent prompts port across the cloud MCP connection and +this in-process layer. Only the tools that exist in every mode are +registered (no folders, search_documents, or get_document_image), and the +guidance strings (tool descriptions) adapt to the local surface the same +way the agent instructions do — they never teach capabilities that only +exist on the cloud. + +Tools never raise for any invocation their signatures accept: every +outcome, including errors, is returned as the same JSON envelope the cloud +emits ({"success": true, ...} / {"error": ...}). Arguments outside a pruned +local signature fail at the Python call boundary; the call_tool path +answers them with the guided error envelope instead. """ from __future__ import annotations @@ -17,7 +21,9 @@ import difflib import json import re +import threading import time +import weakref from typing import Any, Callable, Optional from .errors import PageIndexAPIError @@ -638,10 +644,15 @@ def _browse_documents(client, folder_id: str = "root", recursive: bool = False, if folder_id != "root": return _folder_unsupported("folder_id") if sort not in ("time", "relevance"): - return _failure('sort must be "time" or "relevance"', None, - {"summary": "Invalid sort mode", - "options": ['Use sort="time" or sort="relevance"']}, - "INVALID_INPUT") + return _failure( + 'Invalid sort mode — only the default "time" sort is available ' + "in local mode.", None, + {"summary": "Invalid sort mode", + "options": ['Use sort="time" (newest first) or omit sort', + "Semantic ranking is available on PageIndex cloud " + "(PageIndexCloudClient with an API key)"]}, + "INVALID_INPUT", + ) if sort == "relevance" or query: # Semantic ranking is a cloud capability; like folders, it is not # imitated here. @@ -715,8 +726,10 @@ def _browse_documents(client, folder_id: str = "root", recursive: bool = False, options.append( "Results returned ≠ correct results. Verify these documents match " "the user's actual intent (topic, time period, document type) " - "before proceeding. If they do not match, page through the rest " - "of the library. Do NOT use general knowledge as a substitute." + "before proceeding." + + (" If they do not match, page through the rest of the library." + if has_more else "") + + " Do NOT use general knowledge as a substitute." ) if page_has_processing: options.append("Some documents on this page are still processing. " @@ -1115,15 +1128,19 @@ def _tool_docstring(description: str, properties: dict[str, Any]) -> str: # contract minus the hidden cloud-only parameters, and description strings # adapt to the local surface the same way AGENT_INSTRUCTIONS does — guidance # must not teach capabilities (folders, semantic ranking) or tools -# (search_documents, get_document_image) that do not exist here. Guard tests -# assert both properties; a contract refresh that reintroduces a cloud-only -# reference fails the dead-reference test. +# (search_documents, get_document_image) that do not exist here. Guard +# tests pin structure (contract-minus-hidden equality), tool references +# (the dead-reference test), and capability phrases (the per-docstring +# phrase test) — a contract refresh that reintroduces a cloud-only +# reference fails loudly. #: Cloud-only parameters hidden from the local surface — strict-schema -#: frameworks then make the dead-end calls inexpressible. The -#: implementations still accept them and answer with the guided error -#: envelope, for direct call_tool callers and hosts without schema -#: enforcement. +#: frameworks make the dead-end calls inexpressible, and lenient framework +#: argument models drop them before the call (degrading to the bare call). +#: The call_tool path still answers folder_id/sort/query with the guided +#: error envelope; recursive is simply accepted (flattening a folderless +#: library is the identity). Plain functions reject unknown parameters at +#: the Python call boundary. _LOCAL_HIDDEN_PARAMS: dict[str, tuple[str, ...]] = { "browse_documents": ("folder_id", "recursive", "sort", "query"), "get_document": ("folder_id",), @@ -1143,7 +1160,8 @@ def _tool_docstring(description: str, properties: dict[str, Any]) -> str: "Primary document retrieval tool — first choice for any " "document-related question. Lists your documents newest first with " "names and descriptions; match them against the user's intent and " - "page through with `offset: next_offset` while `has_more` is true. " + "page through with `offset: next_offset` (limit up to 50) while " + "`has_more` is true. " 'Folder browsing and semantic ranking (sort="relevance") are not ' "supported in local mode yet — they work on PageIndex cloud." ), @@ -1262,18 +1280,24 @@ def proxy(**kwargs: Any) -> str: return proxy +_BRIDGES: "weakref.WeakKeyDictionary" = weakref.WeakKeyDictionary() +_BRIDGES_LOCK = threading.Lock() + + def _cloud_bridge(client): - """One bridge per client instance: tool discovery and instructions share - a single MCP session.""" - bridge = getattr(client, "_mcp_bridge", None) - if bridge is None: - from .mcp_bridge import McpBridge - bridge = McpBridge( - f"{client.BASE_URL}/mcp", - {"Authorization": f"Bearer {client.api_key}"}, - ) - client._mcp_bridge = bridge - return bridge + """One bridge per client: tool discovery and instructions share a single + MCP session. Weak-keyed off the instance so clients stay picklable; the + lock closes the check-then-set race under concurrent first calls.""" + with _BRIDGES_LOCK: + bridge = _BRIDGES.get(client) + if bridge is None: + from .mcp_bridge import McpBridge + bridge = McpBridge( + f"{client.BASE_URL}/mcp", + {"Authorization": f"Bearer {client.api_key}"}, + ) + _BRIDGES[client] = bridge + return bridge def _build_cloud_agent_tools(client, include_management: bool) -> list[Callable[..., str]]: @@ -1301,7 +1325,9 @@ def build_agent_tools(client, include_management: bool = False) -> list[Callable Cloud: one function per tool of the live cloud MCP tool set, signatures synthesized from the server's schemas, calls proxied over MCP. Local: the built-in contract tools over the local store. Every function returns - the JSON envelope as a string and never raises. + the JSON envelope as a string and never raises for arguments its + signature accepts (cloud-only parameters are absent from the local + signatures; the call_tool path answers them with the guided envelope). """ if getattr(client, "api_key", None): return _build_cloud_agent_tools(client, include_management) @@ -1394,7 +1420,7 @@ def remove_document(doc_names: list[str]) -> str: PERSISTENCE (before concluding the target document is not in the library): This protocol applies both when results are empty AND when results are returned but none match the user's intent. Do NOT give up after a single discovery attempt. Follow these steps in order: 1. browse_documents() and compare every returned name/description against the user's intent -2. Page through the ENTIRE library with `offset: next_offset` until has_more is false — MANDATORY, must be completed before concluding "not found" +2. Page through the ENTIRE library with `limit: 50` and `offset: next_offset` until has_more is false — MANDATORY, must be completed before concluding "not found" 3. Re-scan for loose matches: synonyms, abbreviations, and partial titles in names/descriptions can identify the target Only after ALL three steps have been tried may you conclude the document is not in the library. Do NOT fall back to general knowledge — if the user's question references their own documents, exhaust every discovery path first.""" @@ -1415,7 +1441,7 @@ def _base_instructions(client) -> str: if not getattr(client, "api_key", None): return AGENT_INSTRUCTIONS instructions = _cloud_bridge(client).instructions() - if not instructions: + if not isinstance(instructions, str) or not instructions.strip(): raise PageIndexAPIError( "The MCP server returned no agent instructions — refusing to " "substitute the SDK's local-subset guidance, which does not " diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 1d8b7e6d5..63189b3ba 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -131,11 +131,20 @@ def test_local_guidance_references_only_local_tools(client): def test_local_guidance_points_cloud_only_capabilities_at_cloud(client): - browse = client.agent_tools()[0].__doc__ + tools = client.agent_tools(include_management=True) + browse = tools[0].__doc__ assert "not supported in local mode yet" in browse assert "PageIndex cloud" in browse - assert "search_documents" not in browse - assert "get_folder_structure" not in browse + # Capability-phrase guard, all docstrings: cloud-only language must not + # drift back in via a contract refresh. browse alone keeps exactly one + # sort="relevance" mention — the sanctioned pointer to the cloud. + for tool in tools: + doc = tool.__doc__ + for phrase in ("shared-with-me", "sub-folder", "get_folder_structure", + "search_documents", "get_document_image"): + assert phrase not in doc, (tool.__name__, phrase) + expected = 1 if tool.__name__ == "browse_documents" else 0 + assert doc.count('sort="relevance"') == expected, tool.__name__ # ── browse_documents ── @@ -170,9 +179,12 @@ def test_browse_documents_pagination(client, store_path): first, _ = run(client, "browse_documents", limit=2) assert [d["name"] for d in first["documents"]] == ["doc2.pdf", "doc1.pdf"] assert first["has_more"] is True and first["next_offset"] == 2 + assert "page through the rest" in json.dumps(first["next_steps"]) second, _ = run(client, "browse_documents", limit=2, offset=2) assert [d["name"] for d in second["documents"]] == ["doc0.pdf"] assert second["has_more"] is False + # No paging advice when there is nothing left to page through. + assert "page through the rest" not in json.dumps(second["next_steps"]) def test_browse_documents_relevance_unsupported(client, store_path): @@ -189,6 +201,9 @@ def test_browse_documents_relevance_unsupported(client, store_path): assert is_error and "not supported in local mode" in stray_query["error"] bad_sort, is_error = run(client, "browse_documents", sort="banana") assert is_error and bad_sort["errorCode"] == "INVALID_INPUT" + # The invalid-sort guidance must not prescribe the cloud-only value. + assert 'Use sort="relevance"' not in json.dumps(bad_sort) + assert "local mode" in bad_sort["error"] def test_browse_documents_empty_and_folder_error(client): @@ -1017,6 +1032,56 @@ def instructions(self): assert len(created) == 1 +def test_cloud_bridge_cache_threadsafe_and_pickle_clean(monkeypatch): + """One bridge per client even under concurrent first calls, and the + bridge lives off the instance so cloud clients stay picklable.""" + import pickle + import threading + import time as time_mod + import pageindex.mcp_bridge as mcp_bridge + created = [] + + class _Bridge(_FakeBridge): + def __init__(self, url, headers): + time_mod.sleep(0.01) # widen the construction window + super().__init__(url, headers) + created.append(self) + + def instructions(self): + return "LIVE" + + monkeypatch.setattr(mcp_bridge, "McpBridge", _Bridge) + cloud = PageIndexCloudClient(api_key="pi-test-key") + workers = ([threading.Thread(target=cloud.agent_tools) for _ in range(4)] + + [threading.Thread(target=cloud.agent_instructions) + for _ in range(4)]) + for worker in workers: + worker.start() + for worker in workers: + worker.join() + assert len(created) == 1 + pickle.dumps(cloud) + + +def test_cloud_agent_instructions_blank_or_nonstring_raises(monkeypatch): + """Whitespace-only or non-string initialize.instructions must hit the + same honest error as a missing one — never a blank system prompt.""" + import pageindex.mcp_bridge as mcp_bridge + + for bad in (" \n\t ", {"not": "a string"}): + class _SilentBridge: + def __init__(self, url, headers): + pass + + def instructions(self, _value=bad): + return _value + + monkeypatch.setattr(mcp_bridge, "McpBridge", _SilentBridge) + cloud = PageIndexCloudClient(api_key="pi-test-key") + with pytest.raises(PageIndexAPIError, match="no agent instructions"): + cloud.agent_instructions() + + def test_cloud_agent_instructions_empty_raises(monkeypatch): """An empty server response must raise, not silently substitute the subset guidance — same posture as the annotation-regression guard.""" From 63b767f70889999e9be7b0a1d61272fe54ac5ab3 Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 12 Aug 2026 00:18:24 +0800 Subject: [PATCH 17/18] chore: keep the demo's doc_id cache file out of the repo --- .gitignore | 1 + examples/documents/attention-residuals.doc_id | 1 - 2 files changed, 1 insertion(+), 1 deletion(-) delete mode 100644 examples/documents/attention-residuals.doc_id diff --git a/.gitignore b/.gitignore index 5193735ca..b5c223b31 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,4 @@ __pycache__ logs/ .pageindex/ dist/ +*.doc_id diff --git a/examples/documents/attention-residuals.doc_id b/examples/documents/attention-residuals.doc_id deleted file mode 100644 index 19003f3ce..000000000 --- a/examples/documents/attention-residuals.doc_id +++ /dev/null @@ -1 +0,0 @@ -pi-c4a794161a904216bde92b3d2c269a19 \ No newline at end of file From 6c9fe2e544c400674daded224b2f399bcf0152c3 Mon Sep 17 00:00:00 2001 From: Ray Date: Wed, 12 Aug 2026 00:29:48 +0800 Subject: [PATCH 18/18] test: live envelope field-parity guard against cloud response drift The frozen contract guards tools/list, but the response envelopes the local tools emit were hand-built to mirror the cloud's and had no drift detector. A key-gated live test now asserts every field local emits exists in the live cloud response for the analogous call (top-level keys, next_steps, document entries, structure nodes, content entries). Guidance wording is deliberately localized and not compared. Verified green against the live server: local and cloud field structures currently match exactly. --- tests/test_agent_tools.py | 57 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 57 insertions(+) diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 63189b3ba..2a3cbc70b 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -967,6 +967,63 @@ def test_live_cloud_contract_parity(): assert (real.get("annotations") or {}).get(key) == value, (name, key) +@pytest.mark.skipif(not LIVE_KEY, reason="PAGEINDEX_API_KEY not set") +def test_live_cloud_envelope_field_parity(tmp_path): + """Response-envelope drift alarm: every field the local tools emit must + exist in the live cloud tool's response for the analogous call — a cloud + rename of a shared field (has_more, next_offset, content, ...) fails + here. Guidance wording is deliberately localized and not compared.""" + from pageindex.mcp_bridge import McpBridge + bridge = McpBridge("https://api.pageindex.ai/mcp", + {"Authorization": f"Bearer {LIVE_KEY}"}) + cloud_browse = json.loads(bridge.call_tool("browse_documents", {"limit": 2})) + assert cloud_browse.get("success") is True and cloud_browse["documents"] + doc_name = cloud_browse["documents"][0]["name"] + cloud = { + "browse_documents": cloud_browse, + "get_document": json.loads(bridge.call_tool( + "get_document", {"doc_name": doc_name})), + "get_document_structure": json.loads(bridge.call_tool( + "get_document_structure", {"doc_name": doc_name})), + "get_page_content": json.loads(bridge.call_tool( + "get_page_content", {"doc_name": doc_name, "pages": "1"})), + } + + store = str(tmp_path / "store") + local_client = PageIndexLocalClient(storage_path=store) + seed_doc(store, "pi-parity", "parity.pdf") + local = { + "browse_documents": run(local_client, "browse_documents")[0], + "get_document": run(local_client, "get_document", + doc_name="parity.pdf")[0], + "get_document_structure": run(local_client, "get_document_structure", + doc_name="parity.pdf")[0], + "get_page_content": run(local_client, "get_page_content", + doc_name="parity.pdf", pages="1")[0], + } + + for name in cloud: + assert cloud[name].get("success") is True, name + missing = set(local[name]) - set(cloud[name]) + assert not missing, (name, missing) + assert (set(local[name]["next_steps"]) + <= set(cloud[name]["next_steps"]) | {"auto_retry"}), name + + local_doc = local["browse_documents"]["documents"][0] + cloud_doc = cloud_browse["documents"][0] + assert set(local_doc) - set(cloud_doc) <= {"metadata"} + + local_nodes = local["get_document_structure"]["structure"] + cloud_nodes = cloud["get_document_structure"]["structure"] + local_node = local_nodes[0] if isinstance(local_nodes, list) else local_nodes + cloud_node = cloud_nodes[0] if isinstance(cloud_nodes, list) else cloud_nodes + assert (set(local_node) + <= set(cloud_node) | {"page_index", "prefix_summary"}) + + assert (set(local["get_page_content"]["content"][0]) + <= set(cloud["get_page_content"]["content"][0])) + + @pytest.mark.skipif(not LIVE_KEY, reason="PAGEINDEX_API_KEY not set") def test_live_cloud_instructions_nonempty(): """The empty-instructions guard raises for cloud clients; the real