feat: finish sync chat_completion stream

Author: jameszyao

taskingai/client/api/inference_api.py

@@ -14,8 +14,6 @@
 import six
 
 from ..api_client import SyncApiClient
-from ..stream import Stream
-from ..models import INFERENCE_CHAT_COMPLETION_STREAM_CAST_MAP
 
 class InferenceApi(object):
 
@@ -34,16 +32,7 @@ def chat_completion(self, body, stream = False, **kwargs):  # noqa: E501
                  returns the request thread.
         """
         kwargs['_return_http_data_only'] = True
-        cast_map = INFERENCE_CHAT_COMPLETION_STREAM_CAST_MAP
-        response = self.chat_completion_with_http_info(body, stream, **kwargs)
-        if not stream:
-            return response
-        else:
-            return Stream(
-                cast_map=cast_map,
-                response=response,
-                client=self.api_client
-            )
+        return self.chat_completion_with_http_info(body, stream, **kwargs)
 
     def chat_completion_with_http_info(self, body, stream, **kwargs):  # noqa: E501
         """Chat Completion  # noqa: E501

taskingai/client/rest.py

@@ -18,6 +18,7 @@
 # python 2 and python 3 compatibility library
 import httpx
 from httpx import HTTPError
+from .stream import Stream
 
 logger = logging.getLogger(__name__)
 
@@ -71,9 +72,21 @@ def __init__(self, configuration, pools_size=4, maxsize=None):
         if configuration.cert_file and configuration.key_file:
             self.client.cert = (configuration.cert_file, configuration.key_file)
 
-    def request(self, method, url, stream = False, query_params=None, headers=None,
+    def _stream_generator(self, method, url, query_params, headers, request_body, _request_timeout):
+        """Generator function for streaming requests."""
+        with self.client.stream(
+                method, url,
+                params=query_params,
+                headers=headers,
+                content=request_body,
+                timeout=_request_timeout
+        ) as response:
+            for line in response.iter_lines():
+                yield line
+
+    def request(self, method, url, stream=False, query_params=None, headers=None,
                 body=None, post_params=None, _preload_content=True,
-                _request_timeout=None) -> RESTResponse | httpx.Response:
+                _request_timeout=None) -> RESTResponse | Stream:
         """
             Perform asynchronous HTTP requests.
 
@@ -110,14 +123,9 @@ def request(self, method, url, stream = False, query_params=None, headers=None,
 
         try:
             if stream:
-                with self.client.stream(
-                        method, url,
-                        params=query_params,
-                        headers=headers,
-                        content=request_body,
-                        timeout=_request_timeout
-                ) as r:
-                    return r
+                return Stream(stream_generator=self._stream_generator(
+                    method, url, query_params, headers, request_body, _request_timeout
+                ))
             else:
                 r = self.client.request(
                     method, url,
@@ -210,6 +218,7 @@ def PATCH(self, url, stream=False, headers=None, query_params=None, post_params=
                             _request_timeout=_request_timeout,
                             body=body)
 
+
 class RESTAsyncClientObject(object):
 
     def __init__(self, configuration, pools_size=4, maxsize=None):
@@ -297,53 +306,53 @@ async def HEAD(self, url, headers=None, query_params=None, _preload_content=True
                                   query_params=query_params)
 
     async def OPTIONS(self, url, headers=None, query_params=None, post_params=None,
-                body=None, _preload_content=True, _request_timeout=None):
+                      body=None, _preload_content=True, _request_timeout=None):
         return await self.request("OPTIONS", url,
-                            headers=headers,
-                            query_params=query_params,
-                            post_params=post_params,
-                            _preload_content=_preload_content,
-                            _request_timeout=_request_timeout,
-                            body=body)
+                                  headers=headers,
+                                  query_params=query_params,
+                                  post_params=post_params,
+                                  _preload_content=_preload_content,
+                                  _request_timeout=_request_timeout,
+                                  body=body)
 
     async def DELETE(self, url, headers=None, query_params=None, body=None,
-               _preload_content=True, _request_timeout=None):
+                     _preload_content=True, _request_timeout=None):
         return await self.request("DELETE", url,
-                            headers=headers,
-                            query_params=query_params,
-                            _preload_content=_preload_content,
-                            _request_timeout=_request_timeout,
-                            body=body)
+                                  headers=headers,
+                                  query_params=query_params,
+                                  _preload_content=_preload_content,
+                                  _request_timeout=_request_timeout,
+                                  body=body)
 
     async def POST(self, url, headers=None, query_params=None, post_params=None,
-             body=None, _preload_content=True, _request_timeout=None):
+                   body=None, _preload_content=True, _request_timeout=None):
         return await self.request("POST", url,
-                            headers=headers,
-                            query_params=query_params,
-                            post_params=post_params,
-                            _preload_content=_preload_content,
-                            _request_timeout=_request_timeout,
-                            body=body)
+                                  headers=headers,
+                                  query_params=query_params,
+                                  post_params=post_params,
+                                  _preload_content=_preload_content,
+                                  _request_timeout=_request_timeout,
+                                  body=body)
 
     async def PUT(self, url, headers=None, query_params=None, post_params=None,
-            body=None, _preload_content=True, _request_timeout=None):
+                  body=None, _preload_content=True, _request_timeout=None):
         return await self.request("PUT", url,
-                            headers=headers,
-                            query_params=query_params,
-                            post_params=post_params,
-                            _preload_content=_preload_content,
-                            _request_timeout=_request_timeout,
-                            body=body)
+                                  headers=headers,
+                                  query_params=query_params,
+                                  post_params=post_params,
+                                  _preload_content=_preload_content,
+                                  _request_timeout=_request_timeout,
+                                  body=body)
 
     async def PATCH(self, url, headers=None, query_params=None, post_params=None,
-              body=None, _preload_content=True, _request_timeout=None):
+                    body=None, _preload_content=True, _request_timeout=None):
         return await self.request("PATCH", url,
-                            headers=headers,
-                            query_params=query_params,
-                            post_params=post_params,
-                            _preload_content=_preload_content,
-                            _request_timeout=_request_timeout,
-                            body=body)
+                                  headers=headers,
+                                  query_params=query_params,
+                                  post_params=post_params,
+                                  _preload_content=_preload_content,
+                                  _request_timeout=_request_timeout,
+                                  body=body)
 
 
 class ApiException(Exception):

taskingai/client/stream.py

@@ -7,52 +7,28 @@
 import httpx
 
 from .models.entity._base import TaskingaiBaseModel
-from .exceptions import ApiException
-
-from .rest import RESTSyncClientObject, RESTAsyncClientObject
 
+from .exceptions import ApiException
 
 class Stream(object):
     """Provides the core interface to iterate over a synchronous stream response."""
 
-    response: httpx.Response
-
-    def __init__(
-        self,
-        *,
-        cast_map: Dict[str, Type[TaskingaiBaseModel]],
-        response: httpx.Response,
-        client: RESTSyncClientObject,
-    ) -> None:
-        if not isinstance(response, httpx.Response):
-            raise TypeError("response must be an httpx.Response object")
-
-        self.response = response
-        self._cast_map = cast_map
-        self._client = client
+    def __init__(self, stream_generator):
+        self._stream_generator = stream_generator
         self._decoder = SSEDecoder()
         self._iterator = self.__stream__()
 
-    def __next__(self) -> TaskingaiBaseModel:
+    def __next__(self):
         return self._iterator.__next__()
 
-    def __iter__(self) -> Iterator[TaskingaiBaseModel]:
+    def __iter__(self):
         for item in self._iterator:
             yield item
 
-    def _iter_events(self) -> Iterator[ServerSentEvent]:
-        yield from self._decoder.iter(self.response.iter_lines())
-
-    def _cast(self, obj_dict, class_type) -> TaskingaiBaseModel:
-        cast_map = self._cast_map
-        if class_type in cast_map:
-            return cast_map[class_type](**obj_dict)
-        else:
-            raise ValueError(f"No class found for type '{class_type}'")
+    def _iter_events(self):
+        yield from self._decoder.iter(self._stream_generator)
 
-    def __stream__(self) -> Iterator[TaskingaiBaseModel]:
-        print("streaming...")
-        response = self.response
+    def __stream__(self):
         iterator = self._iter_events()
 
         for sse in iterator:
@@ -63,14 +39,11 @@ def __stream__(self) -> Iterator[TaskingaiBaseModel]:
                 data = sse.json()
                 if isinstance(data, Dict) and data.get("error"):
                     raise ApiException(
-                        status=response.status_code,
-                        reason="An error ocurred during streaming",
-                        http_resp=response,
+                        status=400,  # or appropriate status code
+                        reason="An error occurred during streaming",
                     )
 
-                object_type = data.get("object")
-                # todo: raise valid format error
-                yield self._cast(data, object_type)
+                yield data
 
         # Ensure the entire stream is consumed
         for sse in iterator:

taskingai/inference/chat_completion.py

@@ -96,12 +96,10 @@ def chat_completion(
         stream=stream
     )
     if not stream:
-        print("not streaming")
         response: ChatCompletionResponse = api_instance.chat_completion(body=body)
         chat_completion_result: ChatCompletion = ChatCompletion(**response["data"])
         return chat_completion_result
     else:
-        print("streaming")
         response: Stream = api_instance.chat_completion(body=body, stream=True)
         return response