mirror of
https://github.com/VectifyAI/PageIndex.git
synced 2026-07-18 21:21:05 +02:00
fix: six P2 correctness/robustness cleanups from the SDK review
- legacy call_llm: open AsyncOpenAI via `async with` so the client (and its HTTP connection pool) is closed instead of leaked. - LocalBackend.add_document: fail fast with CollectionNotFoundError when the collection doesn't exist, before the expensive parse + LLM index (previously the missing FK only tripped at save time, after paying for the LLM work). Also raise builtin FileNotFoundError for a missing path instead of FileTypeError (which now means only "unsupported extension"). - Collection.query(doc_ids=None): the empty-collection guard now always runs — previously it was skipped once PAGEINDEX_EXPERIMENTAL_MULTIDOC was set. A single list_documents call serves both the guard and the multi-doc warning (no separate call just to decide whether to warn). - CloudBackend.query_stream: on early consumer break / GeneratorExit, signal the background SSE thread to stop and force-close the response, so it no longer drains the whole stream in the background. Adds regression tests for each (client closed, fail-fast on unknown collection, FileNotFoundError, empty-check under the multidoc env flag, single list call, early-break thread stop). Claude-Session: https://claude.ai/code/session_01Kx5DgKbhK1N8autqXH8SmS
This commit is contained in:
parent
6c948a332b
commit
fe36e25773
8 changed files with 128 additions and 19 deletions
|
|
@ -344,6 +344,11 @@ class CloudBackend:
|
|||
# Queue carries QueryEvent, an Exception to re-raise, or None (end).
|
||||
queue: asyncio.Queue[QueryEvent | Exception | None] = asyncio.Queue()
|
||||
loop = asyncio.get_running_loop()
|
||||
# Set when the consumer stops early (break / GeneratorExit) so the
|
||||
# background thread stops draining the SSE stream instead of pulling
|
||||
# it to completion in the background.
|
||||
stop = threading.Event()
|
||||
resp_holder: dict[str, requests.Response] = {}
|
||||
|
||||
def _put(item: QueryEvent | Exception | None) -> None:
|
||||
try:
|
||||
|
|
@ -374,6 +379,7 @@ class CloudBackend:
|
|||
stream=True,
|
||||
timeout=120,
|
||||
)
|
||||
resp_holder["resp"] = resp
|
||||
if resp.status_code != 200:
|
||||
body = resp.text[:500] if resp.text else ""
|
||||
raise CloudAPIError(
|
||||
|
|
@ -385,6 +391,8 @@ class CloudBackend:
|
|||
current_tool_args: list[str] = []
|
||||
|
||||
for line in resp.iter_lines(decode_unicode=True):
|
||||
if stop.is_set():
|
||||
return # consumer abandoned the stream
|
||||
if not line or not line.startswith("data: "):
|
||||
continue
|
||||
data_str = line[6:]
|
||||
|
|
@ -439,15 +447,26 @@ class CloudBackend:
|
|||
thread = threading.Thread(target=_stream, daemon=True)
|
||||
thread.start()
|
||||
|
||||
while True:
|
||||
item = await queue.get()
|
||||
if item is None:
|
||||
break
|
||||
if isinstance(item, Exception):
|
||||
raise item
|
||||
yield item
|
||||
|
||||
thread.join(timeout=5)
|
||||
try:
|
||||
while True:
|
||||
item = await queue.get()
|
||||
if item is None:
|
||||
break
|
||||
if isinstance(item, Exception):
|
||||
raise item
|
||||
yield item
|
||||
finally:
|
||||
# On early break / GeneratorExit / raised error: tell the thread to
|
||||
# stop and force-close the response so a read blocked mid-stream
|
||||
# unblocks instead of draining the rest in the background.
|
||||
stop.set()
|
||||
resp = resp_holder.get("resp")
|
||||
if resp is not None:
|
||||
try:
|
||||
resp.close()
|
||||
except Exception:
|
||||
pass
|
||||
thread.join(timeout=5)
|
||||
|
||||
def _get_all_doc_ids(self, collection: str) -> list[str]:
|
||||
"""Get all document IDs in a collection."""
|
||||
|
|
|
|||
|
|
@ -13,7 +13,8 @@ from ..storage.protocol import StorageEngine
|
|||
from ..index.pipeline import build_index
|
||||
from ..index.utils import parse_pages, get_pdf_page_content, get_md_page_content, remove_fields
|
||||
from ..backend.protocol import AgentTools
|
||||
from ..errors import FileTypeError, DocumentNotFoundError, IndexingError, PageIndexError
|
||||
from ..errors import (FileTypeError, DocumentNotFoundError, CollectionNotFoundError,
|
||||
IndexingError, PageIndexError)
|
||||
|
||||
_COLLECTION_NAME_RE = re.compile(r'^[a-zA-Z0-9_-]{1,128}$')
|
||||
|
||||
|
|
@ -79,7 +80,16 @@ class LocalBackend:
|
|||
def add_document(self, collection: str, file_path: str) -> str:
|
||||
file_path = os.path.realpath(file_path)
|
||||
if not os.path.isfile(file_path):
|
||||
raise FileTypeError(f"Not a regular file: {file_path}")
|
||||
# Missing path is a file-not-found error, not an unsupported-type one.
|
||||
raise FileNotFoundError(f"No such file: {file_path}")
|
||||
# Fail fast before the expensive parse + LLM indexing if the collection
|
||||
# doesn't exist — otherwise the FK constraint only trips at save time,
|
||||
# after the LLM work (and its cost) is already spent.
|
||||
if collection not in self._storage.list_collections():
|
||||
raise CollectionNotFoundError(
|
||||
f"Collection '{collection}' does not exist; "
|
||||
f"create it first (e.g. client.collection('{collection}'))."
|
||||
)
|
||||
parser = self._resolve_parser(file_path)
|
||||
|
||||
# Dedup is content-only — same file is reused regardless of IndexConfig
|
||||
|
|
|
|||
|
|
@ -94,14 +94,16 @@ class Collection:
|
|||
raise ValueError(
|
||||
"doc_ids cannot be empty; pass None to query the whole collection"
|
||||
)
|
||||
if doc_ids is None and not _multidoc_acked():
|
||||
if doc_ids is None:
|
||||
# One list_documents call serves both the empty-collection guard
|
||||
# (always) and the multi-doc warning (only when not acknowledged).
|
||||
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:
|
||||
if len(docs) > 1 and not _multidoc_acked():
|
||||
warnings.warn(_MULTIDOC_WARNING, UserWarning, stacklevel=2)
|
||||
if stream:
|
||||
return QueryStream(self._backend, self._name, question, doc_ids)
|
||||
|
|
|
|||
|
|
@ -460,12 +460,12 @@ async def call_llm(prompt, api_key, model="gpt-4.1", temperature=0):
|
|||
"""
|
||||
import openai
|
||||
|
||||
client = openai.AsyncOpenAI(api_key=api_key)
|
||||
response = await client.chat.completions.create(
|
||||
model=model,
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
temperature=temperature,
|
||||
)
|
||||
async with openai.AsyncOpenAI(api_key=api_key) as client:
|
||||
response = await client.chat.completions.create(
|
||||
model=model,
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
temperature=temperature,
|
||||
)
|
||||
return response.choices[0].message.content.strip()
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue