mirror of
https://github.com/VectifyAI/PageIndex.git
synced 2026-05-19 18:35:16 +02:00
feat(collection): doc_ids accepts str|list, design cleanups
- Collection.query and Backend.query/query_stream accept doc_ids as str, list[str] or None. Single str is normalized to [str] inside each backend; bare [] is rejected with ValueError at both layers. - wrap_with_doc_context wraps the scoped doc list in <docs>...</docs> and SCOPED_SYSTEM_PROMPT instructs the agent to treat that block as data, not instructions (defense against prompt injection via auto-generated doc_description). - _require_cloud_api now distinguishes api_key="" from api_key=None; the former gives a targeted error pointing at the empty-string vs fall-back-to-local situation when legacy SDK methods are called. - Legacy PageIndexClient.list_documents docstring spells out the return-shape difference vs collection.list_documents() to flag a silent migration footgun (paginated dict with id/name keys vs plain list[dict] with doc_id/doc_name keys). - Remove dead CloudBackend.get_agent_tools stub (not on the Backend protocol; only ever returned an empty AgentTools()) and the SYSTEM_PROMPT alias (OPEN_/SCOPED_SYSTEM_PROMPT are the explicit names now). - README quick start and streaming example now pass doc_ids; new multi-document section shows both str and list forms. - examples/demo_query_modes.py exercises all five query-mode cases (single-doc, multi-doc with/without env var, scoped single, scoped multi) for manual verification.
This commit is contained in:
parent
d7b36aaf3f
commit
a47c36a3f5
13 changed files with 322 additions and 45 deletions
149
examples/demo_query_modes.py
Normal file
149
examples/demo_query_modes.py
Normal file
|
|
@ -0,0 +1,149 @@
|
|||
"""Demo: exercise Collection.query() in all modes.
|
||||
|
||||
Creates a temp workspace with 2 small markdown docs, then runs:
|
||||
Case 1 — single-doc collection, no doc_ids (open mode, no warning)
|
||||
Case 2 — multi-doc collection, no doc_ids (open mode, UserWarning)
|
||||
Case 2b — same as Case 2 + PAGEINDEX_EXPERIMENTAL_MULTIDOC=1 (warning silenced)
|
||||
Case 3 — scoped: doc_ids=[one_id] (no list_documents call)
|
||||
Case 4 — scoped: doc_ids=[id1, id2] (no list_documents call)
|
||||
|
||||
Requirements:
|
||||
- OPENAI_API_KEY (or any LiteLLM-supported provider key) in env or .env
|
||||
"""
|
||||
import asyncio
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
import warnings
|
||||
from pathlib import Path
|
||||
|
||||
# Load .env if present
|
||||
env_file = Path(__file__).parent.parent / ".env"
|
||||
if env_file.exists():
|
||||
for line in env_file.read_text().splitlines():
|
||||
if "=" in line and not line.strip().startswith("#"):
|
||||
k, v = line.split("=", 1)
|
||||
os.environ.setdefault(k.strip(), v.strip())
|
||||
|
||||
from pageindex import PageIndexClient
|
||||
|
||||
|
||||
def banner(text: str) -> None:
|
||||
print("\n" + "=" * 70)
|
||||
print(text)
|
||||
print("=" * 70)
|
||||
|
||||
|
||||
WORKSPACE = tempfile.mkdtemp(prefix="pi_demo_")
|
||||
print(f"Workspace: {WORKSPACE}")
|
||||
|
||||
docs_dir = Path(WORKSPACE) / "docs"
|
||||
docs_dir.mkdir()
|
||||
alpha_md = docs_dir / "alpha.md"
|
||||
alpha_md.write_text(
|
||||
"# Alpha\n\n"
|
||||
"## Introduction\n"
|
||||
"Alpha is about apples and their nutritional value.\n\n"
|
||||
"## Health benefits\n"
|
||||
"Apples contain fiber and vitamin C, support digestion, and may help "
|
||||
"regulate blood sugar.\n"
|
||||
)
|
||||
beta_md = docs_dir / "beta.md"
|
||||
beta_md.write_text(
|
||||
"# Beta\n\n"
|
||||
"## Introduction\n"
|
||||
"Beta is about bananas and potassium.\n\n"
|
||||
"## Energy\n"
|
||||
"Bananas provide quick energy from natural sugars and are rich in "
|
||||
"potassium, supporting muscle function.\n"
|
||||
)
|
||||
|
||||
client = PageIndexClient(model="gpt-4o-2024-11-20", storage_path=WORKSPACE)
|
||||
|
||||
|
||||
async def stream_and_collect(coro_or_stream) -> list[str]:
|
||||
"""Iterate a QueryStream, print tool calls and answer, return tool-call names."""
|
||||
calls: list[str] = []
|
||||
async for ev in coro_or_stream:
|
||||
if ev.type == "tool_call":
|
||||
calls.append(ev.data["name"])
|
||||
print(f" [tool] {ev.data['name']}({ev.data.get('args','')})")
|
||||
elif ev.type == "answer_done":
|
||||
text = str(ev.data)
|
||||
print(f" [answer] {text[:160]}{'...' if len(text) > 160 else ''}")
|
||||
return calls
|
||||
|
||||
|
||||
try:
|
||||
# ── Case 1 ────────────────────────────────────────────────────────────
|
||||
banner("Case 1: single-doc collection, no doc_ids (no warning expected)")
|
||||
single = client.collection("single_test")
|
||||
d_alpha_solo = single.add(str(alpha_md))
|
||||
print(f"Indexed: {d_alpha_solo}")
|
||||
with warnings.catch_warnings(record=True) as caught:
|
||||
warnings.simplefilter("always")
|
||||
answer = single.query("What is alpha about?")
|
||||
uw = [w for w in caught if issubclass(w.category, UserWarning)]
|
||||
print(f"UserWarning count: {len(uw)} (expected 0)")
|
||||
print(f"Answer: {answer[:160]}{'...' if len(answer) > 160 else ''}")
|
||||
|
||||
# ── Case 2 ────────────────────────────────────────────────────────────
|
||||
banner("Case 2: multi-doc collection, no doc_ids (UserWarning expected)")
|
||||
multi = client.collection("multi_test")
|
||||
d1 = multi.add(str(alpha_md))
|
||||
d2 = multi.add(str(beta_md))
|
||||
print(f"Indexed: {d1}, {d2}")
|
||||
with warnings.catch_warnings(record=True) as caught:
|
||||
warnings.simplefilter("always")
|
||||
answer = multi.query("What are these documents about?")
|
||||
uw = [w for w in caught if issubclass(w.category, UserWarning)]
|
||||
print(f"UserWarning count: {len(uw)} (expected 1)")
|
||||
for w in uw:
|
||||
print(f" ⚠ {str(w.message)[:140]}")
|
||||
print(f"Answer: {answer[:160]}{'...' if len(answer) > 160 else ''}")
|
||||
|
||||
# ── Case 2b ───────────────────────────────────────────────────────────
|
||||
banner("Case 2b: same as Case 2 + PAGEINDEX_EXPERIMENTAL_MULTIDOC=1 (silenced)")
|
||||
prev = os.environ.get("PAGEINDEX_EXPERIMENTAL_MULTIDOC")
|
||||
os.environ["PAGEINDEX_EXPERIMENTAL_MULTIDOC"] = "1"
|
||||
try:
|
||||
with warnings.catch_warnings(record=True) as caught:
|
||||
warnings.simplefilter("always")
|
||||
answer = multi.query("What are these documents about?")
|
||||
uw = [w for w in caught if issubclass(w.category, UserWarning)]
|
||||
print(f"UserWarning count: {len(uw)} (expected 0)")
|
||||
print(f"Answer: {answer[:160]}{'...' if len(answer) > 160 else ''}")
|
||||
finally:
|
||||
if prev is None:
|
||||
del os.environ["PAGEINDEX_EXPERIMENTAL_MULTIDOC"]
|
||||
else:
|
||||
os.environ["PAGEINDEX_EXPERIMENTAL_MULTIDOC"] = prev
|
||||
|
||||
# ── Case 3 ────────────────────────────────────────────────────────────
|
||||
banner(f"Case 3: scoped, doc_ids=[{d1[:8]}…] (no list_documents)")
|
||||
|
||||
async def case3():
|
||||
calls = await stream_and_collect(
|
||||
multi.query("What are apples good for?", doc_ids=[d1], stream=True)
|
||||
)
|
||||
assert "list_documents" not in calls, f"unexpected list_documents call: {calls}"
|
||||
print(f"Tools called: {calls}")
|
||||
asyncio.run(case3())
|
||||
|
||||
# ── Case 4 ────────────────────────────────────────────────────────────
|
||||
banner(f"Case 4: scoped, doc_ids=[{d1[:8]}…, {d2[:8]}…] (no list_documents)")
|
||||
|
||||
async def case4():
|
||||
calls = await stream_and_collect(
|
||||
multi.query("Compare alpha and beta briefly.",
|
||||
doc_ids=[d1, d2], stream=True)
|
||||
)
|
||||
assert "list_documents" not in calls, f"unexpected list_documents call: {calls}"
|
||||
print(f"Tools called: {calls}")
|
||||
asyncio.run(case4())
|
||||
|
||||
print("\nAll cases passed.")
|
||||
|
||||
finally:
|
||||
shutil.rmtree(WORKSPACE, ignore_errors=True)
|
||||
print(f"\nCleaned up {WORKSPACE}")
|
||||
Loading…
Add table
Add a link
Reference in a new issue