diff --git a/pageindex/__init__.py b/pageindex/__init__.py index 0843c677e..83c82dc9c 100644 --- a/pageindex/__init__.py +++ b/pageindex/__init__.py @@ -1,10 +1,11 @@ """PageIndex SDK.""" from typing import TYPE_CHECKING as _TYPE_CHECKING +from .chat_stream import ChatStream from .client import PageIndexClient, PageIndexCloudClient, PageIndexLocalClient from .errors import PageIndexAPIError -from .types import (ChatConfig, CloudIndexConfig, IndexConfig, - LocalIndexConfig) +from .types import (ChatConfig, ChatProcessOptions, CloudIndexConfig, + IndexConfig, LocalIndexConfig) if _TYPE_CHECKING: from .flash import page_index_flash @@ -16,6 +17,7 @@ "PageIndexClient", "PageIndexCloudClient", "PageIndexLocalClient", "PageIndexAPIError", "IndexConfig", "CloudIndexConfig", "LocalIndexConfig", "ChatConfig", + "ChatProcessOptions", "ChatStream", "page_index", "page_index_main", "page_index_flash", "optimize_tree", "md_to_tree", ] @@ -25,10 +27,10 @@ "optimize_tree": ".tree_optimize", "md_to_tree": ".page_index_md", } -_SUBMODULES = {"agent_tools", "client", "cloud_api", "errors", "flash", - "integrations", "local_api", "local_chat", "local_store", - "mcp_bridge", "page_index_classic", "page_index_md", - "tree_optimize", "types", "utils"} +_SUBMODULES = {"agent_tools", "chat_stream", "client", "cloud_api", "errors", + "flash", "integrations", "local_api", "local_chat", + "local_store", "mcp_bridge", "page_index_classic", + "page_index_md", "tree_optimize", "types", "utils"} def __getattr__(name): diff --git a/pageindex/chat_stream.py b/pageindex/chat_stream.py new file mode 100644 index 000000000..caad91948 --- /dev/null +++ b/pageindex/chat_stream.py @@ -0,0 +1,68 @@ +"""chat(stream=True)'s return type: one run, one view — text or events.""" +from __future__ import annotations + +from typing import Any, Iterator, Optional + +from .errors import PageIndexAPIError + + +class ChatStream: + """chat(stream=True)'s stream: iterate it for the answer text pieces + (with show_process, the woven display); read ``.events`` instead for + the typed process event dicts. One underlying run — consume exactly + one view; call chat() again for the other.""" + + def __init__(self, text, events): + self._text = text # () -> Iterator[str] + self._events = events # () -> Iterator[dict], or the refusal text + self._view: Optional[str] = None + self._it: Any = None + self._closed = False + + def _claim(self, view: str) -> None: + if self._view is not None and self._view != view: + raise PageIndexAPIError( + f"This chat stream is being consumed as {self._view}; one " + "run serves one view — call chat() again for the other.") + self._view = view + + def __iter__(self) -> "ChatStream": + return self + + def __next__(self) -> str: + self._claim("text") + if self._it is None: + if self._closed: + raise StopIteration + self._it = self._text() + return next(self._it) + + @property + def events(self) -> Iterator[dict]: + """The run as typed event dicts: {"type": "thinking"|"answer", + "delta": ...}, {"type": "tool_call", "call_id", "name", + "arguments"}, {"type": "tool_result", "call_id", "name", + "output"} — full data, never clipped. Consuming — not merely + reading the attribute — claims the view, so debugger panes and + getattr probing stay side-effect free.""" + def consume(): + if isinstance(self._events, str): + raise PageIndexAPIError(self._events) + self._claim("events") + if self._it is None: + if self._closed: + return + self._it = self._events() + # no `yield from`: a dropped handle must not close the run + for ev in self._it: + yield ev + return consume() + + def close(self) -> None: + """Stop the run: closes the open view, and the stream is dead + afterwards, like a closed generator (own-model chat: a run never + consumed never starts).""" + self._closed = True + close = getattr(self._it, "close", None) + if close is not None: + close() diff --git a/pageindex/client.py b/pageindex/client.py index 8d6aec6fd..3a056fab2 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -6,8 +6,10 @@ import threading import time import warnings -from typing import Any, Callable, Iterator, Mapping, Optional, Union, cast +from typing import (Any, Callable, Iterator, Literal, Mapping, Optional, Union, + cast, overload) +from .chat_stream import ChatStream from .errors import PageIndexAPIError @@ -749,6 +751,32 @@ def get_retrieval(self, retrieval_id: str) -> dict[str, Any]: # ---------- CHAT ---------- + # stream picks the return type: the docstring's `.events` usage must + # type-check for py.typed consumers + @overload + def chat( + self, + messages: Union[str, list[dict[str, str]]], + doc_id: Optional[Union[str, list[str]]] = None, + stream: Literal[False] = False, + model: Optional[str] = None, + reasoning_effort: Optional[str] = None, + show_process: Union[bool, Mapping[str, Any], None] = None, + ) -> str: ... + + @overload + def chat( + self, + messages: Union[str, list[dict[str, str]]], + doc_id: Optional[Union[str, list[str]]] = None, + *, + stream: Literal[True], + model: Optional[str] = None, + reasoning_effort: Optional[str] = None, + show_process: Union[bool, Mapping[str, Any], None] = None, + ) -> ChatStream: ... + + @overload def chat( self, messages: Union[str, list[dict[str, str]]], @@ -756,16 +784,29 @@ def chat( stream: bool = False, model: Optional[str] = None, reasoning_effort: Optional[str] = None, - ) -> Union[str, Iterator[str]]: + show_process: Union[bool, Mapping[str, Any], None] = None, + ) -> Union[str, ChatStream]: ... + + def chat( + self, + messages: Union[str, list[dict[str, str]]], + doc_id: Optional[Union[str, list[str]]] = None, + stream: bool = False, + model: Optional[str] = None, + reasoning_effort: Optional[str] = None, + show_process: Union[bool, Mapping[str, Any], None] = None, + ) -> Union[str, ChatStream]: """ Ask a question about your documents, get the answer. - Thin sugar over ``chat_completions()`` in every mode — same - engine, same wire, minus the envelope. Multi-turn: keep your own - role/content list of the visible conversation (append each answer - as an assistant message) and pass it back. For usage accounting, - streaming metadata, or the tool-use process, use the protocol - surfaces: ``chat_completions()``, ``responses()``, ``messages()``. + Thin sugar over the same engine as ``chat_completions()`` in + every mode — same wire, minus the envelope. Multi-turn: keep your + own role/content list of the visible conversation (append each + answer as an assistant message; join a stream into one only with + ``show_process=False``) and pass it back. For usage + accounting, streaming metadata, or the full protocol transcripts, + use the protocol surfaces: ``chat_completions()``, + ``responses()``, ``messages()``. Args: messages: A question string, or role/content conversation @@ -775,23 +816,81 @@ def chat( documents: also enforced at the tool layer, not just prompted. Cloud documents: the managed chat scopes server-side; own-model chat targets at the prompt level. - stream: Yield the answer as text chunks as it is produced. + stream: Return a ``ChatStream``: iterate it for the answer + as text chunks as they are produced (``show_process`` + is on by default, so the run's process arrives woven + in; ``show_process=False`` gives the bare answer), or + read its ``.events`` property instead for the run as + typed event dicts — thinking/answer deltas, each tool + call and its full result (own-model chat only; never + clipped). One run serves one view. model: Own-model chat only — backend model name (defaults to ``chat_model``). reasoning_effort: Own-model chat only — how hard the model thinks (``"low"`` / ``"medium"`` / ``"high"``; what a backend accepts is its own). Unset sends nothing — the model's default behavior applies. + show_process: Streamed chat — weave the run into + the text stream for display: thinking flows as + "[thinking] " sections, each tool call as a "[tool_call] + name arguments" line with its "[tool_result]" line, and + the answer unlabeled. **On by default**, weaving what the mode + serves: the in-process agent's full run; on a managed + client, the tool calls the endpoint streams (its wire + carries no thinking and no tool results). Pass ``False`` + for the bare answer stream — do that before appending a + streamed answer to the conversation history. + ``True`` shows everything; a dict (typed as + ``pageindex.ChatProcessOptions``) selects the parts — + ``thinking`` / ``tool_call`` / ``tool_result``, bools + defaulting on — and sets ``max_chars``, the per-line + summary cap in characters (default 200). Omitted keys + keep their defaults, so ``{"thinking": False}`` hides + only thinking and ``{}`` equals ``True``. + Thinking appears when the backend streams it + (e.g. Claude models with ``reasoning_effort``; OpenAI + models expose none on the chat protocol). The labels are + not a parse format, and a process stream must not be + appended back as conversation history — for the + machine-readable process use ``.events``, ``responses()`` + or ``messages()``. Returns: - stream=False: the answer string - - stream=True: iterator of text chunks + - stream=True: a ``ChatStream`` — iterating it yields text + chunks (with show_process, on by default, the run's process + woven in as labeled sections); ``.events`` yields typed + event dicts: + ``{"type": "thinking"|"answer", "delta": ...}``, + ``{"type": "tool_call", "call_id", "name", "arguments"}``, + ``{"type": "tool_result", "call_id", "name", "output"}`` """ - result = self.chat_completions(messages, stream=stream, - doc_id=doc_id, model=model, - reasoning_effort=reasoning_effort) + if show_process is not False and show_process is not None: + from .local_chat import _process_options + _process_options(show_process) # a bad value chokes first + if not stream: + raise PageIndexAPIError( + "show_process shows the run as it happens and requires " + "stream=True; only show_process=False (or None) means " + f"off — got {show_process!r}.") if stream: - return cast(Iterator[str], result) + # the default means "on where available" + resolved = True if show_process is None else show_process + if self._local_chat: + from .local_chat import run_chat_stream + return run_chat_stream(self, messages, doc_id=doc_id, + model=model, + reasoning_effort=reasoning_effort, + show_process=resolved) + from .local_chat import run_cloud_chat_stream + chunks = self.chat_completions(messages, stream=True, + stream_metadata=True, + doc_id=doc_id, model=model, + reasoning_effort=reasoning_effort) + return run_cloud_chat_stream( + cast(Iterator[dict[str, Any]], chunks), resolved) + result = self.chat_completions(messages, doc_id=doc_id, model=model, + reasoning_effort=reasoning_effort) envelope = cast(dict[str, Any], result) try: return envelope["choices"][0]["message"]["content"] or "" diff --git a/pageindex/cloud_api.py b/pageindex/cloud_api.py index df3c0088c..a676cf3f0 100644 --- a/pageindex/cloud_api.py +++ b/pageindex/cloud_api.py @@ -272,6 +272,10 @@ def _stream_chat_response(self, response: requests.Response) -> Iterator[str]: try: chunk = json.loads(data) + if chunk.get("error"): + raise PageIndexAPIError( + "Chat completion failed mid-stream: " + f"{chunk['error']}") choices = chunk.get("choices") or [{}] content = choices[0].get("delta", {}).get("content", "") if content: @@ -294,6 +298,10 @@ def _stream_chat_response_raw(self, response: requests.Response) -> Iterator[Dic try: chunk = json.loads(data) + if chunk.get("error"): + raise PageIndexAPIError( + "Chat completion failed mid-stream: " + f"{chunk['error']}") yield chunk except json.JSONDecodeError: continue diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index e32932d2b..0d8e27e16 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -1,4 +1,6 @@ -"""Own-model chat: document-QA agents over the local or cloud agent tools.""" +"""Own-model chat: document-QA agents over the local or cloud agent +tools, and the runs behind both ChatStream views, which also weave the +managed endpoint's chunk stream.""" from __future__ import annotations import asyncio @@ -8,9 +10,10 @@ import threading import time import uuid -from typing import Any, Iterator, Optional, Union +from typing import Any, Iterator, Mapping, Optional, Union from .agent_tools import _base_instructions, doc_targeting_block +from .chat_stream import ChatStream from .errors import PageIndexAPIError CHAT_HEADER = ( @@ -581,6 +584,281 @@ def _responses_usage(raw_responses) -> dict: "total_tokens": prompt + completion} +def _chat_agent(client, messages, doc_id, model, temperature=None, + top_p=None, reasoning_effort=None, extra_body=None, + max_tokens=None, backend=None, extra_headers=None, + ) -> "tuple[Any, list, str]": + """The chat lane's shared prologue: validated history, doc targeting, + and the configured agent. Returns (agent, input items, model name).""" + system_texts, history = _split_chat_messages(messages) + scope = client._local_doc_scope(doc_id) + block = _doc_block(client, doc_id, scoped=scope is not None) + items = ([{"role": "user", "content": block}] if block else []) + history + model_name = model or client.chat_model + managed = _managed_instructions(client, system_texts) + agent = _openai_agent(client, "chat", model_name, managed, + temperature, top_p, doc_ids=scope, + cache_key=_conversation_cache_key( + model_name, managed, doc_id, history), + reasoning_effort=reasoning_effort, + extra_body=extra_body, max_tokens=max_tokens, + backend=_merged_backend(client, backend), + extra_headers=extra_headers) + return agent, items, model_name + + +def _clip(text, cap: int = 200) -> str: + """One display line: whitespace flattened, capped for the terminal.""" + flat = " ".join(str(text).split()) + if len(flat) <= cap: + return flat + return f"{flat[:cap]}... (+{len(flat) - cap} chars)" + + +_PROCESS_DEFAULTS = {"thinking": True, "tool_call": True, + "tool_result": True, "max_chars": 200} + + +def _process_options(show_process) -> dict: + """chat(show_process=...) normalized: True (or {}) is all defaults, a + mapping overrides per key, anything else chokes loudly.""" + if show_process is True: + return dict(_PROCESS_DEFAULTS) + if not isinstance(show_process, Mapping): + raise PageIndexAPIError( + "show_process must be True, False, or a dict with the keys " + "thinking / tool_call / tool_result (bools) and max_chars " + "(int).") + unknown = set(show_process) - set(_PROCESS_DEFAULTS) + if unknown: + raise PageIndexAPIError( + "Unknown show_process keys: " + f"{', '.join(sorted(map(repr, unknown)))} " + "— valid keys: thinking, tool_call, tool_result, max_chars.") + options = {**_PROCESS_DEFAULTS, **show_process} + for key in ("thinking", "tool_call", "tool_result"): + if not isinstance(options[key], bool): + raise PageIndexAPIError(f"show_process[{key!r}] must be a bool.") + cap = options["max_chars"] + if not isinstance(cap, int) or isinstance(cap, bool) or cap < 1: + raise PageIndexAPIError( + "show_process['max_chars'] must be a positive int.") + return options + + +async def _chat_events_agen(client, agent, items, run_kwargs): + """The chat stream's primitive: the run as typed event dicts — + thinking/answer deltas, each tool call (arguments parsed when JSON) + and its full result. Both ChatStream views are built on it.""" + import openai + from agents import Runner + from agents.exceptions import AgentsException, MaxTurnsExceeded + from openai.types.responses import ( + ResponseReasoningSummaryTextDeltaEvent, + ResponseReasoningTextDeltaEvent, ResponseTextDeltaEvent) + streamed = Runner.run_streamed(agent, input=items, **run_kwargs) + completed = False + names = {} # call_id -> tool name, to label results + try: + async for event in streamed.stream_events(): + if event.type == "raw_response_event": + data = event.data + if isinstance(data, ResponseTextDeltaEvent): + if data.delta: + yield {"type": "answer", "delta": data.delta} + elif isinstance(data, ( + ResponseReasoningTextDeltaEvent, + ResponseReasoningSummaryTextDeltaEvent)): + if data.delta: + yield {"type": "thinking", "delta": data.delta} + elif event.type == "run_item_stream_event": + raw = getattr(event.item, "raw_item", None) + if event.name == "tool_called": + name = getattr(raw, "name", None) or "tool" + call_id = getattr(raw, "call_id", None) + if call_id: + names[call_id] = name + arguments = getattr(raw, "arguments", "") or "" + try: + arguments = json.loads(arguments) + except (ValueError, TypeError): + pass + yield {"type": "tool_call", "call_id": call_id, + "name": name, "arguments": arguments} + elif event.name == "tool_output": + call_id = (raw.get("call_id") if isinstance(raw, dict) + else getattr(raw, "call_id", None)) + yield {"type": "tool_result", "call_id": call_id, + "name": names.get(call_id, "tool"), + "output": event.item.output} + completed = True + except (MaxTurnsExceeded, AgentsException, openai.OpenAIError) as exc: + raise _translate_run_error(exc, None, "chat", client) from exc + finally: + if not completed and hasattr(streamed, "cancel"): + streamed.cancel() # abandoned/failed: stop the agent task + await _aclose_backend(agent) + + +def _weave(events, options) -> Iterator[str]: + """Render the typed event stream as display text: a "[thinking] " + section per thinking burst, a "[tool_call] name args" line per call + with its "[tool_result]" line, the answer unlabeled — labels are the + event type names. options=None is the plain answer-only view; + closing this generator closes the source.""" + try: + if options is None: + for ev in events: + if ev["type"] == "answer": + yield ev["delta"] + return + section = None # the open flowing section: thinking/answer/tool + opened = False # anything yielded yet (first section takes no gap) + cap = options["max_chars"] + last_call = None # the [tool_call] line still open for nesting + call_args = {} # call_id -> clipped arguments, to label orphans + + def enter(kind, label: str = "") -> str: + nonlocal section, opened + gap = "\n\n" if opened else "" + opened = True + section = kind + return gap + label + + for ev in events: + kind = ev["type"] + if kind == "answer": + head = enter("answer") if section != "answer" else "" + yield head + ev["delta"] + elif kind == "thinking": + if not options["thinking"]: + continue + head = (enter("thinking", "[thinking] ") + if section != "thinking" else "") + yield head + ev["delta"] + elif kind == "tool_call": + arguments = ev["arguments"] + if not isinstance(arguments, str): + arguments = json.dumps(arguments, ensure_ascii=False) + clipped = _clip(arguments, cap) + call_args[ev["call_id"]] = clipped # even with call lines hidden + if not options["tool_call"]: + continue + last_call = ev["call_id"] + line = f"[tool_call] {ev['name']} {clipped}" + yield enter("tool") + line.rstrip() + elif kind == "tool_result": + if not options["tool_result"]: + continue + out = _clip(ev["output"], cap) + if section == "tool" and ev["call_id"] == last_call: + # directly under its own call line + yield f"\n[tool_result] {ev['name']}: {out}" + else: + # parallel calls, or call lines hidden: standalone, + # arguments echoed to say whose result this is + args = call_args.get(ev["call_id"], "") + head = f"[tool_result] {ev['name']} {args}".rstrip() + gap = ("\n" if section == "tool" and last_call is None + else enter("tool")) + yield f"{gap}{head}: {out}" + last_call = None + finally: + close = getattr(events, "close", None) + if close is not None: + close() # cancel the underlying run on abandonment + + +def _cloud_chunk_events(chunks) -> Iterator[dict]: + """Typed events from the managed endpoint's chunk stream: answer + deltas, and each tool call (name + accumulated arguments) from the + block_metadata tags — the endpoint interleaves tool-argument JSON + into delta.content, distinguished only by those tags. It streams no + thinking and no tool results. Outside a tool block, chunks without + block_metadata (an older server) are answer text.""" + tool = None # [name, argument pieces] while inside a tool_use block + try: + for chunk in chunks: + if not isinstance(chunk, dict): + continue + meta = chunk.get("block_metadata") or {} + kind = meta.get("type") + if kind == "mcp_tool_use_start": + tool = [meta.get("tool_name") or "tool", []] + continue + choices = chunk.get("choices") or [] + delta = (choices[0].get("delta") or {}) if choices else {} + content = delta.get("content") + if kind == "tool_use_stop": + if tool is not None: + name, pieces = tool + arguments = "".join(map(str, pieces)) + try: + arguments = json.loads(arguments) + except ValueError: + pass + yield {"type": "tool_call", "call_id": None, + "name": name, "arguments": arguments} + tool = None + continue + if kind == "tool_use" or tool is not None: + # inside an open block nothing is answer text: argument + # chunks accumulate under any tag rather than leaking + if tool is not None and content: + tool[1].append(content) + continue + if content: + yield {"type": "answer", "delta": content} + finally: + close = getattr(chunks, "close", None) + if close is not None: + close() + + +def run_cloud_chat_stream(chunks, + show_process: Union[bool, Mapping[str, Any]] = True, + ) -> ChatStream: + """chat(stream=True) on a managed client: the text view weaves what + the endpoint serves — tool-call lines from its block_metadata tags + (that wire carries no thinking and no tool results); .events needs + the in-process agent.""" + options = (None if show_process is False + else _process_options(show_process)) + return ChatStream( + text=lambda: _weave(_cloud_chunk_events(chunks), options), + events=("chat events are produced by the in-process agent, " + "which the managed chat endpoint does not serve — " + "construct the client with chat_model=... (or a chat= " + "model) to run the agent in your process.")) + + +def run_chat_stream(client, messages, doc_id=None, model=None, + reasoning_effort=None, + show_process: Union[bool, Mapping[str, Any]] = False, + ) -> ChatStream: + """chat(stream=True): validation and the agent build run here, eagerly; + the run itself starts when the returned stream's chosen view is first + consumed.""" + options = (None if show_process is False or show_process is None + else _process_options(show_process)) + _require_openai_agents("chat") + if isinstance(messages, str): + if not messages.strip(): + raise PageIndexAPIError( + "messages must be a non-empty string or a list of " + "message dicts.") + messages = [{"role": "user", "content": messages}] + agent, items, _ = _chat_agent(client, messages, doc_id, model, + reasoning_effort=reasoning_effort) + run_kwargs = _run_kwargs(None) + + def events(): + return _stream_sync( + lambda: _chat_events_agen(client, agent, items, run_kwargs)) + + return ChatStream(text=lambda: _weave(events(), options), events=events) + + def run_chat_completions(client, messages, stream: bool = False, doc_id=None, temperature: Optional[float] = None, stream_metadata: bool = False, @@ -603,21 +881,12 @@ def run_chat_completions(client, messages, stream: bool = False, "citations need.")) _require_openai_agents("chat_completions") _validate_max_turns(max_turns) - system_texts, history = _split_chat_messages(messages) - scope = client._local_doc_scope(doc_id) - block = _doc_block(client, doc_id, scoped=scope is not None) - items = ([{"role": "user", "content": block}] if block else []) + history - model_name = model or client.chat_model + agent, items, model_name = _chat_agent( + client, messages, doc_id, model, temperature=temperature, + top_p=top_p, reasoning_effort=reasoning_effort, + extra_body=extra_body, max_tokens=max_tokens, backend=backend, + extra_headers=extra_headers) reported_model = _reported_model(model_name) - managed = _managed_instructions(client, system_texts) - agent = _openai_agent(client, "chat", model_name, managed, - temperature, top_p, doc_ids=scope, - cache_key=_conversation_cache_key( - model_name, managed, doc_id, history), - reasoning_effort=reasoning_effort, - extra_body=extra_body, max_tokens=max_tokens, - backend=_merged_backend(client, backend), - extra_headers=extra_headers) recorded: dict = {} _record_chat_finish(agent, recorded) run_kwargs = _run_kwargs(max_turns) diff --git a/pageindex/types.py b/pageindex/types.py index b02962899..cd63a8e10 100644 --- a/pageindex/types.py +++ b/pageindex/types.py @@ -46,4 +46,14 @@ class ChatConfig(TypedDict, total=False): backend: dict +class ChatProcessOptions(TypedDict, total=False): + """``chat(show_process=...)``'s display config. Omitted keys default + on (``max_chars``: 200); ``show_process=True`` is all defaults.""" + + thinking: bool + tool_call: bool + tool_result: bool + max_chars: int + + IndexConfig = Union[CloudIndexConfig, LocalIndexConfig] diff --git a/tests/test_client.py b/tests/test_client.py index ec7ed845e..5602f0c4c 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -1503,6 +1503,30 @@ def test_cloud_chat_stream_parsing(cloud, monkeypatch): assert {"object": "chat.completion.citations", "citations": []} in chunks +def test_cloud_chat_stream_error_chunk_raises(cloud, monkeypatch): + """A server-side failure mid-stream arrives as a final {"error": ...} + chunk after the partial answer: every streaming surface raises on it + instead of ending as a short, seemingly complete answer.""" + client, calls, fake = cloud + lines = [ + b'data: {"choices": [{"delta": {"content": "Partial"}}]}', + b'data: {"error": {"message": "boom", "type": "internal_error"}}', + ] + _patch_requests(monkeypatch, lambda m, url, kw: FakeResponse(lines=lines)) + for stream in ( + lambda: client.chat_completions("q", stream=True), + lambda: client.chat_completions("q", stream=True, + stream_metadata=True), + lambda: client.chat("q", stream=True), + ): + it = stream() + first = next(it) # the partial answer is still delivered + assert first in ("Partial", + {"choices": [{"delta": {"content": "Partial"}}]}) + with pytest.raises(PageIndexAPIError, match="boom"): + list(it) + + def test_cloud_chat_accepts_query_string(cloud): client, calls, fake = cloud fake.payload = {"choices": [{"message": {"content": "ok"}}]} diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index 64d09cdd6..0caa1708e 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -129,6 +129,24 @@ async def stream_response(self, system_instructions, input, self._record(system_instructions, input) output = self.turns.pop(0) sequence = 0 + for index, piece in enumerate(getattr(self, "thinking_pieces", ())): + from openai.types.responses import ( + ResponseReasoningSummaryTextDeltaEvent, + ResponseReasoningTextDeltaEvent) + sequence += 1 + # litellm (every chat-protocol backend) delivers thinking as + # summary deltas; alternate so both accepted variants stay + # covered, production shape first + if index % 2: + yield ResponseReasoningTextDeltaEvent( + type="response.reasoning_text.delta", delta=piece, + content_index=0, item_id="rs_1", output_index=0, + sequence_number=sequence) + else: + yield ResponseReasoningSummaryTextDeltaEvent( + type="response.reasoning_summary_text.delta", + delta=piece, item_id="rs_1", output_index=0, + summary_index=0, sequence_number=sequence) if getattr(self, "emit_created", False): from openai.types.responses import ResponseCreatedEvent sequence += 1 @@ -459,6 +477,441 @@ def fake_cc(**kwargs): assert cloud.chat("q") == "cloud answer" +@needs_agents +def test_chat_process_weaves_thinking_and_tools(client, store_path, + fake_model): + """show_process=True keeps the plain text stream but weaves the run in: + a "[thinking] " section per thinking burst, a "[tool_call] name args" + line per call with its clipped result, and the answer unlabeled.""" + seed_doc(store_path, "pi-a", "report.pdf") + fake = fake_model([ + [_call_item("get_document", {"doc_name": "report.pdf"})], + [_msg_item("The answer")], + ]) + fake.thinking_pieces = ("Need the ", "report") + text = "".join(client.chat("What status?", stream=True, show_process=True)) + assert text.startswith("[thinking] Need the report") + assert '\n\n[tool_call] get_document {"doc_name": "report.pdf"}' in text + assert '"report.pdf"}\n[tool_result] get_document: ' in text + assert text.endswith("\n\nThe answer") + assert "[answer]" not in text + + +def test_chat_process_requires_stream(client): + with pytest.raises(PageIndexAPIError, match="requires stream=True"): + client.chat("q", show_process=True) + # falsy-but-not-False means ON (the ruled falsy-{} trap); the message + # must say how to turn it off, not claim the caller passed True + with pytest.raises(PageIndexAPIError, + match="only show_process=False"): + client.chat("q", show_process={}) + # an invalid value is refused as such, with or without stream=True — + # never told to add stream=True first + for kwargs in ({}, {"stream": True}): + with pytest.raises(PageIndexAPIError, + match="must be True, False, or a dict"): + client.chat("q", show_process=0, **kwargs) + + +def _cloud_chunk(content=None, meta=None, choices=True): + chunk = {"id": "chatcmpl-x", "object": "chat.completion.chunk"} + if choices: + delta = {"content": content} if content is not None else {} + chunk["choices"] = [{"index": 0, "delta": delta, + "finish_reason": None}] + if meta: + chunk["block_metadata"] = meta + return chunk + + +def _managed_cloud_with_tool_stream(monkeypatch): + """A cloud client whose managed stream serves the live wire shape: + block_metadata-tagged text, a tool call, and untagged tail chunks.""" + cloud = PageIndexCloudClient(api_key="pi-test-key") + seen = {} + + def fake_cc(**kwargs): + seen.update(kwargs) + return iter([ + _cloud_chunk("", {"type": "text_block_start", "block_index": 1}), + _cloud_chunk("Let me check.", {"type": "text", "block_index": 1}), + _cloud_chunk(None, {"type": "text_stop", "block_index": 1}), + _cloud_chunk(None, {"type": "mcp_tool_use_start", + "block_index": 2, + "tool_name": "get_document", + "server_name": "pageindex"}), + _cloud_chunk('{"doc_name": ', {"type": "tool_use", + "block_index": 2}), + _cloud_chunk('"report.pdf"}', {"type": "tool_use", + "block_index": 2}), + _cloud_chunk(None, {"type": "tool_use_stop", "block_index": 2}), + _cloud_chunk("", {"type": "text_block_start", "block_index": 3}), + _cloud_chunk("The answer", {"type": "text", "block_index": 3}), + _cloud_chunk(None, {"type": "text_stop", "block_index": 3}), + _cloud_chunk(None), # finish_reason chunk + _cloud_chunk(choices=False), # usage chunk + ]) + + monkeypatch.setattr(cloud._api, "chat_completions", fake_cc) + return cloud, seen + + +def test_chat_process_managed_weaves_what_the_wire_serves(monkeypatch): + """Managed streams weave the endpoint's block_metadata: tool-call + lines with name + accumulated arguments; no thinking, no results.""" + cloud, seen = _managed_cloud_with_tool_stream(monkeypatch) + text = "".join(cloud.chat("What status?", stream=True)) + assert seen["stream_metadata"] is True + assert text == ('Let me check.\n\n' + '[tool_call] get_document {"doc_name": "report.pdf"}\n\n' + 'The answer') + cloud, _ = _managed_cloud_with_tool_stream(monkeypatch) + assert "".join(cloud.chat("q", stream=True, show_process=True)) == text + cloud, _ = _managed_cloud_with_tool_stream(monkeypatch) + no_calls = "".join(cloud.chat("q", stream=True, + show_process={"tool_call": False})) + assert no_calls == "Let me check.The answer" + + +def test_chat_process_managed_off_is_clean_answer(monkeypatch): + """show_process=False on managed: answer text only — the tool-call + JSON the wire interleaves into delta.content must not leak in.""" + cloud, _ = _managed_cloud_with_tool_stream(monkeypatch) + plain = "".join(cloud.chat("q", stream=True, show_process=False)) + assert plain == "Let me check.The answer" + assert "doc_name" not in plain + + +def test_chat_process_managed_open_block_never_leaks(monkeypatch): + """Inside an open tool block nothing is answer text: argument chunks + under an unexpected tag (or none) accumulate into the call instead + of leaking into the show_process=False answer.""" + def cloud_with(tag): + cloud = PageIndexCloudClient(api_key="pi-test-key") + monkeypatch.setattr(cloud._api, "chat_completions", lambda **kw: iter([ + _cloud_chunk(None, {"type": "mcp_tool_use_start", + "tool_name": "get_document"}), + _cloud_chunk('{"doc_name": ', {"type": "tool_use"}), + _cloud_chunk('"report.pdf"}', tag), + _cloud_chunk(None, {"type": "tool_use_stop"}), + _cloud_chunk("The answer"), + ])) + return cloud + + for tag in ({"type": "input_json_delta"}, None): + plain = "".join(cloud_with(tag).chat("q", stream=True, + show_process=False)) + assert plain == "The answer" + woven = "".join(cloud_with(tag).chat("q", stream=True)) + assert '[tool_call] get_document {"doc_name": "report.pdf"}' in woven + + +def test_chat_process_managed_non_string_argument_chunk(monkeypatch): + """A non-string delta.content inside a tool block degrades to a + raw-string argument instead of killing the stream.""" + cloud = PageIndexCloudClient(api_key="pi-test-key") + monkeypatch.setattr(cloud._api, "chat_completions", lambda **kw: iter([ + _cloud_chunk(None, {"type": "mcp_tool_use_start", + "tool_name": "get_document"}), + _cloud_chunk({"partial_json": "{"}, {"type": "tool_use"}), + _cloud_chunk(None, {"type": "tool_use_stop"}), + _cloud_chunk("The answer"), + ])) + text = "".join(cloud.chat("q", stream=True)) + assert "[tool_call] get_document" in text + assert text.endswith("The answer") + + +@needs_agents +def test_chat_process_dict_selects_parts(client, store_path, fake_model): + """Each key hides exactly its own line kind; omitted keys default on, + and {} means all defaults, not "off" (the falsy-dict trap).""" + def run(process): + seed_doc(store_path, "pi-a", "report.pdf") + fake = fake_model([ + [_call_item("get_document", {"doc_name": "report.pdf"})], + [_msg_item("The answer")], + ]) + fake.thinking_pieces = ("Need the report",) + return "".join(client.chat("What status?", stream=True, + show_process=process)) + + no_thinking = run({"thinking": False}) + assert "[thinking]" not in no_thinking + assert "[tool_call] get_document" in no_thinking + assert "[tool_result] get_document: " in no_thinking + + no_calls = run({"tool_call": False}) + assert "[tool_call]" not in no_calls + # results stand alone, each echoing its call's arguments + assert '\n\n[tool_result] get_document {"doc_name": "report.pdf"}: ' in no_calls + assert "[thinking] Need the report" in no_calls + + calls_only = run({"tool_result": False}) + assert "[tool_call] get_document" in calls_only + assert "[tool_result]" not in calls_only + + assert run({}) == run(True) + + +@needs_agents +def test_chat_process_max_chars_caps_lines(client, store_path, fake_model): + seed_doc(store_path, "pi-a", "report.pdf") + fake_model([ + [_call_item("get_document", {"doc_name": "report.pdf"})], + [_msg_item("The answer")], + ]) + text = "".join(client.chat("What status?", stream=True, + show_process={"max_chars": 10})) + line = next(ln for ln in text.splitlines() + if ln.startswith("[tool_result] ")) + body = line.split(": ", 1)[1] + assert body[10:].startswith("... (+") + + +@needs_agents +def test_chat_stream_events_typed_sequence(client, store_path, fake_model): + """.events is the typed view: full data, parsed arguments, no clip.""" + seed_doc(store_path, "pi-a", "report.pdf") + fake = fake_model([ + [_call_item("get_document", {"doc_name": "report.pdf"})], + [_msg_item("The answer")], + ]) + fake.thinking_pieces = ("Need the report",) + events = list(client.chat("What status?", stream=True).events) + kinds = [ev["type"] for ev in events] + assert kinds == ["thinking", "tool_call", "tool_result", + "thinking", "answer", "answer"] + call = events[1] + assert call["name"] == "get_document" + assert call["arguments"] == {"doc_name": "report.pdf"} # parsed + assert call["call_id"] == "call_1" + result = events[2] + assert result["name"] == "get_document" + assert result["call_id"] == "call_1" + assert "... (+" not in str(result["output"]) # never clipped + assert '"next_steps"' in result["output"] + assert "".join(ev["delta"] for ev in events + if ev["type"] == "answer") == "The answer" + + +@needs_agents +def test_chat_stream_supports_next_and_one_view(client, store_path, + fake_model): + fake_model([[_msg_item("The answer")]]) + stream = client.chat("q", stream=True) + assert next(stream) == "The " # iterator protocol survives the wrapper + with pytest.raises(PageIndexAPIError, match="one view"): + next(stream.events) + fake_model([[_msg_item("The answer")]]) + stream = client.chat("q", stream=True) + assert next(stream.events)["type"] == "answer" + with pytest.raises(PageIndexAPIError, match="one view"): + next(stream) + + +@needs_agents +def test_chat_stream_events_read_is_inert(client, store_path, fake_model): + """Reading .events claims nothing — only consuming does. Debugger + panes, hasattr and getattr probing must not poison the text view.""" + fake_model([[_msg_item("The answer")]]) + stream = client.chat("q", stream=True) + assert hasattr(stream, "events") # introspection, not consumption + stream.events + assert "".join(stream) == "The answer" + + +@needs_agents +def test_chat_stream_events_survive_partial_reads(client, store_path, + fake_model): + """Peek at one event, then read the rest: the dropped .events handle + must not close the run underneath.""" + seed_doc(store_path, "pi-a", "report.pdf") + fake_model([ + [_call_item("get_document", {"doc_name": "report.pdf"})], + [_msg_item("The answer")], + ]) + stream = client.chat("What status?", stream=True) + first = next(stream.events) + rest = list(stream.events) + assert first["type"] == "tool_call" + assert [ev["type"] for ev in rest] == ["tool_result", "answer", "answer"] + + +def test_chat_stream_events_refusal_waits_for_consumption(monkeypatch): + """On managed, .events read is inert — getattr(stream, 'events', + None) must not explode — and the refusal raises on first + consumption, leaving the text view usable.""" + cloud = PageIndexCloudClient(api_key="pi-test-key") + monkeypatch.setattr(cloud._api, "chat_completions", + lambda **kwargs: iter([_cloud_chunk("x")])) + stream = cloud.chat("q", stream=True) + events = getattr(stream, "events", None) # the standard probing idiom + assert events is not None + with pytest.raises(PageIndexAPIError, match="managed chat endpoint"): + next(events) + assert "".join(stream) == "x" + + +@needs_agents +def test_chat_stream_shows_process_by_default(client, store_path, + fake_model): + """The text view weaves the process unless show_process=False.""" + seed_doc(store_path, "pi-a", "report.pdf") + fake = fake_model([ + [_call_item("get_document", {"doc_name": "report.pdf"})], + [_msg_item("The answer")], + ]) + fake.thinking_pieces = ("Need the report",) + text = "".join(client.chat("What status?", stream=True)) + assert "[thinking] Need the report" in text + assert "[tool_call] get_document" in text + fake = fake_model([ + [_call_item("get_document", {"doc_name": "report.pdf"})], + [_msg_item("The answer")], + ]) + fake.thinking_pieces = ("Need the report",) + plain = "".join(client.chat("What status?", stream=True, + show_process=False)) + assert plain == "The answer" + + +def test_weave_close_closes_source(): + closed = {} + + def source(): + try: + yield {"type": "answer", "delta": "a"} + yield {"type": "answer", "delta": "b"} + finally: + closed["yes"] = True + + woven = local_chat._weave(source(), None) + assert next(woven) == "a" + woven.close() + assert closed.get("yes") is True + + +@needs_agents +def test_chat_stream_close_stops_the_run(client, store_path, fake_model): + fake_model([[_msg_item("The answer")]]) + stream = client.chat("q", stream=True) + assert next(stream) == "The " + stream.close() + with pytest.raises(StopIteration): + next(stream) + + +@needs_agents +def test_chat_stream_closed_before_consumption_stays_dead(client, store_path, + fake_model): + """close() on an unconsumed stream: the run must never start — a + later next() is StopIteration, .events is empty, the model untouched.""" + fake = fake_model([[_msg_item("The answer")]]) + stream = client.chat("q", stream=True) + stream.close() + with pytest.raises(StopIteration): + next(stream) + assert fake.inputs == [] + fake = fake_model([[_msg_item("The answer")]]) + stream = client.chat("q", stream=True) + stream.close() + assert list(stream.events) == [] + assert fake.inputs == [] + + +def test_chat_stream_managed_cloud(monkeypatch): + """Old-wire chunks (no block_metadata) are answer text under any + show_process; .events needs the in-process agent.""" + cloud = PageIndexCloudClient(api_key="pi-test-key") + monkeypatch.setattr( + cloud._api, "chat_completions", + lambda **kwargs: iter([_cloud_chunk("cloud "), + _cloud_chunk("answer")])) + assert list(cloud.chat("q", stream=True)) == ["cloud ", "answer"] + with pytest.raises(PageIndexAPIError, match="managed chat endpoint"): + next(cloud.chat("q", stream=True).events) + + +def test_chat_process_config_chokes(client): + for bad, match in [ + ({"thinkng": False}, "Unknown show_process key"), + ("thinking", "show_process must be True, False, or a dict"), + ({"max_chars": True}, "positive int"), + ({"max_chars": 0}, "positive int"), + ({"thinking": 1}, "must be a bool"), + ({1: True, "foo": 1}, "Unknown show_process key"), + ]: + with pytest.raises(PageIndexAPIError, match=match): + client.chat("q", stream=True, show_process=bad) + + +def test_chat_process_blank_chat_model_refuses(client): + client.chat_model = None + with pytest.raises(PageIndexAPIError, match="chat_model is empty"): + client.chat("q", stream=True, show_process=True) + + +def test_clip_flattens_and_caps(): + assert local_chat._clip("a\n b\tc") == "a b c" + assert local_chat._clip("x" * 250) == "x" * 200 + "... (+50 chars)" + + +@needs_agents +def test_chat_process_parallel_calls_pair_results(client, store_path, + fake_model): + """A result nests only under its own call line; parallel-call results + stand alone with their call's arguments echoed.""" + seed_doc(store_path, "pi-a", "report.pdf") + seed_doc(store_path, "pi-b", "other.pdf") + fake_model([ + [_call_item("get_document", {"doc_name": "report.pdf"}), + _call_item("get_document", {"doc_name": "other.pdf"}, + call_id="call_2")], + [_msg_item("The answer")], + ]) + text = "".join(client.chat("What status?", stream=True, + show_process=True)) + # no result rides directly under a call that is not its own — the + # bare paired form is absent, every result echoes its arguments + assert "\n[tool_result] get_document: " not in text + assert '\n\n[tool_result] get_document {"doc_name": "report.pdf"}: ' in text + assert '\n[tool_result] get_document {"doc_name": "other.pdf"}: ' in text + assert '\n\n[tool_result] get_document {"doc_name": "other.pdf"}' not in text + for doc in ("report.pdf", "other.pdf"): + line = next(ln for ln in text.splitlines() + if ln.startswith(f'[tool_result] get_document {{"doc_name": ' + f'"{doc}"}}: ')) + assert f'"{doc}"' in line.split("}: ", 1)[1] + + +@needs_agents +def test_chat_stream_drops_empty_deltas(client, store_path, fake_model): + """Mid-stream empty deltas carry nothing and would only flip weave + sections; the event source drops them, like chat_completions.""" + fake = fake_model([[_msg_item("The answer")]]) + fake.pieces = ("The ", "", "answer") + assert list(client.chat("q", stream=True, + show_process=False)) == ["The ", "answer"] + fake = fake_model([[_msg_item("The answer")]]) + fake.pieces = ("The ", "", "answer") + fake.thinking_pieces = ("hm", "", "m") + deltas = [ev["delta"] for ev in client.chat("q", stream=True).events + if ev["type"] in ("answer", "thinking")] + assert "" not in deltas + + +def test_chat_process_managed_validates_before_request(monkeypatch): + """A bad show_process must choke before the billed request is sent.""" + cloud = PageIndexCloudClient(api_key="pi-test-key") + calls = [] + monkeypatch.setattr(cloud._api, "chat_completions", + lambda **kwargs: calls.append(kwargs) or iter(())) + with pytest.raises(PageIndexAPIError, match="Unknown show_process key"): + cloud.chat("q", stream=True, show_process={"thinkng": False}) + assert calls == [] + + # ── responses ── @needs_agents @@ -1550,6 +2003,50 @@ def __init__(self, *args, **kwargs): assert fake.deltas_emitted == 0 # turn 2 never produced output +@needs_agents +def test_chat_stream_abandonment_cancels_pending_turn(client, store_path, + fake_model, + monkeypatch): + """chat(stream=True)'s teardown mirrors the completions lane: closing + the stream mid-run cancels the blocked turn (pump thread exits) + instead of letting it run — and bill — in the background, and the + per-call backend client is closed before its loop ends.""" + import threading + seed_doc(store_path, "pi-a", "report.pdf") + pumps = [] + real_thread = threading.Thread + + class _Tracking(real_thread): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + if getattr(kwargs.get("target"), "__name__", "") == "pump": + pumps.append(self) + + monkeypatch.setattr(threading, "Thread", _Tracking) + fake = fake_model([ + [_call_item("get_document", {"doc_name": "report.pdf"})], + [_msg_item("The answer")], + ]) + fake.block_from = 2 # turn 2 hangs until cancelled + closes = [] + + class _Backend: + _pageindex_caller_http = False + + async def close(self): + closes.append(True) + + fake._client = _Backend() + stream = client.chat("q", stream=True) + assert next(stream).startswith("[tool_call] get_document") + stream.close() + assert len(pumps) == 1 + pumps[0].join(timeout=3.0) + assert not pumps[0].is_alive() + assert fake.deltas_emitted == 0 # turn 2 never produced output + assert closes == [True] # _aclose_backend ran on abandonment + + @needs_anthropic def test_messages_max_tokens_default_resolves_per_model(client, fake_anthropic): """The wire-required budget must not exceed the model's ceiling: the diff --git a/tests/test_package_surface.py b/tests/test_package_surface.py index 8f68cea70..8c60fe902 100644 --- a/tests/test_package_surface.py +++ b/tests/test_package_surface.py @@ -53,8 +53,9 @@ def test_import_pageindex_is_lazy(): probe = ( "import sys; import pageindex; " "heavy = [m for m in ('pageindex.page_index_classic', 'pageindex.flash', " - "'pageindex.utils', 'pageindex.tree_optimize', 'numpy', 'PyPDF2') " - "if m in sys.modules]; " + "'pageindex.utils', 'pageindex.tree_optimize', " + "'pageindex.local_chat', 'numpy', 'PyPDF2', " + "'agents', 'litellm', 'openai', 'anthropic') if m in sys.modules]; " "print(','.join(heavy) or 'clean'); " "print(type(pageindex.page_index_main).__name__)" ) @@ -63,6 +64,23 @@ def test_import_pageindex_is_lazy(): assert out.stdout.split() == ["clean", "function"] +def test_public_method_type_hints_resolve_at_runtime(): + """Tools that introspect signatures at runtime (agents' function_tool, + pydantic, doc generators) evaluate the annotations: every public + method's hints must resolve, ChatStream included.""" + import inspect + import typing + import pageindex + from pageindex import ChatStream, PageIndexClient + hints = {name: typing.get_type_hints(fn) for name, fn + in inspect.getmembers(PageIndexClient, inspect.isfunction) + if not name.startswith("_")} + assert len(hints) > 10, f"public-method walk collapsed: {sorted(hints)}" + assert ChatStream in typing.get_args(hints["chat"]["return"]) + assert pageindex.local_chat.ChatStream is ChatStream, ( + "the import path the class shipped under in 0.2.11-0.2.14") + + def test_sdk_submodules_reachable_and_dunder_probes_stay_lazy(): """The 0.2.10 modules resolve as attributes, and underscore probes (the frequent unknown names: copy/pickle/inspect dunders) raise without