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

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": {

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