Merge pull request #411 from Portkey-AI/fix/remove-custom-chat-completions-stream-handling

chore: refactor chat and text completion methods to custom remove str…

Author: VisargD

portkey_ai/api_resources/apis/chat_complete.py

@@ -75,7 +75,7 @@ def stream_create(  # type: ignore[return]
         **kwargs,
     ) -> Union[ChatCompletions, Iterator[ChatCompletionChunk]]:
         extra_headers = kwargs.get("extra_headers", {})
-        with self.openai_client.with_streaming_response.chat.completions.create(
+        return self.openai_client.chat.completions.create(
             model=model,
             messages=messages,
             stream=stream,
@@ -91,22 +91,7 @@ def stream_create(  # type: ignore[return]
             store=store,
             extra_headers=extra_headers,
             extra_body=kwargs,
-        ) as response:
-            for line in response.iter_lines():
-                json_string = line.replace("data: ", "")
-                json_string = json_string.strip().rstrip("\n")
-                if json_string == "":
-                    continue
-                if json_string.startswith(":"):
-                    continue
-                elif json_string == "[DONE]":
-                    break
-                elif json_string != "":
-                    json_data = json.loads(json_string)
-                    json_data = ChatCompletionChunk(**json_data)
-                    yield json_data
-                else:
-                    return ""
+        )
 
     def normal_create(
         self,
@@ -488,7 +473,7 @@ async def stream_create(
         **kwargs,
     ) -> Union[ChatCompletions, AsyncIterator[ChatCompletionChunk]]:
         extra_headers = kwargs.get("extra_headers", {})
-        async with self.openai_client.with_streaming_response.chat.completions.create(
+        return await self.openai_client.chat.completions.create(
             model=model,
             messages=messages,
             stream=stream,
@@ -504,22 +489,7 @@ async def stream_create(
             store=store,
             extra_headers=extra_headers,
             extra_body=kwargs,
-        ) as response:
-            async for line in response.iter_lines():
-                json_string = line.replace("data: ", "")
-                json_string = json_string.strip().rstrip("\n")
-                if json_string == "":
-                    continue
-                if json_string.startswith(":"):
-                    continue
-                elif json_string == "[DONE]":
-                    break
-                elif json_string != "":
-                    json_data = json.loads(json_string)
-                    json_data = ChatCompletionChunk(**json_data)
-                    yield json_data
-                else:
-                    pass
+        )
 
     async def normal_create(
         self,
@@ -579,7 +549,7 @@ async def create(
         **kwargs,
     ) -> Union[ChatCompletions, AsyncIterator[ChatCompletionChunk]]:
         if stream is True:
-            return self.stream_create(
+            return await self.stream_create(
                 model=model,
                 messages=messages,
                 stream=stream,

portkey_ai/api_resources/apis/complete.py

@@ -39,7 +39,7 @@ def stream_create(  # type: ignore[return]
         stream_options,
         **kwargs,
     ) -> Union[TextCompletion, Iterator[TextCompletionChunk]]:
-        with self.openai_client.with_streaming_response.completions.create(
+        return self.openai_client.completions.create(
             model=model,
             prompt=prompt,
             stream=stream,
@@ -59,20 +59,7 @@ def stream_create(  # type: ignore[return]
             user=user,
             stream_options=stream_options,
             extra_body=kwargs,
-        ) as response:
-            for line in response.iter_lines():
-                json_string = line.replace("data: ", "")
-                json_string = json_string.strip().rstrip("\n")
-                if json_string == "":
-                    continue
-                elif json_string == "[DONE]":
-                    break
-                elif json_string != "":
-                    json_data = json.loads(json_string)
-                    json_data = TextCompletionChunk(**json_data)
-                    yield json_data
-                else:
-                    return ""
+        )
 
     def normal_create(
         self,
@@ -219,7 +206,7 @@ async def stream_create(
         stream_options,
         **kwargs,
     ) -> Union[TextCompletion, AsyncIterator[TextCompletionChunk]]:
-        async with self.openai_client.with_streaming_response.completions.create(
+        return await self.openai_client.completions.create(
             model=model,
             prompt=prompt,
             stream=stream,
@@ -239,20 +226,7 @@ async def stream_create(
             user=user,
             stream_options=stream_options,
             extra_body=kwargs,
-        ) as response:
-            async for line in response.iter_lines():
-                json_string = line.replace("data: ", "")
-                json_string = json_string.strip().rstrip("\n")
-                if json_string == "":
-                    continue
-                elif json_string == "[DONE]":
-                    break
-                elif json_string != "":
-                    json_data = json.loads(json_string)
-                    json_data = TextCompletionChunk(**json_data)
-                    yield json_data
-                else:
-                    pass
+        )
 
     async def normal_create(
         self,
@@ -327,7 +301,7 @@ async def create(
         **kwargs,
     ) -> Union[TextCompletion, AsyncIterator[TextCompletionChunk]]:
         if stream is True:
-            return self.stream_create(
+            return await self.stream_create(
                 model=model,
                 prompt=prompt,
                 stream=stream,