mirror of
https://github.com/VectifyAI/PageIndex.git
synced 2026-07-15 21:11:05 +02:00
d97231b480
- local get_agent_tools: `set(doc_ids) if doc_ids is not None else None` so doc_ids=[] is a scope of nothing (reject all), not open mode. The public query path already guarded []; this hardens direct callers. - legacy get_tree: send summary=true/false (lowercase) instead of Python's capitalized True/False, matching the modern CloudBackend and the API. - markdown parser: track the opening fence character so a ```-fence isn't closed by a ~~~ line (CommonMark), keeping '#'-lines inside it out of headings. - dedupe the cloud base URL: single API_BASE in cloud_api, referenced by CloudBackend and PageIndexClient (was three independent copies). Regression tests for each.
268 lines
11 KiB
Python
268 lines
11 KiB
Python
# tests/sdk/test_local_backend.py
|
|
import asyncio
|
|
import json
|
|
import pytest
|
|
from pathlib import Path
|
|
from pageindex.backend.local import LocalBackend
|
|
from pageindex.storage.sqlite import SQLiteStorage
|
|
from pageindex.errors import FileTypeError, DocumentNotFoundError
|
|
|
|
|
|
@pytest.fixture
|
|
def backend(tmp_path):
|
|
storage = SQLiteStorage(str(tmp_path / "test.db"))
|
|
files_dir = tmp_path / "files"
|
|
return LocalBackend(storage=storage, files_dir=str(files_dir), model="gpt-4o")
|
|
|
|
|
|
def test_collection_lifecycle(backend):
|
|
backend.get_or_create_collection("papers")
|
|
assert "papers" in backend.list_collections()
|
|
backend.delete_collection("papers")
|
|
assert "papers" not in backend.list_collections()
|
|
|
|
|
|
def test_list_documents_empty(backend):
|
|
backend.get_or_create_collection("papers")
|
|
assert backend.list_documents("papers") == []
|
|
|
|
|
|
def test_unsupported_file_type_raises(backend, tmp_path):
|
|
backend.get_or_create_collection("papers")
|
|
bad_file = tmp_path / "test.xyz"
|
|
bad_file.write_text("hello")
|
|
with pytest.raises(FileTypeError):
|
|
backend.add_document("papers", str(bad_file))
|
|
|
|
|
|
def test_add_document_on_empty_markdown_file_does_not_crash(tmp_path):
|
|
"""An empty/whitespace-only .md file used to route into the PDF-oriented
|
|
TOC-detection pipeline (no node ever has 'level' set), wasting an LLM call
|
|
and then raising IndexingError. Must complete instantly with zero LLM
|
|
calls when summary/description are off."""
|
|
from pageindex.config import IndexConfig
|
|
|
|
storage = SQLiteStorage(str(tmp_path / "test.db"))
|
|
backend = LocalBackend(
|
|
storage=storage, files_dir=str(tmp_path / "files"), model="gpt-4o",
|
|
index_config=IndexConfig(if_add_node_summary=False, if_add_doc_description=False),
|
|
)
|
|
backend.get_or_create_collection("papers")
|
|
empty_md = tmp_path / "empty.md"
|
|
empty_md.write_text(" \n\n \n")
|
|
|
|
doc_id = backend.add_document("papers", str(empty_md)) # must not raise
|
|
|
|
assert backend.get_document_structure("papers", doc_id) == []
|
|
|
|
|
|
def test_register_custom_parser(backend):
|
|
from pageindex.parser.protocol import ParsedDocument, ContentNode
|
|
|
|
class TxtParser:
|
|
def supported_extensions(self):
|
|
return [".txt"]
|
|
def parse(self, file_path, **kwargs):
|
|
text = Path(file_path).read_text()
|
|
return ParsedDocument(doc_name="test", nodes=[
|
|
ContentNode(content=text, tokens=len(text.split()), title="Content", index=1, level=1)
|
|
])
|
|
|
|
backend.register_parser(TxtParser())
|
|
# Now .txt should be supported (won't raise FileTypeError)
|
|
assert backend._resolve_parser("test.txt") is not None
|
|
|
|
|
|
# ── Scoped-mode agent tools ──────────────────────────────────────────────────
|
|
|
|
@pytest.fixture
|
|
def populated_backend(backend):
|
|
"""Backend with a 'papers' collection containing two stub docs."""
|
|
backend.get_or_create_collection("papers")
|
|
for did, name, desc in [
|
|
("d1", "alpha.pdf", "About alpha."),
|
|
("d2", "beta.pdf", "About beta."),
|
|
]:
|
|
backend._storage.save_document("papers", did, {
|
|
"doc_name": name, "doc_description": desc,
|
|
"doc_type": "pdf", "file_path": f"/tmp/{name}", "structure": [],
|
|
})
|
|
return backend
|
|
|
|
|
|
def _invoke_tool(tool, args: dict) -> str:
|
|
"""Run a FunctionTool synchronously with a minimal ToolContext."""
|
|
from agents.tool_context import ToolContext
|
|
ctx = ToolContext(context=None, tool_name=tool.name,
|
|
tool_call_id="test", tool_arguments=json.dumps(args))
|
|
return asyncio.run(tool.on_invoke_tool(ctx, json.dumps(args)))
|
|
|
|
|
|
def test_open_mode_includes_list_documents(populated_backend):
|
|
tools = populated_backend.get_agent_tools("papers", doc_ids=None)
|
|
names = {t.name for t in tools.function_tools}
|
|
assert names == {"list_documents", "get_document", "get_document_structure", "get_page_content"}
|
|
|
|
|
|
def test_scoped_mode_excludes_list_documents(populated_backend):
|
|
tools = populated_backend.get_agent_tools("papers", doc_ids=["d1"])
|
|
names = {t.name for t in tools.function_tools}
|
|
assert "list_documents" not in names
|
|
assert names == {"get_document", "get_document_structure", "get_page_content"}
|
|
|
|
|
|
def test_scoped_mode_rejects_out_of_scope_doc_id(populated_backend):
|
|
tools = populated_backend.get_agent_tools("papers", doc_ids=["d1"])
|
|
by_name = {t.name: t for t in tools.function_tools}
|
|
out = json.loads(_invoke_tool(by_name["get_document"], {"doc_id": "d2"}))
|
|
assert "error" in out
|
|
assert "not in scope" in out["error"]
|
|
assert out["allowed_doc_ids"] == ["d1"]
|
|
|
|
|
|
def test_scoped_mode_allows_in_scope_doc_id(populated_backend):
|
|
tools = populated_backend.get_agent_tools("papers", doc_ids=["d1"])
|
|
by_name = {t.name: t for t in tools.function_tools}
|
|
out = json.loads(_invoke_tool(by_name["get_document"], {"doc_id": "d1"}))
|
|
assert out.get("doc_name") == "alpha.pdf"
|
|
|
|
|
|
def test_empty_doc_ids_is_scoped_to_nothing_not_open_mode(populated_backend):
|
|
# doc_ids=[] means "scope to no documents", NOT open mode. It must exclude
|
|
# list_documents and reject every doc_id — otherwise an empty list would
|
|
# collapse to None (truthiness) and silently grant access to the whole
|
|
# collection.
|
|
tools = populated_backend.get_agent_tools("papers", doc_ids=[])
|
|
by_name = {t.name: t for t in tools.function_tools}
|
|
assert "list_documents" not in by_name
|
|
out = json.loads(_invoke_tool(by_name["get_document"], {"doc_id": "d1"}))
|
|
assert "error" in out and "not in scope" in out["error"]
|
|
|
|
|
|
@pytest.mark.parametrize("bad_pages", ["all", "5-", "abc", "3-1"])
|
|
def test_get_page_content_returns_actionable_error_for_bad_page_spec(populated_backend, bad_pages):
|
|
# A malformed page spec must come back as a correctable JSON error (like the
|
|
# legacy retrieval tool), not the agent SDK's generic tool-failure fallback,
|
|
# so the model can retry with a valid range instead of giving up.
|
|
tools = populated_backend.get_agent_tools("papers", doc_ids=["d1"])
|
|
by_name = {t.name: t for t in tools.function_tools}
|
|
out = json.loads(_invoke_tool(by_name["get_page_content"], {"doc_id": "d1", "pages": bad_pages}))
|
|
assert "error" in out
|
|
assert "Invalid pages format" in out["error"]
|
|
|
|
|
|
def test_wrap_with_doc_context_single(populated_backend):
|
|
from pageindex.agent import wrap_with_doc_context
|
|
docs = populated_backend._scoped_docs("papers", ["d1"])
|
|
wrapped = wrap_with_doc_context(docs, "what is this?")
|
|
assert "d1: alpha.pdf — About alpha." in wrapped
|
|
assert "specified the following document" in wrapped
|
|
assert "<docs>" in wrapped and "</docs>" in wrapped
|
|
assert "User question: what is this?" in wrapped
|
|
|
|
|
|
def test_wrap_with_doc_context_multi(populated_backend):
|
|
from pageindex.agent import wrap_with_doc_context
|
|
docs = populated_backend._scoped_docs("papers", ["d1", "d2"])
|
|
wrapped = wrap_with_doc_context(docs, "compare them")
|
|
assert "d1: alpha.pdf — About alpha." in wrapped
|
|
assert "d2: beta.pdf — About beta." in wrapped
|
|
assert "specified the following documents" in wrapped
|
|
assert "<docs>" in wrapped and "</docs>" in wrapped
|
|
assert "User question: compare them" in wrapped
|
|
|
|
|
|
def test_scoped_docs_raises_on_missing(populated_backend):
|
|
with pytest.raises(DocumentNotFoundError, match="nonexistent"):
|
|
populated_backend._scoped_docs("papers", ["d1", "nonexistent"])
|
|
|
|
|
|
def test_normalize_doc_ids():
|
|
assert LocalBackend._normalize_doc_ids("d1") == ["d1"]
|
|
assert LocalBackend._normalize_doc_ids(["d1", "d2"]) == ["d1", "d2"]
|
|
assert LocalBackend._normalize_doc_ids(None) is None
|
|
|
|
|
|
def test_normalize_doc_ids_rejects_empty_list():
|
|
with pytest.raises(ValueError, match="cannot be empty"):
|
|
LocalBackend._normalize_doc_ids([])
|
|
|
|
|
|
# ── error taxonomy: missing docs raise DocumentNotFoundError ─────────────────
|
|
|
|
def test_get_document_missing_raises(backend):
|
|
backend.get_or_create_collection("papers")
|
|
with pytest.raises(DocumentNotFoundError, match="ghost"):
|
|
backend.get_document("papers", "ghost")
|
|
|
|
|
|
def test_delete_document_missing_raises(backend):
|
|
backend.get_or_create_collection("papers")
|
|
with pytest.raises(DocumentNotFoundError, match="ghost"):
|
|
backend.delete_document("papers", "ghost")
|
|
|
|
|
|
def test_delete_collection_rejects_path_traversal(backend, tmp_path):
|
|
# Regression: an unvalidated name like "../.." would rmtree outside files_dir.
|
|
from pageindex.errors import PageIndexError
|
|
canary = tmp_path / "canary.txt"
|
|
canary.write_text("still here")
|
|
with pytest.raises(PageIndexError, match="Invalid collection name"):
|
|
backend.delete_collection("../..")
|
|
assert canary.exists()
|
|
|
|
|
|
@pytest.mark.parametrize("bad_name", ["papers\n", "\npapers", "papers\n\n"])
|
|
def test_get_or_create_collection_rejects_trailing_newline(backend, bad_name):
|
|
# Regression: Python's $ matches just before a final \n, so a $-anchored
|
|
# .match() accepted "papers\n"; get_or_create_collection then hit the SQL
|
|
# CHECK via INSERT OR IGNORE, silently created no row, and handed back a
|
|
# Collection that failed later on add(). .fullmatch() rejects it up front.
|
|
from pageindex.errors import PageIndexError
|
|
with pytest.raises(PageIndexError, match="Invalid collection name"):
|
|
backend.get_or_create_collection(bad_name)
|
|
|
|
|
|
def test_add_document_missing_file_raises_file_not_found(backend, tmp_path):
|
|
backend.get_or_create_collection("papers")
|
|
with pytest.raises(FileNotFoundError):
|
|
backend.add_document("papers", str(tmp_path / "nope.pdf"))
|
|
|
|
|
|
def test_add_document_unknown_collection_fails_fast(backend, tmp_path):
|
|
from pageindex.errors import CollectionNotFoundError
|
|
pdf = tmp_path / "doc.pdf"
|
|
pdf.write_bytes(b"%PDF-1.4")
|
|
# Collection never created -> must raise before any parse/LLM work.
|
|
with pytest.raises(CollectionNotFoundError, match="does not exist"):
|
|
backend.add_document("ghost-collection", str(pdf))
|
|
|
|
|
|
def test_add_document_race_returns_existing_id(backend, tmp_path, monkeypatch):
|
|
"""If the pre-check misses but the INSERT hits UNIQUE (concurrent add),
|
|
add_document must clean up and return the winner's doc_id, not duplicate."""
|
|
import pageindex.backend.local as local_mod
|
|
pdf = tmp_path / "doc.pdf"
|
|
pdf.write_bytes(b"%PDF-1.4 body")
|
|
backend.get_or_create_collection("papers")
|
|
|
|
# Pretend a winning add already stored this content under "winner-id".
|
|
file_hash = backend._file_hash(str(pdf))
|
|
backend._storage.save_document("papers", "winner-id", {
|
|
"doc_name": "doc", "doc_type": "pdf", "file_hash": file_hash, "structure": [],
|
|
})
|
|
# Pre-check misses (returns None) so we reach the INSERT; the post-conflict
|
|
# lookup then returns the winner's id.
|
|
calls = {"n": 0}
|
|
def fake_find(col, h):
|
|
calls["n"] += 1
|
|
return None if calls["n"] == 1 else "winner-id"
|
|
monkeypatch.setattr(backend._storage, "find_document_by_hash", fake_find)
|
|
# avoid real parsing/LLM: stub parser + build_index
|
|
monkeypatch.setattr(backend, "_resolve_parser", lambda p: type("P", (), {
|
|
"parse": lambda self, fp, **k: type("PD", (), {"doc_name": "doc", "nodes": []})()
|
|
})())
|
|
monkeypatch.setattr(local_mod, "build_index", lambda parsed, model=None, opt=None: {"structure": [], "doc_description": ""})
|
|
|
|
result = backend.add_document("papers", str(pdf))
|
|
assert result == "winner-id"
|