scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

Merge remote-tracking branch 'origin/main' into litellm_vertex_identity_labels

Browse source
This commit is contained in:
mateo-berri 2026-10-28 17:24:42 -07:00
commit 8d0512872b
132 changed files with 2560 additions and 566 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
.github/workflows/test-linting.yml vendored
View file

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

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

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

4
litellm-rust/Cargo.lock generated
View file

@ -4393,6 +4393,10 @@ dependencies = [
"rstest",
]
[[package]]
name = "litellm-router-core"
version = "0.1.0"
[[package]]
name = "litellm-router-types"
version = "0.1.0"

33
litellm-rust/crates/model-catalog/src/lib.rs
View file

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

9
litellm-rust/crates/router-core/Cargo.toml Normal file
View file

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

63
litellm-rust/crates/router-core/src/lib.rs Normal file
View file

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

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

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

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

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

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

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

9
litellm/google_genai/streaming_iterator.py
View file

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

5
litellm/interactions/streaming_iterator.py
View file

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

0
litellm/router_utils/batch_utils.py → litellm/litellm_core_utils/batch_utils.py
View file

2
litellm/litellm_core_utils/prompt_templates/common_utils.py
View file

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

57
litellm/litellm_core_utils/streaming_handler.py
View file

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

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

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

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

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

2
litellm/llms/anthropic/prompt_cache_prediction.py
View file

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

28
litellm/model_prices_and_context_window_backup.json
View file

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

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

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

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

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

2
litellm/proxy/common_request_processing.py
View file

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

2
litellm/proxy/guardrails/auto_router_compression.py
View file

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

2
litellm/proxy/health_check.py
View file

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

2
litellm/proxy/hooks/autorouter_baseline_cache.py
View file

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

2
litellm/proxy/hooks/model_max_budget_limiter.py
View file

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

2
litellm/proxy/litellm_pre_call_utils.py
View file

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

12
litellm/proxy/management_endpoints/auto_router_endpoints.py
View file

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

31
litellm/proxy/management_endpoints/model_management_endpoints.py
View file

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

4
litellm/proxy/management_helpers/auto_router_availability.py
View file

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

5
litellm/proxy/management_helpers/auto_router_permissions.py
View file

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

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

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

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

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

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

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

14
litellm/proxy/proxy_server.py
View file

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

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

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

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

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

76
litellm/router.py
View file

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

0
litellm/router_utils/auto_router_model_naming.py → litellm/router_strategy/auto_router/auto_router_model_naming.py
View file

0
litellm/router_utils/auto_router_tuning_baseline.py → litellm/router_strategy/auto_router/auto_router_tuning_baseline.py
View file

0
litellm/router_utils/baseline_request.py → litellm/router_strategy/auto_router/baseline_request.py
View file

2
litellm/router_strategy/budget_limiter.py
View file

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

2
litellm/router_strategy/least_busy.py
View file

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

2
litellm/router_strategy/lowest_cost.py
View file

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

2
litellm/router_strategy/lowest_latency.py
View file

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

2
litellm/router_strategy/lowest_tpm_rpm.py
View file

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

2
litellm/router_strategy/lowest_tpm_rpm_v2.py
View file

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

2
litellm/router_utils/fallback_event_handlers.py
View file

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

28
model_prices_and_context_window.json
View file

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

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

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

4
tests/code_coverage_tests/router_code_coverage.py
View file

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

211
tests/code_coverage_tests/router_surface_allowlist.json Normal file
View file

@ -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"
]
}

235
tests/code_coverage_tests/router_surface_check.py Normal file
View file

@ -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.<name>.
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:]))

2
tests/code_coverage_tests/unbounded_in_baseline.txt
View file

@ -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 <module> raw-sql classifier_type.IN `IN ({_LLM_CLASSIFIER_TYPES_SQL})` 0
litellm/router_strategy/auto_router/auto_router_model_naming.py <module> raw-sql classifier_type.IN `IN ({_LLM_CLASSIFIER_TYPES_SQL})` 0

65
tests/integration/providers/test_gemini_grounding_annotations_wire.py Normal file
View file

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

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

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

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

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

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

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

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

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

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

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

13
tests/unit/router_utils/test_batch_utils.py → tests/unit/litellm_core_utils/test_batch_utils.py
View file

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

Some files were not shown because too many files have changed in this diff Show more