scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-11 08:38:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 280

perf(streaming): reuse the upstream choice for plain litellm HTTP-handler stream chunks (#50786)

Browse source 
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
quietjules 2026-11-10 09:09:18 -08:00 • committed by GitHub
parent 20be20bf45
commit 97b4de224a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 145 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

48
litellm/litellm_core_utils/streaming_handler.py
View file

@ -82,6 +82,8 @@ FUNCTION_CALL_ATTRIBUTE: Final = "function_call"
_SYNC_ITER_EXHAUSTED: Final = object()
_STREAMING_CHOICE_KEYS: Final = frozenset({"finish_reason", "index", "delta", "logprobs", "message"})
_GCHUNK_FIELDS: Final[frozenset] = frozenset(GChunk.__annotations__)
_CompleteStreamingResponse = ModelResponse | TextCompletionResponse | ModelResponseStream | None
_USAGE_COST_HEADER_PROVIDERS: Final[frozenset[str]] = frozenset(
@ -139,11 +141,18 @@ def _entry_len(entries: Mapping[str, object], key: str) -> int:
def _prepend_upstream_choice_extras(
choice: StreamingChoices, upstream_choice: ChunkChoice, extra_keys: frozenset[str]
choice: StreamingChoices, upstream_choice: ChunkChoice | StreamingChoices, extra_keys: frozenset[str]
) -> None:
if not extra_keys:
return
upstream_extras: Final[Mapping[str, object]] = upstream_choice.model_dump(include=set(extra_keys))
stored_extras: Final = {
key: value for key, value in (upstream_choice.__pydantic_extra__ or {}).items() if key in extra_keys
}
upstream_extras: Final[Mapping[str, object]] = (
stored_extras
if _holds_only_scalars(stored_extras.items())
else upstream_choice.model_dump(include=set(extra_keys))
)
choice.__pydantic_extra__ = {**upstream_extras, **(choice.__pydantic_extra__ or {})} # rebind-ok: reuse choice
choice.__pydantic_fields_set__.update(upstream_extras)
@ -439,7 +448,7 @@ class CustomStreamWrapper:
self.completion_stream = completion_stream
self.sent_first_chunk = False
self.sent_last_chunk = False
self._openai_handler_receives_chat_chunks = False
self._openai_handler_receives_chunk_objects = False
self._failure_logging: asyncio.Task[None] | None = None
self._closed = False
self._stream_created_time: float = time.time()
@ -1162,13 +1171,13 @@ class CustomStreamWrapper:
def _prebuilt_plain_delta(self, chunk: object) -> Delta | None:
if (
not self._openai_handler_receives_chat_chunks
or not isinstance(chunk, ChatCompletionChunk)
not self._openai_handler_receives_chunk_objects
or not isinstance(chunk, (ChatCompletionChunk, ModelResponseStream))
or not chunk.choices
):
return None
upstream_delta: Final = cast( # cast-ok: [LIT006] runtime delta may be absent
ChoiceDelta | None,
ChoiceDelta | Delta | None,
chunk.choices[0].delta,
)
if (
@ -1296,22 +1305,23 @@ class CustomStreamWrapper:
upstream_choice: Final = (
original_chunk_for_reuse.choices[0]
if plain_delta_created
and isinstance(original_chunk_for_reuse, ChatCompletionChunk)
and isinstance(original_chunk_for_reuse, (ChatCompletionChunk, ModelResponseStream))
and len(original_chunk_for_reuse.choices) == 1
else None
)
upstream_delta: Final = cast( # cast-ok: [LIT006] runtime delta may be absent
ChoiceDelta | None,
ChoiceDelta | Delta | None,
getattr(upstream_choice, "delta", None),
)
upstream_choice_extra_keys: Final = frozenset(
(upstream_choice.__pydantic_extra__ or {}) if upstream_choice is not None else ()
)
upstream_logprobs: Final[object] = getattr(upstream_choice, "logprobs", None)
can_reuse_choice: Final = (
upstream_choice is not None
and upstream_delta is not None
and "enhancements" not in upstream_choice_extra_keys
and upstream_choice.logprobs is None
and upstream_logprobs is None
and _holds_only_scalars(upstream_delta)
)
if can_reuse_choice and upstream_choice is not None:
@ -1320,7 +1330,7 @@ class CustomStreamWrapper:
_prepend_upstream_choice_extras(
model_response.choices[0],
upstream_choice,
upstream_choice_extra_keys - {"message"},
upstream_choice_extra_keys - _STREAMING_CHOICE_KEYS,
)
else:
choices: Final = []
@ -1737,14 +1747,16 @@ class CustomStreamWrapper:
self._record_finish_choice_fields(cached_chunk)
else: # openai / azure chat model
self.uses_openai_chat_stream = True
if isinstance(chunk, ChatCompletionChunk):
self._openai_handler_receives_chat_chunks = True
if any(str(choice.finish_reason) == "abort" for choice in chunk.choices or ()):
raise litellm.InternalServerError(
message="Upstream generation aborted (finish_reason=abort)",
model=chunk.model,
llm_provider=self.custom_llm_provider,
)
if isinstance(chunk, (ChatCompletionChunk, ModelResponseStream)):
self._openai_handler_receives_chunk_objects = True
if isinstance(chunk, ChatCompletionChunk) and any(
str(choice.finish_reason) == "abort" for choice in chunk.choices or ()
):
raise litellm.InternalServerError(
message="Upstream generation aborted (finish_reason=abort)",
model=chunk.model,
llm_provider=self.custom_llm_provider,
)
if self.custom_llm_provider in [
LlmProviders.AZURE.value,
LlmProviders.AZURE_AI.value,

115
tests/unit/litellm_core_utils/test_streaming_handler.py
View file

@ -7427,6 +7427,121 @@ def test_azure_chunk_creator_streams_choice_content_filter_results_on_later_chun
assert ["content_filter_results" in output.choices[0].model_fields_set for output in outputs] == [True]
def _parsed_openai_like_chunk(
chunk_id: str, delta: Mapping[str, object], finish_reason: str | None = None, **choice_extras: object
) -> ModelResponseStream:
parser: Final = OpenAIChatCompletionStreamingHandler(streaming_response=None, sync_stream=True)
return parser.chunk_parser(
{
"id": chunk_id,
"object": "chat.completion.chunk",
"created": 1,
"model": "served-model",
"choices": [{"index": 0, "delta": dict(delta), "finish_reason": finish_reason, **choice_extras}],
}
)
def _parsed_stream_wrapper(custom_llm_provider: str, logging_obj: Logging) -> CustomStreamWrapper:
wrapper: Final = CustomStreamWrapper(
completion_stream=iter(()),
model="served-model",
custom_llm_provider=custom_llm_provider,
logging_obj=logging_obj,
)
wrapper.chunk_creator(chunk=_parsed_openai_like_chunk("chunk-0", {"role": "assistant", "content": ""}))
return wrapper
@pytest.mark.parametrize("custom_llm_provider", ["hosted_vllm", "deepseek", "groq"])
def test_parsed_openai_like_chunk_creator_builds_one_delta_and_choice_and_no_dump_per_later_chunk(
logging_obj: Logging, custom_llm_provider: str
) -> None:
wrapper: Final = _parsed_stream_wrapper(custom_llm_provider, logging_obj)
later_chunks: Final = (
_parsed_openai_like_chunk("chunk-1", {"content": "hello"}, stop_reason=None),
_parsed_openai_like_chunk("chunk-2", {"reasoning": "thought"}, stop_reason=None),
_parsed_openai_like_chunk("chunk-3", {"content": " world", "reasoning_content": None}, logprobs=None),
)
delta_init: Final = Delta.__init__.__code__
streaming_choices_init: Final = StreamingChoices.__init__.__code__
model_dump: Final = BaseModel.model_dump.__code__
outputs, calls = _chunk_creator_calls(
wrapper, later_chunks, frozenset({delta_init, streaming_choices_init, model_dump})
)
assert sum(call.code is delta_init for call in calls) == len(later_chunks)
assert sum(call.code is streaming_choices_init for call in calls) == len(later_chunks)
assert sum(call.code is model_dump for call in calls) == 0
assert [_wire_dump(output)["choices"] for output in outputs if output is not None] == [
[{"index": 0, "delta": {"content": "hello"}}],
[{"index": 0, "delta": {"reasoning_content": "thought"}}],
[{"index": 0, "delta": {"content": " world"}}],
]
def test_parsed_openai_like_chunk_creator_keeps_choice_extras_ahead_of_choice_keys(logging_obj: Logging) -> None:
wrapper: Final = _parsed_stream_wrapper("hosted_vllm", logging_obj)
later_chunks: Final = (
_parsed_openai_like_chunk("chunk-1", {"content": "hi"}, stop_reason=None, token_ids=[7, 8]),
_parsed_openai_like_chunk("chunk-2", {"content": "!"}, stop_reason=None, matched_stop=2),
)
outputs: Final = tuple(wrapper.chunk_creator(chunk=chunk) for chunk in later_chunks)
choice_dumps: Final = [output.choices[0].model_dump() for output in outputs if output is not None]
assert [list(choice_dump) for choice_dump in choice_dumps] == [
["stop_reason", "token_ids", "finish_reason", "index", "delta", "logprobs"],
["stop_reason", "matched_stop", "finish_reason", "index", "delta", "logprobs"],
]
assert [choice_dump["token_ids"] for choice_dump in choice_dumps[:1]] == [[7, 8]]
assert [choice_dump["matched_stop"] for choice_dump in choice_dumps[1:]] == [2]
assert [choice_dump["delta"]["content"] for choice_dump in choice_dumps] == ["hi", "!"]
def test_parsed_openai_like_stream_still_emits_tool_calls_and_finish_after_reused_chunks(
logging_obj: Logging,
) -> None:
wrapper: Final = _parsed_stream_wrapper("deepseek", logging_obj)
later_chunks: Final = (
_parsed_openai_like_chunk("chunk-1", {"content": "checking"}),
_parsed_openai_like_chunk(
"chunk-2",
{
"tool_calls": [
{"index": 0, "id": "call_1", "type": "function", "function": {"name": "f", "arguments": ""}}
]
},
),
_parsed_openai_like_chunk("chunk-3", {"tool_calls": [{"index": 0, "function": {"arguments": "{}"}}]}),
_parsed_openai_like_chunk("chunk-4", {}, finish_reason="tool_calls"),
)
outputs: Final = tuple(wrapper.chunk_creator(chunk=chunk) for chunk in later_chunks)
assert [_wire_dump(output)["choices"] for output in outputs if output is not None] == [
[{"index": 0, "delta": {"content": "checking"}}],
[
{
"index": 0,
"delta": {
"tool_calls": [
{"id": "call_1", "function": {"arguments": "", "name": "f"}, "type": "function", "index": 0}
]
},
}
],
[
{
"index": 0,
"delta": {"tool_calls": [{"function": {"arguments": "{}"}, "type": "function", "index": 0}]},
}
],
[{"finish_reason": "tool_calls", "index": 0, "delta": {}}],
]
@pytest.mark.parametrize(
"special_delta",
[