diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 7a6752fbff8..772a7a28947 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1183,7 +1183,7 @@ class CustomStreamWrapper: ], ) _streaming_response = StreamingChoices(delta=_delta_obj) - _model_response = ModelResponse(stream=True) + _model_response = ModelResponseStream() _model_response.choices = [_streaming_response] response_obj = {"original_chunk": _model_response} else: diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 1c58a11eebe..6d8b9c5a163 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -558,7 +558,7 @@ class BedrockLLM(BaseAWSLLM): "INSIDE BEDROCK STREAMING TOOL CALLING CONDITION BLOCK" ) # return an iterator - streaming_model_response = ModelResponse(stream=True) + streaming_model_response = ModelResponseStream() streaming_model_response.choices[0].finish_reason = getattr( model_response.choices[0], "finish_reason", "stop" ) @@ -695,7 +695,7 @@ class BedrockLLM(BaseAWSLLM): ) if stream and provider == "ai21": - streaming_model_response = ModelResponse(stream=True) + streaming_model_response = ModelResponseStream() streaming_model_response.choices[0].finish_reason = model_response.choices[ # type: ignore 0 ].finish_reason diff --git a/litellm/llms/codestral/completion/transformation.py b/litellm/llms/codestral/completion/transformation.py index 646c0e8e56c..31d6652f48a 100644 --- a/litellm/llms/codestral/completion/transformation.py +++ b/litellm/llms/codestral/completion/transformation.py @@ -102,7 +102,7 @@ class CodestralTextCompletionConfig(OpenAITextCompletionConfig): "finish_reason": finish_reason, } - original_chunk = litellm.ModelResponse(**chunk_data_dict, stream=True) + original_chunk = litellm.ModelResponseStream(**chunk_data_dict) _choices = chunk_data_dict.get("choices", []) or [] if len(_choices) == 0: return { diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 29d57ee4734..06e0b38f531 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -808,7 +808,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): return response if isinstance(response, ModelResponse) and not isinstance( - response.choices[0], StreamingChoices + response, ModelResponseStream ): # /chat/completions requests if isinstance(response.choices[0].message.content, str): verbose_proxy_logger.debug( @@ -832,7 +832,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): return response # skip streaming here; handled in async_post_call_streaming_iterator_hook - if response.choices and isinstance(response.choices[0], StreamingChoices): + if isinstance(response, ModelResponseStream): return response presidio_config = self.get_presidio_settings_from_request_data( @@ -840,10 +840,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): ) for choice in response.choices: - # Type narrowing: StreamingChoices doesn't have .message attribute - if not hasattr(choice, "message"): - continue - content = getattr(choice.message, "content", None) # type: ignore + content = getattr(choice.message, "content", None) if content is None: continue if isinstance(content, str): diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 2b64231ca76..7ffd20e8278 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -22,9 +22,8 @@ from openai.types.moderation_create_response import Moderation as Moderation from openai.types.moderation_create_response import ( ModerationCreateResponse as ModerationCreateResponse, ) -from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, SerializationInfo, model_validator -from pydantic.functional_serializers import WrapSerializer -from typing_extensions import Annotated, Required, TypedDict +from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, model_validator +from typing_extensions import Required, TypedDict from litellm._uuid import uuid from litellm.types.llms.base import ( @@ -1641,33 +1640,6 @@ class StreamingChatCompletionChunk(OpenAIChatCompletionChunk): super().__init__(**kwargs) -def _serialize_choices_list( - choices: list, handler, info: SerializationInfo -) -> list: - """Serialize each choice individually to avoid Union serializer warnings. - - Pydantic's Union serializer for ``List[Union[Choices, StreamingChoices]]`` - may try the wrong branch first, emitting spurious - ``PydanticSerializationUnexpectedValue`` warnings. By serializing each - item via its own ``model_dump()`` we bypass the Union dispatch entirely. - """ - kwargs: Dict[str, Any] = {} - if info.exclude_none: - kwargs["exclude_none"] = True - if info.exclude_unset: - kwargs["exclude_unset"] = True - if info.exclude_defaults: - kwargs["exclude_defaults"] = True - result = [] - for choice in choices: - if hasattr(choice, "model_dump"): - result.append(choice.model_dump(**kwargs)) - elif isinstance(choice, dict): - result.append(choice) - else: - result.append(choice) - return result - class ModelResponseBase(OpenAIObject): id: str @@ -1777,10 +1749,7 @@ class ModelResponseStream(ModelResponseBase): class ModelResponse(ModelResponseBase): - choices: Annotated[ - List[Union[Choices, StreamingChoices]], - WrapSerializer(_serialize_choices_list, return_type=list), - ] + choices: List[Choices] """The list of completion choices the model generated for the input prompt.""" def __init__( # noqa: PLR0915 @@ -1799,44 +1768,27 @@ class ModelResponse(ModelResponseBase): _response_headers=None, **params, ) -> None: - if stream is not None and stream is True: - object = "chat.completion.chunk" - if choices is not None and isinstance(choices, list): - new_choices = [] - for choice in choices: - _new_choice = None - if isinstance(choice, StreamingChoices): - _new_choice = choice - elif isinstance(choice, dict): - _new_choice = StreamingChoices(**choice) - elif isinstance(choice, BaseModel): - _new_choice = StreamingChoices(**choice.model_dump()) - new_choices.append(_new_choice) - choices = new_choices - else: - choices = [StreamingChoices()] + object = "chat.completion" + if choices is not None and isinstance(choices, list): + new_choices = [] + for choice in choices: + if isinstance(choice, Choices): + _new_choice = choice # type: ignore + elif isinstance(choice, dict): + _new_choice = Choices(**choice) # type: ignore + elif isinstance(choice, BaseModel): + dump = ( + choice.model_dump() + if hasattr(choice, "model_dump") + else choice.dict() + ) + _new_choice = Choices(**dump) # type: ignore + else: + _new_choice = choice + new_choices.append(_new_choice) + choices = new_choices else: - object = "chat.completion" - if choices is not None and isinstance(choices, list): - new_choices = [] - for choice in choices: - if isinstance(choice, Choices): - _new_choice = choice # type: ignore - elif isinstance(choice, dict): - _new_choice = Choices(**choice) # type: ignore - elif isinstance(choice, BaseModel): - dump = ( - choice.model_dump() - if hasattr(choice, "model_dump") - else choice.dict() - ) - _new_choice = Choices(**dump) # type: ignore - else: - _new_choice = choice - new_choices.append(_new_choice) - choices = new_choices - else: - choices = [Choices()] + choices = [Choices()] if id is None: id = _generate_id() else: diff --git a/litellm/utils.py b/litellm/utils.py index 241b9d217b7..de8fa9c1e18 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7369,9 +7369,9 @@ def _get_base_model_from_metadata(model_call_details=None): class ModelResponseIterator: def __init__(self, model_response: ModelResponse, convert_to_delta: bool = False): if convert_to_delta is True: - self.model_response = ModelResponse(stream=True) - _delta = self.model_response.choices[0].delta # type: ignore - _delta.content = model_response.choices[0].message.content # type: ignore + _stream_response = ModelResponseStream() + _stream_response.choices[0].delta.content = model_response.choices[0].message.content # type: ignore + self.model_response: Union[ModelResponse, ModelResponseStream] = _stream_response else: self.model_response = model_response self.is_done = False diff --git a/tests/litellm/test_stream_chunk_builder_images.py b/tests/litellm/test_stream_chunk_builder_images.py index c51a14ede67..92fb0f93aab 100644 --- a/tests/litellm/test_stream_chunk_builder_images.py +++ b/tests/litellm/test_stream_chunk_builder_images.py @@ -72,7 +72,7 @@ def test_stream_chunk_builder_preserves_images(): chunks = [] for chunk in init_chunks: - chunks.append(litellm.ModelResponse(**chunk, stream=True)) + chunks.append(litellm.ModelResponseStream(**chunk)) response = stream_chunk_builder(chunks=chunks) @@ -163,7 +163,7 @@ def test_stream_chunk_builder_preserves_multiple_images(): chunks = [] for chunk in init_chunks: - chunks.append(litellm.ModelResponse(**chunk, stream=True)) + chunks.append(litellm.ModelResponseStream(**chunk)) response = stream_chunk_builder(chunks=chunks) @@ -230,7 +230,7 @@ def test_stream_chunk_builder_no_images(): chunks = [] for chunk in init_chunks: - chunks.append(litellm.ModelResponse(**chunk, stream=True)) + chunks.append(litellm.ModelResponseStream(**chunk)) response = stream_chunk_builder(chunks=chunks) diff --git a/tests/local_testing/test_stream_chunk_builder.py b/tests/local_testing/test_stream_chunk_builder.py index 8224773aa4c..ddb1546097c 100644 --- a/tests/local_testing/test_stream_chunk_builder.py +++ b/tests/local_testing/test_stream_chunk_builder.py @@ -542,7 +542,7 @@ def test_stream_chunk_builder_multiple_tool_calls(): chunks = [] for chunk in init_chunks: - chunks.append(litellm.ModelResponse(**chunk, stream=True)) + chunks.append(litellm.ModelResponseStream(**chunk)) response = stream_chunk_builder(chunks=chunks) print(f"Returned response: {response}") @@ -616,7 +616,7 @@ def test_stream_chunk_builder_openai_prompt_caching(): chunks: List[litellm.ModelResponse] = [] usage_obj = None for chunk in chat_completion: - chunks.append(litellm.ModelResponse(**chunk.model_dump(), stream=True)) + chunks.append(litellm.ModelResponseStream(**chunk.model_dump())) print(f"chunks: {chunks}") @@ -661,7 +661,7 @@ def test_stream_chunk_builder_openai_audio_output_usage(): chunks = [] for chunk in completion: - chunks.append(litellm.ModelResponse(**chunk.model_dump(), stream=True)) + chunks.append(litellm.ModelResponseStream(**chunk.model_dump())) usage_obj: Optional[litellm.Usage] = None diff --git a/tests/local_testing/test_streaming.py b/tests/local_testing/test_streaming.py index ee208b5e0e2..7f7a43095cd 100644 --- a/tests/local_testing/test_streaming.py +++ b/tests/local_testing/test_streaming.py @@ -393,7 +393,7 @@ def test_completion_azure_stream_content_filter_no_delta(): chunk_list = [] for chunk in chunks: - new_chunk = litellm.ModelResponse(stream=True, id=chunk["id"]) + new_chunk = litellm.ModelResponseStream(id=chunk["id"]) if "choices" in chunk and isinstance(chunk["choices"], list): new_choices = [] for choice in chunk["choices"]: @@ -3026,7 +3026,7 @@ def test_unit_test_custom_stream_wrapper(): {"index": 0, "delta": {"content": "How are you?"}, "finish_reason": "stop"} ], } - chunk = litellm.ModelResponse(**chunk, stream=True) + chunk = litellm.ModelResponseStream(**chunk) completion_stream = ModelResponseIterator(model_response=chunk) @@ -3223,7 +3223,7 @@ def test_unit_test_custom_stream_wrapper_openai(): "system_fingerprint": None, "usage": None, } - chunk = litellm.ModelResponse(**chunk, stream=True) + chunk = litellm.ModelResponseStream(**chunk) completion_stream = ModelResponseIterator(model_response=chunk) @@ -3457,7 +3457,7 @@ def test_aamazing_unit_test_custom_stream_wrapper_n(): chunk_list = [] for chunk in chunks: - new_chunk = litellm.ModelResponse(stream=True, id=chunk["id"]) + new_chunk = litellm.ModelResponseStream(id=chunk["id"]) if "choices" in chunk and isinstance(chunk["choices"], list): print("INSIDE CHUNK CHOICES!") new_choices = [] @@ -3541,7 +3541,7 @@ def test_unit_test_custom_stream_wrapper_function_call(): "system_fingerprint": "fp_44709d6fcb", "choices": [{"index": 0, "delta": delta, "finish_reason": "stop"}], } - chunk = litellm.ModelResponse(**chunk, stream=True) + chunk = litellm.ModelResponseStream(**chunk) completion_stream = ModelResponseIterator(model_response=chunk) @@ -3651,7 +3651,7 @@ def test_unit_test_perplexity_citations_chunk(): } ], } - chunk = litellm.ModelResponse(**chunk, stream=True) + chunk = litellm.ModelResponseStream(**chunk) completion_stream = ModelResponseIterator(model_response=chunk) diff --git a/tests/test_litellm/test_model_response_normalization.py b/tests/test_litellm/test_model_response_normalization.py index 52e6b8bc9e6..85b9fc1450f 100644 --- a/tests/test_litellm/test_model_response_normalization.py +++ b/tests/test_litellm/test_model_response_normalization.py @@ -2,7 +2,14 @@ import warnings import pytest -from litellm.types.utils import Choices, Delta, Message, ModelResponse, StreamingChoices +from litellm.types.utils import ( + Choices, + Delta, + Message, + ModelResponse, + ModelResponseStream, + StreamingChoices, +) def test_modelresponse_normalizes_openai_base_models() -> None: @@ -62,10 +69,8 @@ def test_modelresponse_serialization_avoids_pydantic_warnings() -> None: def test_modelresponse_model_dump_json_no_pydantic_warnings() -> None: - """model_dump_json() bypasses the Python model_dump() override and uses - Pydantic's C-level serializer directly. The Union[Choices, StreamingChoices] - field previously triggered PydanticSerializationUnexpectedValue warnings via - this path.""" + """model_dump_json() and model_dump() should not trigger any Pydantic + serialization warnings now that choices is List[Choices] (no Union).""" response = ModelResponse( model="test-model", choices=[ @@ -94,11 +99,10 @@ def test_modelresponse_model_dump_json_no_pydantic_warnings() -> None: ) -def test_streaming_modelresponse_no_pydantic_warnings() -> None: - """Streaming responses use StreamingChoices in the Union field and should - also serialize without warnings.""" - response = ModelResponse( - model="test-model", +def test_streaming_modelresponsestream_no_pydantic_warnings() -> None: + """Streaming responses use ModelResponseStream with List[StreamingChoices] + and should serialize without warnings.""" + response = ModelResponseStream( choices=[ StreamingChoices( finish_reason="stop", @@ -106,7 +110,6 @@ def test_streaming_modelresponse_no_pydantic_warnings() -> None: delta=Delta(content="hello", role="assistant"), ) ], - stream=True, ) with warnings.catch_warnings(record=True) as captured: