scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 200

fix(streaming): keep merged think tags balanced and split SSE lines only on CR/LF (#48755)

Browse source 
* fix(vertex_ai): split Gemini SSE lines on CR/LF only so U+2028 in text keeps the stream

The Gemini and Vertex AI streaming path read the SSE body with httpx
iter_lines()/aiter_lines(). Those follow str.splitlines(), which also breaks
on U+0085, U+2028, U+2029, \x0b, \x0c and \x1c-\x1e. JSON leaves U+0085,
U+2028 and U+2029 unescaped inside strings, so a data: line whose text holds
one of them was cut in two. The tail is not a data: line and the head never
parses, and the accumulate-and-retry fallback then carries the unparsed
prefix forward, so every later event is appended to it and the stream ends
with no more content and no usage.

SSE lines end only at CRLF, LF or CR. The new iter_sse_lines/aiter_sse_lines
split decoded text on exactly those, treating a CR at the end of one chunk and
an LF at the start of the next as one break. The first request and every
long-decode hop now read through them.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(streaming): keep merged think tags balanced and split shared SSE lines only on CR/LF

Close an open <think> block on the first content or tool-call delta and at stream end, reopen it when reasoning returns, and route the shared chat completions stream through the SSE line splitter so U+2028 and U+2029 inside JSON no longer drop chunks.

* fix(streaming): exact-match the SSE [DONE] marker in BaseModelResponseIterator

* fix(streaming): split every SSE reader on CR/LF only and keep the logged think block closed

Switch the remaining httpx stream readers (Anthropic chat, Google GenAI,
Interactions, Codestral, openai_like, Predibase, LangGraph, Vertex Agent
Engine, A2A providers, Azure AI agents, sandboxes, the A2A endpoint, and
others) from iter_lines/aiter_lines to iter_sse_lines/aiter_sse_lines, so a
raw U+2028 or U+2029 inside a JSON string no longer splits an event.

Keep the chunk that closes a think block in the logged and cached response
when a stream ends without a finish_reason, close the block before a legacy
function_call, and make sync Codestral FIM streaming send stream=true and use
the text-completion-codestral stream parser.

Test fakes that only mocked iter_lines now stream real httpx bytes and
assert the content they feed in.

* fix(streaming): split buffered SSE parsers only on CR and LF, shard the codestral tests

Buffered SSE and JSONL parsers used str.splitlines(), which also breaks on
U+2028, U+2029 and U+0085 that JSON allows raw inside strings. They now use
split_sse_lines from sse_lines.py, which also gains a str.split fast path.
Adds tests/unit/llms/codestral to the unit workflow shard.

* fix(streaming): keep U+2028 in restamped Anthropic frames and fold CRLF before splitting SSE lines

* test: type the new stream tests, drop a comment, and benchmark the SSE line splitter

* test: type the iterator class parameters in the line-separator stream tests

* fix(responses): stream tool call arguments opened in a text chunk

* fix(responses): hold a tool call that shares a chunk with text until the text ends

* test: type the a2a SSE handler helper and the sandbox fake responses

---------

Co-authored-by: Tim Vink <vinktim@gmail.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Co-authored-by: michelligabriele <gabriele.michelli@icloud.com>
This commit is contained in:
rj-cairn 2026-10-28 17:00:29 -07:00 • committed by GitHub
parent ab07b7207e
commit d06b889d41
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
80 changed files with 1554 additions and 413 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
.github/workflows/test-unit.yml vendored
View file

@ -245,6 +245,7 @@ jobs:
tests/unit/llms/clarifai
tests/unit/llms/claude_code
tests/unit/llms/cloudflare
tests/unit/llms/codestral
tests/unit/llms/codex
tests/unit/llms/cohere
tests/unit/llms/cometapi

5
litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py
View file

@ -19,6 +19,7 @@ from litellm.a2a_protocol.utils import (
)
from litellm.exceptions import BadRequestError
from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreConfig
from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines
RUNTIME_SESSION_ID_MIN_LENGTH: Final = 33
RUNTIME_SESSION_ID_MAX_LENGTH: Final = 256
@ -50,7 +51,7 @@ _RESERVED_PREFIX_HEADERS: Final[tuple[str, ...]] = (
class _SSELineSource(Protocol):
"""Minimal streaming-response surface used to read SSE lines."""
def aiter_lines(self) -> AsyncIterator[str]: ...
def aiter_text(self) -> AsyncIterator[str]: ...
def _filter_reserved_headers(
@ -244,7 +245,7 @@ class BedrockAgentCoreA2ATransformation:
Yields:
Parsed JSON dicts from SSE data lines
"""
async for line in response.aiter_lines():
async for line in aiter_sse_lines(response.aiter_text()):
line = line.strip()
if not line:
continue

5
litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py
View file

@ -20,6 +20,7 @@ from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
get_async_httpx_client,
)
from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines
from litellm.types.llms.custom_http import httpxSpecialProvider
_IBM_CLOUD_IAM_URL: Final = "https://iam.cloud.ibm.com/identity/token"
@ -74,7 +75,7 @@ class _WXORun(TypedDict, total=False):
class _SSELineSource(Protocol):
def aiter_lines(self) -> AsyncIterator[str]: ...
def aiter_text(self) -> AsyncIterator[str]: ...
class _WXOView(TypedDict, total=False):
@ -231,7 +232,7 @@ class WatsonxOrchestrateHandler:
async def _accumulate_wxo_sse_text(response: _SSELineSource) -> str:
source: Final[_WXOView] = {"sse_source": response}
accumulated_text = ""
async for line in source["sse_source"].aiter_lines():
async for line in aiter_sse_lines(source["sse_source"].aiter_text()):
if not line.startswith("data:"):
continue
data_str = line[5:].strip()

3
litellm/completion_extras/litellm_responses_transformation/transformation.py
View file

@ -34,6 +34,7 @@ from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
from litellm.llms.base_llm.bridges.completion_transformation import (
CompletionTransformationBridge,
)
from litellm.llms.custom_httpx.sse_lines import split_sse_lines
from litellm.responses.litellm_completion_transformation.tool_search import is_anthropic_tool_search_chat_tool
from litellm.responses.sse_output_recovery import (
parse_sse_json_chunk,
@ -1351,7 +1352,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
recovered_output_items: Final[dict[int, dict[str, object]]] = {}
recovered_text_only_items: Final[dict[int, dict[str, object]]] = {}
for chunk in raw_sse.splitlines():
for chunk in split_sse_lines(raw_sse):
parsed_chunk = parse_sse_json_chunk(chunk)
if parsed_chunk is None:
continue

9
litellm/google_genai/streaming_iterator.py
View file

@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, Final
import litellm
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines
from litellm.proxy.pass_through_endpoints.success_handler import (
PassThroughEndpointLogging,
)
@ -131,9 +132,9 @@ class GoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateContent
self.model = model
self.generate_content_provider_config = generate_content_provider_config
self.litellm_metadata = litellm_metadata
# Gemini streamGenerateContent uses SSE line framing; iter_lines keeps
# Gemini streamGenerateContent uses SSE line framing; iter_sse_lines keeps
# large inlineData payloads (e.g. image/jpeg) intact within one event.
self.stream_iterator = response.iter_lines()
self.stream_iterator = iter_sse_lines(response.iter_text())
def __iter__(self):
return self
@ -182,9 +183,9 @@ class AsyncGoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateCo
self.model = model
self.generate_content_provider_config = generate_content_provider_config
self.litellm_metadata = litellm_metadata
# Gemini streamGenerateContent uses SSE line framing; aiter_lines keeps
# Gemini streamGenerateContent uses SSE line framing; aiter_sse_lines keeps
# large inlineData payloads (e.g. image/jpeg) intact within one event.
self.stream_iterator = response.aiter_lines()
self.stream_iterator = aiter_sse_lines(response.aiter_text())
def __aiter__(self):
return self

5
litellm/interactions/streaming_iterator.py
View file

@ -21,6 +21,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base
from litellm.litellm_core_utils.thread_pool_executor import executor
from litellm.llms.base_llm.interactions.transformation import BaseInteractionsAPIConfig
from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines
from litellm.types.interactions import (
InteractionsAPIStreamingResponse,
)
@ -145,7 +146,7 @@ class InteractionsAPIStreamingIterator(BaseInteractionsAPIStreamingIterator):
litellm_metadata=litellm_metadata,
custom_llm_provider=custom_llm_provider,
)
self.stream_iterator = response.aiter_lines()
self.stream_iterator = aiter_sse_lines(response.aiter_text())
def __aiter__(self):
return self
@ -212,7 +213,7 @@ class SyncInteractionsAPIStreamingIterator(BaseInteractionsAPIStreamingIterator)
litellm_metadata=litellm_metadata,
custom_llm_provider=custom_llm_provider,
)
self.stream_iterator = response.iter_lines()
self.stream_iterator = iter_sse_lines(response.iter_text())
def __iter__(self):
return self

57
litellm/litellm_core_utils/streaming_handler.py
View file

@ -1263,7 +1263,7 @@ class CustomStreamWrapper:
)
if _is_delta_empty:
model_response.choices[0].delta = Delta(content=None) # ensure empty delta chunk returned
model_response.choices[0].delta = Delta(content=self._close_open_think_block())
# get any function call arguments
model_response.choices[0].finish_reason = map_finish_reason(
finish_reason=self.received_finish_reason
@ -1292,21 +1292,26 @@ class CustomStreamWrapper:
if self.merge_reasoning_content_in_choices is not True:
return
delta: Final = model_response.choices[0].delta
reasoning_content: Final = getattr(delta, "reasoning_content", None)
reasoning_content: Final = getattr(delta, "reasoning_content", None) or ""
content: Final = delta.content or ""
if reasoning_content:
self._emitted_disqualifying_content = True
closes_thinking: Final = bool(content) and not self.sent_last_thinking_block
opening: Final = "" if self.sent_first_thinking_block else "<think>"
delta.content = opening + reasoning_content + ("</think>" + content if closes_thinking else "")
self.sent_first_thinking_block = True
self.sent_last_thinking_block = self.sent_last_thinking_block or closes_thinking
elif self.sent_first_thinking_block and not self.sent_last_thinking_block and content:
delta.content = "</think>" + content
self.sent_last_thinking_block = True
ends_thinking: Final = bool(content or delta.tool_calls or delta.function_call)
inside_think_block: Final = self.sent_first_thinking_block and not self.sent_last_thinking_block
opening: Final = "<think>" if reasoning_content and not inside_think_block else ""
closing: Final = "</think>" if ends_thinking and (reasoning_content or inside_think_block) else ""
if hasattr(delta, "reasoning_content"):
del delta.reasoning_content
if not reasoning_content and not closing:
return
delta.content = opening + reasoning_content + closing + content
self.sent_first_thinking_block = True
self.sent_last_thinking_block = bool(closing)
self._emitted_disqualifying_content = self._emitted_disqualifying_content or bool(reasoning_content)
def _close_open_think_block(self) -> str | None:
if not self.sent_first_thinking_block or self.sent_last_thinking_block:
return None
self.sent_last_thinking_block = True
return "</think>"
def _dispatch_provider_chunk(
self,
@ -1932,8 +1937,20 @@ class CustomStreamWrapper:
model_response.choices[0].finish_reason == "stop" and self.tool_call
): # don't overwrite for other - potential error finish reasons
model_response.choices[0].finish_reason = "tool_calls"
close_tag: Final = self._close_open_think_block()
if close_tag is not None:
model_response.choices[0].delta.content = close_tag
return model_response
def _final_chunk_for_logging(self, processed_chunk: "ModelResponseStream") -> "ModelResponseStream | None":
if self.received_finish_reason is not None or self.intermittent_finish_reason is not None:
return processed_chunk
if processed_chunk.choices[0].delta.content is None:
return None
return processed_chunk.model_copy(
update={"choices": [processed_chunk.choices[0].model_copy(update={"finish_reason": None})]}
)
def _record_usage_only_chunk(self, model_response: "ModelResponseStream") -> None:
"""
Keep provider usage-only chunks (e.g. OpenRouter's post-finish chunk, which carries a
@ -2173,10 +2190,9 @@ class CustomStreamWrapper:
else:
self.sent_last_chunk = True
processed_chunk: Final = self.finish_reason_handler()
# The logged response is built from self.chunks; keep a finish_reason the provider sent on its
# last content chunk (stripped there), but never add the synthetic "stop" used when it sent none.
if self.received_finish_reason is not None or self.intermittent_finish_reason is not None:
self.chunks.append(processed_chunk)
final_chunk_for_logging: Final = self._final_chunk_for_logging(processed_chunk)
if final_chunk_for_logging is not None:
self.chunks.append(final_chunk_for_logging)
if self.stream_options is None:
self._set_final_chunk_usage(processed_chunk)
## LOGGING
@ -2406,10 +2422,9 @@ class CustomStreamWrapper:
else:
self.sent_last_chunk = True
processed_chunk: Final = self.finish_reason_handler()
# The logged response is built from self.chunks; keep a finish_reason the provider sent on its
# last content chunk (stripped there), but never add the synthetic "stop" used when it sent none.
if self.received_finish_reason is not None or self.intermittent_finish_reason is not None:
self.chunks.append(processed_chunk)
final_chunk_for_logging: Final = self._final_chunk_for_logging(processed_chunk)
if final_chunk_for_logging is not None:
self.chunks.append(final_chunk_for_logging)
if self.stream_options is None:
self._set_final_chunk_usage(processed_chunk)
# see sync __next__'s sibling branch: deliberately do NOT restore

5
litellm/llms/anthropic/chat/handler.py
View file

@ -38,6 +38,7 @@ from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
get_httpx_client,
)
from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines
from litellm.types.llms.anthropic import (
ContentBlockDelta,
ContentBlockStart,
@ -233,7 +234,7 @@ async def make_call(
raise AnthropicError(status_code=500, message=str(e))
completion_stream: Final = ModelResponseIterator(
streaming_response=response.aiter_lines(),
streaming_response=aiter_sse_lines(response.aiter_text()),
sync_stream=False,
json_mode=json_mode,
speed=speed,
@ -304,7 +305,7 @@ def make_sync_call(
)
completion_stream: Final = ModelResponseIterator(
streaming_response=response.iter_lines(),
streaming_response=iter_sse_lines(response.iter_text()),
sync_stream=True,
json_mode=json_mode,
speed=speed,

3
litellm/llms/anthropic/messages/stream_assembly.py
View file

@ -13,6 +13,7 @@ from litellm.llms.anthropic.chat.handler import (
)
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
from litellm.llms.anthropic.messages.native.utils import MESSAGE_DELTA_OPTIONAL_KEYS
from litellm.llms.custom_httpx.sse_lines import split_sse_lines
from litellm.types.utils import (
Choices,
Message,
@ -269,7 +270,7 @@ def build_complete_streaming_response_legacy(
def extract_sse_data(event_str: str) -> dict | None:
"""Parse the JSON object from the ``data:`` line of an Anthropic SSE event."""
for line in event_str.splitlines():
for line in split_sse_lines(event_str):
stripped = line.strip()
if stripped.startswith("data:"):
payload = stripped[len("data:") :].strip()

3
litellm/llms/azure_ai/agents/handler.py
View file

@ -35,6 +35,7 @@ from litellm.llms.azure_ai.agents.transformation import (
AzureAIAgentsConfig,
AzureAIAgentsError,
)
from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines
from litellm.types.llms.openai import (
ChatCompletionAnnotation,
ChatCompletionAnnotationURLCitation,
@ -632,7 +633,7 @@ class AzureAIAgentsHandler:
current_event = None
async for line in response.aiter_lines():
async for line in aiter_sse_lines(response.aiter_text()):
line = line.strip()
if line.startswith("event:"):

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

@ -96,7 +96,7 @@ class BaseModelResponseIterator:
released (and a backend like vLLM aborts generation) when the stream
is abandoned before its natural end.
``streaming_response`` is usually a bare ``aiter_lines()`` generator
``streaming_response`` is usually a bare ``aiter_sse_lines()`` generator
that holds no reference to the response, so the handler that owns the
response attaches it here after construction."""
if self.http_response is not None:

3
litellm/llms/base_llm/guardrail_translation/utils.py
View file

@ -6,6 +6,7 @@ from typing import Final, TypeVar, cast # noqa: TID251 # a rebuilt chat row ha
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
from litellm.llms.custom_httpx.sse_lines import split_sse_lines
from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage
from litellm.types.llms.openai import AllMessageValues, ResponseAPIUsage
@ -24,7 +25,7 @@ def _anthropic_stream_chunk_events(item: object) -> list[dict]:
events: Final[list[dict]] = []
for block in chunk.split("\n\n"):
for line in block.splitlines():
for line in split_sse_lines(block):
stripped = line.strip()
if not stripped.startswith("data:"):
continue

3
litellm/llms/base_llm/sandbox/transformation.py
View file

@ -12,6 +12,7 @@ from typing import Final
import httpx
from pydantic import Field, PrivateAttr
from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines
from litellm.types.llms.base import LiteLLMPydanticObjectBase
SANDBOX_MAX_OUTPUT_BYTES: Final = 10 * 1024 * 1024
@ -99,7 +100,7 @@ class BaseSandboxConfig:
async def _read_capped_lines(self, response: httpx.Response) -> list[str]:
lines: Final[list[str]] = []
total = 0
async for line in response.aiter_lines():
async for line in aiter_sse_lines(response.aiter_text()):
total += len(line.encode("utf-8"))
if total > SANDBOX_MAX_OUTPUT_BYTES:
raise ValueError(

3
litellm/llms/bedrock/files/transformation.py
View file

@ -49,6 +49,7 @@ from litellm.llms.base_llm.files.transformation import (
BaseFilesConfig,
LiteLLMLoggingObj,
)
from litellm.llms.custom_httpx.sse_lines import split_sse_lines
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.llms.bedrock import AwsAuthParams, BedrockBatchRecordKind
from litellm.types.llms.openai import (
@ -400,7 +401,7 @@ def _batch_input_line(line_number: int, line: str) -> _BatchInputLine:
def _non_blank_lines(file_content: str) -> Iterator[tuple[int, str]]:
numbered_lines: Final = enumerate(file_content.splitlines(), start=1)
numbered_lines: Final = enumerate(split_sse_lines(file_content), start=1)
return ((line_number, line) for line_number, line in numbered_lines if line.strip())

3
litellm/llms/bedrock/passthrough/guardrail_translation/handler.py
View file

@ -45,6 +45,7 @@ from litellm.llms.bedrock.passthrough.guardrail_translation.structured_edits imp
structured_edits,
tool_edits,
)
from litellm.llms.custom_httpx.sse_lines import split_sse_lines
from litellm.proxy.guardrails.auto_router_compression import is_compression_guardrail
from litellm.types.utils import GenericGuardrailAPIInputs
@ -947,7 +948,7 @@ def _invoke_frames(items: Sequence[object]) -> tuple[InvokeStreamFrame, ...]:
def _sse_events(sse: bytes) -> tuple[Mapping[str, object], ...]:
return tuple(
_JSON_OBJECT.validate_python(json.loads(line[len("data:") :]))
for line in sse.decode().splitlines()
for line in split_sse_lines(sse.decode())
if line.startswith("data:")
)

5
litellm/llms/bedrock/search/transformation.py
View file

@ -47,6 +47,7 @@ from litellm.llms.base_llm.search.transformation import (
)
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.llms.custom_httpx.sse_lines import split_sse_lines
from litellm.secret_managers.main import get_secret_str
# AgentCore web-search rejects queries longer than 200 characters
@ -146,7 +147,9 @@ def _iter_sse_events(text: str) -> Iterator[Mapping[str, object]]:
progress notifications before the JSON-RPC response.
"""
for chunk in _SSE_EVENT_SEPARATOR.split(text):
payload = "\n".join(line[len("data:") :].lstrip() for line in chunk.splitlines() if line.startswith("data:"))
payload = "\n".join(
line[len("data:") :].lstrip() for line in split_sse_lines(chunk) if line.startswith("data:")
)
if not payload:
continue
try:

3
litellm/llms/chatgpt/responses/transformation.py
View file

@ -11,6 +11,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo
safe_convert_created_field,
)
from litellm.litellm_core_utils.url_utils import join_api_base
from litellm.llms.custom_httpx.sse_lines import split_sse_lines
from litellm.llms.openai.common_utils import OpenAIError
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
from litellm.responses.sse_output_recovery import (
@ -199,7 +200,7 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig):
error_message = None
streamed_output_items: Final[dict[int, dict[str, object]]] = {}
text_only_output_items: Final[dict[int, dict[str, object]]] = {}
for chunk in body_text.splitlines():
for chunk in split_sse_lines(body_text):
parsed_chunk = parse_sse_json_chunk(chunk)
if parsed_chunk is None:
continue

9
litellm/llms/codestral/completion/handler.py
View file

@ -20,6 +20,7 @@ from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
get_async_httpx_client,
)
from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines
from litellm.types.utils import TextChoices
from litellm.utils import CustomStreamWrapper, TextCompletionResponse
@ -109,7 +110,7 @@ async def make_call(
if response.status_code != 200:
raise TextCompletionCodestralError(status_code=response.status_code, message=response.text)
completion_stream: Final = response.aiter_lines()
completion_stream: Final = aiter_sse_lines(response.aiter_text())
# LOGGING
logging_obj.post_call(
input=messages,
@ -339,13 +340,13 @@ class CodestralTextCompletion:
response = litellm.module_level_client.post(
completion_url,
headers=headers,
data=json.dumps(data),
data=json.dumps({**data, "stream": True}),
stream=stream,
)
_response: Final = CustomStreamWrapper(
response.iter_lines(),
iter_sse_lines(response.iter_text()),
model,
custom_llm_provider="codestral",
custom_llm_provider="text-completion-codestral",
logging_obj=logging_obj,
)
return _response

3
litellm/llms/custom_httpx/aiohttp_handler.py
View file

@ -21,6 +21,7 @@ from litellm.llms.custom_httpx.http_handler import (
get_httpx_client,
get_ssl_configuration,
)
from litellm.llms.custom_httpx.sse_lines import iter_sse_lines
from litellm.types.llms.openai import FileTypes
from litellm.types.utils import HttpHandlerRequestFields, ImageResponse, LlmProviders
from litellm.utils import CustomStreamWrapper, ModelResponse, ProviderConfigManager
@ -466,7 +467,7 @@ class BaseLLMAIOHTTPHandler:
)
else:
completion_stream = provider_config.get_model_response_iterator(
streaming_response=response.iter_lines(), sync_stream=True
streaming_response=iter_sse_lines(response.iter_text()), sync_stream=True
)
# LOGGING

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

@ -118,6 +118,7 @@ from litellm.llms.custom_httpx.http_handler import (
get_httpx_client,
header_value,
)
from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines
from litellm.responses.streaming_iterator import (
BaseResponsesAPIStreamingIterator,
MockResponsesAPIStreamingIterator,
@ -1086,7 +1087,7 @@ class BaseLLMHTTPHandler:
completion_stream: object = MockResponseIterator(model_response=model_response, json_mode=json_mode)
else:
completion_stream = provider_config.get_model_response_iterator_for_request(
streaming_response=response.iter_lines(),
streaming_response=iter_sse_lines(response.iter_text()),
sync_stream=True,
json_mode=json_mode,
optional_params=_OPTIONAL_PARAMS.validate_python(optional_params),
@ -1229,7 +1230,7 @@ class BaseLLMHTTPHandler:
completion_stream: object = MockResponseIterator(model_response=model_response, json_mode=json_mode)
else:
completion_stream = provider_config.get_model_response_iterator_for_request(
streaming_response=response.aiter_lines(),
streaming_response=aiter_sse_lines(response.aiter_text()),
sync_stream=False,
json_mode=json_mode,
optional_params=_OPTIONAL_PARAMS.validate_python(optional_params),

74
litellm/llms/custom_httpx/sse_lines.py Normal file
View file

@ -0,0 +1,74 @@
"""
Split a decoded Server-Sent Events body into lines the way the SSE spec does.
httpx's `iter_lines()` / `aiter_lines()` follow `str.splitlines()`, which also breaks on U+0085, U+2028,
U+2029, \\x0b, \\x0c and \\x1c-\\x1e. JSON leaves U+0085, U+2028 and U+2029 unescaped inside strings, so a
`data:` line carrying one of them gets cut in half. SSE lines end only at CRLF, LF or CR
(https://html.spec.whatwg.org/multipage/server-sent-events.html#parsing-an-event-stream)
"""
import io
import re
from collections.abc import AsyncIterable, AsyncIterator, Iterable, Iterator, Sequence
from typing import Final
_SSE_LINE_BREAK: Final = re.compile(r"\r\n|\r|\n")
_SSE_LINE_WITH_BREAK: Final = re.compile(r"[^\r\n]*(?:\r\n|\r|\n)|[^\r\n]+")
def _split_lines(text: str) -> Sequence[str]:
if "\r" not in text:
return text.split("\n")
crlf_folded: Final = text.replace("\r\n", "\n")
if "\r" not in crlf_folded:
return crlf_folded.split("\n")
return _SSE_LINE_BREAK.split(text)
class _SSELineSplitter:
def __init__(self) -> None:
self._partial_line: io.StringIO = io.StringIO()
self._after_cr: bool = False
def feed(self, text: str) -> tuple[str, ...]:
if not text:
return ()
fresh: Final = text[1:] if self._after_cr and text.startswith("\n") else text
self._after_cr = fresh.endswith("\r")
*ended, tail = _split_lines(fresh)
if not ended:
_ = self._partial_line.write(tail)
return ()
first: Final = self._partial_line.getvalue() + ended[0]
self._partial_line = io.StringIO()
_ = self._partial_line.write(tail)
return (first, *ended[1:])
def flush(self) -> tuple[str, ...]:
rest: Final = self._partial_line.getvalue()
self._partial_line = io.StringIO()
return (rest,) if rest else ()
def iter_sse_lines(text_chunks: Iterable[str]) -> Iterator[str]:
splitter: Final = _SSELineSplitter()
for text in text_chunks:
yield from splitter.feed(text)
yield from splitter.flush()
def split_sse_lines(text: str) -> tuple[str, ...]:
return tuple(iter_sse_lines((text,)))
def split_sse_lines_keepends(text: str) -> tuple[str, ...]:
return tuple(_SSE_LINE_WITH_BREAK.findall(text))
async def aiter_sse_lines(text_chunks: AsyncIterable[str]) -> AsyncIterator[str]:
splitter: Final = _SSELineSplitter()
async for text in text_chunks:
for line in splitter.feed(text):
yield line
for line in splitter.flush():
yield line

3
litellm/llms/gemini/google_genai/guardrail_translation/handler.py
View file

@ -23,6 +23,7 @@ from litellm.llms.base_llm.guardrail_translation.base_translation import (
BaseTranslation,
StreamTransformSink,
)
from litellm.llms.custom_httpx.sse_lines import split_sse_lines
from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
@ -121,7 +122,7 @@ def _parse_json_dict_or_none(payload: str) -> Mapping[str, object] | None:
def _sse_payload_texts(sse_text: str) -> tuple[str, ...]:
return tuple(
text
for line in sse_text.splitlines()
for line in split_sse_lines(sse_text)
if line.startswith("data:")
for payload in (line[len("data:") :].strip(),)
if payload and payload != "[DONE]"

3
litellm/llms/huggingface/embedding/transformation.py
View file

@ -17,6 +17,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
)
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
from litellm.llms.custom_httpx.sse_lines import iter_sse_lines
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import Choices, Message, ModelResponse, Usage
@ -357,7 +358,7 @@ class HuggingFaceEmbeddingConfig(BaseConfig):
api_key: str | None = None,
) -> list[dict[str, str]]:
streamed_response: Final = CustomStreamWrapper(
completion_stream=response.iter_lines(),
completion_stream=iter_sse_lines(response.iter_text()),
model=model,
custom_llm_provider="huggingface",
logging_obj=logging_obj,

5
litellm/llms/langgraph/chat/sse_iterator.py
View file

@ -12,6 +12,7 @@ import httpx
from pydantic import TypeAdapter
from litellm._logging import verbose_logger
from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
_DECODED_JSON: Final = TypeAdapter(object)
@ -36,12 +37,12 @@ class LangGraphSSEStreamIterator:
def __iter__(self):
"""Initialize sync iteration."""
self.line_iterator = self.response.iter_lines()
self.line_iterator = iter_sse_lines(self.response.iter_text())
return self
def __aiter__(self):
"""Initialize async iteration."""
self.async_line_iterator = self.response.aiter_lines()
self.async_line_iterator = aiter_sse_lines(self.response.aiter_text())
return self
def _parse_sse_line(self, line: str) -> ModelResponseStream | None:

3
litellm/llms/oobabooga/chat/oobabooga.py
View file

@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Final
import litellm
from litellm.litellm_core_utils.url_utils import join_api_base
from litellm.llms.custom_httpx.http_handler import get_httpx_client
from litellm.llms.custom_httpx.sse_lines import iter_sse_lines
from litellm.utils import EmbeddingResponse, ModelResponse, Usage
from ..common_utils import OobaboogaError
@ -74,7 +75,7 @@ def completion(
stream=optional_params["stream"] if "stream" in optional_params else False,
)
if "stream" in optional_params and optional_params["stream"] is True:
return response.iter_lines()
return iter_sse_lines(response.iter_text())
else:
return oobabooga_config.transform_response(
model=model,

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

@ -16,6 +16,7 @@ from litellm import LlmProviders
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.base_model_iterator import MockResponseIterator
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines
from litellm.llms.databricks.streaming_utils import ModelResponseIterator
from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig
from litellm.llms.openai.openai import OpenAIConfig
@ -65,7 +66,9 @@ async def make_call(
model_response: Final = _fake_streamed_model_response(response.json())
completion_stream = MockResponseIterator(model_response=model_response)
else:
completion_stream = ModelResponseIterator(streaming_response=response.aiter_lines(), sync_stream=False)
completion_stream = ModelResponseIterator(
streaming_response=aiter_sse_lines(response.aiter_text()), sync_stream=False
)
# LOGGING
logging_obj.post_call(
input=messages,
@ -103,7 +106,9 @@ def make_sync_call(
model_response: Final = _fake_streamed_model_response(response.json())
completion_stream = MockResponseIterator(model_response=model_response)
else:
completion_stream = ModelResponseIterator(streaming_response=response.iter_lines(), sync_stream=True)
completion_stream = ModelResponseIterator(
streaming_response=iter_sse_lines(response.iter_text()), sync_stream=True
)
# LOGGING
logging_obj.post_call(

5
litellm/llms/predibase/chat/handler.py
View file

@ -14,6 +14,7 @@ from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
get_async_httpx_client,
)
from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines
from litellm.utils import CustomStreamWrapper, ModelResponse
from ..common_utils import PredibaseError
@ -34,7 +35,7 @@ async def make_call(
if response.status_code != 200:
raise PredibaseError(status_code=response.status_code, message=response.text)
completion_stream: Final = response.aiter_lines()
completion_stream: Final = aiter_sse_lines(response.aiter_text())
# LOGGING
logging_obj.post_call(
input=messages,
@ -163,7 +164,7 @@ class PredibaseChatCompletion:
timeout=timeout,
)
_response: Final = CustomStreamWrapper(
response.iter_lines(),
iter_sse_lines(response.iter_text()),
model,
custom_llm_provider="predibase",
logging_obj=logging_obj,

5
litellm/llms/vertex_ai/agent_engine/transformation.py
View file

@ -22,6 +22,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
convert_content_list_to_str,
)
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines
from litellm.llms.vertex_ai.agent_engine.sse_iterator import (
VertexAgentEngineResponseIterator,
)
@ -357,7 +358,7 @@ class VertexAgentEngineConfig(BaseConfig, VertexBase):
) -> VertexAgentEngineResponseIterator:
"""Return a streaming iterator for SSE responses."""
return VertexAgentEngineResponseIterator(
streaming_response=raw_response.iter_lines(),
streaming_response=iter_sse_lines(raw_response.iter_text()),
sync_stream=True,
)
@ -465,7 +466,7 @@ class VertexAgentEngineConfig(BaseConfig, VertexBase):
# Create iterator for SSE stream (async)
completion_stream: Final = VertexAgentEngineResponseIterator(
streaming_response=response.aiter_lines(),
streaming_response=aiter_sse_lines(response.aiter_text()),
sync_stream=False,
)

9
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
View file

@ -41,6 +41,7 @@ from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
get_httpx_client,
)
from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines
from litellm.types.llms.anthropic import AnthropicThinkingParam
from litellm.types.llms.base import CachedTokensDetails
from litellm.types.llms.gemini import BidiGenerateContentServerMessage
@ -3241,7 +3242,7 @@ async def make_call(
raise
completion_stream: Final = ModelResponseIterator(
streaming_response=response.aiter_lines(),
streaming_response=aiter_sse_lines(response.aiter_text()),
sync_stream=False,
logging_obj=logging_obj,
response_headers=response.headers,
@ -3312,7 +3313,7 @@ def make_sync_call(
raise
completion_stream: Final = ModelResponseIterator(
streaming_response=response.iter_lines(),
streaming_response=iter_sse_lines(response.iter_text()),
sync_stream=True,
logging_obj=logging_obj,
response_headers=response.headers,
@ -4190,7 +4191,7 @@ class ModelResponseIterator:
log_failed_hop(e)
return self._truncated_stream_end(token)
self._start_hop(hop_response, token)
self.streaming_response = hop_response.aiter_lines()
self.streaming_response = aiter_sse_lines(hop_response.aiter_text())
self.async_response_iterator = self.streaming_response
return None
@ -4204,7 +4205,7 @@ class ModelResponseIterator:
log_failed_hop(e)
return self._truncated_stream_end(token)
self._start_hop(hop_response, token)
self.streaming_response = hop_response.iter_lines()
self.streaming_response = iter_sse_lines(hop_response.iter_text())
self.response_iterator = self.streaming_response
return None

3
litellm/proxy/agent_endpoints/a2a_endpoints.py
View file

@ -26,6 +26,7 @@ from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.core_helpers import as_str_mapping
from litellm.litellm_core_utils.url_utils import SSRFError, validate_url
from litellm.llms.a2a.common_utils import resolve_a2a_hop_auth_header
from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.a2a.version_convert import (
A2AVersion,
@ -320,7 +321,7 @@ async def _a2a_sse_event_source(
"error": {"code": -32603, "message": resp.reason_phrase},
}
return
async for line in resp.aiter_lines():
async for line in aiter_sse_lines(resp.aiter_text()):
stripped = line.strip()
if not stripped.startswith("data:"):
continue

4
litellm/proxy/anthropic_endpoints/streaming_model_restamp.py
View file

@ -14,6 +14,8 @@ from typing import Final
from pydantic import TypeAdapter, ValidationError
from litellm.llms.custom_httpx.sse_lines import split_sse_lines_keepends
_MESSAGE_START_EVENT: Final = "message_start"
_MESSAGE_START_MARKER: Final = b"message_start"
_SSE_DATA_FIELD: Final = "data:"
@ -52,7 +54,7 @@ def _restamped_data_line(line: str, requested_model: str) -> str | None:
def _restamped_frame(frame: str, requested_model: str) -> str | None:
lines: Final = frame.splitlines(keepends=True)
lines: Final = split_sse_lines_keepends(frame)
restamped: Final = tuple(_restamped_data_line(line, requested_model) for line in lines)
if all(line is None for line in restamped):
return None

7
litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py
View file

@ -29,6 +29,7 @@ from litellm.llms.anthropic.messages.stream_assembly import (
split_sse_chunk_into_events,
)
from litellm.llms.bedrock.common_utils import BEDROCK_INVOKE_TRACE_FIELD
from litellm.llms.custom_httpx.sse_lines import split_sse_lines
from litellm.proxy._types import PassThroughEndpointLoggingTypedDict
from litellm.proxy.auth.auth_utils import get_end_user_id_from_request_body
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import (
@ -61,7 +62,7 @@ _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_i
def _event_payload(event: str) -> Mapping[str, object] | None:
data_lines: Final = tuple(
line.strip()[len("data:") :].strip() for line in event.splitlines() if line.strip().startswith("data:")
line.strip()[len("data:") :].strip() for line in split_sse_lines(event) if line.strip().startswith("data:")
)
if not data_lines or data_lines[0] in ("", "[DONE]"):
return None
@ -362,7 +363,7 @@ class AnthropicPassthroughLoggingHandler:
) -> str | None:
for raw in all_chunks:
text = raw.decode("utf-8") if isinstance(raw, bytes) else raw
for line in text.splitlines():
for line in split_sse_lines(text):
if not line.startswith("data:"):
continue
try:
@ -402,7 +403,7 @@ class AnthropicPassthroughLoggingHandler:
"""
for raw in reversed(all_chunks):
text = raw.decode("utf-8") if isinstance(raw, bytes) else raw
for line in reversed(text.splitlines()):
for line in reversed(split_sse_lines(text)):
if not line.startswith("data:"):
continue
try:

3
litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py
View file

@ -21,6 +21,7 @@ from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
get_async_httpx_client,
)
from litellm.llms.custom_httpx.sse_lines import split_sse_lines
from litellm.proxy._types import PassThroughEndpointLoggingTypedDict
from litellm.types.llms.custom_http import httpxSpecialProvider
from litellm.types.passthrough_endpoints.tinyfish import (
@ -81,7 +82,7 @@ def sse_poller_spawned(logging_obj: LiteLLMLoggingObj) -> bool:
def run_id_from_sse_frames(frames: bytes) -> str | None:
return _run_id_from_sse_chunks(frames.decode("utf-8", errors="replace").splitlines())
return _run_id_from_sse_chunks(split_sse_lines(frames.decode("utf-8", errors="replace")))
def resolve_tinyfish_agent_api_base() -> str:

3
litellm/proxy/pass_through_endpoints/managed_id_rewriter.py
View file

@ -51,6 +51,7 @@ from litellm.llms.base_llm.managed_resources.isolation import (
can_access_resource,
resolve_resource_owner_id,
)
from litellm.llms.custom_httpx.sse_lines import split_sse_lines
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.batches_endpoints.common_utils import validate_batch_list_limit
from litellm.proxy.common_utils.sse_keepalive import split_complete_sse_frames
@ -830,7 +831,7 @@ _SSE_EVENT_ADAPTER: Final = TypeAdapter(Mapping[str, JsonValue])
def _first_streamed_response(frames: bytes) -> tuple[str, Mapping[str, JsonValue]] | None:
for line in frames.decode("utf-8", errors="replace").splitlines():
for line in split_sse_lines(frames.decode("utf-8", errors="replace")):
if not line.startswith(_SSE_DATA_PREFIX):
continue
try:

3
litellm/proxy/response_polling/background_streaming.py
View file

@ -18,6 +18,7 @@ from starlette.types import Message
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.llms.custom_httpx.sse_lines import split_sse_lines
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.proxy.response_polling.polling_handler import ResponsePollingHandler
@ -81,7 +82,7 @@ class _StreamsBody(Protocol):
def _sse_frame_data(frame: str) -> str | None:
return next((line[6:].strip() for line in frame.splitlines() if line.startswith("data: ")), None)
return next((line[6:].strip() for line in split_sse_lines(frame) if line.startswith("data: ")), None)
async def _never_receive() -> Message:

13
litellm/responses/litellm_completion_transformation/streaming_iterator.py
View file

@ -165,6 +165,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
self._upstream_exhausted: bool = False
self._response_id_primed: bool = False
self._pending_tool_events: list[BaseLiteLLMOpenAIResponseObject] = []
self._tool_calls_held_behind_text: tuple[object, ...] = ()
self._tool_output_index_by_call_id: dict[str, int] = {}
self._tool_args_by_call_id: dict[str, str] = {}
self._tool_item_id_by_call_id: dict[str, str] = {} # mutable-ok: filled per call id as tool call events stream
@ -388,6 +389,15 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
return ""
return self._tool_call_id_by_index.get(index, "")
def _hold_tool_calls_behind_text(self, tool_calls: object) -> None:
if isinstance(tool_calls, list):
self._tool_calls_held_behind_text = (*self._tool_calls_held_behind_text, *tool_calls)
def _release_held_tool_calls(self) -> tuple[object, ...]:
held: Final = self._tool_calls_held_behind_text
self._tool_calls_held_behind_text = ()
return held
def _queue_tool_call_delta_events(self, tool_calls: object) -> None:
"""
Convert chat-completions streaming `tool_calls` deltas into Responses API streaming events.
@ -1365,6 +1375,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
if delta_content:
if not self.sent_message_item_added_event:
self._queue_message_item_added_events()
self._hold_tool_calls_behind_text(getattr(chunk.choices[0].delta, "tool_calls", None))
self._sequence_number += 1
text_delta_event: Final = OutputTextDeltaEvent(
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA,
@ -1379,7 +1390,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
# Priority 3: Handle tool call deltas (if any) -> queue events and emit them
# For each tool call delta, we emit events one at a time to match OpenAI's streaming behavior
if chunk.choices and hasattr(chunk.choices[0].delta, "tool_calls") and chunk.choices[0].delta.tool_calls:
self._queue_tool_call_delta_events(chunk.choices[0].delta.tool_calls)
self._queue_tool_call_delta_events([*self._release_held_tool_calls(), *chunk.choices[0].delta.tool_calls])
# Return one pending tool event at a time
if self._pending_tool_events:
return self._pending_tool_events.pop(0)

27
tests/benchmarks/test_sse_lines_benchmarks.py Normal file
View file

@ -0,0 +1,27 @@
from collections.abc import Sequence
from typing import Final
import pytest
from litellm.llms.custom_httpx.sse_lines import iter_sse_lines
EVENT: Final = 'data: {"choices":[{"delta":{"content":"hello world"}}]}'
def _chunks(line_break: str) -> Sequence[str]:
body: Final = "".join(f"{EVENT}{line_break}{line_break}" for _ in range(2000))
return tuple(body[i : i + 512] for i in range(0, len(body), 512))
LF_CHUNKS: Final = _chunks("\n")
CRLF_CHUNKS: Final = _chunks("\r\n")
@pytest.mark.benchmark
def test_iter_sse_lines_lf_stream() -> None:
assert sum(1 for _ in iter_sse_lines(LF_CHUNKS)) == 4000
@pytest.mark.benchmark
def test_iter_sse_lines_crlf_stream() -> None:
assert sum(1 for _ in iter_sse_lines(CRLF_CHUNKS)) == 4000

32
tests/unit/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py
View file

@ -10,6 +10,7 @@ Verifies that:
"""
import json
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
@ -560,6 +561,37 @@ class TestStreaming:
assert sent_body["params"]["message"]["messageId"] == "msg-001"
assert [event["result"]["kind"] for event in events] == ["task", "status-update"]
@pytest.mark.asyncio
@pytest.mark.usefixtures("httpx_transport")
async def test_streaming_keeps_unicode_line_separators_inside_an_sse_event(self) -> None:
from litellm.a2a_protocol.providers.bedrock_agentcore.config import (
BedrockAgentCoreA2AConfig,
)
event: Final = {
"jsonrpc": "2.0",
"id": "req-001",
"result": {"kind": "artifact-update", "artifact": {"parts": [{"kind": "text", "text": "a\u2028b\u2029c"}]}},
}
with respx.mock(assert_all_called=True) as router:
router.post(url__regex=r".*/invocations.*").mock(
return_value=httpx.Response(
200,
headers={"content-type": "text/event-stream"},
text=f"data: {json.dumps(event, ensure_ascii=False)}\n\n",
)
)
events: Final = [
streamed
async for streamed in BedrockAgentCoreA2AConfig().handle_streaming(
request_id="req-001",
params=SAMPLE_PARAMS,
litellm_params=SAMPLE_LITELLM_PARAMS,
)
]
assert events == [event]
class TestConfigManager:
"""Test that config manager routes 'bedrock' correctly."""

20
tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py
View file

@ -1,6 +1,7 @@
import asyncio
import json
import time
from typing import Final
import httpx
import pytest
@ -36,13 +37,8 @@ class _ShortTtlTokenClient:
return _JsonResponse({"access_token": f"token-{self.calls}", "expires_in": 30})
class _SSELines:
def __init__(self, lines):
self.lines = lines
async def aiter_lines(self):
for line in self.lines:
yield line
def _sse_response(lines: list[str]) -> httpx.Response:
return httpx.Response(200, stream=httpx.ByteStream("".join(f"{line}\n" for line in lines).encode()))
class _InvalidJsonStreamResponse:
@ -218,7 +214,7 @@ def test_cp4d_token_ttl_from_absolute_expiration():
@pytest.mark.asyncio
async def test_accumulate_wxo_sse_text_ignores_non_dict_json_events():
response = _SSELines(
response = _sse_response(
[
"data: null",
"data: true",
@ -230,6 +226,14 @@ async def test_accumulate_wxo_sse_text_ignores_non_dict_json_events():
)
@pytest.mark.asyncio
async def test_accumulate_wxo_sse_text_keeps_unicode_line_separators_inside_an_event() -> None:
separated_text: Final = "before\u2028between\u2029after"
response: Final = _sse_response([f"data: {json.dumps({'results': separated_text}, ensure_ascii=False)}", ""])
assert await WatsonxOrchestrateHandler._accumulate_wxo_sse_text(response) == separated_text
@pytest.mark.asyncio
async def test_short_lived_tokens_are_not_served_from_cache():
client = _ShortTtlTokenClient()

18
tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py
View file

@ -1102,6 +1102,24 @@ def test_recover_output_items_merges_text_only_items_at_distinct_indices():
assert recovered[1]["content"][0]["text"] == "From OUTPUT_TEXT_DONE only"
def test_recover_output_items_keeps_an_item_whose_text_holds_unicode_line_separators() -> None:
from litellm.completion_extras.litellm_responses_transformation.transformation import (
LiteLLMResponsesTransformationHandler,
)
item: Final = {
"type": "message",
"id": "msg_item_0",
"role": "assistant",
"status": "completed",
"content": [{"type": "output_text", "text": "From\u2028OUTPUT_ITEM\u2029DONE", "annotations": []}],
}
event: Final = {"type": "response.output_item.done", "output_index": 0, "item": item}
raw_sse: Final = f"data: {json.dumps(event, ensure_ascii=False)}\ndata: [DONE]\n"
assert LiteLLMResponsesTransformationHandler._recover_output_items_from_raw_sse(raw_sse) == [item]
def test_transform_response_prefers_completed_output_from_raw_sse():
from litellm.completion_extras.litellm_responses_transformation.transformation import (
LiteLLMResponsesTransformationHandler,

83
tests/unit/google_genai/test_google_genai_streaming_iterator.py
View file

@ -62,7 +62,9 @@ def _logging_obj(
],
)
def test_streaming_logging_targets_the_provider_that_served_the_request(
iterator_cls: type,
iterator_cls: (
type[AsyncGoogleGenAIGenerateContentStreamingIterator] | type[GoogleGenAIGenerateContentStreamingIterator]
),
custom_llm_provider: str,
expected_endpoint_type: EndpointType,
):
@ -99,18 +101,15 @@ def _large_inline_data_event() -> str:
return f"data: {json.dumps(payload)}"
def _sse_response(body: str) -> httpx.Response:
return httpx.Response(200, stream=httpx.ByteStream(body.encode()))
@pytest.mark.asyncio
async def test_async_streaming_iterator_yields_complete_sse_events():
"""Large inlineData must not be split across byte-chunk boundaries."""
mock_response = MagicMock()
async def _aiter_lines():
yield _large_inline_data_event()
mock_response.aiter_lines = _aiter_lines
iterator = AsyncGoogleGenAIGenerateContentStreamingIterator(
response=mock_response,
response=_sse_response(_large_inline_data_event() + "\n\n"),
model="gemini-3.1-flash-image-preview",
logging_obj=_logging_obj(),
generate_content_provider_config=MagicMock(),
@ -128,11 +127,8 @@ async def test_async_streaming_iterator_yields_complete_sse_events():
def test_sync_streaming_iterator_yields_complete_sse_events():
mock_response = MagicMock()
mock_response.iter_lines.return_value = iter([_large_inline_data_event()])
iterator = GoogleGenAIGenerateContentStreamingIterator(
response=mock_response,
response=_sse_response(_large_inline_data_event() + "\n\n"),
model="gemini-3.1-flash-image-preview",
logging_obj=_logging_obj(),
generate_content_provider_config=MagicMock(),
@ -150,17 +146,8 @@ def test_sync_streaming_iterator_yields_complete_sse_events():
@pytest.mark.asyncio
async def test_async_streaming_iterator_preserves_multi_field_sse_event():
mock_response = MagicMock()
async def _aiter_lines():
yield "event: message"
yield 'data: {"text":"hi"}'
yield ""
mock_response.aiter_lines = _aiter_lines
iterator = AsyncGoogleGenAIGenerateContentStreamingIterator(
response=mock_response,
response=_sse_response('event: message\ndata: {"text":"hi"}\n\n'),
model="gemini-test",
logging_obj=_logging_obj(),
generate_content_provider_config=MagicMock(),
@ -174,16 +161,8 @@ async def test_async_streaming_iterator_preserves_multi_field_sse_event():
@pytest.mark.asyncio
async def test_async_streaming_iterator_forwards_sse_comment_events():
mock_response = MagicMock()
async def _aiter_lines():
yield ": keepalive"
yield ""
mock_response.aiter_lines = _aiter_lines
iterator = AsyncGoogleGenAIGenerateContentStreamingIterator(
response=mock_response,
response=_sse_response(": keepalive\n\n"),
model="gemini-test",
logging_obj=_logging_obj(),
generate_content_provider_config=MagicMock(),
@ -195,6 +174,46 @@ async def test_async_streaming_iterator_forwards_sse_comment_events():
assert chunk == b": keepalive\n\n"
_LINE_SEPARATED_TEXT: Final = "before\u2028between\u2029after"
async def _first_chunk(
iterator: GoogleGenAIGenerateContentStreamingIterator | AsyncGoogleGenAIGenerateContentStreamingIterator,
) -> bytes:
if isinstance(iterator, AsyncGoogleGenAIGenerateContentStreamingIterator):
return await iterator.__anext__()
return next(iterator)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"iterator_cls",
[AsyncGoogleGenAIGenerateContentStreamingIterator, GoogleGenAIGenerateContentStreamingIterator],
)
async def test_streaming_iterator_keeps_unicode_line_separators_inside_the_event(
iterator_cls: (
type[AsyncGoogleGenAIGenerateContentStreamingIterator] | type[GoogleGenAIGenerateContentStreamingIterator]
),
) -> None:
payload: Final = {"candidates": [{"content": {"parts": [{"text": _LINE_SEPARATED_TEXT}]}}]}
event: Final = f"data: {json.dumps(payload, ensure_ascii=False)}\n\n"
iterator: Final = iterator_cls(
response=_sse_response(event),
model="gemini-3.8-flash",
logging_obj=_logging_obj(),
generate_content_provider_config=MagicMock(),
litellm_metadata={},
custom_llm_provider="gemini",
)
chunk: Final = await _first_chunk(iterator)
assert chunk == event.encode()
assert (
json.loads(chunk[len(b"data: ") : -2])["candidates"][0]["content"]["parts"][0]["text"] == _LINE_SEPARATED_TEXT
)
def _gemini_event(text: str, finish_reason: str | None = None) -> str:
candidate: Final = {"content": {"role": "model", "parts": [{"text": text}]}, "index": 0}
payload: Final = {

45
tests/unit/interactions/test_interactions_streaming_iterator.py
View file

@ -16,7 +16,10 @@ import pytest
import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm.interactions import streaming_iterator as interactions_streaming_iterator_module
from litellm.interactions.streaming_iterator import InteractionsAPIStreamingIterator
from litellm.interactions.streaming_iterator import (
InteractionsAPIStreamingIterator,
SyncInteractionsAPIStreamingIterator,
)
from litellm.litellm_core_utils import thread_pool_executor as thread_pool_executor_module
from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging
from litellm.llms.gemini.interactions.transformation import GoogleAIStudioInteractionsConfig
@ -148,3 +151,43 @@ async def test_stream_logs_its_first_event_time_measured_from_the_call_start():
assert first_event_at is not None
assert recorder.logged == [(call_start, first_event_at)]
_LINE_SEPARATED_TEXT: Final = "before\u2028between\u2029after"
async def _first_event(
iterator: InteractionsAPIStreamingIterator | SyncInteractionsAPIStreamingIterator,
) -> InteractionsAPIStreamingResponse:
if isinstance(iterator, InteractionsAPIStreamingIterator):
return await iterator.__anext__()
return next(iterator)
@pytest.mark.asyncio
@pytest.mark.parametrize("iterator_cls", [InteractionsAPIStreamingIterator, SyncInteractionsAPIStreamingIterator])
async def test_stream_keeps_unicode_line_separators_inside_the_delta_text(
iterator_cls: type[InteractionsAPIStreamingIterator] | type[SyncInteractionsAPIStreamingIterator],
) -> None:
delta: Final = {"type": "text", "text": _LINE_SEPARATED_TEXT}
event: Final = {"event_type": "step.delta", "id": "interaction-1", "delta": delta}
iterator: Final = iterator_cls(
response=httpx.Response(
200, stream=httpx.ByteStream(f"data: {json.dumps(event, ensure_ascii=False)}\n\n".encode())
),
model="gemini/gemini-3.1-pro-preview",
interactions_api_config=GoogleAIStudioInteractionsConfig(),
logging_obj=LitellmLogging(
model="gemini/gemini-3.1-pro-preview",
messages=[{"role": "user", "content": "hi"}],
stream=True,
call_type="ainteraction",
start_time=datetime(2026, 10, 9, 12, 0, 0),
litellm_call_id="interactions-separators",
function_id="interactions-separators",
),
)
first: Final = await _first_event(iterator)
assert first.delta == delta

114
tests/unit/litellm_core_utils/test_streaming_handler.py
View file

@ -506,6 +506,118 @@ def test_multi_chunk_reasoning_and_content(
assert initialized_custom_stream_wrapper.sent_last_thinking_block is True
def _stream_chunk(delta: dict[str, object], finish_reason: str | None = None) -> ModelResponseStream:
return ModelResponseStream(choices=[{"index": 0, "delta": delta, "finish_reason": finish_reason}])
def _content_through_first_call(deltas: Sequence[Delta]) -> str:
first_call: Final = next(
(index for index, delta in enumerate(deltas) if delta.tool_calls or delta.function_call), len(deltas)
)
return "".join(delta.content or "" for delta in deltas[: first_call + 1])
_GET_WEATHER_CALL: Final = {
"index": 0,
"id": "call_1",
"type": "function",
"function": {"name": "get_weather", "arguments": '{"city": "Paris"}'},
}
@pytest.mark.asyncio
@pytest.mark.parametrize("is_async", [True, False])
@pytest.mark.parametrize(
"chunks,expected_content,expected_tool_calls",
[
(
(
_stream_chunk({"reasoning_content": "Check the weather"}),
_stream_chunk({"tool_calls": [_GET_WEATHER_CALL]}),
_stream_chunk({}, "tool_calls"),
),
"<think>Check the weather</think>",
1,
),
(
(
_stream_chunk({"reasoning_content": "Pick a function"}),
_stream_chunk({"function_call": {"name": "get_weather", "arguments": "{}"}}),
_stream_chunk({}, "function_call"),
),
"<think>Pick a function</think>",
0,
),
(
(
_stream_chunk({"reasoning_content": "Step one"}),
_stream_chunk({"reasoning_content": ", step two"}),
_stream_chunk({}, "length"),
),
"<think>Step one, step two</think>",
0,
),
(
(
_stream_chunk({"reasoning_content": "a"}),
_stream_chunk({"content": "b"}),
_stream_chunk({"reasoning_content": "c"}),
_stream_chunk({"content": "d"}),
_stream_chunk({}, "stop"),
),
"<think>a</think>b<think>c</think>d",
0,
),
(
(_stream_chunk({"reasoning_content": "Still thinking"}),),
"<think>Still thinking</think>",
0,
),
(
(_stream_chunk({"content": "hi"}), _stream_chunk({}, "stop")),
"hi",
0,
),
],
ids=[
"reasoning_then_tool_call",
"reasoning_then_function_call",
"reasoning_cut_by_length",
"reasoning_after_answer",
"stream_ends_mid_reasoning",
"answer_without_reasoning",
],
)
async def test_merged_reasoning_stream_keeps_think_tags_balanced(
logging_obj: Logging,
is_async: bool,
chunks: tuple[ModelResponseStream, ...],
expected_content: str,
expected_tool_calls: int,
) -> None:
response: Final = CustomStreamWrapper(
completion_stream=ModelResponseListIterator(model_responses=[chunk.model_copy(deep=True) for chunk in chunks]),
model="deepseek-reasoner",
logging_obj=logging_obj,
)
response.merge_reasoning_content_in_choices = True
deltas: Final = (
tuple([chunk.choices[0].delta async for chunk in response])
if is_async
else tuple(chunk.choices[0].delta for chunk in response)
)
logged: Final = litellm.stream_chunk_builder(chunks=response.chunks)
assert "".join(delta.content or "" for delta in deltas) == expected_content
assert _content_through_first_call(deltas) == expected_content
assert sum(len(delta.tool_calls or ()) for delta in deltas) == expected_tool_calls
assert logged.choices[0].message.content == (expected_content or None)
assert any(chunk.choices[0].finish_reason for chunk in response.chunks if chunk.choices) == any(
chunk.choices[0].finish_reason for chunk in chunks
)
def test_strip_sse_data_from_chunk():
"""Test the static method that strips 'data: ' prefix from SSE chunks"""
# Test with string inputs
@ -7849,7 +7961,7 @@ def gemini_mock_post_streaming(url, **kwargs):
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.headers = {"Content-Type": "text/event-stream"}
mock_response.iter_lines = MagicMock(return_value=stream_response())
mock_response.iter_text = MagicMock(return_value=stream_response())
return mock_response

59
tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py
View file

@ -15,7 +15,7 @@ import litellm
from litellm._uuid import uuid
from litellm.constants import RESPONSE_FORMAT_TOOL_NAME
from litellm.litellm_core_utils.prompt_templates.factory import anthropic_messages_pt
from litellm.llms.anthropic.chat.handler import ModelResponseIterator, _envelope_verdict, make_call
from litellm.llms.anthropic.chat.handler import ModelResponseIterator, _envelope_verdict, make_call, make_sync_call
from litellm.llms.base_llm.base_utils import response_format_schema, unwrap_response_format_envelope
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.types.llms.openai import (
@ -27,24 +27,34 @@ from litellm.types.utils import ChatCompletionDeltaToolCall, ModelResponseStream
from tests.unit.llms.anthropic.anthropic_docs_refusal import ANTHROPIC_DOCS_REFUSAL_MESSAGE
def _text_delta_stream(text: str) -> bytes:
events: Final = (
("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}),
)
return "".join(
f"event: {name}\ndata: {json.dumps(payload, ensure_ascii=False)}\n\n" for name, payload in events
).encode()
def _streamed_text(chunks: list[object]) -> str:
return "".join(chunk.choices[0].delta.content or "" for chunk in chunks if isinstance(chunk, ModelResponseStream))
@pytest.mark.asyncio
async def test_make_call_passes_logging_obj_to_client_post():
"""make_call must pass logging_obj to client.post so track_llm_api_timing can set llm_api_duration_ms for litellm_overhead_time_ms."""
mock_client = AsyncMock()
mock_response = MagicMock()
mock_response.aiter_lines = MagicMock(
return_value=iter([b'data: {"type":"message_start"}\n', b'data: {"type":"message_delta"}\n'])
)
mock_client.post.return_value = mock_response
mock_client.post.return_value = httpx.Response(200, stream=httpx.ByteStream(_text_delta_stream("Hi there")))
logging_obj = MagicMock()
await make_call(
completion_stream, _ = await make_call(
client=mock_client,
api_base="https://api.anthropic.com/v1/messages",
headers={},
data="{}",
model="claude-3-5-haiku",
model="claude-haiku-5-5",
messages=[{"role": "user", "content": "Hi"}],
logging_obj=logging_obj,
timeout=60.0,
@ -54,6 +64,39 @@ async def test_make_call_passes_logging_obj_to_client_post():
mock_client.post.assert_called_once()
call_kwargs = mock_client.post.call_args[1]
assert call_kwargs.get("logging_obj") is logging_obj
assert _streamed_text([chunk async for chunk in completion_stream]) == "Hi there"
_LINE_SEPARATED_TEXT: Final = "before\u2028between\u2029after"
async def _streamed_chunks(sync_stream: bool, text: str) -> list[object]:
transport: Final = httpx.MockTransport(
lambda request: httpx.Response(200, stream=httpx.ByteStream(_text_delta_stream(text)))
)
call_args: Final = {
"api_base": "https://api.anthropic.com/v1/messages",
"headers": {},
"data": "{}",
"model": "claude-haiku-5-5",
"messages": [{"role": "user", "content": "Hi"}],
"logging_obj": MagicMock(),
"timeout": 60.0,
"json_mode": False,
}
if sync_stream:
sync_iterator, _ = make_sync_call(client=HTTPHandler(client=httpx.Client(transport=transport)), **call_args)
return list(sync_iterator)
async_iterator, _ = await make_call(client=AsyncHTTPHandler(transport=transport), **call_args)
return [chunk async for chunk in async_iterator]
@pytest.mark.asyncio
@pytest.mark.parametrize("sync_stream", [True, False])
async def test_streamed_text_keeps_unicode_line_separators(sync_stream: bool) -> None:
chunks: Final = await _streamed_chunks(sync_stream, _LINE_SEPARATED_TEXT)
assert _streamed_text(chunks) == _LINE_SEPARATED_TEXT
def test_anthropic_completion_does_not_send_deployment_default_limits():

24
tests/unit/llms/anthropic/messages/test_stream_assembly.py
View file

@ -144,6 +144,30 @@ def test_anthropic_message_delta_fields_reads_the_fields_the_openai_shape_drops(
assert dict(anthropic_message_delta_fields(GOOGLE_STREAM)) == {}
def test_anthropic_message_delta_fields_reads_a_message_delta_whose_json_holds_unicode_line_separators() -> None:
safeguard_results: Final = [
{
"type": "dangerous_tool_use",
"status": {
"type": "available",
"tool_uses": {
"toolu_01": {"type": "evaluated", "outcome": "flagged", "explanation": "deletes\u2028the\u2029root"}
},
},
}
]
message_delta: Final = {
"type": "message_delta",
"delta": {"stop_reason": "end_turn", "safeguard_results": safeguard_results},
"usage": {"output_tokens": 2},
}
raw_message_delta_frame: Final = f"event: message_delta\ndata: {json.dumps(message_delta, ensure_ascii=False)}\n\n"
fields: Final = anthropic_message_delta_fields((*ANTHROPIC_STREAM[:5], raw_message_delta_frame.encode()))
assert dict(fields) == {"safeguard_results": safeguard_results}
@pytest.mark.parametrize(
"chunks",
[GOOGLE_STREAM, INVALID_UTF8, (ANTHROPIC_ERROR,), ()],

23
tests/unit/llms/azure_ai/agents/test_handler.py
View file

@ -173,7 +173,7 @@ async def test_handler_acompletion_with_thread_id_reuses_that_thread() -> None:
def _sse(event: str, data: dict[str, object] | str) -> str:
return f"event: {event}\ndata: {data if isinstance(data, str) else json.dumps(data)}\n\n"
return f"event: {event}\ndata: {data if isinstance(data, str) else json.dumps(data, ensure_ascii=False)}\n\n"
def _message_delta(text: str) -> dict[str, object]:
@ -222,3 +222,24 @@ async def test_acompletion_stream_yields_agent_message_deltas_then_stop(respx_mo
"thread": {"messages": [{"role": "user", "content": "Hi Agent, what is 10 + 5?"}]},
}
assert chunks == [("10 + 5", None), (" = 15", None), (None, "stop")]
@pytest.mark.asyncio
async def test_acompletion_stream_keeps_unicode_line_separators_inside_a_delta(respx_mock: respx.MockRouter) -> None:
separated_text: Final = "before\u2028between\u2029after"
_ = respx_mock.post(f"{API_BASE}/threads/runs").respond(
200,
text=_sse("thread.message.delta", _message_delta(separated_text)) + _sse("done", "[DONE]"),
headers={"content-type": "text/event-stream"},
)
response: Final = await litellm.acompletion(
model=f"azure_ai/agents/{AGENT_ID}",
messages=[{"role": "user", "content": "Say it with separators"}],
api_base=API_BASE,
api_key="azure-ad-token",
stream=True,
)
chunks: Final = [(chunk.choices[0].delta.content, chunk.choices[0].finish_reason) async for chunk in response]
assert chunks == [(separated_text, None), (None, "stop")]

20
tests/unit/llms/azure_ai/agents/test_transformation.py
View file

@ -1,7 +1,7 @@
import json
from typing import Final
from unittest.mock import MagicMock
import httpx
import pytest
from litellm.llms.azure_ai.agents.transformation import AzureAIAgentsConfig
@ -428,15 +428,10 @@ async def test_azure_ai_agents_streaming_accumulates_annotations_from_multiple_t
"data: [DONE]",
]
async def mock_aiter_lines():
for line in sse_lines:
yield line
mock_response = MagicMock()
mock_response.aiter_lines = MagicMock(return_value=mock_aiter_lines())
sse_response = httpx.Response(200, stream=httpx.ByteStream("\n".join(sse_lines).encode()))
chunks = []
async for chunk in handler._process_sse_stream(mock_response, "azure_ai/agents/asst_123"):
async for chunk in handler._process_sse_stream(sse_response, "azure_ai/agents/asst_123"):
chunks.append(chunk)
final_chunk = chunks[-1]
@ -498,15 +493,10 @@ async def test_azure_ai_agents_streaming_annotations_from_completed_message():
"data: [DONE]",
]
async def mock_aiter_lines():
for line in sse_lines:
yield line
mock_response = MagicMock()
mock_response.aiter_lines = MagicMock(return_value=mock_aiter_lines())
sse_response = httpx.Response(200, stream=httpx.ByteStream("\n".join(sse_lines).encode()))
chunks = []
async for chunk in handler._process_sse_stream(mock_response, "azure_ai/agents/asst_123"):
async for chunk in handler._process_sse_stream(sse_response, "azure_ai/agents/asst_123"):
chunks.append(chunk)
assert len(chunks) >= 1

17
tests/unit/llms/base_llm/guardrail_translation/test_utils.py
View file

@ -1,9 +1,11 @@
import json
from typing import Final
import pytest
import litellm
from litellm.llms.base_llm.guardrail_translation.utils import (
blocked_response_usage,
effective_skip_system_message_for_guardrail,
effective_skip_tool_message_for_guardrail,
)
@ -32,3 +34,18 @@ def test_a_per_guardrail_skip_flag_wins_over_the_global_setting(
assert effective_skip_system_message_for_guardrail(guardrail) is expected
assert effective_skip_tool_message_for_guardrail(guardrail) is expected
def test_blocked_response_usage_reads_a_message_delta_whose_json_holds_unicode_line_separators() -> None:
message_start: Final = {"type": "message_start", "message": {"usage": {"input_tokens": 11, "output_tokens": 1}}}
message_delta: Final = {
"type": "message_delta",
"delta": {"stop_reason": "stop_sequence", "stop_sequence": "END\u2028OF\u2029TURN"},
"usage": {"output_tokens": 7},
}
stream: Final = "".join(
f"event: {event['type']}\ndata: {json.dumps(event, ensure_ascii=False)}\n\n"
for event in (message_start, message_delta)
).encode()
assert blocked_response_usage([stream]) == {"input_tokens": 11, "output_tokens": 7}

14
tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py
View file

@ -4458,6 +4458,20 @@ class TestMalformedBatchJsonl:
_VALID_BATCH_RECORD: Final = json.loads(_VALID_BATCH_LINE)
def test_file_content_conversion_keeps_a_line_whose_json_holds_unicode_line_separators() -> None:
text: Final = "Hello\u2028from\u2029batch"
body: Final = {**_VALID_BATCH_RECORD["body"], "messages": [{"role": "user", "content": text}]}
record: Final = {**_VALID_BATCH_RECORD, "body": body}
file_content: Final = f"{json.dumps(record, ensure_ascii=False)}\n".encode()
bedrock_jsonl, _ = BedrockJsonlFilesTransformation().transform_openai_file_content_to_bedrock_file_content(
file_content
)
(bedrock_record,) = (json.loads(line) for line in bedrock_jsonl.split("\n"))
assert bedrock_record["modelInput"]["messages"] == [{"role": "user", "content": [{"text": text}]}]
def _batch_line(body: object) -> str:
return json.dumps({**_VALID_BATCH_RECORD, "body": body})

14
tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py
View file

@ -18,7 +18,7 @@ import types
from collections.abc import Callable, Mapping, Sequence
from datetime import datetime
from itertools import chain
from typing import ClassVar, Literal, TypeAlias
from typing import ClassVar, Final, Literal, TypeAlias
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
@ -60,6 +60,7 @@ from litellm.llms.bedrock.passthrough.guardrail_translation.handler import (
BedrockPassthroughGuardrailHandler,
_extract_converse_texts,
_is_converse_endpoint,
_sse_events,
_write_back_texts,
)
@ -3840,3 +3841,14 @@ class TestInvokeStreamGuardrailRelay:
assert received_key is not None and received_key.texts == ("In Paris",)
assert rewritten_key == received_key
def test_sse_events_keeps_a_data_line_whose_json_holds_unicode_line_separators() -> None:
event: Final = {
"type": "content_block_delta",
"index": 0,
"delta": {"type": "text_delta", "text": "Card\u2028on\u2029file"},
}
sse: Final = f"event: content_block_delta\ndata: {json.dumps(event, ensure_ascii=False)}\n\n".encode()
assert _sse_events(sse) == (event,)

15
tests/unit/llms/bedrock/search/test_agentcore_search_transformation.py
View file

@ -8,6 +8,7 @@ the sharded CI (coverage collection runs against this tree).
import json
import os
from typing import Final
import httpx
import pytest
@ -682,3 +683,17 @@ def test_transform_search_response_ignores_text_blocks_without_result_objects(bl
body = {"result": {"content": [{"type": "text", "text": block_text}]}}
assert _transform_text(json.dumps(body)).results == []
def test_transform_search_response_keeps_a_result_whose_sse_json_holds_unicode_line_separators() -> None:
snippet: Final = "Snippet\u2028for\u2029result"
results: Final = [{"title": "Result", "url": "https://example.com/1", "text": snippet}]
body: Final = {
"jsonrpc": "2.0",
"id": 1,
"result": {"content": [{"type": "text", "text": json.dumps(results, ensure_ascii=False)}]},
}
response: Final = _transform_text(f"event: message\ndata: {json.dumps(body, ensure_ascii=False)}\n\n")
assert _as_tuples(response) == [("Result", "https://example.com/1", snippet, None, None)]

31
tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py
View file

@ -587,6 +587,37 @@ class TestChatGPTResponsesMultiAccount:
assert headers_default["ChatGPT-Account-Id"] == "acct-default"
def test_chatgpt_non_stream_sse_keeps_an_output_item_whose_text_holds_unicode_line_separators() -> None:
text: Final = "Hello\u2028from\u2029stream"
output_item_done: Final = {
"type": "response.output_item.done",
"output_index": 0,
"item": {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": text}]},
}
completed: Final = {
"type": "response.completed",
"response": {
"id": "resp_test",
"object": "response",
"created_at": 1700000000,
"status": "completed",
"model": "gpt-6.1-sol",
"output": [],
},
}
sse_body: Final = "".join(
f"data: {json.dumps(event, ensure_ascii=False)}\n" for event in (output_item_done, completed)
)
parsed: Final = ChatGPTResponsesAPIConfig().transform_response_api_response(
model="chatgpt/gpt-6.1-sol",
raw_response=httpx.Response(200, headers={"content-type": "text/event-stream"}, text=sse_body),
logging_obj=MagicMock(),
)
assert parsed.output_text == text
@pytest.mark.parametrize("call", ["responses", "completion"])
def test_prompt_cache_key_keeps_one_session_header_across_requests(
call: str, tmp_path: Path, monkeypatch: pytest.MonkeyPatch

47
tests/unit/llms/cloudflare/test_cloudflare_transformation.py
View file

@ -1,6 +1,6 @@
import asyncio
import json
from collections.abc import Callable, Iterator
from collections.abc import AsyncIterator, Callable, Iterator
from typing import Final
from unittest.mock import MagicMock
@ -380,6 +380,35 @@ def test_completion_cloudflare_stream(sync_mode: bool, _cloudflare_httpx_transpo
assert "language" in content.lower()
@pytest.mark.asyncio
@pytest.mark.parametrize("sync_mode", [True, False], ids=["sync", "async"])
async def test_cloudflare_stream_keeps_unicode_line_separators_inside_content(
sync_mode: bool, _cloudflare_httpx_transport: None
) -> None:
streamed_text: Final = ("alpha\u2028", "beta\u2029", "gamma")
base: Final = {"id": "chatcmpl-cf", "object": "chat.completion.chunk", "created": 1, "model": "@cf/glm"}
body: Final = "".join(
f"data: {json.dumps({**base, 'choices': [{'index': 0, 'delta': {'content': text}}]}, ensure_ascii=False)}\n\n"
for text in streamed_text
)
with respx.mock(assert_all_called=True) as api:
api.post(f"{FAKE_API_BASE}/chat/completions").respond(
200, headers={"content-type": "text/event-stream"}, content=(body + "data: [DONE]\n\n").encode()
)
request: Final = {
"model": "cloudflare/@cf/glm",
"messages": [{"role": "user", "content": "repeat"}],
"stream": True,
"api_base": FAKE_API_BASE,
"api_key": FAKE_API_KEY,
}
chunks: Final = (
tuple(completion(**request)) if sync_mode else tuple([c async for c in await acompletion(**request)])
)
assert "".join(c.choices[0].delta.content or "" for c in chunks) == "".join(streamed_text)
def _make_mock_response(json_data: Dict[str, Any]) -> MagicMock:
mock = MagicMock(spec=httpx.Response)
mock.status_code = 200
@ -397,13 +426,13 @@ def test_acompletion_cloudflare_stream(sync_mode, monkeypatch):
if sync_mode:
def _iter_lines():
def _iter_text() -> Iterator[str]:
for chunk in raw_chunks:
yield f"data: {chunk}"
yield "data: [DONE]"
yield f"data: {chunk}\n\n"
yield "data: [DONE]\n\n"
mock_resp = MagicMock()
mock_resp.iter_lines.return_value = _iter_lines()
mock_resp.iter_text.return_value = _iter_text()
mock_resp.status_code = 200
mock_resp.headers = {"content-type": "text/event-stream"}
@ -420,13 +449,13 @@ def test_acompletion_cloudflare_stream(sync_mode, monkeypatch):
mock_post.assert_called_once()
else:
async def _aiter_lines():
async def _aiter_text() -> AsyncIterator[str]:
for chunk in raw_chunks:
yield f"data: {chunk}"
yield "data: [DONE]"
yield f"data: {chunk}\n\n"
yield "data: [DONE]\n\n"
mock_resp = MagicMock()
mock_resp.aiter_lines.return_value = _aiter_lines()
mock_resp.aiter_text.return_value = _aiter_text()
mock_resp.status_code = 200
mock_resp.headers = {"content-type": "text/event-stream"}

0
tests/unit/llms/codestral/__init__.py Normal file
View file

0
tests/unit/llms/codestral/completion/__init__.py Normal file
View file

50
tests/unit/llms/codestral/completion/test_codestral_completion_handler.py Normal file
View file

@ -0,0 +1,50 @@
import json
from typing import Final
import httpx
import pytest
import respx
import litellm
_LINE_SEPARATED_TEXT: Final = "before\u2028between\u2029after"
def _fim_stream_body(text: str) -> bytes:
chunks: Final = (
{"id": "c1", "object": "chat.completion.chunk", "created": 1, "model": "codestral-latest",
"choices": [{"index": 0, "delta": {"role": "assistant", "content": text}, "finish_reason": None}]},
{"id": "c1", "object": "chat.completion.chunk", "created": 1, "model": "codestral-latest",
"choices": [{"index": 0, "delta": {"content": ""}, "finish_reason": "stop"}]},
)
events: Final = "".join(f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n" for chunk in chunks)
return f"{events}data: [DONE]\n\n".encode()
async def _streamed_texts(sync_stream: bool) -> list[str | None]:
request: Final = {
"model": "text-completion-codestral/codestral-latest",
"prompt": "def f(",
"stream": True,
"api_key": "codestral-key",
"max_retries": 0,
}
if sync_stream:
return [chunk.choices[0].text for chunk in litellm.text_completion(**request)]
return [chunk.choices[0].text async for chunk in await litellm.atext_completion(**request)]
@pytest.mark.asyncio
@pytest.mark.parametrize("sync_stream", [True, False])
async def test_fim_stream_keeps_unicode_line_separators_inside_a_delta(
sync_stream: bool, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
route: Final = respx_mock.post("https://codestral.mistral.ai/v1/fim/completions").mock(
return_value=httpx.Response(200, content=_fim_stream_body(_LINE_SEPARATED_TEXT))
)
texts: Final = await _streamed_texts(sync_stream)
assert texts == [_LINE_SEPARATED_TEXT, None]
assert json.loads(route.calls.last.request.content)["stream"] is True

37
tests/unit/llms/custom_httpx/test_aiohttp_handler.py
View file

@ -1,11 +1,17 @@
import json
from typing import Final
from unittest.mock import AsyncMock, Mock, patch
import aiohttp
import httpx
import pytest
import litellm
from litellm.llms.aiohttp_openai.chat.transformation import AiohttpOpenAIChatConfig
from litellm.llms.custom_httpx.aiohttp_handler import BaseLLMAIOHTTPHandler
from litellm.llms.custom_httpx.aiohttp_transport import LiteLLMAiohttpTransport
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.types.utils import ModelResponseStream
class TestBaseLLMAIOHTTPHandler:
@ -445,3 +451,34 @@ class TestBaseLLMAIOHTTPHandler:
# Should use transport, not connector
mock_transport.get_valid_client_session.assert_called_once()
assert result is mock_session_from_transport
def test_make_sync_call_stream_keeps_unicode_line_separators_inside_a_delta() -> None:
text: Final = "before\u2028between\u2029after"
chunk: Final = {
"id": "c1",
"object": "chat.completion.chunk",
"created": 1,
"model": "gpt-5.5",
"choices": [{"index": 0, "delta": {"content": text}, "finish_reason": None}],
}
body: Final = f"data: {json.dumps(chunk, ensure_ascii=False)}\n\ndata: [DONE]\n\n".encode()
transport: Final = httpx.MockTransport(lambda request: httpx.Response(200, stream=httpx.ByteStream(body)))
completion_stream, _ = BaseLLMAIOHTTPHandler().make_sync_call(
provider_config=AiohttpOpenAIChatConfig(),
api_base="https://example.com/v1/chat/completions",
headers={},
data={"model": "gpt-5.5", "messages": []},
model="gpt-5.5",
messages=[],
logging_obj=Mock(),
litellm_params={},
timeout=5.0,
client=HTTPHandler(client=httpx.Client(transport=transport)),
)
contents: Final = [
chunk.choices[0].delta.content for chunk in completion_stream if isinstance(chunk, ModelResponseStream)
]
assert contents == [text]

64
tests/unit/llms/custom_httpx/test_sse_lines.py Normal file
View file

@ -0,0 +1,64 @@
from collections.abc import AsyncIterator, Sequence
from typing import Final
import pytest
from litellm.llms.custom_httpx.sse_lines import (
aiter_sse_lines,
iter_sse_lines,
split_sse_lines,
split_sse_lines_keepends,
)
CASES: Final = (
pytest.param(("data: a\u2028b\r\n\r\n",), ("data: a\u2028b", ""), id="u2028-stays-inside-the-line"),
pytest.param(
("data: \u0085\u2029\x0b\x0c\x1c\x1d\x1e\n",),
("data: \u0085\u2029\x0b\x0c\x1c\x1d\x1e",),
id="no-splitlines-only-separator-breaks-a-line",
),
pytest.param(("a\r\nb\nc\rd\n",), ("a", "b", "c", "d"), id="crlf-lf-and-cr-each-end-a-line"),
pytest.param(("a\r", "\nb\n"), ("a", "b"), id="crlf-split-across-chunks-is-one-break"),
pytest.param(("a\r", "\n\r\n"), ("a", ""), id="crlf-split-then-blank-line"),
pytest.param(("a\r", "", "\nb\n"), ("a", "b"), id="empty-chunk-between-cr-and-lf"),
pytest.param(("a\r", "\r"), ("a", ""), id="cr-then-cr-is-two-breaks"),
pytest.param(("a\r\r\nb\n",), ("a", "", "b"), id="cr-then-crlf-in-one-chunk-is-two-breaks"),
pytest.param(("a\r\nb\r\n\r\n",), ("a", "b", ""), id="crlf-only-chunk"),
pytest.param(("da", "ta: x", "y\n"), ("data: xy",), id="line-spanning-chunks-is-joined"),
pytest.param(("data: x",), ("data: x",), id="unterminated-last-line-is-kept"),
pytest.param(("",), (), id="empty-body"),
)
async def _async_chunks(chunks: Sequence[str]) -> AsyncIterator[str]:
for chunk in chunks:
yield chunk
@pytest.mark.parametrize(("chunks", "expected"), CASES)
def test_iter_sse_lines_splits_only_on_cr_lf_and_crlf(chunks: Sequence[str], expected: Sequence[str]) -> None:
assert tuple(iter_sse_lines(chunks)) == tuple(expected)
@pytest.mark.parametrize(("chunks", "expected"), CASES)
async def test_aiter_sse_lines_splits_only_on_cr_lf_and_crlf(chunks: Sequence[str], expected: Sequence[str]) -> None:
assert tuple([line async for line in aiter_sse_lines(_async_chunks(chunks))]) == tuple(expected)
@pytest.mark.parametrize(("chunks", "expected"), CASES)
def test_split_sse_lines_splits_a_whole_body_only_on_cr_lf_and_crlf(
chunks: Sequence[str], expected: Sequence[str]
) -> None:
assert split_sse_lines("".join(chunks)) == tuple(expected)
@pytest.mark.parametrize(
("text", "expected"),
(
pytest.param("a\r\nb\rc\nd", ("a\r\n", "b\r", "c\n", "d"), id="each-break-stays-on-its-line"),
pytest.param("data: a\u2028b\u2029c\n\n", ("data: a\u2028b\u2029c\n", "\n"), id="u2028-and-u2029-stay-inside"),
pytest.param("", (), id="empty-body"),
),
)
def test_split_sse_lines_keepends_splits_only_on_cr_lf_and_crlf(text: str, expected: Sequence[str]) -> None:
assert split_sse_lines_keepends(text) == tuple(expected)

14
tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py
View file

@ -1,6 +1,6 @@
import copy
import json
from collections.abc import Iterator
from collections.abc import AsyncIterator, Iterator
from typing import Final
from unittest.mock import MagicMock, patch
@ -2160,13 +2160,13 @@ def mock_chat_response_nonanthropic_prompt_caching() -> Dict[str, Any]:
def mock_http_handler_chat_streaming_response() -> MagicMock:
mock_stream_chunks = mock_chat_streaming_response_chunks()
def mock_iter_lines():
def mock_iter_text() -> Iterator[str]:
for chunk in mock_stream_chunks:
for line in chunk.splitlines():
yield line
yield f"{line}\n"
mock_response = MagicMock()
mock_response.iter_lines.side_effect = mock_iter_lines
mock_response.iter_text.side_effect = mock_iter_text
mock_response.status_code = 200
return mock_response
@ -2175,13 +2175,13 @@ def mock_http_handler_chat_streaming_response() -> MagicMock:
def mock_http_handler_chat_async_streaming_response() -> MagicMock:
mock_stream_chunks = mock_chat_streaming_response_chunks()
async def mock_iter_lines():
async def mock_aiter_text() -> AsyncIterator[str]:
for chunk in mock_stream_chunks:
for line in chunk.splitlines():
yield line
yield f"{line}\n"
mock_response = MagicMock()
mock_response.aiter_lines.return_value = mock_iter_lines()
mock_response.aiter_text.return_value = mock_aiter_text()
mock_response.status_code = 200
return mock_response

6
tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py
View file

@ -1329,7 +1329,7 @@ def test_streaming_surfaces_fireworks_response_fields():
raw_response = MagicMock()
raw_response.status_code = 200
raw_response.headers = {}
raw_response.iter_lines = lambda: iter(sse_lines)
raw_response.iter_text = lambda: iter(f"{line}\n" for line in sse_lines)
client = HTTPHandler()
with patch.object(client, "post", return_value=raw_response):
@ -1788,7 +1788,7 @@ def test_streaming_preserves_selected_model_for_private_accounting():
raw_response = MagicMock()
raw_response.status_code = 200
raw_response.headers = {}
raw_response.iter_lines = lambda: iter(sse_lines)
raw_response.iter_text = lambda: iter(f"{line}\n" for line in sse_lines)
client = HTTPHandler()
with patch.object(client, "post", return_value=raw_response):
@ -2275,7 +2275,7 @@ def test_completion_forwards_litellm_user_id_as_user(call_kwargs: dict[str, obje
def test_completion_forwards_litellm_user_id_when_streaming() -> None:
client: Final = _fireworks_chat_client()
client.post.return_value.iter_lines = lambda: iter(())
client.post.return_value.iter_text = lambda: iter(())
litellm.completion(
model="fireworks_ai/accounts/fireworks/models/kimi-k3",
messages=[{"role": "user", "content": "hi"}],

16
tests/unit/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py
View file

@ -4,6 +4,7 @@ Tests for the Google GenAI generateContent guardrail translation handler.
import json
from types import SimpleNamespace
from typing import Final
from unittest.mock import AsyncMock, MagicMock
import pytest
@ -190,6 +191,21 @@ async def test_streaming_raw_sse_chunks_accumulate_text_across_split_frames():
assert guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"] == ["harmful answer"]
@pytest.mark.asyncio
async def test_streaming_raw_sse_text_holding_unicode_line_separators_reaches_the_guardrail_whole() -> None:
guardrail: Final = _mock_guardrail(["clean"])
text: Final = "harmful\u2028answer\u2029here"
payload: Final = {"candidates": [{"content": {"parts": [{"text": text}]}}]}
chunk: Final = f"data: {json.dumps(payload, ensure_ascii=False)}\r\n\r\n"
await GoogleGenAIGenerateContentHandler().process_output_streaming_response(
responses_so_far=[chunk],
guardrail_to_apply=guardrail,
)
assert guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"] == [text]
@pytest.mark.asyncio
async def test_streaming_blocking_guardrail_exception_propagates():
handler = GoogleGenAIGenerateContentHandler()

32
tests/unit/llms/langgraph/chat/test_langgraph_sse_iterator.py Normal file
View file

@ -0,0 +1,32 @@
import json
from typing import Final
import httpx
import pytest
from litellm.llms.langgraph.chat.sse_iterator import LangGraphSSEStreamIterator
_LINE_SEPARATED_TEXT: Final = "before\u2028between\u2029after"
def _messages_event_stream(text: str) -> httpx.Response:
event: Final = ["messages", [[{"type": "AIMessageChunk", "content": text}, {}]]]
return httpx.Response(200, stream=httpx.ByteStream(f"data: {json.dumps(event, ensure_ascii=False)}\n\n".encode()))
async def _streamed_contents(iterator: LangGraphSSEStreamIterator, sync_stream: bool) -> list[str | None]:
if sync_stream:
return [chunk.choices[0].delta.content for chunk in iterator]
return [chunk.choices[0].delta.content async for chunk in iterator]
@pytest.mark.asyncio
@pytest.mark.parametrize("sync_stream", [True, False])
async def test_stream_keeps_unicode_line_separators_inside_a_message_chunk(sync_stream: bool) -> None:
iterator: Final = LangGraphSSEStreamIterator(
response=_messages_event_stream(_LINE_SEPARATED_TEXT), model="langgraph/agent"
)
contents: Final = await _streamed_contents(iterator, sync_stream)
assert contents == [_LINE_SEPARATED_TEXT, None]

29
tests/unit/llms/oobabooga/chat/test_oobabooga.py
View file

@ -1,7 +1,12 @@
import json
from typing import Final
from unittest.mock import MagicMock, patch
import httpx
import respx
import litellm
from litellm.llms.oobabooga.chat import oobabooga
MOCK_COMPLETION_RESPONSE = {
"choices": [{"message": {"role": "assistant", "content": "hi there"}}],
@ -50,3 +55,27 @@ def test_url_valued_model_still_targets_that_url():
mock_post.assert_called_once()
called_url = mock_post.call_args[0][0]
assert called_url == "https://sdk-user.example/v1/chat/completions"
def test_stream_keeps_unicode_line_separators_inside_a_data_line(respx_mock: respx.MockRouter) -> None:
text: Final = "before\u2028between\u2029after"
chunk: Final = {"choices": [{"index": 0, "delta": {"content": text}, "finish_reason": None}]}
data_line: Final = f"data: {json.dumps(chunk, ensure_ascii=False)}"
respx_mock.post("https://ooba.example/v1/chat/completions").mock(
return_value=httpx.Response(200, content=f"{data_line}\n\ndata: [DONE]\n\n".encode())
)
lines: Final = oobabooga.completion(
model="my-model",
messages=[{"role": "user", "content": "hello"}],
api_base="https://ooba.example",
model_response=litellm.ModelResponse(),
print_verbose=print,
encoding=None,
api_key=None,
logging_obj=MagicMock(),
optional_params={"stream": True},
litellm_params={},
)
assert [line for line in lines if line] == [data_line, "data: [DONE]"]

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

@ -4,9 +4,10 @@ from typing import Final
from unittest.mock import Mock
import httpx
import pytest
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.llms.openai_like.chat.handler import make_sync_call
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.openai_like.chat.handler import make_call, make_sync_call
_FAKE_STREAMED_REPLY: Final = {
"id": "chatcmpl-fake-stream",
@ -73,3 +74,35 @@ def test_fake_stream_keeps_reasoning_every_tool_call_and_token_detail():
assert [call.function.name for call in delta.tool_calls] == ["get_weather", "get_time"]
assert chunks[0].usage.prompt_tokens_details.cached_tokens == 64
assert chunks[0].usage.completion_tokens_details.reasoning_tokens == 1920
_LINE_SEPARATED_TEXT: Final = "before\u2028between\u2029after"
def _streamed_reply(request: httpx.Request) -> httpx.Response:
chunk: Final = {"id": "chatcmpl-separators", "choices": [{"index": 0, "delta": {"content": _LINE_SEPARATED_TEXT}}]}
body: Final = f"data: {json.dumps(chunk, ensure_ascii=False)}\n\ndata: [DONE]\n\n"
return httpx.Response(200, stream=httpx.ByteStream(body.encode()))
async def _streamed_chunks(sync_stream: bool) -> list[dict[str, object]]:
call_args: Final = {
"api_base": "https://openai-like.example/v1/chat/completions",
"headers": {},
"data": json.dumps({"model": "openai/gpt-oss-120b", "messages": [], "stream": True}),
"model": "openai/gpt-oss-120b",
"messages": [],
"logging_obj": Mock(),
}
transport: Final = httpx.MockTransport(_streamed_reply)
if sync_stream:
return list(make_sync_call(client=HTTPHandler(client=httpx.Client(transport=transport)), **call_args))
return [chunk async for chunk in await make_call(client=AsyncHTTPHandler(transport=transport), **call_args)]
@pytest.mark.asyncio
@pytest.mark.parametrize("sync_stream", [True, False])
async def test_stream_keeps_unicode_line_separators_inside_the_delta(sync_stream: bool) -> None:
chunks: Final = await _streamed_chunks(sync_stream)
assert "".join(str(chunk["text"]) for chunk in chunks) == _LINE_SEPARATED_TEXT

30
tests/unit/llms/snowflake/chat/test_snowflake_chat_transformation.py
View file

@ -902,20 +902,14 @@ class TestSnowflakeChatCompletion:
@pytest.mark.parametrize("sync_mode", [True, False])
def test_chat_completion_snowflake_stream(self, sync_mode):
raw_chunks = _streaming_chunks()
sse_body = "".join(f"data: {chunk}\n\n" for chunk in _streaming_chunks()) + "data: [DONE]\n\n"
mock_resp = httpx.Response(
200,
headers={"content-type": "text/event-stream"},
stream=httpx.ByteStream(sse_body.encode()),
)
if sync_mode:
def _iter_lines():
for chunk in raw_chunks:
yield f"data: {chunk}"
yield "data: [DONE]"
mock_resp = MagicMock()
mock_resp.iter_lines.return_value = _iter_lines()
mock_resp.status_code = 200
mock_resp.headers = {"content-type": "text/event-stream"}
with patch.object(HTTPHandler, "post", return_value=mock_resp) as mock_post:
response = completion(
model="snowflake/mistral-7b",
@ -930,16 +924,6 @@ class TestSnowflakeChatCompletion:
mock_post.assert_called_once()
else:
async def _aiter_lines():
for chunk in raw_chunks:
yield f"data: {chunk}"
yield "data: [DONE]"
mock_resp = MagicMock()
mock_resp.aiter_lines.return_value = _aiter_lines()
mock_resp.status_code = 200
mock_resp.headers = {"content-type": "text/event-stream"}
async def _run():
with patch.object(
AsyncHTTPHandler,
@ -964,8 +948,8 @@ class TestSnowflakeChatCompletion:
chunks_received = asyncio.run(_run())
assert len(chunks_received) > 0
content = "".join(c.choices[0].delta.content for c in chunks_received if c.choices[0].delta.content)
assert content == "The sky is blue"
async def test_snowflake_claude_async_completion_inlines_remote_images_off_the_event_loop(async_only_image_fetch):

47
tests/unit/llms/test_predibase_transformation.py
View file

@ -1,3 +1,4 @@
import json
from typing import Final
from unittest.mock import AsyncMock, Mock
@ -639,3 +640,49 @@ def test_predibase_completion_passes_existing_config_to_async_completion(monkeyp
assert result == "async-result"
assert captured["async_kwargs"]["predibase_config"] is captured["config_instance"]
_LINE_SEPARATED_TEXT: Final = "before\u2028between\u2029after"
def _generate_stream_body(text: str) -> bytes:
events: Final = (
{"token": {"text": text}},
{"token": {"text": ""}, "generated_text": text, "details": {"finish_reason": "eos_token"}},
)
return "".join(f"data:{json.dumps(event, ensure_ascii=False)}\n\n" for event in events).encode()
async def _streamed_contents(sync_stream: bool) -> list[str | None]:
request: Final = {
"model": "predibase/llama-3-8b-instruct",
"messages": [{"role": "user", "content": "hi"}],
"stream": True,
"predibase_tenant_id": "tenant-1",
"api_key": "pb-key",
"max_retries": 0,
}
if sync_stream:
return [chunk.choices[0].delta.content for chunk in litellm.completion(**request)]
return [chunk.choices[0].delta.content async for chunk in await litellm.acompletion(**request)]
@pytest.mark.asyncio
@pytest.mark.parametrize("sync_stream", [True, False])
async def test_predibase_stream_keeps_unicode_line_separators_inside_a_token(
sync_stream: bool, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_hf_tokenizer_download", True)
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setattr(
litellm,
"custom_prompt_dict",
{"llama-3-8b-instruct": {"roles": {}, "initial_prompt_value": "", "final_prompt_value": ""}},
)
respx_mock.post(
"https://serving.app.predibase.com/tenant-1/deployments/v2/llms/llama-3-8b-instruct/generate_stream"
).mock(return_value=httpx.Response(200, content=_generate_stream_body(_LINE_SEPARATED_TEXT)))
contents: Final = await _streamed_contents(sync_stream)
assert contents == [_LINE_SEPARATED_TEXT, None]

11
tests/unit/llms/triton/test_triton.py
View file

@ -1,5 +1,7 @@
import json
import traceback
from collections.abc import Iterator
from typing import Final
from unittest.mock import MagicMock, patch
import httpx
@ -210,8 +212,8 @@ def test_completion_triton_generate_api(stream, monkeypatch, respx_mock):
mock_response = MagicMock()
if stream:
def mock_iter_lines():
mock_output = "".join(
def mock_iter_text() -> Iterator[str]:
mock_output: Final = "".join(
[
'data: {"model_name":"ensemble","model_version":"1","sequence_end":false,"sequence_id":0,"sequence_start":false,"text_output":"'
+ t
@ -219,10 +221,9 @@ def test_completion_triton_generate_api(stream, monkeypatch, respx_mock):
for t in ["I", " am", " an", " AI", " assistant"]
]
)
for out in mock_output.split("\n"):
yield out
yield mock_output
mock_response.iter_lines = mock_iter_lines
mock_response.iter_text = mock_iter_text
else:
def return_val():

47
tests/unit/llms/vertex_ai/agent_engine/test_transformation.py
View file

@ -14,7 +14,7 @@ import pytest
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler, get_async_httpx_client
from litellm.llms.vertex_ai.agent_engine.sse_iterator import (
VertexAgentEngineResponseIterator,
)
@ -170,3 +170,48 @@ async def test_async_stream_wrapper_without_a_client_posts_through_the_cached_ve
assert type(wrapper) is CustomStreamWrapper
assert [str(request.url) for request in sent_requests] == [api_base]
assert json.loads(sent_requests[0].content) == {"class_method": "stream_query"}
_LINE_SEPARATED_TEXT: Final = "before\u2028between\u2029after"
def _separated_text_reply(request: httpx.Request) -> httpx.Response:
chunk: Final = {"content": {"parts": [{"text": _LINE_SEPARATED_TEXT}], "role": "model"}, "finish_reason": "STOP"}
return httpx.Response(200, stream=httpx.ByteStream(f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n".encode()))
async def _streamed_text(sync_stream: bool) -> str:
messages: Final = [{"role": "user", "content": "hi"}]
call_args: Final = {
"model": "agent_engine/123",
"custom_llm_provider": "vertex_ai",
"logging_obj": Logging(
model="agent_engine/123",
messages=messages,
stream=True,
call_type="completion" if sync_stream else "acompletion",
start_time=datetime(2026, 1, 1),
litellm_call_id="call-separators",
function_id="fn-separators",
),
"api_base": "https://us-central1-aiplatform.googleapis.com/v1/reasoningEngines/123:streamQuery",
"headers": {},
"data": {"class_method": "stream_query"},
"messages": messages,
"litellm_params": {},
}
transport: Final = httpx.MockTransport(_separated_text_reply)
if sync_stream:
sync_wrapper: Final = VertexAgentEngineConfig().get_sync_custom_stream_wrapper(
client=HTTPHandler(client=httpx.Client(transport=transport)), **call_args
)
return "".join(chunk.choices[0].delta.content or "" for chunk in sync_wrapper)
async_wrapper: Final = await VertexAgentEngineConfig().get_async_custom_stream_wrapper(
client=AsyncHTTPHandler(transport=transport), **call_args
)
return "".join([chunk.choices[0].delta.content or "" async for chunk in async_wrapper])
@pytest.mark.parametrize("sync_stream", [True, False])
async def test_stream_keeps_unicode_line_separators_inside_the_text(sync_stream: bool) -> None:
assert await _streamed_text(sync_stream) == _LINE_SEPARATED_TEXT

64
tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
View file

@ -3,8 +3,9 @@ import json
import logging
import re
import threading
from collections.abc import Awaitable, Callable, Iterator, Mapping
from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping
from copy import deepcopy
from itertools import pairwise
from typing import Final, List, Literal, cast, get_args
from unittest.mock import MagicMock, Mock, patch
@ -34,7 +35,7 @@ from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
)
from litellm.types.llms.base import CachedTokensDetails
from litellm.types.llms.vertex_ai import GeminiFinishReason, RequestBody, UsageMetadata
from litellm.types.utils import ChoiceLogprobs, ModelInfo, Usage
from litellm.types.utils import ChoiceLogprobs, ModelInfo, ModelResponseStream, Usage
from litellm.utils import _invalidate_model_cost_lowercase_map, CustomStreamWrapper
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.integrations.custom_logger import CustomLogger
@ -10125,7 +10126,7 @@ def test_vertex_ai_streaming_response_id():
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.iter_lines = MagicMock(return_value=stream_response())
mock_response.iter_text = MagicMock(return_value=stream_response())
return mock_response
logging_obj = MagicMock()
@ -10690,6 +10691,63 @@ async def test_gemini_sends_a_boolean_for_a_gateway_default_written_as_text_or_n
assert sent["includeThoughts"] is expected
class _BodyInChunks(httpx.SyncByteStream, httpx.AsyncByteStream):
def __init__(self, chunks: tuple[bytes, ...]) -> None:
self._chunks: Final = chunks
def __iter__(self) -> Iterator[bytes]:
yield from self._chunks
async def __aiter__(self) -> AsyncIterator[bytes]:
for chunk in self._chunks:
yield chunk
def _split_at(body: bytes, cuts: tuple[int, ...]) -> tuple[bytes, ...]:
return tuple(body[start:end] for start, end in pairwise((0, *sorted(cuts), len(body))))
async def _gemini_stream(is_async: bool, chunks: tuple[bytes, ...]) -> list[ModelResponseStream]:
def handler(request: httpx.Request) -> httpx.Response:
return httpx.Response(200, stream=_BodyInChunks(chunks), headers={"content-type": "text/event-stream"})
kwargs: Final = {
"model": "gemini/gemini-3.8-flash",
"messages": [{"role": "user", "content": "hi"}],
"stream": True,
"stream_options": {"include_usage": True},
"api_key": "fake-gemini-key",
}
if is_async:
stream: Final = await litellm.acompletion(
**kwargs, client=AsyncHTTPHandler(transport=httpx.MockTransport(handler))
)
return [chunk async for chunk in stream]
return list(
litellm.completion(**kwargs, client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler))))
)
@pytest.mark.parametrize("is_async", [True, False])
async def test_gemini_stream_keeps_text_and_usage_when_a_part_contains_a_unicode_line_separator(is_async: bool) -> None:
events: Final = (
{"candidates": [{"content": {"role": "model", "parts": [{"text": "line one\u2028line two"}]}}]},
{
"candidates": [{"content": {"role": "model", "parts": [{"text": " and three"}]}, "finishReason": "STOP"}],
"usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 8, "totalTokenCount": 18},
},
)
body: Final = "".join(f"data: {json.dumps(event, ensure_ascii=False)}\r\n\r\n" for event in events).encode()
cut_inside_separator_and_crlf: Final = (body.index(b"\xe2\x80\xa8") + 1, body.index(b"\r") + 1)
received: Final = await _gemini_stream(is_async, _split_at(body, cut_inside_separator_and_crlf))
usage: Final = next(chunk.usage for chunk in reversed(received) if getattr(chunk, "usage", None))
text: Final = "".join(c.choices[0].delta.content or "" for c in received if c.choices)
assert text == "line one\u2028line two and three"
assert (usage.prompt_tokens, usage.completion_tokens) == (10, 8)
@pytest.mark.parametrize("chunk_indices", [[0, 1], [2, 1, 0]])
def test_grounding_support_preserves_every_web_source(chunk_indices: list[int]) -> None:
metadata: Final = {

8
tests/unit/llms/vertex_ai/test_http_status_201.py
View file

@ -39,8 +39,8 @@ class TestVertexAIHTTPStatus201(unittest.TestCase):
# Create a mock response with status code 201
mock_response = MagicMock()
mock_response.status_code = 201
mock_response.aiter_lines = MagicMock()
mock_response.aiter_lines.return_value = ["test response"]
mock_response.aiter_text = MagicMock()
mock_response.aiter_text.return_value = ["test response"]
# Setup mock client
mock_client = MagicMock()
@ -76,8 +76,8 @@ class TestVertexAIHTTPStatus201(unittest.TestCase):
# Create a mock response with status code 201
mock_response = MagicMock()
mock_response.status_code = 201
mock_response.iter_lines = MagicMock()
mock_response.iter_lines.return_value = ["test response"]
mock_response.iter_text = MagicMock()
mock_response.iter_text.return_value = ["test response"]
mock_post.return_value = mock_response
# Call the make_sync_call function

138
tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py
View file

@ -4,10 +4,11 @@ Mock tests for A2A endpoints.
Tests that invoke_agent_a2a properly integrates with add_litellm_data_to_request.
"""
import asyncio
import json
import socket
import sys
from collections.abc import Awaitable, Callable, Mapping
from collections.abc import AsyncIterator, Awaitable, Callable, Mapping
from contextlib import AbstractContextManager, ExitStack
from dataclasses import dataclass
from typing import Final
@ -24,6 +25,21 @@ from litellm.types.agents import AgentCaller
AddLiteLLMData = Callable[..., Awaitable[dict[str, object]]]
class _SSELineStream(httpx.AsyncByteStream):
def __init__(self, lines: tuple[str, ...], first_line_delay: float) -> None:
self._lines = lines
self._first_line_delay = first_line_delay
async def __aiter__(self) -> AsyncIterator[bytes]:
await asyncio.sleep(self._first_line_delay)
for line in self._lines:
yield f"{line}\n".encode()
def _sse_upstream_response(*lines: str, first_line_delay: float = 0.0) -> httpx.Response:
return httpx.Response(200, stream=_SSELineStream(lines, first_line_delay))
@dataclass(frozen=True, slots=True)
class CapturedAgentCall:
request_id: object
@ -941,13 +957,7 @@ async def test_task_methods_extract_litellm_params_before_forwarding(method: str
mock_http_response.json.return_value = upstream_response
mock_http_response.is_success = True
async def fake_aiter_lines():
yield 'data: {"jsonrpc":"2.0","id":"req-1","result":{"taskId":"task-1"}}'
mock_resp = AsyncMock()
mock_resp.is_success = True
mock_resp.aiter_lines = fake_aiter_lines
mock_resp.aclose = AsyncMock()
mock_resp = _sse_upstream_response('data: {"jsonrpc":"2.0","id":"req-1","result":{"taskId":"task-1"}}')
mock_async_client = MagicMock()
mock_async_client.build_request = MagicMock(return_value=MagicMock())
@ -982,8 +992,8 @@ async def test_task_methods_extract_litellm_params_before_forwarding(method: str
user_api_key_dict=user_api_key_dict,
)
if method == "tasks/resubscribe":
async for _ in response.body_iterator:
pass
streamed = "".join([chunk async for chunk in response.body_iterator])
assert "task-1" in streamed
if method == "tasks/resubscribe":
forwarded_body = mock_async_client.build_request.call_args.kwargs["json"]
@ -1007,14 +1017,7 @@ async def test_subscribe_to_task_returns_sse_stream():
'data: {"jsonrpc":"2.0","id":"req-1","result":{"taskId":"task-1","status":{"state":"completed"}}}',
]
async def fake_aiter_lines():
for line in sse_lines:
yield line
mock_resp = AsyncMock()
mock_resp.is_success = True
mock_resp.aiter_lines = fake_aiter_lines
mock_resp.aclose = AsyncMock()
mock_resp = _sse_upstream_response(*sse_lines)
mock_async_client = MagicMock()
mock_async_client.build_request = MagicMock(return_value=MagicMock())
@ -1064,13 +1067,9 @@ async def test_subscribe_to_task_calls_pre_call_hook():
mock_request = _make_request_mock("tasks/resubscribe", {"id": "task-1"})
user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="t1")
async def fake_aiter_lines():
yield 'data: {"jsonrpc":"2.0","id":"req-1","result":{"taskId":"task-1","status":{"state":"completed"}}}'
mock_resp = AsyncMock()
mock_resp.is_success = True
mock_resp.aiter_lines = fake_aiter_lines
mock_resp.aclose = AsyncMock()
mock_resp = _sse_upstream_response(
'data: {"jsonrpc":"2.0","id":"req-1","result":{"taskId":"task-1","status":{"state":"completed"}}}'
)
mock_async_client = MagicMock()
mock_async_client.build_request = MagicMock(return_value=MagicMock())
@ -1117,9 +1116,9 @@ async def test_subscribe_to_task_calls_pre_call_hook():
)
assert response.media_type == "text/event-stream"
async for _ in response.body_iterator:
pass
streamed = "".join([chunk async for chunk in response.body_iterator])
assert "completed" in streamed
mock_proxy_logging.pre_call_hook.assert_awaited_once()
call_kwargs = mock_proxy_logging.pre_call_hook.await_args.kwargs
assert call_kwargs.get("call_type") == "asend_message"
@ -1149,16 +1148,10 @@ async def test_subscribe_to_task_runs_post_call_streaming_guardrail():
mock_request = _make_request_mock("tasks/resubscribe", {"id": "task-1"})
user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="t1")
async def fake_aiter_lines():
yield (
'data: {"jsonrpc":"2.0","id":"req-1","result":'
'{"kind":"message","parts":[{"kind":"text","text":"resubscribe-secret"}]}}'
)
mock_resp = AsyncMock()
mock_resp.is_success = True
mock_resp.aiter_lines = fake_aiter_lines
mock_resp.aclose = AsyncMock()
mock_resp = _sse_upstream_response(
'data: {"jsonrpc":"2.0","id":"req-1","result":'
'{"kind":"message","parts":[{"kind":"text","text":"resubscribe-secret"}]}}'
)
mock_async_client = MagicMock()
mock_async_client.build_request = MagicMock(return_value=MagicMock())
@ -1600,14 +1593,7 @@ async def test_pascal_method_names_normalize_to_wire_format(pascal_method: str,
mock_http_response.is_success = True
mock_http_response.raise_for_status = MagicMock()
async def _empty_aiter_lines():
return
yield # make it an async generator
mock_sse_resp = AsyncMock()
mock_sse_resp.is_success = True
mock_sse_resp.aiter_lines = _empty_aiter_lines
mock_sse_resp.aclose = AsyncMock()
mock_sse_resp = _sse_upstream_response()
mock_async_client = MagicMock()
mock_async_client.build_request = MagicMock(return_value=MagicMock())
@ -2349,13 +2335,7 @@ async def test_task_methods_forward_trace_header(method: str):
mock_http_response.json.return_value = upstream_response
mock_http_response.is_success = True
async def fake_aiter_lines():
yield 'data: {"jsonrpc":"2.0","id":"req-1","result":{"taskId":"task-1"}}'
mock_resp = AsyncMock()
mock_resp.is_success = True
mock_resp.aiter_lines = fake_aiter_lines
mock_resp.aclose = AsyncMock()
mock_resp = _sse_upstream_response('data: {"jsonrpc":"2.0","id":"req-1","result":{"taskId":"task-1"}}')
mock_async_client = MagicMock()
mock_async_client.build_request = MagicMock(return_value=MagicMock())
@ -2390,8 +2370,8 @@ async def test_task_methods_forward_trace_header(method: str):
user_api_key_dict=user_api_key_dict,
)
if method == "tasks/resubscribe":
async for _ in response.body_iterator:
pass
streamed = "".join([chunk async for chunk in response.body_iterator])
assert "task-1" in streamed
if method == "tasks/resubscribe":
forwarded_headers = mock_async_client.build_request.call_args.kwargs["headers"]
@ -2646,17 +2626,12 @@ def test_served_version_falls_back_to_header_when_unconfigured():
assert _served_version(_agent(None), _request_with_a2a_header(None)) == "0.3"
def _sse_agent_handler(lines):
mock_resp = AsyncMock()
mock_resp.is_success = True
mock_resp.aiter_lines = lines
mock_resp.aclose = AsyncMock()
mock_async_client = MagicMock()
def _sse_agent_handler(upstream_response: httpx.Response) -> MagicMock:
mock_async_client: Final = MagicMock()
mock_async_client.build_request = MagicMock(return_value=MagicMock())
mock_async_client.send = AsyncMock(return_value=mock_resp)
mock_async_client.send = AsyncMock(return_value=upstream_response)
mock_handler = MagicMock()
mock_handler: Final = MagicMock()
mock_handler.client = mock_async_client
return mock_handler
@ -2685,13 +2660,13 @@ async def test_forward_jsonrpc_sse_pings_while_the_upstream_agent_is_still_silen
monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", 0.05)
async def _slow_lines():
await asyncio.sleep(0.3)
yield 'data: {"jsonrpc": "2.0", "id": "req-1", "result": {"kind": "task"}}'
slow_upstream = _sse_upstream_response(
'data: {"jsonrpc": "2.0", "id": "req-1", "result": {"kind": "task"}}', first_line_delay=0.3
)
with patch(
"litellm.llms.custom_httpx.http_handler.get_async_httpx_client",
return_value=_sse_agent_handler(_slow_lines),
return_value=_sse_agent_handler(slow_upstream),
):
response = await _resubscribe_response()
assert response.headers["x-accel-buffering"] == "no"
@ -2713,12 +2688,11 @@ async def test_forward_jsonrpc_sse_is_untouched_while_keepalives_are_unconfigure
monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", None)
async def _lines():
yield 'data: {"jsonrpc": "2.0", "id": "req-1", "result": {"kind": "task"}}'
with patch(
"litellm.llms.custom_httpx.http_handler.get_async_httpx_client",
return_value=_sse_agent_handler(_lines),
return_value=_sse_agent_handler(
_sse_upstream_response('data: {"jsonrpc": "2.0", "id": "req-1", "result": {"kind": "task"}}')
),
):
response = await _resubscribe_response()
assert "x-accel-buffering" not in response.headers
@ -2728,6 +2702,28 @@ async def test_forward_jsonrpc_sse_is_untouched_while_keepalives_are_unconfigure
assert json.loads(chunks[-1].removeprefix("data: "))["result"]["kind"] == "task"
@pytest.mark.asyncio
async def test_forward_jsonrpc_sse_keeps_unicode_line_separators_inside_an_event(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", None)
separated_text: Final = "before\u2028between\u2029after"
event: Final = {
"jsonrpc": "2.0",
"id": "req-1",
"result": {"kind": "message", "parts": [{"kind": "text", "text": separated_text}]},
}
with patch(
"litellm.llms.custom_httpx.http_handler.get_async_httpx_client",
return_value=_sse_agent_handler(_sse_upstream_response(f"data: {json.dumps(event, ensure_ascii=False)}")),
):
response: Final = await _resubscribe_response()
chunks: Final = [chunk async for chunk in response.body_iterator]
assert [json.loads(chunk.removeprefix("data: "))["result"]["parts"][0]["text"] for chunk in chunks] == [
separated_text
]
async def _stream_message_response():
from litellm.proxy.agent_endpoints.a2a_endpoints import _handle_stream_message

23
tests/unit/proxy/anthropic_endpoints/test_streaming_model_restamp.py
View file

@ -4,6 +4,7 @@ Tests for restamping the public model on Anthropic Messages streaming chunks.
import json
from collections.abc import AsyncIterator, Sequence
from typing import Final
import pytest
@ -46,6 +47,28 @@ def test_restamps_sse_bytes_frame():
assert b"event: message_start" in restamped
def test_restamps_sse_bytes_frame_holding_a_line_separator() -> None:
payload: Final = {
"type": "message_start",
"message": {
"id": "msg_1",
"type": "message",
"role": "assistant",
"model": "claude-haiku-4-5-20251001",
"content": [{"type": "text", "text": "alpha\u2028beta"}],
},
}
frame: Final = f"event: message_start\ndata: {json.dumps(payload, ensure_ascii=False)}\n\n".encode()
restamped: Final = restamp_anthropic_stream_chunk_model(frame, "claude-auto-1")
assert isinstance(restamped, bytes)
data_line: Final = next(line for line in restamped.decode().split("\n") if line.startswith("data:"))
message: Final = json.loads(data_line[len("data:") :])["message"]
assert message["model"] == "claude-auto-1"
assert message["content"] == [{"type": "text", "text": "alpha\u2028beta"}]
def test_restamps_event_dict():
chunk = {"type": "message_start", "message": {"id": "msg_1", "model": "claude-sonnet-4-6"}}

48
tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py
View file

@ -2503,6 +2503,33 @@ def test_stream_was_interrupted_skips_data_lines_that_are_not_json_objects(
assert AnthropicPassthroughLoggingHandler._stream_was_interrupted(all_chunks) is interrupted
def test_stream_was_interrupted_reads_a_message_delta_whose_json_holds_unicode_line_separators() -> None:
message_delta: Final = {
"type": "message_delta",
"delta": {"stop_reason": "stop_sequence", "stop_sequence": "END\u2028OF\u2029TURN"},
"usage": {"output_tokens": 3},
}
all_chunks: Final = [
'data: {"type": "content_block_stop", "index": 0}',
f"data: {json.dumps(message_delta, ensure_ascii=False)}",
]
assert AnthropicPassthroughLoggingHandler._stream_was_interrupted(all_chunks) is False
def test_extract_model_reads_a_message_start_whose_json_holds_unicode_line_separators() -> None:
message_start: Final = {
"type": "message_start",
"message": {
"model": "claude-sonnet-5",
"content": [{"type": "text", "text": "Prefill\u2028line\u2029two"}],
},
}
chunk: Final = f"event: message_start\ndata: {json.dumps(message_start, ensure_ascii=False)}\n\n"
assert AnthropicPassthroughLoggingHandler._extract_model_from_anthropic_chunks([chunk]) == "claude-sonnet-5"
@pytest.mark.asyncio
async def test_anthropic_passthrough_handler(mock_httpx_response, mock_response, mock_logging_obj):
"""
@ -3004,3 +3031,24 @@ def test_anthropic_passthrough_stream_bills_the_bedrock_guardrail_trace_of_the_m
(entry,) = _recorded_guardrail_entries(mock_logging_obj)
assert entry["guardrail_usage"] == {"topicPolicyUnits": 3, "contentPolicyUnits": 3}
assert entry["guardrail_cost"] == pytest.approx(6 * 0.00015)
def test_anthropic_passthrough_stream_bills_a_bedrock_guardrail_trace_whose_json_holds_unicode_line_separators(
monkeypatch: pytest.MonkeyPatch, mock_logging_obj: LiteLLMLoggingObj, all_chunks: list[str]
) -> None:
monkeypatch.setattr(litellm, "model_cost", _BEDROCK_TRACE_PRICING)
mock_logging_obj.model_call_details["litellm_params"] = {"metadata": {}}
trace: Final = {"guardrail": {**_BEDROCK_INVOKE_TRACE["guardrail"], "actionReason": "Blocked\u2028by\u2029policy"}}
stop_event: Final = json.dumps({"type": "message_stop", "amazon-bedrock-trace": trace}, ensure_ascii=False)
AnthropicPassthroughLoggingHandler.handle_logging_anthropic_collected_chunks(
litellm_logging_obj=mock_logging_obj,
passthrough_success_handler_obj=MagicMock(),
url_route="/v1/messages",
request_body={"model": "claude", "messages": []},
endpoint_type=MagicMock(),
start_time=datetime(2025, 1, 1),
all_chunks=[*all_chunks, f"event: message_stop\ndata: {stop_event}\n\n"],
end_time=datetime(2025, 1, 1),
)
(entry,) = _recorded_guardrail_entries(mock_logging_obj)
assert entry["guardrail_usage"] == {"topicPolicyUnits": 3, "contentPolicyUnits": 3}

7
tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_tinyfish_passthrough_logging_handler.py
View file

@ -1,6 +1,7 @@
import asyncio
import json
from datetime import datetime
from typing import Final
from unittest.mock import AsyncMock, MagicMock
import httpx
@ -267,6 +268,12 @@ class TestRunIdFromSseFrames:
frames = b'data: {"run_id": "run-7", "event": "INITIALIZED"}\n\ndata: {"run_id": "run-7", "event": "ACTION"}\n\n'
assert run_id_from_sse_frames(frames) == "run-7"
def test_finds_run_id_in_a_frame_whose_json_holds_unicode_line_separators(self) -> None:
event: Final = {"run_id": "run-7", "event": "ACTION", "message": "Clicked\u2028Submit\u2029button"}
frames: Final = f"data: {json.dumps(event, ensure_ascii=False)}\n\n".encode()
assert run_id_from_sse_frames(frames) == "run-7"
def test_skips_frames_without_run_id(self):
frames = b': keepalive\n\ndata: not-json\n\ndata: {"event": "HEARTBEAT"}\n\n'
assert run_id_from_sse_frames(frames) is None

33
tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py
View file

@ -30,7 +30,7 @@ from litellm.llms.base_llm.managed_resources.utils import(
resolve_passthrough_managed_id_provider,
)
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
from typing import Any
from typing import Any, Final
def _user() -> UserAPIKeyAuth:
@ -254,6 +254,37 @@ async def test_streamed_response_with_cr_only_frame_delimiters_is_still_owned_an
assert output == _response_stream_bytes(managed_id).replace(b"\n", b"\r")
def _instructed_response_stream_bytes(raw_id: str, instructions: str) -> bytes:
events: Final = (
{"type": "response.created", "response": {"id": raw_id, "status": "in_progress", "instructions": instructions}},
{"type": "response.completed", "response": {"id": raw_id, "status": "completed", "instructions": instructions}},
)
return b"".join(
f"event: {event['type']}\ndata: {json.dumps(event, ensure_ascii=False)}\n\n".encode() for event in events
)
@pytest.mark.asyncio
async def test_streamed_response_whose_json_holds_unicode_line_separators_is_still_owned_and_rewritten() -> None:
instructions: Final = "Answer\u2028in\u2029verse"
pc: Final = _prisma_client()
output: Final = await _collect(
rewrite_streamed_response_ids(
stream=_chunks(_instructed_response_stream_bytes(RAW_RESPONSE_ID, instructions), 4096),
provider="openai",
method="POST",
route="/openai_passthrough/v1/responses",
user_api_key_dict=_user(),
prisma_client=pc,
)
)
pc.db.litellm_managedobjecttable.upsert.assert_awaited_once()
managed_id: Final = pc.db.litellm_managedobjecttable.upsert.await_args.kwargs["data"]["create"]["unified_object_id"]
assert output == _instructed_response_stream_bytes(managed_id, instructions)
@pytest.mark.asyncio
async def test_streamed_bytes_untouched_on_routes_without_a_response_id():
pc = _prisma_client()

31
tests/unit/proxy/response_polling/test_response_polling_handler.py
View file

@ -17,7 +17,7 @@ following the OpenAI Response API format.
import asyncio
import json
from datetime import datetime, timezone
from typing import Any, Dict, Optional
from typing import Any, Dict, Final, Optional
from unittest.mock import AsyncMock, Mock, patch
import pytest
@ -1400,7 +1400,7 @@ def _make_sse_stream(events: list) -> Mock:
async def _body_iterator():
for event in events:
yield f"data: {json.dumps(event)}"
yield f"data: {json.dumps(event, ensure_ascii=False)}"
yield "data: [DONE]"
mock_response = Mock()
@ -1639,6 +1639,33 @@ class TestBackgroundStreamingTerminalEvents:
assert final_call.kwargs["status"] == "completed"
assert final_call.kwargs["usage"] == {"input_tokens": 10, "output_tokens": 50}
@pytest.mark.asyncio
async def test_response_completed_whose_json_holds_unicode_line_separators_keeps_its_output(self) -> None:
from litellm.proxy.response_polling.background_streaming import (
background_streaming_task,
)
item: Final = {
"id": "item_1",
"type": "message",
"content": [{"type": "output_text", "text": "Line\u2028one\u2029two"}],
}
events: Final = [
{"type": "response.completed", "response": {"id": "resp_123", "status": "completed", "output": [item]}}
]
handler: Final = AsyncMock(spec=ResponsePollingHandler)
with patch(
"litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing"
) as MockProcessor:
MockProcessor.return_value.base_process_llm_request = AsyncMock(return_value=_make_sse_stream(events))
await background_streaming_task(**_make_background_streaming_kwargs("poll_u2028", handler))
calls: Final = handler.update_state.call_args_list
outputs: Final = [call.kwargs["output"] for call in calls if "output" in call.kwargs]
assert outputs[-1] == [item]
assert calls[-1].kwargs["status"] == "completed"
@pytest.mark.asyncio
async def test_fallback_status_derived_from_event_type_when_status_field_missing(
self,

100
tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py
View file

@ -1288,6 +1288,106 @@ async def test_reasoning_item_closes_before_a_tool_call_from_the_same_chunk_open
]
def _text_opening_a_tool_call_chunk(content: str, arguments: str = "") -> ModelResponseStream:
return ModelResponseStream(
id=CHAT_COMPLETION_ID,
created=1748575031,
model="claude-haiku-5-5",
object="chat.completion.chunk",
choices=[
StreamingChoices(
index=0,
delta=Delta(
role="assistant",
content=content,
tool_calls=[
{
"id": "call_weather",
"type": "function",
"function": {"name": "get_weather", "arguments": arguments},
"index": 0,
}
],
),
finish_reason=None,
)
],
)
def _tool_arguments_chunk(arguments: str, finish_reason: str | None = None) -> ModelResponseStream:
return ModelResponseStream(
id=CHAT_COMPLETION_ID,
created=1748575031,
model="claude-haiku-5-5",
object="chat.completion.chunk",
choices=[
StreamingChoices(
index=0,
delta=Delta(
content=None,
tool_calls=[{"type": "function", "function": {"arguments": arguments}, "index": 0}],
),
finish_reason=finish_reason,
)
],
)
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_tool_call_opened_in_a_text_chunk_streams_each_argument_delta(sync_mode: bool) -> None:
argument_pieces: Final = ('{"', "city", '"', ': "', "Paris", '"', "}")
iterator: Final = _build_iterator(
[
_text_opening_a_tool_call_chunk("</think>"),
*(_tool_arguments_chunk(piece) for piece in argument_pieces),
_chunk("", finish_reason="tool_calls"),
]
)
events: Final = await _collect_events(iterator, sync_mode)
assert _event_text(events, ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA) == "</think>"
assert [
event.delta
for event in events
if getattr(event, "type", None) == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA
] == list(argument_pieces)
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_tool_call_sharing_a_chunk_with_text_opens_after_the_whole_message(sync_mode: bool) -> None:
iterator: Final = _build_iterator(
[
_text_opening_a_tool_call_chunk("Checking ", arguments='{"city": "Paris"}'),
_chunk("the "),
_chunk("weather now."),
_chunk("", finish_reason="tool_calls"),
]
)
events: Final = await _collect_events(iterator, sync_mode)
event_types: Final = [getattr(event, "type", None) for event in events]
function_call_added_at: Final = next(
position
for position, event in enumerate(events)
if event_types[position] == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED and event.item.type == "function_call"
)
streamed: Final = [
event
for event in events
if getattr(event, "type", None)
in (ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA)
]
assert ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA not in event_types[function_call_added_at:]
assert [event.sequence_number for event in streamed] == sorted(event.sequence_number for event in streamed)
assert _event_text(events, ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA) == "Checking the weather now."
assert _event_text(events, ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA) == '{"city": "Paris"}'
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_tool_only_stream_puts_the_function_call_at_output_index_zero(sync_mode: bool):

28
tests/unit/sandbox/test_e2b_sandbox.py
View file

@ -7,6 +7,7 @@ integration tests live in tests/integration/sandbox/test_e2b_sandbox.py.
"""
import json
from typing import Final
import httpx
import pytest
@ -19,19 +20,18 @@ from litellm.llms.e2b.sandbox.transformation import (
)
def _line_stream(lines: list[str]) -> httpx.Response:
return httpx.Response(200, stream=httpx.ByteStream("".join(f"{line}\n" for line in lines).encode()))
class FakeResponse:
def __init__(self, *, json_data=None, lines=None, status_code=200):
def __init__(self, *, json_data: object = None, status_code: int = 200) -> None:
self._json = json_data
self._lines = lines or []
self.status_code = status_code
def json(self):
return self._json
async def aiter_lines(self):
for line in self._lines:
yield line
class FakeHTTPClient:
"""Records outbound requests and returns canned responses keyed by URL."""
@ -61,7 +61,7 @@ class FakeHTTPClient:
if url.endswith("/execute"):
if self.execute_raises is not None:
raise self.execute_raises
return FakeResponse(lines=self.execute_lines)
return _line_stream(self.execute_lines)
raise AssertionError(f"unexpected POST {url}")
async def delete(self, url, headers=None, **kwargs):
@ -249,6 +249,20 @@ def test_parse_lines_skips_non_json_lines():
assert result.error is None
@pytest.mark.asyncio
async def test_run_code_keeps_unicode_line_separators_inside_stdout() -> None:
separated_text: Final = "before\u2028between\u2029after\n"
client: Final = FakeHTTPClient(
execute_lines=[json.dumps({"type": "stdout", "text": separated_text, "timestamp": 1}, ensure_ascii=False)]
)
handle: Final = ContainerHandle(id="sbx_sep", provider="e2b", domain="e2b.app")
handle._hidden_params = {"envd_access_token": "tok"}
result: Final = await E2BSandboxConfig().arun_code(container=handle, code="print(text)", client=client)
assert result.stdout == separated_text
@pytest.mark.asyncio
async def test_run_code_aborts_on_output_over_cap():
big_line = "x" * (MAX_OUTPUT_BYTES + 1)

13
tests/unit/sandbox/test_opensandbox_sandbox.py
View file

@ -27,10 +27,13 @@ def sse(data):
return f"data: {json.dumps(data)}"
def _line_stream(lines: list[str]) -> httpx.Response:
return httpx.Response(200, stream=httpx.ByteStream("".join(f"{line}\n" for line in lines).encode()))
class FakeResponse:
def __init__(self, *, json_data=None, lines=None, status_code=200):
def __init__(self, *, json_data: object = None, status_code: int = 200) -> None:
self._json = json_data
self._lines = lines or []
self.status_code = status_code
def json(self):
@ -40,10 +43,6 @@ class FakeResponse:
if self.status_code >= 400:
raise http_status_error(self.status_code)
async def aiter_lines(self):
for line in self._lines:
yield line
class FakeHTTPClient:
def __init__(
@ -93,7 +92,7 @@ class FakeHTTPClient:
if url.endswith("/code"):
if self.execute_raises is not None:
raise self.execute_raises
return FakeResponse(lines=self.execute_lines)
return _line_stream(self.execute_lines)
raise AssertionError(f"unexpected POST {url}")
async def get(self, url, headers=None, params=None, **kwargs):

164
tests/unit/test_groq_streaming_encoding.py
View file

@ -6,148 +6,58 @@ UTF-8 encoded content in streaming responses, specifically fixing
the ASCII encoding error described in issue #12660.
"""
import asyncio
from unittest.mock import AsyncMock, Mock
import json
from typing import Final
from unittest.mock import Mock
import httpx
import pytest
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.openai_like.chat.handler import make_call, make_sync_call
class MockResponse:
"""Mock httpx response for testing UTF-8 handling."""
def __init__(self, test_content: str):
self.test_content = test_content
self.status_code = 200
def iter_text(self, encoding="utf-8"):
"""Mock iter_text that yields content with the specified encoding."""
yield self.test_content
async def aiter_text(self, encoding="utf-8"):
"""Mock aiter_text that yields content with the specified encoding."""
yield self.test_content
def iter_lines(self):
"""Mock iter_lines method for synchronous streaming."""
yield self.test_content
async def aiter_lines(self):
"""Mock aiter_lines method for asynchronous streaming."""
yield self.test_content
def json(self):
return {"choices": [{"delta": {"content": "test"}}]}
_UNICODE_TEXTS: Final = (
"The symbol µ represents micro",
"Testing © character",
"Testing ™ character",
"Testing € character",
"Testing 北京 character",
"Testing 🚀 character",
"Testing Ñoño character",
)
class MockSyncClient:
"""Mock synchronous HTTP client for testing."""
def __init__(self, response_content: str):
self.response_content = response_content
def post(self, *args, **kwargs):
return MockResponse(self.response_content)
def _utf8_stream_transport(text: str) -> httpx.MockTransport:
chunk: Final = {"choices": [{"index": 0, "delta": {"content": text}}]}
body: Final = f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n".encode()
return httpx.MockTransport(lambda request: httpx.Response(200, stream=httpx.ByteStream(body)))
class MockAsyncClient:
"""Mock asynchronous HTTP client for testing."""
def __init__(self, response_content: str):
self.response_content = response_content
async def post(self, *args, **kwargs):
return MockResponse(self.response_content)
def _call_args() -> dict[str, object]:
return {
"api_base": "https://test.com/v1/chat/completions",
"headers": {"Authorization": "Bearer test"},
"data": '{"model": "test", "messages": []}',
"model": "test-model",
"messages": [],
"logging_obj": Mock(),
}
def test_utf8_streaming_sync():
"""Test that synchronous streaming handles UTF-8 characters correctly."""
# Content with the µ character that was causing issues
test_content = (
'data: {"choices":[{"delta":{"content":"The symbol µ represents micro"}}]}\n\n'
@pytest.mark.parametrize("text", _UNICODE_TEXTS)
def test_utf8_streaming_sync(text: str) -> None:
completion_stream: Final = make_sync_call(
client=HTTPHandler(client=httpx.Client(transport=_utf8_stream_transport(text))), **_call_args()
)
mock_client = MockSyncClient(test_content)
mock_logging = Mock()
# This should not raise an ASCII encoding error
completion_stream = make_sync_call(
client=mock_client,
api_base="https://test.com/v1/chat/completions",
headers={"Authorization": "Bearer test"},
data='{"model": "test", "messages": []}',
model="test-model",
messages=[],
logging_obj=mock_logging,
)
# Verify we can iterate through the stream without encoding errors
assert completion_stream is not None
assert "".join(chunk["text"] for chunk in completion_stream) == text
@pytest.mark.asyncio
async def test_utf8_streaming_async():
"""Test that asynchronous streaming handles UTF-8 characters correctly."""
# Content with the µ character that was causing issues
test_content = (
'data: {"choices":[{"delta":{"content":"The symbol µ represents micro"}}]}\n\n'
@pytest.mark.parametrize("text", _UNICODE_TEXTS)
async def test_utf8_streaming_async(text: str) -> None:
completion_stream: Final = await make_call(
client=AsyncHTTPHandler(transport=_utf8_stream_transport(text)), **_call_args()
)
mock_client = MockAsyncClient(test_content)
mock_logging = Mock()
# This should not raise an ASCII encoding error
completion_stream = await make_call(
client=mock_client,
api_base="https://test.com/v1/chat/completions",
headers={"Authorization": "Bearer test"},
data='{"model": "test", "messages": []}',
model="test-model",
messages=[],
logging_obj=mock_logging,
)
# Verify we can iterate through the stream without encoding errors
assert completion_stream is not None
def test_various_unicode_characters():
"""Test streaming with various Unicode characters that could cause issues."""
unicode_test_cases = [
"µ", # Micro symbol (the original issue)
"©", # Copyright symbol
"™", # Trademark symbol
"€", # Euro symbol
"北京", # Chinese characters
"🚀", # Emoji
"Ñoño", # Spanish characters with tildes
]
for unicode_char in unicode_test_cases:
test_content = f'data: {{"choices":[{{"delta":{{"content":"Testing {unicode_char} character"}}}}]}}\n\n'
mock_client = MockSyncClient(test_content)
mock_logging = Mock()
# This should not raise an ASCII encoding error for any Unicode character
completion_stream = make_sync_call(
client=mock_client,
api_base="https://test.com/v1/chat/completions",
headers={"Authorization": "Bearer test"},
data='{"model": "test", "messages": []}',
model="test-model",
messages=[],
logging_obj=mock_logging,
)
assert (
completion_stream is not None
), f"Failed to handle Unicode character: {unicode_char}"
if __name__ == "__main__":
test_utf8_streaming_sync()
asyncio.run(test_utf8_streaming_async())
test_various_unicode_characters()
print("All UTF-8 streaming tests passed!")
assert "".join([chunk["text"] async for chunk in completion_stream]) == text