diff --git a/pageindex/__init__.py b/pageindex/__init__.py index 63085f240..83c82dc9c 100644 --- a/pageindex/__init__.py +++ b/pageindex/__init__.py @@ -1,6 +1,7 @@ """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, ChatProcessOptions, CloudIndexConfig, @@ -8,7 +9,6 @@ if _TYPE_CHECKING: from .flash import page_index_flash - from .local_chat import ChatStream from .page_index_classic import page_index, page_index_main from .page_index_md import md_to_tree from .tree_optimize import optimize_tree @@ -23,15 +23,14 @@ ] _LAZY = { - "ChatStream": ".local_chat", "page_index_flash": ".flash", "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 f9bb16a90..2de7f4da3 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -9,11 +9,9 @@ from typing import (TYPE_CHECKING, Any, Callable, Iterator, Literal, Mapping, Optional, Union, cast, overload) +from .chat_stream import ChatStream from .errors import PageIndexAPIError -if TYPE_CHECKING: - from .local_chat import ChatStream - _litellm_preload_started = False @@ -826,7 +824,7 @@ def chat( backend: Optional[dict[str, Any]] = None, extra_headers: Optional[dict[str, str]] = None, extra_body: Optional[dict[str, Any]] = None, - ) -> "ChatStream": ... + ) -> ChatStream: ... @overload def chat( @@ -898,7 +896,7 @@ def chat( backend: Optional[dict[str, Any]] = None, extra_headers: Optional[dict[str, str]] = None, extra_body: Optional[dict[str, Any]] = None, - ) -> Union[str, "ChatStream"]: ... + ) -> Union[str, ChatStream]: ... @overload def chat( @@ -916,7 +914,7 @@ def chat( backend: Optional[dict[str, Any]] = None, extra_headers: Optional[dict[str, str]] = None, extra_body: Optional[dict[str, Any]] = None, - ) -> Union[str, "ChatStream", dict[str, Any], Iterator[Any]]: ... + ) -> Union[str, ChatStream, dict[str, Any], Iterator[Any]]: ... def chat( self, @@ -933,7 +931,7 @@ def chat( backend: Optional[dict[str, Any]] = None, extra_headers: Optional[dict[str, str]] = None, extra_body: Optional[dict[str, Any]] = None, - ) -> Union[str, "ChatStream", dict[str, Any], Iterator[Any]]: + ) -> Union[str, ChatStream, dict[str, Any], Iterator[Any]]: """ Ask a question about your documents. diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index af8f56804..a23d99439 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -1,6 +1,6 @@ """Own-model chat: document-QA agents over the local or cloud agent -tools, plus the ChatStream views, which also weave the managed -endpoint's chunk stream.""" +tools, and the runs behind both ChatStream views, which also weave the +managed endpoint's chunk stream.""" from __future__ import annotations import asyncio @@ -13,6 +13,7 @@ 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 = ( @@ -774,68 +775,6 @@ def enter(kind, label: str = "") -> str: close() # cancel the underlying run on abandonment -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() - - 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 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