scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-10 03:28:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 85

fix(types): remove StreamingChoices from ModelResponse, use ModelResponseStream

Browse source 
ModelResponse.choices was typed as List[Union[Choices, StreamingChoices]] which
caused Pydantic serialization warnings and false linting errors. Now that
ModelResponseStream exists for streaming, narrow ModelResponse.choices to
List[Choices] and migrate all ModelResponse(stream=True) call sites to use
ModelResponseStream() instead.
This commit is contained in:
Chesars 2026-02-20 17:47:42 -03:00
parent 3564b8d83b
commit 0f20976efa
10 changed files with 59 additions and 107 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/litellm_core_utils/streaming_handler.py
View file

@ -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:

4
litellm/llms/bedrock/chat/invoke_handler.py
View file

@ -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

2
litellm/llms/codestral/completion/transformation.py
View file

@ -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 {

9
litellm/proxy/guardrails/guardrail_hooks/presidio.py
View file

@ -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):

94
litellm/types/utils.py
View file

@ -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:

6
litellm/utils.py
View file

@ -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

6
tests/litellm/test_stream_chunk_builder_images.py
View file

@ -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)

6
tests/local_testing/test_stream_chunk_builder.py
View file

@ -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

12
tests/local_testing/test_streaming.py
View file

@ -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)

25
tests/test_litellm/test_model_response_normalization.py
View file

@ -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: