feat(types): add conversation streaming config and builder

- introduce typed stream callbacks and message chunks for conversation search
- add decorator-based StreamConfigBuilder and wire streaming params into search

Author: tharropoulos

src/typesense/types/document.py

@@ -586,6 +586,162 @@ class NLLanguageParameters(typing.TypedDict):
     nl_query_debug: typing.NotRequired[bool]
 
 
+class MessageChunk(typing.TypedDict):
+    """
+    A single chunk from a conversation stream response.
+
+    Attributes:
+      conversation_id (str): ID of the conversation.
+      message (str): Message content for this chunk.
+    """
+
+    conversation_id: str
+    message: str
+
+
+class StreamConfig(typing.Generic[TDoc], typing.TypedDict, total=False):
+    """
+    Configuration for streaming conversation search responses.
+
+    Attributes:
+      on_chunk: Callback invoked for each streamed chunk (conversation_id, message).
+      on_complete: Callback invoked when the stream completes with the full search response.
+      on_error: Callback invoked if an error occurs during streaming.
+    """
+
+    on_chunk: typing.Callable[[MessageChunk], None]
+    on_complete: "OnCompleteCallback[TDoc]"
+    on_error: typing.Callable[[BaseException], None]
+
+
+OnChunkCallback = typing.Callable[[MessageChunk], None]
+
+
+class OnCompleteCallback(typing.Protocol[TDoc]):
+    def __call__(self, response: "SearchResponse[TDoc]") -> None: ...
+
+
+OnErrorCallback = typing.Callable[[BaseException], None]
+
+
+class StreamConfigBuilder(typing.Generic[TDoc]):
+    """
+    Builder for StreamConfig using decorators.
+
+    Example:
+        >>> stream = StreamConfigBuilder()
+        >>>
+        >>> @stream.on_chunk
+        ... def handle_chunk(chunk: MessageChunk) -> None:
+        ...     print(chunk["message"], end="", flush=True)
+        >>>
+        >>> @stream.on_complete
+        ... def handle_complete(response: dict) -> None:
+        ...     print(f"Done! Found {response.get('found', 0)}")
+        >>>
+        >>> response = await client.collections["docs"].documents.search({
+        ...     "q": "query",
+        ...     "query_by": "content",
+        ...     "conversation_stream": True,
+        ...     "stream_config": stream,
+        ... })
+    """
+
+    def __init__(self) -> None:
+        """Initialize an empty StreamConfigBuilder."""
+        self._on_chunk: OnChunkCallback | None = None
+        self._on_complete: OnCompleteCallback[TDoc] | None = None
+        self._on_error: OnErrorCallback | None = None
+
+    def on_chunk(self, func: OnChunkCallback) -> OnChunkCallback:
+        """
+        Decorator to register an on_chunk callback.
+
+        Args:
+            func: Callback invoked for each streamed message chunk.
+
+        Returns:
+            The original function (unmodified).
+        """
+        self._on_chunk = func
+        return func
+
+    def on_complete(self, func: OnCompleteCallback[TDoc]) -> OnCompleteCallback[TDoc]:
+        """
+        Decorator to register an on_complete callback.
+
+        Args:
+            func: Callback invoked when streaming completes with the full response.
+
+        Returns:
+            The original function (unmodified).
+        """
+        self._on_complete = func
+        return func
+
+    def on_error(self, func: OnErrorCallback) -> OnErrorCallback:
+        """
+        Decorator to register an on_error callback.
+
+        Args:
+            func: Callback invoked if an error occurs during streaming.
+
+        Returns:
+            The original function (unmodified).
+        """
+        self._on_error = func
+        return func
+
+    def build(self) -> StreamConfig[TDoc]:
+        """
+        Build the StreamConfig dictionary.
+
+        Returns:
+            A StreamConfig with the registered callbacks.
+        """
+        config: StreamConfig[TDoc] = {}
+        if self._on_chunk is not None:
+            config["on_chunk"] = self._on_chunk
+        if self._on_complete is not None:
+            config["on_complete"] = self._on_complete
+        if self._on_error is not None:
+            config["on_error"] = self._on_error
+        return config
+
+    def get(
+        self,
+        key: typing.Literal["on_chunk", "on_complete", "on_error"],
+        default: typing.Callable[..., None] | None = None,
+    ) -> typing.Callable[..., None] | None:
+        """
+        Get a callback by key (for compatibility with dict-like access).
+
+        Args:
+            key: The callback name ('on_chunk', 'on_complete', or 'on_error').
+            default: Default value if the callback is not set.
+
+        Returns:
+            The callback function or the default value.
+        """
+        return self.build().get(key, default)
+
+
+class ConversationStreamParameters(typing.Generic[TDoc], typing.TypedDict):
+    """
+    Parameters for conversational search streaming.
+
+    Attributes:
+      conversation_stream (bool): When true, the search response is streamed (SSE).
+      stream_config: Callbacks for stream events. Not sent to the API.
+        Can be a StreamConfig dict or a StreamConfigBuilder instance.
+    """
+
+    conversation_stream: typing.NotRequired[bool]
+    stream_config: typing.NotRequired[
+        typing.Union[StreamConfig[TDoc], StreamConfigBuilder[TDoc]]
+    ]
+
+
 class SearchParameters(
     RequiredSearchParameters,
     QueryParameters,
@@ -598,11 +754,13 @@ class SearchParameters(
     TypoToleranceParameters,
     CachingParameters,
     NLLanguageParameters,
+    ConversationStreamParameters[TDoc],
+    typing.Generic[TDoc],
 ):
     """Parameters for searching documents."""
 
 
-class MultiSearchParameters(SearchParameters):
+class MultiSearchParameters(SearchParameters[TDoc], typing.Generic[TDoc]):
     """
     Parameters for performing a [Federated/Multi-Search](https://typesense.org/docs/26.0/api/federated-multi-search.html#federated-multi-search).
 
@@ -856,7 +1014,7 @@ class LLMResponse(typing.TypedDict):
     model: str
 
 
-class ParsedNLQuery(typing.TypedDict):
+class ParsedNLQuery(typing.Generic[TDoc], typing.TypedDict):
     """
     Schema for a parsed natural language query.
 
@@ -868,8 +1026,8 @@ class ParsedNLQuery(typing.TypedDict):
     """
 
     parse_time_ms: int
-    generated_params: SearchParameters
-    augmented_params: SearchParameters
+    generated_params: SearchParameters[TDoc]
+    augmented_params: SearchParameters[TDoc]
     llm_response: typing.NotRequired[LLMResponse]
 
 
@@ -901,7 +1059,7 @@ class SearchResponse(typing.Generic[TDoc], typing.TypedDict):
     hits: typing.List[Hit[TDoc]]
     grouped_hits: typing.NotRequired[typing.List[GroupedHit[TDoc]]]
     conversation: typing.NotRequired[Conversation]
-    parsed_nl_query: typing.NotRequired[ParsedNLQuery]
+    parsed_nl_query: typing.NotRequired[ParsedNLQuery[TDoc]]
 
 
 class DeleteSingleDocumentParameters(typing.TypedDict):