diff --git a/pageindex/backend/cloud.py b/pageindex/backend/cloud.py index 095b9d9..0c2a0d2 100644 --- a/pageindex/backend/cloud.py +++ b/pageindex/backend/cloud.py @@ -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.""" diff --git a/pageindex/backend/local.py b/pageindex/backend/local.py index 812be1e..a017730 100644 --- a/pageindex/backend/local.py +++ b/pageindex/backend/local.py @@ -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 diff --git a/pageindex/collection.py b/pageindex/collection.py index 053fb43..4cf2f64 100644 --- a/pageindex/collection.py +++ b/pageindex/collection.py @@ -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) diff --git a/pageindex/index/utils.py b/pageindex/index/utils.py index 49e3878..60cb770 100644 --- a/pageindex/index/utils.py +++ b/pageindex/index/utils.py @@ -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() diff --git a/tests/test_cloud_backend.py b/tests/test_cloud_backend.py index 9bb1105..35b8629 100644 --- a/tests/test_cloud_backend.py +++ b/tests/test_cloud_backend.py @@ -219,3 +219,32 @@ def test_query_uses_long_timeout_and_single_attempt(monkeypatch): backend.query("col", "q", doc_ids=["d1"]) assert len(calls) == 1 assert calls[0]["timeout"] == 300 + + +def test_query_stream_early_break_stops_background_thread(monkeypatch): + """Consumer breaking early must signal the SSE thread to stop, not let it + drain the whole stream in the background.""" + import threading + import time as _real_time # autouse fixture stubs cloud_mod.time.sleep, not this + backend = CloudBackend(api_key="pi-test") + drained_all = threading.Event() + + class SlowResponse: + status_code = 200 + text = "" + def iter_lines(self, decode_unicode=True): + for i in range(1000): + yield _sse("text", f"chunk{i} ") + _real_time.sleep(0.002) # pace so the consumer reliably breaks first + drained_all.set() # only reached if the thread was NOT stopped + def close(self): + pass + + monkeypatch.setattr(cloud_mod.requests, "post", lambda *a, **k: SlowResponse()) + + async def _run(): + async for _ in backend.query_stream("col", "q", doc_ids=["d1"]): + break # consume one event, then abandon + + asyncio.run(_run()) + assert not drained_all.is_set() diff --git a/tests/test_collection.py b/tests/test_collection.py index 5f4221d..9a24f65 100644 --- a/tests/test_collection.py +++ b/tests/test_collection.py @@ -94,3 +94,26 @@ def test_query_accepts_str_doc_id(col): def test_query_rejects_empty_list(col): with pytest.raises(ValueError, match="cannot be empty"): col.query("what?", doc_ids=[]) + + +def test_empty_collection_check_runs_even_when_multidoc_acked(monkeypatch): + from unittest.mock import MagicMock + from pageindex.collection import Collection + monkeypatch.setenv("PAGEINDEX_EXPERIMENTAL_MULTIDOC", "1") + backend = MagicMock() + backend.list_documents.return_value = [] + col = Collection(name="papers", backend=backend) + with pytest.raises(ValueError, match="empty"): + col.query("q") # doc_ids=None, collection empty -> must still raise + + +def test_whole_collection_query_lists_documents_once(monkeypatch): + from unittest.mock import MagicMock + from pageindex.collection import Collection + monkeypatch.setenv("PAGEINDEX_EXPERIMENTAL_MULTIDOC", "1") # silence warning path + backend = MagicMock() + backend.list_documents.return_value = [{"doc_id": "d1"}, {"doc_id": "d2"}] + backend.query.return_value = "ans" + col = Collection(name="papers", backend=backend) + col.query("q") + assert backend.list_documents.call_count == 1 # single call at the collection layer diff --git a/tests/test_legacy_utils_contract.py b/tests/test_legacy_utils_contract.py index 2abf5fb..9e67da4 100644 --- a/tests/test_legacy_utils_contract.py +++ b/tests/test_legacy_utils_contract.py @@ -75,6 +75,7 @@ def test_print_tree_keeps_legacy_exclude_fields(capsys): def test_call_llm_keeps_legacy_async_openai_contract(monkeypatch): calls = [] + closed = [] class FakeCompletions: async def create(self, **kwargs): @@ -88,6 +89,15 @@ def test_call_llm_keeps_legacy_async_openai_contract(monkeypatch): self.api_key = api_key self.chat = SimpleNamespace(completions=FakeCompletions()) + # call_llm must open the client as an async context manager so it is + # closed (no leaked HTTP connection pool). + async def __aenter__(self): + return self + + async def __aexit__(self, *exc): + closed.append(True) + return False + fake_openai = SimpleNamespace(AsyncOpenAI=FakeAsyncOpenAI) monkeypatch.setitem(sys.modules, "openai", fake_openai) @@ -99,6 +109,7 @@ def test_call_llm_keeps_legacy_async_openai_contract(monkeypatch): )) assert result == "answer" + assert closed == [True] # client was closed assert calls == [{ "model": "gpt-test", "messages": [{"role": "user", "content": "hello"}], diff --git a/tests/test_local_backend.py b/tests/test_local_backend.py index c4d6c11..b78cf55 100644 --- a/tests/test_local_backend.py +++ b/tests/test_local_backend.py @@ -165,3 +165,18 @@ def test_delete_collection_rejects_path_traversal(backend, tmp_path): with pytest.raises(PageIndexError, match="Invalid collection name"): backend.delete_collection("../..") assert canary.exists() + + +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))