mirror of
https://github.com/VectifyAI/PageIndex.git
synced 2026-07-24 21:41:04 +02:00
feat: restore status and page/line counts in local get_document
Parsers self-report counts via ParsedDocument.metadata (pymupdf page_count; markdown len(lines)); the backend merges them generically into the stored document row at index time. New page_count/line_count columns are added in place for pre-existing DBs. Local status is always "completed" (indexing is synchronous), matching the cloud backend's field. Align the agentic demo's get_document tool docstring (and its col -> collection naming) accordingly.
This commit is contained in:
Ray
parent
43153fecd9
commit
9980fc9414
7 changed files with 55 additions and 38 deletions
|
|
@ -1,28 +1,22 @@
|
|||
"""
|
||||
Agentic Vectorless RAG with PageIndex — Demo
|
||||
Agentic Vectorless RAG with PageIndex - Demo
|
||||
|
||||
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.
|
||||
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.
|
||||
|
||||
Agent tools:
|
||||
- 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
|
||||
- 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
|
||||
|
||||
Steps:
|
||||
1 — Index a PDF and view its tree structure
|
||||
1 — Index a PDF 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 pageindex openai-agents
|
||||
export OPENAI_API_KEY=your-api-key # or any LiteLLM-supported provider
|
||||
Requirements: pip install openai-agents
|
||||
"""
|
||||
import sys
|
||||
import json
|
||||
|
|
@ -49,7 +43,7 @@ WORKSPACE = _EXAMPLES_DIR / "workspace"
|
|||
AGENT_SYSTEM_PROMPT = """
|
||||
You are PageIndex, a document QA assistant.
|
||||
TOOL USE:
|
||||
- Call get_document() first to confirm the document's name and type.
|
||||
- Call get_document() first to confirm status and page/line 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.
|
||||
|
|
@ -67,7 +61,7 @@ def _normalize_model_for_agents_sdk(model: str) -> str:
|
|||
return model
|
||||
|
||||
|
||||
def query_agent(col, doc_id: str, prompt: str, model: str, verbose: bool = False) -> str:
|
||||
def query_agent(collection, 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.
|
||||
|
|
@ -76,15 +70,15 @@ def query_agent(col, doc_id: str, prompt: str, model: str, verbose: bool = False
|
|||
|
||||
@function_tool
|
||||
def get_document() -> str:
|
||||
"""Get document metadata: name, type, and description."""
|
||||
doc = col.get_document(doc_id)
|
||||
"""Get document metadata: status, page count, name, and description."""
|
||||
doc = collection.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 json.dumps(col.get_document_structure(doc_id), ensure_ascii=False)
|
||||
return json.dumps(collection.get_document_structure(doc_id), ensure_ascii=False)
|
||||
|
||||
@function_tool
|
||||
def get_page_content(pages: str) -> str:
|
||||
|
|
@ -93,7 +87,7 @@ def query_agent(col, doc_id: str, prompt: str, model: str, verbose: bool = False
|
|||
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 json.dumps(col.get_page_content(doc_id, pages), ensure_ascii=False)
|
||||
return json.dumps(collection.get_page_content(doc_id, pages), ensure_ascii=False)
|
||||
|
||||
agent = Agent(
|
||||
name="PageIndex",
|
||||
|
|
@ -173,24 +167,24 @@ if __name__ == "__main__":
|
|||
|
||||
# Setup: self-hosted local client + a collection
|
||||
client = LocalClient(storage_path=str(WORKSPACE))
|
||||
col = client.collection("agentic-demo")
|
||||
collection = 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)
|
||||
# Content-hash dedup: re-running reuses the existing doc_id, no re-index.
|
||||
doc_id = col.add(str(PDF_PATH))
|
||||
doc_id = collection.add(str(PDF_PATH))
|
||||
print(f"\ndoc_id: {doc_id}")
|
||||
print("\nTree Structure (top-level sections):")
|
||||
for node in col.get_document_structure(doc_id):
|
||||
for node in collection.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)
|
||||
meta = col.get_document(doc_id)
|
||||
meta = collection.get_document(doc_id)
|
||||
meta.pop("structure", None)
|
||||
print("\n" + json.dumps(meta, ensure_ascii=False, indent=2))
|
||||
|
||||
|
|
@ -200,4 +194,4 @@ if __name__ == "__main__":
|
|||
print("=" * 60)
|
||||
question = "Explain Attention Residuals in simple language."
|
||||
print(f"\nQuestion: '{question}'")
|
||||
query_agent(col, doc_id, question, client.retrieve_model, verbose=True)
|
||||
query_agent(collection, doc_id, question, client.retrieve_model, verbose=True)
|
||||
|
|
|
|||
|
|
@ -134,6 +134,7 @@ class LocalBackend:
|
|||
"file_path": str(managed_path),
|
||||
"file_hash": file_hash,
|
||||
"doc_type": ext.lstrip("."),
|
||||
**(parsed.metadata or {}), # parser-reported, e.g. page_count / line_count
|
||||
"structure": result["structure"],
|
||||
"pages": pages,
|
||||
})
|
||||
|
|
@ -184,6 +185,7 @@ class LocalBackend:
|
|||
use in agent/LLM contexts as it can exhaust the context window.
|
||||
"""
|
||||
doc = self._require_document(collection, doc_id)
|
||||
doc["status"] = "completed" # local indexing is synchronous
|
||||
doc["structure"] = self._storage.get_document_structure(collection, doc_id)
|
||||
if include_text:
|
||||
pages = self._storage.get_pages(collection, doc_id) or []
|
||||
|
|
|
|||
|
|
@ -24,7 +24,8 @@ class MarkdownParser:
|
|||
headers = self._extract_headers(lines)
|
||||
nodes = self._build_nodes(headers, lines, model, doc_title=path.stem)
|
||||
|
||||
return ParsedDocument(doc_name=path.stem, nodes=nodes)
|
||||
return ParsedDocument(doc_name=path.stem, nodes=nodes,
|
||||
metadata={"line_count": len(lines)})
|
||||
|
||||
def _extract_headers(self, lines: list[str]) -> list[dict]:
|
||||
header_pattern = r"^(#{1,6})\s+(.+)$"
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ class PdfParser:
|
|||
nodes = []
|
||||
|
||||
with pymupdf.open(str(path)) as doc:
|
||||
page_count = doc.page_count
|
||||
for i, page in enumerate(doc):
|
||||
page_num = i + 1
|
||||
if images_dir:
|
||||
|
|
@ -35,7 +36,8 @@ class PdfParser:
|
|||
images=images if images else None,
|
||||
))
|
||||
|
||||
return ParsedDocument(doc_name=path.stem, nodes=nodes)
|
||||
return ParsedDocument(doc_name=path.stem, nodes=nodes,
|
||||
metadata={"page_count": page_count})
|
||||
|
||||
@staticmethod
|
||||
def _extract_page_with_images(doc, page, page_num: int,
|
||||
|
|
|
|||
|
|
@ -19,6 +19,9 @@ class ParsedDocument:
|
|||
"""Unified parser output. Always a flat list of ContentNode."""
|
||||
doc_name: str
|
||||
nodes: list[ContentNode]
|
||||
# Doc-level fields merged into the stored document record at index time.
|
||||
# The built-in storage only persists keys it has columns for
|
||||
# (currently page_count / line_count); other keys are dropped.
|
||||
metadata: dict | None = None
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -99,6 +99,8 @@ class SQLiteStorage:
|
|||
file_path TEXT,
|
||||
file_hash TEXT,
|
||||
doc_type TEXT NOT NULL,
|
||||
page_count INTEGER,
|
||||
line_count INTEGER,
|
||||
structure JSON,
|
||||
pages JSON,
|
||||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
|
|
@ -107,6 +109,14 @@ class SQLiteStorage:
|
|||
CREATE INDEX IF NOT EXISTS idx_docs_collection ON documents(collection_name);
|
||||
CREATE INDEX IF NOT EXISTS idx_docs_hash ON documents(collection_name, file_hash);
|
||||
""")
|
||||
# DBs created before the count columns existed: add them in place.
|
||||
cols = {r[1] for r in conn.execute("PRAGMA table_info(documents)")}
|
||||
for name in ("page_count", "line_count"):
|
||||
if name not in cols:
|
||||
try:
|
||||
conn.execute(f"ALTER TABLE documents ADD COLUMN {name} INTEGER")
|
||||
except sqlite3.OperationalError:
|
||||
pass # concurrent open of the same legacy DB already added it
|
||||
conn.commit()
|
||||
|
||||
def create_collection(self, name: str) -> None:
|
||||
|
|
@ -145,10 +155,11 @@ class SQLiteStorage:
|
|||
conn = self._get_conn()
|
||||
conn.execute(
|
||||
"""INSERT INTO documents
|
||||
(doc_id, collection_name, doc_name, doc_description, file_path, file_hash, doc_type, structure, pages)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""",
|
||||
(doc_id, collection_name, doc_name, doc_description, file_path, file_hash, doc_type, page_count, line_count, structure, pages)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
|
||||
(doc_id, collection, doc.get("doc_name"), doc.get("doc_description"),
|
||||
doc.get("file_path"), doc.get("file_hash"), doc["doc_type"],
|
||||
doc.get("page_count"), doc.get("line_count"),
|
||||
json.dumps(doc.get("structure", [])),
|
||||
json.dumps(doc.get("pages")) if doc.get("pages") else None),
|
||||
)
|
||||
|
|
@ -165,13 +176,15 @@ class SQLiteStorage:
|
|||
def get_document(self, collection: str, doc_id: str) -> dict:
|
||||
conn = self._get_conn()
|
||||
row = conn.execute(
|
||||
"SELECT doc_id, doc_name, doc_description, file_path, doc_type FROM documents WHERE doc_id = ? AND collection_name = ?",
|
||||
"SELECT doc_id, doc_name, doc_description, file_path, doc_type, page_count, line_count FROM documents WHERE doc_id = ? AND collection_name = ?",
|
||||
(doc_id, collection),
|
||||
).fetchone()
|
||||
if not row:
|
||||
return {}
|
||||
return {"doc_id": row[0], "doc_name": row[1], "doc_description": row[2],
|
||||
"file_path": row[3], "doc_type": row[4]}
|
||||
doc = {"doc_id": row[0], "doc_name": row[1], "doc_description": row[2],
|
||||
"file_path": row[3], "doc_type": row[4]}
|
||||
doc.update({k: v for k, v in (("page_count", row[5]), ("line_count", row[6])) if v is not None})
|
||||
return doc
|
||||
|
||||
def get_document_structure(self, collection: str, doc_id: str) -> list:
|
||||
conn = self._get_conn()
|
||||
|
|
|
|||
|
|
@ -28,11 +28,13 @@ class _DocumentDetailRequired(DocumentInfo):
|
|||
class DocumentDetail(_DocumentDetailRequired, total=False):
|
||||
"""A document with its tree, as returned by ``get_document()``.
|
||||
|
||||
``structure`` is always present; ``file_path`` is local-only and
|
||||
``status`` is cloud-only, hence total=False for those two only.
|
||||
``structure`` is always present; the remaining fields are
|
||||
backend-specific, hence total=False.
|
||||
"""
|
||||
file_path: str # local backend only
|
||||
status: str # cloud backend only
|
||||
file_path: str # local backend only
|
||||
status: str # local: always "completed" (indexing is synchronous); cloud: server-reported
|
||||
page_count: int # local backend, PDF documents
|
||||
line_count: int # local backend, Markdown documents
|
||||
|
||||
|
||||
class PageContent(TypedDict, total=False):
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue