fix(streaming): convert upstream logprobs so streams with logprobs end with their finish chunk (#33549)

* fix(streaming): normalize provider logprobs to prevent MockValSer serialization crash

* fix(streaming): convert upstream chunk logprobs to litellm ChoiceLogprobs so the finish chunk serializes

A streaming chat completion with logprobs on the OpenAI SDK path attached the
SDK's own ChoiceLogprobs object to the outgoing chunk when the chunk had no
content to rebuild from, which is the provider's finish chunk (empty delta,
logprobs for the end token, finish_reason). The SDK parses stream chunks
without building that nested class, so its pydantic serializer is still a
MockValSer and model_dump_json fails. The proxy then sent an error frame in
place of the finish chunk, and when the same chunk also carried usage the SDK
wrapper itself raised MidStreamFallbackError.

The attach site now validates upstream logprobs through a TypeAdapter that
dumps any non-litellm pydantic model before validating into litellm's
ChoiceLogprobs, and passes litellm instances through unchanged.

* fix(streaming): keep upstream logprobs that don't fit ChoiceLogprobs as plain JSON so the stream still finishes

* fix(streaming): build ChoiceLogprobs' serializer up front and keep nonconforming logprobs as JSON on chunk rebuilds

ChoiceLogprobs inherits the openai SDK's deferred build, so a fresh worker
whose first logprobs arrive on the finish chunk still had no serializer.
Foreign logprobs models now dump in full, so a missing top_logprobs becomes
[] and SDK callers keep attribute access, and chunk rebuilds for content and
usage keep logprobs that don't fit as plain JSON instead of raising.

* fix(streaming): keep only the logprobs keys upstream sent and type refusal logprobs

* fix(types): pass nonconforming non-streaming logprobs through instead of failing the response

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>

litellm/litellm_core_utils/streaming_handler.py

@@ -40,6 +40,7 @@ from litellm.llms.ollama.common_utils import OllamaResponseIdNormalizer
 from litellm.types.llms.base import LiteLLMBaseModel
 from litellm.types.router import GenericLiteLLMParams
 from litellm.types.utils import (
+    CHOICE_LOGPROBS_OR_JSON,
     CacheCreationTokenDetails,
     CompletionTokensDetailsWrapper,
     Delta,
@@ -1201,7 +1202,10 @@ class CustomStreamWrapper:
                                         choice_json.pop(
                                             "finish_reason", None
                                         )  # for mistral etc. which return a value in their last chunk (not-openai compatible).
-                                        choices.append(StreamingChoices(**choice_json))
+                                        choice_json.pop("logprobs", None)
+                                        choices.append(
+                                            StreamingChoices(**choice_json, logprobs=getattr(choice, "logprobs", None))
+                                        )
                                 except Exception:
                                     choices.append(StreamingChoices())
                             setattr(model_response, "choices", choices)
@@ -1605,7 +1609,7 @@ class CustomStreamWrapper:
                     model_response.system_fingerprint = response_obj["original_chunk"].system_fingerprint
                     self.system_fingerprint = response_obj["original_chunk"].system_fingerprint
             if response_obj["logprobs"] is not None:
-                model_response.choices[0].logprobs = response_obj["logprobs"]
+                model_response.choices[0].logprobs = CHOICE_LOGPROBS_OR_JSON.validate_python(response_obj["logprobs"])
 
             if response_obj["usage"] is not None:
                 if isinstance(response_obj["usage"], dict):

litellm/types/utils.py

@@ -7,6 +7,7 @@ from enum import Enum
 from types import MappingProxyType
 from typing import (
     TYPE_CHECKING,
+    Annotated,
     Any,
     Final,
     Literal,
@@ -38,6 +39,7 @@ from openai.types.moderation_create_response import (
 )
 from pydantic import (
     BaseModel,
+    BeforeValidator,
     ConfigDict,
     Field,
     FieldSerializationInfo,
@@ -45,6 +47,7 @@ from pydantic import (
     JsonValue,
     PrivateAttr,
     SkipValidation,
+    TypeAdapter,
     ValidationError,
     field_serializer,
     field_validator,
@@ -1199,7 +1202,7 @@ class ChatCompletionTokenLogprob(OpenAIObject):
     unlikely.
     """
 
-    top_logprobs: list[TopLogprob]
+    top_logprobs: list[TopLogprob] = Field(default_factory=list)
     """List of the most likely tokens and their log probability, at this token
     position.
 
@@ -1236,9 +1239,13 @@ class ChatCompletionTokenLogprob(OpenAIObject):
 
 
 class ChoiceLogprobs(OpenAIObject):
+    model_config = ConfigDict(defer_build=False)
+
     content: list[ChatCompletionTokenLogprob] | None = None
     """A list of message content tokens with log probability information."""
 
+    refusal: Sequence[ChatCompletionTokenLogprob] | None = None
+
     def __contains__(self, key) -> bool:
         # Define custom behavior for the 'in' operator
         return hasattr(self, key)
@@ -1252,6 +1259,21 @@ class ChoiceLogprobs(OpenAIObject):
         return getattr(self, key)
 
 
+def _sdk_logprobs_as_dict(logprobs: object) -> object:
+    if isinstance(logprobs, ChoiceLogprobs) or not isinstance(logprobs, BaseModel):
+        return logprobs
+    return logprobs.model_dump(exclude_unset=True)
+
+
+CHOICE_LOGPROBS_OR_JSON: Final[TypeAdapter[ChoiceLogprobs | JsonValue]] = TypeAdapter(
+    Annotated[
+        ChoiceLogprobs | JsonValue,
+        Field(union_mode="left_to_right"),
+        BeforeValidator(_sdk_logprobs_as_dict),
+    ]
+)
+
+
 class FunctionCall(OpenAIObject):
     arguments: str
     name: str | None = None
@@ -1765,7 +1787,7 @@ class Choices(SafeAttributeModel, OpenAIObject):
                 params["message"] = Message(**dump)
         if logprobs is not None:
             if isinstance(logprobs, dict):
-                params["logprobs"] = ChoiceLogprobs(**logprobs)
+                params["logprobs"] = CHOICE_LOGPROBS_OR_JSON.validate_python(logprobs)
             else:
                 params["logprobs"] = logprobs
         else:
@@ -2131,8 +2153,8 @@ class StreamingChoices(OpenAIObject):
         if enhancements is not None:
             self.enhancements = enhancements
 
-        if logprobs is not None and isinstance(logprobs, dict):
-            self.logprobs = ChoiceLogprobs(**logprobs)
+        if isinstance(logprobs, (dict, BaseModel)):
+            self.logprobs = CHOICE_LOGPROBS_OR_JSON.validate_python(logprobs)
         else:
             self.logprobs = logprobs
 

tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py

@@ -1,6 +1,7 @@
 import asyncio
 import importlib
 import json
+from collections.abc import Mapping
 from datetime import datetime, timedelta
 from typing import Final
 
@@ -1111,6 +1112,44 @@ def test_convert_to_model_response_object_with_logprobs():
     assert result.system_fingerprint is None
     assert result._hidden_params is not None
 
+
+@pytest.mark.parametrize(
+    "upstream_logprobs",
+    [
+        pytest.param({"content": [{"token": "Hi", "logprob": None}]}, id="null-content-logprob"),
+        pytest.param(
+            {
+                "content": [{"token": "Hi", "logprob": -0.1, "top_logprobs": []}],
+                "refusal": [{"token": "I", "logprob": None}],
+            },
+            id="null-refusal-logprob",
+        ),
+    ],
+)
+def test_convert_to_model_response_object_passes_nonconforming_logprobs_through(
+    upstream_logprobs: Mapping[str, object],
+) -> None:
+    result: Final = convert_to_model_response_object(
+        response_object={
+            "id": "chatcmpl-123",
+            "object": "chat.completion",
+            "created": 1702685778,
+            "model": "test-model",
+            "choices": [
+                {
+                    "index": 0,
+                    "message": {"role": "assistant", "content": "Hi"},
+                    "logprobs": upstream_logprobs,
+                    "finish_reason": "stop",
+                }
+            ],
+        },
+        model_response_object=ModelResponse(),
+    )
+
+    assert json.loads(result.model_dump_json())["choices"][0]["logprobs"] == upstream_logprobs
+
+
 def test_convert_to_model_response_object_error():
     """Test error handling for None response object."""
     with pytest.raises(Exception, match="Error in response object format"):

tests/unit/litellm_core_utils/test_streaming_handler.py

@@ -2,12 +2,15 @@ import asyncio
 import datetime
 import itertools
 import json
+import os
 import queue
+import subprocess
 import sys
 import threading
 import time
 from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
+from pathlib import Path
 from types import CodeType, FrameType, MappingProxyType
 from typing import AsyncIterator, Callable, Final, Iterator, Literal, NoReturn, Optional, cast
 from unittest.mock import AsyncMock, MagicMock, Mock, patch
@@ -33,6 +36,7 @@ from litellm.llms.custom_httpx.http_handler import HTTPHandler
 from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler
 from litellm.responses.mcp.chat_completions_handler import _add_mcp_metadata_to_response
 from litellm.types.utils import (
+    ChoiceLogprobs,
     CompletionTokensDetailsWrapper,
     Delta,
     ModelResponse,
@@ -1624,6 +1628,356 @@ async def test_stream_with_null_choices_chunk_finishes_normally(
     assert [choice.finish_reason for choice in streamed_choices if choice.finish_reason] == ["stop"]
 
 
+def _token_logprobs(token: str) -> dict[str, object]:
+    return {
+        "content": [{"token": token, "logprob": -0.5, "bytes": list(token.encode()), "top_logprobs": []}],
+        "refusal": None,
+    }
+
+
+@pytest.mark.parametrize("async_stream", [False, True])
+@pytest.mark.parametrize("transport", ["sdk", "http_handler"])
+@pytest.mark.parametrize(
+    ("finish_chunk_extra", "stream_options", "expected_usage_totals"),
+    [
+        pytest.param({}, None, [], id="finish-chunk-with-logprobs"),
+        pytest.param(
+            {"usage": {"prompt_tokens": 3, "completion_tokens": 3, "total_tokens": 6}},
+            {"include_usage": True},
+            [6],
+            id="finish-chunk-with-logprobs-and-usage",
+        ),
+    ],
+)
+async def test_stream_finish_chunk_with_logprobs_reaches_caller_as_litellm_logprobs(
+    async_stream: bool,
+    transport: str,
+    finish_chunk_extra: dict[str, object],
+    stream_options: dict[str, bool] | None,
+    expected_usage_totals: list[int],
+    logging_obj: Logging,
+) -> None:
+    upstream_tokens: Final = ("Hel", "lo", "<|return|>")
+    upstream_chunks: Final = (
+        {
+            "choices": [
+                {
+                    "index": 0,
+                    "delta": {"role": "assistant", "content": "Hel"},
+                    "logprobs": _token_logprobs("Hel"),
+                    "finish_reason": None,
+                }
+            ]
+        },
+        {
+            "choices": [
+                {"index": 0, "delta": {"content": "lo"}, "logprobs": _token_logprobs("lo"), "finish_reason": None}
+            ]
+        },
+        {
+            "choices": [{"index": 0, "delta": {}, "logprobs": _token_logprobs("<|return|>"), "finish_reason": "stop"}],
+            **finish_chunk_extra,
+        },
+    )
+    body: Final = (
+        "".join(
+            "data: "
+            + json.dumps(
+                {
+                    "id": "chatcmpl-logprobs",
+                    "object": "chat.completion.chunk",
+                    "created": 1,
+                    "model": "test-model",
+                    **chunk,
+                }
+            )
+            + "\n\n"
+            for chunk in upstream_chunks
+        )
+        + "data: [DONE]\n\n"
+    )
+    raw_response: Final = httpx.Response(
+        200,
+        content=body,
+        headers={"content-type": "text/event-stream"},
+        request=httpx.Request("POST", "https://provider.example/v1/chat/completions"),
+    )
+    client: Final = AsyncOpenAI(api_key="unused") if async_stream else OpenAI(api_key="unused")
+    stream: Final = (
+        (
+            AsyncStream(cast_to=ChatCompletionChunk, response=raw_response, client=client)
+            if isinstance(client, AsyncOpenAI)
+            else Stream(cast_to=ChatCompletionChunk, response=raw_response, client=client)
+        )
+        if transport == "sdk"
+        else OpenAIChatCompletionStreamingHandler(
+            streaming_response=raw_response.aiter_lines() if async_stream else raw_response.iter_lines(),
+            sync_stream=not async_stream,
+        )
+    )
+    response: Final = CustomStreamWrapper(
+        completion_stream=stream,
+        model="test-model",
+        logging_obj=logging_obj,
+        custom_llm_provider="openai" if transport == "sdk" else "hosted_vllm",
+        stream_options=stream_options,
+    )
+
+    try:
+        chunks: Final = [chunk async for chunk in response] if async_stream else list(response)
+    finally:
+        if isinstance(client, AsyncOpenAI):
+            await client.close()
+        else:
+            client.close()
+
+    wire_choices: Final = tuple(
+        itertools.chain.from_iterable(
+            json.loads(chunk.model_dump_json(exclude_none=True, exclude_unset=True))["choices"] for chunk in chunks
+        )
+    )
+    finish_choices: Final = [choice for choice in wire_choices if choice.get("finish_reason")]
+    streamed_logprobs: Final = tuple(
+        choice.logprobs
+        for choice in itertools.chain.from_iterable(chunk.choices for chunk in chunks)
+        if choice.logprobs is not None
+    )
+    streamed_tokens: Final = tuple(
+        entry.token for entry in itertools.chain.from_iterable(logprobs.content or () for logprobs in streamed_logprobs)
+    )
+
+    assert [choice["finish_reason"] for choice in finish_choices] == ["stop"]
+    assert [entry["token"] for entry in finish_choices[0]["logprobs"]["content"]] == [upstream_tokens[-1]]
+    assert all(isinstance(logprobs, ChoiceLogprobs) for logprobs in streamed_logprobs)
+    assert streamed_tokens == upstream_tokens
+    assert [chunk.usage.total_tokens for chunk in chunks if getattr(chunk, "usage", None)] == expected_usage_totals
+
+
+async def _sdk_stream_chunks(
+    upstream_chunks: Sequence[Mapping[str, object]],
+    async_stream: bool,
+    logging_obj: Logging,
+    stream_options: dict[str, bool] | None = None,
+) -> list[ModelResponseStream]:
+    body: Final = (
+        "".join(
+            "data: "
+            + json.dumps(
+                {
+                    "id": "chatcmpl-logprobs",
+                    "object": "chat.completion.chunk",
+                    "created": 1,
+                    "model": "test-model",
+                    **chunk,
+                }
+            )
+            + "\n\n"
+            for chunk in upstream_chunks
+        )
+        + "data: [DONE]\n\n"
+    )
+    raw_response: Final = httpx.Response(
+        200,
+        content=body,
+        headers={"content-type": "text/event-stream"},
+        request=httpx.Request("POST", "https://provider.example/v1/chat/completions"),
+    )
+    client: Final = AsyncOpenAI(api_key="unused") if async_stream else OpenAI(api_key="unused")
+    stream: Final = (
+        AsyncStream(cast_to=ChatCompletionChunk, response=raw_response, client=client)
+        if isinstance(client, AsyncOpenAI)
+        else Stream(cast_to=ChatCompletionChunk, response=raw_response, client=client)
+    )
+    response: Final = CustomStreamWrapper(
+        completion_stream=stream,
+        model="test-model",
+        logging_obj=logging_obj,
+        custom_llm_provider="openai",
+        stream_options=stream_options,
+    )
+    try:
+        return [chunk async for chunk in response] if async_stream else list(response)
+    finally:
+        if isinstance(client, AsyncOpenAI):
+            await client.close()
+        else:
+            client.close()
+
+
+@pytest.mark.parametrize("async_stream", [False, True])
+@pytest.mark.parametrize(
+    ("upstream_logprobs", "finish_chunk_extra", "stream_options", "expected_usage_totals"),
+    [
+        pytest.param(
+            {"content": [{"token": "<|return|>", "bytes": None, "logprob": None, "top_logprobs": None}], "refusal": None},
+            {},
+            None,
+            [],
+            id="null-logprob",
+        ),
+        pytest.param(
+            {"content": [{"token": "<|return|>", "logprob": None}]},
+            {},
+            None,
+            [],
+            id="null-logprob-keeps-only-the-keys-upstream-sent",
+        ),
+        pytest.param(-0.25, {}, None, [], id="scalar-logprobs"),
+        pytest.param(
+            {"content": [{"token": "<|return|>", "bytes": None, "logprob": None, "top_logprobs": None}], "refusal": None},
+            {"usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}},
+            {"include_usage": True},
+            [5],
+            id="null-logprob-and-usage",
+        ),
+    ],
+)
+async def test_stream_with_nonstandard_logprobs_passes_them_through_and_finishes(
+    async_stream: bool,
+    upstream_logprobs: object,
+    finish_chunk_extra: dict[str, object],
+    stream_options: dict[str, bool] | None,
+    expected_usage_totals: list[int],
+    logging_obj: Logging,
+) -> None:
+    chunks: Final = await _sdk_stream_chunks(
+        (
+            {
+                "choices": [
+                    {
+                        "index": 0,
+                        "delta": {"role": "assistant", "content": "hi"},
+                        "logprobs": upstream_logprobs,
+                        "finish_reason": None,
+                    }
+                ]
+            },
+            {
+                "choices": [{"index": 0, "delta": {}, "logprobs": upstream_logprobs, "finish_reason": "stop"}],
+                **finish_chunk_extra,
+            },
+        ),
+        async_stream,
+        logging_obj,
+        stream_options,
+    )
+
+    wire_choices: Final = tuple(
+        itertools.chain.from_iterable(
+            json.loads(chunk.model_dump_json(exclude_none=True, exclude_unset=True))["choices"] for chunk in chunks
+        )
+    )
+    logprobs_choices: Final = [choice for choice in wire_choices if "logprobs" in choice]
+
+    assert [choice.get("finish_reason") for choice in logprobs_choices] == [None, "stop"]
+    assert [choice["logprobs"] for choice in logprobs_choices] == [upstream_logprobs, upstream_logprobs]
+    assert [chunk.usage.total_tokens for chunk in chunks if getattr(chunk, "usage", None)] == expected_usage_totals
+
+
+@pytest.mark.parametrize("async_stream", [False, True])
+@pytest.mark.parametrize("logprobs_field", ["content", "refusal"])
+async def test_stream_logprobs_without_top_logprobs_reach_caller_as_litellm_logprobs(
+    async_stream: bool, logprobs_field: Literal["content", "refusal"], logging_obj: Logging
+) -> None:
+    chunks: Final = await _sdk_stream_chunks(
+        tuple(
+            {
+                "choices": [
+                    {
+                        "index": 0,
+                        "delta": delta,
+                        "logprobs": {logprobs_field: [{"token": token, "logprob": -0.5}]},
+                        "finish_reason": finish_reason,
+                    }
+                ]
+            }
+            for token, delta, finish_reason in (
+                ("hi", {"role": "assistant", "content": "hi"}, None),
+                ("<|return|>", {}, "stop"),
+            )
+        ),
+        async_stream,
+        logging_obj,
+    )
+
+    streamed_logprobs: Final = tuple(
+        choice.logprobs
+        for choice in itertools.chain.from_iterable(chunk.choices for chunk in chunks)
+        if choice.logprobs is not None
+    )
+    streamed_entries: Final = tuple(
+        itertools.chain.from_iterable(
+            (logprobs.content if logprobs_field == "content" else logprobs.refusal) or ()
+            for logprobs in streamed_logprobs
+        )
+    )
+
+    assert [(entry.token, entry.logprob, entry.top_logprobs) for entry in streamed_entries] == [
+        ("hi", -0.5, []),
+        ("<|return|>", -0.5, []),
+    ]
+
+
+FINISH_ONLY_LOGPROBS_STREAM_SCRIPT: Final = """
+import json
+import sys
+
+sys.path.insert(0, sys.argv[1])
+
+import httpx
+from openai import OpenAI
+
+import litellm
+
+upstream_choices = (
+    {"index": 0, "delta": {"role": "assistant", "content": "hi"}, "finish_reason": None},
+    {
+        "index": 0,
+        "delta": {},
+        "logprobs": {"content": [{"token": "<|return|>", "logprob": -0.5, "top_logprobs": []}]},
+        "finish_reason": "stop",
+    },
+)
+body = "".join(
+    "data: "
+    + json.dumps({"id": "c", "object": "chat.completion.chunk", "created": 1, "model": "m", "choices": [choice]})
+    + "\\n\\n"
+    for choice in upstream_choices
+) + "data: [DONE]\\n\\n"
+client = OpenAI(
+    api_key="unused",
+    base_url="https://provider.example/v1",
+    http_client=httpx.Client(
+        transport=httpx.MockTransport(
+            lambda request: httpx.Response(200, content=body, headers={"content-type": "text/event-stream"})
+        )
+    ),
+)
+stream = litellm.completion(
+    model="openai/m", messages=[{"role": "user", "content": "hi"}], stream=True, logprobs=True, client=client
+)
+print(json.dumps([json.loads(chunk.model_dump_json(exclude_none=True)) for chunk in stream]))
+"""
+
+
+def test_stream_with_logprobs_only_on_finish_chunk_serializes_in_a_fresh_process() -> None:
+    result: Final = subprocess.run(
+        [sys.executable, "-I", "-c", FINISH_ONLY_LOGPROBS_STREAM_SCRIPT, str(Path(litellm.__file__).parent.parent)],
+        capture_output=True,
+        text=True,
+        timeout=120,
+        env={**os.environ, "LITELLM_LOCAL_MODEL_COST_MAP": "True"},
+    )
+
+    assert result.returncode == 0, result.stdout + result.stderr
+    wire_choices: Final = tuple(
+        itertools.chain.from_iterable(chunk["choices"] for chunk in json.loads(result.stdout.strip().splitlines()[-1]))
+    )
+    finish_choices: Final = [choice for choice in wire_choices if choice.get("finish_reason")]
+    assert [choice["finish_reason"] for choice in finish_choices] == ["stop"]
+    assert [entry["token"] for entry in finish_choices[0]["logprobs"]["content"]] == ["<|return|>"]
+
+
 def _hosted_vllm_stream_wrapper(logging_obj: Logging, error_payload: dict) -> CustomStreamWrapper:
     """A CustomStreamWrapper over the real OpenAI-compatible line iterator,
     fed an HTTP 200 SSE body that carries an in-body error payload the way

ui/litellm-dashboard/src/lib/http/schema.d.ts

@@ -29161,7 +29161,7 @@ export interface components {
             /** Token */
             token: string;
             /** Top Logprobs */
-            top_logprobs: components["schemas"]["TopLogprob"][];
+            top_logprobs?: components["schemas"]["TopLogprob"][];
         } & {
             [key: string]: unknown;
         };
@@ -29292,6 +29292,8 @@ export interface components {
         ChoiceLogprobs: {
             /** Content */
             content?: components["schemas"]["ChatCompletionTokenLogprob"][] | null;
+            /** Refusal */
+            refusal?: components["schemas"]["ChatCompletionTokenLogprob"][] | null;
         } & {
             [key: string]: unknown;
         };