feat(streaming): generate sync client for streaming responses

Author: tharropoulos

src/typesense/sync/api_call.py

@@ -31,6 +31,7 @@
 by other components of the library.
 """
 
+import json
 import sys
 from types import MappingProxyType, TracebackType
 
@@ -51,6 +52,14 @@
 )
 from typesense.node_manager import NodeManager
 from typesense.request_handler import RequestHandler
+from typesense.stream_handlers import (
+    JSONDict,
+    StreamChunk,
+    combine_stream_chunks,
+    is_message_chunk,
+    parse_sse_line,
+)
+from typesense.types.document import StreamConfig
 
 if sys.version_info >= (3, 11):
     import typing
@@ -168,9 +177,9 @@ def __enter__(self) -> "ApiCall":
 
     def __exit__(
         self,
-        exc_type: typing.Optional[typing.Type[BaseException]],
-        exc_val: typing.Optional[BaseException],
-        exc_tb: typing.Optional[TracebackType],
+        exc_type: typing.Type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
     ) -> None:
         """Async context manager exit."""
         self._client.close()
@@ -186,6 +195,8 @@ def get(
         entity_type: typing.Type[TEntityDict],
         as_json: typing.Literal[False],
         params: typing.Union[TParams, None] = None,
+        stream_config: StreamConfig[TEntityDict] | None = None,
+        is_streaming_request: bool = False,
     ) -> str:
         """
         Execute an async GET request to the Typesense API.
@@ -207,6 +218,8 @@ def get(
         entity_type: typing.Type[TEntityDict],
         as_json: typing.Literal[True] = True,
         params: typing.Union[TParams, None] = None,
+        stream_config: StreamConfig[TEntityDict] | None = None,
+        is_streaming_request: bool = False,
     ) -> TEntityDict:
         """
         Execute an async GET request to the Typesense API.
@@ -227,6 +240,8 @@ def get(
         entity_type: typing.Type[TEntityDict],
         as_json: typing.Union[typing.Literal[True], typing.Literal[False]] = True,
         params: typing.Union[TParams, None] = None,
+        stream_config: StreamConfig[TEntityDict] | None = None,
+        is_streaming_request: bool = False,
     ) -> typing.Union[TEntityDict, str]:
         """
         Execute an async GET request to the Typesense API.
@@ -246,6 +261,8 @@ def get(
             entity_type,
             as_json,
             params=params,
+            stream_config=stream_config,
+            is_streaming_request=is_streaming_request,
         )
 
     @typing.overload
@@ -414,6 +431,8 @@ def _execute_request(
         as_json: typing.Literal[True],
         last_exception: typing.Union[None, Exception] = None,
         num_retries: int = 0,
+        stream_config: StreamConfig[TEntityDict] | None = None,
+        is_streaming_request: bool = False,
         **kwargs: typing.Unpack[SessionFunctionKwargs[TParams, TBody]],
     ) -> TEntityDict:
         """Execute an async request with retry logic."""
@@ -427,6 +446,8 @@ def _execute_request(
         as_json: typing.Literal[False],
         last_exception: typing.Union[None, Exception] = None,
         num_retries: int = 0,
+        stream_config: StreamConfig[TEntityDict] | None = None,
+        is_streaming_request: bool = False,
         **kwargs: typing.Unpack[SessionFunctionKwargs[TParams, TBody]],
     ) -> str:
         """Execute an async request with retry logic."""
@@ -439,6 +460,8 @@ def _execute_request(
         as_json: typing.Union[typing.Literal[True], typing.Literal[False]] = True,
         last_exception: typing.Union[None, Exception] = None,
         num_retries: int = 0,
+        stream_config: StreamConfig[TEntityDict] | None = None,
+        is_streaming_request: bool = False,
         **kwargs: typing.Unpack[SessionFunctionKwargs[TParams, TBody]],
     ) -> typing.Union[TEntityDict, str]:
         """
@@ -470,6 +493,10 @@ def _execute_request(
         node, url, request_kwargs = self._prepare_request_params(endpoint, **kwargs)
 
         try:
+            if is_streaming_request and method == "GET":
+                return self._handle_streaming_get(
+                    url, entity_type, stream_config, **request_kwargs
+                )
             return self._make_request_and_process_response(
                 method,
                 url,
@@ -479,13 +506,22 @@ def _execute_request(
             )
         except _SERVER_ERRORS as server_error:
             self.node_manager.set_node_health(node, is_healthy=False)
+            if is_streaming_request and stream_config:
+                on_error = stream_config.get("on_error")
+                if on_error:
+                    try:
+                        on_error(server_error)
+                    except Exception:
+                        pass
             return self._execute_request(
                 method,
                 endpoint,
                 entity_type,
                 as_json,
                 last_exception=server_error,
                 num_retries=num_retries + 1,
+                stream_config=stream_config,
+                is_streaming_request=is_streaming_request,
                 **kwargs,
             )
 
@@ -516,6 +552,73 @@ def _make_request_and_process_response(
             else typing.cast(str, request_response)
         )
 
+    def _handle_streaming_get(
+        self,
+        url: str,
+        entity_type: typing.Type[TEntityDict],
+        stream_config: StreamConfig[TEntityDict] | None,
+        **kwargs: typing.Unpack[SessionFunctionKwargs[TParams, TBody]],
+    ) -> TEntityDict:
+        """Perform an async streaming GET, parse SSE lines, invoke callbacks, return combined result."""
+        headers: typing.Dict[str, str] = {
+            self.request_handler.api_key_header_name: self.config.api_key,
+            "Accept": "text/event-stream",
+        }
+        headers.update(self.config.additional_headers)
+        extra_headers = kwargs.get("headers")
+        if extra_headers:
+            headers.update(extra_headers)
+
+        params = kwargs.get("params")
+        content: typing.Union[str, bytes, None] = None
+        if body := kwargs.get("data"):
+            if isinstance(body, (str, bytes)):
+                content = body
+            else:
+                content = json.dumps(body)
+
+        all_chunks: typing.List[StreamChunk] = []
+        with self._client.stream(
+            "GET",
+            url,
+            params=params,
+            content=content,
+            headers=headers,
+            timeout=self.config.connection_timeout_seconds,
+        ) as response:
+            if response.status_code < 200 or response.status_code >= 300:
+                response.read()
+                error_message = self.request_handler._get_error_message(response)
+                raise self.request_handler._get_exception(response.status_code)(
+                    response.status_code,
+                    error_message,
+                )
+            for line in response.iter_lines():
+                chunk = parse_sse_line(line)
+                if chunk is not None:
+                    all_chunks.append(chunk)
+                    if stream_config and is_message_chunk(chunk):
+                        on_chunk = stream_config.get("on_chunk")
+                        if on_chunk:
+                            try:
+                                on_chunk(chunk)
+                            except Exception:
+                                pass
+
+        self.node_manager.set_node_health(
+            self.node_manager.get_node(),
+            is_healthy=True,
+        )
+        final: JSONDict = combine_stream_chunks(all_chunks)
+        if stream_config:
+            on_complete = stream_config.get("on_complete")
+            if on_complete:
+                try:
+                    on_complete(typing.cast(TEntityDict, final))
+                except Exception:
+                    pass
+        return typing.cast(TEntityDict, final)
+
     def _prepare_request_params(
         self,
         endpoint: str,

src/typesense/sync/documents.py

@@ -43,6 +43,7 @@
     ImportResponseWithId,
     SearchParameters,
     SearchResponse,
+    StreamConfigBuilder,
     UpdateByFilterParameters,
     UpdateByFilterResponse,
 )
@@ -333,16 +334,25 @@ def search(self, search_parameters: SearchParameters) -> SearchResponse[TDoc]:
 
         Args:
             search_parameters (SearchParameters): The search parameters.
+                Use conversation_stream=True and optionally stream_config (on_chunk,
+                on_complete, on_error) for conversational search streaming.
 
         Returns:
             SearchResponse[TDoc]: The search response containing matching documents.
         """
-        stringified_search_params = stringify_search_params(search_parameters)
+        params_for_api = dict(search_parameters)
+        stream_config = params_for_api.pop("stream_config", None)
+        if isinstance(stream_config, StreamConfigBuilder):
+            stream_config = stream_config.build()
+        conversation_stream = params_for_api.get("conversation_stream") is True
+        stringified_search_params = stringify_search_params(params_for_api)
         response: SearchResponse[TDoc] = self.api_call.get(
             self._endpoint_path("search"),
             params=stringified_search_params,
             entity_type=SearchResponse,
             as_json=True,
+            stream_config=stream_config,
+            is_streaming_request=conversation_stream,
         )
         return response