diff --git a/portkey_ai/api_resources/apis/chat_complete.py b/portkey_ai/api_resources/apis/chat_complete.py index 80b943f..fc87b78 100644 --- a/portkey_ai/api_resources/apis/chat_complete.py +++ b/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, diff --git a/portkey_ai/api_resources/apis/complete.py b/portkey_ai/api_resources/apis/complete.py index f0a1cba..0299c4f 100644 --- a/portkey_ai/api_resources/apis/complete.py +++ b/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,