From 37a2849140df1b44c4c408f3f678263b2c347275 Mon Sep 17 00:00:00 2001 From: Sujit Nirmal Date: Tue, 20 Oct 2026 11:09:22 +0000 Subject: [PATCH] fix(streaming): keep chat logprobs when rebuilding a streamed response stream_chunk_builder merged content, tool calls and reasoning from chat chunks but never read choice.logprobs, so the rebuilt response (and a cached streaming replay built from it) came back with no logprobs even when every chunk carried them. Concatenate the per-chunk token logprobs the same way the text completion path already does --- .../streaming_chunk_builder_utils.py | 30 +++++++++- .../test_streaming_chunk_builder_utils.py | 55 +++++++++++++++++++ 2 files changed, 84 insertions(+), 1 deletion(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 39f675099d7..b48f69a980b 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -7,7 +7,7 @@ from itertools import chain, groupby from types import MappingProxyType from typing import TYPE_CHECKING, Final, Optional, Protocol, TypeAlias, TypedDict, TypeVar, Union, cast -from pydantic import BaseModel, TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from typing_extensions import ReadOnly, Required import litellm @@ -30,6 +30,7 @@ from litellm.types.utils import ( ChatCompletionDeltaToolCall, ChatCompletionMessageCustomToolCall, ChatCompletionMessageToolCall, + ChoiceLogprobs, Choices, CompletionTokensDetails, CompletionTokensDetailsWrapper, @@ -248,6 +249,31 @@ def apply_grounding_request_counts( return counted.model_copy(update=updates) +class _ChoiceLogprobsView(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + logprobs: ChoiceLogprobs | None = None + + +class _ChunkLogprobsView(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + choices: tuple[_ChoiceLogprobsView, ...] = () + + +def _first_choice_logprobs(chunk: object) -> ChoiceLogprobs | None: + try: + choices: Final = _ChunkLogprobsView.model_validate(chunk).choices + except ValidationError: + return None + return choices[0].logprobs if choices else None + + +def _streamed_chat_logprobs(chunks: Sequence[object]) -> ChoiceLogprobs | None: + sent: Final = tuple(logprobs for logprobs in map(_first_choice_logprobs, chunks) if logprobs is not None) + if not sent: + return None + return ChoiceLogprobs(content=_joined_or_none(logprobs.content for logprobs in sent)) + + class ChunkProcessor: def __init__(self, chunks: list, messages: Sequence | None = None): self.chunks = self._sort_chunks(chunks) @@ -416,6 +442,7 @@ class ChunkProcessor: model: Final = ChunkProcessor._get_model_from_chunks(chunks, first_chunk_model) system_fingerprint: Final = chunk.get("system_fingerprint", None) service_tier: Final = ChunkProcessor._get_service_tier_from_chunks(chunks) + logprobs: Final = _streamed_chat_logprobs(chunks) role: Final = ChunkProcessor._get_role_from_chunks(chunks) finish_reason = "stop" @@ -452,6 +479,7 @@ class ChunkProcessor: "index": 0, "message": {"role": role, "content": ""}, "finish_reason": finish_reason, + **(MappingProxyType({"logprobs": logprobs}) if logprobs is not None else MappingProxyType({})), } ], "usage": { diff --git a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py index 5c56245cf33..18f78fb131d 100644 --- a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -2517,3 +2517,58 @@ def test_stream_chunk_builder_keeps_reasoning_items_on_text_only_replies(): assert response is not None assert response.choices[0].message.content == "Hello world" assert response.choices[0].message.reasoning_items == reasoning_items + + +def _chat_chunk( + delta: Delta, logprobs: dict[str, object] | None, finish_reason: str | None = None +) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-lp", + created=1, + model="gpt-4o-mini", + object="chat.completion.chunk", + choices=[StreamingChoices(index=0, delta=delta, logprobs=logprobs, finish_reason=finish_reason)], + ) + + +def _token_logprobs(token: str, logprob: float) -> dict[str, object]: + return { + "content": [ + { + "token": token, + "logprob": logprob, + "bytes": list(token.encode()), + "top_logprobs": [{"token": token, "logprob": logprob, "bytes": list(token.encode())}], + } + ] + } + + +@pytest.mark.parametrize("first_delta", [Delta(role="assistant"), Delta(role="assistant", reasoning_content="hmm")]) +@pytest.mark.parametrize("dumped", [False, True]) +def test_stream_chunk_builder_keeps_logprobs_from_every_chat_chunk(first_delta: Delta, dumped: bool) -> None: + chunks: Final = ( + _chat_chunk(first_delta, None), + _chat_chunk(Delta(content="Hello"), _token_logprobs("Hello", -0.1)), + _chat_chunk(Delta(content=" world"), _token_logprobs(" world", -0.2)), + _chat_chunk(Delta(), None, "stop"), + ) + + response: Final = stream_chunk_builder(chunks=[chunk.model_dump() if dumped else chunk for chunk in chunks]) + + assert response is not None + assert response.choices[0].message.content == "Hello world" + token_logprobs: Final = response.choices[0].logprobs.content + assert [(item.token, item.logprob, item.top_logprobs[0].token) for item in token_logprobs] == [ + ("Hello", -0.1, "Hello"), + (" world", -0.2, " world"), + ] + + +def test_stream_chunk_builder_leaves_chat_logprobs_unset_when_no_chunk_sent_any() -> None: + response: Final = stream_chunk_builder( + chunks=[_chat_chunk(Delta(role="assistant", content="Hi"), None), _chat_chunk(Delta(), None, "stop")] + ) + + assert response is not None + assert response.choices[0].model_dump().get("logprobs") is None