mirror of
https://github.com/VectifyAI/PageIndex.git
synced 2026-05-19 18:35:16 +02:00
a47c36a3f5
- 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.
108 lines
3.9 KiB
Python
108 lines
3.9 KiB
Python
# pageindex/collection.py
|
|
from __future__ import annotations
|
|
import os
|
|
import warnings
|
|
from typing import AsyncIterator
|
|
from .events import QueryEvent
|
|
from .backend.protocol import Backend
|
|
|
|
|
|
def _multidoc_acked() -> bool:
|
|
return os.getenv("PAGEINDEX_EXPERIMENTAL_MULTIDOC", "").lower() in ("1", "true", "yes")
|
|
|
|
|
|
_MULTIDOC_WARNING = (
|
|
"Querying the entire collection (no doc_ids) is experimental — a naive "
|
|
"first implementation that lets the agent pick docs from auto-generated "
|
|
"descriptions. Better cross-document retrieval is on the way. Pass "
|
|
"doc_ids=[...] for reliable results, or set "
|
|
"PAGEINDEX_EXPERIMENTAL_MULTIDOC=1 to silence this warning."
|
|
)
|
|
|
|
|
|
class QueryStream:
|
|
"""Wraps backend.query_stream() as an async iterable object."""
|
|
|
|
def __init__(self, backend: Backend, collection: str, question: str,
|
|
doc_ids: list[str] | None = None):
|
|
self._backend = backend
|
|
self._collection = collection
|
|
self._question = question
|
|
self._doc_ids = doc_ids
|
|
|
|
async def stream_events(self) -> AsyncIterator[QueryEvent]:
|
|
async for event in self._backend.query_stream(
|
|
self._collection, self._question, self._doc_ids
|
|
):
|
|
yield event
|
|
|
|
def __aiter__(self):
|
|
return self.stream_events()
|
|
|
|
|
|
class Collection:
|
|
def __init__(self, name: str, backend: Backend):
|
|
self._name = name
|
|
self._backend = backend
|
|
|
|
@property
|
|
def name(self) -> str:
|
|
return self._name
|
|
|
|
def add(self, file_path: str) -> str:
|
|
return self._backend.add_document(self._name, file_path)
|
|
|
|
def list_documents(self) -> list[dict]:
|
|
return self._backend.list_documents(self._name)
|
|
|
|
def get_document(self, doc_id: str, include_text: bool = False) -> dict:
|
|
return self._backend.get_document(self._name, doc_id, include_text=include_text)
|
|
|
|
def get_document_structure(self, doc_id: str) -> list:
|
|
return self._backend.get_document_structure(self._name, doc_id)
|
|
|
|
def get_page_content(self, doc_id: str, pages: str) -> list:
|
|
return self._backend.get_page_content(self._name, doc_id, pages)
|
|
|
|
def delete_document(self, doc_id: str) -> None:
|
|
self._backend.delete_document(self._name, doc_id)
|
|
|
|
def query(self, question: str,
|
|
doc_ids: str | list[str] | None = None,
|
|
stream: bool = False) -> str | QueryStream:
|
|
"""Query documents in this collection.
|
|
|
|
- stream=False: returns answer string (sync)
|
|
- stream=True: returns async iterable of QueryEvent
|
|
|
|
``doc_ids`` can be a single doc id (``str``) or a list. ``None`` queries
|
|
the entire collection (experimental).
|
|
|
|
Usage:
|
|
answer = col.query("question", doc_ids=doc_id) # single
|
|
answer = col.query("question", doc_ids=[d1, d2]) # multi
|
|
async for event in col.query("question", doc_ids=doc_id, stream=True):
|
|
...
|
|
|
|
Passing doc_ids=None queries the entire collection — this is
|
|
experimental; emits a UserWarning unless PAGEINDEX_EXPERIMENTAL_MULTIDOC
|
|
is set.
|
|
"""
|
|
if isinstance(doc_ids, str):
|
|
doc_ids = [doc_ids]
|
|
elif doc_ids == []:
|
|
raise ValueError(
|
|
"doc_ids cannot be empty; pass None to query the whole collection"
|
|
)
|
|
if doc_ids is None and not _multidoc_acked():
|
|
docs = self._backend.list_documents(self._name)
|
|
if not docs:
|
|
raise ValueError(
|
|
f"Cannot query collection '{self._name}': it is empty. "
|
|
"Add documents with col.add(...) first."
|
|
)
|
|
if len(docs) > 1:
|
|
warnings.warn(_MULTIDOC_WARNING, UserWarning, stacklevel=2)
|
|
if stream:
|
|
return QueryStream(self._backend, self._name, question, doc_ids)
|
|
return self._backend.query(self._name, question, doc_ids)
|