diff --git a/README.md b/README.md index 643d990..adb7faf 100644 --- a/README.md +++ b/README.md @@ -179,7 +179,7 @@ import asyncio async def main(): async for ev in col.query("Explain multi-head attention", doc_ids=doc_id, stream=True): - if ev.type == "answer_delta": + if ev.type == "text_delta": print(ev.data, end="", flush=True) elif ev.type == "tool_call": print(f"\n[tool] {ev.data['name']}") @@ -187,7 +187,7 @@ async def main(): asyncio.run(main()) ``` -`ev.type` is one of: `tool_call`, `tool_result`, `answer_delta`, `answer_done`. +`ev.type` is one of: `tool_call`, `tool_result`, `text_delta`, `text_done`. A `text_done` fires each time a text message completes — a local agentic query may emit several as the agent narrates between tool calls; the last `text_done` before the stream ends carries the final answer. ### Multi-document collections (experimental) diff --git a/examples/cloud_demo.py b/examples/cloud_demo.py index cd3344b..6455d7d 100644 --- a/examples/cloud_demo.py +++ b/examples/cloud_demo.py @@ -46,7 +46,7 @@ stream = col.query("What is the main contribution of this paper?", stream=True) async def main(): streamed_text = False async for event in stream: - if event.type == "answer_delta": + if event.type == "text_delta": print(event.data, end="", flush=True) streamed_text = True elif event.type == "tool_call": @@ -55,7 +55,7 @@ async def main(): streamed_text = False args = event.data.get("args", "") print(f"[tool call] {event.data['name']}({args})") - elif event.type == "answer_done": + elif event.type == "text_done": print() streamed_text = False diff --git a/examples/demo_query_modes.py b/examples/demo_query_modes.py index a858ed5..86c620b 100644 --- a/examples/demo_query_modes.py +++ b/examples/demo_query_modes.py @@ -68,7 +68,7 @@ async def stream_and_collect(coro_or_stream) -> list[str]: 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": + elif ev.type == "text_done": text = str(ev.data) print(f" [answer] {text[:160]}{'...' if len(text) > 160 else ''}") return calls diff --git a/examples/local_demo.py b/examples/local_demo.py index f98d25d..fcbc9bb 100644 --- a/examples/local_demo.py +++ b/examples/local_demo.py @@ -51,7 +51,7 @@ stream = col.query( async def main(): streamed_text = False async for event in stream: - if event.type == "answer_delta": + if event.type == "text_delta": print(event.data, end="", flush=True) streamed_text = True elif event.type == "tool_call": @@ -62,7 +62,7 @@ async def main(): elif event.type == "tool_result": preview = str(event.data)[:200] + "..." if len(str(event.data)) > 200 else event.data print(f"[tool output] {preview}") - elif event.type == "answer_done": + elif event.type == "text_done": print() streamed_text = False diff --git a/pageindex/agent.py b/pageindex/agent.py index 4a1a3d0..253d028 100644 --- a/pageindex/agent.py +++ b/pageindex/agent.py @@ -89,7 +89,7 @@ class QueryStream: Usage: stream = col.query("question", stream=True) async for event in stream: - if event.type == "answer_delta": + if event.type == "text_delta": print(event.data, end="", flush=True) """ @@ -117,7 +117,7 @@ class QueryStream: async for event in streamed_run.stream_events(): if isinstance(event, RawResponsesStreamEvent): if isinstance(event.data, ResponseTextDeltaEvent): - yield QueryEvent(type="answer_delta", data=event.data.delta) + yield QueryEvent(type="text_delta", data=event.data.delta) elif isinstance(event, RunItemStreamEvent): item = event.item if item.type == "tool_call_item": @@ -130,7 +130,7 @@ class QueryStream: elif item.type == "message_output_item": text = ItemHelpers.text_message_output(item) if text: - yield QueryEvent(type="answer_done", data=text) + yield QueryEvent(type="text_done", data=text) def __aiter__(self): return self.stream_events() diff --git a/pageindex/backend/cloud.py b/pageindex/backend/cloud.py index bbfd863..417ec6b 100644 --- a/pageindex/backend/cloud.py +++ b/pageindex/backend/cloud.py @@ -256,6 +256,10 @@ class CloudBackend: ) return meta + def _require_documents(self, collection: str, doc_ids: list[str]) -> None: + for doc_id in doc_ids: + self._require_document(collection, doc_id) + def get_document(self, collection: str, doc_id: str, include_text: bool = False) -> dict: if include_text: import warnings @@ -379,6 +383,8 @@ class CloudBackend: raise ValueError( "doc_ids cannot be empty; pass None to query the whole collection" ) + if doc_ids: + self._require_documents(collection, doc_ids) doc_id = doc_ids if doc_ids else self._get_all_doc_ids(collection) if not doc_id: raise ValueError("collection has no documents to query") @@ -414,6 +420,8 @@ class CloudBackend: raise ValueError( "doc_ids cannot be empty; pass None to query the whole collection" ) + if doc_ids: + await asyncio.to_thread(self._require_documents, collection, doc_ids) doc_id = doc_ids if doc_ids else await asyncio.to_thread( self._get_all_doc_ids, collection ) @@ -516,11 +524,11 @@ class CloudBackend: elif block_type == "text" and content: answer_parts.append(content) - _put(QueryEvent(type="answer_delta", data=content)) + _put(QueryEvent(type="text_delta", data=content)) - # Same terminal contract as the local backend: a final - # answer_done event carrying the full answer text. - _put(QueryEvent(type="answer_done", data="".join(answer_parts))) + # The whole cloud answer is one text message, so its text_done + # carries the full answer text. + _put(QueryEvent(type="text_done", data="".join(answer_parts))) except requests.RequestException as e: _put(CloudAPIError(f"Cloud streaming request failed: {e}")) diff --git a/pageindex/events.py b/pageindex/events.py index fc8f304..13dc6d1 100644 --- a/pageindex/events.py +++ b/pageindex/events.py @@ -5,5 +5,5 @@ from typing import Literal, Any @dataclass class QueryEvent: """Event emitted during streaming query.""" - type: Literal["reasoning", "tool_call", "tool_result", "answer_delta", "answer_done"] + type: Literal["reasoning", "tool_call", "tool_result", "text_delta", "text_done"] data: Any diff --git a/tests/test_cloud_backend.py b/tests/test_cloud_backend.py index fcdb081..d023667 100644 --- a/tests/test_cloud_backend.py +++ b/tests/test_cloud_backend.py @@ -331,9 +331,37 @@ def test_doc_ops_skip_guard_when_folders_unavailable(monkeypatch): assert log == [("DELETE", "/doc/d1/")] +def test_query_rejects_doc_from_another_collection(monkeypatch): + backend, log = _membership_backend(monkeypatch, doc_folder="f-other") + with pytest.raises(DocumentNotFoundError, match="collection 'col'"): + backend.query("col", "q", doc_ids=["d1"]) + # guard must fail before the completion request goes out + assert all(method == "GET" and path.endswith("/metadata/") for method, path in log) + + +def test_query_stream_rejects_doc_from_another_collection(monkeypatch): + backend, log = _membership_backend(monkeypatch, doc_folder="f-other") + + async def _run(): + async for _ in backend.query_stream("col", "q", doc_ids=["d1"]): + pass + + with pytest.raises(DocumentNotFoundError, match="collection 'col'"): + asyncio.run(_run()) + assert all(method == "GET" and path.endswith("/metadata/") for method, path in log) + + +def test_query_checks_membership_then_completes(monkeypatch): + backend, log = _membership_backend(monkeypatch) + backend.query("col", "q", doc_ids=["d1"]) + assert ("GET", "/doc/d1/metadata/") in log + assert ("POST", "/chat/completions/") in log + + # ── query_stream: terminal contract and error propagation ─────────────────── def _collect_events(backend, **kwargs): + backend._folder_id_cache.setdefault("col", None) # skip folder lookup async def _run(): events = [] async for ev in backend.query_stream("col", "q", doc_ids=["d1"], **kwargs): @@ -349,7 +377,7 @@ def _sse(block_type, content): }) -def test_query_stream_emits_answer_done(monkeypatch): +def test_query_stream_emits_text_done(monkeypatch): backend = CloudBackend(api_key="pi-test") lines = [_sse("text", "Hello "), _sse("text", "world"), "data: [DONE]"] monkeypatch.setattr( @@ -357,13 +385,14 @@ def test_query_stream_emits_answer_done(monkeypatch): lambda *a, **k: FakeResponse(status_code=200, lines=lines), ) events = _collect_events(backend) - assert [e.type for e in events] == ["answer_delta", "answer_delta", "answer_done"] - # Same contract as the local backend: answer_done carries the full text. + assert [e.type for e in events] == ["text_delta", "text_delta", "text_done"] + # The whole cloud answer is one text message, so text_done carries the full text. assert events[-1].data == "Hello world" def test_query_stream_http_error_raises(monkeypatch): backend = CloudBackend(api_key="pi-test") + backend._folder_id_cache["col"] = None # skip folder lookup monkeypatch.setattr( cloud_mod.requests, "post", lambda *a, **k: FakeResponse(status_code=401, text="unauthorized"), @@ -379,6 +408,7 @@ def test_query_stream_http_error_raises(monkeypatch): def test_query_stream_connect_failure_raises_instead_of_hanging(monkeypatch): backend = CloudBackend(api_key="pi-test") + backend._folder_id_cache["col"] = None # skip folder lookup def fake_post(*a, **k): raise cloud_mod.requests.ConnectionError("dns failure") @@ -397,6 +427,7 @@ def test_query_uses_long_timeout_and_single_attempt(monkeypatch): """Non-streaming chat completion is non-idempotent and slow: it must get a long timeout and must NOT be retried (each retry re-bills the query).""" backend = CloudBackend(api_key="pi-test") + backend._folder_id_cache["col"] = None # skip folder lookup calls = [] def fake_request(method, url, headers=None, **kwargs): @@ -415,6 +446,7 @@ def test_query_stream_early_break_stops_background_thread(monkeypatch): drain the whole stream in the background.""" import threading backend = CloudBackend(api_key="pi-test") + backend._folder_id_cache["col"] = None # skip folder lookup drained_all = threading.Event() class SlowResponse: diff --git a/tests/test_events.py b/tests/test_events.py index 0046130..c097ce9 100644 --- a/tests/test_events.py +++ b/tests/test_events.py @@ -3,13 +3,13 @@ from pageindex.backend.protocol import AgentTools def test_query_event(): - event = QueryEvent(type="answer_delta", data="hello") - assert event.type == "answer_delta" + event = QueryEvent(type="text_delta", data="hello") + assert event.type == "text_delta" assert event.data == "hello" def test_query_event_types(): - for t in ["reasoning", "tool_call", "tool_result", "answer_delta", "answer_done"]: + for t in ["reasoning", "tool_call", "tool_result", "text_delta", "text_done"]: event = QueryEvent(type=t, data="test") assert event.type == t