ChatStream moves to a light module so chat()'s type hints resolve at runtime

pageindex/__init__.py

@@ -1,14 +1,14 @@
 """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,
                     IndexConfig, LocalIndexConfig)
 
 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):

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()

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.
 

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

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