phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

fix(o1_handler.py): remove fake streaming for o1

Browse source 
Closes https://github.com/BerriAI/litellm/issues/6801
This commit is contained in:
Krrish Dholakia 2024-11-21 19:06:36 +05:30
parent 42beb618ae
commit 252c952a05
1 changed files with 2 additions and 34 deletions
Download patch file Download diff file Expand all files Collapse all files

36
litellm/llms/OpenAI/chat/o1_handler.py
View file

@ -17,22 +17,6 @@ from litellm.utils import CustomStreamWrapper
class OpenAIO1ChatCompletion(OpenAIChatCompletion):
async def mock_async_streaming(
self,
response: Any,
model: Optional[str],
logging_obj: Any,
):
model_response = await response
completion_stream = MockResponseIterator(model_response=model_response)
streaming_response = CustomStreamWrapper(
completion_stream=completion_stream,
model=model,
custom_llm_provider="openai",
logging_obj=logging_obj,
)
return streaming_response
def completion(
self,
model_response: ModelResponse,
@ -54,7 +38,7 @@ class OpenAIO1ChatCompletion(OpenAIChatCompletion):
custom_llm_provider: Optional[str] = None,
drop_params: Optional[bool] = None,
):
stream: Optional[bool] = optional_params.pop("stream", False)
# stream: Optional[bool] = optional_params.pop("stream", False)
response = super().completion(
model_response,
timeout,
@ -76,20 +60,4 @@ class OpenAIO1ChatCompletion(OpenAIChatCompletion):
drop_params,
)
if stream is True:
if asyncio.iscoroutine(response):
return self.mock_async_streaming(
response=response, model=model, logging_obj=logging_obj # type: ignore
)
completion_stream = MockResponseIterator(model_response=response)
streaming_response = CustomStreamWrapper(
completion_stream=completion_stream,
model=model,
custom_llm_provider="openai",
logging_obj=logging_obj,
)
return streaming_response
else:
return response
return response