From 98ebd747fd73aa80d0b2efe00a54ecf597772c92 Mon Sep 17 00:00:00 2001 From: Yushu Li Date: Fri, 30 Oct 2026 23:16:05 -0700 Subject: [PATCH] fix(streaming): preserve logprobs in provider fake streams --- litellm/llms/base_llm/base_model_iterator.py | 16 ++- litellm/llms/custom_httpx/llm_http_handler.py | 8 +- litellm/llms/openai_like/chat/handler.py | 4 +- .../llms/base_llm/test_base_model_iterator.py | 101 +++++++++++++++++- .../custom_httpx/test_llm_http_handler.py | 83 ++++++++++++++ .../chat/test_openai_like_chat_handler.py | 63 +++++++++++ 6 files changed, 267 insertions(+), 8 deletions(-) diff --git a/litellm/llms/base_llm/base_model_iterator.py b/litellm/llms/base_llm/base_model_iterator.py index 3afc685b3dd..5ee36cd2ca6 100644 --- a/litellm/llms/base_llm/base_model_iterator.py +++ b/litellm/llms/base_llm/base_model_iterator.py @@ -41,6 +41,8 @@ def _streaming_delta(message: Message) -> Delta: def convert_model_response_to_streaming( model_response: ModelResponse, + *, + preserve_logprobs: bool = False, ) -> ModelResponseStream: """ Convert a ModelResponse to ModelResponseStream. @@ -50,6 +52,7 @@ def convert_model_response_to_streaming( Args: model_response: The ModelResponse to convert + preserve_logprobs: Keep token probabilities only for an unmodified provider response. Returns: ModelResponseStream: A streaming chunk version of the response @@ -65,6 +68,8 @@ def convert_model_response_to_streaming( index=choice.index, delta=_streaming_delta(cast(Choices, choice).message), finish_reason=choice.finish_reason, + # Guardrail replays can redact text while retaining the original token probabilities. + logprobs=getattr(choice, "logprobs", None) if preserve_logprobs else None, ) ) processed_chunk: Final = ModelResponseStream( @@ -235,9 +240,16 @@ class BaseModelResponseIterator: class MockResponseIterator: # for returning ai21 streaming responses - def __init__(self, model_response: ModelResponse, json_mode: bool | None = False): + def __init__( + self, + model_response: ModelResponse, + json_mode: bool | None = False, + *, + preserve_logprobs: bool = False, + ): self.model_response = model_response self.json_mode = json_mode + self.preserve_logprobs = preserve_logprobs self.is_done = False # Sync iterator @@ -245,7 +257,7 @@ class MockResponseIterator: # for returning ai21 streaming responses return self def _chunk_parser(self, chunk_data: ModelResponse) -> ModelResponseStream: - return convert_model_response_to_streaming(chunk_data) + return convert_model_response_to_streaming(chunk_data, preserve_logprobs=self.preserve_logprobs) def __next__(self): if self.is_done: diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 9a2b352dd86..6c0dfb60af4 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1142,7 +1142,9 @@ class BaseLLMHTTPHandler: json_mode=json_mode, ) - completion_stream: object = MockResponseIterator(model_response=model_response, json_mode=json_mode) + completion_stream: object = MockResponseIterator( + model_response=model_response, json_mode=json_mode, preserve_logprobs=True + ) else: completion_stream = provider_config.get_model_response_iterator_for_request( streaming_response=iter_sse_lines(response.iter_text()), @@ -1285,7 +1287,9 @@ class BaseLLMHTTPHandler: json_mode=json_mode, ) - completion_stream: object = MockResponseIterator(model_response=model_response, json_mode=json_mode) + completion_stream: object = MockResponseIterator( + model_response=model_response, json_mode=json_mode, preserve_logprobs=True + ) else: completion_stream = provider_config.get_model_response_iterator_for_request( streaming_response=aiter_sse_lines(response.aiter_text()), diff --git a/litellm/llms/openai_like/chat/handler.py b/litellm/llms/openai_like/chat/handler.py index a450947afdf..33fa08c7cbd 100644 --- a/litellm/llms/openai_like/chat/handler.py +++ b/litellm/llms/openai_like/chat/handler.py @@ -64,7 +64,7 @@ async def make_call( completion_stream = streaming_decoder.aiter_bytes(response.aiter_bytes(chunk_size=1024)) elif fake_stream: model_response: Final = _fake_streamed_model_response(response.json()) - completion_stream = MockResponseIterator(model_response=model_response) + completion_stream = MockResponseIterator(model_response=model_response, preserve_logprobs=True) else: completion_stream = ModelResponseIterator( streaming_response=aiter_sse_lines(response.aiter_text()), sync_stream=False @@ -104,7 +104,7 @@ def make_sync_call( completion_stream = streaming_decoder.iter_bytes(response.iter_bytes(chunk_size=1024)) elif fake_stream: model_response: Final = _fake_streamed_model_response(response.json()) - completion_stream = MockResponseIterator(model_response=model_response) + completion_stream = MockResponseIterator(model_response=model_response, preserve_logprobs=True) else: completion_stream = ModelResponseIterator( streaming_response=iter_sse_lines(response.iter_text()), sync_stream=True diff --git a/tests/unit/llms/base_llm/test_base_model_iterator.py b/tests/unit/llms/base_llm/test_base_model_iterator.py index 2971405885d..163c7623bde 100644 --- a/tests/unit/llms/base_llm/test_base_model_iterator.py +++ b/tests/unit/llms/base_llm/test_base_model_iterator.py @@ -14,10 +14,21 @@ from pydantic import BaseModel from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper -from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator, convert_model_response_to_streaming +from litellm.llms.base_llm.base_model_iterator import ( + BaseModelResponseIterator, + MockResponseIterator, + convert_model_response_to_streaming, +) from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler from litellm.llms.openai.common_utils import OpenAIError -from litellm.types.utils import Choices, GenericStreamingChunk, Message, ModelResponse +from litellm.types.utils import ( + ChoiceLogprobs, + Choices, + GenericStreamingChunk, + Message, + ModelResponse, + ModelResponseStream, +) class TestBaseModelResponseIterator: @@ -435,3 +446,89 @@ def test_error_frame_mentioning_done_raises_instead_of_ending_the_stream() -> No list(OpenAIChatCompletionStreamingHandler(streaming_response=iter((line,)), sync_stream=True)) assert error.value.status_code == 400 + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize("preserve_logprobs", [True, False]) +@pytest.mark.asyncio +async def test_fake_stream_preserves_each_choices_logprobs(sync_mode: bool, preserve_logprobs: bool) -> None: + response: Final = ModelResponse( + id="chatcmpl-logprobs", + created=1_790_000_000, + model="test-model", + choices=[ + Choices( + index=2, + message=Message(content="Yes"), + logprobs={ + "content": [ + { + "token": "Yes", + "logprob": -0.25, + "bytes": [89, 101, 115], + "top_logprobs": [ + {"token": "Yes", "logprob": -0.25, "bytes": [89, 101, 115]}, + {"token": "No", "logprob": -1.75, "bytes": [78, 111]}, + ], + } + ], + "refusal": [], + }, + ), + Choices( + index=7, + message=Message(content=None, refusal="Cannot help"), + logprobs={ + "content": [], + "refusal": [ + {"token": "Cannot", "logprob": -0.5, "bytes": [67, 97, 110, 110, 111, 116], "top_logprobs": []} + ], + }, + ), + Choices(index=9, message=Message(content="No probabilities requested")), + ], + ) + original: Final = response.model_dump_json() + iterator: Final = ( + MockResponseIterator(response, preserve_logprobs=True) if preserve_logprobs else MockResponseIterator(response) + ) + chunks: Final = list(iterator) if sync_mode else [chunk async for chunk in iterator] + + assert len(chunks) == 1 + chunk: Final = chunks[0] + assert isinstance(chunk, ModelResponseStream) + assert [choice.index for choice in chunk.choices] == [2, 7, 9] + for logprobs, source_logprobs in zip( + (choice.logprobs for choice in chunk.choices[:2]), + (choice.logprobs for choice in response.choices[:2]), + strict=True, + ): + if preserve_logprobs: + assert isinstance(logprobs, ChoiceLogprobs) + assert isinstance(source_logprobs, ChoiceLogprobs) + assert logprobs.model_dump() == source_logprobs.model_dump() + else: + assert logprobs is None + assert chunk.choices[2].logprobs is None + assert chunk.choices[0].delta.content == "Yes" + assert getattr(chunk.choices[1].delta, "refusal", None) == "Cannot help" + assert response.model_dump_json() == original + + +def test_response_replay_drops_original_token_probabilities_by_default() -> None: + response: Final = ModelResponse( + choices=[ + Choices( + message=Message(content="[REDACTED]"), + logprobs={ + "content": [{"token": "original-secret", "logprob": -0.25, "bytes": None, "top_logprobs": []}], + "refusal": [], + }, + ) + ] + ) + + for chunk in (convert_model_response_to_streaming(response), next(MockResponseIterator(response))): + assert chunk.choices[0].delta.content == "[REDACTED]" + assert chunk.choices[0].logprobs is None + assert "original-secret" not in chunk.model_dump_json() diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index 23d7feb2557..97ce8764712 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -7,6 +7,7 @@ import threading import time from collections.abc import Awaitable, Callable, Mapping from types import MappingProxyType +from datetime import datetime, timezone from typing import Final from unittest.mock import AsyncMock, Mock, patch @@ -20,6 +21,7 @@ from h2.events import ConnectionTerminated import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.integrations.code_interpreter_interception.handler import ( CodeInterpreterInterceptionLogger, LITELLM_CODE_EXECUTION_TOOL_NAME, @@ -58,6 +60,7 @@ from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude_trans from litellm.llms.bedrock.files.transformation import BedrockFilesConfig from litellm.llms.anthropic.skills.transformation import AnthropicSkillsConfig from litellm.llms.openai.evals.transformation import OpenAIEvalsConfig +from litellm.llms.openai_like.chat.transformation import OpenAILikeChatConfig from litellm.llms.mistral.files.transformation import MistralFilesConfig from litellm.llms.openai.vector_store_files.transformation import OpenAIVectorStoreFilesConfig from litellm.llms.openai.vector_stores.transformation import OpenAIVectorStoreConfig @@ -70,6 +73,7 @@ from litellm.responses.streaming_iterator import ( from litellm.types.llms.openai import HttpxBinaryResponseContent, ResponsesAPIResponse from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( + ChoiceLogprobs, ImageObject, ImageResponse, ModelResponse, @@ -6250,3 +6254,82 @@ async def test_a_stream_answered_with_one_json_body_is_not_replayed_for_a_provid ) with pytest.raises(litellm.exceptions.MidStreamFallbackError, match="before any terminal event"): _ = tuple([event async for event in stream]) + + +@pytest.mark.parametrize("sync_stream", [True, False]) +@pytest.mark.parametrize("include_logprobs", [True, False]) +async def test_fake_chat_stream_keeps_provider_logprobs(sync_stream: bool, include_logprobs: bool) -> None: + expected: Final = ChoiceLogprobs( + content=[ + { + "token": "answer", + "logprob": -0.25, + "bytes": [97, 110, 115, 119, 101, 114], + "top_logprobs": [{"token": "answer", "logprob": -0.25, "bytes": [97, 110, 115, 119, 101, 114]}], + } + ], + refusal=[], + ) + payload: Final = { + "id": "chatcmpl-logprobs", + "created": 1_790_000_000, + "model": "test-model", + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": {"role": "assistant", "content": "answer"}, + "logprobs": expected.model_dump() if include_logprobs else None, + } + ], + } + + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json=payload) + + handler: Final = BaseLLMHTTPHandler() + data: Final = {"model": "test-model", "messages": []} + call_args: Final = { + "provider_config": OpenAILikeChatConfig(), + "api_base": "https://provider.example/v1/chat/completions", + "headers": {}, + "data": data, + "model": "test-model", + "messages": [], + "logging_obj": Logging( + model="test-model", + messages=[], + stream=True, + call_type="completion", + start_time=datetime(2026, 10, 30, tzinfo=timezone.utc), + litellm_call_id="test-logprobs", + function_id="test-function", + ), + "optional_params": {}, + "litellm_params": {}, + "timeout": 5.0, + "fake_stream": True, + } + transport: Final = httpx.MockTransport(respond) + iterator, _ = ( + handler.make_sync_call( + **call_args, + signed_json_body=None, + original_data=data, + client=HTTPHandler(client=httpx.Client(transport=transport)), + ) + if sync_stream + else await handler.make_async_call_stream_helper( + **call_args, custom_llm_provider="openai", client=AsyncHTTPHandler(transport=transport) + ) + ) + chunks: Final = list(iterator) if sync_stream else [chunk async for chunk in iterator] + + assert len(chunks) == 1 + assert chunks[0].choices[0].delta.content == "answer" + actual: Final = chunks[0].choices[0].logprobs + if include_logprobs: + assert isinstance(actual, ChoiceLogprobs) + assert actual == expected + else: + assert actual is None diff --git a/tests/unit/llms/openai_like/chat/test_openai_like_chat_handler.py b/tests/unit/llms/openai_like/chat/test_openai_like_chat_handler.py index 3df1efe311b..5ad3759a83f 100644 --- a/tests/unit/llms/openai_like/chat/test_openai_like_chat_handler.py +++ b/tests/unit/llms/openai_like/chat/test_openai_like_chat_handler.py @@ -1,4 +1,5 @@ import json +from datetime import datetime, timezone from typing import Final from unittest.mock import Mock @@ -6,8 +7,10 @@ from unittest.mock import Mock import httpx import pytest +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.openai_like.chat.handler import make_call, make_sync_call +from litellm.types.utils import ChoiceLogprobs _FAKE_STREAMED_REPLY: Final = { "id": "chatcmpl-fake-stream", @@ -106,3 +109,63 @@ async def test_stream_keeps_unicode_line_separators_inside_the_delta(sync_stream chunks: Final = await _streamed_chunks(sync_stream) assert "".join(str(chunk["text"]) for chunk in chunks) == _LINE_SEPARATED_TEXT + + +@pytest.mark.parametrize("sync_stream", [True, False]) +@pytest.mark.parametrize("include_logprobs", [True, False]) +async def test_fake_stream_preserves_provider_logprobs(sync_stream: bool, include_logprobs: bool) -> None: + expected: Final = ChoiceLogprobs( + content=[ + { + "token": "Paris", + "logprob": -0.25, + "bytes": [80, 97, 114, 105, 115], + "top_logprobs": [ + {"token": "Paris", "logprob": -0.25, "bytes": [80, 97, 114, 105, 115]}, + {"token": "London", "logprob": -1.75, "bytes": [76, 111, 110, 100, 111, 110]}, + ], + } + ], + refusal=[], + ) + payload: Final = { + **_FAKE_STREAMED_REPLY, + "choices": [ + {**_FAKE_STREAMED_REPLY["choices"][0], "logprobs": expected.model_dump() if include_logprobs else None} + ], + } + + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json=payload) + + call_args: Final = { + "api_base": "https://openai-like.example/v1/chat/completions", + "headers": {}, + "data": json.dumps({"model": "test-model", "messages": []}), + "model": "test-model", + "messages": [], + "logging_obj": Logging( + model="test-model", + messages=[], + stream=True, + call_type="completion", + start_time=datetime(2026, 10, 30, tzinfo=timezone.utc), + litellm_call_id="test-logprobs", + function_id="test-function", + ), + "fake_stream": True, + } + transport: Final = httpx.MockTransport(respond) + chunks: Final = ( + list(make_sync_call(client=HTTPHandler(client=httpx.Client(transport=transport)), **call_args)) + if sync_stream + else [chunk async for chunk in await make_call(client=AsyncHTTPHandler(transport=transport), **call_args)] + ) + + assert len(chunks) == 1 + actual: Final = chunks[0].choices[0].logprobs + if include_logprobs: + assert isinstance(actual, ChoiceLogprobs) + assert actual == expected + else: + assert actual is None