From cf7f5ce9bf64f2b49cb5ee223e3ce1cdd6a987ae Mon Sep 17 00:00:00 2001 From: mountain Date: Wed, 8 Jul 2026 18:56:51 +0800 Subject: [PATCH] fix: address PR #272 review findings (directly-fixable items) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Verified against current dev; the compat/behavior decisions (#7 api_key semantics, #10 CLI flags, #11 doc-description default) are deferred. Crashes: - page_index(): snapshot args before importing IndexConfig — locals() was capturing the imported class and IndexConfig(extra='forbid') made every call raise ValidationError. - process_none_page_numbers: pop('page', None) instead of del (a TOC item without 'page' raised KeyError mid-pipeline). - pipeline._run_async: guard only the loop detection, not the run, so a real RuntimeError from the coroutine isn't masked as "asyncio.run() cannot be called from a running event loop". Silent-wrong / robustness: - LocalBackend.get_document_structure and the agent get_document / get_document_structure tools now surface a missing doc (raise / error-JSON) instead of returning empty, matching get_page_content and the cloud backend. - cloud delete_collection drops the cached folder_id. - cloud query raises on an empty collection instead of POSTing doc_id:[]. - LocalClient skips the API-key check for keyless providers (ollama, lm_studio, …) so keyless LiteLLM models aren't rejected at construction. Compat / cleanup: - md_to_tree coerces legacy 'yes'/'no' string flags (a bare 'no' was truthy). - FileTypeError also subclasses ValueError (0.2.x raised ValueError). - _validate_llm_provider no longer mutates global litellm.model_cost_map_url. - __all__ re-includes legacy exports (page_index, md_to_tree, get_*). - Rewrite examples/agentic_vectorless_rag_demo.py to the Collection API and use the in-repo attention.pdf (the old workspace=/client.index/client.documents API no longer exists). Adds tests/test_review_fixes.py (10 regressions). Full suite: 189 passed. Claude-Session: https://claude.ai/code/session_01Kx5DgKbhK1N8autqXH8SmS --- examples/agentic_vectorless_rag_demo.py | 86 ++++++++-------- pageindex/__init__.py | 7 ++ pageindex/backend/cloud.py | 7 ++ pageindex/backend/local.py | 16 ++- pageindex/client.py | 10 +- pageindex/errors.py | 8 +- pageindex/index/page_index.py | 9 +- pageindex/index/page_index_md.py | 14 +++ pageindex/index/pipeline.py | 19 ++-- tests/test_review_fixes.py | 130 ++++++++++++++++++++++++ 10 files changed, 250 insertions(+), 56 deletions(-) create mode 100644 tests/test_review_fixes.py diff --git a/examples/agentic_vectorless_rag_demo.py b/examples/agentic_vectorless_rag_demo.py index b4ed9c2..103a1a0 100644 --- a/examples/agentic_vectorless_rag_demo.py +++ b/examples/agentic_vectorless_rag_demo.py @@ -1,22 +1,28 @@ """ -Agentic Vectorless RAG with PageIndex - Demo +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. +Build 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 lets an agent reason over it for human-like, +context-aware retrieval. + +This demo wires up your OWN agent + tools against the PageIndex Collection API. +For the batteries-included version, just use ``col.query(..., stream=True)`` — +see local_demo.py. 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 + - get_document() — document metadata (name, type, description) + - get_document_structure() — the document's tree-structure index + - get_page_content() — text of specific pages / line ranges Steps: - 1 — Index a PDF and view its tree structure index + 1 — Index a PDF and view its tree structure 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 pageindex openai-agents + export OPENAI_API_KEY=your-api-key # or any LiteLLM-supported provider """ import sys import json @@ -32,19 +38,19 @@ from agents.model_settings import ModelSettings from agents.stream_events import RawResponsesStreamEvent, RunItemStreamEvent from openai.types.responses import ResponseTextDeltaEvent, ResponseReasoningSummaryTextDeltaEvent -from pageindex import PageIndexClient -import pageindex.utils as utils +from pageindex import LocalClient -PDF_URL = "https://arxiv.org/pdf/2603.15031" +PDF_URL = "https://arxiv.org/pdf/1706.03762.pdf" _EXAMPLES_DIR = Path(__file__).parent -PDF_PATH = _EXAMPLES_DIR / "documents" / "attention-residuals.pdf" +PDF_PATH = _EXAMPLES_DIR / "documents" / "attention.pdf" WORKSPACE = _EXAMPLES_DIR / "workspace" +MODEL = "gpt-4o-2024-11-20" # any LiteLLM-supported model 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 the document's name and type. - 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. @@ -52,7 +58,7 @@ 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(col, doc_id: str, prompt: str, model: 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. @@ -61,13 +67,15 @@ 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) + """Get document metadata: name, type, and description.""" + doc = col.get_document(doc_id) + doc.pop("structure", None) # keep tool output small for the LLM context + return json.dumps(doc, ensure_ascii=False) @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(col.get_document_structure(doc_id), ensure_ascii=False) @function_tool def get_page_content(pages: str) -> str: @@ -76,13 +84,13 @@ def query_agent(client: PageIndexClient, doc_id: str, prompt: str, verbose: bool 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(col.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=model, # model_settings=ModelSettings(reasoning={"effort": "low", "summary": "auto"}), # Uncomment to enable reasoning ) @@ -128,12 +136,14 @@ def query_agent(client: PageIndexClient, doc_id: str, prompt: str, verbose: bool print() return "" if not streamed_run.final_output else str(streamed_run.final_output) + # Only the detection is guarded, not the run, so a real error inside _run + # isn't misread as "no running loop". try: asyncio.get_running_loop() - with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: - return pool.submit(asyncio.run, _run()).result() except RuntimeError: return asyncio.run(_run()) + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: + return pool.submit(asyncio.run, _run()).result() if __name__ == "__main__": @@ -152,37 +162,33 @@ if __name__ == "__main__": f.write(chunk) print("Download complete.\n") - # Setup - client = PageIndexClient(workspace=WORKSPACE) + # Setup: self-hosted local client + a collection + client = LocalClient(model=MODEL, storage_path=str(WORKSPACE)) + col = client.collection("agentic-demo") # 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), - None, - ) - if doc_id: - print(f"\nLoaded cached doc_id: {doc_id}") - else: - doc_id = client.index(PDF_PATH) - print(f"\nIndexed. doc_id: {doc_id}") + # Content-hash dedup: re-running reuses the existing doc_id, no re-index. + doc_id = col.add(str(PDF_PATH)) + print(f"\ndoc_id: {doc_id}") print("\nTree Structure (top-level sections):") - structure = json.loads(client.get_document_structure(doc_id)) - utils.print_tree(structure) + for node in col.get_document_structure(doc_id): + print(f" - {node.get('title', '(untitled)')}") # Step 2: View document metadata print("\n" + "=" * 60) print("Step 2: View document metadata") print("=" * 60) - doc_metadata = client.get_document(doc_id) - print(f"\n{doc_metadata}") + meta = col.get_document(doc_id) + meta.pop("structure", None) + print("\n" + json.dumps(meta, ensure_ascii=False, indent=2)) # Step 3: Agent Query print("\n" + "=" * 60) print("Step 3: Agent Query (auto tool-use)") print("=" * 60) - question = "Explain Attention Residuals in simple language." + question = "Explain the Transformer's self-attention in simple language." print(f"\nQuestion: '{question}'") - query_agent(client, doc_id, question, verbose=True) + query_agent(col, doc_id, question, MODEL, verbose=True) diff --git a/pageindex/__init__.py b/pageindex/__init__.py index e05c5b4..e95bce1 100644 --- a/pageindex/__init__.py +++ b/pageindex/__init__.py @@ -61,4 +61,11 @@ __all__ = [ "IndexingError", "CloudAPIError", "FileTypeError", + # Legacy top-level exports (pre-SDK API), kept so `from pageindex import *` + # still binds them. + "page_index", + "md_to_tree", + "get_document", + "get_document_structure", + "get_page_content", ] diff --git a/pageindex/backend/cloud.py b/pageindex/backend/cloud.py index 4f74fce..bd1dba6 100644 --- a/pageindex/backend/cloud.py +++ b/pageindex/backend/cloud.py @@ -162,6 +162,9 @@ class CloudBackend: folder_id = self._get_folder_id(name) if folder_id: self._request("DELETE", f"/folder/{self._enc(folder_id)}/") + # Drop the cached id so a later same-name op re-resolves instead of + # reusing the now-deleted folder_id. + self._folder_id_cache.pop(name, None) # ── Document management ─────────────────────────────────────────────── @@ -307,6 +310,8 @@ class CloudBackend: "doc_ids cannot be empty; pass None to query the whole collection" ) doc_id = doc_ids if doc_ids else self._get_all_doc_ids(collection) + if not doc_id: + raise ValueError("collection has no documents to query") # A non-streaming completion returns nothing until generation # finishes, so it needs far more than the default 30s. retries=1: # retrying this non-idempotent call would redo the full server-side @@ -340,6 +345,8 @@ class CloudBackend: "doc_ids cannot be empty; pass None to query the whole collection" ) doc_id = doc_ids if doc_ids else self._get_all_doc_ids(collection) + if not doc_id: + raise ValueError("collection has no documents to query") headers = self._headers # Queue carries QueryEvent, an Exception to re-raise, or None (end). queue: asyncio.Queue[QueryEvent | Exception | None] = asyncio.Queue() diff --git a/pageindex/backend/local.py b/pageindex/backend/local.py index 8102deb..d6d9676 100644 --- a/pageindex/backend/local.py +++ b/pageindex/backend/local.py @@ -190,6 +190,10 @@ class LocalBackend: LocalBackend._fill_node_text(node["nodes"], page_map) def get_document_structure(self, collection: str, doc_id: str) -> list: + # Parity with get_document / the cloud backend: a missing doc must raise, + # not masquerade as an empty structure. + if not self._storage.get_document(collection, doc_id): + raise DocumentNotFoundError(f"Document {doc_id} not found") return self._storage.get_document_structure(collection, doc_id) def get_page_content(self, collection: str, doc_id: str, pages: str) -> list: @@ -255,7 +259,10 @@ class LocalBackend: rejection = _reject(doc_id) if rejection: return rejection - return json.dumps(storage.get_document(col_name, doc_id)) + doc = storage.get_document(col_name, doc_id) + if not doc: + return json.dumps({"error": f"doc_id '{doc_id}' not found."}) + return json.dumps(doc) @function_tool def get_document_structure(doc_id: str) -> str: @@ -263,6 +270,8 @@ class LocalBackend: rejection = _reject(doc_id) if rejection: return rejection + if not storage.get_document(col_name, doc_id): + return json.dumps({"error": f"doc_id '{doc_id}' not found."}) structure = storage.get_document_structure(col_name, doc_id) return json.dumps(remove_fields(structure, fields=["text"]), ensure_ascii=False) @@ -272,7 +281,10 @@ class LocalBackend: rejection = _reject(doc_id) if rejection: return rejection - result = backend.get_page_content(col_name, doc_id, pages) + try: + result = backend.get_page_content(col_name, doc_id, pages) + except DocumentNotFoundError: + return json.dumps({"error": f"doc_id '{doc_id}' not found."}) return json.dumps(result, ensure_ascii=False) tools = [get_document, get_document_structure, get_page_content] diff --git a/pageindex/client.py b/pageindex/client.py index a143ca8..34883ea 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -113,14 +113,20 @@ class PageIndexClient: @staticmethod def _validate_llm_provider(model: str) -> None: - """Validate model and check API key via litellm. Warns if key seems missing.""" + """Validate the model string and require an API key for providers that + need one. Local / keyless providers (ollama, lm_studio, …) are skipped so + a keyless LiteLLM model isn't rejected at construction time.""" try: import litellm - litellm.model_cost_map_url = "" _, provider, _, _ = litellm.get_llm_provider(model=model) except Exception: return + # LiteLLM providers that run locally and need no API key. + keyless = {"ollama", "ollama_chat", "lm_studio", "hosted_vllm", "vllm"} + if provider in keyless: + return + key = litellm.get_api_key(llm_provider=provider, dynamic_api_key=None) if not key: import os diff --git a/pageindex/errors.py b/pageindex/errors.py index dc4b1c7..b2a7065 100644 --- a/pageindex/errors.py +++ b/pageindex/errors.py @@ -39,6 +39,10 @@ class CloudAPIError(PageIndexAPIError): self.status_code = status_code -class FileTypeError(PageIndexError): - """Unsupported file type.""" +class FileTypeError(PageIndexError, ValueError): + """Unsupported file type. + + Also subclasses ValueError so pre-SDK ``except ValueError`` around indexing + (0.2.x raised ValueError for an unsupported file format) still catches it. + """ pass diff --git a/pageindex/index/page_index.py b/pageindex/index/page_index.py index 9687a28..1982214 100644 --- a/pageindex/index/page_index.py +++ b/pageindex/index/page_index.py @@ -679,11 +679,11 @@ def process_none_page_numbers(toc_items, page_list, start_index=1, model=None): continue item_copy = copy.deepcopy(item) - del item_copy['page'] + item_copy.pop('page', None) result = add_page_number_to_toc(page_contents, item_copy, model) if isinstance(result[0]['physical_index'], str) and result[0]['physical_index'].startswith('').strip()) - del item['page'] + item.pop('page', None) return toc_items @@ -1113,11 +1113,14 @@ def page_index_main(doc, opt=None): def page_index(doc, model=None, toc_check_page_num=None, max_page_num_each_node=None, max_token_num_each_node=None, if_add_node_id=None, if_add_node_summary=None, if_add_doc_description=None, if_add_node_text=None): - from ..config import IndexConfig + # Snapshot the call args BEFORE importing IndexConfig — otherwise the + # imported class would be captured by locals() and rejected by + # IndexConfig(extra="forbid"). user_opt = { arg: value for arg, value in locals().items() if arg != "doc" and value is not None } + from ..config import IndexConfig opt = IndexConfig(**user_opt) return page_index_main(doc, opt) diff --git a/pageindex/index/page_index_md.py b/pageindex/index/page_index_md.py index f9e300a..d187fb4 100644 --- a/pageindex/index/page_index_md.py +++ b/pageindex/index/page_index_md.py @@ -237,7 +237,21 @@ def clean_tree_for_output(tree_nodes): return cleaned_nodes +def _coerce_bool(value): + """Coerce a legacy 'yes'/'no' string flag to bool (a bare 'no' is truthy).""" + if isinstance(value, str): + return value.strip().lower() in ("yes", "true", "1", "y", "on") + return bool(value) + + async def md_to_tree(md_path, if_thinning=False, min_token_threshold=None, if_add_node_summary=False, summary_token_threshold=None, model=None, if_add_doc_description=False, if_add_node_text=False, if_add_node_id=True): + # Accept legacy 'yes'/'no' string flags — a bare 'no' would otherwise be + # truthy and wrongly enable the option. + if_thinning = _coerce_bool(if_thinning) + if_add_node_summary = _coerce_bool(if_add_node_summary) + if_add_doc_description = _coerce_bool(if_add_doc_description) + if_add_node_text = _coerce_bool(if_add_node_text) + if_add_node_id = _coerce_bool(if_add_node_id) with open(md_path, 'r', encoding='utf-8') as f: markdown_content = f.read() line_count = markdown_content.count('\n') + 1 diff --git a/pageindex/index/pipeline.py b/pageindex/index/pipeline.py index 6217a46..355b12f 100644 --- a/pageindex/index/pipeline.py +++ b/pageindex/index/pipeline.py @@ -44,17 +44,22 @@ def _run_async(coro): import asyncio import concurrent.futures import contextvars + # Only the detection is guarded — NOT the run. If the coroutine's own work + # raises RuntimeError, letting it fall into `except RuntimeError` here would + # misfire the "no running loop" branch and mask the real error behind a + # bogus "asyncio.run() cannot be called from a running event loop". try: asyncio.get_running_loop() - # Already inside an event loop -- run in a separate thread. Copy the - # current context so ContextVar-based settings (e.g. the - # max_concurrency_scope override set by build_index) propagate into the - # worker thread instead of silently falling back to the process default. - ctx = contextvars.copy_context() - with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: - return pool.submit(ctx.run, asyncio.run, coro).result() except RuntimeError: + # No running loop -- drive the coroutine directly. return asyncio.run(coro) + # Already inside an event loop -- run in a separate thread so we don't nest + # asyncio.run. Copy the current context so ContextVar-based settings (e.g. + # the max_concurrency_scope override set by build_index) propagate into the + # worker thread; .result() re-raises the worker's real exception unchanged. + ctx = contextvars.copy_context() + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: + return pool.submit(ctx.run, asyncio.run, coro).result() def build_index(parsed: ParsedDocument, model: str = None, opt=None) -> dict: diff --git a/tests/test_review_fixes.py b/tests/test_review_fixes.py new file mode 100644 index 0000000..e447231 --- /dev/null +++ b/tests/test_review_fixes.py @@ -0,0 +1,130 @@ +"""Regression tests for the directly-fixable PR #272 review findings.""" +import asyncio + +import pytest + + +# ── #1: page_index() must not capture the imported IndexConfig into opt ─────── +def test_page_index_wrapper_does_not_capture_indexconfig(monkeypatch): + import pageindex.index.page_index as pi + + captured = {} + + def fake_main(doc, opt): + captured["opt"] = opt + return "ok" + + monkeypatch.setattr(pi, "page_index_main", fake_main) + # Previously raised ValidationError (IndexConfig extra='forbid') because + # locals() captured the just-imported IndexConfig class. + result = pi.page_index("dummy.pdf", model="gpt-4o") + assert result == "ok" + assert captured["opt"].model == "gpt-4o" + + +# ── #2: process_none_page_numbers tolerates items with no 'page' key ────────── +def test_process_none_page_numbers_tolerates_missing_page(monkeypatch): + import pageindex.index.page_index as pi + + monkeypatch.setattr( + pi, "add_page_number_to_toc", + lambda pages, item, model: [{"physical_index": ""}], + ) + toc = [ + {"title": "A", "physical_index": 1}, + {"title": "B"}, # no physical_index AND no 'page' -> used to KeyError + ] + page_list = [("p1", 1), ("p2", 1), ("p3", 1)] + result = pi.process_none_page_numbers(toc, page_list) # must not raise + assert result is toc + assert toc[1]["physical_index"] == 2 + + +# ── P4: a real RuntimeError from the coroutine is not masked ────────────────── +def test_run_async_propagates_worker_runtimeerror(): + from pageindex.index.pipeline import _run_async + + async def boom(): + raise RuntimeError("real indexing error") + + async def outer(): + # Inside a running loop -> _run_async uses the worker-thread path; the + # real error must surface, not a bogus "asyncio.run() cannot be called". + with pytest.raises(RuntimeError, match="real indexing error"): + _run_async(boom()) + + asyncio.run(outer()) + + +# ── #9: FileTypeError also subclasses ValueError ───────────────────────────── +def test_filetypeerror_is_valueerror(): + from pageindex.errors import FileTypeError, PageIndexError + + assert issubclass(FileTypeError, ValueError) + assert issubclass(FileTypeError, PageIndexError) + + +# ── #4: md_to_tree coerces legacy 'yes'/'no' flags ─────────────────────────── +def test_md_coerce_bool(): + from pageindex.index.page_index_md import _coerce_bool + + assert _coerce_bool("no") is False # the whole point: 'no' is NOT truthy + assert _coerce_bool("yes") is True + assert _coerce_bool("YES") is True + assert _coerce_bool(True) is True + assert _coerce_bool(False) is False + + +# ── P6: __all__ includes the legacy top-level exports ──────────────────────── +def test_all_includes_legacy_exports(): + import pageindex + + for name in ("page_index", "md_to_tree", "get_document", + "get_document_structure", "get_page_content"): + assert name in pageindex.__all__, f"{name} missing from __all__" + assert hasattr(pageindex, name), f"{name} not importable" + + +# ── P1: keyless local providers pass validation ────────────────────────────── +def test_validate_llm_provider_skips_keyless_providers(): + from pageindex.client import LocalClient + + # These raised PageIndexError("API key not configured...") before the fix. + LocalClient._validate_llm_provider("ollama/llama3") + LocalClient._validate_llm_provider("lm_studio/some-model") + + +# ── #3/P3: missing doc must raise, not return an empty structure ───────────── +def test_local_get_document_structure_missing_raises(tmp_path): + from pageindex.backend.local import LocalBackend + from pageindex.storage.sqlite import SQLiteStorage + from pageindex.errors import DocumentNotFoundError + + backend = LocalBackend( + storage=SQLiteStorage(str(tmp_path / "t.db")), + files_dir=str(tmp_path / "f"), model="gpt-4o", + ) + backend.get_or_create_collection("c") + with pytest.raises(DocumentNotFoundError): + backend.get_document_structure("c", "ghost") + + +# ── #5: delete_collection drops the cached folder_id ───────────────────────── +def test_cloud_delete_collection_clears_folder_cache(monkeypatch): + from pageindex.backend.cloud import CloudBackend + + backend = CloudBackend(api_key="pi-test") + backend._folder_id_cache["papers"] = "folder-123" + monkeypatch.setattr(backend, "_request", lambda *a, **k: {}) + backend.delete_collection("papers") + assert "papers" not in backend._folder_id_cache + + +# ── #6: querying an empty collection raises instead of sending doc_id:[] ────── +def test_cloud_query_empty_collection_raises(monkeypatch): + from pageindex.backend.cloud import CloudBackend + + backend = CloudBackend(api_key="pi-test") + monkeypatch.setattr(backend, "_get_all_doc_ids", lambda col: []) + with pytest.raises(ValueError, match="no documents"): + backend.query("empty", "q") # doc_ids=None -> resolves to []