diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml
index b9f8d7e7dfb..89ca60ad9d1 100644
--- a/.github/workflows/test-linting.yml
+++ b/.github/workflows/test-linting.yml
@@ -561,6 +561,9 @@ jobs:
- name: router_code_coverage
run: uv run --no-sync python ./tests/code_coverage_tests/router_code_coverage.py
+ - name: router_surface_check
+ run: uv run --no-sync python ./tests/code_coverage_tests/router_surface_check.py
+
- name: test_chat_completion_imports
run: uv run --no-sync python ./tests/code_coverage_tests/test_chat_completion_imports.py
diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml
index 2ddf21bca17..80d5018fe9a 100644
--- a/.github/workflows/test-unit.yml
+++ b/.github/workflows/test-unit.yml
@@ -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
diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock
index 9c45c5c4b1c..10e590b5006 100644
--- a/litellm-rust/Cargo.lock
+++ b/litellm-rust/Cargo.lock
@@ -4393,6 +4393,10 @@ dependencies = [
"rstest",
]
+[[package]]
+name = "litellm-router-core"
+version = "0.1.0"
+
[[package]]
name = "litellm-router-types"
version = "0.1.0"
diff --git a/litellm-rust/crates/model-catalog/src/lib.rs b/litellm-rust/crates/model-catalog/src/lib.rs
index 066c2c83c6b..d42c05e6fe7 100644
--- a/litellm-rust/crates/model-catalog/src/lib.rs
+++ b/litellm-rust/crates/model-catalog/src/lib.rs
@@ -1,3 +1,36 @@
+//! The model catalog: per-model pricing, context limits and capabilities.
+//!
+//! This is the Rust form of `model_prices_and_context_window.json`, the map Python reads as
+//! `litellm.model_cost` and through `litellm.get_model_info`. It answers "what does this model
+//! cost and what can it do"; it knows nothing about configured deployments. The router's view of
+//! deployments (groups, aliases, wildcard patterns, team mappings) is the router registry, defined
+//! in `litellm-router-core`, which reads this catalog through its dynamic config.
+//!
+//! # Structure
+//! - `Catalog`: the parsed map. `Catalog::parse` takes the JSON body and a `Provenance` (where
+//! the body came from: the bundled copy in the repo or the remote copy, with revision and etag),
+//! and keeps every entry in file order with its raw fields.
+//! - `ModelEntry`: one model. `ModelEntry::fields` is the raw JSON object, so unknown fields are
+//! never lost; `ModelEntry::info` is the typed mirror `ModelInfo` (pricing in `pricing`,
+//! modalities and modes in `capabilities`).
+//! - Aliases: an entry may list `aliases`; each alias resolves to its canonical key. Problems
+//! found while indexing them (bad lists, collisions) are kept as `AliasIssue`s instead of
+//! failing the parse.
+//! - `FallbackGeneralizations` / `FallbackRule`: regex rules that map unknown model ids onto known
+//! families. They are data in the file, not entries.
+//! - Validation: `Catalog::validate` checks a catalog against `IntegrityLimits` (a remote copy that
+//! lost most of its models is refused), and `validate_registry` checks every entry against the
+//! checked-in contract.
+//!
+//! # Lookups
+//! `Catalog::lookup` resolves a key exactly first (entry or alias), then case-insensitively, and
+//! returns a `ModelMatch` with both the key that matched and the canonical key of the entry. Provider
+//! prefixes are not stripped here; callers decide which keys to try.
+//!
+//! # Readers
+//! Cost tracking, the router (through `RouterDynamicConfig`) and the gateway. Nothing in this crate
+//! performs I/O; loading and refreshing the file is the caller's job.
+
mod capabilities;
mod catalog;
mod error;
diff --git a/litellm-rust/crates/router-core/Cargo.toml b/litellm-rust/crates/router-core/Cargo.toml
new file mode 100644
index 00000000000..af97bd52c57
--- /dev/null
+++ b/litellm-rust/crates/router-core/Cargo.toml
@@ -0,0 +1,9 @@
+[package]
+name = "litellm-router-core"
+version = "0.1.0"
+edition.workspace = true
+license.workspace = true
+repository.workspace = true
+description = "The router core behind litellm.Router: registry, settings, health, selection and the call loop"
+
+[dependencies]
diff --git a/litellm-rust/crates/router-core/src/lib.rs b/litellm-rust/crates/router-core/src/lib.rs
new file mode 100644
index 00000000000..05d95d02c84
--- /dev/null
+++ b/litellm-rust/crates/router-core/src/lib.rs
@@ -0,0 +1,63 @@
+//! Routing for LiteLLM: decide which deployment serves a request, and drive retries, fallbacks
+//! and cooldowns around the attempts.
+//!
+//! This crate defines the boundary `litellm.Router` is moving behind. Today the Python router in
+//! `litellm/router.py` implements all of it; this crate is where the Rust core will implement the
+//! same interface. Until then it holds this overview and nothing else, so the interface is written
+//! down in code before either side changes.
+//!
+//! # The public surface
+//! Code outside the router uses `litellm.Router` through two kinds of members and nothing else:
+//! - The endpoint methods (`acompletion`, `aembedding`, `aresponses`, `aanthropic_messages`, the
+//! files, batches, assistants and vector store wrappers, 37 in all). Each is a thin wrapper over
+//! `call(op, request)`.
+//! - Eight entry points: `registry()`, `settings()`, `health()`, `select()`, `apply()`,
+//! `observe()`, `register_plugin()` and `retire()`.
+//!
+//! Every other member that outside code uses today has a named replacement among those. The old
+//! names stay available for about a month through a `DeprecatedRouterMembers` mixin, which
+//! implements them over the entry points and warns on use. A CI check
+//! (`tests/code_coverage_tests/router_surface_check.py`) keeps an allowlist of the members outside
+//! code touches; the list may only shrink.
+//!
+//! # Structure
+//! - `Registry`: the configured deployments, indexed by id, model group, alias, wildcard pattern
+//! and team. Immutable and versioned. Seven queries: `deployment(id)`, `deployments(filter)`,
+//! `names(filter)`, `resolve(name, scope)`, `group(name)`, `access_groups()` and
+//! `credentials(id)`. "Registry" is the router's view of what is configured; per-model pricing
+//! and capabilities are the model catalog (`litellm-model-catalog`), which the router reads
+//! through its dynamic config.
+//! - `RouterConfig` / `RouterDynamicConfig`: everything the router needs to know, passed in. The
+//! router never reads process globals. Static config (model list, strategy, retry, timeout and
+//! cooldown policy, Redis settings, plugins) is fixed for the router's life. Dynamic config
+//! (pricing, model aliases, fallbacks, retry counts, callbacks) is a versioned snapshot read
+//! once per call.
+//! - `Settings`: one typed, frozen view of the configuration a caller may read.
+//! - `Health`: cooldowns, blocked deployments and usage, kept in shared cache keys so Python and
+//! Rust workers see the same state. `report(id, outcome)` records outcomes the router did not
+//! see itself (a health check, an interceptor that marks a deployment failed).
+//! - `select(model, request)`: resolve a name, filter eligible deployments, pick one with the
+//! configured strategy, without calling it.
+//! - `apply(changes)`: one `ChangeSet` (upserts, deletes, settings), applied atomically, producing
+//! a new registry version.
+//! - `call(op, request)`: select, run the attempt, and retry or fall back until it succeeds or the
+//! policy is exhausted.
+//!
+//! # Execution
+//! Each attempt resolves an `Executor` after a deployment is picked. A Rust-native operation runs
+//! its `inference-*` route machine in-process; anything else goes through `Host::invoke`. A
+//! started attempt is never replayed on a different executor.
+//!
+//! # Hosts
+//! `Host` is what the embedding process provides: `invoke` for non-native operations, `select`
+//! for strategies and sub-routers not yet ported, `allow_fallback` for access and budget checks,
+//! `sleep`, `observe` for routing events, and `cache`, a DualCache-style store that holds all
+//! routing state under the same keys, encodings and TTLs the Python router uses. `PythonHost`
+//! serves them from the existing Python code when the router runs inside Python; `NativeHost`
+//! serves them in the Rust server, where an operation without a Rust implementation makes the
+//! config fail at build time, never partway through a request.
+//!
+//! # Plugins
+//! Routing strategies, pre-call checks and sub-routers (auto, complexity, quality, adaptive)
+//! register through `register_plugin` and see a `RoutingContext` (registry, settings, health and
+//! the current request), never the router itself.
diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py b/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py
index 1cea0194715..d4400fbb4cc 100644
--- a/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py
+++ b/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py
@@ -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
diff --git a/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py b/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py
index 0a2871d25cc..3937c47ea3c 100644
--- a/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py
+++ b/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py
@@ -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()
diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py
index 47a3a903d82..ee3639ac387 100644
--- a/litellm/completion_extras/litellm_responses_transformation/transformation.py
+++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py
@@ -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
diff --git a/litellm/google_genai/streaming_iterator.py b/litellm/google_genai/streaming_iterator.py
index a22c47cb1e7..08ac7d4f9dc 100644
--- a/litellm/google_genai/streaming_iterator.py
+++ b/litellm/google_genai/streaming_iterator.py
@@ -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
diff --git a/litellm/interactions/streaming_iterator.py b/litellm/interactions/streaming_iterator.py
index f4ad57bf7b4..c7faa3fc510 100644
--- a/litellm/interactions/streaming_iterator.py
+++ b/litellm/interactions/streaming_iterator.py
@@ -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
diff --git a/litellm/router_utils/batch_utils.py b/litellm/litellm_core_utils/batch_utils.py
similarity index 100%
rename from litellm/router_utils/batch_utils.py
rename to litellm/litellm_core_utils/batch_utils.py
diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py
index cf6bdc69556..101f069c0f8 100644
--- a/litellm/litellm_core_utils/prompt_templates/common_utils.py
+++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py
@@ -21,8 +21,8 @@ from pydantic import TypeAdapter
import litellm
from litellm import verbose_logger
+from litellm.litellm_core_utils.batch_utils import InMemoryFile
from litellm.litellm_core_utils.core_helpers import as_str_mapping
-from litellm.router_utils.batch_utils import InMemoryFile
from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionAssistantMessage,
diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py
index 3fb25de43cc..2348d8e57d6 100644
--- a/litellm/litellm_core_utils/streaming_handler.py
+++ b/litellm/litellm_core_utils/streaming_handler.py
@@ -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 ""
- delta.content = opening + reasoning_content + ("" + 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 = "" + 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 = "" if reasoning_content and not inside_think_block else ""
+ closing: Final = "" 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 ""
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
diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py
index 13ceea3e31b..8a82ed98991 100644
--- a/litellm/llms/anthropic/chat/handler.py
+++ b/litellm/llms/anthropic/chat/handler.py
@@ -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,
diff --git a/litellm/llms/anthropic/messages/stream_assembly.py b/litellm/llms/anthropic/messages/stream_assembly.py
index 90bdad8f4bc..c103eb0f975 100644
--- a/litellm/llms/anthropic/messages/stream_assembly.py
+++ b/litellm/llms/anthropic/messages/stream_assembly.py
@@ -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()
diff --git a/litellm/llms/anthropic/prompt_cache_prediction.py b/litellm/llms/anthropic/prompt_cache_prediction.py
index 7a4c707860d..17c178d169e 100644
--- a/litellm/llms/anthropic/prompt_cache_prediction.py
+++ b/litellm/llms/anthropic/prompt_cache_prediction.py
@@ -29,7 +29,7 @@ from litellm.llms.anthropic.messages.native.transformation import (
AnthropicMessagesConfig,
)
from litellm.llms.anthropic.messages.native.utils import AnthropicMessagesRequestUtils, prepare_native_messages
-from litellm.router_utils.baseline_request import (
+from litellm.router_strategy.auto_router.baseline_request import (
BASELINE_PARAMETERS,
capture_baseline_parameters,
)
diff --git a/litellm/llms/azure_ai/agents/handler.py b/litellm/llms/azure_ai/agents/handler.py
index 94f1399f01b..423a09c5c92 100644
--- a/litellm/llms/azure_ai/agents/handler.py
+++ b/litellm/llms/azure_ai/agents/handler.py
@@ -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:"):
diff --git a/litellm/llms/base_llm/base_model_iterator.py b/litellm/llms/base_llm/base_model_iterator.py
index 032951be74a..3afc685b3dd 100644
--- a/litellm/llms/base_llm/base_model_iterator.py
+++ b/litellm/llms/base_llm/base_model_iterator.py
@@ -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:
diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py
index 725146fe894..fb886e42780 100644
--- a/litellm/llms/base_llm/guardrail_translation/utils.py
+++ b/litellm/llms/base_llm/guardrail_translation/utils.py
@@ -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
diff --git a/litellm/llms/base_llm/sandbox/transformation.py b/litellm/llms/base_llm/sandbox/transformation.py
index fe0e3abb87e..9b9b4e0a947 100644
--- a/litellm/llms/base_llm/sandbox/transformation.py
+++ b/litellm/llms/base_llm/sandbox/transformation.py
@@ -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(
diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py
index 8c4702bcf99..0624c2c65db 100644
--- a/litellm/llms/bedrock/files/transformation.py
+++ b/litellm/llms/bedrock/files/transformation.py
@@ -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())
diff --git a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py
index 1b01b0ec6d0..3a65a8928bf 100644
--- a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py
+++ b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py
@@ -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:")
)
diff --git a/litellm/llms/bedrock/search/transformation.py b/litellm/llms/bedrock/search/transformation.py
index 565fc12aa1e..6f72cadad67 100644
--- a/litellm/llms/bedrock/search/transformation.py
+++ b/litellm/llms/bedrock/search/transformation.py
@@ -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:
diff --git a/litellm/llms/chatgpt/responses/transformation.py b/litellm/llms/chatgpt/responses/transformation.py
index 040b7b8c1d8..72b5cde21d1 100644
--- a/litellm/llms/chatgpt/responses/transformation.py
+++ b/litellm/llms/chatgpt/responses/transformation.py
@@ -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
diff --git a/litellm/llms/codestral/completion/handler.py b/litellm/llms/codestral/completion/handler.py
index 6bebfe2d444..6b172e77295 100644
--- a/litellm/llms/codestral/completion/handler.py
+++ b/litellm/llms/codestral/completion/handler.py
@@ -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
diff --git a/litellm/llms/custom_httpx/aiohttp_handler.py b/litellm/llms/custom_httpx/aiohttp_handler.py
index 447d7b685f7..4060d84f7f4 100644
--- a/litellm/llms/custom_httpx/aiohttp_handler.py
+++ b/litellm/llms/custom_httpx/aiohttp_handler.py
@@ -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
diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py
index afa12b58b29..e7da6586b97 100644
--- a/litellm/llms/custom_httpx/llm_http_handler.py
+++ b/litellm/llms/custom_httpx/llm_http_handler.py
@@ -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),
diff --git a/litellm/llms/custom_httpx/sse_lines.py b/litellm/llms/custom_httpx/sse_lines.py
new file mode 100644
index 00000000000..71ea7d621eb
--- /dev/null
+++ b/litellm/llms/custom_httpx/sse_lines.py
@@ -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
diff --git a/litellm/llms/gemini/google_genai/guardrail_translation/handler.py b/litellm/llms/gemini/google_genai/guardrail_translation/handler.py
index 34007bc8fd4..b6b2aa6c28f 100644
--- a/litellm/llms/gemini/google_genai/guardrail_translation/handler.py
+++ b/litellm/llms/gemini/google_genai/guardrail_translation/handler.py
@@ -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]"
diff --git a/litellm/llms/huggingface/embedding/transformation.py b/litellm/llms/huggingface/embedding/transformation.py
index 1b7d6483739..6e6a117227e 100644
--- a/litellm/llms/huggingface/embedding/transformation.py
+++ b/litellm/llms/huggingface/embedding/transformation.py
@@ -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,
diff --git a/litellm/llms/langgraph/chat/sse_iterator.py b/litellm/llms/langgraph/chat/sse_iterator.py
index 575820bceee..6fe355ad98c 100644
--- a/litellm/llms/langgraph/chat/sse_iterator.py
+++ b/litellm/llms/langgraph/chat/sse_iterator.py
@@ -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:
diff --git a/litellm/llms/oobabooga/chat/oobabooga.py b/litellm/llms/oobabooga/chat/oobabooga.py
index 8879c20c5eb..b4dc57ff5ed 100644
--- a/litellm/llms/oobabooga/chat/oobabooga.py
+++ b/litellm/llms/oobabooga/chat/oobabooga.py
@@ -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,
diff --git a/litellm/llms/openai_like/chat/handler.py b/litellm/llms/openai_like/chat/handler.py
index 2c74ef19e7d..a450947afdf 100644
--- a/litellm/llms/openai_like/chat/handler.py
+++ b/litellm/llms/openai_like/chat/handler.py
@@ -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(
diff --git a/litellm/llms/predibase/chat/handler.py b/litellm/llms/predibase/chat/handler.py
index d0a61ea6e00..2ba4a63e08b 100644
--- a/litellm/llms/predibase/chat/handler.py
+++ b/litellm/llms/predibase/chat/handler.py
@@ -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,
diff --git a/litellm/llms/vertex_ai/agent_engine/transformation.py b/litellm/llms/vertex_ai/agent_engine/transformation.py
index 9f592398f74..0339533e1f5 100644
--- a/litellm/llms/vertex_ai/agent_engine/transformation.py
+++ b/litellm/llms/vertex_ai/agent_engine/transformation.py
@@ -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,
)
diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
index f27593dc94e..c506192f253 100644
--- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
+++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
@@ -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
diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json
index dc043dc4d26..8719565a96f 100644
--- a/litellm/model_prices_and_context_window_backup.json
+++ b/litellm/model_prices_and_context_window_backup.json
@@ -28853,6 +28853,7 @@
"mode": "image_generation",
"output_cost_per_image": 0.039,
"output_cost_per_image_token": 3e-05,
+ "output_cost_per_image_token_batches": 1.5e-05,
"output_cost_per_reasoning_token": 2.5e-06,
"output_cost_per_token": 2.5e-06,
"output_cost_per_token_batches": 1.25e-06,
@@ -28910,6 +28911,7 @@
"mode": "image_generation",
"output_cost_per_image": 0.134,
"output_cost_per_image_token": 0.00012,
+ "output_cost_per_image_token_batches": 6e-05,
"output_cost_per_token": 1.2e-05,
"output_cost_per_token_above_200k_tokens": 1.8e-05,
"output_cost_per_token_above_200k_tokens_priority": 3.24e-05,
@@ -28963,12 +28965,13 @@
"mode": "image_generation",
"output_cost_per_image": 0.134,
"output_cost_per_image_token": 0.00012,
+ "output_cost_per_image_token_batches": 6e-05,
"output_cost_per_token": 1.2e-05,
"output_cost_per_token_priority": 2.16e-05,
"output_cost_per_token_above_200k_tokens": 1.8e-05,
"output_cost_per_token_above_200k_tokens_priority": 3.24e-05,
"output_cost_per_token_batches": 6e-06,
- "source": "https://ai.google.dev/gemini-api/docs/pricing",
+ "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/images/generations",
@@ -29014,6 +29017,7 @@
"mode": "image_generation",
"output_cost_per_image": 0.0672,
"output_cost_per_image_token": 6e-05,
+ "output_cost_per_image_token_batches": 3e-05,
"output_cost_per_token": 3e-06,
"output_cost_per_token_batches": 1.5e-06,
"output_cost_per_token_flex": 1.5e-06,
@@ -29059,9 +29063,10 @@
"mode": "image_generation",
"output_cost_per_image": 0.0672,
"output_cost_per_image_token": 6e-05,
+ "output_cost_per_image_token_batches": 3e-05,
"output_cost_per_token": 3e-06,
"output_cost_per_token_batches": 1.5e-06,
- "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models",
+ "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/images/generations",
@@ -29106,6 +29111,7 @@
"mode": "image_generation",
"output_cost_per_image": 0.0336,
"output_cost_per_image_token": 3e-05,
+ "output_cost_per_image_token_batches": 1.5e-05,
"output_cost_per_token": 1.5e-06,
"output_cost_per_token_batches": 7.5e-07,
"output_cost_per_token_flex": 7.5e-07,
@@ -29672,10 +29678,12 @@
},
"search_free_requests": {"count": 10000, "period": "day", "shared_group": "gemini-2.5-pro"},
"google_maps_grounding_cost_per_query": 0.025,
+ "input_cost_per_token_above_200k_tokens_batches": 1.25e-06,
"input_cost_per_token_above_200k_tokens_priority": 4.5e-06,
"input_cost_per_token_batches": 6.25e-07,
"input_cost_per_token_flex": 6.25e-07,
"input_cost_per_token_priority": 2.25e-06,
+ "output_cost_per_token_above_200k_tokens_batches": 7.5e-06,
"output_cost_per_token_above_200k_tokens_priority": 2.7e-05,
"output_cost_per_token_batches": 5e-06,
"output_cost_per_token_flex": 5e-06,
@@ -29690,6 +29698,7 @@
"google_maps_grounding_cost_per_query": 0.014,
"input_cost_per_token": 2e-06,
"input_cost_per_token_above_200k_tokens": 4e-06,
+ "input_cost_per_token_above_200k_tokens_batches": 2e-06,
"input_cost_per_token_above_200k_tokens_priority": 7.2e-06,
"input_cost_per_token_batches": 1e-06,
"input_cost_per_token_flex": 1e-06,
@@ -29701,6 +29710,7 @@
"output_cost_per_image": 0.00012,
"output_cost_per_token": 1.2e-05,
"output_cost_per_token_above_200k_tokens": 1.8e-05,
+ "output_cost_per_token_above_200k_tokens_batches": 9e-06,
"output_cost_per_token_above_200k_tokens_priority": 3.24e-05,
"output_cost_per_token_batches": 6e-06,
"output_cost_per_token_flex": 6e-06,
@@ -30229,6 +30239,7 @@
"google_maps_grounding_cost_per_query": 0.014,
"input_cost_per_token": 2e-06,
"input_cost_per_token_above_200k_tokens": 4e-06,
+ "input_cost_per_token_above_200k_tokens_batches": 2e-06,
"input_cost_per_token_above_200k_tokens_priority": 7.2e-06,
"input_cost_per_token_batches": 1e-06,
"input_cost_per_token_flex": 1e-06,
@@ -30240,6 +30251,7 @@
"output_cost_per_image": 0.00012,
"output_cost_per_token": 1.2e-05,
"output_cost_per_token_above_200k_tokens": 1.8e-05,
+ "output_cost_per_token_above_200k_tokens_batches": 9e-06,
"output_cost_per_token_above_200k_tokens_priority": 3.24e-05,
"output_cost_per_token_batches": 6e-06,
"output_cost_per_token_flex": 6e-06,
@@ -51910,6 +51922,7 @@
"mode": "image_generation",
"output_cost_per_image": 0.039,
"output_cost_per_image_token": 3e-05,
+ "output_cost_per_image_token_batches": 1.5e-05,
"output_cost_per_reasoning_token": 2.5e-06,
"output_cost_per_token": 2.5e-06,
"output_cost_per_token_batches": 1.25e-06,
@@ -51968,6 +51981,7 @@
"mode": "image_generation",
"output_cost_per_image": 0.134,
"output_cost_per_image_token": 0.00012,
+ "output_cost_per_image_token_batches": 6e-05,
"output_cost_per_token": 1.2e-05,
"output_cost_per_token_above_200k_tokens": 1.8e-05,
"output_cost_per_token_above_200k_tokens_priority": 3.24e-05,
@@ -52002,13 +52016,14 @@
"mode": "image_generation",
"output_cost_per_image": 0.134,
"output_cost_per_image_token": 0.00012,
+ "output_cost_per_image_token_batches": 6e-05,
"output_cost_per_token": 1.2e-05,
"output_cost_per_token_priority": 2.16e-05,
"output_cost_per_token_above_200k_tokens": 1.8e-05,
"output_cost_per_token_above_200k_tokens_priority": 3.24e-05,
"output_cost_per_token_batches": 6e-06,
"supports_reasoning": false,
- "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image",
+ "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/images/generations",
@@ -52034,6 +52049,7 @@
"mode": "image_generation",
"output_cost_per_image": 0.0672,
"output_cost_per_image_token": 6e-05,
+ "output_cost_per_image_token_batches": 3e-05,
"output_cost_per_token": 3e-06,
"output_cost_per_token_batches": 1.5e-06,
"output_cost_per_token_flex": 1.5e-06,
@@ -52061,10 +52077,11 @@
"mode": "image_generation",
"output_cost_per_image": 0.0672,
"output_cost_per_image_token": 6e-05,
+ "output_cost_per_image_token_batches": 3e-05,
"output_cost_per_token": 3e-06,
"output_cost_per_token_batches": 1.5e-06,
"supports_reasoning": false,
- "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models",
+ "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/images/generations",
@@ -52089,6 +52106,7 @@
"mode": "image_generation",
"output_cost_per_image": 0.0336,
"output_cost_per_image_token": 3e-05,
+ "output_cost_per_image_token_batches": 1.5e-05,
"output_cost_per_token": 1.5e-06,
"output_cost_per_token_batches": 7.5e-07,
"output_cost_per_token_flex": 7.5e-07,
@@ -74400,6 +74418,7 @@
"mode": "image_generation",
"output_cost_per_image": 0.0336,
"output_cost_per_image_token": 3e-05,
+ "output_cost_per_image_token_batches": 1.5e-05,
"output_cost_per_token": 7.5e-06,
"output_cost_per_token_batches": 3.75e-06,
"output_cost_per_token_priority": 1.35e-05,
@@ -74445,6 +74464,7 @@
"mode": "image_generation",
"output_cost_per_image": 0.0336,
"output_cost_per_image_token": 3e-05,
+ "output_cost_per_image_token_batches": 1.5e-05,
"output_cost_per_token": 7.5e-06,
"output_cost_per_token_batches": 3.75e-06,
"output_cost_per_token_priority": 1.35e-05,
diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py
index 4691c7b0bc5..219ce906340 100644
--- a/litellm/proxy/agent_endpoints/a2a_endpoints.py
+++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py
@@ -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
diff --git a/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py b/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py
index f748a754fe0..32e0f6e8682 100644
--- a/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py
+++ b/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py
@@ -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
diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py
index 70e305f46fb..9bb451495c8 100644
--- a/litellm/proxy/common_request_processing.py
+++ b/litellm/proxy/common_request_processing.py
@@ -139,8 +139,8 @@ from litellm.proxy.utils import ( # noqa: F401 # legacy module exports
)
from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect, spawn_logging_task
from litellm.router import Router
+from litellm.router_strategy.auto_router.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX
from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict
-from litellm.router_utils.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX
from litellm.router_utils.common_utils import resolve_model_group_alias, resolve_served_model
from litellm.types.extraction import ExtractionResponse
from litellm.types.guardrails import GuardrailEventHooks
diff --git a/litellm/proxy/guardrails/auto_router_compression.py b/litellm/proxy/guardrails/auto_router_compression.py
index 1267fb15db8..3071bcbb1dc 100644
--- a/litellm/proxy/guardrails/auto_router_compression.py
+++ b/litellm/proxy/guardrails/auto_router_compression.py
@@ -14,7 +14,7 @@ from typing import TYPE_CHECKING, Final
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket
-from litellm.router_utils.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX
+from litellm.router_strategy.auto_router.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX
from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py
index 0d2f8dc9a0c..be32405154c 100644
--- a/litellm/proxy/health_check.py
+++ b/litellm/proxy/health_check.py
@@ -27,7 +27,7 @@ from litellm.constants import (
HEALTH_CHECK_TIMEOUT_SECONDS,
)
from litellm.litellm_core_utils.health_check_helpers import native_health_check_mode
-from litellm.router_utils.auto_router_model_naming import (
+from litellm.router_strategy.auto_router.auto_router_model_naming import (
StrategyRouterDependency,
classify_strategy_router_model,
strategy_router_dependencies,
diff --git a/litellm/proxy/hooks/autorouter_baseline_cache.py b/litellm/proxy/hooks/autorouter_baseline_cache.py
index c4232609124..1147ce7ec12 100644
--- a/litellm/proxy/hooks/autorouter_baseline_cache.py
+++ b/litellm/proxy/hooks/autorouter_baseline_cache.py
@@ -37,8 +37,8 @@ from litellm.proxy.spend_tracking.savings import (
_effective_model_info, # pyright: ignore[reportPrivateUsage] # existing deployment-price owner
_proxy_llm_router, # pyright: ignore[reportPrivateUsage] # existing optional proxy-router owner
)
+from litellm.router_strategy.auto_router.baseline_request import baseline_request
from litellm.router_strategy.complexity_router.context_compaction import compaction_applied
-from litellm.router_utils.baseline_request import baseline_request
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.router import BaselineRouteStamp
from litellm.types.utils import CallTypes, ModelInfo, Usage
diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py
index 11a47b2d424..1c4a61cfad8 100644
--- a/litellm/proxy/hooks/model_max_budget_limiter.py
+++ b/litellm/proxy/hooks/model_max_budget_limiter.py
@@ -19,6 +19,7 @@ from litellm._internal_context import with_service_target
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
from litellm.integrations.custom_logger import Span
+from litellm.litellm_core_utils.batch_utils import batch_charged_once_marker_key, batch_id_to_charge_once
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
from litellm.llms.bedrock.common_utils import get_bedrock_base_model, get_bedrock_cross_region_inference_regions
@@ -30,7 +31,6 @@ from litellm.proxy.hooks.model_budget_shared_windows import (
SharedWindows,
)
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting, budget_window_seconds
-from litellm.router_utils.batch_utils import batch_charged_once_marker_key, batch_id_to_charge_once
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import BudgetConfig, LlmProviders, StandardLoggingPayload
diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py
index d4528bc85e8..2c05f724374 100644
--- a/litellm/proxy/litellm_pre_call_utils.py
+++ b/litellm/proxy/litellm_pre_call_utils.py
@@ -3290,7 +3290,7 @@ def _per_user_credential_names_for_groups(
GITHUB_COPILOT_AUTH_TYPE_KEY,
GITHUB_COPILOT_PER_USER_AUTH_TYPE,
)
- from litellm.router_utils.auto_router_model_naming import classify_strategy_router_model
+ from litellm.router_strategy.auto_router.auto_router_model_naming import classify_strategy_router_model
names: Final[list[str]] = [] # mutable-ok: accumulates one name per per-user deployment
for group in model_groups:
diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py
index 4f48f3c57b9..3dceab6b7f7 100644
--- a/litellm/proxy/management_endpoints/auto_router_endpoints.py
+++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py
@@ -60,16 +60,16 @@ from litellm.repositories.table_repositories import PrismaTableRepository
from litellm.repositories.team_repository import TeamRepository
from litellm.repositories.user_repository import UserRepository
from litellm.repositories.verification_token_repository import VerificationTokenRepository
+from litellm.router_strategy.auto_router.auto_router_model_naming import (
+ StrategyRouterDependencyRole,
+ classify_strategy_router_model,
+ strategy_router_dependencies,
+)
from litellm.router_strategy.complexity_router import ComplexityRouter
from litellm.router_strategy.complexity_router.request_models import (
AutoRouterRoutingTestRequest,
RequestComplexityRouterConfig,
)
-from litellm.router_utils.auto_router_model_naming import (
- StrategyRouterDependencyRole,
- classify_strategy_router_model,
- strategy_router_dependencies,
-)
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.management_endpoints.auto_router_endpoints import (
SHADOW_EVAL_TURN_VALVE,
@@ -411,7 +411,7 @@ async def validate_complexity_router_config(
auto_router_environment_reference_error(data.complexity_router_config, user_api_key_dict)
)
- from litellm.router_utils.auto_router_model_naming import (
+ from litellm.router_strategy.auto_router.auto_router_model_naming import (
validate_complexity_router_config_write,
)
diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py
index 45006b14df6..8cd3559ba86 100644
--- a/litellm/proxy/management_endpoints/model_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/model_management_endpoints.py
@@ -131,6 +131,23 @@ from litellm.repositories.project_repository import ProjectRepository
from litellm.repositories.table_repositories import ModelTableRepository
from litellm.repositories.team_repository import TeamRepository
from litellm.router import Router
+from litellm.router_strategy.auto_router.auto_router_model_naming import (
+ GATED_AUTO_ROUTER_CAPABILITIES,
+ STRATEGY_ROUTER_PARAM_FIELDS,
+ GatedAutoRouterCapability,
+ capability_limit_violation,
+ carries_complexity_router_settings,
+ count_capability_routers,
+ gated_capabilities_of,
+ is_complexity_router_model,
+ validate_complexity_router_config_placement,
+ validate_complexity_router_config_write,
+ validate_strategy_router_model_write,
+)
+from litellm.router_strategy.auto_router.auto_router_tuning_baseline import (
+ is_mutable_tuned_candidate,
+ tuning_quota_violation,
+)
from litellm.router_strategy.complexity_router import (
DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE,
ClassificationRubric,
@@ -144,20 +161,6 @@ from litellm.router_strategy.complexity_router import (
normalize_classification_prompt,
)
from litellm.router_strategy.complexity_router.config import resolve_complexity_router_config_write
-from litellm.router_utils.auto_router_model_naming import (
- GATED_AUTO_ROUTER_CAPABILITIES,
- STRATEGY_ROUTER_PARAM_FIELDS,
- GatedAutoRouterCapability,
- capability_limit_violation,
- carries_complexity_router_settings,
- count_capability_routers,
- gated_capabilities_of,
- is_complexity_router_model,
- validate_complexity_router_config_placement,
- validate_complexity_router_config_write,
- validate_strategy_router_model_write,
-)
-from litellm.router_utils.auto_router_tuning_baseline import is_mutable_tuned_candidate, tuning_quota_violation
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.llms.bedrock import AwsSessionTag
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
diff --git a/litellm/proxy/management_helpers/auto_router_availability.py b/litellm/proxy/management_helpers/auto_router_availability.py
index 9b511ce0817..ce7859aff2d 100644
--- a/litellm/proxy/management_helpers/auto_router_availability.py
+++ b/litellm/proxy/management_helpers/auto_router_availability.py
@@ -7,14 +7,14 @@ from typing import Final
from pydantic import Json, TypeAdapter, ValidationError
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper
-from litellm.router_utils.auto_router_model_naming import (
+from litellm.router_strategy.auto_router.auto_router_model_naming import (
GATED_AUTO_ROUTER_CAPABILITIES,
capability_limit_violation,
classify_strategy_router_model,
count_capability_routers,
gated_capabilities_of,
)
-from litellm.router_utils.auto_router_tuning_baseline import (
+from litellm.router_strategy.auto_router.auto_router_tuning_baseline import (
is_mutable_tuned_candidate,
mutable_tuned_identities,
tuning_quota_violation,
diff --git a/litellm/proxy/management_helpers/auto_router_permissions.py b/litellm/proxy/management_helpers/auto_router_permissions.py
index 995446db7a6..8597f0d3667 100644
--- a/litellm/proxy/management_helpers/auto_router_permissions.py
+++ b/litellm/proxy/management_helpers/auto_router_permissions.py
@@ -34,13 +34,16 @@ from litellm.repositories.prisma_protocols import DatabaseClient
from litellm.repositories.project_repository import ProjectRepository
from litellm.repositories.table_repositories import TeamMembershipRepository
from litellm.router import Router
+from litellm.router_strategy.auto_router.auto_router_model_naming import (
+ classify_strategy_router_model,
+ strategy_router_dependencies,
+)
from litellm.router_strategy.complexity_router.config import (
ComplexityRouterConfigWrite,
normalize_classifier_config_aliases,
resolve_complexity_router_config_write,
)
from litellm.router_strategy.complexity_router.request_models import RequestComplexityRouterConfig
-from litellm.router_utils.auto_router_model_naming import classify_strategy_router_model, strategy_router_dependencies
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.router import Deployment, updateDeployment
diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py
index 7c7eb51bc5d..5415314f29b 100644
--- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py
+++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py
@@ -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:
diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py
index 2afc1d17bc2..0c4a3f64852 100644
--- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py
+++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py
@@ -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:
diff --git a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py
index 9790c150ae1..cae8ccc93ee 100644
--- a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py
+++ b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py
@@ -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:
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 70e1c25a70a..b263412ce0f 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -168,12 +168,7 @@ from litellm.proxy.common_utils.token_count_settings import (
)
from litellm.proxy.management_helpers.auto_router_availability import AutoRouterCatalogEntry, build_auto_router_catalog
from litellm.proxy.management_helpers.model_group_custom_fields import model_group_custom_fields
-from litellm.router_utils.access_windows import access_windows_config_error
-from litellm.router_utils.add_retry_fallback_headers import (
- get_fallback_errors_from_headers,
- get_hidden_params_dict,
-)
-from litellm.router_utils.auto_router_model_naming import (
+from litellm.router_strategy.auto_router.auto_router_model_naming import (
GATED_AUTO_ROUTER_CAPABILITIES,
STRATEGY_ROUTER_PARAM_FIELDS,
capability_limit_violation,
@@ -181,12 +176,17 @@ from litellm.router_utils.auto_router_model_naming import (
count_capability_routers,
validate_complexity_router_config_placement,
)
-from litellm.router_utils.auto_router_tuning_baseline import (
+from litellm.router_strategy.auto_router.auto_router_tuning_baseline import (
TUNING_BASELINE_PARAM_NAME,
mutable_tuned_identities,
snapshot_tuning_baselines,
tuning_limit_violation,
)
+from litellm.router_utils.access_windows import access_windows_config_error
+from litellm.router_utils.add_retry_fallback_headers import (
+ get_fallback_errors_from_headers,
+ get_hidden_params_dict,
+)
from litellm.router_utils.common_utils import resolve_model_group_alias
from litellm.router_utils.routing_groups import parse_routing_groups
from litellm.types.caching import RedisPipelineIncrementOperation
diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py
index 5f6a7432afb..1eb76f8f558 100644
--- a/litellm/proxy/response_polling/background_streaming.py
+++ b/litellm/proxy/response_polling/background_streaming.py
@@ -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:
diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py
index 219944de521..982158040c4 100644
--- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py
+++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py
@@ -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)
diff --git a/litellm/router.py b/litellm/router.py
index 5ca09499c3b..fc888495e99 100644
--- a/litellm/router.py
+++ b/litellm/router.py
@@ -7,6 +7,29 @@
#
# Thank you ! We ❤️ you! - Krrish & Ishaan
+"""
+litellm.Router: load balancing, retries, fallbacks and cooldowns across deployments.
+
+Public surface. The router is moving behind a narrow interface; code outside
+the router should end up using a Router through two kinds of members only:
+
+- endpoint methods (acompletion, aembedding, aresponses, aanthropic_messages,
+ the files, batches, assistants and vector store wrappers, ...), each a thin
+ wrapper around one call;
+- entry points (being added) that replace direct access to internals: registry() for
+ deployment lookups, settings() for configuration reads, health() for
+ cooldowns, blocks and usage, select() to pick a deployment without calling
+ it, apply() for runtime changes, observe() for routing events,
+ register_plugin() for strategies and sub-routers, and retire() to release
+ the router.
+
+Everything else on the class is internal. The members outside code uses today
+are listed in tests/code_coverage_tests/router_surface_allowlist.json, and a CI
+check (router_surface_check.py next to it) fails on any new outside use, so the
+list only shrinks while callers move to the entry points. The same interface
+is documented for the Rust core in litellm-rust/crates/router-core.
+"""
+
import asyncio
import contextlib
import copy
@@ -85,6 +108,15 @@ from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.otel.routing import routing_decision_attributes
from litellm.integrations.otel.runtime import phase_attributes, phase_event, phase_span
from litellm.litellm_core_utils.asyncify import run_async_function
+from litellm.litellm_core_utils.batch_utils import (
+ batch_create_routing,
+ deployment_supported_endpoints,
+ get_router_metadata_variable_name,
+ is_batch_retrieve_call_type,
+ replace_model_in_jsonl,
+ should_replace_model_in_jsonl,
+ without_deployments_that_cannot_batch,
+)
from litellm.litellm_core_utils.completion_timeout import CompletionTimeout
from litellm.litellm_core_utils.core_helpers import (
coerce_token_limit,
@@ -148,6 +180,15 @@ from litellm.llms.openai_like.model_info import (
get_openai_compatible_model_info,
)
from litellm.responses.bridge_check import keep_registered_mode
+from litellm.router_strategy.auto_router.auto_router_model_naming import (
+ AUTO_ROUTER_MODEL_PREFIX,
+ STRATEGY_ROUTER_PARAM_FIELDS,
+ GatedAutoRouterCapability,
+ capability_limit_violation,
+ claimed_capabilities,
+ classify_strategy_router_model,
+ count_capability_routers,
+)
from litellm.router_strategy.base_routing_strategy import BaseRoutingStrategy
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
from litellm.router_strategy.complexity_router.context_compaction import (
@@ -187,24 +228,6 @@ from litellm.router_utils.add_retry_fallback_headers import (
replace_complexity_router_headers,
response_total_token_count,
)
-from litellm.router_utils.auto_router_model_naming import (
- AUTO_ROUTER_MODEL_PREFIX,
- STRATEGY_ROUTER_PARAM_FIELDS,
- GatedAutoRouterCapability,
- capability_limit_violation,
- claimed_capabilities,
- classify_strategy_router_model,
- count_capability_routers,
-)
-from litellm.router_utils.batch_utils import (
- batch_create_routing,
- deployment_supported_endpoints,
- get_router_metadata_variable_name,
- is_batch_retrieve_call_type,
- replace_model_in_jsonl,
- should_replace_model_in_jsonl,
- without_deployments_that_cannot_batch,
-)
from litellm.router_utils.client_initalization_utils import InitalizeCachedClient, MaxParallelRequestsLimit
from litellm.router_utils.clientside_credential_handler import (
get_dynamic_litellm_params,
@@ -1370,7 +1393,7 @@ class Router:
},
]
- router = Router(model_list=model_list, fallbacks=[{"azure-gpt-3.5-turbo": "openai-gpt-3.5-turbo"}])
+ router = Router(model_list=model_list, fallbacks=[{"azure-gpt-3.5-turbo": ["openai-gpt-3.5-turbo"]}])
```
"""
@@ -1596,6 +1619,10 @@ class Router:
if default_fallbacks is not None or litellm.default_fallbacks is not None:
_fallbacks = default_fallbacks or litellm.default_fallbacks
+ if isinstance(_fallbacks, str):
+ raise TypeError(
+ f"default_fallbacks must be a list of model names, got '{_fallbacks}'. Use ['{_fallbacks}']"
+ )
if self.fallbacks is not None:
self.fallbacks.append({"*": _fallbacks})
else:
@@ -2839,7 +2866,7 @@ class Router:
self._initialize_core_endpoints()
self._initialize_specialized_endpoints()
- def validate_fallbacks(self, fallback_param: list | None):
+ def validate_fallbacks(self, fallback_param: Sequence[Mapping[str, object]] | None):
"""
Validate the fallbacks parameter.
"""
@@ -2852,6 +2879,13 @@ class Router:
raise ValueError(
f"Dictionary '{fallback_dict}' must have exactly one key, but has {len(fallback_dict)} keys."
)
+ ((model_group, targets),) = fallback_dict.items()
+ if isinstance(targets, list) or check_non_standard_fallback_format(fallbacks=[fallback_dict]):
+ continue
+ raise TypeError(
+ f"Fallbacks for '{model_group}' must be a list of model names, got '{targets}'. "
+ f"Use {{'{model_group}': ['{targets}']}}"
+ )
def _add_encrypted_content_affinity_check(self, enable_global_affinity: bool) -> None:
def _move_before_deployment_affinity(
@@ -16260,7 +16294,7 @@ class Router:
to the deployment that actually served the request. Every attempt therefore
writes or clears, never just writes.
"""
- from litellm.router_utils.baseline_request import capture_baseline_parameters
+ from litellm.router_strategy.auto_router.baseline_request import capture_baseline_parameters
from litellm.types.router import BaselineRouteStamp
phase_attributes(routing_decision_attributes(routing_decision))
diff --git a/litellm/router_utils/auto_router_model_naming.py b/litellm/router_strategy/auto_router/auto_router_model_naming.py
similarity index 100%
rename from litellm/router_utils/auto_router_model_naming.py
rename to litellm/router_strategy/auto_router/auto_router_model_naming.py
diff --git a/litellm/router_utils/auto_router_tuning_baseline.py b/litellm/router_strategy/auto_router/auto_router_tuning_baseline.py
similarity index 100%
rename from litellm/router_utils/auto_router_tuning_baseline.py
rename to litellm/router_strategy/auto_router/auto_router_tuning_baseline.py
diff --git a/litellm/router_utils/baseline_request.py b/litellm/router_strategy/auto_router/baseline_request.py
similarity index 100%
rename from litellm/router_utils/baseline_request.py
rename to litellm/router_strategy/auto_router/baseline_request.py
diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py
index eab232eae18..1d502dd8d95 100644
--- a/litellm/router_strategy/budget_limiter.py
+++ b/litellm/router_strategy/budget_limiter.py
@@ -36,12 +36,12 @@ from litellm._logging import verbose_router_logger
from litellm.caching.caching import DualCache
from litellm.caching.redis_cache import RedisCache, RedisPipelineIncrementOperation, log_redis_failure
from litellm.integrations.custom_logger import CustomLogger, Span
+from litellm.litellm_core_utils.batch_utils import batch_charged_once_marker_key, batch_id_to_charge_once
from litellm.litellm_core_utils.core_helpers import (
get_metadata_variable_name_from_kwargs,
)
from litellm.litellm_core_utils.duration_parser import duration_in_seconds_from
from litellm.router_strategy.tag_based_routing import get_tags_from_request_kwargs
-from litellm.router_utils.batch_utils import batch_charged_once_marker_key, batch_id_to_charge_once
from litellm.router_utils.cooldown_callbacks import (
get_prometheus_logger_from_callbacks,
)
diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py
index 6f3e0936641..8e0c3d6a7e7 100644
--- a/litellm/router_strategy/least_busy.py
+++ b/litellm/router_strategy/least_busy.py
@@ -10,7 +10,7 @@ from litellm._logging import verbose_router_logger
from litellm.caching.caching import DualCache
from litellm.caching.redis_cache import log_redis_failure
from litellm.integrations.custom_logger import CustomLogger
-from litellm.router_utils.batch_utils import is_batch_retrieve_call_type
+from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type
IN_FLIGHT_COUNT_TTL_SECONDS: Final = 60 * 60
diff --git a/litellm/router_strategy/lowest_cost.py b/litellm/router_strategy/lowest_cost.py
index 5894da9cb4f..17d03fd4ff8 100644
--- a/litellm/router_strategy/lowest_cost.py
+++ b/litellm/router_strategy/lowest_cost.py
@@ -10,7 +10,7 @@ from litellm._internal_context import with_service_target
from litellm._logging import verbose_router_logger
from litellm.caching.caching import DualCache
from litellm.integrations.custom_logger import CustomLogger
-from litellm.router_utils.batch_utils import is_batch_retrieve_call_type
+from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type
from litellm.router_utils.deployment_pricing import deployment_pricing_models, str_keyed
diff --git a/litellm/router_strategy/lowest_latency.py b/litellm/router_strategy/lowest_latency.py
index 4e994d3fdc6..c352934db6c 100644
--- a/litellm/router_strategy/lowest_latency.py
+++ b/litellm/router_strategy/lowest_latency.py
@@ -13,9 +13,9 @@ from litellm import ModelResponse, token_counter, verbose_logger
from litellm._internal_context import with_service_target
from litellm.caching.caching import DualCache
from litellm.integrations.custom_logger import CustomLogger
+from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type
from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs, safe_divide_seconds
from litellm.responses.utils import ResponseAPILoggingUtils
-from litellm.router_utils.batch_utils import is_batch_retrieve_call_type
from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.types.utils import LiteLLMPydanticObjectBase
diff --git a/litellm/router_strategy/lowest_tpm_rpm.py b/litellm/router_strategy/lowest_tpm_rpm.py
index 2d373e0c266..9636d7f009d 100644
--- a/litellm/router_strategy/lowest_tpm_rpm.py
+++ b/litellm/router_strategy/lowest_tpm_rpm.py
@@ -9,7 +9,7 @@ from litellm._internal_context import with_service_target
from litellm._logging import verbose_router_logger
from litellm.caching.caching import DualCache
from litellm.integrations.custom_logger import CustomLogger
-from litellm.router_utils.batch_utils import is_batch_retrieve_call_type
+from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type
from litellm.types.utils import LiteLLMPydanticObjectBase
from litellm.utils import print_verbose
diff --git a/litellm/router_strategy/lowest_tpm_rpm_v2.py b/litellm/router_strategy/lowest_tpm_rpm_v2.py
index 527ea830d55..0719b30d36d 100644
--- a/litellm/router_strategy/lowest_tpm_rpm_v2.py
+++ b/litellm/router_strategy/lowest_tpm_rpm_v2.py
@@ -15,8 +15,8 @@ from litellm._internal_context import with_service_target
from litellm._logging import verbose_logger, verbose_router_logger
from litellm.caching.caching import DualCache
from litellm.integrations.custom_logger import CustomLogger
+from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type
from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs
-from litellm.router_utils.batch_utils import is_batch_retrieve_call_type
from litellm.types.router import RouterErrors
from litellm.types.utils import LiteLLMPydanticObjectBase, StandardLoggingPayload
from litellm.utils import get_utc_datetime, print_verbose
diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py
index 0a77ad5b9ea..6d1de1f9f6c 100644
--- a/litellm/router_utils/fallback_event_handlers.py
+++ b/litellm/router_utils/fallback_event_handlers.py
@@ -13,6 +13,7 @@ from pydantic import TypeAdapter
import litellm
from litellm._logging import verbose_router_logger
from litellm.integrations.custom_logger import CustomLogger
+from litellm.litellm_core_utils.batch_utils import get_router_metadata_variable_name
from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs, safe_deep_copy
from litellm.litellm_core_utils.credential_accessor import reject_api_key_secret_reference
from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider
@@ -22,7 +23,6 @@ from litellm.router_utils.add_retry_fallback_headers import (
add_fallback_headers_to_response,
get_fallback_error_info,
)
-from litellm.router_utils.batch_utils import get_router_metadata_variable_name
from litellm.router_utils.cooldown_handlers import (
_first_present, # pyright: ignore[reportPrivateUsage] - shared internal helper, used across router_utils
cast_exception_status_to_int,
diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json
index dc043dc4d26..8719565a96f 100644
--- a/model_prices_and_context_window.json
+++ b/model_prices_and_context_window.json
@@ -28853,6 +28853,7 @@
"mode": "image_generation",
"output_cost_per_image": 0.039,
"output_cost_per_image_token": 3e-05,
+ "output_cost_per_image_token_batches": 1.5e-05,
"output_cost_per_reasoning_token": 2.5e-06,
"output_cost_per_token": 2.5e-06,
"output_cost_per_token_batches": 1.25e-06,
@@ -28910,6 +28911,7 @@
"mode": "image_generation",
"output_cost_per_image": 0.134,
"output_cost_per_image_token": 0.00012,
+ "output_cost_per_image_token_batches": 6e-05,
"output_cost_per_token": 1.2e-05,
"output_cost_per_token_above_200k_tokens": 1.8e-05,
"output_cost_per_token_above_200k_tokens_priority": 3.24e-05,
@@ -28963,12 +28965,13 @@
"mode": "image_generation",
"output_cost_per_image": 0.134,
"output_cost_per_image_token": 0.00012,
+ "output_cost_per_image_token_batches": 6e-05,
"output_cost_per_token": 1.2e-05,
"output_cost_per_token_priority": 2.16e-05,
"output_cost_per_token_above_200k_tokens": 1.8e-05,
"output_cost_per_token_above_200k_tokens_priority": 3.24e-05,
"output_cost_per_token_batches": 6e-06,
- "source": "https://ai.google.dev/gemini-api/docs/pricing",
+ "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/images/generations",
@@ -29014,6 +29017,7 @@
"mode": "image_generation",
"output_cost_per_image": 0.0672,
"output_cost_per_image_token": 6e-05,
+ "output_cost_per_image_token_batches": 3e-05,
"output_cost_per_token": 3e-06,
"output_cost_per_token_batches": 1.5e-06,
"output_cost_per_token_flex": 1.5e-06,
@@ -29059,9 +29063,10 @@
"mode": "image_generation",
"output_cost_per_image": 0.0672,
"output_cost_per_image_token": 6e-05,
+ "output_cost_per_image_token_batches": 3e-05,
"output_cost_per_token": 3e-06,
"output_cost_per_token_batches": 1.5e-06,
- "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models",
+ "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/images/generations",
@@ -29106,6 +29111,7 @@
"mode": "image_generation",
"output_cost_per_image": 0.0336,
"output_cost_per_image_token": 3e-05,
+ "output_cost_per_image_token_batches": 1.5e-05,
"output_cost_per_token": 1.5e-06,
"output_cost_per_token_batches": 7.5e-07,
"output_cost_per_token_flex": 7.5e-07,
@@ -29672,10 +29678,12 @@
},
"search_free_requests": {"count": 10000, "period": "day", "shared_group": "gemini-2.5-pro"},
"google_maps_grounding_cost_per_query": 0.025,
+ "input_cost_per_token_above_200k_tokens_batches": 1.25e-06,
"input_cost_per_token_above_200k_tokens_priority": 4.5e-06,
"input_cost_per_token_batches": 6.25e-07,
"input_cost_per_token_flex": 6.25e-07,
"input_cost_per_token_priority": 2.25e-06,
+ "output_cost_per_token_above_200k_tokens_batches": 7.5e-06,
"output_cost_per_token_above_200k_tokens_priority": 2.7e-05,
"output_cost_per_token_batches": 5e-06,
"output_cost_per_token_flex": 5e-06,
@@ -29690,6 +29698,7 @@
"google_maps_grounding_cost_per_query": 0.014,
"input_cost_per_token": 2e-06,
"input_cost_per_token_above_200k_tokens": 4e-06,
+ "input_cost_per_token_above_200k_tokens_batches": 2e-06,
"input_cost_per_token_above_200k_tokens_priority": 7.2e-06,
"input_cost_per_token_batches": 1e-06,
"input_cost_per_token_flex": 1e-06,
@@ -29701,6 +29710,7 @@
"output_cost_per_image": 0.00012,
"output_cost_per_token": 1.2e-05,
"output_cost_per_token_above_200k_tokens": 1.8e-05,
+ "output_cost_per_token_above_200k_tokens_batches": 9e-06,
"output_cost_per_token_above_200k_tokens_priority": 3.24e-05,
"output_cost_per_token_batches": 6e-06,
"output_cost_per_token_flex": 6e-06,
@@ -30229,6 +30239,7 @@
"google_maps_grounding_cost_per_query": 0.014,
"input_cost_per_token": 2e-06,
"input_cost_per_token_above_200k_tokens": 4e-06,
+ "input_cost_per_token_above_200k_tokens_batches": 2e-06,
"input_cost_per_token_above_200k_tokens_priority": 7.2e-06,
"input_cost_per_token_batches": 1e-06,
"input_cost_per_token_flex": 1e-06,
@@ -30240,6 +30251,7 @@
"output_cost_per_image": 0.00012,
"output_cost_per_token": 1.2e-05,
"output_cost_per_token_above_200k_tokens": 1.8e-05,
+ "output_cost_per_token_above_200k_tokens_batches": 9e-06,
"output_cost_per_token_above_200k_tokens_priority": 3.24e-05,
"output_cost_per_token_batches": 6e-06,
"output_cost_per_token_flex": 6e-06,
@@ -51910,6 +51922,7 @@
"mode": "image_generation",
"output_cost_per_image": 0.039,
"output_cost_per_image_token": 3e-05,
+ "output_cost_per_image_token_batches": 1.5e-05,
"output_cost_per_reasoning_token": 2.5e-06,
"output_cost_per_token": 2.5e-06,
"output_cost_per_token_batches": 1.25e-06,
@@ -51968,6 +51981,7 @@
"mode": "image_generation",
"output_cost_per_image": 0.134,
"output_cost_per_image_token": 0.00012,
+ "output_cost_per_image_token_batches": 6e-05,
"output_cost_per_token": 1.2e-05,
"output_cost_per_token_above_200k_tokens": 1.8e-05,
"output_cost_per_token_above_200k_tokens_priority": 3.24e-05,
@@ -52002,13 +52016,14 @@
"mode": "image_generation",
"output_cost_per_image": 0.134,
"output_cost_per_image_token": 0.00012,
+ "output_cost_per_image_token_batches": 6e-05,
"output_cost_per_token": 1.2e-05,
"output_cost_per_token_priority": 2.16e-05,
"output_cost_per_token_above_200k_tokens": 1.8e-05,
"output_cost_per_token_above_200k_tokens_priority": 3.24e-05,
"output_cost_per_token_batches": 6e-06,
"supports_reasoning": false,
- "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image",
+ "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/images/generations",
@@ -52034,6 +52049,7 @@
"mode": "image_generation",
"output_cost_per_image": 0.0672,
"output_cost_per_image_token": 6e-05,
+ "output_cost_per_image_token_batches": 3e-05,
"output_cost_per_token": 3e-06,
"output_cost_per_token_batches": 1.5e-06,
"output_cost_per_token_flex": 1.5e-06,
@@ -52061,10 +52077,11 @@
"mode": "image_generation",
"output_cost_per_image": 0.0672,
"output_cost_per_image_token": 6e-05,
+ "output_cost_per_image_token_batches": 3e-05,
"output_cost_per_token": 3e-06,
"output_cost_per_token_batches": 1.5e-06,
"supports_reasoning": false,
- "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models",
+ "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/images/generations",
@@ -52089,6 +52106,7 @@
"mode": "image_generation",
"output_cost_per_image": 0.0336,
"output_cost_per_image_token": 3e-05,
+ "output_cost_per_image_token_batches": 1.5e-05,
"output_cost_per_token": 1.5e-06,
"output_cost_per_token_batches": 7.5e-07,
"output_cost_per_token_flex": 7.5e-07,
@@ -74400,6 +74418,7 @@
"mode": "image_generation",
"output_cost_per_image": 0.0336,
"output_cost_per_image_token": 3e-05,
+ "output_cost_per_image_token_batches": 1.5e-05,
"output_cost_per_token": 7.5e-06,
"output_cost_per_token_batches": 3.75e-06,
"output_cost_per_token_priority": 1.35e-05,
@@ -74445,6 +74464,7 @@
"mode": "image_generation",
"output_cost_per_image": 0.0336,
"output_cost_per_image_token": 3e-05,
+ "output_cost_per_image_token_batches": 1.5e-05,
"output_cost_per_token": 7.5e-06,
"output_cost_per_token_batches": 3.75e-06,
"output_cost_per_token_priority": 1.35e-05,
diff --git a/tests/benchmarks/test_sse_lines_benchmarks.py b/tests/benchmarks/test_sse_lines_benchmarks.py
new file mode 100644
index 00000000000..1320fae49a6
--- /dev/null
+++ b/tests/benchmarks/test_sse_lines_benchmarks.py
@@ -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
diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py
index 4c9afd26dad..9fb3633d53c 100644
--- a/tests/code_coverage_tests/router_code_coverage.py
+++ b/tests/code_coverage_tests/router_code_coverage.py
@@ -187,13 +187,13 @@ ignored_function_names = [
def main():
router_file = [
"./litellm/router.py",
- "./litellm/router_utils/batch_utils.py",
+ "./litellm/litellm_core_utils/batch_utils.py",
"./litellm/router_utils/pattern_match_deployments.py",
]
# router_file = [
# "../../litellm/router.py",
# "../../litellm/router_utils/pattern_match_deployments.py",
- # "../../litellm/router_utils/batch_utils.py",
+ # "../../litellm/litellm_core_utils/batch_utils.py",
# ] ## LOCAL TESTING
tests_dir = "./tests/" # Update this path if your tests directory is located elsewhere
# tests_dir = "../../tests/" # LOCAL TESTING
diff --git a/tests/code_coverage_tests/router_surface_allowlist.json b/tests/code_coverage_tests/router_surface_allowlist.json
new file mode 100644
index 00000000000..428fd9b4011
--- /dev/null
+++ b/tests/code_coverage_tests/router_surface_allowlist.json
@@ -0,0 +1,211 @@
+{
+ "_doc": "Router members and litellm.router_utils modules used outside the router. Shrink-only: see tests/code_coverage_tests/router_surface_check.py.",
+ "members_outside": [
+ "_get_client",
+ "_zero_cost_cache",
+ "a_add_message",
+ "aadapter_completion",
+ "aanthropic_messages",
+ "abatch_completion",
+ "abatch_completion_fastest_response",
+ "acancel_batch",
+ "acancel_fine_tuning_job",
+ "acompletion",
+ "acreate_assistants",
+ "acreate_batch",
+ "acreate_file",
+ "acreate_fine_tuning_job",
+ "acreate_thread",
+ "adaptive_routers",
+ "adecisions",
+ "adelete_assistant",
+ "aembedding",
+ "afile_content",
+ "afile_delete",
+ "aget_assistants",
+ "aget_messages",
+ "aget_responses",
+ "aget_thread",
+ "alist_batches",
+ "alist_fine_tuning_jobs",
+ "allm_passthrough_route",
+ "amoderation",
+ "aocr",
+ "are_all_deployments_blocked",
+ "arefresh_model_info",
+ "aretrieve_batch",
+ "aretrieve_fine_tuning_job",
+ "arm_routing_read_prefetch",
+ "arun_thread",
+ "async_get_available_deployment",
+ "async_get_fully_unhealthy_model_names",
+ "async_get_healthy_deployments",
+ "async_pre_routing_hook",
+ "auto_routers",
+ "avector_store_create",
+ "avector_store_search",
+ "background_health_check_model_groups",
+ "bind_to_request_logging",
+ "cache",
+ "cache_responses",
+ "cached_deployment_model_info",
+ "cached_model_group_info",
+ "completion",
+ "complexity_routers",
+ "config_deployments",
+ "content_policy_fallbacks",
+ "context_window_fallbacks",
+ "cooldown_time",
+ "default_deployment",
+ "default_litellm_params",
+ "delete_deployment",
+ "deployment_latency_map",
+ "deployment_names",
+ "deployment_usable_by_team",
+ "deployments_for_request",
+ "deployments_named",
+ "discard",
+ "embedding",
+ "enable_health_check_routing",
+ "fallbacks",
+ "generate_model_id",
+ "get_all_deployments",
+ "get_available_deployment_for_pass_through",
+ "get_available_guardrail",
+ "get_candidate_model_ids_for_route",
+ "get_configured_display_name",
+ "get_configured_mode",
+ "get_configured_service_tiers",
+ "get_configured_token_limits",
+ "get_credential_deployment",
+ "get_deployment",
+ "get_deployment_by_model_group_name",
+ "get_deployment_credentials",
+ "get_deployment_credentials_with_provider",
+ "get_deployment_model_info",
+ "get_discovered_model_info",
+ "get_fully_blocked_model_names",
+ "get_model_access_groups",
+ "get_model_from_alias",
+ "get_model_group_info",
+ "get_model_group_usage",
+ "get_model_ids",
+ "get_model_info",
+ "get_model_list",
+ "get_model_list_from_model_alias",
+ "get_model_list_from_routing_groups",
+ "get_model_list_of_routed_group",
+ "get_model_listing_info",
+ "get_model_names",
+ "get_remaining_model_group_usage",
+ "get_routable_deployments",
+ "get_routable_upstream_model",
+ "get_router_deployment_budget_limiter",
+ "get_routing_group",
+ "get_settings",
+ "get_valid_args",
+ "guardrail_list",
+ "has_model_id",
+ "health_check_ignore_transient_errors",
+ "health_state_cache",
+ "is_deployment_blocked",
+ "is_model_access_group_for_wildcard_route",
+ "is_recognized_model",
+ "kubernetes_pod_discovery",
+ "map_team_model",
+ "model_group_alias",
+ "model_list",
+ "model_name_to_deployment_indices",
+ "model_names",
+ "pattern_router",
+ "priority_header_update",
+ "prompt_caching_pin_cache",
+ "provider_budget_config",
+ "quality_routers",
+ "resolve_model_name_from_model_id",
+ "resolve_served_model_name",
+ "resolved_litellm_models",
+ "routable_model_group",
+ "router_general_settings",
+ "routing_groups",
+ "routing_plugins",
+ "schedule_acompletion",
+ "search_tools",
+ "serve_through_request_logging",
+ "should_include_deployment",
+ "split_request_priority",
+ "team_model_to_deployment_indices",
+ "team_pattern_routers",
+ "team_ptu_ceiling",
+ "team_public_model_names",
+ "team_public_model_names_of",
+ "update_redis_cache",
+ "update_settings",
+ "upsert_deployment",
+ "wildcard_route_deployments",
+ "wildcard_routed_models"
+ ],
+ "members_router_family": [
+ "_count_pre_call_check_tokens",
+ "_pre_call_checks",
+ "_update_kwargs_before_fallbacks",
+ "acompletion",
+ "adecisions",
+ "aembedding",
+ "allowed_fails",
+ "allowed_fails_policy",
+ "aresponses",
+ "async_function_with_fallbacks",
+ "async_get_healthy_deployments",
+ "cache",
+ "cached_deployment_model_info",
+ "cooldown_cache",
+ "cooldown_time",
+ "default_max_parallel_requests",
+ "deployments_named",
+ "disable_cooldowns",
+ "embedding",
+ "enable_pre_call_checks",
+ "enable_tag_filtering",
+ "explicit_allowed_fails",
+ "fallback_access_check",
+ "fallback_budget_check",
+ "fallbacks",
+ "get_all_deployments",
+ "get_allowed_fails_from_policy",
+ "get_deployment",
+ "get_model_group",
+ "get_model_ids",
+ "get_model_info",
+ "get_model_list",
+ "get_router_model_info",
+ "log_retry",
+ "model_list",
+ "provider_default_deployment_ids",
+ "routing_group_has_alternatives",
+ "routing_strategy",
+ "search_tools",
+ "set_custom_routing_strategy",
+ "slack_alerting_logger",
+ "tag_filtering_match_any",
+ "tag_routing_prefix",
+ "team_model_has_alternatives"
+ ],
+ "router_utils_imports_outside": [
+ "litellm.router_utils.access_windows",
+ "litellm.router_utils.add_retry_fallback_headers",
+ "litellm.router_utils.clientside_credential_handler",
+ "litellm.router_utils.common_utils",
+ "litellm.router_utils.cooldown_handlers",
+ "litellm.router_utils.deployment_pricing",
+ "litellm.router_utils.fallback_event_handlers",
+ "litellm.router_utils.get_retry_from_policy",
+ "litellm.router_utils.health_state_cache",
+ "litellm.router_utils.pattern_match_deployments",
+ "litellm.router_utils.ptu_shares",
+ "litellm.router_utils.reasoning_effort_capability",
+ "litellm.router_utils.router_callbacks.track_deployment_metrics",
+ "litellm.router_utils.routing_groups",
+ "litellm.router_utils.search_api_router"
+ ]
+}
diff --git a/tests/code_coverage_tests/router_surface_check.py b/tests/code_coverage_tests/router_surface_check.py
new file mode 100644
index 00000000000..0b35bac5240
--- /dev/null
+++ b/tests/code_coverage_tests/router_surface_check.py
@@ -0,0 +1,235 @@
+"""
+Fence for the Router's public surface.
+
+Code outside the router may use only the Router members listed in
+router_surface_allowlist.json, and may import only the litellm.router_utils
+modules listed there. The lists hold what is used today and may only shrink:
+
+- a member or module used outside the router that is not listed fails the
+ check (use an existing entry point instead, or discuss widening the surface);
+- a listed member or module that nothing outside the router uses any more
+ also fails, so the entry is removed in the same change.
+
+"Outside" means every non-test file under litellm/ and enterprise/ except
+litellm/router.py and the router's own modules. Uses from the router family
+(litellm/router_utils, litellm/router_strategy) are listed separately, because
+those modules become private parts of the router core in a later step.
+
+Member uses are found the same way the boundary survey found them: attribute
+access on a receiver whose name looks like a router (llm_router,
+litellm_router_instance, router, ...) or on the Router class, plus
+getattr/setattr/hasattr with a constant name on such a receiver. Members
+include the container file endpoints Router.__init__ registers by name.
+
+Run: python tests/code_coverage_tests/router_surface_check.py [--write]
+--write regenerates the allowlist from the current tree (for shrinking it).
+"""
+
+import ast
+import json
+import os
+import re
+import sys
+from typing import Dict, List, Set
+
+ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
+ALLOWLIST = os.path.join(os.path.dirname(__file__), "router_surface_allowlist.json")
+
+ROUTER_FILES = {
+ "litellm/router.py",
+}
+ROUTER_CLASS_FILES = ["litellm/router.py"]
+ROUTER_CLASS_NAMES = {"Router"}
+
+SCAN_ROOTS = ["litellm", "enterprise"]
+SKIP_DIRS = {
+ ".git",
+ "node_modules",
+ ".venv",
+ "__pycache__",
+ "litellm-rust",
+ "ui",
+ "docs",
+ "cookbook",
+ "tests",
+ "test",
+}
+
+RECEIVER = re.compile(
+ r"(^|[._])(llm_router|router|litellm_router|litellm_router_instance|router_instance"
+ r"|_router|proxy_router|llm_router_instance|router_obj|mock_router|Router)$"
+)
+# Router.__init__ registers these with setattr(self, name, ...), so they never appear as self..
+CONTAINER_FILE_ENDPOINTS = "litellm/containers/endpoints.json"
+DYNAMIC_FUNCS = {"getattr", "setattr", "hasattr"}
+
+
+def _area(rel: str) -> str:
+ if rel.startswith(("litellm/router_utils/", "litellm/router_strategy/")):
+ return "router_family"
+ return "outside"
+
+
+def _is_test(rel: str) -> bool:
+ name = os.path.basename(rel)
+ return "/tests/" in rel or name.startswith("test_") or name.endswith("_test.py")
+
+
+def router_members() -> Set[str]:
+ members: Set[str] = set()
+ for rel in ROUTER_CLASS_FILES:
+ with open(os.path.join(ROOT, rel), encoding="utf-8") as f:
+ tree = ast.parse(f.read())
+ for cls in tree.body:
+ if not isinstance(cls, ast.ClassDef) or cls.name not in ROUTER_CLASS_NAMES:
+ continue
+ for node in cls.body:
+ if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+ members.add(node.name)
+ elif isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name):
+ members.add(node.target.id)
+ elif isinstance(node, ast.Assign):
+ members.update(t.id for t in node.targets if isinstance(t, ast.Name))
+ for node in ast.walk(cls):
+ if (
+ isinstance(node, ast.Attribute)
+ and isinstance(node.value, ast.Name)
+ and node.value.id == "self"
+ and isinstance(node.ctx, ast.Store)
+ ):
+ members.add(node.attr)
+ with open(os.path.join(ROOT, CONTAINER_FILE_ENDPOINTS), encoding="utf-8") as f:
+ for endpoint in json.load(f)["endpoints"]:
+ members.update((endpoint["name"], endpoint["async_name"]))
+ return {m for m in members if not (m.startswith("__") and m.endswith("__"))}
+
+
+def _python_files():
+ for top in SCAN_ROOTS:
+ base = os.path.join(ROOT, top)
+ for dirpath, dirnames, filenames in os.walk(base):
+ dirnames[:] = [d for d in dirnames if d not in SKIP_DIRS]
+ for fn in filenames:
+ if fn.endswith(".py"):
+ path = os.path.join(dirpath, fn)
+ yield os.path.relpath(path, ROOT).replace(os.sep, "/"), path
+
+
+def scan() -> Dict[str, Dict[str, Set[str]]]:
+ """Return {"members": {area: {member: files}}, "router_utils_imports": {module: files}}."""
+ members = router_members()
+ uses: Dict[str, Dict[str, Set[str]]] = {"outside": {}, "router_family": {}}
+ imports: Dict[str, Set[str]] = {}
+ for rel, path in _python_files():
+ if rel in ROUTER_FILES or _is_test(rel):
+ continue
+ with open(path, encoding="utf-8") as f:
+ src = f.read()
+ has_router = "router" in src.lower()
+ if not has_router:
+ continue
+ try:
+ tree = ast.parse(src)
+ except SyntaxError:
+ continue
+ area = _area(rel)
+ for name in member_uses(tree, members):
+ uses[area].setdefault(name, set()).add(rel)
+ if area == "outside":
+ for node in ast.walk(tree):
+ for mod in _router_utils_modules(node):
+ imports.setdefault(mod, set()).add(rel)
+ return {"members": uses, "router_utils_imports": imports}
+
+
+def member_uses(tree: ast.AST, members: Set[str]) -> Set[str]:
+ """The Router members a module uses, on an instance or on the Router class itself."""
+ used: Set[str] = set()
+ for node in ast.walk(tree):
+ if isinstance(node, ast.Attribute) and node.attr in members:
+ if RECEIVER.search(ast.unparse(node.value)):
+ used.add(node.attr)
+ elif (
+ isinstance(node, ast.Call)
+ and isinstance(node.func, ast.Name)
+ and node.func.id in DYNAMIC_FUNCS
+ and len(node.args) >= 2
+ and isinstance(node.args[1], ast.Constant)
+ and node.args[1].value in members
+ and RECEIVER.search(ast.unparse(node.args[0]))
+ ):
+ used.add(node.args[1].value)
+ return used
+
+
+def _router_utils_modules(node: ast.AST) -> List[str]:
+ if isinstance(node, ast.ImportFrom) and node.module:
+ if node.module == "litellm.router_utils":
+ return ["litellm.router_utils." + a.name for a in node.names]
+ if node.module.startswith("litellm.router_utils."):
+ return [node.module]
+ if isinstance(node, ast.Import):
+ return [a.name for a in node.names if a.name.startswith("litellm.router_utils.")]
+ return []
+
+
+def build_allowlist(found) -> dict:
+ return {
+ "_doc": (
+ "Router members and litellm.router_utils modules used outside the router. "
+ "Shrink-only: see tests/code_coverage_tests/router_surface_check.py."
+ ),
+ "members_outside": sorted(found["members"]["outside"]),
+ "members_router_family": sorted(found["members"]["router_family"]),
+ "router_utils_imports_outside": sorted(found["router_utils_imports"]),
+ }
+
+
+def compare(allow: dict, found) -> List[str]:
+ errors: List[str] = []
+ sections = [
+ ("members_outside", found["members"]["outside"], "Router member"),
+ ("members_router_family", found["members"]["router_family"], "Router member (router family)"),
+ ("router_utils_imports_outside", found["router_utils_imports"], "litellm.router_utils import"),
+ ]
+ for key, used, what in sections:
+ listed = set(allow.get(key, []))
+ for name in sorted(set(used) - listed):
+ files = ", ".join(sorted(used[name]))
+ errors.append(
+ f"new {what} outside the router: {name!r} in {files}. The router surface may only "
+ f"shrink: use an entry point instead."
+ )
+ for name in sorted(listed - set(used)):
+ errors.append(
+ f"{what} {name!r} is listed in {os.path.basename(ALLOWLIST)} [{key}] but nothing uses it "
+ f"any more: remove it from the list."
+ )
+ return errors
+
+
+def main(argv: List[str]) -> int:
+ found = scan()
+ if "--write" in argv:
+ with open(ALLOWLIST, "w", encoding="utf-8") as f:
+ json.dump(build_allowlist(found), f, indent=2)
+ f.write("\n")
+ print(f"wrote {ALLOWLIST}")
+ return 0
+ with open(ALLOWLIST, encoding="utf-8") as f:
+ allow = json.load(f)
+ errors = compare(allow, found)
+ for e in errors:
+ print(e)
+ if errors:
+ return 1
+ print(
+ f"router surface ok: {len(allow['members_outside'])} members outside, "
+ f"{len(allow['members_router_family'])} from the router family, "
+ f"{len(allow['router_utils_imports_outside'])} router_utils modules imported outside"
+ )
+ return 0
+
+
+if __name__ == "__main__":
+ sys.exit(main(sys.argv[1:]))
diff --git a/tests/code_coverage_tests/unbounded_in_baseline.txt b/tests/code_coverage_tests/unbounded_in_baseline.txt
index 1da964681a9..d8c6e2241f4 100644
--- a/tests/code_coverage_tests/unbounded_in_baseline.txt
+++ b/tests/code_coverage_tests/unbounded_in_baseline.txt
@@ -129,4 +129,4 @@ litellm/proxy/utils.py prefetch_config_params prisma param_name.in `param_names`
litellm/repositories/daily_activity_repository.py DailyActivityRepository.daily_rows prisma [scope.entity_id_field].in `list(scope.entity_ids)` 0
litellm/repositories/daily_activity_repository.py DailyActivityRepository.daily_rows prisma api_key.in `list(scope.api_keys)` 0
litellm/repositories/daily_activity_repository.py DailyActivityRepository.daily_rows prisma not.in `list(scope.exclude_entity_ids)` 0
-litellm/router_utils/auto_router_model_naming.py raw-sql classifier_type.IN `IN ({_LLM_CLASSIFIER_TYPES_SQL})` 0
+litellm/router_strategy/auto_router/auto_router_model_naming.py raw-sql classifier_type.IN `IN ({_LLM_CLASSIFIER_TYPES_SQL})` 0
diff --git a/tests/integration/providers/test_gemini_grounding_annotations_wire.py b/tests/integration/providers/test_gemini_grounding_annotations_wire.py
new file mode 100644
index 00000000000..11df7c214f3
--- /dev/null
+++ b/tests/integration/providers/test_gemini_grounding_annotations_wire.py
@@ -0,0 +1,65 @@
+import json
+from typing import Final
+
+import pytest
+from integration._support.client import Gateway
+from integration._support.wire import Reply, Request, wire_server
+
+
+def _grounded_reply(request: Request) -> Reply:
+ body: Final = {
+ "candidates": [
+ {
+ "content": {"role": "model", "parts": [{"text": "The passage"}]},
+ "finishReason": "STOP",
+ "groundingMetadata": {
+ "groundingChunks": [
+ {"retrievedContext": {"uri": "gs://documents/source"}},
+ {"web": {"uri": "https://example.com/first", "title": "First"}},
+ {"web": {"uri": "https://example.com/second", "title": "Second"}},
+ ],
+ "groundingSupports": [
+ {"segment": {"startIndex": 4, "endIndex": 11}, "groundingChunkIndices": [0, 99, 2, 1]}
+ ],
+ },
+ }
+ ],
+ "usageMetadata": {"promptTokenCount": 5, "candidatesTokenCount": 2, "totalTokenCount": 7},
+ }
+ assert json.loads(request.body)["contents"] == [{"role": "user", "parts": [{"text": "Find sources"}]}]
+ if ":streamGenerateContent" in request.target:
+ return Reply(content_type="text/event-stream", chunks=(f"data: {json.dumps(body)}\n\n".encode(),))
+ return Reply(body=json.dumps(body).encode())
+
+
+@pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"])
+def test_chat_grounding_annotations_include_every_web_source(gateway: Gateway, stream: bool) -> None:
+ with wire_server(_grounded_reply) as wire, gateway.scenario() as scenario:
+ model: Final = scenario.model(model="gemini/gemini-3.8-flash", api_base=wire.url, api_key="synthetic-key")
+ response: Final = gateway.request(
+ "POST",
+ "/v1/chat/completions",
+ {"model": model, "messages": [{"role": "user", "content": "Find sources"}], "stream": stream},
+ )
+ assert response.status_code == 200, response.text
+ payloads: Final = (
+ [json.loads(line[6:]) for line in response.text.splitlines() if line.startswith("data: {")]
+ if stream
+ else [response.json()]
+ )
+ messages: Final = tuple(payload["choices"][0]["delta" if stream else "message"] for payload in payloads)
+ annotations: Final = tuple(message["annotations"] for message in messages if message.get("annotations"))
+ assert annotations == (
+ [
+ {
+ "type": "url_citation",
+ "url_citation": {
+ "start_index": 4,
+ "end_index": 11,
+ "url": f"https://example.com/{name.lower()}",
+ "title": name,
+ },
+ }
+ for name in ("Second", "First")
+ ],
+ ), response.text
diff --git a/tests/unit/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py b/tests/unit/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py
index 09f09170241..cd3f7622fa2 100644
--- a/tests/unit/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py
+++ b/tests/unit/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py
@@ -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."""
diff --git a/tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py b/tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py
index a300560ae9d..ac3f0390f84 100644
--- a/tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py
+++ b/tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py
@@ -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()
diff --git a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py
index 9aa494c1aab..46f1dbf2e7a 100644
--- a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py
+++ b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py
@@ -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,
diff --git a/tests/unit/google_genai/test_google_genai_streaming_iterator.py b/tests/unit/google_genai/test_google_genai_streaming_iterator.py
index 120c99b6e42..cbb9638b3d9 100644
--- a/tests/unit/google_genai/test_google_genai_streaming_iterator.py
+++ b/tests/unit/google_genai/test_google_genai_streaming_iterator.py
@@ -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 = {
diff --git a/tests/unit/interactions/test_interactions_streaming_iterator.py b/tests/unit/interactions/test_interactions_streaming_iterator.py
index afcd0d559d7..fb1fe92d637 100644
--- a/tests/unit/interactions/test_interactions_streaming_iterator.py
+++ b/tests/unit/interactions/test_interactions_streaming_iterator.py
@@ -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
diff --git a/tests/unit/router_utils/test_batch_utils.py b/tests/unit/litellm_core_utils/test_batch_utils.py
similarity index 97%
rename from tests/unit/router_utils/test_batch_utils.py
rename to tests/unit/litellm_core_utils/test_batch_utils.py
index 79fe86a37c0..cd37d635c43 100644
--- a/tests/unit/router_utils/test_batch_utils.py
+++ b/tests/unit/litellm_core_utils/test_batch_utils.py
@@ -7,7 +7,7 @@ from typing import Dict, Final, List
import pytest
import litellm
-from litellm.router_utils.batch_utils import (
+from litellm.litellm_core_utils.batch_utils import (
InMemoryFile,
get_router_metadata_variable_name,
parse_jsonl_with_embedded_newlines,
@@ -116,7 +116,7 @@ def test_router_metadata_variable_name():
def test_non_json_input():
"""Test that replace_model_in_jsonl returns original content for non-JSON input"""
- from litellm.router_utils.batch_utils import replace_model_in_jsonl
+ from litellm.litellm_core_utils.batch_utils import replace_model_in_jsonl
# Test with non-JSON string
non_json_str = "This is not a JSON string"
@@ -138,7 +138,7 @@ def test_non_json_input():
def test_should_replace_model_in_jsonl():
"""Test that should_replace_model_in_jsonl returns the correct value"""
- from litellm.router_utils.batch_utils import should_replace_model_in_jsonl
+ from litellm.litellm_core_utils.batch_utils import should_replace_model_in_jsonl
assert should_replace_model_in_jsonl(purpose="batch") is True
assert should_replace_model_in_jsonl(purpose="batch", passthrough=True) is False
@@ -340,7 +340,7 @@ def test_replace_model_in_jsonl_with_embedded_newlines():
def test_is_batch_retrieve_call_type_matches_only_batch_retrieves():
- from litellm.router_utils.batch_utils import is_batch_retrieve_call_type
+ from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type
from litellm.types.utils import CallTypes
assert is_batch_retrieve_call_type(CallTypes.aretrieve_batch.value) is True
@@ -389,7 +389,8 @@ def setup_and_teardown():
)
def test_batch_charge_identity_requires_a_cost_bearing_retrieval(call_type: str, cost: float, eligible: bool) -> None:
from openai.types import Batch
- from litellm.router_utils.batch_utils import batch_id_to_charge_once
+
+ from litellm.litellm_core_utils.batch_utils import batch_id_to_charge_once
batch: Final = Batch(
id="batch-budget", completion_window="24h", created_at=1, endpoint="/v1/chat/completions",
@@ -401,7 +402,7 @@ def test_batch_charge_identity_requires_a_cost_bearing_retrieval(call_type: str,
def test_batch_charge_markers_keep_counters_and_batches_independent() -> None:
from litellm.caching.in_memory_cache import InMemoryCache
- from litellm.router_utils.batch_utils import batch_charged_once_marker_key
+ from litellm.litellm_core_utils.batch_utils import batch_charged_once_marker_key
cache: Final = InMemoryCache()
cache.set_cache(batch_charged_once_marker_key("provider-spend", "first-batch"), 1)
diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py
index a1b58521d7e..79978ecfe8e 100644
--- a/tests/unit/litellm_core_utils/test_streaming_handler.py
+++ b/tests/unit/litellm_core_utils/test_streaming_handler.py
@@ -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"),
+ ),
+ "Check the weather",
+ 1,
+ ),
+ (
+ (
+ _stream_chunk({"reasoning_content": "Pick a function"}),
+ _stream_chunk({"function_call": {"name": "get_weather", "arguments": "{}"}}),
+ _stream_chunk({}, "function_call"),
+ ),
+ "Pick a function",
+ 0,
+ ),
+ (
+ (
+ _stream_chunk({"reasoning_content": "Step one"}),
+ _stream_chunk({"reasoning_content": ", step two"}),
+ _stream_chunk({}, "length"),
+ ),
+ "Step one, step two",
+ 0,
+ ),
+ (
+ (
+ _stream_chunk({"reasoning_content": "a"}),
+ _stream_chunk({"content": "b"}),
+ _stream_chunk({"reasoning_content": "c"}),
+ _stream_chunk({"content": "d"}),
+ _stream_chunk({}, "stop"),
+ ),
+ "abcd",
+ 0,
+ ),
+ (
+ (_stream_chunk({"reasoning_content": "Still thinking"}),),
+ "Still thinking",
+ 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
diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py
index 1f4a5aaa256..355d407001c 100644
--- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py
+++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py
@@ -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():
diff --git a/tests/unit/llms/anthropic/messages/test_stream_assembly.py b/tests/unit/llms/anthropic/messages/test_stream_assembly.py
index f7d6b8d4f8c..4993d2fa9bf 100644
--- a/tests/unit/llms/anthropic/messages/test_stream_assembly.py
+++ b/tests/unit/llms/anthropic/messages/test_stream_assembly.py
@@ -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,), ()],
diff --git a/tests/unit/llms/azure_ai/agents/test_handler.py b/tests/unit/llms/azure_ai/agents/test_handler.py
index 71466bb172c..14b6094747f 100644
--- a/tests/unit/llms/azure_ai/agents/test_handler.py
+++ b/tests/unit/llms/azure_ai/agents/test_handler.py
@@ -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")]
diff --git a/tests/unit/llms/azure_ai/agents/test_transformation.py b/tests/unit/llms/azure_ai/agents/test_transformation.py
index 10676438858..413908c8313 100644
--- a/tests/unit/llms/azure_ai/agents/test_transformation.py
+++ b/tests/unit/llms/azure_ai/agents/test_transformation.py
@@ -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
diff --git a/tests/unit/llms/base_llm/guardrail_translation/test_utils.py b/tests/unit/llms/base_llm/guardrail_translation/test_utils.py
index 65208a2bd3b..acd50616f90 100644
--- a/tests/unit/llms/base_llm/guardrail_translation/test_utils.py
+++ b/tests/unit/llms/base_llm/guardrail_translation/test_utils.py
@@ -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}
diff --git a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py
index 7e05ef500f5..9d9e39df944 100644
--- a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py
+++ b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py
@@ -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})
diff --git a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py
index 5aaee1f0ee4..aec1672b335 100644
--- a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py
+++ b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py
@@ -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,)
diff --git a/tests/unit/llms/bedrock/search/test_agentcore_search_transformation.py b/tests/unit/llms/bedrock/search/test_agentcore_search_transformation.py
index 664dbf87fe9..01cca1fc270 100644
--- a/tests/unit/llms/bedrock/search/test_agentcore_search_transformation.py
+++ b/tests/unit/llms/bedrock/search/test_agentcore_search_transformation.py
@@ -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)]
diff --git a/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py b/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py
index a6f7d613f28..635363e77ea 100644
--- a/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py
+++ b/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py
@@ -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
diff --git a/tests/unit/llms/cloudflare/test_cloudflare_transformation.py b/tests/unit/llms/cloudflare/test_cloudflare_transformation.py
index 3a45c83556c..ccbfc02f0e6 100644
--- a/tests/unit/llms/cloudflare/test_cloudflare_transformation.py
+++ b/tests/unit/llms/cloudflare/test_cloudflare_transformation.py
@@ -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"}
diff --git a/tests/unit/llms/codestral/__init__.py b/tests/unit/llms/codestral/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/unit/llms/codestral/completion/__init__.py b/tests/unit/llms/codestral/completion/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/unit/llms/codestral/completion/test_codestral_completion_handler.py b/tests/unit/llms/codestral/completion/test_codestral_completion_handler.py
new file mode 100644
index 00000000000..9ce50d8ceca
--- /dev/null
+++ b/tests/unit/llms/codestral/completion/test_codestral_completion_handler.py
@@ -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
diff --git a/tests/unit/llms/custom_httpx/test_aiohttp_handler.py b/tests/unit/llms/custom_httpx/test_aiohttp_handler.py
index 83c38b65beb..bb66a70da84 100644
--- a/tests/unit/llms/custom_httpx/test_aiohttp_handler.py
+++ b/tests/unit/llms/custom_httpx/test_aiohttp_handler.py
@@ -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]
diff --git a/tests/unit/llms/custom_httpx/test_sse_lines.py b/tests/unit/llms/custom_httpx/test_sse_lines.py
new file mode 100644
index 00000000000..9bb03f6da98
--- /dev/null
+++ b/tests/unit/llms/custom_httpx/test_sse_lines.py
@@ -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)
diff --git a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py
index b3a224209e9..eb2e1d8e347 100644
--- a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py
+++ b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py
@@ -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
diff --git a/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py
index 04b791cf289..61e5545b8bf 100644
--- a/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py
+++ b/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py
@@ -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"}],
diff --git a/tests/unit/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py b/tests/unit/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py
index 4119ce99423..53ddb1135b4 100644
--- a/tests/unit/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py
+++ b/tests/unit/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py
@@ -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()
diff --git a/tests/unit/llms/langgraph/chat/test_langgraph_sse_iterator.py b/tests/unit/llms/langgraph/chat/test_langgraph_sse_iterator.py
new file mode 100644
index 00000000000..06485362162
--- /dev/null
+++ b/tests/unit/llms/langgraph/chat/test_langgraph_sse_iterator.py
@@ -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]
diff --git a/tests/unit/llms/oobabooga/chat/test_oobabooga.py b/tests/unit/llms/oobabooga/chat/test_oobabooga.py
index 395a4fb5715..332b6913890 100644
--- a/tests/unit/llms/oobabooga/chat/test_oobabooga.py
+++ b/tests/unit/llms/oobabooga/chat/test_oobabooga.py
@@ -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]"]
diff --git a/tests/unit/llms/openai_like/chat/test_openai_like_chat_handler.py b/tests/unit/llms/openai_like/chat/test_openai_like_chat_handler.py
index 75b60b2dd6c..3df1efe311b 100644
--- a/tests/unit/llms/openai_like/chat/test_openai_like_chat_handler.py
+++ b/tests/unit/llms/openai_like/chat/test_openai_like_chat_handler.py
@@ -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
diff --git a/tests/unit/llms/snowflake/chat/test_snowflake_chat_transformation.py b/tests/unit/llms/snowflake/chat/test_snowflake_chat_transformation.py
index f9d12b7eeea..36c319d9524 100644
--- a/tests/unit/llms/snowflake/chat/test_snowflake_chat_transformation.py
+++ b/tests/unit/llms/snowflake/chat/test_snowflake_chat_transformation.py
@@ -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):
diff --git a/tests/unit/llms/test_predibase_transformation.py b/tests/unit/llms/test_predibase_transformation.py
index a02840ed4ef..7afbd05b221 100644
--- a/tests/unit/llms/test_predibase_transformation.py
+++ b/tests/unit/llms/test_predibase_transformation.py
@@ -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]
diff --git a/tests/unit/llms/triton/test_triton.py b/tests/unit/llms/triton/test_triton.py
index 7c0de603a3e..238c30816b2 100644
--- a/tests/unit/llms/triton/test_triton.py
+++ b/tests/unit/llms/triton/test_triton.py
@@ -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():
diff --git a/tests/unit/llms/vertex_ai/agent_engine/test_transformation.py b/tests/unit/llms/vertex_ai/agent_engine/test_transformation.py
index 06aaca44e8a..20333fd0465 100644
--- a/tests/unit/llms/vertex_ai/agent_engine/test_transformation.py
+++ b/tests/unit/llms/vertex_ai/agent_engine/test_transformation.py
@@ -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
diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
index 1546d1e9fe0..67c78b62151 100644
--- a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
+++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
@@ -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
@@ -3514,6 +3515,39 @@ def test_vertex_ai_annotation_conversion():
assert "Current time" in time_annotation["url_citation"]["title"]
+@pytest.mark.parametrize(
+ "indices, expected_indices",
+ [([2, 0], [2, 0]), ([1, 99, -1, 2, 0], [2, 0]), ([0], [0]), ([], []), ([1, 99], [])],
+ ids=["multiple-sources", "non-web-first", "single-source", "empty", "no-web"],
+)
+def test_grounding_annotations_preserve_all_web_sources(indices: list[int], expected_indices: list[int]) -> None:
+ chunks: Final = [
+ {"web": {"uri": "https://example.com/first", "title": "First"}},
+ {"retrievedContext": {"uri": "gs://documents/source"}},
+ {"web": {"uri": "https://example.com/second", "title": "Second"}},
+ ]
+ metadata: Final = {
+ "groundingChunks": chunks,
+ "groundingSupports": [
+ {"segment": {"startIndex": 4, "endIndex": 11}, "groundingChunkIndices": indices}
+ ],
+ }
+ annotations: Final = VertexGeminiConfig._convert_grounding_metadata_to_annotations([metadata], "The passage")
+ expected: Final = [
+ {
+ "type": "url_citation",
+ "url_citation": {
+ "start_index": 4,
+ "end_index": 11,
+ "url": chunks[index]["web"]["uri"],
+ "title": chunks[index]["web"]["title"],
+ },
+ }
+ for index in expected_indices
+ ]
+ assert json.dumps(annotations) == json.dumps(expected)
+
+
def test_vertex_ai_annotation_empty_grounding_metadata():
"""
Test handling of empty or missing grounding metadata.
@@ -10092,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()
@@ -10657,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 = {
diff --git a/tests/unit/llms/vertex_ai/test_http_status_201.py b/tests/unit/llms/vertex_ai/test_http_status_201.py
index 34502d5cebe..0ed469acfa7 100644
--- a/tests/unit/llms/vertex_ai/test_http_status_201.py
+++ b/tests/unit/llms/vertex_ai/test_http_status_201.py
@@ -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
diff --git a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py
index a7d7ac5627c..7abe6ad7799 100644
--- a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py
+++ b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py
@@ -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
diff --git a/tests/unit/proxy/anthropic_endpoints/test_streaming_model_restamp.py b/tests/unit/proxy/anthropic_endpoints/test_streaming_model_restamp.py
index 85bb8304b95..5911fcdc073 100644
--- a/tests/unit/proxy/anthropic_endpoints/test_streaming_model_restamp.py
+++ b/tests/unit/proxy/anthropic_endpoints/test_streaming_model_restamp.py
@@ -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"}}
diff --git a/tests/unit/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/unit/proxy/management_endpoints/test_auto_router_endpoints.py
index 748b291d5ff..c13e3d5812b 100644
--- a/tests/unit/proxy/management_endpoints/test_auto_router_endpoints.py
+++ b/tests/unit/proxy/management_endpoints/test_auto_router_endpoints.py
@@ -170,7 +170,7 @@ async def test_preview_and_serving_share_default_model_resolution(
explicit_default: str | None,
expected: str,
):
- from litellm.router_utils.auto_router_model_naming import validate_complexity_router_config_write
+ from litellm.router_strategy.auto_router.auto_router_model_naming import validate_complexity_router_config_write
from litellm.types.management_endpoints.auto_router_endpoints import ComplexityRouterConfigValidationRequest
config: Final = {
@@ -1240,7 +1240,7 @@ class TestAutoRouterBenchmarks:
def test_the_listed_kinds_match_the_router_types_traffic_can_record(self):
from typing import get_args, get_type_hints
- from litellm.router_utils.auto_router_model_naming import StrategyRouterKind
+ from litellm.router_strategy.auto_router.auto_router_model_naming import StrategyRouterKind
from litellm.types.utils import StandardLoggingRoutingDecision
readonly_router_type: Final = get_type_hints(StandardLoggingRoutingDecision, include_extras=True)["router_type"]
diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py
index fc6749aed0f..437d948f8f5 100644
--- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py
+++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py
@@ -1,6 +1,6 @@
-import inspect
import asyncio
import contextlib
+import inspect
import json
import socket
from collections.abc import Iterator, Mapping
@@ -1461,10 +1461,10 @@ class TestDeleteModelClearsRouterRegistry:
add_deployment never restores, so an unguarded pop would make it permanently
unroutable (the same cross-tenant DoS clear_cache was hardened against).
"""
+ from litellm.proxy.management_endpoints.model_management_endpoints import ModelInfoDelete
from litellm.proxy.management_endpoints.model_management_endpoints import (
delete_model as delete_model_endpoint,
)
- from litellm.proxy.management_endpoints.model_management_endpoints import ModelInfoDelete
model_id = "regular-del-1"
admin_user = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
@@ -2739,11 +2739,11 @@ class TestAddAndDeleteModelLifecycle:
- Delete same model again → raises (model not found)
"""
from litellm.proxy.management_endpoints.model_management_endpoints import (
+ ModelInfoDelete,
add_new_model,
- delete_model as delete_model_endpoint,
)
from litellm.proxy.management_endpoints.model_management_endpoints import (
- ModelInfoDelete,
+ delete_model as delete_model_endpoint,
)
model_id = "lifecycle-test-model-123"
@@ -2833,6 +2833,8 @@ class TestDeleteTeamBYOKModelGhost:
async def test_delete_strips_public_name_and_refreshes_cache(self):
from litellm.proxy.management_endpoints.model_management_endpoints import (
ModelInfoDelete,
+ )
+ from litellm.proxy.management_endpoints.model_management_endpoints import (
delete_model as delete_model_endpoint,
)
@@ -2916,6 +2918,8 @@ class TestDeleteTeamBYOKModelGhost:
alias cleanup (delete_team_model_alias), preserving legacy behavior."""
from litellm.proxy.management_endpoints.model_management_endpoints import (
ModelInfoDelete,
+ )
+ from litellm.proxy.management_endpoints.model_management_endpoints import (
delete_model as delete_model_endpoint,
)
@@ -2984,6 +2988,8 @@ class TestDeleteTeamBYOKModelGhost:
team.models when one replica is deleted but a sibling still backs it."""
from litellm.proxy.management_endpoints.model_management_endpoints import (
ModelInfoDelete,
+ )
+ from litellm.proxy.management_endpoints.model_management_endpoints import (
delete_model as delete_model_endpoint,
)
@@ -3065,6 +3071,8 @@ class TestDeleteTeamBYOKModelGhost:
deployment still serves it."""
from litellm.proxy.management_endpoints.model_management_endpoints import (
ModelInfoDelete,
+ )
+ from litellm.proxy.management_endpoints.model_management_endpoints import (
delete_model as delete_model_endpoint,
)
@@ -3145,6 +3153,8 @@ class TestDeleteTeamBYOKModelGhost:
the alias would break routing that works."""
from litellm.proxy.management_endpoints.model_management_endpoints import (
ModelInfoDelete,
+ )
+ from litellm.proxy.management_endpoints.model_management_endpoints import (
delete_model as delete_model_endpoint,
)
@@ -3252,6 +3262,8 @@ class TestDeleteModelTeamAuth:
async def test_proxy_admin_can_delete_model_when_team_deleted(self):
from litellm.proxy.management_endpoints.model_management_endpoints import (
ModelInfoDelete,
+ )
+ from litellm.proxy.management_endpoints.model_management_endpoints import (
delete_model as delete_model_endpoint,
)
@@ -3287,6 +3299,8 @@ class TestDeleteModelTeamAuth:
"""A missing team must never let a non-admin delete the orphan (no fail-open)."""
from litellm.proxy.management_endpoints.model_management_endpoints import (
ModelInfoDelete,
+ )
+ from litellm.proxy.management_endpoints.model_management_endpoints import (
delete_model as delete_model_endpoint,
)
from litellm.proxy.proxy_server import ProxyException
@@ -3322,6 +3336,8 @@ class TestDeleteModelTeamAuth:
"""The auth check must not add a redundant team query on the live-team path."""
from litellm.proxy.management_endpoints.model_management_endpoints import (
ModelInfoDelete,
+ )
+ from litellm.proxy.management_endpoints.model_management_endpoints import (
delete_model as delete_model_endpoint,
)
from litellm.proxy.proxy_server import ProxyException
@@ -4392,7 +4408,6 @@ class TestModelInfoCostMapEchoFilter:
def test_echoed_cost_map_metadata_is_not_persisted(self):
import litellm
-
from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
@@ -4418,7 +4433,6 @@ class TestModelInfoCostMapEchoFilter:
def test_an_edited_value_survives_the_echo_filter(self):
import litellm
-
from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
@@ -4483,7 +4497,6 @@ class TestModelInfoCostMapEchoFilter:
def test_metadata_without_a_cost_map_key_is_persisted(self):
import litellm
-
from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
from litellm.types.utils import echoed_cost_map_fields
@@ -4513,7 +4526,6 @@ class TestModelInfoCostMapEchoFilter:
def test_a_stored_mode_survives_an_echoed_save(self):
import litellm
-
from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
@@ -4536,7 +4548,6 @@ class TestModelInfoCostMapEchoFilter:
def test_resetting_an_override_to_the_cost_map_value_removes_it(self):
import litellm
-
from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
@@ -4560,7 +4571,6 @@ class TestModelInfoCostMapEchoFilter:
def test_reset_is_recognised_after_the_router_registered_the_override(self, monkeypatch: pytest.MonkeyPatch):
import litellm
-
from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
@@ -4586,7 +4596,6 @@ class TestModelInfoCostMapEchoFilter:
from types import MappingProxyType
import litellm
-
from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
@@ -4611,7 +4620,6 @@ class TestModelInfoCostMapEchoFilter:
def test_echo_is_compared_against_the_deployments_lookup_not_the_key(self):
import litellm
-
from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
@@ -4653,7 +4661,6 @@ class TestModelInfoCostMapEchoFilter:
def test_base_model_wins_over_litellm_params_model_for_the_lookup(self):
import litellm
-
from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
@@ -4681,7 +4688,6 @@ class TestModelInfoCostMapEchoFilter:
def test_encrypted_stored_model_is_decrypted_for_the_lookup(self, monkeypatch):
import litellm
-
from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
@@ -6431,7 +6437,7 @@ class TestStrategyRouterWriteValidation:
AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY,
_auto_router_capability_slot,
)
- from litellm.router_utils.auto_router_model_naming import gated_capabilities_of
+ from litellm.router_strategy.auto_router.auto_router_model_naming import gated_capabilities_of
capabilities: Final = gated_capabilities_of(effective_params)
@@ -6568,7 +6574,7 @@ class TestStrategyRouterWriteValidation:
from fastapi import HTTPException
from litellm.proxy.management_endpoints.model_management_endpoints import _auto_router_capability_slot
- from litellm.router_utils.auto_router_tuning_baseline import snapshot_tuning_baselines
+ from litellm.router_strategy.auto_router.auto_router_tuning_baseline import snapshot_tuning_baselines
configs = {
"_TUNED_A": self._TUNED_A,
@@ -6624,7 +6630,7 @@ class TestStrategyRouterWriteValidation:
from litellm.proxy import proxy_server
from litellm.proxy.management_endpoints.model_management_endpoints import _auto_router_capability_slot
- from litellm.router_utils.auto_router_tuning_baseline import snapshot_tuning_baselines
+ from litellm.router_strategy.auto_router.auto_router_tuning_baseline import snapshot_tuning_baselines
fake = self._FakeDb([])
fake.tx_obj.litellm_proxymodeltable.find_many = AsyncMock(
@@ -6669,7 +6675,7 @@ class TestStrategyRouterWriteValidation:
row id before the slot is entered, so a second tuned router is refused before its DB write."""
from litellm.proxy._types import ProxyException
from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model
- from litellm.router_utils.auto_router_tuning_baseline import snapshot_tuning_baselines
+ from litellm.router_strategy.auto_router.auto_router_tuning_baseline import snapshot_tuning_baselines
admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
baselines = snapshot_tuning_baselines([self._db_router_row("a", self._TUNED_A)])
@@ -8980,7 +8986,10 @@ class TestTeamMemberAutoRouterWrites:
from fastapi import HTTPException
from litellm.proxy.management_endpoints.model_management_endpoints import _auto_router_capability_slot
- from litellm.proxy.management_helpers.auto_router_permissions import MemberAutoRouterWrite, validate_member_auto_router_config
+ from litellm.proxy.management_helpers.auto_router_permissions import (
+ MemberAutoRouterWrite,
+ validate_member_auto_router_config,
+ )
row: Final = self._row()
database: Final = self._database(self._team(), row)
@@ -10677,6 +10686,8 @@ class TestFederationGateScopesToWhatTheWriteTouches:
and the team that owns the model must be able to take it off their page."""
from litellm.proxy.management_endpoints.model_management_endpoints import (
ModelInfoDelete,
+ )
+ from litellm.proxy.management_endpoints.model_management_endpoints import (
delete_model as delete_model_endpoint,
)
diff --git a/tests/unit/proxy/management_helpers/test_auto_router_availability.py b/tests/unit/proxy/management_helpers/test_auto_router_availability.py
index a88dcd7f66d..0eaae60b4ea 100644
--- a/tests/unit/proxy/management_helpers/test_auto_router_availability.py
+++ b/tests/unit/proxy/management_helpers/test_auto_router_availability.py
@@ -1,16 +1,15 @@
from collections.abc import Mapping
-from typing import Final
from types import SimpleNamespace
-
-from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper
+from typing import Final
import pytest
+from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper
from litellm.proxy.management_helpers.auto_router_availability import (
auto_router_availability,
build_auto_router_catalog,
)
-from litellm.router_utils.auto_router_tuning_baseline import snapshot_tuning_baselines
+from litellm.router_strategy.auto_router.auto_router_tuning_baseline import snapshot_tuning_baselines
def deployment(
diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py
index e09a179eadf..47b6af9aeb3 100644
--- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py
+++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py
@@ -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}
diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_tinyfish_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_tinyfish_passthrough_logging_handler.py
index 1b83c5140ca..b0f84f2ecbe 100644
--- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_tinyfish_passthrough_logging_handler.py
+++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_tinyfish_passthrough_logging_handler.py
@@ -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
diff --git a/tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py b/tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py
index 49bc2e9c29e..2496aa47d22 100644
--- a/tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py
+++ b/tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py
@@ -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()
diff --git a/tests/unit/proxy/proxy_server/test_lifecycle.py b/tests/unit/proxy/proxy_server/test_lifecycle.py
index 11a44a0ec16..8b1d7c1222a 100644
--- a/tests/unit/proxy/proxy_server/test_lifecycle.py
+++ b/tests/unit/proxy/proxy_server/test_lifecycle.py
@@ -33,9 +33,9 @@ from pydantic import BaseModel
from typing_extensions import TypedDict
import litellm.proxy.proxy_server as ps
+from litellm.proxy.management_helpers.audit_logs import track_audit_task
from litellm.proxy.proxy_server import (
ProxyStartupEvent,
- initialize_shared_aiohttp_session,
_resolve_pydantic_type,
_resolve_typed_dict_type,
cleanup_router_config_variables,
@@ -43,6 +43,7 @@ from litellm.proxy.proxy_server import (
get_litellm_model_info,
initialize,
initialize_from_worker_config,
+ initialize_shared_aiohttp_session,
load_from_azure_key_vault,
proxy_shutdown_event,
proxy_startup_event,
@@ -50,8 +51,6 @@ from litellm.proxy.proxy_server import (
save_worker_config,
)
-from litellm.proxy.management_helpers.audit_logs import track_audit_task
-
from .conftest import normalize
# ---------------------------------------------------------------------------
@@ -1165,7 +1164,7 @@ def test_proxy_startup_event_warns_for_global_budget_without_database():
@pytest.mark.asyncio
async def test_tuning_baseline_v3_is_created_alongside_the_legacy_row():
- from litellm.router_utils.auto_router_tuning_baseline import DEFAULT_TUNING_FINGERPRINT
+ from litellm.router_strategy.auto_router.auto_router_tuning_baseline import DEFAULT_TUNING_FINGERPRINT
prisma_client = MagicMock()
prisma_client.db.litellm_config.find_unique = AsyncMock(return_value=None)
@@ -1186,7 +1185,10 @@ async def test_tuning_baseline_v3_is_created_alongside_the_legacy_row():
@pytest.mark.asyncio
async def test_scorer_baseline_upgrade_preserves_existing_routers_and_is_not_refreshed_on_restart():
- from litellm.router_utils.auto_router_tuning_baseline import mutable_tuned_identities, snapshot_tuning_baselines
+ from litellm.router_strategy.auto_router.auto_router_tuning_baseline import (
+ mutable_tuned_identities,
+ snapshot_tuning_baselines,
+ )
deployments = [
{
diff --git a/tests/unit/proxy/response_polling/test_response_polling_handler.py b/tests/unit/proxy/response_polling/test_response_polling_handler.py
index 63e38ce533a..42acad28488 100644
--- a/tests/unit/proxy/response_polling/test_response_polling_handler.py
+++ b/tests/unit/proxy/response_polling/test_response_polling_handler.py
@@ -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,
diff --git a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py
index ed90acbd2fa..0b971bc1696 100644
--- a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py
+++ b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py
@@ -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(""),
+ *(_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) == ""
+ 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):
diff --git a/tests/unit/router_utils/test_auto_router_model_naming.py b/tests/unit/router_strategy/test_auto_router_model_naming.py
similarity index 99%
rename from tests/unit/router_utils/test_auto_router_model_naming.py
rename to tests/unit/router_strategy/test_auto_router_model_naming.py
index 1cfa3aa5c91..44fcfdaf5be 100644
--- a/tests/unit/router_utils/test_auto_router_model_naming.py
+++ b/tests/unit/router_strategy/test_auto_router_model_naming.py
@@ -3,9 +3,7 @@ from typing import Final
import pytest
-from litellm.router_strategy.complexity_router.fuse_presets import get_fuse_presets
-from litellm.router_strategy.complexity_router.jev_classifier import DEFAULT_JEV_INSTRUCTIONS
-from litellm.router_utils.auto_router_model_naming import (
+from litellm.router_strategy.auto_router.auto_router_model_naming import (
GATED_AUTO_ROUTER_CAPABILITIES,
capability_limit_violation,
carries_complexity_router_settings,
@@ -18,6 +16,8 @@ from litellm.router_utils.auto_router_model_naming import (
validate_complexity_router_config_write,
validate_strategy_router_model_write,
)
+from litellm.router_strategy.complexity_router.fuse_presets import get_fuse_presets
+from litellm.router_strategy.complexity_router.jev_classifier import DEFAULT_JEV_INSTRUCTIONS
COMPLEXITY_FIELDS = frozenset({"complexity_router_config"})
SEMANTIC_FIELDS = frozenset({"auto_router_config", "auto_router_default_model", "auto_router_embedding_model"})
@@ -596,7 +596,7 @@ def test_custom_classifier_prompt_capability(config: Mapping[str, object], expec
],
)
def test_is_complexity_router_model(model: str | None, expected: bool) -> None:
- from litellm.router_utils.auto_router_model_naming import is_complexity_router_model
+ from litellm.router_strategy.auto_router.auto_router_model_naming import is_complexity_router_model
assert is_complexity_router_model(model) is expected
diff --git a/tests/unit/router_utils/test_auto_router_tuning_baseline.py b/tests/unit/router_strategy/test_auto_router_tuning_baseline.py
similarity index 99%
rename from tests/unit/router_utils/test_auto_router_tuning_baseline.py
rename to tests/unit/router_strategy/test_auto_router_tuning_baseline.py
index d0d98ce7042..4ee219a7784 100644
--- a/tests/unit/router_utils/test_auto_router_tuning_baseline.py
+++ b/tests/unit/router_strategy/test_auto_router_tuning_baseline.py
@@ -7,8 +7,7 @@ from typing import Final
import pytest
-from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig
-from litellm.router_utils.auto_router_tuning_baseline import (
+from litellm.router_strategy.auto_router.auto_router_tuning_baseline import (
DEFAULT_TUNING_FINGERPRINT,
HEURISTIC_V1_TUNING_FIELDS,
heuristic_v1_router_fingerprint,
@@ -19,6 +18,7 @@ from litellm.router_utils.auto_router_tuning_baseline import (
tuning_limit_violation,
tuning_quota_violation,
)
+from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig
_TIERS = {"SIMPLE": "cheap", "MEDIUM": "mid", "COMPLEX": "strong", "REASONING": "top"}
_ALT_TIERS = {**_TIERS, "COMPLEX": "other-strong"}
diff --git a/tests/unit/router_utils/test_baseline_request.py b/tests/unit/router_strategy/test_baseline_request.py
similarity index 95%
rename from tests/unit/router_utils/test_baseline_request.py
rename to tests/unit/router_strategy/test_baseline_request.py
index 37d21597752..eb785a41d63 100644
--- a/tests/unit/router_utils/test_baseline_request.py
+++ b/tests/unit/router_strategy/test_baseline_request.py
@@ -1,6 +1,6 @@
from typing import Final
-from litellm.router_utils.baseline_request import baseline_request, capture_baseline_parameters
+from litellm.router_strategy.auto_router.baseline_request import baseline_request, capture_baseline_parameters
def test_baseline_snapshot_owns_nested_caller_settings_and_overrides_routed_settings() -> None:
diff --git a/tests/unit/router_strategy/test_complexity_router.py b/tests/unit/router_strategy/test_complexity_router.py
index ca39c7f5c8f..018f93779a6 100644
--- a/tests/unit/router_strategy/test_complexity_router.py
+++ b/tests/unit/router_strategy/test_complexity_router.py
@@ -24,14 +24,6 @@ from pydantic import TypeAdapter, ValidationError
import litellm
from litellm import Router
-from litellm.integrations.custom_logger import CustomLogger
-from litellm.types.utils import StandardLoggingRoutingDecision
-from litellm.router_utils.auto_router_model_naming import (
- CUSTOMIZATION_CAPABILITY,
- GATED_AUTO_ROUTER_CAPABILITIES,
- HEURISTIC_V2_CAPABILITY,
- count_capability_routers,
-)
from litellm._logging import verbose_router_logger
from litellm.caching.dual_cache import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
@@ -42,7 +34,19 @@ from litellm.constants import (
RETURN_RAW_MODEL_NAME_METADATA_KEY,
SESSION_ID_GENERATED_METADATA_KEY,
)
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.router import as_output_cap, stamped_client_ceiling
+from litellm.router_strategy.auto_router.auto_router_model_naming import (
+ CUSTOMIZATION_CAPABILITY,
+ GATED_AUTO_ROUTER_CAPABILITIES,
+ HEURISTIC_V2_CAPABILITY,
+ count_capability_routers,
+)
+from litellm.router_strategy.complexity_router.capability_classifier import (
+ CAPABILITY_CLASSIFIER_SYSTEM_PROMPT,
+ CapabilityClassifierVerdict,
+)
from litellm.router_strategy.complexity_router.complexity_router import (
_CLASSIFICATION_CURRENT_MESSAGE_ONLY,
_CLASSIFICATION_WITH_CONVERSATION,
@@ -58,18 +62,14 @@ from litellm.router_strategy.complexity_router.complexity_router import (
classification_system_prompt,
custom_tier_classification_prompt,
)
-from litellm.router_strategy.complexity_router.capability_classifier import (
- CAPABILITY_CLASSIFIER_SYSTEM_PROMPT,
- CapabilityClassifierVerdict,
-)
from litellm.router_strategy.complexity_router.config import (
- CapabilityCalibrationConfig,
- CapabilityClassifierConfig,
DEFAULT_CLASSIFICATION_RUBRIC,
DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE,
DEFAULT_COMPLEXITY_CONFIG,
DEFAULT_TECHNICAL_KEYWORDS,
TIER_SEVERITY_ORDER,
+ CapabilityCalibrationConfig,
+ CapabilityClassifierConfig,
ClassificationRubric,
ClassifierLLMConfig,
ComplexityRouterConfig,
@@ -83,10 +83,12 @@ from litellm.router_strategy.complexity_router.jev_classifier import (
JevUsage,
)
from litellm.router_strategy.complexity_router.llm_v2 import LLM_V2_PROMPT_VERSION
+from litellm.router_strategy.complexity_router.request_models import RequestComplexityRouterConfig
from litellm.router_strategy.complexity_router.tier_predictor import (
TierGlobalStatistic,
TrainedTierArtifact,
)
+from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.types.router import (
Deployment,
LiteLLM_Params,
@@ -94,10 +96,7 @@ from litellm.types.router import (
RouterErrors,
TaggedPreRoutingStrategy,
)
-from litellm.types.llms.openai import ResponsesAPIResponse
-from litellm.router_strategy.complexity_router.request_models import RequestComplexityRouterConfig
-from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
-
+from litellm.types.utils import StandardLoggingRoutingDecision
requires_semantic_router = pytest.mark.skipif(
sys.version_info >= (3, 14), reason="The semantic-router extra excludes Python 3.14"
@@ -16299,7 +16298,6 @@ class TestTierHealthFailover:
tests can prove the probe passes a kwargs copy and forwards the prompt arguments.
"""
import litellm as litellm_module
-
from litellm.types.router import RouterRateLimitError
probed_kwargs = []
diff --git a/tests/unit/router_strategy/test_llm_v2.py b/tests/unit/router_strategy/test_llm_v2.py
index 3fd2e8808e7..8c1e9e48661 100644
--- a/tests/unit/router_strategy/test_llm_v2.py
+++ b/tests/unit/router_strategy/test_llm_v2.py
@@ -4,11 +4,12 @@ from typing import Final
from unittest.mock import AsyncMock, MagicMock
import pytest
-import litellm
from pydantic import ValidationError
+import litellm
from litellm import ModelResponse, Router
from litellm.caching.dual_cache import DualCache
+from litellm.router_strategy.auto_router.auto_router_model_naming import strategy_router_dependencies
from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter
from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig, ComplexityTier
from litellm.router_strategy.complexity_router.fuse_presets import get_fuse_presets
@@ -21,7 +22,6 @@ from litellm.router_strategy.complexity_router.llm_v2 import (
LLMV2Verdict,
llm_v2_response_format,
)
-from litellm.router_utils.auto_router_model_naming import strategy_router_dependencies
from litellm.types.llms.openai import ResponsesAPIResponse
diff --git a/tests/unit/router_utils/test_cooldown_handlers.py b/tests/unit/router_utils/test_cooldown_handlers.py
index 24459a5a250..2e1e3acaabb 100644
--- a/tests/unit/router_utils/test_cooldown_handlers.py
+++ b/tests/unit/router_utils/test_cooldown_handlers.py
@@ -3150,8 +3150,6 @@ BENCHED: Final = (litellm.Timeout, litellm.Timeout, RouterRateLimitError, Router
{},
BENCHED,
),
- ({"allowed_fails": 1, "fallbacks": [{"test-model": "backup-model"}]}, {}, BENCHED),
- ({"allowed_fails": 1, "default_fallbacks": "backup-model"}, {}, BENCHED),
],
ids=[
"allowed_fails",
@@ -3160,8 +3158,6 @@ BENCHED: Final = (litellm.Timeout, litellm.Timeout, RouterRateLimitError, Router
"policy_names_timeout",
"allowed_fails_with_a_fallback_group",
"allowed_fails_with_a_dict_fallback_target",
- "allowed_fails_with_a_string_fallback_target",
- "allowed_fails_with_a_string_default_fallbacks",
],
)
@pytest.mark.asyncio
@@ -3207,7 +3203,8 @@ async def test_router_allowed_fails_benches_a_lone_deployment_only_when_the_poli
def test_has_fallback_group_reads_string_targets_and_skips_unparseable_entries(
fallbacks: Sequence[Mapping[str, object]], expected: bool
) -> None:
- router: Final = Router(model_list=[TIMED_OUT_DEPLOYMENT], fallbacks=list(fallbacks))
+ router: Final = Router(model_list=[TIMED_OUT_DEPLOYMENT])
+ router.fallbacks = list(fallbacks)
assert _has_fallback_group(router, "test-model") is expected
diff --git a/tests/unit/sandbox/test_e2b_sandbox.py b/tests/unit/sandbox/test_e2b_sandbox.py
index e01b9120416..e58392ca266 100644
--- a/tests/unit/sandbox/test_e2b_sandbox.py
+++ b/tests/unit/sandbox/test_e2b_sandbox.py
@@ -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)
diff --git a/tests/unit/sandbox/test_opensandbox_sandbox.py b/tests/unit/sandbox/test_opensandbox_sandbox.py
index 2928dea100e..36ae848786c 100644
--- a/tests/unit/sandbox/test_opensandbox_sandbox.py
+++ b/tests/unit/sandbox/test_opensandbox_sandbox.py
@@ -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):
diff --git a/tests/unit/test_groq_streaming_encoding.py b/tests/unit/test_groq_streaming_encoding.py
index ffe87fb9890..37f05272f67 100644
--- a/tests/unit/test_groq_streaming_encoding.py
+++ b/tests/unit/test_groq_streaming_encoding.py
@@ -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
diff --git a/tests/unit/test_private_usage_aliases.py b/tests/unit/test_private_usage_aliases.py
index 1d9c2e87100..d0c5c8247d3 100644
--- a/tests/unit/test_private_usage_aliases.py
+++ b/tests/unit/test_private_usage_aliases.py
@@ -1969,7 +1969,7 @@ ALIAS_CASES: Final = (
("litellm.router_utils.add_retry_fallback_headers", "HiddenParamsAsyncIteratorWrapper", "_inner", "inner", True),
("litellm.router_utils.add_retry_fallback_headers", "", "_HiddenParamsHost", "HiddenParamsHost", False),
(
- "litellm.router_utils.batch_utils",
+ "litellm.litellm_core_utils.batch_utils",
"",
"_get_router_metadata_variable_name",
"get_router_metadata_variable_name",
diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py
index 3c7a222d9a1..300c9cff6ff 100644
--- a/tests/unit/test_router/test_router.py
+++ b/tests/unit/test_router/test_router.py
@@ -31228,3 +31228,51 @@ async def test_anthropic_caller_rate_limit_does_not_block_other_credentials(
finally:
await GLOBAL_LOGGING_WORKER.flush()
await client.close()
+
+
+_FALLBACK_VALIDATION_MODEL_LIST: Final = (
+ {"model_name": "gpt-4", "litellm_params": {"model": "openai/gpt-4", "api_key": "fake-key"}},
+ {"model_name": "gpt-3.5", "litellm_params": {"model": "openai/gpt-3.5-turbo", "api_key": "fake-key"}},
+)
+
+
+@pytest.mark.parametrize(
+ ("router_kwargs", "expected_message"),
+ [
+ ({"fallbacks": [{"gpt-4": "gpt-3.5"}]}, "Fallbacks for 'gpt-4' must be a list of model names, got 'gpt-3.5'"),
+ ({"default_fallbacks": "gpt-3.5"}, "default_fallbacks must be a list of model names, got 'gpt-3.5'"),
+ (
+ {"content_policy_fallbacks": [{"gpt-4": "gpt-3.5"}]},
+ "Fallbacks for 'gpt-4' must be a list of model names, got 'gpt-3.5'",
+ ),
+ ],
+ ids=["string_fallback_target", "string_default_fallbacks", "string_content_policy_fallback_target"],
+)
+def test_router_rejects_a_string_fallback_target_at_init(
+ router_kwargs: Mapping[str, object], expected_message: str
+) -> None:
+ with pytest.raises(TypeError, match="must be a list of model names") as excinfo:
+ Router(model_list=list(_FALLBACK_VALIDATION_MODEL_LIST), **router_kwargs)
+
+ assert expected_message in str(excinfo.value)
+
+
+@pytest.mark.parametrize(
+ ("router_kwargs", "expected_fallbacks"),
+ [
+ ({"fallbacks": [{"gpt-4": ["gpt-3.5"]}]}, [{"gpt-4": ["gpt-3.5"]}]),
+ ({"default_fallbacks": ["gpt-3.5"]}, [{"*": ["gpt-3.5"]}]),
+ (
+ {"fallbacks": [{"gpt-4": [{"model": "gpt-3.5", "api_key": "other-key"}]}]},
+ [{"gpt-4": [{"model": "gpt-3.5", "api_key": "other-key"}]}],
+ ),
+ ({"fallbacks": [{"model": "gpt-3.5"}]}, [{"model": "gpt-3.5"}]),
+ ],
+ ids=["list_fallback_target", "list_default_fallbacks", "dict_override_fallback_target", "direct_target_entry"],
+)
+def test_router_keeps_supported_fallback_shapes(
+ router_kwargs: Mapping[str, object], expected_fallbacks: list[dict[str, object]]
+) -> None:
+ router: Final = Router(model_list=list(_FALLBACK_VALIDATION_MODEL_LIST), **router_kwargs)
+
+ assert router.fallbacks == expected_fallbacks
diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py
index af8e7f53462..513136ae150 100644
--- a/tests/unit/test_router/test_router_helper_utils.py
+++ b/tests/unit/test_router/test_router_helper_utils.py
@@ -4,20 +4,21 @@ import gc
import json
import os
import weakref
-from datetime import datetime, timezone
from collections.abc import Iterator, Mapping
+from datetime import datetime, timezone
from types import MethodType
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
-import litellm
import pytest
+
+import litellm
from litellm import Router
from litellm.caching.dual_cache import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY
-from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
from litellm.integrations.custom_logger import CustomLogger
+from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook
from litellm.types.litellm_params import RoutingStrategyName
@@ -801,10 +802,11 @@ async def test_deployment_callback_on_failure(model_list):
def test_deployment_callback_respects_cooldown_time(model_list):
"""Ensure per-model cooldown_time is honored even when exception headers are present."""
- import httpx
import time
from unittest.mock import patch
+ import httpx
+
router = Router(model_list=model_list)
class FakeException(Exception):
@@ -1811,8 +1813,8 @@ def test_get_allowed_fails_from_policy(model_list, exception_type, exception_nam
def test_initialize_alerting(model_list):
"""Test if the 'initialize_alerting' function is working correctly"""
- from litellm.types.router import AlertingConfig
from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
+ from litellm.types.router import AlertingConfig
router = Router(
model_list=model_list, alerting_config=AlertingConfig(webhook_url="test")
@@ -2245,6 +2247,7 @@ def test_deployments_by_pattern(model_list):
)
def test_pattern_match_deployment_set_model_name(user_request_model, model_name, litellm_model, expected_model):
from re import Match
+
from litellm.router_utils.pattern_match_deployments import PatternMatchRouter
pattern_router = PatternMatchRouter()
@@ -2702,7 +2705,7 @@ def test_handle_clientside_credential_metadata_loading(model_list, function_name
)
def test_handle_clientside_credential_metadata_variable_name(model_list, function_name, metadata_key):
"""Test that _handle_clientside_credential uses the correct metadata variable name based on function name"""
- from litellm.router_utils.batch_utils import get_router_metadata_variable_name
+ from litellm.litellm_core_utils.batch_utils import get_router_metadata_variable_name
router = Router(model_list=model_list)
diff --git a/tests/unit/test_router/test_router_surface_check.py b/tests/unit/test_router/test_router_surface_check.py
new file mode 100644
index 00000000000..a97fd029380
--- /dev/null
+++ b/tests/unit/test_router/test_router_surface_check.py
@@ -0,0 +1,58 @@
+"""
+The router surface fence (tests/code_coverage_tests/router_surface_check.py) must see every way
+outside code reaches a Router member, or new dependencies slip past the shrink-only allowlist.
+"""
+
+import ast
+import importlib.util
+import pathlib
+
+import pytest
+
+_CHECK = pathlib.Path(__file__).parents[2] / "code_coverage_tests" / "router_surface_check.py"
+
+
+@pytest.fixture(scope="module")
+def surface_check():
+ spec = importlib.util.spec_from_file_location("router_surface_check", _CHECK)
+ assert spec is not None and spec.loader is not None
+ module = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(module)
+ return module
+
+
+def _uses(surface_check, src: str) -> set:
+ return surface_check.member_uses(ast.parse(src), surface_check.router_members())
+
+
+def test_instance_and_dynamic_access_are_uses(surface_check):
+ assert _uses(surface_check, "llm_router.get_model_list()\ngetattr(router, 'model_list')") == {
+ "get_model_list",
+ "model_list",
+ }
+
+
+def test_class_level_calls_are_uses(surface_check):
+ assert _uses(surface_check, "litellm.Router.get_valid_args()") == {"get_valid_args"}
+ assert _uses(surface_check, "from litellm import Router\nRouter.get_valid_args()") == {"get_valid_args"}
+
+
+def test_container_file_endpoints_registered_by_name_are_members(surface_check):
+ members = surface_check.router_members()
+ assert {"aupload_container_file", "upload_container_file", "alist_container_files"} <= members
+ assert _uses(surface_check, "await llm_router.aupload_container_file(file=f)") == {"aupload_container_file"}
+
+
+def test_unlisted_class_level_use_fails_the_fence(surface_check):
+ found = {
+ "members": {"outside": {"get_valid_args": {"litellm/new_module.py"}}, "router_family": {}},
+ "router_utils_imports": {},
+ }
+ allow = {"members_outside": [], "members_router_family": [], "router_utils_imports_outside": []}
+ errors = surface_check.compare(allow, found)
+ assert len(errors) == 1
+ assert "'get_valid_args' in litellm/new_module.py" in errors[0]
+
+
+def test_unrelated_receivers_are_not_uses(surface_check):
+ assert _uses(surface_check, "client.get_model_list()\nRouterConfig.model_list") == set()