scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-31 06:56:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 245

Merge 98ebd747fd into a392a39627

Browse source
This commit is contained in:
Yushu Li 2026-10-30 23:17:27 -07:00 • committed by GitHub
commit 89bb0e5f9a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 267 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

16
litellm/llms/base_llm/base_model_iterator.py
View file

@ -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:

8
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -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()),

4
litellm/llms/openai_like/chat/handler.py
View file

@ -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

101
tests/unit/llms/base_llm/test_base_model_iterator.py
View file

@ -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()

83
tests/unit/llms/custom_httpx/test_llm_http_handler.py
View file

@ -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

63
tests/unit/llms/openai_like/chat/test_openai_like_chat_handler.py
View file

@ -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