diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261016000000_background_interaction_settlement_settled_at_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261016000000_background_interaction_settlement_settled_at_index/migration.sql
new file mode 100644
index 00000000000..81d32345794
--- /dev/null
+++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261016000000_background_interaction_settlement_settled_at_index/migration.sql
@@ -0,0 +1,2 @@
+CREATE INDEX CONCURRENTLY IF NOT EXISTS "idx_background_interaction_settlement_settled_at"
+ON "LiteLLM_BackgroundInteractionSettlement" ("settled_at");
diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma
index d40bac380aa..83a0e6f9e05 100644
--- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma
+++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma
@@ -1981,6 +1981,7 @@ model LiteLLM_BackgroundInteractionSettlement {
outcome String?
@@index([claimed_at], map: "idx_background_interaction_settlement_claimed_at")
+ @@index([settled_at], map: "idx_background_interaction_settlement_settled_at")
}
model LiteLLM_Lens {
diff --git a/litellm/__init__.py b/litellm/__init__.py
index ff901ede512..88e56c14187 100644
--- a/litellm/__init__.py
+++ b/litellm/__init__.py
@@ -207,7 +207,7 @@ _async_failure_callback: List[Union[str, Callable[..., object], "CustomLogger"]]
[]
) # internal variable - async custom callbacks are routed here.
pre_call_rules: List[Callable] = []
-post_call_rules: List[Callable] = []
+post_call_rules: List[Callable[..., object]] = []
turn_off_message_logging: Optional[bool] = False
standard_logging_payload_excluded_fields: Optional[List[str]] = (
None # Fields to exclude from StandardLoggingPayload before callbacks receive it
@@ -393,6 +393,7 @@ key_alias_pattern: str | None = None
enable_gemini_default_thinking_level_low: bool = (
False # opt-in: force thinkingLevel low/minimal for Gemini 3 thinking param mapping
)
+gemini_include_thoughts_default: bool = True
####################
logging: bool = True
enable_loadbalancing_on_batch_endpoints: Optional[bool] = None
diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py
index 655131e3c69..5f5b8210096 100644
--- a/litellm/caching/caching.py
+++ b/litellm/caching/caching.py
@@ -25,7 +25,7 @@ from litellm._internal_context import current_service_target, service_target
from litellm._logging import verbose_logger
from litellm.constants import CACHED_STREAMING_CHUNK_DELAY
from litellm.integrations.otel.runtime import phase_span
-from litellm.litellm_core_utils.model_param_helper import ModelParamHelper
+from litellm.litellm_core_utils.model_param_helper import ModelParamHelper, peek_reasoning_summary_aliases
from litellm.types.caching import *
from litellm.types.utils import EmbeddingResponse, is_litellm_owned_kwarg
@@ -376,7 +376,7 @@ class Cache:
)
return "".join(f"{field}: {value}" for field, value in scope_values if value is not None)
- def get_cache_key(self, **kwargs) -> str:
+ def get_cache_key(self, **kwargs: object) -> str:
"""
Get the cache key for the given arguments.
@@ -410,13 +410,16 @@ class Cache:
if litellm.enable_caching_on_provider_specific_optional_params is True: # feature flagged for now
if kwargs[param] is None:
continue # ignore None params
- param_value = kwargs[param]
- cache_key += f"{param}: {param_value}"
+ cache_key += f"{param}: {kwargs[param]}"
if is_semantic_cache:
cache_key += self._get_semantic_cache_tenant_scope(kwargs)
- hashed_cache_key = Cache._get_hashed_cache_key(cache_key)
+ reasoning_summary: Final = peek_reasoning_summary_aliases(kwargs)
+ reasoning_summary_material: Final = (
+ "" if reasoning_summary is None else f"reasoning_summary: {reasoning_summary}"
+ )
+ hashed_cache_key = Cache._get_hashed_cache_key(cache_key + reasoning_summary_material)
hashed_cache_key = self._add_namespace_to_cache_key(hashed_cache_key, **kwargs)
verbose_logger.debug(
"\nCreated cache key: %s (source material length: %d)",
diff --git a/litellm/constants.py b/litellm/constants.py
index 665ab771726..39363cc02e4 100644
--- a/litellm/constants.py
+++ b/litellm/constants.py
@@ -57,6 +57,7 @@ RUNTIME_UPDATABLE_ROUTER_SETTINGS: Final[frozenset[str]] = frozenset(
"model_group_retry_policy",
"model_group_alias",
"enable_weighted_failover",
+ "enable_mid_stream_fallback_continuation",
"enable_tag_filtering",
"tag_routing_prefix",
"optional_pre_call_checks",
@@ -146,6 +147,7 @@ HF_CONFIG_FETCH_TIMEOUT_SECONDS: Final = 10.0
# Streams exceeding this duration are terminated with a Timeout error.
# None (default) = no limit. Set env var to a number of seconds to enable globally.
LITELLM_MAX_STREAMING_DURATION_SECONDS: Final = get_env_float_or_none("LITELLM_MAX_STREAMING_DURATION_SECONDS")
+MID_STREAM_FAILURE_LOG_WAIT_SECONDS: Final = 1.0
# Maximum number of base64 characters to keep in logging payloads.
# Data URIs exceeding this are replaced with a size placeholder.
@@ -1629,6 +1631,7 @@ INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin"
MANAGED_BACKGROUND_RESPONSE_METADATA_KEY: Final = "_managed_background_response"
SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated"
SESSION_ID_OMITTED_METADATA_KEY: Final = "litellm_session_id_omitted"
+CALLER_SESSION_ID_METADATA_KEY: Final = "litellm_caller_session_id"
LITELLM_TRUNCATED_PAYLOAD_FIELD: Final = "litellm_truncated"
LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE: Final = (
"Truncation is a DB storage safeguard. "
@@ -1895,6 +1898,7 @@ _background_interaction_cost_polling_env: Final = os.getenv(
"BACKGROUND_INTERACTION_COST_POLLING_ENABLED", "true"
).lower()
BACKGROUND_INTERACTION_COST_POLLING_ENABLED: Final = _background_interaction_cost_polling_env == "true"
+BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS: Final = 300.0
PROXY_BUDGET_RESCHEDULER_MAX_TIME: Final = get_env_int("PROXY_BUDGET_RESCHEDULER_MAX_TIME", 605)
PROXY_BATCH_WRITE_AT: Final = get_env_int("PROXY_BATCH_WRITE_AT", 10) # in seconds, increased from 10
PROXY_CONFIG_RELOAD_INTERVAL_SECONDS: Final = get_env_int("PROXY_CONFIG_RELOAD_INTERVAL_SECONDS", 30)
diff --git a/litellm/exceptions.py b/litellm/exceptions.py
index 31388b00554..b3a0802db13 100644
--- a/litellm/exceptions.py
+++ b/litellm/exceptions.py
@@ -1217,6 +1217,7 @@ class MidStreamFallbackError(ServiceUnavailableError):
num_retries: int | None = None,
generated_content: str = "",
is_pre_first_chunk: bool = False,
+ emitted_disqualifying_content: bool = False,
):
original_status: Final = getattr(original_exception, "status_code", None)
self.status_code = int(original_status) if original_status is not None else 503
@@ -1229,6 +1230,7 @@ class MidStreamFallbackError(ServiceUnavailableError):
self.num_retries = num_retries
self.generated_content = generated_content
self.is_pre_first_chunk = is_pre_first_chunk
+ self.emitted_disqualifying_content = emitted_disqualifying_content
# Create a response if one wasn't provided
if response is None:
diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py
index 2a2a2a0326d..93dd3d0f5ee 100644
--- a/litellm/integrations/custom_guardrail.py
+++ b/litellm/integrations/custom_guardrail.py
@@ -695,9 +695,7 @@ class CustomGuardrail(CustomLogger):
supported_event_hooks: list[GuardrailEventHooks],
) -> None:
allowed_hooks: Final = frozenset(supported_event_hooks) | (
- frozenset((GuardrailEventHooks.logging_only,))
- if self.uses_apply_guardrail_interface() and not self.use_native_lifecycle_hooks
- else frozenset()
+ frozenset((GuardrailEventHooks.logging_only,)) if self.accepts_logging_only_mode() else frozenset()
)
def _validate_event_hook_list_is_in_supported_event_hooks(
@@ -869,8 +867,13 @@ class CustomGuardrail(CustomLogger):
return True
return False
- def uses_apply_guardrail_interface(self) -> bool:
- return type(self).apply_guardrail is not CustomGuardrail.apply_guardrail
+ @classmethod
+ def uses_apply_guardrail_interface(cls) -> bool:
+ return cls.apply_guardrail is not CustomGuardrail.apply_guardrail
+
+ @classmethod
+ def accepts_logging_only_mode(cls) -> bool:
+ return cls.uses_apply_guardrail_interface() and not cls.use_native_lifecycle_hooks
@classmethod
def supports_logging_only_scope(cls) -> bool:
@@ -1029,7 +1032,7 @@ class CustomGuardrail(CustomLogger):
)
return kwargs, result
litellm_params: Final = kwargs.get("litellm_params") or {}
- scratch_metadata: Final = {
+ scratch_metadata: Final[dict[str, object]] = { # mutable-ok: apply_guardrail records its verdict here
key: value
for key, value in (litellm_params.get("metadata") or {}).items()
if key != "standard_logging_guardrail_information"
@@ -1083,11 +1086,11 @@ class CustomGuardrail(CustomLogger):
async def _scan_logged_call(
self,
- kwargs: dict, # mutable-ok: CustomLogger.async_logging_hook contract
+ kwargs: dict[str, object], # mutable-ok: CustomLogger.async_logging_hook contract
response: object | None,
translation: "BaseTranslation",
output_translation: "BaseTranslation",
- scratch_metadata: dict, # mutable-ok: apply_guardrail records its verdict into request metadata
+ scratch_metadata: dict[str, object], # mutable-ok: apply_guardrail records its verdict into request metadata
) -> None:
scratch_fields: Final = self._copy_scratch_request_fields(kwargs)
scratch_input, scratch_tools = scratch_fields or (None, None)
diff --git a/litellm/integrations/gitlab/gitlab_prompt_manager.py b/litellm/integrations/gitlab/gitlab_prompt_manager.py
index 85b91821f6c..adaacd615a8 100644
--- a/litellm/integrations/gitlab/gitlab_prompt_manager.py
+++ b/litellm/integrations/gitlab/gitlab_prompt_manager.py
@@ -219,7 +219,7 @@ class GitLabTemplateManager:
elif value.replace(".", "").isdigit():
try:
result[key] = float(value)
- except Exception:
+ except ValueError:
result[key] = value
else:
result[key] = value.strip("\"'")
diff --git a/litellm/interactions/background_cost_polling.py b/litellm/interactions/background_cost_polling.py
index 21d438d9621..275eb17787e 100644
--- a/litellm/interactions/background_cost_polling.py
+++ b/litellm/interactions/background_cost_polling.py
@@ -32,19 +32,24 @@ task to settle. A ``BackgroundSettlementStore`` makes the pending settlement
durable across processes: the create registers the request context that
billing needs (never provider credentials), the settlement is claimed
exactly once through the store, and a delete on any replica rebuilds the
-billing context from the store when the poll task is not local. Rows left
-unclaimed by a process that died are resumed at startup. The default store is
-in-memory, which keeps the SDK and single-process behavior unchanged; the
-proxy installs a database-backed one.
+billing context from the store when the poll task is not local. Every row
+still unsettled when a process starts is resumed there, including one whose
+settler died between its claim and its outcome: the database-backed store
+treats a claim as a lease that its settler renews until it records the
+outcome, so such a row is claimable again once the lease runs out. The
+default store is in-memory, which keeps the SDK and single-process behavior
+unchanged; the proxy installs a database-backed one.
"""
import asyncio
+import itertools
from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping, Sequence
from dataclasses import dataclass, field
from datetime import datetime, timezone
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias
+from openai import APIStatusError
from pydantic import ConfigDict, JsonValue, TypeAdapter, ValidationError
from pydantic_core import PydanticSerializationError, to_jsonable_python
@@ -54,10 +59,12 @@ from litellm.constants import (
BACKGROUND_INTERACTION_COST_POLL_MAX_INTERVAL_SECONDS,
BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS,
BACKGROUND_INTERACTION_COST_POLLING_ENABLED,
+ BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS,
)
from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs
from litellm.types.interactions import InteractionsAPIResponse
from litellm.types.llms.base import LiteLLMBaseModel
+from litellm.types.router import CredentialLiteLLMParams
from litellm.types.utils import CustomPricingLiteLLMParams
if TYPE_CHECKING:
@@ -101,6 +108,7 @@ class PendingBackgroundInteraction:
custom_llm_provider: str
create_context: BackgroundInteractionCreateContext
created_at: datetime
+ claim_expires_at: datetime | None = None
class BackgroundSettlementStore(Protocol):
@@ -112,9 +120,11 @@ class BackgroundSettlementStore(Protocol):
async def claim(self, interaction_id: str) -> bool: ...
+ async def renew_claim(self, interaction_id: str) -> None: ...
+
async def record_outcome(self, interaction_id: str, outcome: SettlementOutcome) -> None: ...
- async def unclaimed(self) -> Sequence[PendingBackgroundInteraction]: ...
+ async def unsettled(self) -> Sequence[PendingBackgroundInteraction]: ...
@dataclass(frozen=True, slots=True)
@@ -145,10 +155,13 @@ class InMemoryBackgroundSettlementStore:
self._rows[interaction_id] = None
return True
+ async def renew_claim(self, interaction_id: str) -> None:
+ return None
+
async def record_outcome(self, interaction_id: str, outcome: SettlementOutcome) -> None:
return None
- async def unclaimed(self) -> Sequence[PendingBackgroundInteraction]:
+ async def unsettled(self) -> Sequence[PendingBackgroundInteraction]:
return tuple(row for row in self._rows.values() if row is not None)
@@ -169,11 +182,15 @@ class BackgroundInteractionPollContext:
interaction_id: str
custom_llm_provider: str
logging_obj: "LiteLLMLoggingObj"
- api_key: str | None = None
- api_base: str | None = None
+ credentials: CredentialLiteLLMParams = field(default_factory=CredentialLiteLLMParams)
+ deployment_id: str | None = None
+ model_group: str | None = None
initial_interval_seconds: float = BACKGROUND_INTERACTION_COST_POLL_INITIAL_INTERVAL_SECONDS
max_interval_seconds: float = BACKGROUND_INTERACTION_COST_POLL_MAX_INTERVAL_SECONDS
timeout_seconds: float = BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS
+ start_delay_seconds: float = 0.0
+ outcome_write_timeout_seconds: float = BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS / 2
+ claim_renewal_interval_seconds: float = BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS / 3
store: BackgroundSettlementStore = field(default_factory=InMemoryBackgroundSettlementStore)
resumed: bool = False
@@ -184,15 +201,41 @@ FetchInteraction: TypeAlias = Callable[[BackgroundInteractionPollContext], Await
async def fetch_background_interaction(context: BackgroundInteractionPollContext) -> InteractionsAPIResponse:
from litellm.interactions import aget
+ credentials: Final = context.credentials
return await aget(
interaction_id=context.interaction_id,
custom_llm_provider=context.custom_llm_provider,
- api_key=context.api_key,
- api_base=context.api_base,
+ api_key=credentials.api_key,
+ api_base=credentials.api_base,
+ vertex_project=credentials.vertex_project,
+ vertex_location=credentials.vertex_location,
+ vertex_credentials=credentials.vertex_credentials,
**{"no-log": True},
)
+_FETCH_CREDENTIAL_FIELDS: Final = ("api_key", "api_base", "vertex_project", "vertex_location", "vertex_credentials")
+
+
+def _fetch_credentials(request_kwargs: Mapping[str, object]) -> CredentialLiteLLMParams:
+ try:
+ return CredentialLiteLLMParams.model_validate(
+ {field_name: request_kwargs.get(field_name) for field_name in _FETCH_CREDENTIAL_FIELDS}
+ )
+ except ValidationError:
+ verbose_logger.debug(
+ "Fetching a background interaction with environment credentials: the request's are malformed"
+ )
+ return CredentialLiteLLMParams()
+
+
+_PERMANENT_FETCH_STATUS_CODES: Final = frozenset({400, 401, 403, 404})
+
+
+def _is_permanent_fetch_error(error: Exception) -> bool:
+ return isinstance(error, APIStatusError) and error.status_code in _PERMANENT_FETCH_STATUS_CODES
+
+
def _poll_intervals(initial: float, maximum: float, timeout: float) -> Iterator[float]:
elapsed = 0.0 # rebind-ok: the schedule accumulates the time it has already yielded
interval = initial # rebind-ok: the schedule doubles the interval up to the cap
@@ -231,6 +274,7 @@ _UNCARRIED_METADATA_KEY: Final = "user_api_key_auth"
_JSON_VALUE: Final = TypeAdapter(JsonValue)
_STRING: Final = TypeAdapter(str)
_OBJECT_MAPPING: Final = TypeAdapter(Mapping[str, object])
+_OBJECT_DICT: Final = TypeAdapter(dict[str, object])
def _carries(key: str) -> bool:
@@ -256,10 +300,36 @@ def _as_datetime(start_time: datetime | float) -> datetime:
return start_time if isinstance(start_time, datetime) else datetime.fromtimestamp(start_time, tz=timezone.utc)
+def _spend_metadata(logging_obj: "LiteLLMLoggingObj") -> Mapping[str, object]:
+ """
+ The metadata the creating worker's own spend row is built from, with the
+ request tags computed once, header-derived ones included, because a
+ rebuilt logging object has no request headers to derive them from.
+ """
+ litellm_params: Final = _OBJECT_DICT.validate_python(logging_obj.model_call_details.get("litellm_params") or {})
+ return {
+ **_OBJECT_MAPPING.validate_python(get_litellm_metadata_from_kwargs(kwargs=logging_obj.model_call_details)),
+ **_request_tags(litellm_params),
+ }
+
+
+def _request_tags(litellm_params: Mapping[str, object]) -> Mapping[str, object]:
+ from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
+
+ try:
+ return {
+ "tags": StandardLoggingPayloadSetup.get_request_tags(
+ litellm_params=_OBJECT_DICT.validate_python(litellm_params),
+ proxy_server_request=_OBJECT_DICT.validate_python(litellm_params.get("proxy_server_request") or {}),
+ )
+ }
+ except AttributeError:
+ verbose_logger.debug("Keeping a background interaction's request tags as sent: they are not a list")
+ return {}
+
+
def _create_context(logging_obj: "LiteLLMLoggingObj", custom_llm_provider: str) -> BackgroundInteractionCreateContext:
- metadata: Final = _OBJECT_MAPPING.validate_python(
- get_litellm_metadata_from_kwargs(kwargs=logging_obj.model_call_details)
- )
+ metadata: Final = _spend_metadata(logging_obj)
litellm_params: Final = _OBJECT_MAPPING.validate_python(logging_obj.litellm_params)
model: Final = logging_obj.model_call_details.get("model")
return BackgroundInteractionCreateContext(
@@ -313,6 +383,16 @@ async def _settled_elsewhere(context: BackgroundInteractionPollContext) -> bool:
return False
+async def _settled_elsewhere_before_the_window_ends(
+ context: BackgroundInteractionPollContext, intervals: Iterator[float]
+) -> bool:
+ for interval in intervals:
+ await asyncio.sleep(interval)
+ if await _settled_elsewhere(context):
+ return True
+ return False
+
+
async def _claim(context: BackgroundInteractionPollContext) -> bool | None:
"""
Exactly-once gate between every settler of one interaction, on every
@@ -326,30 +406,118 @@ async def _claim(context: BackgroundInteractionPollContext) -> bool | None:
return None
-async def _record(context: BackgroundInteractionPollContext, outcome: SettlementOutcome) -> SettlementOutcome:
+async def _wrote_outcome(context: BackgroundInteractionPollContext, outcome: SettlementOutcome) -> bool:
try:
await context.store.record_outcome(context.interaction_id, outcome)
- except Exception: # noqa: BLE001 # the outcome is an audit trail; the claim already made the settlement exclusive
- verbose_logger.exception("Could not record the settlement of background interaction %s", context.interaction_id)
+ except Exception as e: # noqa: BLE001 # a failed write is retried while the claim still shuts out takeovers
+ verbose_logger.warning(
+ "Could not record the settlement of background interaction %s: %s", context.interaction_id, e
+ )
+ return False
+ return True
+
+
+async def _recorded(context: BackgroundInteractionPollContext, outcome: SettlementOutcome) -> bool:
+ """
+ Retries a failed write for half the claim lease, while no other settler
+ can take the row over, so a database blip does not leave the row to a
+ takeover or hold the budget reservation that only a recorded release frees.
+ """
+ retries: Final = _poll_intervals(
+ initial=context.initial_interval_seconds,
+ maximum=context.max_interval_seconds,
+ timeout=context.outcome_write_timeout_seconds,
+ )
+ for delay in itertools.chain((0.0,), retries):
+ await asyncio.sleep(delay)
+ if await _wrote_outcome(context, outcome):
+ return True
+ verbose_logger.error(
+ "Could not record the settlement of background interaction %s within %ss; the next process to take "
+ "over its lapsed claim settles it",
+ context.interaction_id,
+ context.outcome_write_timeout_seconds,
+ )
+ return False
+
+
+async def _bill(context: BackgroundInteractionPollContext, response: InteractionsAPIResponse) -> SettlementOutcome:
+ await _bill_settled_interaction(logging_obj=context.logging_obj, response=response)
+ await _recorded(context, "billed")
+ return "billed"
+
+
+async def _release(
+ context: BackgroundInteractionPollContext, outcome: Literal["released", "unsettled"]
+) -> SettlementOutcome:
+ """
+ A release writes no spend row, so a takeover cannot tell whether the
+ settler before it released. The outcome is recorded first and the
+ reservation released only once it is, which leaves a takeover with
+ nothing to release twice.
+ """
+ if await _recorded(context, outcome):
+ await _release_open_budget_reservation(logging_obj=context.logging_obj)
return outcome
+async def _renew_claim_until_cancelled(context: BackgroundInteractionPollContext) -> None:
+ while True:
+ await asyncio.sleep(context.claim_renewal_interval_seconds)
+ try:
+ await context.store.renew_claim(context.interaction_id)
+ except Exception as e: # noqa: BLE001 # a missed renewal is retried next interval, well inside the lease
+ verbose_logger.warning(
+ "Could not renew the settlement claim of background interaction %s: %s", context.interaction_id, e
+ )
+
+
+async def _settle(
+ context: BackgroundInteractionPollContext, response: InteractionsAPIResponse | None
+) -> SettlementOutcome:
+ """
+ Billing runs every success callback in turn, and a slow one can outlast
+ the claim lease, so the claim is renewed until the outcome is recorded.
+ A settler that stops renewing has died, and only then can another take
+ the row over.
+ """
+ renewal: Final = asyncio.create_task(_renew_claim_until_cancelled(context))
+ try:
+ if response is not None and response.status in _TERMINAL_STATUSES and response.usage is not None:
+ return await _bill(context, response)
+ return await _release(context, "released")
+ finally:
+ renewal.cancel()
+
+
async def poll_and_log_background_interaction_cost(
context: BackgroundInteractionPollContext,
fetch_interaction: FetchInteraction = fetch_background_interaction,
) -> SettlementOutcome | None:
last_response: InteractionsAPIResponse | None = None # rebind-ok: the give-up path settles from the last poll
- for interval in _poll_intervals(
+ await asyncio.sleep(context.start_delay_seconds)
+ intervals: Final = _poll_intervals(
initial=context.initial_interval_seconds,
maximum=context.max_interval_seconds,
timeout=context.timeout_seconds,
- ):
+ )
+ for interval in intervals:
await asyncio.sleep(interval)
if await _settled_elsewhere(context):
return None
try:
response = await fetch_interaction(context)
except Exception as e: # noqa: BLE001 # any fetch error must not kill the billing poll loop
+ if context.resumed and _is_permanent_fetch_error(e):
+ verbose_logger.warning(
+ "Stopped fetching resumed background interaction %s: the provider answered %s, so this "
+ "process cannot read it and leaves it to the replica that created it until the poll window ends",
+ context.interaction_id,
+ e,
+ )
+ if await _settled_elsewhere_before_the_window_ends(context, intervals):
+ return None
+ break
verbose_logger.debug(
"Background interaction cost poll for %s failed, will retry: %s",
context.interaction_id,
@@ -363,11 +531,11 @@ async def poll_and_log_background_interaction_cost(
continue
if not claimed:
return None
- return await _record(context, await _settle_terminal(logging_obj=context.logging_obj, response=response))
+ return await _settle(context, response)
if not await _claim(context):
return None
if last_response is not None and last_response.status in _TERMINAL_STATUSES:
- return await _record(context, await _settle_terminal(logging_obj=context.logging_obj, response=last_response))
+ return await _settle(context, last_response)
if last_response is not None and last_response.status not in _POLLABLE_STATUSES:
verbose_logger.error(
"Gave up cost polling for background interaction %s after %ss: its last status %r is in neither "
@@ -379,20 +547,10 @@ async def poll_and_log_background_interaction_cost(
)
else:
verbose_logger.warning(
- "Gave up cost polling for background interaction %s after %ss; its usage will not be tracked",
+ "Gave up cost polling for background interaction %s; its usage will not be tracked",
context.interaction_id,
- context.timeout_seconds,
)
- await _release_open_budget_reservation(logging_obj=context.logging_obj)
- return await _record(context, "unsettled")
-
-
-async def _settle_terminal(logging_obj: "LiteLLMLoggingObj", response: InteractionsAPIResponse) -> SettlementOutcome:
- if response.status in _TERMINAL_STATUSES and response.usage is not None:
- await _bill_settled_interaction(logging_obj=logging_obj, response=response)
- return "billed"
- await _release_open_budget_reservation(logging_obj=logging_obj)
- return "released"
+ return await _release(context, "unsettled")
async def _release_open_budget_reservation(logging_obj: "LiteLLMLoggingObj") -> None:
@@ -405,7 +563,9 @@ async def _release_open_budget_reservation(logging_obj: "LiteLLMLoggingObj") ->
settlement must release the reservation here or the spend counters stay
pinned at the estimated cost.
"""
- metadata: Final = get_litellm_metadata_from_kwargs(kwargs=logging_obj.model_call_details)
+ metadata: Final = _OBJECT_MAPPING.validate_python(
+ get_litellm_metadata_from_kwargs(kwargs=logging_obj.model_call_details)
+ )
budget_reservation: Final = metadata.get("user_api_key_budget_reservation")
if not isinstance(budget_reservation, dict):
return
@@ -518,13 +678,18 @@ class _UnverifiedRegistrationStore:
return False
return await self.local.claim(interaction_id)
+ async def renew_claim(self, interaction_id: str) -> None:
+ if await self.local.is_claimed(interaction_id):
+ return
+ await self.durable.renew_claim(interaction_id)
+
async def record_outcome(self, interaction_id: str, outcome: SettlementOutcome) -> None:
if await self.local.is_claimed(interaction_id):
return
await self.durable.record_outcome(interaction_id, outcome)
- async def unclaimed(self) -> Sequence[PendingBackgroundInteraction]:
- return await self.durable.unclaimed()
+ async def unsettled(self) -> Sequence[PendingBackgroundInteraction]:
+ return await self.durable.unsettled()
async def _registered_store(
@@ -559,8 +724,6 @@ async def maybe_schedule_background_interaction_cost_polling(
logging_obj: Final = create_kwargs.get("litellm_logging_obj")
if not isinstance(logging_obj, Logging):
return None
- api_key: Final = create_kwargs.get("api_key")
- api_base: Final = create_kwargs.get("api_base")
pending: Final = PendingBackgroundInteraction(
interaction_id=response.id,
custom_llm_provider=custom_llm_provider,
@@ -571,8 +734,7 @@ async def maybe_schedule_background_interaction_cost_polling(
interaction_id=response.id,
custom_llm_provider=custom_llm_provider,
logging_obj=logging_obj,
- api_key=api_key if isinstance(api_key, str) else None,
- api_base=api_base if isinstance(api_base, str) else None,
+ credentials=_fetch_credentials(create_kwargs),
store=await _registered_store(store or _STORE.store, pending),
)
return _track_poll(context, fetch_interaction)
@@ -603,10 +765,7 @@ async def _settle_before_delete(
) -> SettlementOutcome | None:
if not await _claim(context):
return None
- if response is None:
- await _release_open_budget_reservation(logging_obj=context.logging_obj)
- return await _record(context, "released")
- return await _record(context, await _settle_terminal(logging_obj=context.logging_obj, response=response))
+ return await _settle(context, response)
async def maybe_settle_background_interaction_before_delete(
@@ -622,14 +781,11 @@ async def maybe_settle_background_interaction_before_delete(
pending: Final = await _pending(settlement_store, interaction_id)
if pending is None:
return None
- api_key: Final = delete_kwargs.get("api_key")
- api_base: Final = delete_kwargs.get("api_base")
context: Final = BackgroundInteractionPollContext(
interaction_id=interaction_id,
custom_llm_provider=pending.custom_llm_provider,
logging_obj=_rebuild_logging_obj(pending.create_context),
- api_key=api_key if isinstance(api_key, str) else None,
- api_base=api_base if isinstance(api_base, str) else None,
+ credentials=_fetch_credentials(delete_kwargs),
store=settlement_store,
)
try:
@@ -644,9 +800,9 @@ async def maybe_settle_background_interaction_before_delete(
return await _settle_before_delete(context, response)
-async def _unclaimed(store: BackgroundSettlementStore) -> Sequence[PendingBackgroundInteraction]:
+async def _unsettled(store: BackgroundSettlementStore) -> Sequence[PendingBackgroundInteraction]:
try:
- return await store.unclaimed()
+ return await store.unsettled()
except Exception: # noqa: BLE001 # an unreadable store at startup leaves its rows for the next boot
verbose_logger.exception("Could not list the unsettled background interactions")
return ()
@@ -662,17 +818,36 @@ class PollSchedule:
DEFAULT_POLL_SCHEDULE: Final = PollSchedule()
+def _deployment_id(create_context: BackgroundInteractionCreateContext) -> str | None:
+ model_info: Final = create_context.metadata.get("model_info")
+ deployment_id: Final = model_info.get("id") if isinstance(model_info, dict) else None
+ return deployment_id if isinstance(deployment_id, str) else None
+
+
+def _model_group(create_context: BackgroundInteractionCreateContext) -> str | None:
+ model_group: Final = create_context.metadata.get("model_group")
+ return model_group if isinstance(model_group, str) else None
+
+
+def _seconds_until(moment: datetime | None) -> float:
+ if moment is None:
+ return 0.0
+ return max((moment - datetime.now(timezone.utc)).total_seconds(), 0.0)
+
+
def _resumed_context(
row: PendingBackgroundInteraction, store: BackgroundSettlementStore, schedule: PollSchedule
) -> BackgroundInteractionPollContext:
- age_seconds: Final = (datetime.now(timezone.utc) - row.created_at).total_seconds()
return BackgroundInteractionPollContext(
interaction_id=row.interaction_id,
custom_llm_provider=row.custom_llm_provider,
logging_obj=_rebuild_logging_obj(row.create_context),
+ deployment_id=_deployment_id(row.create_context),
+ model_group=_model_group(row.create_context),
initial_interval_seconds=schedule.initial_interval_seconds,
max_interval_seconds=schedule.max_interval_seconds,
- timeout_seconds=max(schedule.timeout_seconds - age_seconds, schedule.initial_interval_seconds),
+ timeout_seconds=schedule.timeout_seconds,
+ start_delay_seconds=_seconds_until(row.claim_expires_at),
store=store,
resumed=True,
)
@@ -684,13 +859,15 @@ async def resume_unsettled_background_interactions(
schedule: PollSchedule = DEFAULT_POLL_SCHEDULE,
) -> tuple["asyncio.Task[SettlementOutcome | None]", ...]:
"""
- Pick up every settlement no process has claimed, which is what a replica
- that died mid-poll leaves behind. Each resumed poll keeps the remaining
- share of the original timeout and gets at least one fetch, so a completed
- interaction is still billed however late the resume comes.
+ Pick up every settlement no process has finished, which is what a replica
+ that died mid-poll or mid-settlement leaves behind. A row whose claim is
+ still within its lease waits until the lease runs out, then polls like any
+ other. Each resumed poll gets the full poll window however old its row is,
+ so a fetch that fails right after a late restart is retried rather than
+ ending the row as unsettled.
"""
return tuple(
_track_poll(_resumed_context(row, store, schedule), fetch_interaction)
- for row in await _unclaimed(store)
+ for row in await _unsettled(store)
if row.interaction_id not in _ACTIVE_POLLS
)
diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py
index e766a2be9c3..fa009ca6327 100644
--- a/litellm/litellm_core_utils/litellm_logging.py
+++ b/litellm/litellm_core_utils/litellm_logging.py
@@ -2295,6 +2295,36 @@ class Logging(LiteLLMLoggingBaseClass):
self.model_call_details["combined_usage_object"] = usage
self.model_call_details["response_cost"] = response_cost
+ def record_continued_partial(self, usage: Usage, response_cost: float, response_id: str) -> None:
+ """Bill what a cut stream already consumed on this continuation's success log, under the id the caller saw."""
+ self.model_call_details["continued_partial_usage"] = usage
+ self.model_call_details["continued_partial_cost"] = response_cost
+ self.model_call_details["continued_response_id"] = response_id
+
+ def _add_continued_partial_spend(self, response: object) -> None:
+ partial_usage: Final = self.model_call_details.get("continued_partial_usage")
+ partial_cost: Final = self.model_call_details.get("continued_partial_cost")
+ own_usage: Final = getattr(response, "usage", None)
+ own_cost: Final = self.model_call_details.get("response_cost")
+ if not (isinstance(partial_usage, Usage) and isinstance(partial_cost, float) and isinstance(own_usage, Usage)):
+ return
+ from litellm.cost_calculator import BaseTokenUsageProcessor
+
+ self.model_call_details["combined_usage_object"] = BaseTokenUsageProcessor.combine_usage_objects(
+ usage_objects=[own_usage, partial_usage]
+ )
+ if not isinstance(own_cost, (int, float)):
+ return
+ self.model_call_details["response_cost"] = own_cost + partial_cost
+ if self.cost_breakdown is not None:
+ self.cost_breakdown = self.cost_breakdown | CostBreakdown(
+ total_cost=own_cost + partial_cost,
+ additional_costs={
+ **self.cost_breakdown.get("additional_costs", {}),
+ "mid_stream_partial_cost": partial_cost,
+ },
+ )
+
def record_assembled_response_for_failure(self, assembled: ModelResponse) -> None:
"""Bill a fully streamed response on the failure log when a post-call hook rejects it."""
usage: Final = getattr(assembled, "usage", None)
@@ -2590,6 +2620,15 @@ class Logging(LiteLLMLoggingBaseClass):
return response
return response.model_copy(update={"model": model})
+ def _with_continued_response_id(
+ self,
+ response: ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse,
+ ) -> ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse:
+ response_id: Final = self.model_call_details.get("continued_response_id")
+ if not isinstance(response_id, str) or getattr(response, "id", None) == response_id:
+ return response
+ return response.model_copy(update={"id": response_id})
+
def _success_handler_helper_fn(
self,
result: object = None,
@@ -2908,8 +2947,11 @@ class Logging(LiteLLMLoggingBaseClass):
self.model_call_details["response_cost"] = self.response_cost_calculator(
result=complete_streaming_response
)
+ self._add_continued_partial_spend(complete_streaming_response)
self._merge_hidden_params_from_response_into_metadata(complete_streaming_response)
- logged_streaming_response: Final = self._with_client_facing_stream_model(complete_streaming_response)
+ logged_streaming_response: Final = self._with_continued_response_id(
+ self._with_client_facing_stream_model(complete_streaming_response)
+ )
self.model_call_details["complete_streaming_response"] = logged_streaming_response
## STANDARDIZED LOGGING PAYLOAD
self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload(
@@ -3445,12 +3487,15 @@ class Logging(LiteLLMLoggingBaseClass):
self.model,
)
self.model_call_details["response_cost"] = None
+ self._add_continued_partial_spend(complete_streaming_response)
self._merge_hidden_params_from_response_into_metadata(complete_streaming_response)
await self._prepare_baseline_cache_estimate(complete_streaming_response)
- logged_streaming_response: Final = self._with_client_facing_stream_model(complete_streaming_response)
+ logged_streaming_response: Final = self._with_continued_response_id(
+ self._with_client_facing_stream_model(complete_streaming_response)
+ )
self.model_call_details["async_complete_streaming_response"] = logged_streaming_response
## STANDARDIZED LOGGING PAYLOAD
diff --git a/litellm/litellm_core_utils/llm_response_utils/get_headers.py b/litellm/litellm_core_utils/llm_response_utils/get_headers.py
index b86cc954b0e..3c08f16fa80 100644
--- a/litellm/litellm_core_utils/llm_response_utils/get_headers.py
+++ b/litellm/litellm_core_utils/llm_response_utils/get_headers.py
@@ -1,8 +1,11 @@
from collections.abc import Mapping
from typing import Final
+import httpx
from pydantic import TypeAdapter, ValidationError
+from litellm.constants import PROVIDER_REQUEST_ID_HEADERS
+
_HEADERS: Final = TypeAdapter(Mapping[str, str])
_REQUEST_ID_HEADERS: Final = ("request-id", "x-request-id", "llm_provider-request-id", "llm_provider-x-request-id")
@@ -16,6 +19,16 @@ def get_provider_request_id(headers: object) -> str | None:
return next((normalized[name] for name in _REQUEST_ID_HEADERS if normalized.get(name)), None)
+def provider_request_id_headers(headers: Mapping[str, object] | None) -> httpx.Headers | None:
+ """The provider request id headers in `headers` whose values are ASCII strings, or None when there are none."""
+ request_id_headers: Final = tuple(
+ (name, value)
+ for name, value in (headers or {}).items()
+ if name.lower() in PROVIDER_REQUEST_ID_HEADERS and isinstance(value, str) and value.isascii()
+ )
+ return httpx.Headers(request_id_headers) if request_id_headers else None
+
+
def get_response_headers(_response_headers: Mapping[str, str] | None = None) -> dict:
"""
diff --git a/litellm/litellm_core_utils/model_param_helper.py b/litellm/litellm_core_utils/model_param_helper.py
index 2dc52980875..c9c7a0e89ac 100644
--- a/litellm/litellm_core_utils/model_param_helper.py
+++ b/litellm/litellm_core_utils/model_param_helper.py
@@ -1,3 +1,4 @@
+from collections.abc import Mapping
from functools import lru_cache
from typing import ClassVar, Final
@@ -22,6 +23,24 @@ from litellm.types.llms.anthropic import AnthropicMessagesRequest
from litellm.types.rerank import RerankRequest
+def peek_reasoning_summary_aliases(optional_params: Mapping[str, object]) -> object | None:
+ """Read AI-SDK-style reasoning summary from optional_params or nested extra_body.
+
+ Uses key membership (not ``or`` chains) so falsy values like ``""`` are not skipped.
+ """
+ if "reasoningSummary" in optional_params:
+ return optional_params["reasoningSummary"]
+ if "reasoning_summary" in optional_params:
+ return optional_params["reasoning_summary"]
+ extra_body: Final = optional_params.get("extra_body")
+ if isinstance(extra_body, dict):
+ if "reasoningSummary" in extra_body:
+ return extra_body["reasoningSummary"]
+ if "reasoning_summary" in extra_body:
+ return extra_body["reasoning_summary"]
+ return None
+
+
class ModelParamHelper:
# Cached at class level — deterministic set built from static OpenAI type annotations
relevant_logging_args: ClassVar[frozenset[str]] = frozenset()
diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py
index 81fc865a5bc..044123d8729 100644
--- a/litellm/litellm_core_utils/prompt_templates/common_utils.py
+++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py
@@ -1894,6 +1894,11 @@ def strip_litellm_internal_message_fields(message: AllMessageValues) -> AllMessa
return _message_without_fields(message, LITELLM_INTERNAL_MESSAGE_FIELDS)
+def strip_prefill_prefix_field(message: AllMessageValues) -> AllMessageValues:
+ """Drop the ``prefix`` flag Mistral and DeepSeek read off a prefill, for endpoints that reject unknown keys."""
+ return _message_without_fields(message, frozenset({"prefix"}))
+
+
def filter_value_from_dict(dictionary: dict, key: str, depth: int = 0) -> Any:
"""
Filters a value from a dictionary
diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py
index 0c18ce38b13..87e76fe39f7 100644
--- a/litellm/litellm_core_utils/prompt_templates/factory.py
+++ b/litellm/litellm_core_utils/prompt_templates/factory.py
@@ -9,10 +9,11 @@ import xml.etree.ElementTree as ET
from collections.abc import Container, Iterable, Iterator, Mapping, Sequence
from enum import Enum
from types import MappingProxyType
-from typing import Any, Final, TypeAlias, TypedDict, cast, overload
+from typing import Annotated, Any, Final, TypeAlias, TypedDict, cast, overload
from jinja2.sandbox import ImmutableSandboxedEnvironment
-from pydantic import BaseModel, JsonValue, TypeAdapter
+from pydantic import BaseModel, BeforeValidator, ConfigDict, JsonValue, TypeAdapter, ValidationError, model_validator
+from pydantic_core import ErrorDetails, PydanticCustomError
import litellm
import litellm.types
@@ -2582,6 +2583,48 @@ def _with_native_redacted_data(
_AnthropicMessageList: TypeAlias = list[AllAnthropicPassThroughMessageValues]
+def _client_tool_result(message: AllMessageValues, replaced_server_tool_answers: frozenset[str]) -> bool:
+ return (
+ message_field(message, "role") in ("tool", "function")
+ and message_field(message, "tool_call_id") not in replaced_server_tool_answers
+ )
+
+
+def _tool_result_later_in_turn(
+ messages: Sequence[AllMessageValues], replaced_server_tool_answers: frozenset[str]
+) -> tuple[bool, ...]:
+ from_the_end: Final = itertools.accumulate(
+ reversed(messages),
+ lambda later, message: (
+ message_field(message, "role") in ("user", "tool", "function")
+ and (later or _client_tool_result(message, replaced_server_tool_answers))
+ ),
+ initial=False,
+ )
+ return tuple(from_the_end)[::-1][1:]
+
+
+def _empty_text_placeholder(message: AllMessageValues) -> AnthropicMessagesTextParam:
+ placeholder: Final = AnthropicMessagesTextParam(type="text", text=_EMPTY_TEXT_PLACEHOLDER)
+ add_cache_control_to_content(anthropic_content_element=placeholder, original_content_element=message)
+ return placeholder
+
+
+def _cache_breakpoint_placeholder(
+ message: AllMessageValues, tool_result_follows: bool
+) -> tuple[AnthropicMessagesTextParam, ...]:
+ if tool_result_follows or not message_field(message, "cache_control"):
+ return ()
+ return (_empty_text_placeholder(message),)
+
+
+def _empty_user_turn_placeholder(span: Sequence[AllMessageValues]) -> AnthropicMessagesTextParam | None:
+ empty_user_message: Final = next(
+ (message for message in reversed(span) if message["role"] == "user" and not message.get("content")), None
+ )
+ return None if empty_user_message is None else _empty_text_placeholder(empty_user_message)
+
+
def anthropic_messages_pt(
messages: list[AllMessageValues],
model: str,
@@ -2636,6 +2679,7 @@ def anthropic_messages_pt(
is_vertex_ai = llm_provider.startswith("vertex_ai") if llm_provider else False
force_base64 = is_bedrock_invoke or is_vertex_ai
+ tool_result_later_in_turn: Final = _tool_result_later_in_turn(messages, replaced_server_tool_answers)
msg_i = 0
while msg_i < len(messages):
user_content: list[AnthropicMessagesUserMessageValues] = []
@@ -2648,8 +2692,12 @@ def anthropic_messages_pt(
ChatCompletionToolMessage | ChatCompletionUserMessage | ChatCompletionFunctionMessage
) = messages[msg_i]
if user_message_types_block["role"] == "user":
- if isinstance(user_message_types_block["content"], list):
- for m in user_message_types_block["content"]:
+ if not (user_block_content := user_message_types_block.get("content")):
+ user_content.extend(
+ _cache_breakpoint_placeholder(user_message_types_block, tool_result_later_in_turn[msg_i])
+ )
+ elif isinstance(user_block_content, list):
+ for m in user_block_content:
if m.get("type", "") == "image_url":
m = cast(ChatCompletionImageObject, m)
format = m["image_url"].get("format") if isinstance(m["image_url"], dict) else None
@@ -2720,10 +2768,10 @@ def anthropic_messages_pt(
_file_content_element,
)
)
- elif isinstance(user_message_types_block["content"], str):
+ elif isinstance(user_block_content, str):
_anthropic_content_text_element: AnthropicMessagesTextParam = {
"type": "text",
- "text": user_message_types_block["content"],
+ "text": user_block_content,
}
_content_element = add_cache_control_to_content(
anthropic_content_element=_anthropic_content_text_element,
@@ -2747,6 +2795,8 @@ def anthropic_messages_pt(
if user_content:
new_messages.append({"role": "user", "content": user_content})
+ elif (placeholder := _empty_user_turn_placeholder(messages[init_msg_i:msg_i])) is not None:
+ new_messages.append({"role": "user", "content": [placeholder]})
# Track unique tool IDs in this merge block to avoid duplication
unique_tool_ids: set[str] = set()
@@ -4253,6 +4303,24 @@ def _with_text_when_document_only(message: BedrockMessageBlock) -> BedrockMessag
return BedrockMessageBlock(role="user", content=[*blocks[:cut], placeholder, *blocks[cut:]])
+def _rstrip_text_block(block: BedrockContentBlock) -> BedrockContentBlock:
+ return BedrockContentBlock(text=block["text"].rstrip()) if "text" in block else block
+
+
+def _assistant_turn_rstripped(message: BedrockMessageBlock) -> BedrockMessageBlock:
+ """Claude on Converse rejects a final assistant turn that ends in whitespace, as the Anthropic path strips it."""
+ if message["role"] != "assistant":
+ return message
+ blocks: Final = message["content"]
+ last_text_index: Final = max((index for index, block in enumerate(blocks) if "text" in block), default=None)
+ return BedrockMessageBlock(
+ role="assistant",
+ content=[
+ _rstrip_text_block(block) if index == last_text_index else block for index, block in enumerate(blocks)
+ ],
+ )
+
+
def _ensure_document_messages_have_text(
contents: list[BedrockMessageBlock],
) -> list[BedrockMessageBlock]:
@@ -4844,7 +4912,12 @@ class BedrockConverseMessagesProcessor:
llm_provider=llm_provider,
)
- return _ensure_document_messages_have_text(_rename_duplicate_bedrock_document_names(contents))
+ converse_messages: Final = _ensure_document_messages_have_text(
+ _rename_duplicate_bedrock_document_names(contents)
+ )
+ if not (to_claude and converse_messages):
+ return converse_messages
+ return [*converse_messages[:-1], _assistant_turn_rstripped(converse_messages[-1])]
@staticmethod
def translate_thinking_blocks_to_reasoning_content_blocks(
@@ -5253,7 +5326,10 @@ def _bedrock_converse_messages_pt(
llm_provider=llm_provider,
)
- return _ensure_document_messages_have_text(_rename_duplicate_bedrock_document_names(contents))
+ converse_messages: Final = _ensure_document_messages_have_text(_rename_duplicate_bedrock_document_names(contents))
+ if not (to_claude and converse_messages):
+ return converse_messages
+ return [*converse_messages[:-1], _assistant_turn_rstripped(converse_messages[-1])]
def make_valid_bedrock_tool_name(input_tool_name: str) -> str:
@@ -5516,7 +5592,7 @@ def function_call_prompt(
]
-def response_schema_prompt(model: str, response_schema: dict) -> str:
+def response_schema_prompt(model: str, response_schema: dict, llm_provider: str) -> str:
"""
Decides if a user-defined custom prompt or default needs to be used
@@ -5532,11 +5608,8 @@ def response_schema_prompt(model: str, response_schema: dict) -> str:
custom_prompt_details = litellm.custom_prompt_dict["response_schema_prompt"]
if custom_prompt_details is not None:
- return custom_prompt(
- role_dict=custom_prompt_details["roles"],
- initial_prompt_value=custom_prompt_details["initial_prompt_value"],
- final_prompt_value=custom_prompt_details["final_prompt_value"],
- messages=response_schema_as_message,
+ return render_custom_prompt_template(
+ template=custom_prompt_details, messages=response_schema_as_message, model=model, llm_provider=llm_provider
)
else:
return default_response_schema_prompt(response_schema=response_schema)
@@ -5556,8 +5629,95 @@ def default_response_schema_prompt(response_schema: dict) -> str:
# Custom prompt template
+def _mapping_or_none(value: object) -> object:
+ if value is None or issubclass(type(value), Mapping):
+ return value
+ verbose_logger.warning(
+ "Ignoring a custom prompt template roles value of type %s because it is not a mapping, "
+ "so those messages render without role markers",
+ type(value).__name__,
+ )
+ return None
+
+
+class CustomPromptRole(BaseModel):
+ model_config = ConfigDict(frozen=True)
+
+ pre_message: str | None = None
+ post_message: str | None = None
+
+
+class CustomPromptRoles(BaseModel):
+ model_config = ConfigDict(frozen=True, title="custom prompt template")
+
+ roles: Mapping[str, Annotated[CustomPromptRole | None, BeforeValidator(_mapping_or_none)]]
+
+
+class CustomPromptTemplate(BaseModel):
+ model_config = ConfigDict(frozen=True, title="custom prompt template")
+
+ roles: Annotated[Mapping[object, object] | None, BeforeValidator(_mapping_or_none)] = None
+ initial_prompt_value: str | None = None
+ final_prompt_value: str | None = None
+
+ @model_validator(mode="after")
+ def _sets_a_rendered_field(self) -> "CustomPromptTemplate":
+ if self.roles is None and self.initial_prompt_value is None and self.final_prompt_value is None:
+ raise PydanticCustomError(
+ "empty_custom_prompt_template",
+ "it sets none of roles, initial_prompt_value, or final_prompt_value",
+ )
+ return self
+
+
+def render_custom_prompt_template(
+ template: Mapping[str, object],
+ messages: Sequence[Mapping[str, object]],
+ model: str,
+ llm_provider: str,
+) -> str:
+ try:
+ return _render_custom_prompt_template(template=template, messages=messages)
+ except ValidationError as error:
+ raise litellm.APIConnectionError(
+ message=f"The custom prompt template for {model} is invalid: {_describe_template_errors(error)}",
+ model=model,
+ llm_provider=llm_provider,
+ ) from error
+
+
+def _describe_template_errors(error: ValidationError) -> str:
+ return "; ".join(
+ _describe_template_error(detail) for detail in error.errors(include_url=False, include_input=False)
+ )
+
+
+def _describe_template_error(detail: ErrorDetails) -> str:
+ location: Final = ".".join(str(part) for part in detail["loc"])
+ return f"{location}: {detail['msg']}" if location else detail["msg"]
+
+
+def _render_custom_prompt_template(template: Mapping[str, object], messages: Sequence[Mapping[str, object]]) -> str:
+ validated: Final = CustomPromptTemplate.model_validate(template)
+ message_roles: Final = frozenset(message.get("role") for message in messages)
+ template_roles: Final = validated.roles or {}
+ used_roles: Final = CustomPromptRoles.model_validate(
+ {"roles": {role: entry for role, entry in template_roles.items() if role in message_roles}}
+ )
+ return custom_prompt(
+ role_dict={
+ role: {"pre_message": entry.pre_message or "", "post_message": entry.post_message or ""}
+ for role, entry in used_roles.roles.items()
+ if entry is not None
+ },
+ initial_prompt_value=validated.initial_prompt_value or "",
+ final_prompt_value=validated.final_prompt_value or "",
+ messages=list(messages),
+ )
+
+
def custom_prompt(
- role_dict: dict,
+ role_dict: Mapping[str, Mapping[str, str]],
messages: list,
initial_prompt_value: str = "",
final_prompt_value: str = "",
diff --git a/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py b/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py
index 4d1e77126a9..fd0096f9f4d 100644
--- a/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py
+++ b/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py
@@ -414,10 +414,11 @@ def _assistant_renders(message: object, answered: frozenset[str]) -> bool:
def _renders(message: object, answered: frozenset[str]) -> bool:
"""Whether ``anthropic_messages_pt`` puts a block on the wire for this message.
- A tool message always becomes a ``tool_result`` and a user message with string
- content always becomes a text block (empty text gets a placeholder). A user list
- renders only through parts of a type the converter emits; ``None``, an empty list,
- and a list of other parts vanish. Assistant messages follow ``_assistant_renders``.
+ A tool message always becomes a ``tool_result``, and a user message with string or
+ missing content (``None``, an empty list) always puts its turn on the wire, an empty
+ turn getting a placeholder. A user list renders only through parts of a type the
+ converter emits, so a list of only other parts vanishes. Assistant messages follow
+ ``_assistant_renders``.
"""
role: Final = message_field(message, "role")
if role in _TOOL_ROLES:
@@ -425,8 +426,10 @@ def _renders(message: object, answered: frozenset[str]) -> bool:
if role == "assistant":
return _assistant_renders(message, answered)
content: Final = message_field(message, "content")
- return isinstance(content, str) or any(
- message_field(part, "type") in _RENDERED_PART_TYPES for part in parts_of(content)
+ return (
+ isinstance(content, str)
+ or not content
+ or any(message_field(part, "type") in _RENDERED_PART_TYPES for part in parts_of(content))
)
diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py
index c9b96c0a4c2..e4eb9acc37d 100644
--- a/litellm/litellm_core_utils/realtime_streaming.py
+++ b/litellm/litellm_core_utils/realtime_streaming.py
@@ -1,16 +1,18 @@
import asyncio
import json
import traceback
-from collections.abc import Coroutine, Mapping, Sequence
+from collections.abc import Callable, Coroutine, Mapping, Sequence
from dataclasses import dataclass
from enum import Enum, auto
-from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypedDict, cast
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypeAlias, TypedDict, cast
from pydantic import TypeAdapter
from typing_extensions import ReadOnly
import litellm
from litellm._logging import redact_internal_details_from_client_message, verbose_logger
+from litellm._uuid import uuid
from litellm.constants import (
REALTIME_METERED_USAGE_KEY,
REALTIME_SESSION_FAILURE_LOGGED_KEY,
@@ -116,10 +118,87 @@ def _decode_json_object(payload: str) -> Mapping[str, object]:
return json.loads(payload)
+_GUARDRAIL_TURN_DETECTION_GATE: Final[Mapping[str, object]] = MappingProxyType(
+ {"type": "server_vad", "create_response": False}
+)
+_EMPTY_SESSION: Final[Mapping[str, object]] = MappingProxyType({})
+_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object])
+_GUARDRAIL_EVENT_ID_PREFIX: Final = "litellm_guardrail_"
+
+
+def _json_object(value: object) -> Mapping[str, object] | None:
+ return _JSON_OBJECT.validate_python(value) if isinstance(value, Mapping) else None
+
+
+def _mapping_field(container: Mapping[str, object], key: str) -> Mapping[str, object] | None:
+ return _json_object(container.get(key))
+
+
+def _guardrail_event_id() -> str:
+ return f"{_GUARDRAIL_EVENT_ID_PREFIX}{uuid.uuid4().hex}"
+
+
+def _error_event_id(event: Mapping[str, object]) -> str | None:
+ error: Final = _mapping_field(event, "error")
+ caused_by: Final = None if error is None else error.get("event_id")
+ return caused_by if isinstance(caused_by, str) else None
+
+
+_TurnDetectionGate: TypeAlias = Callable[[object], Mapping[str, object]]
+
+
+def _turn_detection_with_auto_response_off(turn_detection: object) -> Mapping[str, object]:
+ client_turn_detection: Final = _json_object(turn_detection)
+ if client_turn_detection is None:
+ return dict(_GUARDRAIL_TURN_DETECTION_GATE)
+ return {**client_turn_detection, "create_response": False}
+
+
+def _nested_turn_detection_part(session: Mapping[str, object], gate: _TurnDetectionGate) -> Mapping[str, object]:
+ audio: Final = _mapping_field(session, "audio")
+ if audio is None:
+ return {}
+ audio_input: Final = _mapping_field(audio, "input")
+ if audio_input is None or "turn_detection" not in audio_input:
+ return {}
+ gated_input: Final = {
+ **audio_input,
+ "turn_detection": gate(audio_input["turn_detection"]),
+ }
+ return {"audio": {**audio, "input": gated_input}}
+
+
+def _with_create_response_disabled(
+ session: Mapping[str, object], gate: _TurnDetectionGate
+) -> Mapping[str, object] | None:
+ """``session`` with every turn_detection the client sent unable to auto-respond, or None when it sent none."""
+ flat_part: Final = {"turn_detection": gate(session["turn_detection"])} if "turn_detection" in session else {}
+ nested_part: Final = _nested_turn_detection_part(session, gate)
+ if not flat_part and not nested_part:
+ return None
+ return {**session, **flat_part, **nested_part}
+
+
+def _gate_session_update(
+ session: object, gate: _TurnDetectionGate, inject_gate_when_absent: bool
+) -> Mapping[str, object] | None:
+ """The session to forward under a transcript guardrail, or None when the frame needs no change."""
+ client_session: Final = _json_object(session)
+ if client_session is None:
+ return None
+ gated: Final = _with_create_response_disabled(client_session, gate)
+ if gated is not None or not inject_gate_when_absent:
+ return gated
+ return {**client_session, "turn_detection": gate(None)}
+
+
class RealtimeEventNormalizer(Protocol):
def should_drop(self, event: object) -> bool: ...
def normalize(self, event: dict) -> dict: ...
def patch_outgoing_session(self, session: dict) -> dict: ...
+ def gate_turn_detection(self, turn_detection: object, /) -> Mapping[str, object]: ...
+ @property
+ def deletes_conversation_items(self) -> bool: ...
class RealtimeUsageMeter(Protocol):
@@ -197,6 +276,8 @@ class RealTimeStreaming:
# Track whether we have already sent the guardrail turn-detection update
# that disables provider auto-response for transcription guardrails.
self._guardrail_turn_detection_update_sent: bool = False
+ self._session_gate_sent: bool = False
+ self._guardrail_cancel_event_ids: frozenset[str] = frozenset()
# Deferred Gemini Live setup: Pipecat may stream audio before session.update.
# Buffer client audio until the backend acknowledges setup (setupComplete).
self._backend_setup_complete: bool = provider_config is None or provider_config.requires_session_configuration()
@@ -213,6 +294,12 @@ class RealTimeStreaming:
self._is_transcription_session: bool = force_transcription_model is not None
# Optional per-provider GA event normalizer (e.g. XAIRealtimeNormalizer).
self._event_normalizer = event_normalizer
+ self._gate_turn_detection: Final[_TurnDetectionGate] = (
+ _turn_detection_with_auto_response_off if event_normalizer is None else event_normalizer.gate_turn_detection
+ )
+ self._deletes_conversation_items: Final = (
+ event_normalizer is None or event_normalizer.deletes_conversation_items
+ )
self._usage_meter = usage_meter
@property
@@ -734,20 +821,37 @@ class RealTimeStreaming:
except (json.JSONDecodeError, TypeError):
return
+ def _make_turn_detection_update(self, turn_detection: Mapping[str, object]) -> dict[str, object]:
+ session: Final[dict[str, object]] = (
+ {"turn_detection": turn_detection}
+ if self._backend_uses_beta_protocol
+ else {"type": "realtime", "audio": {"input": {"turn_detection": turn_detection}}}
+ )
+ return {"type": "session.update", "session": session}
+
def _make_disable_auto_response_message(self) -> str:
"""Return a session.update that disables VAD auto-response."""
- turn_detection: Final[dict[str, str | bool]] = {
- "type": "server_vad",
- "create_response": False,
- }
- if self._backend_uses_beta_protocol:
- session: dict[str, object] = {"turn_detection": turn_detection}
- else:
- session = {
- "type": "realtime",
- "audio": {"input": {"turn_detection": turn_detection}},
- }
- return json.dumps({"type": "session.update", "session": session})
+ return json.dumps(self._make_turn_detection_update(self._gate_turn_detection(None)))
+
+ async def _send_guardrail_cancel(self) -> None:
+ event_id: Final = _guardrail_event_id()
+ self._guardrail_cancel_event_ids = self._guardrail_cancel_event_ids | {event_id}
+ await self._send_to_backend(json.dumps({"type": "response.cancel", "event_id": event_id}))
+
+ async def _withdraw_blocked_turn(self, item_id: str | None) -> None:
+ await self._send_guardrail_cancel()
+ if item_id is not None and self._deletes_conversation_items:
+ await self._send_to_backend(json.dumps({"type": "conversation.item.delete", "item_id": item_id}))
+
+ def _is_guardrail_cancel_error(self, event: Mapping[str, object]) -> bool:
+ return _error_event_id(event) in self._guardrail_cancel_event_ids
+
+ async def _send_session_gate_once(self) -> None:
+ """Gate the backend's default turn_detection ahead of the first client session.update, which overrides it."""
+ if self._session_gate_sent or not self._should_disable_vad_auto_response():
+ return
+ self._session_gate_sent = True
+ await self._send_to_backend(self._make_disable_auto_response_message())
async def _maybe_send_guardrail_turn_detection_update(self) -> None:
"""Disable provider auto-response once when transcription guardrails are enabled."""
@@ -921,7 +1025,7 @@ class RealTimeStreaming:
if pre_block_backend_message is not None:
await self._send_to_backend(pre_block_backend_message)
if not self._is_transcription_session:
- await self._send_to_backend(json.dumps({"type": "response.cancel"}))
+ await self._withdraw_blocked_turn(item_id)
# Send the policy violation hint (shows as small gray status text in UI).
await self.websocket.send_text(
json.dumps(
@@ -1051,10 +1155,7 @@ class RealTimeStreaming:
self.store_message(event_str)
self._capture_transcription_usage(event)
await self._send_event_to_client(event, event_str)
- blocked = await self.run_realtime_guardrails(
- cast(str, transcript),
- item_id=cast(str | None, event.get("item_id")),
- )
+ blocked = await self.run_realtime_guardrails(cast(str, transcript))
if not blocked and self._should_disable_vad_auto_response():
await self._send_to_backend(json.dumps({"type": "response.create"}))
continue
@@ -1086,7 +1187,7 @@ class RealTimeStreaming:
if event_type == "session.created" and self._should_disable_vad_auto_response():
self.store_message(event_obj)
await self.websocket.send_text(self._event_to_client_json(event_obj))
- await self._send_to_backend(self._make_disable_auto_response_message())
+ await self._send_session_gate_once()
return True
if event_type == "conversation.item.input_audio_transcription.completed":
@@ -1140,6 +1241,10 @@ class RealTimeStreaming:
if self._usage_meter is not None:
self._usage_meter.observe_backend_event(event)
+ if event.get("type") == "error" and self._is_guardrail_cancel_error(event):
+ verbose_logger.debug("Realtime: not forwarding an error caused by a guardrail event: %s", event)
+ continue
+
if await self._handle_raw_backend_message(event, raw_response):
continue
@@ -1472,19 +1577,21 @@ class RealTimeStreaming:
and self.session_configuration_request is None
and not self._guardrail_turn_detection_update_sent
and self._should_disable_vad_auto_response()
- ):
- session: Mapping[str, object] | None = msg_obj.setdefault("session", {})
- if isinstance(session, dict):
- existing_td = session.get("turn_detection")
- if not isinstance(existing_td, dict):
- existing_td = {}
- existing_td["create_response"] = False
- session["turn_detection"] = existing_td
- message = json.dumps(msg_obj)
- guardrail_turn_detection_injected = True
- verbose_logger.debug(
- "Injected turn_detection into first session.update for audio transcription guardrails"
+ and (
+ gated_first := _gate_session_update(
+ client_event.get("session", _EMPTY_SESSION),
+ self._gate_turn_detection,
+ inject_gate_when_absent=True,
)
+ )
+ is not None
+ ):
+ msg_obj["session"] = gated_first
+ message = json.dumps(msg_obj)
+ guardrail_turn_detection_injected = True
+ verbose_logger.debug(
+ "Injected turn_detection into first session.update for audio transcription guardrails"
+ )
## GUARDRAIL: Force ``create_response`` to False in any
# client-provided ``turn_detection`` so a later
@@ -1492,52 +1599,25 @@ class RealTimeStreaming:
# and bypass the transcription guardrail after the
# initial disable. Covers both the flat beta key and the
# nested GA ``audio.input.turn_detection`` shape, since
- # the GA remap below also accepts either form. Skipped
- # when the injection block above already ran for this
- # message, to avoid redundant double-serialization.
+ # the GA remap below also accepts either form. A later
+ # update with no turn_detection is forwarded as is: OpenAI
+ # keeps the earlier disable for an unset field, and xAI
+ # falls back to its manual default, which never auto-responds.
if (
msg_type == "session.update"
and not guardrail_turn_detection_injected
and self._should_disable_vad_auto_response()
+ and (
+ gated_later := _gate_session_update(
+ client_event.get("session"),
+ self._gate_turn_detection,
+ inject_gate_when_absent=False,
+ )
+ )
+ is not None
):
- session = client_event.get("session")
- if isinstance(session, dict):
- td_overridden = False
- flat_td = session.get("turn_detection")
- flat_td_present = flat_td is not None
- if flat_td_present:
- if not isinstance(flat_td, dict):
- flat_td = {}
- if flat_td.get("create_response") is not False:
- flat_td["create_response"] = False
- session["turn_detection"] = flat_td
- td_overridden = True
- nested_td_present = False
- audio = session.get("audio")
- if isinstance(audio, dict):
- audio_input = audio.get("input")
- if isinstance(audio_input, dict):
- nested_td = audio_input.get("turn_detection")
- if nested_td is not None:
- nested_td_present = True
- if not isinstance(nested_td, dict):
- nested_td = {}
- if nested_td.get("create_response") is not False:
- nested_td["create_response"] = False
- audio_input["turn_detection"] = nested_td
- td_overridden = True
- # Symmetric with the first-update injection block:
- # if the client omitted turn_detection entirely on
- # a subsequent session.update, still inject the
- # ``create_response: False`` override so the
- # transcription guardrail cannot be re-enabled by
- # any downstream merge that drops the original
- # disable.
- if not flat_td_present and not nested_td_present:
- session["turn_detection"] = {"create_response": False}
- td_overridden = True
- if td_overridden:
- message = json.dumps(msg_obj)
+ msg_obj["session"] = gated_later
+ message = json.dumps(msg_obj)
# GA compatibility: remap beta-style session fields only when
# the upstream is in GA mode. Beta upstreams expect the flat
@@ -1590,6 +1670,8 @@ class RealTimeStreaming:
# would permanently disable the injection if ``_send_to_backend``
# raised — neither this loop nor
# ``_maybe_send_guardrail_turn_detection_update`` would retry.
+ if msg_type == "session.update" and self.provider_config is None:
+ await self._send_session_gate_once()
sent = await self._send_to_backend(message)
if guardrail_turn_detection_injected and sent:
self._guardrail_turn_detection_update_sent = True
diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py
index 46ad1b70abe..95d062a190c 100644
--- a/litellm/litellm_core_utils/streaming_handler.py
+++ b/litellm/litellm_core_utils/streaming_handler.py
@@ -242,6 +242,7 @@ class CustomStreamWrapper:
self.completion_stream = completion_stream
self.sent_first_chunk = False
self.sent_last_chunk = False
+ self._failure_logging: asyncio.Task[None] | None = None
self._closed = False
self._stream_created_time: float = time.time()
@@ -269,6 +270,7 @@ class CustomStreamWrapper:
self.holding_chunk = ""
self.complete_response = ""
self.response_uptil_now = ""
+ self._emitted_disqualifying_content = False
_model_info: Final[dict] = litellm_params.model_info or {}
_api_base: Final = get_api_base(
@@ -1129,6 +1131,7 @@ class CustomStreamWrapper:
if self.merge_reasoning_content_in_choices is True:
reasoning_content: Final = getattr(model_response.choices[0].delta, "reasoning_content", None)
if reasoning_content:
+ self._emitted_disqualifying_content = True
if self.sent_first_thinking_block is False:
# Ensure content is not None before concatenation
if model_response.choices[0].delta.content is None:
@@ -1865,9 +1868,7 @@ class CustomStreamWrapper:
if response.choices:
choice = response.choices[0]
if isinstance(choice, StreamingChoices):
- self.response_uptil_now += choice.delta.get("content", "") or ""
- else:
- self.response_uptil_now += ""
+ self._accumulate_streamed_delta(choice.delta)
self.rules.post_call_rules(input=self.response_uptil_now, model=self.model)
# HANDLE STREAM OPTIONS
self.chunks.append(response)
@@ -2084,9 +2085,7 @@ class CustomStreamWrapper:
if processed_chunk.choices:
choice = processed_chunk.choices[0]
if isinstance(choice, StreamingChoices):
- self.response_uptil_now += choice.delta.get("content", "") or ""
- else:
- self.response_uptil_now += ""
+ self._accumulate_streamed_delta(choice.delta)
self.rules.post_call_rules(input=self.response_uptil_now, model=self.model)
# Add mcp_list_tools to first chunk if present
if not self.sent_first_chunk and processed_chunk.choices:
@@ -2156,9 +2155,7 @@ class CustomStreamWrapper:
choice = processed_chunk.choices[0]
if isinstance(choice, StreamingChoices):
- self.response_uptil_now += choice.delta.get("content", "") or ""
- else:
- self.response_uptil_now += ""
+ self._accumulate_streamed_delta(choice.delta)
self.rules.post_call_rules(input=self.response_uptil_now, model=self.model)
# RETURN RESULT
self.chunks.append(processed_chunk)
@@ -2172,7 +2169,7 @@ class CustomStreamWrapper:
if self.logging_obj is not None:
self._record_partial_usage_for_failure()
## LOGGING
- asyncio.create_task(
+ self._failure_logging = asyncio.create_task(
self.logging_obj.dispatch_failure_handlers(e, traceback_exception, prefer_async_handlers=True)
)
await self.aclose()
@@ -2305,11 +2302,17 @@ class CustomStreamWrapper:
if self.logging_obj is not None:
self._record_partial_usage_for_failure()
## LOGGING
- asyncio.create_task(
+ self._failure_logging = asyncio.create_task(
self.logging_obj.dispatch_failure_handlers(e, traceback_exception, prefer_async_handlers=True)
)
self._handle_stream_fallback_error(e)
+ async def wait_for_failure_logging(self, timeout: float) -> None:
+ """Let this stream's failure log read its own deployment before a fallback reuses the logging object."""
+ if self._failure_logging is None:
+ return
+ await asyncio.wait((self._failure_logging,), timeout=timeout)
+
def _record_partial_usage_for_failure(self) -> None:
"""
A stream that breaks mid-flight still billed the provider for the chunks
@@ -2346,6 +2349,35 @@ class CustomStreamWrapper:
recover_error,
)
+ # Delta fields a text-only prefill continuation cannot carry, so a stream
+ # that emitted any of them is not eligible for mid-stream continuation.
+ _CONTINUATION_DISQUALIFYING_DELTA_FIELDS: Final = (
+ "tool_calls",
+ "function_call",
+ "reasoning_content",
+ "thinking_blocks",
+ "reasoning_items",
+ "audio",
+ "images",
+ "annotations",
+ )
+
+ @classmethod
+ def _delta_disqualifies_continuation(cls, delta: object) -> bool:
+ get: Final = getattr(delta, "get", None)
+ if not callable(get):
+ return False
+ return any(get(field) for field in cls._CONTINUATION_DISQUALIFYING_DELTA_FIELDS)
+
+ def _accumulate_streamed_delta(self, delta: object) -> None:
+ # Shared by the sync, async, and non-aiohttp iteration sites so answer
+ # text and the disqualifying-content latch stay in step across all three.
+ get: Final = getattr(delta, "get", None)
+ content: Final = get("content") if callable(get) else None
+ self.response_uptil_now += content if isinstance(content, str) else ""
+ if not self._emitted_disqualifying_content and self._delta_disqualifies_continuation(delta):
+ self._emitted_disqualifying_content = True
+
def _handle_stream_fallback_error(self, e: Exception) -> "NoReturn":
"""
Common error handling for both __next__ and __anext__.
@@ -2377,6 +2409,7 @@ class CustomStreamWrapper:
original_exception=e.original_exception,
generated_content=self.response_uptil_now,
is_pre_first_chunk=not self.sent_first_chunk,
+ emitted_disqualifying_content=self._emitted_disqualifying_content or e.emitted_disqualifying_content,
)
# Map to OpenAI exception format. Some providers' mappers (e.g.
@@ -2437,6 +2470,7 @@ class CustomStreamWrapper:
original_exception=mapped_exception,
generated_content=self.response_uptil_now,
is_pre_first_chunk=not self.sent_first_chunk,
+ emitted_disqualifying_content=self._emitted_disqualifying_content,
)
@staticmethod
diff --git a/litellm/llms/anthropic/completion/transformation.py b/litellm/llms/anthropic/completion/transformation.py
index e3ad961ed76..73e94e85f04 100644
--- a/litellm/llms/anthropic/completion/transformation.py
+++ b/litellm/llms/anthropic/completion/transformation.py
@@ -14,8 +14,8 @@ import httpx
import litellm
from litellm.constants import DEFAULT_MAX_TOKENS
from litellm.litellm_core_utils.prompt_templates.factory import (
- custom_prompt,
prompt_factory,
+ render_custom_prompt_template,
)
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
from litellm.llms.base_llm.chat.transformation import (
@@ -244,11 +244,8 @@ class AnthropicTextConfig(BaseConfig):
if model in custom_prompt_dict:
# check if the model has a registered custom prompt
model_prompt_details: Final = custom_prompt_dict[model]
- prompt = custom_prompt(
- role_dict=model_prompt_details["roles"],
- initial_prompt_value=model_prompt_details["initial_prompt_value"],
- final_prompt_value=model_prompt_details["final_prompt_value"],
- messages=messages,
+ prompt = render_custom_prompt_template(
+ template=model_prompt_details, messages=messages, model=model, llm_provider="anthropic"
)
else:
prompt = prompt_factory(model=model, messages=messages, custom_llm_provider="anthropic")
diff --git a/litellm/llms/anthropic/count_tokens/token_counter.py b/litellm/llms/anthropic/count_tokens/token_counter.py
index d78de0037a0..55eaa0d2ac6 100644
--- a/litellm/llms/anthropic/count_tokens/token_counter.py
+++ b/litellm/llms/anthropic/count_tokens/token_counter.py
@@ -6,6 +6,7 @@ from typing import Any, Final
from pydantic import JsonValue
+import litellm
from litellm._logging import verbose_logger
from litellm.exceptions import AuthenticationError
from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler
@@ -61,14 +62,24 @@ class AnthropicTokenCounter(BaseTokenCounter):
try:
auth_header: Final = await AnthropicModelInfo.aget_auth_header(
- api_key=litellm_params.get("api_key"),
+ api_key=litellm_params.get("api_key") or litellm.anthropic_key or litellm.api_key,
api_base=api_base,
litellm_params=litellm_params,
allow_workload_identity=True,
)
if auth_header is None:
- verbose_logger.warning("No Anthropic credential found for token counting")
- return None
+ return TokenCountResponse(
+ total_tokens=0,
+ request_model=request_model,
+ model_used=model_to_use,
+ tokenizer_type="anthropic_api",
+ error=True,
+ error_message=(
+ f"No Anthropic API key is set for model {request_model or model_to_use}, so tokens can't be "
+ "counted. Set api_key on the deployment or ANTHROPIC_API_KEY in the proxy environment."
+ ),
+ status_code=401,
+ )
result: Final = await anthropic_count_tokens_handler.handle_count_tokens_request(
model=model_to_use,
diff --git a/litellm/llms/anthropic/pass_through/adapters/handler.py b/litellm/llms/anthropic/pass_through/adapters/handler.py
index 4c0613777c6..60c63964492 100644
--- a/litellm/llms/anthropic/pass_through/adapters/handler.py
+++ b/litellm/llms/anthropic/pass_through/adapters/handler.py
@@ -1,4 +1,5 @@
from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence
+from dataclasses import dataclass
from typing import (
TYPE_CHECKING,
Final,
@@ -62,6 +63,12 @@ class _CompletionKwargs(TypedDict, total=False, extra_items=object):
custom_llm_provider: str
+@dataclass(frozen=True, slots=True)
+class _ChatTargetReasoning:
+ effort: str
+ summary: str | None = None
+
+
def _messages_have_compaction_block(messages: _AnthropicMessages) -> bool:
"""Return True when any message carries a ``compaction`` content block."""
for msg in messages:
@@ -441,15 +448,33 @@ class LiteLLMMessagesToCompletionTransformationHandler:
completion_kwargs["reasoning_effort"] = updated_reasoning_effort
@staticmethod
- def _plain_effort_for_chat_target(
+ def _requested_reasoning_summary(thinking: Mapping[str, object]) -> str | None:
+ if thinking.get("type") == "disabled":
+ return None
+ summary: Final = thinking.get("summary")
+ if isinstance(summary, str):
+ return summary
+ display: Final = thinking.get("display")
+ return "auto" if display is not None and display != "omitted" else None
+
+ @staticmethod
+ def _chat_target_reasoning(
completion_kwargs: _CompletionKwargs,
*,
thinking: Mapping[str, object] | None,
- ) -> str | None:
- reasoning_effort: Final = completion_kwargs.get("reasoning_effort")
- if not thinking or not isinstance(reasoning_effort, dict) or "summary" not in reasoning_effort:
+ ) -> _ChatTargetReasoning | None:
+ if not thinking:
return None
- effort: Final = reasoning_effort.get("effort")
+ reasoning_effort: Final = completion_kwargs.get("reasoning_effort")
+ wraps_summary: Final = isinstance(reasoning_effort, dict) and "summary" in reasoning_effort
+ requested_summary: Final = LiteLLMMessagesToCompletionTransformationHandler._requested_reasoning_summary(
+ thinking
+ )
+ if not wraps_summary and requested_summary is None:
+ return None
+ effort: Final[object] = (
+ reasoning_effort.get("effort") if isinstance(reasoning_effort, dict) else reasoning_effort
+ )
model: Final = completion_kwargs.get("model")
if not isinstance(effort, str) or not isinstance(model, str) or not model:
return None
@@ -467,7 +492,12 @@ class LiteLLMMessagesToCompletionTransformationHandler:
return None
if resolved_provider == "litellm_proxy":
return None
- from litellm.responses.bridge_check import responses_api_bridge_check
+ from litellm.responses.bridge_check import chat_target_accepts_reasoning_summary, responses_api_bridge_check
+
+ if requested_summary is not None and chat_target_accepts_reasoning_summary(local_model, resolved_provider):
+ return _ChatTargetReasoning(effort=effort, summary=requested_summary)
+ if not wraps_summary or not isinstance(reasoning_effort, dict):
+ return None
web_search_options: Final = completion_kwargs.get("web_search_options")
tools: Final = completion_kwargs.get("tools")
@@ -481,7 +511,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
reasoning_effort=reasoning_effort,
api_base=resolved_api_base,
)
- return None if model_info.get("mode") == "responses" else effort
+ return None if model_info.get("mode") == "responses" else _ChatTargetReasoning(effort=effort)
@staticmethod
def _normalize_reasoning_effort(
@@ -647,12 +677,14 @@ class LiteLLMMessagesToCompletionTransformationHandler:
thinking=thinking,
)
- plain_effort: Final = LiteLLMMessagesToCompletionTransformationHandler._plain_effort_for_chat_target(
+ chat_reasoning: Final = LiteLLMMessagesToCompletionTransformationHandler._chat_target_reasoning(
completion_kwargs,
thinking=thinking,
)
- if plain_effort is not None:
- completion_kwargs["reasoning_effort"] = plain_effort
+ if chat_reasoning is not None:
+ completion_kwargs["reasoning_effort"] = chat_reasoning.effort
+ if chat_reasoning is not None and chat_reasoning.summary is not None:
+ completion_kwargs["reasoning_summary"] = chat_reasoning.summary
return completion_kwargs, tool_name_mapping
diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py
index 0e0b68df449..53734a77f3c 100644
--- a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py
+++ b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py
@@ -30,7 +30,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig):
self,
model: str,
messages: list[dict[str, JsonValue]],
- api_key: str,
+ api_key: str | None,
api_base: str,
litellm_params: dict[str, object] | None = None,
timeout: float | httpx.Timeout | None = None,
@@ -80,6 +80,15 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig):
api_key=api_key,
litellm_params=litellm_params,
)
+ if not {"x-api-key", "Authorization"} & headers.keys():
+ raise AnthropicError(
+ status_code=401,
+ message=(
+ f"No Azure AI API key or Entra ID credential is set for model {model}, so tokens can't be "
+ "counted. Set api_key or an Entra ID credential on the deployment, or AZURE_AI_API_KEY or "
+ "AZURE_API_KEY in the proxy environment."
+ ),
+ )
# Use LiteLLM's async httpx client
async_client: Final = get_async_httpx_client(llm_provider=litellm.LlmProviders.AZURE_AI)
diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py
index 2707d24ee07..0ed432cc394 100644
--- a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py
+++ b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py
@@ -7,6 +7,7 @@ from typing import Any, Final
from pydantic import JsonValue
+import litellm
from litellm._logging import verbose_logger
from litellm.llms.azure_ai.anthropic.count_tokens.handler import (
AzureAIAnthropicCountTokensHandler,
@@ -18,6 +19,21 @@ from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCou
azure_ai_anthropic_count_tokens_handler: Final = AzureAIAnthropicCountTokensHandler()
+def _missing_api_base_response(request_model: str, model_to_use: str) -> TokenCountResponse:
+ return TokenCountResponse(
+ total_tokens=0,
+ request_model=request_model,
+ model_used=model_to_use,
+ tokenizer_type="azure_ai_anthropic_api",
+ error=True,
+ error_message=(
+ f"No Azure AI API base is set for model {request_model or model_to_use}, so tokens can't be counted. "
+ "Set api_base on the deployment or AZURE_AI_API_BASE or AZURE_API_BASE in the proxy environment."
+ ),
+ status_code=500,
+ )
+
+
class AzureAIAnthropicTokenCounter(BaseTokenCounter):
"""Token counter implementation for Azure AI Anthropic provider using the CountTokens API."""
@@ -65,18 +81,15 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter):
if not api_key:
api_key = os.getenv("AZURE_AI_API_KEY")
- # Get API base from deployment config or environment
- api_base = litellm_params.get("api_base")
- if not api_base:
- api_base = os.getenv("AZURE_AI_API_BASE")
-
- if not api_key:
- verbose_logger.warning("No Azure AI API key found for token counting")
- return None
+ api_base: Final = (
+ litellm_params.get("api_base")
+ or os.getenv("AZURE_AI_API_BASE")
+ or os.getenv("AZURE_API_BASE")
+ or litellm.api_base
+ )
if not api_base:
- verbose_logger.warning("No Azure AI API base found for token counting")
- return None
+ return _missing_api_base_response(request_model, model_to_use)
try:
result: Final = await azure_ai_anthropic_count_tokens_handler.handle_count_tokens_request(
diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py
index d7455439151..5946a1f4480 100644
--- a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py
+++ b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py
@@ -24,7 +24,7 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig):
def get_required_headers(
self,
- api_key: str,
+ api_key: str | None,
litellm_params: dict[str, object] | None = None,
) -> dict[str, str]:
"""
@@ -45,12 +45,12 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig):
"Content-Type": "application/json",
"anthropic-version": "2023-06-01",
"anthropic-beta": ANTHROPIC_TOKEN_COUNTING_BETA_VERSION,
- "x-api-key": api_key, # Azure AI Anthropic requires this header
+ **({"x-api-key": api_key} if api_key else {}),
}
# Also set up Azure auth headers for flexibility
litellm_params = litellm_params or {}
- if "api_key" not in litellm_params:
+ if api_key and "api_key" not in litellm_params:
litellm_params["api_key"] = api_key
litellm_params_obj: Final = GenericLiteLLMParams.model_validate(litellm_params)
@@ -58,8 +58,9 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig):
# Get Azure auth headers (api-key or Authorization)
azure_headers = BaseAzureLLM.base_validate_azure_environment(headers={}, litellm_params=litellm_params_obj)
- # Merge Azure auth headers
headers.update(azure_headers)
+ if "api-key" in azure_headers:
+ headers.setdefault("x-api-key", azure_headers["api-key"])
return headers
diff --git a/litellm/llms/base_llm/chat/transformation.py b/litellm/llms/base_llm/chat/transformation.py
index 2b3c644f266..db4d5e0f820 100644
--- a/litellm/llms/base_llm/chat/transformation.py
+++ b/litellm/llms/base_llm/chat/transformation.py
@@ -488,6 +488,10 @@ class BaseConfig(ABC):
"""
return True
+ @property
+ def accepts_reasoning_summary_in_reasoning_effort(self) -> bool:
+ return False
+
def post_stream_processing(self, stream: object) -> object:
"""Hook for providers to post-process streaming responses. Default: pass-through."""
return stream
diff --git a/litellm/llms/base_llm/realtime/transcription_protocol.py b/litellm/llms/base_llm/realtime/transcription_protocol.py
index c3264911d46..b4514fd1da2 100644
--- a/litellm/llms/base_llm/realtime/transcription_protocol.py
+++ b/litellm/llms/base_llm/realtime/transcription_protocol.py
@@ -102,6 +102,16 @@ def json_integer(
return value
+def uses_server_vad(
+ update: TranscriptionSessionUpdate, backend: str, drop_params: bool, error: ProtocolErrorType
+) -> bool:
+ if update.turn_detection_disabled:
+ return False
+ if update.turn_detection_type in (None, "server_vad") or drop_params:
+ return True
+ raise error(f"{backend} supports server_vad turn detection or null. Set drop_params to use server_vad instead")
+
+
def new_event_id() -> str:
return f"event_{uuid.uuid4().hex}"
diff --git a/litellm/llms/base_llm/videos/async_task_utils.py b/litellm/llms/base_llm/videos/async_task_utils.py
new file mode 100644
index 00000000000..d2fb3ed7606
--- /dev/null
+++ b/litellm/llms/base_llm/videos/async_task_utils.py
@@ -0,0 +1,123 @@
+import base64
+from collections.abc import Mapping
+from io import BufferedReader, BytesIO
+from itertools import accumulate, repeat
+from pathlib import Path, PurePath
+from typing import Final
+
+from pydantic import TypeAdapter, ValidationError
+
+from litellm.images.utils import ImageEditRequestUtils
+from litellm.types.videos.utils import decode_video_id_with_provider
+
+_FIELDS_ADAPTER: Final = TypeAdapter(dict[str, object])
+_FILE_TUPLE_ADAPTER: Final = TypeAdapter(tuple[object, ...])
+
+
+def _fields(value: object) -> Mapping[str, object] | None:
+ try:
+ return _FIELDS_ADAPTER.validate_python(value)
+ except ValidationError:
+ return None
+
+
+def _logged_litellm_params(logging_obj: object) -> Mapping[str, object]:
+ return _fields(getattr(logging_obj, "litellm_params", None)) or {}
+
+
+def _id_safe_name(name: object) -> str | None:
+ return name if isinstance(name, str) and name and ";" not in name else None
+
+
+def _stamped_models(metadata: object) -> tuple[str, str | None] | None:
+ fields: Final = _fields(metadata)
+ if fields is None:
+ return None
+ requested: Final = _id_safe_name(fields.get("model_group"))
+ served: Final = _id_safe_name(fields.get("deployment_model_name"))
+ if requested is None:
+ return None
+ return requested, None if served is None or served == requested or "*" in served else served
+
+
+def routed_video_models(logging_obj: object) -> tuple[str, str | None] | None:
+ """
+ A later status or content request checks the key's model access against the model encoded in the id, and
+ routes by it. One upstream model name can be served by several deployments (regions, accounts), so only the
+ group that served the task can send its polls back to the deployment that created it, while a key may be
+ scoped to the alias, routing group or wildcard name it asked for instead. The id carries both names: the one
+ the router was asked for on this hop, and the group that served it when that differs and is not a wildcard.
+ The router stamps both into litellm_metadata, and metadata carries a copy only when the caller sent none.
+ """
+ params: Final = _logged_litellm_params(logging_obj)
+ stamped: Final = (_stamped_models(params.get(key)) for key in ("litellm_metadata", "metadata"))
+ return next((models for models in stamped if models), None)
+
+
+def polled_video_models(logging_obj: object) -> tuple[str | None, str | None]:
+ """
+ The served group survives into the returned id only when the router sent this poll to it, so an edited id
+ cannot write any name it likes back.
+ """
+ video_id: Final = _logged_litellm_params(logging_obj).get("video_id")
+ if not isinstance(video_id, str):
+ return None, None
+ decoded: Final = decode_video_id_with_provider(video_id)
+ served: Final = decoded.get("served_model")
+ routed: Final = routed_video_models(logging_obj)
+ return decoded.get("model_id") or None, served if routed is None or served in routed else None
+
+
+def _tuple_payload(file_obj: object) -> object:
+ if not isinstance(file_obj, tuple):
+ return file_obj
+ file_tuple: Final = _FILE_TUPLE_ADAPTER.validate_python(file_obj)
+ return file_tuple[1] if len(file_tuple) >= 2 else None
+
+
+def _file_content(file_obj: object) -> object:
+ unwrapped: Final = accumulate(repeat(None), lambda current, _: _tuple_payload(current), initial=file_obj)
+ return next(item for item in unwrapped if not isinstance(item, tuple))
+
+
+def input_reference_bytes(file_obj: object) -> bytes:
+ content: Final = _file_content(file_obj)
+ if isinstance(content, PurePath):
+ return Path(content).read_bytes()
+ if isinstance(content, (BytesIO, BufferedReader)):
+ current_position: Final = content.tell()
+ _ = content.seek(0)
+ data: Final = content.read()
+ _ = content.seek(current_position)
+ return data
+ if isinstance(content, bytes):
+ return content
+ if isinstance(content, bytearray):
+ return bytes(content)
+ read: Final = getattr(content, "read", None)
+ if callable(read):
+ read_data: Final[object] = read()
+ if isinstance(read_data, bytes):
+ return read_data
+ raise ValueError("input_reference must be a URL string, bytes, a path, a file object, or a file tuple")
+
+
+def _reference_object_url(reference: Mapping[str, object]) -> str:
+ image_url: Final = reference.get("image_url")
+ if isinstance(image_url, str) and image_url:
+ return image_url
+ raise ValueError("input_reference as an object must carry an image_url, a file_id is not supported")
+
+
+def input_reference_url(image: object) -> str:
+ """A string is a URL the provider fetches itself and is never opened as a local path. OpenAI's object form
+ carries that URL, or a data URL, in ``image_url``."""
+ if isinstance(image, str):
+ return image
+ reference: Final = _fields(image)
+ if reference is not None:
+ return _reference_object_url(reference)
+ content: Final = input_reference_bytes(image)
+ content_type: Final = ImageEditRequestUtils.get_image_content_type(content)
+ encoded: Final = base64.b64encode(content).decode("utf-8")
+ return f"data:{content_type};base64,{encoded}"
diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py
index 848120150cd..f1426aa9375 100644
--- a/litellm/llms/bedrock/base_aws_llm.py
+++ b/litellm/llms/bedrock/base_aws_llm.py
@@ -99,11 +99,16 @@ def _assume_role_params(
)
-_SecureTransportBool = TypedDict("_SecureTransportBool", {"aws:SecureTransport": ReadOnly[Literal["true"]]})
+_SecureTransportBool: Final = TypedDict("_SecureTransportBool", {"aws:SecureTransport": ReadOnly[Literal["true"]]})
+
+_CalledViaLastEquals: Final = TypedDict(
+ "_CalledViaLastEquals", {"aws:CalledViaLast": ReadOnly[Literal["aws-external-anthropic.amazonaws.com"]]}
+)
-class _SecureTransportCondition(TypedDict):
+class _SessionPolicyCondition(TypedDict):
Bool: ReadOnly[_SecureTransportBool]
+ StringEquals: NotRequired[ReadOnly[_CalledViaLastEquals]]
class _SessionPolicyStatement(TypedDict):
@@ -111,7 +116,7 @@ class _SessionPolicyStatement(TypedDict):
Effect: ReadOnly[Literal["Allow"]]
Action: ReadOnly[tuple[str, ...]]
Resource: ReadOnly[Literal["*"]]
- Condition: ReadOnly[_SecureTransportCondition]
+ Condition: ReadOnly[_SessionPolicyCondition]
class WebIdentitySessionPolicy(TypedDict):
@@ -147,11 +152,24 @@ _WEB_IDENTITY_SESSION_POLICY_ACTIONS: Final[Mapping[str, tuple[str, ...]]] = Map
"aws-external-anthropic:Get*",
"aws-external-anthropic:List*",
),
+ "ClaudePlatformOutboundIdentityLiteLLM": (
+ "sts:GetWebIdentityToken",
+ "sts:TagGetWebIdentityToken",
+ ),
"BedrockMantleLiteLLM": ("bedrock-mantle:CreateInference",),
}
)
-_SECURE_TRANSPORT_ONLY: Final = _SecureTransportCondition(Bool=_SecureTransportBool({"aws:SecureTransport": "true"}))
+_SECURE_TRANSPORT_ONLY: Final = _SessionPolicyCondition(Bool=_SecureTransportBool({"aws:SecureTransport": "true"}))
+
+_WEB_IDENTITY_SESSION_POLICY_CONDITIONS: Final[Mapping[str, _SessionPolicyCondition]] = MappingProxyType(
+ {
+ "ClaudePlatformOutboundIdentityLiteLLM": _SessionPolicyCondition(
+ Bool=_SecureTransportBool({"aws:SecureTransport": "true"}),
+ StringEquals=_CalledViaLastEquals({"aws:CalledViaLast": "aws-external-anthropic.amazonaws.com"}),
+ ),
+ }
+)
def build_web_identity_session_policy() -> WebIdentitySessionPolicy:
@@ -163,7 +181,7 @@ def build_web_identity_session_policy() -> WebIdentitySessionPolicy:
Effect="Allow",
Action=actions,
Resource="*",
- Condition=_SECURE_TRANSPORT_ONLY,
+ Condition=_WEB_IDENTITY_SESSION_POLICY_CONDITIONS.get(sid, _SECURE_TRANSPORT_ONLY),
)
for sid, actions in _WEB_IDENTITY_SESSION_POLICY_ACTIONS.items()
),
diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py
index a207bceb418..6ad400661be 100644
--- a/litellm/llms/bedrock/chat/invoke_handler.py
+++ b/litellm/llms/bedrock/chat/invoke_handler.py
@@ -13,6 +13,7 @@ from litellm.caching.caching import InMemoryCache
from litellm.constants import RESPONSE_FORMAT_TOOL_NAME
from litellm.litellm_core_utils.core_helpers import map_finish_reason
from litellm.litellm_core_utils.litellm_logging import Logging
+from litellm.litellm_core_utils.llm_response_utils.get_headers import provider_request_id_headers
from litellm.litellm_core_utils.optional_imports import ensure_optional_import
from litellm.litellm_core_utils.prompt_templates.common_utils import native_redacted_thinking_data
from litellm.llms.anthropic.chat.handler import (
@@ -448,6 +449,7 @@ class _EventStreamTally:
f"x-amzn-requestid={_response_header(response_headers, 'x-amzn-requestid')!r}, "
f"{self.bytes_received} bytes received"
)
+ request_id_headers: Final = provider_request_id_headers(response_headers)
if not self.events:
return BedrockError(
status_code=502,
@@ -455,6 +457,7 @@ class _EventStreamTally:
"Bedrock answered the stream with HTTP 200 but its body decoded to no events "
f"({detail}, first bytes={self.head!r})"
),
+ headers=request_id_headers,
)
if not self.recognized_events:
return BedrockError(
@@ -464,10 +467,12 @@ class _EventStreamTally:
f"event type (event types={sorted(self.unrecognized_event_types)}, {detail}, "
f"first payload={self.unrecognized_head!r})"
),
+ headers=request_id_headers,
)
return BedrockError(
status_code=502,
message=f"Bedrock stream ended with {undecoded} undecoded bytes after {self.events} events ({detail})",
+ headers=request_id_headers,
)
@@ -837,11 +842,12 @@ class AWSEventStreamDecoder:
event_stream_buffer: Final = EventStreamBuffer()
tally: Final = _EventStreamTally(get_bedrock_stream_event_statuses())
+ request_id_headers: Final = provider_request_id_headers(response_headers)
for chunk in iterator:
event_stream_buffer.add_data(chunk)
tally.add_chunk(chunk)
for event in event_stream_buffer:
- message = self._decode_event(event, tally)
+ message = self._decode_event(event, tally, request_id_headers)
if message:
# sse_event = ServerSentEvent(data=message, event="completion")
_data = json.loads(message)
@@ -858,11 +864,12 @@ class AWSEventStreamDecoder:
event_stream_buffer: Final = EventStreamBuffer()
tally: Final = _EventStreamTally(get_bedrock_stream_event_statuses())
+ request_id_headers: Final = provider_request_id_headers(response_headers)
async for chunk in iterator:
event_stream_buffer.add_data(chunk)
tally.add_chunk(chunk)
for event in event_stream_buffer:
- message = self._decode_event(event, tally)
+ message = self._decode_event(event, tally, request_id_headers)
if message:
_data = json.loads(message)
yield self.chunk_parser(chunk_data=_data)
@@ -882,27 +889,34 @@ class AWSEventStreamDecoder:
)
return response_stream_shape
- def _decode_event(self, event: "EventStreamMessage", tally: _EventStreamTally) -> str | None:
+ def _decode_event(
+ self, event: "EventStreamMessage", tally: _EventStreamTally, request_id_headers: httpx.Headers | None
+ ) -> str | None:
response_stream_shape: Final = self._response_stream_shape()
response_dict: Final = bedrock_event_stream_response(event)
tally.add_event(event, response_dict["headers"])
- return self._parse_message_from_response(response_dict, response_stream_shape)
+ return self._parse_message_from_response(response_dict, response_stream_shape, request_id_headers)
def _parse_message_from_event(self, event: "EventStreamMessage") -> str | None:
response_stream_shape: Final = self._response_stream_shape()
return self._parse_message_from_response(bedrock_event_stream_response(event), response_stream_shape)
def _parse_message_from_response(
- self, response_dict: BedrockEventStreamResponseDict, response_stream_shape: "Shape"
+ self,
+ response_dict: BedrockEventStreamResponseDict,
+ response_stream_shape: "Shape",
+ request_id_headers: httpx.Headers | None = None,
) -> str | None:
parsed_response: Final = self.parser.parse(response_dict, response_stream_shape)
if response_dict["status_code"] != 200:
- raise build_bedrock_stream_error(response_dict, response_stream_shape)
+ raise build_bedrock_stream_error(response_dict, response_stream_shape, headers=request_id_headers)
event_type: Final = bedrock_event_stream_header(response_dict["headers"], ":event-type")
event_error_status: Final = bedrock_stream_event_error_status(event_type)
if event_type is not None and event_error_status is not None:
- raise build_bedrock_stream_event_error(event_type, event_error_status, response_dict["body"])
+ raise build_bedrock_stream_event_error(
+ event_type, event_error_status, response_dict["body"], headers=request_id_headers
+ )
if "chunk" in parsed_response:
chunk = parsed_response.get("chunk")
if not chunk:
diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py
index 1b3b0c344df..f026e4f9a99 100644
--- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py
+++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py
@@ -14,9 +14,9 @@ from litellm.litellm_core_utils.get_litellm_params import stored_control_options
from litellm.litellm_core_utils.logging_utils import track_llm_api_timing
from litellm.litellm_core_utils.prompt_templates.factory import (
cohere_message_pt,
- custom_prompt,
deepseek_r1_pt,
prompt_factory,
+ render_custom_prompt_template,
)
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
from litellm.llms.bedrock.chat.invoke_handler import make_call, make_sync_call
@@ -601,7 +601,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
return cast(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL, provider)
return None
- def convert_messages_to_prompt(self, model, messages, provider, custom_prompt_dict) -> tuple[str, list | None]:
+ def convert_messages_to_prompt(self, model: str, messages, provider, custom_prompt_dict) -> tuple[str, list | None]:
# handle anthropic prompts and amazon titan prompts
prompt = ""
chat_history: list | None = None
@@ -609,11 +609,8 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
if model in custom_prompt_dict:
# check if the model has a registered custom prompt
model_prompt_details: Final = custom_prompt_dict[model]
- prompt = custom_prompt(
- role_dict=model_prompt_details["roles"],
- initial_prompt_value=model_prompt_details.get("initial_prompt_value", ""),
- final_prompt_value=model_prompt_details.get("final_prompt_value", ""),
- messages=messages,
+ prompt = render_custom_prompt_template(
+ template=model_prompt_details, messages=messages, model=model, llm_provider="bedrock"
)
return prompt, None
## ELSE
diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py
index e79865ef6bf..5862b05a619 100644
--- a/litellm/llms/bedrock/common_utils.py
+++ b/litellm/llms/bedrock/common_utils.py
@@ -1913,13 +1913,18 @@ def bedrock_event_stream_header(headers: Mapping[str, object], name: str) -> str
return value if isinstance(value, str) else None
-def build_bedrock_stream_event_error(event_type: str, status_code: int, body: bytes) -> BedrockError:
- return BedrockError(status_code=status_code, message=f"{event_type} {body.decode(errors='replace')}")
+def build_bedrock_stream_event_error(
+ event_type: str, status_code: int, body: bytes, headers: httpx.Headers | None = None
+) -> BedrockError:
+ return BedrockError(
+ status_code=status_code, message=f"{event_type} {body.decode(errors='replace')}", headers=headers
+ )
def build_bedrock_stream_error(
response_dict: BedrockEventStreamResponseDict,
response_stream_shape: Shape | None,
+ headers: httpx.Headers | None = None,
) -> BedrockError:
"""Build a BedrockError for a non-200 event-stream error event.
@@ -1929,12 +1934,14 @@ def build_bedrock_stream_error(
"""
exception_type: Final = bedrock_event_stream_header(response_dict["headers"], ":exception-type")
if exception_type is None:
- return BedrockError(status_code=response_dict["status_code"], message=response_dict["body"].decode())
+ return BedrockError(
+ status_code=response_dict["status_code"], message=response_dict["body"].decode(), headers=headers
+ )
member: Final = _structure_members(response_stream_shape).get(exception_type)
modeled_status: Final = None if member is None else _modeled_error_status(member)
status_code: Final = response_dict["status_code"] if modeled_status is None else modeled_status
- return build_bedrock_stream_event_error(exception_type, status_code, response_dict["body"])
+ return build_bedrock_stream_event_error(exception_type, status_code, response_dict["body"], headers=headers)
class BedrockEventStreamDecoderBase:
diff --git a/litellm/llms/codestral/completion/handler.py b/litellm/llms/codestral/completion/handler.py
index f8486d3b274..08799c3a0a5 100644
--- a/litellm/llms/codestral/completion/handler.py
+++ b/litellm/llms/codestral/completion/handler.py
@@ -13,8 +13,8 @@ import litellm
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.litellm_core_utils.logging_utils import track_llm_api_timing
from litellm.litellm_core_utils.prompt_templates.factory import (
- custom_prompt,
prompt_factory,
+ render_custom_prompt_template,
)
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
@@ -261,11 +261,8 @@ class CodestralTextCompletion:
if model in custom_prompt_dict:
# check if the model has a registered custom prompt
model_prompt_details: Final = custom_prompt_dict[model]
- prompt = custom_prompt(
- role_dict=model_prompt_details["roles"],
- initial_prompt_value=model_prompt_details["initial_prompt_value"],
- final_prompt_value=model_prompt_details["final_prompt_value"],
- messages=messages,
+ prompt = render_custom_prompt_template(
+ template=model_prompt_details, messages=messages, model=model, llm_provider="codestral"
)
else:
prompt = prompt_factory(model=model, messages=messages)
diff --git a/litellm/llms/dashscope/videos/transformation.py b/litellm/llms/dashscope/videos/transformation.py
index 17b685a0e1a..1db91cbb127 100644
--- a/litellm/llms/dashscope/videos/transformation.py
+++ b/litellm/llms/dashscope/videos/transformation.py
@@ -3,13 +3,10 @@ DashScope (Alibaba Model Studio) async video task API: create, poll and download
delete or remix endpoint, so those raise.
"""
-import base64
import json
from collections.abc import Mapping
from datetime import datetime
-from io import BufferedReader, BytesIO
from math import log
-from pathlib import Path, PurePath
from types import MappingProxyType
from typing import TYPE_CHECKING, Annotated, Final
@@ -19,9 +16,13 @@ from pydantic import BaseModel, Field, StrictBool, StrictInt, StrictStr, TypeAda
import litellm
from litellm.exceptions import UnsupportedParamsError
-from litellm.images.utils import ImageEditRequestUtils
from litellm.litellm_core_utils.url_utils import encode_url_path_segment
from litellm.llms.base_llm.chat.transformation import BaseLLMException
+from litellm.llms.base_llm.videos.async_task_utils import (
+ input_reference_url,
+ polled_video_models,
+ routed_video_models,
+)
from litellm.llms.base_llm.videos.transformation import BaseVideoConfig
from litellm.llms.custom_httpx.http_handler import (
_get_httpx_client, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # house cached-client factory has no public alias and its stub leaves params untyped
@@ -32,7 +33,6 @@ from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import ModelInfo
from litellm.types.videos.main import VideoCreateOptionalRequestParams, VideoObject
from litellm.types.videos.utils import (
- decode_video_id_with_provider,
encode_video_id_with_provider,
extract_original_video_id,
)
@@ -86,7 +86,6 @@ _EMPTY_PARAMS: Final[dict[str, object]] = {} # mutable-ok: BaseVideoConfig cont
_TASK_RESPONSE_ADAPTER: Final = TypeAdapter(_TaskResponse)
_PARAMETERS_ADAPTER: Final = TypeAdapter(dict[str, object])
-_FILE_TUPLE_ADAPTER: Final = TypeAdapter(tuple[object, ...])
_STATUS_MAP: Final = MappingProxyType(
{
@@ -383,39 +382,6 @@ def _is_priced_tier(resolution: object, priced_tiers: tuple[str, ...]) -> bool:
return isinstance(resolution, str) and resolution.strip().upper() in priced_tiers
-def _read_all_bytes(file_obj: object) -> bytes:
- if isinstance(file_obj, tuple):
- file_tuple: Final = _FILE_TUPLE_ADAPTER.validate_python(file_obj)
- return _read_all_bytes(file_tuple[1] if len(file_tuple) >= 2 else None)
- if isinstance(file_obj, PurePath):
- return Path(file_obj).read_bytes()
- if isinstance(file_obj, (BytesIO, BufferedReader)):
- current_position: Final = file_obj.tell()
- file_obj.seek(0)
- content: Final = file_obj.read()
- file_obj.seek(current_position)
- return content
- if isinstance(file_obj, bytes):
- return file_obj
- if isinstance(file_obj, bytearray):
- return bytes(file_obj)
- read: Final = getattr(file_obj, "read", None)
- if callable(read):
- data: Final = read()
- if isinstance(data, bytes):
- return data
- raise ValueError("input_reference must be a URL string, bytes, a path, a file object, or a file tuple")
-
-
-def _image_url(image: object) -> str:
- if isinstance(image, str):
- return image
- content: Final = _read_all_bytes(image)
- content_type: Final = ImageEditRequestUtils.get_image_content_type(content)
- encoded: Final = base64.b64encode(content).decode("utf-8")
- return f"data:{content_type};base64,{encoded}"
-
-
def _media_reference_type(model: str) -> str | None:
"""Models without the field predate the media array and take the image on the flat ``img_url`` field."""
info: Final = _dashscope_model_info(model)
@@ -558,58 +524,6 @@ def _video_url_from_task(task: _TaskResponse, raw_response: httpx.Response) -> s
raise _video_error(raw_response, 502, "Video URL not found in task response. The task may not have succeeded yet.")
-def _logged_litellm_params(logging_obj: object) -> Mapping[str, object]:
- litellm_params: Final[object] = getattr(logging_obj, "litellm_params", None)
- try:
- return _PARAMETERS_ADAPTER.validate_python(litellm_params)
- except ValidationError:
- return _EMPTY_PARAMS
-
-
-def _polled_models(logging_obj: object) -> tuple[str | None, str | None]:
- """
- The proxy routes by the models encoded in the polled id, so they must survive into the returned id. The served
- group survives only when the router sent this poll to it, so an edited id cannot write any name it likes back.
- """
- video_id: Final = _logged_litellm_params(logging_obj).get("video_id")
- if not isinstance(video_id, str):
- return None, None
- decoded: Final = decode_video_id_with_provider(video_id)
- served: Final = decoded.get("served_model")
- routed: Final = _poll_models(logging_obj)
- return decoded.get("model_id") or None, served if routed is None or served in routed else None
-
-
-def _poll_models(logging_obj: object) -> tuple[str, str | None] | None:
- """
- A later status or content request checks the key's model access against the model encoded in the id, and
- routes by it. DashScope serves the same model name from several regional hosts, so only the group that
- served the task can send its polls back to the region that created it, while a key may be scoped to the
- alias, routing group or wildcard name it asked for instead. The id carries both names: the one the router
- was asked for on this hop, and the group that served it when that differs and is not a wildcard. The router
- stamps both into litellm_metadata, and metadata carries a copy only when the caller sent none.
- """
- params: Final = _logged_litellm_params(logging_obj)
- stamped: Final = (_stamped_poll_models(params.get(key)) for key in ("litellm_metadata", "metadata"))
- return next((models for models in stamped if models), None)
-
-
-def _stamped_poll_models(metadata: object) -> tuple[str, str | None] | None:
- try:
- fields: Final = _PARAMETERS_ADAPTER.validate_python(metadata)
- except ValidationError:
- return None
- requested: Final = _id_safe_name(fields.get("model_group"))
- served: Final = _id_safe_name(fields.get("deployment_model_name"))
- if requested is None:
- return None
- return requested, None if served is None or served == requested or "*" in served else served
-
-
-def _id_safe_name(name: object) -> str | None:
- return name if isinstance(name, str) and name and ";" not in name else None
-
-
class DashScopeVideoConfig(BaseVideoConfig):
def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: BaseVideoConfig contract returns list
return [
@@ -804,7 +718,7 @@ class DashScopeVideoConfig(BaseVideoConfig):
input_reference: Final = video_create_optional_request_params.get("input_reference")
if input_reference is None:
return None, None
- image_url: Final = _image_url(input_reference)
+ image_url: Final = input_reference_url(input_reference)
media_type: Final = _media_reference_type(model)
if media_type is None:
return "img_url", image_url
@@ -823,7 +737,7 @@ class DashScopeVideoConfig(BaseVideoConfig):
requested: Final = self._requested_parameters(request_data)
video_obj: Final = _video_object_from_task(task, raw_response, model=model, requested=requested)
if custom_llm_provider and video_obj.id:
- requested_model, served_model = _poll_models(logging_obj) or (model, None)
+ requested_model, served_model = routed_video_models(logging_obj) or (model, None)
video_obj.id = encode_video_id_with_provider(
video_obj.id, custom_llm_provider, requested_model, served_model
)
@@ -869,7 +783,7 @@ class DashScopeVideoConfig(BaseVideoConfig):
client: "HTTPHandler | None" = None,
) -> VideoObject:
task: Final = _parse_task_response(raw_response)
- model, served_model = _polled_models(logging_obj)
+ model, served_model = polled_video_models(logging_obj)
video_obj: Final = _video_object_from_task(task, raw_response, model=model)
if custom_llm_provider and video_obj.id:
video_obj.id = encode_video_id_with_provider(video_obj.id, custom_llm_provider, model, served_model)
diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py
index 2145c19a71a..74bdd2cd733 100644
--- a/litellm/llms/databricks/chat/transformation.py
+++ b/litellm/llms/databricks/chat/transformation.py
@@ -20,6 +20,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
merge_consecutive_system_messages,
strip_litellm_internal_message_fields,
strip_name_from_message,
+ strip_prefill_prefix_field,
)
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
from litellm.llms.base_llm.base_utils import (
@@ -487,6 +488,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
Databricks does not support:
- 'name' in user message.
- litellm's internal `thinking_blocks` / `reasoning_content` on assistant messages.
+ - the `prefix` prefill flag Mistral and DeepSeek read, which Claude does not need to continue an assistant turn.
- non-Claude json_object "json" checks read only string content, so a single text block is sent as a string.
"""
new_messages = []
@@ -496,7 +498,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
else:
_message = message
_message = strip_name_from_message(_message, allowed_name_roles=["user"])
- _message = strip_litellm_internal_message_fields(_message)
+ _message = strip_prefill_prefix_field(strip_litellm_internal_message_fields(_message))
# Move message-level cache_control into a content block when content is a string.
if "cache_control" in _message and isinstance(_message.get("content"), str):
_message = self._move_cache_control_into_string_content_block(_message)
diff --git a/litellm/llms/gemini/realtime/transformation.py b/litellm/llms/gemini/realtime/transformation.py
index 7c04ebf311f..f4e2941670a 100644
--- a/litellm/llms/gemini/realtime/transformation.py
+++ b/litellm/llms/gemini/realtime/transformation.py
@@ -161,6 +161,10 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
"""Google AI Studio Gemini 3.5+ accepts ``id`` on functionResponses; Vertex AI rejects it."""
return True
+ def _rejects_vad_tuning_with_detection_off(self) -> bool:
+ """Google AI Studio closes the socket (1007) on padding or silence tuning next to ``disabled: true``."""
+ return True
+
@staticmethod
def _usage_detail_alias(details: Mapping[str, int | None] | None, defaults: dict[str, int]) -> dict[str, int]:
if not isinstance(details, dict):
@@ -279,10 +283,15 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
OpenAI ``semantic_vad`` has no Gemini Live equivalent — return an empty
dict so callers omit ``realtimeInputConfig`` (mapping it with
- ``disabled: true`` breaks native-audio sessions).
+ ``disabled: true`` breaks native-audio sessions). An explicit
+ ``create_response: false`` maps to detection off alone where the backend
+ rejects tuning next to ``disabled: true``.
"""
if isinstance(value, dict) and value.get("type") == "semantic_vad" and "create_response" not in value:
return AutomaticActivityDetection()
+ auto_response_off: Final = "create_response" in value and value["create_response"] is False
+ if auto_response_off and self._rejects_vad_tuning_with_detection_off():
+ return AutomaticActivityDetection(disabled=True)
automatic_activity_dection: Final = AutomaticActivityDetection()
if "create_response" in value and isinstance(value["create_response"], bool):
diff --git a/litellm/llms/meta/realtime/transformation.py b/litellm/llms/meta/realtime/transformation.py
index 46231e5e980..892d7ff66de 100644
--- a/litellm/llms/meta/realtime/transformation.py
+++ b/litellm/llms/meta/realtime/transformation.py
@@ -26,6 +26,7 @@ from litellm.llms.base_llm.realtime.transcription_protocol import (
speech_event,
transcription_session,
transcription_session_created_event,
+ uses_server_vad,
)
from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig
from litellm.secret_managers.main import get_secret_str
@@ -230,15 +231,7 @@ def _parse_sample_rate(update: TranscriptionSessionUpdate) -> MuseSampleRate:
return 16_000 if rate == 16_000 else 24_000
-def _parse_mode(update: TranscriptionSessionUpdate) -> MuseMode:
- if update.turn_detection_disabled:
- return "PUSH_TO_TALK"
- if update.turn_detection_type not in (None, "server_vad"):
- raise MuseProtocolError("Muse Voice supports server_vad turn detection or null")
- return "ENDPOINTING"
-
-
-def parse_session_update(payload: str, expected_model: str) -> MuseSessionConfig:
+def parse_session_update(payload: str, expected_model: str, *, drop_params: bool = False) -> MuseSessionConfig:
update: Final = parse_transcription_session_update(payload, MuseProtocolError)
if update.session_type not in (None, "transcription", "realtime"):
raise MuseProtocolError("Muse Voice supports transcription sessions only")
@@ -253,7 +246,7 @@ def parse_session_update(payload: str, expected_model: str) -> MuseSessionConfig
raise MuseProtocolError("realtime session model cannot be changed")
return MuseSessionConfig(
model=normalized_model,
- mode=_parse_mode(update),
+ mode="ENDPOINTING" if uses_server_vad(update, "Muse Voice", drop_params, MuseProtocolError) else "PUSH_TO_TALK",
sample_rate=_parse_sample_rate(update),
language_bias=() if update.language is None else (normalize_language(update.language),),
)
@@ -441,9 +434,11 @@ class MetaRealtimeConfig(BaseRealtimeConfig):
*,
monotonic: Callable[[], float] = time.monotonic,
sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
+ drop_params: bool = False,
) -> None:
self._monotonic: Final = monotonic
self._sleep: Final = sleep
+ self._drop_params: Final = drop_params
self._transformer: Final = MuseEventTransformer()
self._access_token: str | None = None
self._config: MuseSessionConfig | None = None
@@ -553,7 +548,7 @@ class MetaRealtimeConfig(BaseRealtimeConfig):
access_token: Final = self._access_token
if access_token is None:
raise MuseProtocolError("Meta API key was not validated before the session was configured")
- config: Final = parse_session_update(message, model)
+ config: Final = parse_session_update(message, model, drop_params=self._drop_params)
self._config = config
self._transformer.configure(config)
return (json.dumps(config.handshake(access_token), separators=(",", ":")),)
diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py
index 703d45acdf2..ac41310fb61 100644
--- a/litellm/llms/ollama/completion/transformation.py
+++ b/litellm/llms/ollama/completion/transformation.py
@@ -16,8 +16,8 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
)
from litellm.litellm_core_utils.prompt_templates.factory import (
convert_to_ollama_image,
- custom_prompt,
ollama_pt,
+ render_custom_prompt_template,
)
from litellm.litellm_core_utils.prompt_templates.image_handling import (
async_inline_remote_media,
@@ -437,11 +437,8 @@ class OllamaConfig(BaseConfig):
if model in custom_prompt_dict:
# check if the model has a registered custom prompt
model_prompt_details: Final = custom_prompt_dict[model]
- ollama_prompt = custom_prompt(
- role_dict=model_prompt_details["roles"],
- initial_prompt_value=model_prompt_details["initial_prompt_value"],
- final_prompt_value=model_prompt_details["final_prompt_value"],
- messages=messages,
+ ollama_prompt = render_custom_prompt_template(
+ template=model_prompt_details, messages=messages, model=model, llm_provider="ollama"
)
elif text_completion_request: # handle `/completions` requests
ollama_prompt = get_str_from_messages(messages=messages)
diff --git a/litellm/llms/openai/responses/count_tokens/token_counter.py b/litellm/llms/openai/responses/count_tokens/token_counter.py
index 43b3ae9f6f1..7c0697626ca 100644
--- a/litellm/llms/openai/responses/count_tokens/token_counter.py
+++ b/litellm/llms/openai/responses/count_tokens/token_counter.py
@@ -2,20 +2,21 @@
OpenAI Token Counter implementation using the Responses API /input_tokens endpoint.
"""
-import os
from collections.abc import Mapping, Sequence
from typing import Any, Final
+import openai
from pydantic import JsonValue
from typing_extensions import assert_never
from litellm._logging import verbose_logger
+from litellm.litellm_core_utils.asyncify import asyncify
from litellm.llms.anthropic.pass_through.convertible_blocks import unconvertible_blocks_error
from litellm.llms.anthropic.pass_through.responses_adapters.transformation import (
LiteLLMAnthropicToResponsesAPIAdapter,
)
from litellm.llms.base_llm.base_utils import BaseTokenCounter
-from litellm.llms.openai.common_utils import OpenAIError
+from litellm.llms.openai.common_utils import OpenAIError, get_openai_credentials
from litellm.llms.openai.responses.count_tokens.handler import (
OpenAICountTokensHandler,
)
@@ -23,12 +24,26 @@ from litellm.llms.openai.responses.count_tokens.transformation import (
OpenAICountTokensConfig,
ResponsesCountRequestParts,
)
+from litellm.llms.openai.workload_identity import (
+ OpenAIWorkloadIdentityConfig,
+ get_workload_identity_bearer_token,
+ resolve_openai_workload_identity_config,
+)
from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse
# Global handler instance - reuse across all token counting requests
openai_count_tokens_handler: Final = OpenAICountTokensHandler()
+def _workload_identity_config(litellm_params: Mapping[str, object]) -> OpenAIWorkloadIdentityConfig | None:
+ api_base: Final = litellm_params.get("api_base")
+ return resolve_openai_workload_identity_config(
+ api_key=None,
+ api_base=api_base if isinstance(api_base, str) else None,
+ litellm_params=litellm_params,
+ )
+
+
class OpenAITokenCounter(BaseTokenCounter):
"""Token counter implementation for OpenAI provider using the Responses API."""
@@ -107,16 +122,27 @@ class OpenAITokenCounter(BaseTokenCounter):
deployment = deployment or {}
litellm_params: Final = deployment.get("litellm_params", {})
- # Get OpenAI API key from deployment config or environment
- api_key = litellm_params.get("api_key")
- if not api_key:
- api_key = os.getenv("OPENAI_API_KEY")
+ credentials: Final = get_openai_credentials(
+ api_base=litellm_params.get("api_base"), api_key=litellm_params.get("api_key")
+ )
+ credential: Final = credentials.api_key or _workload_identity_config(litellm_params)
- if not api_key:
- verbose_logger.warning("No OpenAI API key found for token counting")
- return None
+ if credential is None:
+ return TokenCountResponse(
+ total_tokens=0,
+ request_model=request_model,
+ model_used=model_to_use,
+ tokenizer_type="openai_api",
+ error=True,
+ error_message=(
+ f"No OpenAI API key is set for model {request_model or model_to_use}, so tokens can't be "
+ "counted. Set api_key or workload identity on the deployment, or OPENAI_API_KEY in the proxy "
+ "environment."
+ ),
+ status_code=401,
+ )
- api_base: Final = litellm_params.get("api_base")
+ api_base: Final = credentials.api_base
input_items, instructions, counted_tools, text = self._responses_request_parts(
model_to_use=model_to_use,
@@ -132,6 +158,11 @@ class OpenAITokenCounter(BaseTokenCounter):
return None
try:
+ api_key: Final = (
+ credential
+ if isinstance(credential, str)
+ else await asyncify(get_workload_identity_bearer_token)(credential)
+ )
result: Final = await openai_count_tokens_handler.handle_count_tokens_request(
model=model_to_use,
input=input_items,
@@ -163,6 +194,17 @@ class OpenAITokenCounter(BaseTokenCounter):
error_message=e.message,
status_code=e.status_code,
)
+ except openai.AuthenticationError as e:
+ verbose_logger.warning("OpenAI workload identity token exchange failed: %s", e.message)
+ return TokenCountResponse(
+ total_tokens=0,
+ request_model=request_model,
+ model_used=model_to_use,
+ tokenizer_type="openai_api",
+ error=True,
+ error_message=e.message,
+ status_code=401,
+ )
except Exception as e:
verbose_logger.warning("Error calling OpenAI CountTokens API: %s", e)
return TokenCountResponse(
diff --git a/litellm/llms/petals/completion/handler.py b/litellm/llms/petals/completion/handler.py
index 425d3729d7c..a2ccd7eaf08 100644
--- a/litellm/llms/petals/completion/handler.py
+++ b/litellm/llms/petals/completion/handler.py
@@ -5,8 +5,8 @@ from typing import Final
import litellm
from litellm.litellm_core_utils.prompt_templates.factory import (
- custom_prompt,
prompt_factory,
+ render_custom_prompt_template,
)
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
@@ -43,11 +43,8 @@ def completion(
if model in litellm.custom_prompt_dict:
# check if the model has a registered custom prompt
model_prompt_details: Final = litellm.custom_prompt_dict[model]
- prompt = custom_prompt(
- role_dict=model_prompt_details["roles"],
- initial_prompt_value=model_prompt_details["initial_prompt_value"],
- final_prompt_value=model_prompt_details["final_prompt_value"],
- messages=messages,
+ prompt = render_custom_prompt_template(
+ template=model_prompt_details, messages=messages, model=model, llm_provider="petals"
)
else:
prompt = prompt_factory(model=model, messages=messages)
diff --git a/litellm/llms/predibase/chat/transformation.py b/litellm/llms/predibase/chat/transformation.py
index 7bd05c6e30b..0dd7e27722e 100644
--- a/litellm/llms/predibase/chat/transformation.py
+++ b/litellm/llms/predibase/chat/transformation.py
@@ -10,8 +10,8 @@ import litellm
from litellm.constants import DEFAULT_MAX_TOKENS
from litellm.litellm_core_utils.core_helpers import map_finish_reason
from litellm.litellm_core_utils.prompt_templates.factory import (
- custom_prompt,
prompt_factory,
+ render_custom_prompt_template,
)
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
from litellm.types.llms.openai import AllMessageValues
@@ -277,11 +277,8 @@ class PredibaseConfig(BaseConfig):
custom_prompt_dict: Final = litellm_params.get("custom_prompt_dict", {})
if model in custom_prompt_dict:
model_prompt_details: Final = custom_prompt_dict[model]
- prompt = custom_prompt(
- role_dict=model_prompt_details["roles"],
- initial_prompt_value=model_prompt_details["initial_prompt_value"],
- final_prompt_value=model_prompt_details["final_prompt_value"],
- messages=messages,
+ prompt = render_custom_prompt_template(
+ template=model_prompt_details, messages=messages, model=model, llm_provider="predibase"
)
else:
prompt = prompt_factory(model=model, messages=messages)
diff --git a/litellm/llms/sagemaker/completion/transformation.py b/litellm/llms/sagemaker/completion/transformation.py
index 274f884de0b..a69846c816f 100644
--- a/litellm/llms/sagemaker/completion/transformation.py
+++ b/litellm/llms/sagemaker/completion/transformation.py
@@ -15,8 +15,8 @@ from httpx._models import Headers, Response
import litellm
from litellm.litellm_core_utils.asyncify import asyncify
from litellm.litellm_core_utils.prompt_templates.factory import (
- custom_prompt,
prompt_factory,
+ render_custom_prompt_template,
)
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
from litellm.types.llms.openai import AllMessageValues
@@ -122,20 +122,14 @@ class SagemakerConfig(BaseConfig):
if model in custom_prompt_dict:
# check if the model has a registered custom prompt
model_prompt_details = custom_prompt_dict[model]
- prompt = custom_prompt(
- role_dict=model_prompt_details.get("roles", None),
- initial_prompt_value=model_prompt_details.get("initial_prompt_value", ""),
- final_prompt_value=model_prompt_details.get("final_prompt_value", ""),
- messages=messages,
+ prompt = render_custom_prompt_template(
+ template=model_prompt_details, messages=messages, model=model, llm_provider="sagemaker"
)
elif hf_model_name in custom_prompt_dict:
# check if the base huggingface model has a registered custom prompt
model_prompt_details = custom_prompt_dict[hf_model_name]
- prompt = custom_prompt(
- role_dict=model_prompt_details.get("roles", None),
- initial_prompt_value=model_prompt_details.get("initial_prompt_value", ""),
- final_prompt_value=model_prompt_details.get("final_prompt_value", ""),
- messages=messages,
+ prompt = render_custom_prompt_template(
+ template=model_prompt_details, messages=messages, model=model, llm_provider="sagemaker"
)
else:
if hf_model_name is None:
diff --git a/litellm/llms/soniox/audio_transcription/handler.py b/litellm/llms/soniox/audio_transcription/handler.py
index b53955ae75d..f86bc142898 100644
--- a/litellm/llms/soniox/audio_transcription/handler.py
+++ b/litellm/llms/soniox/audio_transcription/handler.py
@@ -19,11 +19,14 @@ import asyncio
import math
import time
from collections.abc import Coroutine, Mapping, Sequence
+from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final
import httpx
+from pydantic import TypeAdapter
from typing_extensions import ReadOnly, TypedDict
+import litellm
from litellm.litellm_core_utils.audio_utils.utils import (
get_audio_file_name,
process_audio_file,
@@ -35,7 +38,12 @@ from litellm.llms.custom_httpx.http_handler import (
get_httpx_client,
)
from litellm.llms.soniox.audio_transcription.transformation import (
+ SONIOX_HANDLER_ONLY_PARAMS,
SonioxAudioTranscriptionConfig,
+ SonioxInvalidFormParam,
+ decode_soniox_form_params,
+ raise_soniox_form_error,
+ soniox_language_hints,
)
from litellm.llms.soniox.common_utils import (
SONIOX_DEFAULT_CLEANUP,
@@ -58,6 +66,39 @@ else:
LiteLLMLoggingObj = Any
+_CLEANUP_ENTRIES: Final = TypeAdapter(tuple[object, ...])
+_NOT_FORWARDED_TO_SONIOX: Final[frozenset[str]] = frozenset(
+ (*SONIOX_HANDLER_ONLY_PARAMS, "audio_url", "file_id", "response_format", "language")
+)
+
+
+def _optional_str(value: object) -> str | None:
+ return None if value is None else str(value)
+
+
+def _float_or_default(value: object, default: float) -> float:
+ try:
+ parsed: Final = float(str(value))
+ except ValueError:
+ return default
+ return parsed if math.isfinite(parsed) else default
+
+
+def _int_or_default(value: object, default: int) -> int:
+ try:
+ return int(value) if isinstance(value, (int, float)) else int(str(value))
+ except (ValueError, OverflowError):
+ return default
+
+
+def _cleanup_targets(raw: object) -> tuple[str, ...]:
+ if raw is None:
+ return ()
+ if isinstance(raw, str):
+ return (raw,)
+ return tuple(target for target in _CLEANUP_ENTRIES.validate_python(raw) if isinstance(target, str))
+
+
class _TranscriptionMeta(TypedDict, total=False):
"""Fields the handler reads from a Soniox transcription object."""
@@ -182,7 +223,7 @@ class SonioxAudioTranscriptionHandler:
) -> tuple[
dict[str, str], # auth headers
str, # api_base (no trailing slash)
- dict[str, object], # body for POST /v1/transcriptions (without file_id/audio_url)
+ Mapping[str, object], # body for POST /v1/transcriptions (without file_id/audio_url)
_HandlerOptions, # handler-only options (poll interval, cleanup, ...)
]:
# Validate env -> auth headers.
@@ -198,52 +239,42 @@ class SonioxAudioTranscriptionHandler:
base_url: Final = get_soniox_api_base(api_base)
- # Operate on a local copy so we don't mutate the caller's dict
- # (the caller may reuse `optional_params` for retries or logging).
- params: Final = dict(optional_params)
-
- # Pull handler-only kwargs out of params so they aren't sent
- # to Soniox.
- poll_interval = float(params.pop("soniox_polling_interval", SONIOX_DEFAULT_POLL_INTERVAL))
- try:
- max_attempts = int(params.pop("soniox_max_polling_attempts", SONIOX_DEFAULT_MAX_POLL_ATTEMPTS))
- except (ValueError, OverflowError):
- max_attempts = SONIOX_DEFAULT_MAX_POLL_ATTEMPTS
- cleanup_raw: Final = params.pop("soniox_cleanup", SONIOX_DEFAULT_CLEANUP)
- if cleanup_raw is None:
- cleanup: list[str] = []
- elif isinstance(cleanup_raw, str):
- cleanup = [cleanup_raw]
- else:
- cleanup = list(cleanup_raw)
- filename_override: Final = params.pop("filename", None)
+ decoded: Final = decode_soniox_form_params(
+ optional_params, drop_params=litellm_params.get("drop_params") is True or litellm.drop_params is True
+ )
+ if isinstance(decoded, SonioxInvalidFormParam):
+ raise_soniox_form_error(decoded)
# Server-side clamps. Caller-supplied poll settings (from request kwargs)
# are bounded so an authenticated caller cannot force a worker into a
# tight poll loop (zero interval) or pin it indefinitely (huge attempt
# count). Total polling time is bounded by
# SONIOX_MAX_POLL_ATTEMPTS * SONIOX_MAX_POLL_INTERVAL.
- if not math.isfinite(poll_interval):
- poll_interval = SONIOX_DEFAULT_POLL_INTERVAL
- clamped_poll_interval: Final = max(SONIOX_MIN_POLL_INTERVAL, min(poll_interval, SONIOX_MAX_POLL_INTERVAL))
- clamped_max_attempts: Final = max(1, min(max_attempts, SONIOX_MAX_POLL_ATTEMPTS))
+ poll_interval: Final = _float_or_default(decoded.get("soniox_polling_interval"), SONIOX_DEFAULT_POLL_INTERVAL)
+ max_attempts: Final = _int_or_default(
+ decoded.get("soniox_max_polling_attempts"), SONIOX_DEFAULT_MAX_POLL_ATTEMPTS
+ )
+ cleanup: Final = _cleanup_targets(decoded.get("soniox_cleanup", SONIOX_DEFAULT_CLEANUP))
+ language_hints: Final = soniox_language_hints(decoded.get("language"), decoded.get("language_hints"))
# response_format is handled by LiteLLM post-processing, not Soniox.
handler_opts: Final[_HandlerOptions] = {
- "poll_interval": clamped_poll_interval,
- "max_attempts": clamped_max_attempts,
+ "poll_interval": max(SONIOX_MIN_POLL_INTERVAL, min(poll_interval, SONIOX_MAX_POLL_INTERVAL)),
+ "max_attempts": max(1, min(max_attempts, SONIOX_MAX_POLL_ATTEMPTS)),
"cleanup": cleanup,
- "filename_override": filename_override,
- "audio_url": params.pop("audio_url", None),
- "file_id": params.pop("file_id", None),
- "response_format": params.pop("response_format", None),
+ "filename_override": _optional_str(decoded.get("filename")),
+ "audio_url": _optional_str(decoded.get("audio_url")),
+ "file_id": _optional_str(decoded.get("file_id")),
+ "response_format": _optional_str(decoded.get("response_format")),
}
- # Soniox does not accept `language` directly; map_openai_params should
- # already have translated it, but drop any leftover to be safe.
- params.pop("language", None)
-
- return auth_headers, base_url, params, handler_opts
+ provider_params: Final = MappingProxyType(
+ {
+ **{key: value for key, value in decoded.items() if key not in _NOT_FORWARDED_TO_SONIOX},
+ **({} if language_hints is None else {"language_hints": language_hints}),
+ }
+ )
+ return auth_headers, base_url, provider_params, handler_opts
def _build_create_body(
self,
diff --git a/litellm/llms/soniox/audio_transcription/transformation.py b/litellm/llms/soniox/audio_transcription/transformation.py
index aa1ff239320..30d872f248f 100644
--- a/litellm/llms/soniox/audio_transcription/transformation.py
+++ b/litellm/llms/soniox/audio_transcription/transformation.py
@@ -9,9 +9,14 @@ async API requires multiple HTTP calls and does not fit the single-request
contract of `base_llm_http_handler.audio_transcriptions`.
"""
-from typing import Any, Final
+import re
+from collections.abc import Mapping
+from dataclasses import dataclass
+from types import MappingProxyType
+from typing import Any, Final, NoReturn
from httpx import Headers, Response
+from pydantic import TypeAdapter, ValidationError
from litellm.llms.base_llm.audio_transcription.transformation import (
AudioTranscriptionRequestData,
@@ -26,6 +31,7 @@ from litellm.llms.soniox.common_utils import (
render_soniox_tokens_as_srt,
render_soniox_tokens_as_vtt,
)
+from litellm.llms.soniox.types import SonioxContext
from litellm.types.llms.openai import (
AllMessageValues,
OpenAIAudioTranscriptionOptionalParams,
@@ -58,6 +64,152 @@ SONIOX_HANDLER_ONLY_PARAMS: Final[list[str]] = [
]
+SONIOX_BOOL_PARAMS: Final[frozenset[str]] = frozenset(
+ {"enable_speaker_diarization", "enable_language_identification", "language_hints_strict"}
+)
+_JSON_OBJECT: Final = TypeAdapter(dict[str, object])
+_STRUCTURED_CONTEXT: Final = TypeAdapter(SonioxContext)
+_FORM_BOOL: Final = TypeAdapter(bool)
+_ITEMS: Final = TypeAdapter(tuple[object, ...])
+_BRACKETED_FORM_FIELD: Final = re.compile(
+ r"(?Pcontext|translation)\[(?P\w+)\](?:\[\]\[(?P\w+)\])?"
+)
+_BracketedField = tuple[str, str, str | None, object]
+
+
+@dataclass(frozen=True, slots=True)
+class SonioxInvalidFormParam:
+ key: str
+ value: str
+ expected: str
+
+
+def raise_soniox_form_error(error: SonioxInvalidFormParam) -> NoReturn:
+ raise SonioxException(
+ message=f"`{error.key}` must be {error.expected}, got {error.value!r}", status_code=400, headers=None
+ )
+
+
+def soniox_language_hints(language: object, hints: object) -> object:
+ """Soniox has no `language` field, so the OpenAI `language` becomes `language_hints` when none are set."""
+ if hints is not None or not isinstance(language, str) or not language:
+ return hints
+ return (language,)
+
+
+def _decode_bool(key: str, value: str) -> bool | SonioxInvalidFormParam:
+ try:
+ return _FORM_BOOL.validate_python(value.strip())
+ except ValidationError:
+ return SonioxInvalidFormParam(key=key, value=value, expected="a boolean")
+
+
+def _decode_polling_interval(value: str) -> str | SonioxInvalidFormParam:
+ try:
+ float(value)
+ except ValueError:
+ return SonioxInvalidFormParam(key="soniox_polling_interval", value=value, expected="a number")
+ return value
+
+
+def _decode_context(value: str) -> SonioxContext | str:
+ if not value.lstrip().startswith("{"):
+ return value
+ try:
+ return _STRUCTURED_CONTEXT.validate_json(value)
+ except ValidationError:
+ return value
+
+
+def _decode_translation(value: str) -> Mapping[str, object] | SonioxInvalidFormParam:
+ try:
+ return _JSON_OBJECT.validate_json(value)
+ except ValidationError:
+ return SonioxInvalidFormParam(key="translation", value=value, expected="a JSON object")
+
+
+def _decode_language_hints(value: str) -> tuple[object, ...] | SonioxInvalidFormParam:
+ if not value.lstrip().startswith("["):
+ return tuple(hint.strip() for hint in value.split(",") if hint.strip())
+ try:
+ return _ITEMS.validate_json(value)
+ except ValidationError:
+ return SonioxInvalidFormParam(key="language_hints", value=value, expected="a JSON list")
+
+
+def _decode_form_value(key: str, value: object) -> object | SonioxInvalidFormParam:
+ if not isinstance(value, str):
+ return value
+ if key in SONIOX_BOOL_PARAMS:
+ return _decode_bool(key, value)
+ if key == "soniox_polling_interval":
+ return _decode_polling_interval(value)
+ if key == "context":
+ return _decode_context(value)
+ if key == "translation":
+ return _decode_translation(value)
+ if key == "language_hints":
+ return _decode_language_hints(value)
+ return value
+
+
+def _nested_form_param(param: str, fields: tuple[_BracketedField, ...]) -> Mapping[str, object]:
+ owned: Final = tuple(field for field in fields if field[0] == param)
+ flat: Final = {name: value for _, name, entry_field, value in owned if entry_field is None}
+ listed: Final = {name for _, name, entry_field, _ in owned if entry_field is not None}
+ return {
+ **flat,
+ **{
+ name: (
+ {
+ entry_field: value
+ for _, owner, entry_field, value in owned
+ if owner == name and entry_field is not None
+ },
+ )
+ for name in listed
+ },
+ }
+
+
+def _is_soniox_context(value: object) -> bool:
+ try:
+ _STRUCTURED_CONTEXT.validate_python(value)
+ except ValidationError:
+ return False
+ return True
+
+
+def _nest_bracketed_form_fields(params: Mapping[str, object]) -> Mapping[str, object]:
+ """The OpenAI SDK sends a dict in `extra_body` as `context[terms]` and `context[general][][key]` form keys."""
+ matches: Final = {key: _BRACKETED_FORM_FIELD.fullmatch(key) for key in params}
+ fields: Final = tuple(
+ (match["param"], match["field"], match["entry_field"], params[key]) for key, match in matches.items() if match
+ )
+ nested: Final = {param: _nested_form_param(param, fields) for param in {field[0] for field in fields}}
+ kept: Final = {param: value for param, value in nested.items() if param != "context" or _is_soniox_context(value)}
+ unbracketed: Final = {
+ key: value for key, value in params.items() if (match := matches[key]) is None or match["param"] not in kept
+ }
+ return {**kept, **unbracketed}
+
+
+def decode_soniox_form_params(
+ optional_params: Mapping[str, object],
+ drop_params: bool = False,
+) -> Mapping[str, object] | SonioxInvalidFormParam:
+ """Multipart form fields reach the proxy as strings; restore the JSON types Soniox expects."""
+ decoded: Final = {
+ key: _decode_form_value(key, value) for key, value in _nest_bracketed_form_fields(optional_params).items()
+ }
+ invalid: Final = next((value for value in decoded.values() if isinstance(value, SonioxInvalidFormParam)), None)
+ if invalid is not None and not drop_params:
+ return invalid
+ return MappingProxyType(
+ {key: value for key, value in decoded.items() if not isinstance(value, SonioxInvalidFormParam)}
+ )
+
+
class SonioxAudioTranscriptionConfig(BaseAudioTranscriptionConfig):
"""Configuration for Soniox async speech-to-text transcription."""
@@ -74,25 +226,26 @@ class SonioxAudioTranscriptionConfig(BaseAudioTranscriptionConfig):
model: str,
drop_params: bool,
) -> dict:
- # Translate the OpenAI `language` param into Soniox `language_hints`.
- if "language" in non_default_params and non_default_params["language"]:
- language: Final = non_default_params["language"]
- existing_hints: Final = optional_params.get("language_hints")
- if not existing_hints:
- optional_params["language_hints"] = [language]
- elif language not in existing_hints:
- optional_params["language_hints"] = [language] + list(existing_hints)
-
- # Capture response_format for post-processing (not sent to Soniox API).
- if "response_format" in non_default_params:
- optional_params["response_format"] = non_default_params["response_format"]
-
- # Pass through Soniox-native kwargs unchanged.
- for key in SONIOX_PASSTHROUGH_PARAMS + SONIOX_HANDLER_ONLY_PARAMS:
- if key in non_default_params and non_default_params[key] is not None:
- optional_params[key] = non_default_params[key]
-
- return optional_params
+ language: Final[object] = non_default_params.get("language")
+ language_params: Final[Mapping[str, object]] = (
+ {
+ "language": language,
+ "language_hints": soniox_language_hints(language, optional_params.get("language_hints")),
+ }
+ if language
+ else {}
+ )
+ response_format_params: Final[Mapping[str, object]] = (
+ {"response_format": non_default_params["response_format"]}
+ if "response_format" in non_default_params
+ else {}
+ )
+ passthrough_params: Final[Mapping[str, object]] = {
+ key: non_default_params[key]
+ for key in SONIOX_PASSTHROUGH_PARAMS + SONIOX_HANDLER_ONLY_PARAMS
+ if non_default_params.get(key) is not None
+ }
+ return {**optional_params, **language_params, **response_format_params, **passthrough_params}
def get_error_class(self, error_message: str, status_code: int, headers: dict | Headers) -> BaseLLMException:
return SonioxException(message=error_message, status_code=status_code, headers=headers)
diff --git a/litellm/llms/soniox/types.py b/litellm/llms/soniox/types.py
new file mode 100644
index 00000000000..806d2a8fb0e
--- /dev/null
+++ b/litellm/llms/soniox/types.py
@@ -0,0 +1,24 @@
+from collections.abc import Sequence
+
+from pydantic import ConfigDict, with_config
+from typing_extensions import ReadOnly, TypedDict
+
+
+class SonioxContextGeneralEntry(TypedDict):
+ key: ReadOnly[str]
+ value: ReadOnly[str]
+
+
+class SonioxTranslationTerm(TypedDict):
+ source: ReadOnly[str]
+ target: ReadOnly[str]
+
+
+@with_config(ConfigDict(extra="forbid"))
+class SonioxContext(TypedDict, total=False):
+ """Soniox `context` request field: https://soniox.com/docs/stt/concepts/context"""
+
+ general: ReadOnly[Sequence[SonioxContextGeneralEntry]]
+ text: ReadOnly[str]
+ terms: ReadOnly[Sequence[str]]
+ translation_terms: ReadOnly[Sequence[SonioxTranslationTerm]]
diff --git a/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py b/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py
index 1d4a974fc54..5015678c467 100644
--- a/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py
+++ b/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py
@@ -13,7 +13,6 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.llms.base_llm.realtime.transcription_protocol import (
RealtimeTranscriptionProtocolError,
TranscriptionAudioFormat,
- TranscriptionSessionUpdate,
completed_event,
decode_pcm16_append,
delta_event,
@@ -23,6 +22,7 @@ from litellm.llms.base_llm.realtime.transcription_protocol import (
speech_event,
transcription_session,
transcription_session_created_event,
+ uses_server_vad,
)
from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig, RealtimeBackend
from litellm.llms.vertex_ai.audio_transcription.transformation import (
@@ -125,7 +125,7 @@ def default_session_config(model: str) -> ChirpSessionConfig:
)
-def parse_chirp_session_update(payload: str, expected_model: str) -> ChirpSessionConfig:
+def parse_chirp_session_update(payload: str, expected_model: str, *, drop_params: bool = False) -> ChirpSessionConfig:
update: Final = parse_transcription_session_update(payload, ChirpProtocolError)
if update.session_type not in (None, "transcription", "realtime"):
raise ChirpProtocolError("Speech-to-Text streaming supports transcription sessions only")
@@ -141,7 +141,7 @@ def parse_chirp_session_update(payload: str, expected_model: str) -> ChirpSessio
model=model,
language=None if update.language is None else normalize_transcription_language_to_bcp47(update.language),
sample_rate=_parse_sample_rate(update.audio_format),
- server_vad=_parse_server_vad(update),
+ server_vad=uses_server_vad(update, "Speech-to-Text streaming", drop_params, ChirpProtocolError),
)
@@ -161,14 +161,6 @@ def _parse_sample_rate(audio_format: TranscriptionAudioFormat | None) -> int:
return rate
-def _parse_server_vad(update: TranscriptionSessionUpdate) -> bool:
- if update.turn_detection_disabled:
- return False
- if update.turn_detection_type not in (None, "server_vad"):
- raise ChirpProtocolError("Speech-to-Text streaming supports server_vad turn detection or null")
- return True
-
-
def session_created_event(config: ChirpSessionConfig, session_id: str) -> OpenAIRealtimeTranscriptionSessionCreated:
return transcription_session_created_event(config.openai_session(session_id))
@@ -331,11 +323,13 @@ class VertexChirpRealtimeConfig(BaseRealtimeConfig):
project: str,
location: str | None,
backend_factory: Callable[[SpeechStreamingTarget], RealtimeBackend] = _default_backend_factory,
+ drop_params: bool = False,
) -> None:
self._resolve_access_token: Final = resolve_access_token
self._project: Final = validate_vertex_transcription_project_id(project)
self._location: Final = validate_vertex_transcription_location(location, DEFAULT_SPEECH_TO_TEXT_LOCATION)
self._backend_factory: Final = backend_factory
+ self._drop_params: Final = drop_params
self._transformer: Final = ChirpEventTransformer()
self._config: ChirpSessionConfig | None = None
self._session_id: str | None = None
@@ -423,7 +417,7 @@ class VertexChirpRealtimeConfig(BaseRealtimeConfig):
if self._config is not None:
verbose_logger.debug("Speech-to-Text streaming: ignoring session.update after the stream was configured")
return ()
- config: Final = parse_chirp_session_update(message, model)
+ config: Final = parse_chirp_session_update(message, model, drop_params=self._drop_params)
self._config = config
self._transformer.configure(config, self._session_id or f"sess_{uuid.uuid4().hex}")
return (config.configure_command(),)
diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py
index 98bc679a78e..581e2cb15f7 100644
--- a/litellm/llms/vertex_ai/common_utils.py
+++ b/litellm/llms/vertex_ai/common_utils.py
@@ -738,14 +738,14 @@ def build_vertex_schema(parameters: dict, add_property_ordering: bool = False):
# refs recursively and correctly detects/skips circular references.
unpack_defs(parameters, defs)
+ _convert_schema_types(parameters)
+
# 5. Nullable fields:
# * https://github.com/pydantic/pydantic/issues/1270
# * https://stackoverflow.com/a/58841311
# * https://github.com/pydantic/pydantic/discussions/4872
convert_anyof_null_to_nullable(parameters)
- _convert_schema_types(parameters)
-
# Handle empty strings in enum values - Gemini doesn't accept empty strings in enums
_fix_enum_empty_strings(parameters)
@@ -926,34 +926,28 @@ def convert_anyof_null_to_nullable(schema, depth=0):
)
""" Converts null objects within anyOf by removing them and adding nullable to all remaining objects """
anyof: Final = schema.get("anyOf", None)
- if anyof is not None:
- contains_null = False
+ if isinstance(anyof, list):
+ non_null_branches: Final = [
+ atype for atype in anyof if not (isinstance(atype, dict) and atype.get("type") == "null")
+ ]
+ contains_null: Final = len(non_null_branches) < len(anyof)
+ anyof[:] = non_null_branches or anyof[:1]
for atype in anyof:
- if isinstance(atype, dict) and atype.get("type") == "null":
- # remove null type
- anyof.remove(atype)
- contains_null = True
- elif "type" not in atype and len(atype) == 0:
- # Handle empty object case
+ if not isinstance(atype, dict):
+ continue
+ if not atype:
atype["type"] = "object"
-
- if len(anyof) == 0:
- # Edge case: response schema with only null type present is invalid in Vertex AI
- raise ValueError(
- "Invalid input: AnyOf schema with only null type is not supported. Please provide a non-null type."
- )
-
- if contains_null:
- # set all types to nullable following guidance found here: https://cloud.google.com/vertex-ai/generative-ai/docs/samples/generativeaionvertexai-gemini-controlled-generation-response-schema-3#generativeaionvertexai_gemini_controlled_generation_response_schema_3-python
- # Empty `items: {}` on array branches is left in place; downstream
- # process_items() converts it to {"type": "object"}, which Vertex
- # requires whenever type == "array" (even inside anyOf).
- for atype in anyof:
+ if contains_null:
+ # set all types to nullable following guidance found here: https://cloud.google.com/vertex-ai/generative-ai/docs/samples/generativeaionvertexai-gemini-controlled-generation-response-schema-3#generativeaionvertexai_gemini_controlled_generation_response_schema_3-python
+ # Empty `items: {}` on array branches is left in place; downstream
+ # process_items() converts it to {"type": "object"}, which Vertex
+ # requires whenever type == "array" (even inside anyOf).
atype["nullable"] = True
+ convert_anyof_null_to_nullable(atype, depth=depth + 1)
properties: Final = schema.get("properties", None)
if properties is not None:
- for name, value in properties.items():
+ for value in properties.values():
convert_anyof_null_to_nullable(value, depth=depth + 1)
items: Final = schema.get("items", None)
diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py
index 895d59965b7..7b832b321da 100644
--- a/litellm/llms/vertex_ai/gemini/transformation.py
+++ b/litellm/llms/vertex_ai/gemini/transformation.py
@@ -1154,6 +1154,7 @@ def transform_request_body(
user_response_schema_message: Final = response_schema_prompt(
model=model,
response_schema=optional_params.get("response_schema"),
+ llm_provider=custom_llm_provider,
)
messages.append({"role": "user", "content": user_response_schema_message})
optional_params.pop("response_schema")
diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
index bed7bc1bd74..257fafbb6b8 100644
--- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
+++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
@@ -27,6 +27,7 @@ from litellm.constants import (
DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_PRO,
)
from litellm.exceptions import UnsupportedParamsError
+from litellm.litellm_core_utils.core_helpers import normalize_drop_params
from litellm.litellm_core_utils.json_fragment_accumulator import JSONFragmentAccumulator
from litellm.litellm_core_utils.prompt_templates.factory import (
encode_tool_call_id_with_signature,
@@ -926,11 +927,38 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
# - Requires propertyOrdering
optional_params["response_schema"] = self._map_response_schema(value=schema)
+ @property
+ def accepts_reasoning_summary_in_reasoning_effort(self) -> bool:
+ return True
+
+ @staticmethod
+ def _include_thoughts(requested: bool | None) -> bool:
+ if requested is not None:
+ return requested
+ return normalize_drop_params(litellm.gemini_include_thoughts_default) is not False
+
+ @staticmethod
+ def _include_thoughts_requested_by_reasoning_effort(reasoning_effort: object) -> bool | None:
+ match reasoning_effort:
+ case {"summary": str()}:
+ return True
+ case _:
+ return None
+
+ @staticmethod
+ def _map_reasoning_effort(effort: str, model: str, reasoning_effort: object) -> GeminiThinkingConfig:
+ include_thoughts: Final = VertexGeminiConfig._include_thoughts_requested_by_reasoning_effort(reasoning_effort)
+ if VertexGeminiConfig._is_gemini_3_or_newer(model):
+ return VertexGeminiConfig._map_reasoning_effort_to_thinking_level(effort, model, include_thoughts)
+ return VertexGeminiConfig._map_reasoning_effort_to_thinking_budget(effort, model, include_thoughts)
+
@staticmethod
def _map_reasoning_effort_to_thinking_budget(
reasoning_effort: str,
model: str | None = None,
+ include_thoughts: bool | None = None,
) -> GeminiThinkingConfig:
+ sends_thoughts: Final = VertexGeminiConfig._include_thoughts(include_thoughts)
if reasoning_effort == "minimal":
# Use model-specific minimum thinking budget or fallback
# Check for exact matches first, then partial matches
@@ -945,22 +973,22 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
return {
"thinkingBudget": budget,
- "includeThoughts": True,
+ "includeThoughts": sends_thoughts,
}
elif reasoning_effort == "low":
return {
"thinkingBudget": DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET,
- "includeThoughts": True,
+ "includeThoughts": sends_thoughts,
}
elif reasoning_effort == "medium":
return {
"thinkingBudget": DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET,
- "includeThoughts": True,
+ "includeThoughts": sends_thoughts,
}
elif reasoning_effort == "high":
return {
"thinkingBudget": DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET,
- "includeThoughts": True,
+ "includeThoughts": sends_thoughts,
}
elif reasoning_effort == "disable":
return {
@@ -987,6 +1015,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
def _map_reasoning_effort_to_thinking_level(
reasoning_effort: str,
model: str | None = None,
+ include_thoughts: bool | None = None,
) -> GeminiThinkingConfig:
"""
Map reasoning_effort to thinking_level for Gemini 3+ models.
@@ -998,17 +1027,18 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
GeminiThinkingConfig with thinkingLevel and includeThoughts
"""
supports_minimal: Final = bool(model) and VertexGeminiConfig._supports_minimal_thinking_level(model)
+ sends_thoughts: Final = VertexGeminiConfig._include_thoughts(include_thoughts)
if reasoning_effort == "minimal":
if supports_minimal:
- return {"thinkingLevel": "minimal", "includeThoughts": True}
+ return {"thinkingLevel": "minimal", "includeThoughts": sends_thoughts}
else:
- return {"thinkingLevel": "low", "includeThoughts": True}
+ return {"thinkingLevel": "low", "includeThoughts": sends_thoughts}
elif reasoning_effort == "low":
- return {"thinkingLevel": "low", "includeThoughts": True}
+ return {"thinkingLevel": "low", "includeThoughts": sends_thoughts}
elif reasoning_effort == "medium":
- return {"thinkingLevel": "medium", "includeThoughts": True}
+ return {"thinkingLevel": "medium", "includeThoughts": sends_thoughts}
elif reasoning_effort == "high":
- return {"thinkingLevel": "high", "includeThoughts": True}
+ return {"thinkingLevel": "high", "includeThoughts": sends_thoughts}
elif reasoning_effort in ("disable", "none"):
return {
"thinkingLevel": "minimal" if supports_minimal else "low",
@@ -1061,6 +1091,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
status_code=400,
)
+ @staticmethod
+ def _thinking_asks_for_thoughts(thinking_param: Mapping[str, object]) -> bool:
+ display: Final = thinking_param.get("display")
+ return isinstance(thinking_param.get("summary"), str) or (display is not None and display != "omitted")
+
@staticmethod
def _map_thinking_param(
thinking_param: AnthropicThinkingParam,
@@ -1068,6 +1103,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
) -> GeminiThinkingConfig:
thinking_enabled: Final = thinking_param.get("type") == "enabled"
thinking_budget: Final = thinking_param.get("budget_tokens")
+ asks_for_thoughts: Final = VertexGeminiConfig._thinking_asks_for_thoughts(thinking_param)
+ sends_thoughts: Final = VertexGeminiConfig._include_thoughts(True if asks_for_thoughts else None)
params: Final[GeminiThinkingConfig] = {}
@@ -1077,7 +1114,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
if thinking_budget == 0:
params["includeThoughts"] = False
else:
- params["includeThoughts"] = True
+ params["includeThoughts"] = sends_thoughts
# Follow provider defaults unless explicitly opted into legacy behavior.
if litellm.enable_gemini_default_thinking_level_low is True:
params["thinkingLevel"] = (
@@ -1089,7 +1126,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
else:
# For older Gemini models, use thinkingBudget
if thinking_enabled and not VertexGeminiConfig._is_thinking_budget_zero(thinking_budget):
- params["includeThoughts"] = True
+ params["includeThoughts"] = sends_thoughts
if thinking_budget is not None and isinstance(thinking_budget, int):
params["thinkingBudget"] = thinking_budget
@@ -1304,14 +1341,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
param_name="reasoning_effort",
param_description="thinking_budget",
)
- if VertexGeminiConfig._is_gemini_3_or_newer(model):
- optional_params["thinkingConfig"] = VertexGeminiConfig._map_reasoning_effort_to_thinking_level(
- effort_value, model
- )
- else:
- optional_params["thinkingConfig"] = VertexGeminiConfig._map_reasoning_effort_to_thinking_budget(
- effort_value, model
- )
+ optional_params["thinkingConfig"] = VertexGeminiConfig._map_reasoning_effort(
+ effort=effort_value, model=model, reasoning_effort=value
+ )
elif param == "thinking":
# Validate no conflict with thinking_level
VertexGeminiConfig._validate_thinking_config_conflicts(
diff --git a/litellm/llms/vertex_ai/realtime/transformation.py b/litellm/llms/vertex_ai/realtime/transformation.py
index c14dd535145..55f1cf58cbe 100644
--- a/litellm/llms/vertex_ai/realtime/transformation.py
+++ b/litellm/llms/vertex_ai/realtime/transformation.py
@@ -47,6 +47,9 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig):
def _include_function_response_id(self) -> bool:
return False
+ def _rejects_vad_tuning_with_detection_off(self) -> bool:
+ return False
+
# ------------------------------------------------------------------
# URL
# ------------------------------------------------------------------
@@ -248,9 +251,12 @@ def vertex_realtime_config(
resolve_access_token: Callable[[], Awaitable[str]],
project: str,
location: str | None,
+ drop_params: bool,
) -> VertexAIRealtimeConfig | VertexChirpRealtimeConfig:
if is_vertex_speech_to_text_model(model):
- return VertexChirpRealtimeConfig(resolve_access_token=resolve_access_token, project=project, location=location)
+ return VertexChirpRealtimeConfig(
+ resolve_access_token=resolve_access_token, project=project, location=location, drop_params=drop_params
+ )
return VertexAIRealtimeConfig(
access_token=access_token,
project=project,
diff --git a/litellm/llms/vllm/completion/handler.py b/litellm/llms/vllm/completion/handler.py
index 78e6c74c2f7..f330f0928d0 100644
--- a/litellm/llms/vllm/completion/handler.py
+++ b/litellm/llms/vllm/completion/handler.py
@@ -5,8 +5,8 @@ from typing import Final
import httpx
from litellm.litellm_core_utils.prompt_templates.factory import (
- custom_prompt,
prompt_factory,
+ render_custom_prompt_template,
)
from litellm.utils import ModelResponse, Usage
@@ -56,11 +56,8 @@ def completion(
if model in custom_prompt_dict:
# check if the model has a registered custom prompt
model_prompt_details: Final = custom_prompt_dict[model]
- prompt = custom_prompt(
- role_dict=model_prompt_details["roles"],
- initial_prompt_value=model_prompt_details["initial_prompt_value"],
- final_prompt_value=model_prompt_details["final_prompt_value"],
- messages=messages,
+ prompt = render_custom_prompt_template(
+ template=model_prompt_details, messages=messages, model=model, llm_provider="vllm"
)
else:
prompt = prompt_factory(model=model, messages=messages)
@@ -144,11 +141,8 @@ def batch_completions(model: str, messages: list, optional_params=None, custom_p
# check if the model has a registered custom prompt
model_prompt_details: Final = custom_prompt_dict[model]
for message in messages:
- prompt = custom_prompt(
- role_dict=model_prompt_details["roles"],
- initial_prompt_value=model_prompt_details["initial_prompt_value"],
- final_prompt_value=model_prompt_details["final_prompt_value"],
- messages=message,
+ prompt = render_custom_prompt_template(
+ template=model_prompt_details, messages=message, model=model, llm_provider="vllm"
)
prompts.append(prompt)
else:
diff --git a/litellm/llms/volcengine/videos/__init__.py b/litellm/llms/volcengine/videos/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/litellm/llms/volcengine/videos/transformation.py b/litellm/llms/volcengine/videos/transformation.py
new file mode 100644
index 00000000000..a75a4185186
--- /dev/null
+++ b/litellm/llms/volcengine/videos/transformation.py
@@ -0,0 +1,904 @@
+import json
+from collections.abc import Mapping, Sequence
+from math import log
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Final, TypeVar
+
+import httpx
+from httpx._types import RequestFiles
+from pydantic import BaseModel, StrictBool, StrictInt, StrictStr, TypeAdapter, ValidationError
+
+import litellm
+from litellm.exceptions import UnsupportedParamsError
+from litellm.litellm_core_utils.url_utils import encode_url_path_segment
+from litellm.llms.base_llm.chat.transformation import BaseLLMException
+from litellm.llms.base_llm.videos.async_task_utils import (
+ input_reference_url,
+ polled_video_models,
+ routed_video_models,
+)
+from litellm.llms.base_llm.videos.transformation import BaseVideoConfig
+from litellm.llms.custom_httpx.http_handler import (
+ AsyncHTTPHandler,
+ HTTPHandler,
+ _get_httpx_client, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # house cached-client factory has no public alias and its stub leaves params untyped
+ get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # factory stub leaves params untyped
+)
+from litellm.llms.volcengine.common_utils import get_volcengine_base_url
+from litellm.secret_managers.main import get_secret_str
+from litellm.types.router import GenericLiteLLMParams
+from litellm.types.utils import ModelInfo
+from litellm.types.videos.main import VideoCreateOptionalRequestParams, VideoObject
+from litellm.types.videos.utils import encode_video_id_with_provider, extract_original_video_id
+from litellm.utils import get_model_info
+
+if TYPE_CHECKING:
+ from litellm.litellm_core_utils.litellm_logging import Logging
+
+
+_ParsedT: Final = TypeVar("_ParsedT")
+
+
+class VolcEngineVideoError(BaseLLMException):
+ pass
+
+
+class _TaskError(BaseModel, frozen=True):
+ code: str | None = None
+ message: str | None = None
+
+
+class _TaskContent(BaseModel, frozen=True):
+ video_url: str | None = None
+
+
+class _TaskUsage(BaseModel, frozen=True):
+ total_tokens: int | None = None
+
+
+class _Task(BaseModel, frozen=True):
+ id: str
+ status: str
+ model: str | None = None
+ error: _TaskError | None = None
+ content: _TaskContent | None = None
+ usage: _TaskUsage | None = None
+ created_at: int | None = None
+ updated_at: int | None = None
+ resolution: str | None = None
+ ratio: str | None = None
+ duration: int | float | None = None
+
+
+class _TaskCreateResponse(BaseModel, frozen=True):
+ id: str | None = None
+ error: _TaskError | None = None
+
+
+class _TaskList(BaseModel, frozen=True):
+ items: Sequence[_Task] = ()
+ total: int | None = None
+
+
+class _ContentItem(BaseModel, frozen=True):
+ type: str | None = None
+ role: str | None = None
+
+
+class _KnownFields(BaseModel, frozen=True):
+ resolution: StrictStr | None = None
+ ratio: StrictStr | None = None
+ seed: StrictInt | None = None
+ camera_fixed: StrictBool | None = None
+ watermark: StrictBool | None = None
+ generate_audio: StrictBool | None = None
+
+
+VOLCENGINE_VIDEO_TASKS_PATH: Final = "/api/v3/contents/generations/tasks"
+VOLCENGINE_VIDEO_DEFAULT_DURATION_SECONDS: Final = 4
+VOLCENGINE_VIDEO_DEFAULT_RESOLUTION: Final = "720p"
+_VOLCENGINE_API_MARKER: Final = "/api/"
+_VOLCENGINE_COMPAT_SUFFIXES: Final = ("/api", "/v1")
+_FRAMES_FIELD: Final = "volcengine_video_frames"
+
+_EMPTY_PARAMS: Final[dict[str, object]] = {} # mutable-ok: BaseVideoConfig contract empty params
+
+_TASK_ADAPTER: Final = TypeAdapter(_Task)
+_TASK_CREATE_ADAPTER: Final = TypeAdapter(_TaskCreateResponse)
+_TASK_LIST_ADAPTER: Final = TypeAdapter(_TaskList)
+_FIELDS_ADAPTER: Final = TypeAdapter(dict[str, object])
+_CONTENT_ITEMS_ADAPTER: Final = TypeAdapter(tuple[dict[str, object], ...])
+_FRAMES_ADAPTER: Final = TypeAdapter(dict[str, dict[str, str]])
+_EMPTY_FRAMES: Final[Mapping[str, Mapping[str, str]]] = MappingProxyType({})
+
+_STATUS_MAP: Final = MappingProxyType(
+ {
+ "queued": "queued",
+ "running": "in_progress",
+ "succeeded": "completed",
+ "failed": "failed",
+ "cancelled": "cancelled",
+ "expired": "failed",
+ }
+)
+_PENDING_STATUSES: Final = frozenset({"queued", "running"})
+_TERMINAL_STATUSES: Final = frozenset({"completed", "failed", "cancelled"})
+
+_CREATE_PASSTHROUGH_KEYS: Final = (
+ "resolution",
+ "ratio",
+ "duration",
+ "seed",
+ "camera_fixed",
+ "watermark",
+ "generate_audio",
+ "service_tier",
+ "priority",
+ "output_format",
+ "execution_expires_after",
+ "callback_url",
+ "safety_identifier",
+ "tools",
+ "omni_reference_task_type",
+)
+
+_REQUEST_METADATA: Final = frozenset({"model", "prompt", "user", "extra_headers", "extra_query", "extra_body"})
+_UNSUPPORTED_BY_VOLCENGINE: Final = frozenset({"characters", "image"})
+_MAPPED_ELSEWHERE: Final = frozenset({*_REQUEST_METADATA, *_UNSUPPORTED_BY_VOLCENGINE, "seconds", "size", "parameters"})
+_READ_FROM_THE_REQUEST: Final = _MAPPED_ELSEWHERE - _UNSUPPORTED_BY_VOLCENGINE
+_LITELLM_CONTROLS: Final = frozenset({"drop_params", "additional_drop_params", "organization"})
+_DROP_KEYS_ADAPTER: Final = TypeAdapter(frozenset[str])
+
+# litellm/videos/utils.py merges the caller's raw ``extra_body`` back over the mapped params, which returns the keys
+# this adapter maps itself a second time, as fields of their own. They are already accounted for, so they count as
+# handled rather than unknown.
+_MAPPED_EXTRA_BODY_COPIES: Final = frozenset({"size", "seconds", "user"})
+_CREATE_HANDLED_KEYS: Final = frozenset(
+ {
+ *_CREATE_PASSTHROUGH_KEYS,
+ "content",
+ "input_reference",
+ "frames",
+ *_MAPPED_EXTRA_BODY_COPIES,
+ *_LITELLM_CONTROLS,
+ }
+)
+
+_VIDEO_INPUT_UNSUPPORTED: Final = (
+ "{operation} is not supported through litellm: Ark bills a task with a video input at a different rate per "
+ "token, and litellm bills the video when the task is created, before that cost is known"
+)
+_INPUT_VIDEO_ROLE: Final = "reference_video"
+_INPUT_VIDEO_TYPE: Final = "video_url"
+
+_RATIO_SET: Final = ("16:9", "9:16", "4:3", "3:4", "1:1", "21:9")
+_RESOLUTION_SHORTEST_SIDE: Final = MappingProxyType({"480p": 480, "720p": 720, "1080p": 1080, "4k": 2160})
+_TIER_FRAME_ASPECT: Final = 16 / 9
+
+
+def _video_error(raw_response: httpx.Response, status_code: int, message: str) -> VolcEngineVideoError:
+ """The provider's own error status wins, and a body that failed on a 2xx answers the caller's status instead,
+ since a 2xx error has the proxy retry a call the provider already rejected."""
+ return VolcEngineVideoError(
+ status_code=raw_response.status_code if raw_response.status_code >= 400 else status_code,
+ message=message,
+ headers=raw_response.headers,
+ )
+
+
+def _parse(adapter: TypeAdapter[_ParsedT], raw_response: httpx.Response) -> _ParsedT:
+ try:
+ return adapter.validate_python(raw_response.json())
+ except (json.JSONDecodeError, ValidationError) as e:
+ raise _video_error(raw_response, 502, f"Unexpected Volcengine video response: {e}") from e
+
+
+def _unsupported(model: str, message: str) -> UnsupportedParamsError:
+ return UnsupportedParamsError(message=message, model=model, llm_provider="volcengine")
+
+
+def _droppable(model: str, message: str) -> UnsupportedParamsError:
+ return _unsupported(model, f"{message}, set drop_params=True to ignore it")
+
+
+def _normalized_api_base(api_base: str) -> str:
+ """The Ark credential is shared with the chat, responses and Anthropic surfaces, so its api_base may carry any of
+ their paths, and the video paths carry their own version segment."""
+ trimmed: Final = api_base.rstrip("/")
+ marker_index: Final = trimmed.rfind(_VOLCENGINE_API_MARKER)
+ if marker_index != -1:
+ return trimmed[:marker_index]
+ suffix: Final = next((suffix for suffix in _VOLCENGINE_COMPAT_SUFFIXES if trimmed.endswith(suffix)), None)
+ return trimmed[: -len(suffix)] if suffix is not None else trimmed
+
+
+def _size_dimensions(size: str) -> tuple[int, int] | None:
+ width_str, separator, height_str = size.partition("x")
+ if not separator or not (width_str.isdigit() and height_str.isdigit()):
+ return None
+ if not int(width_str) or not int(height_str):
+ return None
+ return int(width_str), int(height_str)
+
+
+def _ratio_value(ratio: str) -> float:
+ width_str, _, height_str = ratio.partition(":")
+ return int(width_str) / int(height_str)
+
+
+def _ratio_from_size(size: str) -> str | None:
+ """Ark only takes its named ratios, and its own tiers do not reduce onto them (854x480 is 427:240), so the
+ closest ratio wins."""
+ dimensions: Final = _size_dimensions(size)
+ if dimensions is None:
+ return None
+ width, height = dimensions
+ return min(_RATIO_SET, key=lambda candidate: abs(log(width / height / _ratio_value(candidate))))
+
+
+def _resolution_from_size(size: str) -> str | None:
+ """Ark keeps a tier's pixel count across its ratios (960x960 and 1280x720 are both 720p), so the tier whose 16:9
+ frame has the closest pixel count wins."""
+ dimensions: Final = _size_dimensions(size)
+ if dimensions is None:
+ return None
+ width, height = dimensions
+ return min(
+ _RESOLUTION_SHORTEST_SIDE,
+ key=lambda tier: abs(log(width * height / (_RESOLUTION_SHORTEST_SIDE[tier] ** 2 * _TIER_FRAME_ASPECT))),
+ )
+
+
+def _duration_param(seconds: object) -> int | None:
+ if isinstance(seconds, bool):
+ return None
+ if isinstance(seconds, int):
+ return seconds
+ if not isinstance(seconds, (str, float)):
+ return None
+ try:
+ parsed: Final = float(seconds)
+ except ValueError:
+ return None
+ return int(parsed) if parsed.is_integer() else None
+
+
+def _duration_fields(seconds: object, model: str, drop_params: bool) -> Mapping[str, object]:
+ if seconds is None:
+ return _EMPTY_PARAMS
+ duration: Final = _duration_param(seconds)
+ if duration is not None:
+ return MappingProxyType({"duration": duration})
+ if drop_params:
+ return _EMPTY_PARAMS
+ raise _droppable(model, f"seconds must be a whole number of seconds, got {seconds!r}")
+
+
+def _size_fields(size: object, model: str, drop_params: bool) -> Mapping[str, object]:
+ if size is None:
+ return _EMPTY_PARAMS
+ ratio: Final = _ratio_from_size(size) if isinstance(size, str) else None
+ resolution: Final = _resolution_from_size(size) if isinstance(size, str) else None
+ if ratio is not None and resolution is not None:
+ return MappingProxyType({"ratio": ratio, "resolution": resolution})
+ if drop_params:
+ return _EMPTY_PARAMS
+ raise _droppable(model, f"size must be WIDTHxHEIGHT, got {size!r}")
+
+
+def _parameters_fields(parameters: object, model: str, drop_params: bool) -> Mapping[str, object]:
+ if parameters is None:
+ return _EMPTY_PARAMS
+ try:
+ fields: Final = _FIELDS_ADAPTER.validate_python(parameters)
+ except ValidationError:
+ if drop_params:
+ return _EMPTY_PARAMS
+ raise _droppable(
+ model, f"parameters must be an object of Ark request fields, got {type(parameters).__name__}"
+ ) from None
+ misplaced: Final = tuple(sorted(key for key in fields if key in _READ_FROM_THE_REQUEST))
+ if misplaced and not drop_params:
+ raise _droppable(model, f"send {', '.join(misplaced)} as top-level request fields, not inside parameters")
+ return MappingProxyType({key: value for key, value in fields.items() if key not in _READ_FROM_THE_REQUEST})
+
+
+def _user_fields(user: object, model: str, drop_params: bool) -> Mapping[str, object]:
+ if isinstance(user, str) and user:
+ return MappingProxyType({"safety_identifier": user})
+ if user is None or isinstance(user, str) or drop_params:
+ return _EMPTY_PARAMS
+ raise _droppable(model, f"user must be a string, got {type(user).__name__}")
+
+
+def _additional_drop_keys(params: Mapping[str, object]) -> frozenset[str]:
+ try:
+ return _DROP_KEYS_ADAPTER.validate_python(params.get("additional_drop_params") or ())
+ except ValidationError:
+ return frozenset()
+
+
+def _without(params: Mapping[str, object], keys: frozenset[str]) -> Mapping[str, object]:
+ return MappingProxyType({key: value for key, value in params.items() if key not in keys})
+
+
+def _malformed_fields(fields: Mapping[str, object]) -> tuple[tuple[str, str], ...]:
+ try:
+ _ = _KnownFields.model_validate(dict(fields))
+ except ValidationError as e:
+ return tuple((str(error["loc"][0]), error["msg"]) for error in e.errors() if error["loc"])
+ return ()
+
+
+def _drops_params(litellm_params: GenericLiteLLMParams) -> bool:
+ return litellm_params.drop_params is True or litellm.drop_params is True
+
+
+def _model_info(model: str) -> ModelInfo | None:
+ try:
+ return get_model_info(model=model, custom_llm_provider="volcengine")
+ except Exception: # noqa: BLE001 # get_model_info raises bare Exception for unmapped models
+ return None
+
+
+def _pricing_model(model: str, litellm_params: GenericLiteLLMParams) -> str:
+ """A deployment priced through ``model_info.base_model``, like an ``ep-`` endpoint id, bills at that model's
+ tiers."""
+ model_info: Final = _FIELDS_ADAPTER.validate_python(litellm_params["model_info"] or _EMPTY_PARAMS)
+ base_model: Final = model_info.get("base_model")
+ return base_model if isinstance(base_model, str) and base_model else model
+
+
+def _priced_tiers(model: str) -> tuple[str, ...] | None:
+ """A tier the model has no price for is one Ark does not offer it at, and litellm would bill it at the base rate.
+ A model with no tier price listed, missing from the cost map or registered unpriced by the router, has nothing
+ to check against."""
+ info: Final = _model_info(model)
+ if info is None:
+ return None
+ return (
+ tuple(tier for tier in _RESOLUTION_SHORTEST_SIDE if info.get(f"output_cost_per_second_{tier}") is not None)
+ or None
+ )
+
+
+def _delivered_frames(model: str | None) -> Mapping[str, Mapping[str, str]]:
+ info: Final = _model_info(model) if model else None
+ provider_specific: Final[object] = info.get("provider_specific_entry") if info is not None else None
+ frames: Final[object] = provider_specific.get(_FRAMES_FIELD) if isinstance(provider_specific, Mapping) else None
+ try:
+ return _FRAMES_ADAPTER.validate_python(frames) if frames is not None else _EMPTY_FRAMES
+ except ValidationError:
+ return _EMPTY_FRAMES
+
+
+def _size_from_geometry(resolution: str | None, ratio: str | None, model: str | None) -> str | None:
+ """Ark echoes a tier and a ratio rather than pixels, and its rendered frame does not always follow the tier's
+ nominal size, so the frame listed for the model wins over scaling the tier's shortest side up the ratio."""
+ if not resolution or not ratio or ":" not in ratio:
+ return None
+ tier: Final = resolution.strip().lower()
+ normalized_ratio: Final = ratio.strip()
+ listed: Final = _delivered_frames(model).get(tier, MappingProxyType({})).get(normalized_ratio)
+ if listed is not None:
+ return listed
+ shortest_side: Final = _RESOLUTION_SHORTEST_SIDE.get(tier)
+ width_str, _, height_str = normalized_ratio.partition(":")
+ if shortest_side is None or not (width_str.isdigit() and height_str.isdigit()):
+ return None
+ ratio_width: Final = int(width_str)
+ ratio_height: Final = int(height_str)
+ if not ratio_width or not ratio_height:
+ return None
+ shortest_ratio_side: Final = min(ratio_width, ratio_height)
+ return f"{-(-shortest_side * ratio_width // shortest_ratio_side)}x{-(-shortest_side * ratio_height // shortest_ratio_side)}"
+
+
+def _is_input_video(item: _ContentItem) -> bool:
+ return item.type == _INPUT_VIDEO_TYPE or item.role == _INPUT_VIDEO_ROLE
+
+
+def _supplied_items(supplied: object, model: str, drop_params: bool) -> tuple[Mapping[str, object], ...]:
+ try:
+ items: Final = _CONTENT_ITEMS_ADAPTER.validate_python(supplied)
+ entries: Final = tuple(_ContentItem.model_validate(item) for item in items)
+ except ValidationError:
+ if drop_params:
+ return ()
+ raise _droppable(model, "content must be an array of Ark content items") from None
+ if any(entry.type == "text" for entry in entries) and not drop_params:
+ raise _droppable(
+ model,
+ "A text item in content is not supported, put the text in prompt since it is the only text input "
+ "guardrails scan before the request reaches Ark",
+ )
+ if any(_is_input_video(entry) for entry in entries) and not drop_params:
+ raise _droppable(model, _VIDEO_INPUT_UNSUPPORTED.format(operation="Video input (a reference_video item)"))
+ kept: Final = (item for item, entry in zip(items, entries) if entry.type != "text" and not _is_input_video(entry))
+ return tuple(MappingProxyType(item) for item in kept)
+
+
+def _reference_items(params: Mapping[str, object], model: str, drop_params: bool) -> tuple[Mapping[str, object], ...]:
+ """A caller-supplied ``content`` array wins over ``input_reference``, since its roles are the only way to express
+ reference-to-video."""
+ supplied: Final = params.get("content")
+ if supplied is not None:
+ return _supplied_items(supplied, model, drop_params)
+ reference: Final = params.get("input_reference")
+ if reference is None:
+ return ()
+ try:
+ url: Final = input_reference_url(reference)
+ except ValueError as e:
+ if drop_params:
+ return ()
+ raise _droppable(model, str(e)) from None
+ return (MappingProxyType({"type": "image_url", "image_url": {"url": url}, "role": "first_frame"}),)
+
+
+def _sent_content(
+ prompt: str, params: Mapping[str, object], model: str, drop_params: bool
+) -> tuple[Mapping[str, object], ...]:
+ """The prompt always leads, as it is the only text the input guardrails scan."""
+ return (MappingProxyType({"type": "text", "text": prompt}), *_reference_items(params, model, drop_params))
+
+
+def _seconds_label(seconds: object) -> str | None:
+ if isinstance(seconds, float) and seconds.is_integer():
+ return str(int(seconds))
+ return None if seconds is None else str(seconds)
+
+
+def _requested_usage(request_data: Mapping[str, object]) -> Mapping[str, object]:
+ """The create call is the only billed one and Ark's create response carries no usage, so the requested duration
+ and tier are what the video bills at."""
+ duration: Final = _duration_param(request_data.get("duration"))
+ resolution: Final = request_data.get("resolution")
+ return MappingProxyType(
+ {
+ key: value
+ for key, value in (
+ ("duration_seconds", float(duration) if duration is not None and duration > 0 else None),
+ ("video_resolution", resolution.strip().lower() if isinstance(resolution, str) else None),
+ )
+ if value is not None
+ }
+ )
+
+
+def _task_usage(task: _Task) -> Mapping[str, object]:
+ resolution: Final = task.resolution.strip().lower() if task.resolution else None
+ return MappingProxyType(
+ {
+ key: value
+ for key, value in (
+ ("duration_seconds", float(task.duration) if task.duration is not None else None),
+ ("video_resolution", resolution),
+ ("total_tokens", task.usage.total_tokens if task.usage is not None else None),
+ )
+ if value is not None
+ }
+ )
+
+
+def _task_status(task: _Task, raw_response: httpx.Response) -> str:
+ status: Final = _STATUS_MAP.get(task.status.strip().lower())
+ if status is None: # defaulting it would hand back a job the caller polls forever
+ raise _video_error(raw_response, 502, f"Unknown Volcengine task status: {task.status}")
+ return status
+
+
+def _video_object_from_task(task: _Task, raw_response: httpx.Response, model: str | None) -> VideoObject:
+ status: Final = _task_status(task, raw_response)
+ usage: Final = _task_usage(task)
+ seconds: Final = usage.get("duration_seconds")
+ error: Final[dict[str, object] | None] = ( # mutable-ok: VideoObject.error is a dict field
+ {key: value for key, value in (("code", task.error.code), ("message", task.error.message)) if value is not None}
+ if task.error is not None
+ else None
+ )
+ return VideoObject(
+ id=task.id,
+ object="video",
+ status=status,
+ created_at=task.created_at,
+ completed_at=task.updated_at if status in _TERMINAL_STATUSES else None,
+ error=error,
+ seconds=_seconds_label(seconds),
+ size=_size_from_geometry(task.resolution, task.ratio, task.model),
+ model=model or task.model,
+ usage=dict(usage) if usage else None,
+ )
+
+
+def _video_url_from_task(task: _Task, raw_response: httpx.Response) -> str:
+ """A download before the video is ready is the caller polling too early, not a fault of the deployment, so it
+ answers 400, which litellm neither retries nor counts toward a deployment cooldown. A 404 would cool the
+ deployment down, and a 409 or a 5xx would be retried against a job that is simply not finished."""
+ if task.content is not None and task.content.video_url:
+ return task.content.video_url
+ status: Final = task.status.strip().lower()
+ if status in _PENDING_STATUSES:
+ raise _video_error(
+ raw_response, 400, f"Video is still processing (status: {status}). Please wait and try again."
+ )
+ if task.error is not None and (task.error.message or task.error.code):
+ raise _video_error(raw_response, 400, f"Video generation failed: {task.error.message or task.error.code}")
+ if status == "expired":
+ raise _video_error(raw_response, 400, "Task expired. Ark keeps task records for 7 days after creation.")
+ if status in ("failed", "cancelled"):
+ raise _video_error(raw_response, 400, f"Video generation ended as {status} with no video URL.")
+ raise _video_error(raw_response, 502, "Video URL not found in task response.")
+
+
+class VolcEngineVideoConfig(BaseVideoConfig):
+ """Volcengine Ark (Doubao-Seedance) on Ark's async task API: create, poll, download from the signed
+ ``content.video_url``, and list."""
+
+ def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: BaseVideoConfig contract returns list
+ return [
+ "model",
+ "prompt",
+ "input_reference",
+ "seconds",
+ "size",
+ "user",
+ "extra_headers",
+ "content",
+ "parameters",
+ *_CREATE_PASSTHROUGH_KEYS,
+ ]
+
+ def map_openai_params(
+ self,
+ video_create_optional_params: VideoCreateOptionalRequestParams,
+ model: str,
+ drop_params: bool,
+ ) -> dict[str, object]: # mutable-ok: BaseVideoConfig contract
+ requested: Final = _without(video_create_optional_params, _additional_drop_keys(video_create_optional_params))
+ unsupported: Final = next(
+ (key for key, value in requested.items() if key in _UNSUPPORTED_BY_VOLCENGINE and value is not None),
+ None,
+ )
+ if unsupported is not None and not drop_params:
+ raise _droppable(model, f"'{unsupported}' is not supported by Volcengine video generation")
+ passthrough: Final = MappingProxyType(
+ {key: value for key, value in requested.items() if value is not None and key not in _MAPPED_ELSEWHERE}
+ )
+ return {
+ **_user_fields(requested.get("user"), model, drop_params),
+ **_size_fields(requested.get("size"), model, drop_params),
+ **_duration_fields(requested.get("seconds"), model, drop_params),
+ **passthrough,
+ **_parameters_fields(requested.get("parameters"), model, drop_params),
+ }
+
+ def _resolve_api_key(self, api_key: str | None) -> str:
+ resolved_api_key: Final = api_key or get_secret_str("ARK_API_KEY") or get_secret_str("VOLCENGINE_API_KEY")
+ if resolved_api_key is None:
+ raise ValueError(
+ "Volcengine API key is required. Set ARK_API_KEY or VOLCENGINE_API_KEY environment variable or pass "
+ "api_key parameter."
+ )
+ return resolved_api_key
+
+ def validate_environment(
+ self,
+ headers: dict[str, str], # mutable-ok: BaseVideoConfig contract; handler expects a mutable headers dict
+ model: str,
+ api_key: str | None = None,
+ litellm_params: GenericLiteLLMParams | None = None,
+ ) -> dict[str, str]: # mutable-ok: BaseVideoConfig contract
+ resolved_api_key: Final = self._resolve_api_key(
+ api_key
+ or (litellm_params.api_key if litellm_params is not None and litellm_params.api_key else None)
+ or litellm.api_key
+ )
+ headers.update({"Authorization": f"Bearer {resolved_api_key}", "Content-Type": "application/json"})
+ return headers
+
+ def get_complete_url(
+ self,
+ model: str,
+ api_base: str | None,
+ litellm_params: dict[str, object], # mutable-ok: BaseVideoConfig contract
+ ) -> str:
+ return _normalized_api_base(
+ get_volcengine_base_url(api_base or get_secret_str("VOLCENGINE_API_BASE") or get_secret_str("ARK_API_BASE"))
+ )
+
+ def transform_video_create_request(
+ self,
+ model: str,
+ prompt: str,
+ api_base: str,
+ video_create_optional_request_params: dict[str, object], # mutable-ok: BaseVideoConfig contract
+ litellm_params: GenericLiteLLMParams,
+ headers: dict[str, str], # mutable-ok: BaseVideoConfig contract
+ ) -> tuple[dict[str, object], RequestFiles, str]: # mutable-ok: BaseVideoConfig contract
+ """
+ Ark bills a video when the task finishes, while litellm bills it here at create, so every request has to
+ carry a length and a tier whose cost is known now. Ark-picked length (``duration: -1``) and a video input
+ (edits, extensions and reference videos, billed at a different rate per token) are rejected for that reason,
+ and an omitted length is sent as OpenAI's documented default of 4 seconds, since the 2.5 family's own default
+ is ``-1``.
+ """
+ drop_params: Final = _drops_params(litellm_params)
+ dropped_by_config: Final = _additional_drop_keys(video_create_optional_request_params)
+ kept: Final = _without(video_create_optional_request_params, dropped_by_config)
+ params: Final = _without(
+ {
+ **{key: value for key, value in kept.items() if key != "parameters"},
+ **_parameters_fields(kept.get("parameters"), model, drop_params),
+ },
+ dropped_by_config,
+ )
+ content: Final = _sent_content(prompt, params, model, drop_params)
+ if params.get("frames") is not None and not drop_params:
+ raise _droppable(model, "frames is not supported through litellm, pass seconds instead")
+ requested_duration: Final = params.get("duration")
+ parsed_duration: Final = None if requested_duration is None else _duration_param(requested_duration)
+ if requested_duration is not None and (parsed_duration is None or parsed_duration <= 0) and not drop_params:
+ raise _droppable(
+ model,
+ f"duration must be a positive whole number of seconds, got {requested_duration!r}. Ark-picked length "
+ "(-1) is not supported through litellm, since the video is billed when the task is created, before "
+ "Ark has picked its length",
+ )
+ duration: Final = (
+ parsed_duration
+ if parsed_duration is not None and parsed_duration > 0
+ else VOLCENGINE_VIDEO_DEFAULT_DURATION_SECONDS
+ )
+
+ unknown_fields: Final = tuple(
+ sorted(key for key, value in params.items() if value is not None and key not in _CREATE_HANDLED_KEYS)
+ )
+ if unknown_fields and not drop_params:
+ raise _droppable(model, f"litellm does not send {', '.join(unknown_fields)} to Ark video generation")
+ requested_fields: Final = MappingProxyType(
+ {key: params[key] for key in _CREATE_PASSTHROUGH_KEYS if params.get(key) is not None}
+ )
+ malformed: Final = _malformed_fields(requested_fields)
+ if malformed and not drop_params:
+ raise _droppable(model, "; ".join(f"invalid {field}: {message}" for field, message in malformed))
+ resolution: Final = requested_fields.get("resolution")
+ priced_tiers: Final = _priced_tiers(_pricing_model(model, litellm_params))
+ drops_resolution: Final = (
+ priced_tiers is not None and isinstance(resolution, str) and resolution.strip().lower() not in priced_tiers
+ )
+ if drops_resolution and not drop_params:
+ raise _droppable(
+ model, f"{model} is not offered at resolution {resolution!r}, only at {', '.join(priced_tiers or ())}"
+ )
+ dropped: Final = frozenset({*(field for field, _ in malformed), *(("resolution",) if drops_resolution else ())})
+ kept_fields: Final = MappingProxyType(
+ {key: value for key, value in requested_fields.items() if key not in dropped}
+ )
+ default_tier: Final = (
+ MappingProxyType({"resolution": VOLCENGINE_VIDEO_DEFAULT_RESOLUTION})
+ if "resolution" not in kept_fields
+ and priced_tiers is not None
+ and VOLCENGINE_VIDEO_DEFAULT_RESOLUTION in priced_tiers
+ else _EMPTY_PARAMS
+ )
+
+ request_data: Final[dict[str, object]] = { # mutable-ok: request body dict, JSON-serialized by the handler
+ "model": model,
+ "content": [dict(item) for item in content],
+ **default_tier,
+ **kept_fields,
+ "duration": duration,
+ }
+ return request_data, (), f"{api_base}{VOLCENGINE_VIDEO_TASKS_PATH}"
+
+ def transform_video_create_response(
+ self,
+ model: str,
+ raw_response: httpx.Response,
+ logging_obj: "Logging",
+ custom_llm_provider: str | None = None,
+ request_data: dict[str, object] | None = None, # mutable-ok: BaseVideoConfig contract
+ ) -> VideoObject:
+ task: Final = _parse(_TASK_CREATE_ADAPTER, raw_response)
+ if not task.id:
+ message: Final = task.error.message if task.error is not None and task.error.message else "no task id"
+ raise _video_error(raw_response, 400, f"Volcengine did not return a video task id: {message}")
+ requested: Final = request_data or _EMPTY_PARAMS
+ usage: Final = _requested_usage(requested)
+ seconds: Final = usage.get("duration_seconds")
+ requested_model, served_model = routed_video_models(logging_obj) or (model, None)
+ resolution: Final = requested.get("resolution")
+ ratio: Final = requested.get("ratio")
+ return VideoObject(
+ id=encode_video_id_with_provider(task.id, custom_llm_provider, requested_model, served_model)
+ if custom_llm_provider
+ else task.id,
+ object="video",
+ status="queued",
+ model=model,
+ seconds=_seconds_label(seconds),
+ size=_size_from_geometry(
+ resolution if isinstance(resolution, str) else None, ratio if isinstance(ratio, str) else None, model
+ ),
+ usage=dict(usage) if usage else None,
+ )
+
+ @staticmethod
+ def _task_url(video_id: str, api_base: str) -> str:
+ original_task_id: Final = extract_original_video_id(video_id)
+ encoded_task_id: Final = encode_url_path_segment(original_task_id, field_name="video_id")
+ return f"{api_base}{VOLCENGINE_VIDEO_TASKS_PATH}/{encoded_task_id}"
+
+ def transform_video_status_retrieve_request(
+ self,
+ video_id: str,
+ api_base: str,
+ litellm_params: GenericLiteLLMParams,
+ headers: dict[str, str], # mutable-ok: BaseVideoConfig contract
+ ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig contract
+ return self._task_url(video_id, api_base), _EMPTY_PARAMS
+
+ def transform_video_status_retrieve_response(
+ self,
+ raw_response: httpx.Response,
+ logging_obj: "Logging",
+ custom_llm_provider: str | None = None,
+ client: "HTTPHandler | None" = None,
+ ) -> VideoObject:
+ task: Final = _parse(_TASK_ADAPTER, raw_response)
+ model, served_model = polled_video_models(logging_obj)
+ video_obj: Final = _video_object_from_task(task, raw_response, model)
+ if custom_llm_provider:
+ video_obj.id = encode_video_id_with_provider(task.id, custom_llm_provider, model, served_model)
+ return video_obj
+
+ def transform_video_content_request(
+ self,
+ video_id: str,
+ api_base: str,
+ litellm_params: GenericLiteLLMParams,
+ headers: dict[str, str], # mutable-ok: BaseVideoConfig contract
+ variant: str | None = None,
+ ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig contract
+ return self._task_url(video_id, api_base), _EMPTY_PARAMS
+
+ def transform_video_content_response(
+ self,
+ raw_response: httpx.Response,
+ logging_obj: "Logging",
+ client: "HTTPHandler | None" = None,
+ ) -> bytes:
+ video_url: Final = _video_url_from_task(_parse(_TASK_ADAPTER, raw_response), raw_response)
+ httpx_client: Final[HTTPHandler] = client if client is not None else _get_httpx_client()
+ video_response: Final = httpx_client.get(video_url) # pyright: ignore[reportUnknownMemberType] # HTTPHandler.get stub leaves params/headers untyped
+ video_response.raise_for_status()
+ return video_response.content
+
+ async def async_transform_video_content_response(
+ self,
+ raw_response: httpx.Response,
+ logging_obj: "Logging",
+ client: "AsyncHTTPHandler | None" = None,
+ ) -> bytes:
+ video_url: Final = _video_url_from_task(_parse(_TASK_ADAPTER, raw_response), raw_response)
+ async_httpx_client: Final[AsyncHTTPHandler] = (
+ client if client is not None else get_async_httpx_client(llm_provider=litellm.LlmProviders.VOLCENGINE)
+ )
+ video_response: Final = await async_httpx_client.get(video_url) # pyright: ignore[reportUnknownMemberType] # HTTPHandler.get stub leaves params/headers untyped
+ video_response.raise_for_status()
+ return video_response.content
+
+ def transform_video_remix_request(
+ self,
+ video_id: str,
+ prompt: str,
+ api_base: str,
+ litellm_params: GenericLiteLLMParams,
+ headers: dict[str, str], # mutable-ok: BaseVideoConfig contract
+ extra_body: dict[str, object] | None = None, # mutable-ok: BaseVideoConfig contract
+ ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig contract
+ raise _unsupported("", _VIDEO_INPUT_UNSUPPORTED.format(operation="Video remix"))
+
+ def transform_video_remix_response(
+ self,
+ raw_response: httpx.Response,
+ logging_obj: "Logging",
+ custom_llm_provider: str | None = None,
+ ) -> VideoObject:
+ raise NotImplementedError("Video remix is not supported by Volcengine.")
+
+ def transform_video_edit_request(
+ self,
+ prompt: str,
+ video_id: str,
+ api_base: str,
+ litellm_params: GenericLiteLLMParams,
+ headers: dict[str, str], # mutable-ok: BaseVideoConfig contract
+ video_file: object | None = None,
+ extra_body: dict[str, object] | None = None, # mutable-ok: BaseVideoConfig contract
+ prefetched_source_data: dict[str, object] | None = None, # mutable-ok: BaseVideoConfig contract
+ ) -> tuple[str, Mapping[str, object], RequestFiles | None]:
+ raise _unsupported("", _VIDEO_INPUT_UNSUPPORTED.format(operation="Video edit"))
+
+ def transform_video_extension_request(
+ self,
+ prompt: str,
+ video_id: str,
+ seconds: str,
+ api_base: str,
+ litellm_params: GenericLiteLLMParams,
+ headers: dict[str, str], # mutable-ok: BaseVideoConfig contract
+ extra_body: dict[str, object] | None = None, # mutable-ok: BaseVideoConfig contract
+ ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig contract
+ raise _unsupported("", _VIDEO_INPUT_UNSUPPORTED.format(operation="Video extension"))
+
+ def transform_video_list_request(
+ self,
+ api_base: str,
+ litellm_params: GenericLiteLLMParams,
+ headers: dict[str, str], # mutable-ok: BaseVideoConfig contract
+ after: str | None = None,
+ limit: int | None = None,
+ order: str | None = None,
+ extra_query: dict[str, object] | None = None, # mutable-ok: BaseVideoConfig contract
+ ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig contract
+ """Ark pages by page_num/page_size with filters and has no cursor or order."""
+ unsupported: Final = tuple(name for name, value in (("after", after), ("order", order)) if value is not None)
+ if unsupported and not _drops_params(litellm_params):
+ raise _droppable(
+ "", f"Ark's task list pages by page_num in extra_query and has no {' or '.join(unsupported)}"
+ )
+ return f"{api_base}{VOLCENGINE_VIDEO_TASKS_PATH}", {
+ **({"page_size": str(limit)} if limit is not None else {}),
+ **(extra_query or {}),
+ }
+
+ def transform_video_list_response( # pyright: ignore[reportIncompatibleMethodOverride] # base declares dict[str, str] but the payload is a heterogeneous list body
+ self,
+ raw_response: httpx.Response,
+ logging_obj: "Logging",
+ custom_llm_provider: str | None = None,
+ ) -> dict[str, object]: # mutable-ok: OpenAI list body served as JSON by the proxy
+ list_payload: Final = _parse(_TASK_LIST_ADAPTER, raw_response)
+ videos: Final = tuple(_video_object_from_task(task, raw_response, None) for task in list_payload.items)
+ data: Final = [
+ video.model_dump()
+ | (
+ {"id": encode_video_id_with_provider(video.id, custom_llm_provider, video.model)}
+ if custom_llm_provider
+ else {}
+ )
+ for video in videos
+ ]
+ return {
+ "object": "list",
+ "data": data,
+ "total": list_payload.total if list_payload.total is not None else len(data),
+ "has_more": False,
+ **({"first_id": data[0]["id"], "last_id": data[-1]["id"]} if data else {}),
+ }
+
+ def transform_video_delete_request(
+ self,
+ video_id: str,
+ api_base: str,
+ litellm_params: GenericLiteLLMParams,
+ headers: dict[str, str], # mutable-ok: BaseVideoConfig contract
+ ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig contract
+ raise NotImplementedError("Video delete is not supported by Volcengine.")
+
+ def transform_video_delete_response(
+ self,
+ raw_response: httpx.Response,
+ logging_obj: "Logging",
+ ) -> VideoObject:
+ raise NotImplementedError("Video delete is not supported by Volcengine.")
+
+ def get_error_class(
+ self,
+ error_message: str,
+ status_code: int,
+ headers: dict[str, str] | httpx.Headers, # mutable-ok: BaseVideoConfig contract
+ ) -> BaseLLMException:
+ return VolcEngineVideoError(status_code=status_code, message=error_message, headers=headers)
diff --git a/litellm/llms/watsonx/common_utils.py b/litellm/llms/watsonx/common_utils.py
index fb1fa8b9ffb..6d3686e9f85 100644
--- a/litellm/llms/watsonx/common_utils.py
+++ b/litellm/llms/watsonx/common_utils.py
@@ -141,7 +141,7 @@ async def _aconvert_watsonx_messages_core(
model_prompt_dict: Final = custom_prompt_dict[model]
return ptf.custom_prompt(
messages=messages,
- role_dict=model_prompt_dict.get("role_dict", model_prompt_dict.get("roles")),
+ role_dict=model_prompt_dict.get("role_dict", model_prompt_dict.get("roles", {})),
initial_prompt_value=model_prompt_dict.get("initial_prompt_value", ""),
final_prompt_value=model_prompt_dict.get("final_prompt_value", ""),
bos_token=model_prompt_dict.get("bos_token", ""),
@@ -173,7 +173,7 @@ def _convert_watsonx_messages_core(
model_prompt_dict: Final = custom_prompt_dict[model]
return ptf.custom_prompt(
messages=messages,
- role_dict=model_prompt_dict.get("role_dict", model_prompt_dict.get("roles")),
+ role_dict=model_prompt_dict.get("role_dict", model_prompt_dict.get("roles", {})),
initial_prompt_value=model_prompt_dict.get("initial_prompt_value", ""),
final_prompt_value=model_prompt_dict.get("final_prompt_value", ""),
bos_token=model_prompt_dict.get("bos_token", ""),
diff --git a/litellm/llms/xai/realtime/transformation.py b/litellm/llms/xai/realtime/transformation.py
index 5efe125ee60..ffbf814788b 100644
--- a/litellm/llms/xai/realtime/transformation.py
+++ b/litellm/llms/xai/realtime/transformation.py
@@ -16,6 +16,7 @@ construction time (see ``handler.py``) so all normalization is isolated here
and ``RealTimeStreaming`` stays provider-agnostic.
"""
+from collections.abc import Mapping
from typing import Final
@@ -77,6 +78,18 @@ class XAIRealtimeNormalizer:
event = self._normalize_response_usage_event(event, event_type)
return event
+ def gate_turn_detection(self, _turn_detection: object, /) -> Mapping[str, object]:
+ """xAI auto-answers under ``server_vad`` even with ``create_response`` off, so the gate is always untyped."""
+ return {"create_response": False}
+
+ @property
+ def deletes_conversation_items(self) -> bool:
+ """xAI answers ``conversation.item.delete`` with ``Item not found`` until it commits the audio item.
+
+ The guardrail blocks on an interim transcript, before that commit.
+ """
+ return False
+
def patch_outgoing_session(self, session: dict) -> dict:
"""Patch a client ``session.update`` payload before forwarding to xAI.
diff --git a/litellm/main.py b/litellm/main.py
index d8e62ebcdad..e6caa5c25b6 100644
--- a/litellm/main.py
+++ b/litellm/main.py
@@ -52,6 +52,7 @@ from litellm.utils import (
get_litellm_params,
get_optional_params,
peek_reasoning_summary_aliases,
+ reasoning_effort_with_summary_alias,
strip_reasoning_summary_aliases_from_optional_params,
)
@@ -5681,7 +5682,15 @@ def completion(
"allowed_openai_params": allowed_openai_params,
"base_model": base_model,
}
- optional_params = get_optional_params(**optional_param_args, **non_default_params)
+ provider_reasoning_effort: Final = reasoning_effort_with_summary_alias(
+ reasoning_effort,
+ non_default_params,
+ accepts_summary=provider_config is not None
+ and provider_config.accepts_reasoning_summary_in_reasoning_effort,
+ )
+ optional_params = get_optional_params(
+ **{**optional_param_args, "reasoning_effort": provider_reasoning_effort}, **non_default_params
+ )
processed_non_default_params: Final = pre_process_non_default_params(
model=model,
passed_params=optional_param_args,
diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json
index b744032fb6c..49d2ce90af4 100644
--- a/litellm/model_prices_and_context_window_backup.json
+++ b/litellm/model_prices_and_context_window_backup.json
@@ -2773,7 +2773,7 @@
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
- "supports_assistant_prefill": true,
+ "supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
@@ -2808,7 +2808,7 @@
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
- "supports_assistant_prefill": true,
+ "supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
@@ -2843,7 +2843,7 @@
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
- "supports_assistant_prefill": true,
+ "supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
@@ -2878,7 +2878,7 @@
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
- "supports_assistant_prefill": true,
+ "supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
@@ -2913,7 +2913,7 @@
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
- "supports_assistant_prefill": true,
+ "supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
@@ -2948,7 +2948,7 @@
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
- "supports_assistant_prefill": true,
+ "supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
@@ -3743,7 +3743,7 @@
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 1.5e-05,
- "supports_assistant_prefill": true,
+ "supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
@@ -15494,7 +15494,7 @@
},
"supports_adaptive_thinking": true,
"supports_legacy_thinking": true,
- "supports_assistant_prefill": true,
+ "supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
@@ -20832,7 +20832,7 @@
"output_cost_per_token": 2.5000010000000002e-05,
"output_dbu_cost_per_token": 0.000357143,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
- "supports_assistant_prefill": true,
+ "supports_assistant_prefill": false,
"supports_function_calling": true,
"supports_legacy_thinking": true,
"supports_prompt_caching": true,
@@ -21047,7 +21047,7 @@
"output_cost_per_token": 1.5000020000000002e-05,
"output_dbu_cost_per_token": 0.000214286,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
- "supports_assistant_prefill": true,
+ "supports_assistant_prefill": false,
"supports_function_calling": true,
"supports_legacy_thinking": true,
"supports_prompt_caching": true,
@@ -49176,7 +49176,7 @@
"max_tokens": 64000,
"mode": "chat",
"output_cost_per_token": 2.5e-05,
- "supports_assistant_prefill": true,
+ "supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
@@ -51006,7 +51006,7 @@
"mode": "chat",
"output_cost_per_token": 1.5e-05,
"output_cost_per_token_batches": 7.5e-06,
- "supports_assistant_prefill": true,
+ "supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
@@ -51179,7 +51179,7 @@
"supported_regions": [
"global"
],
- "supports_assistant_prefill": true,
+ "supports_assistant_prefill": false,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
@@ -51200,7 +51200,7 @@
"supported_regions": [
"us-central1"
],
- "supports_assistant_prefill": true,
+ "supports_assistant_prefill": false,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
@@ -58822,7 +58822,7 @@
"mode": "chat",
"output_cost_per_token": 1.5e-05,
"output_cost_per_token_batches": 7.5e-06,
- "supports_assistant_prefill": true,
+ "supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
@@ -60467,6 +60467,61 @@
}
]
},
+ "volcengine/doubao-seedance-2-5-260628": {
+ "litellm_provider": "volcengine",
+ "mode": "video_generation",
+ "source": "https://www.volcengine.com/docs/82379/1544106",
+ "output_cost_per_second": 0.216,
+ "output_cost_per_second_480p": 0.096075,
+ "output_cost_per_second_720p": 0.216,
+ "output_cost_per_second_1080p": 0.5346,
+ "supported_modalities": ["text", "image", "video", "audio"],
+ "supported_output_modalities": ["video"],
+ "supported_endpoints": ["/v1/videos"],
+ "provider_specific_entry": {"volcengine_video_frames": {"480p": {"16:9": "854x480", "4:3": "752x560", "1:1": "640x640", "3:4": "560x752", "9:16": "480x854", "21:9": "992x432"}, "720p": {"16:9": "1280x720", "4:3": "1112x834", "1:1": "960x960", "3:4": "834x1112", "9:16": "720x1280", "21:9": "1470x630"}}},
+ "metadata": {"comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s). List 70 CNY per million tokens at 480p and 720p and 77 at 1080p without a video input, matching Ark's per-resolution table and its 1080p 16:9 example of 3.74 CNY/s, converted at 7.0 CNY/USD like the other volcengine entries. Accepts up to 50 reference assets (30 images, 10 videos, 10 audios) and clips of 4-30s; per-asset billing beyond the free allowance is not encoded. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card."}
+ },
+ "volcengine/doubao-seedance-2-0-260128": {
+ "litellm_provider": "volcengine",
+ "mode": "video_generation",
+ "source": "https://www.volcengine.com/docs/82379/1544106",
+ "output_cost_per_second": 0.141942857,
+ "output_cost_per_second_480p": 0.066003429,
+ "output_cost_per_second_720p": 0.141942857,
+ "output_cost_per_second_1080p": 0.354085714,
+ "output_cost_per_second_4k": 0.722057143,
+ "supported_modalities": ["text", "image", "video", "audio"],
+ "supported_output_modalities": ["video"],
+ "supported_endpoints": ["/v1/videos"],
+ "provider_specific_entry": {"volcengine_video_frames": {"480p": {"16:9": "864x496", "4:3": "752x560", "1:1": "640x640", "3:4": "560x752", "9:16": "496x864", "21:9": "992x432"}, "720p": {"16:9": "1280x720", "4:3": "1112x834", "1:1": "960x960", "3:4": "834x1112", "9:16": "720x1280", "21:9": "1470x630"}}},
+ "metadata": {"comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s; the 480p frame is 864x496, 10044 tokens/s, not the tier's nominal 854x480). List 46 CNY per million tokens at 480p and 720p, 51 at 1080p, and 26 at 4k without a video input, matching Ark's per-resolution table and its 1080p and 4k 16:9 examples of 2.48 and 5.05 CNY/s, converted at 7.0 CNY/USD like the other volcengine entries. Clips of 4-15s; 4k is exclusive to this tier. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card."}
+ },
+ "volcengine/doubao-seedance-2-0-fast-260128": {
+ "litellm_provider": "volcengine",
+ "mode": "video_generation",
+ "source": "https://www.volcengine.com/docs/82379/1544106",
+ "output_cost_per_second": 0.114171429,
+ "output_cost_per_second_480p": 0.053089714,
+ "output_cost_per_second_720p": 0.114171429,
+ "supported_modalities": ["text", "image", "video", "audio"],
+ "supported_output_modalities": ["video"],
+ "supported_endpoints": ["/v1/videos"],
+ "provider_specific_entry": {"volcengine_video_frames": {"480p": {"16:9": "864x496", "4:3": "752x560", "1:1": "640x640", "3:4": "560x752", "9:16": "496x864", "21:9": "992x432"}, "720p": {"16:9": "1280x720", "4:3": "1112x834", "1:1": "960x960", "3:4": "834x1112", "9:16": "720x1280", "21:9": "1470x630"}}},
+ "metadata": {"comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s; the 480p frame is 864x496, 10044 tokens/s, not the tier's nominal 854x480). List 37 CNY per million tokens without a video input, converted at 7.0 CNY/USD like the other volcengine entries. Clips of 4-15s at 480p or 720p. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card."}
+ },
+ "volcengine/doubao-seedance-2-0-mini-260615": {
+ "litellm_provider": "volcengine",
+ "mode": "video_generation",
+ "source": "https://www.volcengine.com/docs/82379/1544106",
+ "output_cost_per_second": 0.070971429,
+ "output_cost_per_second_480p": 0.033001714,
+ "output_cost_per_second_720p": 0.070971429,
+ "supported_modalities": ["text", "image"],
+ "supported_output_modalities": ["video"],
+ "supported_endpoints": ["/v1/videos"],
+ "provider_specific_entry": {"volcengine_video_frames": {"480p": {"16:9": "864x496", "4:3": "752x560", "1:1": "640x640", "3:4": "560x752", "9:16": "496x864", "21:9": "992x432"}, "720p": {"16:9": "1280x720", "4:3": "1112x834", "1:1": "960x960", "3:4": "834x1112", "9:16": "720x1280", "21:9": "1470x630"}}},
+ "metadata": {"comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s; the 480p frame is 864x496, 10044 tokens/s, not the tier's nominal 854x480). List 23 CNY per million tokens without a video input, converted at 7.0 CNY/USD like the other volcengine entries. Clips of 4-15s at 480p or 720p. Only text and image inputs are declared; the mini tier's reference video and audio caps were not verified. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card."}
+ },
"bedrock/us-east-1/zai.glm-5": {
"input_cost_per_token": 1e-06,
"output_cost_per_token": 3.2e-06,
@@ -72334,7 +72389,7 @@
"max_output_tokens": 64000,
"max_tokens": 64000,
"mode": "chat",
- "supports_assistant_prefill": true,
+ "supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py
index 4d0eb6a166a..ee1ac79c820 100644
--- a/litellm/proxy/anthropic_endpoints/endpoints.py
+++ b/litellm/proxy/anthropic_endpoints/endpoints.py
@@ -29,6 +29,7 @@ from litellm.proxy.common_request_processing import (
create_response,
log_llm_api_exception,
proxy_exception_from_http_exception,
+ proxy_exception_from_route_error,
resolve_litellm_call_id,
)
from litellm.proxy.common_utils.error_body_call_id import error_body_call_id
@@ -339,10 +340,10 @@ async def count_tokens(
messages: Final = data.get("messages", [])
if not model_name:
- raise HTTPException(status_code=400, detail={"error": "model parameter is required"})
+ raise HTTPException(status_code=400, detail="model parameter is required")
if not messages:
- raise HTTPException(status_code=400, detail={"error": "messages parameter is required"})
+ raise HTTPException(status_code=400, detail="messages parameter is required")
# Create TokenCountRequest for the internal endpoint
from litellm.proxy._types import TokenCountRequest
@@ -360,13 +361,13 @@ async def count_tokens(
messages_format="anthropic",
output_format=_requested_output_format(data),
)
- except HTTPException:
- raise
+ except HTTPException as e:
+ return _anthropic_error_json_response(proxy_exception_from_http_exception(e, {}), request)
except ProxyException as e:
return _anthropic_error_json_response(e, request)
except Exception as e:
log_llm_api_exception(e, litellm_call_id)
- raise HTTPException(status_code=500, detail={"error": f"Internal server error: {e}"})
+ return _anthropic_error_json_response(proxy_exception_from_route_error(e), request)
else:
return {"input_tokens": token_response.total_tokens}
diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py
index 784e3d4719f..d382ff53ffb 100644
--- a/litellm/proxy/google_endpoints/endpoints.py
+++ b/litellm/proxy/google_endpoints/endpoints.py
@@ -1,16 +1,22 @@
from collections.abc import Mapping
from itertools import chain
+from types import MappingProxyType
from typing import Final
from fastapi import APIRouter, Depends, HTTPException, Request, Response
-from fastapi.responses import ORJSONResponse
+from fastapi.responses import JSONResponse, ORJSONResponse
from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError
from litellm.litellm_core_utils.json_validation_rule import normalize_json_schema_types
from litellm.proxy._types import *
-from litellm.proxy._types import TokenCountRequest
+from litellm.proxy._types import ProxyException, TokenCountRequest
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth
-from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
+from litellm.proxy.common_request_processing import (
+ ProxyBaseLLMRequestProcessing,
+ log_llm_api_exception,
+ proxy_exception_from_http_exception,
+ resolve_litellm_call_id,
+)
from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports
_read_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
read_request_body,
@@ -63,6 +69,37 @@ def _input_schema(declaration: _GoogleCountTokensFunctionDeclaration) -> object:
return {"type": "object", "properties": {}}
+_GOOGLE_RPC_STATUS: Final = MappingProxyType(
+ {
+ 400: "INVALID_ARGUMENT",
+ 401: "UNAUTHENTICATED",
+ 403: "PERMISSION_DENIED",
+ 404: "NOT_FOUND",
+ 409: "ABORTED",
+ 429: "RESOURCE_EXHAUSTED",
+ 499: "CANCELLED",
+ 501: "UNIMPLEMENTED",
+ 503: "UNAVAILABLE",
+ 504: "DEADLINE_EXCEEDED",
+ }
+)
+
+
+def _google_error_response(request: Request, error: ProxyException) -> JSONResponse:
+ from litellm.proxy.proxy_server import (
+ _close_dangling_otel_server_span, # pyright: ignore[reportPrivateUsage] # proxy_server keeps the span-close helper private; error JSONResponses returned by the route must stamp the OTel server span like the global ProxyException handler does
+ )
+
+ status_code: Final = int(error.code) if error.code is not None and error.code.isdigit() else 500
+ _close_dangling_otel_server_span(request, status_code, exc=error)
+ status: Final = _GOOGLE_RPC_STATUS.get(status_code, "INTERNAL" if status_code >= 500 else "FAILED_PRECONDITION")
+ return JSONResponse(
+ status_code=status_code,
+ content={"error": {"code": status_code, "message": error.message, "status": status}},
+ headers=error.headers,
+ )
+
+
def _google_count_tokens_request(model_name: str, data: object) -> TokenCountRequest:
from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter
@@ -251,16 +288,26 @@ async def google_count_tokens(request: Request, model_name: str):
data: Final = await read_request_body(request=request)
try:
token_request: Final = _google_count_tokens_request(model_name, data)
- except ValidationError as error:
- raise HTTPException(status_code=400, detail="Invalid Google countTokens request body") from error
+ except ValidationError:
+ return _google_error_response(
+ request,
+ ProxyException(
+ message="Invalid Google countTokens request body", type="invalid_request_error", param=None, code=400
+ ),
+ )
- # Call the internal token counter function with direct request flag set to False
- token_response: Final = await internal_token_counter(
- request=token_request,
- call_endpoint=True,
- )
- if token_response is None:
- return TokenCountDetailsResponse(totalTokens=0, promptTokensDetails=[])
+ try:
+ token_response: Final = await internal_token_counter(request=token_request, call_endpoint=True)
+ except ProxyException as error:
+ return _google_error_response(request, error)
+ except HTTPException as error:
+ return _google_error_response(request, proxy_exception_from_http_exception(error, {}))
+ except Exception as error:
+ log_llm_api_exception(error, resolve_litellm_call_id(request.headers.get("x-litellm-call-id")))
+ return _google_error_response(
+ request,
+ ProxyException(message="Internal server error", type="internal_server_error", param=None, code=500),
+ )
original_response: Final[dict] = token_response.original_response or {}
return TokenCountDetailsResponse(
totalTokens=original_response.get("totalTokens") or token_response.total_tokens or 0,
diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py
index d713a18f470..1d5c530a616 100644
--- a/litellm/proxy/guardrails/guardrail_endpoints.py
+++ b/litellm/proxy/guardrails/guardrail_endpoints.py
@@ -1483,6 +1483,17 @@ async def get_guardrail_info(guardrail_id: str):
raise HTTPException(status_code=500, detail=str(e))
+def _undeclared_logging_only(
+ guardrail_class: type[CustomGuardrail], hooks: Sequence[GuardrailEventHooks]
+) -> tuple[GuardrailEventHooks, ...]:
+ """logging_only when the class accepts it at save time without declaring it in its hooks."""
+ return (
+ (GuardrailEventHooks.logging_only,)
+ if guardrail_class.accepts_logging_only_mode() and GuardrailEventHooks.logging_only not in hooks
+ else ()
+ )
+
+
@router.get(
"/guardrails/ui/add_guardrail_settings",
tags=["Guardrails"],
@@ -1515,7 +1526,7 @@ async def get_guardrail_ui_settings() -> GuardrailUIAddGuardrailSettings:
]
supported_modes_by_provider: Final = {
- provider: [hook.value for hook in hooks]
+ provider: [hook.value for hook in (*hooks, *_undeclared_logging_only(guardrail_class, hooks))]
for provider, guardrail_class in guardrail_class_registry.items()
if (hooks := guardrail_class.get_supported_event_hooks()) is not None
} | MappingProxyType(
diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py
index f959292d567..40e44151e65 100644
--- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py
@@ -6,6 +6,7 @@ to detect and block/mask sensitive content.
"""
import asyncio
+import contextvars
import itertools
import json
import os
@@ -50,6 +51,7 @@ from litellm.types.utils import (
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+ from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.types.guardrails import (
BlockedWord,
@@ -101,6 +103,13 @@ WORD_NUMBER_SEQUENCE_PATTERN: Final = re.compile(
)
WORD_NUMBER_TOKEN_FINDER: Final = re.compile(rf"(?:{WORD_NUMBER_TOKEN_REGEX})", re.IGNORECASE)
+_logging_only_scan_guardrail: Final[contextvars.ContextVar["ContentFilterGuardrail | None"]] = contextvars.ContextVar(
+ "litellm_content_filter_logging_only_scan", default=None
+)
+_logging_only_scan_would_reject: Final[contextvars.ContextVar[bool]] = contextvars.ContextVar(
+ "litellm_content_filter_logging_only_scan_would_reject", default=False
+)
+
class ConditionalCategoryConfig(TypedDict):
identifier_words: Sequence[str]
@@ -1209,6 +1218,18 @@ class ContentFilterGuardrail(CustomGuardrail):
return (keyword, action, description)
return None
+ def _is_logging_only_scan(self) -> bool:
+ return _logging_only_scan_guardrail.get() is self
+
+ def _enforces_block(self, action: ContentFilterAction) -> bool:
+ return action == ContentFilterAction.BLOCK and not self._is_logging_only_scan()
+
+ @staticmethod
+ def _asks_to_block(detection: ContentFilterDetection) -> bool:
+ if detection["type"] == "competitor_intent":
+ return detection["action_hint"] in ("refuse", "reframe")
+ return detection["action"] == ContentFilterAction.BLOCK.value
+
def _handle_conditional_match(
self,
matched_phrase: str,
@@ -1228,7 +1249,7 @@ class ContentFilterGuardrail(CustomGuardrail):
}
detections.append(category_detection)
- if action == ContentFilterAction.BLOCK:
+ if self._enforces_block(action):
error_msg: Final = (
f"Content blocked: {category_name} conditional match '{matched_phrase}' detected (severity: {severity})"
)
@@ -1269,7 +1290,7 @@ class ContentFilterGuardrail(CustomGuardrail):
}
detections.append(category_detection)
- if action == ContentFilterAction.BLOCK:
+ if self._enforces_block(action):
error_msg = f"Content blocked: {category_name} category keyword '{keyword}' detected (severity: {severity})"
verbose_proxy_logger.warning(error_msg)
raise HTTPException(
@@ -1312,7 +1333,7 @@ class ContentFilterGuardrail(CustomGuardrail):
}
detections.append(pattern_detection)
- if action == ContentFilterAction.BLOCK:
+ if self._enforces_block(action):
error_msg: Final = f"Content blocked: {pattern_name} pattern detected"
verbose_proxy_logger.warning(error_msg)
raise HTTPException(
@@ -1346,7 +1367,7 @@ class ContentFilterGuardrail(CustomGuardrail):
}
detections.append(blocked_word_detection)
- if action == ContentFilterAction.BLOCK:
+ if self._enforces_block(action):
error_msg = f"Content blocked: keyword '{keyword}' detected"
if description:
error_msg += f" ({description})"
@@ -1619,7 +1640,7 @@ class ContentFilterGuardrail(CustomGuardrail):
- Number of entities masked
- Whether the guardrail blocked the request (max risk)
"""
- if status == "guardrail_intervened":
+ if status in ("guardrail_intervened", "guardrail_flagged"):
return 10.0
total_masked: Final = sum(masked_entity_count.values()) if masked_entity_count else 0
@@ -1665,6 +1686,8 @@ class ContentFilterGuardrail(CustomGuardrail):
}
detections.append(detection)
+ if self._is_logging_only_scan():
+ return
if action_hint_val == "refuse":
msg = "Content blocked: competitor comparison or ranking intent detected."
if self._competitor_intent_checker and getattr(
@@ -1763,18 +1786,18 @@ class ContentFilterGuardrail(CustomGuardrail):
return raw_name
return None
+ def _reject(self, error: str) -> None:
+ """Raise the 400 an enforcing mode returns, or note it on the leg a logging_only scan is observing."""
+ if self._is_logging_only_scan():
+ _logging_only_scan_would_reject.set(True)
+ return
+ raise HTTPException(status_code=400, detail={"error": error})
+
def _assert_argument_label_clean(
self, text: str, detections: list[ContentFilterDetection], context_label: str
) -> None:
if self._filter_single_text(text, detections=detections) != text:
- raise HTTPException(
- status_code=400,
- detail={
- "error": (
- f"Content blocked: {context_label} argument matched a masking rule on a non-rewritable field"
- )
- },
- )
+ self._reject(f"Content blocked: {context_label} argument matched a masking rule on a non-rewritable field")
def _filter_argument_value(
self,
@@ -1784,10 +1807,8 @@ class ContentFilterGuardrail(CustomGuardrail):
depth: int = 0,
) -> object:
if depth > DEFAULT_MAX_RECURSE_DEPTH:
- raise HTTPException(
- status_code=400,
- detail={"error": f"Content blocked: {context_label} arguments exceed the maximum nesting depth"},
- )
+ self._reject(f"Content blocked: {context_label} arguments exceed the maximum nesting depth")
+ return value
if isinstance(value, str):
return self._filter_single_text(value, detections=detections)
if isinstance(value, (int, float)) and not isinstance(value, bool):
@@ -1874,6 +1895,21 @@ class ContentFilterGuardrail(CustomGuardrail):
if filtered_arguments != arguments:
self._set_tool_call_arguments(tool_call, filtered_arguments)
+ async def _scan_logged_call(
+ self,
+ kwargs: dict[str, object], # mutable-ok: CustomLogger.async_logging_hook contract
+ response: object | None,
+ translation: "BaseTranslation",
+ output_translation: "BaseTranslation",
+ scratch_metadata: dict[str, object], # mutable-ok: apply_guardrail records its verdict into request metadata
+ ) -> None:
+ """Scan the logged call with every BLOCK match recorded instead of raised."""
+ token: Final = _logging_only_scan_guardrail.set(self)
+ try:
+ await super()._scan_logged_call(kwargs, response, translation, output_translation, scratch_metadata)
+ finally:
+ _logging_only_scan_guardrail.reset(token)
+
async def apply_guardrail(
self,
inputs: "GenericGuardrailAPIInputs",
@@ -1904,6 +1940,7 @@ class ContentFilterGuardrail(CustomGuardrail):
masked_entity_count: Final[dict[str, int]] = {}
status: GuardrailStatus = "success"
exception_str: str = ""
+ would_reject_token: Final = _logging_only_scan_would_reject.set(False)
try:
texts: Final = inputs.get("texts", [])
@@ -1947,11 +1984,13 @@ class ContentFilterGuardrail(CustomGuardrail):
exception_str = str(e)
raise e
finally:
+ would_reject: Final = _logging_only_scan_would_reject.get() or any(map(self._asks_to_block, detections))
+ _logging_only_scan_would_reject.reset(would_reject_token)
# Log guardrail information
self._log_guardrail_information(
request_data=request_data,
detections=detections,
- status=status,
+ status="guardrail_flagged" if status == "success" and would_reject else status,
start_time=start_time,
masked_entity_count=masked_entity_count,
exception_str=exception_str,
diff --git a/litellm/proxy/hooks/key_management_event_hooks.py b/litellm/proxy/hooks/key_management_event_hooks.py
index ef93809d47b..ef60a800cdd 100644
--- a/litellm/proxy/hooks/key_management_event_hooks.py
+++ b/litellm/proxy/hooks/key_management_event_hooks.py
@@ -57,13 +57,6 @@ class KeyManagementEventHooks:
)
from litellm.proxy.proxy_server import litellm_proxy_admin_name
- # Send email notification - non-blocking, independent operation
- if data.send_invite_email is True:
- try:
- await KeyManagementEventHooks._send_key_created_email(response.model_dump(exclude_none=True))
- except Exception as e:
- verbose_proxy_logger.warning("Failed to send key created email: %s", e)
-
if is_audit_logging_enabled():
_updated_values: Final = response.model_dump_json(exclude_none=True)
asyncio.create_task(
@@ -86,6 +79,13 @@ class KeyManagementEventHooks:
)
)
+ # Send email notification - non-blocking, independent operation
+ if data.send_invite_email is True:
+ try:
+ await KeyManagementEventHooks._send_key_created_email(response.model_dump(exclude_none=True))
+ except Exception as e:
+ verbose_proxy_logger.warning("Failed to send key created email: %s", e)
+
# Store the generated key in the secret manager - non-blocking, independent operation
try:
await KeyManagementEventHooks._store_virtual_key_in_secret_manager(
@@ -118,16 +118,6 @@ class KeyManagementEventHooks:
)
from litellm.proxy.proxy_server import litellm_proxy_admin_name
- if data.key_alias is not None and data.key_alias != existing_key_row.key_alias:
- try:
- await KeyManagementEventHooks._rename_virtual_key_in_secret_manager(
- current_secret_name=existing_key_row.key_alias or f"virtual-key-{existing_key_row.token}",
- new_secret_name=data.key_alias,
- team_id=existing_key_row.team_id,
- )
- except Exception as e:
- verbose_proxy_logger.warning("Failed to rename virtual key in secret manager: %s", e)
-
if is_audit_logging_enabled():
updated_fields: Final = {
**data.model_dump(exclude_none=True),
@@ -156,6 +146,16 @@ class KeyManagementEventHooks:
)
asyncio.create_task(create_audit_log_for_update(request_data=request_data))
+ if data.key_alias is not None and data.key_alias != existing_key_row.key_alias:
+ try:
+ await KeyManagementEventHooks._rename_virtual_key_in_secret_manager(
+ current_secret_name=existing_key_row.key_alias or f"virtual-key-{existing_key_row.token}",
+ new_secret_name=data.key_alias,
+ team_id=existing_key_row.team_id,
+ )
+ except Exception as e:
+ verbose_proxy_logger.warning("Failed to rename virtual key in secret manager: %s", e)
+
@staticmethod
async def async_key_rotated_hook(
data: RegenerateKeyRequest | None,
@@ -171,6 +171,28 @@ class KeyManagementEventHooks:
)
from litellm.proxy.proxy_server import litellm_proxy_admin_name
+ # store the audit log
+ if is_audit_logging_enabled() and existing_key_row.token is not None:
+ asyncio.create_task(
+ create_audit_log_for_update(
+ request_data=LiteLLM_AuditLogs(
+ id=str(uuid.uuid4()),
+ updated_at=datetime.now(timezone.utc),
+ changed_by=get_audit_log_changed_by(
+ litellm_changed_by=litellm_changed_by,
+ user_api_key_dict=user_api_key_dict,
+ litellm_proxy_admin_name=litellm_proxy_admin_name,
+ ),
+ changed_by_api_key=user_api_key_dict.token,
+ table_name=LitellmTableNames.KEY_TABLE_NAME,
+ object_id=existing_key_row.token,
+ action="rotated",
+ updated_values=response.model_dump_json(exclude_none=True),
+ before_value=existing_key_row.model_dump_json(exclude_none=True),
+ )
+ )
+ )
+
# Store the generated key in the secret manager - non-blocking, independent operation
if response.token_id is not None:
try:
@@ -204,28 +226,6 @@ class KeyManagementEventHooks:
except Exception as e:
verbose_proxy_logger.warning("Failed to send key rotated email: %s", e)
- # store the audit log
- if is_audit_logging_enabled() and existing_key_row.token is not None:
- asyncio.create_task(
- create_audit_log_for_update(
- request_data=LiteLLM_AuditLogs(
- id=str(uuid.uuid4()),
- updated_at=datetime.now(timezone.utc),
- changed_by=get_audit_log_changed_by(
- litellm_changed_by=litellm_changed_by,
- user_api_key_dict=user_api_key_dict,
- litellm_proxy_admin_name=litellm_proxy_admin_name,
- ),
- changed_by_api_key=user_api_key_dict.token,
- table_name=LitellmTableNames.KEY_TABLE_NAME,
- object_id=existing_key_row.token,
- action="rotated",
- updated_values=response.model_dump_json(exclude_none=True),
- before_value=existing_key_row.model_dump_json(exclude_none=True),
- )
- )
- )
-
@staticmethod
async def async_key_deleted_hook(
data: KeyRequest,
diff --git a/litellm/proxy/hooks/user_management_event_hooks.py b/litellm/proxy/hooks/user_management_event_hooks.py
index 6d978929c05..6618af70c0a 100644
--- a/litellm/proxy/hooks/user_management_event_hooks.py
+++ b/litellm/proxy/hooks/user_management_event_hooks.py
@@ -23,6 +23,7 @@ from litellm.proxy._types import (
from litellm.proxy.management_helpers.audit_logs import (
create_audit_log_for_update,
is_audit_logging_enabled,
+ track_audit_task,
)
from litellm.repositories.user_repository import UserRepository
@@ -41,7 +42,18 @@ class UserManagementEventHooks:
- Creating an audit log for the user creation
- Sending a user invitation email to the user
"""
- from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client
+ from litellm.proxy.proxy_server import litellm_proxy_admin_name
+
+ if is_audit_logging_enabled():
+ track_audit_task(
+ asyncio.create_task(
+ UserManagementEventHooks._create_user_created_audit_log(
+ response=response,
+ user_api_key_dict=user_api_key_dict,
+ litellm_proxy_admin_name=litellm_proxy_admin_name,
+ )
+ )
+ )
#########################################################
########## Send User Invitation Email ################
@@ -52,9 +64,14 @@ class UserManagementEventHooks:
user_api_key_dict=user_api_key_dict,
)
- #########################################################
- ########## CREATE AUDIT LOG ################
- #########################################################
+ @staticmethod
+ async def _create_user_created_audit_log(
+ response: NewUserResponse,
+ user_api_key_dict: UserAPIKeyAuth,
+ litellm_proxy_admin_name: str | None,
+ ) -> None:
+ from litellm.proxy.proxy_server import prisma_client
+
try:
if prisma_client is None:
raise Exception(CommonProxyErrors.db_not_connected_error.value)
@@ -63,16 +80,14 @@ class UserManagementEventHooks:
user_row: Final = await UserRepository(prisma_client).find_by_id(response.user_id)
if user_row is None:
raise Exception(f"no user row found for user_id={response.user_id}")
- asyncio.create_task(
- UserManagementEventHooks.create_internal_user_audit_log(
- user_id=user_row.user_id,
- action="created",
- litellm_changed_by=user_api_key_dict.user_id,
- user_api_key_dict=user_api_key_dict,
- litellm_proxy_admin_name=litellm_proxy_admin_name,
- before_value=None,
- after_value=user_row.model_dump_json(exclude_none=True),
- )
+ await UserManagementEventHooks.create_internal_user_audit_log(
+ user_id=user_row.user_id,
+ action="created",
+ litellm_changed_by=user_api_key_dict.user_id,
+ user_api_key_dict=user_api_key_dict,
+ litellm_proxy_admin_name=litellm_proxy_admin_name,
+ before_value=None,
+ after_value=user_row.model_dump_json(exclude_none=True),
)
except Exception as e:
verbose_proxy_logger.warning("Unable to create audit log for user on `/user/new` - %s", e)
diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py
index 6bd31a42e9e..675bf0830f9 100644
--- a/litellm/proxy/litellm_pre_call_utils.py
+++ b/litellm/proxy/litellm_pre_call_utils.py
@@ -27,6 +27,7 @@ from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm._service_logger import ServiceLogging
from litellm._uuid import uuid
from litellm.constants import (
+ CALLER_SESSION_ID_METADATA_KEY,
CLIENT_OUTPUT_CEILING_METADATA_KEY,
CONSUMED_REQUEST_TAGS_METADATA_KEY,
INTERNAL_CALL_ORIGIN_METADATA_KEY,
@@ -187,6 +188,30 @@ def _caller_trace_field(data: Mapping[str, object], metadata_variable_name: str,
return requester_value if isinstance(requester_value, str) and requester_value else None
+def _caller_body_session_id(data: Mapping[str, object], metadata_variable_name: str) -> str | None:
+ top_level: Final = data.get("litellm_session_id")
+ if isinstance(top_level, str) and top_level:
+ return top_level
+ return _caller_trace_field(data, metadata_variable_name, "session_id")
+
+
+def _caller_session_beside_trace_header(
+ normalized_headers: Mapping[str, str], data: Mapping[str, object], metadata_variable_name: str
+) -> str | None:
+ session_headers: Final = {key: value for key, value in normalized_headers.items() if key != "x-litellm-trace-id"}
+ return get_chain_id_from_headers(session_headers) or _caller_body_session_id(data, metadata_variable_name)
+
+
+def _drop_client_caller_session(request_body: Mapping[str, object]) -> None:
+ for metadata_key in ("metadata", "litellm_metadata"):
+ if isinstance(metadata := request_body.get(metadata_key), dict):
+ metadata.pop(CALLER_SESSION_ID_METADATA_KEY, None)
+
+
+def _requester_metadata_snapshot(metadata: Mapping[str, object]) -> Mapping[str, object]:
+ return copy.deepcopy({key: value for key, value in metadata.items() if key != CALLER_SESSION_ID_METADATA_KEY})
+
+
def _stampable_key_hash(user_api_key_dict: UserAPIKeyAuth) -> str | None:
"""Only proxy-validated keys are stamped, proven by the unforgeable
via_virtual_key marker AND a known non-secret shape: the sha256 hex digest
@@ -1623,6 +1648,8 @@ class LiteLLMProxyRequestSetup:
from litellm.proxy._types import LitellmMetadataFromRequestHeaders
metadata_from_headers: Final = LitellmMetadataFromRequestHeaders()
+ request_body: Final = cast(Mapping[str, object], data) # cast-ok: request body is a str-keyed JSON object
+ _drop_client_caller_session(request_body)
spend_logs_metadata: Final = LiteLLMProxyRequestSetup._get_spend_logs_metadata_from_request_headers(headers)
if spend_logs_metadata is not None:
metadata_from_headers["spend_logs_metadata"] = spend_logs_metadata
@@ -1631,6 +1658,9 @@ class LiteLLMProxyRequestSetup:
# Finally update the requests metadata with the `metadata_from_headers`
#########################################################################################
+ normalized_headers: Final = MappingProxyType(
+ {k.lower(): v for k, v in headers.items() if isinstance(k, str) and isinstance(v, str)}
+ )
agent_id_from_header: Final = headers.get("x-litellm-agent-id")
# Explicit litellm headers take precedence; fall back to any x-*-session-id header.
chain_id: Final = get_chain_id_from_headers(dict(headers))
@@ -1639,6 +1669,13 @@ class LiteLLMProxyRequestSetup:
metadata_from_headers["agent_id"] = agent_id_from_header
verbose_proxy_logger.debug("Extracted agent_id from header: %s", agent_id_from_header)
+ spend_log_session_id: Final = (
+ _caller_session_beside_trace_header(normalized_headers, request_body, _metadata_variable_name)
+ if litellm.request_correlation_in_logs
+ and chain_id is not None
+ and chain_id == normalized_headers.get("x-litellm-trace-id")
+ else None
+ )
if chain_id:
metadata_from_headers["trace_id"] = chain_id
metadata_from_headers["session_id"] = chain_id
@@ -1663,7 +1700,6 @@ class LiteLLMProxyRequestSetup:
# - but lets a caller's existing traceparent/baggage headers (from
# real OTel instrumentation) correlate with litellm's own logs instead
# of generating an unrelated trace_id.
- normalized_headers: Final = MappingProxyType({k.lower(): v for k, v in headers.items() if isinstance(k, str)})
if "litellm_trace_id" not in data:
traceparent: Final = normalized_headers.get("traceparent")
if isinstance(traceparent, str):
@@ -1689,11 +1725,7 @@ class LiteLLMProxyRequestSetup:
baggage: Final = normalized_headers.get("baggage")
if isinstance(baggage, str):
session_id_from_baggage: Final = _session_id_from_baggage(baggage)
- caller_session_id: Final = _caller_trace_field(
- cast(Mapping[str, object], data), # cast-ok: request body is a str-keyed JSON object
- _metadata_variable_name,
- "session_id",
- )
+ caller_session_id: Final = _caller_trace_field(request_body, _metadata_variable_name, "session_id")
if session_id_from_baggage and (
caller_session_id is None or caller_session_id == session_id_from_baggage
):
@@ -1702,7 +1734,11 @@ class LiteLLMProxyRequestSetup:
verbose_proxy_logger.debug("Extracted session_id from W3C baggage header")
if isinstance(data[_metadata_variable_name], dict):
- data[_metadata_variable_name].update(metadata_from_headers)
+ data[_metadata_variable_name].update(
+ {**metadata_from_headers, CALLER_SESSION_ID_METADATA_KEY: spend_log_session_id}
+ if spend_log_session_id
+ else metadata_from_headers
+ )
return data
@staticmethod
@@ -2417,7 +2453,7 @@ async def add_litellm_data_to_request(
# them — from leaking into requester_metadata where guardrails and audit
# paths may read from it.
if "metadata" in data and isinstance(data["metadata"], dict):
- data[_metadata_variable_name]["requester_metadata"] = copy.deepcopy(data["metadata"])
+ data[_metadata_variable_name]["requester_metadata"] = _requester_metadata_snapshot(data["metadata"])
if _metadata_variable_name == "litellm_metadata":
data[_metadata_variable_name].update(
_promoted_trace_control_fields(
diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py
index 7877d887409..592c732bfd3 100644
--- a/litellm/proxy/management_endpoints/key_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/key_management_endpoints.py
@@ -7713,6 +7713,11 @@ async def block_key(
code=status.HTTP_404_NOT_FOUND,
)
+ record: Final = await _prisma_table(VerificationTokenRepository(prisma_client)).update(
+ where={"token": hashed_token},
+ data=with_settings_updated_at({"blocked": True}),
+ )
+
if is_audit_logging_enabled():
asyncio.create_task(
create_audit_log_for_update(
@@ -7734,11 +7739,6 @@ async def block_key(
)
)
- record: Final = await _prisma_table(VerificationTokenRepository(prisma_client)).update(
- where={"token": hashed_token},
- data=with_settings_updated_at({"blocked": True}),
- )
-
## UPDATE KEY CACHE - invalidate so next read re-fetches from DB
await delete_cache_key_object(
hashed_token=hashed_token,
@@ -7827,6 +7827,11 @@ async def unblock_key(
code=status.HTTP_404_NOT_FOUND,
)
+ record: Final = await _prisma_table(VerificationTokenRepository(prisma_client)).update(
+ where={"token": hashed_token},
+ data=with_settings_updated_at({"blocked": False}),
+ )
+
if is_audit_logging_enabled():
asyncio.create_task(
create_audit_log_for_update(
@@ -7848,11 +7853,6 @@ async def unblock_key(
)
)
- record: Final = await _prisma_table(VerificationTokenRepository(prisma_client)).update(
- where={"token": hashed_token},
- data=with_settings_updated_at({"blocked": False}),
- )
-
## UPDATE KEY CACHE - invalidate so next read re-fetches from DB
await delete_cache_key_object(
hashed_token=hashed_token,
diff --git a/litellm/proxy/management_helpers/audit_logs.py b/litellm/proxy/management_helpers/audit_logs.py
index ecd6abea3c3..ef95e8e2c0f 100644
--- a/litellm/proxy/management_helpers/audit_logs.py
+++ b/litellm/proxy/management_helpers/audit_logs.py
@@ -5,7 +5,7 @@ Functions to create audit logs for LiteLLM Proxy
import asyncio
import json
from datetime import datetime, timezone
-from typing import Final
+from typing import Final, TypeVar
import litellm
from litellm._logging import verbose_proxy_logger
@@ -134,6 +134,56 @@ def _audit_log_task_done_callback(task: asyncio.Task) -> None:
verbose_proxy_logger.error("Audit log callback task failed: %s", exc, exc_info=exc)
+_T: Final = TypeVar("_T")
+
+
+class AuditTaskRegistry:
+ __slots__ = ("_pending_tasks",)
+
+ def __init__(self) -> None:
+ self._pending_tasks: Final[set[asyncio.Task[object]]] = set() # mutable-ok: background task tracking
+
+ def track(self, task: asyncio.Task[_T]) -> asyncio.Task[_T]:
+ if task.done() or task in self._pending_tasks:
+ return task
+ self._pending_tasks.add(task)
+ task.add_done_callback(self._on_task_done)
+ return task
+
+ def _on_task_done(self, task: asyncio.Task[object]) -> None:
+ self._pending_tasks.discard(task)
+ _audit_log_task_done_callback(task)
+
+ async def drain(self, timeout: float = 5.0) -> None:
+ loop: Final = asyncio.get_running_loop()
+ deadline: Final = loop.time() + timeout
+ while self._pending_tasks:
+ _, pending = await asyncio.wait(tuple(self._pending_tasks), timeout=max(deadline - loop.time(), 0))
+ if pending:
+ self._pending_tasks.difference_update(pending)
+ verbose_proxy_logger.warning(
+ "Timed out draining %d audit tasks on shutdown, dropping: %s",
+ len(pending),
+ ", ".join(sorted(task.get_name() for task in pending)),
+ )
+ return
+
+
+_audit_task_registry: Final = AuditTaskRegistry()
+
+
+def track_audit_task(task: asyncio.Task[_T]) -> asyncio.Task[_T]:
+ return _audit_task_registry.track(task)
+
+
+async def drain_audit_tasks(timeout: float = 5.0) -> None:
+ await _audit_task_registry.drain(timeout=timeout)
+
+
+def _audit_task_name(request_data: LiteLLM_AuditLogs, destination: str) -> str:
+ return f"{request_data.action} {request_data.table_name.value} {request_data.object_id} to {destination}"
+
+
async def _dispatch_audit_log_to_callbacks(
request_data: LiteLLM_AuditLogs,
) -> None:
@@ -153,8 +203,12 @@ async def _dispatch_audit_log_to_callbacks(
continue
if isinstance(resolved, CustomLogger):
- task = asyncio.create_task(resolved.async_log_audit_log_event(payload))
- task.add_done_callback(_audit_log_task_done_callback)
+ track_audit_task(
+ asyncio.create_task(
+ resolved.async_log_audit_log_event(payload),
+ name=_audit_task_name(request_data, type(resolved).__name__),
+ )
+ )
except Exception as e:
verbose_proxy_logger.error("Failed dispatching audit log to callback: %s", e)
@@ -204,6 +258,14 @@ async def create_object_audit_log(
)
+async def _insert_audit_log_row(repository: AuditLogRepository, request_data: LiteLLM_AuditLogs) -> None:
+ try:
+ await repository.table.create(data={**request_data.model_dump(exclude_none=True)})
+ except Exception as e:
+ # [Non-Blocking Exception. Do not allow blocking LLM API call]
+ verbose_proxy_logger.error("Failed Creating audit log %s", e)
+
+
async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs):
"""
Create an audit log for an object.
@@ -231,14 +293,9 @@ async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs):
verbose_proxy_logger.error("prisma_client is None, cannot write audit log to DB")
return
- _request_data: Final = request_data.model_dump(exclude_none=True)
-
- try:
- await AuditLogRepository(prisma_client).table.create(
- data={
- **_request_data,
- }
+ await track_audit_task(
+ asyncio.create_task(
+ _insert_audit_log_row(AuditLogRepository(prisma_client), request_data),
+ name=_audit_task_name(request_data, "the database"),
)
- except Exception as e:
- # [Non-Blocking Exception. Do not allow blocking LLM API call]
- verbose_proxy_logger.error("Failed Creating audit log %s", e)
+ )
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 97d8eda1fad..b56f0e68c4f 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -1200,6 +1200,20 @@ async def _flush_spend_logs_queue_on_shutdown() -> None:
verbose_proxy_logger.exception("Error flushing spend logs queue on shutdown: %s", e)
+async def _drain_audit_tasks_on_shutdown() -> None:
+ try:
+ from litellm.proxy.management_helpers.audit_logs import drain_audit_tasks
+
+ await drain_audit_tasks()
+ except Exception as e: # noqa: BLE001 # shutdown must continue even if the drain fails
+ verbose_proxy_logger.exception("Error draining audit tasks on shutdown: %s", e)
+
+
+async def _flush_request_accumulators(prisma: PrismaClient) -> None:
+ await flush_gateway_requests(prisma, gateway_request_accumulator)
+ await flush_request_errors(prisma, request_error_accumulator)
+
+
async def proxy_shutdown_event(worker_heartbeat: ProxyWorkerHeartbeat | None = None) -> None:
global prisma_client, master_key, user_custom_auth, user_custom_key_generate, user_custom_key_update
verbose_proxy_logger.info("Shutting down LiteLLM Proxy Server")
@@ -1210,10 +1224,11 @@ async def proxy_shutdown_event(worker_heartbeat: ProxyWorkerHeartbeat | None = N
# is lost, and a write attempted after disconnect raises
# ClientNotConnectedError rather than persisting anything. Ordering this
# inside the same guard is what keeps the two from drifting apart.
- await flush_gateway_requests(prisma_client, gateway_request_accumulator)
- await flush_request_errors(prisma_client, request_error_accumulator)
+ await asyncio.gather(_drain_audit_tasks_on_shutdown(), _flush_request_accumulators(prisma_client))
verbose_proxy_logger.debug("Disconnecting from Prisma")
await prisma_client.disconnect()
+ else:
+ await _drain_audit_tasks_on_shutdown()
await drain_post_call_redis_batches()
if litellm.cache is not None:
@@ -1509,7 +1524,6 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState
await asyncio.sleep(5)
asyncio.create_task(_run_agent_grant_id_migration())
- await install_background_interaction_settlement(prisma_client)
## A coordination_redis block saved from the admin UI lives in the database,
## which is only reachable once the prisma client exists. Apply it here, before
@@ -1636,6 +1650,10 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState
## SYNC UI SETTINGS ##
await ProxyStartupEvent.sync_ui_settings_to_general_settings()
+ await install_background_interaction_settlement(
+ prisma_client, router=lambda: llm_router, spend_logs_kept=lambda: disable_spend_logs is False
+ )
+
# Start background health checks AFTER models are loaded and index is built
if use_background_health_checks:
asyncio.create_task(_run_background_health_check()) # start the background health check coroutine.
@@ -14473,6 +14491,16 @@ async def _try_provider_token_count(
return result
+_PROVIDERS_COUNTING_WITHOUT_MESSAGES: Final = frozenset(
+ {litellm.LlmProviders.GEMINI.value, litellm.LlmProviders.VERTEX_AI.value}
+)
+
+
+def _drops_params(deployment: Mapping[str, object] | None) -> bool:
+ litellm_params: Final = as_str_mapping(None if deployment is None else deployment.get("litellm_params"))
+ return litellm.drop_params is True or (litellm_params is not None and litellm_params.get("drop_params") is True)
+
+
def _system_message(system: object) -> ChatCompletionSystemMessage | None:
if not isinstance(system, (str, list)) or not system:
return None
@@ -14539,8 +14567,8 @@ async def count_request_tokens(
litellm_model_name = litellm_model_name.split("/", 1)[1]
model_to_use: str = (
- litellm_model_name or request.model
- ) # use litellm model name, if it's not avalable then fallback to request.model
+ litellm_model_name if isinstance(litellm_model_name, str) and litellm_model_name else request.model
+ )
# Try provider-specific token counting first - only for non-direct requests (from provider endpoints)
provider_counter: BaseTokenCounter | None = None
@@ -14551,13 +14579,23 @@ async def count_request_tokens(
if _model is not None:
model_to_use = _model
+ conversation_empty: Final = not messages and not contents
+ if prompt and not conversation_empty and not _drops_params(deployment):
+ raise ProxyException(
+ message="prompt cannot be sent with messages or contents. Set drop_params to count them without the prompt",
+ type="invalid_request_error",
+ param="prompt",
+ code=400,
+ )
+ prompt_only: Final = bool(prompt) and conversation_empty
if provider_counter is not None:
+ prompt_as_message: Final = prompt_only and custom_llm_provider in _PROVIDERS_COUNTING_WITHOUT_MESSAGES
result: Final = await _try_provider_token_count(
provider_counter=provider_counter,
custom_llm_provider=custom_llm_provider,
model_to_use=model_to_use,
- messages=messages,
- contents=contents,
+ messages=[{"role": "user", "content": prompt}] if prompt_as_message else messages,
+ contents=None if prompt_as_message else contents,
deployment=deployment,
request_model=request.model,
tools=tools,
@@ -14590,8 +14628,10 @@ async def count_request_tokens(
tokenizer_used: Final = str(_tokenizer_used["type"])
system_message: Final = _system_message(system)
+ prompt_turn: Final = prompt_only and not (system_message is None and not tools)
typed_messages: Final = cast( # cast-ok: request messages are raw chat-shaped dicts that token_counter normalizes
- Sequence[AllMessageValues] | None, messages
+ Sequence[AllMessageValues] | None,
+ [{"role": "user", "content": prompt}] if prompt_turn else (None if prompt_only else messages),
)
counted_messages: Final = (
typed_messages if typed_messages is None or system_message is None else (system_message, *typed_messages)
@@ -14601,7 +14641,7 @@ async def count_request_tokens(
)
total_tokens: Final = await offload_token_count(litellm.token_counter)(
model=model_to_use,
- text=prompt,
+ text=prompt if counted_messages is None else None,
messages=counted_messages,
tools=counted_tools,
custom_tokenizer=_tokenizer_used,
@@ -19658,6 +19698,26 @@ async def delete_config_general_settings(
return response
+def _configured_callback_names(value: object) -> tuple[JsonValue, ...]:
+ if isinstance(value, str):
+ return (value,)
+ if isinstance(value, list):
+ return tuple(cast(Iterable[JsonValue], value)) # cast-ok: config.yaml values are untyped
+ return ()
+
+
+def _is_callback_name(entry: object, callback_name: str) -> bool:
+ return isinstance(entry, str) and entry.lower() == callback_name
+
+
+def _callback_lists_to_rewrite(
+ matching_keys: tuple[str, ...], owned_by_config: Callable[[str], bool]
+) -> tuple[str, ...]:
+ """The matching lists the database owns, or all of them when the config file owns every one."""
+ database_owned: Final = tuple(key for key in matching_keys if not owned_by_config(key))
+ return database_owned or matching_keys
+
+
@router.post(
"/config/callback/delete",
tags=["config.yaml"],
@@ -19693,24 +19753,42 @@ async def delete_callback(
try:
# Get current configuration
- config: Final = await proxy_config.get_config()
+ config: Final = cast( # cast-ok: get_config returns an untyped config dict
+ dict[str, object],
+ await proxy_config.get_config(),
+ )
callback_name: Final = data.callback_name.lower()
# Check if callback exists in current configuration
- litellm_settings: Final = config.get("litellm_settings", {})
- success_callbacks: Final = litellm_settings.get("success_callback", [])
+ litellm_settings: Final[Mapping[str, object]] = cast( # cast-ok: the config dict's values are untyped
+ Mapping[str, object],
+ config.get("litellm_settings", {}),
+ )
+ configured_lists: Final = {
+ key: _configured_callback_names(litellm_settings.get(key)) for key in _DB_CONFIG_CALLBACK_EVENT_TYPES
+ }
+ matching_keys: Final = tuple(
+ key
+ for key, names in configured_lists.items()
+ if any(_is_callback_name(entry, callback_name) for entry in names)
+ )
- if callback_name not in success_callbacks:
+ if not matching_keys:
raise HTTPException(
status_code=404,
detail={"error": f"Callback '{callback_name}' not found in active configuration"},
)
- before_success_callbacks: Final = list(success_callbacks)
-
- # Remove callback from success_callback list
- success_callbacks.remove(callback_name)
- config.setdefault("litellm_settings", {})["success_callback"] = success_callbacks
+ rewritten_keys: Final = _callback_lists_to_rewrite(matching_keys, proxy_config.litellm_settings.owned_by_config)
+ remaining_lists: Final = MappingProxyType(
+ {
+ key: tuple(entry for entry in configured_lists[key] if not _is_callback_name(entry, callback_name))
+ for key in rewritten_keys
+ }
+ )
+ before_callbacks: Final[JsonValue] = {key: list(configured_lists[key]) for key in rewritten_keys}
+ after_callbacks: Final[JsonValue] = {key: list(names) for key, names in remaining_lists.items()}
+ config["litellm_settings"] = {**litellm_settings, **after_callbacks}
# Save the updated configuration
await proxy_config.save_config(new_config=config)
@@ -19719,8 +19797,8 @@ async def delete_callback(
create_config_audit_log(
"litellm_settings",
"deleted",
- {"success_callback": before_success_callbacks},
- {"success_callback": success_callbacks},
+ before_callbacks,
+ after_callbacks,
user_api_key_dict,
)
)
@@ -19731,7 +19809,7 @@ async def delete_callback(
return {
"message": f"Successfully deleted callback: {callback_name}",
"removed_callback": callback_name,
- "remaining_callbacks": success_callbacks,
+ "remaining_callbacks": tuple(chain.from_iterable(remaining_lists.values())),
"deleted_at": datetime.now().isoformat(),
}
diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma
index d40bac380aa..83a0e6f9e05 100644
--- a/litellm/proxy/schema.prisma
+++ b/litellm/proxy/schema.prisma
@@ -1981,6 +1981,7 @@ model LiteLLM_BackgroundInteractionSettlement {
outcome String?
@@index([claimed_at], map: "idx_background_interaction_settlement_claimed_at")
+ @@index([settled_at], map: "idx_background_interaction_settlement_settled_at")
}
model LiteLLM_Lens {
diff --git a/litellm/proxy/spend_tracking/background_interaction_settlement.py b/litellm/proxy/spend_tracking/background_interaction_settlement.py
index 165f813cd50..ffa8185619c 100644
--- a/litellm/proxy/spend_tracking/background_interaction_settlement.py
+++ b/litellm/proxy/spend_tracking/background_interaction_settlement.py
@@ -1,21 +1,25 @@
import asyncio
import os
import socket
-from collections.abc import Awaitable, Mapping, Sequence
-from dataclasses import dataclass
-from datetime import datetime, timezone
+from collections.abc import Awaitable, Callable, Mapping, Sequence
+from dataclasses import dataclass, replace
+from datetime import datetime, timedelta, timezone
from itertools import chain
from types import MappingProxyType
-from typing import TYPE_CHECKING, Final, Protocol, TypeVar
+from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeVar
from pydantic import ValidationError
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
-from litellm.constants import BACKGROUND_INTERACTION_COST_POLLING_ENABLED
+from litellm.constants import (
+ BACKGROUND_INTERACTION_COST_POLLING_ENABLED,
+ BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS,
+)
from litellm.interactions.background_cost_polling import (
DEFAULT_POLL_SCHEDULE,
BackgroundInteractionCreateContext,
+ BackgroundInteractionPollContext,
FetchInteraction,
PendingBackgroundInteraction,
PollSchedule,
@@ -24,10 +28,13 @@ from litellm.interactions.background_cost_polling import (
fetch_background_interaction,
resume_unsettled_background_interactions,
)
-from litellm.repositories.table_repositories import BackgroundInteractionSettlementRepository
+from litellm.repositories.table_repositories import BackgroundInteractionSettlementRepository, SpendLogsRepository
+from litellm.types.interactions import InteractionsAPIResponse
+from litellm.types.router import CredentialLiteLLMParams
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
+ from litellm.router import Router
class _SettlementRow(Protocol):
@@ -41,6 +48,8 @@ class _SettlementRow(Protocol):
def created_at(self) -> datetime: ...
@property
def claimed_at(self) -> datetime | None: ...
+ @property
+ def settled_at(self) -> datetime | None: ...
class _NewSettlementRow(TypedDict):
@@ -59,8 +68,24 @@ class _UnclaimedRowKey(TypedDict):
claimed_at: ReadOnly[None]
-class _UnclaimedRows(TypedDict):
- claimed_at: ReadOnly[None]
+class _Before(TypedDict):
+ lt: ReadOnly[datetime]
+
+
+class _ExpiredClaimRowKey(TypedDict):
+ interaction_id: ReadOnly[str]
+ settled_at: ReadOnly[None]
+ claimed_at: ReadOnly[_Before]
+
+
+class _OwnClaimRowKey(TypedDict):
+ interaction_id: ReadOnly[str]
+ claimed_by: ReadOnly[str]
+ settled_at: ReadOnly[None]
+
+
+class _UnsettledRows(TypedDict):
+ settled_at: ReadOnly[None]
class _Claim(TypedDict):
@@ -79,9 +104,19 @@ class _SettlementTableActions(Protocol):
def find_unique(self, *, where: _RowKey) -> Awaitable[_SettlementRow | None]: ...
- def find_many(self, *, where: _UnclaimedRows) -> Awaitable[Sequence[_SettlementRow]]: ...
+ def find_many(self, *, where: _UnsettledRows) -> Awaitable[Sequence[_SettlementRow]]: ...
- def update_many(self, *, data: _Claim | _Outcome, where: _RowKey | _UnclaimedRowKey) -> Awaitable[int]: ...
+ def update_many(
+ self, *, data: _Claim | _Outcome, where: _RowKey | _UnclaimedRowKey | _ExpiredClaimRowKey | _OwnClaimRowKey
+ ) -> Awaitable[int]: ...
+
+
+class _SpendLogKey(TypedDict):
+ request_id: ReadOnly[str]
+
+
+class _SpendLogLookup(Protocol):
+ def find_unique(self, *, where: _SpendLogKey) -> Awaitable[object | None]: ...
_CLEARED_CREATE_CONTEXT: Final[Mapping[str, object]] = MappingProxyType({})
@@ -103,11 +138,7 @@ def _json(data: Mapping[str, object]) -> object:
return Json.keys(**data)
-def _pending_rows(rows: Sequence[_SettlementRow]) -> tuple[PendingBackgroundInteraction, ...]:
- return tuple(chain.from_iterable(_pending_row(row) for row in rows))
-
-
-def _pending_row(row: _SettlementRow) -> tuple[PendingBackgroundInteraction, ...]:
+def _pending_row(row: _SettlementRow, claim_expires_at: datetime | None) -> tuple[PendingBackgroundInteraction, ...]:
try:
create_context: Final = BackgroundInteractionCreateContext.model_validate(row.create_context)
except ValidationError:
@@ -122,14 +153,31 @@ def _pending_row(row: _SettlementRow) -> tuple[PendingBackgroundInteraction, ...
custom_llm_provider=row.custom_llm_provider,
create_context=create_context,
created_at=row.created_at,
+ claim_expires_at=claim_expires_at,
),
)
@dataclass(frozen=True, slots=True)
class PrismaBackgroundSettlementStore:
+ """
+ A claim is a lease, not a settlement: the settler that claims a row bills
+ or releases it and then records the outcome, so a row claimed longer ago
+ than the lease with no outcome belongs to a settler that died in between,
+ and any settler may claim it again. A live settler renews its claim while
+ it settles, so it is never overtaken. The settler before may have billed
+ and then died or failed to record it, so a takeover bills only when the
+ interaction has no spend row yet, and otherwise records it as billed. The
+ spend row is read before the claim is taken, so a failed read takes
+ nothing and the claim is retried. With spend logs off there is no row to
+ tell, so a takeover bills nothing rather than risk billing twice.
+ """
+
table: _SettlementTableActions
+ spend_logs: _SpendLogLookup
+ spend_logs_kept: Callable[[], bool]
claimed_by: str
+ claim_lease_seconds: float = BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS
async def register(self, pending: PendingBackgroundInteraction) -> None:
await self.table.create(
@@ -143,24 +191,71 @@ class PrismaBackgroundSettlementStore:
async def pending(self, interaction_id: str) -> PendingBackgroundInteraction | None:
row: Final = await self._row(interaction_id)
- if row is None or row.claimed_at is not None:
+ if row is None or self._is_owned(row):
return None
- return next(iter(_pending_row(row)), None)
+ return next(iter(_pending_row(row, claim_expires_at=None)), None)
async def is_claimed(self, interaction_id: str) -> bool:
row: Final = await self._row(interaction_id)
- return row is not None and row.claimed_at is not None
+ return row is not None and self._is_owned(row)
async def claim(self, interaction_id: str) -> bool:
+ if await self._claim_where(_UnclaimedRowKey(interaction_id=interaction_id, claimed_at=None)):
+ return True
+ settled_before: Final = await self._outcome_a_dead_settler_left(interaction_id)
+ took_over: Final = await self._claim_where(
+ _ExpiredClaimRowKey(
+ interaction_id=interaction_id,
+ settled_at=None,
+ claimed_at=_Before(lt=datetime.now(timezone.utc) - self._lease),
+ )
+ )
+ if not took_over or settled_before is None:
+ return took_over
+ if settled_before == "unsettled":
+ verbose_proxy_logger.warning(
+ "Background interaction %s lost its settler mid-settle, and with spend logs off nothing shows "
+ "whether it was billed, so it is left unbilled rather than risk billing it twice",
+ interaction_id,
+ )
+ await self.record_outcome(interaction_id, settled_before)
+ return False
+
+ async def renew_claim(self, interaction_id: str) -> None:
+ await self.table.update_many(
+ data=_Claim(claimed_at=datetime.now(timezone.utc), claimed_by=self.claimed_by),
+ where=_OwnClaimRowKey(interaction_id=interaction_id, claimed_by=self.claimed_by, settled_at=None),
+ )
+
+ async def _claim_where(self, where: _UnclaimedRowKey | _ExpiredClaimRowKey) -> bool:
claimed_rows: Final = await _read_from_a_table_that_may_not_exist(
self.table.update_many(
- data=_Claim(claimed_at=datetime.now(timezone.utc), claimed_by=self.claimed_by),
- where=_UnclaimedRowKey(interaction_id=interaction_id, claimed_at=None),
+ data=_Claim(claimed_at=datetime.now(timezone.utc), claimed_by=self.claimed_by), where=where
),
when_missing=0,
)
return claimed_rows == 1
+ async def _outcome_a_dead_settler_left(self, interaction_id: str) -> Literal["billed", "unsettled"] | None:
+ if not self.spend_logs_kept():
+ return "unsettled"
+ if await self.spend_logs.find_unique(where=_SpendLogKey(request_id=interaction_id)) is None:
+ return None
+ return "billed"
+
+ @property
+ def _lease(self) -> timedelta:
+ return timedelta(seconds=self.claim_lease_seconds)
+
+ def _claim_expires_at(self, row: _SettlementRow) -> datetime | None:
+ return None if row.claimed_at is None else row.claimed_at + self._lease
+
+ def _is_owned(self, row: _SettlementRow) -> bool:
+ if row.settled_at is not None:
+ return True
+ claim_expires_at: Final = self._claim_expires_at(row)
+ return claim_expires_at is not None and claim_expires_at > datetime.now(timezone.utc)
+
async def _row(self, interaction_id: str) -> _SettlementRow | None:
return await _read_from_a_table_that_may_not_exist(
self.table.find_unique(where=_RowKey(interaction_id=interaction_id)), when_missing=None
@@ -174,19 +269,78 @@ class PrismaBackgroundSettlementStore:
where=_RowKey(interaction_id=interaction_id),
)
- async def unclaimed(self) -> Sequence[PendingBackgroundInteraction]:
- return _pending_rows(await self.table.find_many(where=_UnclaimedRows(claimed_at=None)))
+ async def unsettled(self) -> Sequence[PendingBackgroundInteraction]:
+ rows: Final = await self.table.find_many(where=_UnsettledRows(settled_at=None))
+ return tuple(chain.from_iterable(_pending_row(row, self._claim_expires_at(row)) for row in rows))
+
+
+def _provider_credentials(router: "Router", deployment_id: str, custom_llm_provider: str) -> Mapping[str, object]:
+ credentials: Final = router.get_deployment_credentials_with_provider(model_id=deployment_id)
+ if credentials is None or credentials.get("custom_llm_provider") != custom_llm_provider:
+ return MappingProxyType({})
+ return MappingProxyType(credentials)
+
+
+def _group_deployment_credentials(router: "Router", context: BackgroundInteractionPollContext) -> Mapping[str, object]:
+ if context.model_group is None:
+ return MappingProxyType({})
+ candidates: Final = tuple(
+ credentials
+ for deployment_id in router.get_model_ids(model_name=context.model_group)
+ if (credentials := _provider_credentials(router, deployment_id, context.custom_llm_provider))
+ )
+ return candidates[0] if len(candidates) == 1 else MappingProxyType({})
+
+
+def _deployment_credentials(
+ router: "Router | None", context: BackgroundInteractionPollContext
+) -> CredentialLiteLLMParams | None:
+ if router is None:
+ return None
+ serving: Final = (
+ _provider_credentials(router, context.deployment_id, context.custom_llm_provider)
+ if context.deployment_id is not None and router.get_deployment(model_id=context.deployment_id) is not None
+ else _group_deployment_credentials(router, context)
+ )
+ return CredentialLiteLLMParams.model_validate(serving) if serving else None
+
+
+def fetch_with_deployment_credentials(
+ router: Callable[[], "Router | None"], fetch_interaction: FetchInteraction = fetch_background_interaction
+) -> FetchInteraction:
+ """
+ A resumed poll has no request to take credentials from, and they are never
+ stored, so it fetches with the credentials of the deployment that served
+ the create, looked up by its id on the live router at each fetch so a
+ deployment loaded from the database after startup is found too. Editing a
+ deployment changes its id, so a missing id falls back to the create's model
+ group when exactly one of its deployments serves the same provider. Anything
+ else falls back to the provider's environment credentials, never to a
+ wildcard route, whose key belongs to whatever provider it routes to.
+ """
+
+ async def fetch(context: BackgroundInteractionPollContext) -> InteractionsAPIResponse:
+ credentials: Final = _deployment_credentials(router(), context)
+ if credentials is None:
+ return await fetch_interaction(context)
+ return await fetch_interaction(replace(context, credentials=credentials))
+
+ return fetch
async def configure_background_interaction_settlement(
table: _SettlementTableActions,
+ spend_logs: _SpendLogLookup,
+ spend_logs_kept: Callable[[], bool],
claimed_by: str,
fetch_interaction: FetchInteraction = fetch_background_interaction,
schedule: PollSchedule = DEFAULT_POLL_SCHEDULE,
) -> tuple["asyncio.Task[SettlementOutcome | None]", ...]:
if not BACKGROUND_INTERACTION_COST_POLLING_ENABLED:
return ()
- store: Final = PrismaBackgroundSettlementStore(table=table, claimed_by=claimed_by)
+ store: Final = PrismaBackgroundSettlementStore(
+ table=table, spend_logs=spend_logs, spend_logs_kept=spend_logs_kept, claimed_by=claimed_by
+ )
configure_background_settlement_store(store)
resumed: Final = await resume_unsettled_background_interactions(store, fetch_interaction, schedule)
if resumed:
@@ -194,11 +348,16 @@ async def configure_background_interaction_settlement(
return resumed
-async def install_background_interaction_settlement(prisma_client: "PrismaClient") -> None:
+async def install_background_interaction_settlement(
+ prisma_client: "PrismaClient", router: Callable[[], "Router | None"], spend_logs_kept: Callable[[], bool]
+) -> None:
try:
await configure_background_interaction_settlement(
table=BackgroundInteractionSettlementRepository(prisma_client).table,
+ spend_logs=SpendLogsRepository(prisma_client).table,
+ spend_logs_kept=spend_logs_kept,
claimed_by=f"{socket.gethostname()}:{os.getpid()}",
+ fetch_interaction=fetch_with_deployment_credentials(router),
)
except Exception as e: # noqa: BLE001 # a boot step must survive any DB error; billing then settles in-process as before
verbose_proxy_logger.warning(
diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py
index 49644f1aa12..9ebea589cc9 100644
--- a/litellm/proxy/spend_tracking/spend_tracking_utils.py
+++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py
@@ -15,6 +15,7 @@ from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.constants import (
+ CALLER_SESSION_ID_METADATA_KEY,
CLI_SESSION_KEY_PREFIX,
EMPTY_MAPPING,
INTERNAL_CALL_ORIGIN_METADATA_KEY,
@@ -33,6 +34,8 @@ from litellm.constants import (
from litellm.constants import (
MAX_STRING_LENGTH_PROMPT_IN_DB as DEFAULT_MAX_STRING_LENGTH_PROMPT_IN_DB,
)
+from litellm.integrations.otel.model.metadata import caller_session_id
+from litellm.integrations.otel.model.trace_controls import TraceControls
from litellm.litellm_core_utils.classifier_logging import (
CLASSIFIER_AUDIT_FIELDS,
classifier_audit_fields,
@@ -570,11 +573,11 @@ def get_logging_payload(
elif isinstance(_usage, dict):
usage = _usage
- # A request that failed mid-stream has no usable response_obj usage, but the
- # streaming handler may have recovered the usage from the chunks already
- # delivered. Honor that override so the partial usage lands in spend tracking.
+ # A request that failed mid-stream has no usable response_obj usage, and a stream
+ # continued on a fallback only has the fallback's. combined_usage_object holds the
+ # usage of every call the client's answer took, so it wins in both cases.
_combined_usage: Final = kwargs.get("combined_usage_object")
- if not usage and isinstance(_combined_usage, litellm.Usage):
+ if (not usage or "continued_partial_usage" in kwargs) and isinstance(_combined_usage, litellm.Usage):
usage = _combined_usage.model_dump()
id = get_spend_logs_id(call_type or "acompletion", response_obj_dict, kwargs, metadata)
@@ -930,6 +933,16 @@ def _get_batch_trace_session_id(call_type: str | None, request_id: str | None) -
return request_id.removesuffix(BATCH_COST_REQUEST_ID_SUFFIX)
+def _caller_session_for_spend_log(kwargs: Mapping[str, object], metadata: Mapping[str, object] | None) -> str | None:
+ beside_trace_header: Final = metadata.get(CALLER_SESSION_ID_METADATA_KEY) if metadata else None
+ caller_session: Final = (
+ beside_trace_header
+ if isinstance(beside_trace_header, str) and beside_trace_header
+ else caller_session_id(kwargs, TraceControls())
+ )
+ return caller_session if caller_session and "\x00" not in caller_session else None
+
+
def _get_session_id_for_spend_log(
kwargs: Mapping[str, object],
metadata: Mapping[str, object] | None,
@@ -937,9 +950,12 @@ def _get_session_id_for_spend_log(
omit_when_missing: bool,
batch_trace_session_id: str | None = None,
) -> str | None:
- """Under `omit` only `metadata.session_id`, the key Langfuse reads, counts as a session; `litellm_session_id` may
- be a copied trace id. Batch call types carry a deterministic session derived from the batch id, which outranks
- the per-request trace ids because those differ between the create call and the cost poller's row."""
+ """Resolve the session id for the spend log row: `omit` honors only metadata.session_id,
+ batch sessions outrank everything, then the conversation id the caller sent (also one an x-litellm-trace-id header
+ displaced from the request's session), then trace ids.
+
+ A session minted by `missing_session_id: generate` or copied from the request's trace id is not a conversation, so
+ it yields to the trace id, which the Responses bridge points at the previous turn's session to keep a chain whole."""
if omit_when_missing:
session_id: Final = metadata.get("session_id") if metadata else None
return str(session_id) if session_id else None
@@ -948,11 +964,14 @@ def _get_session_id_for_spend_log(
if batch_trace_session_id is not None:
return batch_trace_session_id
- if standard_logging_payload is not None and standard_logging_payload.get("trace_id") is not None:
- return str(standard_logging_payload.get("trace_id"))
- if kwargs.get("litellm_trace_id") is not None:
- return str(kwargs.get("litellm_trace_id"))
- return str(uuid.uuid4())
+ if litellm.request_correlation_in_logs and (caller_session := _caller_session_for_spend_log(kwargs, metadata)):
+ return caller_session
+ match standard_logging_payload:
+ case {"trace_id": trace_id} if trace_id is not None:
+ return str(trace_id)
+ case _:
+ litellm_trace_id: Final = kwargs.get("litellm_trace_id")
+ return str(litellm_trace_id) if litellm_trace_id is not None else str(uuid.uuid4())
def _get_request_duration_ms(start_time: datetime, end_time: datetime) -> int | None:
diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py
index b327129793d..4ece88f757d 100644
--- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py
+++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py
@@ -343,7 +343,7 @@ class UISettings(LiteLLMBaseModel):
)
enable_chat_ui: bool = Field(
- default=False,
+ default=True,
description="If true, shows the Chat page in the UI sidebar, letting users chat with an LLM and connect their own MCP server credentials via OAuth.",
)
diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py
index 367d5be7bfc..aeec0bf67c0 100644
--- a/litellm/realtime_api/main.py
+++ b/litellm/realtime_api/main.py
@@ -500,11 +500,13 @@ async def _arealtime(
)
return
+ drop_params: Final = litellm.drop_params is True or normalize_drop_params(kwargs.get("drop_params")) is True
provider_config: BaseRealtimeConfig | None = None
if _custom_llm_provider in LlmProviders._member_map_.values():
provider_config = ProviderConfigManager.get_provider_realtime_config(
model=model,
provider=LlmProviders(_custom_llm_provider),
+ drop_params=drop_params,
)
if provider_config is not None:
await base_llm_http_handler.async_realtime(
@@ -585,7 +587,6 @@ async def _arealtime(
aws_bedrock_runtime_endpoint: Final = kwargs.get("aws_bedrock_runtime_endpoint")
aws_external_id: Final = kwargs.get("aws_external_id")
aws_session_tags: Final = kwargs.get("aws_session_tags")
- drop_params: Final = litellm.drop_params is True or normalize_drop_params(kwargs.get("drop_params")) is True
await bedrock_realtime.async_realtime(
model=model,
@@ -671,6 +672,7 @@ async def _arealtime(
resolve_access_token=resolve_vertex_access_token,
project=resolved_project,
location=vertex_location,
+ drop_params=drop_params,
)
await base_llm_http_handler.async_realtime(
diff --git a/litellm/responses/bridge_check.py b/litellm/responses/bridge_check.py
index c2b4c105199..2552ee8c880 100644
--- a/litellm/responses/bridge_check.py
+++ b/litellm/responses/bridge_check.py
@@ -4,6 +4,7 @@ from urllib.parse import urlsplit
import litellm
from litellm._logging import verbose_logger
+from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.llms.azure_ai.common_utils import (
azure_ai_supports_native_responses,
foundry_chat_rejects_function_tools_while_reasoning,
@@ -11,7 +12,8 @@ from litellm.llms.azure_ai.common_utils import (
from litellm.llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import OpenAIWebSearchOptions
-from litellm.utils import get_model_info_helper
+from litellm.types.utils import LlmProviders
+from litellm.utils import ProviderConfigManager, get_model_info_helper
_OPENAI_DEFAULT_API_BASE: Final = "https://api.openai.com/v1"
_OPENAI_API_HOST: Final = "api.openai.com"
@@ -36,6 +38,18 @@ def _resolve_openai_api_base(api_base: str | None) -> str:
)
+def chat_target_accepts_reasoning_summary(model: str, custom_llm_provider: str | None) -> bool:
+ try:
+ local_model, provider, _, _ = get_llm_provider(model=model, custom_llm_provider=custom_llm_provider)
+ chat_config: Final = ProviderConfigManager.get_provider_chat_config(
+ model=local_model, provider=LlmProviders(provider)
+ )
+ except Exception as e: # noqa: BLE001 # a capability probe must never fail the request it probes for
+ verbose_logger.debug("chat target reasoning summary check failed: %s", e)
+ return False
+ return chat_config is not None and chat_config.accepts_reasoning_summary_in_reasoning_effort
+
+
def responses_api_bridge_check(
model: str,
custom_llm_provider: str,
diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py
index 732c32ff2cb..e78ec5e5d3e 100644
--- a/litellm/responses/litellm_completion_transformation/session_handler.py
+++ b/litellm/responses/litellm_completion_transformation/session_handler.py
@@ -281,7 +281,13 @@ class ResponsesSessionHandler:
SQL query
- SELECT session_id FROM spend_logs WHERE response_id = previous_response_id, SELECT * FROM spend_logs WHERE session_id = session_id
+ SELECT session_id, api_key, "user", team_id FROM spend_logs WHERE request_id = previous_response_id, then
+ SELECT * FROM spend_logs with that session_id and the same owner
+
+ Only rows with the referenced response's owner are returned: the same user and team, and for a key with
+ neither, the same key. The session id can be caller supplied, so matching on it alone would mix owners. A
+ user's or team's regenerated or second key keeps the history; a key with neither does not, and turns written
+ by another owner (an admin continuing the chain) are not replayed.
A just-finished turn gets a short second chance: the worker that served it may
still be writing its spend log when the follow-up arrives, and an empty result
@@ -303,13 +309,24 @@ class ResponsesSessionHandler:
query: Final = """
WITH matching_session AS (
- SELECT session_id
+ SELECT session_id, api_key, "user", team_id
FROM "LiteLLM_SpendLogs"
WHERE request_id = $1
)
SELECT *
- FROM "LiteLLM_SpendLogs"
- WHERE session_id IN (SELECT session_id FROM matching_session)
+ FROM "LiteLLM_SpendLogs" AS logs
+ WHERE EXISTS (
+ SELECT 1
+ FROM matching_session
+ WHERE matching_session.session_id = logs.session_id
+ AND matching_session."user" IS NOT DISTINCT FROM logs."user"
+ AND matching_session.team_id IS NOT DISTINCT FROM logs.team_id
+ AND (
+ COALESCE(matching_session."user", '') <> ''
+ OR COALESCE(matching_session.team_id, '') <> ''
+ OR matching_session.api_key = logs.api_key
+ )
+ )
ORDER BY "endTime" ASC;
"""
diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py
index 7f5d143e98e..587cdfcee25 100644
--- a/litellm/responses/litellm_completion_transformation/transformation.py
+++ b/litellm/responses/litellm_completion_transformation/transformation.py
@@ -39,7 +39,11 @@ from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_logger
from litellm.caching import InMemoryCache
-from litellm.constants import REDACTED_BY_LITELLM, REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER
+from litellm.constants import (
+ CALLER_SESSION_ID_METADATA_KEY,
+ REDACTED_BY_LITELLM,
+ REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER,
+)
from litellm.integrations.anthropic_cache_control_hook import MAX_CACHE_CONTROL_BLOCKS, AnthropicCacheControlHook
from litellm.litellm_core_utils.core_helpers import normalize_drop_params
from litellm.litellm_core_utils.get_supported_openai_params import (
@@ -56,6 +60,7 @@ from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.bedrock.common_utils import bedrock_declared_cache_point_support
from litellm.llms.openrouter.chat.transformation import openrouter_model_supports_cache_control
from litellm.llms.snowflake.chat.transformation import is_snowflake_claude_model
+from litellm.responses.bridge_check import chat_target_accepts_reasoning_summary
from litellm.responses.litellm_completion_transformation.session_handler import (
ResponsesSessionHandler,
)
@@ -562,6 +567,14 @@ class ChatCompletionSession(TypedDict, total=False):
########### End of Initialize Classes used for Responses API ###########
+def _chain_session_metadata(chain_session: str, request_kwargs: Mapping[str, object]) -> Mapping[str, object]:
+ key: Final = "litellm_metadata" if request_kwargs.get("litellm_metadata") else "metadata"
+ metadata: Final = request_kwargs.get(key)
+ if not isinstance(metadata, dict):
+ return MappingProxyType({})
+ return {key: {**_STR_KEY_DICT_ADAPTER.validate_python(metadata), CALLER_SESSION_ID_METADATA_KEY: chain_session}}
+
+
class LiteLLMCompletionResponsesConfig:
@staticmethod
def get_supported_openai_params(model: str) -> list:
@@ -785,8 +798,8 @@ class LiteLLMCompletionResponsesConfig:
provider either rejects it or silently drops it, and dropping it turns reasoning off while
still billing for the turn. ``summary`` has no chat equivalent, so it rides the
``reasoning_summary`` alias, which ``litellm.completion`` reassembles into ``{effort,
- summary}`` when it bridges the model back onto the Responses API, and is sent to nothing
- else.
+ summary}`` when it bridges the model back onto the Responses API or targets a chat provider
+ whose config maps a summary ask (Gemini's thought summaries), and is sent to nothing else.
"""
if not reasoning_param:
return ResponsesReasoningChatForm(effort=None, summary=None)
@@ -807,7 +820,10 @@ class LiteLLMCompletionResponsesConfig:
reasoning_summary=summary,
api_base=api_base,
)
- return ResponsesReasoningChatForm(effort=effort, summary=summary if bridges_back else None)
+ forwards_summary: Final = bridges_back or (
+ isinstance(effort, str) and chat_target_accepts_reasoning_summary(model, custom_llm_provider)
+ )
+ return ResponsesReasoningChatForm(effort=effort, summary=summary if forwards_summary else None)
@staticmethod
def transform_responses_api_request_to_chat_completion_request(
@@ -1074,8 +1090,6 @@ class LiteLLMCompletionResponsesConfig:
else:
# Both are empty - this likely means function_call_output had empty/invalid call_id
# Provide a helpful error message
- import litellm
-
raise litellm.BadRequestError(
message=(
f"Unable to create messages for completion request. "
@@ -1090,8 +1104,14 @@ class LiteLLMCompletionResponsesConfig:
)
litellm_completion_request["messages"] = combined_messages
- litellm_completion_request["litellm_trace_id"] = chat_completion_session.get("litellm_session_id")
- return litellm_completion_request
+ chain_session: Final = chat_completion_session.get("litellm_session_id")
+ litellm_completion_request["litellm_trace_id"] = chain_session
+ chain_session_metadata: Final[Mapping[str, object]] = (
+ _chain_session_metadata(chain_session, request_kwargs)
+ if chain_session and litellm.request_correlation_in_logs
+ else MappingProxyType({})
+ )
+ return {**litellm_completion_request, **chain_session_metadata}
@staticmethod
def _transform_response_input_param_to_chat_completion_message(
diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py
index 7cf68332c07..89c32aed1fb 100644
--- a/litellm/responses/streaming_iterator.py
+++ b/litellm/responses/streaming_iterator.py
@@ -31,6 +31,7 @@ from litellm.litellm_core_utils.core_helpers import process_response_headers
from litellm.litellm_core_utils.internal_call_metadata import without_budget_reservation
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base
+from litellm.litellm_core_utils.llm_response_utils.get_headers import provider_request_id_headers
from litellm.litellm_core_utils.llm_response_utils.response_metadata import (
update_response_metadata,
)
@@ -262,7 +263,9 @@ def stream_error_status_and_message(error_obj: object) -> tuple[int, str]:
return _status_code_for_error_fields(error_type, error_code), message
-def _map_stream_error_to_exception(error_obj: object, model: str, custom_llm_provider: str) -> Exception:
+def _map_stream_error_to_exception(
+ error_obj: object, model: str, custom_llm_provider: str, headers: httpx.Headers | None = None
+) -> Exception:
from litellm.llms.base_llm.chat.transformation import BaseLLMException
error_message, error_type, error_code = _error_event_fields(error_obj)
@@ -272,6 +275,7 @@ def _map_stream_error_to_exception(error_obj: object, model: str, custom_llm_pro
status_code=status_code,
message=f"Error code: {status_code} - {{'error': {error_body}}}",
body=error_body,
+ headers=headers,
)
try:
return litellm.exception_type(
@@ -667,7 +671,12 @@ class BaseResponsesAPIStreamingIterator:
)
def _map_error_event_exception(self, error_obj: object) -> Exception:
- return _map_stream_error_to_exception(error_obj, self.model or "", self.custom_llm_provider or "")
+ return _map_stream_error_to_exception(
+ error_obj,
+ self.model or "",
+ self.custom_llm_provider or "",
+ headers=provider_request_id_headers(self._raw_response_headers),
+ )
def _maybe_raise_for_error_event(self, result: object) -> None:
chunk_type: Final = getattr(result, "type", None)
@@ -683,11 +692,17 @@ class BaseResponsesAPIStreamingIterator:
mapped_exception: Final = self._map_error_event_exception(error_obj)
if not _mid_stream_fallback_eligible(mapped_exception):
raise mapped_exception
+ mapped_status: Final = getattr(mapped_exception, "status_code", None)
raise MidStreamFallbackError(
message=str(mapped_exception),
model=self.model or "",
llm_provider=self.custom_llm_provider or "",
original_exception=mapped_exception,
+ response=httpx.Response(
+ status_code=int(mapped_status) if mapped_status is not None else 503,
+ headers=provider_request_id_headers(self._raw_response_headers),
+ request=httpx.Request(method="POST", url=f"https://{self.custom_llm_provider or ''}.com/v1/"),
+ ),
generated_content=self._generated_content,
is_pre_first_chunk=not self._yielded_first_chunk,
)
diff --git a/litellm/router.py b/litellm/router.py
index a549bb06e7d..3c969610e83 100644
--- a/litellm/router.py
+++ b/litellm/router.py
@@ -38,7 +38,7 @@ from dataclasses import dataclass
from datetime import datetime, timezone
from functools import lru_cache, partial
from types import MappingProxyType
-from typing import TYPE_CHECKING, Any, Final, Generic, Literal, Optional, TypeAlias, TypeVar, Union, cast
+from typing import TYPE_CHECKING, Any, Final, Generic, Literal, NoReturn, Optional, TypeAlias, TypeVar, Union, cast
import anyio
import httpx
@@ -68,6 +68,7 @@ from litellm.constants import (
DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER,
DEFAULT_MAX_LRU_CACHE_SIZE,
INTERNAL_CALL_ORIGIN_METADATA_KEY,
+ MID_STREAM_FAILURE_LOG_WAIT_SECONDS,
OUTPUT_TOKEN_CEILING_PARAMS,
PROMPT_CACHE_PIN_CACHE_MAX_SIZE,
ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY,
@@ -75,7 +76,7 @@ from litellm.constants import (
RUNTIME_UPDATABLE_ROUTER_SETTINGS,
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
)
-from litellm.integrations.custom_guardrail import is_guardrail_intervention
+from litellm.integrations.custom_guardrail import CustomGuardrail, is_guardrail_intervention
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
@@ -128,6 +129,7 @@ from litellm.llms.base_llm.vector_store.transformation import (
RouterVectorStoreEmbeddingExecutor,
vector_store_request_metadata,
)
+from litellm.llms.bedrock.common_utils import get_bedrock_base_model
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client
from litellm.llms.openai_like.json_loader import JSONProviderRegistry
from litellm.llms.openai_like.model_info import (
@@ -193,6 +195,7 @@ from litellm.router_utils.clientside_credential_handler import (
is_clientside_credential,
)
from litellm.router_utils.common_utils import (
+ MID_STREAM_CONTINUATION_KWARG,
filter_team_based_models,
filter_web_search_deployments,
format_fallback_outcome_message,
@@ -232,7 +235,10 @@ from litellm.router_utils.fallback_event_handlers import (
check_non_standard_fallback_format,
clear_pre_routing_selection,
committed_retry_budget_for_request,
+ continuation_fallback_targets,
fallback_lookup_groups,
+ fallback_target_name,
+ fallback_target_overrides_request_params,
fallbacks_disabled_for_request,
get_fallback_model_group_for_lookup_groups,
get_pre_routing_selection,
@@ -257,6 +263,17 @@ from litellm.router_utils.handle_error import (
send_llm_exception_alert,
)
from litellm.router_utils.health_state_cache import DeploymentHealthCache
+from litellm.router_utils.pre_call_checks.continuation_prefill_check import (
+ MID_STREAM_CONTINUATION_MARKER,
+ ContinuationOutput,
+ ContinuationPrefillDeploymentCheck,
+ ContinuedPartial,
+ constrains_continuation,
+ continuation_deployments,
+ continuation_output_ceilings,
+ continued_text,
+ deployment_can_continue,
+)
from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
DeploymentAffinityCheck,
warn_on_unknown_model_group_affinity_flags,
@@ -302,6 +319,8 @@ from litellm.router_utils.routing_groups import (
)
from litellm.router_utils.routing_read_batch import ROUTER_USAGE_TARGET, RoutingPrefetch, RoutingReadBatch
from litellm.scheduler import FlowItem, Scheduler
+from litellm.secret_managers.main import str_to_bool
+from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.litellm_params import RoutingStrategyName
from litellm.types.llms.openai import (
AllMessageValues,
@@ -328,6 +347,7 @@ from litellm.types.router import (
FallbackBudgetCheck,
GuardrailTypedDict,
LiteLLM_Params,
+ LiteLLMParamsTypedDict,
MockRouterTestingParams,
ModelGroupInfo,
OptionalPreCallChecks,
@@ -369,11 +389,13 @@ from litellm.utils import (
EmbeddingResponse,
ModelResponse,
Rules,
+ TextCompletionResponse,
function_setup,
get_llm_provider,
get_non_default_completion_params,
get_secret,
get_utc_datetime,
+ is_generalized_model_info,
is_region_allowed,
provider_rejectable_params,
set_live_deployment_replay,
@@ -484,6 +506,7 @@ _CLAUDE_CODE_SESSION_ID_RE: Final = re.compile(r"^[a-zA-Z0-9_\-]{8,}$")
_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS: Final = 3600
CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET: Final = "claude_code_session_router_binding"
+_PRE_CALL_CHECK_FILTERED_PARAMS: Final = frozenset({"response_format"})
_RUNTIME_TOGGLEABLE_PRE_CALL_CHECKS: Final[Mapping[str, type[CustomLogger]]] = MappingProxyType(
{
"prompt_caching": PromptCachingDeploymentCheck,
@@ -523,6 +546,15 @@ _ROUTING_KWARGS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeA
_FALLBACK_HOP_ADAPTER: Final = TypeAdapter(Mapping[str, object])
_OPAQUE_VALUE_ADAPTER: Final = TypeAdapter(object)
_DEPLOYMENT_SELECTED_EVENT: Final = "litellm.request.deployment_selected"
+_CONTENT_BLOCKS_ADAPTER: Final = TypeAdapter(list[dict[str, object]])
+
+
+def _extend_trailing_text_block(blocks: Sequence[Mapping[str, object]], text: str) -> Sequence[Mapping[str, object]]:
+ """Continue the last text block in place so the prefill reads as the text the caller already received."""
+ last: Final = blocks[-1] if blocks else None
+ if last is not None and last.get("type") == "text" and isinstance(last_text := last.get("text"), str):
+ return [*blocks[:-1], {**last, "text": f"{last_text}{text}"}]
+ return [*blocks, {"type": "text", "text": text}]
def _is_fallback_hop(request_kwargs: Mapping[str, object]) -> bool:
@@ -991,6 +1023,49 @@ def as_output_cap(value: object) -> int | None:
return cap if cap >= 0 else None
+def _exact_base_model_max_input_tokens(base_model: str) -> int | None:
+ try:
+ model_info: Final = litellm.get_model_info(model=base_model)
+ except Exception: # noqa: BLE001 # get_model_info raises ModelNotMappedError or a bare Exception for an unmapped model
+ return None
+ if is_generalized_model_info(model_info):
+ return None
+ return model_info.get("max_input_tokens")
+
+
+def _provider_model_id(model_name: str, custom_llm_provider: str) -> str:
+ model_id: Final = model_name.rsplit("/", 1)[-1]
+ return get_bedrock_base_model(model_id) if custom_llm_provider.startswith("bedrock") else model_id
+
+
+def _names_route_model(base_model: str, route_key: str, custom_llm_provider: str) -> bool:
+ route_model_id: Final = _provider_model_id(route_key, custom_llm_provider)
+ base_model_id: Final = _provider_model_id(base_model, custom_llm_provider)
+ return route_model_id in (base_model_id, base_model_id.split("@", 1)[0])
+
+
+def _declared_base_model(user_model_info: Mapping[str, object]) -> str | None:
+ match user_model_info:
+ case {"base_model": str() as base_model}:
+ return base_model
+ case _:
+ return None
+
+
+def _base_model_widened_window(
+ route_model_info: ModelMapInfo, base_model: str | None, custom_llm_provider: str
+) -> Mapping[str, int]:
+ route_max_input_tokens: Final = route_model_info.get("max_input_tokens")
+ if base_model is None or route_max_input_tokens is None:
+ return MappingProxyType({})
+ if not _names_route_model(base_model, route_model_info["key"], custom_llm_provider):
+ return MappingProxyType({})
+ base_max_input_tokens: Final = _exact_base_model_max_input_tokens(base_model)
+ if base_max_input_tokens is None or base_max_input_tokens <= route_max_input_tokens:
+ return MappingProxyType({})
+ return MappingProxyType({"max_input_tokens": base_max_input_tokens})
+
+
class Router:
@property
def _routing_groups(self) -> dict[str, RoutingGroup]:
@@ -1072,6 +1147,7 @@ class Router:
health_check_ignore_transient_errors: bool = False,
background_health_check_model_groups: Sequence[str] | None = None,
enable_weighted_failover: bool = False,
+ enable_mid_stream_fallback_continuation: bool = False,
fallback_access_check: FallbackAccessCheck | None = None,
fallback_budget_check: FallbackBudgetCheck | None = None,
auto_router_capability_limit: AutoRouterCapabilityLimit | None = None,
@@ -1113,6 +1189,7 @@ class Router:
deployment_affinity_ttl_seconds (int): TTL for user-key -> deployment affinity mapping. Defaults to 3600.
ignore_invalid_deployments (bool): Ignores invalid deployments, and continues with other deployments. Default is to raise an error.
enable_weighted_failover (bool): When True and the routing strategy is "simple-shuffle", a retryable failure on one deployment causes the request to re-pick (weighted) across the other deployments in the same model group before any cross-group fallback runs. Bounded by `max_fallbacks`. Async-only: currently honored by `router.acompletion()` and other async entrypoints. The sync `router.completion()` path falls back to the regular fallback flow. Defaults to False.
+ enable_mid_stream_fallback_continuation (bool): When True, a chat-completions stream that breaks after plain assistant text has been delivered continues on a fallback deployment via assistant prefill instead of surfacing the error. Only deployments whose model supports assistant prefill are eligible, so the partial text is continued, not regenerated; if none is, the original error is surfaced. Streams that emitted tool calls, thinking blocks, audio/images, or constrained (JSON / forced tool_choice) output are never continued. Async-only. Defaults to False.
fallback_access_check (Optional[FallbackAccessCheck]): Awaited before each cross-model-group fallback attempt on the async path; a fallback target it rejects is skipped. Defaults to None (every configured fallback is attempted).
fallback_budget_check (Optional[FallbackBudgetCheck]): Awaited before each cross-model-group fallback attempt on the async path; a fallback target it rejects as over budget is skipped. Defaults to None (budget is not re-checked on fallback).
Returns:
@@ -1308,6 +1385,7 @@ class Router:
self.fallback_on_unsupported_params = fallback_on_unsupported_params
self.enable_health_check_routing = enable_health_check_routing
self.enable_weighted_failover = enable_weighted_failover
+ self.enable_mid_stream_fallback_continuation = self._is_flag_enabled(enable_mid_stream_fallback_continuation)
self.health_check_ignore_transient_errors = health_check_ignore_transient_errors
self.background_health_check_model_groups: frozenset[str] | None = (
frozenset(background_health_check_model_groups)
@@ -1502,6 +1580,14 @@ class Router:
default_pre_call_checks: Final[OptionalPreCallChecks] = []
self.add_optional_pre_call_checks(default_pre_call_checks)
+ if self.enable_mid_stream_fallback_continuation:
+ self._register_continuation_prefill_check()
+
+ @staticmethod
+ def _register_continuation_prefill_check() -> None:
+ # Process-global on purpose: Router.discard() must not drop a filter another router still needs.
+ litellm.logging_callback_manager.add_litellm_callback(ContinuationPrefillDeploymentCheck())
+
def discard(self):
"""
Pseudo-destructor to be invoked to clean up global data structures when router is no longer used.
@@ -3096,15 +3182,53 @@ class Router:
usage: Final = cast(Usage | None, getattr(fallback_item, "usage", None))
usage_objects: Final = [usage] if usage is not None else []
- if (
- complete_response_object_usage is not None
- and hasattr(complete_response_object_usage, "usage")
- and complete_response_object_usage.usage is not None
- ):
+ if complete_response_object_usage is not None:
usage_objects.append(complete_response_object_usage)
combined_usage: Final = BaseTokenUsageProcessor.combine_usage_objects(usage_objects=usage_objects)
setattr(fallback_item, "usage", combined_usage)
+ @staticmethod
+ def _price_cut_partial(
+ cut_stream: CustomStreamWrapper,
+ partial_response: ModelResponse | TextCompletionResponse | None,
+ ) -> "ContinuedPartial | None":
+ """Price the cut stream's partial, plus any partial it inherited, before a fallback can reuse its logging."""
+ from litellm.cost_calculator import BaseTokenUsageProcessor
+
+ cut_logging: Final = cut_stream.logging_obj
+ partial_usage: Final = getattr(partial_response, "usage", None)
+ if partial_response is None or not isinstance(partial_usage, Usage):
+ return None
+ inherited_usage: Final = cut_logging.model_call_details.get("continued_partial_usage")
+ inherited_cost: Final = cut_logging.model_call_details.get("continued_partial_cost")
+ inherited_id: Final = cut_logging.model_call_details.get("continued_response_id")
+ priced_partial: Final = partial_response.model_copy(
+ update={"model": cut_stream.model} if cut_stream.model else {}
+ )
+ priced_partial._hidden_params = { # pyright: ignore[reportPrivateUsage] # no public accessor
+ key: value for key, value in get_hidden_params_dict(partial_response).items() if key != "response_cost"
+ }
+ return ContinuedPartial(
+ usage=BaseTokenUsageProcessor.combine_usage_objects(
+ usage_objects=[partial_usage, inherited_usage]
+ if isinstance(inherited_usage, Usage)
+ else [partial_usage]
+ ),
+ response_cost=(cut_logging.response_cost_calculator(result=priced_partial) or 0.0)
+ + (inherited_cost if isinstance(inherited_cost, float) else 0.0),
+ response_id=inherited_id if isinstance(inherited_id, str) else partial_response.id,
+ )
+
+ @staticmethod
+ def _carry_partial_spend_to_fallback(fallback_response: object, cut_partial: "ContinuedPartial | None") -> None:
+ """Bill what the cut legs consumed on the continuation's success log."""
+ fallback_logging: Final = getattr(fallback_response, "logging_obj", None)
+ if cut_partial is None or not isinstance(fallback_logging, LiteLLMLogging):
+ return
+ fallback_logging.record_continued_partial(
+ usage=cut_partial.usage, response_cost=cut_partial.response_cost, response_id=cut_partial.response_id
+ )
+
@staticmethod
def _prepare_fallback_hidden_params(
fallback_response: object,
@@ -3154,7 +3278,7 @@ class Router:
self,
model_response: CustomStreamWrapper,
messages: list[dict[str, str]],
- initial_kwargs: dict,
+ initial_kwargs: dict[str, Any],
deployment_slot: contextlib.AsyncExitStack | None = None,
first_chunk: _FirstChunk[AsyncIterator[object]] | None = None,
) -> CustomStreamWrapper:
@@ -3217,16 +3341,22 @@ class Router:
with anyio.CancelScope(shield=True):
await close_model_response()
await held_slot.aclose()
- if not e.is_pre_first_chunk and (
- e.generated_content or _stream_chunks_have_generated_content(model_response.chunks)
- ):
- if e.original_exception is not None:
- raise e.original_exception from e
- raise
+ committed: Final = bool(
+ not e.is_pre_first_chunk
+ and (e.generated_content or _stream_chunks_have_generated_content(model_response.chunks))
+ )
+ cut_hidden_params: Final = get_hidden_params(model_response) or {}
+ continue_after_content: Final = committed and self._mid_stream_continuation_eligible(
+ e=e, request_kwargs=initial_kwargs, cut_deployment_id=cut_hidden_params.get("model_id")
+ )
+ if committed and not continue_after_content:
+ self._raise_original_mid_stream_error(e)
from litellm.main import stream_chunk_builder
- complete_response_object: Final = stream_chunk_builder(chunks=model_response.chunks)
+ complete_response_object: Final = stream_chunk_builder(
+ chunks=model_response.chunks, messages=messages if continue_after_content else None
+ )
complete_response_object_usage: Final = cast(
Usage | None,
getattr(complete_response_object, "usage", None),
@@ -3242,7 +3372,41 @@ class Router:
"content_policy_fallbacks", self.content_policy_fallbacks
)
initial_kwargs["original_function"] = self._acompletion
- initial_kwargs["messages"] = messages
+ continuation_messages: Final = (
+ self._build_completion_continuation_input(messages, e.generated_content)
+ if continue_after_content
+ else None
+ )
+ continuation_output: Final = (
+ ContinuationOutput(
+ e.generated_content,
+ continuation_messages[-1],
+ model_group,
+ complete_response_object,
+ checks_post_call_rules=initial_kwargs.get(MID_STREAM_CONTINUATION_KWARG)
+ is not MID_STREAM_CONTINUATION_MARKER,
+ )
+ if continuation_messages is not None
+ else None
+ )
+ cut_partial: Final = (
+ Router._price_cut_partial(model_response, complete_response_object)
+ if continuation_output is not None
+ else None
+ )
+ if continue_after_content:
+ emitted_tokens: Final = int(
+ getattr(complete_response_object_usage, "completion_tokens", 0) or 0
+ )
+ reduced_ceilings: Final = continuation_output_ceilings(initial_kwargs, emitted_tokens)
+ if reduced_ceilings is None or continuation_messages is None:
+ self._raise_original_mid_stream_error(e)
+ await model_response.wait_for_failure_logging(timeout=MID_STREAM_FAILURE_LOG_WAIT_SECONDS)
+ initial_kwargs.update(reduced_ceilings)
+ initial_kwargs["messages"] = continuation_messages
+ initial_kwargs[MID_STREAM_CONTINUATION_KWARG] = MID_STREAM_CONTINUATION_MARKER
+ else:
+ initial_kwargs["messages"] = messages
self._update_kwargs_before_fallbacks(model=model_group, kwargs=initial_kwargs)
fallback_response = await self.async_function_with_fallbacks_common_utils(
e=e,
@@ -3258,11 +3422,17 @@ class Router:
# If fallback returns a streaming response, iterate over it
if hasattr(fallback_response, "__aiter__"):
+ Router._carry_partial_spend_to_fallback(fallback_response, cut_partial)
prepared_fallback_hidden_params = Router._adopt_fallback_response_headers(
wrapper_ref, fallback_response
)
fallback_headers_are_settled = False
- async for fallback_item in fallback_response:
+ fallback_items: Final = (
+ continuation_output.splice(fallback_response)
+ if continuation_output is not None
+ else fallback_response
+ )
+ async for fallback_item in fallback_items:
if not fallback_headers_are_settled:
fallback_headers_are_settled = True
# a fallback that failed over again only repoints itself once it yields
@@ -3275,7 +3445,10 @@ class Router:
and isinstance(fallback_item, ModelResponseStream)
and hasattr(fallback_item, "usage")
):
- self._combine_fallback_usage(fallback_item, complete_response_object_usage)
+ self._combine_fallback_usage(
+ fallback_item,
+ complete_response_object_usage if continuation_output is not None else None,
+ )
yield fallback_item
else:
# If fallback returns a non-streaming response, yield None
@@ -3432,6 +3605,97 @@ class Router:
total_tokens=(partial_usage.total_tokens or 0) + (fb.total_tokens or 0),
)
+ @staticmethod
+ def _raise_original_mid_stream_error(e: "MidStreamFallbackError") -> NoReturn:
+ """Surface the provider error the stream wrapper carried instead of the internal MidStreamFallbackError."""
+ if e.original_exception is not None:
+ raise e.original_exception from e
+ raise e
+
+ def _params_sent_to_deployment(
+ self, request_kwargs: Mapping[str, object], deployment: Deployment | None
+ ) -> Mapping[str, object]:
+ """The request as the deployment received it, merged in the order _update_kwargs_with_deployment applies."""
+ deployment_params: Final = deployment.litellm_params.model_dump(exclude_none=True) if deployment else {}
+ defaults: Final = {key: value for key, value in self.default_litellm_params.items() if value is not None}
+ deployment_tool_choice: Final = (
+ {"tool_choice": deployment_params["tool_choice"]}
+ if deployment_params.get("tool_choice") is not None
+ else {}
+ )
+ return MappingProxyType({**deployment_params, **defaults, **deployment_tool_choice, **request_kwargs})
+
+ def _names_deployment_that_cannot_continue(self, target: object) -> bool:
+ name: Final = fallback_target_name(target)
+ if name is None or name in self.model_names or not self.has_model_id(name):
+ return False
+ deployment: Final = self.get_deployment(model_id=name)
+ return deployment is not None and not deployment_can_continue(deployment.model_dump(exclude_none=True))
+
+ def _mid_stream_continuation_eligible(
+ self,
+ e: "MidStreamFallbackError",
+ request_kwargs: Mapping[str, object],
+ cut_deployment_id: object = None,
+ ) -> bool:
+ """Whether a stream that broke after plain assistant text may continue via prefill."""
+ if not self.enable_mid_stream_fallback_continuation:
+ return False
+ if not e.generated_content or e.emitted_disqualifying_content:
+ return False
+ if isinstance(e.original_exception, litellm.ContentPolicyViolationError):
+ return False
+ cut_deployment: Final = (
+ self.get_deployment(model_id=cut_deployment_id) if isinstance(cut_deployment_id, str) else None
+ )
+ if isinstance(cut_deployment_id, str) and cut_deployment is None:
+ return False
+ if constrains_continuation(self._params_sent_to_deployment(request_kwargs, cut_deployment)):
+ return False
+ requested_group: Final = request_kwargs.get("model")
+ lookup_groups: Final = fallback_lookup_groups(
+ request_kwargs, requested_group if isinstance(requested_group, str) else None
+ )
+ targets: Final = continuation_fallback_targets(request_kwargs.get("fallbacks", self.fallbacks), lookup_groups)
+ if any(fallback_target_overrides_request_params(target) for target in targets):
+ return False
+ if any(self._names_deployment_that_cannot_continue(target) for target in targets):
+ return False
+ return not any(
+ isinstance(callback, CustomGuardrail)
+ and callback._event_hook_is_event_type( # pyright: ignore[reportPrivateUsage] # the only check of a guardrail's hooks
+ GuardrailEventHooks.post_call
+ )
+ for callback in litellm.callbacks
+ )
+
+ @staticmethod
+ def _build_completion_continuation_input(
+ messages: Sequence[Mapping[str, object]],
+ generated_content: str,
+ ) -> Sequence[Mapping[str, object]] | None:
+ """Append the partial output as an assistant prefill, or extend a trailing `prefix: True` assistant turn in
+ place so the request never ends in two assistant messages. None when the request ends in a plain assistant
+ turn, in a prefill turn whose content has a shape this cannot extend, or when the cut came mid-echo of the
+ prefill."""
+ last: Final = messages[-1] if messages else None
+ if last is None or last.get("role") != "assistant":
+ return [*messages, {"role": "assistant", "content": generated_content, "prefix": True}]
+ if last.get("prefix") is not True:
+ return None
+ new_text: Final = continued_text(last, generated_content)
+ if new_text is None:
+ return None
+ content: Final = last.get("content")
+ if content is None or isinstance(content, str):
+ return [*messages[:-1], {**last, "content": f"{content or ''}{new_text}", "prefix": True}]
+ try:
+ blocks: Final = _CONTENT_BLOCKS_ADAPTER.validate_python(content)
+ except ValidationError:
+ return None
+ extended: Final = _extend_trailing_text_block(blocks, new_text)
+ return [*messages[:-1], {**last, "content": extended, "prefix": True}]
+
@staticmethod
def _build_responses_continuation_input(
input_val: Union[str, "ResponseInputParam"] | None,
@@ -3833,13 +4097,6 @@ class Router:
raise e.original_exception from e
raise
- from litellm.main import stream_chunk_builder
-
- complete_response_object: Final = stream_chunk_builder(chunks=model_response.chunks)
- complete_response_object_usage: Final = cast(
- Usage | None,
- getattr(complete_response_object, "usage", None),
- )
try:
model_group: Final = cast(str, initial_kwargs.get("model"))
fallbacks: Final[list | None] = initial_kwargs.get("fallbacks", router_self.fallbacks)
@@ -3885,7 +4142,7 @@ class Router:
and isinstance(fallback_item, ModelResponseStream)
and hasattr(fallback_item, "usage")
):
- router_self._combine_fallback_usage(fallback_item, complete_response_object_usage)
+ router_self._combine_fallback_usage(fallback_item, None)
yield fallback_item
else:
yield None
@@ -11343,6 +11600,7 @@ class Router:
Augment litellm info with additional params set in `model_info`.
For azure models, ignore the `model:`. Only set max tokens, cost values if base_model is set.
+ For every other provider, a base_model naming the same model as `model:` can only widen its context window.
Returns
- ModelInfo - If found -> typed dict with max tokens, input cost, etc.
@@ -11439,20 +11697,23 @@ class Router:
## CHECK USER SET MODEL INFO
raw_user_model_info: Final = deployment.get("model_info")
- user_model_info: Final = (
+ user_model_info: Final = _MODEL_INFO_ADAPTER.validate_python(
raw_user_model_info.model_dump(exclude_none=True)
if isinstance(raw_user_model_info, BaseModel)
- else raw_user_model_info
+ else raw_user_model_info or {}
)
# get_model_info() hands back an lru_cache'd dict, so merge into a copy; unset
# values are skipped or Deployment's None pricing defaults would erase the map's
merged_model_info: Final[ModelMapInfo] = {
**copy.deepcopy(model_info),
- **self.get_discovered_model_info((deployment.get("model_info") or {}).get("id")),
- **MappingProxyType(
- {key: value for key, value in (user_model_info or MappingProxyType({})).items() if value is not None}
+ **_base_model_widened_window(
+ model_info,
+ _declared_base_model(user_model_info) or _declared_base_model(litellm_params.model_extra or {}),
+ custom_llm_provider,
),
+ **self.get_discovered_model_info((deployment.get("model_info") or {}).get("id")),
+ **MappingProxyType({key: value for key, value in user_model_info.items() if value is not None}),
}
return merged_model_info
@@ -12642,16 +12903,34 @@ class Router:
str(deployment_params.get("model") or ""), deployment_params.get("custom_llm_provider")
):
return True
+ return Router._deployment_supports_param(deployment, group, param)
+
+ @staticmethod
+ def _capability_model_and_provider(deployment_params: LiteLLMParamsTypedDict, group: str) -> tuple[str, str]:
+ model_label: Final = str(deployment_params.get("model") or group)
+ authenticating_provider: Final = declared_authenticating_provider(
+ model_label, deployment_params.get("custom_llm_provider")
+ )
+ if authenticating_provider is not None:
+ return model_label.removeprefix(f"{authenticating_provider}/"), authenticating_provider
+ model, custom_llm_provider, _, _ = litellm.get_llm_provider(
+ model=model_label,
+ custom_llm_provider=deployment_params.get("custom_llm_provider"),
+ litellm_params=LiteLLM_Params.model_validate(deployment_params),
+ )
+ return model, custom_llm_provider
+
+ @staticmethod
+ def _deployment_supports_param(deployment: DeploymentTypedDict, group: str, param: str) -> bool:
+ deployment_params: Final = deployment.get("litellm_params")
+ if not deployment_params:
+ return True
deployment_model_info: Final = deployment.get("model_info")
base_model: Final = (
deployment_model_info.get("base_model") if deployment_model_info else None
) or deployment_params.get("base_model")
try:
- model, custom_llm_provider, _, _ = litellm.get_llm_provider(
- model=deployment_params.get("model") or group,
- custom_llm_provider=deployment_params.get("custom_llm_provider"),
- litellm_params=LiteLLM_Params.model_validate(deployment_params),
- )
+ model, custom_llm_provider = Router._capability_model_and_provider(deployment_params, group)
supported: Final = litellm.get_supported_openai_params(
model=model,
custom_llm_provider=custom_llm_provider,
@@ -12659,7 +12938,7 @@ class Router:
)
except Exception as e: # noqa: BLE001 # best-effort filter: an unresolvable provider must not narrow the request
verbose_router_logger.debug(
- "litellm.router.py::_deployment_accepts_param: keeping %s for model=%s. Got - %s", param, group, e
+ "litellm.router.py::_deployment_supports_param: keeping %s for model=%s. Got - %s", param, group, e
)
return True
return supported is None or param in supported
@@ -12935,6 +13214,7 @@ class Router:
"retry_policy",
"model_group_alias",
"enable_weighted_failover",
+ "enable_mid_stream_fallback_continuation",
"enable_tag_filtering",
"tag_routing_prefix",
"fallback_on_unsupported_params",
@@ -12991,12 +13271,10 @@ class Router:
rebuild_routing_groups = True
elif var == "optional_pre_call_checks":
self.set_optional_pre_call_checks(kwargs[var])
+ elif var == "enable_mid_stream_fallback_continuation":
+ self._set_mid_stream_fallback_continuation(kwargs[var])
elif var == "retry_policy":
- value = kwargs[var]
- if isinstance(value, dict):
- value = RetryPolicy(**value)
- if value is None or isinstance(value, RetryPolicy):
- setattr(self, var, value)
+ self._set_retry_policy(kwargs[var])
else:
value = kwargs[var]
# only run routing strategy init if it has changed
@@ -13023,6 +13301,20 @@ class Router:
self._routing_groups_input = routing_groups_input
verbose_router_logger.debug("Updated Router settings: %s", self.get_settings())
+ @staticmethod
+ def _is_flag_enabled(value: object) -> bool:
+ return value is True or (isinstance(value, str) and str_to_bool(value) is True)
+
+ def _set_mid_stream_fallback_continuation(self, enabled: object) -> None:
+ self.enable_mid_stream_fallback_continuation = self._is_flag_enabled(enabled)
+ if self.enable_mid_stream_fallback_continuation:
+ self._register_continuation_prefill_check()
+
+ def _set_retry_policy(self, value: object) -> None:
+ policy: Final = RetryPolicy.model_validate(value) if isinstance(value, dict) else value
+ if policy is None or isinstance(policy, RetryPolicy):
+ self.retry_policy = policy
+
def _get_client(self, deployment, kwargs, client_type=None):
"""
Returns the appropriate client based on the given deployment, kwargs, and client_type.
@@ -13156,7 +13448,8 @@ class Router:
Returns None when no deployment limits its context window, and when counting fails. The
caller pairs this with ``skip_inline_token_count`` so neither case puts the count back on
- the loop: a failed count leaves the deployments unfiltered, exactly as before.
+ the loop: a failed count only skips the context-window check - RPM, region and
+ invalid-params checks still run for every deployment.
"""
if messages is None and input is None:
return None
@@ -13174,6 +13467,66 @@ class Router:
)
return None
+ def _resolve_pre_call_check_input_tokens(
+ self,
+ messages: Sequence[Mapping[str, str]] | None,
+ input: str | Sequence[object] | None,
+ request_kwargs: Mapping[str, object] | None,
+ input_token_count: int | None,
+ count_inline: bool,
+ ) -> int | None:
+ if input_token_count is not None:
+ return input_token_count
+ if not count_inline:
+ return None
+ try:
+ return self._count_pre_call_check_tokens(
+ messages=cast(list[dict[str, str]] | None, messages), # cast-ok: forwarded to the sync counter
+ input=cast(str | list[object] | None, input), # cast-ok: forwarded to the sync counter
+ request_kwargs=request_kwargs,
+ )
+ except Exception as e:
+ verbose_router_logger.error(
+ "litellm.router.py::_pre_call_checks: failed to count tokens. "
+ "Skipping the context-window check only; RPM/region/param checks still apply. Got - %s",
+ e,
+ )
+ return None
+
+ @staticmethod
+ def _params_neither_dropped_nor_allowed(params: Mapping[str, object], candidates: frozenset[str]) -> frozenset[str]:
+ if params.get("drop_params") is True:
+ return frozenset()
+ dropped: Final = params.get("additional_drop_params")
+ dropped_names: Final = (
+ frozenset(entry for entry in dropped if isinstance(entry, str))
+ if isinstance(dropped, list)
+ else frozenset()
+ )
+ return candidates - Router._declared_param_allowlist(params) - dropped_names
+
+ @staticmethod
+ def _deployment_lacks_params(deployment: DeploymentTypedDict, group: str, params: frozenset[str]) -> bool:
+ return any(not Router._deployment_supports_param(deployment, group, param) for param in params)
+
+ @staticmethod
+ def _deployment_rejects_params(
+ deployment: DeploymentTypedDict, group: str, params: frozenset[str], request_kwargs: Mapping[str, object]
+ ) -> bool:
+ call_settings: Final[Mapping[str, object]] = {**(deployment.get("litellm_params") or {}), **request_kwargs}
+ return Router._deployment_lacks_params(
+ deployment, group, Router._params_neither_dropped_nor_allowed(call_settings, params)
+ )
+
+ @staticmethod
+ def _pre_call_check_requested_chat_params(request_kwargs: Mapping[str, object]) -> frozenset[str]:
+ requested: Final = litellm.utils.get_non_default_params(passed_params=dict(request_kwargs))
+ return frozenset(
+ param
+ for param, value in requested.items()
+ if param in _PRE_CALL_CHECK_FILTERED_PARAMS and not isinstance(value, str)
+ )
+
def _pre_call_checks(
self,
model: str,
@@ -13183,6 +13536,7 @@ class Router:
request_kwargs: dict | None = None,
input_token_count: int | None = None,
skip_inline_token_count: bool = False,
+ request_deployments: Sequence[DeploymentTypedDict] | None = None,
):
"""
Filter out model in model group, if:
@@ -13201,13 +13555,6 @@ class Router:
invalid_model_indices: Final = set() # Use set for O(1) membership checks
- # Token counting (tiktoken) is the dominant on-loop cost for large prompts.
- # Only count when a deployment actually declares max_input_tokens, and count
- # at most once; for model groups with no context-window limit it is skipped.
- # Async callers pass the count in, already computed off the event loop, and set
- # skip_inline_token_count so a failed off-loop count is not retried back on the loop.
- input_tokens: int | None = input_token_count
-
_context_window_error = False
_potential_error_str = ""
_rate_limit_error = False
@@ -13216,6 +13563,24 @@ class Router:
has_countable_input: Final = (messages is not None or input is not None) and not compaction_pending(
request_kwargs
)
+ deployment_max_input_tokens: Final = tuple(
+ self._deployment_max_input_tokens(model, deployment) for deployment in _returned_deployments
+ )
+ input_tokens: Final = self._resolve_pre_call_check_input_tokens(
+ messages=messages,
+ input=cast(str | list[object] | None, input), # cast-ok: forwarded to the resolver
+ request_kwargs=cast(Mapping[str, object] | None, request_kwargs), # cast-ok: forwarded to the resolver
+ input_token_count=input_token_count,
+ count_inline=has_countable_input
+ and not skip_inline_token_count
+ and any(limit is not None for limit in deployment_max_input_tokens),
+ )
+ request_settings: Final[Mapping[str, object]] = request_kwargs if request_kwargs is not None else {}
+ requested_params: Final[frozenset[str]] = (
+ self._pre_call_check_requested_chat_params(request_settings)
+ if request_kwargs is not None and litellm.drop_params is False
+ else frozenset()
+ )
## get model group RPM ##
dt: Final = get_utc_datetime()
@@ -13229,39 +13594,22 @@ class Router:
_litellm_params = deployment.get("litellm_params", {})
_model_info = deployment.get("model_info", {})
- # see if we have the info for this model
- _deployment_model = None # per-deployment model name (avoids overwriting the outer `model` group name)
- try:
- base_model = _model_info.get("base_model", None)
- if base_model is None:
- base_model = _litellm_params.get("base_model", None)
- _deployment_model = base_model or _litellm_params.get("model", None)
- model_info = self.get_router_model_info(deployment=deployment, received_model_name=model)
-
- max_input_tokens = model_info.get("max_input_tokens") if isinstance(model_info, dict) else None
- if isinstance(max_input_tokens, int) and has_countable_input:
- if input_tokens is None:
- if skip_inline_token_count:
- return _returned_deployments
- try:
- input_tokens = self._count_pre_call_check_tokens(
- messages=messages, input=input, request_kwargs=request_kwargs
- )
- except Exception as e:
- verbose_router_logger.error(
- "litellm.router.py::_pre_call_checks: failed to count tokens. Returning initial list of deployments. Got - %s",
- e,
- )
- return _returned_deployments
- if input_tokens > max_input_tokens:
- invalid_model_indices.add(idx)
- _context_window_error = True
- _potential_error_str += (
- f"Model={_deployment_model}, Max Input Tokens={max_input_tokens}, Got={input_tokens}"
- )
- continue
- except Exception as e:
- verbose_router_logger.exception("An error occurs - %s", e)
+ max_input_tokens = deployment_max_input_tokens[idx]
+ if (
+ max_input_tokens is not None
+ and has_countable_input
+ and input_tokens is not None
+ and input_tokens > max_input_tokens
+ ):
+ invalid_model_indices.add(idx)
+ _context_window_error = True
+ deployment_model = (
+ _model_info.get("base_model") or _litellm_params.get("base_model") or _litellm_params.get("model")
+ )
+ _potential_error_str += (
+ f"Model={deployment_model}, Max Input Tokens={max_input_tokens}, Got={input_tokens}"
+ )
+ continue
model_id = _model_info.get("id", "")
## RPM CHECK ##
@@ -13293,53 +13641,57 @@ class Router:
invalid_model_indices.add(idx)
continue
- ## INVALID PARAMS ## -> catch 'gpt-3.5-turbo-16k' not supporting 'response_format' param
- if request_kwargs is not None and litellm.drop_params is False:
- # get supported params — use per-deployment model to avoid overwriting the outer model group name
- _dep_model_for_params: str = _deployment_model or model
- try:
- (
- _dep_model_for_params,
- custom_llm_provider,
- _,
- _,
- ) = litellm.get_llm_provider(
- model=_dep_model_for_params,
- litellm_params=LiteLLM_Params.model_validate(_litellm_params),
- )
- except Exception as e: # noqa: BLE001 # best-effort filter: an unresolvable provider must not fail the request
- verbose_router_logger.debug(
- "litellm.router.py::_pre_call_checks: skipping supported-params check for model=%s. Got - %s",
- _dep_model_for_params,
- e,
- )
- continue
+ typed_deployments: Final = tuple(
+ cast(DeploymentTypedDict, deployment) # cast-ok: router deployments are plain dicts
+ for deployment in _returned_deployments
+ )
+ param_candidates: Final = tuple(
+ (idx, deployment) for idx, deployment in enumerate(typed_deployments) if idx not in invalid_model_indices
+ )
+ candidates_lacking_params: Final = tuple(
+ self._deployment_lacks_params(deployment, model, requested_params) for _, deployment in param_candidates
+ )
+ no_deployment_takes_params: Final = bool(candidates_lacking_params) and all(candidates_lacking_params)
+ candidates_failing_params: Final = (
+ tuple(
+ self._deployment_rejects_params(deployment, model, requested_params, request_settings)
+ for _, deployment in param_candidates
+ )
+ if no_deployment_takes_params
+ else candidates_lacking_params
+ )
+ excluded_indices: Final = frozenset(invalid_model_indices) | frozenset(
+ idx for (idx, _), fails in zip(param_candidates, candidates_failing_params, strict=True) if fails
+ )
- supported_openai_params = litellm.get_supported_openai_params(
- model=_dep_model_for_params,
- custom_llm_provider=custom_llm_provider,
- )
-
- if supported_openai_params is None:
- continue
- else:
- # check the non-default openai params in request kwargs
- non_default_params = litellm.utils.get_non_default_params(passed_params=request_kwargs)
- special_params = ["response_format"]
- # check if all params are supported
- for k, v in non_default_params.items():
- if k not in supported_openai_params and k in special_params:
- # if not -> invalid model
- verbose_router_logger.debug("INVALID MODEL INDEX @ REQUEST KWARG FILTERING, k=%s", k)
- invalid_model_indices.add(idx)
-
- if len(invalid_model_indices) == len(_returned_deployments):
+ if len(excluded_indices) == len(_returned_deployments):
"""
- no healthy deployments available b/c context window checks or rate limit error
- First check for rate limit errors (if this is true, it means the model passed the context window check but failed the rate limit check)
"""
+ group_deployments: Final = (
+ tuple(request_deployments)
+ if request_deployments is not None
+ else (tuple(self.get_all_deployments(model_name=model)) or typed_deployments)
+ )
+ if (
+ not _context_window_error
+ and group_deployments
+ and all(
+ self._deployment_rejects_params(deployment, model, requested_params, request_settings)
+ for deployment in group_deployments
+ )
+ ):
+ raise litellm.UnsupportedParamsError(
+ message=(
+ f"no deployment of model={model} supports {sorted(requested_params)}. "
+ "To drop these, set `litellm.drop_params=True` or for proxy `litellm_settings: drop_params: true`, "
+ f"or send allowed_openai_params={sorted(requested_params)} to pass them through"
+ ),
+ model=model,
+ )
if _rate_limit_error is True: # allow generic fallback logic to take place
raise RouterRateLimitErrorBasic(
model=model,
@@ -13351,9 +13703,9 @@ class Router:
model=model,
llm_provider="",
)
- if len(invalid_model_indices) > 0:
+ if len(excluded_indices) > 0:
# Single-pass filter using set for O(1) lookups (avoids O(n^2) from repeated pops)
- _returned_deployments = [d for i, d in enumerate(_returned_deployments) if i not in invalid_model_indices]
+ _returned_deployments = [d for i, d in enumerate(_returned_deployments) if i not in excluded_indices]
return _returned_deployments
@@ -13776,6 +14128,10 @@ class Router:
)
return healthy_deployments
+ request_deployments: Final = cast( # cast-ok: router deployments are plain dicts
+ Sequence[DeploymentTypedDict], healthy_deployments
+ )
+
# Health-check-based filtering (before cooldown)
healthy_deployments = await self._async_filter_health_check_unhealthy_deployments(
healthy_deployments=healthy_deployments,
@@ -13835,6 +14191,7 @@ class Router:
request_kwargs=request_kwargs,
),
skip_inline_token_count=True,
+ request_deployments=request_deployments,
)
# check if user wants to do tag based routing
healthy_deployments = await get_deployments_for_tag(
@@ -14985,9 +15342,8 @@ class Router:
request_kwargs=request_kwargs,
)
strategy, strategy_selector = self._get_routing_context(model, request_kwargs)
- pick_attributes: Final = _deployment_pick_attributes(
- model, _ROUTING_KWARGS_ADAPTER.validate_python(request_kwargs)
- )
+ routing_kwargs: Final = _ROUTING_KWARGS_ADAPTER.validate_python(request_kwargs)
+ pick_attributes: Final = _deployment_pick_attributes(model, routing_kwargs)
if isinstance(healthy_deployments, dict):
if (healthy_deployments.get("model_info") or {}).get("blocked") is True:
@@ -15001,6 +15357,9 @@ class Router:
return healthy_deployments
parent_otel_span: Final[Span | None] = get_parent_otel_span_from_kwargs(request_kwargs)
+ request_deployments: Final = cast( # cast-ok: router deployments are plain dicts
+ Sequence[DeploymentTypedDict], healthy_deployments
+ )
# Health-check-based filtering (before cooldown)
healthy_deployments = self._filter_health_check_unhealthy_deployments(
@@ -15022,7 +15381,9 @@ class Router:
)
healthy_deployments = _pre_cooldown_deployments
- healthy_deployments = self._filter_blocked_deployments(healthy_deployments)
+ healthy_deployments = continuation_deployments(
+ self._filter_blocked_deployments(healthy_deployments), routing_kwargs
+ )
# filter pre-call checks
if self.enable_pre_call_checks and (messages is not None or input is not None):
@@ -15032,6 +15393,7 @@ class Router:
messages=messages,
input=input,
request_kwargs=request_kwargs,
+ request_deployments=request_deployments,
)
## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2)
@@ -15195,6 +15557,9 @@ class Router:
# 4. Apply health-check and cooldown filtering
parent_otel_span: Final[Span | None] = get_parent_otel_span_from_kwargs(request_kwargs)
+ request_deployments: Final = cast( # cast-ok: router deployments are plain dicts
+ Sequence[DeploymentTypedDict], pass_through_deployments
+ )
pass_through_deployments = self._filter_health_check_unhealthy_deployments(
healthy_deployments=pass_through_deployments,
parent_otel_span=parent_otel_span,
@@ -15216,6 +15581,7 @@ class Router:
messages=messages,
input=input,
request_kwargs=request_kwargs,
+ request_deployments=request_deployments,
)
if self._is_priority_routing_group(model):
diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py
index b20f71cd51f..1622a30b95f 100644
--- a/litellm/router_strategy/budget_limiter.py
+++ b/litellm/router_strategy/budget_limiter.py
@@ -405,7 +405,8 @@ class RouterBudgetLimiting(CustomLogger):
Runs once the budget start time exists in Redis Cache (on the 2nd and subsequent requests to the same provider)
- Increments the spend in memory cache (so spend instantly updated in memory)
- - Queues the increment operation to Redis Pipeline (using batched pipeline to optimize performance. Using Redis for multi instance environment of LiteLLM)
+ - Queues the increment for the batched Redis pipeline when Redis is configured, since only a multi instance
+ deployment has anything to sync
"""
increment_op: Final = RedisPipelineIncrementOperation(
key=spend_key,
@@ -418,6 +419,8 @@ class RouterBudgetLimiting(CustomLogger):
value=response_cost,
ttl=ttl,
)
+ if self.dual_cache.redis_cache is None:
+ return
self.redis_increment_operation_queue.append(increment_op)
def _get_redis_increment_queue_lock(self) -> asyncio.Lock:
diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py
index e12304c6ed1..c4e4d9ccc0b 100644
--- a/litellm/router_utils/common_utils.py
+++ b/litellm/router_utils/common_utils.py
@@ -27,7 +27,8 @@ from litellm.types.utils import LlmProviders, ProviderSpecificModelInfo
_V = TypeVar("_V")
-ROUTER_ONLY_CALL_KWARGS: Final = frozenset({"silent_model", "include_fallback_errors"})
+MID_STREAM_CONTINUATION_KWARG: Final = "_mid_stream_continuation"
+ROUTER_ONLY_CALL_KWARGS: Final = frozenset({"silent_model", "include_fallback_errors", MID_STREAM_CONTINUATION_KWARG})
_OBJECT_MAPPING_ADAPTER: Final[TypeAdapter[Mapping[object, object]]] = TypeAdapter(Mapping[object, object])
diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py
index 1f37c6c8cb1..06bd3329cc1 100644
--- a/litellm/router_utils/fallback_event_handlers.py
+++ b/litellm/router_utils/fallback_event_handlers.py
@@ -615,6 +615,32 @@ def get_fallback_model_group(
return fallback_model_group, generic_fallback_idx
+def _resolved_fallback_chain(entries: Sequence[object], lookup_groups: tuple[str, ...]) -> Sequence[object]:
+ group_entries: Final[Sequence[Mapping[str, list[str]]]] = tuple( # mutable-ok: mirrors the fallbacks config shape
+ entry for entry in entries if isinstance(entry, Mapping) and entry
+ )
+ chain, _ = get_fallback_model_group_for_lookup_groups(fallbacks=group_entries, lookup_groups=lookup_groups)
+ return chain if isinstance(chain, list) else ()
+
+
+def continuation_fallback_targets(fallbacks: object, lookup_groups: tuple[str, ...]) -> Sequence[object]:
+ entries: Final[Sequence[object]] = fallbacks if isinstance(fallbacks, list) else ()
+ return (
+ entries
+ if check_non_standard_fallback_format(fallbacks=entries)
+ else _resolved_fallback_chain(entries, lookup_groups)
+ )
+
+
+def fallback_target_overrides_request_params(target: object) -> bool:
+ return isinstance(target, Mapping) and "model" in target and len(target) > 1
+
+
+def fallback_target_name(target: object) -> str | None:
+ name: Final = target.get("model") if isinstance(target, Mapping) else target
+ return name if isinstance(name, str) else None
+
+
PROVIDER_SCOPED_RESOURCE_KEYS: Final = ("input_file_id", "training_file", "batch_id", "file_id", "fine_tuning_job_id")
PROVIDER_SCOPED_RESOURCE_FUNCTION_NAMES: Final = frozenset(
{
diff --git a/litellm/router_utils/pre_call_checks/continuation_prefill_check.py b/litellm/router_utils/pre_call_checks/continuation_prefill_check.py
new file mode 100644
index 00000000000..8f211ef67db
--- /dev/null
+++ b/litellm/router_utils/pre_call_checks/continuation_prefill_check.py
@@ -0,0 +1,264 @@
+"""Keeps a mid-stream continuation on deployments whose model supports assistant prefill and whose own params leave
+the output unconstrained."""
+
+from collections.abc import AsyncIterable, AsyncIterator, Mapping, Sequence
+from dataclasses import dataclass
+from itertools import accumulate
+from os.path import commonprefix
+from types import MappingProxyType
+from typing import Final
+
+from pydantic import TypeAdapter, ValidationError
+
+import litellm
+from litellm.integrations.custom_logger import CustomLogger, Span
+from litellm.litellm_core_utils.rules import Rules
+from litellm.router_utils.common_utils import MID_STREAM_CONTINUATION_KWARG
+from litellm.types.llms.openai import AllMessageValues
+from litellm.types.utils import ModelResponse, ModelResponseStream, Usage
+from litellm.utils import supports_assistant_prefill
+
+
+class _ContinuationMarker:
+ """Only the router constructs one, so a JSON request body cannot forge the continuation flag."""
+
+
+MID_STREAM_CONTINUATION_MARKER: Final = _ContinuationMarker()
+
+_STR_KEYED_DICT_ADAPTER: Final = TypeAdapter(dict[str, object])
+_NO_FIELDS: Final[Mapping[str, object]] = MappingProxyType({})
+_UNCONSTRAINED_RESPONSE_FORMAT: Final[Mapping[str, str]] = MappingProxyType({"type": "text"})
+_OUTPUT_CEILING_KEYS: Final = ("max_tokens", "max_completion_tokens")
+
+
+def _str_keyed(value: object) -> Mapping[str, object]:
+ try:
+ return _STR_KEYED_DICT_ADAPTER.validate_python(value)
+ except ValidationError:
+ return _NO_FIELDS
+
+
+def deployment_supports_prefill(deployment: object) -> bool:
+ deployment_map: Final = _str_keyed(deployment)
+ model_info: Final = _str_keyed(deployment_map.get("model_info"))
+ declared: Final = model_info.get("supports_assistant_prefill")
+ if isinstance(declared, bool):
+ return declared
+ litellm_params: Final = _str_keyed(deployment_map.get("litellm_params"))
+ candidates: Final = (model_info.get("base_model"), litellm_params.get("base_model"), litellm_params.get("model"))
+ model: Final = next((candidate for candidate in candidates if isinstance(candidate, str) and candidate), None)
+ return model is not None and supports_assistant_prefill(model=model)
+
+
+def constrains_continuation(params: Mapping[str, object]) -> bool:
+ sent: Final = {**params, **_str_keyed(params.get("extra_body"))}
+ choice_count: Final = sent.get("n")
+ response_format: Final = sent.get("response_format")
+ tool_choice: Final = sent.get("tool_choice")
+ return (
+ (isinstance(choice_count, int) and choice_count > 1)
+ or (response_format is not None and response_format != _UNCONSTRAINED_RESPONSE_FORMAT)
+ or tool_choice == "required"
+ or isinstance(tool_choice, Mapping)
+ or sent.get("merge_reasoning_content_in_choices") is True
+ )
+
+
+def deployment_can_continue(deployment: object) -> bool:
+ litellm_params: Final = _str_keyed(_str_keyed(deployment).get("litellm_params"))
+ return deployment_supports_prefill(deployment) and not constrains_continuation(litellm_params)
+
+
+def continuation_deployments(
+ healthy_deployments: list[dict[str, object]], # mutable-ok: router deployment-list contract
+ request_kwargs: Mapping[str, object] | None,
+) -> list[dict[str, object]]: # mutable-ok: returns a mutable deployment list
+ marker: Final = (request_kwargs or {}).get(MID_STREAM_CONTINUATION_KWARG)
+ if not isinstance(marker, _ContinuationMarker):
+ return healthy_deployments
+ return [deployment for deployment in healthy_deployments if deployment_can_continue(deployment)]
+
+
+class ContinuationPrefillDeploymentCheck(CustomLogger):
+ async def async_filter_deployments(
+ self,
+ model: str,
+ healthy_deployments: list[dict[str, object]], # mutable-ok: CustomLogger deployment-list contract
+ messages: Sequence[AllMessageValues] | None,
+ request_kwargs: Mapping[str, object] | None = None,
+ parent_otel_span: Span | None = None,
+ ) -> list[dict[str, object]]: # mutable-ok: returns a mutable deployment list
+ return continuation_deployments(healthy_deployments, request_kwargs)
+
+
+def _delta_text(chunk: ModelResponseStream) -> str:
+ if not chunk.choices:
+ return ""
+ content: Final = chunk.choices[0].delta.content
+ return content if isinstance(content, str) else ""
+
+
+_CONTENT_BLOCKS_ADAPTER: Final = TypeAdapter(list[dict[str, object]])
+
+
+def _prefill_text(prefill_message: Mapping[str, object]) -> str:
+ content: Final = prefill_message.get("content")
+ if isinstance(content, str):
+ return content
+ try:
+ blocks: Final = _CONTENT_BLOCKS_ADAPTER.validate_python(content)
+ except ValidationError:
+ return ""
+ return "".join(
+ text for block in blocks if block.get("type") == "text" and isinstance(text := block.get("text"), str)
+ )
+
+
+def _resendable_whitespace(prefill: str) -> str:
+ return prefill[len(prefill.rstrip()) :] if prefill.strip() else ""
+
+
+def continued_text(prefill_message: Mapping[str, object], generated_content: str) -> str | None:
+ prefill: Final = _prefill_text(prefill_message)
+ if generated_content and len(generated_content) < len(prefill) and prefill.startswith(generated_content):
+ return None
+ echo: Final = len(prefill) if prefill and generated_content.startswith(prefill) else 0
+ after_echo: Final = generated_content[echo:]
+ return after_echo[len(commonprefix((after_echo, _resendable_whitespace(prefill)))) :]
+
+
+def _reduced_ceilings(params: Mapping[str, object], emitted_tokens: int) -> Mapping[str, int]:
+ return MappingProxyType(
+ {key: value - emitted_tokens for key in _OUTPUT_CEILING_KEYS if isinstance(value := params.get(key), int)}
+ )
+
+
+def continuation_output_ceilings(
+ request_kwargs: Mapping[str, object], emitted_tokens: int
+) -> Mapping[str, object] | None:
+ """The caller's output ceilings, top-level and in extra_body, minus the tokens already emitted, or None once
+ nothing is left."""
+ extra_body: Final = _str_keyed(request_kwargs.get("extra_body"))
+ top_level: Final = _reduced_ceilings(request_kwargs, emitted_tokens)
+ in_extra_body: Final = _reduced_ceilings(extra_body, emitted_tokens)
+ if any(ceiling <= 0 for ceiling in (*top_level.values(), *in_extra_body.values())):
+ return None
+ if not in_extra_body:
+ return top_level
+ return MappingProxyType({**top_level, "extra_body": {**extra_body, **in_extra_body}})
+
+
+def _with_first_delta(chunk: ModelResponseStream, update: Mapping[str, object]) -> ModelResponseStream:
+ choice: Final = chunk.choices[0]
+ delta: Final = choice.delta.model_copy(update=update)
+ return chunk.model_copy(update={"choices": [choice.model_copy(update={"delta": delta}), *chunk.choices[1:]]})
+
+
+def _without_leading_text(chunk: ModelResponseStream, count: int) -> ModelResponseStream:
+ if count <= 0:
+ return chunk
+ return _with_first_delta(chunk, {"content": _delta_text(chunk)[count:]})
+
+
+def _without_role(chunk: ModelResponseStream) -> ModelResponseStream:
+ if not chunk.choices or chunk.choices[0].delta.role is None:
+ return chunk
+ return _with_first_delta(chunk, {"role": None})
+
+
+def _stream_identity(primary_response: object) -> Mapping[str, object]:
+ if not isinstance(primary_response, ModelResponse):
+ return _NO_FIELDS
+ return MappingProxyType({"id": primary_response.id, "created": primary_response.created})
+
+
+def _drop_leading_text(chunks: tuple[ModelResponseStream, ...], count: int) -> tuple[ModelResponseStream, ...]:
+ lengths: Final = tuple(len(_delta_text(chunk)) for chunk in chunks)
+ offsets: Final = tuple(accumulate(lengths, initial=0))[:-1]
+ return tuple(
+ _without_leading_text(chunk, min(length, count - offset))
+ for chunk, offset, length in zip(chunks, offsets, lengths, strict=True)
+ )
+
+
+@dataclass(frozen=True, slots=True)
+class ContinuedPartial:
+ usage: Usage
+ response_cost: float
+ response_id: str
+
+
+class ContinuationOutput:
+ """Splices a fallback's continuation onto the stream the client already received. Drops a leading echo of the
+ prefill, since some providers (Mistral) return the prefix inside the completion, and the prefill's trailing
+ whitespace sent again by a provider that strips it (Anthropic), keeps the primary's chunk id, created time and single role delta, and
+ runs litellm.post_call_rules over the joined text, since each stream wrapper only sees its own fragment. A
+ continuation nested in another leaves the rules to the outermost one, the only one holding the whole text."""
+
+ def __init__(
+ self,
+ emitted_text: str,
+ prefill_message: Mapping[str, object],
+ model: str,
+ primary_response: object,
+ checks_post_call_rules: bool = True,
+ ) -> None:
+ self._text = emitted_text
+ self._checks_post_call_rules: Final = checks_post_call_rules
+ self._prefill_text: Final = _prefill_text(prefill_message)
+ self._model: Final = model
+ self._identity: Final = _stream_identity(primary_response)
+ self._rules: Final = Rules()
+ self._held: tuple[ModelResponseStream, ...] = ()
+ self._echo_settled = not self._prefill_text
+ self._seam = _resendable_whitespace(self._prefill_text)
+
+ async def splice(self, stream: AsyncIterable[object]) -> AsyncIterator[object]:
+ async for item in stream:
+ if not isinstance(item, ModelResponseStream):
+ yield item
+ continue
+ for released in self._observed(self._past_seam(self._released(item))):
+ yield self._in_primary_stream(released)
+ for released in self._observed(self._past_seam(self._flushed())):
+ yield self._in_primary_stream(released)
+
+ def _in_primary_stream(self, chunk: ModelResponseStream) -> ModelResponseStream:
+ return _without_role(chunk.model_copy(update=self._identity) if self._identity else chunk)
+
+ def _flushed(self) -> tuple[ModelResponseStream, ...]:
+ held: Final = self._held
+ self._held = ()
+ self._echo_settled = True
+ return held
+
+ def _released(self, chunk: ModelResponseStream) -> tuple[ModelResponseStream, ...]:
+ if self._echo_settled:
+ return (chunk,)
+ held: Final = (*self._held, chunk)
+ text: Final = "".join(_delta_text(held_chunk) for held_chunk in held)
+ if len(text) < len(self._prefill_text) and self._prefill_text.startswith(text):
+ self._held = held
+ return ()
+ self._held = ()
+ self._echo_settled = True
+ if not text.startswith(self._prefill_text):
+ return held
+ return _drop_leading_text(held, len(self._prefill_text))
+
+ def _past_seam(self, chunks: tuple[ModelResponseStream, ...]) -> tuple[ModelResponseStream, ...]:
+ return tuple(self._without_seam(chunk) for chunk in chunks)
+
+ def _without_seam(self, chunk: ModelResponseStream) -> ModelResponseStream:
+ text: Final = _delta_text(chunk)
+ matched: Final = len(commonprefix((text, self._seam)))
+ self._seam = self._seam[matched:] if matched == len(text) else ""
+ return _without_leading_text(chunk, matched)
+
+ def _observed(self, chunks: tuple[ModelResponseStream, ...]) -> tuple[ModelResponseStream, ...]:
+ if not self._checks_post_call_rules or not litellm.post_call_rules:
+ return chunks
+ for chunk in chunks:
+ self._text += _delta_text(chunk)
+ self._rules.post_call_rules(input=self._text, model=self._model)
+ return chunks
diff --git a/litellm/types/litellm_params.py b/litellm/types/litellm_params.py
index 096ad1c30dd..8387afa22bb 100644
--- a/litellm/types/litellm_params.py
+++ b/litellm/types/litellm_params.py
@@ -260,6 +260,7 @@ class PromptOptions:
prompt_label: str | None = None
litellm_system_prompt: str | None = None
custom_prompt_dict: Mapping[str, object] | None = None
+ initial_prompt_value: str | None = None
roles: Mapping[str, object] | None = None
final_prompt_value: str | None = None
bos_token: str | None = None
diff --git a/litellm/types/utils.py b/litellm/types/utils.py
index 0bcb732a0f5..7df8cc7452f 100644
--- a/litellm/types/utils.py
+++ b/litellm/types/utils.py
@@ -4045,7 +4045,9 @@ class CustomPricingLiteLLMParams(MirroredPricingParams):
DEPLOYMENT_SCOPED_PRICING_FIELDS: Final[frozenset[str]] = frozenset({"off_peak_pricing"})
-DEPLOYMENT_SCOPED_CAPABILITY_FIELDS: Final[frozenset[str]] = frozenset({"supports_regex_lookaround"})
+DEPLOYMENT_SCOPED_CAPABILITY_FIELDS: Final[frozenset[str]] = frozenset(
+ {"supports_regex_lookaround", "supports_assistant_prefill"}
+)
SHARED_BACKEND_MODEL_INFO_FIELDS: Final[frozenset[str]] = (
frozenset(ModelInfoBase.__required_keys__ | ModelInfoBase.__optional_keys__)
diff --git a/litellm/utils.py b/litellm/utils.py
index c0c80f09e4a..27669d8fd4a 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -89,6 +89,7 @@ from litellm.litellm_core_utils.fallback_generalizations import (
match_capability_generalizations,
match_fill_missing_generalizations,
)
+from litellm.litellm_core_utils.model_param_helper import peek_reasoning_summary_aliases
from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload
from litellm.litellm_core_utils.tokenizer import strip_special_tokens
from litellm.rust_bridge import tokenizer as tokenizer_dispatch
@@ -3059,6 +3060,20 @@ def supports_prompt_cache_breakpoint(model: str, custom_llm_provider: str | None
)
+def supports_assistant_prefill(model: str, custom_llm_provider: str | None = None) -> bool:
+ """
+ Whether the model can continue a prefilled assistant message (Anthropic's
+ ``prefix: True`` trick and equivalents). Missing metadata reads as False,
+ so a mid-stream fallback continuation is only routed to a model known to
+ support it.
+ """
+ return _supports_factory(
+ model=model,
+ custom_llm_provider=custom_llm_provider,
+ key="supports_assistant_prefill",
+ )
+
+
def supports_thinking_cache_preservation(model: str, custom_llm_provider: str | None = None) -> bool:
return supports_factory(
model=model,
@@ -9932,6 +9947,10 @@ class ProviderConfigManager:
return MinimaxVideoConfig()
elif LlmProviders.EDENAI == provider:
return litellm.EdenAIVideoConfig()
+ elif LlmProviders.VOLCENGINE == provider:
+ from litellm.llms.volcengine.videos.transformation import VolcEngineVideoConfig
+
+ return VolcEngineVideoConfig()
return None
@staticmethod
@@ -9956,6 +9975,7 @@ class ProviderConfigManager:
def get_provider_realtime_config(
model: str,
provider: LlmProviders,
+ drop_params: bool = False,
) -> BaseRealtimeConfig | None:
if LlmProviders.GEMINI == provider:
from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig
@@ -9964,7 +9984,7 @@ class ProviderConfigManager:
if LlmProviders.META == provider:
from litellm.llms.meta.realtime.transformation import MetaRealtimeConfig
- return MetaRealtimeConfig()
+ return MetaRealtimeConfig(drop_params=drop_params)
return None
@staticmethod
@@ -10446,22 +10466,21 @@ def get_non_default_completion_params(kwargs: Mapping[str, object]) -> dict[str,
return filter_out_litellm_params(kwargs, excluding=litellm.OPENAI_CHAT_COMPLETION_PARAMS)
-def peek_reasoning_summary_aliases(optional_params: dict) -> object | None:
- """Read AI-SDK-style reasoning summary from optional_params or nested extra_body.
-
- Uses key membership (not ``or`` chains) so falsy values like ``""`` are not skipped.
- """
- if "reasoningSummary" in optional_params:
- return optional_params["reasoningSummary"]
- if "reasoning_summary" in optional_params:
- return optional_params["reasoning_summary"]
- extra_body: Final = optional_params.get("extra_body")
- if isinstance(extra_body, dict):
- if "reasoningSummary" in extra_body:
- return extra_body["reasoningSummary"]
- if "reasoning_summary" in extra_body:
- return extra_body["reasoning_summary"]
- return None
+def reasoning_effort_with_summary_alias(
+ reasoning_effort: object, params: Mapping[str, object], accepts_summary: bool
+) -> object:
+ summary: Final = peek_reasoning_summary_aliases(params)
+ if not accepts_summary or not isinstance(summary, str):
+ return reasoning_effort
+ match reasoning_effort:
+ case str():
+ return {"effort": reasoning_effort, "summary": summary}
+ case {"summary": _}:
+ return reasoning_effort
+ case {"effort": str() as effort}:
+ return {"effort": effort, "summary": summary}
+ case _:
+ return reasoning_effort
def strip_reasoning_summary_aliases_from_optional_params(
diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json
index b744032fb6c..49d2ce90af4 100644
--- a/model_prices_and_context_window.json
+++ b/model_prices_and_context_window.json
@@ -2773,7 +2773,7 @@
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
- "supports_assistant_prefill": true,
+ "supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
@@ -2808,7 +2808,7 @@
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
- "supports_assistant_prefill": true,
+ "supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
@@ -2843,7 +2843,7 @@
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
- "supports_assistant_prefill": true,
+ "supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
@@ -2878,7 +2878,7 @@
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
- "supports_assistant_prefill": true,
+ "supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
@@ -2913,7 +2913,7 @@
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
- "supports_assistant_prefill": true,
+ "supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
@@ -2948,7 +2948,7 @@
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
- "supports_assistant_prefill": true,
+ "supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
@@ -3743,7 +3743,7 @@
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 1.5e-05,
- "supports_assistant_prefill": true,
+ "supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
@@ -15494,7 +15494,7 @@
},
"supports_adaptive_thinking": true,
"supports_legacy_thinking": true,
- "supports_assistant_prefill": true,
+ "supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
@@ -20832,7 +20832,7 @@
"output_cost_per_token": 2.5000010000000002e-05,
"output_dbu_cost_per_token": 0.000357143,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
- "supports_assistant_prefill": true,
+ "supports_assistant_prefill": false,
"supports_function_calling": true,
"supports_legacy_thinking": true,
"supports_prompt_caching": true,
@@ -21047,7 +21047,7 @@
"output_cost_per_token": 1.5000020000000002e-05,
"output_dbu_cost_per_token": 0.000214286,
"source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving",
- "supports_assistant_prefill": true,
+ "supports_assistant_prefill": false,
"supports_function_calling": true,
"supports_legacy_thinking": true,
"supports_prompt_caching": true,
@@ -49176,7 +49176,7 @@
"max_tokens": 64000,
"mode": "chat",
"output_cost_per_token": 2.5e-05,
- "supports_assistant_prefill": true,
+ "supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
@@ -51006,7 +51006,7 @@
"mode": "chat",
"output_cost_per_token": 1.5e-05,
"output_cost_per_token_batches": 7.5e-06,
- "supports_assistant_prefill": true,
+ "supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
@@ -51179,7 +51179,7 @@
"supported_regions": [
"global"
],
- "supports_assistant_prefill": true,
+ "supports_assistant_prefill": false,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
@@ -51200,7 +51200,7 @@
"supported_regions": [
"us-central1"
],
- "supports_assistant_prefill": true,
+ "supports_assistant_prefill": false,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
@@ -58822,7 +58822,7 @@
"mode": "chat",
"output_cost_per_token": 1.5e-05,
"output_cost_per_token_batches": 7.5e-06,
- "supports_assistant_prefill": true,
+ "supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
@@ -60467,6 +60467,61 @@
}
]
},
+ "volcengine/doubao-seedance-2-5-260628": {
+ "litellm_provider": "volcengine",
+ "mode": "video_generation",
+ "source": "https://www.volcengine.com/docs/82379/1544106",
+ "output_cost_per_second": 0.216,
+ "output_cost_per_second_480p": 0.096075,
+ "output_cost_per_second_720p": 0.216,
+ "output_cost_per_second_1080p": 0.5346,
+ "supported_modalities": ["text", "image", "video", "audio"],
+ "supported_output_modalities": ["video"],
+ "supported_endpoints": ["/v1/videos"],
+ "provider_specific_entry": {"volcengine_video_frames": {"480p": {"16:9": "854x480", "4:3": "752x560", "1:1": "640x640", "3:4": "560x752", "9:16": "480x854", "21:9": "992x432"}, "720p": {"16:9": "1280x720", "4:3": "1112x834", "1:1": "960x960", "3:4": "834x1112", "9:16": "720x1280", "21:9": "1470x630"}}},
+ "metadata": {"comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s). List 70 CNY per million tokens at 480p and 720p and 77 at 1080p without a video input, matching Ark's per-resolution table and its 1080p 16:9 example of 3.74 CNY/s, converted at 7.0 CNY/USD like the other volcengine entries. Accepts up to 50 reference assets (30 images, 10 videos, 10 audios) and clips of 4-30s; per-asset billing beyond the free allowance is not encoded. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card."}
+ },
+ "volcengine/doubao-seedance-2-0-260128": {
+ "litellm_provider": "volcengine",
+ "mode": "video_generation",
+ "source": "https://www.volcengine.com/docs/82379/1544106",
+ "output_cost_per_second": 0.141942857,
+ "output_cost_per_second_480p": 0.066003429,
+ "output_cost_per_second_720p": 0.141942857,
+ "output_cost_per_second_1080p": 0.354085714,
+ "output_cost_per_second_4k": 0.722057143,
+ "supported_modalities": ["text", "image", "video", "audio"],
+ "supported_output_modalities": ["video"],
+ "supported_endpoints": ["/v1/videos"],
+ "provider_specific_entry": {"volcengine_video_frames": {"480p": {"16:9": "864x496", "4:3": "752x560", "1:1": "640x640", "3:4": "560x752", "9:16": "496x864", "21:9": "992x432"}, "720p": {"16:9": "1280x720", "4:3": "1112x834", "1:1": "960x960", "3:4": "834x1112", "9:16": "720x1280", "21:9": "1470x630"}}},
+ "metadata": {"comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s; the 480p frame is 864x496, 10044 tokens/s, not the tier's nominal 854x480). List 46 CNY per million tokens at 480p and 720p, 51 at 1080p, and 26 at 4k without a video input, matching Ark's per-resolution table and its 1080p and 4k 16:9 examples of 2.48 and 5.05 CNY/s, converted at 7.0 CNY/USD like the other volcengine entries. Clips of 4-15s; 4k is exclusive to this tier. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card."}
+ },
+ "volcengine/doubao-seedance-2-0-fast-260128": {
+ "litellm_provider": "volcengine",
+ "mode": "video_generation",
+ "source": "https://www.volcengine.com/docs/82379/1544106",
+ "output_cost_per_second": 0.114171429,
+ "output_cost_per_second_480p": 0.053089714,
+ "output_cost_per_second_720p": 0.114171429,
+ "supported_modalities": ["text", "image", "video", "audio"],
+ "supported_output_modalities": ["video"],
+ "supported_endpoints": ["/v1/videos"],
+ "provider_specific_entry": {"volcengine_video_frames": {"480p": {"16:9": "864x496", "4:3": "752x560", "1:1": "640x640", "3:4": "560x752", "9:16": "496x864", "21:9": "992x432"}, "720p": {"16:9": "1280x720", "4:3": "1112x834", "1:1": "960x960", "3:4": "834x1112", "9:16": "720x1280", "21:9": "1470x630"}}},
+ "metadata": {"comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s; the 480p frame is 864x496, 10044 tokens/s, not the tier's nominal 854x480). List 37 CNY per million tokens without a video input, converted at 7.0 CNY/USD like the other volcengine entries. Clips of 4-15s at 480p or 720p. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card."}
+ },
+ "volcengine/doubao-seedance-2-0-mini-260615": {
+ "litellm_provider": "volcengine",
+ "mode": "video_generation",
+ "source": "https://www.volcengine.com/docs/82379/1544106",
+ "output_cost_per_second": 0.070971429,
+ "output_cost_per_second_480p": 0.033001714,
+ "output_cost_per_second_720p": 0.070971429,
+ "supported_modalities": ["text", "image"],
+ "supported_output_modalities": ["video"],
+ "supported_endpoints": ["/v1/videos"],
+ "provider_specific_entry": {"volcengine_video_frames": {"480p": {"16:9": "864x496", "4:3": "752x560", "1:1": "640x640", "3:4": "560x752", "9:16": "496x864", "21:9": "992x432"}, "720p": {"16:9": "1280x720", "4:3": "1112x834", "1:1": "960x960", "3:4": "834x1112", "9:16": "720x1280", "21:9": "1470x630"}}},
+ "metadata": {"comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s; the 480p frame is 864x496, 10044 tokens/s, not the tier's nominal 854x480). List 23 CNY per million tokens without a video input, converted at 7.0 CNY/USD like the other volcengine entries. Clips of 4-15s at 480p or 720p. Only text and image inputs are declared; the mini tier's reference video and audio caps were not verified. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card."}
+ },
"bedrock/us-east-1/zai.glm-5": {
"input_cost_per_token": 1e-06,
"output_cost_per_token": 3.2e-06,
@@ -72334,7 +72389,7 @@
"max_output_tokens": 64000,
"max_tokens": 64000,
"mode": "chat",
- "supports_assistant_prefill": true,
+ "supports_assistant_prefill": false,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json
index 45104d40018..0325fe1f796 100644
--- a/provider_endpoints_support.json
+++ b/provider_endpoints_support.json
@@ -2766,7 +2766,8 @@
"batches": false,
"rerank": false,
"a2a": true,
- "interactions": true
+ "interactions": true,
+ "video_generations": true
}
},
"voyage": {
diff --git a/schema.prisma b/schema.prisma
index d40bac380aa..83a0e6f9e05 100644
--- a/schema.prisma
+++ b/schema.prisma
@@ -1981,6 +1981,7 @@ model LiteLLM_BackgroundInteractionSettlement {
outcome String?
@@index([claimed_at], map: "idx_background_interaction_settlement_claimed_at")
+ @@index([settled_at], map: "idx_background_interaction_settlement_settled_at")
}
model LiteLLM_Lens {
diff --git a/tests/benchmarks/test_prompt_template_benchmarks.py b/tests/benchmarks/test_prompt_template_benchmarks.py
new file mode 100644
index 00000000000..c7858771d35
--- /dev/null
+++ b/tests/benchmarks/test_prompt_template_benchmarks.py
@@ -0,0 +1,21 @@
+from collections.abc import Mapping
+from typing import Final
+
+import pytest
+
+from litellm.litellm_core_utils.prompt_templates.factory import anthropic_messages_pt
+
+LONG_RUN_OF_EMPTY_CACHED_USER_MESSAGES: Final[tuple[Mapping[str, object], ...]] = (
+ {"role": "user", "content": "hi"},
+ {"role": "assistant", "content": "hello"},
+ *({"role": "user", "content": [], "cache_control": {"type": "ephemeral"}} for _ in range(1_000)),
+)
+
+
+@pytest.mark.benchmark
+def test_anthropic_messages_pt_long_run_of_empty_cached_user_messages() -> None:
+ result: Final = anthropic_messages_pt(
+ messages=list(LONG_RUN_OF_EMPTY_CACHED_USER_MESSAGES), model="claude-opus-5", llm_provider="vertex_ai"
+ )
+
+ assert len(result[-1]["content"]) == 1_000
diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py
index ba2913bcfb7..6762e4cd35c 100644
--- a/tests/code_coverage_tests/router_code_coverage.py
+++ b/tests/code_coverage_tests/router_code_coverage.py
@@ -114,6 +114,24 @@ ignored_function_names = [
"_deployment_that_just_failed", # Tested through the same-boundary hop test in test_router_order_fallback.py
"_collect_ptu_share_teams", # Cached behind team_ptu_ceiling, tested through the live model list test in test_router.py
"_resolve_team_ptu_ceiling", # Cached behind team_ptu_ceiling, tested through the live model list test in test_router.py
+ "_register_continuation_prefill_check", # Tested through the continuation tests in test_router.py
+ "_set_mid_stream_fallback_continuation", # Tested through the update_settings continuation test in test_router.py
+ "_is_flag_enabled", # Tested through the string-flag continuation tests in test_router.py
+ "_params_sent_to_deployment", # Tested through the configured-constraint continuation tests in test_router.py
+ "_names_deployment_that_cannot_continue", # Tested through the deployment-id continuation tests in test_router.py
+ "_set_retry_policy", # Tested through update_settings in test_router_retry_policy_update.py
+ "_resolve_pre_call_check_input_tokens", # Tested through the token-count failure pre-call check tests in test_router.py
+ "_params_neither_dropped_nor_allowed", # Tested through the allowed-or-dropped param pre-call check tests in test_router.py
+ "_deployment_rejects_params", # Tested through the unsupported-param pre-call check tests in test_router.py
+ "_deployment_lacks_params", # Tested through the unsupported-param pre-call check tests in test_router.py
+ "_pre_call_check_requested_chat_params", # Tested through the unsupported-param pre-call check tests in test_router.py
+ "_deployment_supports_param", # Tested through the deployment-capability pre-call check tests in test_router.py
+ "_capability_model_and_provider", # Tested through the authenticating-provider pre-call check test in test_router.py
+ "_exact_base_model_max_input_tokens", # Tested through get_router_model_info in test_router_model_cost_isolation.py
+ "_base_model_widened_window", # Tested through get_router_model_info in test_router_model_cost_isolation.py
+ "_provider_model_id", # Tested through get_router_model_info in test_router_model_cost_isolation.py
+ "_names_route_model", # Tested through get_router_model_info in test_router_model_cost_isolation.py
+ "_declared_base_model", # Tested through get_router_model_info in test_router_model_cost_isolation.py
]
diff --git a/tests/integration/_support/upstream.py b/tests/integration/_support/upstream.py
index d2b67c5438e..186bc345fab 100644
--- a/tests/integration/_support/upstream.py
+++ b/tests/integration/_support/upstream.py
@@ -9,9 +9,10 @@ import struct
import uuid
import zlib
from collections import deque
-from collections.abc import AsyncIterator, Mapping
+from collections.abc import AsyncIterator, Callable, Iterator, Mapping
from dataclasses import dataclass, field
from queue import SimpleQueue
+from types import MappingProxyType
from typing import Final, cast
import httpx
@@ -42,7 +43,7 @@ from starlette.datastructures import UploadFile as StarletteUploadFile
from starlette.requests import Request
from starlette.responses import JSONResponse, Response, StreamingResponse
from starlette.routing import Route, WebSocketRoute
-from starlette.websockets import WebSocket, WebSocketDisconnect
+from starlette.websockets import WebSocket, WebSocketDisconnect, WebSocketState
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
INTERNAL_FIELDS: Final = frozenset(
@@ -81,6 +82,25 @@ class Observation:
api_key: str = ""
+@dataclass(frozen=True, slots=True)
+class _RealtimeConnection:
+ websocket: WebSocket
+ response: RealtimeResponse
+ session: dict[str, JsonValue]
+ scenario_id: str
+ authorization: str
+ api_key: str
+ events: Iterator[dict[str, JsonValue]]
+ typed_turn_detection: deque[JsonValue]
+
+ def auto_answers(self) -> bool:
+ return (
+ self.response.dialect == "xai"
+ and len(self.typed_turn_detection) > 0
+ and _is_typed_server_vad(self.typed_turn_detection[-1])
+ )
+
+
class InteractionState(BaseModel):
"""What the scripted Interactions API answers for one interaction id until a DELETE drops it."""
@@ -92,6 +112,12 @@ class InteractionState(BaseModel):
delay_seconds: float = 0
+class _ReceivedFrame(BaseModel):
+ model_config = ConfigDict(frozen=True)
+ type: str
+ text: str | None = None
+
+
class _ScenarioRegistration(BaseModel):
scenario_id: str
response: StoredResponse
@@ -415,23 +441,47 @@ class Provider:
) -> None:
await websocket.accept()
session: Final = _realtime_session(response, scenario_id, websocket.query_params.get("model", ""))
+ frames: Final = _received_texts(websocket)
+ held: Final = await anext(frames, None) if response.created_after_first_frame else None
for _ in range(response.created_repeats):
await websocket.send_json({"type": response.created_event, "event_id": _event_id(), "session": session})
- event_index: Final = iter(response.events)
- async for message in websocket.iter_json():
- payload: Final = JSON_OBJECT.validate_python(message)
- self.observations.put(Observation(websocket.url.path, authorization, payload, "WEBSOCKET_FRAME", api_key))
- if payload.get("type") == "session.update":
- await websocket.send_json(
- _realtime_update_reply(payload.get("session"), session, response.session_type)
+ connection: Final = _RealtimeConnection(
+ websocket, response, session, scenario_id, authorization, api_key, iter(response.events), deque(maxlen=1)
+ )
+ if held is not None:
+ await self._realtime_frame(connection, held)
+ async for message in frames:
+ await self._realtime_frame(connection, message)
+
+ async def _realtime_frame(self, connection: _RealtimeConnection, message: str) -> None:
+ websocket: Final = connection.websocket
+ response: Final = connection.response
+ payload: Final = JSON_OBJECT.validate_json(message)
+ self.observations.put(
+ Observation(websocket.url.path, connection.authorization, payload, "WEBSOCKET_FRAME", connection.api_key)
+ )
+ frame_type: Final = payload.get("type")
+ if frame_type == "session.update":
+ reply: Final = _realtime_update_reply(
+ payload.get("session"), connection.session, response.session_type, response.dialect
+ )
+ if reply["type"] == "session.updated":
+ connection.typed_turn_detection.extend(
+ filter(_declares_a_type, _declared_turn_detection(payload.get("session")))
)
- continue
- if payload.get("type") not in _REALTIME_TRIGGERS:
- continue
- event: Final = next(event_index, None)
- if event is None:
- continue
- await websocket.send_json(_rendered_realtime_event(event, scenario_id))
+ await _reply_while_open(websocket, reply)
+ return
+ if frame_type == "response.cancel":
+ await _reply_while_open(websocket, _cancel_not_active(payload.get("event_id"), response.dialect))
+ return
+ if frame_type not in _REALTIME_TRIGGERS:
+ return
+ event: Final = next(connection.events, None)
+ if event is None:
+ return
+ await _reply_while_open(websocket, _rendered_realtime_event(event, connection.scenario_id))
+ if frame_type == "input_audio_buffer.commit" and connection.auto_answers():
+ await _reply_while_open(websocket, _rendered_realtime_event(dict(_AUTO_ANSWER), connection.scenario_id))
async def muse_realtime(self, websocket: WebSocket) -> None:
await websocket.accept()
@@ -454,8 +504,23 @@ class Provider:
await websocket.send_json(_rendered_realtime_event(pending.popleft(), scenario_id))
async def gemini_live(self, websocket: WebSocket) -> None:
- authorization: Final = websocket.headers.get("authorization", "")
- scenario_id: Final = websocket.headers.get("x-goog-user-project", "")
+ await self._gemini_live_session(
+ websocket,
+ websocket.headers.get("x-goog-user-project", ""),
+ websocket.headers.get("authorization", ""),
+ _accepts_every_setup,
+ )
+
+ async def gemini_ai_studio_live(self, websocket: WebSocket) -> None:
+ await self._gemini_live_session(websocket, websocket.query_params.get("key", ""), "", _ai_studio_setup_refusal)
+
+ async def _gemini_live_session(
+ self,
+ websocket: WebSocket,
+ scenario_id: str,
+ authorization: str,
+ setup_refusal: Callable[[Mapping[str, JsonValue]], str | None],
+ ) -> None:
self.observations.put(
Observation(
websocket.url.path,
@@ -477,6 +542,9 @@ class Provider:
Observation(websocket.url.path, authorization, payload, "WEBSOCKET_FRAME", scenario_id)
)
if "setup" in payload:
+ if (refusal := setup_refusal(payload)) is not None:
+ await websocket.close(code=_INVALID_ARGUMENT_CLOSE, reason=refusal[:_CLOSE_REASON_BYTES])
+ return
await websocket.send_json({"setupComplete": {}})
continue
if not _GEMINI_LIVE_TRIGGERS.intersection(payload):
@@ -587,6 +655,7 @@ class Provider:
WebSocketRoute("/openai/realtime", self.realtime),
WebSocketRoute("/v1/asr/realtime", self.muse_realtime),
WebSocketRoute(GEMINI_LIVE_PATH, self.gemini_live),
+ WebSocketRoute(GEMINI_AI_STUDIO_LIVE_PATH, self.gemini_ai_studio_live),
]
)
@@ -612,6 +681,19 @@ def _interaction_body(interaction_id: str, state: InteractionState) -> dict[str,
_REALTIME_TRIGGERS: Final = frozenset({"response.create", "input_audio_buffer.commit"})
_GEMINI_LIVE_TRIGGERS: Final = frozenset({"clientContent", "realtimeInput"})
GEMINI_LIVE_PATH: Final = "/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent"
+GEMINI_AI_STUDIO_LIVE_PATH: Final = "/ws/google.ai.generativelanguage.v1beta.GenerativeService.BidiGenerateContent"
+_INVALID_ARGUMENT_CLOSE: Final = 1007
+_CLOSE_REASON_BYTES: Final = 123
+_AI_STUDIO_TUNING_FIELDS: Final = MappingProxyType(
+ {"prefixPaddingMs": "prefix_padding_ms", "silenceDurationMs": "silence_duration_ms"}
+)
+_AUTO_ANSWER: Final[Mapping[str, JsonValue]] = MappingProxyType(
+ {
+ "type": "response.done",
+ "event_id": "evt_$UNIQUE_ID",
+ "response": {"id": "resp_auto_$UNIQUE_ID", "object": "realtime.response", "status": "completed", "output": []},
+ }
+)
_TRANSCRIPTION_UPDATE_REFUSED: Final = "Passing a realtime session update to a transcription session is not allowed."
_REALTIME_UPDATE_REFUSED: Final = "Passing a transcription session update to a realtime session is not allowed."
_NESTED_TURN_DETECTION_TYPE: Final = "session.audio.input.turn_detection.type"
@@ -651,11 +733,11 @@ def _missing_turn_detection_type(session: Mapping[str, JsonValue]) -> str | None
def _realtime_update_reply(
- session: JsonValue, created: Mapping[str, JsonValue], session_type: str | None
+ session: JsonValue, created: Mapping[str, JsonValue], session_type: str | None, dialect: str
) -> dict[str, JsonValue]:
if not isinstance(session, dict):
return _realtime_error("invalid_value", "Invalid value for 'session': expected an object.", "session")
- missing: Final = _missing_turn_detection_type(session)
+ missing: Final = None if dialect == "xai" else _missing_turn_detection_type(session)
if missing is not None:
return _realtime_error("missing_required_parameter", f"Missing required parameter: '{missing}'.", missing)
declared: Final = session.get("type")
@@ -666,6 +748,82 @@ def _realtime_update_reply(
return {"type": "session.updated", "event_id": _event_id(), "session": {**created, **session}}
+async def _received_texts(websocket: WebSocket) -> AsyncIterator[str]:
+ while (frame := _ReceivedFrame.model_validate(await websocket.receive())).type == "websocket.receive":
+ if frame.text is not None:
+ yield frame.text
+
+
+async def _reply_while_open(websocket: WebSocket, frame: Mapping[str, JsonValue]) -> None:
+ if websocket.application_state is not WebSocketState.CONNECTED:
+ return
+ try:
+ await websocket.send_json(frame)
+ except WebSocketDisconnect:
+ return
+
+
+def _cancel_not_active(cancel_event_id: JsonValue, dialect: str) -> dict[str, JsonValue]:
+ message: Final = "Cancellation failed: no active response found"
+ if dialect == "xai":
+ own_id: Final = _event_id()
+ return {
+ "type": "error",
+ "event_id": own_id,
+ "error": {
+ "message": message,
+ "type": "invalid_request_error",
+ "code": "invalid_request_error",
+ "event_id": own_id,
+ },
+ }
+ return {
+ "type": "error",
+ "event_id": _event_id(),
+ "error": {
+ "type": "invalid_request_error",
+ "code": "response_cancel_not_active",
+ "message": message,
+ "param": None,
+ "event_id": cancel_event_id,
+ },
+ }
+
+
+def _declared_turn_detection(session: JsonValue) -> tuple[JsonValue, ...]:
+ if not isinstance(session, dict):
+ return ()
+ audio: Final = session.get("audio")
+ audio_input: Final = audio.get("input") if isinstance(audio, dict) else None
+ if isinstance(audio_input, dict) and "turn_detection" in audio_input:
+ return (audio_input["turn_detection"],)
+ return (session["turn_detection"],) if "turn_detection" in session else ()
+
+
+def _declares_a_type(turn_detection: JsonValue) -> bool:
+ return isinstance(turn_detection, dict) and "type" in turn_detection
+
+
+def _is_typed_server_vad(turn_detection: JsonValue) -> bool:
+ return isinstance(turn_detection, dict) and turn_detection.get("type") == "server_vad"
+
+
+def _accepts_every_setup(_payload: Mapping[str, JsonValue]) -> str | None:
+ return None
+
+
+def _ai_studio_setup_refusal(payload: Mapping[str, JsonValue]) -> str | None:
+ setup: Final = payload.get("setup")
+ config: Final = setup.get("realtimeInputConfig") if isinstance(setup, dict) else None
+ detection: Final = config.get("automaticActivityDetection") if isinstance(config, dict) else None
+ if not isinstance(detection, dict) or detection.get("disabled") is not True:
+ return None
+ tuning: Final = tuple(name for key, name in _AI_STUDIO_TUNING_FIELDS.items() if key in detection)
+ if not tuning:
+ return None
+ return f"automatic_activity_detection.disabled is true, but the following fields were also set: {', '.join(tuning)}"
+
+
def _rendered_realtime_event(event: Mapping[str, JsonValue], scenario_id: str) -> dict[str, JsonValue]:
return JSON_OBJECT.validate_json(
json.dumps(event, separators=(",", ":"))
diff --git a/tests/integration/configuration/test_callback_delete.py b/tests/integration/configuration/test_callback_delete.py
new file mode 100644
index 00000000000..d58dd252593
--- /dev/null
+++ b/tests/integration/configuration/test_callback_delete.py
@@ -0,0 +1,351 @@
+import json
+import os
+import uuid
+from collections.abc import Generator, Iterator, Mapping
+from contextlib import contextmanager
+from dataclasses import dataclass
+from pathlib import Path
+from types import MappingProxyType
+from typing import Final
+
+import pytest
+import yaml
+from integration._support.client import JSON_OBJECT, Gateway, JsonValue, eventually, list_value, object_value
+from integration._support.database import read_rows, scratch_database
+from integration._support.process import owned_proxy
+from integration._support.wire import Reply, Request, Wire, wire_server
+
+CALLBACK: Final = "generic_api"
+SIBLING: Final = "langsmith"
+LIST_TYPES: Final = MappingProxyType(
+ {
+ "success_callback": "success",
+ "failure_callback": "failure",
+ "callbacks": "success_and_failure",
+ }
+)
+SINK_ENVIRONMENT: Final = MappingProxyType(
+ {
+ "GENERIC_LOGGER_ENDPOINT": "http://127.0.0.1:9/unused",
+ "LANGSMITH_API_KEY": "synthetic-langsmith-key",
+ "LANGSMITH_BASE_URL": "http://127.0.0.1:9",
+ }
+)
+NOT_FOUND: Final = f"Callback '{CALLBACK}' not found in active configuration"
+WORKER_SYNC_SECONDS: Final = 40
+
+
+@dataclass(frozen=True, slots=True)
+class Rig:
+ proxy: Gateway
+ database_url: str
+
+
+@contextmanager
+def _rig(gateway: Gateway, tmp_path: Path, file_settings: Mapping[str, JsonValue]) -> Generator[Rig]:
+ config: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
+ path: Final = tmp_path / f"callback_delete_{uuid.uuid4().hex}.yaml"
+ path.write_text(
+ yaml.safe_dump({**config, "litellm_settings": {**object_value(config["litellm_settings"]), **file_settings}})
+ )
+ with (
+ scratch_database() as database_url,
+ owned_proxy(
+ gateway,
+ tmp_path,
+ {"DATABASE_URL": database_url, **SINK_ENVIRONMENT},
+ config=path,
+ remove_environment=("DATABASE_URL_READ_REPLICA",),
+ workers=2,
+ ) as proxy,
+ ):
+ yield Rig(proxy, database_url)
+
+
+@pytest.fixture
+def rig(gateway: Gateway, tmp_path: Path) -> Iterator[Rig]:
+ with _rig(gateway, tmp_path, {}) as created:
+ yield created
+
+
+def _configured(proxy: Gateway) -> tuple[tuple[str, str], ...]:
+ rows: Final = proxy.get("/get/config/callbacks")["callbacks"]
+ assert isinstance(rows, list)
+ return tuple(
+ sorted(
+ (str(row["name"]), str(row["type"]))
+ for row in (object_value(entry) for entry in rows)
+ if not row.get("read_only")
+ )
+ )
+
+
+def _listed(proxy: Gateway, expected: tuple[tuple[str, str], ...]) -> None:
+ eventually(lambda: _configured(proxy), lambda rows: rows == expected, seconds=WORKER_SYNC_SECONDS)
+
+
+def _stored_settings(database_url: str) -> Mapping[str, JsonValue]:
+ rows: Final = read_rows(
+ 'SELECT param_value FROM "LiteLLM_Config" WHERE param_name = %s',
+ ("litellm_settings",),
+ database_url=database_url,
+ )
+ assert len(rows) == 1, rows
+ value: Final = rows[0]["param_value"]
+ return object_value(json.loads(value) if isinstance(value, str) else value)
+
+
+def _configure(
+ proxy: Gateway, settings: Mapping[str, JsonValue], environment: Mapping[str, JsonValue] | None = None
+) -> None:
+ body: Final[Mapping[str, JsonValue]] = {
+ "litellm_settings": dict(settings),
+ "environment_variables": None if environment is None else dict(environment),
+ }
+ assert proxy.post("/config/update", body) == {"message": "Config updated successfully"}
+
+
+def _delete(proxy: Gateway, name: str) -> tuple[int, Mapping[str, JsonValue]]:
+ response: Final = proxy.request("POST", "/config/callback/delete", {"callback_name": name})
+ return response.status_code, JSON_OBJECT.validate_json(response.content)
+
+
+@pytest.mark.parametrize("list_key", tuple(LIST_TYPES))
+def test_callback_listed_under_any_settings_list_is_deleted_from_the_listing_and_the_store(
+ rig: Rig, list_key: str
+) -> None:
+ _configure(rig.proxy, {list_key: [CALLBACK]})
+ _listed(rig.proxy, ((CALLBACK, LIST_TYPES[list_key]),))
+
+ status, body = _delete(rig.proxy, CALLBACK)
+
+ assert status == 200, body
+ assert (body["removed_callback"], body["remaining_callbacks"]) == (CALLBACK, [])
+ assert _stored_settings(rig.database_url)[list_key] == []
+ _listed(rig.proxy, ())
+
+
+def test_callback_in_both_success_and_failure_lists_is_removed_from_both_in_one_delete(rig: Rig) -> None:
+ _configure(rig.proxy, {"success_callback": [CALLBACK], "failure_callback": [CALLBACK]})
+ _listed(rig.proxy, ((CALLBACK, "failure"), (CALLBACK, "success")))
+
+ status, body = _delete(rig.proxy, CALLBACK)
+
+ assert status == 200, body
+ stored: Final = _stored_settings(rig.database_url)
+ assert (stored["success_callback"], stored["failure_callback"]) == ([], [])
+ _listed(rig.proxy, ())
+
+
+def test_delete_keeps_other_callbacks_and_reports_what_is_left_in_the_list_it_edited(rig: Rig) -> None:
+ _configure(rig.proxy, {"callbacks": [SIBLING, CALLBACK], "failure_callback": [SIBLING]})
+ _listed(rig.proxy, ((CALLBACK, "success_and_failure"), (SIBLING, "failure"), (SIBLING, "success_and_failure")))
+
+ status, body = _delete(rig.proxy, CALLBACK)
+
+ assert status == 200, body
+ assert body["remaining_callbacks"] == [SIBLING]
+ stored: Final = _stored_settings(rig.database_url)
+ assert (stored["callbacks"], stored["failure_callback"]) == ([SIBLING], [SIBLING])
+ _listed(rig.proxy, ((SIBLING, "failure"), (SIBLING, "success_and_failure")))
+
+
+def test_delete_matches_the_callback_name_case_insensitively(rig: Rig) -> None:
+ _configure(rig.proxy, {"failure_callback": [CALLBACK.upper()], "callbacks": [CALLBACK]})
+ _listed(rig.proxy, ((CALLBACK.upper(), "failure"), (CALLBACK, "success_and_failure")))
+
+ status, body = _delete(rig.proxy, CALLBACK.upper())
+
+ assert status == 200, body
+ assert body["removed_callback"] == CALLBACK
+ _listed(rig.proxy, ())
+
+
+def test_delete_of_an_unconfigured_callback_is_a_404_and_changes_nothing(rig: Rig) -> None:
+ _configure(rig.proxy, {"callbacks": [SIBLING]})
+ _listed(rig.proxy, ((SIBLING, "success_and_failure"),))
+ before: Final = _stored_settings(rig.database_url)
+
+ status, body = _delete(rig.proxy, CALLBACK)
+
+ assert (status, body) == (404, {"detail": {"error": NOT_FOUND}})
+ assert _stored_settings(rig.database_url) == before
+ _listed(rig.proxy, ((SIBLING, "success_and_failure"),))
+
+
+@pytest.mark.parametrize(
+ ("file_key", "database_key"), (("failure_callback", "success_callback"), ("success_callback", "failure_callback"))
+)
+def test_delete_removes_the_database_copy_and_keeps_the_copy_the_config_file_sets(
+ gateway: Gateway, tmp_path: Path, file_key: str, database_key: str
+) -> None:
+ with _rig(gateway, tmp_path, {file_key: [CALLBACK]}) as file_owned:
+ _configure(file_owned.proxy, {database_key: [CALLBACK]})
+ _listed(file_owned.proxy, ((CALLBACK, "failure"), (CALLBACK, "success")))
+
+ status, body = _delete(file_owned.proxy, CALLBACK)
+
+ assert status == 200, body
+ assert body["remaining_callbacks"] == []
+ assert _stored_settings(file_owned.database_url)[database_key] == []
+ _listed(file_owned.proxy, ((CALLBACK, LIST_TYPES[file_key]),))
+
+
+def test_delete_of_a_callback_only_the_config_file_lists_is_refused_and_names_the_file(
+ gateway: Gateway, tmp_path: Path
+) -> None:
+ with _rig(gateway, tmp_path, {"callbacks": [CALLBACK]}) as file_owned:
+ _listed(file_owned.proxy, ((CALLBACK, "success_and_failure"),))
+
+ status, body = _delete(file_owned.proxy, CALLBACK)
+
+ assert status == 400, body
+ detail: Final = object_value(body["detail"])
+ assert (detail["keys"], detail["section"]) == (["callbacks"], "litellm_settings")
+ assert "set in the config file" in str(detail["error"])
+ _listed(file_owned.proxy, ((CALLBACK, "success_and_failure"),))
+
+
+def _completion(request: Request) -> Reply:
+ message: Final = object_value(list_value(JSON_OBJECT.validate_json(request.body)["messages"])[0])
+ if str(message["content"]).endswith("-failure"):
+ return Reply(
+ status=400,
+ body=json.dumps(
+ {"error": {"type": "invalid_request_error", "code": "synthetic", "message": "synthetic failure"}}
+ ).encode(),
+ )
+ return Reply(
+ body=json.dumps(
+ {
+ "id": "chatcmpl-" + uuid.uuid4().hex,
+ "object": "chat.completion",
+ "created": 1,
+ "model": "gpt-4o-mini",
+ "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
+ "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4},
+ }
+ ).encode()
+ )
+
+
+@dataclass(slots=True)
+class Sink:
+ wire: Wire
+ bodies: list[bytes] # mutable-ok: drain() consumes the queue, later reads keep earlier batches
+
+ def received(self, markers: tuple[str, ...]) -> bool:
+ self.bodies.extend(request.body for request in self.wire.drain())
+ return all(any(marker.encode() in body for body in self.bodies) for marker in markers)
+
+
+@dataclass(frozen=True, slots=True)
+class Pair:
+ first: Gateway
+ peer: Gateway
+ database_url: str
+ provider: Wire
+ generic: Sink
+ langsmith: Sink
+
+ def running(self, name: str) -> tuple[bool, bool]:
+ return (_runs(self.first, name), _runs(self.peer, name))
+
+
+def _runs(proxy: Gateway, name: str) -> bool:
+ active: Final = list_value(proxy.get("/active/callbacks")["all_litellm_callbacks"])
+ return any(name in str(entry).lower() for entry in active)
+
+
+@contextmanager
+def _pair(gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Generator[Pair]:
+ with (
+ scratch_database() as database_url,
+ wire_server(_completion) as provider,
+ wire_server(lambda _: Reply()) as generic,
+ wire_server(lambda _: Reply()) as langsmith,
+ ):
+ monkeypatch.setenv("DATABASE_URL", database_url)
+ environment: Final = {
+ "DATABASE_URL": database_url,
+ "DEFAULT_FLUSH_INTERVAL_SECONDS": "1",
+ "LITELLM_STORE_AUDIT_LOGS": "True",
+ }
+ removed: Final = (
+ "DATABASE_URL_READ_REPLICA",
+ "GENERIC_LOGGER_ENDPOINT",
+ "GENERIC_LOGGER_HEADERS",
+ "LANGSMITH_API_KEY",
+ "LANGSMITH_BASE_URL",
+ )
+ with (
+ owned_proxy(gateway, tmp_path, environment, remove_environment=removed) as first,
+ owned_proxy(gateway, tmp_path, environment, remove_environment=removed) as peer,
+ ):
+ yield Pair(first, peer, database_url, provider, Sink(generic, []), Sink(langsmith, []))
+
+
+def _send_on_both(pair: Pair, model: str, key: str, marker: str, outcome: str) -> tuple[str, str]:
+ markers: Final = (f"{marker}-first-{outcome}", f"{marker}-peer-{outcome}")
+ for proxy, text in zip((pair.first, pair.peer), markers, strict=True):
+ response = proxy.request(
+ "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": text}]}, key=key
+ )
+ assert response.status_code == (400 if outcome == "failure" else 200), response.text
+ return markers
+
+
+@pytest.mark.parametrize(("list_key", "outcome"), (("failure_callback", "failure"), ("callbacks", "success")))
+def test_deleted_callback_stops_receiving_events_on_both_processes(
+ gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, list_key: str, outcome: str
+) -> None:
+ assert os.environ.get("LITELLM_LICENSE"), "LITELLM_LICENSE must be set: audit logs are an enterprise feature"
+ marker: Final = "callback-delete-" + uuid.uuid4().hex
+ with _pair(gateway, tmp_path, monkeypatch) as pair, pair.first.scenario() as scenario:
+ _configure(
+ pair.first,
+ {list_key: [SIBLING, CALLBACK]},
+ {
+ "GENERIC_LOGGER_ENDPOINT": pair.generic.wire.url,
+ "LANGSMITH_API_KEY": "langsmith-key-" + marker,
+ "LANGSMITH_BASE_URL": pair.langsmith.wire.url,
+ },
+ )
+ eventually(lambda: pair.running(SIBLING), lambda running: running == (True, True), seconds=WORKER_SYNC_SECONDS)
+ model: Final = scenario.model(api_base=pair.provider.url + "/v1")
+ key: Final = scenario.key(models=[model])
+ before: Final = _send_on_both(pair, model, key, marker + "-before", outcome)
+ eventually(lambda: pair.generic.received(before), bool, seconds=25)
+ eventually(lambda: pair.langsmith.received(before), bool, seconds=25)
+
+ status, body = _delete(pair.first, SIBLING)
+
+ assert status == 200, body
+ assert (body["removed_callback"], body["remaining_callbacks"]) == (SIBLING, [CALLBACK])
+ assert _stored_settings(pair.database_url)[list_key] == [CALLBACK]
+ audit: Final = eventually(
+ lambda: read_rows(
+ 'SELECT before_value, updated_values FROM "LiteLLM_AuditLog" WHERE object_id = %s AND action = %s',
+ ("litellm_settings", "deleted"),
+ database_url=pair.database_url,
+ ),
+ bool,
+ seconds=10,
+ )
+ assert len(audit) == 1, audit
+ before_value: Final = list_value(object_value(audit[0]["before_value"])[list_key])
+ assert sorted(str(name) for name in before_value) == sorted([SIBLING, CALLBACK]), audit
+ assert audit[0]["updated_values"] == {list_key: [CALLBACK]}, audit
+ eventually(
+ lambda: pair.running(SIBLING), lambda running: running == (False, False), seconds=WORKER_SYNC_SECONDS
+ )
+
+ after: Final = _send_on_both(pair, model, key, marker + "-after", outcome)
+ eventually(lambda: pair.generic.received(after), bool, seconds=25)
+ _configure(pair.first, {list_key: [CALLBACK, SIBLING]})
+ eventually(lambda: pair.running(SIBLING), lambda running: running == (True, True), seconds=WORKER_SYNC_SECONDS)
+ barrier: Final = _send_on_both(pair, model, key, marker + "-barrier", outcome)
+ eventually(lambda: pair.langsmith.received(barrier), bool, seconds=25)
+
+ leaked: Final = tuple(text for text in after if any(text.encode() in body for body in pair.langsmith.bodies))
+ assert leaked == (), leaked
diff --git a/tests/integration/cost_calculation/cost_tracking_case.py b/tests/integration/cost_calculation/cost_tracking_case.py
index 038c8d6b9e3..208d95e6a14 100644
--- a/tests/integration/cost_calculation/cost_tracking_case.py
+++ b/tests/integration/cost_calculation/cost_tracking_case.py
@@ -235,6 +235,8 @@ class RealtimeResponse(BaseModel):
session_type: str | None = None
created_event: Literal["session.created", "transcription_session.created"] = "session.created"
created_repeats: int = 1
+ created_after_first_frame: bool = False
+ dialect: Literal["openai", "xai"] = "openai"
StoredResponse: TypeAlias = Annotated[
diff --git a/tests/integration/observability/test_bedrock_error_request_id.py b/tests/integration/observability/test_bedrock_error_request_id.py
index 230619a2863..436dcc572d5 100644
--- a/tests/integration/observability/test_bedrock_error_request_id.py
+++ b/tests/integration/observability/test_bedrock_error_request_id.py
@@ -1,13 +1,208 @@
+import asyncio
+import base64
import json
import uuid
-from typing import Final
+from collections.abc import Mapping
+from concurrent.futures import ThreadPoolExecutor
+from types import MappingProxyType
+from typing import Final, Literal, NamedTuple
+from urllib.parse import unquote
-from integration._support.client import Gateway, eventually, object_value
+import anthropic
+import httpx
+import openai
+import pytest
+from integration._support.client import Gateway, Scenario, eventually, object_value
from integration._support.database import read_rows
+from integration._support.responses_stream import (
+ created,
+ delta,
+ error_event,
+ failed,
+ frame,
+ healthy_stream,
+ openai_error_event,
+ spend_row,
+)
+from integration._support.upstream import _aws_event_frame, aws_event_stream_frame
from integration._support.wire import Reply, Request, wire_server
+from pydantic import JsonValue
+
+import litellm
+from litellm.exceptions import MidStreamFallbackError
+from litellm.litellm_core_utils.litellm_logging import _get_provider_request_id
_MODEL: Final = "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0"
+_BEDROCK_RESPONSES_MODEL: Final = "bedrock/us.openai.gpt-5.6-sol"
_TOKEN: Final = "synthetic-bedrock-bearer"
+_SHAPES: Final = ("nested-error", "top-level-error", "response.failed", "error-before-output")
+_ERROR_CODES: Final = ("server_error", "rate_limit_exceeded", "invalid_request_error", "content_policy_violation")
+_INVOKE_MODEL_ID: Final = "anthropic.claude-3-haiku-20240307-v1:0"
+_BEDROCK_AWS: Final[Mapping[str, JsonValue]] = MappingProxyType(
+ {
+ "aws_access_key_id": "AKIASCRIPTEDPROVIDER",
+ "aws_secret_access_key": "scripted-secret",
+ "aws_region_name": "us-east-1",
+ }
+)
+_EVENT_STREAM: Final = "application/vnd.amazon.eventstream"
+_BEDROCK_STREAM_TARGETS: Final = frozenset(
+ (
+ "/model/anthropic.claude-sonnet-4-5-20250929-v1:0/converse-stream",
+ f"/model/{_INVOKE_MODEL_ID}/invoke-with-response-stream",
+ )
+)
+
+BedrockFailure = Literal["exception-event", "exception-message", "no-known-event", "truncated-frame"]
+
+
+class _FailureShape(NamedTuple):
+ bedrock_status: int
+ converse_class: str
+ after_output: tuple[bool, ...]
+
+
+BedrockEndpoint = Literal["chat", "responses", "messages", "messages-invoke"]
+_BEDROCK_ENDPOINTS: Final[tuple[BedrockEndpoint, ...]] = ("chat", "responses", "messages", "messages-invoke")
+_BEDROCK_FAILURES: Final[Mapping[BedrockFailure, _FailureShape]] = MappingProxyType(
+ {
+ "exception-event": _FailureShape(500, "ServiceUnavailableError", (False, True)),
+ "exception-message": _FailureShape(429, "RateLimitError", (False, True)),
+ "no-known-event": _FailureShape(502, "BadGatewayError", (False,)),
+ "truncated-frame": _FailureShape(502, "BadGatewayError", (True,)),
+ }
+)
+_BEDROCK_CELLS: Final = tuple(
+ pytest.param(
+ endpoint, failure, after_output, id=f"{endpoint}-{failure}-{'after' if after_output else 'before'}-output"
+ )
+ for endpoint in _BEDROCK_ENDPOINTS
+ for failure, shape in _BEDROCK_FAILURES.items()
+ for after_output in shape.after_output
+)
+
+
+def _stream(identity: str, shape: str, error_code: str, message: str) -> tuple[bytes, ...]:
+ opening: Final = (frame(created(identity)), frame(delta(identity, "partial output")))
+ nested: Final = frame(error_event({"type": error_code, "code": error_code, "message": message}))
+ match shape:
+ case "nested-error":
+ return (*opening, nested)
+ case "top-level-error":
+ return (*opening, frame(openai_error_event(error_code, message)))
+ case "response.failed":
+ return (*opening, frame(failed(identity, error_code, message)))
+ case "error-before-output":
+ return (nested,)
+ raise AssertionError(shape)
+
+
+def _bedrock_frame(event_type: str, payload: dict[str, JsonValue]) -> bytes:
+ return _aws_event_frame(event_type, payload, "sc", "u")
+
+
+def _invoke_chunk(event: dict[str, JsonValue]) -> bytes:
+ return _bedrock_frame("chunk", {"bytes": base64.b64encode(json.dumps(event).encode()).decode()})
+
+
+def _bedrock_output(invoke: bool, text: str = "partial output") -> tuple[bytes, ...]:
+ if not invoke:
+ return (
+ _bedrock_frame("messageStart", {"role": "assistant"}),
+ _bedrock_frame("contentBlockDelta", {"delta": {"text": text}, "contentBlockIndex": 0}),
+ )
+ started: Final[dict[str, JsonValue]] = {
+ "id": f"msg_{uuid.uuid4().hex}",
+ "type": "message",
+ "role": "assistant",
+ "content": [],
+ "model": _INVOKE_MODEL_ID,
+ "stop_reason": None,
+ "usage": {"input_tokens": 3, "output_tokens": 1},
+ }
+ return (
+ _invoke_chunk({"type": "message_start", "message": started}),
+ _invoke_chunk({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
+ _invoke_chunk({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}),
+ )
+
+
+def _bedrock_failure_frame(failure: BedrockFailure, message: str) -> bytes:
+ match failure:
+ case "exception-event":
+ return _bedrock_frame("internalServerException", {"message": message})
+ case "exception-message":
+ headers: Final = {
+ ":exception-type": "throttlingException",
+ ":content-type": "application/json",
+ ":message-type": "exception",
+ }
+ return aws_event_stream_frame(headers, json.dumps({"message": message}).encode())
+ case "no-known-event":
+ return _bedrock_frame("somethingBedrockAddedLater", {"message": message})
+ case "truncated-frame":
+ return _bedrock_frame("contentBlockDelta", {"delta": {"text": message}, "contentBlockIndex": 0})[:-7]
+
+
+def _bedrock_failure_text(failure: BedrockFailure, message: str) -> str:
+ return "undecoded bytes after" if failure == "truncated-frame" else message
+
+
+def _bedrock_stream_request(
+ endpoint: BedrockEndpoint, model: str, prompt: str, *, stream: bool = True
+) -> tuple[str, dict[str, JsonValue]]:
+ extra: Final[dict[str, JsonValue]] = {"stream": stream, "num_retries": 0, "cache": {"no-cache": True}}
+ turn: Final[list[JsonValue]] = [{"role": "user", "content": prompt}]
+ match endpoint:
+ case "chat":
+ return "/v1/chat/completions", {"model": model, "messages": turn, **extra}
+ case "responses":
+ return "/v1/responses", {"model": model, "input": prompt, **extra}
+ case "messages" | "messages-invoke":
+ return "/v1/messages", {"model": model, "messages": turn, "max_tokens": 16, **extra}
+
+
+def _failure_information(call_id: str) -> dict[str, JsonValue]:
+ row: Final = spend_row(call_id)
+ assert row["status"] == "failure", row
+ metadata: Final = row["metadata"]
+ parsed: Final = json.loads(metadata) if isinstance(metadata, str) else object_value(metadata)
+ return object_value(parsed["error_information"])
+
+
+def _streamed_response(gateway: Gateway, model: str, prompt: str) -> httpx.Response:
+ response: Final = gateway.request(
+ "POST",
+ "/v1/responses",
+ {"model": model, "input": prompt, "stream": True, "num_retries": 0, "cache": {"no-cache": True}},
+ )
+ assert response.status_code == 200, response.text
+ return response
+
+
+async def _consume_sdk_stream(asynchronous: bool, api_base: str, identity: str) -> None:
+ if asynchronous:
+ async_stream: Final = await litellm.aresponses(
+ model="openai/gpt-6",
+ input=identity,
+ stream=True,
+ api_base=api_base,
+ api_key="synthetic-provider-key",
+ num_retries=0,
+ )
+ async for _ in async_stream:
+ pass
+ return
+ tuple(
+ litellm.responses(
+ model="openai/gpt-6",
+ input=identity,
+ stream=True,
+ api_base=api_base,
+ api_key="synthetic-provider-key",
+ num_retries=0,
+ )
+ )
def test_bedrock_500_keeps_amzn_request_id_on_error_headers_and_failure_log(gateway: Gateway) -> None:
@@ -54,3 +249,556 @@ def test_bedrock_500_keeps_amzn_request_id_on_error_headers_and_failure_log(gate
parsed: Final = json.loads(metadata) if isinstance(metadata, str) else object_value(metadata)
error_information: Final = object_value(parsed["error_information"])
assert error_information["error_provider_request_id"] == amzn_request_id, error_information
+
+
+@pytest.mark.parametrize("shape", _SHAPES)
+@pytest.mark.parametrize("error_code", _ERROR_CODES)
+def test_responses_stream_failure_log_retains_provider_request_id(
+ gateway: Gateway, shape: str, error_code: str
+) -> None:
+ identity: Final = f"resp_{uuid.uuid4().hex}"
+ provider_request_id: Final = str(uuid.uuid4())
+ message: Final = f"stream failure {identity}"
+
+ def respond(request: Request) -> Reply:
+ assert request.method == "POST" and request.target == "/responses", request.target
+ assert json.loads(request.body)["stream"] is True
+ return Reply(
+ headers={"x-amzn-RequestId": provider_request_id},
+ content_type="text/event-stream",
+ chunks=_stream(identity, shape, error_code, message),
+ )
+
+ with wire_server(respond) as wire, gateway.scenario() as scenario:
+ model: Final = scenario.model(model="openai/gpt-6", api_base=wire.url, num_retries=0)
+ response: Final = _streamed_response(gateway, model, identity)
+ assert message in response.text, response.text
+ assert ("partial output" in response.text) is (shape != "error-before-output"), response.text
+ assert response.headers.get("llm_provider-x-amzn-requestid") == provider_request_id, dict(response.headers)
+ assert len(wire.drain()) == 1
+ information: Final = _failure_information(response.headers["x-litellm-call-id"])
+ assert information["error_provider_request_id"] == provider_request_id, information
+
+
+@pytest.mark.parametrize("shape", _SHAPES)
+def test_responses_stream_failure_without_provider_request_id_logs_none(gateway: Gateway, shape: str) -> None:
+ identity: Final = f"resp_{uuid.uuid4().hex}"
+ message: Final = f"stream failure {identity}"
+
+ def respond(request: Request) -> Reply:
+ assert request.method == "POST" and request.target == "/responses", request.target
+ return Reply(content_type="text/event-stream", chunks=_stream(identity, shape, "server_error", message))
+
+ with wire_server(respond) as wire, gateway.scenario() as scenario:
+ model: Final = scenario.model(model="openai/gpt-6", api_base=wire.url, num_retries=0)
+ response: Final = _streamed_response(gateway, model, identity)
+ assert message in response.text, response.text
+ assert "llm_provider-x-amzn-requestid" not in response.headers, dict(response.headers)
+ assert len(wire.drain()) == 1
+ information: Final = _failure_information(response.headers["x-litellm-call-id"])
+ assert information["error_class"] == "InternalServerError", information
+ assert information["error_provider_request_id"] is None, information
+
+
+@pytest.mark.parametrize("shape", ("nested-error", "response.failed"))
+def test_bedrock_responses_stream_server_error_logs_amzn_request_id(gateway: Gateway, shape: str) -> None:
+ identity: Final = f"resp_{uuid.uuid4().hex}"
+ amzn_request_id: Final = str(uuid.uuid4())
+ message: Final = f"bedrock stream failure {identity}"
+
+ def respond(request: Request) -> Reply:
+ assert request.method == "POST" and request.target == "/openai/v1/responses", request.target
+ assert request.headers["authorization"] == f"Bearer {_TOKEN}", request.headers
+ assert json.loads(request.body)["stream"] is True
+ return Reply(
+ headers={"x-amzn-RequestId": amzn_request_id},
+ content_type="text/event-stream",
+ chunks=_stream(identity, shape, "server_error", message),
+ )
+
+ with wire_server(respond) as wire, gateway.scenario() as scenario:
+ model: Final = scenario.model(
+ model=_BEDROCK_RESPONSES_MODEL,
+ api_key=_TOKEN,
+ api_base=wire.url,
+ aws_region_name="us-east-1",
+ num_retries=0,
+ )
+ response: Final = _streamed_response(gateway, model, identity)
+ assert "partial output" in response.text and message in response.text, response.text
+ assert response.headers.get("llm_provider-x-amzn-requestid") == amzn_request_id, dict(response.headers)
+ assert len(wire.drain()) == 1
+ information: Final = _failure_information(response.headers["x-litellm-call-id"])
+ assert information["llm_provider"] == "bedrock", information
+ assert information["error_provider_request_id"] == amzn_request_id, information
+
+
+def test_responses_stream_failure_log_reads_openai_x_request_id(gateway: Gateway) -> None:
+ identity: Final = f"resp_{uuid.uuid4().hex}"
+ provider_request_id: Final = f"req_{uuid.uuid4().hex}"
+ message: Final = f"stream failure {identity}"
+
+ def respond(request: Request) -> Reply:
+ assert request.method == "POST" and request.target == "/responses", request.target
+ return Reply(
+ headers={"x-request-id": provider_request_id},
+ content_type="text/event-stream",
+ chunks=_stream(identity, "top-level-error", "server_error", message),
+ )
+
+ with wire_server(respond) as wire, gateway.scenario() as scenario:
+ model: Final = scenario.model(model="openai/gpt-6", api_base=wire.url, num_retries=0)
+ response: Final = _streamed_response(gateway, model, identity)
+ assert message in response.text, response.text
+ assert len(wire.drain()) == 1
+ information: Final = _failure_information(response.headers["x-litellm-call-id"])
+ assert information["error_provider_request_id"] == provider_request_id, information
+
+
+@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async"))
+@pytest.mark.parametrize("event_type", ("error", "response.failed"))
+async def test_responses_sdk_stream_error_retains_provider_request_id(asynchronous: bool, event_type: str) -> None:
+ identity: Final = f"resp_{uuid.uuid4().hex}"
+ provider_request_id: Final = str(uuid.uuid4())
+ shape: Final = "nested-error" if event_type == "error" else "response.failed"
+
+ def respond(request: Request) -> Reply:
+ assert request.method == "POST" and request.target == "/responses", request.target
+ return Reply(
+ headers={"x-amzn-RequestId": provider_request_id, "retry-after": "7"},
+ content_type="text/event-stream",
+ chunks=_stream(identity, shape, "server_error", identity),
+ )
+
+ with wire_server(respond) as wire:
+ with pytest.raises(MidStreamFallbackError) as raised:
+ await _consume_sdk_stream(asynchronous, wire.url, identity)
+ assert raised.value.status_code == 500
+ assert raised.value.response.status_code == 500
+ assert str(raised.value.response.request.url) == "https://openai.com/v1/"
+ assert raised.value.generated_content == "partial output"
+ assert dict(raised.value.response.headers) == {"x-amzn-requestid": provider_request_id}
+ assert _get_provider_request_id(raised.value) == provider_request_id
+ assert raised.value.original_exception is not None
+ assert _get_provider_request_id(raised.value.original_exception) == provider_request_id
+ assert len(wire.drain()) == 1
+
+
+def _concurrent_reply(request: Request, failing: tuple[str, ...], recovery: str) -> Reply:
+ prompt: Final = json.loads(request.body)["input"]
+ assert prompt in (*failing, recovery)
+ chunks: Final = (
+ healthy_stream(f"resp_{prompt}", "recovered")
+ if prompt == recovery
+ else _stream(f"resp_{prompt}", "nested-error", "server_error", prompt)
+ )
+ return Reply(content_type="text/event-stream", headers={"x-amzn-requestid": prompt}, chunks=chunks)
+
+
+def test_concurrent_responses_failures_keep_distinct_request_ids_and_recover(gateway: Gateway) -> None:
+ identity: Final = uuid.uuid4().hex
+ prompts: Final = tuple(f"{identity}-{index}" for index in range(4))
+ recovery: Final = f"{identity}-recovery"
+
+ with (
+ wire_server(lambda request: _concurrent_reply(request, prompts, recovery)) as wire,
+ gateway.scenario() as scenario,
+ ):
+ model: Final = scenario.model(model="openai/gpt-6", api_base=wire.url, num_retries=0)
+
+ def call(prompt: str) -> str:
+ response: Final = _streamed_response(gateway, model, prompt)
+ assert (prompt if prompt != recovery else "recovered") in response.text, response.text
+ return response.headers["x-litellm-call-id"]
+
+ with ThreadPoolExecutor(max_workers=4) as workers:
+ call_ids: Final = tuple(workers.map(call, prompts))
+ recovered_call_id: Final = call(recovery)
+ assert len(wire.drain()) == 5
+ for prompt, call_id in zip(prompts, call_ids, strict=True):
+ assert _failure_information(call_id)["error_provider_request_id"] == prompt
+ assert spend_row(recovered_call_id)["status"] == "success"
+
+
+def _bedrock_logged_error(endpoint: BedrockEndpoint, failure: BedrockFailure) -> tuple[str, str]:
+ shape: Final = _BEDROCK_FAILURES[failure]
+ if endpoint == "messages-invoke":
+ return "BedrockError", str(shape.bedrock_status)
+ return shape.converse_class, "503" if failure == "exception-event" else str(shape.bedrock_status)
+
+
+def _bedrock_client_status(endpoint: BedrockEndpoint, failure: BedrockFailure, after_output: bool) -> int:
+ logged_code: Final = int(_bedrock_logged_error(endpoint, failure)[1])
+ if after_output or endpoint in ("responses", "messages"):
+ return 200
+ if endpoint == "messages-invoke" and failure == "no-known-event":
+ return 200
+ return logged_code
+
+
+@pytest.mark.parametrize(("endpoint", "failure", "after_output"), _BEDROCK_CELLS)
+def test_bedrock_event_stream_failure_log_retains_amzn_request_id(
+ gateway: Gateway, endpoint: BedrockEndpoint, failure: BedrockFailure, after_output: bool
+) -> None:
+ amzn_request_id: Final = str(uuid.uuid4())
+ message: Final = f"bedrock stream failure {amzn_request_id}"
+ invoke: Final = endpoint == "messages-invoke"
+ frames: Final = (*(_bedrock_output(invoke) if after_output else ()), _bedrock_failure_frame(failure, message))
+
+ def respond(request: Request) -> Reply:
+ assert request.method == "POST" and unquote(request.target) in _BEDROCK_STREAM_TARGETS, request.target
+ return Reply(
+ headers={"x-amzn-RequestId": amzn_request_id, "retry-after": "7"},
+ content_type=_EVENT_STREAM,
+ body=b"".join(frames),
+ )
+
+ with wire_server(respond) as wire, gateway.scenario() as scenario:
+ model: Final = scenario.model(
+ model=f"bedrock/invoke/{_INVOKE_MODEL_ID}" if invoke else _MODEL,
+ api_base=wire.url,
+ num_retries=0,
+ **_BEDROCK_AWS,
+ )
+ path, body = _bedrock_stream_request(endpoint, model, message)
+ with gateway.client.stream(
+ "POST", path, json=body, headers={"Authorization": f"Bearer {gateway.key}"}
+ ) as response:
+ text: Final = "\n".join(line for line in response.iter_lines() if line)
+ status: Final = response.status_code
+ response_headers: Final = dict(response.headers)
+ assert status == _bedrock_client_status(endpoint, failure, after_output), text
+ assert _bedrock_failure_text(failure, message) in text and ("partial output" in text) is after_output, text
+ assert len(wire.drain()) == 1
+ information: Final = _failure_information(response_headers["x-litellm-call-id"])
+ assert (information["error_class"], information["error_code"]) == _bedrock_logged_error(endpoint, failure)
+ assert information["error_provider_request_id"] == amzn_request_id, information
+ if status != 200 and not invoke:
+ assert response_headers.get("llm_provider-x-amzn-requestid") == amzn_request_id, response_headers
+ assert "llm_provider-retry-after" not in response_headers, response_headers
+
+
+def _bedrock_model(scenario: Scenario, endpoint: BedrockEndpoint, api_base: str) -> str:
+ return scenario.model(
+ model=f"bedrock/invoke/{_INVOKE_MODEL_ID}" if endpoint == "messages-invoke" else _MODEL,
+ api_base=api_base,
+ num_retries=0,
+ **_BEDROCK_AWS,
+ )
+
+
+def _bedrock_stream(
+ gateway: Gateway, endpoint: BedrockEndpoint, model: str, prompt: str
+) -> tuple[int, str, dict[str, str]]:
+ path, body = _bedrock_stream_request(endpoint, model, prompt)
+ with gateway.client.stream("POST", path, json=body, headers={"Authorization": f"Bearer {gateway.key}"}) as response:
+ text: Final = "\n".join(line for line in response.iter_lines() if line)
+ return response.status_code, text, dict(response.headers)
+
+
+@pytest.mark.parametrize("after_output", (False, True), ids=("before-output", "after-output"))
+@pytest.mark.parametrize("endpoint", _BEDROCK_ENDPOINTS)
+def test_bedrock_event_stream_failure_without_amzn_request_id_logs_none(
+ gateway: Gateway, endpoint: BedrockEndpoint, after_output: bool
+) -> None:
+ message: Final = f"bedrock stream failure {uuid.uuid4()}"
+ frames: Final = (
+ *(_bedrock_output(endpoint == "messages-invoke") if after_output else ()),
+ _bedrock_failure_frame("exception-event", message),
+ )
+
+ def respond(request: Request) -> Reply:
+ assert unquote(request.target) in _BEDROCK_STREAM_TARGETS, request.target
+ return Reply(headers={"retry-after": "7"}, content_type=_EVENT_STREAM, body=b"".join(frames))
+
+ with wire_server(respond) as wire, gateway.scenario() as scenario:
+ status, text, response_headers = _bedrock_stream(
+ gateway, endpoint, _bedrock_model(scenario, endpoint, wire.url), message
+ )
+ assert status == _bedrock_client_status(endpoint, "exception-event", after_output), text
+ assert message in text, text
+ assert not any("amzn" in name for name in response_headers), response_headers
+ assert len(wire.drain()) == 1
+ information: Final = _failure_information(response_headers["x-litellm-call-id"])
+ assert (information["error_class"], information["error_code"]) == _bedrock_logged_error(
+ endpoint, "exception-event"
+ )
+ assert information["error_provider_request_id"] is None, information
+
+
+@pytest.mark.parametrize("endpoint", ("responses", "messages", "messages-invoke"))
+def test_bedrock_non_stream_500_logs_amzn_request_id_on_every_endpoint(
+ gateway: Gateway, endpoint: BedrockEndpoint
+) -> None:
+ amzn_request_id: Final = str(uuid.uuid4())
+ targets: Final = frozenset(
+ (
+ "/model/anthropic.claude-sonnet-4-5-20250929-v1:0/converse",
+ f"/model/{_INVOKE_MODEL_ID}/invoke",
+ )
+ )
+
+ def respond(request: Request) -> Reply:
+ assert request.method == "POST" and unquote(request.target) in targets, request.target
+ return Reply(
+ status=500,
+ headers={"x-amzn-RequestId": amzn_request_id},
+ body=b'{"message":"synthetic bedrock failure"}',
+ )
+
+ with wire_server(respond) as wire, gateway.scenario() as scenario:
+ path, body = _bedrock_stream_request(
+ endpoint, _bedrock_model(scenario, endpoint, wire.url), amzn_request_id, stream=False
+ )
+ response: Final = gateway.request("POST", path, body)
+ assert response.status_code == 503, response.text
+ if endpoint == "responses":
+ assert response.headers.get("llm_provider-x-amzn-requestid") == amzn_request_id, dict(response.headers)
+ assert len(wire.drain()) == 1
+ information: Final = _failure_information(response.headers["x-litellm-call-id"])
+ assert information["error_provider_request_id"] == amzn_request_id, information
+
+
+@pytest.mark.parametrize("amzn_request_id", ("", "a" * 5000, "dup-1, dup-1"), ids=("empty", "5kb", "joined-duplicate"))
+def test_bedrock_stream_failure_logs_odd_request_ids_like_the_non_stream_error(
+ gateway: Gateway, amzn_request_id: str
+) -> None:
+ stream_body: Final = b"".join((*_bedrock_output(False), _bedrock_failure_frame("exception-event", "odd id")))
+
+ def respond(request: Request) -> Reply:
+ if unquote(request.target) in _BEDROCK_STREAM_TARGETS:
+ return Reply(headers={"x-amzn-RequestId": amzn_request_id}, content_type=_EVENT_STREAM, body=stream_body)
+ return Reply(status=500, headers={"x-amzn-RequestId": amzn_request_id}, body=b'{"message":"odd id"}')
+
+ with wire_server(respond) as wire, gateway.scenario() as scenario:
+ model: Final = _bedrock_model(scenario, "chat", wire.url)
+ status, text, stream_headers = _bedrock_stream(gateway, "chat", model, "odd id")
+ path, body = _bedrock_stream_request("chat", model, "odd id", stream=False)
+ non_stream: Final = gateway.request("POST", path, body)
+ assert status == 200 and "odd id" in text, text
+ assert non_stream.status_code == 503, non_stream.text
+ assert len(wire.drain()) == 2
+ streamed: Final = _failure_information(stream_headers["x-litellm-call-id"])
+ unstreamed: Final = _failure_information(non_stream.headers["x-litellm-call-id"])
+ assert streamed["error_provider_request_id"] == unstreamed["error_provider_request_id"], (streamed, unstreamed)
+
+
+_BURST_CALLS: Final = tuple(
+ (endpoint, failure, after_output)
+ for endpoint in _BEDROCK_ENDPOINTS
+ for failure, shape in _BEDROCK_FAILURES.items()
+ for after_output in shape.after_output
+)
+
+
+def _healthy_bedrock_stream(invoke: bool) -> bytes:
+ closing: Final = (
+ (
+ _invoke_chunk({"type": "content_block_stop", "index": 0}),
+ _invoke_chunk(
+ {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 1}}
+ ),
+ _invoke_chunk({"type": "message_stop"}),
+ )
+ if invoke
+ else (
+ _bedrock_frame("contentBlockStop", {"contentBlockIndex": 0}),
+ _bedrock_frame("messageStop", {"stopReason": "end_turn"}),
+ _bedrock_frame("metadata", {"usage": {"inputTokens": 3, "outputTokens": 1, "totalTokens": 4}}),
+ )
+ )
+ return b"".join((*_bedrock_output(invoke, "recovered"), *closing))
+
+
+def test_concurrent_bedrock_stream_failures_on_every_endpoint_keep_distinct_request_ids(gateway: Gateway) -> None:
+ identity: Final = uuid.uuid4().hex
+ failing: Final = {f"{identity}-{index:02d}": call for index, call in enumerate(_BURST_CALLS)}
+ recoveries: Final = {f"{identity}-r{index}": endpoint for index, endpoint in enumerate(_BEDROCK_ENDPOINTS)}
+
+ def respond(request: Request) -> Reply:
+ (token,) = (token for token in (*failing, *recoveries) if token in request.body.decode())
+ if token in recoveries:
+ body = _healthy_bedrock_stream(recoveries[token] == "messages-invoke")
+ else:
+ endpoint, failure, after_output = failing[token]
+ opening = _bedrock_output(endpoint == "messages-invoke") if after_output else ()
+ body = b"".join((*opening, _bedrock_failure_frame(failure, token)))
+ return Reply(headers={"x-amzn-RequestId": token, "retry-after": "7"}, content_type=_EVENT_STREAM, body=body)
+
+ with wire_server(respond) as wire, gateway.scenario() as scenario:
+ models: Final = {endpoint: _bedrock_model(scenario, endpoint, wire.url) for endpoint in _BEDROCK_ENDPOINTS}
+
+ def call(token: str, endpoint: BedrockEndpoint) -> tuple[int, str, dict[str, str]]:
+ return _bedrock_stream(gateway, endpoint, models[endpoint], token)
+
+ with ThreadPoolExecutor(max_workers=8) as workers:
+ outcomes: Final = dict(
+ zip(failing, workers.map(call, failing, (endpoint for endpoint, _, _ in failing.values())), strict=True)
+ )
+ recovered: Final = {token: call(token, endpoint) for token, endpoint in recoveries.items()}
+ assert len(wire.drain()) == len(failing) + len(recoveries)
+ for token, (endpoint, failure, after_output) in failing.items():
+ status, text, response_headers = outcomes[token]
+ assert status == _bedrock_client_status(endpoint, failure, after_output), (token, text)
+ assert _bedrock_failure_text(failure, token) in text, (token, text)
+ information = _failure_information(response_headers["x-litellm-call-id"])
+ assert information["error_provider_request_id"] == token, (token, information)
+ for token, (status, text, response_headers) in recovered.items():
+ assert status == 200 and "recovered" in text, (token, text)
+ assert spend_row(response_headers["x-litellm-call-id"])["status"] == "success"
+
+
+def _sdk_stream_text(
+ client_kind: str, asynchronous: bool, gateway: Gateway, model: str, prompt: str
+) -> tuple[str, str]:
+ extra_body: Final[dict[str, JsonValue]] = {"num_retries": 0, "cache": {"no-cache": True}}
+ turn: Final = [{"role": "user", "content": prompt}]
+
+ def open_stream(client: openai.OpenAI | openai.AsyncOpenAI | anthropic.Anthropic | anthropic.AsyncAnthropic):
+ match client_kind:
+ case "openai-chat":
+ return client.chat.completions.with_raw_response.create(
+ model=model, messages=turn, stream=True, extra_body=extra_body
+ )
+ case "openai-responses":
+ return client.responses.with_raw_response.create(
+ model=model, input=prompt, stream=True, extra_body=extra_body
+ )
+ case "anthropic-messages":
+ return client.messages.with_raw_response.create(
+ model=model, messages=turn, max_tokens=16, stream=True, extra_body=extra_body
+ )
+ raise AssertionError(client_kind)
+
+ seen: Final[list[str]] = []
+ base_url: Final = str(gateway.client.base_url).rstrip("/")
+ if not asynchronous:
+ with httpx.Client(timeout=15, trust_env=False) as transport:
+ client = (
+ anthropic.Anthropic(api_key=gateway.key, base_url=base_url, max_retries=0, http_client=transport)
+ if client_kind == "anthropic-messages"
+ else openai.OpenAI(api_key=gateway.key, base_url=f"{base_url}/v1", max_retries=0, http_client=transport)
+ )
+ raw = open_stream(client)
+ try:
+ seen.extend(repr(event) for event in raw.parse())
+ except (openai.APIError, anthropic.APIError) as error:
+ seen.append(str(error))
+ return raw.headers["x-litellm-call-id"], "\n".join(seen)
+
+ async def run() -> str:
+ async with httpx.AsyncClient(timeout=15, trust_env=False) as transport:
+ client = (
+ anthropic.AsyncAnthropic(api_key=gateway.key, base_url=base_url, max_retries=0, http_client=transport)
+ if client_kind == "anthropic-messages"
+ else openai.AsyncOpenAI(
+ api_key=gateway.key, base_url=f"{base_url}/v1", max_retries=0, http_client=transport
+ )
+ )
+ raw = await open_stream(client)
+ try:
+ async for event in raw.parse():
+ seen.append(repr(event))
+ except (openai.APIError, anthropic.APIError) as error:
+ seen.append(str(error))
+ return raw.headers["x-litellm-call-id"]
+
+ return asyncio.run(run()), "\n".join(seen)
+
+
+@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async"))
+@pytest.mark.parametrize("client_kind", ("openai-chat", "openai-responses", "anthropic-messages"))
+def test_sdk_clients_streaming_bedrock_failure_through_proxy_log_amzn_request_id(
+ gateway: Gateway, client_kind: str, asynchronous: bool
+) -> None:
+ amzn_request_id: Final = str(uuid.uuid4())
+ message: Final = f"bedrock stream failure {amzn_request_id}"
+ frames: Final = (*_bedrock_output(False), _bedrock_failure_frame("exception-event", message))
+
+ def respond(request: Request) -> Reply:
+ assert unquote(request.target) in _BEDROCK_STREAM_TARGETS, request.target
+ return Reply(headers={"x-amzn-RequestId": amzn_request_id}, content_type=_EVENT_STREAM, body=b"".join(frames))
+
+ with wire_server(respond) as wire, gateway.scenario() as scenario:
+ call_id, seen = _sdk_stream_text(
+ client_kind, asynchronous, gateway, _bedrock_model(scenario, "chat", wire.url), message
+ )
+ assert "partial output" in seen and message in seen, seen
+ assert len(wire.drain()) == 1
+ information: Final = _failure_information(call_id)
+ assert information["error_provider_request_id"] == amzn_request_id, information
+
+
+async def _consume_bedrock_sdk_stream(api: Literal["chat", "responses"], asynchronous: bool, api_base: str) -> str:
+ common: Final[dict[str, JsonValue]] = {
+ "model": _MODEL,
+ "stream": True,
+ "api_base": api_base,
+ "num_retries": 0,
+ **_BEDROCK_AWS,
+ }
+ turn: Final[list[JsonValue]] = [{"role": "user", "content": "probe"}]
+ match api, asynchronous:
+ case "chat", True:
+ return "".join([repr(chunk) async for chunk in await litellm.acompletion(messages=turn, **common)])
+ case "chat", False:
+ return "".join(repr(chunk) for chunk in litellm.completion(messages=turn, **common))
+ case "responses", True:
+ return "".join([repr(event) async for event in await litellm.aresponses(input="probe", **common)])
+ case "responses", False:
+ return "".join(repr(event) for event in litellm.responses(input="probe", **common))
+
+
+@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async"))
+@pytest.mark.parametrize("api", ("chat", "responses"))
+async def test_bedrock_sdk_stream_exception_carries_amzn_request_id(
+ api: Literal["chat", "responses"], asynchronous: bool
+) -> None:
+ amzn_request_id: Final = str(uuid.uuid4())
+ message: Final = f"bedrock stream failure {uuid.uuid4()}"
+ frames: Final = (*_bedrock_output(False), _bedrock_failure_frame("exception-event", message))
+
+ def respond(request: Request) -> Reply:
+ assert unquote(request.target) in _BEDROCK_STREAM_TARGETS, request.target
+ return Reply(
+ headers={"x-amzn-RequestId": amzn_request_id, "retry-after": "7"},
+ content_type=_EVENT_STREAM,
+ body=b"".join(frames),
+ )
+
+ with wire_server(respond) as wire:
+ with pytest.raises(MidStreamFallbackError) as raised:
+ await _consume_bedrock_sdk_stream(api, asynchronous, wire.url)
+ assert len(wire.drain()) == 1
+ provider_error: Final = raised.value.original_exception
+ assert isinstance(provider_error, litellm.ServiceUnavailableError), provider_error
+ assert raised.value.generated_content == "partial output"
+ assert message in str(provider_error)
+ assert _get_provider_request_id(provider_error) == amzn_request_id
+ assert "retry-after" not in provider_error.response.headers, provider_error.response.headers
+
+
+@pytest.mark.parametrize("api", ("chat", "responses"))
+@pytest.mark.parametrize("healthy", (True, False), ids=("healthy", "failing"))
+async def test_sync_bedrock_sdk_stream_with_non_ascii_request_id_header_keeps_streaming(
+ api: Literal["chat", "responses"], healthy: bool
+) -> None:
+ body: Final = (
+ _healthy_bedrock_stream(False)
+ if healthy
+ else b"".join((*_bedrock_output(False), _bedrock_failure_frame("exception-event", "non-ascii id failure")))
+ )
+
+ def respond(request: Request) -> Reply:
+ assert unquote(request.target) in _BEDROCK_STREAM_TARGETS, request.target
+ return Reply(headers={"x-amzn-RequestId": "req-\u00fc-1"}, content_type=_EVENT_STREAM, body=body)
+
+ with wire_server(respond) as wire:
+ if healthy:
+ assert "recovered" in await _consume_bedrock_sdk_stream(api, False, wire.url)
+ else:
+ with pytest.raises(MidStreamFallbackError) as raised:
+ await _consume_bedrock_sdk_stream(api, False, wire.url)
+ assert "non-ascii id failure" in str(raised.value.original_exception)
+ assert _get_provider_request_id(raised.value.original_exception) is None
+ assert len(wire.drain()) == 1
diff --git a/tests/integration/observability/test_content_filter_logging_only.py b/tests/integration/observability/test_content_filter_logging_only.py
new file mode 100644
index 00000000000..ab1779c497d
--- /dev/null
+++ b/tests/integration/observability/test_content_filter_logging_only.py
@@ -0,0 +1,712 @@
+"""`litellm_content_filter` in `mode: logging_only`, driven through a real proxy.
+
+The provider is the only double, an owned wire peer. The proxy, its guardrail registry, Postgres and Redis run
+for real with two workers. Every check reads the spend row, which is written after the logging-only scans finish
+"""
+
+from __future__ import annotations
+
+import json
+import uuid
+from collections.abc import Iterator
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from hashlib import sha256
+from pathlib import Path
+from typing import Final
+
+import anthropic
+import openai
+import pytest
+import yaml
+from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment, object_value
+from integration._support.database import read_rows
+from integration._support.mcp import EntryPoint, McpCaller, echo_tool, register_mcp, scripted_peer, tool_calls
+from integration._support.otlp_sink import AuditConfigWriter, SpanSinks, recorded_spans
+from integration._support.process import owned_proxy_process
+from integration._support.wire import Reply, Request, Wire, wire_server
+
+AWS_KEY: Final = "AKIA" + "SYNTHETIC0000000"
+GITHUB_TOKEN: Final = "ghp_" + "a1" * 18
+FORBIDDEN_WORD: Final = "synthetic-forbidden-word"
+LEAK_MARK: Final = "SYNTHETIC-LEAK-TOKEN"
+OUTAGE_MARK: Final = "SYNTHETIC-PROVIDER-OUTAGE"
+OBSERVE_BLOCK: Final = "cf-observe-block"
+OBSERVE_MASK: Final = "cf-observe-mask"
+ENFORCE_BLOCK: Final = "cf-enforce-block"
+ENFORCE_MASK: Final = "cf-enforce-mask"
+OBSERVE_CATEGORY: Final = "cf-observe-category"
+OBSERVE_COMPETITOR: Final = "cf-observe-competitor"
+SECRETS: Final = f"AWS_ACCESS_KEY_ID={AWS_KEY}\nGITHUB_TOKEN={GITHUB_TOKEN}"
+
+
+def _answer(marker: str) -> str:
+ return "synthetic answer " + marker
+
+
+def _provider_answer(body: bytes) -> str:
+ text: Final = body.decode()
+ start: Final = text.find("mark-")
+ assert start >= 0, text
+ answer: Final = _answer(text[start : start + 37])
+ return answer + " " + GITHUB_TOKEN if LEAK_MARK in text else answer
+
+
+def _sse(events: tuple[dict[str, object], ...]) -> tuple[bytes, ...]:
+ return tuple(("data: " + json.dumps(event) + "\n\n").encode() for event in events)
+
+
+def _chat_reply(answer: str, streaming: bool) -> Reply:
+ if not streaming:
+ return Reply(
+ body=json.dumps(
+ {
+ "id": "chatcmpl-" + uuid.uuid4().hex,
+ "object": "chat.completion",
+ "created": 1,
+ "model": "gpt-5.4-mini",
+ "choices": [
+ {"index": 0, "message": {"role": "assistant", "content": answer}, "finish_reason": "stop"}
+ ],
+ "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
+ }
+ ).encode()
+ )
+ identity: Final = "chatcmpl-" + uuid.uuid4().hex
+
+ def chunk(delta: dict[str, object], finish: str | None) -> dict[str, object]:
+ return {
+ "id": identity,
+ "object": "chat.completion.chunk",
+ "created": 1,
+ "model": "gpt-5.4-mini",
+ "choices": [{"index": 0, "delta": delta, "finish_reason": finish}],
+ }
+
+ events: Final = (chunk({"role": "assistant", "content": answer[:9]}, None), chunk({"content": answer[9:]}, "stop"))
+ return Reply(chunks=(*_sse(events), b"data: [DONE]\n\n"), content_type="text/event-stream")
+
+
+def _messages_reply(answer: str) -> Reply:
+ return Reply(
+ body=json.dumps(
+ {
+ "id": "msg_" + uuid.uuid4().hex,
+ "type": "message",
+ "role": "assistant",
+ "model": "claude-haiku-5-5",
+ "content": [{"type": "text", "text": answer}],
+ "stop_reason": "end_turn",
+ "stop_sequence": None,
+ "usage": {"input_tokens": 10, "output_tokens": 5},
+ }
+ ).encode()
+ )
+
+
+def _responses_reply(answer: str) -> Reply:
+ return Reply(
+ body=json.dumps(
+ {
+ "id": "resp_" + uuid.uuid4().hex,
+ "object": "response",
+ "created_at": 1,
+ "status": "completed",
+ "model": "gpt-5.4-mini",
+ "output": [
+ {
+ "type": "message",
+ "id": "msgo_" + uuid.uuid4().hex,
+ "status": "completed",
+ "role": "assistant",
+ "content": [{"type": "output_text", "text": answer, "annotations": []}],
+ }
+ ],
+ "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
+ }
+ ).encode()
+ )
+
+
+def _provider(request: Request) -> Reply:
+ if OUTAGE_MARK.encode() in request.body:
+ return Reply(status=503, body=b'{"error": {"message": "synthetic provider outage", "type": "server_error"}}')
+ answer: Final = _provider_answer(request.body)
+ if request.target.endswith("/v1/messages"):
+ return _messages_reply(answer)
+ if request.target.endswith("/v1/responses"):
+ return _responses_reply(answer)
+ assert request.target.endswith("/v1/chat/completions"), request.target
+ return _chat_reply(answer, bool(json.loads(request.body).get("stream")))
+
+
+def _content_filter(name: str, mode: str, action: str, *, default_on: bool) -> dict[str, object]:
+ return {
+ "guardrail_name": name,
+ "litellm_params": {
+ "guardrail": "litellm_content_filter",
+ "mode": mode,
+ "default_on": default_on,
+ "patterns": [
+ {"pattern_type": "regex", "pattern": "AKIA[0-9A-Z]{16}", "name": "aws_access_key", "action": action},
+ {"pattern_type": "regex", "pattern": "ghp_[A-Za-z0-9]{36}", "name": "github_pat", "action": action},
+ ],
+ "blocked_words": [{"keyword": FORBIDDEN_WORD, "action": action}],
+ },
+ }
+
+
+def _observe_only(name: str, **rules: object) -> dict[str, object]:
+ return {
+ "guardrail_name": name,
+ "litellm_params": {"guardrail": "litellm_content_filter", "mode": "logging_only", "default_on": False, **rules},
+ }
+
+
+def _rig_config(root: Path) -> Path:
+ config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
+ config["litellm_settings"]["cache"] = False
+ config["general_settings"]["store_prompts_in_spend_logs"] = True
+ config["guardrails"] = [
+ _content_filter(OBSERVE_BLOCK, "logging_only", "BLOCK", default_on=True),
+ _content_filter(OBSERVE_MASK, "logging_only", "MASK", default_on=True),
+ _content_filter(ENFORCE_BLOCK, "pre_call", "BLOCK", default_on=False),
+ _content_filter(ENFORCE_MASK, "pre_call", "MASK", default_on=False),
+ _observe_only(
+ OBSERVE_CATEGORY,
+ categories=[
+ {"category": "harmful_violence", "enabled": True, "action": "BLOCK", "severity_threshold": "medium"}
+ ],
+ blocked_words=[{"keyword": FORBIDDEN_WORD, "action": "BLOCK"}],
+ ),
+ _observe_only(
+ OBSERVE_COMPETITOR,
+ competitor_intent_config={
+ "brand_self": ["emirates"],
+ "competitors": ["qatar airways"],
+ "domain_words": ["airline", "flight"],
+ "policy": {"competitor_comparison": "refuse"},
+ "threshold_high": 0.5,
+ },
+ ),
+ ]
+ path: Final = root / "content-filter-logging-only.yaml"
+ path.write_text(yaml.safe_dump(config))
+ return path
+
+
+@dataclass(frozen=True, slots=True)
+class Rig:
+ proxy: Gateway
+ provider: Wire
+ chat_model: str
+ claude_model: str
+ responses_model: str
+ scenario: Scenario
+
+ def base(self) -> str:
+ return str(self.proxy.client.base_url).rstrip("/")
+
+ def provider_bodies(self, marker: str) -> tuple[str, ...]:
+ return tuple(request.body.decode() for request in self.provider.drain() if marker.encode() in request.body)
+
+ def entries(self, response_id: str, name: str) -> tuple[dict[str, object], ...]:
+ return _named_entries(
+ 'SELECT request_id, status, metadata FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (response_id,), name
+ )
+
+
+def _named_entries(query: str, parameters: tuple[str, ...], name: str) -> tuple[dict[str, object], ...]:
+ rows: Final = eventually(lambda: read_rows(query, parameters), lambda values: len(values) >= 1, seconds=70)
+ assert len(rows) == 1, rows
+ recorded: Final = object_value(rows[0]["metadata"]).get("guardrail_information")
+ assert isinstance(recorded, list), rows[0]
+ return tuple(entry for entry in (object_value(item) for item in recorded) if entry.get("guardrail_name") == name)
+
+
+def _findings(entry: dict[str, object]) -> tuple[str, ...]:
+ response: Final = entry.get("guardrail_response")
+ assert isinstance(response, list), entry
+ return tuple(
+ str(detection.get("pattern_name") or detection.get("keyword") or detection.get("type"))
+ for detection in (object_value(item) for item in response)
+ )
+
+
+def _verdicts(entries: tuple[dict[str, object], ...]) -> list[tuple[object, object, tuple[str, ...]]]:
+ return [(entry["guardrail_mode"], entry["guardrail_status"], _findings(entry)) for entry in entries]
+
+
+@pytest.fixture(scope="module")
+def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
+ root: Final = tmp_path_factory.mktemp("content-filter-logging-only")
+ with (
+ gateway_from_environment() as gateway,
+ wire_server(_provider) as provider,
+ owned_proxy_process(gateway, root, {}, config=_rig_config(root), workers=2) as owned,
+ owned.gateway.scenario() as scenario,
+ ):
+ chat: Final = scenario.model(
+ model="openai/gpt-5.4-mini", api_base=provider.url + "/v1", api_key="synthetic-openai-key"
+ )
+ claude: Final = scenario.model(
+ model="anthropic/claude-haiku-5-5", api_base=provider.url, api_key="synthetic-anthropic-key"
+ )
+ responses: Final = scenario.model(
+ model="openai/gpt-5.4-mini", api_base=provider.url + "/v1", api_key="synthetic-openai-key"
+ )
+ yield Rig(owned.gateway, provider, chat, claude, responses, scenario)
+
+
+def _marker() -> str:
+ return "mark-" + uuid.uuid4().hex
+
+
+def _sync_chat(rig: Rig, text: str) -> tuple[str, str]:
+ client: Final = openai.OpenAI(base_url=rig.base() + "/v1", api_key=rig.proxy.key, max_retries=0)
+ completion: Final = client.chat.completions.create(
+ model=rig.chat_model, messages=[{"role": "user", "content": text}]
+ )
+ return completion.id, completion.choices[0].message.content or ""
+
+
+async def _async_chat(rig: Rig, text: str) -> tuple[str, str]:
+ client: Final = openai.AsyncOpenAI(base_url=rig.base() + "/v1", api_key=rig.proxy.key, max_retries=0)
+ completion: Final = await client.chat.completions.create(
+ model=rig.chat_model, messages=[{"role": "user", "content": text}]
+ )
+ return completion.id, completion.choices[0].message.content or ""
+
+
+def _streamed_chat(rig: Rig, text: str) -> tuple[str, str]:
+ client: Final = openai.OpenAI(base_url=rig.base() + "/v1", api_key=rig.proxy.key, max_retries=0)
+ chunks: Final = tuple(
+ client.chat.completions.create(model=rig.chat_model, messages=[{"role": "user", "content": text}], stream=True)
+ )
+ return chunks[0].id, "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices)
+
+
+def _messages(rig: Rig, text: str) -> tuple[str, str]:
+ client: Final = anthropic.Anthropic(base_url=rig.base(), api_key=rig.proxy.key, max_retries=0)
+ message: Final = client.messages.create(
+ model=rig.claude_model, max_tokens=64, messages=[{"role": "user", "content": text}]
+ )
+ assert message.content[0].type == "text"
+ return message.id, message.content[0].text
+
+
+def _responses(rig: Rig, text: str) -> tuple[str, str]:
+ client: Final = openai.OpenAI(base_url=rig.base() + "/v1", api_key=rig.proxy.key, max_retries=0)
+ result: Final = client.responses.create(model=rig.responses_model, input=text)
+ return result.id, result.output_text
+
+
+SYNC_CALLS: Final = {"chat": _sync_chat, "chat_stream": _streamed_chat, "messages": _messages, "responses": _responses}
+EVERY_FINDING: Final = ("aws_access_key", "github_pat", FORBIDDEN_WORD)
+
+
+@pytest.mark.parametrize("endpoint", list(SYNC_CALLS))
+def test_logging_only_block_records_every_match_without_touching_the_call(rig: Rig, endpoint: str) -> None:
+ marker: Final = _marker()
+ prompt: Final = f"{marker} what did the file contain?\n{SECRETS}\n{FORBIDDEN_WORD}"
+ response_id, answer = SYNC_CALLS[endpoint](rig, prompt)
+ assert answer == _answer(marker)
+ sent: Final = rig.provider_bodies(marker)
+ assert len(sent) == 1 and AWS_KEY in sent[0] and GITHUB_TOKEN in sent[0] and FORBIDDEN_WORD in sent[0], sent
+
+ assert _verdicts(rig.entries(response_id, OBSERVE_BLOCK)) == [
+ ("logging_only", "guardrail_flagged", EVERY_FINDING),
+ ("logging_only", "success", ()),
+ ]
+
+
+@pytest.mark.asyncio
+async def test_logging_only_block_records_every_match_for_an_async_client(rig: Rig) -> None:
+ marker: Final = _marker()
+ response_id, answer = await _async_chat(rig, f"{marker} {SECRETS}")
+ assert answer == _answer(marker)
+
+ assert _verdicts(rig.entries(response_id, OBSERVE_BLOCK)) == [
+ ("logging_only", "guardrail_flagged", ("aws_access_key", "github_pat")),
+ ("logging_only", "success", ()),
+ ]
+
+
+def test_logging_only_block_still_scans_the_response_after_a_request_match(rig: Rig) -> None:
+ marker: Final = _marker()
+ response_id, answer = _sync_chat(rig, f"{marker} {LEAK_MARK} key {AWS_KEY}")
+ assert answer == _answer(marker) + " " + GITHUB_TOKEN
+
+ assert _verdicts(rig.entries(response_id, OBSERVE_BLOCK)) == [
+ ("logging_only", "guardrail_flagged", ("aws_access_key",)),
+ ("logging_only", "guardrail_flagged", ("github_pat",)),
+ ]
+
+
+def _tool_call_history(text: str, arguments: str) -> list[dict[str, object]]:
+ return [
+ {"role": "user", "content": text},
+ {
+ "role": "assistant",
+ "content": None,
+ "tool_calls": [
+ {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": arguments}}
+ ],
+ },
+ {"role": "tool", "tool_call_id": "call_1", "content": "done"},
+ ]
+
+
+@pytest.mark.parametrize(
+ ("arguments", "request_findings"),
+ (
+ pytest.param(json.dumps({AWS_KEY: 1}), ("aws_access_key",), id="masked-argument-key"),
+ pytest.param("[" * 150 + "1" + "]" * 150, (), id="too-deep"),
+ ),
+)
+def test_logging_only_flags_tool_arguments_an_enforcing_mode_rejects_and_keeps_scanning(
+ rig: Rig, arguments: str, request_findings: tuple[str, ...]
+) -> None:
+ marker: Final = _marker()
+ messages: Final = _tool_call_history(f"{marker} {LEAK_MARK}", arguments)
+ headers: Final = {"Authorization": f"Bearer {rig.proxy.key}"}
+ enforced: Final = rig.proxy.client.post(
+ "/v1/chat/completions",
+ json={"model": rig.chat_model, "messages": messages, "guardrails": [ENFORCE_MASK]},
+ headers=headers,
+ )
+ assert enforced.status_code == 400, enforced.text
+ assert rig.provider_bodies(marker) == ()
+
+ observed: Final = rig.proxy.client.post(
+ "/v1/chat/completions", json={"model": rig.chat_model, "messages": messages}, headers=headers
+ )
+ assert observed.status_code == 200, observed.text
+ completion: Final = object_value(observed.json())
+ assert len(rig.provider_bodies(marker)) == 1
+
+ assert _verdicts(rig.entries(str(completion["id"]), OBSERVE_MASK)) == [
+ ("logging_only", "guardrail_flagged", request_findings),
+ ("logging_only", "success", ("github_pat",)),
+ ]
+
+
+@pytest.mark.parametrize("entry", ("mcp", "rest"))
+def test_logging_only_block_records_every_match_on_an_mcp_tool_call(rig: Rig, entry: EntryPoint) -> None:
+ alias: Final = "cfmcp" + uuid.uuid4().hex[:8]
+ arguments: Final = {"note": f"{_marker()} {SECRETS}"}
+ with scripted_peer(echo_tool("echo")) as peer:
+ server_id: Final = register_mcp(rig.scenario, peer, alias)
+ key: Final = rig.scenario.key(object_permission={"mcp_servers": [server_id]})
+ outcome: Final = McpCaller(rig.proxy, key, entry, alias).call(f"{alias}-echo", arguments, server_id=server_id)
+ assert (outcome.error, outcome.text) == (None, json.dumps(arguments, sort_keys=True)), outcome.raw
+ assert [call["body"]["params"]["arguments"] for call in tool_calls(peer.drain())] == [arguments]
+
+ entries: Final = _named_entries(
+ 'SELECT metadata FROM "LiteLLM_SpendLogs" WHERE api_key = %s AND call_type = %s',
+ (sha256(key.encode()).hexdigest(), "call_mcp_tool"),
+ OBSERVE_BLOCK,
+ )
+ assert _verdicts(entries) == [("logging_only", "guardrail_flagged", ("aws_access_key", "github_pat"))], entries
+
+
+@pytest.mark.parametrize(
+ ("guardrail", "prompt", "findings"),
+ (
+ pytest.param(
+ OBSERVE_CATEGORY,
+ f"how to kill someone {FORBIDDEN_WORD}",
+ ("how to kill someone", FORBIDDEN_WORD),
+ id="category-keyword",
+ ),
+ pytest.param(
+ OBSERVE_COMPETITOR, "Is Qatar Airways better than Emirates?", ("competitor_intent",), id="competitor-intent"
+ ),
+ ),
+)
+def test_logging_only_block_flags_category_and_competitor_matches_without_blocking(
+ rig: Rig, guardrail: str, prompt: str, findings: tuple[str, ...]
+) -> None:
+ marker: Final = _marker()
+ response: Final = rig.proxy.client.post(
+ "/v1/chat/completions",
+ json={
+ "model": rig.chat_model,
+ "messages": [{"role": "user", "content": f"{marker} {prompt}"}],
+ "guardrails": [guardrail],
+ },
+ headers={"Authorization": f"Bearer {rig.proxy.key}"},
+ )
+ assert response.status_code == 200, response.text
+ completion: Final = object_value(response.json())
+ assert len(rig.provider_bodies(marker)) == 1
+
+ assert _verdicts(rig.entries(str(completion["id"]), guardrail)) == [
+ ("logging_only", "guardrail_flagged", findings),
+ ("logging_only", "success", ()),
+ ]
+
+
+def test_logging_only_mask_records_every_match(rig: Rig) -> None:
+ marker: Final = _marker()
+ response_id, answer = _sync_chat(rig, f"{marker} {SECRETS}\n{FORBIDDEN_WORD}")
+ assert answer == _answer(marker)
+
+ entries: Final = rig.entries(response_id, OBSERVE_MASK)
+ assert _verdicts(entries) == [
+ ("logging_only", "success", EVERY_FINDING),
+ ("logging_only", "success", ()),
+ ]
+ assert entries[0]["masked_entity_count"] == {"aws_access_key": 1, "github_pat": 1, "blocked_word": 1}, entries
+
+
+def test_logging_only_records_nothing_for_a_clean_call(rig: Rig) -> None:
+ marker: Final = _marker()
+ response_id, answer = _sync_chat(rig, f"{marker} nothing secret here")
+ assert answer == _answer(marker)
+
+ assert _verdicts(rig.entries(response_id, OBSERVE_BLOCK)) == [
+ ("logging_only", "success", ()),
+ ("logging_only", "success", ()),
+ ]
+
+
+def test_pre_call_block_still_rejects_the_request(rig: Rig) -> None:
+ marker: Final = _marker()
+ response: Final = rig.proxy.client.post(
+ "/v1/chat/completions",
+ json={
+ "model": rig.chat_model,
+ "messages": [{"role": "user", "content": f"{marker} {SECRETS}"}],
+ "guardrails": [ENFORCE_BLOCK],
+ },
+ headers={"Authorization": f"Bearer {rig.proxy.key}"},
+ )
+ assert response.status_code == 400, response.text
+ assert "aws_access_key" in response.text
+ assert rig.provider_bodies(marker) == ()
+
+
+def test_apply_guardrail_endpoint_still_blocks_with_a_logging_only_content_filter(rig: Rig) -> None:
+ response: Final = rig.proxy.client.post(
+ "/guardrails/apply_guardrail",
+ json={"guardrail_name": OBSERVE_BLOCK, "text": f"key {AWS_KEY}"},
+ headers={"Authorization": f"Bearer {rig.proxy.key}"},
+ )
+ assert response.status_code == 400, response.text
+ assert "aws_access_key" in response.text
+
+
+def _advertised_modes(rig: Rig) -> dict[str, object]:
+ response: Final = rig.proxy.client.get(
+ "/guardrails/ui/add_guardrail_settings", headers={"Authorization": f"Bearer {rig.proxy.key}"}
+ )
+ assert response.status_code == 200, response.text
+ return object_value(response.json()["supported_modes_by_provider"])
+
+
+@pytest.mark.parametrize("provider", ["litellm_content_filter", "generic_guardrail_api", "bedrock", "akto"])
+def test_dashboard_offers_logging_only_where_the_config_accepts_it(rig: Rig, provider: str) -> None:
+ modes: Final = _advertised_modes(rig)
+ offered: Final = modes[provider]
+ assert isinstance(offered, list) and "logging_only" in offered, modes
+ assert offered.count("logging_only") == 1, offered
+
+
+@pytest.mark.parametrize("provider", ["azure/prompt_shield", "lakera_v2"])
+def test_dashboard_does_not_offer_logging_only_where_the_config_rejects_it(rig: Rig, provider: str) -> None:
+ offered: Final = _advertised_modes(rig)[provider]
+ assert isinstance(offered, list) and "logging_only" not in offered, offered
+
+
+def test_a_logging_only_content_filter_created_through_the_api_records_every_match(rig: Rig) -> None:
+ name: Final = "cf-observe-api-" + uuid.uuid4().hex[:8]
+ created: Final = rig.proxy.client.post(
+ "/guardrails",
+ json={"guardrail": {"guardrail_name": name, **_content_filter(name, "logging_only", "BLOCK", default_on=True)}},
+ headers={"Authorization": f"Bearer {rig.proxy.key}"},
+ )
+ assert created.status_code == 200, created.text
+ guardrail_id: Final = str(object_value(created.json())["guardrail_id"])
+ try:
+ marker: Final = _marker()
+ response_id: Final = eventually(
+ lambda: _sync_chat(rig, f"{marker} {SECRETS}")[0],
+ lambda identity: bool(rig.entries(identity, name)),
+ seconds=70,
+ )
+ assert _verdicts(rig.entries(response_id, name))[0] == (
+ "logging_only",
+ "guardrail_flagged",
+ ("aws_access_key", "github_pat"),
+ )
+ finally:
+ rig.proxy.client.delete(f"/guardrails/{guardrail_id}", headers={"Authorization": f"Bearer {rig.proxy.key}"})
+
+
+def test_guardrails_monitor_counts_observe_only_matches_as_flagged_not_blocked(rig: Rig) -> None:
+ name: Final = "cf-observe-usage-" + uuid.uuid4().hex[:8]
+ headers: Final = {"Authorization": f"Bearer {rig.proxy.key}"}
+ created: Final = rig.proxy.client.post(
+ "/guardrails",
+ json={"guardrail": {"guardrail_name": name, **_content_filter(name, "logging_only", "BLOCK", default_on=True)}},
+ headers=headers,
+ )
+ assert created.status_code == 200, created.text
+ guardrail_id: Final = str(object_value(created.json())["guardrail_id"])
+ try:
+ eventually(
+ lambda: _sync_chat(rig, f"{_marker()} {SECRETS}")[0],
+ lambda identity: bool(rig.entries(identity, name)),
+ seconds=70,
+ )
+ today: Final = datetime.now(timezone.utc).date().isoformat()
+ detail: Final = eventually(
+ lambda: object_value(
+ rig.proxy.client.get(
+ f"/guardrails/usage/detail/{guardrail_id}",
+ params={"start_date": today, "end_date": today},
+ headers=headers,
+ ).json()
+ ),
+ lambda body: int(str(body.get("requestsEvaluated", 0))) >= 1,
+ seconds=60,
+ )
+ assert (detail["failRate"], detail["status"]) == (0.0, "healthy"), detail
+ finally:
+ rig.proxy.client.delete(f"/guardrails/{guardrail_id}", headers=headers)
+
+
+@pytest.mark.parametrize(
+ ("environment", "name_key", "status_key"),
+ (
+ pytest.param({}, "guardrail_name", "guardrail_status", id="otel-v1"),
+ pytest.param({"LITELLM_OTEL_V2": "1"}, "litellm.guardrail.name", "litellm.guardrail.status", id="otel-v2"),
+ ),
+)
+def test_otel_guardrail_spans_carry_the_flagged_status_without_an_error(
+ gateway: Gateway,
+ audit_sinks: SpanSinks,
+ otel_audit_config: AuditConfigWriter,
+ tmp_path: Path,
+ environment: dict[str, str],
+ name_key: str,
+ status_key: str,
+) -> None:
+ config_path: Final = otel_audit_config(tmp_path, {"cache": False})
+ config: Final = yaml.safe_load(config_path.read_text())
+ config["guardrails"] = [_content_filter(OBSERVE_BLOCK, "logging_only", "BLOCK", default_on=True)]
+ config_path.write_text(yaml.safe_dump(config))
+ overrides: Final = {"OTEL_EXPORTER": "http/json", "OTEL_ENDPOINT": audit_sinks.operator, **environment}
+ with (
+ wire_server(_provider) as provider,
+ owned_proxy_process(gateway, tmp_path, overrides, config=config_path, workers=2) as owned,
+ owned.gateway.scenario() as scenario,
+ ):
+ model: Final = scenario.model(
+ model="openai/gpt-5.4-mini", api_base=provider.url + "/v1", api_key="synthetic-openai-key"
+ )
+ start, _ = recorded_spans(audit_sinks.operator)
+ marker: Final = _marker()
+ client: Final = openai.OpenAI(
+ base_url=str(owned.gateway.client.base_url).rstrip("/") + "/v1", api_key=owned.gateway.key, max_retries=0
+ )
+ completion: Final = client.chat.completions.create(
+ model=model, messages=[{"role": "user", "content": f"{marker} {SECRETS}"}]
+ )
+ assert completion.choices[0].message.content == _answer(marker)
+
+ spans: Final = eventually(
+ lambda: tuple(
+ span
+ for span in recorded_spans(audit_sinks.operator, start)[1]
+ if span["attributes"].get(name_key) == OBSERVE_BLOCK
+ ),
+ lambda found: len(found) >= 2,
+ seconds=60,
+ return_last_on_timeout=True,
+ )
+ assert sorted(
+ (str(span["attributes"].get(status_key)), "error.type" in span["attributes"]) for span in spans
+ ) == [("guardrail_flagged", False), ("success", False)], spans
+
+
+def _burst_call(rig: Rig, index: int) -> tuple[str, int, str, str]:
+ marker: Final = _marker()
+ if index % 4 == 0:
+ enforced: Final = rig.proxy.client.post(
+ "/v1/chat/completions",
+ json={
+ "model": rig.chat_model,
+ "messages": [{"role": "user", "content": f"{marker} {SECRETS}"}],
+ "guardrails": [ENFORCE_BLOCK],
+ },
+ headers={"Authorization": f"Bearer {rig.proxy.key}"},
+ )
+ return "enforce", enforced.status_code, marker, ""
+ if index % 4 == 1:
+ failed: Final = rig.proxy.client.post(
+ "/v1/chat/completions",
+ json={
+ "model": rig.chat_model,
+ "messages": [{"role": "user", "content": f"{marker} {OUTAGE_MARK} {SECRETS}"}],
+ },
+ headers={"Authorization": f"Bearer {rig.proxy.key}"},
+ )
+ return "outage", failed.status_code, marker, failed.headers.get("x-litellm-call-id", "")
+ endpoint: Final = list(SYNC_CALLS)[index % len(SYNC_CALLS)]
+ response_id, answer = SYNC_CALLS[endpoint](rig, f"{marker} {SECRETS}\n{FORBIDDEN_WORD}")
+ assert answer == _answer(marker), (endpoint, answer)
+ return "observe:" + endpoint, 200, marker, response_id
+
+
+def _failure_row(call_id: str) -> tuple[object, tuple[dict[str, object], ...]]:
+ rows: Final = eventually(
+ lambda: read_rows('SELECT status, metadata FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (call_id,)),
+ lambda values: len(values) >= 1,
+ seconds=70,
+ )
+ assert len(rows) == 1, rows
+ recorded: Final = object_value(rows[0]["metadata"]).get("guardrail_information") or []
+ assert isinstance(recorded, list), rows[0]
+ return rows[0]["status"], tuple(
+ entry for entry in (object_value(item) for item in recorded) if entry.get("guardrail_name") == OBSERVE_BLOCK
+ )
+
+
+def test_concurrent_burst_with_a_provider_outage_records_every_observed_call_once(rig: Rig) -> None:
+ def call(index: int) -> tuple[str, int, str, str]:
+ return _burst_call(rig, index)
+
+ with ThreadPoolExecutor(max_workers=12) as pool:
+ results: Final = tuple(pool.map(call, range(32)))
+
+ assert sorted({(kind.split(":")[0], status) for kind, status, _, _ in results}) == [
+ ("enforce", 400),
+ ("observe", 200),
+ ("outage", 503),
+ ], results
+ sent: Final = rig.provider_bodies("mark-")
+ for kind, _, marker, _ in results:
+ if kind == "enforce":
+ assert not any(marker in body for body in sent), marker
+ outages: Final = tuple(call_id for kind, _, _, call_id in results if kind == "outage")
+ assert len(outages) == 8 and all(outages), results
+ assert [_failure_row(call_id) for call_id in outages] == [("failure", ())] * 8
+ observed: Final = tuple((kind, response_id) for kind, _, _, response_id in results if kind.startswith("observe"))
+ assert len(observed) == 16
+ for kind, response_id in observed:
+ assert _verdicts(rig.entries(response_id, OBSERVE_BLOCK)) == [
+ ("logging_only", "guardrail_flagged", EVERY_FINDING),
+ ("logging_only", "success", ()),
+ ], (kind, response_id)
+
+ marker: Final = _marker()
+ response_id, answer = _sync_chat(rig, f"{marker} {SECRETS}")
+ assert answer == _answer(marker)
+ assert _verdicts(rig.entries(response_id, OBSERVE_BLOCK))[0] == (
+ "logging_only",
+ "guardrail_flagged",
+ ("aws_access_key", "github_pat"),
+ )
diff --git a/tests/integration/observability/test_logging_only_scope_config.py b/tests/integration/observability/test_logging_only_scope_config.py
index c8c765a68a2..354e3b464fe 100644
--- a/tests/integration/observability/test_logging_only_scope_config.py
+++ b/tests/integration/observability/test_logging_only_scope_config.py
@@ -179,16 +179,16 @@ def test_native_content_filter_scope_logs_without_blocking(
(
identity,
"logging_only",
- "guardrail_intervened" if direction == blocked_side else "success",
+ "guardrail_flagged" if direction == blocked_side else "success",
)
for direction in expected_directions
), entries
- intervened_entries: Final = tuple(
- entry for entry in entries if entry["guardrail_status"] == "guardrail_intervened"
+ flagged_entries: Final = tuple(
+ entry for entry in entries if entry["guardrail_status"] == "guardrail_flagged"
)
- assert len(intervened_entries) == 1, entries
- assert intervened_entries[0]["guardrail_response"] is not None, entries
- assert intervened_entries[0]["guardrail_response"] == "REDACTED_BY_LITELM", entries
+ assert len(flagged_entries) == 1, entries
+ assert flagged_entries[0]["guardrail_response"] is not None, entries
+ assert flagged_entries[0]["guardrail_response"] == "REDACTED_BY_LITELM", entries
finally:
delete_scenario(upstream_handle)
diff --git a/tests/integration/observability/test_otel_excluded_services_matrix.py b/tests/integration/observability/test_otel_excluded_services_matrix.py
index 07646a0b189..8a7236f8626 100644
--- a/tests/integration/observability/test_otel_excluded_services_matrix.py
+++ b/tests/integration/observability/test_otel_excluded_services_matrix.py
@@ -358,6 +358,10 @@ def _has_root(spans: tuple[Span, ...]) -> bool:
return any(span["kind"] == SERVER for span in spans)
+def _has_model_span(spans: tuple[Span, ...]) -> bool:
+ return any("gen_ai.operation.name" in span["attributes"] for span in spans)
+
+
def _trace_of_call(sink: str, call_id: str, since: int) -> tuple[Span, ...]:
_, spans = recorded_spans(sink, since)
traces: Final = {span["trace_id"] for span in spans if span["attributes"].get("litellm.call_id") == call_id}
@@ -367,7 +371,7 @@ def _trace_of_call(sink: str, call_id: str, since: int) -> tuple[Span, ...]:
def _operator_trace(rig: Rig, sent: Sent, cursors: Cursors) -> tuple[Span, ...]:
trace: Final = eventually(
lambda: _trace_of_call(rig.sinks.operator, sent.call_id, cursors.operator),
- lambda spans: _has_root(spans) and "redis" in _db_systems(spans),
+ lambda spans: _has_root(spans) and "redis" in _db_systems(spans) and _has_model_span(spans),
seconds=40,
)
assert len({span["trace_id"] for span in trace}) == 1, _names(trace)
@@ -405,7 +409,7 @@ def _assert_tenant_mirrors(rig: Rig, operator: tuple[Span, ...], cursors: Cursor
def _assert_withheld(rig: Rig, sent: Sent, cursors: Cursors) -> tuple[Span, ...]:
tenant: Final = _assert_tenant_mirrors(rig, _operator_trace(rig, sent, cursors), cursors)
- assert any("gen_ai.operation.name" in span["attributes"] for span in tenant), _names(tenant)
+ assert _has_model_span(tenant), _names(tenant)
return tenant
@@ -719,7 +723,7 @@ def _assert_tenant_kept(
return (
sum(1 for span in spans if span["kind"] == SERVER) == 1
and "redis" in _db_systems(spans)
- and (not needs_model_span or any("gen_ai.operation.name" in span["attributes"] for span in spans))
+ and (not needs_model_span or _has_model_span(spans))
)
tenant: Final = eventually(
@@ -730,7 +734,7 @@ def _assert_tenant_kept(
)
assert sum(1 for span in tenant if span["kind"] == SERVER) == 1, _names(tenant)
assert "redis" in _db_systems(tenant), f"redis spans missing at the tenant: {_names(tenant)}"
- assert not needs_model_span or any("gen_ai.operation.name" in span["attributes"] for span in tenant), _names(tenant)
+ assert not needs_model_span or _has_model_span(tenant), _names(tenant)
return tenant
diff --git a/tests/integration/providers/test_anthropic_empty_user_turn_wire.py b/tests/integration/providers/test_anthropic_empty_user_turn_wire.py
new file mode 100644
index 00000000000..654a88c60bb
--- /dev/null
+++ b/tests/integration/providers/test_anthropic_empty_user_turn_wire.py
@@ -0,0 +1,216 @@
+import json
+from collections.abc import Iterator, Mapping, Sequence
+from dataclasses import dataclass
+from pathlib import Path
+from types import MappingProxyType
+from typing import Final
+
+import pytest
+import yaml
+from integration._support.client import Gateway, gateway_from_environment, object_value
+from integration._support.process import graceful_stop_seconds, owned_proxy
+from integration._support.wire import Reply, Request, Wire, wire_server
+from integration.providers._cache_control_marks_support import (
+ anthropic_deployment,
+ anthropic_peer,
+ final_text,
+ new_marker,
+)
+from pydantic import JsonValue, TypeAdapter
+
+pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
+
+_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
+_MODEL: Final = "empty-user-turn"
+_PLACEHOLDER: Final = "[System: Empty message content sanitised to satisfy protocol]"
+_REPLY: Final = "Noted."
+_EMPTY_ROWS: Final[Mapping[str, Mapping[str, JsonValue]]] = MappingProxyType(
+ {
+ "list": MappingProxyType({"role": "user", "content": []}),
+ "null": MappingProxyType({"role": "user", "content": None}),
+ "missing": MappingProxyType({"role": "user"}),
+ "empty-string": MappingProxyType({"role": "user", "content": ""}),
+ }
+)
+_CHAT_EMPTY_KINDS: Final = ("list", "null", "missing")
+
+Turn = tuple[str, tuple[str, ...]]
+
+
+@dataclass(frozen=True, slots=True)
+class _Rig:
+ gateway: Gateway
+ provider: Wire
+
+
+def _empty_row(kind: str) -> dict[str, JsonValue]:
+ return dict(_EMPTY_ROWS[kind])
+
+
+def _guardrail_peer(request: Request) -> Reply:
+ kind: Final = request.target.split("/")[1]
+ rows: Final = _JSON_OBJECT.validate_json(request.body)["structured_messages"]
+ assert isinstance(rows, list), rows
+ verdict: Final = {"action": "GUARDRAIL_INTERVENED", "structured_messages": [*rows, _empty_row(kind)]}
+ return Reply(body=json.dumps(verdict).encode())
+
+
+def _guardrail(kind: str, api_base: str) -> dict[str, JsonValue]:
+ return {
+ "guardrail_name": f"append-{kind}",
+ "litellm_params": {
+ "guardrail": "generic_guardrail_api",
+ "mode": "pre_call",
+ "default_on": False,
+ "api_base": f"{api_base}/{kind}",
+ "api_key": "synthetic-guardrail-key",
+ },
+ }
+
+
+def _config(directory: Path, provider_url: str, guardrail_url: str) -> Path:
+ stock: Final = _JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
+ config: Final[dict[str, JsonValue]] = {
+ **stock,
+ "model_list": [anthropic_deployment(_MODEL, provider_url)],
+ "router_settings": {**object_value(stock["router_settings"]), "num_retries": 0},
+ "guardrails": [_guardrail(kind, guardrail_url) for kind in _EMPTY_ROWS],
+ }
+ path: Final = directory / "empty-user-turn.yaml"
+ path.write_text(yaml.safe_dump(config))
+ return path
+
+
+@pytest.fixture(scope="module")
+def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]:
+ directory: Final = tmp_path_factory.mktemp("empty-user-turn")
+ with (
+ wire_server(anthropic_peer) as provider,
+ wire_server(_guardrail_peer) as guardrail,
+ gateway_from_environment() as gateway,
+ ):
+ config: Final = _config(directory, provider.url, guardrail.url)
+ with owned_proxy(gateway, directory, {}, config=config, workers=2) as owned:
+ yield _Rig(owned, provider)
+
+
+def _block(block: JsonValue) -> str:
+ part: Final = object_value(block)
+ return str(part["text"]) if part.get("type") == "text" else str(part.get("type"))
+
+
+def _turn(message: JsonValue) -> Turn:
+ row: Final = object_value(message)
+ content: Final = row["content"]
+ if isinstance(content, str):
+ return (str(row["role"]), (content,))
+ assert isinstance(content, list), row
+ return (str(row["role"]), tuple(_block(block) for block in content))
+
+
+def _send(rig: _Rig, path: str, body: Mapping[str, JsonValue]) -> None:
+ response: Final = rig.gateway.client.post(
+ path, json=dict(body), headers={"Authorization": f"Bearer {rig.gateway.key}"}
+ )
+ assert response.status_code == 200, response.text
+
+
+def _provider_turns(rig: _Rig, marker: str) -> tuple[Turn, ...]:
+ needle: Final = marker.encode()
+ requests: Final = [request for request in rig.provider.drain() if needle in request.body]
+ assert len(requests) == 1, requests
+ messages: Final = _JSON_OBJECT.validate_json(requests[0].body)["messages"]
+ assert isinstance(messages, list), messages
+ return tuple(_turn(message) for message in messages)
+
+
+def _ask_row(ask: str) -> dict[str, JsonValue]:
+ return {"role": "user", "content": ask}
+
+
+def _chat(messages: Sequence[JsonValue]) -> dict[str, JsonValue]:
+ return {"model": _MODEL, "max_tokens": 64, "messages": list(messages)}
+
+
+@pytest.mark.parametrize("kind", _CHAT_EMPTY_KINDS)
+def test_chat_trailing_empty_user_message_keeps_a_user_turn(rig: _Rig, kind: str) -> None:
+ marker: Final = new_marker()
+ ask: Final = final_text(marker)
+ _send(
+ rig,
+ "/v1/chat/completions",
+ _chat([_ask_row(ask), {"role": "assistant", "content": _REPLY}, _empty_row(kind)]),
+ )
+ assert _provider_turns(rig, marker) == (("user", (ask,)), ("assistant", (_REPLY,)), ("user", (_PLACEHOLDER,)))
+
+
+@pytest.mark.parametrize("kind", _CHAT_EMPTY_KINDS)
+def test_chat_empty_user_message_beside_text_sends_only_the_text(rig: _Rig, kind: str) -> None:
+ marker: Final = new_marker()
+ ask: Final = final_text(marker)
+ _send(rig, "/v1/chat/completions", _chat([_ask_row(ask), _empty_row(kind)]))
+ assert _provider_turns(rig, marker) == (("user", (ask,)),)
+
+
+def _system_then_empty_turns(rig: _Rig, kind: str) -> tuple[Turn, ...]:
+ marker: Final = new_marker()
+ ask: Final = final_text(marker)
+ _send(
+ rig,
+ "/v1/chat/completions",
+ _chat(
+ [
+ _ask_row(ask),
+ {"role": "assistant", "content": _REPLY},
+ {"role": "system", "content": "Answer in French."},
+ _empty_row(kind),
+ ]
+ ),
+ )
+ turns: Final = _provider_turns(rig, marker)
+ assert turns[0] == ("user", (ask,)), turns
+ return turns[1:]
+
+
+@pytest.mark.parametrize("kind", _CHAT_EMPTY_KINDS)
+def test_chat_system_message_before_an_empty_user_message_is_placed_like_an_empty_string(rig: _Rig, kind: str) -> None:
+ assert _system_then_empty_turns(rig, kind) == _system_then_empty_turns(rig, "empty-string")
+
+
+def test_responses_trailing_empty_input_message_keeps_a_user_turn(rig: _Rig) -> None:
+ marker: Final = new_marker()
+ ask: Final = final_text(marker)
+ body: Final[dict[str, JsonValue]] = {
+ "model": _MODEL,
+ "max_output_tokens": 64,
+ "input": [
+ _ask_row(ask),
+ {"role": "assistant", "content": _REPLY},
+ {"role": "user", "content": []},
+ ],
+ }
+ _send(rig, "/v1/responses", body)
+ assert _provider_turns(rig, marker) == (("user", (ask,)), ("assistant", (_REPLY,)), ("user", (_PLACEHOLDER,)))
+
+
+@pytest.mark.parametrize("kind", tuple(_EMPTY_ROWS))
+def test_messages_guardrail_empty_user_row_is_sent_like_an_empty_string(rig: _Rig, kind: str) -> None:
+ marker: Final = new_marker()
+ ask: Final = final_text(marker)
+ body: Final[dict[str, JsonValue]] = {
+ "model": _MODEL,
+ "max_tokens": 64,
+ "messages": [_ask_row(ask)],
+ "guardrails": [f"append-{kind}"],
+ }
+ _send(rig, "/v1/messages", body)
+ assert _provider_turns(rig, marker) == (("user", (ask,)), ("user", (_PLACEHOLDER,)))
+
+
+def test_a_long_run_of_empty_user_messages_is_sent_as_one_text_turn(rig: _Rig) -> None:
+ marker: Final = new_marker()
+ ask: Final = final_text(marker)
+ _send(
+ rig, "/v1/chat/completions", _chat([*({"role": "user", "content": []} for _ in range(10_000)), _ask_row(ask)])
+ )
+ assert _provider_turns(rig, marker) == (("user", (ask,)),)
diff --git a/tests/integration/providers/test_count_tokens_error_shapes_wire.py b/tests/integration/providers/test_count_tokens_error_shapes_wire.py
new file mode 100644
index 00000000000..200f5e81673
--- /dev/null
+++ b/tests/integration/providers/test_count_tokens_error_shapes_wire.py
@@ -0,0 +1,205 @@
+import json
+from collections.abc import Iterator
+from typing import Final
+
+import anthropic
+import openai
+import pytest
+from integration._support.client import Gateway, gateway_from_environment
+from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
+from integration._support.wire import Reply, Request, wire_server
+from pydantic import JsonValue, TypeAdapter
+
+pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
+
+_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
+_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": "Count this synthetic request."}]
+_PROVIDER_KEY_VARIABLES: Final = (
+ "ANTHROPIC_API_KEY",
+ "ANTHROPIC_AUTH_TOKEN",
+ "AZURE_AI_API_KEY",
+ "AZURE_API_KEY",
+ "AZURE_OPENAI_API_KEY",
+ "AZURE_AD_TOKEN",
+ "AZURE_CLIENT_ID",
+ "AZURE_CLIENT_SECRET",
+ "AZURE_TENANT_ID",
+ "OPENAI_API_KEY",
+)
+
+
+@pytest.fixture(scope="module")
+def keyless_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedProxy]:
+ with (
+ gateway_from_environment() as gateway,
+ owned_proxy_process(
+ gateway,
+ tmp_path_factory.mktemp("count-tokens-error-shapes"),
+ {},
+ remove_environment=_PROVIDER_KEY_VARIABLES,
+ workers=1,
+ ) as owned,
+ ):
+ yield owned
+
+
+@pytest.fixture(scope="module")
+def azure_environment_key_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedProxy]:
+ with (
+ gateway_from_environment() as gateway,
+ owned_proxy_process(
+ gateway,
+ tmp_path_factory.mktemp("count-tokens-azure-environment-key"),
+ {"AZURE_API_KEY": "synthetic-azure-env-key"},
+ remove_environment=tuple(name for name in _PROVIDER_KEY_VARIABLES if name != "AZURE_API_KEY"),
+ workers=1,
+ ) as owned,
+ ):
+ yield owned
+
+
+def _anthropic(gateway: Gateway) -> anthropic.Anthropic:
+ return anthropic.Anthropic(
+ api_key=gateway.key, base_url=str(gateway.client.base_url), max_retries=0, http_client=gateway.client
+ )
+
+
+def _never_called(request: Request) -> Reply:
+ raise AssertionError(f"count reached the provider without a key: {request.target}")
+
+
+@pytest.mark.parametrize(
+ ("deployment", "missing"),
+ (
+ ({"model": "anthropic/claude-sonnet-4-6"}, "No Anthropic API key is set for model "),
+ ({"model": "azure_ai/claude-sonnet-4-6"}, "No Azure AI API key or Entra ID credential is set for model "),
+ ),
+ ids=("anthropic", "azure_ai_anthropic"),
+)
+def test_anthropic_count_without_a_provider_credential_is_an_anthropic_401(
+ keyless_proxy: OwnedProxy, deployment: dict[str, JsonValue], missing: str
+) -> None:
+ gateway: Final = keyless_proxy.gateway
+ with wire_server(_never_called) as wire, gateway.scenario() as scenario:
+ model: Final = scenario.model(**deployment, api_base=wire.url, api_key=None, num_retries=0)
+ with pytest.raises(anthropic.AuthenticationError) as raised:
+ _anthropic(gateway).messages.count_tokens(model=model, messages=_MESSAGES)
+ assert raised.value.body == {
+ "type": "error",
+ "error": {"type": "authentication_error", "message": raised.value.body["error"]["message"]},
+ }
+ assert missing in raised.value.message, raised.value.message
+ assert wire.drain() == ()
+
+
+def test_azure_ai_anthropic_count_sends_the_azure_environment_key_as_x_api_key(
+ azure_environment_key_proxy: OwnedProxy,
+) -> None:
+ gateway: Final = azure_environment_key_proxy.gateway
+
+ def count(request: Request) -> Reply:
+ assert request.headers.get("x-api-key") == "synthetic-azure-env-key", request.headers
+ return Reply(body=json.dumps({"input_tokens": 58}).encode())
+
+ with wire_server(count) as wire, gateway.scenario() as scenario:
+ model: Final = scenario.model(
+ model="azure_ai/claude-sonnet-4-6",
+ api_base=wire.url,
+ api_key=None,
+ num_retries=0,
+ )
+ result: Final = _anthropic(gateway).messages.count_tokens(model=model, messages=_MESSAGES)
+ assert result.input_tokens == 58
+
+
+def test_azure_ai_anthropic_count_authenticates_with_an_entra_token(keyless_proxy: OwnedProxy) -> None:
+ gateway: Final = keyless_proxy.gateway
+
+ def count(request: Request) -> Reply:
+ assert request.headers.get("authorization") == "Bearer synthetic-entra-token", request.headers
+ assert "x-api-key" not in request.headers, request.headers
+ return Reply(body=json.dumps({"input_tokens": 61}).encode())
+
+ with wire_server(count) as wire, gateway.scenario() as scenario:
+ model: Final = scenario.model(
+ model="azure_ai/claude-sonnet-4-6",
+ api_base=wire.url,
+ api_key=None,
+ azure_ad_token="synthetic-entra-token",
+ num_retries=0,
+ )
+ result: Final = _anthropic(gateway).messages.count_tokens(model=model, messages=_MESSAGES)
+ assert result.input_tokens == 61
+ assert len(wire.drain()) == 1
+
+
+def test_anthropic_count_request_errors_use_the_anthropic_envelope(keyless_proxy: OwnedProxy) -> None:
+ gateway: Final = keyless_proxy.gateway
+ upstream_error: Final = {"type": "error", "error": {"type": "invalid_request_error", "message": "scripted reject"}}
+ with (
+ wire_server(lambda _request: Reply(status=400, body=json.dumps(upstream_error).encode())) as wire,
+ gateway.scenario() as scenario,
+ ):
+ model: Final = scenario.model(
+ model="anthropic/claude-sonnet-4-6", api_base=wire.url, api_key="synthetic-key", num_retries=0
+ )
+ missing_messages: Final = gateway.request("POST", "/v1/messages/count_tokens", {"model": model})
+ assert missing_messages.status_code == 400, missing_messages.text
+ assert _JSON_OBJECT.validate_json(missing_messages.content) == {
+ "type": "error",
+ "error": {"type": "invalid_request_error", "message": "messages parameter is required"},
+ }
+
+ rejected: Final = gateway.request(
+ "POST", "/v1/messages/count_tokens", {"model": model, "messages": _MESSAGES}
+ )
+ assert rejected.status_code == 400, rejected.text
+ body: Final = _JSON_OBJECT.validate_json(rejected.content)
+ assert set(body) <= {"type", "error", "request_id"}, body
+ assert body["error"] == upstream_error["error"], body
+ assert len(wire.drain()) == 1
+
+
+def test_openai_input_tokens_without_a_provider_key_is_an_openai_401(keyless_proxy: OwnedProxy) -> None:
+ gateway: Final = keyless_proxy.gateway
+ with wire_server(_never_called) as wire, gateway.scenario() as scenario:
+ model: Final = scenario.model(
+ model="openai/gpt-4o-mini", api_base=f"{wire.url}/v1", api_key=None, num_retries=0
+ )
+ client: Final = openai.OpenAI(
+ api_key=gateway.key, base_url=str(gateway.client.base_url), max_retries=0, http_client=gateway.client
+ )
+ with pytest.raises(openai.AuthenticationError) as raised:
+ client.responses.input_tokens.count(model=model, input="Count this synthetic request.")
+ assert f"No OpenAI API key is set for model {model}" in raised.value.message
+ assert wire.drain() == ()
+
+
+def test_google_count_tokens_errors_use_the_google_envelope(keyless_proxy: OwnedProxy) -> None:
+ gateway: Final = keyless_proxy.gateway
+ upstream_error: Final = {
+ "error": {"code": 400, "message": "API key not valid. Please pass a valid API key.", "status": "INVALID_ARGUMENT"}
+ }
+ with (
+ wire_server(lambda _request: Reply(status=400, body=json.dumps(upstream_error).encode())) as wire,
+ gateway.scenario() as scenario,
+ ):
+ model: Final = scenario.model(
+ model="gemini/gemini-2.5-flash", api_base=wire.url, api_key="synthetic-key", num_retries=0
+ )
+ rejected: Final = gateway.request(
+ "POST",
+ f"/v1beta/models/{model}:countTokens",
+ {"contents": [{"role": "user", "parts": [{"text": "Count this synthetic request."}]}]},
+ )
+ assert rejected.status_code == 400, rejected.text
+ error: Final = _JSON_OBJECT.validate_json(rejected.content)["error"]
+ assert isinstance(error, dict) and error["code"] == 400 and error["status"] == "INVALID_ARGUMENT", error
+ assert "API key not valid" in str(error["message"]), error
+ assert len(wire.drain()) == 1
+
+ malformed: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", {"contents": "not-a-list"})
+ assert malformed.status_code == 400, malformed.text
+ assert _JSON_OBJECT.validate_json(malformed.content) == {
+ "error": {"code": 400, "message": "Invalid Google countTokens request body", "status": "INVALID_ARGUMENT"}
+ }
diff --git a/tests/integration/providers/test_custom_prompt_template_chaos.py b/tests/integration/providers/test_custom_prompt_template_chaos.py
new file mode 100644
index 00000000000..c0da612af23
--- /dev/null
+++ b/tests/integration/providers/test_custom_prompt_template_chaos.py
@@ -0,0 +1,542 @@
+import asyncio
+import json
+import re
+import threading
+import uuid
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from pathlib import Path
+from queue import SimpleQueue
+from types import MappingProxyType
+from typing import Final, Literal
+from urllib.parse import urlsplit
+
+import httpx
+import psutil
+import pytest
+import yaml
+from integration._support.client import Gateway, eventually
+from integration._support.database import read_rows
+from integration._support.process import graceful_stop_seconds, owned_proxy_process
+from integration._support.wire import Reply, Request, Wire, wire_server
+from pydantic import JsonValue, TypeAdapter
+
+_BACKEND: Final = "llama3-custom-template-chaos"
+_API_KEY: Final = "synthetic-ollama-key"
+_CONFIG_MODEL: Final = "ollama-custom-template-chaos"
+_SYSTEM: Final = "Answer in one word."
+_ROLES: Final[dict[str, JsonValue]] = {
+ "system": {"pre_message": "", "post_message": ""},
+ "user": {"pre_message": "", "post_message": ""},
+}
+_REJECTION: Final = "initial_prompt_value: Input should be a valid string"
+_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
+_MARKER: Final = re.compile(r"marker-([0-9a-f]{32})")
+_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]")
+_OWNED_PROXY_CELL_SECONDS: Final = 2 * graceful_stop_seconds() + 120
+
+Endpoint = Literal["chat", "messages", "responses"]
+
+
+@dataclass(frozen=True, slots=True)
+class _Call:
+ endpoint: Endpoint
+ stream: bool
+ marker: str
+ rejected: bool = False
+
+
+@dataclass(frozen=True, slots=True)
+class _Served:
+ call: _Call
+ status: int
+ text: str
+
+
+def _question(marker: str) -> str:
+ return f"What colour is the sky? marker-{marker}"
+
+
+def _prompt(marker: str) -> str:
+ return f"{_SYSTEM}{_question(marker)}"
+
+
+def _path(endpoint: Endpoint) -> str:
+ match endpoint:
+ case "chat":
+ return "/v1/chat/completions"
+ case "messages":
+ return "/v1/messages"
+ case "responses":
+ return "/v1/responses"
+
+
+def _body(model: str, call: _Call) -> dict[str, JsonValue]:
+ common: Final[dict[str, JsonValue]] = {
+ "model": model,
+ "stream": call.stream,
+ "num_retries": 0,
+ "cache": {"no-cache": True},
+ **({"initial_prompt_value": 5} if call.rejected else {}),
+ }
+ match call.endpoint:
+ case "chat":
+ return {
+ **common,
+ "messages": [
+ {"role": "system", "content": _SYSTEM},
+ {"role": "user", "content": _question(call.marker)},
+ ],
+ }
+ case "messages":
+ return {
+ **common,
+ "max_tokens": 64,
+ "system": _SYSTEM,
+ "messages": [{"role": "user", "content": _question(call.marker)}],
+ }
+ case "responses":
+ return {**common, "store": False, "instructions": _SYSTEM, "input": _question(call.marker)}
+
+
+def _generate_reply(marker: str, stream: bool, drop_connection: bool = False) -> Reply:
+ done: Final = {
+ "model": _BACKEND,
+ "created_at": "2026-10-15T00:00:00Z",
+ "response": "",
+ "done": True,
+ "done_reason": "stop",
+ "prompt_eval_count": 30,
+ "eval_count": 12,
+ }
+ if not stream:
+ return Reply(
+ body=json.dumps({**done, "response": f"blue marker-{marker}"}).encode(), drop_connection=drop_connection
+ )
+ frames: Final = (
+ *(
+ {"model": _BACKEND, "created_at": "2026-10-15T00:00:00Z", "response": piece, "done": False}
+ for piece in ("blue ", f"marker-{marker}")
+ ),
+ done,
+ )
+ return Reply(
+ content_type="application/x-ndjson",
+ chunks=tuple(json.dumps(frame).encode() + b"\n" for frame in frames),
+ drop_connection=drop_connection,
+ )
+
+
+def _is_generate(request: Request) -> bool:
+ return (request.method, request.target) == ("POST", "/api/generate")
+
+
+def _marker_of(request: Request) -> str:
+ found: Final = _MARKER.search(request.body.decode())
+ assert found is not None, request.body
+ return found.group(1)
+
+
+def _echo(request: Request) -> Reply:
+ if not _is_generate(request):
+ return Reply(body=b"{}")
+ stream: Final = _JSON_OBJECT.validate_json(request.body).get("stream") is True
+ return _generate_reply(_marker_of(request), stream)
+
+
+def _assert_each_prompt_rendered_once(received: tuple[Request, ...], markers: frozenset[str]) -> None:
+ generates: Final = tuple(request for request in received if _is_generate(request))
+ assert sorted(_marker_of(request) for request in generates) == sorted(markers)
+ for request in generates:
+ body: Final = _JSON_OBJECT.validate_json(request.body)
+ assert body["prompt"] == _prompt(_marker_of(request)), body["prompt"]
+ assert "roles" not in body and "initial_prompt_value" not in body, sorted(body)
+
+
+def _response_id(served: _Served) -> str | None:
+ if served.call.endpoint == "responses":
+ return None
+ if not served.call.stream:
+ identity: Final = _JSON_OBJECT.validate_json(served.text)["id"]
+ assert isinstance(identity, str)
+ return identity
+ for line in served.text.splitlines():
+ if not line.startswith("data: ") or line == "data: [DONE]":
+ continue
+ payload: Final = _JSON_OBJECT.validate_json(line.removeprefix("data: "))
+ if served.call.endpoint == "chat":
+ first: Final = payload["id"]
+ assert isinstance(first, str)
+ return first
+ if payload.get("type") == "message_start":
+ message: Final = payload["message"]
+ assert isinstance(message, dict) and isinstance(message["id"], str)
+ return message["id"]
+ raise AssertionError(served.text)
+
+
+def _spend_statuses(model: str, expected: int) -> MappingProxyType[str, str]:
+ rows: Final = eventually(
+ lambda: read_rows('SELECT request_id, status FROM "LiteLLM_SpendLogs" WHERE model_group=%s', (model,)),
+ lambda found: len(found) >= expected,
+ seconds=70,
+ )
+ statuses: Final = MappingProxyType({str(row["request_id"]): str(row["status"]) for row in rows})
+ assert len(statuses) == len(rows) == expected, rows
+ return statuses
+
+
+def _successes(statuses: Mapping[str, str]) -> frozenset[str]:
+ return frozenset(identity for identity, status in statuses.items() if status == "success")
+
+
+async def _send(client: httpx.AsyncClient, key: str, model: str, call: _Call) -> _Served:
+ async with client.stream(
+ "POST",
+ _path(call.endpoint),
+ json=_body(model, call),
+ headers={"Authorization": f"Bearer {key}", "anthropic-version": "2023-06-01"},
+ ) as response:
+ raw: Final = await response.aread()
+ return _Served(call=call, status=response.status_code, text=raw.decode())
+
+
+async def _burst(
+ base_url: str, key: str, model: str, calls: tuple[_Call, ...], *, tolerate_transport_errors: bool = False
+) -> tuple[_Served, ...]:
+ async with httpx.AsyncClient(base_url=base_url, timeout=60, trust_env=False) as client:
+ results: Final = await asyncio.gather(
+ *(_send(client, key, model, call) for call in calls), return_exceptions=tolerate_transport_errors
+ )
+ for result in results:
+ assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result)
+ return tuple(result for result in results if isinstance(result, _Served))
+
+
+def _calls(count: int, endpoints: tuple[Endpoint, ...], stream: Callable[[int], bool]) -> tuple[_Call, ...]:
+ return tuple(
+ _Call(endpoint=endpoints[index % len(endpoints)], stream=stream(index), marker=uuid.uuid4().hex)
+ for index in range(count)
+ )
+
+
+def _assert_answered_with_its_own_marker(served: _Served) -> None:
+ assert served.status == 200, served.text
+ assert set(_MARKER.findall(served.text)) == {served.call.marker}, served.text
+
+
+def _assert_rejected_by_name(served: _Served) -> None:
+ assert served.status == 500, served.text
+ assert _REJECTION in served.text and "marker-" not in served.text, served.text
+
+
+async def test_concurrent_burst_renders_each_template_and_rejects_only_the_bad_requests(gateway: Gateway) -> None:
+ calls: Final = tuple(
+ _Call(endpoint=call.endpoint, stream=call.stream, marker=call.marker, rejected=index % 5 == 4)
+ for index, call in enumerate(_calls(30, ("chat", "messages", "responses"), lambda index: index % 2 == 0))
+ )
+ accepted: Final = tuple(call for call in calls if not call.rejected)
+ with wire_server(_echo) as wire, gateway.scenario() as scenario:
+ model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY, roles=_ROLES)
+ served: Final = await _burst(str(gateway.client.base_url), gateway.key, model, calls)
+ assert len(served) == 30
+ for item in served:
+ if item.call.rejected:
+ _assert_rejected_by_name(item)
+ else:
+ _assert_answered_with_its_own_marker(item)
+ _assert_each_prompt_rendered_once(wire.drain(), frozenset(call.marker for call in accepted))
+ answered: Final = tuple(item for item in served if not item.call.rejected)
+ known: Final = frozenset(identity for identity in map(_response_id, answered) if identity is not None)
+ assert len(known) == 16, known
+ statuses: Final = _spend_statuses(model, 30)
+ assert len(_successes(statuses)) == len(accepted), statuses
+ assert known <= _successes(statuses)
+
+
+async def test_dropped_ollama_connections_fail_their_callers_and_the_rest_keep_their_template(
+ gateway: Gateway,
+) -> None:
+ calls: Final = _calls(12, ("chat",), lambda index: index % 2 == 1)
+ dropped: Final = frozenset(call.marker for index, call in enumerate(calls) if index % 3 == 0)
+
+ def respond(request: Request) -> Reply:
+ if not _is_generate(request):
+ return Reply(body=b"{}")
+ marker: Final = _marker_of(request)
+ stream: Final = _JSON_OBJECT.validate_json(request.body).get("stream") is True
+ return _generate_reply(marker, stream, drop_connection=marker in dropped)
+
+ with wire_server(respond) as wire, gateway.scenario() as scenario:
+ model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY, roles=_ROLES)
+ served: Final = await _burst(str(gateway.client.base_url), gateway.key, model, calls)
+ assert len(served) == 12
+ for item in served:
+ if item.call.marker in dropped:
+ assert item.status == 500, item.text
+ assert "APIConnectionError" in item.text and "marker-" not in item.text, item.text
+ else:
+ _assert_answered_with_its_own_marker(item)
+ recovery: Final = _Call(endpoint="chat", stream=False, marker=uuid.uuid4().hex)
+ (recovered,) = await _burst(str(gateway.client.base_url), gateway.key, model, (recovery,))
+ _assert_answered_with_its_own_marker(recovered)
+ _assert_each_prompt_rendered_once(wire.drain(), frozenset(call.marker for call in (*calls, recovery)))
+ answered: Final = tuple(item for item in served if item.call.marker not in dropped)
+ survivors: Final = frozenset(
+ identity for identity in map(_response_id, (*answered, recovered)) if identity is not None
+ )
+ assert len(survivors) == 9, survivors
+ statuses: Final = _spend_statuses(model, 13)
+ assert _successes(statuses) == survivors, statuses
+ assert sum(status == "failure" for status in statuses.values()) == len(dropped), statuses
+
+
+def _owned_config(
+ wires: tuple[Wire, ...],
+ tmp_path: Path,
+ template: Mapping[str, JsonValue],
+ litellm_settings: Mapping[str, JsonValue],
+ router_settings: Mapping[str, JsonValue] = MappingProxyType({}),
+) -> Path:
+ base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
+ deployments: Final = [
+ {
+ "model_name": _CONFIG_MODEL,
+ "litellm_params": {
+ "model": f"ollama/{_BACKEND}",
+ "api_base": wire.url,
+ "api_key": _API_KEY,
+ "order": order,
+ **template,
+ },
+ }
+ for order, wire in enumerate(wires, start=1)
+ ]
+ config: Final = {
+ **base,
+ "model_list": deployments,
+ "litellm_settings": {**base.get("litellm_settings", {}), **litellm_settings},
+ "router_settings": {**base.get("router_settings", {}), **router_settings},
+ }
+ path: Final = tmp_path / "ollama-custom-template-chaos.yaml"
+ path.write_text(yaml.safe_dump(config))
+ return path
+
+
+def _open_upstream_connections(pid: int, upstream: str) -> int:
+ port: Final = urlsplit(upstream).port
+ return sum(
+ 1
+ for connection in psutil.Process(pid).net_connections(kind="tcp")
+ if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port
+ )
+
+
+@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS)
+async def test_worker_sigkill_mid_burst_leaves_the_sibling_rendering_templates(
+ gateway: Gateway, tmp_path: Path
+) -> None:
+ calls: Final = _calls(20, ("chat",), lambda _: False)
+ release: Final = threading.Event()
+ held_markers: Final[SimpleQueue[str]] = SimpleQueue()
+
+ def held(request: Request) -> Reply:
+ if not _is_generate(request):
+ return Reply(body=b"{}")
+ held_markers.put(_marker_of(request))
+ assert release.wait(timeout=60), "The burst was never released"
+ return _echo(request)
+
+ with wire_server(held) as wire:
+ path: Final = _owned_config((wire,), tmp_path, {"roles": _ROLES}, {})
+ with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
+ candidate: Final = owned.gateway
+ workers: Final = eventually(
+ lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())),
+ lambda pids: len(pids) == 2,
+ seconds=30,
+ )
+ burst: Final = asyncio.create_task(
+ _burst(
+ str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, calls, tolerate_transport_errors=True
+ )
+ )
+ await asyncio.to_thread(eventually, held_markers.qsize, lambda size: size == 20, 60)
+ held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers})
+ assert sum(held_by.values()) == 20, held_by
+ victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__)
+ with owned.replacing_worker(victim_pid) as kill:
+ kill.freeze()
+ await asyncio.to_thread(kill)
+ release.set()
+ served: Final = await burst
+ assert held_by[survivor_pid] >= 10, held_by
+ assert len(served) == held_by[survivor_pid], (held_by, len(served))
+ for item in served:
+ _assert_answered_with_its_own_marker(item)
+ follow_up: Final = _Call(endpoint="chat", stream=False, marker=uuid.uuid4().hex)
+ (answered,) = await _burst(str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, (follow_up,))
+ _assert_answered_with_its_own_marker(answered)
+ _assert_each_prompt_rendered_once(wire.drain(), frozenset(call.marker for call in (*calls, follow_up)))
+
+
+@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS)
+async def test_rejected_templates_never_cool_down_the_preferred_deployment(gateway: Gateway, tmp_path: Path) -> None:
+ rejected: Final = tuple(
+ _Call(endpoint="chat", stream=False, marker=uuid.uuid4().hex, rejected=True) for _ in range(4)
+ )
+ valid: Final = _calls(8, ("chat",), lambda _: False)
+ with wire_server(_echo) as preferred, wire_server(_echo) as standby:
+ path: Final = _owned_config(
+ (preferred, standby), tmp_path, {"roles": _ROLES}, {}, {"disable_cooldowns": False, "allowed_fails": 0}
+ )
+ with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
+ base_url: Final = str(owned.gateway.client.base_url)
+ for served in await _burst(base_url, owned.gateway.key, _CONFIG_MODEL, rejected):
+ _assert_rejected_by_name(served)
+ for served in await _burst(base_url, owned.gateway.key, _CONFIG_MODEL, valid):
+ _assert_answered_with_its_own_marker(served)
+ _assert_each_prompt_rendered_once(preferred.drain(), frozenset(call.marker for call in valid))
+ assert not any(_is_generate(request) for request in standby.drain())
+
+
+def _cached_chat(candidate: Gateway, marker: str, field: str, value: str) -> httpx.Response:
+ return candidate.request(
+ "POST",
+ "/v1/chat/completions",
+ {
+ "model": _CONFIG_MODEL,
+ "messages": [{"role": "system", "content": _SYSTEM}, {"role": "user", "content": _question(marker)}],
+ field: value,
+ },
+ )
+
+
+def _served_from_cache(candidate: Gateway, field: str) -> bool:
+ marker: Final = uuid.uuid4().hex
+ first: Final = _cached_chat(candidate, marker, field, "")
+ assert first.status_code == 200, first.text
+ assert "x-litellm-cache-key" not in first.headers, sorted(first.headers)
+ eventually(
+ lambda: _cached_chat(candidate, marker, field, ""),
+ lambda found: "x-litellm-cache-key" in found.headers,
+ seconds=20,
+ )
+ changed: Final = _cached_chat(candidate, marker, field, "")
+ assert changed.status_code == 200, changed.text
+ return "x-litellm-cache-key" in changed.headers
+
+
+@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS)
+def test_initial_and_final_prompt_values_key_the_provider_specific_cache_alike(
+ gateway: Gateway, tmp_path: Path
+) -> None:
+ template: Final[dict[str, JsonValue]] = {
+ "roles": _ROLES,
+ "initial_prompt_value": "",
+ "final_prompt_value": "",
+ }
+ with wire_server(_echo) as wire:
+ path: Final = _owned_config(
+ (wire,), tmp_path, template, {"enable_caching_on_provider_specific_optional_params": True}
+ )
+ with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
+ initial: Final = _served_from_cache(owned.gateway, "initial_prompt_value")
+ final: Final = _served_from_cache(owned.gateway, "final_prompt_value")
+ assert initial == final, {"initial_prompt_value": initial, "final_prompt_value": final}
+
+
+def _gemini_reply(_: Request) -> Reply:
+ candidate: Final[dict[str, JsonValue]] = {
+ "content": {"role": "model", "parts": [{"text": '{"colour": "blue"}'}]},
+ "finishReason": "STOP",
+ }
+ usage: Final[dict[str, JsonValue]] = {"promptTokenCount": 30, "candidatesTokenCount": 12, "totalTokenCount": 42}
+ return Reply(body=json.dumps({"candidates": [candidate], "usageMetadata": usage}).encode())
+
+
+def _texts(contents: JsonValue) -> tuple[str, ...]:
+ assert isinstance(contents, list), contents
+ parts: Final = tuple(part for content in contents if isinstance(content, dict) for part in _parts(content))
+ return tuple(text for part in parts if isinstance(text := part.get("text"), str))
+
+
+def _parts(content: dict[str, JsonValue]) -> tuple[dict[str, JsonValue], ...]:
+ parts: Final = content.get("parts")
+ assert isinstance(parts, list), content
+ return tuple(part for part in parts if isinstance(part, dict))
+
+
+def _anthropic_text_reply(_: Request) -> Reply:
+ completion: Final[dict[str, JsonValue]] = {
+ "type": "completion",
+ "id": "compl_scripted",
+ "completion": "blue",
+ "stop_reason": "stop_sequence",
+ }
+ return Reply(body=json.dumps(completion).encode())
+
+
+def _boot_deployment(name: str, model: str, wire: Wire) -> dict[str, JsonValue]:
+ return {"model_name": name, "litellm_params": {"model": model, "api_base": wire.url, "api_key": _API_KEY}}
+
+
+def _boot_chat(candidate: Gateway, model: str, extra: Mapping[str, JsonValue]) -> None:
+ response: Final = candidate.request(
+ "POST",
+ "/v1/chat/completions",
+ {
+ "model": model,
+ "messages": [{"role": "system", "content": _SYSTEM}, {"role": "user", "content": _question(model)}],
+ "cache": {"no-cache": True},
+ **extra,
+ },
+ )
+ assert response.status_code == 200, response.text
+
+
+def _anthropic_text_prompts(wire: Wire) -> MappingProxyType[str, JsonValue]:
+ bodies: Final = tuple(_JSON_OBJECT.validate_json(request.body) for request in wire.drain())
+ return MappingProxyType({str(body["model"]): body["prompt"] for body in bodies})
+
+
+@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS)
+def test_boot_time_custom_prompt_dict_renders_partial_templates(gateway: Gateway, tmp_path: Path) -> None:
+ templates: Final[dict[str, JsonValue]] = {
+ "response_schema_prompt": {"roles": {"user": {"pre_message": "", "post_message": ""}}},
+ "claude-2.1": {"roles": _ROLES},
+ "claude-instant-1.2": {"initial_prompt_value": "", "final_prompt_value": ""},
+ }
+ schema: Final[dict[str, JsonValue]] = {
+ "type": "object",
+ "properties": {"colour": {"type": "string"}},
+ "required": ["colour"],
+ }
+ with wire_server(_gemini_reply) as gemini, wire_server(_anthropic_text_reply) as anthropic_text:
+ base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
+ deployments: Final = [
+ _boot_deployment("gemini-schema", "gemini/scripted-template-model", gemini),
+ _boot_deployment("anthropic-text-roles", "anthropic_text/claude-2.1", anthropic_text),
+ _boot_deployment("anthropic-text-values", "anthropic_text/claude-instant-1.2", anthropic_text),
+ ]
+ settings: Final = {**base.get("litellm_settings", {}), "custom_prompt_dict": templates}
+ path: Final = tmp_path / "boot-time-custom-prompt-dict.yaml"
+ path.write_text(yaml.safe_dump({**base, "model_list": deployments, "litellm_settings": settings}))
+ with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
+ _boot_chat(
+ owned.gateway,
+ "gemini-schema",
+ {"response_format": {"type": "json_schema", "json_schema": {"name": "sky", "schema": schema}}},
+ )
+ posts: Final = tuple(request for request in gemini.drain() if request.method == "POST")
+ assert len(posts) == 1, [request.target for request in posts]
+ texts: Final = _texts(_JSON_OBJECT.validate_json(posts[0].body)["contents"])
+ wrapped: Final = tuple(text for text in texts if text.startswith(""))
+ assert len(wrapped) == 1, texts
+ assert wrapped[0].endswith("") and "colour" in wrapped[0], wrapped[0]
+ _boot_chat(owned.gateway, "anthropic-text-roles", {})
+ _boot_chat(owned.gateway, "anthropic-text-values", {})
+ assert _anthropic_text_prompts(anthropic_text) == {
+ "claude-2.1": f"{_SYSTEM}{_question('anthropic-text-roles')}",
+ "claude-instant-1.2": f"{_SYSTEM}{_question('anthropic-text-values')}",
+ }
diff --git a/tests/integration/providers/test_custom_prompt_template_wire.py b/tests/integration/providers/test_custom_prompt_template_wire.py
new file mode 100644
index 00000000000..49a47a8c5e5
--- /dev/null
+++ b/tests/integration/providers/test_custom_prompt_template_wire.py
@@ -0,0 +1,701 @@
+import json
+import uuid
+from collections.abc import Callable, Iterator, Mapping
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import Final
+
+import anthropic
+import openai
+import pytest
+from integration._support.client import Gateway, Scenario, eventually, object_value, string_value
+from integration._support.database import read_rows
+from integration._support.wire import Reply, Request, Wire, wire_server
+from pydantic import JsonValue, TypeAdapter
+
+_BACKEND: Final = "llama3-custom-template"
+_API_KEY: Final = "synthetic-ollama-key"
+_SYSTEM: Final = "Answer in one word."
+_QUESTION: Final = "What colour is the sky?"
+_ANSWER: Final = "blue"
+_INITIAL: Final = ""
+_FINAL: Final = ""
+_ROLES: Final[dict[str, JsonValue]] = {
+ "system": {"pre_message": "", "post_message": ""},
+ "user": {"pre_message": "", "post_message": ""},
+ "assistant": {"pre_message": "", "post_message": ""},
+}
+_ROLES_ONLY: Final[dict[str, JsonValue]] = {"roles": _ROLES}
+_VALUES_ONLY: Final[dict[str, JsonValue]] = {"initial_prompt_value": _INITIAL, "final_prompt_value": _FINAL}
+_FULL: Final[dict[str, JsonValue]] = {**_ROLES_ONLY, **_VALUES_ONLY}
+_ROLES_PROMPT: Final = f"{_SYSTEM}{_QUESTION}"
+_VALUES_PROMPT: Final = f"{_INITIAL}{_SYSTEM}{_QUESTION}{_FINAL}"
+_FULL_PROMPT: Final = f"{_INITIAL}{_ROLES_PROMPT}{_FINAL}"
+_TEMPLATES: Final[Mapping[str, tuple[dict[str, JsonValue], str]]] = {
+ "roles_only": (_ROLES_ONLY, _ROLES_PROMPT),
+ "values_only": (_VALUES_ONLY, _VALUES_PROMPT),
+}
+_TEMPLATE_FIELDS: Final = frozenset({"roles", "initial_prompt_value", "final_prompt_value", "bos_token", "eos_token"})
+_AWS: Final[dict[str, JsonValue]] = {
+ "aws_access_key_id": "AKIASYNTHETICKEY0001",
+ "aws_secret_access_key": "synthetic-secret-key-for-testing",
+ "aws_region_name": "us-east-1",
+}
+_NO_CACHE: Final[dict[str, JsonValue]] = {"no-cache": True}
+_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
+_JSON_LIST: Final = TypeAdapter(list[JsonValue])
+
+
+def _generate_reply(text: str) -> Reply:
+ return Reply(
+ body=json.dumps(
+ {
+ "model": _BACKEND,
+ "created_at": "2026-10-15T00:00:00Z",
+ "response": text,
+ "done": True,
+ "done_reason": "stop",
+ "prompt_eval_count": 30,
+ "eval_count": 12,
+ }
+ ).encode()
+ )
+
+
+def _streamed_reply(text: str) -> Reply:
+ frames: Final = (
+ {"model": _BACKEND, "created_at": "2026-10-15T00:00:00Z", "response": text, "done": False},
+ {
+ "model": _BACKEND,
+ "created_at": "2026-10-15T00:00:00Z",
+ "response": "",
+ "done": True,
+ "done_reason": "stop",
+ "prompt_eval_count": 30,
+ "eval_count": 12,
+ },
+ )
+ return Reply(
+ content_type="application/x-ndjson", chunks=tuple(json.dumps(frame).encode() + b"\n" for frame in frames)
+ )
+
+
+def _is_generate(request: Request) -> bool:
+ return (request.method, request.target) == ("POST", "/api/generate")
+
+
+def _echo_answer(request: Request) -> Reply:
+ stream: Final = _JSON_OBJECT.validate_json(request.body).get("stream") is True
+ return _streamed_reply(_ANSWER) if stream else _generate_reply(_ANSWER)
+
+
+@contextmanager
+def _ollama_server(respond: Callable[[Request], Reply] = _echo_answer) -> Iterator[Wire]:
+ with wire_server(lambda request: respond(request) if _is_generate(request) else Reply(body=b"{}")) as wire:
+ yield wire
+
+
+def _generate_bodies(wire: Wire) -> tuple[dict[str, JsonValue], ...]:
+ return tuple(_JSON_OBJECT.validate_json(request.body) for request in wire.drain() if _is_generate(request))
+
+
+def _assert_rendered(body: dict[str, JsonValue], prompt: str, *, stream: bool) -> None:
+ assert body["model"] == _BACKEND
+ assert body["prompt"] == prompt, body["prompt"]
+ assert body["stream"] is stream
+ options: Final = body.get("options", {})
+ assert isinstance(options, dict)
+ assert not _TEMPLATE_FIELDS & (set(body) | set(options)), sorted(body)
+
+
+def _assert_only_rendered(wire: Wire, prompt: str, *, stream: bool = False) -> None:
+ bodies: Final = _generate_bodies(wire)
+ assert len(bodies) == 1, bodies
+ _assert_rendered(bodies[0], prompt, stream=stream)
+
+
+def _ollama_model(scenario: Scenario, wire: Wire, template: Mapping[str, JsonValue]) -> str:
+ return scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY, **template)
+
+
+def _chat_messages() -> list[dict[str, JsonValue]]:
+ return [{"role": "system", "content": _SYSTEM}, {"role": "user", "content": _QUESTION}]
+
+
+def _spend_row(identity: str) -> dict[str, JsonValue]:
+ rows: Final = eventually(
+ lambda: read_rows(
+ 'SELECT model_group, status, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
+ (identity,),
+ ),
+ lambda found: len(found) == 1,
+ seconds=70,
+ )
+ return rows[0]
+
+
+def _model_spend_rows(model: str, expected: int) -> list[dict[str, JsonValue]]:
+ rows: Final = eventually(
+ lambda: read_rows(
+ 'SELECT request_id, status, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE model_group=%s',
+ (model,),
+ ),
+ lambda found: len(found) >= expected,
+ seconds=70,
+ )
+ assert len({row["request_id"] for row in rows}) == len(rows) == expected, rows
+ return rows
+
+
+def _billed(model: str) -> dict[str, JsonValue]:
+ return {"model_group": model, "status": "success", "prompt_tokens": 30, "completion_tokens": 12}
+
+
+def _openai_client(gateway: Gateway) -> openai.OpenAI:
+ return openai.OpenAI(base_url=str(gateway.client.base_url) + "/v1", api_key=gateway.key, max_retries=0)
+
+
+def _async_openai_client(gateway: Gateway) -> openai.AsyncOpenAI:
+ return openai.AsyncOpenAI(base_url=str(gateway.client.base_url) + "/v1", api_key=gateway.key, max_retries=0)
+
+
+def _anthropic_client(gateway: Gateway) -> anthropic.Anthropic:
+ return anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0)
+
+
+def _async_anthropic_client(gateway: Gateway) -> anthropic.AsyncAnthropic:
+ return anthropic.AsyncAnthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0)
+
+
+def _post(gateway: Gateway, path: str, body: Mapping[str, JsonValue]) -> tuple[int, str]:
+ response: Final = gateway.request("POST", path, {**body, "cache": _NO_CACHE})
+ return response.status_code, response.text
+
+
+def _chat(gateway: Gateway, model: str, **extra: JsonValue) -> tuple[int, str]:
+ return _post(gateway, "/v1/chat/completions", {"model": model, "messages": _chat_messages(), **extra})
+
+
+def _chat_content(text: str) -> str:
+ choices: Final = _JSON_LIST.validate_python(_JSON_OBJECT.validate_json(text)["choices"])
+ message: Final = object_value(object_value(choices[0])["message"])
+ return string_value(message["content"])
+
+
+def test_openai_sdk_chat_renders_a_roles_only_template(gateway: Gateway) -> None:
+ with _ollama_server() as wire, gateway.scenario() as scenario:
+ model: Final = _ollama_model(scenario, wire, _ROLES_ONLY)
+ completion: Final = _openai_client(gateway).chat.completions.create(
+ model=model,
+ messages=_chat_messages(), # pyright: ignore[reportArgumentType] # plain JSON messages
+ extra_body={"cache": _NO_CACHE},
+ )
+ assert completion.choices[0].message.content == _ANSWER
+ assert completion.choices[0].finish_reason == "stop"
+ assert completion.usage is not None
+ assert (completion.usage.prompt_tokens, completion.usage.completion_tokens) == (30, 12)
+ _assert_only_rendered(wire, _ROLES_PROMPT)
+ assert _spend_row(completion.id) == _billed(model)
+
+
+async def test_async_openai_sdk_chat_stream_renders_a_roles_only_template(gateway: Gateway) -> None:
+ with _ollama_server() as wire, gateway.scenario() as scenario:
+ model: Final = _ollama_model(scenario, wire, _ROLES_ONLY)
+ stream: Final = await _async_openai_client(gateway).chat.completions.create(
+ model=model,
+ messages=_chat_messages(), # pyright: ignore[reportArgumentType] # plain JSON messages
+ stream=True,
+ stream_options={"include_usage": True},
+ extra_body={"cache": _NO_CACHE},
+ )
+ chunks: Final = [chunk async for chunk in stream]
+ assert {chunk.id for chunk in chunks} == {chunks[0].id}
+ assert "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) == _ANSWER
+ usages: Final = [chunk.usage for chunk in chunks if chunk.usage is not None]
+ assert [(usage.prompt_tokens, usage.completion_tokens) for usage in usages] == [(30, 12)]
+ _assert_only_rendered(wire, _ROLES_PROMPT, stream=True)
+ assert _spend_row(chunks[0].id) == _billed(model)
+
+
+def test_openai_sdk_completions_render_a_roles_only_template(gateway: Gateway) -> None:
+ with _ollama_server() as wire, gateway.scenario() as scenario:
+ model: Final = _ollama_model(scenario, wire, _ROLES_ONLY)
+ completion: Final = _openai_client(gateway).completions.create(
+ model=model, prompt=_QUESTION, extra_body={"cache": _NO_CACHE}
+ )
+ assert completion.choices[0].text == _ANSWER
+ _assert_only_rendered(wire, f"{_QUESTION}")
+ assert _spend_row(completion.id) == _billed(model)
+
+
+def test_anthropic_sdk_messages_render_a_roles_only_template(gateway: Gateway) -> None:
+ with _ollama_server() as wire, gateway.scenario() as scenario:
+ model: Final = _ollama_model(scenario, wire, _ROLES_ONLY)
+ message: Final = _anthropic_client(gateway).messages.create(
+ model=model,
+ max_tokens=64,
+ system=_SYSTEM,
+ messages=[{"role": "user", "content": _QUESTION}],
+ extra_body={"cache": _NO_CACHE},
+ )
+ assert message.stop_reason == "end_turn"
+ assert [(block.type, getattr(block, "text", None)) for block in message.content] == [("text", _ANSWER)]
+ assert (message.usage.input_tokens, message.usage.output_tokens) == (30, 12)
+ _assert_only_rendered(wire, _ROLES_PROMPT)
+ assert _spend_row(message.id) == _billed(model)
+
+
+async def test_async_anthropic_sdk_messages_stream_render_a_roles_only_template(gateway: Gateway) -> None:
+ with _ollama_server() as wire, gateway.scenario() as scenario:
+ model: Final = _ollama_model(scenario, wire, _ROLES_ONLY)
+ async with _async_anthropic_client(gateway).messages.stream(
+ model=model,
+ max_tokens=64,
+ system=_SYSTEM,
+ messages=[{"role": "user", "content": _QUESTION}],
+ extra_body={"cache": _NO_CACHE},
+ ) as stream:
+ text: Final = "".join([piece async for piece in stream.text_stream])
+ final: Final = await stream.get_final_message()
+ assert text == _ANSWER
+ assert final.stop_reason == "end_turn"
+ _assert_only_rendered(wire, _ROLES_PROMPT, stream=True)
+ assert _spend_row(final.id) == _billed(model)
+
+
+def test_openai_sdk_responses_render_a_roles_only_template(gateway: Gateway) -> None:
+ with _ollama_server() as wire, gateway.scenario() as scenario:
+ model: Final = _ollama_model(scenario, wire, _ROLES_ONLY)
+ response: Final = _openai_client(gateway).responses.create(
+ model=model, instructions=_SYSTEM, input=_QUESTION, store=False, extra_body={"cache": _NO_CACHE}
+ )
+ assert response.output_text == _ANSWER
+ assert response.usage is not None and (response.usage.input_tokens, response.usage.output_tokens) == (30, 12)
+ _assert_only_rendered(wire, _ROLES_PROMPT)
+ rows: Final = _model_spend_rows(model, 1)
+ assert (rows[0]["status"], rows[0]["prompt_tokens"], rows[0]["completion_tokens"]) == ("success", 30, 12)
+
+
+async def test_async_openai_sdk_responses_stream_render_a_roles_only_template(gateway: Gateway) -> None:
+ with _ollama_server() as wire, gateway.scenario() as scenario:
+ model: Final = _ollama_model(scenario, wire, _ROLES_ONLY)
+ stream: Final = await _async_openai_client(gateway).responses.create(
+ model=model,
+ instructions=_SYSTEM,
+ input=_QUESTION,
+ store=False,
+ stream=True,
+ extra_body={"cache": _NO_CACHE},
+ )
+ events: Final = [event async for event in stream]
+ assert "".join(event.delta for event in events if event.type == "response.output_text.delta") == _ANSWER
+ completed: Final = [event for event in events if event.type == "response.completed"]
+ assert len(completed) == 1 and completed[0].response.output_text == _ANSWER
+ _assert_only_rendered(wire, _ROLES_PROMPT, stream=True)
+ assert _model_spend_rows(model, 1)[0]["status"] == "success"
+
+
+def test_prompt_values_only_template_wraps_the_bare_messages(gateway: Gateway) -> None:
+ with _ollama_server() as wire, gateway.scenario() as scenario:
+ model: Final = _ollama_model(scenario, wire, _VALUES_ONLY)
+ code, text = _chat(gateway, model)
+ assert code == 200, text
+ assert _chat_content(text) == _ANSWER
+ _assert_only_rendered(wire, _VALUES_PROMPT)
+
+
+def test_full_template_renders_every_part(gateway: Gateway) -> None:
+ with _ollama_server() as wire, gateway.scenario() as scenario:
+ model: Final = _ollama_model(scenario, wire, _FULL)
+ code, text = _chat(gateway, model)
+ assert code == 200, text
+ assert _chat_content(text) == _ANSWER
+ _assert_only_rendered(wire, _FULL_PROMPT)
+
+
+@dataclass(frozen=True, slots=True)
+class _Member:
+ deployment: Callable[[str], dict[str, JsonValue]]
+ reply: Callable[[str], Reply]
+ prompt_field: str
+ target: Callable[[str], bool]
+
+
+def _codestral_reply(text: str) -> Reply:
+ return Reply(
+ body=json.dumps(
+ {
+ "id": "cmpl-scripted",
+ "object": "chat.completion",
+ "model": "codestral-latest",
+ "created": 1,
+ "choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}],
+ "usage": {"prompt_tokens": 30, "completion_tokens": 12, "total_tokens": 42},
+ }
+ ).encode()
+ )
+
+
+def _predibase_reply(text: str) -> Reply:
+ return Reply(body=json.dumps({"generated_text": text, "details": {"finish_reason": "eos_token"}}).encode())
+
+
+def _sagemaker_reply(text: str) -> Reply:
+ return Reply(body=json.dumps([{"generated_text": text}]).encode())
+
+
+def _bedrock_meta_reply(text: str) -> Reply:
+ return Reply(
+ body=json.dumps(
+ {"generation": text, "prompt_token_count": 30, "generation_token_count": 12, "stop_reason": "stop"}
+ ).encode()
+ )
+
+
+_FAMILY: Final[Mapping[str, _Member]] = {
+ "codestral": _Member(
+ deployment=lambda url: {
+ "model": "text-completion-codestral/codestral-latest",
+ "api_base": f"{url}/v1/fim/completions",
+ "api_key": "scripted-key",
+ },
+ reply=_codestral_reply,
+ prompt_field="prompt",
+ target=lambda target: target == "/v1/fim/completions",
+ ),
+ "predibase": _Member(
+ deployment=lambda url: {
+ "model": "predibase/llama-3-8b-instruct",
+ "api_base": url,
+ "api_key": "scripted-key",
+ "predibase_tenant_id": "scripted-tenant",
+ },
+ reply=_predibase_reply,
+ prompt_field="inputs",
+ target=lambda target: target.startswith("/scripted-tenant/deployments/v2/llms/llama-3-8b-instruct/generate"),
+ ),
+ "sagemaker": _Member(
+ deployment=lambda url: {
+ "model": "sagemaker/scripted-endpoint",
+ "sagemaker_base_url": f"{url}/endpoints/scripted-endpoint/invocations",
+ **_AWS,
+ },
+ reply=_sagemaker_reply,
+ prompt_field="inputs",
+ target=lambda target: target == "/endpoints/scripted-endpoint/invocations",
+ ),
+ "bedrock_invoke": _Member(
+ deployment=lambda url: {"model": "bedrock/invoke/meta.llama3-8b-instruct-v1:0", "api_base": url, **_AWS},
+ reply=_bedrock_meta_reply,
+ prompt_field="prompt",
+ target=lambda target: target.endswith("/invoke"),
+ ),
+}
+
+
+@pytest.mark.parametrize("template", sorted(_TEMPLATES))
+@pytest.mark.parametrize("member", sorted(_FAMILY))
+def test_family_member_renders_a_partial_template(gateway: Gateway, member: str, template: str) -> None:
+ spec: Final = _FAMILY[member]
+ fields, prompt = _TEMPLATES[template]
+ with wire_server(lambda _: spec.reply(_ANSWER)) as wire, gateway.scenario() as scenario:
+ model: Final = scenario.model(**spec.deployment(wire.url), **fields)
+ code, text = _chat(gateway, model)
+ received: Final = tuple(request for request in wire.drain() if request.method == "POST")
+ assert code == 200, text
+ assert _chat_content(text) == _ANSWER
+ assert len(received) == 1, [request.target for request in received]
+ assert spec.target(received[0].target), received[0].target
+ body: Final = _JSON_OBJECT.validate_json(received[0].body)
+ assert body[spec.prompt_field] == prompt, body
+ assert not _TEMPLATE_FIELDS & set(body), sorted(body)
+
+
+@pytest.mark.parametrize("provider", ("vllm", "petals"))
+def test_in_process_family_member_fails_alone_and_the_rig_keeps_serving(gateway: Gateway, provider: str) -> None:
+ with _ollama_server() as wire, gateway.scenario() as scenario:
+ broken: Final = scenario.model(model=f"{provider}/scripted-model", **_ROLES_ONLY)
+ healthy: Final = _ollama_model(scenario, wire, _ROLES_ONLY)
+ failed, failure = _chat(gateway, broken)
+ assert failed >= 400, failure
+ code, text = _chat(gateway, healthy)
+ assert code == 200, text
+ _assert_only_rendered(wire, _ROLES_PROMPT)
+
+
+def _watsonx_reply(text: str) -> Reply:
+ result: Final[dict[str, JsonValue]] = {
+ "generated_text": text,
+ "generated_token_count": 12,
+ "input_token_count": 30,
+ "stop_reason": "eos_token",
+ }
+ return Reply(body=json.dumps({"model_id": "ibm/granite", "results": [result]}).encode())
+
+
+def test_watsonx_text_deployment_with_a_template_still_answers(gateway: Gateway) -> None:
+ with wire_server(lambda _: _watsonx_reply(_ANSWER)) as wire, gateway.scenario() as scenario:
+ model: Final = scenario.model(
+ model="watsonx_text/ibm/granite-chat-scripted",
+ api_base=wire.url,
+ token="scripted-token",
+ project_id="p",
+ **_ROLES_ONLY,
+ )
+ code, text = _chat(gateway, model)
+ assert code == 200, text
+ assert _chat_content(text) == _ANSWER
+ posts: Final = tuple(request for request in wire.drain() if request.method == "POST")
+ assert len(posts) == 1, [request.target for request in posts]
+ body: Final = _JSON_OBJECT.validate_json(posts[0].body)
+ assert _QUESTION in string_value(body["input"]), body
+ assert not _TEMPLATE_FIELDS & set(body), sorted(body)
+
+
+def _openai_chat_reply(text: str) -> Reply:
+ return Reply(
+ body=json.dumps(
+ {
+ "id": "chatcmpl-scripted",
+ "object": "chat.completion",
+ "created": 1,
+ "model": "scripted-compat",
+ "choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}],
+ "usage": {"prompt_tokens": 30, "completion_tokens": 12, "total_tokens": 42},
+ }
+ ).encode()
+ )
+
+
+def _converse_reply(text: str) -> Reply:
+ return Reply(
+ body=json.dumps(
+ {
+ "output": {"message": {"role": "assistant", "content": [{"text": text}]}},
+ "stopReason": "end_turn",
+ "usage": {"inputTokens": 30, "outputTokens": 12, "totalTokens": 42},
+ "metrics": {"latencyMs": 1},
+ }
+ ).encode()
+ )
+
+
+_FORWARDING: Final[Mapping[str, tuple[Callable[[str], dict[str, JsonValue]], Callable[[str], Reply]]]] = {
+ "openai_compatible": (
+ lambda url: {"model": "openai/scripted-compat", "api_base": f"{url}/v1", "api_key": "scripted-key"},
+ _openai_chat_reply,
+ ),
+ "bedrock_converse": (
+ lambda url: {"model": "bedrock/converse/meta.llama3-8b-instruct-v1:0", "api_base": url, **_AWS},
+ _converse_reply,
+ ),
+}
+
+
+def _keys_anywhere(value: JsonValue) -> frozenset[str]:
+ if isinstance(value, dict):
+ nested: Final = frozenset(key for child in value.values() for key in _keys_anywhere(child))
+ return frozenset(value) | nested
+ if isinstance(value, list):
+ return frozenset(key for child in value for key in _keys_anywhere(child))
+ return frozenset()
+
+
+@pytest.mark.parametrize("source", ("deployment", "request"))
+@pytest.mark.parametrize("member", sorted(_FORWARDING))
+def test_template_fields_never_reach_a_provider_that_does_not_render_them(
+ gateway: Gateway, member: str, source: str
+) -> None:
+ deployment, reply = _FORWARDING[member]
+ with wire_server(lambda _: reply(_ANSWER)) as wire, gateway.scenario() as scenario:
+ model: Final = scenario.model(**deployment(wire.url), **(_FULL if source == "deployment" else {}))
+ code, text = _chat(gateway, model, **(_FULL if source == "request" else {}))
+ assert code == 200, text
+ assert _chat_content(text) == _ANSWER
+ posts: Final = tuple(request for request in wire.drain() if request.method == "POST")
+ assert len(posts) == 1, [request.target for request in posts]
+ body: Final = _JSON_OBJECT.validate_json(posts[0].body)
+ assert not _TEMPLATE_FIELDS & _keys_anywhere(body), body
+ assert _INITIAL not in posts[0].body.decode() and _FINAL not in posts[0].body.decode(), body
+
+
+@pytest.mark.parametrize("value", (5, [_INITIAL]), ids=("int", "list"))
+@pytest.mark.parametrize("source", ("deployment", "request"))
+def test_non_string_initial_prompt_value_is_rejected_by_name(gateway: Gateway, source: str, value: JsonValue) -> None:
+ with _ollama_server() as wire, gateway.scenario() as scenario:
+ template: Final[dict[str, JsonValue]] = {**_FULL, "initial_prompt_value": value}
+ broken: Final = _ollama_model(scenario, wire, template if source == "deployment" else _FULL)
+ healthy: Final = _ollama_model(scenario, wire, _FULL)
+ failed, failure = _chat(gateway, broken, **({"initial_prompt_value": value} if source == "request" else {}))
+ assert failed == 500, failure
+ assert "initial_prompt_value: Input should be a valid string" in failure, failure
+ assert _generate_bodies(wire) == ()
+ code, text = _chat(gateway, healthy)
+ assert code == 200, text
+ _assert_only_rendered(wire, _FULL_PROMPT)
+
+
+@pytest.mark.parametrize("value", ("", None), ids=("empty", "null"))
+def test_blank_initial_prompt_value_in_the_request_drops_the_deployment_value(
+ gateway: Gateway, value: JsonValue
+) -> None:
+ with _ollama_server() as wire, gateway.scenario() as scenario:
+ model: Final = _ollama_model(scenario, wire, _FULL)
+ code, text = _chat(gateway, model, initial_prompt_value=value)
+ assert code == 200, text
+ _assert_only_rendered(wire, f"{_ROLES_PROMPT}{_FINAL}")
+
+
+def test_request_initial_prompt_value_overrides_the_deployment_value(gateway: Gateway) -> None:
+ with _ollama_server() as wire, gateway.scenario() as scenario:
+ model: Final = _ollama_model(scenario, wire, _FULL)
+ code, text = _chat(gateway, model, initial_prompt_value="")
+ assert code == 200, text
+ _assert_only_rendered(wire, f"{_ROLES_PROMPT}{_FINAL}")
+
+
+def test_five_kilobyte_initial_prompt_value_leads_the_prompt(gateway: Gateway) -> None:
+ large: Final = "x" * 5120
+ with _ollama_server() as wire, gateway.scenario() as scenario:
+ model: Final = _ollama_model(scenario, wire, _ROLES_ONLY)
+ code, text = _chat(gateway, model, initial_prompt_value=large)
+ assert code == 200, text
+ _assert_only_rendered(wire, f"{large}{_ROLES_PROMPT}")
+
+
+def test_identical_requests_each_render_and_bill_once(gateway: Gateway) -> None:
+ with _ollama_server() as wire, gateway.scenario() as scenario:
+ model: Final = _ollama_model(scenario, wire, _ROLES_ONLY)
+ first, first_text = _chat(gateway, model)
+ second, second_text = _chat(gateway, model)
+ assert (first, second) == (200, 200), (first_text, second_text)
+ bodies: Final = _generate_bodies(wire)
+ assert len(bodies) == 2, bodies
+ for body in bodies:
+ _assert_rendered(body, _ROLES_PROMPT, stream=False)
+ identities: Final = {string_value(_JSON_OBJECT.validate_json(text)["id"]) for text in (first_text, second_text)}
+ assert len(identities) == 2
+ for identity in identities:
+ assert _spend_row(identity) == _billed(model)
+
+
+def test_template_with_only_a_bos_token_names_the_fields_it_lacks(gateway: Gateway) -> None:
+ with _ollama_server() as wire, gateway.scenario() as scenario:
+ model: Final = _ollama_model(scenario, wire, {"bos_token": ""})
+ code, text = _chat(gateway, model)
+ assert code == 500, text
+ assert "it sets none of roles, initial_prompt_value, or final_prompt_value" in text, text
+ assert _generate_bodies(wire) == ()
+
+
+def test_roles_that_are_not_a_mapping_render_without_role_markers(gateway: Gateway) -> None:
+ with _ollama_server() as wire, gateway.scenario() as scenario:
+ model: Final = _ollama_model(scenario, wire, {**_VALUES_ONLY, "roles": json.dumps(_ROLES)})
+ code, text = _chat(gateway, model)
+ assert code == 200, text
+ _assert_only_rendered(wire, _VALUES_PROMPT)
+
+
+def test_provider_error_reaches_the_caller_after_one_attempt(gateway: Gateway) -> None:
+ missing: Final = Reply(status=404, body=json.dumps({"error": f"model '{_BACKEND}' not found"}).encode())
+ with _ollama_server(lambda _: missing) as wire, gateway.scenario() as scenario:
+ model: Final = _ollama_model(scenario, wire, _ROLES_ONLY)
+ code, text = _chat(gateway, model)
+ assert code == 404, text
+ assert f"model '{_BACKEND}' not found" in text, text
+ bodies: Final = _generate_bodies(wire)
+ assert len(bodies) == 1, bodies
+ _assert_rendered(bodies[0], _ROLES_PROMPT, stream=False)
+
+
+def _template_keys(fields: Mapping[str, JsonValue]) -> tuple[bool, bool]:
+ return "initial_prompt_value" in fields, "final_prompt_value" in fields
+
+
+def test_pass_through_body_treats_initial_and_final_prompt_values_alike(gateway: Gateway) -> None:
+ with wire_server(lambda _: _generate_reply(_ANSWER)) as wire, gateway.scenario() as scenario:
+ path: Final = f"/integration-template-passthrough-{uuid.uuid4().hex}"
+ created: Final = gateway.post(
+ "/config/pass_through_endpoint", {"path": path, "target": f"{wire.url}/api/generate"}
+ )
+ endpoint: Final = object_value(_JSON_LIST.validate_python(created["endpoints"])[0])
+ scenario.cleanups.callback(
+ lambda: gateway.request(
+ "DELETE", "/config/pass_through_endpoint", params={"endpoint_id": str(endpoint["id"])}
+ )
+ )
+ response: Final = eventually(
+ lambda: gateway.request("POST", path, {"model": _BACKEND, "prompt": _QUESTION, "stream": False, **_FULL}),
+ lambda answered: answered.status_code != 404,
+ seconds=30,
+ )
+ assert response.status_code == 200, response.text
+ posts: Final = tuple(request for request in wire.drain() if request.method == "POST")
+ forwarded: Final = _JSON_OBJECT.validate_json(posts[-1].body)
+ assert (forwarded["model"], forwarded["prompt"]) == (_BACKEND, _QUESTION), forwarded
+ initial, final = _template_keys(forwarded)
+ assert initial == final, sorted(forwarded)
+
+
+def _a2a_reply(request: Request) -> Reply:
+ result: Final[dict[str, JsonValue]] = {
+ "kind": "message",
+ "role": "agent",
+ "messageId": f"{uuid.uuid4().hex}-out",
+ "parts": [{"kind": "text", "text": _ANSWER}],
+ }
+ outbound: Final = _JSON_OBJECT.validate_json(request.body)
+ return Reply(body=json.dumps({"jsonrpc": "2.0", "id": outbound["id"], "result": result}).encode())
+
+
+def test_a2a_params_treat_initial_and_final_prompt_values_alike(gateway: Gateway) -> None:
+ marker: Final = uuid.uuid4().hex
+ agent_name: Final = f"template-agent-{marker}"
+
+ def peer(request: Request) -> Reply:
+ if request.method == "GET":
+ return Reply(body=json.dumps(_agent_card(wire.url, agent_name)).encode())
+ return _a2a_reply(request)
+
+ with wire_server(peer) as wire, gateway.scenario() as scenario:
+ created: Final = gateway.post(
+ "/v1/agents", {"agent_name": agent_name, "agent_card_params": _agent_card(wire.url, agent_name)}
+ )
+ agent_id: Final = string_value(created["agent_id"])
+ scenario.cleanups.callback(lambda: gateway.request("DELETE", f"/v1/agents/{agent_id}"))
+ message: Final[dict[str, JsonValue]] = {
+ "kind": "message",
+ "role": "user",
+ "messageId": f"{marker}-in",
+ "parts": [{"kind": "text", "text": _QUESTION}],
+ }
+ invoked: Final = gateway.request(
+ "POST",
+ f"/a2a/{agent_id}",
+ {
+ "jsonrpc": "2.0",
+ "id": marker,
+ "method": "message/send",
+ "params": {"message": message, "initial_prompt_value": _INITIAL, "final_prompt_value": _FINAL},
+ },
+ )
+ assert invoked.status_code == 200, invoked.text
+ assert object_value(invoked.json()["result"])["parts"] == [{"kind": "text", "text": _ANSWER}], invoked.text
+ posts: Final = tuple(request for request in wire.drain() if request.method == "POST")
+ assert len(posts) == 1, [request.target for request in posts]
+ params: Final = object_value(_JSON_OBJECT.validate_json(posts[0].body)["params"])
+ assert _template_keys(params) == (False, False), sorted(params)
+
+
+def _agent_card(url: str, name: str) -> dict[str, JsonValue]:
+ return {
+ "protocolVersion": "0.3",
+ "name": name,
+ "description": "Scripted A2A peer",
+ "version": "1.0.0",
+ "url": url + "/",
+ "capabilities": {"streaming": True},
+ "defaultInputModes": ["text"],
+ "defaultOutputModes": ["text"],
+ "skills": [],
+ }
diff --git a/tests/integration/providers/test_gemini_anyof_null_branches_wire.py b/tests/integration/providers/test_gemini_anyof_null_branches_wire.py
new file mode 100644
index 00000000000..9fa99cb69f8
--- /dev/null
+++ b/tests/integration/providers/test_gemini_anyof_null_branches_wire.py
@@ -0,0 +1,681 @@
+import asyncio
+import base64
+import json
+import re
+import threading
+import uuid
+from collections.abc import AsyncIterator, Iterator, Mapping
+from concurrent.futures import ThreadPoolExecutor
+from itertools import product
+from types import MappingProxyType
+from typing import Final, Literal
+
+import anthropic
+import httpx
+import openai
+import pytest
+import websockets
+from integration._support.client import (
+ JSON_OBJECT,
+ Gateway,
+ Scenario,
+ gateway_from_environment,
+ object_value,
+ string_value,
+)
+from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
+from integration._support.vertex import service_account_json
+from integration._support.wire import Reply, Request, Wire, wire_server
+from integration._support.upstream import GEMINI_LIVE_PATH
+from integration.providers._gemini_tts_voice_support import live_deployment, live_scenario
+from integration.providers._vertex_batch_rows import encoded, gcs_peer, is_gcs_upload, jsonl, upload, vertex_model
+from pydantic import JsonValue, TypeAdapter
+
+Provider = Literal["vertex_ai", "gemini"]
+Endpoint = Literal["chat", "messages", "responses"]
+
+_BACKEND: Final = "gemini-3.8-flash"
+_NO_JSON_SCHEMA_BACKEND: Final = "gemini-flash-latest"
+_IMAGE_BACKEND: Final = "gemini-3.1-flash-image"
+_PROJECT: Final = "scripted-project"
+_LOCATION: Final = "us-central1"
+_GEMINI_KEY: Final = "sk-scripted-gemini"
+_COUNT: Final = 777
+_FAIL_MARKER: Final = "scripted-invalid-argument"
+_STRING_ANYOF: Final = re.compile(rb'"anyOf":\s*"')
+_SCALAR_BRANCH: Final = re.compile(rb'"anyOf":\s*\[\s*(?:true|false|-?\d)')
+_PNG: Final = base64.b64encode(b"\x89PNG\r\n\x1a\nscripted").decode()
+_JSON_LIST: Final = TypeAdapter(list[JsonValue])
+
+_NULL: Final[Mapping[str, JsonValue]] = {"type": "null"}
+_STRING: Final[Mapping[str, JsonValue]] = {"type": "string"}
+_INTEGER: Final[Mapping[str, JsonValue]] = {"type": "integer"}
+_NULLABLE_STRING: Final[JsonValue] = {"anyOf": [{"type": "string", "nullable": True}]}
+_NULLABLE_NULL: Final[JsonValue] = {"anyOf": [{"type": "null", "nullable": True}]}
+
+_UNIONS: Final[tuple[tuple[str, JsonValue, JsonValue], ...]] = (
+ ("null-null-string", {"anyOf": [_NULL, _NULL, _STRING]}, _NULLABLE_STRING),
+ ("null-only", {"anyOf": [_NULL]}, _NULLABLE_NULL),
+ ("two-nulls-only", {"anyOf": [_NULL, _NULL]}, _NULLABLE_NULL),
+ ("string-null", {"anyOf": [_STRING, _NULL]}, _NULLABLE_STRING),
+ (
+ "integer-null-null-string",
+ {"anyOf": [_INTEGER, _NULL, _NULL, _STRING]},
+ {"anyOf": [{"type": "integer", "nullable": True}, {"type": "string", "nullable": True}]},
+ ),
+ ("null-empty", {"anyOf": [_NULL, {}]}, {"anyOf": [{"type": "object", "nullable": True}]}),
+ ("empty-only", {"anyOf": [{}]}, {"anyOf": [{"type": "object"}]}),
+ ("empty-list", {"anyOf": []}, {"anyOf": []}),
+ ("object-anyof", {"anyOf": _STRING}, {"anyOf": _STRING}),
+ ("null-anyof", {"anyOf": None}, {"anyOf": None}),
+ ("five-hundred-nulls-then-string", {"anyOf": [*([_NULL] * 500), _STRING]}, _NULLABLE_STRING),
+ ("string-null-type-array", {"type": ["string", "null"]}, _NULLABLE_STRING),
+ (
+ "null-inside-a-union-branch",
+ {"anyOf": [{"type": "object", "properties": {"x": {"anyOf": [_INTEGER, _NULL]}}}, _NULL]},
+ {
+ "anyOf": [
+ {
+ "type": "object",
+ "properties": {"x": {"anyOf": [{"type": "integer", "nullable": True}]}},
+ "nullable": True,
+ }
+ ]
+ },
+ ),
+)
+_FANOUT_UNIONS: Final = tuple(case for case in _UNIONS if case[0] in {"null-null-string", "null-only", "object-anyof"})
+
+
+def _text_reply() -> Mapping[str, JsonValue]:
+ return {
+ "candidates": [{"content": {"role": "model", "parts": [{"text": "OK"}]}, "finishReason": "STOP"}],
+ "usageMetadata": {"promptTokenCount": 9, "candidatesTokenCount": 1, "totalTokenCount": 10},
+ "modelVersion": _BACKEND,
+ }
+
+
+def _image_reply() -> Mapping[str, JsonValue]:
+ return {
+ "candidates": [
+ {"content": {"role": "model", "parts": [{"inlineData": {"mimeType": "image/png", "data": _PNG}}]}}
+ ],
+ "usageMetadata": {"promptTokenCount": 9, "candidatesTokenCount": 1, "totalTokenCount": 10},
+ }
+
+
+def _rejection(message: str) -> Reply:
+ error: Final = {"code": 400, "message": message, "status": "INVALID_ARGUMENT"}
+ return Reply(status=400, body=json.dumps({"error": error}).encode())
+
+
+def _google(request: Request) -> Reply:
+ if is_gcs_upload(request):
+ return gcs_peer(request)
+ if _FAIL_MARKER.encode() in request.body:
+ return _rejection(f"{_FAIL_MARKER}: the scripted tool schema was refused")
+ if _STRING_ANYOF.search(request.body):
+ return _rejection("Invalid value at 'tools[0].function_declarations[0].parameters' (anyOf), not a list")
+ if _SCALAR_BRANCH.search(request.body):
+ return _rejection(
+ "Invalid value at 'tools[0].function_declarations[0].parameters.properties[0].value.any_of[0]'"
+ )
+ path: Final = request.target.split("?")[0]
+ if path.endswith(":countTokens"):
+ return Reply(body=json.dumps({"totalTokens": _COUNT}).encode())
+ if path.endswith(":streamGenerateContent"):
+ return Reply(content_type="text/event-stream", chunks=(f"data: {json.dumps(_text_reply())}\n\n".encode(),))
+ if path.endswith(f"{_IMAGE_BACKEND}:generateContent"):
+ return Reply(body=json.dumps(_image_reply()).encode())
+ if path.endswith(":generateContent"):
+ return Reply(body=json.dumps(_text_reply()).encode())
+ return Reply(status=404, body=json.dumps({"error": {"code": 404, "message": request.target}}).encode())
+
+
+def _deployment(gateway: Gateway, scenario: Scenario, wire: Wire, provider: Provider, backend: str = _BACKEND) -> str:
+ if provider == "gemini":
+ return scenario.model(model=f"gemini/{backend}", api_base=wire.url, api_key=_GEMINI_KEY)
+ return scenario.model(
+ model=f"vertex_ai/{backend}",
+ api_base=wire.url,
+ api_key=None,
+ vertex_project=_PROJECT,
+ vertex_location=_LOCATION,
+ vertex_credentials=service_account_json(_PROJECT, gateway.upstream_url.rstrip("/")),
+ )
+
+
+def _parameters(value: JsonValue) -> dict[str, JsonValue]:
+ return {"type": "object", "properties": {"value": value}, "required": ["value"]}
+
+
+def _declaration(description: str, value: JsonValue) -> dict[str, JsonValue]:
+ return {"name": "record", "description": description, "parameters": _parameters(value)}
+
+
+def _tool(endpoint: Endpoint, description: str, parameters: JsonValue) -> dict[str, JsonValue]:
+ if endpoint == "messages":
+ return {"name": "record", "description": description, "input_schema": parameters}
+ if endpoint == "responses":
+ return {"type": "function", "name": "record", "description": description, "parameters": parameters}
+ return {"type": "function", "function": {"name": "record", "description": description, "parameters": parameters}}
+
+
+def _request_body(
+ endpoint: Endpoint, model: str, description: str, parameters: JsonValue, stream: bool
+) -> tuple[str, dict[str, JsonValue]]:
+ prompt: Final = f"Record the value hello by calling the record tool. {description}"
+ tools: Final[list[JsonValue]] = [_tool(endpoint, description, parameters)]
+ streaming: Final[dict[str, JsonValue]] = {"stream": True} if stream else {}
+ if endpoint == "messages":
+ messages_body: Final[dict[str, JsonValue]] = {
+ "model": model,
+ "max_tokens": 16,
+ "messages": [{"role": "user", "content": prompt}],
+ "tools": tools,
+ }
+ return "/v1/messages", {**messages_body, **streaming}
+ if endpoint == "responses":
+ return "/v1/responses", {"model": model, "input": prompt, "tools": tools, **streaming}
+ chat_body: Final[dict[str, JsonValue]] = {
+ "model": model,
+ "max_tokens": 16,
+ "messages": [{"role": "user", "content": prompt}],
+ "tools": tools,
+ }
+ return "/v1/chat/completions", {**chat_body, **streaming}
+
+
+def _body(request: Request) -> dict[str, JsonValue]:
+ return JSON_OBJECT.validate_json(request.body)
+
+
+def _json(response: httpx.Response) -> dict[str, JsonValue]:
+ return JSON_OBJECT.validate_json(response.content)
+
+
+def _generate_requests(wire: Wire) -> tuple[Request, ...]:
+ return tuple(
+ request
+ for request in wire.drain()
+ if request.target.split("?")[0].endswith((":generateContent", ":streamGenerateContent"))
+ )
+
+
+def _sent_declarations(body: Mapping[str, JsonValue]) -> JsonValue:
+ tools: Final = _JSON_LIST.validate_python(body["tools"])
+ assert len(tools) == 1, tools
+ tool: Final = object_value(tools[0])
+ assert len(tool) == 1, tool
+ return next(iter(tool.values()))
+
+
+def _description() -> str:
+ return f"Record a value {uuid.uuid4().hex[:12]}"
+
+
+def _assert_caller_saw_ok(endpoint: Endpoint, response: httpx.Response, stream: bool) -> None:
+ assert response.status_code == 200, response.text
+ if stream:
+ assert "OK" in response.text, response.text
+ return
+ payload: Final = _json(response)
+ assert string_value(payload["id"]), payload
+ if endpoint == "chat":
+ choice: Final = object_value(_JSON_LIST.validate_python(payload["choices"])[0])
+ assert object_value(choice["message"])["content"] == "OK", payload
+ return
+ assert "OK" in response.text, payload
+
+
+@pytest.mark.parametrize("provider", ("vertex_ai", "gemini"))
+@pytest.mark.parametrize(("union", "expected"), [case[1:] for case in _UNIONS], ids=[case[0] for case in _UNIONS])
+def test_chat_tool_union_reaches_google_with_null_branches_folded_into_nullable(
+ gateway: Gateway, provider: Provider, union: JsonValue, expected: JsonValue
+) -> None:
+ with wire_server(_google) as wire, gateway.scenario() as scenario:
+ model: Final = _deployment(gateway, scenario, wire, provider)
+ description: Final = _description()
+ path, body = _request_body("chat", model, description, _parameters(union), False)
+ response: Final = gateway.request("POST", path, body)
+ _assert_caller_saw_ok("chat", response, False)
+ requests: Final = _generate_requests(wire)
+ assert len(requests) == 1, [request.target for request in requests]
+ assert requests[0].target.split("?")[0].endswith(f"{_BACKEND}:generateContent"), requests[0].target
+ assert _sent_declarations(_body(requests[0])) == [_declaration(description, expected)]
+
+
+@pytest.mark.parametrize("stream", (False, True), ids=("whole", "stream"))
+@pytest.mark.parametrize("endpoint", ("chat", "messages", "responses"))
+@pytest.mark.parametrize("provider", ("vertex_ai", "gemini"))
+@pytest.mark.parametrize(
+ ("union", "expected"), [case[1:] for case in _FANOUT_UNIONS], ids=[case[0] for case in _FANOUT_UNIONS]
+)
+def test_every_unified_endpoint_sends_the_same_folded_union(
+ gateway: Gateway, provider: Provider, endpoint: Endpoint, stream: bool, union: JsonValue, expected: JsonValue
+) -> None:
+ with wire_server(_google) as wire, gateway.scenario() as scenario:
+ model: Final = _deployment(gateway, scenario, wire, provider)
+ description: Final = _description()
+ path, body = _request_body(endpoint, model, description, _parameters(union), stream)
+ response: Final = gateway.request("POST", path, body)
+ _assert_caller_saw_ok(endpoint, response, stream)
+ requests: Final = _generate_requests(wire)
+ assert len(requests) == 1, [request.target for request in requests]
+ method: Final = "streamGenerateContent" if stream else "generateContent"
+ assert requests[0].target.split("?")[0].endswith(f"{_BACKEND}:{method}"), requests[0].target
+ assert _sent_declarations(_body(requests[0])) == [_declaration(description, expected)]
+
+
+def test_nested_item_and_property_unions_are_folded_at_every_depth(gateway: Gateway) -> None:
+ parameters: Final[dict[str, JsonValue]] = {
+ "type": "object",
+ "properties": {
+ "values": {"type": "array", "items": {"anyOf": [_NULL, _NULL, _INTEGER]}},
+ "nested": {"type": "object", "properties": {"inner": {"anyOf": [_NULL, _NULL, _STRING]}}},
+ },
+ "required": ["values", "nested"],
+ }
+ with wire_server(_google) as wire, gateway.scenario() as scenario:
+ model: Final = _deployment(gateway, scenario, wire, "vertex_ai")
+ description: Final = _description()
+ path, body = _request_body("chat", model, description, parameters, False)
+ response: Final = gateway.request("POST", path, body)
+ _assert_caller_saw_ok("chat", response, False)
+ requests: Final = _generate_requests(wire)
+ assert len(requests) == 1, [request.target for request in requests]
+ assert _sent_declarations(_body(requests[0])) == [
+ {
+ "name": "record",
+ "description": description,
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "values": {"type": "array", "items": {"anyOf": [{"type": "integer", "nullable": True}]}},
+ "nested": {"type": "object", "properties": {"inner": _NULLABLE_STRING}},
+ },
+ "required": ["values", "nested"],
+ },
+ }
+ ]
+
+
+def test_a_string_anyof_reaches_google_and_its_400_reaches_the_caller(gateway: Gateway) -> None:
+ with wire_server(_google) as wire, gateway.scenario() as scenario:
+ model: Final = _deployment(gateway, scenario, wire, "gemini")
+ description: Final = _description()
+ path, body = _request_body("chat", model, description, _parameters({"anyOf": "string"}), False)
+ response: Final = gateway.request("POST", path, body)
+ assert response.status_code == 400, response.text
+ assert "not a list" in response.text, response.text
+ requests: Final = _generate_requests(wire)
+ assert len(requests) == 1, [request.target for request in requests]
+ assert _sent_declarations(_body(requests[0])) == [_declaration(description, {"anyOf": "string"})]
+
+
+@pytest.mark.parametrize("provider", ("vertex_ai", "gemini"))
+def test_a_google_rejection_of_a_folded_union_reaches_the_caller(gateway: Gateway, provider: Provider) -> None:
+ with wire_server(_google) as wire, gateway.scenario() as scenario:
+ model: Final = _deployment(gateway, scenario, wire, provider)
+ description: Final = f"{_FAIL_MARKER} {uuid.uuid4().hex[:12]}"
+ path, body = _request_body("chat", model, description, _parameters({"anyOf": [_NULL, _NULL, _STRING]}), False)
+ response: Final = gateway.request("POST", path, body)
+ assert response.status_code == 400, response.text
+ assert _FAIL_MARKER in response.text, response.text
+ requests: Final = _generate_requests(wire)
+ assert len(requests) == 1, [request.target for request in requests]
+ assert _sent_declarations(_body(requests[0])) == [_declaration(description, _NULLABLE_STRING)]
+
+
+@pytest.mark.parametrize(
+ ("union", "sent"),
+ (({"anyOf": [True, _NULL]}, {"anyOf": [True]}), ({"anyOf": [_NULL, 5]}, {"anyOf": [5]})),
+ ids=("true-null", "null-int"),
+)
+@pytest.mark.parametrize("provider", ("vertex_ai", "gemini"))
+def test_a_union_branch_that_is_not_an_object_reaches_google_and_its_400_reaches_the_caller(
+ gateway: Gateway, provider: Provider, union: JsonValue, sent: JsonValue
+) -> None:
+ with wire_server(_google) as wire, gateway.scenario() as scenario:
+ model: Final = _deployment(gateway, scenario, wire, provider)
+ rejected_description: Final = _description()
+ path, broken = _request_body("chat", model, rejected_description, _parameters(union), False)
+ failed: Final = gateway.request("POST", path, broken)
+ assert failed.status_code == 400, failed.text
+ assert "any_of[0]" in failed.text, failed.text
+ rejected: Final = _generate_requests(wire)
+ assert len(rejected) == 1, [request.target for request in rejected]
+ assert _sent_declarations(_body(rejected[0])) == [_declaration(rejected_description, sent)]
+ description: Final = _description()
+ _, valid = _request_body("chat", model, description, _parameters({"anyOf": [_NULL, _NULL, _STRING]}), False)
+ recovered: Final = gateway.request("POST", path, valid)
+ _assert_caller_saw_ok("chat", recovered, False)
+ requests: Final = _generate_requests(wire)
+ assert len(requests) == 1, [request.target for request in requests]
+ assert _sent_declarations(_body(requests[0])) == [_declaration(description, _NULLABLE_STRING)]
+
+
+def test_a_repeated_identical_request_sends_the_same_folded_union_each_time(gateway: Gateway) -> None:
+ with wire_server(_google) as wire, gateway.scenario() as scenario:
+ model: Final = _deployment(gateway, scenario, wire, "vertex_ai")
+ description: Final = _description()
+ path, body = _request_body("chat", model, description, _parameters({"anyOf": [_NULL, _NULL, _STRING]}), False)
+ first: Final = gateway.request("POST", path, body)
+ second: Final = gateway.request("POST", path, body)
+ _assert_caller_saw_ok("chat", first, False)
+ _assert_caller_saw_ok("chat", second, False)
+ assert _json(first)["id"] != _json(second)["id"]
+ requests: Final = _generate_requests(wire)
+ assert len(requests) == 2, [request.target for request in requests]
+ assert [_sent_declarations(_body(request)) for request in requests] == [
+ [_declaration(description, _NULLABLE_STRING)]
+ ] * 2
+
+
+def test_openai_sdk_chat_sends_the_folded_union(gateway: Gateway) -> None:
+ with wire_server(_google) as wire, gateway.scenario() as scenario:
+ model: Final = _deployment(gateway, scenario, wire, "gemini")
+ description: Final = _description()
+ client: Final = openai.OpenAI(
+ base_url=f"{str(gateway.client.base_url).rstrip('/')}/v1", api_key=gateway.key, max_retries=0
+ )
+ completion: Final = client.chat.completions.create(
+ model=model,
+ max_tokens=16,
+ messages=[{"role": "user", "content": "Record the value hello by calling the record tool."}],
+ tools=[_tool("chat", description, _parameters({"anyOf": [_NULL, _NULL, _STRING]}))], # pyright: ignore[reportArgumentType] # SDK TypedDicts reject a JsonValue mapping
+ )
+ assert completion.choices[0].message.content == "OK"
+ requests: Final = _generate_requests(wire)
+ assert len(requests) == 1, [request.target for request in requests]
+ assert _sent_declarations(_body(requests[0])) == [_declaration(description, _NULLABLE_STRING)]
+
+
+async def _async_openai_stream_text(gateway: Gateway, model: str, description: str) -> str:
+ client: Final = openai.AsyncOpenAI(
+ base_url=f"{str(gateway.client.base_url).rstrip('/')}/v1", api_key=gateway.key, max_retries=0
+ )
+ stream: Final = await client.chat.completions.create(
+ model=model,
+ max_tokens=16,
+ stream=True,
+ messages=[{"role": "user", "content": "Record the value hello by calling the record tool."}],
+ tools=[_tool("chat", description, _parameters({"anyOf": [_NULL]}))], # pyright: ignore[reportArgumentType] # SDK TypedDicts reject a JsonValue mapping
+ )
+ return "".join([chunk.choices[0].delta.content or "" async for chunk in stream if chunk.choices])
+
+
+def test_async_openai_sdk_stream_sends_the_null_only_union_as_one_nullable_null(gateway: Gateway) -> None:
+ with wire_server(_google) as wire, gateway.scenario() as scenario:
+ model: Final = _deployment(gateway, scenario, wire, "vertex_ai")
+ description: Final = _description()
+ assert asyncio.run(_async_openai_stream_text(gateway, model, description)) == "OK"
+ requests: Final = _generate_requests(wire)
+ assert len(requests) == 1, [request.target for request in requests]
+ assert requests[0].target.split("?")[0].endswith(":streamGenerateContent"), requests[0].target
+ assert _sent_declarations(_body(requests[0])) == [_declaration(description, _NULLABLE_NULL)]
+
+
+def test_anthropic_sdk_messages_sends_the_folded_union(gateway: Gateway) -> None:
+ with wire_server(_google) as wire, gateway.scenario() as scenario:
+ model: Final = _deployment(gateway, scenario, wire, "vertex_ai")
+ description: Final = _description()
+ client: Final = anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0)
+ message: Final = client.messages.create(
+ model=model,
+ max_tokens=16,
+ messages=[{"role": "user", "content": "Record the value hello by calling the record tool."}],
+ tools=[_tool("messages", description, _parameters({"anyOf": [_NULL, _NULL]}))], # pyright: ignore[reportArgumentType] # SDK TypedDicts reject a JsonValue mapping
+ )
+ assert [block.type for block in message.content] == ["text"]
+ requests: Final = _generate_requests(wire)
+ assert len(requests) == 1, [request.target for request in requests]
+ assert _sent_declarations(_body(requests[0])) == [_declaration(description, _NULLABLE_NULL)]
+
+
+@pytest.mark.parametrize(
+ ("union", "expected"), [case[1:] for case in _FANOUT_UNIONS], ids=[case[0] for case in _FANOUT_UNIONS]
+)
+def test_vertex_count_tokens_sends_the_folded_union(gateway: Gateway, union: JsonValue, expected: JsonValue) -> None:
+ with wire_server(_google) as wire, gateway.scenario() as scenario:
+ model: Final = _deployment(gateway, scenario, wire, "vertex_ai")
+ description: Final = _description()
+ response: Final = gateway.request(
+ "POST",
+ "/v1/messages/count_tokens",
+ {
+ "model": model,
+ "messages": [{"role": "user", "content": "Record the value hello by calling the record tool."}],
+ "tools": [_tool("messages", description, _parameters(union))],
+ },
+ )
+ assert response.status_code == 200, response.text
+ assert _json(response)["input_tokens"] == _COUNT, response.text
+ requests: Final = tuple(request for request in wire.drain() if request.target.endswith(":countTokens"))
+ assert len(requests) == 1, [request.target for request in requests]
+ assert _sent_declarations(_body(requests[0])) == [_declaration(description, expected)]
+
+
+@pytest.mark.parametrize(
+ ("union", "expected"), [case[1:] for case in _FANOUT_UNIONS], ids=[case[0] for case in _FANOUT_UNIONS]
+)
+def test_gemini_token_counter_sends_the_folded_union(gateway: Gateway, union: JsonValue, expected: JsonValue) -> None:
+ with wire_server(_google) as wire, gateway.scenario() as scenario:
+ model: Final = _deployment(gateway, scenario, wire, "gemini")
+ description: Final = _description()
+ response: Final = gateway.request(
+ "POST",
+ "/utils/token_counter",
+ {
+ "model": model,
+ "contents": [
+ {"role": "user", "parts": [{"text": "Record the value hello by calling the record tool."}]}
+ ],
+ "tools": [_tool("chat", description, _parameters(union))],
+ },
+ params={"call_endpoint": "true"},
+ )
+ assert response.status_code == 200, response.text
+ assert _json(response)["total_tokens"] == _COUNT, response.text
+ requests: Final = tuple(request for request in wire.drain() if ":countTokens" in request.target)
+ assert len(requests) == 1, [request.target for request in requests]
+ sent: Final = object_value(_body(requests[0])["generateContentRequest"])
+ assert _sent_declarations(sent) == [_declaration(description, expected)]
+
+
+def test_response_format_on_a_model_without_json_schema_support_folds_the_union(gateway: Gateway) -> None:
+ with wire_server(_google) as wire, gateway.scenario() as scenario:
+ model: Final = _deployment(gateway, scenario, wire, "gemini", _NO_JSON_SCHEMA_BACKEND)
+ response: Final = gateway.request(
+ "POST",
+ "/v1/chat/completions",
+ {
+ "model": model,
+ "messages": [{"role": "user", "content": "Return the value as JSON."}],
+ "response_format": {
+ "type": "json_schema",
+ "json_schema": {"name": "record", "schema": _parameters({"anyOf": [_NULL, _NULL, _STRING]})},
+ },
+ },
+ )
+ _assert_caller_saw_ok("chat", response, False)
+ requests: Final = _generate_requests(wire)
+ assert len(requests) == 1, [request.target for request in requests]
+ generation_config: Final = object_value(_body(requests[0])["generationConfig"])
+ assert generation_config["response_mime_type"] == "application/json", generation_config
+ assert generation_config["response_schema"] == {**_parameters(_NULLABLE_STRING), "propertyOrdering": ["value"]}
+
+
+def test_native_generate_content_response_schema_folds_the_union(gateway: Gateway) -> None:
+ with wire_server(_google) as wire, gateway.scenario() as scenario:
+ model: Final = _deployment(gateway, scenario, wire, "gemini", _NO_JSON_SCHEMA_BACKEND)
+ response: Final = gateway.request(
+ "POST",
+ f"/v1beta/models/{model}:generateContent",
+ {
+ "contents": [{"role": "user", "parts": [{"text": "Return the value as JSON."}]}],
+ "generationConfig": {
+ "responseMimeType": "application/json",
+ "responseSchema": _parameters({"anyOf": [_NULL, _NULL, _STRING]}),
+ },
+ },
+ )
+ assert response.status_code == 200, response.text
+ requests: Final = _generate_requests(wire)
+ assert len(requests) == 1, [request.target for request in requests]
+ generation_config: Final = object_value(_body(requests[0])["generationConfig"])
+ assert generation_config["responseSchema"] == {**_parameters(_NULLABLE_STRING), "propertyOrdering": ["value"]}
+
+
+@pytest.mark.parametrize(
+ ("union", "expected"), [case[1:] for case in _FANOUT_UNIONS], ids=[case[0] for case in _FANOUT_UNIONS]
+)
+def test_image_generation_tools_send_the_folded_union(gateway: Gateway, union: JsonValue, expected: JsonValue) -> None:
+ with wire_server(_google) as wire, gateway.scenario() as scenario:
+ model: Final = _deployment(gateway, scenario, wire, "gemini", _IMAGE_BACKEND)
+ description: Final = _description()
+ response: Final = gateway.request(
+ "POST",
+ "/v1/images/generations",
+ {
+ "model": model,
+ "prompt": "A small red square on a white background",
+ "tools": [_tool("chat", description, _parameters(union))],
+ },
+ )
+ assert response.status_code == 200, response.text
+ assert len(_JSON_LIST.validate_python(_json(response)["data"])) == 1, response.text
+ requests: Final = _generate_requests(wire)
+ assert len(requests) == 1, [request.target for request in requests]
+ assert _sent_declarations(_body(requests[0])) == [_declaration(description, expected)]
+
+
+def test_vertex_batch_row_is_stored_with_the_folded_union(gateway: Gateway) -> None:
+ with wire_server(_google) as wire, gateway.scenario() as scenario:
+ model: Final = vertex_model(gateway, scenario, wire.url, _BACKEND)
+ description: Final = _description()
+ row: Final[dict[str, JsonValue]] = {
+ "custom_id": "anyof-row",
+ "method": "POST",
+ "url": "/v1/chat/completions",
+ "body": {
+ "model": model,
+ "messages": [{"role": "user", "content": "Record the value hello by calling the record tool."}],
+ "tools": [_tool("chat", description, _parameters({"anyOf": [_NULL, _NULL, _STRING]}))],
+ },
+ }
+ response: Final = upload(gateway, "model", model, jsonl(encoded(row)))
+ assert response.status_code == 200, response.text
+ uploads: Final = tuple(request for request in wire.drain() if is_gcs_upload(request))
+ assert len(uploads) == 1, [request.target for request in uploads]
+ stored: Final = tuple(
+ object_value(JSON_OBJECT.validate_json(line)["request"]) for line in uploads[0].body.splitlines() if line
+ )
+ assert len(stored) == 1, stored
+ assert _sent_declarations(stored[0]) == [_declaration(description, _NULLABLE_STRING)]
+
+
+async def _event_types_until_done(socket: websockets.ClientConnection) -> AsyncIterator[str]:
+ async for message in socket:
+ kind: Final = string_value(JSON_OBJECT.validate_json(message)["type"])
+ yield kind
+ if kind in {"response.done", "error"}:
+ return
+
+
+async def _realtime_session_with_tools(gateway: Gateway, model: str, tools: list[JsonValue]) -> tuple[str, ...]:
+ ws_base: Final = str(gateway.client.base_url).rstrip("/").replace("http://", "ws://")
+ headers: Final = {"Authorization": f"Bearer {gateway.key}"}
+ async with websockets.connect(f"{ws_base}/v1/realtime?model={model}", additional_headers=headers) as socket:
+ first: Final = JSON_OBJECT.validate_json(await socket.recv())
+ await socket.send(json.dumps({"type": "session.update", "session": {"tools": tools}}))
+ item: Final = {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "say the line"}]}
+ await socket.send(json.dumps({"type": "conversation.item.create", "item": item}))
+ events: Final = [kind async for kind in _event_types_until_done(socket)]
+ return (string_value(first["type"]), *events)
+
+
+def _realtime_types(gateway: Gateway, model: str, tools: list[JsonValue]) -> tuple[str, ...]:
+ async def bounded() -> tuple[str, ...]:
+ try:
+ return await asyncio.wait_for(_realtime_session_with_tools(gateway, model, tools), 30)
+ except TimeoutError:
+ return ("timeout",)
+
+ return asyncio.run(bounded())
+
+
+def _live_frames(gateway: Gateway, scenario_id: str) -> tuple[dict[str, JsonValue], ...]:
+ with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream:
+ observed: Final = tuple(
+ map(object_value, _JSON_LIST.validate_python(_json(upstream.get("/__observations"))["requests"]))
+ )
+ return tuple(
+ object_value(request["body"])
+ for request in observed
+ if request["api_key"] == scenario_id
+ and request["method"] == "WEBSOCKET_FRAME"
+ and request["path"] == GEMINI_LIVE_PATH
+ )
+
+
+@pytest.fixture(scope="module")
+def deferred_live(tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedProxy]:
+ directory: Final = tmp_path_factory.mktemp("anyof-live")
+ with gateway_from_environment() as gateway:
+ with owned_proxy_process(gateway, directory, {"LITELLM_GEMINI_LIVE_DEFER_SETUP": "true"}, workers=2) as proxy:
+ yield proxy
+
+
+@pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
+def test_deferred_realtime_setup_carries_session_tools_with_the_folded_union(deferred_live: OwnedProxy) -> None:
+ gateway: Final = deferred_live.gateway
+ with gateway.scenario() as scenario:
+ handle: Final = live_scenario(scenario)
+ model: Final = live_deployment(scenario, handle.scenario_id, gateway.upstream_url)
+ description: Final = _description()
+ tools: Final[list[JsonValue]] = [_tool("chat", description, _parameters({"anyOf": [_NULL, _NULL, _STRING]}))]
+ types: Final = _realtime_types(gateway, model, tools)
+ assert types[0] == "session.created", types
+ assert types[-1] == "response.done", types
+ frames: Final = _live_frames(gateway, handle.scenario_id)
+ setups: Final = tuple(object_value(frame["setup"]) for frame in frames if "setup" in frame)
+ assert len(setups) == 1, [sorted(frame) for frame in frames]
+ assert _sent_declarations(setups[0]) == [_declaration(description, _NULLABLE_STRING)]
+
+
+_PROVIDERS: Final[tuple[Provider, ...]] = ("vertex_ai", "gemini")
+_ENDPOINTS: Final[tuple[Endpoint, ...]] = ("chat", "messages", "responses")
+_BURST: Final[tuple[tuple[Provider, Endpoint, bool, bool], ...]] = tuple(
+ (provider, endpoint, stream, index % 4 == 0)
+ for index, (provider, endpoint, stream, _) in enumerate(product(_PROVIDERS, _ENDPOINTS, (False, True), range(2)))
+)
+
+
+def test_a_concurrent_mixed_burst_folds_every_union_once_and_keeps_each_status(gateway: Gateway) -> None:
+ with wire_server(_google) as wire, gateway.scenario() as scenario:
+ models: Final = MappingProxyType(
+ {provider: _deployment(gateway, scenario, wire, provider) for provider in _PROVIDERS}
+ )
+ shots: Final = tuple(
+ (provider, endpoint, stream, f"{_FAIL_MARKER if refused else 'burst'} {uuid.uuid4().hex}")
+ for provider, endpoint, stream, refused in _BURST
+ )
+
+ def send(shot: tuple[Provider, Endpoint, bool, str]) -> httpx.Response:
+ provider, endpoint, stream, description = shot
+ path, body = _request_body(
+ endpoint, models[provider], description, _parameters({"anyOf": [_NULL, _NULL, _STRING]}), stream
+ )
+ return gateway.request("POST", path, body)
+
+ with ThreadPoolExecutor(max_workers=len(shots)) as pool:
+ responses: Final = tuple(pool.map(send, shots))
+ expected_statuses: Final = [400 if _FAIL_MARKER in shot[3] else 200 for shot in shots]
+ assert [response.status_code for response in responses] == expected_statuses, [
+ response.text[:120] for response in responses if response.status_code not in {200, 400}
+ ]
+ requests: Final = _generate_requests(wire)
+ assert len(requests) == len(shots), len(requests)
+ sent: Final = sorted((json.dumps(_sent_declarations(_body(request)), sort_keys=True) for request in requests))
+ assert sent == sorted(json.dumps([_declaration(shot[3], _NULLABLE_STRING)], sort_keys=True) for shot in shots)
diff --git a/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py b/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py
index f12254b8d1d..ba2261e5a1f 100644
--- a/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py
+++ b/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py
@@ -8,10 +8,10 @@ while a client frame alone never flips a voice session into one. Every row runs
which answers the injected and rewritten updates the way the vendors do.
A blocked transcript on a transcription session reaches the client as a ``guardrail_violation`` error and
-nothing else: the ``response.cancel``, the voiced block prompt and the ``response.create`` a voice session gets
-never go to a backend that cannot speak, while ``on_violation: end_session`` and ``end_session_after_n_fails``
-still close the session. A guardrail that raises anything but a block closes the session the way it does on a
-voice session.
+nothing else: the ``response.cancel``, the blocked item's deletion, the voiced block prompt and the
+``response.create`` a voice session gets never go to a backend that cannot speak, while
+``on_violation: end_session`` and ``end_session_after_n_fails`` still close the session. A guardrail that raises
+anything but a block closes the session the way it does on a voice session.
"""
from __future__ import annotations
@@ -51,6 +51,7 @@ from tests.integration._support.database import read_rows
from tests.integration._support.process import (
OwnedProxy,
UpstreamSlot,
+ graceful_stop_seconds,
owned_proxy_process,
owned_upstream,
stop_root_process,
@@ -60,7 +61,7 @@ from tests.integration.cost_calculation.cost_tracking_case import RealtimeRespon
RecordProperty = Callable[[str, object], None]
-pytestmark: Final = pytest.mark.timeout(180)
+pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
WORKERS: Final = 2
BURST: Final = 20
@@ -83,7 +84,7 @@ CREATED_TYPES: Final = frozenset({"session.created", "transcription_session.crea
TRANSCRIPT_COMPLETED: Final = "conversation.item.input_audio_transcription.completed"
RESPONSE_DONE: Final = "response.done"
GUARDRAIL_VIOLATION: Final = ("guardrail_violation", "content_policy_violation")
-MISSING_TURN_DETECTION_TYPE: Final = ("invalid_request_error", "missing_required_parameter")
+SESSION_TYPE_MISMATCH: Final = ("invalid_request_error", "invalid_parameter")
INVALID_SESSION_VALUE: Final = ("invalid_request_error", "invalid_value")
FIVE_KB: Final = "x" * 5120
UNAUTHENTICATED_STATUS: Final = 403
@@ -104,10 +105,9 @@ BETA_TRANSCRIPTION_UPDATE: Final[dict[str, JsonValue]] = {
"input_audio_transcription": {"model": TRANSCRIBE_MODEL},
}
VOICE_UPDATE: Final[dict[str, JsonValue]] = {"type": "realtime", "instructions": "answer briefly"}
-VOICE_UPDATE_FORWARDED: Final[dict[str, JsonValue]] = {
- "type": "realtime",
- "instructions": "answer briefly",
- "audio": {"input": {"turn_detection": {"create_response": False}}},
+VOICE_UPDATE_GATED: Final[dict[str, JsonValue]] = {
+ **VOICE_UPDATE,
+ "audio": {"input": {"turn_detection": {"type": "server_vad", "create_response": False}}},
}
RE_ENABLE_UPDATE: Final[dict[str, JsonValue]] = {
"type": "realtime",
@@ -123,14 +123,32 @@ CLIENT_DECLARED_TRANSCRIPTION: Final[dict[str, JsonValue]] = {
}
CLIENT_DECLARED_TRANSCRIPTION_FORWARDED: Final[dict[str, JsonValue]] = {
"type": "transcription",
- "audio": {"input": {"turn_detection": {"create_response": False}}},
+ "audio": {"input": {"turn_detection": {"type": "server_vad", "create_response": False}}},
}
-TURN_DETECTION_TYPE_MISSING: Final = "Missing required parameter: 'session.audio.input.turn_detection.type'."
+REALTIME_UPDATE_REFUSED: Final = "Passing a transcription session update to a realtime session is not allowed."
GA_INJECTED_UPDATE: Final[dict[str, JsonValue]] = {
"type": "realtime",
"audio": {"input": {"turn_detection": {"type": "server_vad", "create_response": False}}},
}
BETA_INJECTED_UPDATE: Final[dict[str, JsonValue]] = {"turn_detection": {"type": "server_vad", "create_response": False}}
+XAI_VOICE_MODEL: Final = "grok-voice"
+GEMINI_LIVE_MODEL: Final = "gemini-3.8-live"
+VERTEX_LIVE_LOCATION: Final = "us-central1"
+UNTYPED_GATE: Final[dict[str, JsonValue]] = {"create_response": False}
+SEMANTIC_VAD: Final[dict[str, JsonValue]] = {"type": "semantic_vad", "eagerness": "low"}
+TUNED_SERVER_VAD: Final[dict[str, JsonValue]] = {
+ "type": "server_vad",
+ "prefix_padding_ms": 300,
+ "silence_duration_ms": 500,
+}
+GEMINI_TUNED_DETECTION: Final[dict[str, JsonValue]] = {
+ "disabled": True,
+ "prefixPaddingMs": 300,
+ "silenceDurationMs": 500,
+}
+CLIENT_CANCEL_ID: Final = "client_cancel_1"
+CANCEL_NOT_ACTIVE: Final = ("invalid_request_error", "response_cancel_not_active")
+AUTO_ANSWER_PREFIX: Final = "resp_auto_"
COMMIT: Final[dict[str, JsonValue]] = {"type": "input_audio_buffer.commit"}
APPEND: Final[dict[str, JsonValue]] = {"type": "input_audio_buffer.append", "audio": PCM_200MS_24KHZ}
PUSH_TO_TALK: Final = "PUSH_TO_TALK"
@@ -172,7 +190,13 @@ GUARDRAIL_END_CLOSE: Final = 1000
PROXY_FAILURE_CLOSE: Final = 1011
PROXY_FAILURE_REASON: Final = "proxy failed while relaying the upstream websocket"
TOOL_OUTPUT_BLOCKED: Final = json.dumps({"error": "Tool output blocked by content policy"})
-VOICE_BLOCK_FRAMES: Final = ("response.cancel", "conversation.item.create", "response.create")
+VOICE_BLOCK_FRAMES: Final = (
+ "response.cancel",
+ "conversation.item.delete",
+ "conversation.item.create",
+ "response.create",
+)
+XAI_VOICE_BLOCK_FRAMES: Final = ("response.cancel", "conversation.item.create", "response.create")
CREATED_SECONDS: Final = 60.0
STEP_SECONDS: Final = 15.0
SDK_SECONDS: Final = 45.0
@@ -318,6 +342,30 @@ def _voice_scenario(*transcripts: JsonValue) -> RealtimeResponse:
return RealtimeResponse(content_type="application/x-realtime", events=tuple(_voice_turns(transcripts)))
+def _xai_voice_scenario(transcript: str) -> RealtimeResponse:
+ return RealtimeResponse(
+ content_type="application/x-realtime", events=tuple(_voice_turns((transcript,))), dialect="xai"
+ )
+
+
+def _voice_scenario_created_after_the_first_frame(transcript: str) -> RealtimeResponse:
+ return RealtimeResponse(
+ content_type="application/x-realtime", events=tuple(_voice_turns((transcript,))), created_after_first_frame=True
+ )
+
+
+def _live_scenario() -> RealtimeResponse:
+ return RealtimeResponse(content_type="application/x-realtime", events=())
+
+
+def _realtime_update(turn_detection: JsonValue) -> dict[str, JsonValue]:
+ return {"type": "realtime", "audio": {"input": {"turn_detection": turn_detection}}}
+
+
+def _without_auto_response(turn_detection: dict[str, JsonValue]) -> dict[str, JsonValue]:
+ return {**turn_detection, "create_response": False}
+
+
def _muse_scenario(transcript: str) -> RealtimeResponse:
return RealtimeResponse(
content_type="application/x-realtime",
@@ -358,8 +406,41 @@ def _with_transcription_model(update: dict[str, JsonValue], model: str) -> dict[
return {**update, "audio": {"input": {**audio_input, "transcription": {"model": model}}}}
-def _muse_deployment(scenario: Scenario, scenario_id: str, api_base: str) -> str:
- return scenario.model(model=f"meta/{MUSE_MODEL}", api_key=scenario_id, api_base=api_base)
+def _muse_deployment(scenario: Scenario, scenario_id: str, api_base: str, *, drop_params: bool) -> str:
+ return scenario.model(model=f"meta/{MUSE_MODEL}", api_key=scenario_id, api_base=api_base, drop_params=drop_params)
+
+
+def _gemini_studio_deployment(scenario: Scenario, scenario_id: str, upstream_url: str) -> str:
+ return scenario.model(
+ model=f"gemini/{GEMINI_LIVE_MODEL}",
+ api_key=scenario_id,
+ api_base=upstream_url.rstrip("/"),
+ model_info={"mode": "realtime"},
+ )
+
+
+def _vertex_credentials(upstream_url: str) -> str:
+ return json.dumps(
+ {
+ "type": "external_account",
+ "audience": "synthetic-vertex-audience",
+ "subject_token_type": "urn:ietf:params:oauth:token-type:jwt",
+ "token_url": f"{upstream_url}/_oauth/token",
+ "credential_source": {"url": f"{upstream_url}/health"},
+ }
+ )
+
+
+def _vertex_live_deployment(scenario: Scenario, project: str, upstream_url: str) -> str:
+ upstream: Final = upstream_url.rstrip("/")
+ return scenario.model(
+ model=f"vertex_ai/{GEMINI_LIVE_MODEL}",
+ vertex_project=project,
+ vertex_credentials=_vertex_credentials(upstream),
+ vertex_location=VERTEX_LIVE_LOCATION,
+ api_base=upstream,
+ model_info={"mode": "realtime"},
+ )
def _named_deployment_id(creator: Gateway, scenario: Scenario, name: str, model: str, scenario_id: str) -> str:
@@ -549,6 +630,55 @@ def _transcribe(
return asyncio.run(_session(ws_base, path, query, key, frames, until=until, headers=headers))
+async def _received_after_sending(
+ socket: ClientConnection, frames: tuple[dict[str, JsonValue], ...], until: str
+) -> AsyncIterator[dict[str, JsonValue]]:
+ try:
+ for frame in frames:
+ await socket.send(json.dumps(frame))
+ async for event in _until(socket, until, STEP_SECONDS):
+ yield event
+ except ConnectionClosed as closed:
+ yield {"type": "closed", "code": _close_code(closed)}
+
+
+def _is_proxy_answer(event: dict[str, JsonValue]) -> bool:
+ if event.get("type") != RESPONSE_DONE:
+ return False
+ return not string_value(object_value(event["response"])["id"]).startswith(AUTO_ANSWER_PREFIX)
+
+
+async def _until_the_proxy_answer(
+ socket: ClientConnection, frames: tuple[dict[str, JsonValue], ...]
+) -> AsyncIterator[dict[str, JsonValue]]:
+ try:
+ yield JSON_OBJECT.validate_json(await asyncio.wait_for(socket.recv(), CREATED_SECONDS))
+ for frame in frames:
+ await socket.send(json.dumps(frame))
+ deadline: Final = asyncio.get_running_loop().time() + STEP_SECONDS
+ while (event := await _next_event(socket, deadline)) is not None:
+ yield event
+ if _is_proxy_answer(event):
+ return
+ yield {"type": "timeout"}
+ except ConnectionClosed as closed:
+ yield {"type": "closed", "code": _close_code(closed)}
+
+
+async def _talked_until_the_proxy_answer(
+ ws_base: str, query: str, key: str, frames: tuple[dict[str, JsonValue], ...]
+) -> Session:
+ headers: Final = {"Authorization": f"Bearer {key}"}
+ async with websockets.connect(f"{ws_base}/v1/realtime?{query}", additional_headers=headers) as socket:
+ return Session(tuple([event async for event in _until_the_proxy_answer(socket, frames)]), None)
+
+
+async def _talked_at_connect(ws_base: str, query: str, key: str, frames: tuple[dict[str, JsonValue], ...]) -> Session:
+ headers: Final = {"Authorization": f"Bearer {key}"}
+ async with websockets.connect(f"{ws_base}/v1/realtime?{query}", additional_headers=headers) as socket:
+ return Session(tuple([event async for event in _received_after_sending(socket, frames, RESPONSE_DONE)]), None)
+
+
def _talk(
ws_base: str,
query: str,
@@ -617,10 +747,43 @@ def _session_updates(observed: tuple[dict[str, JsonValue], ...]) -> tuple[JsonVa
return tuple(frame.get("session") for frame in _sent(observed) if frame.get("type") == "session.update")
+def _assert_voice_updates(observed: tuple[dict[str, JsonValue], ...], *later: dict[str, JsonValue]) -> None:
+ updates: Final = _session_updates(observed)
+ assert updates[0] == GA_INJECTED_UPDATE, updates
+ assert updates[1] == VOICE_UPDATE_GATED, updates
+ assert updates[2:] == later, updates
+
+
def _query(upgrade: dict[str, JsonValue]) -> JsonValue:
return upgrade["query"]
+def _declared_turn_detection(update: JsonValue) -> JsonValue:
+ session: Final = object_value(update)
+ audio_input: Final = object_value(object_value(session.get("audio", {})).get("input", {}))
+ return audio_input["turn_detection"] if "turn_detection" in audio_input else session.get("turn_detection")
+
+
+def _upstream_turn_detections(observed: tuple[dict[str, JsonValue], ...]) -> tuple[JsonValue, ...]:
+ return tuple(_declared_turn_detection(update) for update in _session_updates(observed))
+
+
+def _response_ids(session: Session) -> tuple[str, ...]:
+ done: Final = tuple(object_value(event["response"]) for event in session.events if event["type"] == RESPONSE_DONE)
+ return tuple(string_value(response["id"]) for response in done)
+
+
+def _error_event_ids(session: Session) -> tuple[JsonValue, ...]:
+ errors: Final = tuple(object_value(event["error"]) for event in session.events if event["type"] == "error")
+ return tuple(error.get("event_id") for error in errors)
+
+
+def _setup_detection(observed: tuple[dict[str, JsonValue], ...]) -> JsonValue:
+ setups: Final = tuple(object_value(frame["setup"]) for frame in _sent(observed) if "setup" in frame)
+ assert len(setups) == 1, _sent(observed)
+ return object_value(setups[0].get("realtimeInputConfig", {})).get("automaticActivityDetection")
+
+
def _spend_rows(key: str, count: int) -> list[dict[str, JsonValue]]:
return eventually(
lambda: read_rows(SPEND_SQL, (sha256(key.encode()).hexdigest(),)),
@@ -704,6 +867,7 @@ def _write_config(directory: Path, ca_bundle: Path) -> Path:
"litellm_settings": {
"ssl_verify": str(ca_bundle),
"enable_redis_auth_cache": True,
+ "gemini_live_defer_setup": True,
"cache": True,
"cache_params": {"type": "redis", "host": "os.environ/REDIS_HOST", "port": "os.environ/REDIS_PORT"},
},
@@ -918,19 +1082,18 @@ def test_voice_session_keeps_the_injected_update_and_the_clean_transcript_flows(
assert session.types == (
"session.created",
"session.updated",
- "error",
+ "session.updated",
TRANSCRIPT_COMPLETED,
RESPONSE_DONE,
), session
- assert session.errors == (MISSING_TURN_DETECTION_TYPE,), session
- assert session.error_messages == (TURN_DETECTION_TYPE_MISSING,), session
+ assert session.errors == (), session
assert _sent_types(observed) == (
"session.update",
"session.update",
"input_audio_buffer.commit",
"response.create",
), _sent(observed)
- assert _session_updates(observed) == (GA_INJECTED_UPDATE, VOICE_UPDATE_FORWARDED), _session_updates(observed)
+ _assert_voice_updates(observed)
rows: Final = _spend_rows(key, 1)
assert rows[0]["call_type"] == "_arealtime", rows
@@ -945,17 +1108,15 @@ def test_voice_session_blocked_transcript_is_refused_through_the_backend(guardra
)
observed: Final = _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id)
assert session.transcripts == (BLOCKED_TRANSCRIPT,), session
- assert session.errors == (MISSING_TURN_DETECTION_TYPE, GUARDRAIL_VIOLATION), session
+ assert session.errors == (GUARDRAIL_VIOLATION,), session
assert session.types[-1] == RESPONSE_DONE, session
assert _sent_types(observed) == (
"session.update",
"session.update",
"input_audio_buffer.commit",
- "response.cancel",
- "conversation.item.create",
- "response.create",
+ *VOICE_BLOCK_FRAMES,
), _sent(observed)
- assert _session_updates(observed) == (GA_INJECTED_UPDATE, VOICE_UPDATE_FORWARDED), _session_updates(observed)
+ _assert_voice_updates(observed)
def test_voice_session_cannot_re_enable_auto_response_with_a_later_update(guardrail_proxy: OwnedProxy) -> None:
@@ -967,12 +1128,8 @@ def test_voice_session_cannot_re_enable_auto_response_with_a_later_update(guardr
session: Final = _talk(_ws_base(_owned_url(guardrail_proxy)), f"model={model}", key, frames)
observed: Final = _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id)
assert session.types[-1] == RESPONSE_DONE, session
- assert session.errors == (MISSING_TURN_DETECTION_TYPE,), session
- assert _session_updates(observed) == (
- GA_INJECTED_UPDATE,
- VOICE_UPDATE_FORWARDED,
- RE_ENABLE_FORWARDED,
- ), _session_updates(observed)
+ assert session.errors == (), session
+ _assert_voice_updates(observed, RE_ENABLE_FORWARDED)
def test_client_declared_transcription_type_does_not_bypass_the_voice_guardrail(guardrail_proxy: OwnedProxy) -> None:
@@ -984,16 +1141,14 @@ def test_client_declared_transcription_type_does_not_bypass_the_voice_guardrail(
session: Final = _talk(_ws_base(_owned_url(guardrail_proxy)), f"model={model}", key, frames, seconds=20)
observed: Final = _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id)
assert session.transcripts == (BLOCKED_TRANSCRIPT,), session
- assert session.errors == (MISSING_TURN_DETECTION_TYPE, GUARDRAIL_VIOLATION), session
- assert session.error_messages[0] == TURN_DETECTION_TYPE_MISSING, session
+ assert session.errors == (SESSION_TYPE_MISMATCH, GUARDRAIL_VIOLATION), session
+ assert session.error_messages[0] == REALTIME_UPDATE_REFUSED, session
assert session.types[-1] == RESPONSE_DONE, session
assert _sent_types(observed) == (
"session.update",
"session.update",
"input_audio_buffer.commit",
- "response.cancel",
- "conversation.item.create",
- "response.create",
+ *VOICE_BLOCK_FRAMES,
), _sent(observed)
assert _session_updates(observed) == (
GA_INJECTED_UPDATE,
@@ -1014,6 +1169,202 @@ def test_backend_session_created_typed_transcription_skips_the_injection_on_the_
assert [_query(upgrade) for upgrade in _upgrades(observed)] == [[["model", TRANSCRIBE_MODEL]]]
+def test_voice_session_typed_turn_detection_keeps_its_type_and_tuning(guardrail_proxy: OwnedProxy) -> None:
+ with guardrail_proxy.gateway.scenario() as scenario:
+ handle: Final = _scripted(scenario, _voice_scenario(CLEAN_TRANSCRIPT))
+ key: Final = scenario.key()
+ model: Final = _openai_deployment(scenario, handle.scenario_id, model=VOICE_MODEL)
+ frames: Final = _voice_frames(_realtime_update(SEMANTIC_VAD))
+ session: Final = _talk(_ws_base(_owned_url(guardrail_proxy)), f"model={model}", key, frames)
+ observed: Final = _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id)
+ assert session.errors == (), session
+ assert session.types[-1] == RESPONSE_DONE, session
+ assert _session_updates(observed) == (
+ GA_INJECTED_UPDATE,
+ _realtime_update(_without_auto_response(SEMANTIC_VAD)),
+ ), _session_updates(observed)
+
+
+@dataclass(frozen=True, slots=True)
+class LaterUpdateCase:
+ headers: Mapping[str, str]
+ first: dict[str, JsonValue]
+ later: dict[str, JsonValue]
+ forwarded: tuple[dict[str, JsonValue], ...]
+
+
+LATER_UPDATE_CASES: Final = (
+ pytest.param(
+ LaterUpdateCase(
+ {},
+ VOICE_UPDATE,
+ {"type": "realtime", "instructions": "later"},
+ (GA_INJECTED_UPDATE, VOICE_UPDATE_GATED),
+ ),
+ id="ga",
+ ),
+ pytest.param(
+ LaterUpdateCase(
+ BETA_HEADERS,
+ {"turn_detection": {"type": "server_vad"}},
+ {"instructions": "later"},
+ (BETA_INJECTED_UPDATE, BETA_INJECTED_UPDATE),
+ ),
+ id="beta",
+ ),
+)
+
+
+@pytest.mark.parametrize("case", LATER_UPDATE_CASES)
+def test_later_update_without_turn_detection_is_forwarded_verbatim(
+ guardrail_proxy: OwnedProxy, case: LaterUpdateCase
+) -> None:
+ with guardrail_proxy.gateway.scenario() as scenario:
+ handle: Final = _scripted(scenario, _voice_scenario(CLEAN_TRANSCRIPT))
+ key: Final = scenario.key()
+ model: Final = _openai_deployment(scenario, handle.scenario_id, model=VOICE_MODEL)
+ session: Final = _drive(
+ _ws_base(_owned_url(guardrail_proxy)),
+ f"model={model}",
+ key,
+ (_step(_update_frame(case.first), _update_frame(case.later), COMMIT, until=RESPONSE_DONE),),
+ headers=case.headers,
+ )
+ observed: Final = _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id)
+ assert session.errors == (), session
+ assert session.types[-1] == RESPONSE_DONE, session
+ assert _session_updates(observed) == (*case.forwarded, case.later), _session_updates(observed)
+
+
+def test_session_update_sent_before_session_created_follows_the_gate(guardrail_proxy: OwnedProxy) -> None:
+ with guardrail_proxy.gateway.scenario() as scenario:
+ handle: Final = _scripted(scenario, _voice_scenario_created_after_the_first_frame(CLEAN_TRANSCRIPT))
+ key: Final = scenario.key()
+ model: Final = _openai_deployment(scenario, handle.scenario_id, model=VOICE_MODEL)
+ frames: Final = _voice_frames(_realtime_update(TUNED_SERVER_VAD))
+ session: Final = asyncio.run(
+ _talked_at_connect(_ws_base(_owned_url(guardrail_proxy)), f"model={model}", key, frames)
+ )
+ observed: Final = _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id)
+ assert session.types == (
+ "session.created",
+ "session.updated",
+ "session.updated",
+ TRANSCRIPT_COMPLETED,
+ RESPONSE_DONE,
+ ), session
+ assert session.errors == (), session
+ assert _session_updates(observed) == (
+ GA_INJECTED_UPDATE,
+ _realtime_update(_without_auto_response(TUNED_SERVER_VAD)),
+ ), _session_updates(observed)
+
+
+def test_client_response_cancel_error_still_reaches_the_client(guardrail_proxy: OwnedProxy) -> None:
+ with guardrail_proxy.gateway.scenario() as scenario:
+ handle: Final = _scripted(scenario, _voice_scenario(CLEAN_TRANSCRIPT))
+ key: Final = scenario.key()
+ model: Final = _openai_deployment(scenario, handle.scenario_id, model=VOICE_MODEL)
+ cancel: Final[dict[str, JsonValue]] = {"type": "response.cancel", "event_id": CLIENT_CANCEL_ID}
+ session: Final = _drive(
+ _ws_base(_owned_url(guardrail_proxy)), f"model={model}", key, (_step(cancel, until="error"),)
+ )
+ assert session.errors == (CANCEL_NOT_ACTIVE,), session
+ assert _error_event_ids(session) == (CLIENT_CANCEL_ID,), session
+
+
+def _xai_blocked_session(
+ guardrail_proxy: OwnedProxy, scenario: Scenario
+) -> tuple[Session, tuple[dict[str, JsonValue], ...]]:
+ handle: Final = _scripted(scenario, _xai_voice_scenario(BLOCKED_TRANSCRIPT))
+ key: Final = scenario.key()
+ model: Final = scenario.model(model=f"xai/{XAI_VOICE_MODEL}", api_key=handle.scenario_id)
+ frames: Final = _voice_frames(_realtime_update(TUNED_SERVER_VAD))
+ session: Final = asyncio.run(
+ _talked_until_the_proxy_answer(_ws_base(_owned_url(guardrail_proxy)), f"model={model}", key, frames)
+ )
+ return session, _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id)
+
+
+def test_xai_voice_session_gets_untyped_gates_and_no_unguarded_answer(guardrail_proxy: OwnedProxy) -> None:
+ with guardrail_proxy.gateway.scenario() as scenario:
+ session, observed = _xai_blocked_session(guardrail_proxy, scenario)
+ turn_detections: Final = _upstream_turn_detections(observed)
+ assert turn_detections == (UNTYPED_GATE, UNTYPED_GATE), turn_detections
+ assert session.transcripts == (BLOCKED_TRANSCRIPT,), session
+ assert GUARDRAIL_VIOLATION in session.errors, session
+ assert not any(response.startswith(AUTO_ANSWER_PREFIX) for response in _response_ids(session)), session
+
+
+def test_xai_voice_session_block_sends_xai_no_item_delete(guardrail_proxy: OwnedProxy) -> None:
+ with guardrail_proxy.gateway.scenario() as scenario:
+ session, observed = _xai_blocked_session(guardrail_proxy, scenario)
+ assert session.transcripts == (BLOCKED_TRANSCRIPT,), session
+ assert _sent_types(observed) == (
+ "session.update",
+ "session.update",
+ "input_audio_buffer.commit",
+ *XAI_VOICE_BLOCK_FRAMES,
+ ), _sent(observed)
+
+
+@pytest.mark.xfail(
+ strict=True,
+ raises=AssertionError,
+ reason=(
+ "xAI answers the guardrail's response.cancel with an error carrying its own event_id, so the proxy cannot "
+ "tell it apart from an error for a client event and forwards it, as the merge base does"
+ ),
+)
+def test_xai_guardrail_cancel_error_does_not_reach_the_client(guardrail_proxy: OwnedProxy) -> None:
+ with guardrail_proxy.gateway.scenario() as scenario:
+ session, _ = _xai_blocked_session(guardrail_proxy, scenario)
+ assert session.errors == (GUARDRAIL_VIOLATION,), session
+
+
+TUNED_VOICE_UPDATES: Final = (
+ pytest.param(_realtime_update(TUNED_SERVER_VAD), id="nested"),
+ pytest.param({"turn_detection": TUNED_SERVER_VAD}, id="flat"),
+)
+
+
+@pytest.mark.parametrize("update", TUNED_VOICE_UPDATES)
+def test_gemini_ai_studio_voice_session_with_vad_tuning_connects_under_the_guardrail(
+ guardrail_proxy: OwnedProxy, update: dict[str, JsonValue]
+) -> None:
+ with guardrail_proxy.gateway.scenario() as scenario:
+ handle: Final = _scripted(scenario, _live_scenario())
+ key: Final = scenario.key()
+ upstream_url: Final = guardrail_proxy.gateway.upstream_url
+ model: Final = _gemini_studio_deployment(scenario, handle.scenario_id, upstream_url)
+ session: Final = _drive(
+ _ws_base(_owned_url(guardrail_proxy)),
+ f"model={model}",
+ key,
+ (_step(_update_frame(update), until="session.updated"),),
+ )
+ observed: Final = _observed(upstream_url, handle.scenario_id)
+ assert session.types == ("session.created", "session.updated"), session
+ assert _setup_detection(observed) == {"disabled": True}, _sent(observed)
+
+
+def test_vertex_live_voice_session_keeps_vad_tuning_under_the_guardrail(guardrail_proxy: OwnedProxy) -> None:
+ with guardrail_proxy.gateway.scenario() as scenario:
+ handle: Final = _scripted(scenario, _live_scenario())
+ key: Final = scenario.key()
+ upstream_url: Final = guardrail_proxy.gateway.upstream_url
+ model: Final = _vertex_live_deployment(scenario, handle.scenario_id, upstream_url)
+ session: Final = _drive(
+ _ws_base(_owned_url(guardrail_proxy)),
+ f"model={model}",
+ key,
+ (_step(_update_frame(_realtime_update(TUNED_SERVER_VAD)), until="session.updated"),),
+ )
+ observed: Final = _observed(upstream_url, handle.scenario_id)
+ assert session.types == ("session.created", "session.updated"), session
+ assert _setup_detection(observed) == GEMINI_TUNED_DETECTION, _sent(observed)
+
+
MUSE_PUSH_TO_TALK_FRAMES: Final[tuple[dict[str, JsonValue], ...]] = (
{"binary_bytes": MUSE_PACKET_BYTES},
{"binary_bytes": MUSE_PACKET_BYTES},
@@ -1041,10 +1392,11 @@ def _muse_session(
turn_detection: JsonValue,
*,
until: str | None = None,
+ drop_params: bool = False,
) -> tuple[Session, tuple[dict[str, JsonValue], ...]]:
handle: Final = _scripted(scenario, _muse_scenario(transcript), control_url=tls_upstream)
key: Final = scenario.key()
- model: Final = _muse_deployment(scenario, handle.scenario_id, tls_upstream)
+ model: Final = _muse_deployment(scenario, handle.scenario_id, tls_upstream, drop_params=drop_params)
verdict: Final = TRANSCRIPT_COMPLETED if BLOCKED_WORD not in transcript else "error"
session: Final = _talk(
_ws_base(_owned_url(guardrail_proxy)),
@@ -1120,6 +1472,40 @@ def test_muse_server_vad_transcription_session_keeps_endpointing(
), _muse_audio_frames(observed)
+def test_muse_semantic_vad_under_drop_params_falls_back_to_endpointing(
+ guardrail_proxy: OwnedProxy, tls_upstream: str
+) -> None:
+ with guardrail_proxy.gateway.scenario() as scenario:
+ session, observed = _muse_session(
+ guardrail_proxy,
+ scenario,
+ tls_upstream,
+ CLEAN_TRANSCRIPT,
+ f"&{TRANSCRIPTION_QUERY}",
+ {"type": "semantic_vad"},
+ drop_params=True,
+ )
+ assert session.transcripts == (CLEAN_TRANSCRIPT,), session
+ assert session.errors == (), session
+ assert _muse_handshake(observed)["mode"] == ENDPOINTING, _muse_handshake(observed)
+
+
+def test_muse_semantic_vad_without_drop_params_is_refused(guardrail_proxy: OwnedProxy, tls_upstream: str) -> None:
+ with guardrail_proxy.gateway.scenario() as scenario:
+ session, observed = _muse_session(
+ guardrail_proxy,
+ scenario,
+ tls_upstream,
+ CLEAN_TRANSCRIPT,
+ "",
+ {"type": "semantic_vad"},
+ until=TRANSCRIPT_COMPLETED,
+ )
+ assert session.types[-1] == "closed", session
+ assert session.transcripts == (), session
+ assert _upgrades(observed) == (), observed
+
+
@pytest.mark.xfail(
strict=True,
reason=(
@@ -1229,8 +1615,8 @@ def test_opt_in_guardrail_gates_a_voice_session_on_an_opted_out_key(guardrail_pr
)
observed: Final = _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id)
assert session.transcripts == (BLOCKED_TRANSCRIPT,), session
- assert session.errors == (MISSING_TURN_DETECTION_TYPE, GUARDRAIL_VIOLATION), session
- assert _session_updates(observed) == (GA_INJECTED_UPDATE, VOICE_UPDATE_FORWARDED), _session_updates(observed)
+ assert session.errors == (GUARDRAIL_VIOLATION,), session
+ _assert_voice_updates(observed)
def test_opt_in_guardrail_leaves_a_transcription_session_alone_on_an_opted_out_key(guardrail_proxy: OwnedProxy) -> None:
@@ -1717,12 +2103,11 @@ def test_voice_session_guardrail_raising_runtime_error_closes_with_the_same_prox
assert session.types == (
"session.created",
"session.updated",
- "error",
+ "session.updated",
TRANSCRIPT_COMPLETED,
"error",
"closed",
), session
- assert session.errors[0] == MISSING_TURN_DETECTION_TYPE, session
_assert_closed_by_a_proxy_failure(session)
assert _sent_types(observed) == ("session.update", "session.update", "input_audio_buffer.commit"), _sent(
observed
@@ -1740,8 +2125,8 @@ def test_voice_session_guardrail_raising_value_error_is_voiced_through_the_backe
(_step(VOICE_UPDATE_FRAME, COMMIT, until=RESPONSE_DONE),),
)
assert session.transcripts == (blocked,), session
- assert session.errors == (MISSING_TURN_DETECTION_TYPE, GUARDRAIL_VIOLATION), session
- assert session.error_messages[1] == f"{VALUE_ERROR_WORD} is not allowed", session
+ assert session.errors == (GUARDRAIL_VIOLATION,), session
+ assert session.error_messages == (f"{VALUE_ERROR_WORD} is not allowed",), session
assert session.types[-1] == RESPONSE_DONE, session
assert _sent_types(observed) == (
"session.update",
@@ -1764,7 +2149,7 @@ def test_voice_session_second_violation_under_end_session_after_n_fails_closes_a
(_step(VOICE_UPDATE_FRAME, COMMIT, until=RESPONSE_DONE), _step(COMMIT, until="closed")),
)
assert session.transcripts == (blocked, blocked), session
- assert session.errors == (MISSING_TURN_DETECTION_TYPE, GUARDRAIL_VIOLATION, GUARDRAIL_VIOLATION), session
+ assert session.errors == (GUARDRAIL_VIOLATION, GUARDRAIL_VIOLATION), session
assert session.types[-1] == "closed", session
assert session.close_code == GUARDRAIL_END_CLOSE, session
assert _sent_types(observed) == (
@@ -1948,7 +2333,7 @@ def test_non_string_transcript_field_closes_a_voice_session_with_the_same_proxy_
assert session.types == (
"session.created",
"session.updated",
- "error",
+ "session.updated",
TRANSCRIPT_COMPLETED,
"error",
"closed",
diff --git a/tests/integration/providers/test_realtime_unguarded_voice_turns.py b/tests/integration/providers/test_realtime_unguarded_voice_turns.py
index 4dd7fd9aeef..c7489685273 100644
--- a/tests/integration/providers/test_realtime_unguarded_voice_turns.py
+++ b/tests/integration/providers/test_realtime_unguarded_voice_turns.py
@@ -114,12 +114,10 @@ TRANSCRIPT_FIELDS: Final = (
pytest.param({"transcript": {"text": "a"}}, id="object"),
)
VOICE_UPDATE: Final[dict[str, JsonValue]] = {"type": "realtime", "instructions": "answer briefly"}
-VOICE_UPDATE_FORWARDED: Final[dict[str, JsonValue]] = {
- "type": "realtime",
- "instructions": "answer briefly",
- "audio": {"input": {"turn_detection": {"create_response": False}}},
+VOICE_UPDATE_GATED: Final[dict[str, JsonValue]] = {
+ **VOICE_UPDATE,
+ "audio": {"input": {"turn_detection": {"type": "server_vad", "create_response": False}}},
}
-TURN_DETECTION_TYPE_MISSING: Final = "Missing required parameter: 'session.audio.input.turn_detection.type'."
INJECTED_UPDATE: Final[dict[str, JsonValue]] = {
"type": "realtime",
"audio": {"input": {"turn_detection": {"type": "server_vad", "create_response": False}}},
@@ -982,15 +980,15 @@ def test_list_mode_transcript_guardrail_still_turns_auto_response_off_and_answer
assert session.types == (
"session.created",
"session.updated",
- "error",
+ "session.updated",
TRANSCRIPT_COMPLETED,
RESPONSE_DONE,
), session
assert session.transcript_fields == ({"transcript": TRANSCRIPT},), session
- assert session.error_messages == (TURN_DETECTION_TYPE_MISSING,), session
+ assert session.error_messages == (), session
assert _sent(observed) == (
_update_frame(INJECTED_UPDATE),
- _update_frame(VOICE_UPDATE_FORWARDED),
+ _update_frame(VOICE_UPDATE_GATED),
COMMIT,
PROXY_RESPONSE_CREATE,
), _sent(observed)
diff --git a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py
index 5fe82d37e9e..8de93256f6c 100644
--- a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py
+++ b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py
@@ -65,7 +65,7 @@ class _GeminiPeer:
return requests
-def _deployment(gateway: Gateway, scenario: Scenario, api_base: str) -> str:
+def _deployment(gateway: Gateway, scenario: Scenario, api_base: str, **parameters: JsonValue) -> str:
return scenario.model(
model=f"vertex_ai/{_MODEL}",
api_base=api_base,
@@ -73,6 +73,7 @@ def _deployment(gateway: Gateway, scenario: Scenario, api_base: str) -> str:
vertex_project=_PROJECT,
vertex_location=_LOCATION,
vertex_credentials=service_account_json(_PROJECT, gateway.upstream_url.rstrip("/")),
+ **parameters,
)
@@ -220,8 +221,8 @@ _CONTENT_BLOCK_BODIES: Final[Mapping[str, Mapping[str, JsonValue]]] = MappingPro
)
-def _ai_studio_deployment(scenario: Scenario, api_base: str) -> str:
- return scenario.model(model=f"gemini/{_MODEL}", api_base=api_base, api_key="scripted-gemini-key")
+def _ai_studio_deployment(scenario: Scenario, api_base: str, **parameters: JsonValue) -> str:
+ return scenario.model(model=f"gemini/{_MODEL}", api_base=api_base, api_key="scripted-gemini-key", **parameters)
def _counted_contents(count_body: Mapping[str, JsonValue]) -> JsonValue:
@@ -594,3 +595,110 @@ def test_concurrent_count_waves_recover_after_vertex_count_outage(gateway: Gatew
tuple(_json_body(request) for request in requests)
== ({"contents": [{"role": "user", "parts": [{"text": _USER}]}]},) * 9
)
+
+
+_CONTENTS: Final[JsonValue] = [{"role": "user", "parts": [{"text": _USER}]}]
+_OPENAI_TOOLS: Final[JsonValue] = [
+ {
+ "type": "function",
+ "function": {"name": "noop", "description": "Return no output", "parameters": {"type": "object"}},
+ }
+]
+
+
+def _provider_deployment(
+ gateway: Gateway, scenario: Scenario, provider: str, api_base: str, **parameters: JsonValue
+) -> str:
+ if provider == "vertex_ai":
+ return _deployment(gateway, scenario, api_base, **parameters)
+ return _ai_studio_deployment(scenario, api_base, **parameters)
+
+
+def _count(gateway: Gateway, body: Mapping[str, JsonValue], call_endpoint: bool = True) -> httpx.Response:
+ return gateway.request(
+ "POST", "/utils/token_counter", body, params={"call_endpoint": "true" if call_endpoint else "false"}
+ )
+
+
+def _total_tokens(response: httpx.Response) -> JsonValue:
+ return _JSON_OBJECT.validate_json(response.content)["total_tokens"]
+
+
+@pytest.mark.parametrize(
+ ("provider", "empty_conversation"),
+ (("vertex_ai", {}), ("gemini", {}), ("gemini", {"messages": []}), ("gemini", {"contents": []})),
+ ids=("vertex-prompt-only", "ai-studio-prompt-only", "ai-studio-empty-messages", "ai-studio-empty-contents"),
+)
+def test_utils_token_counter_sends_a_prompt_as_the_user_turn_of_its_messages_form(
+ gateway: Gateway, provider: str, empty_conversation: Mapping[str, JsonValue]
+) -> None:
+ peer: Final = _GeminiPeer()
+ with wire_server(peer) as wire, gateway.scenario() as scenario:
+ model: Final = _provider_deployment(gateway, scenario, provider, wire.url)
+ prompted: Final = _count(gateway, {"model": model, "prompt": _USER, **empty_conversation})
+ prompted_requests: Final = _count_requests(wire)
+ as_message: Final = _count(gateway, {"model": model, "messages": _MESSAGES})
+ message_requests: Final = _count_requests(wire)
+ assert prompted.status_code == 200, prompted.text
+ assert as_message.status_code == 200, as_message.text
+ assert _total_tokens(prompted) == _COUNT, prompted.text
+ assert len(prompted_requests) == 1
+ assert len(message_requests) == 1
+ assert _counted_contents(_json_body(prompted_requests[0])) == _counted_contents(_json_body(message_requests[0]))
+
+
+@pytest.mark.parametrize(
+ ("conversation", "call_endpoint"),
+ (({"messages": _MESSAGES}, True), ({"contents": _CONTENTS}, True), ({"messages": _MESSAGES}, False)),
+ ids=("provider-count-messages", "provider-count-contents", "local-count-messages"),
+)
+def test_utils_token_counter_rejects_a_prompt_next_to_a_conversation_without_calling_the_provider(
+ gateway: Gateway, conversation: Mapping[str, JsonValue], call_endpoint: bool
+) -> None:
+ peer: Final = _GeminiPeer()
+ with wire_server(peer) as wire, gateway.scenario() as scenario:
+ model: Final = _ai_studio_deployment(scenario, wire.url)
+ response: Final = _count(gateway, {"model": model, "prompt": "hi", **conversation}, call_endpoint)
+ requests: Final = _count_requests(wire)
+ assert response.status_code == 400, response.text
+ assert _JSON_OBJECT.validate_python(_JSON_OBJECT.validate_json(response.content)["error"])["param"] == "prompt"
+ assert requests == ()
+
+
+@pytest.mark.parametrize("call_endpoint", (True, False), ids=("provider-count", "local-count"))
+def test_utils_token_counter_counts_only_the_conversation_when_drop_params_drops_the_prompt(
+ gateway: Gateway, call_endpoint: bool
+) -> None:
+ peer: Final = _GeminiPeer()
+ with wire_server(peer) as wire, gateway.scenario() as scenario:
+ model: Final = _ai_studio_deployment(scenario, wire.url, drop_params=True)
+ mixed: Final = _count(gateway, {"model": model, "prompt": "hi", "messages": _MESSAGES}, call_endpoint)
+ mixed_requests: Final = _count_requests(wire)
+ alone: Final = _count(gateway, {"model": model, "messages": _MESSAGES}, call_endpoint)
+ alone_requests: Final = _count_requests(wire)
+ assert mixed.status_code == 200, mixed.text
+ assert alone.status_code == 200, alone.text
+ assert _total_tokens(mixed) == _total_tokens(alone)
+ assert [_counted_contents(_json_body(request)) for request in mixed_requests] == [
+ _counted_contents(_json_body(request)) for request in alone_requests
+ ]
+
+
+@pytest.mark.parametrize(
+ "extra",
+ ({"system": "Follow the system instruction."}, {"tools": _OPENAI_TOOLS}),
+ ids=("system", "tools"),
+)
+def test_utils_token_counter_counts_a_local_prompt_with_system_or_tools_like_its_messages_form(
+ gateway: Gateway, extra: Mapping[str, JsonValue]
+) -> None:
+ peer: Final = _GeminiPeer()
+ with wire_server(peer) as wire, gateway.scenario() as scenario:
+ model: Final = _ai_studio_deployment(scenario, wire.url)
+ prompted: Final = _count(gateway, {"model": model, "prompt": _USER, **extra}, call_endpoint=False)
+ as_message: Final = _count(gateway, {"model": model, "messages": _MESSAGES, **extra}, call_endpoint=False)
+ requests: Final = _count_requests(wire)
+ assert prompted.status_code == 200, prompted.text
+ assert as_message.status_code == 200, as_message.text
+ assert _total_tokens(prompted) == _total_tokens(as_message)
+ assert requests == ()
diff --git a/tests/integration/routing/test_base_model_context_window_owned_proxy.py b/tests/integration/routing/test_base_model_context_window_owned_proxy.py
new file mode 100644
index 00000000000..d09ec868c78
--- /dev/null
+++ b/tests/integration/routing/test_base_model_context_window_owned_proxy.py
@@ -0,0 +1,171 @@
+import json
+import uuid
+from collections.abc import Iterator, Mapping
+from dataclasses import dataclass
+from typing import Final
+
+import pytest
+from integration._support.client import gateway_from_environment
+from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
+from integration._support.wire import Reply, Request, Wire, wire_server
+from integration.providers._cache_control_marks_support import marker_of, owned_config
+from pydantic import JsonValue
+
+_TOKEN: Final = "synthetic-bedrock-bearer"
+_KIMI: Final = "moonshotai.kimi-k2.5"
+_SONNET: Final = "global.anthropic.claude-sonnet-4-6"
+_KIMI_BASE: Final = f"bedrock/us-east-1/{_KIMI}"
+_KIMI_TARGET: Final = f"/model/{_KIMI}/converse"
+_SONNET_TARGET: Final = f"/model/{_SONNET}/converse"
+_PAST_ROUTE_ROW_LINES: Final = 16250
+_PAST_ESCALATION_BUFFER_LINES: Final = 15480
+_SMALL_LINES: Final = 1300
+_REPLY: Final = json.dumps(
+ {
+ "output": {"message": {"role": "assistant", "content": [{"text": "8115"}]}},
+ "stopReason": "end_turn",
+ "usage": {"inputTokens": 9, "outputTokens": 2, "totalTokens": 11},
+ "metrics": {"latencyMs": 1},
+ }
+).encode()
+
+
+@dataclass(frozen=True, slots=True)
+class _Rig:
+ owned: OwnedProxy
+ wire: Wire
+
+
+def _converse_peer(request: Request) -> Reply:
+ assert request.headers["authorization"] == f"Bearer {_TOKEN}", request.headers
+ return Reply(body=_REPLY)
+
+
+def _deployment(
+ name: str,
+ model: str,
+ api_base: str,
+ *,
+ model_info: Mapping[str, JsonValue] | None = None,
+ **params: JsonValue,
+) -> dict[str, JsonValue]:
+ return {
+ "model_name": name,
+ "litellm_params": {
+ "model": f"bedrock/converse/{model}",
+ "api_key": _TOKEN,
+ "aws_region_name": "us-east-1",
+ "aws_bedrock_runtime_endpoint": api_base,
+ **params,
+ },
+ "model_info": {"id": f"{name}-{uuid.uuid4().hex}", **(model_info or {})},
+ }
+
+
+@pytest.fixture(scope="module")
+def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]:
+ directory: Final = tmp_path_factory.mktemp("base-model-window")
+ with gateway_from_environment() as gateway, wire_server(_converse_peer) as wire:
+ config: Final = owned_config(
+ directory,
+ [
+ _deployment("base-in-model-info", _KIMI, wire.url, model_info={"base_model": _KIMI_BASE}),
+ _deployment("base-in-litellm-params", _KIMI, wire.url, base_model=_KIMI_BASE),
+ _deployment("no-base", _KIMI, wire.url),
+ _deployment("other-model-base", _KIMI, wire.url, model_info={"base_model": _SONNET}),
+ _deployment(
+ "declared-window",
+ _KIMI,
+ wire.url,
+ model_info={"base_model": _KIMI_BASE, "max_input_tokens": 257000},
+ ),
+ _deployment(
+ "unmapped-base", _KIMI, wire.url, model_info={"base_model": f"{_KIMI_BASE}-not-in-cost-map"}
+ ),
+ _deployment(
+ "smaller-other-model-base",
+ _KIMI,
+ wire.url,
+ model_info={"base_model": "bedrock/us-east-1/meta.llama3-8b-instruct-v1:0"},
+ ),
+ _deployment("one-million", _SONNET, wire.url),
+ {
+ "model_name": "complexity",
+ "litellm_params": {
+ "model": "auto_router/complexity_router",
+ "complexity_router_config": {
+ "tiers": {
+ "SIMPLE": "base-in-model-info",
+ "MEDIUM": "base-in-model-info",
+ "COMPLEX": "one-million",
+ "REASONING": "one-million",
+ },
+ "enable_context_window_escalation": True,
+ },
+ },
+ },
+ ],
+ litellm_settings={"cache": False},
+ router_settings={"enable_pre_call_checks": True},
+ )
+ with owned_proxy_process(gateway, directory, {}, config=config, workers=2, client_timeout=120) as owned:
+ yield _Rig(owned, wire)
+
+
+def _prompt(lines: int, marker: str) -> str:
+ listing: Final = "\n".join(f"Line {i}: the value of item {i} is {(i * 7) % 1000}." for i in range(1, lines + 1))
+ return f"{listing}\n\nWhat is the value of item 12345? Answer with the number only. marker-{marker}"
+
+
+def _request(surface: str, model: str, text: str) -> tuple[str, dict[str, JsonValue]]:
+ if surface == "messages":
+ return "/v1/messages", {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": text}]}
+ if surface == "responses":
+ return "/v1/responses", {"model": model, "max_output_tokens": 16, "input": text}
+ return "/v1/chat/completions", {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": text}]}
+
+
+def _send(rig: _Rig, surface: str, model: str, lines: int) -> tuple[int, str, tuple[str, ...]]:
+ marker: Final = uuid.uuid4().hex
+ path, body = _request(surface, model, _prompt(lines, marker))
+ response: Final = rig.owned.gateway.request("POST", path, body)
+ targets: Final = tuple(request.target for request in rig.wire.drain() if marker_of(request) == marker)
+ return response.status_code, response.text, targets
+
+
+@pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
+@pytest.mark.parametrize("surface", ["chat", "messages", "responses"])
+def test_a_same_model_base_model_row_widens_the_window_pre_call_checks_enforce(rig: _Rig, surface: str) -> None:
+ status, text, targets = _send(rig, surface, "base-in-model-info", _PAST_ROUTE_ROW_LINES)
+ assert status == 200, text
+ assert targets == (_KIMI_TARGET,)
+
+
+@pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
+def test_a_same_model_base_model_in_litellm_params_widens_the_window_too(rig: _Rig) -> None:
+ status, text, targets = _send(rig, "chat", "base-in-litellm-params", _PAST_ROUTE_ROW_LINES)
+ assert status == 200, text
+ assert targets == (_KIMI_TARGET,)
+
+
+@pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
+@pytest.mark.parametrize("model", ["no-base", "other-model-base", "declared-window", "unmapped-base"])
+def test_a_base_model_that_cannot_widen_keeps_the_route_window_and_calls_no_upstream(rig: _Rig, model: str) -> None:
+ status, text, targets = _send(rig, "chat", model, _PAST_ROUTE_ROW_LINES)
+ assert status == 400, text
+ assert "ContextWindowExceededError" in text, text
+ assert targets == ()
+
+
+@pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
+def test_a_smaller_other_model_base_model_never_narrows_the_route_window(rig: _Rig) -> None:
+ status, text, targets = _send(rig, "chat", "smaller-other-model-base", _SMALL_LINES)
+ assert status == 200, text
+ assert targets == (_KIMI_TARGET,)
+
+
+@pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
+def test_complexity_router_keeps_a_prompt_on_the_tier_whose_base_model_row_holds_it(rig: _Rig) -> None:
+ status, text, targets = _send(rig, "chat", "complexity", _PAST_ESCALATION_BUFFER_LINES)
+ assert status == 200, text
+ assert targets == (_KIMI_TARGET,)
diff --git a/tests/integration/routing/test_end_user_region_routing.py b/tests/integration/routing/test_end_user_region_routing.py
index 16638d04bd8..d8ba2da042a 100644
--- a/tests/integration/routing/test_end_user_region_routing.py
+++ b/tests/integration/routing/test_end_user_region_routing.py
@@ -1,5 +1,5 @@
import uuid
-from collections.abc import Iterator
+from collections.abc import Callable, Iterator
from pathlib import Path
from typing import Final
@@ -8,12 +8,17 @@ import pytest
import yaml
from integration._support.client import Gateway, gateway_from_environment
from integration._support.process import owned_proxy
+from pydantic import JsonValue
pytestmark: Final = pytest.mark.timeout(180)
MODEL: Final = "regional-model"
UPSTREAM_BY_REGION: Final = {"eu": "regional-eu-upstream", "us": "regional-us-upstream"}
CALLS: Final = 5
+UNCOUNTABLE_MODEL: Final = "regional-uncountable-model"
+PINNED_REGION: Final = "eu"
+WEIGHT_BY_REGION: Final = {"eu": 1, "us": 99}
+SILENT_WAV: Final = "UklGRiQAAABXQVZFZm10IBAAAAABAAEAQB8AAIA+AAACABAAZGF0YQAAAAA="
@pytest.fixture(scope="module")
@@ -22,16 +27,32 @@ def candidate(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]:
with gateway_from_environment() as base:
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config["model_list"] = [
- {
- "model_name": MODEL,
- "litellm_params": {
- "model": f"openai/{upstream}",
- "api_base": f"{base.upstream_url}/v1",
- "api_key": "synthetic-region-key",
- "region_name": region,
- },
- }
- for region, upstream in UPSTREAM_BY_REGION.items()
+ *(
+ {
+ "model_name": MODEL,
+ "litellm_params": {
+ "model": f"openai/{upstream}",
+ "api_base": f"{base.upstream_url}/v1",
+ "api_key": "synthetic-region-key",
+ "region_name": region,
+ },
+ }
+ for region, upstream in UPSTREAM_BY_REGION.items()
+ ),
+ *(
+ {
+ "model_name": UNCOUNTABLE_MODEL,
+ "litellm_params": {
+ "model": f"openai/{upstream}",
+ "api_base": f"{base.upstream_url}/v1",
+ "api_key": "synthetic-region-key",
+ "region_name": region,
+ "weight": WEIGHT_BY_REGION[region],
+ },
+ "model_info": {"id": f"{UNCOUNTABLE_MODEL}-{region}", "max_input_tokens": 128000},
+ }
+ for region, upstream in UPSTREAM_BY_REGION.items()
+ ),
]
config["router_settings"] = {**config["router_settings"], "enable_pre_call_checks": True}
path: Final = directory / "region-routing.yaml"
@@ -70,3 +91,55 @@ def test_an_end_users_allowed_region_pins_every_call_to_that_regions_deployment(
assert [response.headers.get("x-litellm-model-region") for response in responses] == [region] * CALLS
observed: Final = upstream.get("/__observations").json()["requests"]
assert [request["body"]["model"] for request in observed] == [UPSTREAM_BY_REGION[region]] * CALLS
+
+
+def _uncountable_chat(marker: str, end_user: str) -> tuple[str, dict[str, JsonValue]]:
+ return (
+ "/v1/chat/completions",
+ {
+ "model": UNCOUNTABLE_MODEL,
+ "user": end_user,
+ "messages": [
+ {
+ "role": "user",
+ "content": [
+ {"type": "text", "text": f"region {marker}"},
+ {"type": "input_audio", "input_audio": {"data": SILENT_WAV, "format": "wav"}},
+ ],
+ }
+ ],
+ },
+ )
+
+
+def _uncountable_embedding(marker: str, end_user: str) -> tuple[str, dict[str, JsonValue]]:
+ return "/v1/embeddings", {"model": UNCOUNTABLE_MODEL, "user": end_user, "input": [f"region {marker}", marker]}
+
+
+@pytest.mark.parametrize(
+ ("build", "recorded_chat_calls"),
+ [
+ pytest.param(_uncountable_chat, CALLS, id="chat-input-audio"),
+ pytest.param(_uncountable_embedding, 0, id="embeddings-list-input"),
+ ],
+)
+def test_an_uncountable_prompt_still_pins_an_end_user_to_their_region(
+ candidate: Gateway, build: Callable[[str, str], tuple[str, dict[str, JsonValue]]], recorded_chat_calls: int
+) -> None:
+ with (
+ candidate.scenario() as scenario,
+ httpx.Client(base_url=candidate.upstream_url, timeout=5, trust_env=False) as upstream,
+ ):
+ end_user: Final = f"integration-end-user-{uuid.uuid4().hex}"
+ candidate.post("/end_user/new", {"user_id": end_user, "allowed_model_region": PINNED_REGION})
+ scenario.cleanups.callback(candidate.post, "/end_user/delete", {"user_ids": [end_user]})
+ key: Final = scenario.key(models=[UNCOUNTABLE_MODEL])
+ upstream.get("/__observations").raise_for_status()
+ requests: Final = tuple(build(uuid.uuid4().hex, end_user) for _ in range(CALLS))
+ responses: Final = tuple(candidate.request("POST", path, body, key=key) for path, body in requests)
+ assert [response.status_code for response in responses] == [200] * CALLS, [r.text for r in responses]
+ pinned_deployment: Final = f"{UNCOUNTABLE_MODEL}-{PINNED_REGION}"
+ assert [response.headers.get("x-litellm-model-id") for response in responses] == [pinned_deployment] * CALLS
+ observed: Final = upstream.get("/__observations").json()["requests"]
+ pinned_upstream: Final = UPSTREAM_BY_REGION[PINNED_REGION]
+ assert [request["body"]["model"] for request in observed] == [pinned_upstream] * recorded_chat_calls
diff --git a/tests/integration/routing/test_pre_call_response_format_filter.py b/tests/integration/routing/test_pre_call_response_format_filter.py
new file mode 100644
index 00000000000..24a4434fb5a
--- /dev/null
+++ b/tests/integration/routing/test_pre_call_response_format_filter.py
@@ -0,0 +1,161 @@
+import uuid
+from collections.abc import Iterator, Mapping
+from pathlib import Path
+from typing import Final
+
+import httpx
+import pytest
+import yaml
+from integration._support.client import Gateway, gateway_from_environment
+from integration._support.process import owned_proxy
+from pydantic import JsonValue
+
+pytestmark: Final = pytest.mark.timeout(180)
+
+LACKS_JSON_MODE: Final = "gpt-4"
+TAKES_JSON_MODE: Final = "gpt-4o-mini"
+UNSUPPORTED: Final = "json-mode-unsupported"
+DROPPING: Final = "json-mode-dropped"
+MIXED: Final = "json-mode-mixed"
+FALLBACK_SOURCE: Final = "json-mode-fallback-source"
+SMALL_WINDOW: Final = "json-mode-small-window"
+FALLBACK_TARGET: Final = "json-mode-fallback-target"
+JSON_MODE: Final[Mapping[str, JsonValue]] = {"type": "json_object"}
+CALLS: Final = 4
+SMALL_WINDOW_TOKENS: Final = 50
+
+
+def _deployment(group: str, upstream_model: str, upstream_url: str, **params: JsonValue) -> dict[str, JsonValue]:
+ return {
+ "model_name": group,
+ "litellm_params": {
+ "model": f"openai/{upstream_model}",
+ "api_base": f"{upstream_url}/v1",
+ "api_key": "synthetic-json-mode-key",
+ **params,
+ },
+ }
+
+
+@pytest.fixture(scope="module")
+def candidate(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]:
+ directory: Final = tmp_path_factory.mktemp("json-mode-filter")
+ with gateway_from_environment() as base:
+ config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
+ url: Final = base.upstream_url
+ config["model_list"] = [
+ _deployment(UNSUPPORTED, LACKS_JSON_MODE, url),
+ _deployment(DROPPING, LACKS_JSON_MODE, url, drop_params=True),
+ _deployment(MIXED, LACKS_JSON_MODE, url),
+ _deployment(MIXED, TAKES_JSON_MODE, url),
+ _deployment(FALLBACK_SOURCE, LACKS_JSON_MODE, url),
+ {
+ **_deployment(SMALL_WINDOW, LACKS_JSON_MODE, url),
+ "model_info": {"max_input_tokens": SMALL_WINDOW_TOKENS},
+ },
+ _deployment(FALLBACK_TARGET, TAKES_JSON_MODE, url),
+ ]
+ config["router_settings"] = {
+ **config["router_settings"],
+ "enable_pre_call_checks": True,
+ "fallbacks": [{FALLBACK_SOURCE: [FALLBACK_TARGET]}],
+ "context_window_fallbacks": [{SMALL_WINDOW: [FALLBACK_TARGET]}],
+ }
+ path: Final = directory / "json-mode-filter.yaml"
+ path.write_text(yaml.safe_dump(config))
+ with owned_proxy(base, directory, {}, config=path) as proxy:
+ yield proxy
+
+
+def _json_mode_chat(group: str, text: str, **extra: JsonValue) -> dict[str, JsonValue]:
+ return {
+ "model": group,
+ "response_format": dict(JSON_MODE),
+ "messages": [{"role": "user", "content": f"Answer in JSON. {text}"}],
+ **extra,
+ }
+
+
+def _drain(upstream: httpx.Client) -> list[dict[str, JsonValue]]:
+ response: Final = upstream.get("/__observations")
+ response.raise_for_status()
+ return [request["body"] for request in response.json()["requests"]]
+
+
+@pytest.fixture
+def upstream(candidate: Gateway) -> Iterator[httpx.Client]:
+ with httpx.Client(base_url=candidate.upstream_url, timeout=5, trust_env=False) as client:
+ _drain(client)
+ yield client
+
+
+@pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"])
+def test_json_mode_answers_400_without_an_upstream_call_when_no_deployment_takes_it(
+ candidate: Gateway, upstream: httpx.Client, stream: bool
+) -> None:
+ response: Final = candidate.request(
+ "POST", "/v1/chat/completions", _json_mode_chat(UNSUPPORTED, uuid.uuid4().hex, stream=stream)
+ )
+
+ assert response.status_code == 400, response.text
+ assert _drain(upstream) == []
+
+
+@pytest.mark.parametrize(
+ ("group", "extra", "forwarded_format"),
+ [
+ pytest.param(DROPPING, {}, None, id="deployment-drop-params"),
+ pytest.param(UNSUPPORTED, {"allowed_openai_params": ["response_format"]}, JSON_MODE, id="request-allowlist"),
+ pytest.param(UNSUPPORTED, {"drop_params": True}, None, id="request-drop-params"),
+ ],
+)
+def test_json_mode_reaches_a_deployment_that_lacks_it_once_drop_or_allow_is_set(
+ candidate: Gateway,
+ upstream: httpx.Client,
+ group: str,
+ extra: Mapping[str, JsonValue],
+ forwarded_format: Mapping[str, JsonValue] | None,
+) -> None:
+ response: Final = candidate.request(
+ "POST", "/v1/chat/completions", _json_mode_chat(group, uuid.uuid4().hex, **extra)
+ )
+
+ assert response.status_code == 200, response.text
+ received: Final = _drain(upstream)
+ assert [body["model"] for body in received] == [LACKS_JSON_MODE]
+ assert received[0].get("response_format") == forwarded_format
+
+
+def test_json_mode_in_a_mixed_group_always_reaches_the_deployment_that_takes_it(
+ candidate: Gateway, upstream: httpx.Client
+) -> None:
+ responses: Final = tuple(
+ candidate.request("POST", "/v1/chat/completions", _json_mode_chat(MIXED, uuid.uuid4().hex))
+ for _ in range(CALLS)
+ )
+
+ assert [response.status_code for response in responses] == [200] * CALLS, [r.text for r in responses]
+ received: Final = _drain(upstream)
+ assert [(body["model"], body.get("response_format")) for body in received] == [
+ (TAKES_JSON_MODE, dict(JSON_MODE))
+ ] * CALLS
+
+
+@pytest.mark.parametrize(
+ ("group", "text"),
+ [
+ pytest.param(FALLBACK_SOURCE, "short", id="generic-fallback"),
+ pytest.param(SMALL_WINDOW, "padding " * (4 * SMALL_WINDOW_TOKENS), id="context-window-fallback"),
+ ],
+)
+def test_json_mode_falls_back_to_a_group_that_takes_it(
+ candidate: Gateway, upstream: httpx.Client, group: str, text: str
+) -> None:
+ response: Final = candidate.request(
+ "POST", "/v1/chat/completions", _json_mode_chat(group, f"{uuid.uuid4().hex} {text}")
+ )
+
+ assert response.status_code == 200, response.text
+ assert [(body["model"], body.get("response_format")) for body in _drain(upstream)] == [
+ (TAKES_JSON_MODE, dict(JSON_MODE))
+ ]
diff --git a/tests/integration/sdk/test_router_budget_limiter_wire.py b/tests/integration/sdk/test_router_budget_limiter_wire.py
index cf4fda41891..0863ee6d6cc 100644
--- a/tests/integration/sdk/test_router_budget_limiter_wire.py
+++ b/tests/integration/sdk/test_router_budget_limiter_wire.py
@@ -1,7 +1,9 @@
from __future__ import annotations
+import asyncio
import json
import os
+import time
import uuid
from collections.abc import Iterator
from typing import Final
@@ -222,3 +224,95 @@ async def test_tag_budget_blocks_exhausted_tag_but_allows_another(
assert response.choices[0].message.content == "openai"
assert len(requests) == 1
assert json.loads(requests[0].body)["model"] == _MODEL
+
+
+
+_ISOLATED_SPEND_KEY: Final = "provider_spend:hosted_vllm:1d"
+_INPUT_COST_PER_TOKEN: Final = 0.001
+_OUTPUT_COST_PER_TOKEN: Final = 0.002
+
+
+def _isolated_provider_deployment(api_base: str) -> dict[str, JsonValue]:
+ """Only these two tests route to hosted_vllm, so success callbacks still draining from earlier tests cannot land here."""
+ return {
+ "model_name": "isolated-budget",
+ "litellm_params": {
+ "model": "hosted_vllm/llama-3",
+ "api_base": api_base,
+ "api_key": _API_KEY,
+ "input_cost_per_token": _INPUT_COST_PER_TOKEN,
+ "output_cost_per_token": _OUTPUT_COST_PER_TOKEN,
+ },
+ "model_info": {"id": "vllm-deployment"},
+ }
+
+
+def _isolated_budget_router(api_base: str, redis_client: Redis | None) -> Router:
+ redis_params: Final = (
+ {}
+ if redis_client is None
+ else {"redis_host": os.environ["REDIS_HOST"], "redis_port": int(os.environ["REDIS_PORT"])}
+ )
+ return Router(
+ model_list=[_isolated_provider_deployment(api_base)],
+ provider_budget_config={"hosted_vllm": BudgetConfig(budget_duration="1d", max_budget=10.0)},
+ num_retries=0,
+ disable_cooldowns=True,
+ **redis_params,
+ )
+
+
+async def _complete_and_wait_for_spend(router: Router, previous: float) -> float:
+ response: Final = await router.acompletion(
+ model="isolated-budget",
+ messages=[{"role": "user", "content": "budget tracking"}],
+ )
+ cost: Final = (
+ response.usage.prompt_tokens * _INPUT_COST_PER_TOKEN + response.usage.completion_tokens * _OUTPUT_COST_PER_TOKEN
+ )
+ deadline: Final = time.monotonic() + 5.0
+ while time.monotonic() < deadline:
+ spend: float | None = router.cache.in_memory_cache.get_cache(_ISOLATED_SPEND_KEY)
+ if spend is not None and spend > previous:
+ assert spend == pytest.approx(previous + cost)
+ return spend
+ await asyncio.sleep(0.05)
+ raise AssertionError(f"{_ISOLATED_SPEND_KEY} never rose above {previous}")
+
+
+@pytest.mark.asyncio
+async def test_router_without_redis_tracks_spend_but_queues_no_redis_increments() -> None:
+ """A Router with no Redis never flushes, so anything it queues for Redis is retained for the life of the process."""
+ with wire_server(_peer) as wire:
+ router: Final = _isolated_budget_router(f"{wire.url}/openai/v1", redis_client=None)
+ first: Final = await _complete_and_wait_for_spend(router, previous=0.0)
+ second: Final = await _complete_and_wait_for_spend(router, previous=first)
+ third: Final = await _complete_and_wait_for_spend(router, previous=second)
+ wire.drain()
+
+ assert third > second > first > 0.0
+ assert router.router_budget_logger is not None
+ assert router.router_budget_logger.redis_increment_operation_queue == []
+
+
+async def _wait_for_redis_spend(redis_client: Redis, at_least: float) -> float:
+ deadline: Final = time.monotonic() + 15.0
+ while time.monotonic() < deadline:
+ raw: str | None = redis_client.get(_ISOLATED_SPEND_KEY)
+ if raw is not None and float(raw) >= at_least:
+ return float(raw)
+ await asyncio.sleep(0.1)
+ raise AssertionError(f"{_ISOLATED_SPEND_KEY} in Redis never reached {at_least}")
+
+
+@pytest.mark.asyncio
+async def test_router_with_redis_flushes_queued_spend_increments(redis_client: Redis) -> None:
+ redis_client.delete(_ISOLATED_SPEND_KEY)
+ with wire_server(_peer) as wire:
+ router: Final = _isolated_budget_router(f"{wire.url}/openai/v1", redis_client=redis_client)
+ first: Final = await _complete_and_wait_for_spend(router, previous=0.0)
+ second: Final = await _complete_and_wait_for_spend(router, previous=first)
+ redis_spend: Final = await _wait_for_redis_spend(redis_client, at_least=second)
+ wire.drain()
+
+ assert redis_spend == pytest.approx(second)
diff --git a/tests/integration/spend/test_background_interaction_settlement.py b/tests/integration/spend/test_background_interaction_settlement.py
index 16eb0a3434e..4408bd147f9 100644
--- a/tests/integration/spend/test_background_interaction_settlement.py
+++ b/tests/integration/spend/test_background_interaction_settlement.py
@@ -1,3 +1,4 @@
+import json
import math
import socket
import time
@@ -17,6 +18,8 @@ from integration._support.client import (
Gateway,
eventually,
gateway_from_environment,
+ held,
+ list_value,
object_value,
string_value,
)
@@ -54,6 +57,9 @@ _CUSTOM_INPUT_RATE: Final = 2e-06
_CUSTOM_OUTPUT_RATE: Final = 4e-05
_ENV_KEY: Final = "integration-gemini-env-key"
_DEPLOYMENT_KEY: Final = "integration-gemini-deployment-key"
+_DATABASE_DEPLOYMENT_KEY: Final = "integration-gemini-database-deployment-key"
+_ROTATED_DEPLOYMENT_KEY: Final = "integration-gemini-rotated-deployment-key"
+_TAG: Final = "integration-settlement-team"
_CREATOR_POLL: Final = {"BACKGROUND_INTERACTION_COST_POLL_INITIAL_INTERVAL_SECONDS": "300"}
_SETTLER_POLL: Final = {
"BACKGROUND_INTERACTION_COST_POLL_INITIAL_INTERVAL_SECONDS": "1",
@@ -65,8 +71,10 @@ _RESUMER_POLL: Final = {
"BACKGROUND_INTERACTION_COST_POLL_MAX_INTERVAL_SECONDS": "1",
"BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS": "120",
}
+_EAGER_RESUMER_POLL: Final = {**_RESUMER_POLL, "BACKGROUND_INTERACTION_COST_POLL_INITIAL_INTERVAL_SECONDS": "0.01"}
+_LATE_RESUMER_POLL: Final = {**_RESUMER_POLL, "BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS": "30"}
_SPEND_QUERY: Final = (
- "SELECT request_id, spend, call_type, status, model, prompt_tokens, completion_tokens "
+ "SELECT request_id, spend, call_type, status, model, prompt_tokens, completion_tokens, request_tags "
'FROM "LiteLLM_SpendLogs" WHERE request_id = %s'
)
_SETTLEMENT_QUERY: Final = (
@@ -74,6 +82,16 @@ _SETTLEMENT_QUERY: Final = (
'settled_at IS NOT NULL AS settled, create_context FROM "LiteLLM_BackgroundInteractionSettlement" '
"WHERE interaction_id = %s"
)
+_ABANDONED_CLAIM: Final = (
+ 'UPDATE "LiteLLM_BackgroundInteractionSettlement" SET claimed_by = %s, '
+ "claimed_at = (now() AT TIME ZONE 'UTC') - interval '10 minutes', settled_at = NULL, outcome = NULL, "
+ "create_context = %s::jsonb WHERE interaction_id = %s"
+)
+_DEAD_SETTLER: Final = "dead-settler:1"
+_CREATED_AN_HOUR_AGO: Final = (
+ 'UPDATE "LiteLLM_BackgroundInteractionSettlement" '
+ "SET created_at = (now() AT TIME ZONE 'UTC') - interval '1 hour' WHERE interaction_id = %s"
+)
_SETTLEMENT_TABLE_PRESENT_QUERY: Final = "SELECT to_regclass(%s) IS NOT NULL AS present"
_SETTLEMENT_TABLE: Final = '"LiteLLM_BackgroundInteractionSettlement"'
_SETTLEMENT_BY_CALL_QUERY: Final = (
@@ -179,6 +197,40 @@ def _write_config(directory: Path, upstream_url: str, models: Deployments) -> Pa
return path
+def _variant_config(rig: Rig, name: str, document: Mapping[str, JsonValue]) -> Path:
+ path: Final = rig.directory / name
+ path.write_text(yaml.safe_dump(dict(document)))
+ return path
+
+
+def _rig_config(rig: Rig) -> dict[str, JsonValue]:
+ return JSON_OBJECT.validate_python(yaml.safe_load(rig.config.read_text()))
+
+
+def _config_without_spend_logs(rig: Rig) -> Path:
+ config: Final = _rig_config(rig)
+ general_settings: Final = {**object_value(config["general_settings"]), "disable_spend_logs": True}
+ return _variant_config(
+ rig, "settlement_config_without_spend_logs.yaml", {**config, "general_settings": general_settings}
+ )
+
+
+def _config_with_a_rotated_key(rig: Rig) -> Path:
+ config: Final = _rig_config(rig)
+ model_list: Final = [_with_rotated_key(entry, rig.models.in_progress) for entry in list_value(config["model_list"])]
+ return _variant_config(rig, "settlement_config_rotated_key.yaml", {**config, "model_list": model_list})
+
+
+def _with_rotated_key(entry: JsonValue, model_name: str) -> JsonValue:
+ deployment: Final = object_value(entry)
+ if deployment["model_name"] != model_name:
+ return deployment
+ return {
+ **deployment,
+ "litellm_params": {**object_value(deployment["litellm_params"]), "api_key": _ROTATED_DEPLOYMENT_KEY},
+ }
+
+
def _interaction(status: str, usage: dict[str, JsonValue] | None, http_status: int = 200) -> JsonResponse:
return JsonResponse(
content_type="application/json",
@@ -206,12 +258,14 @@ def _create(
path: str = "/v1beta/interactions",
background: bool = True,
text: str | None = None,
+ headers: Mapping[str, str] | None = None,
) -> str:
response: Final = replica.request(
"POST",
path,
{"model": model, "input": text or f"settle {uuid.uuid4().hex}", "background": background},
key=key,
+ headers=headers,
)
assert response.status_code == 200, response.text
return string_value(JSON_OBJECT.validate_json(response.content)["id"])
@@ -252,6 +306,16 @@ def _settlement(interaction_id: str) -> dict[str, JsonValue] | None:
return rows[0] if rows else None
+def _registered(interaction_id: str) -> dict[str, JsonValue]:
+ row: Final = eventually(lambda: _settlement(interaction_id), lambda value: value is not None)
+ assert row is not None
+ return row
+
+
+def _abandon_claim(interaction_id: str, create_context: JsonValue) -> None:
+ write_rows(_ABANDONED_CLAIM, (_DEAD_SETTLER, json.dumps(create_context), interaction_id))
+
+
def _settlement_table_present() -> bool:
return read_rows(_SETTLEMENT_TABLE_PRESENT_QUERY, (_SETTLEMENT_TABLE,))[0]["present"] is True
@@ -315,6 +379,11 @@ def _assert_billed(row: Mapping[str, JsonValue], rates: tuple[float, float]) ->
return spend
+def _assert_billed_with_tag(row: Mapping[str, JsonValue], rates: tuple[float, float]) -> None:
+ _assert_billed(row, rates)
+ assert _TAG in list_value(row["request_tags"]), row
+
+
def _key_spend(replica: Gateway, key: str) -> float:
spend: Final = object_value(replica.get("/key/info", {"key": key})["info"])["spend"]
assert isinstance(spend, float | int), spend
@@ -333,12 +402,12 @@ def _drain(rig: Rig) -> list[JsonValue]:
return requests
-def _calls(rig: Rig, interaction_id: str) -> tuple[tuple[str, str], ...]:
- suffix: Final = f"/v1beta/interactions/{interaction_id}"
+def _calls(rig: Rig, *interaction_ids: str) -> tuple[tuple[str, str], ...]:
+ suffixes: Final = tuple(f"/v1beta/interactions/{interaction_id}" for interaction_id in interaction_ids)
return tuple(
(string_value(object_value(entry)["method"]), string_value(object_value(entry)["api_key"]))
for entry in _drain(rig)
- if string_value(object_value(entry)["path"]).endswith(suffix)
+ if string_value(object_value(entry)["path"]).endswith(suffixes)
)
@@ -447,7 +516,7 @@ def test_a_replica_booting_later_resumes_and_bills_unclaimed_interactions(rig: R
model: Final = rig.models.in_progress
key: Final = scenario.key()
created_after: Final = time.time()
- created: Final = tuple(_create(rig.creator, model, key) for _ in range(3))
+ created: Final = tuple(_create(rig.creator, model, key, headers={"x-litellm-tags": _TAG}) for _ in range(3))
for item in created:
_state(rig, item, _completed())
rates: Final = _rates(rig.creator, model)
@@ -457,13 +526,113 @@ def test_a_replica_booting_later_resumes_and_bills_unclaimed_interactions(rig: R
pids: Final = _worker_pids(resumer.process.pid)
assert len(pids) == 2, pids
for item in created:
- _assert_billed(_await_spend_row(item, seconds=90), rates)
+ _assert_billed_with_tag(_await_spend_row(item, seconds=90), rates)
claimer: Final = _claimer_pid(_await_outcome(item, "billed"))
assert claimer in pids or _booted_after(claimer, created_after), (claimer, pids)
+ assert {key for _, key in _calls(rig, *created)} == {_DEPLOYMENT_KEY}
for item in created:
assert len(_spend_rows(item)) == 1
+def test_a_replica_booting_later_resumes_a_database_model_with_its_own_key(rig: Rig) -> None:
+ with rig.creator.scenario() as scenario:
+ model: Final = scenario.model(
+ model=_MODEL, api_key=_DATABASE_DEPLOYMENT_KEY, api_base=f"{rig.upstream.url}/{rig.models.in_progress}"
+ )
+ created: Final = _create(rig.creator, model, scenario.key())
+ _state(rig, created, _completed())
+ with owned_proxy_process(
+ rig.gateway, rig.directory, rig.environment(**_EAGER_RESUMER_POLL), config=rig.config, workers=2
+ ):
+ _await_spend_row(created, seconds=90)
+ _await_outcome(created, "billed")
+ assert {key for _, key in _calls(rig, created)} == {_DATABASE_DEPLOYMENT_KEY}
+
+
+def test_a_replica_booting_later_takes_over_lapsed_claims_and_bills_only_the_unbilled(rig: Rig) -> None:
+ with rig.creator.scenario() as scenario:
+ model: Final = rig.models.in_progress
+ key: Final = scenario.key()
+ unbilled: Final = _create(rig.creator, model, key)
+ billed: Final = _create(rig.settler, model, key)
+ contexts: Final = {item: _registered(item)["create_context"] for item in (unbilled, billed)}
+ for item in (unbilled, billed):
+ _state(rig, item, _completed())
+ rates: Final = _rates(rig.creator, model)
+ spend: Final = _assert_billed(_await_spend_row(billed), rates)
+ _await_outcome(billed, "billed")
+ for item, create_context in contexts.items():
+ _abandon_claim(item, create_context)
+ with owned_proxy_process(
+ rig.gateway, rig.directory, rig.environment(**_EAGER_RESUMER_POLL), config=rig.config, workers=2
+ ):
+ _assert_billed(_await_spend_row(unbilled, seconds=90), rates)
+ assert _await_outcome(unbilled, "billed")["claimed_by"] != _DEAD_SETTLER
+ assert _await_outcome(billed, "billed", seconds=90)["claimed_by"] != _DEAD_SETTLER
+ held(
+ lambda: _key_spend(rig.creator, key),
+ lambda observed: math.isclose(observed, 2 * spend, rel_tol=1e-9),
+ holding=5,
+ seconds=30,
+ )
+ for item in (unbilled, billed):
+ assert len(_spend_rows(item)) == 1
+
+
+def test_a_lapsed_claim_with_spend_logs_off_is_recorded_unsettled_and_never_billed(rig: Rig) -> None:
+ with rig.creator.scenario() as scenario:
+ model: Final = rig.models.in_progress
+ key: Final = scenario.key()
+ created: Final = _create(rig.creator, model, key)
+ _abandon_claim(created, _registered(created)["create_context"])
+ _state(rig, created, _completed())
+ with owned_proxy_process(
+ rig.gateway,
+ rig.directory,
+ rig.environment(**_EAGER_RESUMER_POLL),
+ config=_config_without_spend_logs(rig),
+ workers=2,
+ ):
+ assert _await_outcome(created, "unsettled", seconds=90)["claimed_by"] != _DEAD_SETTLER
+ assert _spend_rows(created) == []
+ assert _key_spend(rig.creator, key) == 0
+
+
+def test_a_resumed_interaction_whose_deployment_key_rotated_fetches_with_the_new_key(rig: Rig) -> None:
+ with rig.creator.scenario() as scenario:
+ model: Final = rig.models.in_progress
+ created: Final = _create(rig.creator, model, scenario.key())
+ _state(rig, created, _completed())
+ with owned_proxy_process(
+ rig.gateway,
+ rig.directory,
+ rig.environment(**_EAGER_RESUMER_POLL),
+ config=_config_with_a_rotated_key(rig),
+ workers=2,
+ ) as resumer:
+ assert _model_info(resumer.gateway, model)["id"] != _model_info(rig.creator, model)["id"]
+ _assert_billed(_await_spend_row(created, seconds=90), _rates(rig.creator, model))
+ _await_outcome(created, "billed")
+ assert {key for _, key in _calls(rig, created)} == {_ROTATED_DEPLOYMENT_KEY}
+
+
+def test_a_late_restart_retries_a_failed_fetch_within_a_full_poll_window(rig: Rig) -> None:
+ with rig.creator.scenario() as scenario:
+ model: Final = rig.models.in_progress
+ created: Final = _create(rig.creator, model, scenario.key())
+ _registered(created)
+ write_rows(_CREATED_AN_HOUR_AGO, (created,))
+ _state(rig, created, InteractionState(status="completed", usage=_USAGE, get_status=503))
+ _drain(rig)
+ with owned_proxy_process(
+ rig.gateway, rig.directory, rig.environment(**_LATE_RESUMER_POLL), config=rig.config, workers=2
+ ):
+ eventually(lambda: _calls(rig, created), lambda calls: calls != (), seconds=90)
+ _state(rig, created, _completed())
+ _assert_billed(_await_spend_row(created), _rates(rig.creator, model))
+ _await_outcome(created, "billed")
+
+
def test_deletes_on_the_creating_proxy_bill_each_interaction_once(rig: Rig) -> None:
with rig.settler.scenario() as scenario:
model: Final = rig.models.in_progress
diff --git a/tests/integration/spend/test_spend_log_caller_session.py b/tests/integration/spend/test_spend_log_caller_session.py
new file mode 100644
index 00000000000..213583d21a1
--- /dev/null
+++ b/tests/integration/spend/test_spend_log_caller_session.py
@@ -0,0 +1,783 @@
+import itertools
+import json
+import re
+import threading
+import uuid
+from urllib.parse import urlsplit
+from collections import deque
+from collections.abc import Iterator, Mapping, Sequence
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import ExitStack
+from dataclasses import dataclass, replace
+from datetime import datetime, timedelta, timezone
+from hashlib import sha256
+from pathlib import Path
+from queue import SimpleQueue
+from types import MappingProxyType
+from typing import Final
+
+import anthropic
+import httpx
+import openai
+import psutil
+import pytest
+import yaml
+from integration._support.client import (
+ JSON_OBJECT,
+ Gateway,
+ Scenario,
+ eventually,
+ gateway_from_environment,
+ list_value,
+ object_value,
+ string_value,
+)
+from integration._support.budget import readiness_seconds
+from integration._support.database import read_rows
+from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
+from integration._support.wire import Reply, Request, Wire, wire_server
+from pydantic import JsonValue
+
+pytestmark: Final = pytest.mark.timeout(2 * readiness_seconds(2) + 2 * graceful_stop_seconds() + 240)
+
+MARKER: Final = re.compile(rb"callersess-[0-9a-f]{32}")
+CALLER_SESSION_KEY: Final = "litellm_caller_session_id"
+PATHS: Final = MappingProxyType(
+ {"chat": "/v1/chat/completions", "messages": "/v1/messages", "responses": "/v1/responses"}
+)
+ENDPOINT_STREAMS: Final = tuple(itertools.product(PATHS, (False, True)))
+Headers = Mapping[str, str] | Sequence[tuple[str, str]]
+
+
+def _marker() -> str:
+ return "callersess-" + uuid.uuid4().hex
+
+
+def _token(label: str) -> str:
+ return f"{label}-{uuid.uuid4().hex}"
+
+
+def _frame(payload: Mapping[str, JsonValue]) -> bytes:
+ return b"data: " + json.dumps(payload).encode() + b"\n\n"
+
+
+def _chat_reply(text: str, stream: bool) -> Reply:
+ identity: Final = "chatcmpl-" + uuid.uuid4().hex
+ usage: Final[dict[str, JsonValue]] = {"prompt_tokens": 7, "completion_tokens": 2, "total_tokens": 9}
+ if not stream:
+ return Reply(
+ body=json.dumps(
+ {
+ "id": identity,
+ "object": "chat.completion",
+ "created": 1,
+ "model": "gpt-4o-mini",
+ "choices": [
+ {"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}
+ ],
+ "usage": usage,
+ }
+ ).encode()
+ )
+ chunk: Final[dict[str, JsonValue]] = {
+ "id": identity,
+ "object": "chat.completion.chunk",
+ "created": 1,
+ "model": "gpt-4o-mini",
+ }
+ return Reply(
+ content_type="text/event-stream",
+ chunks=(
+ _frame({**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": text}}]}),
+ _frame({**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], "usage": usage}),
+ b"data: [DONE]\n\n",
+ ),
+ )
+
+
+def _upstream(request: Request) -> Reply:
+ markers: Final = tuple(found.group(0).decode() for found in MARKER.finditer(request.body))
+ if request.method != "POST" or not request.target.endswith("/chat/completions") or not markers:
+ return Reply(status=404, body=b'{"error":"unexpected request"}')
+ return _chat_reply(f"answer {markers[-1]}", JSON_OBJECT.validate_json(request.body).get("stream") is True)
+
+
+def _user_markers(body: Mapping[str, JsonValue]) -> tuple[str, ...]:
+ user_contents: Final = tuple(
+ json.dumps(object_value(message)["content"])
+ for message in list_value(body.get("messages", []))
+ if object_value(message)["role"] == "user"
+ )
+ return tuple(found.group(0).decode() for content in user_contents if (found := MARKER.search(content.encode())))
+
+
+def _spans(batch: Request) -> Iterator[dict[str, JsonValue]]:
+ for resource in list_value(JSON_OBJECT.validate_json(batch.body)["resourceSpans"]):
+ for scope in list_value(object_value(resource)["scopeSpans"]):
+ yield from (object_value(span) for span in list_value(object_value(scope)["spans"]))
+
+
+def _attributes(span: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
+ return {
+ string_value(object_value(attribute)["key"]): object_value(attribute)["value"]
+ for attribute in list_value(span.get("attributes", []))
+ }
+
+
+def _batch_attributes(batches: Sequence[Request]) -> Iterator[dict[str, JsonValue]]:
+ return (_attributes(span) for span in itertools.chain.from_iterable(_spans(batch) for batch in batches))
+
+
+@dataclass(frozen=True, slots=True)
+class Collector:
+ wire: Wire
+ accepted: deque[Request]
+ guard: threading.Lock
+
+ def span_attributes(self, response_id: str) -> tuple[dict[str, JsonValue], ...]:
+ with self.guard:
+ batches: Final = tuple(self.accepted)
+ attributes: Final = tuple(_batch_attributes(batches))
+ return tuple(
+ found
+ for found in attributes
+ if object_value(found.get("gen_ai.response.id", {})).get("stringValue") == response_id
+ )
+
+ def conversation_ids(self, conversations: frozenset[str]) -> tuple[str, ...]:
+ with self.guard:
+ batches: Final = tuple(self.accepted)
+ responded: Final = (
+ object_value(found.get("gen_ai.conversation.id", {})).get("stringValue")
+ for found in _batch_attributes(batches)
+ if "gen_ai.response.id" in found
+ )
+ return tuple(
+ conversation
+ for conversation in responded
+ if isinstance(conversation, str) and conversation in conversations
+ )
+
+
+@dataclass(frozen=True, slots=True)
+class Rig:
+ proxy: Gateway
+ scenario: Scenario
+ model: str
+ upstream: Wire
+ process: OwnedProxy
+
+ def body(self, endpoint: str, marker: str, extra: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
+ match endpoint:
+ case "chat":
+ return {"model": self.model, "messages": [{"role": "user", "content": marker}], **extra}
+ case "messages":
+ return {
+ "model": self.model,
+ "max_tokens": 16,
+ "messages": [{"role": "user", "content": marker}],
+ **extra,
+ }
+ case _:
+ return {"model": self.model, "input": marker, **extra}
+
+ def post(
+ self,
+ endpoint: str,
+ marker: str,
+ fields: Mapping[str, JsonValue] = MappingProxyType({}),
+ *,
+ stream: bool = False,
+ headers: Headers = (),
+ key: str | None = None,
+ ) -> httpx.Response:
+ payload: Final = self.body(endpoint, marker, fields) | ({"stream": True} if stream else {})
+ request_headers: Final = (
+ ("Authorization", f"Bearer {self.proxy.key if key is None else key}"),
+ *(headers.items() if isinstance(headers, Mapping) else headers),
+ )
+ with self.proxy.client.stream("POST", PATHS[endpoint], json=payload, headers=request_headers) as response:
+ response.read()
+ return response
+
+ def send(
+ self,
+ endpoint: str,
+ marker: str,
+ fields: Mapping[str, JsonValue] = MappingProxyType({}),
+ *,
+ stream: bool = False,
+ headers: Headers = (),
+ key: str | None = None,
+ ) -> httpx.Response:
+ response: Final = self.post(endpoint, marker, fields, stream=stream, headers=headers, key=key)
+ assert response.status_code == 200, response.text
+ return response
+
+ def upstream_bodies(self, marker: str) -> tuple[dict[str, JsonValue], ...]:
+ bodies: Final = tuple(JSON_OBJECT.validate_json(request.body) for request in self.upstream.drain())
+ return tuple(body for body in bodies if _user_markers(body)[-1:] == (marker,))
+
+
+@dataclass(frozen=True, slots=True)
+class Rigs:
+ flag_off: Rig
+ flag_on: Rig
+ sink: Collector
+
+
+def _spend_rows(call_id: str) -> list[dict[str, JsonValue]]:
+ return read_rows(
+ 'SELECT session_id, request_id, metadata::text AS metadata_text FROM "LiteLLM_SpendLogs" '
+ "WHERE litellm_call_id=%s",
+ (call_id,),
+ )
+
+
+def _spend_row(response: httpx.Response) -> dict[str, JsonValue]:
+ call_id: Final = response.headers["x-litellm-call-id"]
+ rows: Final = eventually(lambda: _spend_rows(call_id), lambda values: len(values) >= 1, seconds=70)
+ assert len(rows) == 1, rows
+ return rows[0]
+
+
+def _session(response: httpx.Response) -> JsonValue:
+ return _spend_row(response)["session_id"]
+
+
+def _model(scenario: Scenario, upstream: Wire) -> str:
+ return scenario.model(
+ model="hosted_vllm/gpt-4o-mini",
+ api_base=f"{upstream.url}/v1",
+ use_chat_completions_api=True,
+ input_cost_per_token=0.001,
+ output_cost_per_token=0.002,
+ )
+
+
+def _rig_config(
+ directory: Path, name: str, litellm_settings: Mapping[str, JsonValue], extra: Mapping[str, JsonValue]
+) -> Path:
+ shared: Final = object_value(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
+ config: Final = {
+ **shared,
+ "general_settings": {**object_value(shared["general_settings"]), "store_prompts_in_spend_logs": True},
+ "litellm_settings": {**object_value(shared["litellm_settings"]), **litellm_settings},
+ **extra,
+ }
+ path: Final = directory / f"{name}.yaml"
+ path.write_text(yaml.safe_dump(config))
+ return path
+
+
+@pytest.fixture(scope="module")
+def rigs(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rigs]:
+ accepted: Final[deque[Request]] = deque() # mutable-ok: sink thread appends each accepted OTLP batch
+ guard: Final = threading.Lock()
+
+ def sink(request: Request) -> Reply:
+ with guard:
+ accepted.append(request)
+ return Reply()
+
+ directory: Final = tmp_path_factory.mktemp("caller-session")
+ with ExitStack() as stack:
+ upstream: Final = stack.enter_context(wire_server(_upstream))
+ collector: Final = Collector(stack.enter_context(wire_server(sink)), accepted, guard)
+ shared: Final = stack.enter_context(gateway_from_environment())
+ flag_off: Final = stack.enter_context(
+ owned_proxy_process(shared, directory, {}, config=_rig_config(directory, "flag-off", {}, {}), workers=2)
+ )
+ flag_on: Final = stack.enter_context(
+ owned_proxy_process(
+ shared,
+ directory,
+ {"LITELLM_OTEL_V2": "1", "OTEL_BSP_SCHEDULE_DELAY": "300"},
+ config=_rig_config(
+ directory,
+ "flag-on",
+ {"request_correlation_in_logs": True, "callbacks": ["otel"]},
+ {
+ "callback_settings": {
+ "otel": {"exporter": "http/json", "endpoint": collector.wire.url, "mapper_names": ["genai"]}
+ }
+ },
+ ),
+ workers=2,
+ )
+ )
+ off_scenario: Final = stack.enter_context(flag_off.gateway.scenario())
+ on_scenario: Final = stack.enter_context(flag_on.gateway.scenario())
+ yield Rigs(
+ Rig(flag_off.gateway, off_scenario, _model(off_scenario, upstream), upstream, flag_off),
+ Rig(flag_on.gateway, on_scenario, _model(on_scenario, upstream), upstream, flag_on),
+ collector,
+ )
+
+
+@pytest.mark.parametrize(("endpoint", "stream"), ENDPOINT_STREAMS)
+def test_spend_log_session_is_the_body_session_beside_a_trace_header(rigs: Rigs, endpoint: str, stream: bool) -> None:
+ marker: Final = _marker()
+ session: Final = _token("convo")
+ response: Final = rigs.flag_on.send(
+ endpoint,
+ marker,
+ stream=stream,
+ headers={"x-litellm-trace-id": _token("trace")},
+ fields={"litellm_session_id": session},
+ )
+
+ row: Final = _spend_row(response)
+ upstream: Final = rigs.flag_on.upstream_bodies(marker)
+
+ assert row["session_id"] == session, row
+ assert CALLER_SESSION_KEY not in string_value(row["metadata_text"]), row
+ assert len(upstream) == 1, upstream
+ assert CALLER_SESSION_KEY not in json.dumps(upstream[0]), upstream
+
+
+def test_sdk_clients_record_the_body_session_beside_a_trace_header(rigs: Rigs) -> None:
+ rig: Final = rigs.flag_on
+ base_url: Final = str(rig.proxy.client.base_url)
+ traces: Final = tuple(_token("trace") for _ in range(3))
+ sessions: Final = tuple(_token("convo") for _ in range(3))
+ markers: Final = tuple(_marker() for _ in range(3))
+ chat_client: Final = openai.OpenAI(base_url=base_url + "/v1", api_key=rig.proxy.key, max_retries=0)
+ messages_client: Final = anthropic.Anthropic(base_url=base_url, api_key=rig.proxy.key, max_retries=0)
+
+ with chat_client.chat.completions.with_streaming_response.create(
+ model=rig.model,
+ messages=[{"role": "user", "content": markers[0]}],
+ stream=True,
+ extra_headers={"x-litellm-trace-id": traces[0]},
+ extra_body={"litellm_session_id": sessions[0]},
+ ) as chat_stream:
+ chat_call: Final = chat_stream.headers["x-litellm-call-id"]
+ chat_text: Final = "".join(
+ chunk.choices[0].delta.content or "" for chunk in chat_stream.parse() if chunk.choices
+ )
+ message: Final = messages_client.messages.with_raw_response.create(
+ model=rig.model,
+ max_tokens=16,
+ messages=[{"role": "user", "content": markers[1]}],
+ extra_headers={"x-litellm-trace-id": traces[1]},
+ extra_body={"litellm_session_id": sessions[1]},
+ )
+ with chat_client.responses.with_streaming_response.create(
+ model=rig.model,
+ input=markers[2],
+ stream=True,
+ extra_headers={"x-litellm-trace-id": traces[2]},
+ extra_body={"litellm_session_id": sessions[2]},
+ ) as responses_stream:
+ responses_call: Final = responses_stream.headers["x-litellm-call-id"]
+ completed: Final = tuple(event for event in responses_stream.parse() if event.type == "response.completed")
+
+ first_block: Final = message.parse().content[0]
+ assert chat_text == f"answer {markers[0]}", chat_text
+ assert isinstance(first_block, anthropic.types.TextBlock) and first_block.text == f"answer {markers[1]}", message
+ assert len(completed) == 1, completed
+ calls: Final = (chat_call, message.headers["x-litellm-call-id"], responses_call)
+ landed: Final = tuple(
+ eventually(lambda call=call: _spend_rows(call), lambda values: len(values) == 1, seconds=70)[0]["session_id"]
+ for call in calls
+ )
+ assert landed == sessions, (landed, calls)
+
+
+def test_conversation_span_keeps_the_trace_header_and_no_internal_key(rigs: Rigs) -> None:
+ trace: Final = _token("trace")
+ response: Final = rigs.flag_on.send(
+ "chat", _marker(), headers={"x-litellm-trace-id": trace}, fields={"litellm_session_id": _token("convo")}
+ )
+ response_id: Final = string_value(JSON_OBJECT.validate_json(response.content)["id"])
+
+ spans: Final = eventually(lambda: rigs.sink.span_attributes(response_id), lambda found: len(found) == 1, seconds=30)
+
+ assert spans[0]["gen_ai.conversation.id"] == {"stringValue": trace}, spans
+ assert CALLER_SESSION_KEY not in json.dumps(spans[0]), spans
+
+
+def test_a_continued_chains_conversation_spans_keep_each_turns_trace_header(rigs: Rigs) -> None:
+ traces: Final = (_token("trace"), _token("trace"))
+ first: Final = _turn(rigs.flag_on, _marker(), headers={"x-litellm-trace-id": traces[0]})
+ _turn(
+ rigs.flag_on,
+ _marker(),
+ first,
+ headers={"x-litellm-trace-id": traces[1]},
+ fields={"litellm_session_id": _token("convo")},
+ )
+
+ conversations: Final = eventually(
+ lambda: rigs.sink.conversation_ids(frozenset(traces)), lambda found: len(found) == 2, seconds=30
+ )
+
+ assert sorted(conversations) == sorted(traces), conversations
+
+
+@pytest.mark.parametrize("endpoint", list(PATHS))
+def test_flag_off_keeps_the_trace_id_beside_a_body_session(rigs: Rigs, endpoint: str) -> None:
+ trace: Final = _token("trace")
+
+ response: Final = rigs.flag_off.send(
+ endpoint, _marker(), headers={"x-litellm-trace-id": trace}, fields={"litellm_session_id": _token("convo")}
+ )
+
+ assert _session(response) == trace
+
+
+@pytest.mark.parametrize("endpoint", list(PATHS))
+def test_response_cache_hit_records_the_body_session_beside_a_trace_header(rigs: Rigs, endpoint: str) -> None:
+ marker: Final = _marker()
+ sessions: Final = (_token("convo"), _token("convo"))
+ responses: Final = tuple(
+ rigs.flag_on.send(
+ endpoint, marker, headers={"x-litellm-trace-id": _token("trace")}, fields={"litellm_session_id": session}
+ )
+ for session in sessions
+ )
+
+ assert tuple(_session(response) for response in responses) == sessions
+ assert len(rigs.flag_on.upstream_bodies(marker)) == 1
+
+
+def test_metadata_session_beside_a_new_trace_per_turn_groups_both_turns(rigs: Rigs) -> None:
+ session: Final = _token("convo")
+ responses: Final = tuple(
+ rigs.flag_on.send(
+ "chat",
+ _marker(),
+ headers={"x-litellm-trace-id": _token("trace")},
+ fields={"metadata": {"session_id": session}},
+ )
+ for _ in range(2)
+ )
+
+ assert tuple(_session(response) for response in responses) == (session, session)
+ now: Final = datetime.now(timezone.utc)
+ page: Final = rigs.flag_on.proxy.get(
+ "/spend/logs/ui",
+ params={
+ "session_id": session,
+ "start_date": (now - timedelta(days=1)).strftime("%Y-%m-%d %H:%M:%S"),
+ "end_date": (now + timedelta(days=1)).strftime("%Y-%m-%d %H:%M:%S"),
+ },
+ )
+ listed: Final = sorted(string_value(object_value(row)["request_id"]) for row in list_value(page["data"]))
+ assert listed == sorted(
+ string_value(JSON_OBJECT.validate_json(response.content)["id"]) for response in responses
+ ), page
+
+
+def _shape(shape: str, trace: str, session: str, header: str) -> tuple[Headers, dict[str, JsonValue], str]:
+ traced: Final = {"x-litellm-trace-id": trace}
+ match shape:
+ case "session-header-wins":
+ return {"x-litellm-session-id": header}, {"litellm_session_id": session}, header
+ case "trace-header-only":
+ return traced, {}, trace
+ case "body-session-only":
+ return {}, {"litellm_session_id": session}, session
+ case "int-session":
+ return traced, {"litellm_session_id": 7}, trace
+ case "list-session":
+ return traced, {"litellm_session_id": [session]}, trace
+ case "empty-session":
+ return traced, {"litellm_session_id": ""}, trace
+ case "null-session":
+ return traced, {"litellm_session_id": None}, trace
+ case "nul-session":
+ return traced, {"litellm_session_id": session + "\x00x"}, trace
+ case "two-trace-headers":
+ return (
+ (("x-litellm-trace-id", trace), ("x-litellm-trace-id", header)),
+ {"litellm_session_id": session},
+ session,
+ )
+ case _:
+ return traced, {"litellm_session_id": session + "s" * 5120}, session + "s" * 5120
+
+
+@pytest.mark.parametrize(
+ "shape",
+ [
+ "session-header-wins",
+ "trace-header-only",
+ "body-session-only",
+ "int-session",
+ "list-session",
+ "empty-session",
+ "null-session",
+ "nul-session",
+ "two-trace-headers",
+ "5kb-session",
+ ],
+)
+def test_spend_log_session_for_each_session_shape(rigs: Rigs, shape: str) -> None:
+ headers, fields, expected = _shape(shape, _token("trace"), _token("convo"), _token("header"))
+ neighbor_trace: Final = _token("neighbor")
+
+ response: Final = rigs.flag_on.send("chat", _marker(), fields, headers=headers)
+ neighbor: Final = rigs.flag_on.send("chat", _marker(), headers={"x-litellm-trace-id": neighbor_trace})
+
+ assert _session(response) == expected
+ assert _session(neighbor) == neighbor_trace
+
+
+@pytest.mark.parametrize("endpoint", ["chat", "messages"])
+def test_client_sent_internal_session_key_is_ignored(rigs: Rigs, endpoint: str) -> None:
+ forged: Final = _token("forged")
+ trace: Final = _token("trace")
+ forged_metadata: Final[dict[str, JsonValue]] = {
+ "metadata" if endpoint == "chat" else "litellm_metadata": {CALLER_SESSION_KEY: forged}
+ }
+
+ untraced: Final = rigs.flag_on.send(endpoint, _marker(), forged_metadata)
+ traced: Final = rigs.flag_on.send(endpoint, _marker(), forged_metadata, headers={"x-litellm-trace-id": trace})
+
+ assert _session(untraced) not in (forged, None)
+ assert _session(traced) == trace
+
+
+def test_identical_requests_each_log_once_with_the_body_session(rigs: Rigs) -> None:
+ session: Final = _token("convo")
+ trace: Final = _token("trace")
+ marker: Final = _marker()
+
+ responses: Final = tuple(
+ rigs.flag_on.send("chat", marker, headers={"x-litellm-trace-id": trace}, fields={"litellm_session_id": session})
+ for _ in range(3)
+ )
+
+ assert len({response.headers["x-litellm-call-id"] for response in responses}) == 3
+ assert tuple(_session(response) for response in responses) == (session, session, session)
+
+
+def _pass_through_route(rig: Rig, target: str) -> str:
+ path: Final = f"/caller-session-{uuid.uuid4().hex}"
+ created: Final = rig.proxy.post("/config/pass_through_endpoint", {"path": path, "target": target, "auth": True})
+ endpoint_id: Final = string_value(object_value(list_value(created["endpoints"])[0])["id"])
+ rig.scenario.cleanups.callback(
+ lambda: rig.proxy.request("DELETE", "/config/pass_through_endpoint", params={"endpoint_id": endpoint_id})
+ )
+ return path
+
+
+@pytest.mark.parametrize("field", ["litellm_session_id", "metadata"])
+def test_pass_through_spend_log_session_is_the_body_session_beside_a_trace_header(rigs: Rigs, field: str) -> None:
+ marker: Final = _marker()
+ session: Final = _token("convo")
+ rig: Final = rigs.flag_on
+ path: Final = _pass_through_route(rig, f"{rig.upstream.url}/v1/chat/completions")
+ key: Final = rig.scenario.key(allowed_passthrough_routes=[path])
+ fields: Final[dict[str, JsonValue]] = (
+ {"litellm_session_id": session} if field == "litellm_session_id" else {"metadata": {"session_id": session}}
+ )
+
+ response: Final = rig.proxy.request(
+ "POST",
+ path,
+ rig.body("chat", marker, fields),
+ key=key,
+ headers={"x-litellm-trace-id": _token("trace")},
+ )
+
+ rows: Final = eventually(
+ lambda: read_rows(
+ 'SELECT session_id FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (sha256(key.encode()).hexdigest(),)
+ ),
+ lambda found: len(found) >= 1,
+ seconds=70,
+ )
+ assert response.status_code == 200, response.text
+ assert f"answer {marker}" in response.text
+ assert len(rig.upstream_bodies(marker)) == 1
+ assert rows == [{"session_id": session}]
+
+
+def _turn(
+ rig: Rig,
+ marker: str,
+ previous: str | None = None,
+ *,
+ key: str | None = None,
+ headers: Headers = (),
+ fields: Mapping[str, JsonValue] = MappingProxyType({}),
+) -> str:
+ chained: Final[dict[str, JsonValue]] = {"previous_response_id": previous} if previous else {}
+ response: Final = rig.send("responses", marker, {**chained, **fields}, headers=headers, key=key)
+ _spend_row(response)
+ return string_value(JSON_OBJECT.validate_json(response.content)["id"])
+
+
+def _turn_fields(shape: str, session: str) -> tuple[dict[str, JsonValue], dict[str, JsonValue]]:
+ match shape:
+ case "session-every-turn":
+ return {"litellm_session_id": session}, {"litellm_session_id": session}
+ case "session-from-turn-2":
+ return {}, {"litellm_session_id": session}
+ case _:
+ return {"litellm_session_id": _token("convo")}, {"litellm_session_id": _token("convo")}
+
+
+@pytest.mark.parametrize(
+ ("correlated", "shape", "replays_turn_1"),
+ [
+ (True, "session-every-turn", True),
+ (True, "session-from-turn-2", True),
+ (True, "new-session-per-turn", True),
+ (False, "session-from-turn-2", False),
+ ],
+)
+def test_bridged_chain_history(rigs: Rigs, correlated: bool, shape: str, replays_turn_1: bool) -> None:
+ rig: Final = rigs.flag_on if correlated else rigs.flag_off
+ first_fields, second_fields = _turn_fields(shape, _token("convo"))
+ markers: Final = (_marker(), _marker(), _marker())
+ first: Final = _turn(rig, markers[0], headers={"x-litellm-trace-id": _token("trace")}, fields=first_fields)
+ second: Final = _turn(rig, markers[1], first, headers={"x-litellm-trace-id": _token("trace")}, fields=second_fields)
+
+ _turn(rig, markers[2], second)
+
+ replayed: Final = rig.upstream_bodies(markers[2])
+ assert len(replayed) == 1, replayed
+ assert _user_markers(replayed[0]) == (markers if replays_turn_1 else markers[1:])
+
+
+def test_bridged_chains_sharing_a_metadata_session_replay_each_others_turns(rigs: Rigs) -> None:
+ shared: Final[dict[str, JsonValue]] = {"metadata": {"session_id": _token("convo")}}
+ markers: Final = (_marker(), _marker(), _marker())
+ first: Final = _turn(rigs.flag_on, markers[0], fields=shared)
+ _turn(rigs.flag_on, markers[1], fields=shared)
+
+ _turn(rigs.flag_on, markers[2], first)
+
+ replayed: Final = rigs.flag_on.upstream_bodies(markers[2])
+ assert len(replayed) == 1, replayed
+ assert _user_markers(replayed[0]) == markers
+
+
+@dataclass(frozen=True, slots=True)
+class Owners:
+ first: str
+ second: str
+
+
+def _owners(scenario: Scenario, model: str, shape: str) -> Owners:
+ user: Final = _token("user")
+ match shape:
+ case "other-user":
+ return Owners(
+ scenario.key(user_id=user, models=[model]), scenario.key(user_id=_token("user"), models=[model])
+ )
+ case "same-user":
+ return Owners(scenario.key(user_id=user, models=[model]), scenario.key(user_id=user, models=[model]))
+ case "same-team":
+ team: Final = scenario.team()
+ return Owners(scenario.key(team_id=team, models=[model]), scenario.key(team_id=team, models=[model]))
+ case _:
+ return Owners(scenario.key(models=[model]), scenario.key(models=[model]))
+
+
+@pytest.mark.parametrize(
+ ("shape", "replays_first_owner"),
+ [("other-user", False), ("same-user", True), ("same-team", True), ("ownerless-keys", False)],
+)
+def test_bridged_history_follows_the_previous_responses_owner(
+ rigs: Rigs, shape: str, replays_first_owner: bool
+) -> None:
+ rig: Final = rigs.flag_off
+ owners: Final = _owners(rig.scenario, rig.model, shape)
+ session: Final = _token("convo")
+ markers: Final = (_marker(), _marker(), _marker())
+ _turn(rig, markers[0], key=owners.first, fields={"litellm_session_id": session})
+ second: Final = _turn(rig, markers[1], key=owners.second, fields={"litellm_session_id": session})
+
+ _turn(rig, markers[2], second, key=owners.second)
+
+ replayed: Final = rig.upstream_bodies(markers[2])
+ assert len(replayed) == 1, replayed
+ assert _user_markers(replayed[0]) == (markers if replays_first_owner else markers[1:])
+
+
+BURST: Final = ENDPOINT_STREAMS * 4
+
+
+def test_concurrent_burst_across_endpoints_logs_each_session_once(rigs: Rigs) -> None:
+ sessions: Final = tuple(_token("convo") for _ in BURST)
+
+ def one(index: int) -> httpx.Response:
+ endpoint, stream = BURST[index]
+ return rigs.flag_on.send(
+ endpoint,
+ _marker(),
+ stream=stream,
+ headers={"x-litellm-trace-id": _token("trace")},
+ fields={"litellm_session_id": sessions[index]},
+ )
+
+ with ThreadPoolExecutor(max_workers=12) as pool:
+ responses: Final = tuple(pool.map(one, range(len(BURST))))
+
+ assert tuple(_session(response) for response in responses) == sessions
+
+
+def _open_upstream_connections(pid: int, upstream: Wire) -> int:
+ port: Final = urlsplit(upstream.url).port
+ return sum(
+ 1
+ for connection in psutil.Process(pid).net_connections(kind="tcp")
+ if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port
+ )
+
+
+def _post_or_none(rig: Rig, session: str) -> httpx.Response | None:
+ try:
+ return rig.post(
+ "chat", _marker(), headers={"x-litellm-trace-id": _token("trace")}, fields={"litellm_session_id": session}
+ )
+ except httpx.HTTPError:
+ return None
+
+
+def test_killing_one_worker_mid_burst_keeps_logging_the_body_session(rigs: Rigs) -> None:
+ release: Final = threading.Event()
+ arrived: Final[SimpleQueue[bytes]] = SimpleQueue()
+
+ def held(request: Request) -> Reply:
+ arrived.put(request.body)
+ assert release.wait(timeout=60), "the burst was never released"
+ return _upstream(request)
+
+ process: Final = rigs.flag_on.process
+ sessions: Final = tuple(_token("convo") for _ in range(20))
+ after_session: Final = _token("convo")
+ with wire_server(held) as wire:
+ rig: Final = replace(rigs.flag_on, model=_model(rigs.flag_on.scenario, wire), upstream=wire)
+ workers: Final = eventually(process.worker_pids, lambda found: len(found) == 2, seconds=30)
+ with ThreadPoolExecutor(max_workers=len(sessions)) as pool:
+ pending: Final = tuple(pool.submit(_post_or_none, rig, session) for session in sessions)
+ eventually(arrived.qsize, lambda size: size == len(sessions), seconds=60)
+ held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire) for pid in workers})
+ victim, survivor = sorted(workers, key=held_by.__getitem__)
+ with process.replacing_worker(victim) as kill:
+ kill.freeze()
+ kill()
+ release.set()
+ results: Final = tuple(future.result() for future in pending)
+ after: Final = rig.send(
+ "chat",
+ _marker(),
+ headers={"x-litellm-trace-id": _token("trace")},
+ fields={"litellm_session_id": after_session},
+ )
+
+ served: Final = tuple(
+ (response, session)
+ for response, session in zip(results, sessions, strict=True)
+ if response is not None and response.status_code == 200
+ )
+ assert sum(held_by.values()) == len(sessions), held_by
+ assert len(served) == held_by[survivor] >= len(sessions) // 2, (held_by, len(served))
+ assert _session(after) == after_session
+ assert tuple(_session(response) for response, _ in served) == tuple(session for _, session in served)
diff --git a/tests/integration/streaming/test_mid_stream_fallback_continuation.py b/tests/integration/streaming/test_mid_stream_fallback_continuation.py
new file mode 100644
index 00000000000..8ec52f26dde
--- /dev/null
+++ b/tests/integration/streaming/test_mid_stream_fallback_continuation.py
@@ -0,0 +1,749 @@
+from __future__ import annotations
+
+import functools
+import hashlib
+import itertools
+import json
+import math
+import operator
+import threading
+import uuid
+from collections.abc import Callable, Generator, Iterable, Iterator, Mapping, Sequence
+from contextlib import ExitStack, contextmanager
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Final
+
+import openai
+import pytest
+import yaml
+from integration._support.client import JSON_OBJECT, Gateway, eventually, gateway_from_environment, object_value
+from integration._support.database import read_rows, scratch_database
+from integration._support.process import graceful_stop_seconds, owned_proxy
+from integration._support.wire import Reply, Request, Wire, wire_server
+from openai.types.chat import ChatCompletionChunk
+from pydantic import JsonValue, TypeAdapter
+
+pytestmark: Final = pytest.mark.timeout(4 * graceful_stop_seconds() + 120)
+
+SETTING: Final = "enable_mid_stream_fallback_continuation"
+PARTIAL: Final = ("Alpha beta", " gamma")
+JOINED_PARTIAL: Final = "".join(PARTIAL)
+RELAY_TEXT: Final = " delta"
+BACKUP_TEXT: Final = " omega."
+PLAIN_TEXT: Final = "plain answer"
+TOOL_PREAMBLE: Final = "Let me check"
+BACKUP_USAGE: Final = (40, 6)
+MAX_TOKENS: Final = 200
+PRIMARY_PRICES: Final = (0.001, 0.002)
+BACKUP_PRICES: Final = (0.003, 0.005)
+SHARED_BACKEND: Final = "openai/gpt-4o-mini"
+SEPARATE_BACKEND: Final = "openai/gpt-4o"
+WEATHER_TOOL: Final[Mapping[str, JsonValue]] = {
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
+ },
+}
+TOOL_CALL_DELTA: Final[Mapping[str, JsonValue]] = {
+ "index": 0,
+ "id": "call_1",
+ "type": "function",
+ "function": {"name": "get_weather", "arguments": ""},
+}
+NOTHING_SEEN: Final[tuple[ChatCompletionChunk, ...]] = ()
+ALLOW_EVERYTHING: Final = "def apply_guardrail(inputs, request_data, input_type):\n return allow()\n"
+EVENT_BATCH: Final = TypeAdapter(list[dict[str, JsonValue]])
+
+
+@functools.cache
+def cut_gate(marker: str, hop: str, attempt: int) -> threading.Event:
+ return threading.Event()
+
+
+@functools.cache
+def attempts(marker: str, hop: str) -> Iterator[int]:
+ return itertools.count()
+
+
+def frame(identity: str, delta: Mapping[str, JsonValue], finish: str | None = None) -> bytes:
+ body: Final = {
+ "id": identity,
+ "object": "chat.completion.chunk",
+ "created": 1700000000,
+ "model": "gpt-4o-mini",
+ "choices": [{"index": 0, "delta": dict(delta), "finish_reason": finish}],
+ }
+ return b"data: " + json.dumps(body).encode() + b"\n\n"
+
+
+def complete_stream(identity: str, text: str, usage: tuple[int, int]) -> Reply:
+ usage_frame: Final[dict[str, JsonValue]] = {
+ "id": identity,
+ "object": "chat.completion.chunk",
+ "created": 1700000000,
+ "model": "gpt-4o-mini",
+ "choices": [],
+ "usage": {"prompt_tokens": usage[0], "completion_tokens": usage[1], "total_tokens": sum(usage)},
+ }
+ return Reply(
+ content_type="text/event-stream",
+ chunks=(
+ frame(identity, {"role": "assistant", "content": text}),
+ frame(identity, {}, finish="stop"),
+ b"data: " + json.dumps(usage_frame).encode() + b"\n\n",
+ b"data: [DONE]\n\n",
+ ),
+ )
+
+
+def cut_after_caller_sees(marker: str, hop: str, *frames: bytes) -> Reply:
+ return Reply(
+ content_type="text/event-stream",
+ chunks=(b"".join(frames), b"data: never sent\n\n"),
+ abort_after=1,
+ gate_after_first=cut_gate(marker, hop, next(attempts(marker, hop))),
+ )
+
+
+def first_user_text(request: Request) -> str:
+ messages: Final = JSON_OBJECT.validate_json(request.body)["messages"]
+ assert isinstance(messages, list), messages
+ first: Final = object_value(messages[0])["content"]
+ assert isinstance(first, str), first
+ return first
+
+
+def marker_of(request: Request) -> str:
+ return first_user_text(request).split()[-1]
+
+
+def is_chat(request: Request) -> bool:
+ return request.target.endswith("/chat/completions")
+
+
+def model_list() -> Reply:
+ return Reply(body=b'{"object": "list", "data": []}')
+
+
+def primary(request: Request) -> Reply:
+ if not is_chat(request):
+ return model_list()
+ scenario: Final = first_user_text(request).split()[0]
+ marker: Final = marker_of(request)
+ identity: Final = f"chatcmpl-{marker}-primary"
+ if scenario == "reject":
+ return Reply(status=500, body=b'{"error": {"message": "primary is down", "type": "server_error"}}')
+ if scenario == "tool-cut":
+ return cut_after_caller_sees(
+ marker,
+ "primary",
+ frame(identity, {"role": "assistant", "content": TOOL_PREAMBLE}),
+ frame(identity, {"tool_calls": [dict(TOOL_CALL_DELTA)]}),
+ )
+ return cut_after_caller_sees(
+ marker,
+ "primary",
+ frame(identity, {"role": "assistant", "content": PARTIAL[0]}),
+ frame(identity, {"content": PARTIAL[1]}),
+ )
+
+
+def relay(request: Request) -> Reply:
+ if not is_chat(request):
+ return model_list()
+ marker: Final = marker_of(request)
+ return cut_after_caller_sees(
+ marker, "relay", frame(f"chatcmpl-{marker}-relay", {"role": "assistant", "content": RELAY_TEXT})
+ )
+
+
+def backup(request: Request) -> Reply:
+ if not is_chat(request):
+ return model_list()
+ return complete_stream(f"chatcmpl-{marker_of(request)}-backup", BACKUP_TEXT, BACKUP_USAGE)
+
+
+def no_prefill(request: Request) -> Reply:
+ if not is_chat(request):
+ return model_list()
+ return complete_stream(f"chatcmpl-{marker_of(request)}-plain", PLAIN_TEXT, (9, 2))
+
+
+def accept_events(request: Request) -> Reply:
+ return Reply()
+
+
+def registered_guardrail(name: str, mode: str) -> dict[str, JsonValue]:
+ return {
+ "guardrail_name": name,
+ "litellm_params": {
+ "guardrail": "custom_code",
+ "mode": mode,
+ "default_on": False,
+ "custom_code": ALLOW_EVERYTHING,
+ },
+ }
+
+
+@dataclass(frozen=True, slots=True)
+class Rig:
+ gateway: Gateway
+ database_url: str
+ primary: Wire
+ relay: Wire
+ backup: Wire
+ no_prefill: Wire
+ events: Wire
+
+
+def joined_text(chunks: Iterable[ChatCompletionChunk]) -> str:
+ return "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices)
+
+
+def has_tool_call(chunks: Iterable[ChatCompletionChunk]) -> bool:
+ return any(choice.delta.tool_calls for chunk in chunks for choice in chunk.choices)
+
+
+@dataclass(frozen=True, slots=True)
+class Release:
+ gate: threading.Event
+ ready: Callable[[tuple[ChatCompletionChunk, ...]], bool]
+
+
+def after_text(marker: str, hop: str, text: str, attempt: int = 0) -> Release:
+ return Release(cut_gate(marker, hop, attempt), lambda seen: joined_text(seen) == text)
+
+
+@dataclass(frozen=True, slots=True)
+class Streamed:
+ chunks: tuple[ChatCompletionChunk, ...]
+ error: openai.APIError | None
+
+ @property
+ def text(self) -> str:
+ return joined_text(self.chunks)
+
+ @property
+ def ids(self) -> frozenset[str]:
+ return frozenset(chunk.id for chunk in self.chunks)
+
+ @property
+ def role_deltas(self) -> int:
+ return sum(1 for chunk in self.chunks for choice in chunk.choices if choice.delta.role is not None)
+
+ @property
+ def finish_reasons(self) -> tuple[str, ...]:
+ return tuple(choice.finish_reason for chunk in self.chunks for choice in chunk.choices if choice.finish_reason)
+
+ @property
+ def usage(self) -> tuple[int, int]:
+ reported: Final = tuple(chunk.usage for chunk in self.chunks if chunk.usage is not None)
+ assert reported, self.chunks
+ return (reported[-1].prompt_tokens, reported[-1].completion_tokens)
+
+
+def deployment(
+ name: str,
+ wire: Wire,
+ prices: tuple[float, float],
+ *,
+ backend: str = SHARED_BACKEND,
+ prefill: bool | None = None,
+) -> dict[str, JsonValue]:
+ return {
+ "model_name": name,
+ "litellm_params": {
+ "model": backend,
+ "api_base": wire.url + "/v1",
+ "api_key": "synthetic-continuation-key",
+ "input_cost_per_token": prices[0],
+ "output_cost_per_token": prices[1],
+ },
+ **({"model_info": {"supports_assistant_prefill": prefill}} if prefill is not None else {}),
+ }
+
+
+def write_config(
+ directory: Path,
+ model_entries: Sequence[Mapping[str, JsonValue]],
+ fallbacks: Sequence[Mapping[str, JsonValue]],
+ router_settings: Mapping[str, JsonValue],
+ guardrails: Sequence[Mapping[str, JsonValue]],
+) -> Path:
+ base: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
+ config: Final = {
+ **base,
+ "model_list": [dict(entry) for entry in model_entries],
+ "litellm_settings": {
+ **object_value(base["litellm_settings"]),
+ "callbacks": ["generic_api"],
+ "DEFAULT_FLUSH_INTERVAL_SECONDS": 1,
+ },
+ "guardrails": [dict(entry) for entry in guardrails],
+ "router_settings": {
+ **object_value(base["router_settings"]),
+ "num_retries": 0,
+ "fallbacks": [dict(entry) for entry in fallbacks],
+ **router_settings,
+ },
+ }
+ path: Final = directory / f"continuation-{uuid.uuid4().hex[:8]}.yaml"
+ path.write_text(yaml.safe_dump(config))
+ return path
+
+
+@contextmanager
+def booted_rig(
+ directory: Path,
+ fallbacks: Sequence[Mapping[str, JsonValue]],
+ router_settings: Mapping[str, JsonValue],
+ workers: int,
+ guardrails: Sequence[Mapping[str, JsonValue]] = (),
+) -> Generator[Rig]:
+ with ExitStack() as stack:
+ shared: Final = stack.enter_context(gateway_from_environment())
+ database_url: Final = stack.enter_context(scratch_database())
+ wires: Final = tuple(
+ stack.enter_context(wire_server(respond)) for respond in (primary, relay, backup, no_prefill, accept_events)
+ )
+ primary_wire, relay_wire, backup_wire, no_prefill_wire, events_wire = wires
+ config: Final = write_config(
+ directory,
+ (
+ *(
+ deployment(name, primary_wire, PRIMARY_PRICES)
+ for name in ("primary", "primary-skip", "primary-none", "primary-relay", "primary-shared")
+ ),
+ deployment("relay", relay_wire, PRIMARY_PRICES, prefill=True),
+ deployment("backup", backup_wire, BACKUP_PRICES, prefill=True),
+ deployment("no-prefill", no_prefill_wire, PRIMARY_PRICES, backend=SEPARATE_BACKEND),
+ deployment("no-prefill-shared", no_prefill_wire, PRIMARY_PRICES),
+ ),
+ fallbacks,
+ router_settings,
+ guardrails,
+ )
+ gateway: Final = stack.enter_context(
+ owned_proxy(
+ shared,
+ directory,
+ {"DATABASE_URL": database_url, "GENERIC_LOGGER_ENDPOINT": events_wire.url},
+ config=config,
+ remove_environment=("DATABASE_URL_READ_REPLICA",),
+ workers=workers,
+ )
+ )
+ yield Rig(gateway, database_url, primary_wire, relay_wire, backup_wire, no_prefill_wire, events_wire)
+
+
+@pytest.fixture(scope="module")
+def configured(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
+ fallbacks: Final[tuple[Mapping[str, JsonValue], ...]] = (
+ {"primary": ["backup"]},
+ {"primary-skip": ["no-prefill", "backup"]},
+ {"primary-none": ["no-prefill"]},
+ {"primary-relay": ["relay"]},
+ {"relay": ["backup"]},
+ {"primary-shared": ["no-prefill-shared"]},
+ )
+ with booted_rig(
+ tmp_path_factory.mktemp("continuation-configured"),
+ fallbacks,
+ {SETTING: True},
+ 2,
+ (registered_guardrail("input-check", "pre_call"),),
+ ) as rig:
+ yield rig
+
+
+@pytest.fixture(scope="module")
+def output_guarded(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
+ with booted_rig(
+ tmp_path_factory.mktemp("continuation-output-guarded"),
+ ({"primary": ["backup"]},),
+ {SETTING: True},
+ 2,
+ (registered_guardrail("output-check", "post_call"),),
+ ) as rig:
+ yield rig
+
+
+@pytest.fixture(scope="module")
+def unset(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
+ with booted_rig(
+ tmp_path_factory.mktemp("continuation-unset"), ({"primary": ["no-prefill", "backup"]},), {}, 1
+ ) as rig:
+ yield rig
+
+
+def new_key(rig: Rig) -> str:
+ key: Final = rig.gateway.post("/key/generate", {})["key"]
+ assert isinstance(key, str), key
+ return key
+
+
+def appended(seen: tuple[ChatCompletionChunk, ...], chunk: ChatCompletionChunk) -> tuple[ChatCompletionChunk, ...]:
+ return (*seen, chunk)
+
+
+def observed(
+ response: Iterable[ChatCompletionChunk], releases: Sequence[Release]
+) -> Iterator[ChatCompletionChunk | openai.APIError]:
+ try:
+ for seen in itertools.accumulate(response, appended, initial=NOTHING_SEEN):
+ for release in releases:
+ if release.ready(seen):
+ release.gate.set()
+ yield from seen[-1:]
+ except openai.APIError as error:
+ yield error
+
+
+def stream(
+ rig: Rig,
+ key: str,
+ model: str,
+ prompt: str,
+ releases: Sequence[Release] = (),
+ extra: Mapping[str, JsonValue] | None = None,
+) -> Streamed:
+ base_url: Final = f"{rig.gateway.client.base_url}".rstrip("/") + "/v1"
+ try:
+ with openai.OpenAI(base_url=base_url, api_key=key, max_retries=0, timeout=60) as client:
+ response: Final = client.chat.completions.create(
+ model=model,
+ messages=[{"role": "user", "content": prompt}],
+ stream=True,
+ stream_options={"include_usage": True},
+ max_tokens=MAX_TOKENS,
+ extra_body=dict(extra or {}),
+ )
+ outcome: Final = tuple(observed(response, releases))
+ finally:
+ for release in releases:
+ release.gate.set()
+ errors: Final = tuple(item for item in outcome if isinstance(item, openai.APIError))
+ return Streamed(
+ tuple(item for item in outcome if isinstance(item, ChatCompletionChunk)), errors[0] if errors else None
+ )
+
+
+def chat_bodies(wire: Wire, marker: str) -> tuple[dict[str, JsonValue], ...]:
+ return tuple(
+ JSON_OBJECT.validate_json(request.body)
+ for request in wire.drain()
+ if is_chat(request) and marker_of(request) == marker
+ )
+
+
+def spend_rows(rig: Rig, key: str) -> list[dict[str, JsonValue]]:
+ return eventually(
+ lambda: read_rows(
+ 'SELECT request_id, status, spend, prompt_tokens, completion_tokens, metadata FROM "LiteLLM_SpendLogs" '
+ "WHERE api_key=%s",
+ (hashlib.sha256(key.encode()).hexdigest(),),
+ database_url=rig.database_url,
+ ),
+ lambda rows: len(rows) >= 1,
+ seconds=70,
+ )
+
+
+def amount(value: JsonValue) -> float:
+ assert isinstance(value, (int, float)), value
+ return float(value)
+
+
+def key_spend(rig: Rig, key: str) -> float:
+ return amount(object_value(rig.gateway.get("/key/info", {"key": key})["info"])["spend"])
+
+
+def close(observed_value: float, expected: float) -> bool:
+ return math.isclose(observed_value, expected, rel_tol=1e-9, abs_tol=1e-12)
+
+
+def cost(tokens: tuple[int, int], prices: tuple[float, float]) -> float:
+ return tokens[0] * prices[0] + tokens[1] * prices[1]
+
+
+def prefill_turn(text: str) -> dict[str, JsonValue]:
+ return {"role": "assistant", "content": text, "prefix": True}
+
+
+@dataclass(frozen=True, slots=True)
+class Billed:
+ rows: list[dict[str, JsonValue]]
+ partial: tuple[int, int]
+
+
+def billed_once(rig: Rig, key: str, received: Streamed) -> Billed:
+ rows: Final = spend_rows(rig, key)
+ assert [row["status"] for row in rows] == ["success"], rows
+ row_tokens: Final = (rows[0]["prompt_tokens"], rows[0]["completion_tokens"])
+ assert isinstance(row_tokens[0], int) and isinstance(row_tokens[1], int), rows
+ assert received.usage == row_tokens, (received.usage, row_tokens)
+ partial: Final = (row_tokens[0] - BACKUP_USAGE[0], row_tokens[1] - BACKUP_USAGE[1])
+ assert partial[0] > 0 and partial[1] > 0, (row_tokens, BACKUP_USAGE)
+ return Billed(rows, partial)
+
+
+def assert_spend_prices_every_token(rig: Rig, key: str, billed: Billed) -> None:
+ partial_cost: Final = cost(billed.partial, PRIMARY_PRICES)
+ breakdown: Final = object_value(object_value(billed.rows[0]["metadata"])["cost_breakdown"])
+ spend: Final = amount(billed.rows[0]["spend"])
+ reported_partial_cost: Final = amount(object_value(breakdown["additional_costs"])["mid_stream_partial_cost"])
+ assert close(spend, cost(BACKUP_USAGE, BACKUP_PRICES) + partial_cost) and close(
+ reported_partial_cost, partial_cost
+ ), f"partial={billed.partial} spend={spend} breakdown={json.dumps(breakdown)}"
+ eventually(lambda: key_spend(rig, key), lambda value: close(value, spend), seconds=70)
+ assert spend_rows(rig, key) == billed.rows
+
+
+def test_text_cut_continues_on_the_prefill_fallback_and_bills_both_calls_once(configured: Rig) -> None:
+ marker: Final = uuid.uuid4().hex
+ key: Final = new_key(configured)
+ prompt: Final = f"text-cut {marker}"
+ received: Final = stream(configured, key, "primary", prompt, (after_text(marker, "primary", JOINED_PARTIAL),))
+ assert received.error is None, received.error
+ assert received.text == JOINED_PARTIAL + BACKUP_TEXT, received.chunks
+ assert received.ids == {f"chatcmpl-{marker}-primary"}, received.ids
+ assert received.role_deltas == 1, received.chunks
+ assert received.finish_reasons == ("stop",), received.chunks
+ assert len(chat_bodies(configured.primary, marker)) == 1
+ continued: Final = chat_bodies(configured.backup, marker)
+ assert len(continued) == 1, continued
+ assert continued[0]["messages"] == [{"role": "user", "content": prompt}, prefill_turn(JOINED_PARTIAL)]
+ assert "_mid_stream_continuation" not in continued[0], continued[0]
+ billed: Final = billed_once(configured, key, received)
+ assert continued[0]["max_tokens"] == MAX_TOKENS - billed.partial[1], (continued[0], billed.partial)
+ assert_spend_prices_every_token(configured, key, billed)
+
+
+def test_spend_row_of_a_continued_stream_carries_the_id_the_caller_received(configured: Rig) -> None:
+ marker: Final = uuid.uuid4().hex
+ key: Final = new_key(configured)
+ received: Final = stream(
+ configured, key, "primary", f"text-cut {marker}", (after_text(marker, "primary", JOINED_PARTIAL),)
+ )
+ assert received.error is None, received.error
+ assert [row["request_id"] for row in spend_rows(configured, key)] == [received.chunks[0].id]
+
+
+def test_tool_call_before_the_cut_keeps_the_original_error(configured: Rig) -> None:
+ marker: Final = uuid.uuid4().hex
+ received: Final = stream(
+ configured,
+ new_key(configured),
+ "primary",
+ f"tool-cut {marker}",
+ (Release(cut_gate(marker, "primary", 0), has_tool_call),),
+ {"tools": [dict(WEATHER_TOOL)]},
+ )
+ assert received.error is not None, received.chunks
+ assert (received.text, has_tool_call(received.chunks)) == (TOOL_PREAMBLE, True), received.chunks
+ assert len(chat_bodies(configured.primary, marker)) == 1
+ assert chat_bodies(configured.backup, marker) == ()
+
+
+def test_continuation_skips_a_fallback_without_prefill_support(configured: Rig) -> None:
+ marker: Final = uuid.uuid4().hex
+ prompt: Final = f"text-cut {marker}"
+ received: Final = stream(
+ configured, new_key(configured), "primary-skip", prompt, (after_text(marker, "primary", JOINED_PARTIAL),)
+ )
+ assert received.error is None, received.error
+ assert received.text == JOINED_PARTIAL + BACKUP_TEXT, received.chunks
+ assert chat_bodies(configured.no_prefill, marker) == ()
+ assert [body["messages"] for body in chat_bodies(configured.backup, marker)] == [
+ [{"role": "user", "content": prompt}, prefill_turn(JOINED_PARTIAL)]
+ ]
+
+
+def test_no_prefill_capable_fallback_keeps_the_original_error(configured: Rig) -> None:
+ marker: Final = uuid.uuid4().hex
+ received: Final = stream(
+ configured,
+ new_key(configured),
+ "primary-none",
+ f"text-cut {marker}",
+ (after_text(marker, "primary", JOINED_PARTIAL),),
+ )
+ assert received.error is not None, received.chunks
+ assert "No deployments available" not in str(received.error), received.error
+ assert received.text == JOINED_PARTIAL, received.chunks
+ assert len(chat_bodies(configured.primary, marker)) == 1
+ assert chat_bodies(configured.no_prefill, marker) == ()
+
+
+def test_undeclared_fallback_sharing_a_backend_with_a_prefill_deployment_keeps_the_original_error(
+ configured: Rig,
+) -> None:
+ marker: Final = uuid.uuid4().hex
+ received: Final = stream(
+ configured,
+ new_key(configured),
+ "primary-shared",
+ f"text-cut {marker}",
+ (after_text(marker, "primary", JOINED_PARTIAL),),
+ )
+ assert (received.error is not None, received.text) == (True, JOINED_PARTIAL), received.chunks
+ assert chat_bodies(configured.no_prefill, marker) == ()
+
+
+def test_failure_before_content_still_falls_back_to_a_fallback_without_prefill_support(configured: Rig) -> None:
+ marker: Final = uuid.uuid4().hex
+ prompt: Final = f"reject {marker}"
+ received: Final = stream(configured, new_key(configured), "primary-none", prompt)
+ assert received.error is None, received.error
+ assert received.text == PLAIN_TEXT, received.chunks
+ assert [body["messages"] for body in chat_bodies(configured.no_prefill, marker)] == [
+ [{"role": "user", "content": prompt}]
+ ]
+
+
+def stream_with_two_cuts(rig: Rig, key: str, marker: str) -> Streamed:
+ return stream(
+ rig,
+ key,
+ "primary-relay",
+ f"text-cut {marker}",
+ (after_text(marker, "primary", JOINED_PARTIAL), after_text(marker, "relay", JOINED_PARTIAL + RELAY_TEXT)),
+ )
+
+
+def test_second_cut_continues_with_both_partials_in_one_prefill(configured: Rig) -> None:
+ marker: Final = uuid.uuid4().hex
+ key: Final = new_key(configured)
+ received: Final = stream_with_two_cuts(configured, key, marker)
+ assert received.error is None, received.error
+ assert received.text == JOINED_PARTIAL + RELAY_TEXT + BACKUP_TEXT, received.chunks
+ assert received.role_deltas == 1, received.chunks
+ user_turn: Final = {"role": "user", "content": f"text-cut {marker}"}
+ relayed: Final = chat_bodies(configured.relay, marker)
+ continued: Final = chat_bodies(configured.backup, marker)
+ assert [body["messages"] for body in relayed] == [[user_turn, prefill_turn(JOINED_PARTIAL)]]
+ assert [body["messages"] for body in continued] == [[user_turn, prefill_turn(JOINED_PARTIAL + RELAY_TEXT)]]
+ partial: Final = billed_once(configured, key, received).partial
+ relay_ceiling: Final = relayed[0]["max_tokens"]
+ assert isinstance(relay_ceiling, int), relayed[0]
+ assert 0 < MAX_TOKENS - relay_ceiling < partial[1], (relayed[0], partial)
+ assert continued[0]["max_tokens"] == MAX_TOKENS - partial[1], (continued[0], partial)
+
+
+def test_second_cut_bills_every_call_on_one_row(configured: Rig) -> None:
+ key: Final = new_key(configured)
+ received: Final = stream_with_two_cuts(configured, key, uuid.uuid4().hex)
+ assert received.error is None, received.error
+ assert_spend_prices_every_token(configured, key, billed_once(configured, key, received))
+
+
+def test_repeating_a_continued_request_never_returns_only_the_continuation(configured: Rig) -> None:
+ marker: Final = uuid.uuid4().hex
+ key: Final = new_key(configured)
+ prompt: Final = f"text-cut {marker}"
+ answers: Final = tuple(
+ stream(configured, key, "primary", prompt, (after_text(marker, "primary", JOINED_PARTIAL, attempt),))
+ for attempt in range(2)
+ )
+ assert tuple((answer.error, answer.text) for answer in answers) == ((None, JOINED_PARTIAL + BACKUP_TEXT),) * 2
+
+
+def test_without_the_setting_a_text_cut_keeps_the_original_error(unset: Rig) -> None:
+ marker: Final = uuid.uuid4().hex
+ received: Final = stream(
+ unset, new_key(unset), "primary", f"text-cut {marker}", (after_text(marker, "primary", JOINED_PARTIAL),)
+ )
+ assert received.error is not None, received.chunks
+ assert received.text == JOINED_PARTIAL, received.chunks
+ assert len(chat_bodies(unset.primary, marker)) == 1
+ assert chat_bodies(unset.no_prefill, marker) == ()
+ assert chat_bodies(unset.backup, marker) == ()
+
+
+def router_setting(rig: Rig) -> JsonValue:
+ return object_value(rig.gateway.get("/router/settings")["current_values"]).get(SETTING)
+
+
+def test_turning_the_setting_on_at_runtime_continues_on_a_prefill_fallback(unset: Rig) -> None:
+ rig: Final = unset
+ rig.gateway.post("/config/update", {"router_settings": {SETTING: True}})
+ try:
+ assert router_setting(rig) is True
+ marker: Final = uuid.uuid4().hex
+ prompt: Final = f"text-cut {marker}"
+ received: Final = stream(rig, new_key(rig), "primary", prompt, (after_text(marker, "primary", JOINED_PARTIAL),))
+ assert received.error is None, received.error
+ assert received.text == JOINED_PARTIAL + BACKUP_TEXT, received.chunks
+ assert chat_bodies(rig.no_prefill, marker) == ()
+ assert [body["messages"] for body in chat_bodies(rig.backup, marker)] == [
+ [{"role": "user", "content": prompt}, prefill_turn(JOINED_PARTIAL)]
+ ]
+ finally:
+ rig.gateway.post("/config/update", {"router_settings": {SETTING: False}})
+
+
+def test_a_registered_output_guardrail_keeps_the_original_error(output_guarded: Rig) -> None:
+ marker: Final = uuid.uuid4().hex
+ received: Final = stream(
+ output_guarded,
+ new_key(output_guarded),
+ "primary",
+ f"text-cut {marker}",
+ (after_text(marker, "primary", JOINED_PARTIAL),),
+ )
+ assert (received.error is not None, received.text) == (True, JOINED_PARTIAL), received.chunks
+ assert len(chat_bodies(output_guarded.primary, marker)) == 1
+ assert chat_bodies(output_guarded.backup, marker) == ()
+
+
+def test_a_request_fallback_bringing_its_own_messages_keeps_the_original_error(configured: Rig) -> None:
+ marker: Final = uuid.uuid4().hex
+ received: Final = stream(
+ configured,
+ new_key(configured),
+ "primary",
+ f"text-cut {marker}",
+ (after_text(marker, "primary", JOINED_PARTIAL),),
+ {"fallbacks": [{"model": "backup", "messages": [{"role": "user", "content": f"start over {marker}"}]}]},
+ )
+ assert (received.error is not None, received.text) == (True, JOINED_PARTIAL), received.chunks
+ assert len(chat_bodies(configured.primary, marker)) == 1
+ assert chat_bodies(configured.backup, marker) == ()
+
+
+def carries_tag(event: Mapping[str, JsonValue], marker: str) -> bool:
+ tags: Final = event.get("request_tags")
+ return isinstance(tags, list) and marker in tags
+
+
+def tagged_events(batches: Iterable[Request], marker: str) -> tuple[dict[str, JsonValue], ...]:
+ events: Final = itertools.chain.from_iterable(EVENT_BATCH.validate_json(batch.body) for batch in batches)
+ return tuple(event for event in events if carries_tag(event, marker))
+
+
+def logged_events(sink: Wire, marker: str) -> Iterator[tuple[dict[str, JsonValue], ...]]:
+ return itertools.accumulate((tagged_events(sink.drain(), marker) for _ in itertools.count()), operator.add)
+
+
+def test_the_cut_call_logs_its_failure_under_its_own_deployment_and_price(configured: Rig) -> None:
+ marker: Final = uuid.uuid4().hex
+ received: Final = stream(
+ configured,
+ new_key(configured),
+ "primary",
+ f"text-cut {marker}",
+ (after_text(marker, "primary", JOINED_PARTIAL),),
+ {"metadata": {"tags": [marker]}},
+ )
+ assert received.error is None, received.error
+ polls: Final = logged_events(configured.events, marker)
+ events: Final = eventually(lambda: next(polls), lambda logged: len(logged) >= 2, seconds=30)
+ assert sorted(str(event["status"]) for event in events) == ["failure", "success"], events
+ failure: Final = next(event for event in events if event["status"] == "failure")
+ assert failure["model_group"] == "primary", failure
+ assert str(failure["api_base"]).startswith(configured.primary.url), failure
+ tokens: Final = (failure["prompt_tokens"], failure["completion_tokens"])
+ assert isinstance(tokens[0], int) and isinstance(tokens[1], int) and tokens[0] > 0 and tokens[1] > 0, failure
+ assert close(amount(failure["response_cost"]), cost((tokens[0], tokens[1]), PRIMARY_PRICES)), failure
diff --git a/tests/integration/translation/chat_completions/reasoning/__init__.py b/tests/integration/translation/chat_completions/reasoning/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/integration/translation/chat_completions/reasoning/test_chat_completions_reasoning_gemini.py b/tests/integration/translation/chat_completions/reasoning/test_chat_completions_reasoning_gemini.py
new file mode 100644
index 00000000000..194e6972453
--- /dev/null
+++ b/tests/integration/translation/chat_completions/reasoning/test_chat_completions_reasoning_gemini.py
@@ -0,0 +1,1041 @@
+import json
+from collections.abc import Mapping
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import replace
+from pathlib import Path
+from typing import Final
+from unittest.mock import ANY
+from uuid import uuid4
+
+import pytest
+from integration._support.client import Gateway, eventually, list_value, object_value, string_value
+from integration._support.process import graceful_stop_seconds
+from integration._support.provider import SharedProvider
+from integration._support.wire import Reply
+from integration.caching.response_cache_case import CACHE_KEY, redis_store
+from integration.translation.case import SSEFrame, TranslationTestCase
+from integration.translation.chat_completions.bases.gemini import GEMINI_3_8_FLASH_TEST_CASE
+from integration.translation.chat_completions.bases.vertex_ai import (
+ GEMINI_3_8_FLASH_TEST_CASE as VERTEX_GEMINI_3_8_FLASH_TEST_CASE,
+)
+from integration.translation.gemini_thoughts import (
+ LONG_SUMMARY,
+ gemini_thoughts_proxy,
+ sent_thinking_config,
+ thoughts_off,
+ with_thinking_config,
+)
+from integration.translation.messages.bases.gemini import (
+ GEMINI_3_8_FLASH_TEST_CASE as MESSAGES_GEMINI_3_8_FLASH_TEST_CASE,
+)
+from integration.translation.responses.bases.gemini import (
+ GEMINI_3_8_FLASH_TEST_CASE as RESPONSES_GEMINI_3_8_FLASH_TEST_CASE,
+)
+from integration.translation.runner import (
+ assert_provider_request,
+ assert_translation,
+ encode_sse_frame,
+ parse_sse_frames,
+)
+from pydantic import JsonValue
+
+__all__ = ("thoughts_off",)
+
+pytestmark = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
+
+GEMINI_3_8_FLASH_THINKING_SHOWN_THOUGHT: Final = "**My Inspection Strategy: Validation First**\n\nOkay, here's how I'm approaching this project review. The instruction specifically mentions considering zero, negative, and non-integer quantities. To efficiently address that, I'm going to start with `shop/validation.py`. My reasoning is simple: the validation logic is the first line of defense. Understanding how the system *expects* to handle these edge cases will inform everything else. It's the logical place to begin my investigation to understand the code. I'll use the `read_file` tool to examine that file. I am **not** changing any files. I will call the `chat_completions` client for this project review. I am now proceeding with reading `shop/validation.py`.\n\n\n"
+
+GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE: Final = (
+ "EowMCokMAWkUfRPuOEs7DW4IXhoH/SNh9qYezwNSqhXT2tKoI1mcXDG+yx7R+12sMmzvjHh2OSghAlsvmR6welOzIM1LMw6c"
+ "QrSX2IPRo/xbtUvDLlo6aqufzRtini5EF0yxQdmyn9r9xznzDyuEf4eqrbK0KQAArEVM13rm3CTDIjyDRQPawLALfUh/e61/"
+ "XVcWrr1PZLR5Jj46ImyrobGGLmJuoytm3ezTEihTb2b0q3/vHOW9cO/gJzMcbCI8koexQLuBB1KLSHUCaZZQYPj4jLx/g0de"
+ "2l3H/DYR7Cdroh0P/R3mewkhu+3D2ELnHKt85InrdQOqdXp/XhICsivPgsIEWHMRr6kQ3ZLIUJvJbJ7zo7/4bYnHOm+JqmoX"
+ "E+VMobI2YZ+ff9GOYg4vsJRTeF7ff9kZHYBWRM9XvTsZMbeCZHRbCOLNmSSJPfn2ToDOgMxU1u/GqHYmVSV97ybT3Y0mxwtf"
+ "2WcVlPUCkFX0GW3GP5IHVkEISo0TznJSI3/yYb+5UmQFMuuIedDdhKa74eU1/K59VuTnd3DzVX/lfhpFbGLDkOSurRsDsViN"
+ "69zq+6VbLi+5ARi6SJqUU+1U02QaSIQVp8kxtrhEIsPedjCIYHQQNxR8hFPMhnEssR51THNELWH6sVccxpJIbFfGjl0TZS9d"
+ "rfYL2c+IBX56lrH8IX2gz1iMOsESKwTyAECmqJIL7d/HcVdTk404Hd9SYycB+2MP2GN2XA7ivoO/zfkkFcf+PGmXPed8dHlW"
+ "bAxrQcOAz89p8kPzb1bou5QUT3feeRG06aJL31gX7BSYj19mvLeo0AmMzx+Wb91NFlwUvlJYlhJo4NfZWg/ymcGclgM8r7iE"
+ "MmWZa8wul8g4hbJOg3GE9NV70zvKqEpQrJLTa9jHca0DBvj8N3C+XvpeEg+Dys0cxE7Ak0VVYbvu9r0OHKSWtC7NDYrJsF2n"
+ "VFmdavSxpyPEqPaAdWdiewDH0nhP5GCR/ciikokWFQMF5iClWRfL72qEbFx3cQkctqEf+kK2CevIWKPf8bvkrKVDdPqbjk0G"
+ "9NWuud+YgNxl4dSCkAWoMrAFlxXozEqMmanMAtZk/7nPeOVFbYEL/xOpSCfxaEKS1zEmMwTb/UNVAogsNY9YdBSFYVNSrS8F"
+ "N/IrHBvsP66+YEFuR46Di6cWuuZQdSOfB2uQe1eSC8IgrQ6XUhQHWeXKUXfbcb/1abPOi5PUMYyLBQ7r52UKCEfOBc1j0oNL"
+ "wYzI/vWAz74gq8NblQhHB2VlrInJ9rRVY1MehInqUH0kcG6ysQDQLiLJiS30TV0bZh05I86RfcjaN1twbRir82YchRjXG8/1"
+ "6iFRJ54Sk7oVzbdnbFOugDea8qIblFmgUTq0+SRLGBFcDMvAqYLNcR2ZO7o0xinMG7IuNRwoSMahnmreWYr9saVJAScd3All"
+ "i6RQuxjGIdxG0s0UNuoWiSNIf1Z2ngoJrvZHTc0iZK+A2zv3CBqWAaS8WGwSfcItnlc6hMB0bTWasoCelmRjfv+1jQnKzKMj"
+ "DzZlO6skMkv5x9GzLrAvpYoulLYa/wp9+exw/Ll7MxSB22aeAGfai9C7Vj310y8nApsG3QxaoUUO4jg1eCjrrluBLgoVY8cE"
+ "/O95IAUa2rFzVrOWbHuob7TEI/8+BwieGeyeHdPGn4fNeKYMkRjFeVJ1GA/2+iCHCjCbSNm4sdzh9188/Es3q4acRfQUiqwR"
+ "iQzW2E1fy9mXRDY9y5fdXuzjHtuCkH8MqR0cioN1p4OGZzPi6xXFaqV576NLjySBpfpQEYeXRA3j5NVwTwbeYKnGOot4bEtJ"
+ "JQfelwkXLXeAfxzCTD0/vGtdhilzEK8wSYM9a+8llr5D4XwtEsu/68YbXJ2jwzsS3rE/Frulrf6paz9fS45ZXj98+OT5EpFO"
+ "o5XHfTXKb5Brzr0tKfPSB2Mx6fXbYP3lZRwfJIaUq8Alwdo7g2jnHQUhCNBB8hkb/IkQrnhsAGS6f0ZEjwf6ErVOa8DkfizN"
+ "y4WTdS1D3gOUJxOcszm8J5b085tsfFxEJVgA0VjI60DBY0owWn2I"
+)
+
+GEMINI_3_8_FLASH_THOUGHTS_OFF_RESPONSE: Final[Mapping[str, JsonValue]] = {
+ "candidates": [
+ {
+ "content": {
+ "parts": [
+ {
+ "text": "To understand how zero, negative, and non-integer quantities (such as floats or invalid types) are currently handled and where the checkout rejection is enforced, we need to inspect `shop/validation.py` next.",
+ },
+ {
+ "functionCall": {
+ "name": "read_file",
+ "args": {
+ "path": "shop/validation.py",
+ "end_line": 50,
+ "start_line": 1,
+ },
+ "id": "call_330625",
+ },
+ "thoughtSignature": GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE,
+ },
+ ],
+ "role": "model",
+ },
+ "finishReason": "STOP",
+ "index": 0,
+ "finishMessage": "Model generated function call(s).",
+ },
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 358,
+ "candidatesTokenCount": 77,
+ "totalTokenCount": 793,
+ "promptTokensDetails": [
+ {
+ "modality": "TEXT",
+ "tokenCount": 358,
+ },
+ ],
+ "thoughtsTokenCount": 358,
+ "serviceTier": "standard",
+ },
+ "modelVersion": "gemini-3.8-flash",
+ "responseId": "xy3RapaNEbT4-sAPhMrt6AU",
+}
+
+GEMINI_3_8_FLASH_THINKING_OMITTED_THOUGHT_SIGNATURE: Final = (
+ "EscECsQEAWkUfROlWYT6hX0DWOKSCg4MSuLzOV7u43EgJDn0dfS2R8v/MVGGcNolGA9rteb/cV3dyCSQ7D2EO2CB5whDkblj"
+ "fGAAw6NBhqnidYNQiXbsg7oqynxS116ZQnk302S4KzvAbIvo4gpRbahPVClXKkP8PCQWA5eRQ/e9oQeHr97gbBB7+PLsl0NK"
+ "m7YPV6lP3fAXx/r3EADbdNL1diHylKd6tCWUKRea/x2QIOBQZdY5jHSU73z/pnPOotVwggEiasDLSLA+uUR1KllgGc8WKZ4i"
+ "uRM3L1gDr1QZ4Rht3quiF8AwfopWE5m2ON+FD78aHUDG/lgDs0vXnJluzri9LpPaexfbpvwMPtAx8Bxrq1/dCeekR2VnVx1c"
+ "qZ7TSoWMqk4x98mPpyUFz1hSBRnXaOGvQl5Wk7Iart49tH2LoIBYoV+PkCgNZ2MOi2uzc7OlN/8gvOeBM05xbGAplsLiuKd4"
+ "jkgv7Nz8Jby3v/hxtQT9+Ow3mxwSaiIfEkuaxkPiK1cX1nwT92gFzpmhVsrGqpiCUdAfos39Wd2/rcrFbtl5iMAVcHG5JGjw"
+ "8AjubFqiWDMqWPC58FfsX9TLC2BMUbPf2xi3hrw+1QWDtSuNl9s+dofkCogXFdLA9lAFc0oRj3MVhzYUhn5eixApJdggDsJF"
+ "iJaQtRihs4dstbXX9QXYFnx22Ql3x9TWJibiH8jz4fJVOf0wreMWZ2xVXBiw2xKdy5dB7xRDnrQD/sdcpTpCWURYXVCXglJ9"
+ "DD7e7/OCNMLCPw=="
+)
+
+GEMINI_3_8_FLASH_THINKING_OMITTED_RESPONSE: Final[Mapping[str, JsonValue]] = {
+ "candidates": [
+ {
+ "content": {
+ "parts": [
+ {
+ "text": "To understand how the checkout validates quantities\u2014including zero, negative, and non-integer quantities\u2014I need to inspect `shop/validation.py` next.",
+ },
+ {
+ "functionCall": {
+ "name": "read_file",
+ "args": {
+ "start_line": 1,
+ "path": "shop/validation.py",
+ "end_line": 40,
+ },
+ "id": "call_745953",
+ },
+ "thoughtSignature": GEMINI_3_8_FLASH_THINKING_OMITTED_THOUGHT_SIGNATURE,
+ },
+ ],
+ "role": "model",
+ },
+ "finishReason": "STOP",
+ "index": 0,
+ "finishMessage": "Model generated function call(s).",
+ },
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 358,
+ "candidatesTokenCount": 65,
+ "totalTokenCount": 531,
+ "promptTokensDetails": [
+ {
+ "modality": "TEXT",
+ "tokenCount": 358,
+ },
+ ],
+ "thoughtsTokenCount": 108,
+ "serviceTier": "standard",
+ },
+ "modelVersion": "gemini-3.8-flash",
+ "responseId": "zS3RasS-Lp2-jMcPoay98As",
+}
+
+GEMINI_3_8_FLASH_THINKING_SHOWN_THOUGHT_SIGNATURE: Final = (
+ "EqUECqIEAWkUfRNZEW6qD3+V/1LAc1riV0fvqkg4Bldd7YYDfln+ZHNg6Lbw2T31UZHTum0qYvrAOor848+zSkFG7qBDfdZ9"
+ "69FGl/dQGq2llr1yrv23blnbwC7tO0uevr8225EZ/4Pv+hbfP0teFcSZsasE4LQangtxYkewghc6cH1Acz293SIz1cqF/3CM"
+ "Ta+FtJ5x106oauWy79+z1VbV9TNyHpNtsrUy3ADR2PFNBjXKf2TvUWygzR4w2RJg581+dfx0sI97O3UlZrhSKRe+p/6siUP7"
+ "qnysURAia7k+ew5f05Nuo1aDwD2G20G2s9v7L0PxnCCAay5L4YTub5LXxvUGBXMhSHVb9MmO8+2rNPqDtPWh0bOQdKbkWAV/"
+ "dMckMLTKANDbnzBjtCK1HbN7Rw0sIBSnBVy59z7h+0ierHFgA8Vntfp/nvAy0wDpuZu09agdB3T6FO2HCvWCFG2Y+0B4lGJV"
+ "ErOknvn50gtn9Tpn1QSGv4VDRrrcv6nCOyohoL9CsBeO1Bs+RItyZSAoTA3bVB5IjcF2kKK27+/+VRJjN0cO4Cm5RD/ETb4+"
+ "kI101AFWJ3CYwKEu91SAVTtXLewQYd5LBiJ3Btt68m8d5uUp0a1kAn2cH6g2C7TOfoJ7BgU5+vU81mS8hk9QOFGb9RYUXs8Y"
+ "cJl21l6WfMMZri5bMI9qCAFn+YZdBBeIbGq+JJcIgjr/R2tULmU584+5beoEydkS"
+)
+
+GEMINI_3_8_FLASH_THINKING_SHOWN_RESPONSE: Final[Mapping[str, JsonValue]] = {
+ "candidates": [
+ {
+ "content": {
+ "parts": [
+ {
+ "text": GEMINI_3_8_FLASH_THINKING_SHOWN_THOUGHT,
+ "thought": True,
+ },
+ {
+ "text": "I need to inspect `shop/validation.py` next to examine how zero, negative, and non-integer quantities are currently handled in `validate_quantity`.",
+ },
+ {
+ "functionCall": {
+ "name": "read_file",
+ "args": {
+ "end_line": 40,
+ "start_line": 1,
+ "path": "shop/validation.py",
+ },
+ "id": "call_386884",
+ },
+ "thoughtSignature": GEMINI_3_8_FLASH_THINKING_SHOWN_THOUGHT_SIGNATURE,
+ },
+ ],
+ "role": "model",
+ },
+ "finishReason": "STOP",
+ "index": 0,
+ "finishMessage": "Model generated function call(s).",
+ },
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 358,
+ "candidatesTokenCount": 66,
+ "totalTokenCount": 529,
+ "promptTokensDetails": [
+ {
+ "modality": "TEXT",
+ "tokenCount": 358,
+ },
+ ],
+ "thoughtsTokenCount": 105,
+ "serviceTier": "standard",
+ },
+ "modelVersion": "gemini-3.8-flash",
+ "responseId": "0C3RavXPIJWX9MoPtorosAs",
+}
+
+GEMINI_2_5_FLASH_THOUGHTS_OFF_RESPONSE: Final[Mapping[str, JsonValue]] = {
+ "candidates": [
+ {
+ "content": {
+ "parts": [
+ {
+ "text": "I will read the `shop/validation.py` file to understand the quantity validation logic.",
+ },
+ {
+ "functionCall": {
+ "name": "read_file",
+ "args": {
+ "start_line": 1,
+ "path": "shop/validation.py",
+ "end_line": 20,
+ },
+ },
+ },
+ ],
+ "role": "model",
+ },
+ "finishReason": "STOP",
+ "index": 0,
+ "finishMessage": "Model generated function call(s).",
+ },
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 334,
+ "candidatesTokenCount": 53,
+ "totalTokenCount": 387,
+ "promptTokensDetails": [
+ {
+ "modality": "TEXT",
+ "tokenCount": 334,
+ },
+ ],
+ "serviceTier": "standard",
+ },
+ "modelVersion": "gemini-2.5-flash",
+ "responseId": "0y3Rar76LO_K-8YPh6Kd0Ak",
+}
+
+VERTEX_GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE: Final = (
+ "AY89a18W+yVUlI0b2XsXQujlNKztnSpKXMf0+1JFiprxfp7kc0bG2g9WoPITuNkCqkuBaO5SkZaKDrKQTlJFJKj+H6SnETLk"
+ "JrR3SAzp2GTjwM2QRsJjI5ky54oylVT93Jh51PWCMp6vl8tZstMJrois8b1i4/eTOyKaVprypLW8DOIImqN3Cbkea8nEDe9r"
+ "XsdpdvsR7i+NvD0eiCW+Z4ZIflx5NVENbyEezjyJCL3lg3bV2Mfb0ytsACHXuqgGQEcH9HBKQKGkF5n7mA0GNLw6wAqlk8FE"
+ "JKttP6NgS+HzS/U+xaqxci6B50l0hIMw3pTOHu8Gg5fRKxf/dzvks+s709/ODCGBDrhMzpf/ilnPAMrB67cQrHvQOhnI2pGk"
+ "cvkZv+ZGwAQxjJHs/r0aXDGU2mwN9RY+/rSEv8y4MWd5c1KIs7xa1GytDBaaG8Zwh3kZzaSOYjVKIk7cy31mPJ01jCNdW9Pg"
+ "1p2r2iwbzJrBHo5Ku3s1yK4rADeY1WOU63sM6evVkWJEtu6KPxAqE9JkQKQoAPOFtoP89RtiTSl0nZZaO0stbqG7G+AkE2iO"
+ "5WkKCB1o55R7EEZDalDgFPput8khce5hPrhIGikN3yo11sHsAA2PAuDcHo4L3eWaYPExFiZtB+gVTbBwx42ST3UuBB9lrHtk"
+ "GCrp7tlOBHJJY+RlXj1G0TJsSxUZDSEswCB+gjG+1tdEchKGRRKOBUjWGARCozBsIEuBOR0F1ugSCLYn7CumS2bLwmMhKY7+"
+ "mfKXEgF2U84MkmwDEOBSTubNbRhy0qNhVPcLT8pJvGvtTvLny1C90TBTHErVzOKb3Z095fFpUyW7MmFQj5J9rX3A160W0SNQ"
+ "RKXf2Yr6ikyVhLBSFNvQ00n81iXlyg7bp++ChbbCJBFCUcrGF5WnPLk3+isixodP+G0b+pNtCNo8xMkrvqGqL5AMlmGJ9Qm4"
+ "a9S6QgDRaw5AzcV/W90a7/BcaIeqkIly3Z9M0+d3Eep4igYj6AK0JzX1BhchMogq9g/dquTTZZk4edVKflBmQcwg4qGGvioI"
+ "X/ae4tKJVqB2U0Ed+A6t3JZ4Ggdpopdnuy29XvdEyTmUg+fF2sTCgAdK07hueKYO2pkhqj86mSy88iHV0/secWQs0ZuHtcSW"
+ "D9lAG57yvpMf6RIwccLTyVmM4jcO0A2OB0sTnNIVKYYtswPtUqbrPMTR5saBfEfnULipI521kpb3BseODBKar+G1FISLrb0Q"
+ "OjMhNnCWCoCSVDk1FbDmtEvSjfiUzeJtFYoqM8PIXyGiBZnsmjGaOw3QMNvLfkoy4gqE6nvFVgwIk3+J6Z1GdFecCPSSCatz"
+ "LRwwQ8B5ZdlJ6+xjTDWIuyaOByD5BsB98e0YHyuLSYP17w6+0pSGXf3M6AdBJXK7hcL/vu3sGzftYARFHgRY1PQsYCqW8RVy"
+ "uJL/GwEAT4EgzkDDWufizfUHJEAUVBpQKkxZW7eTbyHyPS/Uiu5a74f5L8lYyep1pAdqt3iK4lvmQ3T9h3DAQIQ15T8PLV78"
+ "FjzJvZDQ2uMryvoRS1MGEcK6dxc4VvcWalQTjyczE4jQIxEEiN3U5Yo5CaM6FFGDnjIJBiy0fUD16P8e4ZaEaRFQCY4QZ0sA"
+ "bhwSHWhnU7QWvA/CVnqOVni5TpRsY1ShtEaQXKTslUuz96lB6XwEeaDkPqsPdw+mWY9QmmLqXnXERnTq13rQY1s="
+)
+
+VERTEX_GEMINI_3_8_FLASH_THOUGHTS_OFF_RESPONSE: Final[Mapping[str, JsonValue]] = {
+ "candidates": [
+ {
+ "content": {
+ "role": "model",
+ "parts": [
+ {
+ "text": "To understand how zero, negative, and non-integer quantities are handled during checkout validation, we need to inspect `shop/validation.py` next.",
+ },
+ {
+ "functionCall": {
+ "name": "read_file",
+ "args": {
+ "start_line": 1,
+ "path": "shop/validation.py",
+ "end_line": 50,
+ },
+ "id": "call_1272842",
+ },
+ "thoughtSignature": VERTEX_GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE,
+ },
+ ],
+ },
+ "finishReason": "STOP",
+ },
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 246,
+ "candidatesTokenCount": 64,
+ "totalTokenCount": 583,
+ "trafficType": "ON_DEMAND",
+ "promptTokensDetails": [
+ {
+ "modality": "TEXT",
+ "tokenCount": 246,
+ },
+ ],
+ "candidatesTokensDetails": [
+ {
+ "modality": "TEXT",
+ "tokenCount": 64,
+ },
+ ],
+ "thoughtsTokenCount": 273,
+ },
+ "modelVersion": "gemini-3.8-flash",
+ "createTime": "2026-10-15T20:02:49.484984Z",
+ "responseId": "aTHRavjMHYvi6tkPt6OroAo",
+}
+
+GEMINI_3_8_FLASH_THOUGHTS_OFF_STREAM_THOUGHT_SIGNATURE: Final = (
+ "ErkLCrYLAWkUfRO3TXF/xF4wcqKCnAFmIW9MEkIMZeaFFG/BmstoEFBw+0v0jari9LOu/TkcHKNFvewXanY1T+6rm512Fsvm"
+ "3cbxU5C7oP8kGPeEHIBOs2BIeocBFx3ZXSmIn25pOLd2YOqDaTzu/De/QX2JUkno3r2Ozbvuy0wvS5EsD3816rYYnaYc7+4E"
+ "iZRQhf5r6notr70qcHJ3DuKIP8ghDEz0gNSN1bAUa8N5Hvjh7ITbHLoSkeKO0pOjXaSWfQ86xv9F2zhpMZfW+muxcm1+g9tT"
+ "5QkyJxMtn+cnHi/w4MbQy+9tbJWdkkLo2H6z81pBsmjvIZAQnjjF0zP8aL91/+S9Po8HNV0LsrGs4P4kUgRR33I6UPMCVnHf"
+ "B7IxmVSEIGeFQ8gg1mJRDj7yZJFBgf61ya97bVWAhiAm6t/EHdnMi9xIPFEYSrY6ed8ZGUnf05x/GRMVOfrmeCePZIHXpftk"
+ "UfI2W4+qIM1YGR8+GD7KOtQEkcsrjRGHvhmcv1GW3NlLVpZ4LWIqX6ZtmZrOjDFvgnCg/jYdzDHL5j2vQiiSlcpsHJ4UZWRj"
+ "NmqUyUY0OZC77erTqensEeWJggUfyizdp3InLZtWM/C5B6rhLGku3LgGwAxdSf5aY25i0f+Ayl/xE/nnCkXS81D64z/+Xqdn"
+ "W4jsZn2tg8TGb11T6GncBK82aE/T9zVUHT7mbUtEljdinxNc/NS+9mWEkhtLSOAQqQmGF24oxMFlCcejVhGZz5QnFpjnu2Uk"
+ "COv2AZ7S9iYJQa/xT1LyQ18IibvAUWQPRBpjwRnOf0wi0beuE9Q/h0Y2itFjCt+mz5JlFRyQRJcpEOupOsfF47ULZCLeG0Uz"
+ "Af8JQzHMUWL3ESlOs30cLQQBJwtHW3QD3ZA1uwMCZnJcpBvIr7GhRdF/VhKZ7C8aYZLMOk32M1i+iGe5lYOM2WuoEzbprTGz"
+ "aVQeQBWEIdyT+RhXjpvLrrsWOIWI/OAw1Y8VEtouLt246yDjDrVuKUdJfXPz9VfnJ/7dKmsYWAaOCBnQ6edZ9HY8vzEKyuma"
+ "mcaJWSrwuq1mIouQ1Pw4fvazi6JE2xx17hEuOpS5sLD8E89V9uUbkcDne4XFoHNF9mm3nVgTplNFqo5J4kSiEZJ6qJveiMyF"
+ "sw3F7SbuU0f95AE/HVwMfDhBbnT3j0R5s2zejxLW6mWuASodcn+J9WGbRfMDLZV+aozCPy+rrXrIfeVQw8nLfjQbwY/iQLFM"
+ "rQCM4m9F5MLrIM0jGKePUeLRmkhJ92LiqEmDMkOZct09jVaZpSW6dRO+ANGfxqJ0y8KppBA7jgWy/PDKMHfwQFWh3Goxc86m"
+ "otxTUbF2i15mhhPGN2Z5xsRe1lsYRsFV7LNIkqLXcUVQ7uQGmlzi7wCGiMasvxDgzTLdTGAJ6FsaWzr+GAHH/RFSyRA9a6Qr"
+ "/Z/h9tm5sx5l6hTQrYspK3qK+cYI7BJCuiYYiIMWEAV6qb5SNhMZ88UlTayryuoolQ7pBco59MBV+P6oAoWAXFh69D4yDAe8"
+ "L2iNa6UTLt2BSMpAgqE5HdrsifGPfwOwNrwtfv/t4eXIBpBVgqC4csZSrN4zUsP6fH3mPZ4oPgwkDFaLu/x1S/PdxKvF0mYg"
+ "e3ORBogDHyO00QzHPP/9sl2PEgzIQmQv+0tkmpmAeTUfU8FMlRouFsgbrO2wUx1P8KoaFO4ZToQHeNaoqhI2yZ+GB9EquuBp"
+ "Xpu46qaU1Wrf3w7vtI0aLUPDKnudfZ3b/me3LVm1cgutomL5o33J86U0v62TwJz8xnLB0881oHQVXX4VfaPV4BVk74QCMcqK"
+ "/SRt3GpK9V1OGQh47ubs5MUFBnqauOnOKmttTxerI5QNi/t5vSwIvdy2MBcXnCXPJSjt3MpZjSZo34PQHK+1iuR1BCbmmHQk"
+ "v2UCtcYUZfnJ4FuKNXaZ/AFDT7VxbdnOaS8HVA=="
+)
+
+GEMINI_3_8_FLASH_THOUGHTS_OFF_STREAM_FRAMES: Final = (
+ SSEFrame(
+ data={
+ "candidates": [
+ {
+ "content": {
+ "parts": [
+ {
+ "text": "To understand how zero, negative, and non-integer quantities are handled during checkout, I need to inspect `shop/validation.py",
+ },
+ ],
+ "role": "model",
+ },
+ "index": 0,
+ },
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 358,
+ "candidatesTokenCount": 27,
+ "totalTokenCount": 673,
+ "promptTokensDetails": [
+ {
+ "modality": "TEXT",
+ "tokenCount": 358,
+ },
+ ],
+ "thoughtsTokenCount": 288,
+ "serviceTier": "standard",
+ },
+ "modelVersion": "gemini-3.8-flash",
+ "responseId": "yi3RaufbH6bXjMcPt5mF2Q0",
+ }
+ ),
+ SSEFrame(
+ data={
+ "candidates": [
+ {
+ "content": {
+ "parts": [
+ {
+ "text": "` next.",
+ },
+ ],
+ "role": "model",
+ },
+ "index": 0,
+ },
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 358,
+ "candidatesTokenCount": 30,
+ "totalTokenCount": 676,
+ "promptTokensDetails": [
+ {
+ "modality": "TEXT",
+ "tokenCount": 358,
+ },
+ ],
+ "thoughtsTokenCount": 288,
+ "serviceTier": "standard",
+ },
+ "modelVersion": "gemini-3.8-flash",
+ "responseId": "yi3RaufbH6bXjMcPt5mF2Q0",
+ }
+ ),
+ SSEFrame(
+ data={
+ "candidates": [
+ {
+ "content": {
+ "parts": [
+ {
+ "functionCall": {
+ "name": "read_file",
+ "args": {
+ "end_line": 50,
+ "path": "shop/validation.py",
+ "start_line": 1,
+ },
+ "id": "call_317677",
+ },
+ "thoughtSignature": GEMINI_3_8_FLASH_THOUGHTS_OFF_STREAM_THOUGHT_SIGNATURE,
+ },
+ ],
+ "role": "model",
+ },
+ "index": 0,
+ },
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 358,
+ "candidatesTokenCount": 63,
+ "totalTokenCount": 709,
+ "promptTokensDetails": [
+ {
+ "modality": "TEXT",
+ "tokenCount": 358,
+ },
+ ],
+ "thoughtsTokenCount": 288,
+ "serviceTier": "standard",
+ },
+ "modelVersion": "gemini-3.8-flash",
+ "responseId": "yi3RaufbH6bXjMcPt5mF2Q0",
+ }
+ ),
+ SSEFrame(
+ data={
+ "candidates": [
+ {
+ "content": {
+ "parts": [
+ {
+ "text": "",
+ },
+ ],
+ "role": "model",
+ },
+ "finishReason": "STOP",
+ "index": 0,
+ },
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 358,
+ "candidatesTokenCount": 63,
+ "totalTokenCount": 709,
+ "promptTokensDetails": [
+ {
+ "modality": "TEXT",
+ "tokenCount": 358,
+ },
+ ],
+ "thoughtsTokenCount": 288,
+ "serviceTier": "standard",
+ },
+ "modelVersion": "gemini-3.8-flash",
+ "responseId": "yi3RaufbH6bXjMcPt5mF2Q0",
+ }
+ ),
+)
+
+GEMINI_2_5_FLASH_REQUEST: Final = with_thinking_config(
+ {
+ **GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request,
+ "contents": [
+ list_value(GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request["contents"])[0],
+ {
+ "role": "model",
+ "parts": [
+ {"text": "I will locate the quantity validation and its tests."},
+ {
+ "function_call": {
+ "name": "search_files",
+ "args": {"query": "quantity", "paths": ["shop/", "tests/"]},
+ }
+ },
+ ],
+ },
+ {
+ "role": "user",
+ "parts": [
+ {
+ "function_response": {
+ "name": "search_files",
+ "response": {
+ "matches": [
+ {
+ "path": "shop/validation.py",
+ "line": 18,
+ "text": "def validate_quantity(quantity):",
+ },
+ {
+ "path": "tests/test_validation.py",
+ "line": 7,
+ "text": "def test_zero_quantity():",
+ },
+ ],
+ },
+ },
+ },
+ ],
+ },
+ list_value(GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request["contents"])[3],
+ ],
+ },
+ {"thinkingBudget": 4096, "includeThoughts": False},
+)
+
+
+def gemini_chat_response(
+ *,
+ content: str,
+ tool_call: Mapping[str, JsonValue],
+ message_fields: Mapping[str, JsonValue],
+ usage: Mapping[str, JsonValue],
+) -> Mapping[str, JsonValue]:
+ return {
+ **GEMINI_3_8_FLASH_TEST_CASE.expected_litellm_response,
+ "choices": [
+ {
+ "finish_reason": "tool_calls",
+ "index": 0,
+ "message": {
+ "content": content,
+ "role": "assistant",
+ "tool_calls": [dict(tool_call)],
+ "images": [],
+ **message_fields,
+ },
+ "provider_specific_fields": {"native_finish_reason": "STOP"},
+ },
+ ],
+ "usage": dict(usage),
+ }
+
+
+def signed_tool_call(call_id: str, arguments: str, signature: str) -> Mapping[str, JsonValue]:
+ return {
+ "index": 0,
+ "provider_specific_fields": {"thought_signature": signature},
+ "function": {"arguments": arguments, "name": "read_file"},
+ "id": f"{call_id}__thought__{signature}",
+ "type": "function",
+ }
+
+
+THINKING_REQUEST: Final[Mapping[str, JsonValue]] = {
+ key: value for key, value in GEMINI_3_8_FLASH_TEST_CASE.litellm_request.items() if key != "reasoning_effort"
+}
+
+GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE: Final = replace(
+ GEMINI_3_8_FLASH_TEST_CASE,
+ scenario="thoughts_off",
+ expected_provider_request=with_thinking_config(
+ GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request, {"thinkingLevel": "high", "includeThoughts": False}
+ ),
+ mock_provider_response=GEMINI_3_8_FLASH_THOUGHTS_OFF_RESPONSE,
+ expected_litellm_response=gemini_chat_response(
+ content=(
+ "To understand how zero, negative, and non-integer quantities (such as floats or invalid types) are "
+ "currently handled and where the checkout rejection is enforced, we need to inspect "
+ "`shop/validation.py` next."
+ ),
+ tool_call=signed_tool_call(
+ "call_330625",
+ '{"path": "shop/validation.py", "end_line": 50, "start_line": 1}',
+ GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE,
+ ),
+ message_fields={
+ "thinking_blocks": [],
+ "provider_specific_fields": {"thought_signatures": [GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE]},
+ },
+ usage={
+ "completion_tokens": 435,
+ "prompt_tokens": 358,
+ "total_tokens": 793,
+ "completion_tokens_details": {"reasoning_tokens": 358, "text_tokens": 77},
+ "prompt_tokens_details": {"text_tokens": 358},
+ },
+ ),
+)
+
+GEMINI_2_5_FLASH_THOUGHTS_OFF_TEST_CASE: Final = replace(
+ GEMINI_3_8_FLASH_TEST_CASE,
+ scenario="thoughts_off",
+ litellm_request={**GEMINI_3_8_FLASH_TEST_CASE.litellm_request, "model": "gemini/gemini-2.5-flash"},
+ expected_provider_endpoint="/models/gemini-2.5-flash:generateContent",
+ expected_provider_request=GEMINI_2_5_FLASH_REQUEST,
+ mock_provider_response=GEMINI_2_5_FLASH_THOUGHTS_OFF_RESPONSE,
+ expected_litellm_response={
+ **gemini_chat_response(
+ content="I will read the `shop/validation.py` file to understand the quantity validation logic.",
+ tool_call={
+ "index": 0,
+ "function": {
+ "arguments": '{"start_line": 1, "path": "shop/validation.py", "end_line": 20}',
+ "name": "read_file",
+ },
+ "id": ANY,
+ "type": "function",
+ },
+ message_fields={"thinking_blocks": []},
+ usage={
+ "completion_tokens": 53,
+ "prompt_tokens": 334,
+ "total_tokens": 387,
+ "completion_tokens_details": {"text_tokens": 53},
+ "prompt_tokens_details": {"text_tokens": 334},
+ },
+ ),
+ "model": "gemini/gemini-2.5-flash",
+ },
+)
+
+VERTEX_GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE: Final = replace(
+ VERTEX_GEMINI_3_8_FLASH_TEST_CASE,
+ scenario="thoughts_off",
+ expected_provider_request=with_thinking_config(
+ VERTEX_GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request, {"thinkingLevel": "high", "includeThoughts": False}
+ ),
+ mock_provider_response=VERTEX_GEMINI_3_8_FLASH_THOUGHTS_OFF_RESPONSE,
+ expected_litellm_response={
+ **gemini_chat_response(
+ content=(
+ "To understand how zero, negative, and non-integer quantities are handled during checkout validation, "
+ "we need to inspect `shop/validation.py` next."
+ ),
+ tool_call=signed_tool_call(
+ "call_1272842",
+ '{"start_line": 1, "path": "shop/validation.py", "end_line": 50}',
+ VERTEX_GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE,
+ ),
+ message_fields={
+ "thinking_blocks": [],
+ "provider_specific_fields": {
+ "thought_signatures": [VERTEX_GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE]
+ },
+ },
+ usage={
+ "completion_tokens": 337,
+ "prompt_tokens": 246,
+ "total_tokens": 583,
+ "completion_tokens_details": {"reasoning_tokens": 273, "text_tokens": 64},
+ "prompt_tokens_details": {"text_tokens": 246},
+ },
+ ),
+ "model": "vertex_ai/gemini-3.8-flash",
+ },
+)
+
+GEMINI_3_8_FLASH_THINKING_OMITTED_TEST_CASE: Final = replace(
+ GEMINI_3_8_FLASH_TEST_CASE,
+ scenario="thinking_omitted",
+ litellm_request={
+ **THINKING_REQUEST,
+ "thinking": {"type": "enabled", "budget_tokens": 1024, "display": "omitted"},
+ },
+ expected_provider_request=with_thinking_config(
+ GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request, {"includeThoughts": False}
+ ),
+ mock_provider_response=GEMINI_3_8_FLASH_THINKING_OMITTED_RESPONSE,
+ expected_litellm_response=gemini_chat_response(
+ content=(
+ "To understand how the checkout validates quantities—including zero, negative, and non-integer "
+ "quantities—I need to inspect `shop/validation.py` next."
+ ),
+ tool_call=signed_tool_call(
+ "call_745953",
+ '{"start_line": 1, "path": "shop/validation.py", "end_line": 40}',
+ GEMINI_3_8_FLASH_THINKING_OMITTED_THOUGHT_SIGNATURE,
+ ),
+ message_fields={
+ "thinking_blocks": [],
+ "provider_specific_fields": {"thought_signatures": [GEMINI_3_8_FLASH_THINKING_OMITTED_THOUGHT_SIGNATURE]},
+ },
+ usage={
+ "completion_tokens": 173,
+ "prompt_tokens": 358,
+ "total_tokens": 531,
+ "completion_tokens_details": {"reasoning_tokens": 108, "text_tokens": 65},
+ "prompt_tokens_details": {"text_tokens": 358},
+ },
+ ),
+)
+
+GEMINI_3_8_FLASH_THINKING_SHOWN_TEST_CASE: Final = replace(
+ GEMINI_3_8_FLASH_TEST_CASE,
+ scenario="thinking_shown",
+ litellm_request={
+ **THINKING_REQUEST,
+ "thinking": {"type": "enabled", "budget_tokens": 1024, "display": "summarized"},
+ },
+ expected_provider_request=with_thinking_config(
+ GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request, {"includeThoughts": True}
+ ),
+ mock_provider_response=GEMINI_3_8_FLASH_THINKING_SHOWN_RESPONSE,
+ expected_litellm_response=gemini_chat_response(
+ content=(
+ "I need to inspect `shop/validation.py` next to examine how zero, negative, and non-integer quantities "
+ "are currently handled in `validate_quantity`."
+ ),
+ tool_call=signed_tool_call(
+ "call_386884",
+ '{"end_line": 40, "start_line": 1, "path": "shop/validation.py"}',
+ GEMINI_3_8_FLASH_THINKING_SHOWN_THOUGHT_SIGNATURE,
+ ),
+ message_fields={
+ "reasoning_content": GEMINI_3_8_FLASH_THINKING_SHOWN_THOUGHT,
+ "thinking_blocks": [{"type": "thinking", "thinking": GEMINI_3_8_FLASH_THINKING_SHOWN_THOUGHT}],
+ "provider_specific_fields": {"thought_signatures": [GEMINI_3_8_FLASH_THINKING_SHOWN_THOUGHT_SIGNATURE]},
+ },
+ usage={
+ "completion_tokens": 171,
+ "prompt_tokens": 358,
+ "total_tokens": 529,
+ "completion_tokens_details": {"reasoning_tokens": 105, "text_tokens": 66},
+ "prompt_tokens_details": {"text_tokens": 358},
+ },
+ ),
+)
+
+
+def with_request(case: TranslationTestCase, scenario: str, **fields: JsonValue) -> TranslationTestCase:
+ return replace(case, scenario=scenario, litellm_request={**case.litellm_request, **fields})
+
+
+def with_thinking(case: TranslationTestCase, scenario: str, display: JsonValue) -> TranslationTestCase:
+ return with_request(case, scenario, thinking={"type": "enabled", "budget_tokens": 1024, "display": display})
+
+
+def with_summary(case: TranslationTestCase, scenario: str, summary: JsonValue) -> TranslationTestCase:
+ return with_request(case, scenario, reasoning_effort={"effort": "high", "summary": summary})
+
+
+@pytest.mark.parametrize(
+ "case",
+ (
+ GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE,
+ GEMINI_2_5_FLASH_THOUGHTS_OFF_TEST_CASE,
+ VERTEX_GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE,
+ ),
+ ids=(
+ GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE.id,
+ GEMINI_2_5_FLASH_THOUGHTS_OFF_TEST_CASE.id,
+ VERTEX_GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE.id,
+ ),
+)
+def test_gemini_keeps_thinking_and_billing_reasoning_without_thought_text_when_the_gateway_turns_it_off(
+ case: TranslationTestCase, thoughts_off: Gateway, provider: SharedProvider
+) -> None:
+ assert_translation(case, thoughts_off, provider)
+
+
+SUMMARY_ASKS: Final = (
+ with_summary(GEMINI_3_8_FLASH_TEST_CASE, "summary_auto", "auto"),
+ with_summary(GEMINI_3_8_FLASH_TEST_CASE, "summary_empty", ""),
+ with_summary(GEMINI_3_8_FLASH_TEST_CASE, "summary_long", LONG_SUMMARY),
+ with_summary(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, "summary_int", 7),
+ with_summary(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, "summary_list", ["auto"]),
+ with_summary(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, "summary_null", None),
+ with_request(GEMINI_3_8_FLASH_TEST_CASE, "summary_alias", reasoning_summary="auto"),
+ with_request(GEMINI_3_8_FLASH_TEST_CASE, "summary_alias_camel", reasoningSummary="detailed"),
+ with_request(
+ GEMINI_3_8_FLASH_TEST_CASE,
+ "summary_alias_effort_dict",
+ reasoning_effort={"effort": "high"},
+ reasoning_summary="auto",
+ ),
+ with_request(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, "summary_alias_null", reasoning_summary=None),
+ with_request(
+ GEMINI_3_8_FLASH_THINKING_SHOWN_TEST_CASE,
+ "thinking_summary",
+ thinking={"type": "enabled", "budget_tokens": 1024, "summary": "auto"},
+ ),
+)
+
+
+@pytest.mark.parametrize("case", SUMMARY_ASKS, ids=tuple(case.id for case in SUMMARY_ASKS))
+def test_only_a_string_reasoning_summary_brings_thought_text_back_over_the_gateway_default(
+ case: TranslationTestCase, thoughts_off: Gateway, provider: SharedProvider
+) -> None:
+ assert_translation(case, thoughts_off, provider)
+
+
+def without_closed_schema(tool: JsonValue) -> JsonValue:
+ function: Final = object_value(object_value(tool)["function"])
+ parameters: Final = object_value(function["parameters"])
+ open_parameters: Final = {name: value for name, value in parameters.items() if name != "additionalProperties"}
+ return {**object_value(tool), "function": {**function, "parameters": open_parameters}}
+
+
+def cacheable_run(case: TranslationTestCase, marker: str) -> TranslationTestCase:
+ messages: Final = list_value(case.litellm_request["messages"])
+ system: Final = object_value(messages[0])
+ instruction: Final = object_value(case.expected_provider_request["system_instruction"])
+ (part,) = list_value(instruction["parts"])
+ text: Final = f"{string_value(system['content'])} Review run {marker}."
+ return replace(
+ case,
+ litellm_request={
+ **{name: value for name, value in case.litellm_request.items() if name != "cache"},
+ "messages": [{**system, "content": text}, messages[1]],
+ "tools": [without_closed_schema(tool) for tool in list_value(case.litellm_request["tools"])],
+ },
+ expected_provider_request={
+ **case.expected_provider_request,
+ "contents": [list_value(case.expected_provider_request["contents"])[0]],
+ "system_instruction": {**instruction, "parts": [{**object_value(part), "text": text}]},
+ },
+ )
+
+
+def response_cache_keys() -> frozenset[str]:
+ keys: Final = (raw.decode() for raw in redis_store().scan_iter())
+ return frozenset(key for key in keys if CACHE_KEY.match(key))
+
+
+def stored_reply_key(before: frozenset[str]) -> str:
+ (key,) = eventually(lambda: response_cache_keys() - before, lambda added: len(added) == 1)
+ return key
+
+
+GEMINI_3_8_FLASH_SUMMARY_ALIAS_TEST_CASE: Final = with_request(
+ GEMINI_3_8_FLASH_TEST_CASE, "summary_alias", reasoning_summary="auto"
+)
+
+
+@pytest.mark.parametrize(
+ ("first", "second"),
+ (
+ (GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, GEMINI_3_8_FLASH_SUMMARY_ALIAS_TEST_CASE),
+ (GEMINI_3_8_FLASH_SUMMARY_ALIAS_TEST_CASE, GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE),
+ ),
+ ids=("plain_then_summary_alias", "summary_alias_then_plain"),
+)
+def test_a_summary_alias_ask_and_a_plain_request_never_share_a_cached_reply(
+ first: TranslationTestCase, second: TranslationTestCase, thoughts_off: Gateway, provider: SharedProvider
+) -> None:
+ marker: Final = uuid4().hex
+ before: Final = response_cache_keys()
+ assert_translation(cacheable_run(first, marker), thoughts_off, provider)
+ first_key: Final = stored_reply_key(before)
+ assert_translation(cacheable_run(second, marker), thoughts_off, provider)
+ second_key: Final = stored_reply_key(before | {first_key})
+ repeat: Final = thoughts_off.request("POST", second.litellm_endpoint, cacheable_run(second, marker).litellm_request)
+ assert repeat.status_code == 200, repeat.text
+ assert repeat.headers.get("x-litellm-cache-key") == second_key, sorted(repeat.headers)
+ assert provider.received() == ()
+
+
+DISPLAY_ASKS: Final = (
+ GEMINI_3_8_FLASH_THINKING_SHOWN_TEST_CASE,
+ with_thinking(GEMINI_3_8_FLASH_THINKING_SHOWN_TEST_CASE, "display_updates", "updates"),
+ with_thinking(GEMINI_3_8_FLASH_THINKING_SHOWN_TEST_CASE, "display_int", 7),
+ with_thinking(GEMINI_3_8_FLASH_THINKING_SHOWN_TEST_CASE, "display_empty", ""),
+ with_thinking(GEMINI_3_8_FLASH_THINKING_SHOWN_TEST_CASE, "display_list", ["omitted"]),
+ with_thinking(GEMINI_3_8_FLASH_THINKING_SHOWN_TEST_CASE, "display_long", LONG_SUMMARY),
+ GEMINI_3_8_FLASH_THINKING_OMITTED_TEST_CASE,
+)
+
+
+@pytest.mark.parametrize("case", DISPLAY_ASKS, ids=tuple(case.id for case in DISPLAY_ASKS))
+def test_any_thinking_display_but_omitted_brings_thought_text_back_over_the_gateway_default(
+ case: TranslationTestCase, thoughts_off: Gateway, provider: SharedProvider
+) -> None:
+ assert_translation(case, thoughts_off, provider)
+
+
+DEFAULT_GATEWAY_DISPLAYS: Final = (
+ with_thinking(GEMINI_3_8_FLASH_THINKING_SHOWN_TEST_CASE, "display_omitted_default", "omitted"),
+ GEMINI_3_8_FLASH_THINKING_SHOWN_TEST_CASE,
+)
+
+
+@pytest.mark.parametrize("case", DEFAULT_GATEWAY_DISPLAYS, ids=tuple(case.id for case in DEFAULT_GATEWAY_DISPLAYS))
+def test_thinking_display_never_turns_thought_text_off_on_the_default_gateway(
+ case: TranslationTestCase, gateway: Gateway, provider: SharedProvider
+) -> None:
+ assert_translation(case, gateway, provider)
+
+
+def test_a_streamed_turn_keeps_thinking_and_billing_reasoning_without_thought_text(
+ thoughts_off: Gateway, provider: SharedProvider
+) -> None:
+ provider.expect(
+ Reply(
+ content_type="text/event-stream",
+ chunks=tuple(encode_sse_frame(frame) for frame in GEMINI_3_8_FLASH_THOUGHTS_OFF_STREAM_FRAMES),
+ )
+ )
+ response: Final = thoughts_off.request(
+ "POST",
+ "/v1/chat/completions",
+ {**GEMINI_3_8_FLASH_TEST_CASE.litellm_request, "stream": True, "stream_options": {"include_usage": True}},
+ )
+ assert_provider_request(
+ provider.received(),
+ "/models/gemini-3.8-flash:streamGenerateContent?alt=sse",
+ GEMINI_3_8_FLASH_TEST_CASE.expected_provider_headers,
+ GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE.expected_provider_request,
+ )
+ assert response.status_code == 200, response.text
+ frames: Final = tuple(frame.data for frame in parse_sse_frames(response.text) if frame.data != "[DONE]")
+ deltas: Final = tuple(
+ object_value(object_value(choice)["delta"])
+ for frame in frames
+ for choice in list_value(frame.get("choices", []))
+ )
+ assert "".join(str(delta.get("content") or "") for delta in deltas) == (
+ "To understand how zero, negative, and non-integer quantities are handled during checkout, I need to "
+ "inspect `shop/validation.py` next."
+ )
+ assert [delta for delta in deltas if delta.get("reasoning_content") or delta.get("thinking_blocks")] == []
+ usage: Final = [object_value(frame["usage"]) for frame in frames if frame.get("usage")][-1]
+ assert {field: value for field, value in usage.items() if field != "cost"} == {
+ "completion_tokens": 351,
+ "prompt_tokens": 358,
+ "total_tokens": 709,
+ "completion_tokens_details": {"reasoning_tokens": 288, "text_tokens": 63},
+ "prompt_tokens_details": {"text_tokens": 358},
+ }
+ assert isinstance(usage["cost"], float) and usage["cost"] > 0
+
+
+@pytest.mark.parametrize(
+ ("configured", "case"),
+ (("false", GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE), (None, GEMINI_3_8_FLASH_TEST_CASE)),
+ ids=("string-false", "null"),
+)
+def test_a_flag_written_as_text_or_left_null_still_sends_gemini_a_boolean(
+ configured: JsonValue, case: TranslationTestCase, gateway: Gateway, provider: SharedProvider, tmp_path: Path
+) -> None:
+ with gemini_thoughts_proxy(gateway, tmp_path, configured) as owned:
+ assert_translation(case, owned.gateway, provider)
+
+
+BURST_ASKS: Final[tuple[tuple[str, Mapping[str, JsonValue], Mapping[str, JsonValue]], ...]] = (
+ (
+ "/v1/chat/completions",
+ {**GEMINI_3_8_FLASH_TEST_CASE.litellm_request, "reasoning_effort": {"effort": "low", "summary": "auto"}},
+ {"thinkingLevel": "low", "includeThoughts": True},
+ ),
+ (
+ "/v1/chat/completions",
+ GEMINI_3_8_FLASH_TEST_CASE.litellm_request,
+ {"thinkingLevel": "high", "includeThoughts": False},
+ ),
+ (
+ "/v1/messages",
+ MESSAGES_GEMINI_3_8_FLASH_TEST_CASE.litellm_request,
+ {"thinkingLevel": "high", "includeThoughts": True},
+ ),
+ (
+ "/v1/messages",
+ {**MESSAGES_GEMINI_3_8_FLASH_TEST_CASE.litellm_request, "thinking": {"type": "adaptive"}},
+ {"thinkingLevel": "high", "includeThoughts": False},
+ ),
+ (
+ "/v1/responses",
+ {**RESPONSES_GEMINI_3_8_FLASH_TEST_CASE.litellm_request, "reasoning": {"effort": "high", "summary": "auto"}},
+ {"thinkingLevel": "high", "includeThoughts": True},
+ ),
+ (
+ "/v1/responses",
+ RESPONSES_GEMINI_3_8_FLASH_TEST_CASE.litellm_request,
+ {"thinkingLevel": "high", "includeThoughts": False},
+ ),
+)
+
+
+def _assert_each_request_follows_its_own_ask(gateway: Gateway, provider: SharedProvider) -> None:
+ burst: Final = BURST_ASKS * 4
+ for _ in burst:
+ provider.expect(Reply(body=json.dumps(GEMINI_3_8_FLASH_THOUGHTS_OFF_RESPONSE).encode()))
+ with ThreadPoolExecutor(max_workers=12) as pool:
+ responses: Final = tuple(pool.map(lambda ask: gateway.request("POST", ask[0], ask[1]), burst))
+ assert [response.status_code for response in responses] == [200] * len(burst), [
+ response.text for response in responses if response.status_code != 200
+ ]
+ sent: Final = sorted(json.dumps(sent_thinking_config(request), sort_keys=True) for request in provider.received())
+ assert sent == sorted(json.dumps(expected, sort_keys=True) for _, _, expected in burst)
+
+
+def test_every_worker_honors_the_flag_on_every_endpoint_under_a_burst_and_after_a_worker_respawns(
+ gateway: Gateway, provider: SharedProvider, tmp_path: Path
+) -> None:
+ with gemini_thoughts_proxy(gateway, tmp_path, False, workers=2) as owned:
+ _assert_each_request_follows_its_own_ask(owned.gateway, provider)
+ owned.kill_worker(min(owned.worker_pids()))
+ _assert_each_request_follows_its_own_ask(owned.gateway, provider)
diff --git a/tests/integration/translation/gemini_thoughts.py b/tests/integration/translation/gemini_thoughts.py
new file mode 100644
index 00000000000..e0d89c3d43e
--- /dev/null
+++ b/tests/integration/translation/gemini_thoughts.py
@@ -0,0 +1,63 @@
+from collections.abc import Iterator, Mapping
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Final
+
+import pytest
+import yaml
+from integration._support.client import JSON_OBJECT, Gateway, list_value, object_value
+from integration._support.process import OwnedProxy, owned_proxy_process
+from integration._support.wire import Request
+from pydantic import JsonValue
+
+GEMINI_2_5_FLASH_DEPLOYMENT: Final[Mapping[str, JsonValue]] = {
+ "model_name": "gemini/gemini-2.5-flash",
+ "litellm_params": {
+ "model": "gemini/gemini-2.5-flash",
+ "api_base": "http://127.0.0.1:8191",
+ "api_key": "synthetic-gemini-key",
+ },
+}
+LONG_SUMMARY: Final = "detailed " * 600
+
+
+def with_thinking_config(
+ request: Mapping[str, JsonValue], thinking_config: Mapping[str, JsonValue]
+) -> Mapping[str, JsonValue]:
+ return {
+ **request,
+ "generationConfig": {**object_value(request["generationConfig"]), "thinkingConfig": dict(thinking_config)},
+ }
+
+
+def sent_thinking_config(request: Request) -> Mapping[str, JsonValue]:
+ body: Final = JSON_OBJECT.validate_json(request.body)
+ return object_value(object_value(body["generationConfig"])["thinkingConfig"])
+
+
+@contextmanager
+def gemini_thoughts_proxy(
+ gateway: Gateway, directory: Path, default: JsonValue, *, workers: int = 1
+) -> Iterator[OwnedProxy]:
+ stock: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
+ config: Final = directory / "gemini-thoughts.yaml"
+ config.write_text(
+ yaml.safe_dump(
+ {
+ **stock,
+ "model_list": [*list_value(stock["model_list"]), dict(GEMINI_2_5_FLASH_DEPLOYMENT)],
+ "litellm_settings": {
+ **object_value(stock["litellm_settings"]),
+ "gemini_include_thoughts_default": default,
+ },
+ }
+ )
+ )
+ with owned_proxy_process(gateway, directory, {}, config=config, workers=workers) as owned:
+ yield owned
+
+
+@pytest.fixture(scope="module")
+def thoughts_off(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]:
+ with gemini_thoughts_proxy(session_gateway, tmp_path_factory.mktemp("gemini-thoughts-off"), False) as owned:
+ yield owned.gateway
diff --git a/tests/integration/translation/messages/reasoning/test_messages_reasoning_gemini.py b/tests/integration/translation/messages/reasoning/test_messages_reasoning_gemini.py
new file mode 100644
index 00000000000..c350d6e2a10
--- /dev/null
+++ b/tests/integration/translation/messages/reasoning/test_messages_reasoning_gemini.py
@@ -0,0 +1,373 @@
+from collections.abc import Mapping
+from dataclasses import replace
+from typing import Final
+
+import pytest
+from integration._support.client import Gateway, object_value
+from integration._support.process import graceful_stop_seconds
+from integration._support.provider import SharedProvider
+from integration._support.wire import Reply
+from integration.translation.case import SSEFrame, TranslationTestCase
+from integration.translation.gemini_thoughts import LONG_SUMMARY, thoughts_off, with_thinking_config
+from integration.translation.messages.bases.gemini import GEMINI_3_8_FLASH_TEST_CASE, GEMINI_MESSAGES_RESPONSE
+from integration.translation.runner import (
+ assert_provider_request,
+ assert_translation,
+ encode_sse_frame,
+ parse_sse_frames,
+)
+from pydantic import JsonValue
+
+__all__ = ("thoughts_off",)
+
+pytestmark = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
+
+GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE: Final = (
+ "EscLCsQLAWkUfRP4jtXWjNAugx5FeAwxasL/kKzH9gXUHHB2cTK13h4C8C6tKVYMKP5rXoPq4yhoXSrsDiIjUEpuAjdQYYCz"
+ "Wvx21tGAN4iBA4VELBqPOi92vo9YVGpczJ/XVREq9KNRFIkyMzWyEzR1lBbJFA4detKHzSVJ7VQ1iRBn9jk6xvknTCxQhbnl"
+ "Pct8TfMo4XluFsM2PY3qAqoX1DjkrjQNe4nrN9Yoglq8Bwr2ayyc3DUNkNJDu0DQN20W6HU7fbeqSizEOsvGIvJVEUZhMOFF"
+ "0r8NyzUGkNiiAZ+TEupIV+G0eTr6sVlMPLS1N74ryQ55sevdj6irenHlP19IwkiDNhy6ApRIoblY5I1uJVt4lj4tQyACGMAE"
+ "VOTg3Gb78hPgJ+8cF3ukmRF0H1bpYH3n/mpnkqiggfK3vm7fVnsCy03W7Keyl7xd7Gfto/HL7vPCta/sO6ZmZm28JvdCr/Ce"
+ "X/W82nHuEqcjiR8fNPk413rL/geE4Lppru1H9ESWBICXzwl6pLxaBBNr8BD/jGMvlYA/ZEcBfKWtYyzow3+yp+4M0fsY1Ah7"
+ "pxVCcCbivMRv2X5J5m9btU1IIm2REJicjMSWCxh4lQMEuuk2FaPzge+wXROC/Zx0HwNw2iPw6VpziTOByACQxXOE7RdWULq1"
+ "3rkbyXCMkuAfG281SiylOnJ3I3tkbRcFoTaY5DcMBKYCCmoJJZPO1HKJ2DMw5Spz1HDftQ7lg9df5FDoZbmehUZKfY+j/2LN"
+ "//Wr5OEUmuBpojC5ht3YYgXaxclo/pa70KnLHlQWHPZ3nqdB2ru0tqgOaXSv9KZNhtvXgagLsmffQ/rm5QiRoEeG4oZ5E6Ga"
+ "VLQB9OvWCDNdzgNVzss0tveNcG/cttJBnkSd5/azHSvzkCDQYh4J/GT3zDVvOyOayWTqzGSLgIoFGpTES7Vsbpph62apefvO"
+ "+ZAlxPNNo/3tHncHzGc8Pku07fF1yhiKdGRV1OjYndfQv6UuDc3G7n2oU90iCKWRoCqOYAotzw1yCh4lbsSpEWO2iB9kWW93"
+ "mJsdqQBdk6e9UiuHdJIOfQ+cOVNOrltXXtZZ+T2mKFTFaWaN1F7MMH+xIX10dcnsWzq5o0Qz/tpwEfFOaQ5j43ez5BJdE5yD"
+ "wj4sDIdY57jAQ4141jU2oHpJCjY5Obhd4692k/FRtN7IMWMn5rQS+uTctWOh/27qatdy+QeAMPZseT4Yoi3W6W//Um9+sJY4"
+ "RjJ0Q27Fjb/wEqBkGsqx1sIidQc19sZWa2jQxVgG3T8xHYCMbgkqT/OHJaHGa4gxo78ajhdQnGYNWtvuRlGDpdUsxcNfF+6z"
+ "kbqWgcxKleRnQ1jKKZKS9GeiLoJnBRk8wIgmnrleblMB0GRj8i124HKLTfAm3WHlqy3xpACI3AIxPbk1GrATdkrb496eecVG"
+ "ko9YTrsdgOn5RV6Ggy/Rxr5wAAvPKBrqx/u1Whyr05pge9AZcGAtFRIZMWRCBnNg+oHE8M1EVipx3wQQXFDFO27Jpr3IandO"
+ "GgpbjbpwcW0fnDKLw7Ndt8SCyyVUEZwZ4Vp4tXJCOv9rzIMHJi+68Upwp5ILirrBFT5vKs//ub6l5hHzEfZty66+DvXyrpqT"
+ "VLYRnYMEO10Te24QyYVCO4qSnbJzfwTv5Kp4ODP2ONdSuJgkyYy78MzZ926H1/8rOlTEbWpAYb+QJwrt1nBJBkWRNDi2QNrl"
+ "UlFKTXZNLNXTP/rHHqfw1us8bhwe4Up2BbECPO+8F0tgSifXusnShtsRwpW7o+z/iXpnwuWePH92YiORo43mdoG4MJ54Emwf"
+ "B7x9ftguHpq1ZDOBoa3frE8l8a5RIEHemkzQsLlPyZSB+eQgnXrToXwUeEfR3EHG0vSCfokMLk0HinhWGa6hgIyJZQAdUIg4"
+ "GALvjSsrGEHwJaeLNNn/9EscJOEGyFQFMEy+sVIkL6Pj2+iZuG0OKahk"
+)
+
+GEMINI_3_8_FLASH_THOUGHTS_OFF_RESPONSE: Final[Mapping[str, JsonValue]] = {
+ "candidates": [
+ {
+ "content": {
+ "parts": [
+ {
+ "text": "I need to inspect `shop/validation.py` next to examine how `validate_quantity` handles zero, negative, and non-integer quantities.",
+ },
+ {
+ "functionCall": {
+ "name": "read_file",
+ "args": {
+ "end_line": 40,
+ "path": "shop/validation.py",
+ "start_line": 1,
+ },
+ "id": "call_412671",
+ },
+ "thoughtSignature": GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE,
+ },
+ ],
+ "role": "model",
+ },
+ "finishReason": "STOP",
+ "index": 0,
+ "finishMessage": "Model generated function call(s).",
+ },
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 354,
+ "candidatesTokenCount": 64,
+ "totalTokenCount": 739,
+ "promptTokensDetails": [
+ {
+ "modality": "TEXT",
+ "tokenCount": 354,
+ },
+ ],
+ "thoughtsTokenCount": 321,
+ "serviceTier": "standard",
+ },
+ "modelVersion": "gemini-3.8-flash",
+ "responseId": "1C3RatW-Jtuc_uMPhaLroA0",
+}
+
+GEMINI_3_8_FLASH_THOUGHTS_OFF_STREAM_THOUGHT_SIGNATURE: Final = (
+ "Eo0LCooLAWkUfRMuNip0QePQVa4NCyMrc0dwJR2qI/PSRcyScXfwpZhjBJVTorCs5ei6D4xzVf9attfTHx8bOBCX9wXynQ6V"
+ "JxU6uN562z5VAyrgoGv2nVGJk/mT7SnzIFiOc4RN+xebNGO0BIDf95xyOMeYP+ZHXquqtx1SnQXw7V6p81wTA4cGMnqjNE7F"
+ "Vs7OqlcpvzeFUmE8gl/Cy1a+u5SgLHwfnvOR3uzIX3u8kdk//dzvCKqXYaRStDf0dsrJICex/9WoOBfZKCujLBMK7CCjKEWF"
+ "nqERtFxXayTUKz1dvvBkcZ+MbhtrG3RFkBegqokQn9lus3icwBAfkZ9b3ps42j/mYw7DfjR+3Dg/Pyz8mgJVRC312QRxt+70"
+ "yOCU8POm/GTHJAZvvxdm9U4eiYBs3789F5oV34YzZl1+pcE4WVRwiFkB0vi17xy5XJFYTtzuk9d9ySr5Fz1de6nTKWk6TpuX"
+ "PM4q8bhX51kJWPCr1W4oVn1GhKpCtRyrXzKhuMsWcHXu/yKHUVjxyEFY+cTsVtHfrMHsIhMHyNnBqz2JclGXtFZScaq6AZ8g"
+ "0JZPqqAXgfgNyx2RaxoaCtITznOlBwMC1hTAZCrtjhHQNvGsAEdaRBAnZRFgeDHkuL7r7D9tr38UdxaUf/R5vfl5eO0ksj4X"
+ "bdeUG/LAnWP2CLUVjIoF50o4ZHKRy5UT/Mhr45986ePmtWnIN9t6dXOCB5I4qoEaljwkDn+llwZMbPq1orHdSrUAR6flcl3h"
+ "qVutnVOFgN3NeHvNkrkjUCvckiQL/ZBiG5YeDKfMP+XECDWIoTS2/DCJ7SX6d44B2yJ/PO2PH3nm7hr81U7r7xYKly4/2YFX"
+ "8bAsFJt4lGihfDHXFOr/VbQQd/CobZtLV9UZP7m1wUGXfL0knP31bdFWDKmSp/2giddbpIsGQf6j7h+YiE3M6QEfvRi7kaDX"
+ "CQ16O/4XJzG85/BxPw8ssR1cN+xLM95pZkDQO2MN1rxLVogF+T+Xc0M2SSjwbahM4nwYmisOqXv4Wuy0V7hQYJUySCvHabdH"
+ "nmTi+3mIfs9/zzHT3QUGZvAumBCZYfz1kqyz6mytXSwQbhhktmDLgIaj1abGm2gSn10fxl8pDm/mTq8/owS9aNPyDYChNBMr"
+ "NGYEBMXA4DtZy9TjybzMXKGPi4d/M7EZM5sC33WnCl2JovveyeZUNPWkZTV0+TjSSqk8C488huCOzv+PgJMvtBq6XiQW9ZV1"
+ "HQWk8XRMK7VMMhuv9sRqy1Lh/h8kIlSWtaCCNwpqIgUlZ+xybbY9dRGdsWq5FoxCU/OLmlSv7BcqXAamxxfIK9inXRKd/xkI"
+ "bR7OCeIaFq97Z5M08LfUpu9WjcXGgvtsppOUB4Rik2uCBVSzA/OqQdY1/N3q+BPYCvhOVVEBOJEEAvH9iAQOtCw4peaCtSLL"
+ "7vLh7s/lIZGBYj5I4zjoOfoSuT9g1x9uvv9h2sJUeP1HsN9Mp95weKBbF3HUN5F/lVRTan2I+Wb2+WnMMbashS1hoVIfk/Bc"
+ "426M3xPy7JDsczXLUP0PAyhAc2+xiFnkI120l4IHT3JK53AArx7svpTxzLOvPffVIcodkyWXHeJeC6eQbD6Elt/wWbx5Ig1L"
+ "mJQP7egSoyqmzu0zpuIC0xVsmDdgnumSi7acTjp1mIsLHHLLaWPcpxaNeiwBQf/7FfKsGqJ6p+I+lzt/rdRDOR4apkwKRwnr"
+ "d5beSf38TJd7L8Zj1qXp36Q2rnvtKubSm7f451w9ZbKsqbMDrwoagNWsfG9kc0CR0wtraSRzBHD6GUfoapPbplOLsSyjzEvX"
+ "i48aG814wRmQk932cd8ETx1a9BIbN/bsOr6stttugYakpvD3b8ckhfXW7eXfpKywBBzu1s1KLrU="
+)
+
+GEMINI_3_8_FLASH_THOUGHTS_OFF_STREAM_FRAMES: Final = (
+ SSEFrame(
+ data={
+ "candidates": [
+ {
+ "content": {
+ "parts": [
+ {
+ "text": "I need to read `shop/validation.py` next to see how `validate_quantity` currently handles zero, negative, and",
+ },
+ ],
+ "role": "model",
+ },
+ "index": 0,
+ },
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 354,
+ "candidatesTokenCount": 27,
+ "totalTokenCount": 676,
+ "promptTokensDetails": [
+ {
+ "modality": "TEXT",
+ "tokenCount": 354,
+ },
+ ],
+ "thoughtsTokenCount": 295,
+ "serviceTier": "standard",
+ },
+ "modelVersion": "gemini-3.8-flash",
+ "responseId": "2C3RavD_A77W-sAP9brxyQk",
+ }
+ ),
+ SSEFrame(
+ data={
+ "candidates": [
+ {
+ "content": {
+ "parts": [
+ {
+ "text": " non-integer quantities.",
+ },
+ ],
+ "role": "model",
+ },
+ "index": 0,
+ },
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 354,
+ "candidatesTokenCount": 32,
+ "totalTokenCount": 681,
+ "promptTokensDetails": [
+ {
+ "modality": "TEXT",
+ "tokenCount": 354,
+ },
+ ],
+ "thoughtsTokenCount": 295,
+ "serviceTier": "standard",
+ },
+ "modelVersion": "gemini-3.8-flash",
+ "responseId": "2C3RavD_A77W-sAP9brxyQk",
+ }
+ ),
+ SSEFrame(
+ data={
+ "candidates": [
+ {
+ "content": {
+ "parts": [
+ {
+ "functionCall": {
+ "name": "read_file",
+ "args": {
+ "start_line": 1,
+ "end_line": 50,
+ "path": "shop/validation.py",
+ },
+ "id": "call_654547",
+ },
+ "thoughtSignature": GEMINI_3_8_FLASH_THOUGHTS_OFF_STREAM_THOUGHT_SIGNATURE,
+ },
+ ],
+ "role": "model",
+ },
+ "index": 0,
+ },
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 354,
+ "candidatesTokenCount": 65,
+ "totalTokenCount": 714,
+ "promptTokensDetails": [
+ {
+ "modality": "TEXT",
+ "tokenCount": 354,
+ },
+ ],
+ "thoughtsTokenCount": 295,
+ "serviceTier": "standard",
+ },
+ "modelVersion": "gemini-3.8-flash",
+ "responseId": "2C3RavD_A77W-sAP9brxyQk",
+ }
+ ),
+ SSEFrame(
+ data={
+ "candidates": [
+ {
+ "content": {
+ "parts": [
+ {
+ "text": "",
+ },
+ ],
+ "role": "model",
+ },
+ "finishReason": "STOP",
+ "index": 0,
+ },
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 354,
+ "candidatesTokenCount": 65,
+ "totalTokenCount": 714,
+ "promptTokensDetails": [
+ {
+ "modality": "TEXT",
+ "tokenCount": 354,
+ },
+ ],
+ "thoughtsTokenCount": 295,
+ "serviceTier": "standard",
+ },
+ "modelVersion": "gemini-3.8-flash",
+ "responseId": "2C3RavD_A77W-sAP9brxyQk",
+ }
+ ),
+)
+
+ADAPTIVE_THINKING: Final[Mapping[str, JsonValue]] = {"type": "adaptive"}
+
+GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE: Final = replace(
+ GEMINI_3_8_FLASH_TEST_CASE,
+ scenario="thoughts_off",
+ litellm_request={**GEMINI_3_8_FLASH_TEST_CASE.litellm_request, "thinking": dict(ADAPTIVE_THINKING)},
+ expected_provider_request=with_thinking_config(
+ GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request, {"thinkingLevel": "high", "includeThoughts": False}
+ ),
+ mock_provider_response=GEMINI_3_8_FLASH_THOUGHTS_OFF_RESPONSE,
+ expected_litellm_response={
+ **GEMINI_MESSAGES_RESPONSE,
+ "id": "1C3RatW-Jtuc_uMPhaLroA0",
+ "model": "gemini/gemini-3.8-flash",
+ "content": [
+ {
+ "type": "text",
+ "text": (
+ "I need to inspect `shop/validation.py` next to examine how `validate_quantity` handles zero, "
+ "negative, and non-integer quantities."
+ ),
+ },
+ {
+ "type": "tool_use",
+ "id": "call_412671",
+ "name": "read_file",
+ "input": {"path": "shop/validation.py", "start_line": 1, "end_line": 40},
+ "provider_specific_fields": {"signature": GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE},
+ },
+ ],
+ "usage": {"input_tokens": 354, "output_tokens": 385},
+ },
+)
+
+
+def with_thinking(case: TranslationTestCase, scenario: str, field: str, value: JsonValue) -> TranslationTestCase:
+ return replace(
+ case,
+ scenario=scenario,
+ litellm_request={**case.litellm_request, "thinking": {**ADAPTIVE_THINKING, field: value}},
+ )
+
+
+THINKING_ASKS: Final = (
+ with_thinking(GEMINI_3_8_FLASH_TEST_CASE, "display_summarized", "display", "summarized"),
+ with_thinking(GEMINI_3_8_FLASH_TEST_CASE, "display_updates", "display", "updates"),
+ with_thinking(GEMINI_3_8_FLASH_TEST_CASE, "display_int", "display", 7),
+ with_thinking(GEMINI_3_8_FLASH_TEST_CASE, "display_empty", "display", ""),
+ with_thinking(GEMINI_3_8_FLASH_TEST_CASE, "display_list", "display", ["omitted"]),
+ with_thinking(GEMINI_3_8_FLASH_TEST_CASE, "display_long", "display", LONG_SUMMARY),
+ with_thinking(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, "display_omitted", "display", "omitted"),
+ with_thinking(GEMINI_3_8_FLASH_TEST_CASE, "summary_auto", "summary", "auto"),
+ with_thinking(GEMINI_3_8_FLASH_TEST_CASE, "summary_empty", "summary", ""),
+ with_thinking(GEMINI_3_8_FLASH_TEST_CASE, "summary_long", "summary", LONG_SUMMARY),
+ with_thinking(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, "summary_int", "summary", 7),
+ with_thinking(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, "summary_list", "summary", ["auto"]),
+ with_thinking(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, "summary_null", "summary", None),
+)
+
+
+def test_gemini_keeps_thinking_and_billing_reasoning_without_thought_text_when_the_gateway_turns_it_off(
+ thoughts_off: Gateway, provider: SharedProvider
+) -> None:
+ assert_translation(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, thoughts_off, provider)
+
+
+@pytest.mark.parametrize("case", THINKING_ASKS, ids=tuple(case.id for case in THINKING_ASKS))
+def test_a_thinking_display_or_summary_ask_brings_thought_text_back_over_the_gateway_default(
+ case: TranslationTestCase, thoughts_off: Gateway, provider: SharedProvider
+) -> None:
+ assert_translation(case, thoughts_off, provider)
+
+
+def test_display_omitted_alone_keeps_thought_text_on_the_default_gateway(
+ gateway: Gateway, provider: SharedProvider
+) -> None:
+ assert_translation(
+ with_thinking(GEMINI_3_8_FLASH_TEST_CASE, "display_omitted_default", "display", "omitted"), gateway, provider
+ )
+
+
+def test_a_streamed_turn_keeps_thinking_and_billing_reasoning_without_thought_text(
+ thoughts_off: Gateway, provider: SharedProvider
+) -> None:
+ provider.expect(
+ Reply(
+ content_type="text/event-stream",
+ chunks=tuple(encode_sse_frame(frame) for frame in GEMINI_3_8_FLASH_THOUGHTS_OFF_STREAM_FRAMES),
+ )
+ )
+ response: Final = thoughts_off.request(
+ "POST", "/v1/messages", {**GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE.litellm_request, "stream": True}
+ )
+ assert_provider_request(
+ provider.received(),
+ "/models/gemini-3.8-flash:streamGenerateContent?alt=sse",
+ GEMINI_3_8_FLASH_TEST_CASE.expected_provider_headers,
+ GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE.expected_provider_request,
+ )
+ assert response.status_code == 200, response.text
+ events: Final = tuple(object_value(frame.data) for frame in parse_sse_frames(response.text))
+ started: Final = tuple(
+ object_value(event["content_block"])["type"] for event in events if event.get("type") == "content_block_start"
+ )
+ assert started == ("text", "tool_use")
+ deltas: Final = tuple(
+ object_value(event["delta"]) for event in events if event.get("type") == "content_block_delta"
+ )
+ assert "".join(str(delta.get("text", "")) for delta in deltas) == (
+ "I need to read `shop/validation.py` next to see how `validate_quantity` currently handles zero, negative, "
+ "and non-integer quantities."
+ )
+ assert [object_value(event["usage"]) for event in events if event.get("type") == "message_delta"][-1] == {
+ "input_tokens": 354,
+ "output_tokens": 360,
+ }
diff --git a/tests/integration/translation/responses/reasoning/__init__.py b/tests/integration/translation/responses/reasoning/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/integration/translation/responses/reasoning/test_responses_reasoning_gemini.py b/tests/integration/translation/responses/reasoning/test_responses_reasoning_gemini.py
new file mode 100644
index 00000000000..984f7d2a8c4
--- /dev/null
+++ b/tests/integration/translation/responses/reasoning/test_responses_reasoning_gemini.py
@@ -0,0 +1,400 @@
+from collections.abc import Mapping
+from dataclasses import replace
+from typing import Final
+from unittest.mock import ANY
+
+import pytest
+from integration._support.client import Gateway, list_value, object_value
+from integration._support.process import graceful_stop_seconds
+from integration._support.provider import SharedProvider
+from integration._support.wire import Reply
+from integration.translation.case import SSEFrame, TranslationTestCase
+from integration.translation.gemini_thoughts import LONG_SUMMARY, thoughts_off, with_thinking_config
+from integration.translation.responses.bases.gemini import GEMINI_3_8_FLASH_TEST_CASE, GEMINI_RESPONSES_RESPONSE
+from integration.translation.runner import (
+ assert_provider_request,
+ assert_translation,
+ encode_sse_frame,
+ parse_sse_frames,
+)
+from pydantic import JsonValue
+
+__all__ = ("thoughts_off",)
+
+pytestmark = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
+
+GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE: Final = (
+ "EvQLCvELAWkUfRM8dwxvoWqyX5HLMaxcqWN9o4LPX6M6MZM/OD962KfMSvacFOxwOEb2B3rQM7QgoyT58YrItR4zpDSZJUxA"
+ "N4shekmtkSRpwrqE0rmH1v9AE8w5lR0Jewj8xmiiUXUrK0qvpyxhbVJRmIgdt9CqW+B4bJHYpp1VFAuHMk8eyBRoVNXakpsG"
+ "xjvHWGDomXbyJ7XuG1uHzvA5QtV2lgKUwRNvQqbFgbMhKN4dc8Imm2kwCLQONRN6wYIKJnYIi6ib+MhsFY+hZ6sPysMzVyrk"
+ "nGhc6pZwOKHbUjAraY/X98HO7UEL/q86LKUIvKTS4FfZ8keuyIQQBMdX68fdZPqurR4fI2dr5pFSDbaJWEVaE9XAY9wOXhOE"
+ "uC5SxY0U7nPU39eycNqAoAf+O7i/uHn1eCfLMhvAZyoH0z6vMfofnVn6L6alZtnRFQ3rEdJWEZHFmsLOk7G9SHlLAYKQjEdj"
+ "5aQHsKqBG1SFODPWSeKduI5w/kgFwgrw1ye1rnFTd+FDBs3DsQUlwh+jbPg8tr3KnCL5f2bxoGShUanPcGRAJ4b/IWGM3+wf"
+ "Whcx+3JY/OfuklkCnw86uRlcJiEjWd9XZ+jAby3ECZxBuXsHSE+Ynzpd0mp5g/S1orzqlmUJ/PlOcbDR/AsFduZvZwG0n4hq"
+ "hbXIIPszIpEyGDhkC//vxXsWdo5gLcsnEoUlLCESN3Sc12goKneId6u8pVWfzvHwRa0pK0LIWo4GcWZepZpEoYARlUZN6Dat"
+ "FNJ0nFW1sGaNb4CEYdSSp3m+GVobHpdz8qsGJxWrzzc6ZYmO7LbDJ25dJtyCl9QcfyzGbkFVDdAaTBEM23EdyZMFc2IIT24K"
+ "OVx5esCx+VjDZOmOzue/ZR0xrfaPXVa0CMuzGA8WwiphpWTbD8xZUnqEUZQNwKrjU91+SIhaBKcTVM8F1Wj0t9vlQkkMlzth"
+ "bJxIw+Ccu51gh/s8/awIDsBwC2Ax6eYIhUoWxUWYcphX7r3lBKhIuzv7FyO8tmaRkIOy6Eb1lXHB4fHAIOimZ/n4ClgK/KOZ"
+ "XVLosaGMKsmjYAdrxxPaNqs+say+pQ08bec/c+mqWeDbTOJ2JrCKM/Q+EcyD+82QqWlO68boeJyyt0fIvTCAdn57xN5Th7r1"
+ "WGhT9YHeSB1Y77P9FaJLQqAOjr+yogG7suiCP4QHg1arBlVzvaxAdnlBeXSv6pbQJuuvsFV9+SfwKHVXaBv68bcjSYKm8YZa"
+ "vfK57jweJUxvknCUpzrDRBdVZX28H60eQ3X6S0cV/696AO8NHoKuMv9zsCyu4WXWnpYEObtVD4Sn5hdSn1tnOkyUTLmN+bHV"
+ "sRGe5WTmrgnTKbrGDqpeqFpYLeMftkKWX+DHZj/LvZiNH1pHJid4nbEvcUKHciculasWiLZXDhEDKA1uzvSeG7CHjtHziqRi"
+ "b0dY7pQ0iEJUrjIwMVPnFZoyrywYZVPYgk+jna9FJypscuCz5XzdV+N2b2yrKryU6NhOQ5u/lzMNEFsyYxuXp6LZ5CVgdH1T"
+ "k+1rUf3sWPTO+CyDj+GzWvo5ljP7HihqyitafZJeKOSfaPfrJpBq+R8jnwDOVLCS1lnkhsof12PFoQHAE0YzGDtVscrstW4O"
+ "OaEdtNyNh30fhwwaZ/XGM4O9yJ8l7CgUvrgY/t8/f1hB/ZIc980HGs3G/bmVBYTxeZ18rqanrURtVuNtveyGief6eiEVIFjo"
+ "HWkAFXM+sn1zsJI6l4N9TPMbkoPxLCw2KaOaI/7mQjsH6VktpGuuqpHw3S5tiw3OGwh0lEwydKT0bk4hUhWqoBHqFpCrmKqA"
+ "g3z0kEHxQjQ1j8GmHusdPp4odrCWphw44m1VeWYjF4VjQ16Vik3SSrC8g3v9sHhtKM+3RL6z8zzukqgcQkXf7RXWvmFttckj"
+ "yYWM+vfQm/UvCzQ4dKvFHMqgmwuyEISPO9CD68/518wNpbQD1iB/RnxKwyYoD4u9YQ9TzXCzk2eP0NBB6pRmGg5nBo2SP0PB"
+ "BRZqlxnFX4XfQiS9pxNL"
+)
+
+GEMINI_3_8_FLASH_THOUGHTS_OFF_RESPONSE: Final[Mapping[str, JsonValue]] = {
+ "candidates": [
+ {
+ "content": {
+ "parts": [
+ {
+ "text": "To see how zero, negative, and non-integer quantities are currently handled and rejected by the validation logic, the next file I need to inspect is `shop/validation.py`.",
+ },
+ {
+ "functionCall": {
+ "name": "read_file",
+ "args": {
+ "path": "shop/validation.py",
+ "start_line": 1,
+ "end_line": 50,
+ },
+ "id": "call_584229",
+ },
+ "thoughtSignature": GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE,
+ },
+ ],
+ "role": "model",
+ },
+ "finishReason": "STOP",
+ "index": 0,
+ "finishMessage": "Model generated function call(s).",
+ },
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 354,
+ "candidatesTokenCount": 70,
+ "totalTokenCount": 790,
+ "promptTokensDetails": [
+ {
+ "modality": "TEXT",
+ "tokenCount": 354,
+ },
+ ],
+ "thoughtsTokenCount": 366,
+ "serviceTier": "standard",
+ },
+ "modelVersion": "gemini-3.8-flash",
+ "responseId": "2y3RasgqkpyMxw-RxpeQBg",
+}
+
+GEMINI_3_8_FLASH_THOUGHTS_OFF_STREAM_THOUGHT_SIGNATURE: Final = (
+ "EvYSCvMSAWkUfROtzl1tHuKrIElCi7xsI6n+7cPuqsM/qZSYIpK3gONKCCo0E4moadzMQ3sN09ll3gRjbcBg6rKSck+JTBhd"
+ "mrLsP/sqfxiYfMx4lWzDeq/Pw/IPhT2g5/2EvSU4FDsUczorTsD55gKvgRT8v05YTFMGF/np8r3VT52Q2yAzwRK+31qn8pYS"
+ "c25if5coTZJrvmJZDlZ9/64aibghRps9zN07yRVXpUjtLyKMaTDH0loGgiQn0I7Cxfhw9oUyZels7Un7pB7omA4RKW5xp5rx"
+ "Y+WgSeAEyyxNcPIkRvaxHaY1p2QaR3m17a0L7joJYNB+gNg7R/4Mfwpr8WuSOyLfBaKTuC5IiYWLaRm0ibASYzd59pPUyArt"
+ "mIx7pPfWxAr8/FT8TVQTSSeHXF88x/xJ2WSTSzQ3+FPw6Wm/LpfeGwQPX/g1+8f+KVcj1qAwT/tUobqOh7DdWl68CnNFReX4"
+ "SwBqRjTzZ276P7l15fhWmnQkiWXmLVsufJi5VXidEVrAsuEyVe7JmpExi3Y+0pn82D9x7lyv/TrOCVQgKwC0Zoq9NfuQxWa0"
+ "RvqmAv6toMfI8owTyzVoItnoGdhcYamwr7AyECKTuqnzMuylzOx5ttZnRc12/Ey2GQR3UxhyxMNLk01/UDicLnEXFl8Zn56i"
+ "rj51atVum7mJc48qekcKXQ5Lc3/t277v1d+eLmuiHebshpQwg+qQPHnnfky0WL98w343XDb29wwK3+veqQSjrffv5rQY9Tvq"
+ "3PMuN1pjU7Z3rL6fZmt5UrhqQQnTf4dA7X7n/aLkpA8SpY6ICDDr2x1z0yW5kMnHfxxk+N+j5eWwoTJuxxy94rLaYj3CKT57"
+ "XhMthQAoF1DGNpGxhbfC/tVuFE9FkxEBiYXqDz5fy2AU+S2zatT7/9mkFaHLY6hkOCS48X9WBVobxsQZZUnkbxplCzg4Huin"
+ "gPnl1bUDEprLHfTYgZHcnn3FgWFPwzxNBebrc4xjqXciMgReu4kOLP18JZbIJSjhP/RZAnQJEEOI8OrVGZlmiC7/JJx+8PTX"
+ "XYMTKtiVU/R3gOhOeVNN9WXruUoQzuy+JfJGg8ZQ2ni0Ugkx0sLO5rDsTOZ67bA+B99FQnQcBLyb40EPTOzrLrcM9wzXL+YZ"
+ "JUwNtdaQGxh/aUeKex/8of7GUH9ZsCZW0YN5xuEs8e2Q+vi2Msluckvki/wcyx7oOxmjw8bEbUkPBFq/ZEOcQqtetPLBzNPt"
+ "smJ4EfOU/kwBBLxplt6JrF+0ILeWJuH7OGpvpseXEosLpJ9HV4oIHL9co8EWGH/oAUzcuvhymR8FRziVE2IgVkreryIfHAs/"
+ "umsJgk1dXZk4A/R7KtjHKau0yxLAYJh6jWLvLsLv5N44VfKCQtyWFdUOSGbqn6E4YuOYVc+x+TGQcSMniobTekwm+UtqzVH8"
+ "mmtd8nwzwVkDgqZEAHrrnJOALLPeIbPuQ0/zjuUJmQovnnSqUyyCxQDpaxrlLD3iIZp9/9hPTY+lYHlam7VfChpQ7dPQ075t"
+ "ECiPA5TgMzh1qNQBF+rgKtrcrPo1ZDFqlseAybmSKqnCdvD5vdfdeVL8cmQIQao9qHDVh/aR0uwOsVKKzavPIYH2mhomg/cz"
+ "XB1nJG2zXB99g08Bd14c/zHHPNeqkqmwURCTxzPh8aeoRPcC8uL/eyawL66FfyUo40og9uridBDA4qx4ZMI+P1dBKFnMPM7o"
+ "Blpo+Sg9m6XZNPNVbDqRKKmX9OWCng7/PAKX+wa8EzRetNbyh73eMviMZaiN9YAyoKiS+nR5FqrPgNcFMqa8dHDEyY6O3Vuh"
+ "5Er5R3iGNACXMQvLRDVy/1UmFY1mzwew9lrDcUkZ0MNKBrbtDC87XEFGMejqvoJgF5Yl+q1mJE6gMBDTo6UEP1SOctxHyX23"
+ "mvqhCK9wOvK91UGnNyd+8ZQtQPP/L99d2y7hHmE6VEA/nUa2qkkO99vEQmJV3ilax+wZknABwk5sFdh7vWLwhEvxkdGYUh6c"
+ "HnA70VX5Kh0twJPKwE7gNASf/rvE5kYdrei5RtZNKMpVulXMPVho3eRBW6Qf6Ka0JjDZF5Dr8IXXsjP2LFTmnLNcnxu0FXx8"
+ "4xRuw1+yN3cP5R8LFlSGJmqPf5kFwAwpzT5FAXETiBaj7lg3WyG+OaDpab2NG+BvHbdfrDAuFKSsTDZ3Vt58fj6fuhS8tFd4"
+ "YzWL6X30pj+r7QRAefk+2CQDKM1TnbAJbyBcG2DVfeWpgmNjtzoV+RXNykIFIy6rSeOnUBAJCXwbAe7Scp5w/lx03a5aooRU"
+ "CP7x7CAmAU+YuGerG7A/3bbWDqN4EHGawkloK9UXmlOVX7uzLy0q3fJZb/6pSvPX+DR2rjFq0CSKnNFyBXYQRHx+DekAHWJS"
+ "VGUw7RLAPd4DE2HRiYTss7vtXalkP0hQ/vnt8lS9bjyBCi8ly+4FjFpo/SeKx0TKFulbFp1GPpsNERT+7kEkTsNZQHXWXPRB"
+ "PH0qmpFWZjI98fzWIGXGF376X0bcg6g67C/K/+YtH+QIU82ScAM0ozzQ1uzb0ci6wnM4JjNKA3ctVX6zuPR20Tb7VLlq5rwR"
+ "coOKjGQmBGbUiW5aKAAsMnIDhxf2lPCcwwb51C0Zb6hfOG+ka7LLcstgq4fmXHQ11BbTFu9ctAhy0E5V3aYSJAilb6VXZdD4"
+ "kSl3I0D2Xc/pTLPAzsxg6mEHL7J/DD1lRvEtSh4U34RCwxMk+pQbfEo6y/lkIVhXuE8p7d4mhPtsDJkls/3vnl+rR5ycebjH"
+ "g/zC01CjHhbVglb+HVEC5FMTrBNy/bweTE6c7qNHV6EWRN64nqvs4IW/jEVoaI2qt9WAJZ1br21rCE01iQagQUs1/ZH2KpQB"
+ "jTe3ju/BotwiE4TdqHcDL/VP62yq9KLeiBNXY2ctb3NwJzzfLrM2KHv2b4YmWuZ/VLABBfh9r3OzSmnFBvmWsZjZOHHEQXMI"
+ "hI+WDpRGB54CWmnWs6YzQK70IKrZUtq4BfPQJGf8OqO/ZeCJA4kinQbCaTTqQOJLsONKbPTqm8ImXz/QiML+KFvY3Z3ReW3T"
+ "NsJ3CGW7NRY1cfBIkhHuKHxPK1Ua5H3b7yP+h8OKrbqyjYPw2kE56TOH75HnIp/Ib6w+7uZpMVIbM4fYCRWlvQ7cUtaRq4fk"
+ "qqgHykE8bJYoU9eJFmPJNj4gt051sXJoslr2C6mvHYdSZgKxNa1ksd95cPKS3yURYQ=="
+)
+
+GEMINI_3_8_FLASH_THOUGHTS_OFF_STREAM_FRAMES: Final = (
+ SSEFrame(
+ data={
+ "candidates": [
+ {
+ "content": {
+ "parts": [
+ {
+ "text": "To understand how zero, negative, and non-integer quantities are checked during validation, I need to inspect",
+ },
+ ],
+ "role": "model",
+ },
+ "index": 0,
+ },
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 354,
+ "candidatesTokenCount": 21,
+ "totalTokenCount": 847,
+ "promptTokensDetails": [
+ {
+ "modality": "TEXT",
+ "tokenCount": 354,
+ },
+ ],
+ "thoughtsTokenCount": 472,
+ "serviceTier": "standard",
+ },
+ "modelVersion": "gemini-3.8-flash",
+ "responseId": "3S3Rat6pNonj_uMPv9i78Qw",
+ }
+ ),
+ SSEFrame(
+ data={
+ "candidates": [
+ {
+ "content": {
+ "parts": [
+ {
+ "text": " `shop/validation.py` next.",
+ },
+ ],
+ "role": "model",
+ },
+ "index": 0,
+ },
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 354,
+ "candidatesTokenCount": 30,
+ "totalTokenCount": 856,
+ "promptTokensDetails": [
+ {
+ "modality": "TEXT",
+ "tokenCount": 354,
+ },
+ ],
+ "thoughtsTokenCount": 472,
+ "serviceTier": "standard",
+ },
+ "modelVersion": "gemini-3.8-flash",
+ "responseId": "3S3Rat6pNonj_uMPv9i78Qw",
+ }
+ ),
+ SSEFrame(
+ data={
+ "candidates": [
+ {
+ "content": {
+ "parts": [
+ {
+ "functionCall": {
+ "name": "read_file",
+ "args": {
+ "end_line": 50,
+ "path": "shop/validation.py",
+ "start_line": 1,
+ },
+ "id": "call_5018",
+ },
+ "thoughtSignature": GEMINI_3_8_FLASH_THOUGHTS_OFF_STREAM_THOUGHT_SIGNATURE,
+ },
+ ],
+ "role": "model",
+ },
+ "index": 0,
+ },
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 354,
+ "candidatesTokenCount": 63,
+ "totalTokenCount": 889,
+ "promptTokensDetails": [
+ {
+ "modality": "TEXT",
+ "tokenCount": 354,
+ },
+ ],
+ "thoughtsTokenCount": 472,
+ "serviceTier": "standard",
+ },
+ "modelVersion": "gemini-3.8-flash",
+ "responseId": "3S3Rat6pNonj_uMPv9i78Qw",
+ }
+ ),
+ SSEFrame(
+ data={
+ "candidates": [
+ {
+ "content": {
+ "parts": [
+ {
+ "text": "",
+ },
+ ],
+ "role": "model",
+ },
+ "finishReason": "STOP",
+ "index": 0,
+ },
+ ],
+ "usageMetadata": {
+ "promptTokenCount": 354,
+ "candidatesTokenCount": 63,
+ "totalTokenCount": 889,
+ "promptTokensDetails": [
+ {
+ "modality": "TEXT",
+ "tokenCount": 354,
+ },
+ ],
+ "thoughtsTokenCount": 472,
+ "serviceTier": "standard",
+ },
+ "modelVersion": "gemini-3.8-flash",
+ "responseId": "3S3Rat6pNonj_uMPv9i78Qw",
+ }
+ ),
+)
+
+GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE: Final = replace(
+ GEMINI_3_8_FLASH_TEST_CASE,
+ scenario="thoughts_off",
+ expected_provider_request=with_thinking_config(
+ GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request, {"thinkingLevel": "high", "includeThoughts": False}
+ ),
+ mock_provider_response=GEMINI_3_8_FLASH_THOUGHTS_OFF_RESPONSE,
+ expected_litellm_response={
+ **GEMINI_RESPONSES_RESPONSE,
+ "model": "gemini/gemini-3.8-flash",
+ "output": [
+ {
+ "type": "message",
+ "id": ANY,
+ "status": "completed",
+ "role": "assistant",
+ "content": [
+ {
+ "type": "output_text",
+ "text": (
+ "To see how zero, negative, and non-integer quantities are currently handled and rejected "
+ "by the validation logic, the next file I need to inspect is `shop/validation.py`."
+ ),
+ "annotations": [],
+ },
+ ],
+ "phase": None,
+ },
+ {
+ "arguments": '{"path": "shop/validation.py", "start_line": 1, "end_line": 50}',
+ "call_id": "call_584229__thought__" + GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE,
+ "name": "read_file",
+ "type": "function_call",
+ "id": "fc_call_584229__thought__" + GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE,
+ "namespace": None,
+ "status": "completed",
+ "provider_specific_fields": {"thought_signature": GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE},
+ },
+ ],
+ "usage": {
+ "input_tokens": 354,
+ "input_tokens_details": {
+ "audio_tokens": None,
+ "cached_tokens": 0,
+ "cached_tokens_details": None,
+ "image_tokens": None,
+ "text_tokens": 354,
+ "video_tokens": None,
+ },
+ "output_tokens": 436,
+ "output_tokens_details": {"audio_tokens": None, "reasoning_tokens": 366, "text_tokens": 70},
+ "total_tokens": 790,
+ "cost": None,
+ },
+ },
+)
+
+
+def with_summary(case: TranslationTestCase, scenario: str, summary: JsonValue) -> TranslationTestCase:
+ reasoning: Final[Mapping[str, JsonValue]] = {"effort": "high", "summary": summary}
+ return replace(
+ case,
+ scenario=scenario,
+ litellm_request={**case.litellm_request, "reasoning": dict(reasoning)},
+ expected_litellm_response={**case.expected_litellm_response, "reasoning": dict(reasoning)},
+ )
+
+
+SUMMARY_ASKS: Final = (
+ with_summary(GEMINI_3_8_FLASH_TEST_CASE, "summary_auto", "auto"),
+ with_summary(GEMINI_3_8_FLASH_TEST_CASE, "summary_concise", "concise"),
+ with_summary(GEMINI_3_8_FLASH_TEST_CASE, "summary_empty", ""),
+ with_summary(GEMINI_3_8_FLASH_TEST_CASE, "summary_long", LONG_SUMMARY),
+ with_summary(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, "summary_int", 7),
+ with_summary(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, "summary_list", ["auto"]),
+ with_summary(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, "summary_null", None),
+)
+
+
+def test_gemini_keeps_thinking_and_billing_reasoning_without_thought_text_when_the_gateway_turns_it_off(
+ thoughts_off: Gateway, provider: SharedProvider
+) -> None:
+ assert_translation(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, thoughts_off, provider)
+
+
+@pytest.mark.parametrize("case", SUMMARY_ASKS, ids=tuple(case.id for case in SUMMARY_ASKS))
+def test_a_reasoning_summary_ask_brings_thought_text_back_over_the_gateway_default(
+ case: TranslationTestCase, thoughts_off: Gateway, provider: SharedProvider
+) -> None:
+ assert_translation(case, thoughts_off, provider)
+
+
+def test_a_streamed_turn_keeps_thinking_and_billing_reasoning_without_thought_text(
+ thoughts_off: Gateway, provider: SharedProvider
+) -> None:
+ provider.expect(
+ Reply(
+ content_type="text/event-stream",
+ chunks=tuple(encode_sse_frame(frame) for frame in GEMINI_3_8_FLASH_THOUGHTS_OFF_STREAM_FRAMES),
+ )
+ )
+ response: Final = thoughts_off.request(
+ "POST", "/v1/responses", {**GEMINI_3_8_FLASH_TEST_CASE.litellm_request, "stream": True}
+ )
+ assert_provider_request(
+ provider.received(),
+ "/models/gemini-3.8-flash:streamGenerateContent?alt=sse",
+ GEMINI_3_8_FLASH_TEST_CASE.expected_provider_headers,
+ GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE.expected_provider_request,
+ )
+ assert response.status_code == 200, response.text
+ events: Final = tuple(
+ object_value(frame.data) for frame in parse_sse_frames(response.text) if frame.data != "[DONE]"
+ )
+ assert [event["type"] for event in events if str(event["type"]).startswith("response.reasoning")] == []
+ assert "".join(str(event["delta"]) for event in events if event["type"] == "response.output_text.delta") == (
+ "To understand how zero, negative, and non-integer quantities are checked during validation, I need to "
+ "inspect `shop/validation.py` next."
+ )
+ completed: Final = object_value(
+ next(event["response"] for event in events if event["type"] == "response.completed")
+ )
+ assert [object_value(item)["type"] for item in list_value(completed["output"])] == ["message", "function_call"]
+ assert object_value(completed["usage"])["output_tokens_details"] == {
+ "reasoning_tokens": 472,
+ "text_tokens": 63,
+ }
diff --git a/tests/unit/caching/test_caching.py b/tests/unit/caching/test_caching.py
index b84529de0d6..91083921130 100644
--- a/tests/unit/caching/test_caching.py
+++ b/tests/unit/caching/test_caching.py
@@ -291,6 +291,37 @@ def test_exact_cache_key_includes_anthropic_messages_params(anthropic_param):
assert baseline != cache.get_cache_key(model="claude-sonnet-4-5", messages=messages, **anthropic_param)
+@pytest.mark.parametrize(
+ "summary_ask",
+ [
+ {"reasoning_summary": "auto"},
+ {"reasoningSummary": "auto"},
+ {"extra_body": {"reasoning_summary": "auto"}},
+ {"extra_body": {"reasoningSummary": "auto"}},
+ ],
+ ids=["top_level", "top_level_camel", "extra_body", "extra_body_camel"],
+)
+def test_exact_cache_key_includes_the_reasoning_summary_alias(summary_ask: dict[str, object]) -> None:
+ """The AI SDK summary alias decides whether thought summaries come back, wherever the request
+ carries it, so a reply cached without it must not answer a request that asks for one."""
+ cache: Final = Cache(type=LiteLLMCacheType.LOCAL)
+ messages: Final = [{"role": "user", "content": "how much is the ball?"}]
+ no_ask: Final = cache.get_cache_key(model="gemini/gemini-3.8-flash", messages=messages, reasoning_effort="medium")
+ ask: Final = cache.get_cache_key(
+ model="gemini/gemini-3.8-flash", messages=messages, reasoning_effort="medium", **summary_ask
+ )
+ assert no_ask != ask
+ assert ask == cache.get_cache_key(
+ model="gemini/gemini-3.8-flash", messages=messages, reasoning_effort="medium", **summary_ask
+ )
+ assert ask != cache.get_cache_key(
+ model="gemini/gemini-3.8-flash",
+ messages=messages,
+ reasoning_effort="medium",
+ reasoning_summary="detailed",
+ )
+
+
@pytest.mark.asyncio
async def test_embedding_cache_skips_write_when_one_input_yields_many_embeddings(monkeypatch):
"""A cross-encoder behind /embeddings returns one score per document for a single
@@ -561,10 +592,7 @@ def test_caching_v2(): # test in memory cache
litellm.cache = None # disable cache
litellm.success_callback = []
litellm._async_success_callback = []
- if (
- response2["choices"][0]["message"]["content"]
- != response1["choices"][0]["message"]["content"]
- ):
+ if response2["choices"][0]["message"]["content"] != response1["choices"][0]["message"]["content"]:
print(f"response1: {response1}")
print(f"response2: {response2}")
pytest.fail(f"Error occurred:")
@@ -595,10 +623,7 @@ def test_caching_with_ttl():
litellm.cache = None # disable cache
litellm.success_callback = []
litellm._async_success_callback = []
- assert (
- response2["choices"][0]["message"]["content"]
- != response1["choices"][0]["message"]["content"]
- )
+ assert response2["choices"][0]["message"]["content"] != response1["choices"][0]["message"]["content"]
except Exception as e:
print(f"error occurred: {traceback.format_exc()}")
pytest.fail(f"Error occurred: {e}")
@@ -632,9 +657,7 @@ def test_caching_with_default_ttl():
def test_caching_with_models_v2():
- messages = [
- {"role": "user", "content": "who is ishaan CTO of litellm from litellm 2023"}
- ]
+ messages = [{"role": "user", "content": "who is ishaan CTO of litellm from litellm 2023"}]
litellm.cache = Cache()
print("test2 for caching")
litellm.set_verbose = True
@@ -657,18 +680,12 @@ def test_caching_with_models_v2():
litellm.cache = None
litellm.success_callback = []
litellm._async_success_callback = []
- if (
- response3["choices"][0]["message"]["content"]
- == response2["choices"][0]["message"]["content"]
- ):
+ if response3["choices"][0]["message"]["content"] == response2["choices"][0]["message"]["content"]:
# if models are different, it should not return cached response
print(f"response2: {response2}")
print(f"response3: {response3}")
pytest.fail(f"Error occurred:")
- if (
- response1["choices"][0]["message"]["content"]
- != response2["choices"][0]["message"]["content"]
- ):
+ if response1["choices"][0]["message"]["content"] != response2["choices"][0]["message"]["content"]:
print(f"response1: {response1}")
print(f"response2: {response2}")
pytest.fail(f"Error occurred:")
@@ -708,9 +725,7 @@ def test_get_cache_key():
cache_key = cache_instance.get_cache_key(
**{
"model": "gpt-3.5-turbo",
- "messages": [
- {"role": "user", "content": "write a one sentence poem about: 7510"}
- ],
+ "messages": [{"role": "user", "content": "write a one sentence poem about: 7510"}],
"max_tokens": 40,
"temperature": 0.2,
"stream": True,
@@ -721,9 +736,7 @@ def test_get_cache_key():
cache_key_2 = cache_instance.get_cache_key(
**{
"model": "gpt-3.5-turbo",
- "messages": [
- {"role": "user", "content": "write a one sentence poem about: 7510"}
- ],
+ "messages": [{"role": "user", "content": "write a one sentence poem about: 7510"}],
"max_tokens": 40,
"temperature": 0.2,
"stream": True,
@@ -736,9 +749,9 @@ def test_get_cache_key():
# Hexadecimal representation of the hash
hash_hex = hash_object.hexdigest()
assert cache_key == hash_hex
- assert (
- cache_key_2 == hash_hex
- ), f"{cache_key} != {cache_key_2}. The same kwargs should have the same cache key across runs"
+ assert cache_key_2 == hash_hex, (
+ f"{cache_key} != {cache_key_2}. The same kwargs should have the same cache key across runs"
+ )
embedding_cache_key = cache_instance.get_cache_key(
**{
@@ -756,15 +769,13 @@ def test_get_cache_key():
print(embedding_cache_key)
- embedding_cache_key_str = (
- "model: azure/text-embedding-ada-002input: ['hi who is ishaan']"
- )
+ embedding_cache_key_str = "model: azure/text-embedding-ada-002input: ['hi who is ishaan']"
hash_object = hashlib.sha256(embedding_cache_key_str.encode())
# Hexadecimal representation of the hash
hash_hex = hash_object.hexdigest()
- assert (
- embedding_cache_key == hash_hex
- ), f"{embedding_cache_key} != 'model: azure/text-embedding-ada-002input: ['hi who is ishaan']'. The same kwargs should have the same cache key across runs"
+ assert embedding_cache_key == hash_hex, (
+ f"{embedding_cache_key} != 'model: azure/text-embedding-ada-002input: ['hi who is ishaan']'. The same kwargs should have the same cache key across runs"
+ )
# Proxy - embedding cache, test if embedding key, gets model_group and not model
embedding_cache_key_2 = cache_instance.get_cache_key(
@@ -817,9 +828,7 @@ def test_get_cache_key():
)
print(embedding_cache_key_2)
- embedding_cache_key_str_2 = (
- "model: EMBEDDING_MODEL_GROUPinput: ['hi who is ishaan']"
- )
+ embedding_cache_key_str_2 = "model: EMBEDDING_MODEL_GROUPinput: ['hi who is ishaan']"
hash_object = hashlib.sha256(embedding_cache_key_str_2.encode())
# Hexadecimal representation of the hash
hash_hex = hash_object.hexdigest()
@@ -852,9 +861,7 @@ def test_redis_caching_multiple_namespaces():
# Mock the Redis client creation from the _redis module
with (
patch("litellm._redis.get_redis_client") as mock_get_redis_client,
- patch(
- "litellm._redis.get_redis_connection_pool"
- ) as mock_get_redis_connection_pool,
+ patch("litellm._redis.get_redis_connection_pool") as mock_get_redis_connection_pool,
):
# Create a mock Redis client that simulates real Redis behavior
mock_redis_client = MagicMock()
@@ -927,18 +934,10 @@ def test_redis_caching_multiple_namespaces():
mock_response="Response without namespace",
)
- print(
- f"Response 1 type: {type(response_1)} - ID: {getattr(response_1, 'id', 'N/A')}"
- )
- print(
- f"Response 2 type: {type(response_2)} - ID: {getattr(response_2, 'id', 'N/A')}"
- )
- print(
- f"Response 3 type: {type(response_3)} - Cache hit: {isinstance(response_3, str)}"
- )
- print(
- f"Response 4 type: {type(response_4)} - ID: {getattr(response_4, 'id', 'N/A')}"
- )
+ print(f"Response 1 type: {type(response_1)} - ID: {getattr(response_1, 'id', 'N/A')}")
+ print(f"Response 2 type: {type(response_2)} - ID: {getattr(response_2, 'id', 'N/A')}")
+ print(f"Response 3 type: {type(response_3)} - Cache hit: {isinstance(response_3, str)}")
+ print(f"Response 4 type: {type(response_4)} - ID: {getattr(response_4, 'id', 'N/A')}")
print(f"Redis storage keys: {list(redis_storage.keys())}")
@@ -947,10 +946,7 @@ def test_redis_caching_multiple_namespaces():
namespace_1_keys = [k for k in cache_keys if k.startswith(f"{namespace_1}:")]
namespace_2_keys = [k for k in cache_keys if k.startswith(f"{namespace_2}:")]
no_namespace_keys = [
- k
- for k in cache_keys
- if not k.startswith(f"{namespace_1}:")
- and not k.startswith(f"{namespace_2}:")
+ k for k in cache_keys if not k.startswith(f"{namespace_1}:") and not k.startswith(f"{namespace_2}:")
]
print(f"Namespace 1 keys: {namespace_1_keys}")
@@ -963,9 +959,7 @@ def test_redis_caching_multiple_namespaces():
assert len(no_namespace_keys) > 0, "Should have cache keys for no namespace"
# The main test: response 3 should be a cache hit (string) because it uses same namespace as response 1
- assert isinstance(
- response_3, str
- ), "Response 3 should be a cache hit (string) for same namespace"
+ assert isinstance(response_3, str), "Response 3 should be a cache hit (string) for same namespace"
# response 1 & 2 should be ModelResponse objects (cache misses)
assert hasattr(response_1, "id"), "Response 1 should be a ModelResponse object"
@@ -973,14 +967,14 @@ def test_redis_caching_multiple_namespaces():
assert hasattr(response_4, "id"), "Response 4 should be a ModelResponse object"
# response 1 & 2 should have different IDs (different namespaces)
- assert (
- response_1.id != response_2.id
- ), f"Expected different response ID for different namespace. Got {response_1.id} and {response_2.id}"
+ assert response_1.id != response_2.id, (
+ f"Expected different response ID for different namespace. Got {response_1.id} and {response_2.id}"
+ )
# response 1 & 4 should have different IDs (different namespaces)
- assert (
- response_1.id != response_4.id
- ), f"Expected different response ID for no namespace vs namespaced. Got {response_1.id} and {response_4.id}"
+ assert response_1.id != response_4.id, (
+ f"Expected different response ID for no namespace vs namespaced. Got {response_1.id} and {response_4.id}"
+ )
_TOOL_TURN_ITEM: Final = {"role": "user", "content": "hi"}
@@ -1305,12 +1299,8 @@ async def test_embedding_caching_individual_items(monkeypatch: pytest.MonkeyPatc
monkeypatch.setattr(litellm, "cache", Cache(type=LiteLLMCacheType.LOCAL))
item: Final = f"individual embedding {uuid.uuid4().hex}"
- first: Final = await aembedding(
- model="text-embedding-3-small", input=item, caching=True, mock_response="0.1,0.2"
- )
- second: Final = await aembedding(
- model="text-embedding-3-small", input=item, caching=True, mock_response="0.3,0.4"
- )
+ first: Final = await aembedding(model="text-embedding-3-small", input=item, caching=True, mock_response="0.1,0.2")
+ second: Final = await aembedding(model="text-embedding-3-small", input=item, caching=True, mock_response="0.3,0.4")
assert second.data[0].embedding == first.data[0].embedding
assert second._hidden_params["cache_hit"] is True
@@ -1448,9 +1438,7 @@ async def test_caching_with_cache_controls(monkeypatch: pytest.MonkeyPatch) -> N
@pytest.mark.asyncio
-async def test_embedding_caching_redis_ttl(
- monkeypatch: pytest.MonkeyPatch, redis_client_without_network: None
-) -> None:
+async def test_embedding_caching_redis_ttl(monkeypatch: pytest.MonkeyPatch, redis_client_without_network: None) -> None:
mock_pipeline: Final = AsyncMock()
mock_set: Final = AsyncMock()
mock_pipeline.__aenter__.return_value.set = mock_set
@@ -1506,18 +1494,14 @@ async def test_logging_turn_off_message_logging_streaming(
assert cached_response["choices"][0]["message"]["content"] == "cached"
-def test_redis_caching_default_ttl(
- monkeypatch: pytest.MonkeyPatch, redis_client_without_network: None
-) -> None:
+def test_redis_caching_default_ttl(monkeypatch: pytest.MonkeyPatch, redis_client_without_network: None) -> None:
monkeypatch.setattr(litellm, "default_redis_ttl", 120)
cache: Final = RedisCache(host="127.0.0.1", port="6379")
assert cache.default_ttl == 120
-def test_redis_caching_llm_caching_ttl(
- monkeypatch: pytest.MonkeyPatch, redis_client_without_network: None
-) -> None:
+def test_redis_caching_llm_caching_ttl(monkeypatch: pytest.MonkeyPatch, redis_client_without_network: None) -> None:
monkeypatch.setattr(litellm, "default_redis_ttl", 120)
cache: Final = RedisCache(host="127.0.0.1", port="6379")
set_cache: Final = MagicMock()
@@ -1529,9 +1513,7 @@ def test_redis_caching_llm_caching_ttl(
@pytest.mark.asyncio
-async def test_redis_caching_ttl_pipeline(
- monkeypatch: pytest.MonkeyPatch, redis_client_without_network: None
-) -> None:
+async def test_redis_caching_ttl_pipeline(monkeypatch: pytest.MonkeyPatch, redis_client_without_network: None) -> None:
monkeypatch.setattr(litellm, "default_redis_ttl", 120)
cache: Final = RedisCache(host="127.0.0.1", port="6379")
pipeline: Final = AsyncMock()
@@ -1551,9 +1533,7 @@ async def test_redis_caching_ttl_pipeline(
@pytest.mark.asyncio
-async def test_redis_caching_ttl_sadd(
- monkeypatch: pytest.MonkeyPatch, redis_client_without_network: None
-) -> None:
+async def test_redis_caching_ttl_sadd(monkeypatch: pytest.MonkeyPatch, redis_client_without_network: None) -> None:
monkeypatch.setattr(litellm, "default_redis_ttl", 120)
cache: Final = RedisCache(host="127.0.0.1", port="6379")
redis_client: Final = AsyncMock()
diff --git a/tests/unit/integration_support/test_upstream.py b/tests/unit/integration_support/test_upstream.py
index 0cbcf1b7ee1..9e114e3b0b7 100644
--- a/tests/unit/integration_support/test_upstream.py
+++ b/tests/unit/integration_support/test_upstream.py
@@ -21,7 +21,11 @@ async def test_realtime_disconnect_after_session_creation_keeps_unexpected_error
response: Final = RealtimeResponse(content_type="application/x-realtime", events=({"type": "response.done"},))
provider.scenario_store.put("owned-scenario", response)
incoming: Final[deque[Message]] = deque(
- ({"type": "websocket.connect"}, {"type": "websocket.receive", "text": '{"type":"response.create"}'})
+ (
+ {"type": "websocket.connect"},
+ {"type": "websocket.receive", "text": '{"type":"response.create"}'},
+ {"type": "websocket.disconnect", "code": 1006},
+ )
)
outgoing: Final[deque[Message]] = deque()
diff --git a/tests/unit/integrations/test_custom_guardrail.py b/tests/unit/integrations/test_custom_guardrail.py
index 9bd41f4a602..614ff97553a 100644
--- a/tests/unit/integrations/test_custom_guardrail.py
+++ b/tests/unit/integrations/test_custom_guardrail.py
@@ -2790,7 +2790,7 @@ class TestLoggingOnlyApplyGuardrail:
assert out_kwargs["messages"] == kwargs["messages"]
assert (
out_kwargs["standard_logging_object"]["guardrail_information"][0]["guardrail_status"]
- == "guardrail_intervened"
+ == "guardrail_flagged"
)
@pytest.mark.asyncio
diff --git a/tests/unit/interactions/test_background_cost_polling.py b/tests/unit/interactions/test_background_cost_polling.py
index 9dc71a4fa1b..d55994dfb19 100644
--- a/tests/unit/interactions/test_background_cost_polling.py
+++ b/tests/unit/interactions/test_background_cost_polling.py
@@ -1,11 +1,14 @@
import asyncio
+import json
import time
-from datetime import datetime, timezone
+from datetime import datetime, timedelta, timezone
from itertools import islice
from typing import Optional
+import httpx
import pytest
+import litellm
from litellm.interactions.background_cost_polling import (
_create_context,
_poll_intervals,
@@ -21,6 +24,7 @@ from litellm.interactions.background_cost_polling import (
)
from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs
from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging
+from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload
from litellm.types.interactions import InteractionsAPIResponse
USAGE_BLOCK = {
@@ -74,6 +78,7 @@ def _context(
logging_obj: LitellmLogging,
timeout_seconds: float = 1.0,
store: Optional[InMemoryBackgroundSettlementStore] = None,
+ resumed: bool = False,
) -> BackgroundInteractionPollContext:
return BackgroundInteractionPollContext(
interaction_id="interactions/bg-abc",
@@ -82,7 +87,9 @@ def _context(
initial_interval_seconds=0.001,
max_interval_seconds=0.002,
timeout_seconds=timeout_seconds,
+ outcome_write_timeout_seconds=0.01,
store=store if store is not None else InMemoryBackgroundSettlementStore(),
+ resumed=resumed,
)
@@ -255,6 +262,180 @@ async def test_poller_retries_after_fetch_error_and_still_bills():
assert logging_obj.model_call_details["response_cost"] > 0
+def _provider_error_response(status_code: int) -> httpx.Response:
+ return httpx.Response(status_code, request=httpx.Request("GET", "https://generativelanguage.googleapis.com"))
+
+
+@pytest.mark.parametrize(
+ "error",
+ [
+ litellm.BadRequestError(message="invalid interaction id", model="gemini-2.5-flash", llm_provider="gemini"),
+ litellm.AuthenticationError(message="API key not valid", model="gemini-2.5-flash", llm_provider="gemini"),
+ litellm.PermissionDeniedError(
+ message="permission denied",
+ model="gemini-2.5-flash",
+ llm_provider="gemini",
+ response=_provider_error_response(403),
+ ),
+ litellm.NotFoundError(message="interaction not found", model="gemini-2.5-flash", llm_provider="gemini"),
+ ],
+)
+@pytest.mark.asyncio
+async def test_a_resumed_poll_stops_fetching_at_the_first_error_no_retry_changes_and_releases_the_reservation(error):
+ """
+ The regression: a resumed fetch the provider refuses for good, like an
+ unknown id or a key rotated away, was retried every interval until the
+ poll timed out.
+ """
+ reservation = _reservation()
+ logging_obj = _logging_obj_with_reservation(reservation)
+ fetch, calls = _fetch_sequence(error)
+
+ outcome = await poll_and_log_background_interaction_cost(
+ _context(logging_obj, timeout_seconds=0.05, resumed=True), fetch_interaction=fetch
+ )
+
+ assert outcome == "unsettled"
+ assert len(calls) == 1
+ assert reservation["finalized"] is True
+
+
+@pytest.mark.asyncio
+async def test_a_resumed_poll_the_provider_refuses_leaves_the_bill_to_the_replica_still_polling_it():
+ """
+ The regression: a replica that booted with credentials unable to read the
+ interaction settled it as unsettled on its first 404, and the replica that
+ created it, still polling, then found it settled and never billed it.
+ """
+ store = InMemoryBackgroundSettlementStore()
+ creator_logging_obj = _logging_obj()
+ resumed_fetched = asyncio.Event()
+
+ async def creator_fetch(context):
+ if resumed_fetched.is_set():
+ return _response("completed", with_usage=True)
+ return _response("in_progress", with_usage=False)
+
+ async def resumed_fetch(context):
+ resumed_fetched.set()
+ raise litellm.NotFoundError(message="interaction not found", model="gemini-2.5-flash", llm_provider="gemini")
+
+ creator_outcome, resumed_outcome = await asyncio.gather(
+ poll_and_log_background_interaction_cost(
+ _context(creator_logging_obj, store=store), fetch_interaction=creator_fetch
+ ),
+ poll_and_log_background_interaction_cost(
+ _context(_logging_obj(), store=store, resumed=True), fetch_interaction=resumed_fetch
+ ),
+ )
+
+ assert (creator_outcome, resumed_outcome) == ("billed", None)
+ assert creator_logging_obj.model_call_details["response_cost"] > 0
+
+
+@pytest.mark.parametrize(
+ "error",
+ [
+ litellm.RateLimitError(message="quota exceeded", model="gemini-2.5-flash", llm_provider="gemini"),
+ litellm.ServiceUnavailableError(message="overloaded", model="gemini-2.5-flash", llm_provider="gemini"),
+ litellm.InternalServerError(message="internal error", model="gemini-2.5-flash", llm_provider="gemini"),
+ litellm.Timeout(message="timed out", model="gemini-2.5-flash", llm_provider="gemini"),
+ litellm.BadRequestError(message="invalid interaction id", model="gemini-2.5-flash", llm_provider="gemini"),
+ litellm.AuthenticationError(message="API key not valid", model="gemini-2.5-flash", llm_provider="gemini"),
+ litellm.NotFoundError(message="interaction not found", model="gemini-2.5-flash", llm_provider="gemini"),
+ ],
+)
+@pytest.mark.asyncio
+async def test_the_creating_poll_retries_through_any_provider_error(error):
+ logging_obj = _logging_obj()
+ fetch, calls = _fetch_sequence(error, _response("completed", with_usage=True))
+
+ outcome = await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch)
+
+ assert outcome == "billed"
+ assert len(calls) == 2
+
+
+class _OutcomeWritesFail:
+ def __init__(self):
+ self.store = InMemoryBackgroundSettlementStore()
+
+ async def register(self, pending):
+ await self.store.register(pending)
+
+ async def pending(self, interaction_id):
+ return await self.store.pending(interaction_id)
+
+ async def is_claimed(self, interaction_id):
+ return await self.store.is_claimed(interaction_id)
+
+ async def claim(self, interaction_id):
+ return await self.store.claim(interaction_id)
+
+ async def record_outcome(self, interaction_id, outcome):
+ raise RuntimeError("database unavailable")
+
+ async def unsettled(self):
+ return await self.store.unsettled()
+
+
+@pytest.mark.parametrize(
+ "last_status, timeout_seconds",
+ [("failed", 1.0), ("in_progress", 0.01)],
+)
+@pytest.mark.asyncio
+async def test_a_release_it_could_not_record_leaves_the_reservation_to_the_settler_that_takes_over(
+ last_status, timeout_seconds
+):
+ """
+ The regression: a settler released the reservation and then failed to
+ record it, so the replica that took over the lapsed claim released the
+ same reservation again and the key's budget lost that amount twice.
+ """
+ reservation = _reservation()
+ logging_obj = _logging_obj_with_reservation(reservation)
+ fetch, _ = _fetch_sequence(_response(last_status, with_usage=False))
+
+ await poll_and_log_background_interaction_cost(
+ _context(logging_obj, timeout_seconds=timeout_seconds, store=_OutcomeWritesFail()),
+ fetch_interaction=fetch,
+ )
+
+ assert reservation["finalized"] is False
+
+
+class _FirstOutcomeWriteFails(_OutcomeWritesFail):
+ def __init__(self):
+ super().__init__()
+ self.failed_once = False
+
+ async def record_outcome(self, interaction_id, outcome):
+ if not self.failed_once:
+ self.failed_once = True
+ raise RuntimeError("connection reset")
+ await self.store.record_outcome(interaction_id, outcome)
+
+
+@pytest.mark.asyncio
+async def test_a_release_whose_first_outcome_write_failed_still_frees_the_reservation():
+ """
+ The regression: one failed outcome write left the key's budget
+ reservation held until a later boot or delete took the row over, and an
+ active key's spend counter never expired it in the meantime.
+ """
+ reservation = _reservation()
+ store = _FirstOutcomeWriteFails()
+ fetch, _ = _fetch_sequence(_response("failed", with_usage=False))
+
+ outcome = await poll_and_log_background_interaction_cost(
+ _context(_logging_obj_with_reservation(reservation), store=store), fetch_interaction=fetch
+ )
+
+ assert outcome == "released"
+ assert reservation["finalized"] is True
+ assert await store.store.unsettled() == ()
+
+
@pytest.mark.asyncio
async def test_schedule_creates_poll_task_for_in_progress_create():
logging_obj = _logging_obj()
@@ -302,6 +483,59 @@ async def test_schedule_registers_an_agent_only_create_that_names_no_model():
assert _rebuild_logging_obj(pending.create_context).model is None
+@pytest.mark.asyncio
+async def test_schedule_registers_a_create_whose_tags_are_a_string_and_keeps_them_as_sent():
+ """
+ The regression: storing the create's context read string tags as a list
+ and raised, so a create sent with a tags string answered with an error
+ after the provider had already started the interaction.
+ """
+ logging_obj = _logging_obj(litellm_params={"metadata": _create_metadata(tags="team-a,nightly-batch")})
+ store = InMemoryBackgroundSettlementStore()
+
+ task = await maybe_schedule_background_interaction_cost_polling(
+ response=_response("in_progress", with_usage=False),
+ create_kwargs={"litellm_logging_obj": logging_obj},
+ custom_llm_provider="gemini",
+ store=store,
+ )
+
+ assert isinstance(task, asyncio.Task)
+ task.cancel()
+ with pytest.raises(asyncio.CancelledError):
+ await task
+ pending = await store.pending("interactions/bg-abc")
+ assert pending is not None
+ rebuilt = _rebuild_logging_obj(pending.create_context)
+ assert get_litellm_metadata_from_kwargs(kwargs=rebuilt.model_call_details)["tags"] == "team-a,nightly-batch"
+
+
+@pytest.mark.asyncio
+async def test_delete_fetches_with_the_vertex_project_location_and_credentials_the_request_used():
+ store = InMemoryBackgroundSettlementStore()
+ await _create_on_a_replica_that_then_dies(_logging_obj(litellm_params={"metadata": _create_metadata()}), store)
+ fetch, captured = _capturing_fetch(_response("completed", with_usage=True))
+
+ outcome = await maybe_settle_background_interaction_before_delete(
+ interaction_id="interactions/bg-abc",
+ delete_kwargs={
+ "vertex_project": "request-project",
+ "vertex_location": "europe-west4",
+ "vertex_credentials": {"type": "service_account", "project_id": "request-project"},
+ },
+ fetch_interaction=fetch,
+ store=store,
+ )
+
+ assert outcome == "billed"
+ credentials = captured[0].credentials
+ assert (credentials.vertex_project, credentials.vertex_location, credentials.vertex_credentials) == (
+ "request-project",
+ "europe-west4",
+ {"type": "service_account", "project_id": "request-project"},
+ )
+
+
@pytest.mark.asyncio
@pytest.mark.parametrize(
"response,create_kwargs",
@@ -781,6 +1015,105 @@ async def test_restart_resumes_only_the_rows_no_replica_claimed():
assert await store.is_claimed("interactions/bg-orphaned")
+@pytest.mark.asyncio
+async def test_a_restart_long_after_the_create_still_retries_a_failed_fetch_before_giving_up():
+ """
+ The regression: a resumed poll kept only what was left of the original
+ window, so a row older than the window got a single fetch, and one
+ transient error right after the restart ended it unbilled.
+ """
+ store = InMemoryBackgroundSettlementStore()
+ await store.register(
+ PendingBackgroundInteraction(
+ interaction_id="interactions/bg-abc",
+ custom_llm_provider="gemini",
+ create_context=_create_context(_logging_obj(litellm_params={"metadata": _create_metadata()}), "gemini"),
+ created_at=datetime.now(timezone.utc) - timedelta(seconds=FAST_SCHEDULE.timeout_seconds * 100),
+ )
+ )
+ fetch, calls = _fetch_sequence(
+ RuntimeError("transient network error"),
+ _response("completed", with_usage=True),
+ )
+
+ (resumed,) = await resume_unsettled_background_interactions(store, fetch, schedule=FAST_SCHEDULE)
+
+ assert await asyncio.wait_for(resumed, timeout=5) == "billed"
+ assert len(calls) == 2
+
+
+def _create_through_the_proxy_router(tags: list) -> LitellmLogging:
+ logging_obj = LitellmLogging(
+ model="gemini-2.5-flash",
+ messages=[],
+ stream=False,
+ call_type="acreate_interaction",
+ start_time=time.time(),
+ litellm_call_id="bg-interactions-call-id",
+ function_id="bg-interactions-fn-id",
+ )
+ logging_obj.update_from_kwargs(
+ kwargs={
+ "metadata": _create_metadata(
+ tags=tags, spend_logs_metadata={"job": "nightly"}, model_group="client-chosen-group"
+ ),
+ "litellm_metadata": {"model_info": {"id": "deployment-1"}, "model_group": "gemini-background"},
+ },
+ model="gemini-2.5-flash",
+ optional_params={},
+ litellm_params={"litellm_call_id": "bg-interactions-call-id"},
+ custom_llm_provider="gemini",
+ )
+ return logging_obj
+
+
+def _create_with_header_tags(tags: list) -> LitellmLogging:
+ return _logging_obj(
+ litellm_params={
+ "metadata": _create_metadata(tags=tags, spend_logs_metadata={"job": "nightly"}),
+ "proxy_server_request": {"headers": {"user-agent": "claude-cli/2.1.0"}},
+ }
+ )
+
+
+async def _billed_spend_row(logging_obj: LitellmLogging) -> dict:
+ response = _response("completed", with_usage=True)
+ fetch, _ = _fetch_sequence(response)
+ assert await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch) == "billed"
+ settled_at = datetime.now(timezone.utc)
+ row = get_logging_payload(
+ kwargs={**logging_obj.model_call_details, "completion_start_time": settled_at},
+ response_obj=response,
+ start_time=settled_at,
+ end_time=settled_at,
+ )
+ return {
+ "request_tags": json.loads(row["request_tags"]),
+ "model_id": row["model_id"],
+ "model_group": row["model_group"],
+ "spend_logs_metadata": json.loads(row["metadata"]).get("spend_logs_metadata"),
+ }
+
+
+@pytest.mark.parametrize("create", [_create_through_the_proxy_router, _create_with_header_tags])
+@pytest.mark.asyncio
+async def test_a_bill_settled_off_the_creating_worker_matches_the_creating_worker_spend_row(create):
+ """
+ The regressions: the stored create context kept only the router's metadata
+ on the proxy path, so a row billed after a restart or by another replica
+ lost the request tags that tag-based spend reports group by; then it let
+ the client's request metadata override the router's, so a client-sent
+ model group relabeled the bill another replica wrote.
+ """
+ creator_row = await _billed_spend_row(create(["team-a", "nightly-batch"]))
+ stored = _create_context(create(["team-a", "nightly-batch"]), "gemini")
+
+ rebuilt_row = await _billed_spend_row(_rebuild_logging_obj(stored))
+
+ assert "team-a" in creator_row["request_tags"]
+ assert rebuilt_row == creator_row
+
+
class _ClaimAnswersOnlyAfterTheLastFetch:
def __init__(self):
self.store = InMemoryBackgroundSettlementStore()
@@ -809,8 +1142,8 @@ class _ClaimAnswersOnlyAfterTheLastFetch:
async def record_outcome(self, interaction_id, outcome):
return None
- async def unclaimed(self):
- return await self.store.unclaimed()
+ async def unsettled(self):
+ return await self.store.unsettled()
@pytest.mark.asyncio
@@ -844,7 +1177,7 @@ class _DownStore:
async def record_outcome(self, interaction_id, outcome):
raise RuntimeError("database unavailable")
- async def unclaimed(self):
+ async def unsettled(self):
raise RuntimeError("database unavailable")
@@ -868,8 +1201,8 @@ class _RegistersThenRaises:
async def record_outcome(self, interaction_id, outcome):
return None
- async def unclaimed(self):
- return await self.store.unclaimed()
+ async def unsettled(self):
+ return await self.store.unsettled()
@pytest.mark.asyncio
@@ -898,7 +1231,7 @@ async def test_create_whose_registration_raised_after_landing_still_claims_the_s
assert outcome == "billed"
assert await store.is_claimed("interactions/bg-abc")
- assert await store.unclaimed() == ()
+ assert await store.unsettled() == ()
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
@@ -907,7 +1240,7 @@ async def test_create_whose_registration_raised_after_landing_still_claims_the_s
@pytest.mark.asyncio
async def test_delete_on_a_worker_that_resumed_the_poll_fails_when_it_cannot_fetch():
"""
- After a restart every worker resumes the unclaimed rows, so none of them
+ After a restart every worker resumes the unsettled rows, so none of them
is the creator whose delete may release and delete on a failed fetch. A
resumed worker's delete fails like any other replica's, and its own poll
still bills the interaction once it completes.
@@ -965,9 +1298,9 @@ class _LandsThenGoesDown:
async def record_outcome(self, interaction_id, outcome):
return None
- async def unclaimed(self):
+ async def unsettled(self):
self._answer()
- return await self.store.unclaimed()
+ return await self.store.unsettled()
def _answer(self):
if self.down:
@@ -992,7 +1325,7 @@ class _TableLessStore:
async def record_outcome(self, interaction_id, outcome):
raise RuntimeError("the settlement table does not exist")
- async def unclaimed(self):
+ async def unsettled(self):
raise RuntimeError("the settlement table does not exist")
@@ -1027,7 +1360,7 @@ async def test_create_whose_registration_and_read_back_both_failed_bills_once_th
assert (while_down, recovered) == (None, "billed")
assert len(calls) == 2
assert await store.is_claimed("interactions/bg-abc")
- assert await store.unclaimed() == ()
+ assert await store.unsettled() == ()
assert await resume_unsettled_background_interactions(store, poll_fetch, schedule=FAST_SCHEDULE) == ()
task.cancel()
with pytest.raises(asyncio.CancelledError):
diff --git a/tests/unit/litellm_core_utils/llm_response_utils/test_get_headers.py b/tests/unit/litellm_core_utils/llm_response_utils/test_get_headers.py
index f52e837136a..1eccf5010ec 100644
--- a/tests/unit/litellm_core_utils/llm_response_utils/test_get_headers.py
+++ b/tests/unit/litellm_core_utils/llm_response_utils/test_get_headers.py
@@ -8,6 +8,7 @@ from litellm.litellm_core_utils.llm_response_utils.get_headers import (
_get_llm_provider_headers,
get_response_headers,
get_provider_request_id,
+ provider_request_id_headers,
)
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
@@ -89,6 +90,31 @@ def test_invalid_provider_request_ids_remain_absent(headers: object) -> None:
assert get_provider_request_id(headers) is None
+@pytest.mark.parametrize(
+ ("headers", "expected"),
+ (
+ (
+ {"X-Amzn-RequestId": "amzn-1", "retry-after": "7", "content-type": "text/event-stream"},
+ {"x-amzn-requestid": "amzn-1"},
+ ),
+ (
+ {"x-request-id": "req_1", "cf-ray": "ray-1", "x-ratelimit-remaining-requests": "0"},
+ {"x-request-id": "req_1", "cf-ray": "ray-1"},
+ ),
+ ({"retry-after": "7"}, None),
+ ({"x-request-id": 7, "request-id": None, "x-amzn-requestid": "amzn-2"}, {"x-amzn-requestid": "amzn-2"}),
+ ({"x-amzn-requestid": "amzn-\u00fc", "x-request-id": "req_2"}, {"x-request-id": "req_2"}),
+ ({}, None),
+ (None, None),
+ ),
+)
+def test_provider_request_id_headers_keeps_only_request_id_headers(
+ headers: dict[str, object] | None, expected: dict[str, str] | None
+) -> None:
+ result = provider_request_id_headers(headers)
+ assert (None if result is None else dict(result)) == expected
+
+
@pytest.fixture(autouse=True)
def _vcr_outcome_gate(request, vcr):
install_live_call_probe(request, vcr)
diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py
index dba1730208d..bf977325a85 100644
--- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py
+++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py
@@ -3,10 +3,12 @@ import json
import logging
import os
import re
+import sys
import time
-from collections.abc import Callable
+import types
+from collections.abc import Callable, Mapping, Sequence
from copy import deepcopy
-from typing import Final, List
+from typing import Final, List, TypeAlias
from unittest.mock import MagicMock, patch
import httpx
@@ -16,6 +18,7 @@ import respx
import litellm
from litellm.litellm_core_utils.prompt_templates.factory import (
BEDROCK_DOCUMENT_PLACEHOLDER_TEXT,
+ _EMPTY_TEXT_PLACEHOLDER,
BedrockConverseMessagesProcessor,
BedrockImageProcessor,
_bedrock_converse_messages_pt,
@@ -68,6 +71,17 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
from litellm.llms.vertex_ai.gemini.transformation import gemini_convert_messages_with_history
from litellm.types.llms.bedrock import BedrockConverseReasoningContentBlock
from litellm.types.llms.openai import AllMessageValues
+from litellm.litellm_core_utils.prompt_templates.factory import response_schema_prompt
+from litellm.llms.anthropic.completion.transformation import AnthropicTextConfig
+from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import AmazonInvokeConfig
+from litellm.llms.codestral.completion.handler import CodestralTextCompletion
+from litellm.llms.ollama.completion.transformation import OllamaConfig
+from litellm.llms.petals.completion import handler as petals_handler
+from litellm.llms.predibase.chat.transformation import PredibaseConfig
+from litellm.llms.sagemaker.completion.transformation import SagemakerConfig
+from litellm.llms.vllm.completion import handler as vllm_handler
+from litellm.llms.watsonx.common_utils import _convert_watsonx_messages_core
+from litellm.types.utils import ModelResponse, TextCompletionResponse
READ_FILE: Final[list[object]] = [{"name": "read_file", "description": "Read the contents of a file."}]
@@ -628,6 +642,64 @@ async def test_bedrock_converse_replays_reasoning_to_claude_only_as_base64_of_cl
assert assistant_content == [*sent, {"text": "hello"}]
+@pytest.mark.asyncio
+@pytest.mark.parametrize("use_async", [False, True])
+@pytest.mark.parametrize(
+ "model, final_text",
+ [("us.anthropic.claude-sonnet-4-6", "Leaves fall"), ("us.amazon.nova-pro-v1:0", "Leaves fall \n")],
+ ids=["claude", "nova"],
+)
+async def test_bedrock_converse_strips_trailing_whitespace_from_a_final_assistant_turn_for_claude_only(
+ use_async: bool, model: str, final_text: str
+):
+ """Claude on Converse 400s on a final assistant turn ending in whitespace, which a cut stream's prefill often has."""
+ messages: Final = [
+ {"role": "user", "content": "Write a haiku."},
+ {"role": "assistant", "content": "Sure. "},
+ {"role": "user", "content": "About autumn."},
+ {"role": "assistant", "content": "Leaves fall \n"},
+ ]
+
+ result: Final = (
+ await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async(
+ messages=messages, model=model, llm_provider="bedrock"
+ )
+ if use_async
+ else _bedrock_converse_messages_pt(messages=messages, model=model, llm_provider="bedrock")
+ )
+
+ assert [m["content"] for m in result if m["role"] == "assistant"] == [
+ [{"text": "Sure. "}],
+ [{"text": final_text}],
+ ]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("use_async", [False, True])
+async def test_bedrock_converse_keeps_whitespace_between_text_blocks_of_a_final_claude_assistant_turn(
+ use_async: bool,
+):
+ messages: Final = [
+ {"role": "user", "content": "Write a haiku."},
+ {
+ "role": "assistant",
+ "content": [{"type": "text", "text": "Leaves fall "}, {"type": "text", "text": "softly \n"}],
+ },
+ ]
+
+ result: Final = (
+ await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async(
+ messages=messages, model="us.anthropic.claude-sonnet-4-6", llm_provider="bedrock"
+ )
+ if use_async
+ else _bedrock_converse_messages_pt(
+ messages=messages, model="us.anthropic.claude-sonnet-4-6", llm_provider="bedrock"
+ )
+ )
+
+ assert result[-1]["content"] == [{"text": "Leaves fall "}, {"text": "softly"}]
+
+
def test_translate_thinking_blocks_keeps_dropping_redacted_data_when_called_the_old_way():
"""Outside callers that only pass the thinking blocks keep the old behavior, with no redactedContent sent."""
assert BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks(
@@ -795,6 +867,138 @@ def test_anthropic_messages_pt_keeps_signed_thinking_block():
assert thinking_blocks[0]["thinking"] == "genuine anthropic reasoning"
+EMPTY_USER_MESSAGES: Final = pytest.mark.parametrize(
+ "empty_user_message",
+ [{"role": "user", "content": None}, {"role": "user", "content": []}, {"role": "user"}],
+ ids=["null", "empty-list", "content-key-dropped"],
+)
+
+Message: TypeAlias = Mapping[str, object]
+
+
+def _convert(messages: Sequence[Message]) -> Sequence[Message]:
+ return anthropic_messages_pt(messages=list(messages), model="claude-opus-5", llm_provider="vertex_ai")
+
+
+@EMPTY_USER_MESSAGES
+def test_anthropic_messages_pt_preserves_trailing_user_message_with_empty_content(empty_user_message: Message) -> None:
+ result: Final = _convert(
+ [
+ {"role": "user", "content": "hi"},
+ {"role": "assistant", "content": "hello"},
+ empty_user_message,
+ ]
+ )
+
+ assert [message["role"] for message in result] == ["user", "assistant", "user"]
+ assert result[-1]["content"] == [{"type": "text", "text": _EMPTY_TEXT_PLACEHOLDER}]
+
+
+@EMPTY_USER_MESSAGES
+@pytest.mark.parametrize(
+ "cache_control", [{}, {"cache_control": {"type": "ephemeral"}}], ids=["no-cache-control", "with-cache-control"]
+)
+def test_anthropic_messages_pt_empty_user_message_beside_tool_result_keeps_tool_result_only(
+ empty_user_message: Message, cache_control: Message
+) -> None:
+ result: Final = _convert(
+ [
+ {"role": "user", "content": "hi"},
+ {
+ "role": "assistant",
+ "content": None,
+ "tool_calls": [{"id": "toolu_1", "type": "function", "function": {"name": "f", "arguments": "{}"}}],
+ },
+ {**empty_user_message, **cache_control},
+ {"role": "tool", "tool_call_id": "toolu_1", "content": "ok"},
+ ]
+ )
+
+ assert [message["role"] for message in result] == ["user", "assistant", "user"]
+ assert [block["type"] for block in result[-1]["content"]] == ["tool_result"]
+
+
+@EMPTY_USER_MESSAGES
+def test_anthropic_messages_pt_empty_user_message_keeps_cache_control_like_an_empty_string(
+ empty_user_message: Message,
+) -> None:
+ def last_turn(last_user_message: Message) -> Message:
+ return _convert(
+ [
+ {"role": "user", "content": "hi"},
+ {"role": "assistant", "content": "hello"},
+ {**last_user_message, "cache_control": {"type": "ephemeral"}},
+ ]
+ )[-1]
+
+ assert last_turn(empty_user_message)["content"][0]["cache_control"] == {"type": "ephemeral"}
+ assert last_turn(empty_user_message) == last_turn({"role": "user", "content": ""})
+
+
+TOOL_CALL_AND_RESULT: Final[tuple[Message, ...]] = (
+ {"role": "user", "content": "hi"},
+ {
+ "role": "assistant",
+ "content": None,
+ "tool_calls": [{"id": "toolu_1", "type": "function", "function": {"name": "f", "arguments": "{}"}}],
+ },
+ {"role": "tool", "tool_call_id": "toolu_1", "content": "ok"},
+)
+
+
+@EMPTY_USER_MESSAGES
+@pytest.mark.parametrize(
+ "messages_before",
+ [
+ [
+ {"role": "user", "content": "hi"},
+ {"role": "assistant", "content": "hello"},
+ {"role": "user", "content": "q"},
+ ],
+ [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}, {"role": "user", "content": []}],
+ TOOL_CALL_AND_RESULT,
+ ],
+ ids=["after-user-text", "after-another-empty-message", "after-tool-result"],
+)
+def test_anthropic_messages_pt_empty_user_message_beside_other_content_sends_nothing_extra(
+ empty_user_message: Message, messages_before: Sequence[Message]
+) -> None:
+ assert _convert([*messages_before, empty_user_message]) == _convert(messages_before)
+
+
+@EMPTY_USER_MESSAGES
+def test_anthropic_messages_pt_empty_user_message_with_cache_control_after_text_keeps_its_breakpoint(
+ empty_user_message: Message,
+) -> None:
+ def last_turn(last_user_message: Message) -> Message:
+ return _convert(
+ [
+ {"role": "user", "content": "hi"},
+ {"role": "assistant", "content": "hello"},
+ {"role": "user", "content": "long document"},
+ {**last_user_message, "cache_control": {"type": "ephemeral"}},
+ ]
+ )[-1]
+
+ assert last_turn(empty_user_message)["content"][-1]["cache_control"] == {"type": "ephemeral"}
+ assert last_turn(empty_user_message) == last_turn({"role": "user", "content": ""})
+
+
+def test_anthropic_messages_pt_keeps_every_breakpoint_of_a_long_run_of_empty_user_messages() -> None:
+ result: Final = _convert(
+ [
+ {"role": "user", "content": "hi"},
+ {"role": "assistant", "content": "hello"},
+ *[{"role": "user", "content": [], "cache_control": {"type": "ephemeral"}}] * 10_000,
+ ]
+ )
+
+ assert (
+ result[-1]["content"]
+ == [{"type": "text", "text": _EMPTY_TEXT_PLACEHOLDER, "cache_control": {"type": "ephemeral"}}] * 10_000
+ )
+
+
def test_convert_to_azure_openai_messages():
"""Test coverting image_url to azure_openai spec"""
@@ -4933,7 +5137,9 @@ def test_bedrock_converse_messages_pt_user_message_without_content_adds_no_block
assert _bedrock_converse_messages_pt(
messages=with_message, model="anthropic.claude-haiku-4-5", llm_provider="bedrock"
- ) == _bedrock_converse_messages_pt(messages=without_message, model="anthropic.claude-haiku-4-5", llm_provider="bedrock")
+ ) == _bedrock_converse_messages_pt(
+ messages=without_message, model="anthropic.claude-haiku-4-5", llm_provider="bedrock"
+ )
@pytest.mark.asyncio
@@ -5063,6 +5269,7 @@ def _vcr_outcome_gate(request, vcr):
yield
record_vcr_outcome(request, vcr)
+
@pytest.fixture(scope="function")
def isolate_litellm_state():
"""
@@ -5115,6 +5322,7 @@ def isolate_litellm_state():
setattr(litellm, attr, original_value)
_invalidate_model_cost_lowercase_map()
+
_SCALAR_DEFAULTS = {
"num_retries": getattr(litellm, "num_retries", None),
"num_retries_per_request": getattr(litellm, "num_retries_per_request", None),
@@ -5135,6 +5343,7 @@ _SCALAR_DEFAULTS = {
"api_key": getattr(litellm, "api_key", None),
}
+
@pytest.fixture(scope="module")
def setup_and_teardown():
"""
@@ -5157,6 +5366,7 @@ def setup_and_teardown():
litellm.in_memory_llm_clients_cache.flush_cache()
yield
+
@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown")
def test_function_call_non_openai_model():
try:
@@ -5185,6 +5395,7 @@ def test_function_call_non_openai_model():
print(e)
pass
+
@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown")
def test_empty_content():
"""
@@ -5203,6 +5414,7 @@ def test_empty_content():
litellm_call_id=str(uuid.uuid4()),
)
+
@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown")
def test_thought_signature_removal_for_non_gemini():
"""
@@ -5251,6 +5463,7 @@ def test_thought_signature_removal_for_non_gemini():
assert THOUGHT_SIGNATURE_SEPARATOR not in processed_messages[1]["tool_calls"][0]["id"]
assert THOUGHT_SIGNATURE_SEPARATOR not in processed_messages[2]["tool_call_id"]
+
@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown")
def test_thought_signature_preserved_for_gemini():
"""
@@ -5297,6 +5510,7 @@ def test_thought_signature_preserved_for_gemini():
assert THOUGHT_SIGNATURE_SEPARATOR in processed_messages[1]["tool_calls"][0]["id"]
assert THOUGHT_SIGNATURE_SEPARATOR in processed_messages[2]["tool_call_id"]
+
@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown")
def test_thought_signature_removal_with_multiple_tool_calls():
"""
@@ -5351,6 +5565,7 @@ def test_thought_signature_removal_with_multiple_tool_calls():
assert processed_messages[2]["tool_call_id"] == "call_1"
assert processed_messages[3]["tool_call_id"] == "call_2"
+
@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown")
def test_messages_without_tool_calls_unchanged():
"""
@@ -7727,3 +7942,334 @@ def test_anthropic_messages_pt_merges_a_long_assistant_run_in_linear_time():
growth: Final = _merged_assistant_run_cost(4000) / _merged_assistant_run_cost(1000)
assert growth < 8.0, f"cost grew {growth:.1f}x for 4x the assistant run"
+
+
+TEMPLATE_MODEL = "partial-template-model"
+TEMPLATE_ROLES = {"user": {"pre_message": "[INST] ", "post_message": " [/INST]"}}
+TEMPLATE_MESSAGES = [{"role": "user", "content": "hi"}]
+TEMPLATE_RENDERED_MESSAGES = "[INST] hi [/INST]"
+TEMPLATE_BARE_MESSAGES = "hi"
+
+TEMPLATE_RESPONSE_SCHEMA = {"type": "object"}
+TEMPLATE_RENDERED_SCHEMA = f"[INST] {TEMPLATE_RESPONSE_SCHEMA} [/INST]"
+TEMPLATE_BARE_SCHEMA = str(TEMPLATE_RESPONSE_SCHEMA)
+
+TEMPLATE_INITIAL = ""
+TEMPLATE_FINAL = ""
+
+PARTIAL_TEMPLATE = {"roles": TEMPLATE_ROLES}
+FULL_TEMPLATE = {"roles": TEMPLATE_ROLES, "initial_prompt_value": TEMPLATE_INITIAL, "final_prompt_value": TEMPLATE_FINAL}
+ROLELESS_TEMPLATE = {"initial_prompt_value": TEMPLATE_INITIAL, "final_prompt_value": TEMPLATE_FINAL}
+
+
+class _PromptCaptured(Exception):
+ def __init__(self, prompt):
+ super().__init__(prompt)
+ self.prompt = prompt
+
+
+class _CapturingLogging:
+ def pre_call(self, input, api_key, additional_args=None, **kwargs):
+ raise _PromptCaptured(input)
+
+ def post_call(self, *args, **kwargs):
+ pass
+
+
+def _captured_prompt(call):
+ with pytest.raises(_PromptCaptured) as excinfo:
+ call()
+ return excinfo.value.prompt
+
+
+def _stub_vllm_import(monkeypatch):
+ class _SamplingParams:
+ def __init__(self, **kwargs):
+ self.kwargs = kwargs
+
+ class _LLM:
+ def __init__(self, model):
+ self.model = model
+
+ def generate(self, prompts, sampling_params):
+ raise _PromptCaptured(prompts)
+
+ stub = types.ModuleType("vllm")
+ stub.LLM = _LLM
+ stub.SamplingParams = _SamplingParams
+
+ monkeypatch.setitem(sys.modules, "vllm", stub)
+ monkeypatch.setattr(vllm_handler, "llm", None)
+
+
+def _render_ollama(template, monkeypatch, messages=TEMPLATE_MESSAGES):
+ request = OllamaConfig().transform_request(
+ model=TEMPLATE_MODEL,
+ messages=messages,
+ optional_params={},
+ litellm_params={"custom_prompt_dict": {TEMPLATE_MODEL: template}},
+ headers={},
+ )
+ return request["prompt"]
+
+
+def _render_predibase(template, monkeypatch):
+ request = PredibaseConfig().transform_request(
+ model=TEMPLATE_MODEL,
+ messages=TEMPLATE_MESSAGES,
+ optional_params={},
+ litellm_params={"custom_prompt_dict": {TEMPLATE_MODEL: template}},
+ headers={},
+ )
+ return request["inputs"]
+
+
+def _render_anthropic_text(template, monkeypatch):
+ monkeypatch.setattr(litellm, "custom_prompt_dict", {TEMPLATE_MODEL: template})
+ return AnthropicTextConfig()._get_anthropic_text_prompt_from_messages(
+ messages=TEMPLATE_MESSAGES, model=TEMPLATE_MODEL
+ )
+
+
+def _render_response_schema_prompt(template, monkeypatch):
+ monkeypatch.setattr(litellm, "custom_prompt_dict", {f"{TEMPLATE_MODEL}/response_schema_prompt": template})
+ return response_schema_prompt(
+ model=TEMPLATE_MODEL, response_schema=TEMPLATE_RESPONSE_SCHEMA, llm_provider="vertex_ai"
+ )
+
+
+def _render_petals(template, monkeypatch):
+ monkeypatch.setattr(litellm, "custom_prompt_dict", {TEMPLATE_MODEL: template})
+ return _captured_prompt(
+ lambda: petals_handler.completion(
+ model=TEMPLATE_MODEL,
+ messages=TEMPLATE_MESSAGES,
+ api_base="http://petals.invalid",
+ model_response=ModelResponse(),
+ print_verbose=lambda *args, **kwargs: None,
+ encoding=None,
+ logging_obj=_CapturingLogging(),
+ optional_params={},
+ )
+ )
+
+
+def _render_codestral(template, monkeypatch):
+ return _captured_prompt(
+ lambda: CodestralTextCompletion().completion(
+ model=TEMPLATE_MODEL,
+ messages=TEMPLATE_MESSAGES,
+ api_base="http://codestral.invalid",
+ custom_prompt_dict={TEMPLATE_MODEL: template},
+ model_response=TextCompletionResponse(),
+ print_verbose=lambda *args, **kwargs: None,
+ encoding=None,
+ api_key="sk-not-used",
+ logging_obj=_CapturingLogging(),
+ optional_params={},
+ timeout=1.0,
+ )
+ )
+
+
+def _render_vllm_completion(template, monkeypatch):
+ _stub_vllm_import(monkeypatch)
+ return _captured_prompt(
+ lambda: vllm_handler.completion(
+ model=TEMPLATE_MODEL,
+ messages=TEMPLATE_MESSAGES,
+ model_response=ModelResponse(),
+ print_verbose=lambda *args, **kwargs: None,
+ encoding=None,
+ logging_obj=_CapturingLogging(),
+ optional_params={},
+ custom_prompt_dict={TEMPLATE_MODEL: template},
+ )
+ )
+
+
+def _render_vllm_batch(template, monkeypatch):
+ _stub_vllm_import(monkeypatch)
+ prompts = _captured_prompt(
+ lambda: vllm_handler.batch_completions(
+ model=TEMPLATE_MODEL,
+ messages=[TEMPLATE_MESSAGES],
+ optional_params={},
+ custom_prompt_dict={TEMPLATE_MODEL: template},
+ )
+ )
+ assert len(prompts) == 1
+ return prompts[0]
+
+
+def _render_bedrock_invoke(template, monkeypatch):
+ prompt, _ = AmazonInvokeConfig().convert_messages_to_prompt(
+ model=TEMPLATE_MODEL, messages=TEMPLATE_MESSAGES, provider="meta", custom_prompt_dict={TEMPLATE_MODEL: template}
+ )
+ return prompt
+
+
+def _render_sagemaker(template, monkeypatch):
+ return SagemakerConfig()._transform_prompt(
+ model=TEMPLATE_MODEL,
+ messages=TEMPLATE_MESSAGES,
+ custom_prompt_dict={TEMPLATE_MODEL: template},
+ hf_model_name=None,
+ )
+
+
+def _render_watsonx(template, monkeypatch):
+ return _convert_watsonx_messages_core(
+ model=TEMPLATE_MODEL,
+ messages=TEMPLATE_MESSAGES,
+ provider="ibm",
+ custom_prompt_dict={TEMPLATE_MODEL: template},
+ apply_template_fn=None,
+ )
+
+
+TEMPLATE_CALL_SITES = [
+ pytest.param(_render_ollama, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="ollama"),
+ pytest.param(_render_petals, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="petals"),
+ pytest.param(_render_vllm_completion, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="vllm-completion"),
+ pytest.param(_render_vllm_batch, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="vllm-batch-completions"),
+ pytest.param(_render_predibase, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="predibase"),
+ pytest.param(_render_codestral, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="codestral"),
+ pytest.param(_render_anthropic_text, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="anthropic-text"),
+ pytest.param(
+ _render_response_schema_prompt, TEMPLATE_RENDERED_SCHEMA, TEMPLATE_BARE_SCHEMA, id="response-schema-prompt"
+ ),
+ pytest.param(_render_bedrock_invoke, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="bedrock-invoke"),
+ pytest.param(_render_sagemaker, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="sagemaker"),
+ pytest.param(_render_watsonx, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="watsonx"),
+]
+
+
+@pytest.mark.parametrize(("render", "rendered_body", "bare_body"), TEMPLATE_CALL_SITES)
+def test_template_without_the_optional_prompt_values_renders(render, rendered_body, bare_body, monkeypatch):
+ assert render(PARTIAL_TEMPLATE, monkeypatch) == rendered_body
+
+
+@pytest.mark.parametrize(("render", "rendered_body", "bare_body"), TEMPLATE_CALL_SITES)
+def test_template_with_the_optional_prompt_values_still_applies_them(render, rendered_body, bare_body, monkeypatch):
+ assert render(FULL_TEMPLATE, monkeypatch) == TEMPLATE_INITIAL + rendered_body + TEMPLATE_FINAL
+
+
+@pytest.mark.parametrize(("render", "rendered_body", "bare_body"), TEMPLATE_CALL_SITES)
+def test_template_without_roles_renders_messages_unwrapped(render, rendered_body, bare_body, monkeypatch):
+ assert render(ROLELESS_TEMPLATE, monkeypatch) == TEMPLATE_INITIAL + bare_body + TEMPLATE_FINAL
+
+
+def test_template_wraps_only_the_roles_it_defines_across_a_conversation(monkeypatch):
+ messages = [
+ {"role": "system", "content": "Be terse."},
+ {"role": "user", "content": "hi"},
+ {"role": "assistant", "content": "hello"},
+ {"role": "user", "content": "sky colour?"},
+ ]
+
+ rendered = _render_ollama(FULL_TEMPLATE, monkeypatch, messages=messages)
+
+ assert rendered == "Be terse.[INST] hi [/INST]hello[INST] sky colour? [/INST]"
+
+
+@pytest.mark.parametrize(
+ "template",
+ [
+ pytest.param({"roles": {**TEMPLATE_ROLES, "tool": None}}, id="null-role-entry"),
+ pytest.param({"roles": {"user": {**TEMPLATE_ROLES["user"], "priority": 1}}}, id="extra-role-key"),
+ pytest.param({"roles": {**TEMPLATE_ROLES, "tool": {"pre_message": 5}}}, id="unused-role-wrong-type"),
+ pytest.param({"roles": {**TEMPLATE_ROLES, 1: {"pre_message": "x"}}}, id="unused-non-string-role"),
+ pytest.param({"roles": TEMPLATE_ROLES, "initial_prompt_value": None, "bos_token": 7}, id="null-and-unread-keys"),
+ ],
+)
+def test_template_shapes_main_rendered_keep_rendering(template, monkeypatch):
+ assert _render_ollama(template, monkeypatch) == TEMPLATE_RENDERED_MESSAGES
+
+
+@pytest.mark.parametrize(
+ "template",
+ [
+ pytest.param({"roles": {"user": "[INST]"}}, id="string-role-entry"),
+ pytest.param({"roles": {"user": ["[INST]"]}}, id="list-role-entry"),
+ pytest.param({"roles": ["user"]}, id="list-roles"),
+ pytest.param({"roles": '{"user": {"pre_message": "[INST]"}}'}, id="json-string-roles"),
+ ],
+)
+def test_roles_that_are_not_a_mapping_render_unwrapped_with_a_warning(template, monkeypatch, caplog):
+ with caplog.at_level(logging.WARNING, logger="LiteLLM"):
+ rendered = _render_ollama({**template, "final_prompt_value": TEMPLATE_FINAL}, monkeypatch)
+
+ assert rendered == TEMPLATE_BARE_MESSAGES + TEMPLATE_FINAL
+ assert "not a mapping" in caplog.text
+
+
+def test_null_role_entry_renders_that_role_unwrapped_without_a_warning(monkeypatch, caplog):
+ with caplog.at_level(logging.WARNING, logger="LiteLLM"):
+ rendered = _render_ollama({"roles": {"user": None}}, monkeypatch)
+
+ assert rendered == TEMPLATE_BARE_MESSAGES
+ assert "not a mapping" not in caplog.text
+
+
+@pytest.mark.parametrize(
+ ("template", "field"),
+ [
+ pytest.param({"initial_prompt_value": 5}, "initial_prompt_value", id="initial-prompt-value"),
+ pytest.param({"final_prompt_value": ["x"]}, "final_prompt_value", id="final-prompt-value"),
+ pytest.param({"roles": {"user": {"pre_message": 5}}}, "roles.user.pre_message", id="pre-message"),
+ ],
+)
+def test_wrong_type_template_field_fails_as_a_server_error_naming_it(template, field):
+ with pytest.raises(litellm.APIConnectionError) as excinfo:
+ litellm.completion(
+ model=f"ollama/{TEMPLATE_MODEL}",
+ messages=TEMPLATE_MESSAGES,
+ api_base="http://ollama.invalid",
+ **template,
+ )
+
+ assert excinfo.value.status_code == 500
+ assert f"{field}: Input should be a valid string" in excinfo.value.message
+
+
+@pytest.mark.parametrize(
+ "template",
+ [
+ pytest.param({"bos_token": ""}, id="bos-token-only"),
+ pytest.param({"eos_token": ""}, id="eos-token-only"),
+ pytest.param({"roles": '{"user": {"pre_message": "[INST]"}}'}, id="json-string-roles-only"),
+ ],
+)
+def test_template_that_sets_nothing_the_prompt_uses_fails_as_a_server_error(template):
+ with pytest.raises(litellm.APIConnectionError) as excinfo:
+ litellm.completion(
+ model=f"ollama/{TEMPLATE_MODEL}",
+ messages=TEMPLATE_MESSAGES,
+ api_base="http://ollama.invalid",
+ **template,
+ )
+
+ assert excinfo.value.status_code == 500
+ assert "is invalid: it sets none of roles, initial_prompt_value, or final_prompt_value" in excinfo.value.message
+
+
+def test_wrong_type_template_field_in_requests_leaves_the_deployments_in_rotation():
+ router = litellm.Router(
+ model_list=[
+ {
+ "model_name": "template-group",
+ "litellm_params": {"model": f"ollama/{TEMPLATE_MODEL}", "api_base": "http://ollama.invalid"},
+ }
+ for _ in range(2)
+ ],
+ num_retries=0,
+ allowed_fails=0,
+ )
+
+ for _ in range(2):
+ with pytest.raises(litellm.APIConnectionError):
+ router.completion(model="template-group", messages=TEMPLATE_MESSAGES, initial_prompt_value=5)
+ response = router.completion(model="template-group", messages=TEMPLATE_MESSAGES, mock_response="Blue")
+
+ assert response.choices[0].message.content == "Blue"
diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py
index 6bc1178f64b..91d7acd331e 100644
--- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py
+++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py
@@ -29,11 +29,8 @@ def _texts(message: dict) -> list[str]:
return [block["text"] for block in message["content"]]
-SENDS_NOTHING = pytest.mark.parametrize(
- "empty_content",
- [[], None, [{"type": "input_audio", "input_audio": {"data": "AAAA", "format": "wav"}}]],
- ids=["empty-list", "none", "unsupported-part-only"],
-)
+EMPTY_USER_CONTENT: Final = pytest.mark.parametrize("empty_content", [[], None], ids=["empty-list", "none"])
+UNSUPPORTED_PART_ONLY: Final = ({"type": "input_audio", "input_audio": {"data": "AAAA", "format": "wav"}},)
def test_split_leading_system_run_keeps_later_system_messages_in_the_conversation():
@@ -213,9 +210,9 @@ def test_flagged_placement_of_an_earlier_run_does_not_move_when_later_turns_are_
assert _roles(placed_n_plus_one) == ["user", "assistant", "user", "assistant", "user"]
-@SENDS_NOTHING
-def test_flagged_placement_converts_a_run_whose_preceding_user_turn_sends_nothing(empty_content):
- placed = place_mid_conversation_system(
+@EMPTY_USER_CONTENT
+def test_flagged_placement_anchors_a_run_on_a_preceding_empty_user_turn(empty_content: list[object] | None) -> None:
+ placed: Final = place_mid_conversation_system(
[
{"role": "user", "content": empty_content},
{"role": "system", "content": "reminder"},
@@ -225,18 +222,48 @@ def test_flagged_placement_converts_a_run_whose_preceding_user_turn_sends_nothin
supports_mid_conversation_system=True,
)
+ assert _roles(placed) == ["user", "system", "assistant", "user"]
+ assert _texts(placed[1]) == ["reminder"]
+
+
+@EMPTY_USER_CONTENT
+def test_flagged_placement_anchors_a_run_on_a_following_empty_user_turn(empty_content: list[object] | None) -> None:
+ placed: Final = place_mid_conversation_system(
+ [
+ {"role": "user", "content": "q1"},
+ {"role": "assistant", "content": "a1"},
+ {"role": "system", "content": "reminder"},
+ {"role": "user", "content": empty_content},
+ ],
+ supports_mid_conversation_system=True,
+ )
+
+ assert _roles(placed) == ["user", "assistant", "user", "system"]
+ assert _texts(placed[3]) == ["reminder"]
+
+
+def test_flagged_placement_converts_a_run_whose_preceding_user_turn_sends_nothing() -> None:
+ placed: Final = place_mid_conversation_system(
+ [
+ {"role": "user", "content": list(UNSUPPORTED_PART_ONLY)},
+ {"role": "system", "content": "reminder"},
+ {"role": "assistant", "content": "a1"},
+ {"role": "user", "content": "q2"},
+ ],
+ supports_mid_conversation_system=True,
+ )
+
assert _roles(placed) == ["user", "user", "assistant", "user"]
assert _texts(placed[1]) == [CONVERTED_SYSTEM_NOTE, "reminder"]
-@SENDS_NOTHING
-def test_flagged_placement_converts_a_run_whose_following_user_turn_sends_nothing(empty_content):
- placed = place_mid_conversation_system(
+def test_flagged_placement_converts_a_run_whose_following_user_turn_sends_nothing() -> None:
+ placed: Final = place_mid_conversation_system(
[
{"role": "user", "content": "q1"},
{"role": "assistant", "content": "a1"},
{"role": "system", "content": "reminder"},
- {"role": "user", "content": empty_content},
+ {"role": "user", "content": list(UNSUPPORTED_PART_ONLY)},
],
supports_mid_conversation_system=True,
)
diff --git a/tests/unit/litellm_core_utils/test_realtime_streaming.py b/tests/unit/litellm_core_utils/test_realtime_streaming.py
index 8de5e27d99c..9b2c42450ae 100644
--- a/tests/unit/litellm_core_utils/test_realtime_streaming.py
+++ b/tests/unit/litellm_core_utils/test_realtime_streaming.py
@@ -2351,8 +2351,7 @@ async def test_guardrail_turn_detection_injection_tolerates_non_dict_value(
injected_turn_detection = session_obj.get("turn_detection") or session_obj.get("audio", {}).get("input", {}).get(
"turn_detection"
)
- assert isinstance(injected_turn_detection, dict)
- assert injected_turn_detection["create_response"] is False
+ assert injected_turn_detection == {"type": "server_vad", "create_response": False}
assert streaming._guardrail_turn_detection_update_sent is True
@@ -2417,8 +2416,7 @@ async def test_subsequent_session_update_cannot_reenable_vad_when_guardrails_act
forwarded_turn_detection = session_obj.get("turn_detection") or session_obj.get("audio", {}).get("input", {}).get(
"turn_detection"
)
- assert isinstance(forwarded_turn_detection, dict)
- assert forwarded_turn_detection["create_response"] is False
+ assert forwarded_turn_detection == {"type": "server_vad", "create_response": False}
@pytest.mark.asyncio
@@ -3927,3 +3925,317 @@ async def test_transcription_session_guardrail_block_only_reports_violation(
sent_to_backend: Final = [call.args[0] for call in backend_ws.send.await_args_list]
assert sent_to_backend == [], sent_to_backend
assert backend_ws.close.await_count == (1 if expect_session_closed else 0)
+
+
+_GA_GUARDRAIL_GATE: Final = {
+ "type": "realtime",
+ "audio": {"input": {"turn_detection": {"type": "server_vad", "create_response": False}}},
+}
+
+
+async def _forwarded_ga_session_updates(
+ monkeypatch: pytest.MonkeyPatch,
+ *client_sessions: dict[str, object],
+ event_normalizer: XAIRealtimeNormalizer | None = None,
+) -> tuple[dict[str, object], ...]:
+ """Run a GA client's session.updates through a transcript-guardrailed proxy and return the forwarded sessions."""
+ monkeypatch.setattr(litellm, "callbacks", [_transcription_guardrail()])
+ client_ws: Final = AsyncMock()
+ client_ws.scope = {"headers": []}
+ client_ws.receive_text = AsyncMock(
+ side_effect=[
+ *(json.dumps({"type": "session.update", "session": session}) for session in client_sessions),
+ ConnectionClosed(None, None),
+ ]
+ )
+ backend_ws: Final = MagicMock()
+ backend_ws.send = AsyncMock()
+ logging_obj: Final = MagicMock()
+ logging_obj.async_success_handler = AsyncMock()
+ logging_obj.success_handler = MagicMock()
+ await RealTimeStreaming(client_ws, backend_ws, logging_obj, event_normalizer=event_normalizer).client_ack_messages()
+ return tuple(json.loads(call.args[0])["session"] for call in backend_ws.send.await_args_list)
+
+
+@pytest.mark.asyncio
+async def test_first_ga_session_update_keeps_client_turn_detection_when_guardrail_disables_auto_response(
+ monkeypatch: pytest.MonkeyPatch,
+):
+ """A transcript guardrail only flips create_response: the client's nested GA
+ turn_detection, its type included, must reach the backend intact, or OpenAI
+ rejects the whole update and the client's session config is lost."""
+ forwarded: Final = await _forwarded_ga_session_updates(
+ monkeypatch,
+ {
+ "type": "realtime",
+ "instructions": "hi",
+ "audio": {
+ "input": {
+ "turn_detection": {"type": "server_vad", "silence_duration_ms": 700, "create_response": True},
+ "transcription": {"model": "gpt-4o-transcribe"},
+ }
+ },
+ },
+ )
+
+ gate, client_update = forwarded
+ assert gate == _GA_GUARDRAIL_GATE
+ assert "turn_detection" not in client_update
+ assert client_update["instructions"] == "hi"
+ assert client_update["audio"]["input"] == {
+ "turn_detection": {"type": "server_vad", "silence_duration_ms": 700, "create_response": False},
+ "transcription": {"model": "gpt-4o-transcribe"},
+ }
+
+
+@pytest.mark.asyncio
+async def test_first_ga_session_update_without_turn_detection_gets_a_typed_gate(monkeypatch: pytest.MonkeyPatch):
+ forwarded: Final = await _forwarded_ga_session_updates(monkeypatch, {"type": "realtime", "instructions": "hi"})
+
+ gate, client_update = forwarded
+ assert gate == _GA_GUARDRAIL_GATE
+ assert "turn_detection" not in client_update
+ assert client_update["audio"]["input"]["turn_detection"] == {"type": "server_vad", "create_response": False}
+
+
+@pytest.mark.asyncio
+async def test_later_ga_session_update_without_turn_detection_is_forwarded_without_one(
+ monkeypatch: pytest.MonkeyPatch,
+):
+ """OpenAI keeps the earlier create_response=false for a field left unset, while a
+ bare create_response carries no type and gets the whole update rejected."""
+ forwarded: Final = await _forwarded_ga_session_updates(
+ monkeypatch,
+ {"type": "realtime", "audio": {"input": {"turn_detection": {"type": "server_vad", "create_response": True}}}},
+ {"type": "realtime", "instructions": "later"},
+ )
+
+ gate, first, later = forwarded
+ assert gate == _GA_GUARDRAIL_GATE
+ assert first["audio"]["input"]["turn_detection"] == {"type": "server_vad", "create_response": False}
+ assert later == {"type": "realtime", "instructions": "later"}
+
+
+_XAI_TYPED_VAD_SESSION: Final = {
+ "type": "realtime",
+ "audio": {"input": {"turn_detection": {"type": "server_vad", "silence_duration_ms": 700}}},
+}
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "client_sessions",
+ [
+ pytest.param(({"type": "realtime", "instructions": "hi"},), id="no-turn-detection"),
+ pytest.param(({"type": "realtime", "audio": {"input": {"turn_detection": None}}},), id="push-to-talk"),
+ pytest.param((_XAI_TYPED_VAD_SESSION,), id="typed-server-vad"),
+ pytest.param(
+ ({"type": "realtime", "instructions": "hi"}, _XAI_TYPED_VAD_SESSION), id="later-typed-server-vad"
+ ),
+ ],
+)
+async def test_xai_session_update_gate_leaves_out_the_vad_type(
+ monkeypatch: pytest.MonkeyPatch, client_sessions: tuple[dict[str, object], ...]
+):
+ """xAI answers on its own under a typed turn_detection even with create_response off,
+ so every turn_detection the proxy forwards to xAI under the guardrail must carry no type."""
+ forwarded: Final = await _forwarded_ga_session_updates(
+ monkeypatch, *client_sessions, event_normalizer=XAIRealtimeNormalizer()
+ )
+
+ assert [session["audio"]["input"]["turn_detection"] for session in forwarded] == [
+ {"create_response": False}
+ ] * (len(client_sessions) + 1)
+
+
+@pytest.mark.asyncio
+async def test_xai_session_created_gate_leaves_out_the_vad_type(monkeypatch: pytest.MonkeyPatch):
+ monkeypatch.setattr(litellm, "callbacks", [_transcription_guardrail()])
+ client_ws: Final = MagicMock()
+ client_ws.send_text = AsyncMock()
+ backend_ws: Final = MagicMock()
+ backend_ws.send = AsyncMock()
+ backend_ws.recv = AsyncMock(
+ side_effect=[
+ json.dumps({"type": "session.created", "session": {"turn_detection": {"type": None}}}).encode(),
+ ConnectionClosed(None, None),
+ ]
+ )
+ logging_obj: Final = _mock_logging()
+ logging_obj.async_success_handler = AsyncMock()
+ logging_obj.success_handler = MagicMock()
+ streaming: Final = RealTimeStreaming(client_ws, backend_ws, logging_obj, event_normalizer=XAIRealtimeNormalizer())
+
+ await streaming.backend_to_client_send_messages()
+
+ sent: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list]
+ assert sent == [
+ {
+ "type": "session.update",
+ "session": {"type": "realtime", "audio": {"input": {"turn_detection": {"create_response": False}}}},
+ }
+ ]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "client_turn_detection",
+ [
+ pytest.param({"type": "semantic_vad", "eagerness": "low"}, id="semantic-vad"),
+ pytest.param({"type": "server_vad", "silence_duration_ms": 700}, id="tuned-server-vad"),
+ ],
+)
+async def test_guardrail_gate_goes_ahead_of_a_session_update_sent_before_session_created(
+ monkeypatch: pytest.MonkeyPatch, client_turn_detection: dict[str, object]
+):
+ """OpenAI applies events in order, so the gate must precede a client update sent on connect:
+ the client's turn_detection then wins, and when OpenAI refuses that update the gate still holds."""
+ monkeypatch.setattr(litellm, "callbacks", [_transcription_guardrail()])
+ client_ws: Final = AsyncMock()
+ client_ws.scope = {"headers": []}
+ client_ws.receive_text = AsyncMock(
+ side_effect=[
+ json.dumps(
+ {
+ "type": "session.update",
+ "session": {"type": "realtime", "audio": {"input": {"turn_detection": client_turn_detection}}},
+ }
+ ),
+ ConnectionClosed(None, None),
+ ]
+ )
+ backend_ws: Final = MagicMock()
+ backend_ws.send = AsyncMock()
+ backend_ws.recv = AsyncMock(
+ side_effect=[json.dumps({"type": "session.created", "session": {}}).encode(), ConnectionClosed(None, None)]
+ )
+ logging_obj: Final = _mock_logging()
+ logging_obj.async_success_handler = AsyncMock()
+ logging_obj.success_handler = MagicMock()
+ streaming: Final = RealTimeStreaming(client_ws, backend_ws, logging_obj)
+
+ await streaming.client_ack_messages()
+ await streaming.backend_to_client_send_messages()
+
+ sent: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list]
+ assert [event["session"]["audio"]["input"]["turn_detection"] for event in sent] == [
+ {"type": "server_vad", "create_response": False},
+ {**client_turn_detection, "create_response": False},
+ ]
+
+
+@pytest.mark.asyncio
+async def test_client_never_sees_the_error_caused_by_the_guardrail_cancel(monkeypatch: pytest.MonkeyPatch):
+ """OpenAI answers the guardrail's response.cancel with response_cancel_not_active when nothing
+ is answering yet, an error about an event the client never sent."""
+
+ class _BlockingGuardrail(CustomGuardrail):
+ async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
+ raise ValueError("blocked")
+
+ monkeypatch.setattr(
+ litellm,
+ "callbacks",
+ [
+ _BlockingGuardrail(
+ guardrail_name="blocking", event_hook=GuardrailEventHooks.realtime_input_transcription, default_on=True
+ )
+ ],
+ )
+ client_ws: Final = MagicMock()
+ client_ws.send_text = AsyncMock()
+ backend_ws: Final = MagicMock()
+ backend_ws.send = AsyncMock()
+
+ def cancel_not_active(event_id: str) -> bytes:
+ return json.dumps(
+ {"type": "error", "error": {"code": "response_cancel_not_active", "event_id": event_id}}
+ ).encode()
+
+ def proxy_cancel_event_id() -> str:
+ sent = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list]
+ return next(event["event_id"] for event in sent if event["type"] == "response.cancel")
+
+ backend_frames: Final = iter(
+ (
+ lambda: _make_transcript_event("blocked words"),
+ lambda: cancel_not_active(proxy_cancel_event_id()),
+ lambda: cancel_not_active("evt_from_client"),
+ lambda: cancel_not_active("litellm_guardrail_from_client"),
+ )
+ )
+
+ async def recv(**_kwargs: object) -> bytes:
+ frame = next(backend_frames, None)
+ if frame is None:
+ raise ConnectionClosed(None, None)
+ return frame()
+
+ backend_ws.recv = recv
+ logging_obj: Final = _mock_logging()
+ logging_obj.async_success_handler = AsyncMock()
+ logging_obj.success_handler = MagicMock()
+
+ await RealTimeStreaming(client_ws, backend_ws, logging_obj).backend_to_client_send_messages()
+
+ client_errors: Final = [
+ event["error"] for call in client_ws.send_text.await_args_list if (event := json.loads(call.args[0]))["type"] == "error"
+ ]
+ assert [(error["code"], error.get("event_id")) for error in client_errors] == [
+ ("content_policy_violation", None),
+ ("response_cancel_not_active", "evt_from_client"),
+ ("response_cancel_not_active", "litellm_guardrail_from_client"),
+ ]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ ("event_normalizer", "expected_frames"),
+ [
+ (None, ["response.cancel", "conversation.item.delete", "conversation.item.create", "response.create"]),
+ (XAIRealtimeNormalizer(), ["response.cancel", "conversation.item.create", "response.create"]),
+ ],
+ ids=["deletes_the_blocked_item", "xai_keeps_the_item_it_cannot_delete"],
+)
+async def test_blocked_transcript_is_removed_from_the_conversation_before_the_refusal_turn(
+ monkeypatch: pytest.MonkeyPatch,
+ event_normalizer: XAIRealtimeNormalizer | None,
+ expected_frames: list[str],
+):
+ class _BlockingGuardrail(CustomGuardrail):
+ async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
+ raise ValueError("blocked")
+
+ monkeypatch.setattr(
+ litellm,
+ "callbacks",
+ [
+ _BlockingGuardrail(
+ guardrail_name="blocking", event_hook=GuardrailEventHooks.realtime_input_transcription, default_on=True
+ )
+ ],
+ )
+ client_ws: Final = MagicMock()
+ client_ws.send_text = AsyncMock()
+ backend_ws: Final = MagicMock()
+ backend_ws.send = AsyncMock()
+ backend_frames: Final = iter((_make_transcript_event("blocked words", item_id="item_blocked"),))
+
+ async def recv(**_kwargs: object) -> bytes:
+ frame = next(backend_frames, None)
+ if frame is None:
+ raise ConnectionClosed(None, None)
+ return frame
+
+ backend_ws.recv = recv
+ logging_obj: Final = _mock_logging()
+ logging_obj.async_success_handler = AsyncMock()
+ logging_obj.success_handler = MagicMock()
+
+ await RealTimeStreaming(
+ client_ws, backend_ws, logging_obj, event_normalizer=event_normalizer
+ ).backend_to_client_send_messages()
+
+ sent: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list]
+ assert [event["type"] for event in sent] == expected_frames
+ assert all(event["item_id"] == "item_blocked" for event in sent if event["type"] == "conversation.item.delete")
diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py
index 4347a6639bd..fa649b48a51 100644
--- a/tests/unit/litellm_core_utils/test_streaming_handler.py
+++ b/tests/unit/litellm_core_utils/test_streaming_handler.py
@@ -1019,6 +1019,87 @@ async def test_bridged_stream_mid_stream_fallback_error_is_rebuilt_around_the_pr
assert excinfo.value.generated_content == ""
+@pytest.mark.asyncio
+@pytest.mark.parametrize("merge_reasoning", [True, False], ids=["merged_into_content", "kept_as_reasoning"])
+async def test_mid_stream_break_after_reasoning_disqualifies_continuation(logging_obj: Logging, merge_reasoning: bool):
+ import httpx
+
+ from litellm.exceptions import MidStreamFallbackError
+
+ class _ReasoningThenCut:
+ def __init__(self) -> None:
+ self.chunks = iter(
+ (
+ ModelResponseStream(choices=[{"index": 0, "delta": {"reasoning_content": "Let me think"}}]),
+ ModelResponseStream(choices=[{"index": 0, "delta": {"content": "The answer"}}]),
+ )
+ )
+
+ def __aiter__(self) -> "_ReasoningThenCut":
+ return self
+
+ async def __anext__(self) -> ModelResponseStream:
+ next_chunk: Final = next(self.chunks, None)
+ if next_chunk is None:
+ raise httpx.ReadTimeout("connection dropped")
+ return next_chunk
+
+ response: Final = CustomStreamWrapper(
+ completion_stream=_ReasoningThenCut(), model="deepseek-reasoner", logging_obj=logging_obj
+ )
+ response.merge_reasoning_content_in_choices = merge_reasoning
+
+ with pytest.raises(MidStreamFallbackError) as excinfo:
+ async for _ in response:
+ pass
+
+ assert excinfo.value.is_pre_first_chunk is False
+ assert excinfo.value.emitted_disqualifying_content is True
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "outer_delta,inner_flag,expected",
+ [
+ ({"tool_calls": [{"index": 0}]}, False, True),
+ ({"content": "Let me check."}, True, True),
+ ({"content": "Let me check."}, False, False),
+ ],
+ ids=["bridge_wrapper_sent_a_tool_call", "inner_stream_sent_a_tool_call", "plain_text_only"],
+)
+async def test_bridged_stream_rebuild_keeps_the_continuation_disqualifier(
+ logging_obj: Logging, outer_delta: dict[str, object], inner_flag: bool, expected: bool
+):
+ from litellm.exceptions import InternalServerError, MidStreamFallbackError
+
+ provider_error: Final = InternalServerError(message="upstream 500", llm_provider="azure", model="gpt-6.1-sol")
+ inner_error: Final = MidStreamFallbackError(
+ message=str(provider_error),
+ model="gpt-6.1-sol",
+ llm_provider="azure",
+ original_exception=provider_error,
+ is_pre_first_chunk=False,
+ emitted_disqualifying_content=inner_flag,
+ )
+
+ async def _raise_inner_error(**kwargs: object) -> NoReturn:
+ raise inner_error
+
+ response: Final = CustomStreamWrapper(
+ completion_stream=None,
+ model="gpt-6.1-sol",
+ logging_obj=logging_obj,
+ custom_llm_provider="azure",
+ make_call=_raise_inner_error,
+ )
+ response._accumulate_streamed_delta(outer_delta)
+
+ with pytest.raises(MidStreamFallbackError) as excinfo:
+ await response.__anext__()
+
+ assert excinfo.value.emitted_disqualifying_content is expected
+
+
@pytest.mark.asyncio
@pytest.mark.parametrize("sync_mode", [True, False])
async def test_streaming_rate_limit_midstream_fallback_keeps_provider_headers(logging_obj: Logging, sync_mode: bool):
@@ -6230,6 +6311,59 @@ class TestStableStreamingResponseId:
assert wrapper.model_response_creator().id == "chatcmpl-from-provider"
+class TestContinuationDisqualifiers:
+ """The mid-stream continuation eligibility hinges on classifying which deltas
+ carry output a text-only prefill cannot represent."""
+
+ @pytest.mark.parametrize(
+ "field",
+ [
+ "tool_calls",
+ "function_call",
+ "reasoning_content",
+ "thinking_blocks",
+ "reasoning_items",
+ "audio",
+ "images",
+ "annotations",
+ ],
+ )
+ def test_disqualifying_fields_flagged(self, field):
+ # reasoning_content included: it reaches the caller as visible reasoning a
+ # text-only prefill cannot carry, so a reasoning fallback would re-derive it
+ assert CustomStreamWrapper._delta_disqualifies_continuation({field: [{"x": 1}]}) is True
+
+ @pytest.mark.parametrize("delta", [{"content": "hi"}, {}, {"role": "assistant"}])
+ def test_plain_text_not_flagged(self, delta):
+ assert CustomStreamWrapper._delta_disqualifies_continuation(delta) is False
+
+ def test_non_mapping_delta_is_safe(self):
+ assert CustomStreamWrapper._delta_disqualifies_continuation(object()) is False
+
+ def test_accumulate_grows_text_and_leaves_flag_clear_for_plain_text(self):
+ wrapper = object.__new__(CustomStreamWrapper)
+ wrapper.response_uptil_now = ""
+ wrapper._emitted_disqualifying_content = False
+
+ wrapper._accumulate_streamed_delta({"content": "Hel"})
+ wrapper._accumulate_streamed_delta({"content": "lo"})
+
+ assert wrapper.response_uptil_now == "Hello"
+ assert wrapper._emitted_disqualifying_content is False
+
+ def test_accumulate_latches_flag_on_disqualifying_delta(self):
+ wrapper = object.__new__(CustomStreamWrapper)
+ wrapper.response_uptil_now = ""
+ wrapper._emitted_disqualifying_content = False
+
+ wrapper._accumulate_streamed_delta({"content": "Hi"})
+ wrapper._accumulate_streamed_delta({"tool_calls": [{"index": 0}]})
+ wrapper._accumulate_streamed_delta({"content": "there"})
+
+ assert wrapper.response_uptil_now == "Hithere"
+ assert wrapper._emitted_disqualifying_content is True
+
+
@pytest.mark.asyncio
async def test_async_stream_without_usage_counts_tokens_off_the_event_loop():
from tests.large_text import text
diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py
index 0c093700822..724e137c835 100644
--- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py
+++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py
@@ -6763,15 +6763,11 @@ def test_chat_flagged_model_converts_system_with_no_following_user_turn(local_mo
assert _texts(result["messages"][2]) == [CONVERTED_SYSTEM_NOTE, REMINDER_TEXT]
-@pytest.mark.parametrize(
- "empty_content",
- [[], None, [{"type": "input_audio", "input_audio": {"data": "AAAA", "format": "wav"}}]],
- ids=["empty-list", "none", "unsupported-part-only"],
-)
-def test_chat_flagged_model_converts_a_system_behind_a_user_turn_that_sends_nothing(
- local_model_cost_map, empty_content
-):
- messages = [
+@pytest.mark.parametrize("empty_content", [[], None], ids=["empty-list", "none"])
+def test_chat_flagged_model_anchors_a_system_on_an_empty_user_turn(
+ local_model_cost_map: None, empty_content: list[object] | None
+) -> None:
+ messages: Final = [
{"role": "system", "content": "You are terse."},
{"role": "user", "content": empty_content},
{"role": "system", "content": REMINDER_TEXT},
@@ -6779,7 +6775,22 @@ def test_chat_flagged_model_converts_a_system_behind_a_user_turn_that_sends_noth
{"role": "user", "content": "Second question"},
]
- result = _chat_request(AnthropicConfig(), FLAGGED_CLAUDE, messages)
+ result: Final = _chat_request(AnthropicConfig(), FLAGGED_CLAUDE, messages)
+
+ assert [m["role"] for m in result["messages"]] == ["user", "system", "assistant", "user"]
+ assert _texts(result["messages"][1]) == [REMINDER_TEXT]
+
+
+def test_chat_flagged_model_converts_a_system_behind_a_user_turn_that_sends_nothing(local_model_cost_map: None) -> None:
+ messages: Final = [
+ {"role": "system", "content": "You are terse."},
+ {"role": "user", "content": [{"type": "input_audio", "input_audio": {"data": "AAAA", "format": "wav"}}]},
+ {"role": "system", "content": REMINDER_TEXT},
+ {"role": "assistant", "content": "First answer"},
+ {"role": "user", "content": "Second question"},
+ ]
+
+ result: Final = _chat_request(AnthropicConfig(), FLAGGED_CLAUDE, messages)
assert [m["role"] for m in result["messages"]] == ["user", "assistant", "user"]
assert _texts(result["messages"][0]) == [CONVERTED_SYSTEM_NOTE, REMINDER_TEXT]
diff --git a/tests/unit/llms/anthropic/pass_through/adapters/test_handler_reasoning_effort_normalization.py b/tests/unit/llms/anthropic/pass_through/adapters/test_handler_reasoning_effort_normalization.py
index 0dfe4a93649..4b96d2e7991 100644
--- a/tests/unit/llms/anthropic/pass_through/adapters/test_handler_reasoning_effort_normalization.py
+++ b/tests/unit/llms/anthropic/pass_through/adapters/test_handler_reasoning_effort_normalization.py
@@ -6,13 +6,17 @@ regression they guard is the one a caller sees: a tier the proxy advertises has
leaves the adapter, in the shape the target expects.
"""
+from collections.abc import Mapping
from typing import Final
import pytest
+import litellm
from litellm.llms.anthropic.pass_through.adapters.handler import (
LiteLLMMessagesToCompletionTransformationHandler,
)
+from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig
+from litellm.utils import reasoning_effort_with_summary_alias
MESSAGES = [{"role": "user", "content": "hello"}]
@@ -38,14 +42,14 @@ def _reasoning_effort_sent(model: str, provider: str, reasoning_effort: object)
return completion_kwargs.get("reasoning_effort")
-def _reasoning_effort_sent_for_thinking(
+def _completion_kwargs_for_thinking(
model: str,
provider: str | None,
thinking: dict[str, object],
*,
tools: list[dict[str, object]] | None = None,
api_base: str | None = None,
-) -> object:
+) -> Mapping[str, object]:
extra_kwargs: Final = {
key: value for key, value in (("custom_llm_provider", provider), ("api_base", api_base)) if value is not None
}
@@ -66,7 +70,20 @@ def _reasoning_effort_sent_for_thinking(
output_format=None,
extra_kwargs=extra_kwargs,
)
- return completion_kwargs.get("reasoning_effort")
+ return completion_kwargs
+
+
+def _reasoning_effort_sent_for_thinking(
+ model: str,
+ provider: str | None,
+ thinking: dict[str, object],
+ *,
+ tools: list[dict[str, object]] | None = None,
+ api_base: str | None = None,
+) -> object:
+ return _completion_kwargs_for_thinking(model, provider, thinking, tools=tools, api_base=api_base).get(
+ "reasoning_effort"
+ )
SUMMARIZED_THINKING = {"type": "enabled", "budget_tokens": 4096, "summary": "auto"}
@@ -78,6 +95,66 @@ MULTIPLY_TOOL = {
}
+class TestAThinkingSummaryAskReachesGemini:
+ """Gemini maps a summary ask to ``includeThoughts``, so ``gemini_include_thoughts_default`` only
+ decides for a ``/v1/messages`` request that asks for nothing, the same as on chat completions."""
+
+ @pytest.mark.parametrize(
+ "model, provider",
+ [("gemini/gemini-3.1-pro-preview", "gemini"), ("vertex_ai/gemini-2.5-flash", "vertex_ai")],
+ )
+ @pytest.mark.parametrize(
+ "thinking, sends_thoughts",
+ [
+ (SUMMARIZED_THINKING, True),
+ ({**PLAIN_THINKING, "display": "summarized"}, True),
+ ({**PLAIN_THINKING, "display": "updates"}, True),
+ ({**PLAIN_THINKING, "display": "omitted"}, False),
+ (PLAIN_THINKING, False),
+ ({"type": "disabled", "display": "summarized"}, False),
+ ],
+ )
+ def test_the_request_ask_beats_the_gateway_default(
+ self,
+ local_model_cost_map: None,
+ monkeypatch: pytest.MonkeyPatch,
+ model: str,
+ provider: str,
+ thinking: dict[str, object],
+ sends_thoughts: bool,
+ ) -> None:
+ monkeypatch.setattr(litellm, "gemini_include_thoughts_default", False)
+
+ sent: Final = _completion_kwargs_for_thinking(model, provider, thinking)
+ gemini: Final = VertexGeminiConfig()
+ reasoning_effort: Final = reasoning_effort_with_summary_alias(
+ sent.get("reasoning_effort"), sent, accepts_summary=gemini.accepts_reasoning_summary_in_reasoning_effort
+ )
+ mapped: Final = gemini.map_openai_params(
+ {"reasoning_effort": reasoning_effort}, {}, model.split("/", 1)[1], True
+ )
+
+ assert mapped["thinkingConfig"]["includeThoughts"] is sends_thoughts
+ assert not isinstance(sent.get("reasoning_effort"), dict)
+
+ def test_auto_summary_leaves_the_gateway_default_in_charge(
+ self, local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch
+ ) -> None:
+ monkeypatch.setattr(litellm, "gemini_include_thoughts_default", False)
+ monkeypatch.setenv("LITELLM_REASONING_AUTO_SUMMARY", "true")
+
+ sent: Final = _reasoning_effort_sent_for_thinking("gemini/gemini-3.1-pro-preview", "gemini", PLAIN_THINKING)
+
+ assert sent == "high"
+
+ def test_a_display_ask_leaves_other_chat_targets_on_the_plain_tier(self, local_model_cost_map: None) -> None:
+ sent: Final = _reasoning_effort_sent_for_thinking(
+ "databricks/databricks-qwen35-122b-a10b", "databricks", {**PLAIN_THINKING, "display": "summarized"}
+ )
+
+ assert sent == "high"
+
+
class TestTheSummaryWrappingOnlyRidesTheResponsesBridge:
"""Only the Responses API takes ``reasoning_effort`` as a dict. Databricks answered the wrapped
``{"effort", "summary"}`` with ``field 'reasoning_effort' expects input with json type 'string'
@@ -100,7 +177,9 @@ class TestTheSummaryWrappingOnlyRidesTheResponsesBridge:
) -> None:
monkeypatch.setenv("LITELLM_REASONING_AUTO_SUMMARY", "true")
- sent = _reasoning_effort_sent_for_thinking("databricks/databricks-qwen35-122b-a10b", "databricks", PLAIN_THINKING)
+ sent: Final = _reasoning_effort_sent_for_thinking(
+ "databricks/databricks-qwen35-122b-a10b", "databricks", PLAIN_THINKING
+ )
assert sent == "high"
@@ -113,7 +192,7 @@ class TestTheSummaryWrappingOnlyRidesTheResponsesBridge:
],
)
def test_a_bridged_target_keeps_the_summary(self, local_model_cost_map: None, model: str, provider: str) -> None:
- sent = _reasoning_effort_sent_for_thinking(model, provider, SUMMARIZED_THINKING)
+ sent: Final = _reasoning_effort_sent_for_thinking(model, provider, SUMMARIZED_THINKING)
assert sent == {"effort": "high", "summary": "auto"}
@@ -122,7 +201,7 @@ class TestTheSummaryWrappingOnlyRidesTheResponsesBridge:
) -> None:
monkeypatch.setenv("LITELLM_REASONING_AUTO_SUMMARY", "true")
- sent = _reasoning_effort_sent_for_thinking("azure/responses/gpt-5-mini", "azure", PLAIN_THINKING)
+ sent: Final = _reasoning_effort_sent_for_thinking("azure/responses/gpt-5-mini", "azure", PLAIN_THINKING)
assert sent == {"effort": "high", "summary": "detailed"}
@@ -139,21 +218,21 @@ class TestTheSummaryWrappingOnlyRidesTheResponsesBridge:
"""``completion()`` keeps a gpt-5.5 deployment with function tools on Foundry's chat route when
its ``api_base`` is a Foundry OpenAI host, and bridges it to Responses when the base makes it an
Azure OpenAI deployment. The adapter has to read the same ``api_base`` to land on the same call."""
- sent = _reasoning_effort_sent_for_thinking(
+ sent: Final = _reasoning_effort_sent_for_thinking(
"azure_ai/gpt-5.5", "azure_ai", SUMMARIZED_THINKING, tools=[MULTIPLY_TOOL], api_base=api_base
)
assert sent == expected
def test_a_provider_resolved_from_the_api_base_gets_the_plain_tier(self, local_model_cost_map: None) -> None:
- sent = _reasoning_effort_sent_for_thinking(
+ sent: Final = _reasoning_effort_sent_for_thinking(
"kimi-k3", None, SUMMARIZED_THINKING, api_base="https://api.together.xyz/v1"
)
assert sent == "high"
def test_a_chained_gateway_keeps_the_dict_for_its_own_bridge(self, local_model_cost_map: None) -> None:
- sent = _reasoning_effort_sent_for_thinking("litellm_proxy/gpt-5.4", "litellm_proxy", SUMMARIZED_THINKING)
+ sent: Final = _reasoning_effort_sent_for_thinking("litellm_proxy/gpt-5.4", "litellm_proxy", SUMMARIZED_THINKING)
assert sent == {"effort": "high", "summary": "auto"}
@@ -192,7 +271,7 @@ class TestTheNormalizedTierIsTheTierSent:
def test_the_dict_form_normalizes_effort_and_keeps_its_siblings(
self, local_model_cost_map, model, provider, expected
):
- sent = _reasoning_effort_sent(model, provider, {"effort": "max", "summary": "detailed"})
+ sent: Final = _reasoning_effort_sent(model, provider, {"effort": "max", "summary": "detailed"})
assert sent == {"effort": expected, "summary": "detailed"}
@@ -215,8 +294,8 @@ class TestTheNormalizedTierIsTheTierSent:
],
)
def test_an_azure_hosted_astra_deployment_drops_to_the_tier_it_accepts(
- self, local_model_cost_map, model, provider
- ):
+ self, local_model_cost_map: None, model: str, provider: str
+ ) -> None:
"""The deployment answers ``max`` with a 400 naming ``none`` through ``xhigh``, so the rows
say so and the adapter sends the tier below instead of the rejected one."""
assert _reasoning_effort_sent(model, provider, "max") == "xhigh"
diff --git a/tests/unit/llms/anthropic/test_count_tokens_oauth.py b/tests/unit/llms/anthropic/test_count_tokens_oauth.py
index 7f717091754..16ec513916c 100644
--- a/tests/unit/llms/anthropic/test_count_tokens_oauth.py
+++ b/tests/unit/llms/anthropic/test_count_tokens_oauth.py
@@ -197,6 +197,30 @@ class TestCountTokensUsesWorkloadIdentity:
assert ANTHROPIC_TOKEN_COUNTING_BETA_VERSION in betas
assert ANTHROPIC_OAUTH_BETA_HEADER not in betas
+ @pytest.mark.asyncio
+ async def test_a_global_anthropic_key_counts_like_chat_would_authenticate(
+ self, monkeypatch, httpx_transport_clients
+ ):
+ for name in ("ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"):
+ monkeypatch.delenv(name, raising=False)
+ monkeypatch.setattr(litellm, "anthropic_key", FAKE_REGULAR_KEY)
+
+ with respx.mock(assert_all_called=True) as router:
+ route = router.post("https://api.anthropic.com/v1/messages/count_tokens").mock(
+ return_value=httpx.Response(200, json={"input_tokens": 13})
+ )
+ result = await AnthropicTokenCounter().count_tokens(
+ model_to_use="claude-sonnet-4-5",
+ messages=[{"role": "user", "content": "hi"}],
+ contents=None,
+ deployment={"litellm_params": {"model": "anthropic/claude-sonnet-4-5"}},
+ request_model="claude-sonnet-4-5",
+ )
+
+ assert result is not None
+ assert result.total_tokens == 13
+ assert route.calls.last.request.headers["x-api-key"] == FAKE_REGULAR_KEY
+
@pytest.mark.asyncio
async def test_a_failed_mint_degrades_like_an_anthropic_error(self, monkeypatch):
from litellm.llms.anthropic.count_tokens import token_counter as token_counter_module
diff --git a/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py
index 29380c6fc53..86ad6c8d9d9 100644
--- a/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py
+++ b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py
@@ -186,3 +186,27 @@ async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_
"status_code": 500,
}
)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("set_base", ("env", "global"), ids=("AZURE_API_BASE", "litellm.api_base"))
+async def test_count_uses_the_shared_api_base_chat_would_use(monkeypatch: pytest.MonkeyPatch, set_base: str) -> None:
+ monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+ monkeypatch.delenv("AZURE_AI_API_BASE", raising=False)
+ monkeypatch.delenv("AZURE_API_BASE", raising=False)
+ if set_base == "env":
+ monkeypatch.setenv("AZURE_API_BASE", API_BASE)
+ else:
+ monkeypatch.setattr(litellm, "api_base", API_BASE)
+ with respx.mock:
+ route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 6}))
+ result: Final = await AzureAIAnthropicTokenCounter().count_tokens(
+ model_to_use="claude-sonnet-4-5",
+ messages=[{"role": "user", "content": "hi"}],
+ contents=None,
+ deployment={"litellm_params": {"api_key": "azure-key"}},
+ request_model="azure_ai/claude-sonnet-4-5",
+ )
+ assert result is not None
+ assert result.total_tokens == 6
+ assert route.called
diff --git a/tests/unit/llms/bedrock/chat/test_invoke_handler.py b/tests/unit/llms/bedrock/chat/test_invoke_handler.py
index 0609f122949..069438e25c0 100644
--- a/tests/unit/llms/bedrock/chat/test_invoke_handler.py
+++ b/tests/unit/llms/bedrock/chat/test_invoke_handler.py
@@ -16,6 +16,7 @@ import litellm
from litellm import AmazonInvokeConfig
from litellm.exceptions import MidStreamFallbackError
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+from litellm.litellm_core_utils.litellm_logging import _get_provider_request_id
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.llms.bedrock.chat.invoke_handler import (
AmazonOpenAICompatibleStreamDecoder,
@@ -864,14 +865,16 @@ async def test_async_invoke_streaming_non_200_forwards_bedrock_response_headers(
def _event_stream_frame(event_type: str, payload: bytes) -> bytes:
+ return _typed_event_stream_frame(
+ {":event-type": event_type, ":content-type": "application/json", ":message-type": "event"}, payload
+ )
+
+
+def _typed_event_stream_frame(frame_headers: Mapping[str, str], payload: bytes) -> bytes:
def header(name: str, value: str) -> bytes:
return bytes([len(name)]) + name.encode() + bytes([7]) + struct.pack(">H", len(value)) + value.encode()
- headers: Final = (
- header(":event-type", event_type)
- + header(":content-type", "application/json")
- + header(":message-type", "event")
- )
+ headers: Final = b"".join(header(name, value) for name, value in frame_headers.items())
prelude: Final = struct.pack(">II", 12 + len(headers) + len(payload) + 4, len(headers))
body: Final = prelude + struct.pack(">I", binascii.crc32(prelude)) + headers + payload
return body + struct.pack(">I", binascii.crc32(body))
@@ -1367,6 +1370,87 @@ def test_iter_bytes_keeps_a_stream_whose_unknown_event_sits_beside_known_frames(
assert "stop" in finish_reasons
+_STREAM_FAILURE_BODIES: Final = (
+ pytest.param(
+ _converse_event_frame("internalServerException", {"message": _UPSTREAM_REJECTION}), 500, id="exception-event"
+ ),
+ pytest.param(
+ _typed_event_stream_frame(
+ {
+ ":exception-type": "throttlingException",
+ ":content-type": "application/json",
+ ":message-type": "exception",
+ },
+ json.dumps({"message": _UPSTREAM_REJECTION}).encode(),
+ ),
+ 429,
+ id="exception-message",
+ ),
+ pytest.param(_unknown_event_frame(), 502, id="no-known-event"),
+ pytest.param(b"", 502, id="no-events"),
+ pytest.param(_CONVERSE_CONTENT_FRAMES[0] + _truncated_frame(), 502, id="undecoded-tail"),
+)
+
+
+def _stream_headers(request_id: str | None) -> httpx.Headers:
+ request_id_header: Final = {} if request_id is None else {"x-amzn-RequestId": request_id}
+ return httpx.Headers(
+ {"content-type": "application/vnd.amazon.eventstream", "retry-after": "7", **request_id_header}
+ )
+
+
+def _assert_carries_only_the_request_id(error: BedrockError, status: int, request_id: str | None) -> None:
+ assert error.status_code == status
+ assert _get_provider_request_id(error) == request_id
+ if request_id is None:
+ assert error.headers is None
+ return
+ assert dict(error.headers or {}) == {"x-amzn-requestid": request_id}
+ assert dict(error.response.headers) == {"x-amzn-requestid": request_id}
+
+
+@pytest.mark.parametrize("request_id", ("req-stream-1", None))
+@pytest.mark.parametrize(("body", "status"), _STREAM_FAILURE_BODIES)
+def test_iter_bytes_stream_error_carries_only_the_provider_request_id_header(
+ body: bytes, status: int, request_id: str | None
+) -> None:
+ decoder: Final = AWSEventStreamDecoder(model="us.moonshotai.kimi-k3")
+
+ with pytest.raises(BedrockError) as exc_info:
+ list(decoder.iter_bytes(iter([body]), response_headers=_stream_headers(request_id)))
+
+ _assert_carries_only_the_request_id(exc_info.value, status, request_id)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(("body", "status"), _STREAM_FAILURE_BODIES)
+async def test_aiter_bytes_stream_error_carries_only_the_provider_request_id_header(body: bytes, status: int) -> None:
+ async def _chunks() -> AsyncIterator[bytes]:
+ yield body
+
+ decoder: Final = AWSEventStreamDecoder(model="us.moonshotai.kimi-k3")
+
+ with pytest.raises(BedrockError) as exc_info:
+ _ = [chunk async for chunk in decoder.aiter_bytes(_chunks(), response_headers=_stream_headers("req-stream-2"))]
+
+ _assert_carries_only_the_request_id(exc_info.value, status, "req-stream-2")
+
+
+@pytest.mark.parametrize("request_id", ("req-\u00fc-1", "req-\udcfc-1"), ids=("latin-1", "surrogate-escaped"))
+def test_iter_bytes_streams_and_fails_normally_with_a_non_ascii_provider_request_id(request_id: str) -> None:
+ decoder: Final = AWSEventStreamDecoder(model="us.moonshotai.kimi-k3")
+ headers: Final = {"content-type": "application/vnd.amazon.eventstream", "x-amzn-RequestId": request_id}
+
+ chunks: Final = list(decoder.iter_bytes(iter(_CONVERSE_CONTENT_FRAMES), response_headers=headers))
+ with pytest.raises(BedrockError) as exc_info:
+ list(decoder.iter_bytes(iter([_unknown_event_frame()]), response_headers=headers))
+
+ texts: Final = [chunk.choices[0].delta.content for chunk in chunks if isinstance(chunk, ModelResponseStream)]
+ assert "".join(text or "" for text in texts) == "hi"
+ assert exc_info.value.status_code == 502
+ assert exc_info.value.headers is None
+
+
def _assert_exception_event_surfaced_with_its_modeled_status(error: BaseException, event_type: str) -> None:
assert not isinstance(error, litellm.BadGatewayError)
assert getattr(error, "status_code", None) == _modeled_status(event_type)
diff --git a/tests/unit/llms/bedrock/test_web_identity_session_policy.py b/tests/unit/llms/bedrock/test_web_identity_session_policy.py
index cbb69b4ceed..332e1258e9c 100644
--- a/tests/unit/llms/bedrock/test_web_identity_session_policy.py
+++ b/tests/unit/llms/bedrock/test_web_identity_session_policy.py
@@ -176,6 +176,44 @@ class TestClaudePlatformActionsCovered:
)
+class TestClaudePlatformOutboundIdentityCovered:
+ """Claude Platform on AWS calls ``sts:GetWebIdentityToken`` with the
+ caller's session before serving inference, and tags the token it
+ requests (``sts:TagGetWebIdentityToken``). Without either in the ceiling,
+ every claude_platform request via OIDC auth denies with "no session
+ policy allows the sts: action" even when the role's identity
+ policy grants it."""
+
+ @pytest.mark.parametrize("action", ["sts:GetWebIdentityToken", "sts:TagGetWebIdentityToken"])
+ def test_outbound_identity_action_present(self, action: str):
+ assert action in _granted_actions(_captured_policy()), (
+ f"{action} missing from session policy, so bedrock/claude_platform/* requests will 403 on OIDC auth"
+ )
+
+ @pytest.mark.parametrize("action", ["sts:GetWebIdentityToken", "sts:TagGetWebIdentityToken"])
+ def test_outbound_identity_action_is_granted_only_via_claude_platform(self, action: str):
+ granting: Final = tuple(stmt for stmt in _captured_policy()["Statement"] if action in stmt["Action"])
+ assert granting, f"{action} missing from session policy"
+ assert all(
+ stmt["Condition"].get("StringEquals") == {"aws:CalledViaLast": "aws-external-anthropic.amazonaws.com"}
+ and stmt["Condition"]["Bool"] == {"aws:SecureTransport": "true"}
+ for stmt in granting
+ ), f"{action} must only be usable over TLS when Claude Platform on AWS calls STS for the session"
+
+ def test_outbound_identity_statement_is_scoped(self):
+ policy = _captured_policy()
+ stmt = _statement_by_sid(policy, "ClaudePlatformOutboundIdentityLiteLLM")
+ assert stmt["Effect"] == "Allow"
+ assert tuple(stmt["Action"]) == ("sts:GetWebIdentityToken", "sts:TagGetWebIdentityToken"), (
+ "the outbound identity statement must grant only the two outbound identity actions"
+ )
+
+ def test_no_sts_wildcard(self):
+ assert "sts:*" not in _granted_actions(_captured_policy()), (
+ "session policy must not grant sts:*, only the two outbound identity actions are needed"
+ )
+
+
class TestBedrockMantleActionsCovered:
"""LIT-3859: bedrock_mantle inference authorizes against the
``bedrock-mantle`` action namespace, so the session-policy ceiling
@@ -306,6 +344,15 @@ class TestPolicyTransportConditions:
"to keep parity with the bedrock statement"
)
+ def test_only_the_outbound_identity_statement_requires_a_claude_platform_caller(self):
+ directly_called: Final = tuple(
+ stmt for stmt in _captured_policy()["Statement"] if stmt["Sid"] != "ClaudePlatformOutboundIdentityLiteLLM"
+ )
+ assert directly_called
+ assert all("StringEquals" not in stmt["Condition"] for stmt in directly_called), (
+ "statements LiteLLM calls directly must not require aws:CalledViaLast, or OIDC auth 403s on those routes"
+ )
+
_STS_SESSION_POLICY_PLAINTEXT_LIMIT: Final = 2048
diff --git a/tests/unit/llms/dashscope/test_dashscope_video_transformation.py b/tests/unit/llms/dashscope/test_dashscope_video_transformation.py
index 65eefcab1af..15b77aaee3d 100644
--- a/tests/unit/llms/dashscope/test_dashscope_video_transformation.py
+++ b/tests/unit/llms/dashscope/test_dashscope_video_transformation.py
@@ -194,8 +194,9 @@ class TestDashScopeVideoInputReference:
lambda tmp_path: ("first.png", PNG_BYTES, "image/png"),
lambda tmp_path: ("first.png", io.BytesIO(PNG_BYTES)),
lambda tmp_path: _write_png(tmp_path),
+ lambda tmp_path: ("first.png", ("inner.png", PNG_BYTES)),
],
- ids=["filename-bytes-content-type", "filename-file", "path"],
+ ids=["filename-bytes-content-type", "filename-file", "path", "nested-tuple"],
)
def test_file_tuple_and_path_references_are_encoded(self, make_reference, tmp_path):
"""Regression: tuples and paths are valid FileTypes but raised before reaching DashScope."""
diff --git a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py
index 650517c92ca..1e353a7a82a 100644
--- a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py
+++ b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py
@@ -461,6 +461,21 @@ def test_transform_messages_sanitizes_empty_content():
assert result[1]["content"] == "Hi"
+def test_transform_messages_drops_the_prefill_prefix_flag():
+ """Databricks rejects unknown message keys, so a continuation prefill reaches Claude without ``prefix``."""
+ messages = [
+ {"role": "user", "content": "How do transformers work?"},
+ {"role": "assistant", "content": "Transformers use", "prefix": True},
+ ]
+
+ result = DatabricksConfig().transform_messages(
+ messages=messages, model="databricks-claude-sonnet-4-6", is_async=False
+ )
+
+ assert result[-1] == {"role": "assistant", "content": "Transformers use"}
+ assert messages[-1]["prefix"] is True
+
+
def test_transform_request_preserves_unity_model_service_name():
config = DatabricksConfig()
result = config.transform_request(
diff --git a/tests/unit/llms/gemini/realtime/test_gemini_realtime_transformation.py b/tests/unit/llms/gemini/realtime/test_gemini_realtime_transformation.py
index bcd5f3d8d19..017f279d6c4 100644
--- a/tests/unit/llms/gemini/realtime/test_gemini_realtime_transformation.py
+++ b/tests/unit/llms/gemini/realtime/test_gemini_realtime_transformation.py
@@ -1,6 +1,6 @@
import json
from collections.abc import Mapping
-from typing import cast
+from typing import Final, cast
from unittest.mock import MagicMock
import pytest
@@ -1323,6 +1323,40 @@ def test_gemini_realtime_pipecat_semantic_vad_omits_realtime_input_config():
assert setup["tools"][0]["function_declarations"][0]["name"] == "terminate_call"
+@pytest.mark.parametrize(
+ ("create_response", "expected_detection"),
+ [
+ (False, {"disabled": True}),
+ (True, {"disabled": False, "prefixPaddingMs": 200, "silenceDurationMs": 700}),
+ ],
+)
+def test_gemini_turn_detection_tuning_is_dropped_when_auto_response_is_off(
+ create_response: bool, expected_detection: dict[str, object]
+):
+ session_update: Final = {
+ "type": "session.update",
+ "session": {
+ "audio": {
+ "input": {
+ "turn_detection": {
+ "type": "server_vad",
+ "create_response": create_response,
+ "prefix_padding_ms": 200,
+ "silence_duration_ms": 700,
+ }
+ },
+ },
+ },
+ }
+ messages: Final = GeminiRealtimeConfig().transform_realtime_request(
+ json.dumps(session_update),
+ "gemini-live-2.5-flash-native-audio",
+ session_configuration_request=None,
+ )
+ setup: Final = json.loads(messages[0])["setup"]
+ assert setup["realtimeInputConfig"]["automaticActivityDetection"] == expected_detection
+
+
def test_gemini_input_audio_buffer_commit_maps_to_audio_stream_end():
config = GeminiRealtimeConfig()
setup = {
diff --git a/tests/unit/llms/meta/realtime/test_meta_realtime_transformation.py b/tests/unit/llms/meta/realtime/test_meta_realtime_transformation.py
index a5d7e47fb65..8b46ea53713 100644
--- a/tests/unit/llms/meta/realtime/test_meta_realtime_transformation.py
+++ b/tests/unit/llms/meta/realtime/test_meta_realtime_transformation.py
@@ -20,6 +20,8 @@ from litellm.llms.meta.realtime.transformation import (
session_created_event,
)
from litellm.types.llms.meta import MuseMode
+from litellm.types.utils import LlmProviders
+from litellm.utils import ProviderConfigManager
from litellm.types.realtime import RealtimeResponseTransformInput
EMPTY_TRANSFORM_INPUT: Final[RealtimeResponseTransformInput] = {
@@ -164,6 +166,22 @@ def test_session_rejects_unsupported_audio_model_and_hints(session: dict[str, ob
parse_session_update(_event("session.update", session={"type": "transcription", **session}), MUSE_MODEL)
+def _handshake(turn_detection: str, drop_params: bool) -> tuple[str | bytes, ...]:
+ config = ProviderConfigManager.get_provider_realtime_config(MUSE_MODEL, LlmProviders.META, drop_params=drop_params)
+ assert config is not None
+ config.validate_environment({}, MUSE_MODEL, api_key="secret-token")
+ return tuple(config.transform_realtime_request(_ga_session_update(turn_detection=turn_detection), MUSE_MODEL))
+
+
+def test_semantic_vad_falls_back_to_server_vad_under_drop_params():
+ assert _handshake("semantic_vad", drop_params=True) == _handshake("server_vad", drop_params=False)
+
+
+def test_semantic_vad_without_drop_params_names_the_drop_params_fallback():
+ with pytest.raises(MuseProtocolError, match="Set drop_params to use server_vad instead"):
+ _handshake("semantic_vad", drop_params=False)
+
+
def test_session_created_event_exposes_openai_transcription_shape():
config = parse_session_update(
_event(
diff --git a/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py b/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py
index 976ecff01ea..befdc5ee9fd 100644
--- a/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py
+++ b/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py
@@ -388,3 +388,44 @@ async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_
"status_code": 500,
}
)
+
+
+@pytest.mark.asyncio
+async def test_count_uses_the_global_openai_key_chat_would_use(monkeypatch: pytest.MonkeyPatch) -> None:
+ monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+ monkeypatch.setattr(litellm, "openai_key", "sk-from-litellm-settings")
+ monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+ with respx.mock:
+ route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 9}))
+ result: Final = await OpenAITokenCounter().count_tokens(
+ model_to_use="gpt-4o",
+ messages=[{"role": "user", "content": "hi"}],
+ contents=None,
+ deployment={"litellm_params": {}},
+ request_model="openai/gpt-4o",
+ )
+ assert result is not None
+ assert result.total_tokens == 9
+ assert route.calls.last.request.headers["authorization"] == "Bearer sk-from-litellm-settings"
+
+
+@pytest.mark.asyncio
+async def test_count_sends_the_global_key_only_to_the_global_api_base(monkeypatch: pytest.MonkeyPatch) -> None:
+ monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+ monkeypatch.setattr(litellm, "openai_key", "sk-private-gateway")
+ monkeypatch.setattr(litellm, "api_base", "https://gateway.internal.example/v1")
+ monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+ with respx.mock:
+ route: Final = respx.post("https://gateway.internal.example/v1/responses/input_tokens").mock(
+ return_value=httpx.Response(200, json={"input_tokens": 4})
+ )
+ result: Final = await OpenAITokenCounter().count_tokens(
+ model_to_use="gpt-4o",
+ messages=[{"role": "user", "content": "hi"}],
+ contents=None,
+ deployment={"litellm_params": {}},
+ request_model="openai/gpt-4o",
+ )
+ assert result is not None
+ assert result.total_tokens == 4
+ assert route.calls.last.request.headers["authorization"] == "Bearer sk-private-gateway"
diff --git a/tests/unit/llms/openai/test_openai_workload_identity.py b/tests/unit/llms/openai/test_openai_workload_identity.py
index d0cf52d0e28..2236de1eb5f 100644
--- a/tests/unit/llms/openai/test_openai_workload_identity.py
+++ b/tests/unit/llms/openai/test_openai_workload_identity.py
@@ -648,3 +648,56 @@ class TestClientsideBaseOverride:
)
is None
)
+
+
+class TestCountTokens:
+ @pytest.mark.asyncio
+ @respx.mock
+ async def test_count_authenticates_with_the_workload_identity_bearer(
+ self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch
+ ) -> None:
+ from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter
+
+ monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+ monkeypatch.setattr(litellm, "api_key", None)
+ monkeypatch.setattr(litellm, "openai_key", None)
+ mock_token_exchange()
+ count_route: Final = respx.post("https://api.openai.com/v1/responses/input_tokens").mock(
+ return_value=httpx.Response(200, json={"input_tokens": 12})
+ )
+ result: Final = await OpenAITokenCounter().count_tokens(
+ model_to_use="gpt-4o",
+ messages=[{"role": "user", "content": "hi"}],
+ contents=None,
+ deployment={"litellm_params": {}},
+ request_model="openai/gpt-4o",
+ )
+ assert result is not None
+ assert result.error is False
+ assert result.total_tokens == 12
+ assert count_route.calls.last.request.headers["authorization"] == "Bearer exchanged-bearer-token"
+
+ @pytest.mark.asyncio
+ @respx.mock
+ async def test_failed_mint_becomes_a_count_error_response(
+ self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch
+ ) -> None:
+ from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter
+
+ monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+ monkeypatch.setattr(litellm, "api_key", None)
+ monkeypatch.setattr(litellm, "openai_key", None)
+ respx.post(TOKEN_EXCHANGE_URL).mock(return_value=httpx.Response(400, json={"error": "invalid_grant"}))
+ count_route: Final = respx.post("https://api.openai.com/v1/responses/input_tokens")
+ result: Final = await OpenAITokenCounter().count_tokens(
+ model_to_use="gpt-4o",
+ messages=[{"role": "user", "content": "hi"}],
+ contents=None,
+ deployment={"litellm_params": {}},
+ request_model="openai/gpt-4o",
+ )
+ assert result is not None
+ assert result.error is True
+ assert result.status_code == 401
+ assert result.total_tokens == 0
+ assert not count_route.called
diff --git a/tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_handler.py b/tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_handler.py
index d2d7d2247f1..9c70ebe7ece 100644
--- a/tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_handler.py
+++ b/tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_handler.py
@@ -339,6 +339,65 @@ class TestPollLimitsClamping:
a worker on tight poll loops.
"""
+ def test_should_raise_400_when_form_bool_param_is_invalid(self):
+ handler = SonioxAudioTranscriptionHandler()
+ with pytest.raises(SonioxException) as exc_info:
+ handler._prepare(
+ audio_file=None,
+ optional_params={
+ "enable_speaker_diarization": "maybe",
+ "audio_url": "https://example.com/a.wav",
+ },
+ litellm_params={},
+ api_key="sk-test",
+ api_base=None,
+ provider_config=SonioxAudioTranscriptionConfig(),
+ headers={},
+ )
+ assert exc_info.value.status_code == 400
+ assert "enable_speaker_diarization" in str(exc_info.value)
+
+ def test_should_raise_400_when_poll_interval_is_not_a_number(self):
+ with pytest.raises(SonioxException) as exc_info:
+ SonioxAudioTranscriptionHandler()._prepare(
+ audio_file=None,
+ optional_params={"soniox_polling_interval": "abc", "audio_url": "https://example.com/a.wav"},
+ litellm_params={},
+ api_key="sk-test",
+ api_base=None,
+ provider_config=SonioxAudioTranscriptionConfig(),
+ headers={},
+ )
+ assert exc_info.value.status_code == 400
+ assert "soniox_polling_interval" in str(exc_info.value)
+
+ def test_should_use_default_poll_interval_for_bad_value_when_drop_params_is_set(self):
+ from litellm.llms.soniox.common_utils import SONIOX_DEFAULT_POLL_INTERVAL
+
+ _, _, _, opts = SonioxAudioTranscriptionHandler()._prepare(
+ audio_file=None,
+ optional_params={"soniox_polling_interval": "abc", "audio_url": "https://example.com/a.wav"},
+ litellm_params={"drop_params": True},
+ api_key="sk-test",
+ api_base=None,
+ provider_config=SonioxAudioTranscriptionConfig(),
+ headers={},
+ )
+ assert opts["poll_interval"] == SONIOX_DEFAULT_POLL_INTERVAL
+
+ @pytest.mark.parametrize("raw, expected", [(10.0, 10), (10, 10), ("10", 10)])
+ def test_should_keep_numeric_max_polling_attempts(self, raw: object, expected: int):
+ _, _, _, opts = SonioxAudioTranscriptionHandler()._prepare(
+ audio_file=None,
+ optional_params={"soniox_max_polling_attempts": raw, "audio_url": "https://example.com/a.wav"},
+ litellm_params={},
+ api_key="sk-test",
+ api_base=None,
+ provider_config=SonioxAudioTranscriptionConfig(),
+ headers={},
+ )
+ assert opts["max_attempts"] == expected
+
def test_should_clamp_poll_interval_to_minimum(self):
from litellm.llms.soniox.common_utils import SONIOX_MIN_POLL_INTERVAL
@@ -707,6 +766,112 @@ class TestCleanupNormalization:
assert "https://api.soniox.com/v1/transcriptions/tx_1" in deletes
+ def test_should_skip_non_string_cleanup_entries(self):
+ _, _, _, opts = SonioxAudioTranscriptionHandler()._prepare(
+ audio_file=None,
+ optional_params={"soniox_cleanup": ["transcription", 1, None], "audio_url": "https://example.com/a.wav"},
+ litellm_params={},
+ api_key="sk-test",
+ api_base=None,
+ provider_config=SonioxAudioTranscriptionConfig(),
+ headers={},
+ )
+ assert opts["cleanup"] == ("transcription",)
+
+ @pytest.mark.parametrize(
+ "language_hints, expected",
+ [("en,es", ["en", "es"]), (["es"], ["es"]), (None, ["de"])],
+ )
+ def test_should_use_language_only_without_language_hints(self, language_hints, expected, monkeypatch):
+ import litellm
+
+ monkeypatch.setattr("time.sleep", lambda *_: None)
+ client = _MockSyncClient(
+ {
+ "POST https://api.soniox.com/v1/transcriptions": [_make_response({"id": "tx_1"})],
+ "GET https://api.soniox.com/v1/transcriptions/tx_1": [_make_response({"status": "completed"})],
+ "GET https://api.soniox.com/v1/transcriptions/tx_1/transcript": [
+ _make_response({"text": "hallo", "tokens": []})
+ ],
+ "DELETE https://api.soniox.com/v1/transcriptions/tx_1": [_make_response({})],
+ }
+ )
+ litellm.transcription(
+ model="soniox/stt-async-v4",
+ file=("a.wav", b"RIFF", "audio/wav"),
+ language="de",
+ api_key="sk-test",
+ client=client,
+ audio_url="https://example.com/a.wav",
+ **({} if language_hints is None else {"language_hints": language_hints}),
+ )
+ body = next(c for c in client.calls if c["method"] == "POST")["json"]
+ assert list(body["language_hints"]) == expected
+ assert "language" not in body
+
+
+class TestProxyFormFields:
+ def test_should_send_soniox_types_for_fields_the_proxy_parsed_from_an_sdk_form(self, monkeypatch):
+ import litellm
+ from starlette.requests import Request
+
+ from litellm.proxy.common_utils.http_parsing_utils import get_form_data
+
+ monkeypatch.setattr("time.sleep", lambda *_: None)
+ multipart = httpx.Request(
+ "POST",
+ "http://proxy/v1/audio/transcriptions",
+ data={
+ "model": "soniox/stt-async-v4",
+ "audio_url": "https://example.com/a.wav",
+ "context[terms][]": ["Celebrex"],
+ "context[general][][key]": "domain",
+ "context[general][][value]": "Healthcare",
+ "translation[type]": "one_way",
+ "translation[target_language]": "de",
+ "enable_speaker_diarization": "true",
+ "language_hints[]": ["en", "es"],
+ },
+ files={"file": ("a.wav", b"RIFF", "audio/wav")},
+ )
+ multipart_body = multipart.read()
+
+ async def receive():
+ return {"type": "http.request", "body": multipart_body, "more_body": False}
+
+ request = Request(
+ {
+ "type": "http",
+ "method": "POST",
+ "path": "/v1/audio/transcriptions",
+ "headers": [(b"content-type", multipart.headers["content-type"].encode())],
+ },
+ receive,
+ )
+ form = asyncio.run(get_form_data(request))
+ client = _MockSyncClient(
+ {
+ "POST https://api.soniox.com/v1/transcriptions": [_make_response({"id": "tx_1"})],
+ "GET https://api.soniox.com/v1/transcriptions/tx_1": [_make_response({"status": "completed"})],
+ "GET https://api.soniox.com/v1/transcriptions/tx_1/transcript": [
+ _make_response({"text": "hello", "tokens": []})
+ ],
+ "DELETE https://api.soniox.com/v1/transcriptions/tx_1": [_make_response({})],
+ }
+ )
+ litellm.transcription(
+ file=("a.wav", b"RIFF", "audio/wav"),
+ api_key="sk-test",
+ client=client,
+ **{key: value for key, value in form.items() if key != "file"},
+ )
+ body = json.loads(json.dumps(next(c for c in client.calls if c["method"] == "POST")["json"]))
+ assert body["context"] == {"terms": ["Celebrex"], "general": [{"key": "domain", "value": "Healthcare"}]}
+ assert body["translation"] == {"type": "one_way", "target_language": "de"}
+ assert body["enable_speaker_diarization"] is True
+ assert body["language_hints"] == ["en", "es"]
+
+
class TestErrorResponses:
def test_should_raise_on_4xx_during_create_with_json_error(self, monkeypatch):
monkeypatch.setattr("time.sleep", lambda *_: None)
@@ -780,6 +945,51 @@ class TestPassthroughBodyBuilding:
post_call = next(c for c in client.calls if c["method"] == "POST")
assert "context" not in post_call["json"]
+ def test_should_send_form_string_params_to_soniox_as_json_types(self, monkeypatch):
+ """Simulates the proxy path, where every multipart form field is a string."""
+ monkeypatch.setattr("time.sleep", lambda *_: None)
+ responses = {
+ "POST https://api.soniox.com/v1/transcriptions": [
+ _make_response({"id": "tx_1"}),
+ ],
+ "GET https://api.soniox.com/v1/transcriptions/tx_1": [
+ _make_response({"status": "completed"}),
+ ],
+ "GET https://api.soniox.com/v1/transcriptions/tx_1/transcript": [
+ _make_response({"text": "ok", "tokens": []}),
+ ],
+ "DELETE https://api.soniox.com/v1/transcriptions/tx_1": [
+ _make_response({}),
+ ],
+ }
+ client = _MockSyncClient(responses)
+ form_params = {
+ "audio_url": "https://example.com/a.wav",
+ "context": '{"terms": ["Celebrex"], "general": [{"key": "domain", "value": "Healthcare"}]}',
+ "translation": '{"type": "one_way", "target_language": "es"}',
+ "language_hints": "en,es",
+ "language_hints_strict": "true",
+ "enable_speaker_diarization": "true",
+ "enable_language_identification": "false",
+ "soniox_polling_interval": "0.5",
+ }
+ SonioxAudioTranscriptionHandler().audio_transcriptions(
+ audio_file=None,
+ optional_params=form_params,
+ litellm_params={},
+ atranscription=False,
+ **_common_call_kwargs(client),
+ )
+ body = next(c for c in client.calls if c["method"] == "POST")["json"]
+ assert body["context"] == {"terms": ["Celebrex"], "general": [{"key": "domain", "value": "Healthcare"}]}
+ assert body["translation"] == {"type": "one_way", "target_language": "es"}
+ assert list(body["language_hints"]) == ["en", "es"]
+ assert body["language_hints_strict"] is True
+ assert body["enable_speaker_diarization"] is True
+ assert body["enable_language_identification"] is False
+ assert "soniox_polling_interval" not in body
+ assert form_params["context"].startswith("{"), "caller's params must not be mutated"
+
class TestSecretRedaction:
"""Secret-bearing fields must be redacted before reaching logging callbacks.
diff --git a/tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py b/tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py
index c99b28adc49..e3e1c9ab493 100644
--- a/tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py
+++ b/tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py
@@ -9,8 +9,16 @@ import pytest
from litellm.llms.soniox.audio_transcription.transformation import (
SonioxAudioTranscriptionConfig,
+ SonioxInvalidFormParam,
+ decode_soniox_form_params,
+ raise_soniox_form_error,
)
from litellm.llms.soniox.common_utils import SonioxException
+from litellm.llms.soniox.types import (
+ SonioxContext,
+ SonioxContextGeneralEntry,
+ SonioxTranslationTerm,
+)
from litellm.types.utils import TranscriptionResponse
@@ -40,9 +48,9 @@ class TestMapOpenAIParams:
model="stt-async-v4",
drop_params=False,
)
- assert result["language_hints"] == ["en"]
+ assert list(result["language_hints"]) == ["en"]
- def test_should_prepend_language_to_existing_hints(self):
+ def test_should_keep_explicit_hints_over_language(self):
cfg = SonioxAudioTranscriptionConfig()
result = cfg.map_openai_params(
non_default_params={"language": "en"},
@@ -50,7 +58,7 @@ class TestMapOpenAIParams:
model="stt-async-v4",
drop_params=False,
)
- assert result["language_hints"] == ["en", "fr"]
+ assert list(result["language_hints"]) == ["fr"]
def test_should_not_duplicate_language_already_in_hints(self):
cfg = SonioxAudioTranscriptionConfig()
@@ -60,7 +68,7 @@ class TestMapOpenAIParams:
model="stt-async-v4",
drop_params=False,
)
- assert result["language_hints"] == ["en", "fr"]
+ assert list(result["language_hints"]) == ["en", "fr"]
def test_should_passthrough_soniox_native_kwargs(self):
cfg = SonioxAudioTranscriptionConfig()
@@ -243,6 +251,21 @@ class TestTransformAudioTranscriptionResponse:
assert "Speaker 1:" in resp.text
assert "Speaker 2:" in resp.text
+ def test_should_keep_source_transcript_for_translation_without_tags(self):
+ cfg = SonioxAudioTranscriptionConfig()
+ payload = {
+ "transcript": {
+ "text": "Hello there.",
+ "tokens": [
+ {"text": "Hello", "translation_status": "original"},
+ {"text": " there.", "translation_status": "original"},
+ {"text": "Hallo", "translation_status": "translation"},
+ {"text": " zusammen.", "translation_status": "translation"},
+ ],
+ }
+ }
+ assert cfg.build_response_from_payload(payload).text == "Hello there."
+
def test_should_set_language_when_all_tokens_share_one(self):
cfg = SonioxAudioTranscriptionConfig()
payload = {
@@ -678,3 +701,175 @@ class TestGetErrorClass:
err = cfg.get_error_class(error_message="boom", status_code=500, headers={})
assert isinstance(err, SonioxException)
assert err.status_code == 500
+
+
+class TestDecodeSonioxFormParams:
+ """Proxy multipart form fields arrive as strings; Soniox needs real JSON types."""
+
+ def test_should_decode_json_string_context_with_all_four_sections(self):
+ context = {
+ "general": [{"key": "domain", "value": "Healthcare"}],
+ "text": "Follow-up visit for a patient on blood thinners.",
+ "terms": ["Celebrex", "Zyrtec"],
+ "translation_terms": [{"source": "Mr. Smith", "target": "Sr. Smith"}],
+ }
+ result = decode_soniox_form_params({"context": json.dumps(context)})
+ assert result["context"] == context
+
+ def test_should_keep_plain_text_context_as_string(self):
+ result = decode_soniox_form_params({"context": "medical conversation"})
+ assert result["context"] == "medical conversation"
+
+ def test_should_leave_sdk_typed_values_untouched(self):
+ context = SonioxContext(
+ general=[SonioxContextGeneralEntry(key="domain", value="Healthcare")],
+ text="Follow-up visit notes.",
+ terms=["Celebrex"],
+ translation_terms=[SonioxTranslationTerm(source="Mr. Smith", target="Sr. Smith")],
+ )
+ translation = {"type": "one_way", "target_language": "es"}
+ result = decode_soniox_form_params(
+ {
+ "context": context,
+ "translation": translation,
+ "enable_speaker_diarization": True,
+ "language_hints": ["en", "es"],
+ }
+ )
+ assert result["context"] is context
+ assert result["translation"] is translation
+ assert result["enable_speaker_diarization"] is True
+ assert result["language_hints"] == ["en", "es"]
+
+ @pytest.mark.parametrize("raw", ["[inaudible] consultation", "{unbalanced", '{"terms": ["Celebrex"'])
+ def test_should_keep_bracket_prefixed_free_form_context_as_string(self, raw: str):
+ assert decode_soniox_form_params({"context": raw})["context"] == raw
+
+ def test_should_decode_json_string_translation(self):
+ result = decode_soniox_form_params({"translation": '{"type": "one_way", "target_language": "es"}'})
+ assert result["translation"] == {"type": "one_way", "target_language": "es"}
+
+ @pytest.mark.parametrize(
+ "raw, expected",
+ [("true", True), ("True", True), ("1", True), ("false", False), ("FALSE", False), ("0", False)],
+ )
+ def test_should_decode_boolean_strings(self, raw: str, expected: bool):
+ result = decode_soniox_form_params(
+ {
+ "enable_speaker_diarization": raw,
+ "enable_language_identification": raw,
+ "language_hints_strict": raw,
+ }
+ )
+ assert result == {
+ "enable_speaker_diarization": expected,
+ "enable_language_identification": expected,
+ "language_hints_strict": expected,
+ }
+
+ @pytest.mark.parametrize("raw, expected", [("on", True), ("yes", True), ("off", False), ("no", False)])
+ def test_should_decode_html_style_boolean_strings(self, raw: str, expected: bool):
+ assert decode_soniox_form_params({"enable_speaker_diarization": raw})["enable_speaker_diarization"] is expected
+
+ def test_should_return_error_value_on_non_boolean_string(self):
+ result = decode_soniox_form_params({"context": '{"general": []}', "enable_speaker_diarization": "maybe"})
+ assert result == SonioxInvalidFormParam(key="enable_speaker_diarization", value="maybe", expected="a boolean")
+
+ def test_should_drop_non_boolean_string_when_drop_params_is_set(self):
+ result = decode_soniox_form_params(
+ {"enable_speaker_diarization": "maybe", "language_hints_strict": "true"}, drop_params=True
+ )
+ assert result == {"language_hints_strict": True}
+
+ def test_should_return_error_value_on_non_numeric_polling_interval(self):
+ result = decode_soniox_form_params({"soniox_polling_interval": "abc"})
+ assert result == SonioxInvalidFormParam(key="soniox_polling_interval", value="abc", expected="a number")
+
+ def test_should_drop_non_numeric_polling_interval_when_drop_params_is_set(self):
+ assert decode_soniox_form_params({"soniox_polling_interval": "abc"}, drop_params=True) == {}
+
+ def test_should_map_invalid_form_value_error_to_400(self):
+ with pytest.raises(SonioxException) as exc_info:
+ raise_soniox_form_error(
+ SonioxInvalidFormParam(key="enable_speaker_diarization", value="maybe", expected="a boolean")
+ )
+ assert exc_info.value.status_code == 400
+ assert "enable_speaker_diarization" in str(exc_info.value)
+ assert "maybe" in str(exc_info.value)
+
+ @pytest.mark.parametrize(
+ "raw", ['["Celebrex", "Zyrtec"]', '{"note": "follow-up visit"}', '{"terms": "Celebrex"}', '"quoted"']
+ )
+ def test_should_keep_json_context_soniox_cannot_take_as_an_object_as_string(self, raw: str):
+ assert decode_soniox_form_params({"context": raw})["context"] == raw
+
+ def test_should_nest_openai_sdk_bracketed_form_fields(self):
+ result = decode_soniox_form_params(
+ {
+ "context[terms]": ["Celebrex", "Zyrtec"],
+ "context[text]": "Follow-up visit.",
+ "translation[type]": "one_way",
+ "translation[target_language]": "de",
+ "client_reference_id": "ref-1",
+ }
+ )
+ assert result == {
+ "context": {"terms": ["Celebrex", "Zyrtec"], "text": "Follow-up visit."},
+ "translation": {"type": "one_way", "target_language": "de"},
+ "client_reference_id": "ref-1",
+ }
+
+ def test_should_prefer_explicit_field_over_bracketed_fields(self):
+ result = decode_soniox_form_params({"context": "plain text", "context[terms]": ["Celebrex"]})
+ assert result == {"context": "plain text"}
+
+ @pytest.mark.parametrize(
+ "raw, expected",
+ [('["en", "es"]', ["en", "es"]), ("en,es", ["en", "es"]), ("en, es ,", ["en", "es"]), ("en", ["en"])],
+ )
+ def test_should_decode_language_hints_strings(self, raw: str, expected: list):
+ assert list(decode_soniox_form_params({"language_hints": raw})["language_hints"]) == expected
+
+ @pytest.mark.parametrize("raw", ['{"type": "one_way"', "de", '["one_way"]'])
+ def test_should_return_error_value_on_translation_that_is_not_a_json_object(self, raw: str):
+ assert decode_soniox_form_params({"translation": raw}) == SonioxInvalidFormParam(
+ key="translation", value=raw, expected="a JSON object"
+ )
+
+ def test_should_return_error_value_on_malformed_json_language_hints(self):
+ assert decode_soniox_form_params({"language_hints": "[en, es]"}) == SonioxInvalidFormParam(
+ key="language_hints", value="[en, es]", expected="a JSON list"
+ )
+
+ def test_should_drop_malformed_json_params_when_drop_params_is_set(self):
+ result = decode_soniox_form_params(
+ {"translation": '{"type": "one_way"', "language_hints": "[en, es]", "language_hints_strict": "true"},
+ drop_params=True,
+ )
+ assert result == {"language_hints_strict": True}
+
+ def test_should_rebuild_openai_sdk_list_of_object_context_fields(self):
+ result = decode_soniox_form_params(
+ {
+ "context[terms]": ["Celebrex"],
+ "context[general][][key]": "domain",
+ "context[general][][value]": "Healthcare",
+ "context[translation_terms][][source]": "Mr. Smith",
+ "context[translation_terms][][target]": "Sr. Smith",
+ }
+ )
+ assert result == {
+ "context": {
+ "terms": ["Celebrex"],
+ "general": ({"key": "domain", "value": "Healthcare"},),
+ "translation_terms": ({"source": "Mr. Smith", "target": "Sr. Smith"},),
+ }
+ }
+
+ def test_should_leave_bracketed_context_soniox_cannot_take_unnested(self):
+ params = {"context[terms]": ["Celebrex"], "context[note]": "follow-up"}
+ assert decode_soniox_form_params(params) == params
+
+ def test_should_not_json_decode_unrelated_string_params(self):
+ result = decode_soniox_form_params({"client_reference_id": '["not", "json-decoded"]'})
+ assert result["client_reference_id"] == '["not", "json-decoded"]'
diff --git a/tests/unit/llms/test_predibase_transformation.py b/tests/unit/llms/test_predibase_transformation.py
index d6425488be3..20d547c3b63 100644
--- a/tests/unit/llms/test_predibase_transformation.py
+++ b/tests/unit/llms/test_predibase_transformation.py
@@ -39,14 +39,9 @@ def test_predibase_transform_request_non_stream():
assert "stream" not in request_data["parameters"]
-def test_predibase_transform_request_custom_prompt(monkeypatch):
+def test_predibase_transform_request_custom_prompt():
config = PredibaseConfig()
- monkeypatch.setattr(
- "litellm.llms.predibase.chat.transformation.custom_prompt",
- lambda **kwargs: "custom-prompt",
- )
-
request_data = config.transform_request(
model="predibase-model",
messages=[{"role": "user", "content": "hello"}],
@@ -54,16 +49,16 @@ def test_predibase_transform_request_custom_prompt(monkeypatch):
litellm_params={
"custom_prompt_dict": {
"predibase-model": {
- "roles": {},
- "initial_prompt_value": "",
- "final_prompt_value": "",
+ "roles": {"user": {"pre_message": "", "post_message": ""}},
+ "initial_prompt_value": "",
+ "final_prompt_value": "",
}
}
},
headers={},
)
- assert request_data["inputs"] == "custom-prompt"
+ assert request_data["inputs"] == "hello"
def test_predibase_get_complete_url_stream_and_non_stream():
diff --git a/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py b/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py
index b50d6e19458..5c7099832ba 100644
--- a/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py
+++ b/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py
@@ -17,6 +17,7 @@ from litellm.llms.vertex_ai.audio_transcription.realtime_transformation import (
parse_chirp_session_update,
)
from litellm.llms.vertex_ai.common_utils import VertexAIError
+from litellm.llms.vertex_ai.realtime.transformation import vertex_realtime_config
from litellm.types.llms.vertex_ai_speech_to_text import (
VertexSpeechStreamingConfigured,
VertexSpeechStreamingResponse,
@@ -173,6 +174,26 @@ def test_unsupported_session_settings_are_rejected(payload: str, message: str):
parse_chirp_session_update(payload, MODEL)
+def _configure_command(turn_detection: str, drop_params: bool) -> list[object]:
+ config = vertex_realtime_config(
+ MODEL, access_token="token", resolve_access_token=_token, project="proj-1", location="us", drop_params=drop_params
+ )
+ config.transform_session_created_event(MODEL, "sess_1")
+ return [
+ json.loads(command) if isinstance(command, str) else command
+ for command in config.transform_realtime_request(_ga_session_update(turn_detection=turn_detection), MODEL)
+ ]
+
+
+def test_semantic_vad_falls_back_to_server_vad_under_drop_params():
+ assert _configure_command("semantic_vad", drop_params=True) == _configure_command("server_vad", drop_params=False)
+
+
+def test_semantic_vad_without_drop_params_names_the_drop_params_fallback():
+ with pytest.raises(ChirpProtocolError, match="Set drop_params to use server_vad instead"):
+ _configure_command("semantic_vad", drop_params=False)
+
+
def test_session_update_configures_once_and_later_updates_are_ignored():
config = _config()
config.transform_session_created_event(MODEL, "sess_1")
diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
index 40ede74b98b..5c0e9f1a305 100644
--- a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
+++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
@@ -1,6 +1,7 @@
import asyncio, importlib, os, uuid
import json
import re
+from collections.abc import Awaitable, Callable, Mapping
from copy import deepcopy
from typing import Final, List, Literal, cast, get_args
from unittest.mock import MagicMock, Mock, patch
@@ -25,6 +26,7 @@ from litellm.types.llms.vertex_ai import GeminiFinishReason, UsageMetadata
from litellm.types.utils import ChoiceLogprobs, Usage
from litellm.utils import _invalidate_model_cost_lowercase_map, CustomStreamWrapper
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
+from litellm.integrations.custom_logger import CustomLogger
from litellm.llms.vertex_ai.gemini.transformation import (
gemini_convert_messages_with_history,
)
@@ -38,62 +40,62 @@ from litellm.llms.vertex_ai.context_caching.transformation import (
import tempfile
GEMINI_3_IMAGE_SIZE_MAPPINGS: Final[tuple[tuple[str, str, str], ...]] = (
- ('512x512', '1:1', '512'),
- ('1024x1024', '1:1', '1K'),
- ('2048x2048', '1:1', '2K'),
- ('4096x4096', '1:1', '4K'),
- ('256x1024', '1:4', '512'),
- ('512x2048', '1:4', '1K'),
- ('1024x4096', '1:4', '2K'),
- ('2048x8192', '1:4', '4K'),
- ('192x1536', '1:8', '512'),
- ('384x3072', '1:8', '1K'),
- ('768x6144', '1:8', '2K'),
- ('1536x12288', '1:8', '4K'),
- ('424x632', '2:3', '512'),
- ('848x1264', '2:3', '1K'),
- ('1696x2528', '2:3', '2K'),
- ('3392x5056', '2:3', '4K'),
- ('632x424', '3:2', '512'),
- ('1264x848', '3:2', '1K'),
- ('2528x1696', '3:2', '2K'),
- ('5056x3392', '3:2', '4K'),
- ('448x600', '3:4', '512'),
- ('896x1200', '3:4', '1K'),
- ('1792x2400', '3:4', '2K'),
- ('3584x4800', '3:4', '4K'),
- ('1024x256', '4:1', '512'),
- ('2048x512', '4:1', '1K'),
- ('4096x1024', '4:1', '2K'),
- ('8192x2048', '4:1', '4K'),
- ('600x448', '4:3', '512'),
- ('1200x896', '4:3', '1K'),
- ('2400x1792', '4:3', '2K'),
- ('4800x3584', '4:3', '4K'),
- ('464x576', '4:5', '512'),
- ('928x1152', '4:5', '1K'),
- ('1856x2304', '4:5', '2K'),
- ('3712x4608', '4:5', '4K'),
- ('576x464', '5:4', '512'),
- ('1152x928', '5:4', '1K'),
- ('2304x1856', '5:4', '2K'),
- ('4608x3712', '5:4', '4K'),
- ('1536x192', '8:1', '512'),
- ('3072x384', '8:1', '1K'),
- ('6144x768', '8:1', '2K'),
- ('12288x1536', '8:1', '4K'),
- ('384x688', '9:16', '512'),
- ('768x1376', '9:16', '1K'),
- ('1536x2752', '9:16', '2K'),
- ('3072x5504', '9:16', '4K'),
- ('688x384', '16:9', '512'),
- ('1376x768', '16:9', '1K'),
- ('2752x1536', '16:9', '2K'),
- ('5504x3072', '16:9', '4K'),
- ('792x336', '21:9', '512'),
- ('1584x672', '21:9', '1K'),
- ('3168x1344', '21:9', '2K'),
- ('6336x2688', '21:9', '4K'),
+ ("512x512", "1:1", "512"),
+ ("1024x1024", "1:1", "1K"),
+ ("2048x2048", "1:1", "2K"),
+ ("4096x4096", "1:1", "4K"),
+ ("256x1024", "1:4", "512"),
+ ("512x2048", "1:4", "1K"),
+ ("1024x4096", "1:4", "2K"),
+ ("2048x8192", "1:4", "4K"),
+ ("192x1536", "1:8", "512"),
+ ("384x3072", "1:8", "1K"),
+ ("768x6144", "1:8", "2K"),
+ ("1536x12288", "1:8", "4K"),
+ ("424x632", "2:3", "512"),
+ ("848x1264", "2:3", "1K"),
+ ("1696x2528", "2:3", "2K"),
+ ("3392x5056", "2:3", "4K"),
+ ("632x424", "3:2", "512"),
+ ("1264x848", "3:2", "1K"),
+ ("2528x1696", "3:2", "2K"),
+ ("5056x3392", "3:2", "4K"),
+ ("448x600", "3:4", "512"),
+ ("896x1200", "3:4", "1K"),
+ ("1792x2400", "3:4", "2K"),
+ ("3584x4800", "3:4", "4K"),
+ ("1024x256", "4:1", "512"),
+ ("2048x512", "4:1", "1K"),
+ ("4096x1024", "4:1", "2K"),
+ ("8192x2048", "4:1", "4K"),
+ ("600x448", "4:3", "512"),
+ ("1200x896", "4:3", "1K"),
+ ("2400x1792", "4:3", "2K"),
+ ("4800x3584", "4:3", "4K"),
+ ("464x576", "4:5", "512"),
+ ("928x1152", "4:5", "1K"),
+ ("1856x2304", "4:5", "2K"),
+ ("3712x4608", "4:5", "4K"),
+ ("576x464", "5:4", "512"),
+ ("1152x928", "5:4", "1K"),
+ ("2304x1856", "5:4", "2K"),
+ ("4608x3712", "5:4", "4K"),
+ ("1536x192", "8:1", "512"),
+ ("3072x384", "8:1", "1K"),
+ ("6144x768", "8:1", "2K"),
+ ("12288x1536", "8:1", "4K"),
+ ("384x688", "9:16", "512"),
+ ("768x1376", "9:16", "1K"),
+ ("1536x2752", "9:16", "2K"),
+ ("3072x5504", "9:16", "4K"),
+ ("688x384", "16:9", "512"),
+ ("1376x768", "16:9", "1K"),
+ ("2752x1536", "16:9", "2K"),
+ ("5504x3072", "16:9", "4K"),
+ ("792x336", "21:9", "512"),
+ ("1584x672", "21:9", "1K"),
+ ("3168x1344", "21:9", "2K"),
+ ("6336x2688", "21:9", "4K"),
)
@@ -314,9 +316,7 @@ def test_vertex_ai_response_json_schema_preserves_refs_for_gemini_2():
# $defs and $ref should be preserved (not unpacked)
assert "response_json_schema" in transformed_request
result_schema = transformed_request["response_json_schema"]
- assert (
- "$defs" in result_schema
- ), "responseJsonSchema should preserve $defs for Gemini 2.0+"
+ assert "$defs" in result_schema, "responseJsonSchema should preserve $defs for Gemini 2.0+"
def test_vertex_ai_get_json_schema_preserves_refs_for_nested_pydantic():
@@ -396,22 +396,14 @@ def test_vertex_ai_response_json_schema_for_gemini_2():
# Types should be lowercase (standard JSON Schema format)
assert transformed_request["response_json_schema"]["type"] == "object"
- assert (
- transformed_request["response_json_schema"]["properties"]["name"]["type"]
- == "string"
- )
- assert (
- transformed_request["response_json_schema"]["properties"]["age"]["type"]
- == "integer"
- )
+ assert transformed_request["response_json_schema"]["properties"]["name"]["type"] == "string"
+ assert transformed_request["response_json_schema"]["properties"]["age"]["type"] == "integer"
# Should NOT have propertyOrdering (not needed for responseJsonSchema)
assert "propertyOrdering" not in transformed_request["response_json_schema"]
# additionalProperties should be preserved (supported by responseJsonSchema)
- assert (
- transformed_request["response_json_schema"].get("additionalProperties") == False
- )
+ assert transformed_request["response_json_schema"].get("additionalProperties") == False
def test_vertex_ai_response_schema_for_old_models():
@@ -648,7 +640,9 @@ def test_response_has_search_grounding_detection():
is False
)
assert (
- VertexGeminiConfig._response_has_search_grounding({"candidates": [{"groundingMetadata": {"webSearchQueries": []}}]})
+ VertexGeminiConfig._response_has_search_grounding(
+ {"candidates": [{"groundingMetadata": {"webSearchQueries": []}}]}
+ )
is False
)
assert VertexGeminiConfig._response_has_search_grounding({"candidates": []}) is False
@@ -789,9 +783,7 @@ def test_vertex_ai_url_context_tool_use_tokens_billed_as_input_tokens():
"candidates": [
{
"urlContextMetadata": {"urlMetadata": []},
- "groundingMetadata": {
- "groundingChunks": [{"web": {"uri": "https://example.com", "title": "Example"}}]
- },
+ "groundingMetadata": {"groundingChunks": [{"web": {"uri": "https://example.com", "title": "Example"}}]},
}
],
"usageMetadata": UsageMetadata(
@@ -823,9 +815,7 @@ def test_url_context_tool_use_tokens_are_priced_as_input():
"candidates": [
{
"urlContextMetadata": {"urlMetadata": []},
- "groundingMetadata": {
- "groundingChunks": [{"web": {"uri": "https://example.com", "title": "Example"}}]
- },
+ "groundingMetadata": {"groundingChunks": [{"web": {"uri": "https://example.com", "title": "Example"}}]},
}
],
"usageMetadata": UsageMetadata(
@@ -886,9 +876,7 @@ def test_missing_prompt_tokens_details_still_prices_the_full_prompt():
"candidates": [
{
"urlContextMetadata": {"urlMetadata": []},
- "groundingMetadata": {
- "groundingChunks": [{"web": {"uri": "https://example.com", "title": "Example"}}]
- },
+ "groundingMetadata": {"groundingChunks": [{"web": {"uri": "https://example.com", "title": "Example"}}]},
}
],
"usageMetadata": UsageMetadata(
@@ -925,9 +913,7 @@ def test_streaming_chunk_includes_reasoning_tokens():
"thoughtsTokenCount": 3,
},
}
- iterator = ModelResponseIterator(
- streaming_response=[], sync_stream=True, logging_obj=litellm_logging
- )
+ iterator = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=litellm_logging)
streaming_chunk = iterator.chunk_parser(chunk)
assert streaming_chunk.usage is not None
assert streaming_chunk.usage.prompt_tokens == 5
@@ -963,17 +949,12 @@ def test_streaming_chunk_includes_reasoning_content():
"usageMetadata": {},
}
- iterator = ModelResponseIterator(
- streaming_response=[], sync_stream=True, logging_obj=litellm_logging
- )
+ iterator = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=litellm_logging)
streaming_chunk = iterator.chunk_parser(chunk)
# The text content should be empty and reasoning_content should be populated
assert streaming_chunk.choices[0].delta.content is None
- assert (
- streaming_chunk.choices[0].delta.reasoning_content
- == "I'm thinking through the problem..."
- )
+ assert streaming_chunk.choices[0].delta.reasoning_content == "I'm thinking through the problem..."
def test_streaming_chunk_with_tool_calls_and_thought_includes_reasoning_content():
@@ -1018,24 +999,16 @@ def test_streaming_chunk_with_tool_calls_and_thought_includes_reasoning_content(
},
}
- iterator = ModelResponseIterator(
- streaming_response=[], sync_stream=True, logging_obj=litellm_logging
- )
+ iterator = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=litellm_logging)
streaming_chunk = iterator.chunk_parser(chunk)
# Verify reasoning_content comes from the thought: true part
- assert (
- streaming_chunk.choices[0].delta.reasoning_content
- == "Let me think about how to get the time..."
- )
+ assert streaming_chunk.choices[0].delta.reasoning_content == "Let me think about how to get the time..."
# Verify tool calls are also present
assert streaming_chunk.choices[0].delta.tool_calls is not None
assert len(streaming_chunk.choices[0].delta.tool_calls) == 1
- assert (
- streaming_chunk.choices[0].delta.tool_calls[0].function.name
- == "get_current_time"
- )
+ assert streaming_chunk.choices[0].delta.tool_calls[0].function.name == "get_current_time"
def test_streaming_chunk_with_tool_calls_no_thought_no_reasoning_content():
@@ -1077,9 +1050,7 @@ def test_streaming_chunk_with_tool_calls_no_thought_no_reasoning_content():
},
}
- iterator = ModelResponseIterator(
- streaming_response=[], sync_stream=True, logging_obj=litellm_logging
- )
+ iterator = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=litellm_logging)
streaming_chunk = iterator.chunk_parser(chunk)
# reasoning_content should be None - thoughtSignature alone does NOT mean reasoning
@@ -1286,11 +1257,7 @@ def test_vertex_ai_usage_metadata_with_image_tokens_in_prompt():
# Verify the math: prompt_tokens = text + image
# 533 = 6 (text) + 527 (image)
- assert (
- result.prompt_tokens_details.text_tokens
- + result.prompt_tokens_details.image_tokens
- == result.prompt_tokens
- )
+ assert result.prompt_tokens_details.text_tokens + result.prompt_tokens_details.image_tokens == result.prompt_tokens
def test_map_response_modalities_video():
@@ -1480,9 +1447,7 @@ def test_vertex_ai_streaming_usage_calculation():
}
# Create iterator and parse chunk
- iterator = ModelResponseIterator(
- streaming_response=[], sync_stream=True, logging_obj=MagicMock()
- )
+ iterator = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=MagicMock())
iterator.chunk_parser(chunk)
# Verify calculate_usage was called with correct parameters
@@ -1554,9 +1519,7 @@ def test_vertex_ai_streaming_usage_web_search_calculation():
}
# Create iterator and parse chunk
- iterator = ModelResponseIterator(
- streaming_response=[], sync_stream=True, logging_obj=MagicMock()
- )
+ iterator = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=MagicMock())
completed_response = iterator.chunk_parser(chunk)
usage: Usage = completed_response.usage
@@ -1775,7 +1738,9 @@ async def test_gemini_image_generation_async_stream(respx_mock: respx.MockRouter
"modelVersion": "gemini-2.5-flash-image",
},
{
- "candidates": [{"content": {"role": "model", "parts": [{"text": ""}]}, "finishReason": "STOP", "index": 0}],
+ "candidates": [
+ {"content": {"role": "model", "parts": [{"text": ""}]}, "finishReason": "STOP", "index": 0}
+ ],
"usageMetadata": {"promptTokenCount": 13, "candidatesTokenCount": 1290, "totalTokenCount": 1303},
"modelVersion": "gemini-2.5-flash-image",
},
@@ -2029,9 +1994,14 @@ def test_gemini_openai_web_search_tool_to_google_search(respx_mock: respx.MockRo
GEMINI_GROUNDING_METADATA: Final = {
"webSearchQueries": ["capital of France"],
- "searchEntryPoint": {"renderedContent": ""},
+ "searchEntryPoint": {"renderedContent": ''},
"groundingChunks": [
- {"web": {"uri": "https://vertexaisearch.cloud.google.com/grounding-api-redirect/AbF9w", "title": "wikipedia.org"}}
+ {
+ "web": {
+ "uri": "https://vertexaisearch.cloud.google.com/grounding-api-redirect/AbF9w",
+ "title": "wikipedia.org",
+ }
+ }
],
"groundingSupports": [
{
@@ -2188,9 +2158,7 @@ def test_vertex_ai_transform_parts():
# Test case 2: Tool call mode (is_function_call=False) - Single message with multiple tool calls
parts_with_multiple_functions = [
- HttpxPartType(
- functionCall={"name": "get_current_weather", "args": {"location": "Boston"}}
- ),
+ HttpxPartType(functionCall={"name": "get_current_weather", "args": {"location": "Boston"}}),
HttpxPartType(
functionCall={
"name": "get_forecast",
@@ -2220,9 +2188,7 @@ def test_vertex_ai_transform_parts():
# Test case 3: Simulating multiple messages - cumulative indexing across messages
# First message with 2 tool calls (starting from index 0)
first_message_parts = [
- HttpxPartType(
- functionCall={"name": "get_weather", "args": {"location": "Boston"}}
- ),
+ HttpxPartType(functionCall={"name": "get_weather", "args": {"location": "Boston"}}),
HttpxPartType(functionCall={"name": "get_time", "args": {"timezone": "EST"}}),
]
@@ -2239,9 +2205,7 @@ def test_vertex_ai_transform_parts():
# Second message with 1 tool call (continuing from previous index)
second_message_parts = [
- HttpxPartType(
- functionCall={"name": "send_email", "args": {"to": "user@example.com"}}
- ),
+ HttpxPartType(functionCall={"name": "send_email", "args": {"to": "user@example.com"}}),
]
function, tools, updated_idx = VertexGeminiConfig._transform_parts(
@@ -2259,9 +2223,7 @@ def test_vertex_ai_transform_parts():
# Third message with 2 more tool calls (continuing from previous index)
third_message_parts = [
- HttpxPartType(
- functionCall={"name": "create_calendar_event", "args": {"title": "Meeting"}}
- ),
+ HttpxPartType(functionCall={"name": "create_calendar_event", "args": {"title": "Meeting"}}),
HttpxPartType(functionCall={"name": "set_reminder", "args": {"time": "10:00"}}),
]
@@ -2308,9 +2270,7 @@ def test_vertex_ai_transform_parts():
assert updated_idx == 10 # Index should remain unchanged
# Test case 6: Function call with empty args
- parts_with_empty_args = [
- HttpxPartType(functionCall={"name": "simple_function", "args": {}})
- ]
+ parts_with_empty_args = [HttpxPartType(functionCall={"name": "simple_function", "args": {}})]
function, tools, updated_idx = VertexGeminiConfig._transform_parts(
parts=parts_with_empty_args, cumulative_tool_call_idx=0, is_function_call=True
@@ -2326,13 +2286,9 @@ def test_vertex_ai_transform_parts():
# Test case 7: Mixed content with function calls - ensuring tool call IDs are unique
mixed_parts = [
HttpxPartType(text="Before function call"),
- HttpxPartType(
- functionCall={"name": "function_a", "args": {"param": "value_a"}}
- ),
+ HttpxPartType(functionCall={"name": "function_a", "args": {"param": "value_a"}}),
HttpxPartType(text="Between function calls"),
- HttpxPartType(
- functionCall={"name": "function_b", "args": {"param": "value_b"}}
- ),
+ HttpxPartType(functionCall={"name": "function_b", "args": {"param": "value_b"}}),
HttpxPartType(text="After function calls"),
]
@@ -2524,25 +2480,17 @@ def test_vertex_ai_tool_call_id_format():
tool_id = tool["id"]
# Should start with 'call_'
- assert tool_id.startswith(
- "call_"
- ), f"ID should start with 'call_', got: {tool_id}"
+ assert tool_id.startswith("call_"), f"ID should start with 'call_', got: {tool_id}"
# Should have exactly 33 total characters (call_ + 28 hex chars)
- assert (
- len(tool_id) == 33
- ), f"ID should be 33 characters long, got {len(tool_id)}: {tool_id}"
+ assert len(tool_id) == 33, f"ID should be 33 characters long, got {len(tool_id)}: {tool_id}"
# The part after 'call_' should be 28 hex characters
hex_part = tool_id[5:] # Remove 'call_' prefix
- assert (
- len(hex_part) == 28
- ), f"Hex part should be 28 characters, got {len(hex_part)}: {hex_part}"
+ assert len(hex_part) == 28, f"Hex part should be 28 characters, got {len(hex_part)}: {hex_part}"
# Should only contain valid hex characters
- assert re.match(
- r"^[0-9a-f]{28}$", hex_part
- ), f"Should contain only lowercase hex chars, got: {hex_part}"
+ assert re.match(r"^[0-9a-f]{28}$", hex_part), f"Should contain only lowercase hex chars, got: {hex_part}"
# Verify IDs are unique
assert tools[0]["id"] != tools[1]["id"], "Tool call IDs should be unique"
@@ -2559,9 +2507,7 @@ def test_vertex_ai_tool_call_id_format():
ids_generated.add(test_tools[0]["id"])
# All generated IDs should be unique
- assert (
- len(ids_generated) == 10
- ), f"All 10 IDs should be unique, got {len(ids_generated)} unique IDs"
+ assert len(ids_generated) == 10, f"All 10 IDs should be unique, got {len(ids_generated)} unique IDs"
def test_vertex_ai_map_google_maps_tool_simple():
@@ -2675,19 +2621,16 @@ def test_vertex_ai_penalty_parameters_validation():
for model, should_support in test_cases:
# Test _supports_penalty_parameters method
- assert (
- v._supports_penalty_parameters(model) == should_support
- ), f"Model {model} penalty support should be {should_support}"
+ assert v._supports_penalty_parameters(model) == should_support, (
+ f"Model {model} penalty support should be {should_support}"
+ )
# Test get_supported_openai_params method
supported_params = v.get_supported_openai_params(model)
- has_penalty_params = (
- "frequency_penalty" in supported_params
- and "presence_penalty" in supported_params
+ has_penalty_params = "frequency_penalty" in supported_params and "presence_penalty" in supported_params
+ assert has_penalty_params == should_support, (
+ f"Model {model} should {'include' if should_support else 'exclude'} penalty params in supported list"
)
- assert (
- has_penalty_params == should_support
- ), f"Model {model} should {'include' if should_support else 'exclude'} penalty params in supported list"
# Test parameter mapping for unsupported model
model = "gemini-2.5-pro-preview-06-05"
@@ -2707,12 +2650,8 @@ def test_vertex_ai_penalty_parameters_validation():
)
# Penalty parameters should be filtered out for unsupported models
- assert (
- "frequency_penalty" not in result
- ), "frequency_penalty should be filtered out for unsupported model"
- assert (
- "presence_penalty" not in result
- ), "presence_penalty should be filtered out for unsupported model"
+ assert "frequency_penalty" not in result, "frequency_penalty should be filtered out for unsupported model"
+ assert "presence_penalty" not in result, "presence_penalty should be filtered out for unsupported model"
# Other parameters should still be included
assert "temperature" in result, "temperature should still be included"
@@ -2741,18 +2680,18 @@ def test_vertex_ai_gemini_3_penalty_parameters_unsupported():
for model in gemini_3_models:
# Test _supports_penalty_parameters method
- assert (
- v._supports_penalty_parameters(model) == False
- ), f"Gemini 3 model {model} should not support penalty parameters"
+ assert v._supports_penalty_parameters(model) == False, (
+ f"Gemini 3 model {model} should not support penalty parameters"
+ )
# Test get_supported_openai_params method
supported_params = v.get_supported_openai_params(model)
- assert (
- "frequency_penalty" not in supported_params
- ), f"frequency_penalty should not be in supported params for {model}"
- assert (
- "presence_penalty" not in supported_params
- ), f"presence_penalty should not be in supported params for {model}"
+ assert "frequency_penalty" not in supported_params, (
+ f"frequency_penalty should not be in supported params for {model}"
+ )
+ assert "presence_penalty" not in supported_params, (
+ f"presence_penalty should not be in supported params for {model}"
+ )
# Test parameter mapping - penalty params should be filtered out
non_default_params = {
@@ -2771,36 +2710,28 @@ def test_vertex_ai_gemini_3_penalty_parameters_unsupported():
)
# Penalty parameters should be filtered out for Gemini 3 models
- assert (
- "frequency_penalty" not in result
- ), f"frequency_penalty should be filtered out for Gemini 3 model {model}"
- assert (
- "presence_penalty" not in result
- ), f"presence_penalty should be filtered out for Gemini 3 model {model}"
+ assert "frequency_penalty" not in result, f"frequency_penalty should be filtered out for Gemini 3 model {model}"
+ assert "presence_penalty" not in result, f"presence_penalty should be filtered out for Gemini 3 model {model}"
# Other parameters should still be included
- assert (
- "temperature" in result
- ), f"temperature should still be included for Gemini 3 model {model}"
- assert (
- "max_output_tokens" in result
- ), f"max_output_tokens should still be included for Gemini 3 model {model}"
+ assert "temperature" in result, f"temperature should still be included for Gemini 3 model {model}"
+ assert "max_output_tokens" in result, f"max_output_tokens should still be included for Gemini 3 model {model}"
assert result["temperature"] == 0.7
assert result["max_output_tokens"] == 100
# Test that non-Gemini 3 models still support penalty parameters (if they're not in the unsupported list)
non_gemini_3_model = "gemini-2.5-pro"
- assert (
- v._supports_penalty_parameters(non_gemini_3_model) == True
- ), f"Non-Gemini 3 model {non_gemini_3_model} should support penalty parameters"
+ assert v._supports_penalty_parameters(non_gemini_3_model) == True, (
+ f"Non-Gemini 3 model {non_gemini_3_model} should support penalty parameters"
+ )
supported_params = v.get_supported_openai_params(non_gemini_3_model)
- assert (
- "frequency_penalty" in supported_params
- ), f"frequency_penalty should be in supported params for {non_gemini_3_model}"
- assert (
- "presence_penalty" in supported_params
- ), f"presence_penalty should be in supported params for {non_gemini_3_model}"
+ assert "frequency_penalty" in supported_params, (
+ f"frequency_penalty should be in supported params for {non_gemini_3_model}"
+ )
+ assert "presence_penalty" in supported_params, (
+ f"presence_penalty should be in supported params for {non_gemini_3_model}"
+ )
def test_vertex_ai_annotation_streaming_events():
@@ -2823,14 +2754,10 @@ def test_vertex_ai_annotation_streaming_events():
chunk_with_annotations = {
"candidates": [
{
- "content": {
- "parts": [{"text": "The weather in San Francisco today is clear."}]
- },
+ "content": {"parts": [{"text": "The weather in San Francisco today is clear."}]},
"groundingMetadata": {
"webSearchQueries": ["weather San Francisco today"],
- "searchEntryPoint": {
- "renderedContent": "Search results
"
- },
+ "searchEntryPoint": {"renderedContent": "Search results
"},
"groundingChunks": [
{
"web": {
@@ -2862,9 +2789,7 @@ def test_vertex_ai_annotation_streaming_events():
}
# Create iterator and parse chunk
- iterator = ModelResponseIterator(
- streaming_response=[], sync_stream=True, logging_obj=litellm_logging
- )
+ iterator = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=litellm_logging)
streaming_chunk = iterator.chunk_parser(chunk_with_annotations)
# Verify the chunk was parsed correctly
@@ -2986,12 +2911,8 @@ def test_vertex_ai_annotation_conversion():
}
# Convert grounding metadata to annotations
- content_text = (
- "The weather in San Francisco is currently 72°F and the time is 2:30 PM"
- )
- annotations = VertexGeminiConfig._convert_grounding_metadata_to_annotations(
- [grounding_metadata], content_text
- )
+ content_text = "The weather in San Francisco is currently 72°F and the time is 2:30 PM"
+ annotations = VertexGeminiConfig._convert_grounding_metadata_to_annotations([grounding_metadata], content_text)
# Verify annotations were created
assert len(annotations) == 3 # One for each grounding support
@@ -3033,9 +2954,7 @@ def test_vertex_ai_annotation_empty_grounding_metadata():
# Test with empty grounding metadata
empty_metadata = {}
- annotations = VertexGeminiConfig._convert_grounding_metadata_to_annotations(
- [empty_metadata], "test content"
- )
+ annotations = VertexGeminiConfig._convert_grounding_metadata_to_annotations([empty_metadata], "test content")
assert len(annotations) == 0
# Test with missing groundingSupports
@@ -3043,9 +2962,7 @@ def test_vertex_ai_annotation_empty_grounding_metadata():
"webSearchQueries": ["test query"],
"groundingChunks": [{"web": {"uri": "https://example.com", "title": "Test"}}],
}
- annotations = VertexGeminiConfig._convert_grounding_metadata_to_annotations(
- [metadata_no_supports], "test content"
- )
+ annotations = VertexGeminiConfig._convert_grounding_metadata_to_annotations([metadata_no_supports], "test content")
assert len(annotations) == 0
# Test with empty groundingSupports
@@ -3408,9 +3325,7 @@ def test_reasoning_effort_maps_to_thinking_level_gemini_3():
def test_gemini_37_38_flash_floor_minimal_thinking_level(
local_model_cost_map, model, reasoning_effort, include_thoughts
):
- result = VertexGeminiConfig._map_reasoning_effort_to_thinking_level(
- reasoning_effort, model
- )
+ result = VertexGeminiConfig._map_reasoning_effort_to_thinking_level(reasoning_effort, model)
assert result["thinkingLevel"] == "low"
assert result["includeThoughts"] is include_thoughts
@@ -3434,9 +3349,7 @@ def test_gemini_37_38_flash_floor_minimal_thinking_level(
def test_gemini_flash_minimal_thinking_support(
local_model_cost_map, model, reasoning_effort, expected_level, include_thoughts
):
- result = VertexGeminiConfig._map_reasoning_effort_to_thinking_level(
- reasoning_effort, model
- )
+ result = VertexGeminiConfig._map_reasoning_effort_to_thinking_level(reasoning_effort, model)
assert result["thinkingLevel"] == expected_level
assert result["includeThoughts"] is include_thoughts
@@ -3446,12 +3359,8 @@ def test_gemini_38_flash_feature_flag_uses_low_thinking_level(local_model_cost_m
monkeypatch.setattr(litellm, "enable_gemini_default_thinking_level_low", True)
thinking_param = {"type": "enabled", "budget_tokens": 1024}
- result_38 = VertexGeminiConfig._map_thinking_param(
- thinking_param, model="gemini-3.8-flash"
- )
- result_36 = VertexGeminiConfig._map_thinking_param(
- thinking_param, model="gemini-3.6-flash"
- )
+ result_38 = VertexGeminiConfig._map_thinking_param(thinking_param, model="gemini-3.8-flash")
+ result_36 = VertexGeminiConfig._map_thinking_param(thinking_param, model="gemini-3.6-flash")
assert result_38["thinkingLevel"] == "low"
assert result_36["thinkingLevel"] == "minimal"
@@ -3933,9 +3842,7 @@ def test_gemini_image_models_excluded_from_thinking():
)
# None of these should have thinkingConfig
- assert (
- "thinkingConfig" not in result
- ), f"Model {model} should not have thinkingConfig"
+ assert "thinkingConfig" not in result, f"Model {model} should not have thinkingConfig"
def test_partial_json_chunk_after_first_chunk():
@@ -3966,9 +3873,7 @@ def test_partial_json_chunk_after_first_chunk():
first_chunk = '{"candidates": [{"content": {"parts": [{"text": "Hello"}]}}]}'
result1 = iterator.handle_valid_json_chunk(first_chunk)
assert result1 is not None, "First complete chunk should parse OK"
- assert (
- iterator.sent_first_chunk is True
- ), "sent_first_chunk should be True after first chunk"
+ assert iterator.sent_first_chunk is True, "sent_first_chunk should be True after first chunk"
# Later chunk arrives PARTIAL (simulating network fragmentation)
partial_chunk = '{"candidates": [{"content":'
@@ -3976,9 +3881,7 @@ def test_partial_json_chunk_after_first_chunk():
# Should switch to accumulation mode instead of crashing
assert result2 is None, "Partial chunk should return None while accumulating"
- assert (
- iterator.chunk_type == "accumulated_json"
- ), "Should switch to accumulated_json mode"
+ assert iterator.chunk_type == "accumulated_json", "Should switch to accumulated_json mode"
def test_partial_json_chunk_on_first_chunk():
@@ -3998,9 +3901,7 @@ def test_partial_json_chunk_on_first_chunk():
result = iterator.handle_valid_json_chunk(partial)
assert result is None, "Partial first chunk should return None"
- assert (
- iterator.chunk_type == "accumulated_json"
- ), "Should switch to accumulated_json mode"
+ assert iterator.chunk_type == "accumulated_json", "Should switch to accumulated_json mode"
def test_accumulated_json_does_not_reparse_every_fragment():
@@ -4032,16 +3933,12 @@ def test_accumulated_json_does_not_reparse_every_fragment():
iterator.chunk_type = "accumulated_json"
text = "x" * 200_000 # no braces/brackets so only the final fragment closes
- blob = json.dumps(
- {"candidates": [{"content": {"role": "model", "parts": [{"text": text}]}}]}
- )
+ blob = json.dumps({"candidates": [{"content": {"role": "model", "parts": [{"text": text}]}}]})
fragments = [blob[i : i + 4096] for i in range(0, len(blob), 4096)]
assert len(fragments) > 10, "need a multi-fragment payload to exercise the bug"
parsed = None
- with patch.object(
- json.JSONDecoder, "raw_decode", autospec=True, side_effect=json.JSONDecoder.raw_decode
- ) as spy:
+ with patch.object(json.JSONDecoder, "raw_decode", autospec=True, side_effect=json.JSONDecoder.raw_decode) as spy:
for fragment in fragments:
out = iterator.handle_accumulated_json_chunk(chunk=fragment)
if out is not None:
@@ -4073,9 +3970,7 @@ def test_accumulated_json_partial_fragment_returns_none_without_parsing():
)
iterator.chunk_type = "accumulated_json"
- with patch.object(
- json.JSONDecoder, "raw_decode", autospec=True, side_effect=json.JSONDecoder.raw_decode
- ) as spy:
+ with patch.object(json.JSONDecoder, "raw_decode", autospec=True, side_effect=json.JSONDecoder.raw_decode) as spy:
result = iterator.handle_accumulated_json_chunk(
chunk='{"candidates": [{"content": {"parts": [{"text": "partial'
)
@@ -4174,17 +4069,11 @@ def test_vertex_ai_multiple_tool_types_separate_objects():
tool_types_in_first = [k for k in tools[0].keys()]
tool_types_in_second = [k for k in tools[1].keys()]
- assert (
- len(tool_types_in_first) == 1
- ), f"First Tool should have exactly 1 type, got {tool_types_in_first}"
- assert (
- len(tool_types_in_second) == 1
- ), f"Second Tool should have exactly 1 type, got {tool_types_in_second}"
+ assert len(tool_types_in_first) == 1, f"First Tool should have exactly 1 type, got {tool_types_in_first}"
+ assert len(tool_types_in_second) == 1, f"Second Tool should have exactly 1 type, got {tool_types_in_second}"
# Verify the correct tool types are present
- assert (
- "enterpriseWebSearch" in tools[0]
- ), "First Tool should contain enterpriseWebSearch"
+ assert "enterpriseWebSearch" in tools[0], "First Tool should contain enterpriseWebSearch"
assert "url_context" in tools[1], "Second Tool should contain url_context"
@@ -4646,9 +4535,7 @@ def test_vertex_ai_multiple_function_declarations_grouped():
)
# Should have only 1 Tool object (function declarations grouped)
- assert (
- len(tools) == 1
- ), f"Expected 1 Tool object for grouped functions, got {len(tools)}"
+ assert len(tools) == 1, f"Expected 1 Tool object for grouped functions, got {len(tools)}"
# Should contain function_declarations with 2 functions
assert "function_declarations" in tools[0]
@@ -5070,9 +4957,7 @@ def test_chunk_parser_handles_prompt_feedback_block():
logging_obj = Mock()
logging_obj.optional_params = {}
- streaming_obj = ModelResponseIterator(
- streaming_response=iter([]), sync_stream=True, logging_obj=logging_obj
- )
+ streaming_obj = ModelResponseIterator(streaming_response=iter([]), sync_stream=True, logging_obj=logging_obj)
# Act
result = streaming_obj.chunk_parser(blocked_chunk)
@@ -5080,9 +4965,9 @@ def test_chunk_parser_handles_prompt_feedback_block():
# Assert
assert result is not None, "Result should not be None"
assert len(result.choices) == 1, "Should have exactly one choice"
- assert (
- result.choices[0].finish_reason == "content_filter"
- ), f"finish_reason should be content_filter, got {result.choices[0].finish_reason}"
+ assert result.choices[0].finish_reason == "content_filter", (
+ f"finish_reason should be content_filter, got {result.choices[0].finish_reason}"
+ )
assert result.choices[0].delta.content is None, "content should be None"
@@ -5106,9 +4991,7 @@ def test_chunk_parser_handles_prompt_feedback_safety_block():
logging_obj = Mock()
logging_obj.optional_params = {}
- streaming_obj = ModelResponseIterator(
- streaming_response=iter([]), sync_stream=True, logging_obj=logging_obj
- )
+ streaming_obj = ModelResponseIterator(streaming_response=iter([]), sync_stream=True, logging_obj=logging_obj)
# Act
result = streaming_obj.chunk_parser(blocked_chunk)
@@ -5145,9 +5028,7 @@ def test_chunk_parser_handles_prompt_feedback_block_with_usage():
logging_obj = Mock()
logging_obj.optional_params = {}
- streaming_obj = ModelResponseIterator(
- streaming_response=iter([]), sync_stream=True, logging_obj=logging_obj
- )
+ streaming_obj = ModelResponseIterator(streaming_response=iter([]), sync_stream=True, logging_obj=logging_obj)
# Act
result = streaming_obj.chunk_parser(blocked_chunk)
@@ -5155,23 +5036,17 @@ def test_chunk_parser_handles_prompt_feedback_block_with_usage():
# Assert - 验证 content_filter 响应和 usage 都被正确处理
assert result is not None, "Result should not be None"
assert len(result.choices) == 1, "Should have exactly one choice"
- assert (
- result.choices[0].finish_reason == "content_filter"
- ), f"finish_reason should be content_filter, got {result.choices[0].finish_reason}"
+ assert result.choices[0].finish_reason == "content_filter", (
+ f"finish_reason should be content_filter, got {result.choices[0].finish_reason}"
+ )
assert result.choices[0].delta.content is None, "content should be None"
# 验证 usage 信息被正确提取
assert hasattr(result, "usage"), "result should have usage attribute"
assert result.usage is not None, "usage should not be None"
- assert (
- result.usage.prompt_tokens == 8175
- ), f"prompt_tokens should be 8175, got {result.usage.prompt_tokens}"
- assert (
- result.usage.completion_tokens == 0
- ), f"completion_tokens should be 0, got {result.usage.completion_tokens}"
- assert (
- result.usage.total_tokens == 8175
- ), f"total_tokens should be 8175, got {result.usage.total_tokens}"
+ assert result.usage.prompt_tokens == 8175, f"prompt_tokens should be 8175, got {result.usage.prompt_tokens}"
+ assert result.usage.completion_tokens == 0, f"completion_tokens should be 0, got {result.usage.completion_tokens}"
+ assert result.usage.total_tokens == 8175, f"total_tokens should be 8175, got {result.usage.total_tokens}"
def test_vertex_ai_traffic_type_preserved_in_hidden_params_streaming():
@@ -5190,14 +5065,10 @@ def test_vertex_ai_traffic_type_preserved_in_hidden_params_streaming():
},
}
- iterator = ModelResponseIterator(
- streaming_response=[], sync_stream=True, logging_obj=MagicMock()
- )
+ iterator = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=MagicMock())
result = iterator.chunk_parser(chunk)
- assert (
- result._hidden_params["provider_specific_fields"]["traffic_type"] == "ON_DEMAND"
- )
+ assert result._hidden_params["provider_specific_fields"]["traffic_type"] == "ON_DEMAND"
def test_vertex_ai_traffic_type_preserved_in_hidden_params_non_streaming():
@@ -5236,10 +5107,7 @@ def test_vertex_ai_traffic_type_preserved_in_hidden_params_non_streaming():
encoding=None,
)
- assert (
- result._hidden_params["provider_specific_fields"]["traffic_type"]
- == "PROVISIONED_THROUGHPUT"
- )
+ assert result._hidden_params["provider_specific_fields"]["traffic_type"] == "PROVISIONED_THROUGHPUT"
def test_vertex_ai_service_tier_streaming():
@@ -5323,9 +5191,7 @@ def test_vertex_ai_traffic_type_surfaced_in_responses_api():
from litellm.types.utils import Choices, Message
model_response = ModelResponse()
- model_response._hidden_params["provider_specific_fields"] = {
- "traffic_type": "ON_DEMAND"
- }
+ model_response._hidden_params["provider_specific_fields"] = {"traffic_type": "ON_DEMAND"}
model_response.choices = [
Choices(
message=Message(content="Hello", role="assistant"),
@@ -5334,15 +5200,15 @@ def test_vertex_ai_traffic_type_surfaced_in_responses_api():
)
]
- responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response(
- request_input="test",
- chat_completion_response=model_response,
- responses_api_request={},
+ responses_api_response = (
+ LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response(
+ request_input="test",
+ chat_completion_response=model_response,
+ responses_api_request={},
+ )
)
- assert (
- responses_api_response.provider_specific_fields["traffic_type"] == "ON_DEMAND"
- )
+ assert responses_api_response.provider_specific_fields["traffic_type"] == "ON_DEMAND"
def test_vertex_ai_web_search_options_parameter():
@@ -5416,12 +5282,8 @@ def test_vertex_ai_web_search_options_in_map_openai_params():
assert "tools" in optional_params, "tools should be added to optional_params"
assert len(optional_params["tools"]) == 1, "Should have exactly one tool"
assert "googleSearch" in optional_params["tools"][0], "Tool should be googleSearch"
- assert (
- optional_params["tools"][0]["googleSearch"] == {}
- ), "googleSearch should be empty config"
- assert (
- "web_search_options" not in optional_params
- ), "web_search_options should be removed after transformation"
+ assert optional_params["tools"][0]["googleSearch"] == {}, "googleSearch should be empty config"
+ assert "web_search_options" not in optional_params, "web_search_options should be removed after transformation"
def test_vertex_ai_service_tier_in_map_openai_params():
@@ -5509,24 +5371,16 @@ def test_vertex_ai_usage_metadata_with_video_tokens_in_prompt():
# Verify prompt token details include video tokens
assert result.prompt_tokens_details is not None
- assert (
- result.prompt_tokens_details.video_tokens == 10240
- ), "Prompt video tokens should be 10240"
- assert (
- result.prompt_tokens_details.text_tokens == 9
- ), "Prompt text tokens should be 9"
- assert (
- result.prompt_tokens_details.audio_tokens == 200
- ), "Prompt audio tokens should be 200"
+ assert result.prompt_tokens_details.video_tokens == 10240, "Prompt video tokens should be 10240"
+ assert result.prompt_tokens_details.text_tokens == 9, "Prompt text tokens should be 9"
+ assert result.prompt_tokens_details.audio_tokens == 200, "Prompt audio tokens should be 200"
# Verify completion token details
assert result.completion_tokens_details is not None
- assert (
- result.completion_tokens_details.text_tokens == 79
- ), "Completion text tokens should be 79"
- assert (
- result.completion_tokens_details.video_tokens is None
- ), "Completion video tokens should be None (text-only response)"
+ assert result.completion_tokens_details.text_tokens == 79, "Completion text tokens should be 79"
+ assert result.completion_tokens_details.video_tokens is None, (
+ "Completion video tokens should be None (text-only response)"
+ )
def test_vertex_ai_usage_metadata_with_video_tokens_in_candidates():
@@ -5667,9 +5521,9 @@ def test_vertex_ai_usage_metadata_with_document_tokens_in_prompt():
# DOCUMENT tokens should be included in text_tokens: 8 (TEXT) + 774 (DOCUMENT) = 782
assert result.prompt_tokens_details is not None
- assert (
- result.prompt_tokens_details.text_tokens == 782
- ), "DOCUMENT modality tokens should be added to text_tokens (8 TEXT + 774 DOCUMENT = 782)"
+ assert result.prompt_tokens_details.text_tokens == 782, (
+ "DOCUMENT modality tokens should be added to text_tokens (8 TEXT + 774 DOCUMENT = 782)"
+ )
# Verify completion token details
assert result.completion_tokens_details is not None
@@ -6166,9 +6020,7 @@ def test_mid_stream_429_error_raises_during_iteration():
{
"content": {
"role": "model",
- "parts": [
- {"text": "Let me think about this...", "thought": True}
- ],
+ "parts": [{"text": "Let me think about this...", "thought": True}],
},
"index": 0,
}
@@ -6188,9 +6040,7 @@ def test_mid_stream_429_error_raises_during_iteration():
{
"content": {
"role": "model",
- "parts": [
- {"text": "I'll generate the image now.", "thought": True}
- ],
+ "parts": [{"text": "I'll generate the image now.", "thought": True}],
},
"index": 0,
}
@@ -6227,6 +6077,7 @@ def test_mid_stream_429_error_raises_during_iteration():
# Iterate the stream: first chunks should succeed, then 429 error should be raised
results = []
+
def _drain():
for chunk in streaming_obj:
if chunk is not None:
@@ -6236,9 +6087,7 @@ def test_mid_stream_429_error_raises_during_iteration():
_drain()
# Verify: received normal chunks before the error
- assert (
- len(results) >= 1
- ), "Should have received at least 1 normal chunk before the error"
+ assert len(results) >= 1, "Should have received at least 1 normal chunk before the error"
# Verify: 429 error is properly raised
assert exc_info.value.status_code == 429
@@ -6427,9 +6276,7 @@ def _accumulating_gemini_iterator():
ModelResponseIterator,
)
- iterator = ModelResponseIterator(
- streaming_response=[], sync_stream=True, logging_obj=MagicMock()
- )
+ iterator = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=MagicMock())
iterator.chunk_type = "accumulated_json"
return iterator
@@ -6978,6 +6825,7 @@ def _vcr_outcome_gate(request, vcr):
yield
record_vcr_outcome(request, vcr)
+
@pytest.fixture(scope="function")
def isolate_litellm_state():
"""
@@ -7030,6 +6878,7 @@ def isolate_litellm_state():
setattr(litellm, attr, original_value)
_invalidate_model_cost_lowercase_map()
+
_SCALAR_DEFAULTS = {
"num_retries": getattr(litellm, "num_retries", None),
"num_retries_per_request": getattr(litellm, "num_retries_per_request", None),
@@ -7050,6 +6899,7 @@ _SCALAR_DEFAULTS = {
"api_key": getattr(litellm, "api_key", None),
}
+
@pytest.fixture(scope="module")
def setup_and_teardown():
"""
@@ -7072,6 +6922,7 @@ def setup_and_teardown():
litellm.in_memory_llm_clients_cache.flush_cache()
yield
+
@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown")
def test_thought_true_creates_thinking_block():
"""
@@ -7086,6 +6937,7 @@ def test_thought_true_creates_thinking_block():
assert block["thinking"] == "Some thinking"
assert block["signature"] == "sig-1"
+
@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown")
def test_thought_true_with_empty_text_creates_block():
"""
@@ -7097,6 +6949,7 @@ def test_thought_true_with_empty_text_creates_block():
assert len(thinking_blocks) == 1
assert thinking_blocks[0]["thinking"] == ""
+
@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown")
def test_thought_signature_without_thought_does_not_create_block():
"""
@@ -7109,6 +6962,7 @@ def test_thought_signature_without_thought_does_not_create_block():
thinking_blocks = config._extract_thinking_blocks_from_parts(parts)
assert thinking_blocks == []
+
@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown")
def test_extract_thought_signatures_from_regular_parts():
"""
@@ -7128,6 +6982,7 @@ def test_extract_thought_signatures_from_regular_parts():
assert len(signatures) == 1
assert signatures[0] == "sig-regular-123"
+
@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown")
def test_extract_multiple_thought_signatures():
"""
@@ -7146,6 +7001,7 @@ def test_extract_multiple_thought_signatures():
assert signatures[0] == "sig-1"
assert signatures[1] == "sig-2"
+
@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown")
def test_round_trip_thought_signature_in_conversation():
"""
@@ -7180,6 +7036,7 @@ def test_round_trip_thought_signature_in_conversation():
assert "thoughtSignature" in text_part
assert text_part["thoughtSignature"] == "sig-round-trip-abc"
+
@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown")
def test_round_trip_without_thought_signature_still_works():
"""
@@ -7418,18 +7275,10 @@ def test_gemini_flash_image_preview_models(model_name: str):
mock_response = ImageResponse()
mock_response.data = [ImageObject(b64_json="test_base64_data", url=None)]
- with patch(
- "litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post"
- ) as mock_post:
+ with patch("litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post") as mock_post:
mock_http_response = MagicMock()
mock_http_response.json.return_value = {
- "candidates": [
- {
- "content": {
- "parts": [{"inlineData": {"data": "test_base64_image_data"}}]
- }
- }
- ]
+ "candidates": [{"content": {"parts": [{"inlineData": {"data": "test_base64_image_data"}}]}}]
}
mock_http_response.status_code = 200
mock_post.return_value = mock_http_response
@@ -7447,9 +7296,7 @@ def test_gemini_flash_image_preview_models(model_name: str):
mock_post.assert_called_once()
call_args = mock_post.call_args
- called_url = (
- call_args[0][0] if call_args[0] else call_args.kwargs.get("url", "")
- )
+ called_url = call_args[0][0] if call_args[0] else call_args.kwargs.get("url", "")
assert ":generateContent" in called_url
assert model_name.split("/", 1)[1] in called_url
@@ -7481,23 +7328,13 @@ def test_gemini_flash_image_preview_models(model_name: str):
),
],
)
-def test_gemini_image_generation_forwards_image_config(
- model: str, kwargs: dict, expected_image_config: dict
-):
+def test_gemini_image_generation_forwards_image_config(model: str, kwargs: dict, expected_image_config: dict):
from unittest.mock import patch, MagicMock
- with patch(
- "litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post"
- ) as mock_post:
+ with patch("litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post") as mock_post:
mock_http_response = MagicMock()
mock_http_response.json.return_value = {
- "candidates": [
- {
- "content": {
- "parts": [{"inlineData": {"data": "test_base64_image_data"}}]
- }
- }
- ]
+ "candidates": [{"content": {"parts": [{"inlineData": {"data": "test_base64_image_data"}}]}}]
}
mock_http_response.status_code = 200
mock_post.return_value = mock_http_response
@@ -7558,9 +7395,7 @@ def test_gemini_image_generation_openai_size_maps_to_google_table(
map_openai_size_to_gemini_image_config,
)
- assert map_openai_size_to_gemini_image_config(
- size, "gemini-3-pro-image-preview"
- ) == {
+ assert map_openai_size_to_gemini_image_config(size, "gemini-3-pro-image-preview") == {
"aspectRatio": expected_aspect_ratio,
"imageSize": expected_image_size,
}
@@ -7584,9 +7419,7 @@ def test_gemini_image_generation_openai_size_snaps_to_nearest_option(
map_openai_size_to_gemini_image_config,
)
- assert map_openai_size_to_gemini_image_config(
- size, "gemini-3-pro-image-preview"
- ) == {
+ assert map_openai_size_to_gemini_image_config(size, "gemini-3-pro-image-preview") == {
"aspectRatio": expected_aspect_ratio,
"imageSize": expected_image_size,
}
@@ -7598,9 +7431,7 @@ def test_gemini_image_generation_openai_size_auto_uses_google_defaults(size: str
map_openai_size_to_gemini_image_config,
)
- assert map_openai_size_to_gemini_image_config(
- size, "gemini-3-pro-image-preview"
- ) is None
+ assert map_openai_size_to_gemini_image_config(size, "gemini-3-pro-image-preview") is None
def test_gemini_imagen_models_use_predict_endpoint():
@@ -7610,13 +7441,9 @@ def test_gemini_imagen_models_use_predict_endpoint():
from unittest.mock import patch, MagicMock
from litellm.types.utils import ImageResponse, ImageObject
- with patch(
- "litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post"
- ) as mock_post:
+ with patch("litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post") as mock_post:
mock_http_response = MagicMock()
- mock_http_response.json.return_value = {
- "predictions": [{"bytesBase64Encoded": "test_base64_image_data"}]
- }
+ mock_http_response.json.return_value = {"predictions": [{"bytesBase64Encoded": "test_base64_image_data"}]}
mock_http_response.status_code = 200
mock_post.return_value = mock_http_response
@@ -7632,9 +7459,7 @@ def test_gemini_imagen_models_use_predict_endpoint():
mock_post.assert_called_once()
call_args = mock_post.call_args
- called_url = (
- call_args[0][0] if call_args[0] else call_args.kwargs.get("url", "")
- )
+ called_url = call_args[0][0] if call_args[0] else call_args.kwargs.get("url", "")
assert ":predict" in called_url
assert "imagen-3.0-generate-001" in called_url
@@ -7797,11 +7622,7 @@ async def test_claude_tool_use_with_gemini():
if "data: " in chunk_str:
try:
- data_line = [
- line
- for line in chunk_str.split("\n")
- if line.startswith("data: ")
- ][0]
+ data_line = [line for line in chunk_str.split("\n") if line.startswith("data: ")][0]
json_str = data_line[6:]
chunk_data = json.loads(json_str)
@@ -7814,24 +7635,15 @@ async def test_claude_tool_use_with_gemini():
if (
chunk_data.get("type") == "message_delta"
- and chunk_data.get("delta", {}).get("stop_reason")
- is not None
+ and chunk_data.get("delta", {}).get("stop_reason") is not None
and "usage" in chunk_data
):
has_usage_in_message_delta = True
usage = chunk_data["usage"]
- assert (
- "input_tokens" in usage
- ), "input_tokens should be present in usage"
- assert (
- "output_tokens" in usage
- ), "output_tokens should be present in usage"
- assert isinstance(
- usage["input_tokens"], int
- ), "input_tokens should be an integer"
- assert isinstance(
- usage["output_tokens"], int
- ), "output_tokens should be an integer"
+ assert "input_tokens" in usage, "input_tokens should be present in usage"
+ assert "output_tokens" in usage, "output_tokens should be present in usage"
+ assert isinstance(usage["input_tokens"], int), "input_tokens should be an integer"
+ assert isinstance(usage["output_tokens"], int), "output_tokens should be an integer"
print(f"Found usage in message_delta: {usage}")
except (json.JSONDecodeError, IndexError) as e:
@@ -7846,9 +7658,7 @@ async def test_claude_tool_use_with_gemini():
assert is_content_block_tool_use, "content_block_tool_use should be present"
assert is_partial_json, "partial_json should be present"
- assert (
- has_usage_in_message_delta
- ), "Usage should be present in message_delta with stop_reason"
+ assert has_usage_in_message_delta, "Usage should be present in message_delta with stop_reason"
assert is_content_block_stop, "is_content_block_stop should be present"
@@ -7877,28 +7687,22 @@ def test_gemini_reasoning_effort_minimal():
)
request_body = raw_request["raw_request_body"]
- assert (
- "generationConfig" in request_body
- ), f"Model {model} should have generationConfig"
+ assert "generationConfig" in request_body, f"Model {model} should have generationConfig"
generation_config = request_body["generationConfig"]
- assert (
- "thinkingConfig" in generation_config
- ), f"Model {model} should have thinkingConfig"
+ assert "thinkingConfig" in generation_config, f"Model {model} should have thinkingConfig"
thinking_config = generation_config["thinkingConfig"]
- assert (
- "thinkingBudget" in thinking_config
- ), f"Model {model} should have thinkingBudget"
+ assert "thinkingBudget" in thinking_config, f"Model {model} should have thinkingBudget"
actual_budget = thinking_config["thinkingBudget"]
- assert (
- actual_budget == expected_min_budget
- ), f"Model {model} should map 'minimal' to {expected_min_budget} tokens, got {actual_budget}"
+ assert actual_budget == expected_min_budget, (
+ f"Model {model} should map 'minimal' to {expected_min_budget} tokens, got {actual_budget}"
+ )
- assert thinking_config.get(
- "includeThoughts", True
- ), f"Model {model} should have includeThoughts=True for minimal reasoning effort"
+ assert thinking_config.get("includeThoughts", True), (
+ f"Model {model} should have includeThoughts=True for minimal reasoning effort"
+ )
try:
raw_request = return_raw_request(
@@ -7913,9 +7717,7 @@ def test_gemini_reasoning_effort_minimal():
request_body = raw_request["raw_request_body"]
generation_config = request_body["generationConfig"]
thinking_config = generation_config["thinkingConfig"]
- assert (
- thinking_config["thinkingBudget"] == 128
- ), "Unknown model should use generic fallback of 128 tokens"
+ assert thinking_config["thinkingBudget"] == 128, "Unknown model should use generic fallback of 128 tokens"
except Exception as e:
print(f"Note: Unknown model test skipped due to: {e}")
pass
@@ -7938,9 +7740,7 @@ def test_gemini_exception_message_format():
mock_response.text = "Invalid API key provided"
mock_response.headers = {}
- mock_exception = httpx.HTTPStatusError(
- message="Bad Request", request=Mock(), response=mock_response
- )
+ mock_exception = httpx.HTTPStatusError(message="Bad Request", request=Mock(), response=mock_response)
mock_exception.response = mock_response
mock_exception.status_code = 400
@@ -7960,9 +7760,9 @@ def test_gemini_exception_message_format():
f"Expected 'GeminiException' in error message, got: {error_message}. "
f"This test should fail before the fix is implemented."
)
- assert (
- "VertexAIException" not in error_message
- ), f"Should not contain 'VertexAIException' in error message, got: {error_message}"
+ assert "VertexAIException" not in error_message, (
+ f"Should not contain 'VertexAIException' in error message, got: {error_message}"
+ )
def test_reasoning_effort_none_mapping():
@@ -8016,17 +7816,11 @@ def test_gemini_function_args_preserve_unicode():
parsed_args = json.loads(arguments_str)
assert parsed_args["message"] == "やあ", "Japanese characters should be preserved"
- assert (
- parsed_args["recipient"] == "たけし"
- ), "Japanese characters should be preserved"
+ assert parsed_args["recipient"] == "たけし", "Japanese characters should be preserved"
assert "\\u" not in arguments_str, "Should not contain Unicode escape sequences"
- assert (
- "やあ" in arguments_str
- ), "Original Japanese characters should be in the string"
- assert (
- "たけし" in arguments_str
- ), "Original Japanese characters should be in the string"
+ assert "やあ" in arguments_str, "Original Japanese characters should be in the string"
+ assert "たけし" in arguments_str, "Original Japanese characters should be in the string"
parts_spanish = [
{
@@ -8079,12 +7873,8 @@ def test_anthropic_thinking_param_to_gemini_3_provider_defaults():
model="gemini-3-flash",
)
- assert (
- "thinkingLevel" not in result
- ), "Should not force thinkingLevel for Gemini 3"
- assert (
- "thinkingBudget" not in result
- ), "Should NOT have thinkingBudget for Gemini 3"
+ assert "thinkingLevel" not in result, "Should not force thinkingLevel for Gemini 3"
+ assert "thinkingBudget" not in result, "Should NOT have thinkingBudget for Gemini 3"
assert result["includeThoughts"] is True
thinking_param_disabled: AnthropicThinkingParam = {
@@ -8098,10 +7888,7 @@ def test_anthropic_thinking_param_to_gemini_3_provider_defaults():
)
assert result_disabled.get("includeThoughts") is False
- assert (
- "thinkingLevel" not in result_disabled
- or result_disabled.get("thinkingLevel") is None
- )
+ assert "thinkingLevel" not in result_disabled or result_disabled.get("thinkingLevel") is None
thinking_param_zero: AnthropicThinkingParam = {
"type": "enabled",
@@ -8114,10 +7901,7 @@ def test_anthropic_thinking_param_to_gemini_3_provider_defaults():
)
assert result_zero["includeThoughts"] is False
- assert (
- "thinkingLevel" not in result_zero
- or result_zero.get("thinkingLevel") is None
- )
+ assert "thinkingLevel" not in result_zero or result_zero.get("thinkingLevel") is None
result_gemini3flashpreview = VertexGeminiConfig._map_thinking_param(
thinking_param=thinking_param,
@@ -8203,9 +7987,7 @@ def test_anthropic_thinking_param_to_gemini_2_thinkingBudget():
)
assert "thinkingBudget" in result_gemini2, "Should have thinkingBudget for Gemini 2"
- assert (
- "thinkingLevel" not in result_gemini2
- ), "Should NOT have thinkingLevel for Gemini 2"
+ assert "thinkingLevel" not in result_gemini2, "Should NOT have thinkingLevel for Gemini 2"
assert result_gemini2["includeThoughts"] is True
assert result_gemini2["thinkingBudget"] == 10000
@@ -8241,12 +8023,8 @@ def test_anthropic_thinking_param_via_map_openai_params():
assert "thinkingConfig" in result, "Should have thinkingConfig in optional_params"
thinking_config = result["thinkingConfig"]
- assert (
- "thinkingLevel" not in thinking_config
- ), "Should not force thinkingLevel for Gemini 3 by default"
- assert (
- "thinkingBudget" not in thinking_config
- ), "Should NOT have thinkingBudget for Gemini 3"
+ assert "thinkingLevel" not in thinking_config, "Should not force thinkingLevel for Gemini 3 by default"
+ assert "thinkingBudget" not in thinking_config, "Should NOT have thinkingBudget for Gemini 3"
assert thinking_config["includeThoughts"] is True
optional_params_2 = {}
@@ -8259,12 +8037,8 @@ def test_anthropic_thinking_param_via_map_openai_params():
assert "thinkingConfig" in result_2, "Should have thinkingConfig in optional_params"
thinking_config_2 = result_2["thinkingConfig"]
- assert (
- "thinkingBudget" in thinking_config_2
- ), "Should have thinkingBudget for Gemini 2"
- assert (
- "thinkingLevel" not in thinking_config_2
- ), "Should NOT have thinkingLevel for Gemini 2"
+ assert "thinkingBudget" in thinking_config_2, "Should have thinkingBudget for Gemini 2"
+ assert "thinkingLevel" not in thinking_config_2, "Should NOT have thinkingLevel for Gemini 2"
assert thinking_config_2["includeThoughts"] is True
assert thinking_config_2["thinkingBudget"] == 10000
@@ -8284,9 +8058,9 @@ def test_gemini_31_flash_lite_reasoning_effort_minimal():
reasoning_effort="minimal",
model="gemini-3.1-flash-lite-preview",
)
- assert (
- result["thinkingLevel"] == "minimal"
- ), f"Expected thinkingLevel='minimal' for gemini-3.1-flash-lite-preview, got '{result['thinkingLevel']}'"
+ assert result["thinkingLevel"] == "minimal", (
+ f"Expected thinkingLevel='minimal' for gemini-3.1-flash-lite-preview, got '{result['thinkingLevel']}'"
+ )
assert result["includeThoughts"] is True
from litellm.utils import return_raw_request
@@ -8302,12 +8076,12 @@ def test_gemini_31_flash_lite_reasoning_effort_minimal():
)
generation_config = raw_request["raw_request_body"]["generationConfig"]
thinking_config = generation_config["thinkingConfig"]
- assert (
- thinking_config.get("thinkingLevel") == "minimal"
- ), f"Expected thinkingLevel='minimal' via full flow, got {thinking_config}"
- assert (
- "thinkingBudget" not in thinking_config
- ), "gemini-3.1-flash-lite-preview should use thinkingLevel, not thinkingBudget"
+ assert thinking_config.get("thinkingLevel") == "minimal", (
+ f"Expected thinkingLevel='minimal' via full flow, got {thinking_config}"
+ )
+ assert "thinkingBudget" not in thinking_config, (
+ "gemini-3.1-flash-lite-preview should use thinkingLevel, not thinkingBudget"
+ )
@pytest.mark.usefixtures("fake_provider_credentials")
@@ -8855,10 +8629,7 @@ async def test_gemini_context_caching_disabled_flag(sync_mode):
"content": [
{
"type": "text",
- "text": "Here is the full text of a complex legal agreement {}".format(
- uuid.uuid4()
- )
- * 4000,
+ "text": "Here is the full text of a complex legal agreement {}".format(uuid.uuid4()) * 4000,
"cache_control": {"type": "ephemeral"},
}
],
@@ -8894,9 +8665,7 @@ async def test_gemini_context_caching_disabled_flag(sync_mode):
else:
client = AsyncHTTPHandler(concurrent_limit=1)
- with patch.object(
- client, "post", side_effect=mock_gemini_request
- ) as mock_client:
+ with patch.object(client, "post", side_effect=mock_gemini_request) as mock_client:
try:
if sync_mode:
response = litellm.completion(
@@ -8919,26 +8688,22 @@ async def test_gemini_context_caching_disabled_flag(sync_mode):
print(e)
# When caching is disabled, should only make 1 call (no separate cache creation call)
- assert (
- mock_client.call_count == 1
- ), f"Expected 1 call when caching is disabled, got {mock_client.call_count}"
+ assert mock_client.call_count == 1, (
+ f"Expected 1 call when caching is disabled, got {mock_client.call_count}"
+ )
first_call_args = mock_client.call_args_list[0].kwargs
first_call_positional_args = mock_client.call_args_list[0].args
print(f"first_call_args with caching disabled: {first_call_args}")
- print(
- f"first_call_positional_args with caching disabled: {first_call_positional_args}"
- )
+ print(f"first_call_positional_args with caching disabled: {first_call_positional_args}")
# Assert that cachedContents is NOT in the URL when caching is disabled
url = first_call_args.get(
"url",
first_call_positional_args[0] if first_call_positional_args else "",
)
- assert (
- "cachedContents" not in url
- ), "cachedContents should not be in URL when caching is disabled"
+ assert "cachedContents" not in url, "cachedContents should not be in URL when caching is disabled"
finally:
# Restore original flag value
@@ -9031,9 +8796,7 @@ def test_gemini_function_call_parameter_in_messages():
"contents": [
{
"role": "user",
- "parts": [
- {"text": "search for weather in boston (use `search`)"}
- ],
+ "parts": [{"text": "search for weather in boston (use `search`)"}],
},
{
"role": "model",
@@ -9052,17 +8815,13 @@ def test_gemini_function_call_parameter_in_messages():
{
"function_response": {
"name": "search",
- "response": {
- "content": "The current weather in Boston is 22°F."
- },
+ "response": {"content": "The current weather in Boston is 22°F."},
}
}
],
},
],
- "system_instruction": {
- "parts": [{"text": "Use search for most queries."}]
- },
+ "system_instruction": {"parts": [{"text": "Use search for most queries."}]},
"tools": [
{
"function_declarations": [
@@ -9137,9 +8896,7 @@ def test_gemini_function_call_parameter_in_messages_2():
{
"function_response": {
"name": "search",
- "response": {
- "content": "The weather in Boston is 100 degrees."
- },
+ "response": {"content": "The weather in Boston is 100 degrees."},
}
}
],
@@ -9517,9 +9274,7 @@ def test_vertex_ai_cache_creation_kwarg_cannot_change_usage_or_cost() -> None:
mock_response.status_code = 200
mock_response.headers = {"Content-Type": "application/json"}
mock_response.json.return_value = {
- "candidates": [
- {"content": {"parts": [{"text": "hi"}], "role": "model"}, "finishReason": "STOP", "index": 0}
- ],
+ "candidates": [{"content": {"parts": [{"text": "hi"}], "role": "model"}, "finishReason": "STOP", "index": 0}],
"usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 5, "totalTokenCount": 15},
}
client = HTTPHandler()
@@ -9585,14 +9340,51 @@ def load_vertex_ai_credentials():
# Gemini prebuilt TTS voices per https://ai.google.dev/gemini-api/docs/speech-generation (page dated 2026-10-01, read 2026-10-07)
GEMINI_TTS_PREBUILT_VOICES: Final[frozenset[str]] = frozenset(
{
- "Zephyr", "Puck", "Charon", "Kore", "Fenrir", "Leda", "Orus", "Aoede", "Callirrhoe", "Autonoe",
- "Enceladus", "Iapetus", "Umbriel", "Algieba", "Despina", "Erinome", "Algenib", "Rasalgethi",
- "Laomedeia", "Achernar", "Alnilam", "Schedar", "Gacrux", "Pulcherrima", "Achird", "Zubenelgenubi",
- "Vindemiatrix", "Sadachbia", "Sadaltager", "Sulafat",
+ "Zephyr",
+ "Puck",
+ "Charon",
+ "Kore",
+ "Fenrir",
+ "Leda",
+ "Orus",
+ "Aoede",
+ "Callirrhoe",
+ "Autonoe",
+ "Enceladus",
+ "Iapetus",
+ "Umbriel",
+ "Algieba",
+ "Despina",
+ "Erinome",
+ "Algenib",
+ "Rasalgethi",
+ "Laomedeia",
+ "Achernar",
+ "Alnilam",
+ "Schedar",
+ "Gacrux",
+ "Pulcherrima",
+ "Achird",
+ "Zubenelgenubi",
+ "Vindemiatrix",
+ "Sadachbia",
+ "Sadaltager",
+ "Sulafat",
}
)
OPENAI_TTS_VOICES: Final[tuple[str, ...]] = (
- "alloy", "ash", "ballad", "cedar", "coral", "echo", "fable", "marin", "onyx", "sage", "shimmer", "verse",
+ "alloy",
+ "ash",
+ "ballad",
+ "cedar",
+ "coral",
+ "echo",
+ "fable",
+ "marin",
+ "onyx",
+ "sage",
+ "shimmer",
+ "verse",
)
@@ -9678,3 +9470,289 @@ def test_streaming_usage_reads_cache_creation_tokens_recorded_on_logging_obj():
assert chunk.usage.prompt_tokens == 1536 + 1280
assert chunk.usage.cache_creation_input_tokens == 1280
+
+
+THOUGHT_SUMMARY_REQUESTS: Final = (
+ pytest.param("gemini-2.5-flash", {"reasoning_effort": "minimal"}, id="budget-minimal"),
+ pytest.param("gemini-2.5-flash", {"reasoning_effort": "low"}, id="budget-low"),
+ pytest.param("gemini-2.5-flash", {"reasoning_effort": "medium"}, id="budget-medium"),
+ pytest.param("gemini-2.5-flash", {"reasoning_effort": "high"}, id="budget-high"),
+ pytest.param("gemini-3.6-flash", {"reasoning_effort": "minimal"}, id="level-minimal"),
+ pytest.param("gemini-3.8-flash", {"reasoning_effort": "minimal"}, id="level-minimal-mapped-to-low"),
+ pytest.param("gemini-3.8-flash", {"reasoning_effort": "low"}, id="level-low"),
+ pytest.param("gemini-3.8-flash", {"reasoning_effort": "medium"}, id="level-medium"),
+ pytest.param("gemini-3.8-flash", {"reasoning_effort": "high"}, id="level-high"),
+ pytest.param("gemini-2.5-flash", {"thinking": {"type": "enabled", "budget_tokens": 2048}}, id="budget-thinking"),
+ pytest.param("gemini-3.8-flash", {"thinking": {"type": "enabled", "budget_tokens": 2048}}, id="level-thinking"),
+)
+
+
+async def _sent_thinking_config(model: str, thinking_params: dict[str, object]) -> dict[str, object]:
+ captured: Final[dict[str, dict]] = {}
+
+ def upstream(request: httpx.Request) -> httpx.Response:
+ captured["body"] = json.loads(request.content)
+ return httpx.Response(
+ 200,
+ json={
+ "candidates": [{"content": {"parts": [{"text": "hi"}], "role": "model"}, "finishReason": "STOP"}],
+ "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2},
+ },
+ request=request,
+ )
+
+ client: Final = AsyncHTTPHandler()
+ client.client = httpx.AsyncClient(transport=httpx.MockTransport(upstream))
+ await litellm.acompletion(
+ model=f"gemini/{model}",
+ messages=[{"role": "user", "content": "hi"}],
+ api_key="fake-gemini-key",
+ client=client,
+ **thinking_params,
+ )
+ return captured["body"]["generationConfig"]["thinkingConfig"]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(("model", "thinking_params"), THOUGHT_SUMMARY_REQUESTS)
+async def test_gemini_thought_summaries_off_keeps_thinking_config_otherwise_unchanged(
+ model: str,
+ thinking_params: dict[str, object],
+ local_model_cost_map: None,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ summaries_on: Final = await _sent_thinking_config(model, thinking_params)
+ monkeypatch.setattr(litellm, "gemini_include_thoughts_default", False)
+ summaries_off: Final = await _sent_thinking_config(model, thinking_params)
+
+ assert summaries_on["includeThoughts"] is True
+ assert summaries_off == {**summaries_on, "includeThoughts": False}
+
+
+PER_REQUEST_THOUGHT_SUMMARY_ASKS: Final = (
+ pytest.param(
+ "gemini-3.8-flash",
+ {"thinking": {"type": "enabled", "budget_tokens": 2048, "display": "summarized"}},
+ False,
+ True,
+ id="level-display-summarized-beats-off",
+ ),
+ pytest.param(
+ "gemini-3.8-flash",
+ {"thinking": {"type": "enabled", "budget_tokens": 2048, "display": "omitted"}},
+ True,
+ True,
+ id="level-display-omitted-leaves-the-default-on",
+ ),
+ pytest.param(
+ "gemini-2.5-flash",
+ {"thinking": {"type": "enabled", "budget_tokens": 2048, "display": "summarized"}},
+ False,
+ True,
+ id="budget-display-summarized-beats-off",
+ ),
+ pytest.param(
+ "gemini-2.5-flash",
+ {"thinking": {"type": "enabled", "budget_tokens": 2048, "display": "omitted"}},
+ True,
+ True,
+ id="budget-display-omitted-leaves-the-default-on",
+ ),
+ pytest.param(
+ "gemini-3.8-flash",
+ {"thinking": {"type": "enabled", "budget_tokens": 2048, "display": "updates"}},
+ False,
+ True,
+ id="level-display-updates-beats-off",
+ ),
+ pytest.param(
+ "gemini-3.8-flash",
+ {"reasoning_effort": {"effort": "medium", "summary": "auto"}},
+ False,
+ True,
+ id="level-effort-summary-beats-off",
+ ),
+ pytest.param(
+ "gemini-2.5-flash",
+ {"reasoning_effort": {"effort": "medium", "summary": "detailed"}},
+ False,
+ True,
+ id="budget-effort-summary-beats-off",
+ ),
+ pytest.param(
+ "gemini-3.8-flash",
+ {"thinking": {"type": "enabled", "budget_tokens": 2048, "summary": "auto"}},
+ False,
+ True,
+ id="level-thinking-summary-beats-off",
+ ),
+ pytest.param(
+ "gemini-2.5-flash",
+ {"thinking": {"type": "enabled", "budget_tokens": 2048, "summary": "detailed"}},
+ False,
+ True,
+ id="budget-thinking-summary-beats-off",
+ ),
+ pytest.param(
+ "gemini-3.8-flash",
+ {"reasoning_effort": "medium", "reasoning_summary": "auto"},
+ False,
+ True,
+ id="level-reasoning-summary-alias-beats-off",
+ ),
+ pytest.param(
+ "gemini-2.5-flash",
+ {"reasoning_effort": "medium", "reasoningSummary": "detailed"},
+ False,
+ True,
+ id="budget-reasoning-summary-camel-alias-beats-off",
+ ),
+ pytest.param(
+ "gemini-3.8-flash",
+ {"reasoning_effort": "none", "reasoning_summary": "auto"},
+ True,
+ False,
+ id="effort-none-ignores-summary-alias",
+ ),
+ pytest.param(
+ "gemini-3.8-flash",
+ {"reasoning_effort": {"effort": "medium"}},
+ False,
+ False,
+ id="effort-without-summary-keeps-off",
+ ),
+ pytest.param(
+ "gemini-3.8-flash",
+ {"reasoning_effort": {"effort": "medium"}, "reasoning_summary": "auto"},
+ False,
+ True,
+ id="effort-dict-takes-the-summary-alias",
+ ),
+ pytest.param(
+ "gemini-3.8-flash",
+ {"reasoning_effort": {"effort": "none", "summary": "auto"}},
+ True,
+ False,
+ id="effort-none-ignores-summary",
+ ),
+)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(("model", "thinking_params", "gateway_default", "expected"), PER_REQUEST_THOUGHT_SUMMARY_ASKS)
+async def test_gemini_per_request_thought_summary_ask_beats_gateway_default(
+ model: str,
+ thinking_params: dict[str, object],
+ gateway_default: bool,
+ expected: bool,
+ local_model_cost_map: None,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ monkeypatch.setattr(litellm, "gemini_include_thoughts_default", gateway_default)
+
+ sent: Final = await _sent_thinking_config(model, thinking_params)
+
+ assert sent["includeThoughts"] is expected
+
+
+def _includes_thoughts_client(sent: list[bool]) -> AsyncHTTPHandler:
+ def upstream(request: httpx.Request) -> httpx.Response:
+ sent.append(json.loads(request.content)["generationConfig"]["thinkingConfig"]["includeThoughts"])
+ return httpx.Response(
+ 200,
+ json={
+ "candidates": [{"content": {"parts": [{"text": "hi"}], "role": "model"}, "finishReason": "STOP"}],
+ "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2},
+ },
+ request=request,
+ )
+
+ client: Final = AsyncHTTPHandler()
+ client.client = httpx.AsyncClient(transport=httpx.MockTransport(upstream))
+ return client
+
+
+async def _chat_summary_ask(client: AsyncHTTPHandler, asks: bool) -> None:
+ await litellm.acompletion(
+ model="gemini/gemini-3.8-flash",
+ messages=[{"role": "user", "content": "hi"}],
+ api_key="fake-gemini-key",
+ client=client,
+ reasoning_effort="medium",
+ **({"reasoning_summary": "auto"} if asks else {}),
+ )
+
+
+async def _messages_summary_ask(client: AsyncHTTPHandler, asks: bool) -> None:
+ await litellm.anthropic.messages.acreate(
+ model="gemini/gemini-3.8-flash",
+ messages=[{"role": "user", "content": "hi"}],
+ max_tokens=64,
+ api_key="fake-gemini-key",
+ client=client,
+ thinking={"type": "enabled", "budget_tokens": 2048, **({"display": "summarized"} if asks else {})},
+ )
+
+
+async def _responses_summary_ask(client: AsyncHTTPHandler, asks: bool) -> None:
+ await litellm.aresponses(
+ model="gemini/gemini-3.8-flash",
+ input="hi",
+ api_key="fake-gemini-key",
+ client=client,
+ reasoning={"effort": "medium", **({"summary": "auto"} if asks else {})},
+ )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "send",
+ (_chat_summary_ask, _messages_summary_ask, _responses_summary_ask),
+ ids=("chat-reasoning-summary-alias", "messages-display-summarized", "responses-reasoning-summary"),
+)
+async def test_gemini_summary_ask_turns_thoughts_on_and_logs_the_same_reasoning_effort(
+ send: Callable[[AsyncHTTPHandler, bool], Awaitable[None]],
+ local_model_cost_map: None,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ logged: Final[list[object]] = []
+
+ class ReasoningEffortLogger(CustomLogger):
+ async def async_log_success_event(
+ self,
+ kwargs: Mapping[str, Mapping[str, Mapping[str, object]]],
+ response_obj: object,
+ start_time: object,
+ end_time: object,
+ ) -> None:
+ logged.append(kwargs["standard_logging_object"]["model_parameters"].get("reasoning_effort"))
+
+ monkeypatch.setattr(litellm, "gemini_include_thoughts_default", False)
+ monkeypatch.setattr(litellm, "callbacks", [ReasoningEffortLogger()])
+ sent: Final[list[bool]] = []
+ client: Final = _includes_thoughts_client(sent)
+
+ await send(client, False)
+ await send(client, True)
+ await asyncio.sleep(0)
+ await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0)
+
+ assert sent == [False, True]
+ assert len(logged) == 2
+ assert isinstance(logged[0], str)
+ assert logged[1] == logged[0]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ ("gateway_default", "expected"),
+ (("false", False), (" Off ", False), ("0", False), ("true", True), (None, True), ("not-a-flag", True)),
+)
+async def test_gemini_sends_a_boolean_for_a_gateway_default_written_as_text_or_null(
+ gateway_default: object, expected: bool, local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch
+) -> None:
+ monkeypatch.setattr(litellm, "gemini_include_thoughts_default", gateway_default)
+
+ sent: Final = await _sent_thinking_config("gemini-3.8-flash", {"reasoning_effort": "high"})
+
+ assert sent["includeThoughts"] is expected
diff --git a/tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py b/tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py
index a705c169782..0bfff870d15 100644
--- a/tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py
+++ b/tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py
@@ -223,6 +223,36 @@ def test_vertex_session_update_normalizes_ga_remapped_fields(
)
+def test_vertex_keeps_turn_detection_tuning_when_auto_response_is_off():
+ session_update: Final = {
+ "type": "session.update",
+ "session": {
+ "audio": {
+ "input": {
+ "turn_detection": {
+ "type": "server_vad",
+ "create_response": False,
+ "prefix_padding_ms": 200,
+ "silence_duration_ms": 700,
+ }
+ },
+ },
+ },
+ }
+ cfg: Final = VertexAIRealtimeConfig(access_token="tok", project="my-proj", location="us-central1")
+ messages: Final = cfg.transform_realtime_request(
+ json.dumps(session_update),
+ _NATIVE_AUDIO_MODEL,
+ session_configuration_request=None,
+ )
+ setup: Final = json.loads(messages[0])["setup"]
+ assert setup["realtimeInputConfig"]["automaticActivityDetection"] == {
+ "disabled": True,
+ "prefixPaddingMs": 200,
+ "silenceDurationMs": 700,
+ }
+
+
# ---------------------------------------------------------------------------
# Round-trip test: text-in / text-out via RealTimeStreaming
# ---------------------------------------------------------------------------
diff --git a/tests/unit/llms/vertex_ai/test_vertex.py b/tests/unit/llms/vertex_ai/test_vertex.py
index 044a07859f9..04cf391c700 100644
--- a/tests/unit/llms/vertex_ai/test_vertex.py
+++ b/tests/unit/llms/vertex_ai/test_vertex.py
@@ -1814,8 +1814,8 @@ def test_gemini_nullable_object_tool_schema_httpx(
"type": "object",
"required": ["source"],
"properties": {"source": {"type": "string"}},
+ "nullable": True,
},
- {"type": "null"},
]
diff --git a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py
index 8528724a24e..d244fe6bd90 100644
--- a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py
+++ b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py
@@ -1,3 +1,4 @@
+import copy
import time
from collections.abc import Iterator
from contextlib import contextmanager
@@ -14,6 +15,7 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler
from litellm.llms.vertex_ai.common_utils import (
VertexAITokenCounter,
+ build_vertex_schema,
convert_anyof_null_to_nullable,
convert_vertex_datetime_to_openai_datetime,
get_vertex_location_from_url,
@@ -234,6 +236,147 @@ def test_nested_anyof_conversion():
assert schema == expected
+def test_anyof_conversion_with_multiple_null_branches():
+ """Adjacent null branches must all be removed."""
+ schema: Final = {
+ "type": "object",
+ "properties": {
+ "example": {
+ "anyOf": [
+ {"type": "null"},
+ {"type": "null"},
+ {"type": "string"},
+ ]
+ }
+ },
+ }
+
+ convert_anyof_null_to_nullable(schema)
+
+ expected: Final = {
+ "type": "object",
+ "properties": {"example": {"anyOf": [{"type": "string", "nullable": True}]}},
+ }
+ assert schema == expected
+
+
+def test_anyof_conversion_with_null_before_empty_object():
+ """An empty-object branch following a null branch still gets a type."""
+ schema: Final = {
+ "type": "object",
+ "properties": {"example": {"anyOf": [{"type": "null"}, {}]}},
+ }
+
+ convert_anyof_null_to_nullable(schema)
+
+ expected: Final = {
+ "type": "object",
+ "properties": {"example": {"anyOf": [{"type": "object", "nullable": True}]}},
+ }
+ assert schema == expected
+
+
+@pytest.mark.parametrize("null_count", [1, 2, 4])
+def test_anyof_conversion_with_only_null_branches_keeps_one_nullable_null(null_count: int):
+ """Vertex accepts a NULL type, so a null-only union is forwarded as one nullable null branch."""
+ schema: Final = {
+ "type": "object",
+ "properties": {"example": {"anyOf": [{"type": "null"} for _ in range(null_count)]}},
+ }
+
+ convert_anyof_null_to_nullable(schema)
+
+ expected: Final = {
+ "type": "object",
+ "properties": {"example": {"anyOf": [{"type": "null", "nullable": True}]}},
+ }
+ assert schema == expected
+
+
+def test_anyof_conversion_leaves_empty_anyof_for_the_provider_to_judge():
+ """An empty union has no null branch to strip, so it is forwarded unchanged."""
+ schema: Final = {"type": "object", "properties": {"example": {"anyOf": []}}}
+
+ convert_anyof_null_to_nullable(schema)
+
+ assert schema == {"type": "object", "properties": {"example": {"anyOf": []}}}
+
+
+@pytest.mark.parametrize("malformed", ["string", {"type": "string"}], ids=["string", "object"])
+def test_anyof_conversion_forwards_a_non_list_anyof_unchanged(malformed: JsonValue):
+ """A malformed union is the provider's to judge, as it was before null branches were stripped in one pass."""
+ schema: Final = {"type": "object", "properties": {"example": {"anyOf": copy.deepcopy(malformed)}}}
+
+ convert_anyof_null_to_nullable(schema)
+
+ assert schema == {"type": "object", "properties": {"example": {"anyOf": malformed}}}
+
+
+@pytest.mark.parametrize(
+ ("branches", "expected"),
+ [
+ ([{"type": "null"}, 5], [5]),
+ ([True, {"type": "null"}], [True]),
+ ([{"type": "null"}, 5, {"type": "string"}], [5, {"type": "string", "nullable": True}]),
+ ],
+ ids=["null-int", "true-null", "null-int-string"],
+)
+def test_anyof_conversion_forwards_a_non_object_branch_for_the_provider_to_judge(
+ branches: list[JsonValue], expected: list[JsonValue]
+):
+ schema: Final = {"type": "object", "properties": {"example": {"anyOf": branches}}}
+
+ convert_anyof_null_to_nullable(schema)
+
+ assert schema == {"type": "object", "properties": {"example": {"anyOf": expected}}}
+
+
+
+@pytest.mark.parametrize(
+ ("field", "expected"),
+ [
+ ({"type": ["integer", "null"]}, {"anyOf": [{"type": "integer", "nullable": True}]}),
+ (
+ {
+ "anyOf": [
+ {"type": "object", "properties": {"x": {"anyOf": [{"type": "integer"}, {"type": "null"}]}}},
+ {"type": "null"},
+ ]
+ },
+ {
+ "anyOf": [
+ {
+ "type": "object",
+ "properties": {"x": {"anyOf": [{"type": "integer", "nullable": True}]}},
+ "nullable": True,
+ }
+ ]
+ },
+ ),
+ (
+ {"type": ["object", "null"], "properties": {"x": {"type": ["string", "null"]}}},
+ {
+ "anyOf": [
+ {
+ "type": "object",
+ "properties": {"x": {"anyOf": [{"type": "string", "nullable": True}]}},
+ "nullable": True,
+ }
+ ]
+ },
+ ),
+ (
+ {"type": "array", "items": {"type": ["number", "null"]}},
+ {"type": "array", "items": {"anyOf": [{"type": "number", "nullable": True}]}},
+ ),
+ ],
+ ids=["type-array", "null-inside-a-union-branch", "nullable-object-with-nullable-property", "nullable-items"],
+)
+def test_vertex_schema_folds_every_null_branch_into_nullable(field: JsonValue, expected: JsonValue):
+ schema: Final = build_vertex_schema({"type": "object", "properties": {"value": field}})
+
+ assert schema["properties"]["value"] == expected
+
def test_anyof_with_excessive_nesting():
"""Test conversion with excessive nesting > max levels +1 deep."""
# generate a schema with excessive nesting
diff --git a/tests/unit/llms/volcengine/test_volcengine_video_transformation.py b/tests/unit/llms/volcengine/test_volcengine_video_transformation.py
new file mode 100644
index 00000000000..1fcf49eeb13
--- /dev/null
+++ b/tests/unit/llms/volcengine/test_volcengine_video_transformation.py
@@ -0,0 +1,1131 @@
+import base64
+import io
+import json
+from pathlib import Path
+from typing import Final
+from unittest.mock import Mock
+
+import httpx
+import pytest
+import respx
+
+import litellm
+from litellm import Router
+from litellm.exceptions import UnsupportedParamsError
+from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
+from litellm.llms.openai.cost_calculation import video_generation_cost
+from litellm.llms.volcengine.videos.transformation import VolcEngineVideoConfig, VolcEngineVideoError
+from litellm.types.router import GenericLiteLLMParams
+from litellm.types.videos.utils import decode_video_id_with_provider, encode_video_id_with_provider
+from litellm.utils import ProviderConfigManager
+from litellm.videos.main import avideo_content, avideo_generation
+
+PNG_BYTES: Final = base64.b64decode(
+ "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg=="
+)
+PNG_DATA_URL: Final = f"data:image/png;base64,{base64.b64encode(PNG_BYTES).decode()}"
+
+API_BASE: Final = "https://ark.cn-beijing.volces.com"
+TASKS_URL: Final = f"{API_BASE}/api/v3/contents/generations/tasks"
+SEEDANCE_2_5: Final = "doubao-seedance-2-5-260628"
+SEEDANCE_2_0: Final = "doubao-seedance-2-0-260128"
+SEEDANCE_2_0_FAST: Final = "doubao-seedance-2-0-fast-260128"
+TASK_ID: Final = "cgt-2026-abc"
+
+
+def _response(payload: object, status_code: int = 200) -> httpx.Response:
+ return httpx.Response(status_code, json=payload, request=httpx.Request("GET", TASKS_URL))
+
+
+def _text_response(text: str, status_code: int = 200) -> httpx.Response:
+ return httpx.Response(status_code, text=text, request=httpx.Request("GET", TASKS_URL))
+
+
+def _create_body(
+ params: dict[str, object], model: str = SEEDANCE_2_5, drop_params: bool | None = None
+) -> dict[str, object]:
+ body, files, url = VolcEngineVideoConfig().transform_video_create_request(
+ model=model,
+ prompt="a cat on a roof",
+ api_base=API_BASE,
+ video_create_optional_request_params=params,
+ litellm_params=GenericLiteLLMParams(drop_params=drop_params),
+ headers={},
+ )
+ assert (files, url) == ((), TASKS_URL)
+ return body
+
+
+def _created(request_data: dict[str, object], model: str = SEEDANCE_2_5, logging_obj: object = None):
+ return VolcEngineVideoConfig().transform_video_create_response(
+ model=model,
+ raw_response=_response({"id": TASK_ID}),
+ logging_obj=logging_obj, # pyright: ignore[reportArgumentType] # the transform only reads litellm_params off it
+ custom_llm_provider="volcengine",
+ request_data=request_data,
+ )
+
+
+def _polled(payload: object, logging_obj: object = None):
+ return VolcEngineVideoConfig().transform_video_status_retrieve_response(
+ raw_response=_response(payload),
+ logging_obj=logging_obj, # pyright: ignore[reportArgumentType] # the transform only reads litellm_params off it
+ custom_llm_provider="volcengine",
+ )
+
+
+def _task(**fields: object) -> dict[str, object]:
+ return {"id": TASK_ID, "model": SEEDANCE_2_5, "status": "succeeded", **fields}
+
+
+@pytest.mark.parametrize(
+ "input_reference",
+ [
+ PNG_BYTES,
+ io.BytesIO(PNG_BYTES),
+ ("frame.png", PNG_BYTES, "image/png"),
+ ("frame.png", io.BytesIO(PNG_BYTES)),
+ ],
+ ids=["bytes", "file", "tuple", "tuple-with-file"],
+)
+def test_a_file_input_reference_becomes_a_data_uri_first_frame(input_reference: object):
+ body = _create_body({"input_reference": input_reference, "duration": 4})
+
+ assert body == {
+ "model": SEEDANCE_2_5,
+ "content": [
+ {"type": "text", "text": "a cat on a roof"},
+ {"type": "image_url", "image_url": {"url": PNG_DATA_URL}, "role": "first_frame"},
+ ],
+ "resolution": "720p",
+ "duration": 4,
+ }
+
+
+def test_a_path_input_reference_is_read_from_disk(tmp_path: Path):
+ frame = tmp_path / "frame.png"
+ frame.write_bytes(PNG_BYTES)
+
+ body = _create_body({"input_reference": frame, "duration": 4})
+
+ assert body["content"] == [
+ {"type": "text", "text": "a cat on a roof"},
+ {"type": "image_url", "image_url": {"url": PNG_DATA_URL}, "role": "first_frame"},
+ ]
+
+
+def test_a_url_input_reference_is_passed_as_is_and_never_opened_as_a_path():
+ body = _create_body({"input_reference": "https://x/first.png", "duration": 4})
+
+ assert body["content"] == [
+ {"type": "text", "text": "a cat on a roof"},
+ {"type": "image_url", "image_url": {"url": "https://x/first.png"}, "role": "first_frame"},
+ ]
+
+
+def test_openais_object_input_reference_sends_its_image_url():
+ body = _create_body({"input_reference": {"image_url": "https://x/first.png"}, "duration": 4})
+
+ assert body["content"] == [
+ {"type": "text", "text": "a cat on a roof"},
+ {"type": "image_url", "image_url": {"url": "https://x/first.png"}, "role": "first_frame"},
+ ]
+
+
+def test_a_file_id_input_reference_is_a_bad_request_unless_drop_params():
+ with pytest.raises(UnsupportedParamsError, match="file_id is not supported"):
+ _ = _create_body({"input_reference": {"file_id": "file-abc"}, "duration": 4})
+
+ body = _create_body({"input_reference": {"file_id": "file-abc"}, "duration": 4}, drop_params=True)
+
+ assert body["content"] == [{"type": "text", "text": "a cat on a roof"}]
+
+
+def test_a_caller_content_array_wins_over_input_reference_and_the_prompt_leads():
+ body = _create_body(
+ {
+ "content": [{"type": "image_url", "image_url": {"url": "https://x/a.png"}, "role": "reference_image"}],
+ "input_reference": "https://x/ignored.png",
+ "duration": 4,
+ }
+ )
+
+ assert body["content"] == [
+ {"type": "text", "text": "a cat on a roof"},
+ {"type": "image_url", "image_url": {"url": "https://x/a.png"}, "role": "reference_image"},
+ ]
+
+
+IMAGE_ITEM: Final = {"type": "image_url", "image_url": {"url": "https://x/a.png"}, "role": "reference_image"}
+
+
+def test_a_text_item_in_content_is_rejected_because_input_guardrails_only_scan_the_prompt():
+ with pytest.raises(UnsupportedParamsError, match="put the text in prompt"):
+ _ = _create_body({"content": [{"type": "text", "text": "text the guardrail never saw"}], "duration": 4})
+
+
+def test_a_text_item_in_content_is_dropped_under_drop_params_so_ark_never_sees_unscanned_text():
+ body = _create_body(
+ {"content": [{"type": "text", "text": "text the guardrail never saw"}, IMAGE_ITEM], "duration": 4},
+ drop_params=True,
+ )
+
+ assert body["content"] == [{"type": "text", "text": "a cat on a roof"}, IMAGE_ITEM]
+
+
+@pytest.mark.parametrize(
+ "item",
+ [
+ {"type": "video_url", "video_url": {"url": "https://x/v.mp4"}, "role": "reference_video"},
+ {"type": "video_url", "video_url": {"url": "https://x/v.mp4"}},
+ {"type": "image_url", "image_url": {"url": "https://x/v.png"}, "role": "reference_video"},
+ ],
+ ids=["reference-video", "untagged-video", "mistagged-role"],
+)
+def test_a_video_input_is_rejected_because_its_cost_is_not_known_at_create(item: dict[str, object]):
+ with pytest.raises(UnsupportedParamsError, match="Video input"):
+ _ = _create_body({"content": [item], "duration": 4})
+
+ body = _create_body({"content": [item, IMAGE_ITEM], "duration": 4}, drop_params=True)
+
+ assert body["content"] == [{"type": "text", "text": "a cat on a roof"}, IMAGE_ITEM]
+
+
+@pytest.mark.parametrize(
+ "content",
+ ["a string", [{"type": "image_url", "role": 7}], [{"type": 7}]],
+ ids=["not-an-array", "non-string-role", "non-string-type"],
+)
+def test_a_content_value_that_is_not_an_array_of_items_is_rejected_unless_drop_params(content: object):
+ with pytest.raises(UnsupportedParamsError, match="content must be an array"):
+ _ = _create_body({"content": content, "duration": 4})
+
+ body = _create_body({"content": content, "duration": 4}, drop_params=True)
+
+ assert body["content"] == [{"type": "text", "text": "a cat on a roof"}]
+
+
+def test_an_omitted_duration_and_tier_are_sent_as_openais_four_seconds_at_720p_so_ark_cannot_pick_what_is_billed():
+ body = _create_body({})
+
+ assert body == {
+ "model": SEEDANCE_2_5,
+ "content": [{"type": "text", "text": "a cat on a roof"}],
+ "resolution": "720p",
+ "duration": 4,
+ }
+
+
+@pytest.mark.usefixtures("local_model_cost_map")
+def test_an_endpoint_id_priced_through_base_model_gets_the_base_models_default_tier():
+ body, _, _ = VolcEngineVideoConfig().transform_video_create_request(
+ model="ep-20261015-abc",
+ prompt="a cat on a roof",
+ api_base=API_BASE,
+ video_create_optional_request_params={"duration": 4},
+ litellm_params=GenericLiteLLMParams(model_info={"base_model": f"volcengine/{SEEDANCE_2_5}"}),
+ headers={},
+ )
+
+ assert (body["model"], body["resolution"]) == ("ep-20261015-abc", "720p")
+
+
+@pytest.mark.parametrize("field", ["return_last_frame", "draft", "frame_count"])
+def test_a_field_litellm_does_not_send_to_ark_is_rejected(field: str):
+ with pytest.raises(UnsupportedParamsError, match=f"does not send {field} to Ark"):
+ _ = _create_body({"duration": 4, field: True})
+
+
+def test_a_field_litellm_does_not_send_to_ark_is_dropped_under_drop_params():
+ body = _create_body({"duration": 4, "return_last_frame": True, "seed": 7}, drop_params=True)
+
+ assert "return_last_frame" not in body
+ assert body["seed"] == 7
+
+
+@pytest.mark.parametrize("duration", [-1, "-1", 0, "abc", 4.5])
+def test_a_duration_ark_would_pick_or_that_is_not_whole_seconds_is_rejected(duration: object):
+ with pytest.raises(UnsupportedParamsError, match="duration must be a positive whole number"):
+ _ = _create_body({"duration": duration})
+
+ assert _create_body({"duration": duration}, drop_params=True)["duration"] == 4
+
+
+def test_frames_is_rejected_because_it_bypasses_the_billed_duration():
+ with pytest.raises(UnsupportedParamsError, match="frames"):
+ _ = _create_body({"frames": 121, "duration": 6})
+
+ body = _create_body({"frames": 121, "duration": 6}, drop_params=True)
+
+ assert ("frames" not in body, body["duration"]) == (True, 6)
+
+
+def test_ark_fields_reach_the_body_and_false_flags_are_kept():
+ body = _create_body(
+ {
+ "duration": 4,
+ "resolution": "480p",
+ "ratio": "9:16",
+ "seed": 7,
+ "watermark": False,
+ "generate_audio": False,
+ "service_tier": "flex",
+ }
+ )
+
+ assert body == {
+ "model": SEEDANCE_2_5,
+ "content": [{"type": "text", "text": "a cat on a roof"}],
+ "resolution": "480p",
+ "ratio": "9:16",
+ "seed": 7,
+ "watermark": False,
+ "generate_audio": False,
+ "service_tier": "flex",
+ "duration": 4,
+ }
+
+
+@pytest.mark.usefixtures("local_model_cost_map")
+def test_a_resolution_the_model_has_no_price_for_is_rejected():
+ with pytest.raises(UnsupportedParamsError, match="not offered at resolution '1080p', only at 480p, 720p"):
+ _ = _create_body({"duration": 4, "resolution": "1080p"}, model=SEEDANCE_2_0_FAST)
+
+
+@pytest.mark.usefixtures("local_model_cost_map")
+def test_an_unpriced_resolution_is_dropped_under_drop_params():
+ body = _create_body(
+ {"duration": 4, "resolution": "1080p", "ratio": "16:9"}, model=SEEDANCE_2_0_FAST, drop_params=True
+ )
+
+ assert body == {
+ "model": SEEDANCE_2_0_FAST,
+ "content": [{"type": "text", "text": "a cat on a roof"}],
+ "resolution": "720p",
+ "ratio": "16:9",
+ "duration": 4,
+ }
+
+
+@pytest.mark.usefixtures("local_model_cost_map")
+def test_a_model_the_router_registered_without_prices_keeps_its_resolution():
+ endpoint_id: Final = "ep-20261015-unpriced"
+ _ = Router(
+ model_list=[
+ {"model_name": "ark", "litellm_params": {"model": f"volcengine/{endpoint_id}", "api_key": "sk-test"}}
+ ]
+ )
+
+ body = _create_body({"duration": 4, "resolution": "480p"}, model=endpoint_id)
+ body_without_a_tier = _create_body({"duration": 4}, model=endpoint_id)
+
+ assert body["resolution"] == "480p"
+ assert "resolution" not in body_without_a_tier
+
+
+def test_a_mistyped_ark_field_is_rejected():
+ with pytest.raises(UnsupportedParamsError, match="invalid seed"):
+ _ = _create_body({"duration": 4, "seed": "abc", "watermark": False})
+
+
+def test_a_mistyped_ark_field_is_dropped_under_drop_params():
+ body = _create_body({"duration": 4, "seed": "abc", "watermark": False}, drop_params=True)
+
+ assert body == {
+ "model": SEEDANCE_2_5,
+ "content": [{"type": "text", "text": "a cat on a roof"}],
+ "resolution": "720p",
+ "watermark": False,
+ "duration": 4,
+ }
+
+
+@pytest.mark.parametrize(
+ "size,expected",
+ [
+ ("1920x1080", {"ratio": "16:9", "resolution": "1080p"}),
+ ("854x480", {"ratio": "16:9", "resolution": "480p"}),
+ ("864x496", {"ratio": "16:9", "resolution": "480p"}),
+ ("720x1280", {"ratio": "9:16", "resolution": "720p"}),
+ ("1024x768", {"ratio": "4:3", "resolution": "720p"}),
+ ("640x640", {"ratio": "1:1", "resolution": "480p"}),
+ ("960x960", {"ratio": "1:1", "resolution": "720p"}),
+ ("1440x1440", {"ratio": "1:1", "resolution": "1080p"}),
+ ("1112x834", {"ratio": "4:3", "resolution": "720p"}),
+ ("1470x630", {"ratio": "21:9", "resolution": "720p"}),
+ ("3840x2160", {"ratio": "16:9", "resolution": "4k"}),
+ ],
+)
+def test_size_maps_onto_an_ark_ratio_and_tier(size: str, expected: dict[str, str]):
+ assert VolcEngineVideoConfig().map_openai_params({"size": size}, SEEDANCE_2_5, False) == expected
+
+
+def test_explicit_ratio_and_resolution_beat_the_size_derived_ones():
+ mapped = VolcEngineVideoConfig().map_openai_params(
+ {"size": "1280x720", "seconds": "4", "ratio": "1:1", "resolution": "480p", "parameters": {"seed": 3}},
+ SEEDANCE_2_5,
+ False,
+ )
+
+ assert mapped == {"ratio": "1:1", "resolution": "480p", "duration": 4, "seed": 3}
+
+
+@pytest.mark.parametrize(
+ "params,message",
+ [
+ ({"seconds": "4.5"}, "seconds must be a whole number"),
+ ({"size": "big"}, "size must be WIDTHxHEIGHT"),
+ ({"size": "16:9"}, "size must be WIDTHxHEIGHT"),
+ ({"parameters": "seed=3"}, "parameters must be an object"),
+ ({"characters": [{"id": "c"}]}, "'characters' is not supported"),
+ ({"image": "https://x/a.png"}, "'image' is not supported"),
+ ],
+)
+def test_a_value_ark_cannot_take_raises_unless_drop_params(params: dict[str, object], message: str):
+ with pytest.raises(UnsupportedParamsError, match=message):
+ _ = VolcEngineVideoConfig().map_openai_params(params, SEEDANCE_2_5, False)
+
+ assert VolcEngineVideoConfig().map_openai_params(params, SEEDANCE_2_5, True) == {}
+
+
+def test_request_metadata_is_not_mapped_into_the_body_and_user_becomes_the_safety_identifier():
+ mapped = VolcEngineVideoConfig().map_openai_params(
+ {"user": "u", "extra_headers": {"x": "y"}, "seconds": 4}, SEEDANCE_2_5, False
+ )
+
+ assert mapped == {"safety_identifier": "u", "duration": 4}
+
+
+def test_an_explicit_safety_identifier_wins_over_user():
+ mapped = VolcEngineVideoConfig().map_openai_params(
+ {"user": "u", "safety_identifier": "s", "seconds": 4}, SEEDANCE_2_5, False
+ )
+
+ assert mapped["safety_identifier"] == "s"
+
+
+@pytest.mark.usefixtures("local_model_cost_map")
+@pytest.mark.parametrize(
+ "model,request_data,seconds,size,usage",
+ [
+ (
+ SEEDANCE_2_5,
+ {"duration": 4, "resolution": "480p", "ratio": "16:9"},
+ "4",
+ "854x480",
+ {"duration_seconds": 4.0, "video_resolution": "480p"},
+ ),
+ (
+ SEEDANCE_2_0,
+ {"duration": 5, "resolution": "480p", "ratio": "16:9"},
+ "5",
+ "864x496",
+ {"duration_seconds": 5.0, "video_resolution": "480p"},
+ ),
+ (SEEDANCE_2_5, {"duration": 5}, "5", None, {"duration_seconds": 5.0}),
+ ],
+ ids=["2.5-480p", "2.0-480p-renders-its-own-frame", "no-geometry"],
+)
+def test_the_create_response_reports_what_the_task_bills(
+ model: str, request_data: dict[str, object], seconds: str, size: str | None, usage: dict[str, object]
+):
+ video = _created(request_data, model=model)
+
+ assert video.model_dump(include={"status", "seconds", "size", "usage", "model"}) == {
+ "status": "queued",
+ "seconds": seconds,
+ "size": size,
+ "usage": usage,
+ "model": model,
+ }
+
+
+@pytest.mark.parametrize(
+ "usage,expected_cost",
+ [
+ ({"duration_seconds": 4.0, "video_resolution": "480p"}, 4 * 0.1),
+ ({"duration_seconds": 5.0, "video_resolution": "720p"}, 5 * 0.2),
+ ({"duration_seconds": 5.0}, 5 * 0.2),
+ ],
+)
+def test_the_create_usage_bills_its_duration_at_its_tier_rate(usage: dict[str, object], expected_cost: float):
+ model_info = {
+ "output_cost_per_second": 0.2,
+ "output_cost_per_second_480p": 0.1,
+ "output_cost_per_second_720p": 0.2,
+ }
+ video = _created({"duration": int(usage["duration_seconds"]), "resolution": usage.get("video_resolution")})
+
+ cost = video_generation_cost(
+ model=SEEDANCE_2_5,
+ duration_seconds=video.usage["duration_seconds"],
+ model_info=model_info, # pyright: ignore[reportArgumentType] # a hand-written rate card stands in for ModelInfo
+ video_resolution=video.usage.get("video_resolution"),
+ )
+
+ assert cost == pytest.approx(expected_cost)
+
+
+def test_the_created_id_carries_the_requested_and_served_model_groups():
+ logging_obj = Mock(
+ litellm_params={"litellm_metadata": {"model_group": "seedance", "deployment_model_name": "seedance-account-b"}}
+ )
+
+ video = _created({"duration": 5}, logging_obj=logging_obj)
+
+ assert decode_video_id_with_provider(video.id) == {
+ "custom_llm_provider": "volcengine",
+ "model_id": "seedance",
+ "video_id": TASK_ID,
+ "served_model": "seedance-account-b",
+ }
+
+
+def test_a_create_without_router_metadata_carries_the_model():
+ video = _created({"duration": 5}, logging_obj=Mock(litellm_params={}))
+
+ assert decode_video_id_with_provider(video.id) == {
+ "custom_llm_provider": "volcengine",
+ "model_id": SEEDANCE_2_5,
+ "video_id": TASK_ID,
+ }
+
+
+@pytest.mark.parametrize(
+ "raw_response,status_code,message",
+ [
+ (_response({"error": {"code": "InvalidParameter", "message": "bad ratio"}}), 400, "bad ratio"),
+ (_response({"error": {"code": "Forbidden", "message": "no access"}}, 403), 403, "no access"),
+ (_text_response("bad gateway"), 502, "Unexpected Volcengine video response"),
+ ],
+ ids=["error-on-200", "error-status-kept", "not-json"],
+)
+def test_a_create_response_without_a_task_id_raises_instead_of_queueing_an_empty_video(
+ raw_response: httpx.Response, status_code: int, message: str
+):
+ with pytest.raises(VolcEngineVideoError, match=message) as raised:
+ _ = VolcEngineVideoConfig().transform_video_create_response(
+ model=SEEDANCE_2_5,
+ raw_response=raw_response,
+ logging_obj=None, # pyright: ignore[reportArgumentType] # the transform only reads litellm_params off it
+ custom_llm_provider="volcengine",
+ request_data={"duration": 5},
+ )
+
+ assert raised.value.status_code == status_code
+
+
+@pytest.mark.parametrize(
+ "ark_status,expected",
+ [
+ ("queued", "queued"),
+ ("running", "in_progress"),
+ ("succeeded", "completed"),
+ ("failed", "failed"),
+ ("cancelled", "cancelled"),
+ ("expired", "failed"),
+ ],
+)
+def test_the_ark_status_maps_onto_the_openai_vocabulary(ark_status: str, expected: str):
+ assert _polled(_task(status=ark_status)).status == expected
+
+
+@pytest.mark.usefixtures("local_model_cost_map")
+def test_a_completed_task_reports_its_delivered_geometry_duration_and_tokens():
+ video = _polled(
+ _task(
+ created_at=1791989938,
+ updated_at=1791990031,
+ resolution="480p",
+ ratio="16:9",
+ duration=4,
+ usage={"completion_tokens": 40594, "total_tokens": 40594},
+ model=SEEDANCE_2_0,
+ content={"video_url": "https://tos/v.mp4"},
+ )
+ )
+
+ assert video.model_dump(exclude={"id"}) | {"id": decode_video_id_with_provider(video.id)["video_id"]} == {
+ "id": TASK_ID,
+ "object": "video",
+ "status": "completed",
+ "created_at": 1791989938,
+ "completed_at": 1791990031,
+ "expires_at": None,
+ "error": None,
+ "progress": None,
+ "remixed_from_video_id": None,
+ "seconds": "4",
+ "size": "864x496",
+ "model": SEEDANCE_2_0,
+ "usage": {"duration_seconds": 4.0, "video_resolution": "480p", "total_tokens": 40594},
+ }
+
+
+@pytest.mark.parametrize(
+ "resolution,ratio,expected",
+ [
+ ("720p", "9:16", "720x1280"),
+ ("1080p", "16:9", "1920x1080"),
+ ("480p", "adaptive", None),
+ ],
+)
+def test_a_frame_the_cost_map_does_not_list_is_scaled_from_the_tier(resolution: str, ratio: str, expected: str | None):
+ assert _polled(_task(resolution=resolution, ratio=ratio, model="doubao-seedance-unlisted")).size == expected
+
+
+def test_a_failed_task_carries_its_error_block():
+ video = _polled(_task(status="failed", error={"code": "OutputVideoSensitiveContentDetected", "message": "flagged"}))
+
+ assert (video.status, video.error, video.completed_at) == (
+ "failed",
+ {"code": "OutputVideoSensitiveContentDetected", "message": "flagged"},
+ None,
+ )
+
+
+@pytest.mark.parametrize(
+ "payload,message",
+ [
+ (_task(status="paused"), "Unknown Volcengine task status: paused"),
+ ({"status": "running"}, "Unexpected Volcengine video response"),
+ ({"id": TASK_ID}, "Unexpected Volcengine video response"),
+ ],
+ ids=["unmapped-status", "no-id", "no-status"],
+)
+def test_a_poll_litellm_cannot_read_raises_instead_of_returning_a_job_to_poll_forever(payload: object, message: str):
+ with pytest.raises(VolcEngineVideoError, match=message) as raised:
+ _ = _polled(payload)
+
+ assert raised.value.status_code == 502
+
+
+def test_a_poll_keeps_the_ids_requested_and_served_model_when_the_router_served_it():
+ video_id = encode_video_id_with_provider(TASK_ID, "volcengine", "seedance", served_model="seedance-account-b")
+ logging_obj = Mock(
+ litellm_params={
+ "video_id": video_id,
+ "litellm_metadata": {"model_group": "seedance", "deployment_model_name": "seedance-account-b"},
+ }
+ )
+
+ video = _polled(_task(status="running"), logging_obj=logging_obj)
+
+ assert video.id == video_id
+ assert video.model == "seedance"
+
+
+def test_a_poll_drops_a_served_model_the_router_did_not_route_to():
+ edited = encode_video_id_with_provider(
+ TASK_ID, "volcengine", "seedance-account-a", served_model="seedance-account-b"
+ )
+ logging_obj = Mock(litellm_params={"video_id": edited, "litellm_metadata": {"model_group": "seedance-account-a"}})
+
+ video = _polled(_task(status="running"), logging_obj=logging_obj)
+
+ assert video.id == encode_video_id_with_provider(TASK_ID, "volcengine", "seedance-account-a")
+
+
+async def _download(task_payload: dict[str, object], video_bytes: bytes = b"mp4") -> bytes:
+ def respond(request: httpx.Request) -> httpx.Response:
+ if request.url.host == "tos.example":
+ return httpx.Response(200, content=video_bytes)
+ return httpx.Response(200, json=task_payload)
+
+ async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client:
+ handler: Final = AsyncHTTPHandler()
+ await handler.close()
+ handler.client = http_client
+ return await avideo_content(
+ video_id=TASK_ID,
+ custom_llm_provider="volcengine",
+ model=f"volcengine/{SEEDANCE_2_5}",
+ api_key="sk-test",
+ api_base=API_BASE,
+ client=handler,
+ )
+
+
+@pytest.mark.asyncio
+async def test_the_download_fetches_the_signed_url_through_the_callers_client():
+ content = await _download(_task(content={"video_url": "https://tos.example/v.mp4"}), video_bytes=b"the-video")
+
+ assert content == b"the-video"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "payload,message",
+ [
+ (_task(status="queued"), "still processing"),
+ (_task(status="running"), "still processing"),
+ (_task(status="failed", error={"message": "flagged"}), "flagged"),
+ (_task(status="failed"), "ended as failed"),
+ (_task(status="cancelled"), "ended as cancelled"),
+ (_task(status="expired"), "Task expired"),
+ ],
+)
+async def test_a_download_with_no_video_yet_is_a_bad_request_so_it_is_neither_retried_nor_cools_the_deployment(
+ payload: dict[str, object], message: str
+):
+ with pytest.raises(litellm.BadRequestError, match=message) as raised:
+ _ = await _download(payload)
+
+ assert raised.value.status_code == 400
+
+
+def test_the_list_maps_ark_paging_and_filters():
+ url, params = VolcEngineVideoConfig().transform_video_list_request(
+ api_base=API_BASE,
+ litellm_params=GenericLiteLLMParams(),
+ headers={},
+ limit=20,
+ extra_query={"filter.status": "succeeded"},
+ )
+
+ assert (url, params) == (TASKS_URL, {"page_size": "20", "filter.status": "succeeded"})
+
+
+@pytest.mark.parametrize("paging", [{"after": "video_abc"}, {"order": "asc"}])
+def test_a_list_cursor_or_order_is_rejected_unless_drop_params(paging: dict[str, str]):
+ config = VolcEngineVideoConfig()
+
+ with pytest.raises(UnsupportedParamsError, match="pages by page_num"):
+ _ = config.transform_video_list_request(
+ api_base=API_BASE, litellm_params=GenericLiteLLMParams(), headers={}, **paging
+ )
+ _, params = config.transform_video_list_request(
+ api_base=API_BASE, litellm_params=GenericLiteLLMParams(drop_params=True), headers={}, limit=5, **paging
+ )
+
+ assert params == {"page_size": "5"}
+
+
+def test_the_list_response_is_an_openai_list_body():
+ response = VolcEngineVideoConfig().transform_video_list_response(
+ raw_response=_response({"total": 1, "items": [_task(status="running")]}),
+ logging_obj=Mock(), # pyright: ignore[reportArgumentType] # unused by the transform
+ custom_llm_provider="volcengine",
+ )
+
+ encoded = encode_video_id_with_provider(TASK_ID, "volcengine", SEEDANCE_2_5)
+ assert {key: value for key, value in response.items() if key != "data"} == {
+ "object": "list",
+ "total": 1,
+ "has_more": False,
+ "first_id": encoded,
+ "last_id": encoded,
+ }
+ assert [(item["id"], item["status"]) for item in response["data"]] == [(encoded, "in_progress")]
+
+
+def test_an_empty_list_still_carries_the_object_type():
+ response = VolcEngineVideoConfig().transform_video_list_response(
+ raw_response=_response({"total": 0, "items": []}),
+ logging_obj=Mock(), # pyright: ignore[reportArgumentType] # unused by the transform
+ )
+
+ assert response == {"object": "list", "data": [], "total": 0, "has_more": False}
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("call", ["avideo_edit", "avideo_extension", "avideo_remix"])
+async def test_edit_extension_and_remix_are_a_bad_request_because_a_video_input_cannot_be_billed_at_create(call: str):
+ with pytest.raises(litellm.BadRequestError, match="Ark bills a task with a video input") as raised:
+ _ = await getattr(litellm, call)(
+ video_id="https://x/source.mp4",
+ prompt="make it brighter",
+ seconds="5",
+ model=f"volcengine/{SEEDANCE_2_5}",
+ custom_llm_provider="volcengine",
+ api_key="sk-test",
+ api_base=API_BASE,
+ )
+
+ assert raised.value.status_code == 400
+
+
+@pytest.mark.parametrize(
+ "api_base",
+ [
+ None,
+ API_BASE,
+ f"{API_BASE}/",
+ f"{API_BASE}/api/v3",
+ f"{API_BASE}/api/compatible",
+ f"{API_BASE}/api/v3/responses",
+ f"{API_BASE}/v1",
+ ],
+)
+def test_any_ark_surface_on_the_shared_credential_normalizes_to_the_host_root(
+ api_base: str | None, monkeypatch: pytest.MonkeyPatch
+):
+ monkeypatch.delenv("VOLCENGINE_API_BASE", raising=False)
+ monkeypatch.delenv("ARK_API_BASE", raising=False)
+
+ assert VolcEngineVideoConfig().get_complete_url("m", api_base, {}) == API_BASE
+
+
+@pytest.mark.parametrize("env_var", ["VOLCENGINE_API_BASE", "ARK_API_BASE"])
+def test_an_api_base_env_var_routes_a_deployment_without_api_base(env_var: str, monkeypatch: pytest.MonkeyPatch):
+ monkeypatch.delenv("VOLCENGINE_API_BASE", raising=False)
+ monkeypatch.delenv("ARK_API_BASE", raising=False)
+ monkeypatch.setenv(env_var, "https://ark.ap-southeast.bytepluses.com/api/v3")
+
+ assert VolcEngineVideoConfig().get_complete_url("m", None, {}) == "https://ark.ap-southeast.bytepluses.com"
+ assert VolcEngineVideoConfig().get_complete_url("m", f"{API_BASE}/api/v3", {}) == API_BASE
+
+
+@pytest.mark.parametrize(
+ "api_base,expected",
+ [
+ ("https://gw.internal/ark/api/v3", "https://gw.internal/ark"),
+ ("https://gw.internal/api/ark/api/v3", "https://gw.internal/api/ark"),
+ ],
+)
+def test_a_path_prefix_before_the_api_marker_survives(api_base: str, expected: str):
+ assert VolcEngineVideoConfig().get_complete_url("m", api_base, {}) == expected
+
+
+def test_a_missing_api_key_raises(monkeypatch: pytest.MonkeyPatch):
+ monkeypatch.delenv("ARK_API_KEY", raising=False)
+ monkeypatch.delenv("VOLCENGINE_API_KEY", raising=False)
+
+ with pytest.raises(ValueError, match="Volcengine API key is required"):
+ _ = VolcEngineVideoConfig().validate_environment({}, SEEDANCE_2_5)
+
+
+def test_volcengine_resolves_to_the_video_config():
+ config = ProviderConfigManager.get_provider_video_config(
+ model=SEEDANCE_2_5, provider=litellm.LlmProviders.VOLCENGINE
+ )
+
+ assert type(config) is VolcEngineVideoConfig
+
+
+ARK_VIDEO_TIERS: Final = {
+ SEEDANCE_2_5: {"480p": ((854, 480), 70), "720p": ((1280, 720), 70), "1080p": ((1920, 1080), 77)},
+ SEEDANCE_2_0: {"480p": ((864, 496), 46), "720p": ((1280, 720), 46), "1080p": ((1920, 1080), 51), "4k": ((3840, 2160), 26)},
+ SEEDANCE_2_0_FAST: {"480p": ((864, 496), 37), "720p": ((1280, 720), 37)},
+ "doubao-seedance-2-0-mini-260615": {"480p": ((864, 496), 23), "720p": ((1280, 720), 23)},
+}
+CNY_PER_USD: Final = 7.0
+
+
+@pytest.mark.usefixtures("local_model_cost_map")
+@pytest.mark.parametrize("model", sorted(ARK_VIDEO_TIERS))
+def test_each_tier_rate_is_its_rendered_pixels_at_arks_price_for_that_tier(model: str):
+ """Ark prices 1080p and 4k per token above and below the 480p/720p rate, so a tier is not the base
+ rate scaled by pixels. Prices and frames are Ark's published per-resolution table, read 2026-10-16
+ from https://www.volcengine.com/docs/82379/1544106, whose own 2.0 examples agree: 1080p 16:9 5s bills
+ 2.48 CNY/s and 4k bills 5.05 CNY/s."""
+ info = litellm.get_model_info(model=model, custom_llm_provider="volcengine")
+ frames = info["provider_specific_entry"]["volcengine_video_frames"]
+
+ assert {tier: info[f"output_cost_per_second_{tier}"] for tier in ARK_VIDEO_TIERS[model]} == {
+ tier: pytest.approx(width * height * 24 / 1024 * cny_per_million / 1e6 / CNY_PER_USD, rel=1e-5)
+ for tier, ((width, height), cny_per_million) in ARK_VIDEO_TIERS[model].items()
+ }
+
+ assert {
+ tier: tuple(int(side) for side in frames[tier]["16:9"].split("x"))
+ for tier in ARK_VIDEO_TIERS[model]
+ if tier in frames
+ } == {tier: frame for tier, (frame, _) in ARK_VIDEO_TIERS[model].items() if tier in frames}
+
+
+async def _wire_body(**kwargs: object) -> dict[str, object]:
+ sent: Final[list[httpx.Request]] = []
+
+ def respond(request: httpx.Request) -> httpx.Response:
+ sent.append(request)
+ return httpx.Response(200, json={"id": TASK_ID})
+
+ async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client:
+ handler: Final = AsyncHTTPHandler()
+ await handler.close()
+ handler.client = http_client
+ _ = await avideo_generation(api_key="sk-test", api_base=API_BASE, client=handler, **kwargs)
+
+ (request,) = sent
+ assert (str(request.url), request.headers["authorization"]) == (TASKS_URL, "Bearer sk-test")
+ return json.loads(request.content)
+
+
+@pytest.mark.asyncio
+async def test_seconds_size_and_extra_body_content_reach_the_wire():
+ body = await _wire_body(
+ model=f"volcengine/{SEEDANCE_2_5}",
+ prompt="use these",
+ seconds="4",
+ size="854x480",
+ extra_body={
+ "content": [{"type": "image_url", "image_url": {"url": "https://x/a.png"}, "role": "reference_image"}]
+ },
+ )
+
+ assert body == {
+ "model": SEEDANCE_2_5,
+ "content": [
+ {"type": "text", "text": "use these"},
+ {"type": "image_url", "image_url": {"url": "https://x/a.png"}, "role": "reference_image"},
+ ],
+ "ratio": "16:9",
+ "resolution": "480p",
+ "duration": 4,
+ }
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "fields",
+ [{"parameters": {"seed": 3}}, {"extra_body": {"parameters": {"seed": 3}}}],
+ ids=["top_level", "extra_body"],
+)
+async def test_parameters_reach_the_wire_unpacked_whether_sent_top_level_or_in_extra_body(fields: dict[str, object]):
+ body = await _wire_body(model=f"volcengine/{SEEDANCE_2_5}", prompt="p", seconds="4", **fields)
+
+ assert body == {
+ "model": SEEDANCE_2_5,
+ "content": [{"type": "text", "text": "p"}],
+ "resolution": "720p",
+ "duration": 4,
+ "seed": 3,
+ }
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "extra_body,expected",
+ [
+ ({"size": "1920x1080"}, {"resolution": "1080p"}),
+ ({"seconds": "8"}, {"duration": 8}),
+ ({"user": "someone"}, {"safety_identifier": "someone"}),
+ ],
+ ids=["size", "seconds", "user"],
+)
+async def test_an_extra_body_field_litellm_maps_itself_is_honoured_and_its_leftover_copy_ignored(
+ extra_body: dict[str, object], expected: dict[str, object]
+):
+ """The shared request builder copies extra_body onto the top-level params, where this adapter maps it, and then
+ merges the raw extra_body back over the mapped params, returning the same key a second time as a field of its
+ own. The second copy has to be ignored rather than rejected as a field litellm never sends."""
+ body = await _wire_body(
+ model=f"volcengine/{SEEDANCE_2_5}",
+ prompt="p",
+ seconds="4",
+ size="1280x720",
+ extra_body=extra_body,
+ )
+
+ assert body == {
+ "model": SEEDANCE_2_5,
+ "content": [{"type": "text", "text": "p"}],
+ "ratio": "16:9",
+ "resolution": "720p",
+ "duration": 4,
+ **expected,
+ }
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "fields,misplaced",
+ [
+ ({"parameters": {"seconds": "8"}}, "seconds"),
+ ({"extra_body": {"parameters": {"seconds": "8"}}}, "seconds"),
+ ({"extra_body": {"parameters": {"size": "1920x1080"}}}, "size"),
+ ({"extra_body": {"parameters": {"user": "someone"}}}, "user"),
+ ({"extra_body": {"parameters": {"prompt": "a dog"}}}, "prompt"),
+ ({"extra_body": {"parameters": {"parameters": {"seed": 3}}}}, "parameters"),
+ ],
+ ids=["top_level_seconds", "seconds", "size", "user", "prompt", "nested_parameters"],
+)
+async def test_a_request_field_litellm_maps_itself_is_rejected_inside_parameters(
+ fields: dict[str, object], misplaced: str
+):
+ with pytest.raises(UnsupportedParamsError, match=f"send {misplaced} as top-level request fields"):
+ _ = await _wire_body(model=f"volcengine/{SEEDANCE_2_5}", prompt="p", **fields)
+
+
+@pytest.mark.asyncio
+async def test_a_request_field_inside_parameters_is_dropped_under_drop_params_and_the_rest_is_sent():
+ body = await _wire_body(
+ model=f"volcengine/{SEEDANCE_2_5}",
+ prompt="p",
+ drop_params=True,
+ extra_body={"parameters": {"seconds": "8", "seed": 3}},
+ )
+
+ assert body == {
+ "model": SEEDANCE_2_5,
+ "content": [{"type": "text", "text": "p"}],
+ "resolution": "720p",
+ "duration": 4,
+ "seed": 3,
+ }
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "extra_body,field",
+ [
+ ({"prompt": "a dog"}, "prompt"),
+ ({"model": SEEDANCE_2_0_FAST}, "model"),
+ ({"extra_headers": {"x-a": "1"}}, "extra_headers"),
+ ],
+ ids=["prompt", "model", "extra_headers"],
+)
+async def test_an_extra_body_copy_of_a_request_field_litellm_does_not_map_is_rejected(
+ extra_body: dict[str, object], field: str
+):
+ with pytest.raises(UnsupportedParamsError, match=f"does not send {field} to Ark"):
+ _ = await _wire_body(model=f"volcengine/{SEEDANCE_2_5}", prompt="p", extra_body=extra_body)
+
+
+@pytest.mark.asyncio
+async def test_a_user_that_is_not_a_string_is_rejected_rather_than_ignored():
+ with pytest.raises(UnsupportedParamsError, match="user must be a string, got int"):
+ _ = await _wire_body(model=f"volcengine/{SEEDANCE_2_5}", prompt="p", extra_body={"user": 5})
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "control",
+ [{"drop_params": False}, {"organization": "org-1"}],
+ ids=["drop_params", "organization"],
+)
+async def test_a_litellm_control_field_is_not_mistaken_for_an_ark_field(control: dict[str, object]):
+ body = await _wire_body(model=f"volcengine/{SEEDANCE_2_5}", prompt="p", seconds="4", **control)
+
+ assert body == {
+ "model": SEEDANCE_2_5,
+ "content": [{"type": "text", "text": "p"}],
+ "resolution": "720p",
+ "duration": 4,
+ }
+
+
+@pytest.mark.asyncio
+async def test_additional_drop_params_removes_the_named_ark_fields_from_the_wire():
+ body = await _wire_body(
+ model=f"volcengine/{SEEDANCE_2_5}",
+ prompt="p",
+ seconds="4",
+ additional_drop_params=["seed"],
+ extra_body={"seed": 3, "watermark": False},
+ )
+
+ assert body == {
+ "model": SEEDANCE_2_5,
+ "content": [{"type": "text", "text": "p"}],
+ "resolution": "720p",
+ "duration": 4,
+ "watermark": False,
+ }
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ ("dropped", "request_fields"),
+ [
+ pytest.param("user", {"user": "u-1"}, id="user"),
+ pytest.param("seconds", {"seconds": "8"}, id="seconds"),
+ pytest.param("size", {"size": "1920x1080"}, id="size"),
+ pytest.param("parameters", {"parameters": {"seed": 3}}, id="parameters"),
+ pytest.param("parameters", {"extra_body": {"parameters": {"seed": 3}}}, id="extra_body_parameters"),
+ pytest.param("image", {"image": {"gcsUri": "gs://b/a.png"}}, id="unsupported_image"),
+ ],
+)
+async def test_additional_drop_params_naming_a_request_field_sends_what_omitting_it_sends(
+ dropped: str, request_fields: dict[str, object]
+):
+ request: Final = {"model": f"volcengine/{SEEDANCE_2_5}", "prompt": "p"}
+
+ with_drop = await _wire_body(**request, **request_fields, additional_drop_params=[dropped])
+
+ assert with_drop == await _wire_body(**request)
+
+
+@pytest.mark.asyncio
+async def test_an_extra_body_duration_of_minus_one_is_rejected_before_anything_is_sent():
+ with pytest.raises(UnsupportedParamsError, match="Ark-picked length"):
+ _ = await _wire_body(model=f"volcengine/{SEEDANCE_2_5}", prompt="p", extra_body={"duration": -1})
+
+
+ACCOUNT_A_HOST: Final = "https://ark-a.example"
+ACCOUNT_B_HOST: Final = "https://ark-b.example"
+
+
+def _two_account_router(**router_kwargs: object) -> Router:
+ return Router(
+ model_list=[
+ {
+ "model_name": "seedance-a",
+ "litellm_params": {
+ "model": f"volcengine/{SEEDANCE_2_5}",
+ "api_key": "key-a",
+ "api_base": ACCOUNT_A_HOST,
+ },
+ },
+ {
+ "model_name": "seedance-b",
+ "litellm_params": {
+ "model": f"volcengine/{SEEDANCE_2_5}",
+ "api_key": "key-b",
+ "api_base": ACCOUNT_B_HOST,
+ },
+ },
+ ],
+ **router_kwargs, # pyright: ignore[reportArgumentType] # forwarded Router options
+ )
+
+
+def _proxy_poll_model(router: Router, video_id: str) -> str | None:
+ decoded: Final = decode_video_id_with_provider(video_id)
+ requested: Final = router.resolve_model_name_from_model_id(
+ decoded.get("model_id"), decoded.get("custom_llm_provider")
+ )
+ served: Final = decoded.get("served_model")
+ return router.resolve_served_model_name(requested, served) if requested and served else requested
+
+
+@pytest.mark.parametrize("caller_metadata", [None, {"tags": ["proxy-request"]}])
+@respx.mock
+def test_a_video_is_polled_on_the_ark_account_that_created_it(caller_metadata: dict[str, object] | None):
+ router = _two_account_router()
+ _ = respx.post(f"{ACCOUNT_B_HOST}/api/v3/contents/generations/tasks").mock(
+ return_value=httpx.Response(200, json={"id": TASK_ID})
+ )
+ task_get = respx.get(f"{ACCOUNT_B_HOST}/api/v3/contents/generations/tasks/{TASK_ID}").mock(
+ return_value=httpx.Response(200, json=_task(status="running"))
+ )
+
+ created = router.video_generation(model="seedance-b", prompt="a cat", seconds="4", metadata=caller_metadata)
+ _ = router.video_status(model=_proxy_poll_model(router, created.id), video_id=created.id)
+
+ assert [call.request.headers["authorization"] for call in task_get.calls] == ["Bearer key-b"]
+
+
+@respx.mock
+def test_a_video_created_through_a_routing_group_is_polled_on_the_member_that_created_it():
+ router = _two_account_router(
+ routing_groups=[
+ {"group_name": "seedance", "models": ["seedance-a", "seedance-b"], "routing_strategy": "simple-shuffle"}
+ ]
+ )
+ for host in (ACCOUNT_A_HOST, ACCOUNT_B_HOST):
+ _ = respx.post(f"{host}/api/v3/contents/generations/tasks").mock(
+ return_value=httpx.Response(200, json={"id": TASK_ID})
+ )
+ _ = respx.get(f"{host}/api/v3/contents/generations/tasks/{TASK_ID}").mock(
+ return_value=httpx.Response(200, json=_task(status="running"))
+ )
+
+ for _ in range(8):
+ created = router.video_generation(model="seedance", prompt="a cat", seconds="4")
+ polled = router.video_status(model=_proxy_poll_model(router, created.id), video_id=created.id)
+ assert polled.id == created.id
+
+ hosts = [call.request.url.host for call in respx.calls]
+ assert hosts[1::2] == hosts[0::2]
diff --git a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py
index 69d182667ec..e4129ebbc39 100644
--- a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py
+++ b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py
@@ -440,8 +440,6 @@ class TestErrorLogCarriesCallId:
@pytest.mark.asyncio
async def test_count_tokens_failure_log_carries_callers_call_id(self, caplog: pytest.LogCaptureFixture):
- from fastapi import HTTPException
-
import litellm.proxy.anthropic_endpoints.endpoints as ep
import litellm.proxy.proxy_server as proxy_server
from litellm.proxy._types import UserAPIKeyAuth
@@ -460,11 +458,10 @@ class TestErrorLogCarriesCallId:
proxy_server, "count_request_tokens", new=AsyncMock(side_effect=RuntimeError("tokenizer down"))
), # test-quality-ok: module global imported at call time; the test targets the endpoint's except block
caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"),
- pytest.raises(HTTPException) as raised,
):
- await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth())
+ response = await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth())
- assert raised.value.status_code == 500
+ assert response.status_code == 500
record = self._error_record(caplog)
assert record.litellm_call_id == call_id
assert call_id in record.getMessage()
diff --git a/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py b/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py
index 451792109de..114bbb184c6 100644
--- a/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py
+++ b/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py
@@ -182,6 +182,29 @@ def test_google_count_tokens_unchanged():
assert body["totalTokens"] == 7
+def test_google_count_tokens_unexpected_failure_uses_the_google_error_envelope():
+ try:
+ client = _build_test_client()
+ except ImportError as e:
+ pytest.skip(f"Skipping test due to missing dependency: {e}")
+
+ with patch(
+ "litellm.proxy.proxy_server.token_counter",
+ new_callable=AsyncMock,
+ side_effect=RuntimeError("tokenizer unavailable"),
+ ):
+ response = client.post(
+ "/v1beta/models/test-model:countTokens",
+ json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]},
+ )
+
+ assert response.status_code == 500
+ error = response.json()["error"]
+ assert error["code"] == 500
+ assert error["status"] == "INTERNAL"
+ assert "tokenizer unavailable" not in error["message"]
+
+
def test_google_count_tokens_uses_normalized_total_when_provider_shape_differs():
"""A provider that counts in Anthropic shape (Bedrock, Anthropic) carries no totalTokens key, so the normalized total must win over the raw response's missing field."""
try:
diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/litellm_content_filter/test_content_filter.py b/tests/unit/proxy/guardrails/guardrail_hooks/litellm_content_filter/test_content_filter.py
index 5d706a65a79..abcbb714186 100644
--- a/tests/unit/proxy/guardrails/guardrail_hooks/litellm_content_filter/test_content_filter.py
+++ b/tests/unit/proxy/guardrails/guardrail_hooks/litellm_content_filter/test_content_filter.py
@@ -9,7 +9,11 @@ Tests 40 different sentences to validate the conditional matching logic:
"""
import importlib
+import json
import os
+from collections.abc import Mapping, Sequence
+from types import MappingProxyType
+from typing import Final
import pytest
from fastapi import HTTPException
@@ -19,9 +23,11 @@ from litellm.proxy.guardrails.content_filter_data import POLICY_TEMPLATES_DIR
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import (
ContentFilterGuardrail,
)
+from litellm.types.guardrails import BlockedWord, ContentFilterAction, ContentFilterPattern, GuardrailEventHooks
from litellm.types.proxy.guardrails.guardrail_hooks.litellm_content_filter import (
ContentFilterCategoryConfig,
)
+from litellm.types.utils import Choices, Message, ModelResponse
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
# Test cases: (sentence, expected_result, reason)
@@ -1793,3 +1799,237 @@ class TestSGPDPAPerformance:
if __name__ == "__main__":
pytest.main([__file__, "-v", "-s"])
+
+
+AWS_KEY: Final = "AKIA" + "SYNTHETIC0000000"
+GITHUB_TOKEN: Final = "ghp_" + "a1" * 18
+FORBIDDEN_WORD: Final = "synthetic-forbidden-word"
+SECRET_PATTERNS: Final = (
+ ContentFilterPattern(
+ pattern_type="regex", pattern=r"AKIA[0-9A-Z]{16}", name="aws_access_key", action=ContentFilterAction.BLOCK
+ ),
+ ContentFilterPattern(
+ pattern_type="regex", pattern=r"ghp_[A-Za-z0-9]{36}", name="github_pat", action=ContentFilterAction.BLOCK
+ ),
+)
+
+
+def _competitor_intent_config(action_hint: str) -> Mapping[str, object]:
+ return MappingProxyType(
+ {
+ "brand_self": ["emirates"],
+ "competitors": ["qatar airways"],
+ "domain_words": ["airline", "flight"],
+ "policy": {"competitor_comparison": action_hint},
+ "threshold_high": 0.5,
+ }
+ )
+
+
+def _observing_filter(
+ patterns: tuple[ContentFilterPattern, ...] = (),
+ categories: tuple[ContentFilterCategoryConfig, ...] = (),
+ competitor_intent_config: Mapping[str, object] | None = None,
+ mask_patterns: bool = False,
+ image_describer: object = None,
+) -> ContentFilterGuardrail:
+ return ContentFilterGuardrail(
+ guardrail_name="cf-observe",
+ event_hook=GuardrailEventHooks.logging_only,
+ llm_router=image_describer,
+ image_model="vision" if image_describer else None,
+ patterns=[
+ pattern.model_copy(update={"action": ContentFilterAction.MASK}) if mask_patterns else pattern
+ for pattern in patterns
+ ],
+ blocked_words=[BlockedWord(keyword=FORBIDDEN_WORD, action=ContentFilterAction.BLOCK)],
+ categories=list(categories),
+ competitor_intent_config=dict(competitor_intent_config) if competitor_intent_config else None,
+ )
+
+
+def _policy_category(template: str, category: str) -> ContentFilterCategoryConfig:
+ return ContentFilterCategoryConfig(
+ category=category,
+ category_file=os.path.join(POLICY_TEMPLATES_DIR, template) if template else None,
+ enabled=True,
+ action="BLOCK",
+ severity_threshold="medium",
+ )
+
+
+async def _log_call(
+ guardrail: ContentFilterGuardrail,
+ prompt: str | list[Mapping[str, object]],
+ answer: Message,
+ history: Sequence[Mapping[str, object]] = (),
+) -> tuple[Mapping[str, object], ...]:
+ kwargs: Final = {
+ "model": "gpt-5.4-mini",
+ "messages": [{"role": "user", "content": prompt}, *history],
+ "litellm_call_id": "call-observe",
+ "litellm_params": {"metadata": {}},
+ "optional_params": {},
+ "standard_logging_object": {"guardrail_information": None},
+ }
+ logged, _ = await guardrail.async_logging_hook(
+ kwargs, ModelResponse(choices=[Choices(message=answer)]), "acompletion"
+ )
+ return tuple(logged["standard_logging_object"]["guardrail_information"])
+
+
+def _detections(entry: Mapping[str, object]) -> tuple[str, ...]:
+ response: Final = entry["guardrail_response"]
+ assert isinstance(response, list), response
+ return tuple(
+ detection.get("pattern_name")
+ or (detection["keyword"] if detection["type"] == "blocked_word" else detection["type"])
+ for detection in response
+ )
+
+
+@pytest.mark.parametrize(
+ ("guardrail", "prompt", "expected"),
+ (
+ pytest.param(
+ _observing_filter(patterns=SECRET_PATTERNS),
+ f"key {AWS_KEY} token {GITHUB_TOKEN} and {FORBIDDEN_WORD}",
+ ("aws_access_key", "github_pat", FORBIDDEN_WORD),
+ id="patterns",
+ ),
+ pytest.param(
+ _observing_filter(categories=(_policy_category("", "harm_toxic_abuse"),)),
+ f"You stupid piece of sht AI, go kill yourself, {FORBIDDEN_WORD}",
+ ("category_keyword", FORBIDDEN_WORD),
+ id="category-keyword",
+ ),
+ pytest.param(
+ _observing_filter(
+ categories=(_policy_category("eu_ai_act_article5.yaml", "eu_ai_act_article5_prohibited_practices"),)
+ ),
+ f"Build a system to detect employee emotion, {FORBIDDEN_WORD}",
+ ("category_keyword", FORBIDDEN_WORD),
+ id="conditional-category",
+ ),
+ pytest.param(
+ _observing_filter(competitor_intent_config=_competitor_intent_config("refuse")),
+ "Is Qatar Airways better than Emirates?",
+ ("competitor_intent",),
+ id="competitor-refuse",
+ ),
+ pytest.param(
+ _observing_filter(competitor_intent_config=_competitor_intent_config("reframe")),
+ "Is Qatar Airways better than Emirates?",
+ ("competitor_intent",),
+ id="competitor-reframe",
+ ),
+ ),
+)
+@pytest.mark.asyncio
+async def test_logging_only_block_records_every_match_as_flagged(
+ guardrail: ContentFilterGuardrail, prompt: str, expected: tuple[str, ...]
+) -> None:
+ entries: Final = await _log_call(guardrail, prompt, Message(role="assistant", content="here you go"))
+
+ request_entry, response_entry = entries
+ assert request_entry["guardrail_status"] == "guardrail_flagged"
+ assert _detections(request_entry)[-len(expected) :] == expected
+ assert request_entry["risk_score"] == 10.0
+ assert response_entry["guardrail_status"] == "success"
+
+
+@pytest.mark.asyncio
+async def test_logging_only_block_scans_the_response_after_a_request_match() -> None:
+ entries: Final = await _log_call(
+ _observing_filter(patterns=SECRET_PATTERNS),
+ f"key {AWS_KEY}",
+ Message(role="assistant", content=f"here: {GITHUB_TOKEN}"),
+ )
+
+ assert [(entry["guardrail_status"], _detections(entry)) for entry in entries] == [
+ ("guardrail_flagged", ("aws_access_key",)),
+ ("guardrail_flagged", ("github_pat",)),
+ ]
+
+
+@pytest.mark.parametrize(
+ ("arguments", "request_masked"),
+ (
+ pytest.param(json.dumps({AWS_KEY: 1}), {"aws_access_key": 1}, id="masked-argument-key"),
+ pytest.param(json.dumps(json.loads("[" * 150 + "1" + "]" * 150)), {}, id="too-deep"),
+ ),
+)
+@pytest.mark.asyncio
+async def test_logging_only_flags_tool_arguments_enforcement_rejects_and_keeps_scanning(
+ arguments: str, request_masked: Mapping[str, int]
+) -> None:
+ history: Final = (
+ {
+ "role": "assistant",
+ "content": None,
+ "tool_calls": [
+ {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": arguments}}
+ ],
+ },
+ {"role": "tool", "tool_call_id": "call_1", "content": "done"},
+ )
+
+ entries: Final = await _log_call(
+ _observing_filter(patterns=SECRET_PATTERNS, mask_patterns=True),
+ "hi",
+ Message(role="assistant", content=f"here: {GITHUB_TOKEN}"),
+ history,
+ )
+
+ assert [(entry["guardrail_status"], entry["masked_entity_count"]) for entry in entries] == [
+ ("guardrail_flagged", request_masked),
+ ("success", {"github_pat": 1}),
+ ]
+
+
+@pytest.mark.asyncio
+async def test_logging_only_scan_leaves_enforcement_of_the_same_guardrail_blocking() -> None:
+ guardrail: Final = _observing_filter(patterns=SECRET_PATTERNS)
+ await _log_call(guardrail, f"key {AWS_KEY}", Message(role="assistant", content="ok"))
+ request_data: Final = {"metadata": {}}
+
+ with pytest.raises(HTTPException) as exc_info:
+ await guardrail.apply_guardrail(
+ inputs={"texts": [f"key {AWS_KEY} token {GITHUB_TOKEN}"]}, request_data=request_data, input_type="request"
+ )
+
+ assert exc_info.value.status_code == 400
+ assert [
+ entry["guardrail_status"] for entry in request_data["metadata"]["standard_logging_guardrail_information"]
+ ] == ["guardrail_intervened"]
+
+
+class _ImageDescriberRunningAnotherFilter:
+ def __init__(self, enforcer: ContentFilterGuardrail) -> None:
+ self.enforcer = enforcer
+ self.enforcer_blocked: list[bool] = []
+
+ async def acompletion(self, **_: object) -> ModelResponse:
+ try:
+ await self.enforcer.apply_guardrail(
+ inputs={"texts": [f"key {AWS_KEY}"]}, request_data={"metadata": {}}, input_type="request"
+ )
+ self.enforcer_blocked.append(False)
+ except HTTPException:
+ self.enforcer_blocked.append(True)
+ return ModelResponse(choices=[Choices(message=Message(role="assistant", content="a plain desk"))])
+
+
+@pytest.mark.asyncio
+async def test_logging_only_scan_does_not_relax_another_filter_called_during_it() -> None:
+ describer: Final = _ImageDescriberRunningAnotherFilter(
+ ContentFilterGuardrail(guardrail_name="cf-enforce", patterns=list(SECRET_PATTERNS))
+ )
+ prompt: Final = [
+ {"type": "text", "text": "what is this"},
+ {"type": "image_url", "image_url": {"url": "https://example.com/desk.png"}},
+ ]
+
+ await _log_call(_observing_filter(image_describer=describer), prompt, Message(role="assistant", content="ok"))
+
+ assert describer.enforcer_blocked == [True]
diff --git a/tests/unit/proxy/guardrails/test_guardrail_endpoints.py b/tests/unit/proxy/guardrails/test_guardrail_endpoints.py
index 4401e3e95c0..908a49946bb 100644
--- a/tests/unit/proxy/guardrails/test_guardrail_endpoints.py
+++ b/tests/unit/proxy/guardrails/test_guardrail_endpoints.py
@@ -41,6 +41,8 @@ from litellm.proxy.guardrails.content_filter_data import DATA_ROOTS
from litellm.proxy.guardrails.guardrail_endpoints import (
test_custom_code_guardrail as run_custom_code_test_endpoint,
)
+from litellm.proxy.guardrails.guardrail_hooks.azure.prompt_shield import AzureContentSafetyPromptShieldGuardrail
+from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail
MOCK_ADMIN_USER = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
from litellm.proxy.guardrails.guardrail_registry import (
@@ -3143,7 +3145,36 @@ async def test_ui_settings_map_matches_runtime_supported_event_hooks():
continue
assert provider in result.supported_modes_by_provider, provider
- assert result.supported_modes_by_provider[provider] == [hook.value for hook in declared], provider
+ offered = result.supported_modes_by_provider[provider]
+ assert offered[: len(declared)] == [hook.value for hook in declared], provider
+ assert set(offered) - {hook.value for hook in declared} <= {"logging_only"}, provider
+ assert offered.count("logging_only") == (
+ 1 if guardrail_class.accepts_logging_only_mode() or GuardrailEventHooks.logging_only in declared else 0
+ ), provider
+
+
+@pytest.mark.asyncio
+async def test_ui_settings_offer_logging_only_where_the_generic_scan_runs():
+ modes_by_provider = (await get_guardrail_ui_settings()).supported_modes_by_provider
+
+ assert modes_by_provider["litellm_content_filter"] == [
+ *(hook.value for hook in ContentFilterGuardrail.get_supported_event_hooks()),
+ "logging_only",
+ ]
+ assert modes_by_provider["akto"].count("logging_only") == 1
+ assert modes_by_provider["azure/prompt_shield"] == [
+ hook.value for hook in AzureContentSafetyPromptShieldGuardrail.get_supported_event_hooks()
+ ]
+ with pytest.raises(ValueError, match="not in the supported event hooks"):
+ AzureContentSafetyPromptShieldGuardrail(
+ guardrail_name="native-logging-only",
+ api_key="k",
+ api_base="https://example.invalid",
+ event_hook="logging_only",
+ )
+ assert (
+ ContentFilterGuardrail(guardrail_name="cf-logging-only", event_hook="logging_only").event_hook == "logging_only"
+ )
def test_content_filter_runtime_rejects_unsupported_mcp_hook():
diff --git a/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py b/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py
index 33904637366..c6369a1f3d7 100644
--- a/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py
+++ b/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py
@@ -20,7 +20,6 @@ from litellm.proxy.hooks.model_max_budget_limiter import (
model_budget_start_time_cache_key,
resolve_model_budget,
)
-from litellm.types.caching import RedisPipelineIncrementOperation
async def _reset_at_of_open_window(
@@ -500,16 +499,15 @@ async def test_async_log_success_event_pushes_redis_increments_when_redis_config
@pytest.mark.asyncio
-async def test_model_budget_limiter_initializes_redis_increment_queue_lock():
+async def test_model_budget_limiter_without_redis_tracks_spend_locally_and_queues_nothing():
dual_cache = DualCache()
limiter = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache)
spend_key = "virtual_key_spend:test-key:gpt-4:1d"
await limiter._increment_spend_in_current_window(spend_key=spend_key, response_cost=0.01, ttl=86400)
- assert limiter.redis_increment_operation_queue == [
- RedisPipelineIncrementOperation(key=spend_key, increment_value=0.01, ttl=86400)
- ]
+ assert await dual_cache.async_get_cache(spend_key) == 0.01
+ assert limiter.redis_increment_operation_queue == []
@pytest.mark.asyncio
diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py
index 1af9def3afb..57e0a4fc47c 100644
--- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py
+++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py
@@ -3407,6 +3407,43 @@ async def test_block_key_existing_key_succeeds(monkeypatch):
assert result == mock_updated_record
+@pytest.mark.asyncio
+@pytest.mark.parametrize("endpoint_name", ["block_key", "unblock_key"])
+async def test_block_and_unblock_write_no_audit_row_when_the_key_update_fails(monkeypatch, endpoint_name):
+ from litellm.proxy._types import BlockKeyRequest
+ from litellm.proxy.management_endpoints import key_management_endpoints
+ from litellm.proxy.management_helpers.audit_logs import drain_audit_tasks
+
+ hashed_token = "a1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd"
+ existing_key_record = MagicMock()
+ existing_key_record.token = hashed_token
+ existing_key_record.model_dump_json.return_value = f'{{"token": "{hashed_token}"}}'
+
+ mock_prisma_client = AsyncMock()
+ mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=existing_key_record)
+ mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock(side_effect=RuntimeError("update failed"))
+ audit_writer = AsyncMock()
+
+ monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+ monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())
+ monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
+ monkeypatch.setattr("litellm.proxy.proxy_server.create_audit_log_for_update", audit_writer)
+ monkeypatch.setattr("litellm.store_audit_logs", True)
+
+ with pytest.raises(RuntimeError, match="update failed"):
+ await getattr(key_management_endpoints, endpoint_name)(
+ data=BlockKeyRequest(key=hashed_token),
+ http_request=MagicMock(),
+ user_api_key_dict=UserAPIKeyAuth(
+ user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin_user"
+ ),
+ litellm_changed_by=None,
+ )
+ await drain_audit_tasks(timeout=1.0)
+
+ assert audit_writer.call_count == 0
+
+
@pytest.mark.asyncio
async def test_validate_key_team_change_with_member_permissions():
"""
diff --git a/tests/unit/proxy/management_helpers/test_audit_logs_proxy.py b/tests/unit/proxy/management_helpers/test_audit_logs_proxy.py
index 1ea608cfc17..325275ab391 100644
--- a/tests/unit/proxy/management_helpers/test_audit_logs_proxy.py
+++ b/tests/unit/proxy/management_helpers/test_audit_logs_proxy.py
@@ -1,14 +1,18 @@
import os
import traceback
from litellm._uuid import uuid
-from datetime import datetime
+from datetime import datetime, timezone
+from collections.abc import Coroutine
+from typing import Any, Final
from dotenv import load_dotenv
from fastapi import Request
from fastapi.routing import APIRoute
+import contextlib
import io
+import logging
import time
# this file is to test litellm/proxy
@@ -38,11 +42,16 @@ from starlette.datastructures import URL
from litellm.proxy.management_helpers.audit_logs import (
create_audit_log_for_update,
+ drain_audit_tasks,
get_audit_log_changed_by,
+ track_audit_task,
)
from litellm.proxy._types import LiteLLM_AuditLogs, LitellmTableNames, UserAPIKeyAuth
+from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
-from unittest.mock import patch, AsyncMock
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.types.utils import StandardAuditLogPayload
+from unittest.mock import AsyncMock, MagicMock, patch
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
import json
@@ -174,7 +183,6 @@ async def test_create_audit_log_for_update_premium_user():
patch("litellm.store_audit_logs", True),
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
):
-
mock_prisma.db.litellm_auditlog.create = AsyncMock()
request_data = LiteLLM_AuditLogs(
@@ -215,9 +223,7 @@ def prisma_client():
os.environ["DATABASE_URL"] = modified_url
# Assuming PrismaClient is a class that needs to be instantiated
- prisma_client = PrismaClient(
- database_url=os.environ["DATABASE_URL"], proxy_logging_obj=proxy_logging_obj
- )
+ prisma_client = PrismaClient(database_url=os.environ["DATABASE_URL"], proxy_logging_obj=proxy_logging_obj)
return prisma_client
@@ -252,10 +258,482 @@ async def test_create_audit_log_in_db(prisma_client):
await asyncio.sleep(1)
# now read the last log from the db
- last_log = await prisma_client.db.litellm_auditlog.find_first(
- where={"id": audit_log_id}
- )
+ last_log = await prisma_client.db.litellm_auditlog.find_first(where={"id": audit_log_id})
assert last_log.id == audit_log_id
setattr(litellm, "store_audit_logs", False)
+
+
+@pytest.mark.asyncio
+async def test_track_audit_task_lifecycle():
+ task_completed: Final = asyncio.Event()
+
+ async def _sample_coroutine() -> None:
+ await asyncio.sleep(0)
+ task_completed.set()
+
+ task: Final = track_audit_task(asyncio.create_task(_sample_coroutine()))
+ assert not task.done()
+
+ await drain_audit_tasks(timeout=1.0)
+ assert task_completed.is_set()
+ assert task.done()
+
+ completed_task: Final = asyncio.create_task(asyncio.sleep(0))
+ await completed_task
+ assert track_audit_task(completed_task) is completed_task
+
+
+@pytest.mark.asyncio
+async def test_drain_audit_tasks_waits_for_all_tasks():
+ event1: Final = asyncio.Event()
+ event2: Final = asyncio.Event()
+
+ async def _worker(event: asyncio.Event) -> None:
+ await asyncio.sleep(0)
+ event.set()
+
+ task1: Final = track_audit_task(asyncio.create_task(_worker(event1)))
+ task2: Final = track_audit_task(asyncio.create_task(_worker(event2)))
+
+ assert not task1.done()
+ assert not task2.done()
+
+ await drain_audit_tasks(timeout=1.0)
+
+ assert event1.is_set()
+ assert event2.is_set()
+ assert task1.done()
+ assert task2.done()
+
+
+@pytest.mark.asyncio
+async def test_drain_audit_tasks_handles_failing_task_cleanly():
+ async def _failing_worker() -> None:
+ await asyncio.sleep(0)
+ raise RuntimeError("test audit log write failure")
+
+ task: Final = track_audit_task(asyncio.create_task(_failing_worker()))
+ assert not task.done()
+
+ await drain_audit_tasks(timeout=1.0)
+ assert task.done()
+ assert isinstance(task.exception(), RuntimeError)
+
+
+@pytest.mark.asyncio
+async def test_drain_audit_tasks_timeout_does_not_cancel_pending_tasks():
+ async def _long_running_worker() -> None:
+ await asyncio.Event().wait()
+
+ task: Final = track_audit_task(asyncio.create_task(_long_running_worker()))
+ try:
+ assert not task.done()
+ await drain_audit_tasks(timeout=0)
+ assert not task.cancelled()
+ assert not task.done()
+ finally:
+ task.cancel()
+ try:
+ await task
+ except asyncio.CancelledError:
+ pass
+
+
+@pytest.mark.asyncio
+async def test_drain_audit_tasks_persists_management_audit_write_before_db_disconnect():
+ mock_prisma: Final = MagicMock()
+ writes: Final[list[dict[str, Any]]] = []
+ is_disconnected: Final = [False]
+
+ async def _mock_create(data: dict[str, Any]) -> None:
+ if is_disconnected[0]:
+ raise RuntimeError("Database already disconnected")
+ await asyncio.sleep(0)
+ writes.append(data)
+
+ mock_prisma.db.litellm_auditlog.create = AsyncMock(side_effect=_mock_create)
+
+ with (
+ patch("litellm.store_audit_logs", True),
+ patch("litellm.proxy.proxy_server.premium_user", True),
+ patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
+ ):
+ request_data: Final = LiteLLM_AuditLogs(
+ id=str(uuid.uuid4()),
+ updated_at=datetime.now(timezone.utc),
+ changed_by="admin-user",
+ changed_by_api_key="sk-admin",
+ table_name=LitellmTableNames.KEY_TABLE_NAME,
+ object_id="token-123",
+ action="blocked",
+ updated_values="{}",
+ before_value='{"key": "test"}',
+ )
+
+ task: Final = track_audit_task(asyncio.create_task(create_audit_log_for_update(request_data=request_data)))
+ assert not task.done()
+
+ await drain_audit_tasks(timeout=2.0)
+ is_disconnected[0] = True
+
+ assert task.done()
+ assert len(writes) == 1
+ assert writes[0]["object_id"] == "token-123"
+ assert writes[0]["action"] == "blocked"
+
+
+@pytest.mark.asyncio
+async def test_create_audit_log_auto_registers_in_drain():
+ mock_prisma: Final = MagicMock()
+ writes: Final[list[dict[str, Any]]] = []
+
+ async def _mock_create(data: dict[str, Any]) -> None:
+ await asyncio.sleep(0)
+ writes.append(data)
+
+ mock_prisma.db.litellm_auditlog.create = AsyncMock(side_effect=_mock_create)
+
+ with (
+ patch("litellm.store_audit_logs", True),
+ patch("litellm.proxy.proxy_server.premium_user", True),
+ patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
+ ):
+ request_data: Final = LiteLLM_AuditLogs(
+ id=str(uuid.uuid4()),
+ updated_at=datetime.now(timezone.utc),
+ changed_by="admin-user",
+ changed_by_api_key="sk-admin",
+ table_name=LitellmTableNames.KEY_TABLE_NAME,
+ object_id="token-456",
+ action="deleted",
+ updated_values="{}",
+ before_value='{"key": "test-deleted"}',
+ )
+
+ task: Final = asyncio.create_task(create_audit_log_for_update(request_data=request_data))
+ await asyncio.sleep(0)
+
+ await drain_audit_tasks(timeout=2.0)
+
+ assert task.done()
+ assert len(writes) == 1
+ assert writes[0]["object_id"] == "token-456"
+ assert writes[0]["action"] == "deleted"
+
+
+@pytest.mark.asyncio
+async def test_drain_audit_tasks_discards_timed_out_tasks_on_timeout():
+ async def _hung_worker() -> None:
+ await asyncio.Event().wait()
+
+ hung_task: Final = track_audit_task(asyncio.create_task(_hung_worker()))
+ try:
+ await drain_audit_tasks(timeout=0)
+ assert not hung_task.done()
+
+ second_drain: Final = asyncio.create_task(drain_audit_tasks(timeout=60))
+ for _ in range(3):
+ await asyncio.sleep(0)
+ assert second_drain.done()
+ finally:
+ hung_task.cancel()
+ try:
+ await hung_task
+ except asyncio.CancelledError:
+ pass
+
+
+@pytest.mark.asyncio
+async def test_drain_audit_tasks_captures_tasks_queued_during_drain():
+ finished_tasks: Final[list[str]] = []
+
+ async def _first_worker() -> None:
+ await asyncio.sleep(0)
+ track_audit_task(asyncio.create_task(_second_worker()))
+ finished_tasks.append("first")
+
+ async def _second_worker() -> None:
+ await asyncio.sleep(0)
+ finished_tasks.append("second")
+
+ track_audit_task(asyncio.create_task(_first_worker()))
+ await drain_audit_tasks(timeout=2.0)
+
+ assert finished_tasks == ["first", "second"]
+
+
+class _AuditCallbackBlockedUntilReleased(CustomLogger):
+ def __init__(self, release: asyncio.Event) -> None:
+ super().__init__()
+ self.release: Final = release
+ self.logged_object_ids: Final[list[str]] = []
+
+ async def async_log_audit_log_event(self, audit_log: StandardAuditLogPayload) -> None:
+ await self.release.wait()
+ self.logged_object_ids.append(audit_log["object_id"])
+
+
+class _WarningRecorder(logging.Handler):
+ def __init__(self) -> None:
+ super().__init__(level=logging.WARNING)
+ self.messages: Final[list[str]] = []
+
+ def emit(self, record: logging.LogRecord) -> None:
+ self.messages.append(record.getMessage())
+
+
+def _deleted_key_audit_log(object_id: str) -> LiteLLM_AuditLogs:
+ return LiteLLM_AuditLogs(
+ id=str(uuid.uuid4()),
+ updated_at=datetime.now(timezone.utc),
+ changed_by="admin-user",
+ table_name=LitellmTableNames.KEY_TABLE_NAME,
+ object_id=object_id,
+ action="deleted",
+ updated_values="{}",
+ before_value="{}",
+ )
+
+
+@pytest.mark.asyncio
+async def test_drain_waits_for_a_slow_audit_log_callback():
+ release: Final = asyncio.Event()
+ callback: Final = _AuditCallbackBlockedUntilReleased(release)
+
+ with (
+ patch("litellm.store_audit_logs", True),
+ patch("litellm.proxy.proxy_server.premium_user", True),
+ patch("litellm.proxy.proxy_server.prisma_client", None),
+ patch("litellm.audit_log_callbacks", [callback]),
+ ):
+ await create_audit_log_for_update(request_data=_deleted_key_audit_log("token-callback"))
+ drain: Final = asyncio.create_task(drain_audit_tasks(timeout=60))
+ for _ in range(3):
+ await asyncio.sleep(0)
+ drain_returned_before_the_callback_finished: Final = drain.done()
+ release.set()
+ await drain
+
+ assert not drain_returned_before_the_callback_finished
+ assert callback.logged_object_ids == ["token-callback"]
+
+
+@pytest.mark.asyncio
+async def test_drain_timeout_warning_names_each_dropped_audit_row():
+ mock_prisma: Final = MagicMock()
+ mock_prisma.db.litellm_auditlog.create = AsyncMock(side_effect=_io_that_never_finishes)
+ recorder: Final = _WarningRecorder()
+ verbose_proxy_logger.addHandler(recorder)
+
+ with (
+ patch("litellm.store_audit_logs", True),
+ patch("litellm.proxy.proxy_server.premium_user", True),
+ patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
+ ):
+ writer: Final = asyncio.create_task(
+ create_audit_log_for_update(request_data=_deleted_key_audit_log("token-dropped"))
+ )
+ try:
+ for _ in range(3):
+ await asyncio.sleep(0)
+ await drain_audit_tasks(timeout=0.05)
+ finally:
+ verbose_proxy_logger.removeHandler(recorder)
+ writer.cancel()
+ with contextlib.suppress(asyncio.CancelledError):
+ await writer
+
+ assert recorder.messages == [
+ "Timed out draining 1 audit tasks on shutdown, dropping: "
+ f"deleted {LitellmTableNames.KEY_TABLE_NAME.value} token-dropped to the database"
+ ]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ ("store_audit_logs", "premium_user", "expected_write_count"),
+ [(True, True, 1), (False, True, 0), (True, False, 0)],
+ ids=["row_written", "logging_disabled", "not_premium"],
+)
+async def test_drain_does_not_wait_for_code_that_runs_after_an_inline_audit_write(
+ store_audit_logs: bool, premium_user: bool, expected_write_count: int
+):
+ release_caller: Final = asyncio.Event()
+ mock_prisma: Final = MagicMock()
+ writes: Final[list[dict[str, Any]]] = []
+ mock_prisma.db.litellm_auditlog.create = AsyncMock(side_effect=lambda data: writes.append(data))
+
+ async def _caller_that_outlives_the_audit_call() -> None:
+ await create_audit_log_for_update(
+ request_data=LiteLLM_AuditLogs(
+ id=str(uuid.uuid4()),
+ updated_at=datetime.now(timezone.utc),
+ changed_by="test-user",
+ table_name=LitellmTableNames.KEY_TABLE_NAME,
+ object_id="test-obj",
+ action="updated",
+ updated_values="{}",
+ before_value="{}",
+ )
+ )
+ await release_caller.wait()
+
+ with (
+ patch("litellm.store_audit_logs", store_audit_logs),
+ patch("litellm.proxy.proxy_server.premium_user", premium_user),
+ patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
+ ):
+ caller: Final = asyncio.create_task(_caller_that_outlives_the_audit_call())
+ for _ in range(3):
+ await asyncio.sleep(0)
+ drain: Final = asyncio.create_task(drain_audit_tasks(timeout=60))
+ for _ in range(3):
+ await asyncio.sleep(0)
+ drain_returned_while_caller_was_running: Final = drain.done()
+ release_caller.set()
+ await caller
+ await drain
+
+ assert drain_returned_while_caller_was_running
+ assert len(writes) == expected_write_count
+
+
+async def _io_that_never_finishes(*args: object, **kwargs: object) -> None:
+ await asyncio.Event().wait()
+
+
+async def _drain_while_hook_io_hangs(hook: Coroutine[object, object, None]) -> None:
+ hook_task: Final = asyncio.create_task(hook)
+ for _ in range(3):
+ await asyncio.sleep(0)
+ await drain_audit_tasks(timeout=2.0)
+ hook_task.cancel()
+
+
+@pytest.mark.asyncio
+async def test_key_generated_audit_write_is_drained_while_the_invite_email_hangs():
+ from litellm.proxy._types import GenerateKeyRequest, GenerateKeyResponse
+ from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks
+
+ mock_prisma: Final = MagicMock()
+ writes: Final[list[dict[str, Any]]] = []
+ mock_prisma.db.litellm_auditlog.create = AsyncMock(side_effect=lambda data: writes.append(data))
+
+ with (
+ patch("litellm.store_audit_logs", True),
+ patch("litellm.proxy.proxy_server.premium_user", True),
+ patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
+ patch.object(KeyManagementEventHooks, "_send_key_created_email", side_effect=_io_that_never_finishes),
+ ):
+ await _drain_while_hook_io_hangs(
+ KeyManagementEventHooks.async_key_generated_hook(
+ data=GenerateKeyRequest(send_invite_email=True),
+ response=GenerateKeyResponse(key="sk-new", token_id="hashed-new"),
+ user_api_key_dict=UserAPIKeyAuth(api_key="sk-admin", user_id="admin"),
+ )
+ )
+
+ assert [(write["object_id"], write["action"]) for write in writes] == [("hashed-new", "created")]
+
+
+@pytest.mark.asyncio
+async def test_key_updated_audit_write_is_drained_while_the_secret_manager_rename_hangs():
+ from litellm.proxy._types import LiteLLM_VerificationToken, UpdateKeyRequest
+ from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks
+
+ mock_prisma: Final = MagicMock()
+ writes: Final[list[dict[str, Any]]] = []
+ mock_prisma.db.litellm_auditlog.create = AsyncMock(side_effect=lambda data: writes.append(data))
+
+ with (
+ patch("litellm.store_audit_logs", True),
+ patch("litellm.proxy.proxy_server.premium_user", True),
+ patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
+ patch.object(
+ KeyManagementEventHooks, "_rename_virtual_key_in_secret_manager", side_effect=_io_that_never_finishes
+ ),
+ ):
+ await _drain_while_hook_io_hangs(
+ KeyManagementEventHooks.async_key_updated_hook(
+ data=UpdateKeyRequest(key="sk-existing", key_alias="new-alias"),
+ existing_key_row=LiteLLM_VerificationToken(token="hashed-existing", key_alias="old-alias"),
+ response=None,
+ user_api_key_dict=UserAPIKeyAuth(api_key="sk-admin", user_id="admin"),
+ )
+ )
+
+ assert [(write["object_id"], write["action"]) for write in writes] == [(hash_token("sk-existing"), "updated")]
+
+
+@pytest.mark.asyncio
+async def test_key_rotated_audit_write_is_drained_while_the_secret_manager_rotation_hangs():
+ from litellm.proxy._types import GenerateKeyResponse, LiteLLM_VerificationToken
+ from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks
+
+ mock_prisma: Final = MagicMock()
+ writes: Final[list[dict[str, Any]]] = []
+ mock_prisma.db.litellm_auditlog.create = AsyncMock(side_effect=lambda data: writes.append(data))
+
+ with (
+ patch("litellm.store_audit_logs", True),
+ patch("litellm.proxy.proxy_server.premium_user", True),
+ patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
+ patch.object(
+ KeyManagementEventHooks, "_rotate_virtual_key_in_secret_manager", side_effect=_io_that_never_finishes
+ ),
+ ):
+ await _drain_while_hook_io_hangs(
+ KeyManagementEventHooks.async_key_rotated_hook(
+ data=None,
+ existing_key_row=LiteLLM_VerificationToken(token="hashed-old"),
+ response=GenerateKeyResponse(key="sk-rotated", token_id="hashed-new"),
+ user_api_key_dict=UserAPIKeyAuth(api_key="sk-admin", user_id="admin"),
+ )
+ )
+
+ assert [(write["object_id"], write["action"]) for write in writes] == [("hashed-old", "rotated")]
+
+
+@pytest.mark.asyncio
+async def test_user_created_audit_write_is_drained_while_its_user_read_is_still_running():
+ from litellm.proxy._types import NewUserRequest, NewUserResponse
+ from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks
+
+ mock_prisma: Final = MagicMock()
+ writes: Final[list[dict[str, Any]]] = []
+ release_user_read: Final = asyncio.Event()
+
+ async def _slow_user_read(*args: object, **kwargs: object) -> dict[str, str]:
+ await release_user_read.wait()
+ return {"user_id": "new-user"}
+
+ mock_prisma.db.litellm_auditlog.create = AsyncMock(side_effect=lambda data: writes.append(data))
+ mock_prisma.db.litellm_usertable.find_unique = AsyncMock(side_effect=_slow_user_read)
+
+ with (
+ patch("litellm.store_audit_logs", True),
+ patch("litellm.proxy.proxy_server.premium_user", True),
+ patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
+ patch.object(UserManagementEventHooks, "async_send_user_invitation_email", side_effect=_io_that_never_finishes),
+ ):
+ hook: Final = asyncio.create_task(
+ UserManagementEventHooks.async_user_created_hook(
+ data=NewUserRequest(user_id="new-user", send_invite_email=True),
+ response=NewUserResponse(key="sk-new", user_id="new-user"),
+ user_api_key_dict=UserAPIKeyAuth(api_key="sk-admin", user_id="admin"),
+ )
+ )
+ for _ in range(3):
+ await asyncio.sleep(0)
+ drain: Final = asyncio.create_task(drain_audit_tasks(timeout=60))
+ for _ in range(3):
+ await asyncio.sleep(0)
+ drain_waited_for_the_user_read: Final = not drain.done()
+ release_user_read.set()
+ await drain
+ hook.cancel()
+
+ assert drain_waited_for_the_user_read
+ assert [(write["object_id"], write["action"]) for write in writes] == [("new-user", "created")]
diff --git a/tests/unit/proxy/proxy_server/test_lifecycle.py b/tests/unit/proxy/proxy_server/test_lifecycle.py
index 69f9c72c9c8..71daa3bb873 100644
--- a/tests/unit/proxy/proxy_server/test_lifecycle.py
+++ b/tests/unit/proxy/proxy_server/test_lifecycle.py
@@ -17,12 +17,14 @@ Pins covered:
from __future__ import annotations
+import asyncio
import inspect
import json
import logging
import os
import subprocess
from collections.abc import Awaitable, Callable
+from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@@ -48,6 +50,8 @@ from litellm.proxy.proxy_server import (
save_worker_config,
)
+from litellm.proxy.management_helpers.audit_logs import track_audit_task
+
from .conftest import normalize
# ---------------------------------------------------------------------------
@@ -223,6 +227,46 @@ async def test_proxy_shutdown_drains_gateway_requests_before_disconnecting(monke
assert calls == ["flush", "disconnect"]
+@pytest.mark.asyncio
+async def test_proxy_shutdown_drains_audit_writes_beside_the_gateway_flush_before_disconnecting(monkeypatch):
+ """
+ A management call returns before its audit row is written, so shutdown waits
+ for that write while prisma is still connected. It waits beside the gateway
+ flush rather than before it, so the two waits overlap inside the teardown budget.
+ """
+ calls: list[str] = [] # mutable-ok: records call order, which is the assertion
+ gateway_flush_started: Final = asyncio.Event()
+
+ fake_prisma = MagicMock()
+ fake_prisma.disconnect = AsyncMock(side_effect=lambda: calls.append("disconnect"))
+ monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
+
+ async def _record_flush(client, accumulator):
+ gateway_flush_started.set()
+ calls.append("flush")
+
+ async def _audit_write_that_finishes_after_the_flush_starts() -> None:
+ await gateway_flush_started.wait()
+ calls.append("audit")
+
+ monkeypatch.setattr(ps, "flush_gateway_requests", _record_flush, raising=False)
+
+ fake_jwt = MagicMock()
+ fake_jwt.close = AsyncMock()
+ monkeypatch.setattr(ps, "jwt_handler", fake_jwt, raising=False)
+ monkeypatch.setattr(ps, "db_writer_client", None, raising=False)
+
+ import litellm
+
+ monkeypatch.setattr(litellm, "cache", None, raising=False)
+ monkeypatch.setattr(litellm, "success_callback", [], raising=False)
+
+ track_audit_task(asyncio.create_task(_audit_write_that_finishes_after_the_flush_starts()))
+ await asyncio.wait_for(proxy_shutdown_event(), timeout=2.0)
+
+ assert calls == ["flush", "audit", "disconnect"]
+
+
@pytest.mark.asyncio
async def test_proxy_shutdown_skips_gateway_flush_without_a_database(monkeypatch):
"""No prisma client means nothing to drain to, and no attempt is made."""
@@ -245,6 +289,34 @@ async def test_proxy_shutdown_skips_gateway_flush_without_a_database(monkeypatch
assert flush.await_count == 0
+@pytest.mark.asyncio
+async def test_proxy_shutdown_without_a_database_drains_audit_callbacks_before_the_rest_of_teardown(monkeypatch):
+ calls: list[str] = [] # mutable-ok: records call order, which is the assertion
+ monkeypatch.setattr(ps, "prisma_client", None, raising=False)
+ monkeypatch.setattr(
+ ps, "drain_post_call_redis_batches", AsyncMock(side_effect=lambda: calls.append("redis")), raising=False
+ )
+
+ fake_jwt = MagicMock()
+ fake_jwt.close = AsyncMock()
+ monkeypatch.setattr(ps, "jwt_handler", fake_jwt, raising=False)
+ monkeypatch.setattr(ps, "db_writer_client", None, raising=False)
+
+ import litellm
+
+ monkeypatch.setattr(litellm, "cache", None, raising=False)
+ monkeypatch.setattr(litellm, "success_callback", [], raising=False)
+
+ async def _audit_callback_still_sending() -> None:
+ await asyncio.sleep(0.05)
+ calls.append("audit")
+
+ track_audit_task(asyncio.create_task(_audit_callback_still_sending()))
+ await asyncio.wait_for(proxy_shutdown_event(), timeout=2.0)
+
+ assert calls == ["audit", "redis"]
+
+
@pytest.mark.asyncio
async def test_proxy_shutdown_event_prisma_disconnect_raises_error(monkeypatch):
fake_prisma = MagicMock()
diff --git a/tests/unit/proxy/proxy_server/test_routes_config.py b/tests/unit/proxy/proxy_server/test_routes_config.py
index 4c8265b82c3..2d86ff6a7c2 100644
--- a/tests/unit/proxy/proxy_server/test_routes_config.py
+++ b/tests/unit/proxy/proxy_server/test_routes_config.py
@@ -14,10 +14,15 @@ Routes covered:
from __future__ import annotations
import asyncio
+import contextlib
import json
+from collections.abc import Callable
+from typing import Final
from unittest.mock import AsyncMock, MagicMock
+import httpx
import pytest
+from fastapi.testclient import TestClient
from .conftest import VOLATILE_KEYS, normalize
@@ -944,6 +949,228 @@ def test_config_callback_delete_not_found(client, auth_as, mock_prisma, monkeypa
assert "langfuse" in str(response.json()).lower() or "not found" in str(response.json()).lower()
+def _delete_callback_roundtrip(
+ client: TestClient,
+ auth_as: Callable[..., contextlib.AbstractContextManager[None]],
+ mock_prisma: MagicMock,
+ monkeypatch: pytest.MonkeyPatch,
+ litellm_settings: dict[str, object],
+ callback_name: str,
+ owned_by_config: frozenset[str] = frozenset(),
+) -> tuple[httpx.Response, dict[str, object]]:
+ """POST /config/callback/delete as admin against a stubbed config, and
+ return (response, saved litellm_settings) so each case pins the full body
+ plus the exact config persisted."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ _install_litellm_config(mock_prisma)
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+ monkeypatch.setattr(ps, "store_model_in_db", True)
+
+ fake_proxy_config: Final = MagicMock()
+ fake_proxy_config.get_config = AsyncMock(return_value={"litellm_settings": litellm_settings})
+ fake_proxy_config.save_config = AsyncMock()
+ fake_proxy_config.add_deployment = AsyncMock()
+ fake_proxy_config.litellm_settings.owned_by_config = owned_by_config.__contains__
+ monkeypatch.setattr(ps, "proxy_config", fake_proxy_config)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response: Final = client.post("/config/callback/delete", json={"callback_name": callback_name})
+ saved: Final = fake_proxy_config.save_config.await_args.kwargs["new_config"]["litellm_settings"]
+ return response, saved
+
+
+def test_config_callback_delete_from_failure_callback(
+ client: TestClient,
+ auth_as: Callable[..., contextlib.AbstractContextManager[None]],
+ mock_prisma: MagicMock,
+ monkeypatch: pytest.MonkeyPatch,
+):
+ """A callback configured only under litellm_settings.failure_callback is
+ listed by GET /get/config/callbacks, so the delete route must remove it."""
+ response, saved = _delete_callback_roundtrip(
+ client,
+ auth_as,
+ mock_prisma,
+ monkeypatch,
+ {"success_callback": ["slack"], "failure_callback": ["datadog", "sentry"]},
+ "datadog",
+ )
+ assert response.status_code == 200
+ assert response.json()["removed_callback"] == "datadog"
+ assert saved == {"success_callback": ["slack"], "failure_callback": ["sentry"]}
+ assert response.json()["remaining_callbacks"] == ["sentry"]
+
+
+def test_config_callback_delete_from_callbacks_list(
+ client: TestClient,
+ auth_as: Callable[..., contextlib.AbstractContextManager[None]],
+ mock_prisma: MagicMock,
+ monkeypatch: pytest.MonkeyPatch,
+):
+ """A callback configured only under litellm_settings.callbacks is deleted
+ while unrelated lists pass through untouched."""
+ response, saved = _delete_callback_roundtrip(
+ client,
+ auth_as,
+ mock_prisma,
+ monkeypatch,
+ {"success_callback": ["langfuse"], "callbacks": ["datadog"], "drop_params": True},
+ "datadog",
+ )
+ assert response.status_code == 200
+ assert response.json()["removed_callback"] == "datadog"
+ assert saved == {"success_callback": ["langfuse"], "callbacks": [], "drop_params": True}
+ assert response.json()["remaining_callbacks"] == []
+
+
+def test_config_callback_delete_from_multiple_lists(
+ client: TestClient,
+ auth_as: Callable[..., contextlib.AbstractContextManager[None]],
+ mock_prisma: MagicMock,
+ monkeypatch: pytest.MonkeyPatch,
+):
+ """A name present in both success_callback and callbacks is removed from
+ both, and remaining_callbacks reports the leftovers of the lists it left."""
+ response, saved = _delete_callback_roundtrip(
+ client,
+ auth_as,
+ mock_prisma,
+ monkeypatch,
+ {
+ "success_callback": ["datadog", "slack"],
+ "failure_callback": ["sentry"],
+ "callbacks": ["langfuse", "datadog"],
+ },
+ "datadog",
+ )
+ assert response.status_code == 200
+ assert saved == {
+ "success_callback": ["slack"],
+ "failure_callback": ["sentry"],
+ "callbacks": ["langfuse"],
+ }
+ assert response.json()["remaining_callbacks"] == ["slack", "langfuse"]
+
+
+def test_config_callback_delete_skips_lists_the_config_file_owns(
+ client: TestClient,
+ auth_as: Callable[..., contextlib.AbstractContextManager[None]],
+ mock_prisma: MagicMock,
+ monkeypatch: pytest.MonkeyPatch,
+):
+ """A name in a database-owned list and a config-file-owned list is removed
+ from the database-owned one only, so the save is not refused."""
+ response, saved = _delete_callback_roundtrip(
+ client,
+ auth_as,
+ mock_prisma,
+ monkeypatch,
+ {"success_callback": ["datadog", "slack"], "failure_callback": ["datadog"]},
+ "datadog",
+ owned_by_config=frozenset({"failure_callback"}),
+ )
+ assert response.status_code == 200
+ assert saved == {"success_callback": ["slack"], "failure_callback": ["datadog"]}
+ assert response.json()["remaining_callbacks"] == ["slack"]
+
+
+def test_config_callback_delete_removes_database_copy_when_config_file_owns_success_callback(
+ client: TestClient,
+ auth_as: Callable[..., contextlib.AbstractContextManager[None]],
+ mock_prisma: MagicMock,
+ monkeypatch: pytest.MonkeyPatch,
+):
+ """A name the config file sets in success_callback is left there while its
+ database-owned failure_callback copy is removed."""
+ response, saved = _delete_callback_roundtrip(
+ client,
+ auth_as,
+ mock_prisma,
+ monkeypatch,
+ {"success_callback": ["datadog"], "failure_callback": ["datadog", "sentry"]},
+ "datadog",
+ owned_by_config=frozenset({"success_callback"}),
+ )
+ assert response.status_code == 200
+ assert saved == {"success_callback": ["datadog"], "failure_callback": ["sentry"]}
+ assert response.json()["remaining_callbacks"] == ["sentry"]
+
+
+def test_config_callback_delete_from_string_valued_list(
+ client: TestClient,
+ auth_as: Callable[..., contextlib.AbstractContextManager[None]],
+ mock_prisma: MagicMock,
+ monkeypatch: pytest.MonkeyPatch,
+):
+ """A list written as a single string is matched and saved back as an empty list."""
+ response, saved = _delete_callback_roundtrip(
+ client,
+ auth_as,
+ mock_prisma,
+ monkeypatch,
+ {"failure_callback": "datadog"},
+ "datadog",
+ )
+ assert response.status_code == 200
+ assert saved == {"failure_callback": []}
+
+
+def test_config_callback_delete_case_insensitive(
+ client: TestClient,
+ auth_as: Callable[..., contextlib.AbstractContextManager[None]],
+ mock_prisma: MagicMock,
+ monkeypatch: pytest.MonkeyPatch,
+):
+ """A mixed-case entry is deleted by the name GET /get/config/callbacks shows for it."""
+ response, saved = _delete_callback_roundtrip(
+ client,
+ auth_as,
+ mock_prisma,
+ monkeypatch,
+ {"callbacks": ["DataDog", "otel"]},
+ "DataDog",
+ )
+ assert response.status_code == 200
+ assert response.json()["removed_callback"] == "datadog"
+ assert saved == {"callbacks": ["otel"]}
+
+
+def test_config_callback_delete_saves_against_the_baseline_it_read(
+ client: TestClient,
+ auth_as: Callable[..., contextlib.AbstractContextManager[None]],
+ mock_prisma: MagicMock,
+ monkeypatch: pytest.MonkeyPatch,
+):
+ """save_config gets the object get_config returned, so it diffs against the
+ settings as they were read and writes only the edited callback list."""
+ from litellm.proxy import proxy_server as ps
+ from litellm.proxy._types import LitellmUserRoles
+
+ _install_litellm_config(mock_prisma)
+ monkeypatch.setattr(ps, "prisma_client", mock_prisma)
+ monkeypatch.setattr(ps, "store_model_in_db", True)
+
+ read_time: Final = {"litellm_settings": {"failure_callback": ["datadog", "sentry"], "max_budget": 10}}
+ loaded: Final = ps._ConfigWithBaseline(read_time)
+ fake_proxy_config: Final = MagicMock()
+ fake_proxy_config.get_config = AsyncMock(return_value=loaded)
+ fake_proxy_config.save_config = AsyncMock()
+ fake_proxy_config.add_deployment = AsyncMock()
+ fake_proxy_config.litellm_settings.owned_by_config = frozenset[str]().__contains__
+ monkeypatch.setattr(ps, "proxy_config", fake_proxy_config)
+
+ with auth_as(LitellmUserRoles.PROXY_ADMIN):
+ response: Final = client.post("/config/callback/delete", json={"callback_name": "datadog"})
+
+ assert response.status_code == 200
+ saved: Final = fake_proxy_config.save_config.await_args.kwargs["new_config"]
+ assert saved is loaded
+ assert saved.baseline == read_time
+ assert saved["litellm_settings"] == {"failure_callback": ["sentry"], "max_budget": 10}
+
+
# ---------------------------------------------------------------------------
# GET /get/config/callbacks
# ---------------------------------------------------------------------------
diff --git a/tests/unit/proxy/proxy_server/test_routes_utils.py b/tests/unit/proxy/proxy_server/test_routes_utils.py
index 295b9e5ae8b..118b80da31e 100644
--- a/tests/unit/proxy/proxy_server/test_routes_utils.py
+++ b/tests/unit/proxy/proxy_server/test_routes_utils.py
@@ -390,8 +390,7 @@ def test_token_counter_fallback_counts_tools_system_and_anthropic_blocks(client,
assert full > bare
-def test_token_counter_fallback_prompt_with_tools_does_not_500(client, auth_as, monkeypatch):
- """Regression: a ``prompt`` request carrying ``tools`` but no ``messages`` still counts, because the fallback attaches tools only when counting messages (``token_counter`` rejects tools on the text path)."""
+def test_token_counter_fallback_counts_a_prompt_with_tools_like_its_messages_form(client, auth_as, monkeypatch):
monkeypatch.setattr(proxy_server, "llm_router", None)
monkeypatch.setattr(litellm, "disable_token_counter", False, raising=False)
prompt = "count the tokens in this sentence please"
@@ -412,4 +411,6 @@ def test_token_counter_fallback_prompt_with_tools_does_not_500(client, auth_as,
)
assert response.status_code == 200, response.text
- assert response.json()["total_tokens"] == litellm.token_counter(model="claude-fable-5", text=prompt)
+ assert response.json()["total_tokens"] == litellm.token_counter(
+ model="claude-fable-5", messages=[{"role": "user", "content": prompt}], tools=tools
+ )
diff --git a/tests/unit/proxy/spend_tracking/test_background_interaction_settlement.py b/tests/unit/proxy/spend_tracking/test_background_interaction_settlement.py
index 21384279bbb..eea7f8a81aa 100644
--- a/tests/unit/proxy/spend_tracking/test_background_interaction_settlement.py
+++ b/tests/unit/proxy/spend_tracking/test_background_interaction_settlement.py
@@ -1,22 +1,28 @@
import asyncio
import time
from dataclasses import dataclass
-from datetime import datetime, timezone
+from datetime import datetime, timedelta, timezone
from typing import Optional
import pytest
import litellm.interactions.background_cost_polling as bg
+from litellm import Router
+from litellm.constants import BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS
from litellm.interactions.background_cost_polling import (
_create_context,
+ BackgroundInteractionPollContext,
configure_background_settlement_store,
maybe_settle_background_interaction_before_delete,
PendingBackgroundInteraction,
+ poll_and_log_background_interaction_cost,
PollSchedule,
+ resume_unsettled_background_interactions,
)
from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging
from litellm.proxy.spend_tracking.background_interaction_settlement import (
configure_background_interaction_settlement,
+ fetch_with_deployment_credentials,
install_background_interaction_settlement,
PrismaBackgroundSettlementStore,
)
@@ -78,7 +84,42 @@ class _FakeSettlementTable:
return len(matched)
def _matching(self, where) -> list:
- return [row for row in self.rows.values() if all(getattr(row, column) == value for column, value in where.items())]
+ return [row for row in self.rows.values() if _matches(row, where)]
+
+
+def _matches(row: "_Row", where: dict) -> bool:
+ return all(_satisfies(row, column, condition) for column, condition in where.items())
+
+
+def _satisfies(row: "_Row", column: str, condition) -> bool:
+ value = getattr(row, column)
+ if isinstance(condition, dict):
+ return value is not None and value < condition["lt"]
+ return value == condition
+
+
+class _FakeSpendLogs:
+ """The spend log table, holding only the request ids of the rows written so far."""
+
+ def __init__(self, request_ids: tuple[str, ...] = ()):
+ self.request_ids = request_ids
+
+ async def find_unique(self, *, where):
+ return where["request_id"] if where["request_id"] in self.request_ids else None
+
+
+def _store(
+ table,
+ claimed_by: str,
+ claim_lease_seconds: float = BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS,
+) -> PrismaBackgroundSettlementStore:
+ return PrismaBackgroundSettlementStore(
+ table=table,
+ spend_logs=_FakeSpendLogs(),
+ spend_logs_kept=lambda: True,
+ claimed_by=claimed_by,
+ claim_lease_seconds=claim_lease_seconds,
+ )
def _logging_obj(metadata: Optional[dict] = None) -> LitellmLogging:
@@ -110,7 +151,12 @@ def _pending(interaction_id: str) -> PendingBackgroundInteraction:
)
-def _stored_row(interaction_id: str, claimed: bool = False, create_context: Optional[object] = None) -> _Row:
+def _stored_row(
+ interaction_id: str,
+ claimed_at: Optional[datetime] = None,
+ settled: bool = False,
+ create_context: Optional[object] = None,
+) -> _Row:
return _Row(
interaction_id=interaction_id,
custom_llm_provider="gemini",
@@ -120,11 +166,21 @@ def _stored_row(interaction_id: str, claimed: bool = False, create_context: Opti
else _create_context(_logging_obj(), "gemini").model_dump(mode="json")
),
created_at=datetime.now(timezone.utc),
- claimed_at=datetime.now(timezone.utc) if claimed else None,
- claimed_by="replica-a:1" if claimed else None,
+ claimed_at=claimed_at,
+ claimed_by="replica-a:1" if claimed_at is not None else None,
+ settled_at=claimed_at if settled else None,
+ outcome="billed" if settled else None,
)
+def _minutes_ago(minutes: float) -> datetime:
+ return datetime.now(timezone.utc) - timedelta(minutes=minutes)
+
+
+def _past_the_default_lease() -> datetime:
+ return datetime.now(timezone.utc) - timedelta(seconds=BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS + 60)
+
+
def _completed(interaction_id: str) -> InteractionsAPIResponse:
return InteractionsAPIResponse(
id=interaction_id, model="gemini-2.5-flash", status="completed", steps=[], usage=dict(USAGE_BLOCK)
@@ -144,20 +200,20 @@ def _capturing_fetch():
@pytest.mark.asyncio
async def test_registered_row_reads_back_as_the_same_pending_interaction():
table = _FakeSettlementTable()
- store = PrismaBackgroundSettlementStore(table=table, claimed_by="replica-a:1")
+ store = _store(table, claimed_by="replica-a:1")
pending = _pending("interactions/bg-1")
await store.register(pending)
assert await store.pending("interactions/bg-1") == pending
- assert await store.unclaimed() == (pending,)
+ assert await store.unsettled() == (pending,)
@pytest.mark.asyncio
async def test_claim_is_won_by_exactly_one_settler():
table = _FakeSettlementTable()
- replica_a = PrismaBackgroundSettlementStore(table=table, claimed_by="replica-a:1")
- replica_b = PrismaBackgroundSettlementStore(table=table, claimed_by="replica-b:1")
+ replica_a = _store(table, claimed_by="replica-a:1")
+ replica_b = _store(table, claimed_by="replica-b:1")
await replica_a.register(_pending("interactions/bg-1"))
assert await replica_b.claim("interactions/bg-1") is True
@@ -200,7 +256,7 @@ class _MissingSettlementTable:
async def test_a_missing_table_holds_no_rows_and_takes_no_registration():
from prisma.errors import TableNotFoundError
- store = PrismaBackgroundSettlementStore(table=_MissingSettlementTable(), claimed_by="replica-a:1")
+ store = _store(_MissingSettlementTable(), claimed_by="replica-a:1")
with pytest.raises(TableNotFoundError):
await store.register(_pending("interactions/bg-1"))
@@ -208,29 +264,62 @@ async def test_a_missing_table_holds_no_rows_and_takes_no_registration():
assert await store.is_claimed("interactions/bg-1") is False
assert await store.claim("interactions/bg-1") is False
with pytest.raises(TableNotFoundError):
- await store.unclaimed()
+ await store.unsettled()
@pytest.mark.asyncio
-async def test_unclaimed_skips_claimed_and_unreadable_rows():
+async def test_unsettled_lists_claimed_rows_with_no_outcome_and_when_their_lease_ends():
+ claimed_at = _minutes_ago(1)
table = _FakeSettlementTable(
rows=(
_stored_row("interactions/bg-orphaned"),
- _stored_row("interactions/bg-settled", claimed=True),
+ _stored_row("interactions/bg-mid-settlement", claimed_at=claimed_at),
+ _stored_row("interactions/bg-settled", claimed_at=claimed_at, settled=True),
_stored_row("interactions/bg-from-the-future", create_context={"schema": "unknown"}),
)
)
- store = PrismaBackgroundSettlementStore(table=table, claimed_by="replica-b:1")
+ store = _store(table, claimed_by="replica-b:1", claim_lease_seconds=300)
- unclaimed = await store.unclaimed()
+ unsettled = await store.unsettled()
- assert [row.interaction_id for row in unclaimed] == ["interactions/bg-orphaned"]
+ assert [(row.interaction_id, row.claim_expires_at) for row in unsettled] == [
+ ("interactions/bg-orphaned", None),
+ ("interactions/bg-mid-settlement", claimed_at + timedelta(seconds=300)),
+ ]
+
+
+@pytest.mark.asyncio
+async def test_a_claim_whose_settler_died_is_taken_over_once_its_lease_runs_out():
+ """
+ The regression: a settler that died between its claim and its outcome left
+ the row claimed forever, so no replica, restart, or delete ever billed it.
+ """
+ table = _FakeSettlementTable(
+ rows=(
+ _stored_row("interactions/bg-dead-settler", claimed_at=_minutes_ago(10)),
+ _stored_row("interactions/bg-settled", claimed_at=_minutes_ago(10), settled=True),
+ )
+ )
+ replica_b = _store(table, claimed_by="replica-b:1", claim_lease_seconds=300)
+ replica_c = _store(table, claimed_by="replica-c:1", claim_lease_seconds=300)
+
+ assert await replica_b.is_claimed("interactions/bg-dead-settler") is False
+ assert await replica_b.pending("interactions/bg-dead-settler") is not None
+ assert await replica_b.claim("interactions/bg-dead-settler") is True
+ assert await replica_c.claim("interactions/bg-dead-settler") is False
+ assert await replica_c.is_claimed("interactions/bg-dead-settler") is True
+ assert table.rows["interactions/bg-dead-settler"].claimed_by == "replica-b:1"
+
+ assert await replica_b.is_claimed("interactions/bg-settled") is True
+ assert await replica_b.pending("interactions/bg-settled") is None
+ assert await replica_b.claim("interactions/bg-settled") is False
+ assert table.rows["interactions/bg-settled"].claimed_by == "replica-a:1"
@pytest.mark.asyncio
async def test_record_outcome_keeps_the_audit_trail_and_drops_the_stored_request_context():
table = _FakeSettlementTable()
- store = PrismaBackgroundSettlementStore(table=table, claimed_by="replica-a:1")
+ store = _store(table, claimed_by="replica-a:1")
await store.register(_pending("interactions/bg-1"))
assert await store.claim("interactions/bg-1")
assert table.rows["interactions/bg-1"].create_context
@@ -244,23 +333,198 @@ async def test_record_outcome_keeps_the_audit_trail_and_drops_the_stored_request
assert row.create_context == {}
+@pytest.mark.asyncio
+async def test_a_takeover_records_the_bill_a_settler_wrote_before_dying_instead_of_billing_again():
+ """
+ The regression: a settler that billed and then died, or failed to record
+ the outcome, left its claim to lapse with no outcome, and the replica that
+ took the claim over billed the interaction a second time.
+ """
+ table = _FakeSettlementTable(
+ rows=(
+ _stored_row("interactions/bg-billed-then-died", claimed_at=_past_the_default_lease()),
+ _stored_row("interactions/bg-died-before-billing", claimed_at=_past_the_default_lease()),
+ )
+ )
+ fetch, _ = _capturing_fetch()
+ previous_store = bg._STORE.store
+ try:
+ resumed = await configure_background_interaction_settlement(
+ table=table,
+ spend_logs=_FakeSpendLogs(request_ids=("interactions/bg-billed-then-died",)),
+ spend_logs_kept=lambda: True,
+ claimed_by="replica-b:1",
+ fetch_interaction=fetch,
+ schedule=FAST_SCHEDULE,
+ )
+
+ assert [await asyncio.wait_for(task, timeout=5) for task in resumed] == [None, "billed"]
+ assert [(row.claimed_by, row.outcome) for row in table.rows.values()] == [
+ ("replica-b:1", "billed"),
+ ("replica-b:1", "billed"),
+ ]
+ assert await bg._STORE.store.unsettled() == ()
+ finally:
+ configure_background_settlement_store(previous_store)
+
+
+@pytest.mark.asyncio
+async def test_with_spend_logs_off_a_takeover_leaves_the_interaction_unbilled_instead_of_billing_it_twice():
+ """
+ The regression: with spend logs off a settler that billed and then died
+ leaves no spend row to find, so the replica that took over its lapsed
+ claim billed the interaction a second time.
+ """
+ table = _FakeSettlementTable(
+ rows=(
+ _stored_row("interactions/bg-claimed-then-died", claimed_at=_past_the_default_lease()),
+ _stored_row("interactions/bg-orphaned"),
+ )
+ )
+ fetch, _ = _capturing_fetch()
+ previous_store = bg._STORE.store
+ try:
+ resumed = await configure_background_interaction_settlement(
+ table=table,
+ spend_logs=_FakeSpendLogs(),
+ spend_logs_kept=lambda: False,
+ claimed_by="replica-b:1",
+ fetch_interaction=fetch,
+ schedule=FAST_SCHEDULE,
+ )
+
+ assert [await asyncio.wait_for(task, timeout=5) for task in resumed] == [None, "billed"]
+ assert [(row.claimed_by, row.outcome) for row in table.rows.values()] == [
+ ("replica-b:1", "unsettled"),
+ ("replica-b:1", "billed"),
+ ]
+ finally:
+ configure_background_settlement_store(previous_store)
+
+
+class _SpendLogsWhoseFirstReadFails(_FakeSpendLogs):
+ def __init__(self):
+ super().__init__()
+ self.reads = 0
+
+ async def find_unique(self, *, where):
+ self.reads += 1
+ if self.reads == 1:
+ raise RuntimeError("read replica unavailable")
+ return await super().find_unique(where=where)
+
+
+@pytest.mark.asyncio
+async def test_a_takeover_whose_spend_log_read_failed_retries_and_bills_instead_of_stranding_the_row():
+ """
+ The regression: a takeover wrote its claim before reading the spend log,
+ so a failed read left the row under a fresh claim of its own, which the
+ poll then took for another settler's and left unbilled until a reboot.
+ """
+ table = _FakeSettlementTable(
+ rows=(_stored_row("interactions/bg-dead-settler", claimed_at=_past_the_default_lease()),)
+ )
+ fetch, _ = _capturing_fetch()
+ previous_store = bg._STORE.store
+ try:
+ (resumed,) = await configure_background_interaction_settlement(
+ table=table,
+ spend_logs=_SpendLogsWhoseFirstReadFails(),
+ spend_logs_kept=lambda: True,
+ claimed_by="replica-b:1",
+ fetch_interaction=fetch,
+ schedule=FAST_SCHEDULE,
+ )
+
+ assert await asyncio.wait_for(resumed, timeout=5) == "billed"
+ assert table.rows["interactions/bg-dead-settler"].outcome == "billed"
+ finally:
+ configure_background_settlement_store(previous_store)
+
+
+def _logging_obj_whose_bill_takes(seconds: float, bills: list) -> LitellmLogging:
+ logging_obj = _logging_obj()
+
+ async def bill(result):
+ bills.append(result.id)
+ await asyncio.sleep(seconds)
+
+ logging_obj.async_log_background_interaction_completion = bill
+ return logging_obj
+
+
+@pytest.mark.asyncio
+async def test_a_settle_that_outlasts_the_lease_is_not_taken_over_and_billed_twice():
+ """
+ The regression: a success callback that held the settle past the claim
+ lease let a booting replica take the live claim over and bill the same
+ interaction a second time.
+ """
+ lease_seconds = 0.6
+ table = _FakeSettlementTable()
+ creator = _store(table, claimed_by="replica-a:1", claim_lease_seconds=lease_seconds)
+ resumer = _store(table, claimed_by="replica-b:1", claim_lease_seconds=lease_seconds)
+ await creator.register(_pending("interactions/bg-slow-settle"))
+ fetch, _ = _capturing_fetch()
+ bills = []
+
+ def poll(store, resumed: bool, start_delay_seconds: float):
+ return poll_and_log_background_interaction_cost(
+ BackgroundInteractionPollContext(
+ interaction_id="interactions/bg-slow-settle",
+ custom_llm_provider="gemini",
+ logging_obj=_logging_obj_whose_bill_takes(lease_seconds * 2.5, bills),
+ initial_interval_seconds=0.001,
+ max_interval_seconds=0.002,
+ timeout_seconds=lease_seconds * 3,
+ start_delay_seconds=start_delay_seconds,
+ claim_renewal_interval_seconds=lease_seconds / 10,
+ store=store,
+ resumed=resumed,
+ ),
+ fetch,
+ )
+
+ outcomes = await asyncio.gather(
+ poll(creator, resumed=False, start_delay_seconds=0.0),
+ poll(resumer, resumed=True, start_delay_seconds=lease_seconds * 1.5),
+ )
+
+ assert outcomes == ["billed", None]
+ assert bills == ["interactions/bg-slow-settle"]
+ assert table.rows["interactions/bg-slow-settle"].claimed_by == "replica-a:1"
+
+
@pytest.mark.asyncio
async def test_configure_installs_the_store_and_resumes_the_orphaned_rows():
table = _FakeSettlementTable(
- rows=(_stored_row("interactions/bg-orphaned"), _stored_row("interactions/bg-settled", claimed=True))
+ rows=(
+ _stored_row("interactions/bg-orphaned"),
+ _stored_row("interactions/bg-dead-settler", claimed_at=_past_the_default_lease()),
+ _stored_row("interactions/bg-settled", claimed_at=_minutes_ago(1), settled=True),
+ )
)
fetch, captured = _capturing_fetch()
previous_store = bg._STORE.store
try:
resumed = await configure_background_interaction_settlement(
- table=table, claimed_by="replica-b:1", fetch_interaction=fetch, schedule=FAST_SCHEDULE
+ table=table,
+ spend_logs=_FakeSpendLogs(),
+ spend_logs_kept=lambda: True,
+ claimed_by="replica-b:1",
+ fetch_interaction=fetch,
+ schedule=FAST_SCHEDULE,
)
- assert len(resumed) == 1
- assert await asyncio.wait_for(resumed[0], timeout=5) == "billed"
- assert [context.interaction_id for context in captured] == ["interactions/bg-orphaned"]
- assert table.rows["interactions/bg-orphaned"].claimed_by == "replica-b:1"
- assert table.rows["interactions/bg-orphaned"].outcome == "billed"
+ assert [await asyncio.wait_for(task, timeout=5) for task in resumed] == ["billed", "billed"]
+ assert sorted(context.interaction_id for context in captured) == [
+ "interactions/bg-dead-settler",
+ "interactions/bg-orphaned",
+ ]
+ for interaction_id in ("interactions/bg-orphaned", "interactions/bg-dead-settler"):
+ assert table.rows[interaction_id].claimed_by == "replica-b:1"
+ assert table.rows[interaction_id].outcome == "billed"
+ assert table.rows["interactions/bg-settled"].claimed_by == "replica-a:1"
await table.create(
data={
@@ -288,7 +552,9 @@ class _PrismaClientWithoutSettlementTable:
async def test_install_keeps_booting_when_the_settlement_table_is_unreachable():
previous_store = bg._STORE.store
- await install_background_interaction_settlement(_PrismaClientWithoutSettlementTable())
+ await install_background_interaction_settlement(
+ _PrismaClientWithoutSettlementTable(), router=lambda: None, spend_logs_kept=lambda: True
+ )
assert bg._STORE.store is previous_store
@@ -296,3 +562,166 @@ async def test_install_keeps_booting_when_the_settlement_table_is_unreachable():
@dataclass(frozen=True)
class _JsonLike:
data: object
+
+
+@pytest.mark.asyncio
+async def test_a_resumed_row_still_under_a_live_claim_waits_for_the_lease_before_settling():
+ claimed_at = datetime.now(timezone.utc)
+ table = _FakeSettlementTable(rows=(_stored_row("interactions/bg-mid-settlement", claimed_at=claimed_at),))
+ store = _store(table, claimed_by="replica-b:1", claim_lease_seconds=0.2)
+ fetched_at = []
+
+ async def fetch(context):
+ fetched_at.append(datetime.now(timezone.utc))
+ return _completed(context.interaction_id)
+
+ (resumed,) = await resume_unsettled_background_interactions(store, fetch, schedule=FAST_SCHEDULE)
+
+ assert await asyncio.wait_for(resumed, timeout=5) == "billed"
+ assert fetched_at[0] >= claimed_at + timedelta(seconds=0.2)
+ assert table.rows["interactions/bg-mid-settlement"].claimed_by == "replica-b:1"
+
+
+def _router_with_a_deployment_scoped_key() -> Router:
+ return Router(
+ model_list=[
+ {
+ "model_name": "gemini-background",
+ "litellm_params": {
+ "model": "gemini/gemini-2.5-flash",
+ "api_key": "deployment-scoped-key",
+ "api_base": "https://deployment-scoped.example",
+ },
+ "model_info": {"id": "deployment-1"},
+ }
+ ]
+ )
+
+
+def _router_with_a_catch_all_route() -> Router:
+ return Router(model_list=[{"model_name": "*", "litellm_params": {"model": "gemini/*", "api_key": "catch-all-key"}}])
+
+
+def _router_with_two_deployments_in_the_group() -> Router:
+ return Router(
+ model_list=[
+ {
+ "model_name": "gemini-background",
+ "litellm_params": {"model": "gemini/gemini-2.5-flash", "api_key": f"deployment-key-{index}"},
+ "model_info": {"id": f"deployment-{index}"},
+ }
+ for index in (2, 3)
+ ]
+ )
+
+
+def _router_with_a_vertex_deployment() -> Router:
+ return Router(
+ model_list=[
+ {
+ "model_name": "vertex-background",
+ "litellm_params": {
+ "model": "vertex_ai/gemini-2.5-flash",
+ "vertex_project": "deployment-project",
+ "vertex_location": "europe-west4",
+ "vertex_credentials": "/secrets/deployment-service-account.json",
+ },
+ "model_info": {"id": "vertex-deployment-1"},
+ }
+ ]
+ )
+
+
+def _create_context_served_by(deployment_id: str) -> dict:
+ metadata = {
+ "user_api_key": "0123456789abcdef" * 4,
+ "model_info": {"id": deployment_id},
+ "model_group": "gemini-background",
+ }
+ return _create_context(_logging_obj(metadata), "gemini").model_dump(mode="json")
+
+
+@pytest.mark.parametrize("deployment_id", ["deployment-1", "deployment-1-before-an-edit"])
+@pytest.mark.asyncio
+async def test_a_resumed_poll_fetches_with_the_credentials_of_the_deployment_that_served_the_create(deployment_id):
+ """
+ The regression: a resumed poll has no request to take credentials from, so
+ it fetched with the environment's key, which a deployment-scoped key is
+ not, and every fetch failed until the row was given up unbilled. Editing
+ the deployment since gives it a new id, and its model group still names it.
+ """
+ table = _FakeSettlementTable(
+ rows=(_stored_row("interactions/bg-1", create_context=_create_context_served_by(deployment_id)),)
+ )
+ router = _router_with_a_deployment_scoped_key()
+ fetch, captured = _capturing_fetch()
+ previous_store = bg._STORE.store
+ try:
+ (resumed,) = await configure_background_interaction_settlement(
+ table=table,
+ spend_logs=_FakeSpendLogs(),
+ spend_logs_kept=lambda: True,
+ claimed_by="replica-b:1",
+ fetch_interaction=fetch_with_deployment_credentials(lambda: router, fetch),
+ schedule=FAST_SCHEDULE,
+ )
+
+ assert await asyncio.wait_for(resumed, timeout=5) == "billed"
+ assert (captured[0].credentials.api_key, captured[0].credentials.api_base) == (
+ "deployment-scoped-key",
+ "https://deployment-scoped.example",
+ )
+ finally:
+ configure_background_settlement_store(previous_store)
+
+
+@pytest.mark.parametrize(
+ "router, deployment_id, model_group",
+ [
+ (_router_with_a_deployment_scoped_key, "deployment-removed-since", None),
+ (_router_with_a_deployment_scoped_key, None, None),
+ (lambda: None, "deployment-1", "gemini-background"),
+ (_router_with_a_catch_all_route, "deployment-removed-since", None),
+ (_router_with_two_deployments_in_the_group, "deployment-removed-since", "gemini-background"),
+ (_router_with_a_vertex_deployment, "deployment-removed-since", "vertex-background"),
+ (_router_with_a_vertex_deployment, "vertex-deployment-1", None),
+ ],
+)
+@pytest.mark.asyncio
+async def test_a_resumed_poll_without_a_known_deployment_fetches_with_the_environment_credentials(
+ router, deployment_id, model_group
+):
+ live_router = router()
+ fetch, captured = _capturing_fetch()
+ context = BackgroundInteractionPollContext(
+ interaction_id="interactions/bg-1",
+ custom_llm_provider="gemini",
+ logging_obj=_logging_obj(),
+ deployment_id=deployment_id,
+ model_group=model_group,
+ )
+
+ await fetch_with_deployment_credentials(lambda: live_router, fetch)(context)
+
+ assert captured == [context]
+
+
+@pytest.mark.asyncio
+async def test_a_resumed_vertex_poll_fetches_with_the_project_location_and_credentials_of_its_deployment():
+ router = _router_with_a_vertex_deployment()
+ fetch, captured = _capturing_fetch()
+ context = BackgroundInteractionPollContext(
+ interaction_id="interactions/bg-1",
+ custom_llm_provider="vertex_ai",
+ logging_obj=_logging_obj(),
+ deployment_id="vertex-deployment-1",
+ )
+
+ await fetch_with_deployment_credentials(lambda: router, fetch)(context)
+
+ credentials = captured[0].credentials
+ assert (credentials.vertex_project, credentials.vertex_location, credentials.vertex_credentials) == (
+ "deployment-project",
+ "europe-west4",
+ "/secrets/deployment-service-account.json",
+ )
diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py
index 7291f49af99..de2676bafaa 100644
--- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py
+++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py
@@ -14,8 +14,8 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict
import litellm
import litellm.constants as litellm_constants
import litellm.proxy.spend_tracking.spend_tracking_utils as spend_tracking_utils
-from litellm.constants import LITELLM_TRUNCATED_PAYLOAD_FIELD, LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, LITTELM_CLI_SERVICE_ACCOUNT_NAME, LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, MAX_SPEND_LOG_MODEL_NAME_LENGTH, REDACTED_BY_LITELM_STRING, SESSION_ID_OMITTED_METADATA_KEY, UNKNOWN_MODEL_SPEND_LOG_MODEL
-from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY
+from litellm.constants import LITELLM_TRUNCATED_PAYLOAD_FIELD, LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, LITTELM_CLI_SERVICE_ACCOUNT_NAME, LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, MAX_SPEND_LOG_MODEL_NAME_LENGTH, REDACTED_BY_LITELM_STRING, SESSION_ID_GENERATED_METADATA_KEY, SESSION_ID_OMITTED_METADATA_KEY, UNKNOWN_MODEL_SPEND_LOG_MODEL
+from litellm.constants import CALLER_SESSION_ID_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
from litellm.llms.base_llm.ocr.transformation import OCRResponse, OCRUsageInfo
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
@@ -243,6 +243,82 @@ def test_batch_session_outranks_the_per_request_trace_id():
assert session_id == "batch-uid-1"
+def _chat_session_id(
+ monkeypatch: pytest.MonkeyPatch,
+ litellm_params: Mapping[str, object],
+ correlation_in_logs: bool = True,
+ batch_trace_session_id: str | None = None,
+ metadata: Mapping[str, object] = MappingProxyType({}),
+) -> str | None:
+ monkeypatch.setattr(litellm, "request_correlation_in_logs", correlation_in_logs)
+ return _get_session_id_for_spend_log(
+ kwargs={"litellm_trace_id": "trace-abc", "litellm_params": litellm_params},
+ metadata=metadata,
+ standard_logging_payload=_TRACE_ONLY_STANDARD_LOGGING,
+ omit_when_missing=False,
+ batch_trace_session_id=batch_trace_session_id,
+ )
+
+
+_CALLER_SESSION_PARAMS: Final = MappingProxyType({"litellm_session_id": "sess-1", "metadata": {"trace_id": "otel-1"}})
+
+
+def test_caller_session_wins_over_the_per_request_trace_id(monkeypatch: pytest.MonkeyPatch):
+ assert _chat_session_id(monkeypatch, _CALLER_SESSION_PARAMS) == "sess-1"
+
+
+def test_batch_session_outranks_a_caller_session(monkeypatch: pytest.MonkeyPatch):
+ assert _chat_session_id(monkeypatch, _CALLER_SESSION_PARAMS, batch_trace_session_id="batch-uid-1") == "batch-uid-1"
+
+
+def test_caller_session_keeps_the_legacy_trace_id_with_correlation_logging_off(monkeypatch: pytest.MonkeyPatch):
+ assert _chat_session_id(monkeypatch, _CALLER_SESSION_PARAMS, correlation_in_logs=False) == "trace-abc"
+
+
+@pytest.mark.parametrize(("correlation_in_logs", "expected_session_id"), [(True, "convo-1"), (False, "trace-abc")])
+def test_session_displaced_by_a_trace_header_wins_only_with_correlation_logging(
+ monkeypatch: pytest.MonkeyPatch, correlation_in_logs: bool, expected_session_id: str
+):
+ header_chain: Final = {"litellm_session_id": "trace-abc", "metadata": {"trace_id": "trace-abc", "session_id": "trace-abc"}}
+ session_id: Final = _chat_session_id(
+ monkeypatch,
+ header_chain,
+ correlation_in_logs=correlation_in_logs,
+ metadata={CALLER_SESSION_ID_METADATA_KEY: "convo-1"},
+ )
+ assert session_id == expected_session_id
+
+
+@pytest.mark.parametrize(
+ ("litellm_params", "metadata"),
+ [
+ pytest.param({"litellm_session_id": "convo\x00x"}, {}, id="body-session"),
+ pytest.param({"litellm_session_id": "trace-abc"}, {CALLER_SESSION_ID_METADATA_KEY: "convo\x00x"}, id="beside-trace-header"),
+ ],
+)
+def test_caller_session_postgres_cannot_store_keeps_the_trace_id(
+ monkeypatch: pytest.MonkeyPatch, litellm_params: Mapping[str, object], metadata: Mapping[str, object]
+):
+ assert _chat_session_id(monkeypatch, litellm_params, metadata=metadata) == "trace-abc"
+
+
+@pytest.mark.parametrize(
+ "litellm_params",
+ [
+ pytest.param({"metadata": {}}, id="no-session"),
+ pytest.param(
+ {"litellm_session_id": "minted-1", "metadata": {"session_id": "minted-1", SESSION_ID_GENERATED_METADATA_KEY: True}},
+ id="minted-by-generate-policy",
+ ),
+ pytest.param({"litellm_session_id": "job-7", "metadata": {"trace_id": "job-7"}}, id="copied-from-trace-id"),
+ ],
+)
+def test_session_that_is_not_a_caller_conversation_keeps_the_trace_id(
+ monkeypatch: pytest.MonkeyPatch, litellm_params: Mapping[str, object]
+):
+ assert _chat_session_id(monkeypatch, litellm_params) == "trace-abc"
+
+
def test_omit_policy_still_suppresses_batch_sessions():
session_id: Final = _get_session_id_for_spend_log(
kwargs={},
diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py
index 96c11be2650..93202d93566 100644
--- a/tests/unit/proxy/test_litellm_pre_call_utils.py
+++ b/tests/unit/proxy/test_litellm_pre_call_utils.py
@@ -17,6 +17,7 @@ from starlette.datastructures import Headers
import litellm
from litellm.constants import (
+ CALLER_SESSION_ID_METADATA_KEY,
ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY,
SERVER_STREAMING_CLASSIFICATION_KEY,
SERVER_STREAMING_CLASSIFICATION_MARKER,
@@ -3308,6 +3309,93 @@ def test_add_litellm_metadata_from_request_headers_x_litellm_trace_id_sets_chain
assert data["litellm_trace_id"] == "foo"
+@pytest.mark.parametrize(("correlation_in_logs", "expected_spend_log_session"), [(True, "convo-1"), (False, None)])
+@pytest.mark.parametrize("session_in_metadata", [False, True], ids=["top-level", "metadata"])
+def test_x_litellm_trace_id_keeps_the_body_session_for_spend_logs_only(
+ monkeypatch: pytest.MonkeyPatch,
+ correlation_in_logs: bool,
+ expected_spend_log_session: str | None,
+ session_in_metadata: bool,
+):
+ monkeypatch.setattr(litellm, "request_correlation_in_logs", correlation_in_logs)
+ metadata: Final[dict[str, str]] = {"session_id": "convo-1"} if session_in_metadata else {}
+ data: Final[dict[str, object]] = (
+ {"metadata": metadata} if session_in_metadata else {"metadata": metadata, "litellm_session_id": "convo-1"}
+ )
+ LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
+ headers={"x-litellm-trace-id": "req-trace-1"}, data=data, _metadata_variable_name="metadata"
+ )
+ assert metadata.get(CALLER_SESSION_ID_METADATA_KEY) == expected_spend_log_session
+ assert data["litellm_session_id"] == "req-trace-1"
+ assert metadata["session_id"] == "req-trace-1"
+ assert data["litellm_trace_id"] == "req-trace-1"
+
+
+@pytest.mark.parametrize(
+ "session_header", ["x-litellm-session-id", "x-claude-code-session-id", "x-session-id"]
+)
+@pytest.mark.parametrize("with_trace_header", [False, True], ids=["session-header-only", "beside-trace-header"])
+def test_session_headers_still_beat_the_body_session_with_correlation_logging(
+ monkeypatch: pytest.MonkeyPatch, session_header: str, with_trace_header: bool
+):
+ monkeypatch.setattr(litellm, "request_correlation_in_logs", True)
+ metadata: Final[dict[str, str]] = {}
+ data: Final[dict[str, object]] = {"metadata": metadata, "litellm_session_id": "convo-1"}
+ trace_header: Final = {"x-litellm-trace-id": "req-trace-1"} if with_trace_header else {}
+ LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
+ headers={**trace_header, session_header: "hdr-session-1"}, data=data, _metadata_variable_name="metadata"
+ )
+ spend_log_session: Final = metadata.get(CALLER_SESSION_ID_METADATA_KEY) or data["litellm_session_id"]
+ assert spend_log_session == "hdr-session-1"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("forged_in", ["metadata", "litellm_metadata"])
+@pytest.mark.parametrize(
+ "headers",
+ [{"x-litellm-session-id": "hdr-session-1"}, {}],
+ ids=["session-header", "no-headers"],
+)
+async def test_client_metadata_cannot_set_the_spend_log_caller_session(
+ monkeypatch: pytest.MonkeyPatch, forged_in: str, headers: dict[str, str]
+):
+ monkeypatch.setattr(litellm, "request_correlation_in_logs", True)
+ updated: Final = await add_litellm_data_to_request(
+ data={
+ "model": "gpt-5.4-mini",
+ "messages": [{"role": "user", "content": "hi"}],
+ forged_in: {CALLER_SESSION_ID_METADATA_KEY: "forged-session"},
+ },
+ request=_make_request_mock("/v1/chat/completions", headers),
+ user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
+ proxy_config=MagicMock(),
+ general_settings={},
+ version="test-version",
+ )
+ assert CALLER_SESSION_ID_METADATA_KEY not in updated["metadata"]
+
+
+@pytest.mark.asyncio
+async def test_spend_log_caller_session_stays_out_of_requester_metadata(monkeypatch: pytest.MonkeyPatch):
+ monkeypatch.setattr(litellm, "request_correlation_in_logs", True)
+ updated: Final = await add_litellm_data_to_request(
+ data={
+ "model": "gpt-5.4-mini",
+ "messages": [{"role": "user", "content": "hi"}],
+ "litellm_session_id": "convo-1",
+ "metadata": {"team_label": "kept"},
+ },
+ request=_make_request_mock("/v1/chat/completions", {"x-litellm-trace-id": "req-trace-1"}),
+ user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
+ proxy_config=MagicMock(),
+ general_settings={},
+ version="test-version",
+ )
+ assert updated["metadata"][CALLER_SESSION_ID_METADATA_KEY] == "convo-1"
+ assert updated["metadata"]["requester_metadata"]["team_label"] == "kept"
+ assert CALLER_SESSION_ID_METADATA_KEY not in updated["metadata"]["requester_metadata"]
+
+
def test_add_litellm_metadata_from_request_headers_x_litellm_session_id_sets_chain_id():
"""x-litellm-session-id sets both metadata and top-level litellm_session_id/litellm_trace_id for call chaining."""
headers = {"x-litellm-session-id": "bar"}
diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py
index 0ecebc7560c..e824c5017c6 100644
--- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py
+++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py
@@ -12612,6 +12612,7 @@ def test_dump_redacted_config_serializes_non_json_native_values():
async def test_update_config_general_settings_emits_audit_log(monkeypatch):
import litellm.proxy.proxy_server as proxy_server_module
from litellm.proxy._types import ConfigFieldUpdate
+ from litellm.proxy.management_helpers.audit_logs import drain_audit_tasks
from litellm.proxy.proxy_server import update_config_general_settings
existing = {"max_parallel_requests": 5, "some_api_key": "sk-stored-secret"}
@@ -12635,6 +12636,7 @@ async def test_update_config_general_settings_emits_audit_log(monkeypatch):
)
# Audit is scheduled via asyncio.create_task; yield so it runs.
await asyncio.sleep(0)
+ await drain_audit_tasks()
fake.db.litellm_auditlog.create.assert_awaited_once()
written = fake.db.litellm_auditlog.create.call_args.kwargs["data"]
@@ -13004,6 +13006,7 @@ async def test_update_config_general_settings_applies_ssrf_globals(monkeypatch):
async def test_delete_config_general_settings_emits_deleted_audit_log(monkeypatch):
import litellm.proxy.proxy_server as proxy_server_module
from litellm.proxy._types import ConfigFieldDelete
+ from litellm.proxy.management_helpers.audit_logs import drain_audit_tasks
from litellm.proxy.proxy_server import delete_config_general_settings
existing = {"max_parallel_requests": 5}
@@ -13023,6 +13026,7 @@ async def test_delete_config_general_settings_emits_deleted_audit_log(monkeypatc
)
# Audit is scheduled via asyncio.create_task; yield so it runs.
await asyncio.sleep(0)
+ await drain_audit_tasks()
fake.db.litellm_auditlog.create.assert_awaited_once()
written = fake.db.litellm_auditlog.create.call_args.kwargs["data"]
diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py
index 9bed719227f..7d70e058396 100644
--- a/tests/unit/proxy/test_proxy_token_counter.py
+++ b/tests/unit/proxy/test_proxy_token_counter.py
@@ -9,7 +9,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
from dotenv import load_dotenv
-from pydantic import JsonValue
+from pydantic import BaseModel, ConfigDict, JsonValue
load_dotenv()
@@ -21,10 +21,13 @@ from fastapi.responses import JSONResponse
import litellm
from litellm import Router
+from litellm.llms.anthropic.count_tokens import handler as anthropic_count_tokens_handler
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter
from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
+from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler
+from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import VertexAIPartnerModelsTokenCounter
from litellm.proxy._types import ProxyException, TokenCountRequest
from litellm.proxy.anthropic_endpoints.endpoints import (
count_tokens as anthropic_count_tokens,
@@ -327,22 +330,16 @@ async def test_internal_token_counter_anthropic_provider_detection():
setattr(litellm.proxy.proxy_server, "llm_router", llm_router)
- # Test with is_direct_request=False but non-Anthropic provider
- response = await token_counter(
- request=TokenCountRequest(
- model="gpt-test",
- messages=[{"role": "user", "content": "hello"}],
- ),
- call_endpoint=True,
- )
-
- print("Non-Anthropic provider test response:", response)
-
- # Verify response structure
- assert response.request_model == "gpt-test"
- assert response.model_used == "gpt-4"
- assert response.total_tokens > 0
- assert response.tokenizer_type == "openai_tokenizer" # Should use LiteLLM tokenizer
+ with pytest.raises(ProxyException) as missing_key:
+ await token_counter(
+ request=TokenCountRequest(
+ model="gpt-test",
+ messages=[{"role": "user", "content": "hello"}],
+ ),
+ call_endpoint=True,
+ )
+ assert missing_key.value.code == "401"
+ assert "No OpenAI API key is set for model gpt-test" in missing_key.value.message
@pytest.mark.asyncio
@@ -352,12 +349,14 @@ async def test_anthropic_endpoint_error_handling():
"""
from unittest.mock import MagicMock
- from fastapi import HTTPException, Request
+ from fastapi import Request
from litellm.proxy.anthropic_endpoints.endpoints import count_tokens
# Mock request object
mock_request = MagicMock(spec=Request)
+ mock_request.headers = {}
+ mock_request.scope = {}
mock_user_api_key_dict = MagicMock()
# Test missing model parameter
@@ -375,12 +374,13 @@ async def test_anthropic_endpoint_error_handling():
anthropic_endpoints.read_request_body = mock_read_request_body
try:
- # Should raise HTTPException for missing model
- with pytest.raises(HTTPException) as exc_info:
- await count_tokens(mock_request, mock_user_api_key_dict)
+ response = await count_tokens(mock_request, mock_user_api_key_dict)
- assert exc_info.value.status_code == 400
- assert "model parameter is required" in str(exc_info.value.detail)
+ assert response.status_code == 400
+ assert json.loads(response.body) == {
+ "type": "error",
+ "error": {"type": "invalid_request_error", "message": "model parameter is required"},
+ }
print("✅ Error handling test passed!")
@@ -471,7 +471,7 @@ async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter():
mock_router.model_list = [
{
"model_name": "gpt-4",
- "litellm_params": {"model": "openai/gpt-4"},
+ "litellm_params": {"model": "databricks/gpt-4"},
"model_info": {},
}
]
@@ -480,7 +480,7 @@ async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter():
mock_router.async_get_available_deployment = AsyncMock(
return_value={
"model_name": "gpt-4",
- "litellm_params": {"model": "openai/gpt-4"},
+ "litellm_params": {"model": "databricks/gpt-4"},
"model_info": {},
}
)
@@ -1177,3 +1177,277 @@ async def test_anthropic_endpoint_429_rate_limit_error_format():
finally:
anthropic_endpoints.read_request_body = original_read_request_body
proxy_server.count_request_tokens = original_token_counter
+
+
+class _GeminiPart(BaseModel):
+ model_config = ConfigDict(frozen=True)
+ text: str = ""
+
+
+class _GeminiContent(BaseModel):
+ model_config = ConfigDict(frozen=True)
+ parts: tuple[_GeminiPart, ...] = ()
+
+
+class _GeminiCountTokensBody(BaseModel):
+ model_config = ConfigDict(frozen=True)
+ contents: tuple[_GeminiContent, ...] = ()
+
+
+def _word_count(content: _GeminiContent) -> int:
+ return sum(len(part.text.split()) for part in content.parts)
+
+
+def _gemini_word_counting_peer(request: httpx.Request) -> httpx.Response:
+ body: Final = _GeminiCountTokensBody.model_validate_json(request.content)
+ total: Final = sum(_word_count(content) for content in body.contents)
+ return httpx.Response(200, json={"totalTokens": total}, request=request)
+
+
+def _unavailable_peer(request: httpx.Request) -> httpx.Response:
+ return httpx.Response(503, json={"error": {"message": "unavailable"}}, request=request)
+
+
+def _gemini_router(drop_params: bool = False) -> Router:
+ return Router(
+ model_list=[
+ {
+ "model_name": "gemini-counted",
+ "litellm_params": {
+ "model": "gemini/gemini-3.8-flash",
+ "api_key": "test-key",
+ "api_base": "https://count-tokens.test",
+ "drop_params": drop_params,
+ },
+ }
+ ]
+ )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "text", ["hello there", "The quick brown fox jumps over the lazy dog. " * 50], ids=["short", "long"]
+)
+async def test_call_endpoint_counts_a_prompt_like_the_same_text_sent_as_a_user_message(
+ text: str, monkeypatch: pytest.MonkeyPatch
+) -> None:
+ monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router())
+ client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer))
+
+ with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client):
+ from_prompt: Final = await token_counter(
+ request=TokenCountRequest(model="gemini-counted", prompt=text), call_endpoint=True
+ )
+ from_message: Final = await token_counter(
+ request=TokenCountRequest(model="gemini-counted", messages=[{"role": "user", "content": text}]),
+ call_endpoint=True,
+ )
+
+ assert from_prompt.error is False
+ assert from_prompt.total_tokens == from_message.total_tokens == len(text.split())
+
+
+class _PartnerMessage(BaseModel):
+ model_config = ConfigDict(frozen=True)
+ content: str = ""
+
+
+class _PartnerCountRequest(BaseModel):
+ model_config = ConfigDict(frozen=True)
+ messages: tuple[_PartnerMessage, ...] = ()
+
+
+async def _vertex_partner_word_count(
+ model: str, request_data: JsonValue, litellm_params: JsonValue
+) -> dict[str, JsonValue]:
+ body: Final = _PartnerCountRequest.model_validate(request_data)
+ return {"input_tokens": sum(len(message.content.split()) for message in body.messages)}
+
+
+@pytest.mark.asyncio
+async def test_call_endpoint_counts_a_prompt_on_a_vertex_ai_deployment(monkeypatch: pytest.MonkeyPatch) -> None:
+ router: Final = Router(
+ model_list=[
+ {
+ "model_name": "vertex-counted",
+ "litellm_params": {
+ "model": "vertex_ai/claude-sonnet-5-5",
+ "vertex_project": "test-project",
+ "vertex_location": "us-east5",
+ },
+ }
+ ]
+ )
+ monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router)
+ text: Final = "The quick brown fox jumps over the lazy dog."
+
+ with patch.object(
+ VertexAIPartnerModelsTokenCounter, "handle_count_tokens_request", side_effect=_vertex_partner_word_count
+ ):
+ from_prompt: Final = await token_counter(
+ request=TokenCountRequest(model="vertex-counted", prompt=text), call_endpoint=True
+ )
+
+ assert (from_prompt.error, from_prompt.total_tokens) == (False, 9)
+
+
+_FOX: Final = "The quick brown fox jumps over the lazy dog."
+
+
+_CONVERSATIONS: Final = (
+ {"messages": [{"role": "user", "content": _FOX}]},
+ {"contents": [{"role": "user", "parts": [{"text": _FOX}]}]},
+)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("conversation", _CONVERSATIONS, ids=["messages", "contents"])
+@pytest.mark.parametrize("call_endpoint", [True, False], ids=["provider", "local"])
+async def test_prompt_next_to_a_conversation_answers_400_without_drop_params(
+ conversation: dict[str, JsonValue], call_endpoint: bool, monkeypatch: pytest.MonkeyPatch
+) -> None:
+ monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router())
+ client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer))
+
+ with (
+ patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client),
+ pytest.raises(ProxyException) as rejected,
+ ):
+ await token_counter(
+ request=TokenCountRequest(model="gemini-counted", prompt="hi", **conversation), call_endpoint=call_endpoint
+ )
+
+ assert (rejected.value.code, rejected.value.param) == ("400", "prompt")
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ ("conversation", "call_endpoint", "global_drop_params", "deployment_drop_params"),
+ [
+ (_CONVERSATIONS[0], True, False, True),
+ (_CONVERSATIONS[1], True, False, True),
+ (_CONVERSATIONS[0], False, False, True),
+ (_CONVERSATIONS[0], False, True, False),
+ ],
+ ids=[
+ "provider_messages",
+ "provider_contents",
+ "local_messages",
+ "local_messages_global_setting",
+ ],
+)
+async def test_drop_params_counts_a_conversation_without_the_prompt(
+ conversation: dict[str, JsonValue],
+ call_endpoint: bool,
+ global_drop_params: bool,
+ deployment_drop_params: bool,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router(drop_params=deployment_drop_params))
+ monkeypatch.setattr(litellm, "drop_params", global_drop_params)
+ client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer))
+
+ with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client):
+ with_prompt: Final = await token_counter(
+ request=TokenCountRequest(model="gemini-counted", prompt="hi", **conversation), call_endpoint=call_endpoint
+ )
+ conversation_only: Final = await token_counter(
+ request=TokenCountRequest(model="gemini-counted", **conversation), call_endpoint=call_endpoint
+ )
+
+ assert with_prompt.error is False
+ assert with_prompt.total_tokens == conversation_only.total_tokens > 0
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("empty_conversation", [{"messages": []}, {"contents": []}], ids=["messages", "contents"])
+async def test_call_endpoint_prompt_next_to_an_empty_conversation_counts_the_prompt(
+ empty_conversation: dict[str, JsonValue], monkeypatch: pytest.MonkeyPatch
+) -> None:
+ monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router())
+ client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer))
+
+ with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client):
+ counted: Final = await token_counter(
+ request=TokenCountRequest(model="gemini-counted", prompt=_FOX, **empty_conversation), call_endpoint=True
+ )
+
+ assert (counted.error, counted.total_tokens) == (False, 9)
+
+
+@pytest.mark.asyncio
+async def test_local_count_of_a_prompt_next_to_empty_messages_counts_the_prompt(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router())
+
+ with_empty: Final = await token_counter(request=TokenCountRequest(model="gemini-counted", prompt=_FOX, messages=[]))
+ alone: Final = await token_counter(request=TokenCountRequest(model="gemini-counted", prompt=_FOX))
+
+ assert with_empty.error is False
+ assert with_empty.total_tokens == alone.total_tokens > 0
+
+
+_WEATHER_TOOL: Final = {
+ "type": "function",
+ "function": {
+ "name": "get_weather",
+ "description": "Get the current weather for a city",
+ "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
+ },
+}
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "chat_fields",
+ [{"system": "You answer in one short sentence. " * 20}, {"tools": [_WEATHER_TOOL]}],
+ ids=["system", "tools"],
+)
+@pytest.mark.parametrize("empty_conversation", [{}, {"messages": []}], ids=["no_messages", "empty_messages"])
+async def test_local_count_of_a_prompt_with_system_or_tools_counts_them_like_the_message_form(
+ chat_fields: dict[str, JsonValue], empty_conversation: dict[str, JsonValue], monkeypatch: pytest.MonkeyPatch
+) -> None:
+ monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router())
+
+ from_prompt: Final = await token_counter(
+ request=TokenCountRequest(model="gemini-counted", prompt=_FOX, **empty_conversation, **chat_fields)
+ )
+ from_message: Final = await token_counter(
+ request=TokenCountRequest(model="gemini-counted", messages=[{"role": "user", "content": _FOX}], **chat_fields)
+ )
+ prompt_alone: Final = await token_counter(request=TokenCountRequest(model="gemini-counted", prompt=_FOX))
+
+ assert from_prompt.error is False
+ assert from_prompt.total_tokens == from_message.total_tokens > prompt_alone.total_tokens
+
+
+@pytest.mark.asyncio
+async def test_call_endpoint_prompt_keeps_the_local_count_on_a_counter_that_needs_messages(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ router: Final = Router(
+ model_list=[
+ {
+ "model_name": "anthropic-counted",
+ "litellm_params": {
+ "model": "anthropic/claude-sonnet-5-5",
+ "api_key": "test-key",
+ "api_base": "https://count-tokens.test",
+ },
+ }
+ ]
+ )
+ monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router)
+ client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_unavailable_peer))
+ request: Final = TokenCountRequest(model="anthropic-counted", prompt="The quick brown fox jumps over the lazy dog.")
+
+ with patch.object(anthropic_count_tokens_handler, "get_async_httpx_client", return_value=client):
+ with_endpoint: Final = await token_counter(request=request, call_endpoint=True)
+ local: Final = await token_counter(request=request, call_endpoint=False)
+
+ assert (with_endpoint.error, with_endpoint.total_tokens, with_endpoint.tokenizer_type) == (
+ False,
+ local.total_tokens,
+ local.tokenizer_type,
+ )
diff --git a/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py
index 6705e5ac8a5..da799a307a6 100644
--- a/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py
+++ b/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py
@@ -3771,6 +3771,24 @@ class TestPtuCostAttributionUISetting:
assert not mock_prisma.db.litellm_uisettings.upsert.called
+class TestEnableChatUISetting:
+ def test_chat_ui_is_on_when_nothing_is_stored(self, mock_auth, monkeypatch):
+ TestPtuCostAttributionUISetting._mock_prisma(monkeypatch)
+
+ response = client.get("/get/ui_settings")
+
+ assert response.status_code == 200
+ assert response.json()["values"]["enable_chat_ui"] is True
+
+ def test_admin_can_still_turn_chat_ui_off(self, mock_auth, monkeypatch):
+ TestPtuCostAttributionUISetting._mock_prisma(monkeypatch, stored={"enable_chat_ui": False})
+
+ response = client.get("/get/ui_settings")
+
+ assert response.status_code == 200
+ assert response.json()["values"]["enable_chat_ui"] is False
+
+
class TestApplyUserBudgetToTeamKeysUISetting:
"""``apply_user_budget_to_team_keys`` mirrors general_settings on every GET.
diff --git a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py
index 0d55fc35e1d..cc5d5c62739 100644
--- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py
+++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py
@@ -17,6 +17,7 @@ from openai.types.responses.response_function_web_search import (
from pydantic import TypeAdapter
import litellm
+from litellm.constants import CALLER_SESSION_ID_METADATA_KEY
from litellm.integrations.anthropic_cache_control_hook import MAX_CACHE_CONTROL_BLOCKS, AnthropicCacheControlHook
from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR
from litellm.litellm_core_utils.prompt_templates.factory import (
@@ -3109,7 +3110,6 @@ class TestToolTransformation:
("bedrock/converse/global.anthropic.claude-sonnet-5", "bedrock_converse"),
("anthropic.claude-sonnet-4-5-20250929-v1:0", "bedrock"),
("claude-sonnet-5", "vertex_ai"),
- ("gemini-3.1-pro-preview", "vertex_ai"),
("moonshotai.kimi-k2-thinking", "bedrock_mantle"),
],
)
@@ -3117,8 +3117,8 @@ class TestToolTransformation:
"""
A Responses request carrying ``reasoning.summary`` must still reach a chat provider as a
plain ``reasoning_effort`` string. ``summary`` is Responses-only, and forwarding the whole
- object turns reasoning off: Bedrock Converse and Vertex silently discard a non-string
- ``reasoning_effort``, and Bedrock Mantle rejects the request outright.
+ object turns reasoning off: Bedrock Converse and Claude on Vertex silently discard a
+ non-string ``reasoning_effort``, and Bedrock Mantle rejects the request outright.
"""
responses_api_request = {"reasoning": {"effort": "medium", "summary": "auto"}}
@@ -3131,6 +3131,42 @@ class TestToolTransformation:
assert result["reasoning_effort"] == "medium"
+ @pytest.mark.parametrize(
+ "model, custom_llm_provider",
+ [("gemini-3.1-pro-preview", "vertex_ai"), ("gemini-2.5-flash", "gemini")],
+ )
+ @pytest.mark.parametrize(
+ "reasoning, sends_thoughts",
+ [({"effort": "medium", "summary": "auto"}, True), ({"effort": "medium"}, False)],
+ )
+ def test_a_summary_ask_turns_gemini_thoughts_on_over_the_gateway_default(
+ self,
+ monkeypatch: pytest.MonkeyPatch,
+ model: str,
+ custom_llm_provider: str,
+ reasoning: dict[str, str],
+ sends_thoughts: bool,
+ ) -> None:
+ from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig
+ from litellm.utils import reasoning_effort_with_summary_alias
+
+ monkeypatch.setattr(litellm, "gemini_include_thoughts_default", False)
+ bridged: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request(
+ model=model,
+ input="hi",
+ responses_api_request={"reasoning": reasoning},
+ custom_llm_provider=custom_llm_provider,
+ )
+
+ gemini: Final = VertexGeminiConfig()
+ reasoning_effort: Final = reasoning_effort_with_summary_alias(
+ bridged["reasoning_effort"], bridged, accepts_summary=gemini.accepts_reasoning_summary_in_reasoning_effort
+ )
+ mapped: Final = gemini.map_openai_params({"reasoning_effort": reasoning_effort}, {}, model, True)
+
+ assert mapped["thinkingConfig"]["includeThoughts"] is sends_thoughts
+ assert bridged["reasoning_effort"] == "medium"
+
@pytest.mark.parametrize(
"model, custom_llm_provider",
[
@@ -6464,3 +6500,69 @@ def test_transform_request_reads_service_tier_from_responses_api_request() -> No
)
assert result["service_tier"] == "priority"
+
+
+def _continued_chain_request(request_kwargs: Mapping[str, object]) -> Mapping[str, object]:
+ async def session_lookup(previous_response_id: str) -> ChatCompletionSession:
+ assert previous_response_id == "resp_turn_1"
+ return ChatCompletionSession(
+ messages=[{"role": "user", "content": "turn 1"}], litellm_session_id="chain-session"
+ )
+
+ return asyncio.run(
+ LiteLLMCompletionResponsesConfig.async_responses_api_session_handler(
+ previous_response_id="resp_turn_1",
+ litellm_completion_request={
+ "model": _CLAUDE_ON_BEDROCK,
+ "custom_llm_provider": "bedrock",
+ "messages": [{"role": "user", "content": "turn 2"}],
+ "metadata": request_kwargs.get("metadata"),
+ },
+ instructions=None,
+ session_lookup=session_lookup,
+ request_kwargs=request_kwargs,
+ )
+ )
+
+
+@pytest.mark.parametrize(
+ ("request_kwargs", "logged_metadata_key"),
+ [
+ (
+ {
+ "metadata": {"tag": "client"},
+ "litellm_metadata": {CALLER_SESSION_ID_METADATA_KEY: "turn-session", "user_api_key": "hashed"},
+ },
+ "litellm_metadata",
+ ),
+ ({"metadata": {"user_api_key": "hashed"}}, "metadata"),
+ ],
+)
+def test_a_continued_chain_logs_its_spend_under_the_chain_session_with_request_correlation(
+ monkeypatch: pytest.MonkeyPatch, request_kwargs: Mapping[str, object], logged_metadata_key: str
+):
+ monkeypatch.setattr(litellm, "request_correlation_in_logs", True)
+ sent_kwargs: Final = deepcopy(request_kwargs)
+
+ request: Final = _continued_chain_request(request_kwargs)
+
+ assert request["litellm_trace_id"] == "chain-session"
+ assert "litellm_session_id" not in request
+ assert request[logged_metadata_key] == {
+ **TypeAdapter(dict[str, object]).validate_python(request_kwargs[logged_metadata_key]),
+ CALLER_SESSION_ID_METADATA_KEY: "chain-session",
+ }
+ assert request_kwargs == sent_kwargs
+
+
+def test_a_continued_chain_leaves_the_turns_metadata_alone_without_request_correlation(
+ monkeypatch: pytest.MonkeyPatch,
+):
+ monkeypatch.setattr(litellm, "request_correlation_in_logs", False)
+ caller_metadata: Final = {CALLER_SESSION_ID_METADATA_KEY: "turn-session", "user_api_key": "hashed"}
+
+ request: Final = _continued_chain_request({"metadata": caller_metadata, "litellm_metadata": caller_metadata})
+
+ assert request["litellm_trace_id"] == "chain-session"
+ assert "litellm_metadata" not in request
+ assert request["metadata"] == {CALLER_SESSION_ID_METADATA_KEY: "turn-session", "user_api_key": "hashed"}
diff --git a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py
index 002a595a5e7..5db563be4b0 100644
--- a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py
+++ b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py
@@ -1,4 +1,5 @@
import json
+import sqlite3
from typing import Final
from unittest.mock import AsyncMock, patch
@@ -787,3 +788,80 @@ async def test_message_history_replays_real_key_named_tool_payloads() -> None:
tool_message: Final = result["messages"][2]
assert json.loads(assistant_message["tool_calls"][0]["function"]["arguments"]) == function_arguments
assert json.loads(tool_message["content"]) == function_output
+
+
+class _SqliteBackedPrismaDB:
+ def __init__(self, rows: list[tuple[str, str, str, str, str | None, str | None]]) -> None:
+ self._connection = sqlite3.connect(":memory:")
+ self._connection.row_factory = sqlite3.Row
+ self._connection.execute(
+ 'CREATE TABLE "LiteLLM_SpendLogs" '
+ '(request_id TEXT, api_key TEXT, session_id TEXT, "endTime" TEXT, "user" TEXT, team_id TEXT)'
+ )
+ self._connection.executemany('INSERT INTO "LiteLLM_SpendLogs" VALUES (?, ?, ?, ?, ?, ?)', rows)
+
+ async def query_raw(self, query: str, *args: object) -> list[dict[str, object]]:
+ return [dict(row) for row in self._connection.execute(query.replace("$1", "?"), args)]
+
+
+@pytest.mark.asyncio
+async def test_session_lookup_does_not_return_another_keys_rows_for_a_shared_session_id():
+ """A caller-supplied litellm_session_id can collide with another key's session; history must stay per key."""
+ sqlite_db: Final = _SqliteBackedPrismaDB(
+ [
+ ("victim-1", "victim-key", "shared-session", "2026-01-01T00:00:01", "", ""),
+ ("victim-2", "victim-key", "shared-session", "2026-01-01T00:00:02", "", ""),
+ ("attacker-1", "attacker-key", "shared-session", "2026-01-01T00:00:03", "", ""),
+ ("attacker-2", "attacker-key", "other-session", "2026-01-01T00:00:04", "", ""),
+ ]
+ )
+ fake_prisma_client = _FakePrismaClient(results=[])
+ fake_prisma_client.db = sqlite_db
+
+ with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client):
+ spend_logs = await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id("attacker-1")
+ victim_logs = await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id("victim-1")
+
+ assert [row["request_id"] for row in spend_logs] == ["attacker-1"]
+ assert [row["request_id"] for row in victim_logs] == ["victim-1", "victim-2"]
+
+
+async def _history_request_ids(sqlite_db: _SqliteBackedPrismaDB, response_id: str) -> list[str]:
+ fake_prisma_client: Final = _FakePrismaClient(results=[])
+ fake_prisma_client.db = sqlite_db
+ with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client):
+ spend_logs: Final = await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id(response_id)
+ return [row["request_id"] for row in spend_logs]
+
+
+@pytest.mark.asyncio
+async def test_session_lookup_keeps_jwt_users_and_teams_apart_when_they_share_an_empty_api_key():
+ sqlite_db: Final = _SqliteBackedPrismaDB(
+ [
+ ("alice-1", "", "shared-session", "2026-01-01T00:00:01", "alice", "team-a"),
+ ("bob-1", "", "shared-session", "2026-01-01T00:00:02", "bob", "team-a"),
+ ("service-a-1", "", "shared-session", "2026-01-01T00:00:03", "", "team-a"),
+ ("service-b-1", "", "shared-session", "2026-01-01T00:00:04", "", "team-b"),
+ ("alice-2", "", "shared-session", "2026-01-01T00:00:05", "alice", "team-a"),
+ ]
+ )
+
+ assert await _history_request_ids(sqlite_db, "alice-1") == ["alice-1", "alice-2"]
+ assert await _history_request_ids(sqlite_db, "bob-1") == ["bob-1"]
+ assert await _history_request_ids(sqlite_db, "service-a-1") == ["service-a-1"]
+ assert await _history_request_ids(sqlite_db, "service-b-1") == ["service-b-1"]
+
+
+@pytest.mark.asyncio
+async def test_session_lookup_follows_the_owner_across_keys_and_keeps_ownerless_keys_apart():
+ sqlite_db: Final = _SqliteBackedPrismaDB(
+ [
+ ("alice-old-key", "hashed-old", "shared-session", "2026-01-01T00:00:01", "alice", "team-a"),
+ ("ownerless-1", "hashed-other", "shared-session", "2026-01-01T00:00:02", "", ""),
+ ("alice-new-key", "hashed-new", "shared-session", "2026-01-01T00:00:03", "alice", "team-a"),
+ ("ownerless-2", "hashed-another", "shared-session", "2026-01-01T00:00:04", "", ""),
+ ]
+ )
+
+ assert await _history_request_ids(sqlite_db, "alice-new-key") == ["alice-old-key", "alice-new-key"]
+ assert await _history_request_ids(sqlite_db, "ownerless-2") == ["ownerless-2"]
diff --git a/tests/unit/responses/test_streaming_iterator_error_events.py b/tests/unit/responses/test_streaming_iterator_error_events.py
index c37e57a43ca..001be7e1fbd 100644
--- a/tests/unit/responses/test_streaming_iterator_error_events.py
+++ b/tests/unit/responses/test_streaming_iterator_error_events.py
@@ -20,14 +20,18 @@ Pydantic ValidationError (previously typed as Optional[str]).
"""
import json
+from datetime import datetime
from importlib import import_module
+from typing import Final
from unittest.mock import Mock, patch
+import httpx
import pytest
import litellm
from litellm.exceptions import MidStreamFallbackError
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+from litellm.litellm_core_utils.litellm_logging import _get_provider_request_id
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
from litellm.responses.streaming_iterator import (
_ERROR_CODE_HTTP_STATUS,
@@ -64,6 +68,52 @@ def _make_error_chunk(error_type: str, code: str, message: str = "err") -> Error
return ErrorEvent(type=ResponsesAPIStreamEvents.ERROR, sequence_number=0, error=error_obj)
+@pytest.mark.parametrize("provider", ("openai", "azure", "bedrock"))
+@pytest.mark.parametrize("header", ("x-amzn-RequestId", "x-request-id", None))
+@pytest.mark.parametrize(
+ "code,status", (("server_error", 500), ("rate_limit_exceeded", 429), ("invalid_request_error", 400))
+)
+def test_stream_error_preserves_provider_request_id(provider: str, header: str | None, code: str, status: int) -> None:
+ stream_headers: Final = {"retry-after": "7", "content-type": "text/event-stream"}
+ response: Final = httpx.Response(
+ 200, headers={**stream_headers, header: "provider-request-id"} if header else stream_headers
+ )
+ logging: Final = LiteLLMLoggingObj(
+ model="test-model",
+ messages=[],
+ stream=True,
+ call_type="responses",
+ start_time=datetime(2026, 1, 1),
+ litellm_call_id="test-call",
+ function_id="test-function",
+ )
+ iterator: Final = BaseResponsesAPIStreamingIterator(
+ response=response,
+ model="test-model",
+ responses_api_provider_config=None,
+ logging_obj=logging,
+ custom_llm_provider=provider,
+ )
+ expected_id: Final = "provider-request-id" if header else None
+ expected_headers: Final = {header.lower(): "provider-request-id"} if header else {}
+ mapped: Final = iterator._map_error_event_exception({"type": code, "code": code, "message": "failure"})
+ assert _get_provider_request_id(mapped) == expected_id
+ assert dict(getattr(mapped, "litellm_response_headers", None) or {}) == expected_headers
+ expected_status: Final = 503 if provider == "bedrock" and status == 500 else status
+ with pytest.raises((litellm.BadRequestError, MidStreamFallbackError)) as raised:
+ iterator._maybe_raise_for_error_event(_make_error_chunk(code, code, "failure"))
+ assert raised.value.status_code == expected_status
+ assert _get_provider_request_id(raised.value) == expected_id
+ assert response.status_code == 200
+ if status in (429, 500):
+ assert isinstance(raised.value, MidStreamFallbackError)
+ assert dict(raised.value.response.headers) == expected_headers
+ assert raised.value.response.status_code == expected_status
+ assert str(raised.value.response.request.url) == f"https://{provider}.com/v1/"
+ assert raised.value.generated_content == ""
+ assert raised.value.is_pre_first_chunk is True
+
+
def test_maybe_raise_for_error_event_wraps_unknown_error_in_mid_stream_fallback():
iterator = _make_iterator()
chunk = _make_error_chunk("server_error", "internal_error", "something went wrong")
diff --git a/tests/unit/router_strategy/test_budget_limiter.py b/tests/unit/router_strategy/test_budget_limiter.py
index 685ea46547e..380aa573b3c 100644
--- a/tests/unit/router_strategy/test_budget_limiter.py
+++ b/tests/unit/router_strategy/test_budget_limiter.py
@@ -14,6 +14,7 @@ import pytest
from litellm import Router
from litellm.caching.caching import DualCache
+from litellm.caching.in_memory_cache import InMemoryCache
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
from litellm.types.utils import BudgetConfig
import sys, os, asyncio, time, random
@@ -232,7 +233,11 @@ async def test_a_monthly_router_budget_ends_a_calendar_month_after_it_opens(
"""
monkeypatch.setattr(litellm, "tag_budget_config", tag_budget_config)
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
- limiter = RouterBudgetLimiting(dual_cache=DualCache(), clock=charged_at.timestamp, **limiter_config)
+ limiter = RouterBudgetLimiting(
+ dual_cache=DualCache(in_memory_cache=InMemoryCache(clock=charged_at.timestamp)),
+ clock=charged_at.timestamp,
+ **limiter_config,
+ )
await limiter.dual_cache.async_set_cache(key=start_time_key, value=opened_at.timestamp())
await _log_success(
@@ -243,9 +248,8 @@ async def test_a_monthly_router_budget_ends_a_calendar_month_after_it_opens(
},
)
- assert limiter.redis_increment_operation_queue == [
- {"key": spend_key, "increment_value": 0.25, "ttl": int((closes_at - charged_at).total_seconds())}
- ]
+ assert await limiter.dual_cache.async_get_cache(spend_key) == 0.25
+ assert limiter.dual_cache.in_memory_cache.ttl_dict[spend_key] == closes_at.timestamp()
@pytest.mark.asyncio
@@ -362,24 +366,18 @@ async def test_new_budget_window_sets_spend_and_start_time() -> None:
@pytest.mark.asyncio
-async def test_spend_increment_updates_memory_and_queues_redis_operation() -> None:
- provider_budget: Final = RouterBudgetLimiting(dual_cache=DualCache(), provider_budget_config={})
- await provider_budget.dual_cache.async_set_cache(key="provider_spend:openai:1d", value=1.0, ttl=86400)
-
- await provider_budget._increment_spend_in_current_window(
- spend_key="provider_spend:openai:1d",
- response_cost=0.5,
- ttl=86400,
+async def test_redis_less_limiter_tracks_spend_without_queueing_redis_increments(disable_budget_sync) -> None:
+ """With no Redis there is nothing to flush to, so queued increments would be retained for the life of the process."""
+ limiter: Final = RouterBudgetLimiting(
+ dual_cache=DualCache(),
+ provider_budget_config={"openai": {"budget_limit": 10.0, "time_period": "1d"}},
)
- assert await provider_budget.dual_cache.async_get_cache("provider_spend:openai:1d") == 1.5
- assert provider_budget.redis_increment_operation_queue == [
- {
- "key": "provider_spend:openai:1d",
- "increment_value": 0.5,
- "ttl": 86400,
- }
- ]
+ for _ in range(3):
+ await _log_success(limiter, _success_kwargs(provider_in_litellm_params="openai", provider_in_payload="openai"))
+
+ assert await limiter.dual_cache.async_get_cache("provider_spend:openai:1d") == pytest.approx(0.75)
+ assert limiter.redis_increment_operation_queue == []
@pytest.mark.asyncio
diff --git a/tests/unit/router_utils/pre_call_checks/test_continuation_prefill_check.py b/tests/unit/router_utils/pre_call_checks/test_continuation_prefill_check.py
new file mode 100644
index 00000000000..8fa1c799863
--- /dev/null
+++ b/tests/unit/router_utils/pre_call_checks/test_continuation_prefill_check.py
@@ -0,0 +1,288 @@
+import pytest
+
+import litellm
+from litellm.router_utils.pre_call_checks.continuation_prefill_check import (
+ MID_STREAM_CONTINUATION_KWARG,
+ MID_STREAM_CONTINUATION_MARKER,
+ ContinuationOutput,
+ ContinuationPrefillDeploymentCheck,
+ continuation_output_ceilings,
+ deployment_supports_prefill,
+)
+
+PREFILL_MODEL = "anthropic/prefill-capable-test-model"
+NON_PREFILL_MODEL = "openai/prefill-unknown-test-model"
+UNMAPPED_MODEL = "openai/unmapped-test-model"
+
+
+@pytest.fixture(autouse=True)
+def _synthetic_cost_map_entries() -> None:
+ litellm.register_model(
+ {
+ PREFILL_MODEL: {"litellm_provider": "anthropic", "mode": "chat", "supports_assistant_prefill": True},
+ NON_PREFILL_MODEL: {"litellm_provider": "openai", "mode": "chat"},
+ },
+ persist_across_reloads=False,
+ )
+ litellm.get_model_info.cache_clear()
+
+
+def _deployment(model: str, dep_id: str) -> dict[str, object]:
+ return {"litellm_params": {"model": model}, "model_info": {"id": dep_id}}
+
+
+def test_deployment_supports_prefill_reads_capability():
+ assert deployment_supports_prefill(_deployment(PREFILL_MODEL, "a")) is True
+ assert deployment_supports_prefill(_deployment(NON_PREFILL_MODEL, "b")) is False
+ assert deployment_supports_prefill({"litellm_params": {"model": PREFILL_MODEL}}) is True
+ assert deployment_supports_prefill({"litellm_params": {"model": PREFILL_MODEL}, "model_info": "bogus"}) is True
+
+
+def test_deployment_model_info_override_wins_over_cost_map():
+ # model_info True opts in a model that is not in the cost map
+ assert (
+ deployment_supports_prefill(
+ {"litellm_params": {"model": "vendor/custom-model"}, "model_info": {"supports_assistant_prefill": True}}
+ )
+ is True
+ )
+ # model_info False opts out a model the cost map would otherwise allow
+ assert (
+ deployment_supports_prefill(
+ {"litellm_params": {"model": PREFILL_MODEL}, "model_info": {"supports_assistant_prefill": False}}
+ )
+ is False
+ )
+ # model_info without the key falls through to the cost map
+ assert deployment_supports_prefill(_deployment(PREFILL_MODEL, "z")) is True
+
+
+@pytest.mark.parametrize(
+ "deployment,expected",
+ [
+ ({"litellm_params": {"model": "bedrock/arn:custom-import", "base_model": PREFILL_MODEL}}, True),
+ ({"litellm_params": {"model": "azure/my-deployment"}, "model_info": {"base_model": PREFILL_MODEL}}, True),
+ (
+ {
+ "litellm_params": {"model": "azure/my-deployment", "base_model": NON_PREFILL_MODEL},
+ "model_info": {"base_model": PREFILL_MODEL},
+ },
+ True,
+ ),
+ (
+ {
+ "litellm_params": {"model": PREFILL_MODEL, "base_model": NON_PREFILL_MODEL},
+ "model_info": {"base_model": NON_PREFILL_MODEL},
+ },
+ False,
+ ),
+ ],
+ ids=[
+ "litellm_params_base_model",
+ "model_info_base_model",
+ "model_info_base_model_wins",
+ "base_model_wins_over_model",
+ ],
+)
+def test_deployment_supports_prefill_reads_the_cost_map_key_of_the_base_model(deployment, expected):
+ assert deployment_supports_prefill(deployment) is expected
+
+
+def test_deployment_supports_prefill_rejects_malformed_deployments():
+ assert deployment_supports_prefill({}) is False
+ assert deployment_supports_prefill({"litellm_params": {}}) is False
+ assert deployment_supports_prefill("not-a-dict") is False
+
+
+@pytest.mark.asyncio
+async def test_filter_is_noop_without_continuation_marker():
+ """A normal (non-continuation) request must be passed through untouched, even
+ if some deployments cannot prefill."""
+ check = ContinuationPrefillDeploymentCheck()
+ deployments = [_deployment(PREFILL_MODEL, "a"), _deployment(NON_PREFILL_MODEL, "b")]
+
+ for request_kwargs in ({}, None, {MID_STREAM_CONTINUATION_KWARG: False}):
+ result = await check.async_filter_deployments(
+ model="group", healthy_deployments=deployments, messages=None, request_kwargs=request_kwargs
+ )
+ assert result == deployments
+
+
+@pytest.mark.asyncio
+async def test_filter_ignores_forged_client_flag():
+ """A client cannot steer routing: a plain truthy value under the marker key
+ (which the proxy could forward from the request body) is not the internal
+ sentinel, so the filter leaves the deployment list untouched."""
+ check = ContinuationPrefillDeploymentCheck()
+ deployments = [_deployment(PREFILL_MODEL, "a"), _deployment(NON_PREFILL_MODEL, "b")]
+
+ for forged in (True, "true", 1, {"any": "json"}):
+ result = await check.async_filter_deployments(
+ model="group",
+ healthy_deployments=deployments,
+ messages=None,
+ request_kwargs={MID_STREAM_CONTINUATION_KWARG: forged},
+ )
+ assert result == deployments
+
+
+@pytest.mark.asyncio
+async def test_filter_keeps_only_prefill_capable_on_continuation():
+ check = ContinuationPrefillDeploymentCheck()
+ deployments = [_deployment(PREFILL_MODEL, "a"), _deployment(NON_PREFILL_MODEL, "b")]
+
+ result = await check.async_filter_deployments(
+ model="group",
+ healthy_deployments=deployments,
+ messages=None,
+ request_kwargs={MID_STREAM_CONTINUATION_KWARG: MID_STREAM_CONTINUATION_MARKER},
+ )
+ assert [d["model_info"]["id"] for d in result] == ["a"]
+
+
+@pytest.mark.asyncio
+async def test_filter_empties_group_when_no_prefill_capable_deployment():
+ """No prefill-capable deployment -> empty result, so the router advances the
+ fallback chain and ultimately surfaces the original error."""
+ check = ContinuationPrefillDeploymentCheck()
+ deployments = [_deployment(NON_PREFILL_MODEL, "b"), _deployment(UNMAPPED_MODEL, "c")]
+
+ result = await check.async_filter_deployments(
+ model="group",
+ healthy_deployments=deployments,
+ messages=None,
+ request_kwargs={MID_STREAM_CONTINUATION_KWARG: MID_STREAM_CONTINUATION_MARKER},
+ )
+ assert result == []
+
+
+@pytest.mark.asyncio
+async def test_filter_drops_deployments_whose_own_params_constrain_the_continuation():
+ check = ContinuationPrefillDeploymentCheck()
+ plain = _deployment(PREFILL_MODEL, "plain")
+ constrained = [
+ {"litellm_params": {"model": PREFILL_MODEL, **params}, "model_info": {"id": dep_id}}
+ for dep_id, params in (
+ ("json", {"response_format": {"type": "json_object"}}),
+ ("two_choices", {"extra_body": {"n": 2}}),
+ ("forced_tool", {"tool_choice": {"type": "function", "function": {"name": "f"}}}),
+ ("merged_reasoning", {"merge_reasoning_content_in_choices": True}),
+ )
+ ]
+
+ result = await check.async_filter_deployments(
+ model="group",
+ healthy_deployments=[plain, *constrained],
+ messages=None,
+ request_kwargs={MID_STREAM_CONTINUATION_KWARG: MID_STREAM_CONTINUATION_MARKER},
+ )
+ assert [d["model_info"]["id"] for d in result] == ["plain"]
+
+
+def _text_chunk(text: str) -> litellm.ModelResponseStream:
+ return litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": text}}])
+
+
+@pytest.mark.parametrize(
+ "prefill_content,fallback_texts,expected_texts",
+ [
+ ("Hello", ("", "Hel", "lo wor", "ld"), ("", "", " wor", "ld")),
+ ("Hello", ("Hello",), ("",)),
+ ("Hello", (" world",), (" world",)),
+ ("Hello", ("Hel", "p me"), ("Hel", "p me")),
+ ("Hello", ("Hel",), ("Hel",)),
+ ([{"type": "text", "text": "Sure: "}, {"type": "text", "text": "Hello"}], ("Sure: Hello", " world"), ("", " world")),
+ ],
+ ids=[
+ "echo_split_across_chunks_is_dropped",
+ "exact_echo_then_end_is_dropped",
+ "continuation_without_echo_passes_through",
+ "near_echo_is_kept_in_order",
+ "stream_ending_mid_match_is_kept",
+ "content_block_prefill_echo_is_dropped",
+ ],
+)
+@pytest.mark.asyncio
+async def test_continuation_output_drops_only_a_full_echo_of_the_prefill(
+ prefill_content, fallback_texts, expected_texts
+):
+ output = ContinuationOutput(
+ "Hello", {"role": "assistant", "content": prefill_content, "prefix": True}, "group", ()
+ )
+
+ async def fallback_stream():
+ for text in fallback_texts:
+ yield _text_chunk(text)
+
+ released = [chunk async for chunk in output.splice(fallback_stream())]
+
+ assert tuple(chunk.choices[0].delta.content for chunk in released) == expected_texts
+
+
+@pytest.mark.parametrize(
+ "prefill,fallback_texts,expected_texts",
+ [
+ ("Transformers use ", ("self-attention",), ("self-attention",)),
+ ("Hello ", ("Hello ", " world"), ("", "world")),
+ ("Hello \n", ("\nworld",), ("\nworld",)),
+ ("A.\n\n", ("\n Hi", "\nthere"), (" Hi", "\nthere")),
+ ("\n\n", ("\nHi",), ("\nHi",)),
+ ("\nHello\n", ("\n",), ("",)),
+ ],
+ ids=[
+ "continuation_without_whitespace_passes_through",
+ "echo_then_re_emitted_space_are_both_dropped",
+ "different_whitespace_is_kept",
+ "seam_ends_at_the_first_divergence",
+ "whitespace_only_prefill_has_no_seam",
+ "re_emitted_whitespace_at_stream_end_is_dropped",
+ ],
+)
+@pytest.mark.asyncio
+async def test_continuation_output_drops_only_whitespace_the_client_already_has(
+ prefill, fallback_texts, expected_texts
+):
+ output = ContinuationOutput(prefill, {"role": "assistant", "content": prefill, "prefix": True}, "group", ())
+
+ async def fallback_stream():
+ for text in fallback_texts:
+ yield _text_chunk(text)
+
+ released = [chunk async for chunk in output.splice(fallback_stream())]
+
+ assert tuple(chunk.choices[0].delta.content for chunk in released) == expected_texts
+
+
+@pytest.mark.parametrize(
+ "request_kwargs,expected",
+ [
+ ({"max_tokens": 100}, {"max_tokens": 70}),
+ ({"max_tokens": 100, "max_completion_tokens": 40}, {"max_tokens": 70, "max_completion_tokens": 10}),
+ ({}, {}),
+ (
+ {"max_tokens": 100, "extra_body": {"max_tokens": 50, "top_k": 3}},
+ {"max_tokens": 70, "extra_body": {"max_tokens": 20, "top_k": 3}},
+ ),
+ ({"extra_body": {"max_completion_tokens": 45}}, {"extra_body": {"max_completion_tokens": 15}}),
+ ({"extra_body": {"top_k": 3}}, {}),
+ ],
+ ids=["max_tokens", "both_ceilings", "no_ceiling", "extra_body_max_tokens", "extra_body_only", "extra_body_no_ceiling"],
+)
+def test_continuation_output_ceilings_take_the_emitted_tokens_off_every_ceiling_the_request_sends(
+ request_kwargs, expected
+):
+ assert continuation_output_ceilings(request_kwargs, 30) == expected
+
+
+@pytest.mark.parametrize(
+ "request_kwargs",
+ [
+ {"max_tokens": 30},
+ {"max_tokens": 20},
+ {"max_tokens": 100, "max_completion_tokens": 30},
+ {"max_tokens": 100, "extra_body": {"max_tokens": 30}},
+ ],
+ ids=["max_tokens_met", "max_tokens_exceeded", "one_of_two_met", "extra_body_met"],
+)
+def test_continuation_output_ceilings_decline_once_a_ceiling_has_no_tokens_left(request_kwargs):
+ assert continuation_output_ceilings(request_kwargs, 30) is None
diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py
index 5a6da70d756..6bc52ab3dfc 100644
--- a/tests/unit/test_router/test_router.py
+++ b/tests/unit/test_router/test_router.py
@@ -1,4 +1,5 @@
import ast, asyncio, importlib, time
+import contextlib
import copy
import functools
import gc
@@ -11,7 +12,8 @@ import threading
import warnings
from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Mapping
from datetime import datetime, timedelta, timezone
-from types import SimpleNamespace
+from pathlib import Path
+from types import MappingProxyType, SimpleNamespace
from typing import TYPE_CHECKING, Final, Literal, cast
from unittest.mock import AsyncMock, MagicMock, Mock, patch
@@ -82,6 +84,7 @@ from litellm.types.router import (
ModelInfo,
PreRoutingHookResponse,
RetryPolicy,
+ RouterRateLimitError,
RoutingContext,
)
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
@@ -2932,6 +2935,1216 @@ async def test_acompletion_streaming_iterator():
print("\n=== All tests passed! ===")
+def _make_midstream_source(error, chunks=None, first_chunk=None):
+ from unittest.mock import MagicMock
+
+ opening_chunk = (
+ first_chunk if first_chunk is not None else MagicMock(choices=[MagicMock(delta=MagicMock(content="Hello"))])
+ )
+
+ class _Source:
+ def __init__(self):
+ self.index = 0
+ self.chunks = chunks if chunks is not None else []
+ self.model = "gpt-4"
+ self.custom_llm_provider = "openai"
+ self.logging_obj = MagicMock()
+
+ def __aiter__(self):
+ return self
+
+ async def __anext__(self):
+ if self.index == 0:
+ self.index += 1
+ return opening_chunk
+ raise error
+
+ async def wait_for_failure_logging(self, timeout):
+ return None
+
+ return _Source()
+
+
+class _FakeFallbackStream:
+ def __init__(self, item):
+ self._item = item
+ self._done = False
+
+ def __aiter__(self):
+ return self
+
+ async def __anext__(self):
+ if self._done:
+ raise StopAsyncIteration
+ self._done = True
+ return self._item
+
+
+@pytest.mark.asyncio
+async def test_acompletion_streaming_iterator_continues_after_content_when_eligible():
+ """Flag on + plain-text break: the router re-enters the fallback chain with an
+ assistant-prefill continuation and the mid-stream marker, then streams the
+ fallback's output instead of re-raising."""
+ from unittest.mock import AsyncMock, MagicMock, patch
+
+ from litellm.exceptions import MidStreamFallbackError
+ from litellm.router_utils.pre_call_checks.continuation_prefill_check import (
+ MID_STREAM_CONTINUATION_KWARG,
+ MID_STREAM_CONTINUATION_MARKER,
+ )
+
+ router = litellm.Router(
+ model_list=[
+ {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}},
+ {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}},
+ ],
+ fallbacks=[{"gpt-4": ["backup"]}],
+ enable_mid_stream_fallback_continuation=True,
+ )
+
+ error = MidStreamFallbackError(
+ message="Connection lost",
+ model="gpt-4",
+ llm_provider="openai",
+ generated_content="Hello",
+ is_pre_first_chunk=False,
+ emitted_disqualifying_content=False,
+ )
+ source = _make_midstream_source(error)
+ fallback_chunk = litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "world"}}])
+ initial_kwargs = {"model": "gpt-4", "stream": True}
+
+ with patch.object(
+ router,
+ "async_function_with_fallbacks_common_utils",
+ new=AsyncMock(return_value=_FakeFallbackStream(fallback_chunk)),
+ ) as mock_fallback:
+ result = await router._acompletion_streaming_iterator(
+ model_response=source, messages=[{"role": "user", "content": "Hi"}], initial_kwargs=initial_kwargs
+ )
+ collected = [chunk async for chunk in result]
+
+ mock_fallback.assert_awaited_once()
+ passed_kwargs = mock_fallback.await_args.kwargs["kwargs"]
+ assert passed_kwargs[MID_STREAM_CONTINUATION_KWARG] is MID_STREAM_CONTINUATION_MARKER
+ assert passed_kwargs["messages"][-1] == {"role": "assistant", "content": "Hello", "prefix": True}
+ assert fallback_chunk in collected
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "continuation_text,joined_output_trips_rule",
+ [("6789", True), ("6788", False)],
+ ids=["rule_trips_only_on_the_joined_text", "clean_continuation_streams"],
+)
+async def test_acompletion_streaming_iterator_continuation_runs_post_call_rules_on_the_joined_text(
+ monkeypatch, continuation_text, joined_output_trips_rule
+):
+ from unittest.mock import AsyncMock, patch
+
+ from litellm.exceptions import MidStreamFallbackError
+
+ monkeypatch.setattr(litellm, "post_call_rules", [lambda output: "123-45-6789" not in output])
+ router = litellm.Router(
+ model_list=[
+ {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}},
+ {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}},
+ ],
+ fallbacks=[{"gpt-4": ["backup"]}],
+ enable_mid_stream_fallback_continuation=True,
+ )
+ error = MidStreamFallbackError(
+ message="Connection lost", model="gpt-4", llm_provider="openai", generated_content="123-45-",
+ is_pre_first_chunk=False, emitted_disqualifying_content=False,
+ )
+ fallback_chunk = litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": continuation_text}}])
+
+ with patch.object(
+ router,
+ "async_function_with_fallbacks_common_utils",
+ new=AsyncMock(return_value=_FakeFallbackStream(fallback_chunk)),
+ ):
+ result = await router._acompletion_streaming_iterator(
+ model_response=_make_midstream_source(error),
+ messages=[{"role": "user", "content": "Hi"}],
+ initial_kwargs={"model": "gpt-4", "stream": True},
+ )
+ if joined_output_trips_rule:
+ with pytest.raises(litellm.APIResponseValidationError):
+ async for _ in result:
+ pass
+ else:
+ assert fallback_chunk in [chunk async for chunk in result]
+
+
+@pytest.mark.asyncio
+async def test_a_second_cut_runs_post_call_rules_only_on_the_text_the_caller_received(monkeypatch):
+ from unittest.mock import AsyncMock, patch
+
+ from litellm.exceptions import MidStreamFallbackError
+ from litellm.router_utils.pre_call_checks.continuation_prefill_check import (
+ MID_STREAM_CONTINUATION_KWARG,
+ MID_STREAM_CONTINUATION_MARKER,
+ )
+
+ checked_texts = []
+ monkeypatch.setattr(litellm, "post_call_rules", [lambda output: checked_texts.append(output) is None])
+ router = litellm.Router(
+ model_list=[
+ {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}},
+ {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}},
+ ],
+ fallbacks=[{"gpt-4": ["backup"]}],
+ enable_mid_stream_fallback_continuation=True,
+ )
+
+ def cut(generated_content):
+ return MidStreamFallbackError(
+ message="Connection lost", model="gpt-4", llm_provider="openai", generated_content=generated_content,
+ is_pre_first_chunk=False, emitted_disqualifying_content=False,
+ )
+
+ messages = [{"role": "user", "content": "Hi"}]
+ middle_leg = await router._acompletion_streaming_iterator(
+ model_response=_make_midstream_source(
+ cut(" there"),
+ first_chunk=litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": " there"}}]),
+ ),
+ messages=[*messages, {"role": "assistant", "content": "Hello", "prefix": True}],
+ initial_kwargs={"model": "gpt-4", "stream": True, MID_STREAM_CONTINUATION_KWARG: MID_STREAM_CONTINUATION_MARKER},
+ )
+ backup_leg = _FakeFallbackStream(
+ litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": " friend"}}])
+ )
+
+ with patch.object(
+ router, "async_function_with_fallbacks_common_utils", new=AsyncMock(side_effect=[middle_leg, backup_leg])
+ ):
+ result = await router._acompletion_streaming_iterator(
+ model_response=_make_midstream_source(cut("Hello")),
+ messages=messages,
+ initial_kwargs={"model": "gpt-4", "stream": True},
+ )
+ _ = [chunk async for chunk in result]
+
+ assert checked_texts == ["Hello there", "Hello there friend"]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "fallback_texts",
+ [("Hel", "lo", " world"), (" world",)],
+ ids=["fallback_echoes_the_prefill", "fallback_continues_after_the_prefill"],
+)
+async def test_acompletion_streaming_iterator_continuation_never_repeats_the_partial_output(fallback_texts):
+ from unittest.mock import AsyncMock, patch
+
+ from litellm.exceptions import MidStreamFallbackError
+
+ router = litellm.Router(
+ model_list=[
+ {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}},
+ {"model_name": "backup", "litellm_params": {"model": "mistral/mistral-small-latest", "api_key": "k2"}},
+ ],
+ fallbacks=[{"gpt-4": ["backup"]}],
+ enable_mid_stream_fallback_continuation=True,
+ )
+ error = MidStreamFallbackError(
+ message="Connection lost", model="gpt-4", llm_provider="openai", generated_content="Hello",
+ is_pre_first_chunk=False, emitted_disqualifying_content=False,
+ )
+
+ async def fallback_stream():
+ for text in fallback_texts:
+ yield litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": text}}])
+
+ with patch.object(
+ router, "async_function_with_fallbacks_common_utils", new=AsyncMock(return_value=fallback_stream())
+ ):
+ result = await router._acompletion_streaming_iterator(
+ model_response=_make_midstream_source(error),
+ messages=[{"role": "user", "content": "Hi"}],
+ initial_kwargs={"model": "gpt-4", "stream": True},
+ )
+ collected = [chunk async for chunk in result]
+
+ continuation = [chunk for chunk in collected if isinstance(chunk, litellm.ModelResponseStream)]
+ assert "".join(chunk.choices[0].delta.content for chunk in continuation) == " world"
+
+
+@pytest.mark.asyncio
+async def test_acompletion_streaming_iterator_continuation_stays_in_the_primary_stream():
+ from unittest.mock import AsyncMock, patch
+
+ from litellm.exceptions import MidStreamFallbackError
+
+ router = litellm.Router(
+ model_list=[
+ {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}},
+ {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}},
+ ],
+ fallbacks=[{"gpt-4": ["backup"]}],
+ enable_mid_stream_fallback_continuation=True,
+ )
+ error = MidStreamFallbackError(
+ message="Connection lost", model="gpt-4", llm_provider="openai", generated_content="Hello",
+ is_pre_first_chunk=False, emitted_disqualifying_content=False,
+ )
+ primary_chunk = litellm.ModelResponseStream(
+ id="chatcmpl-primary", created=111, choices=[{"index": 0, "delta": {"role": "assistant", "content": "Hello"}}]
+ )
+
+ async def fallback_stream():
+ yield litellm.ModelResponseStream(
+ id="msg-fallback", created=222, choices=[{"index": 0, "delta": {"role": "assistant", "content": " wor"}}]
+ )
+ yield litellm.ModelResponseStream(
+ id="msg-fallback", created=222, choices=[{"index": 0, "delta": {"content": "ld"}}]
+ )
+
+ with patch.object(
+ router, "async_function_with_fallbacks_common_utils", new=AsyncMock(return_value=fallback_stream())
+ ):
+ result = await router._acompletion_streaming_iterator(
+ model_response=_make_midstream_source(error, chunks=[primary_chunk]),
+ messages=[{"role": "user", "content": "Hi"}],
+ initial_kwargs={"model": "gpt-4", "stream": True},
+ )
+ collected = [chunk async for chunk in result]
+
+ continuation = [chunk for chunk in collected if isinstance(chunk, litellm.ModelResponseStream)]
+ assert [(chunk.id, chunk.created, chunk.choices[0].delta.role) for chunk in continuation] == [
+ ("chatcmpl-primary", 111, None),
+ ("chatcmpl-primary", 111, None),
+ ]
+ assert "".join(chunk.choices[0].delta.content for chunk in continuation) == " world"
+
+
+def _continuation_router():
+ return litellm.Router(
+ model_list=[
+ {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}},
+ {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}},
+ ],
+ fallbacks=[{"gpt-4": ["backup"]}],
+ enable_mid_stream_fallback_continuation=True,
+ )
+
+
+def _usage_chunk(prompt_tokens, completion_tokens):
+ from litellm.types.utils import Usage
+
+ return litellm.ModelResponseStream(
+ model="gpt-4",
+ choices=[],
+ usage=Usage(
+ prompt_tokens=prompt_tokens,
+ completion_tokens=completion_tokens,
+ total_tokens=prompt_tokens + completion_tokens,
+ ),
+ )
+
+
+async def _drain_with_fallback(router, error, primary_chunks, messages, fallback_items):
+ from unittest.mock import AsyncMock, patch
+
+ async def fallback_stream():
+ for item in fallback_items:
+ yield item
+
+ with patch.object(
+ router, "async_function_with_fallbacks_common_utils", new=AsyncMock(return_value=fallback_stream())
+ ) as mock_fallback:
+ result = await router._acompletion_streaming_iterator(
+ model_response=_make_midstream_source(error, chunks=primary_chunks),
+ messages=messages,
+ initial_kwargs={"model": "gpt-4", "stream": True},
+ )
+ collected = [chunk async for chunk in result]
+ streamed = [chunk for chunk in collected if isinstance(chunk, litellm.ModelResponseStream)]
+ return mock_fallback.await_args.kwargs["kwargs"], streamed
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("primary_reports_usage", [True, False], ids=["usage_reported", "usage_not_reported"])
+async def test_acompletion_streaming_iterator_continuation_usage_sums_both_legs(primary_reports_usage):
+ from litellm.exceptions import MidStreamFallbackError
+
+ messages = [{"role": "user", "content": "Explain transformers in one paragraph."}]
+ partial = "Transformers use "
+ error = MidStreamFallbackError(
+ message="Connection lost", model="gpt-4", llm_provider="openai", generated_content=partial,
+ is_pre_first_chunk=False, emitted_disqualifying_content=False,
+ )
+ content_chunk = litellm.ModelResponseStream(
+ model="gpt-4", choices=[{"index": 0, "delta": {"role": "assistant", "content": partial}}]
+ )
+ primary_chunks = [content_chunk, _usage_chunk(21, 3)] if primary_reports_usage else [content_chunk]
+ primary_usage = litellm.stream_chunk_builder(chunks=list(primary_chunks), messages=messages).usage
+
+ _, streamed = await _drain_with_fallback(
+ _continuation_router(), error, primary_chunks, messages, [_usage_chunk(91, 5)]
+ )
+
+ assert primary_usage.prompt_tokens > 0
+ assert streamed[-1].usage.prompt_tokens == primary_usage.prompt_tokens + 91
+ assert streamed[-1].usage.completion_tokens == primary_usage.completion_tokens + 5
+
+
+class _CutSSE(httpx.AsyncByteStream):
+ def __init__(self, content: str) -> None:
+ self._content: Final = content
+
+ async def __aiter__(self) -> AsyncIterator[bytes]:
+ yield f'data: {json.dumps({"id": "cut", "object": "chat.completion.chunk", "created": 1700000000, "model": "m", "choices": [{"index": 0, "delta": {"role": "assistant", "content": self._content}, "finish_reason": None}]})}\n\n'.encode()
+ raise httpx.RemoteProtocolError("peer closed connection without sending complete message body")
+
+
+_BACKUP_SSE: Final = "".join(
+ f"data: {json.dumps(chunk)}\n\n"
+ for chunk in (
+ {"id": "b", "object": "chat.completion.chunk", "created": 1700000000, "model": "mistral-large-latest",
+ "choices": [{"index": 0, "delta": {"role": "assistant", "content": " self-attention."}, "finish_reason": None}]},
+ {"id": "b", "object": "chat.completion.chunk", "created": 1700000000, "model": "mistral-large-latest",
+ "choices": [{"index": 0, "delta": {"content": ""}, "finish_reason": "stop"}],
+ "usage": {"prompt_tokens": 91, "completion_tokens": 5, "total_tokens": 96}},
+ )
+) + "data: [DONE]\n\n"
+
+
+class _SpendRecorder(CustomLogger):
+ def __init__(self) -> None:
+ super().__init__()
+ self.successes: Final[list[dict[str, object]]] = []
+ self.failures: Final[list[dict[str, object]]] = []
+ self.spend_rows: Final[list[dict[str, object]]] = []
+ self.succeeded: Final = asyncio.Event()
+
+ async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None:
+ from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload
+
+ self.successes.append(kwargs["standard_logging_object"])
+ self.spend_rows.append(get_logging_payload(kwargs, response_obj, start_time, end_time))
+ self.succeeded.set()
+
+ async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None:
+ self.failures.append(kwargs["standard_logging_object"])
+
+
+class _SlowFailureRecorder(_SpendRecorder):
+ def __init__(self) -> None:
+ super().__init__()
+ self.failed_deployment_ids: Final[list[str]] = []
+
+ async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None:
+ await asyncio.sleep(0.05)
+ self.failed_deployment_ids.append(kwargs["litellm_params"]["metadata"]["model_info"]["id"])
+ await super().async_log_failure_event(kwargs, response_obj, start_time, end_time)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("cut_groups", [("primary",), ("primary", "middle")], ids=["one_break", "two_breaks"])
+async def test_continued_stream_success_log_bills_every_cut_leg(
+ monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, cut_groups: tuple[str, ...]
+) -> None:
+ monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+ recorder: Final = _SpendRecorder()
+ monkeypatch.setattr(litellm, "callbacks", [recorder])
+ respx_mock.post("https://api.openai.com/v1/chat/completions").mock(
+ return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=_CutSSE("Transformers use"))
+ )
+ respx_mock.post(host="api.deepseek.com").mock(
+ return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=_CutSSE(" learned"))
+ )
+ respx_mock.post(host="api.mistral.ai").mock(
+ return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, text=_BACKUP_SSE)
+ )
+ router: Final = Router(
+ model_list=[
+ {
+ "model_name": "primary",
+ "litellm_params": {
+ "model": "openai/gpt-4o",
+ "api_key": "k1",
+ "input_cost_per_token": 0.001,
+ "output_cost_per_token": 0.002,
+ },
+ },
+ {
+ "model_name": "middle",
+ "litellm_params": {
+ "model": "deepseek/deepseek-chat",
+ "api_key": "k2",
+ "input_cost_per_token": 0.003,
+ "output_cost_per_token": 0.004,
+ },
+ },
+ {"model_name": "backup", "litellm_params": {"model": "mistral/mistral-large-latest", "api_key": "k3"}},
+ ],
+ fallbacks=[{"primary": ["middle" if "middle" in cut_groups else "backup"]}, {"middle": ["backup"]}],
+ enable_mid_stream_fallback_continuation=True,
+ )
+
+ response: Final = await router.acompletion(
+ model="primary",
+ messages=[{"role": "user", "content": "How do transformers work?"}],
+ stream=True,
+ stream_options={"include_usage": True},
+ )
+ chunks: Final = [chunk async for chunk in response]
+ await asyncio.wait_for(recorder.succeeded.wait(), timeout=10)
+ await asyncio.sleep(0.5)
+
+ backup_prompt_cost, backup_completion_cost = litellm.cost_per_token(
+ model="mistral/mistral-large-latest", prompt_tokens=91, completion_tokens=5
+ )
+ success: Final = recorder.successes[-1]
+ assert len(recorder.failures) == len(cut_groups)
+ assert [failure["response_cost"] for failure in recorder.failures] == pytest.approx(
+ [
+ failure["prompt_tokens"] * input_price + failure["completion_tokens"] * output_price
+ for failure, (input_price, output_price) in zip(recorder.failures, ((0.001, 0.002), (0.003, 0.004)))
+ ]
+ )
+ assert success["response_cost"] == pytest.approx(
+ sum(failure["response_cost"] for failure in recorder.failures) + backup_prompt_cost + backup_completion_cost
+ )
+ assert success["prompt_tokens"] == sum(failure["prompt_tokens"] for failure in recorder.failures) + 91
+ assert success["completion_tokens"] == sum(failure["completion_tokens"] for failure in recorder.failures) + 5
+ assert success["cost_breakdown"]["total_cost"] == pytest.approx(success["response_cost"])
+ spend_row: Final = recorder.spend_rows[-1]
+ assert (spend_row["spend"], spend_row["prompt_tokens"], spend_row["completion_tokens"]) == (
+ success["response_cost"],
+ success["prompt_tokens"],
+ success["completion_tokens"],
+ )
+ client_usage: Final = next(chunk.usage for chunk in reversed(chunks) if getattr(chunk, "usage", None))
+ assert (client_usage.prompt_tokens, client_usage.completion_tokens) == (
+ success["prompt_tokens"],
+ success["completion_tokens"],
+ )
+ assert {chunk.id for chunk in chunks} == {spend_row["request_id"]}
+
+
+@pytest.mark.asyncio
+async def test_continued_stream_sharing_the_proxy_logging_object_bills_the_cut_to_its_own_deployment(
+ monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
+) -> None:
+ monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+ recorder: Final = _SlowFailureRecorder()
+ monkeypatch.setattr(litellm, "callbacks", [recorder])
+ respx_mock.post("https://api.openai.com/v1/chat/completions").mock(
+ return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=_CutSSE("Transformers use"))
+ )
+ respx_mock.post(host="api.mistral.ai").mock(
+ return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, text=_BACKUP_SSE)
+ )
+ router: Final = Router(
+ model_list=[
+ {
+ "model_name": "primary",
+ "litellm_params": {
+ "model": "openai/gpt-4o",
+ "api_key": "k1",
+ "input_cost_per_token": 0.001,
+ "output_cost_per_token": 0.002,
+ },
+ "model_info": {"id": "primary-deployment"},
+ },
+ {"model_name": "backup", "litellm_params": {"model": "mistral/mistral-large-latest", "api_key": "k3"}},
+ ],
+ fallbacks=[{"primary": ["backup"]}],
+ enable_mid_stream_fallback_continuation=True,
+ )
+ messages: Final = [{"role": "user", "content": "How do transformers work?"}]
+ shared_logging, _ = litellm.utils.function_setup(
+ original_function="acompletion",
+ rules_obj=litellm.utils.Rules(),
+ start_time=datetime.now(),
+ model="primary",
+ messages=messages,
+ stream=True,
+ litellm_call_id="shared-call",
+ )
+
+ response: Final = await router.acompletion(
+ model="primary",
+ messages=messages,
+ stream=True,
+ stream_options={"include_usage": True},
+ litellm_call_id="shared-call",
+ litellm_logging_obj=shared_logging,
+ )
+ _ = [chunk async for chunk in response]
+ await asyncio.wait_for(recorder.succeeded.wait(), timeout=10)
+ await asyncio.sleep(0.5)
+
+ backup_prompt_cost, backup_completion_cost = litellm.cost_per_token(
+ model="mistral/mistral-large-latest", prompt_tokens=91, completion_tokens=5
+ )
+ (failure,) = recorder.failures
+ assert (failure["model_group"], failure["model_id"]) == ("primary", "primary-deployment")
+ assert recorder.failed_deployment_ids == ["primary-deployment"]
+ assert failure["response_cost"] == pytest.approx(
+ failure["prompt_tokens"] * 0.001 + failure["completion_tokens"] * 0.002
+ )
+ assert recorder.successes[-1]["response_cost"] == pytest.approx(
+ failure["response_cost"] + backup_prompt_cost + backup_completion_cost
+ )
+
+
+def test_carry_partial_spend_to_fallback_adds_the_partial_the_cut_stream_inherited():
+ from litellm.types.utils import Usage
+
+ cut_logging: Final = LiteLLMLogging(
+ model="deepseek/deepseek-chat", messages=[], stream=True, call_type="acompletion",
+ start_time=datetime.now(), litellm_call_id="cut", function_id="cut",
+ )
+ cut_logging.model_call_details["response_cost"] = 0.5
+ cut_logging.update_environment_variables(
+ litellm_params={"custom_llm_provider": "deepseek"}, optional_params={}, model="deepseek/deepseek-chat"
+ )
+ cut_logging.record_continued_partial(
+ usage=Usage(prompt_tokens=40, completion_tokens=7, total_tokens=47),
+ response_cost=0.001,
+ response_id="chatcmpl-primary",
+ )
+ fallback_logging: Final = LiteLLMLogging(
+ model="mistral-large-latest", messages=[], stream=True, call_type="acompletion",
+ start_time=datetime.now(), litellm_call_id="fallback", function_id="fallback",
+ )
+
+ Router._carry_partial_spend_to_fallback(
+ SimpleNamespace(logging_obj=fallback_logging),
+ Router._price_cut_partial(
+ cast(CustomStreamWrapper, SimpleNamespace(logging_obj=cut_logging, model="deepseek/deepseek-chat")),
+ litellm.ModelResponse(
+ id="chatcmpl-relay", usage=Usage(prompt_tokens=60, completion_tokens=3, total_tokens=63)
+ ),
+ ),
+ )
+
+ relay_prompt_cost, relay_completion_cost = litellm.cost_per_token(
+ model="deepseek/deepseek-chat", prompt_tokens=60, completion_tokens=3
+ )
+ carried_usage: Final = fallback_logging.model_call_details["continued_partial_usage"]
+ assert (carried_usage.prompt_tokens, carried_usage.completion_tokens) == (100, 10)
+ assert fallback_logging.model_call_details["continued_partial_cost"] == pytest.approx(
+ relay_prompt_cost + relay_completion_cost + 0.001
+ )
+ assert fallback_logging.model_call_details["continued_response_id"] == "chatcmpl-primary"
+
+
+@pytest.mark.asyncio
+async def test_acompletion_streaming_iterator_plain_fallback_keeps_the_fallback_usage():
+ from litellm.exceptions import MidStreamFallbackError
+
+ error = MidStreamFallbackError(
+ message="Connection lost", model="gpt-4", llm_provider="anthropic", generated_content="",
+ is_pre_first_chunk=True,
+ )
+ role_chunk = litellm.ModelResponseStream(
+ model="gpt-4", choices=[{"index": 0, "delta": {"role": "assistant", "content": ""}}]
+ )
+
+ _, streamed = await _drain_with_fallback(
+ _continuation_router(),
+ error,
+ [role_chunk, _usage_chunk(21, 0)],
+ [{"role": "user", "content": "Hi"}],
+ [_usage_chunk(91, 5)],
+ )
+
+ assert (streamed[-1].usage.prompt_tokens, streamed[-1].usage.completion_tokens) == (91, 5)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "partial,fallback_texts,expected_answer",
+ [
+ ("Transformers use ", (" self", "-attention"), "Transformers use self-attention"),
+ ("def f():\n ", ("\n", " return 1"), "def f():\n return 1"),
+ ("First.\n", ("\n\nSecond",), "First.\n\nSecond"),
+ ],
+ ids=["space", "indentation", "paragraph_break"],
+)
+async def test_acompletion_streaming_iterator_continuation_does_not_double_trailing_whitespace(
+ partial, fallback_texts, expected_answer
+):
+ from litellm.exceptions import MidStreamFallbackError
+
+ error = MidStreamFallbackError(
+ message="Connection lost", model="gpt-4", llm_provider="openai", generated_content=partial,
+ is_pre_first_chunk=False, emitted_disqualifying_content=False,
+ )
+ fallback_items = [
+ litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": text}}]) for text in fallback_texts
+ ]
+
+ passed_kwargs, streamed = await _drain_with_fallback(
+ _continuation_router(), error, [], [{"role": "user", "content": "Hi"}], fallback_items
+ )
+
+ assert passed_kwargs["messages"][-1]["content"] == partial
+ assert partial + "".join(chunk.choices[0].delta.content for chunk in streamed) == expected_answer
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("enabled", [True, "true", " True "])
+async def test_update_settings_turns_on_mid_stream_continuation(monkeypatch, enabled):
+ from unittest.mock import AsyncMock, patch
+
+ from litellm.exceptions import MidStreamFallbackError
+ from litellm.router_utils.pre_call_checks.continuation_prefill_check import (
+ MID_STREAM_CONTINUATION_KWARG,
+ MID_STREAM_CONTINUATION_MARKER,
+ ContinuationPrefillDeploymentCheck,
+ )
+
+ monkeypatch.setattr(litellm, "callbacks", [])
+ router = litellm.Router(
+ model_list=[
+ {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}},
+ {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}},
+ ],
+ fallbacks=[{"gpt-4": ["backup"]}],
+ )
+ router.update_settings(enable_mid_stream_fallback_continuation=enabled)
+
+ error = MidStreamFallbackError(
+ message="Connection lost", model="gpt-4", llm_provider="openai", generated_content="Hello",
+ is_pre_first_chunk=False, emitted_disqualifying_content=False,
+ )
+ fallback_chunk = litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": " world"}}])
+ with patch.object(
+ router,
+ "async_function_with_fallbacks_common_utils",
+ new=AsyncMock(return_value=_FakeFallbackStream(fallback_chunk)),
+ ) as mock_fallback:
+ result = await router._acompletion_streaming_iterator(
+ model_response=_make_midstream_source(error),
+ messages=[{"role": "user", "content": "Hi"}],
+ initial_kwargs={"model": "gpt-4", "stream": True},
+ )
+ collected = [chunk async for chunk in result]
+
+ assert mock_fallback.await_args.kwargs["kwargs"][MID_STREAM_CONTINUATION_KWARG] is MID_STREAM_CONTINUATION_MARKER
+ assert fallback_chunk in collected
+ assert any(isinstance(callback, ContinuationPrefillDeploymentCheck) for callback in litellm.callbacks)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "error_kwargs,request_kwargs",
+ [
+ ({"emitted_disqualifying_content": True}, {}),
+ ({"emitted_disqualifying_content": False}, {"response_format": {"type": "json_object"}}),
+ ({"emitted_disqualifying_content": False}, {"tool_choice": "required"}),
+ ({"emitted_disqualifying_content": False}, {"merge_reasoning_content_in_choices": True}),
+ ({"emitted_disqualifying_content": False}, {"n": 2}),
+ ],
+ ids=["tool_or_thinking_emitted", "json_mode", "forced_tool_choice", "merged_reasoning", "multiple_choices"],
+)
+async def test_acompletion_streaming_iterator_declines_ineligible_after_content(error_kwargs, request_kwargs):
+ """Flag on but the break is not continuation-safe: the router re-raises and
+ never enters the fallback chain, so no duplicated/rejected request is sent."""
+ from unittest.mock import AsyncMock, patch
+
+ from litellm.exceptions import MidStreamFallbackError
+
+ router = litellm.Router(
+ model_list=[
+ {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}},
+ {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}},
+ ],
+ fallbacks=[{"gpt-4": ["backup"]}],
+ enable_mid_stream_fallback_continuation=True,
+ )
+ error = MidStreamFallbackError(
+ message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello",
+ is_pre_first_chunk=False, **error_kwargs,
+ )
+ source = _make_midstream_source(error)
+ initial_kwargs = {"model": "gpt-4", "stream": True, **request_kwargs}
+
+ with patch.object(router, "async_function_with_fallbacks_common_utils", new=AsyncMock()) as mock_fallback:
+ result = await router._acompletion_streaming_iterator(
+ model_response=source, messages=[{"role": "user", "content": "Hi"}], initial_kwargs=initial_kwargs
+ )
+ with pytest.raises(MidStreamFallbackError):
+ async for _ in result:
+ pass
+
+ mock_fallback.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_acompletion_streaming_iterator_raises_content_filter_instead_of_continuing():
+ """A content filter that stops a stream mid-answer reaches the client as ContentPolicyViolationError,
+ never as a continuation that lets a backup finish the filtered answer."""
+ from unittest.mock import AsyncMock, patch
+
+ from litellm.exceptions import MidStreamFallbackError
+
+ router = litellm.Router(
+ model_list=[
+ {"model_name": "gpt-4", "litellm_params": {"model": "azure/gpt-4", "api_key": "k1"}},
+ {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}},
+ ],
+ fallbacks=[{"gpt-4": ["backup"]}],
+ enable_mid_stream_fallback_continuation=True,
+ )
+ content_filter = litellm.ContentPolicyViolationError(message="filtered", model="gpt-4", llm_provider="azure")
+ error = MidStreamFallbackError(
+ message="filtered", model="gpt-4", llm_provider="azure", generated_content="Step one is to obtain the",
+ is_pre_first_chunk=False, emitted_disqualifying_content=False, original_exception=content_filter,
+ )
+
+ with patch.object(router, "async_function_with_fallbacks_common_utils", new=AsyncMock()) as mock_fallback:
+ result = await router._acompletion_streaming_iterator(
+ model_response=_make_midstream_source(error),
+ messages=[{"role": "user", "content": "Hi"}],
+ initial_kwargs={"model": "gpt-4", "stream": True},
+ )
+ with pytest.raises(litellm.ContentPolicyViolationError):
+ async for _ in result:
+ pass
+
+ mock_fallback.assert_not_awaited()
+
+
+async def _continues_after_content(
+ initial_kwargs: dict[str, object],
+ primary_params: Mapping[str, object] = MappingProxyType({}),
+ default_litellm_params: dict[str, object] | None = None,
+ cut_model_id: str = "primary-deployment",
+) -> bool:
+ from unittest.mock import AsyncMock, patch
+
+ from litellm.exceptions import MidStreamFallbackError
+
+ router: Final = litellm.Router(
+ model_list=[
+ {
+ "model_name": "gpt-4",
+ "litellm_params": {"model": "gpt-4", "api_key": "k1", **primary_params},
+ "model_info": {"id": "primary-deployment"},
+ },
+ {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}},
+ {
+ "model_name": "openai-backup",
+ "litellm_params": {"model": "gpt-4o", "api_key": "k3"},
+ "model_info": {"id": "no-prefill-deployment", "supports_assistant_prefill": False},
+ },
+ {
+ "model_name": "claude-backup",
+ "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k4"},
+ "model_info": {"id": "prefill-deployment", "supports_assistant_prefill": True},
+ },
+ {
+ "model_name": "json-claude-backup",
+ "litellm_params": {
+ "model": "anthropic/claude-3-opus-20240229",
+ "api_key": "k5",
+ "response_format": {"type": "json_object"},
+ },
+ "model_info": {"id": "json-prefill-deployment", "supports_assistant_prefill": True},
+ },
+ ],
+ fallbacks=[{"gpt-4": ["backup"]}],
+ default_litellm_params=default_litellm_params,
+ enable_mid_stream_fallback_continuation=True,
+ )
+ error: Final = MidStreamFallbackError(
+ message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello",
+ is_pre_first_chunk=False, emitted_disqualifying_content=False,
+ )
+ fallback_chunk: Final = litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": " world"}}])
+ cut_stream: Final = _make_midstream_source(error)
+ cut_stream._hidden_params = {"model_id": cut_model_id}
+ with patch.object(
+ router,
+ "async_function_with_fallbacks_common_utils",
+ new=AsyncMock(return_value=_FakeFallbackStream(fallback_chunk)),
+ ) as mock_fallback:
+ result: Final = await router._acompletion_streaming_iterator(
+ model_response=cut_stream,
+ messages=[{"role": "user", "content": "Hi"}],
+ initial_kwargs={"model": "gpt-4", "stream": True, **initial_kwargs},
+ )
+ with contextlib.suppress(MidStreamFallbackError):
+ _ = [chunk async for chunk in result]
+ return mock_fallback.await_count == 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "event_hook,continues",
+ [("post_call", False), (None, False), ("pre_call", True)],
+ ids=["post_call_guardrail", "every_hook_guardrail", "input_only_guardrail"],
+)
+async def test_mid_stream_continuation_declines_while_an_output_guardrail_is_registered(
+ monkeypatch: pytest.MonkeyPatch, event_hook: str | None, continues: bool
+) -> None:
+ from litellm.types.guardrails import GuardrailEventHooks
+
+ guardrail: Final = CustomGuardrail(
+ guardrail_name="output-check", event_hook=GuardrailEventHooks(event_hook) if event_hook else None
+ )
+ monkeypatch.setattr(litellm, "callbacks", [guardrail])
+
+ assert await _continues_after_content({}) is continues
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "fallbacks,continues",
+ [
+ ([{"model": "backup", "messages": [{"role": "user", "content": "Say hi"}]}], False),
+ ([{"model": "backup", "stream": False}], False),
+ ([{"model": "backup", "max_tokens": 300}], False),
+ ([{"model": "backup"}], True),
+ ([{"gpt-4": [{"model": "backup", "stream": False}]}], False),
+ ([{"*": [{"model": "backup", "max_tokens": 300}]}], False),
+ ([{"gpt-4": [{"model": "backup"}, "backup"]}], True),
+ ([{"gpt-4": ["backup"]}, {"other-group": [{"model": "backup", "stream": False}]}], True),
+ ([{}], True),
+ ([{"gpt-4": 5}], True),
+ ],
+ ids=[
+ "entry_with_its_own_messages",
+ "entry_turning_off_streaming",
+ "entry_with_its_own_max_tokens",
+ "entry_with_model_only",
+ "group_target_turning_off_streaming",
+ "generic_group_target_with_its_own_max_tokens",
+ "group_targets_with_model_only",
+ "override_only_in_another_groups_chain",
+ "empty_entry",
+ "non_list_chain",
+ ],
+)
+async def test_mid_stream_continuation_declines_a_request_fallback_that_overrides_request_params(
+ fallbacks: list[object], continues: bool
+) -> None:
+ assert await _continues_after_content({"fallbacks": fallbacks}) is continues
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "fallbacks,continues",
+ [
+ ([{"gpt-4": ["no-prefill-deployment"]}], False),
+ ([{"gpt-4": ["backup", "no-prefill-deployment"]}], False),
+ ([{"model": "no-prefill-deployment"}], False),
+ ([{"gpt-4": ["prefill-deployment"]}], True),
+ ([{"gpt-4": ["openai-backup"]}], True),
+ ([{"gpt-4": ["json-prefill-deployment"]}], False),
+ ],
+ ids=[
+ "deployment_id_without_prefill",
+ "deployment_id_without_prefill_later_in_chain",
+ "flat_entry_naming_a_deployment_id_without_prefill",
+ "deployment_id_with_prefill",
+ "group_without_prefill_left_to_the_deployment_filter",
+ "deployment_id_with_its_own_response_format",
+ ],
+)
+async def test_mid_stream_continuation_declines_a_fallback_naming_a_deployment_that_cannot_continue(
+ fallbacks: list[object], continues: bool
+) -> None:
+ assert await _continues_after_content({"fallbacks": fallbacks}) is continues
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "primary_params,default_litellm_params,continues",
+ [
+ ({"response_format": {"type": "json_object"}}, None, False),
+ ({"n": 2}, None, False),
+ ({"tool_choice": "required"}, None, False),
+ ({"merge_reasoning_content_in_choices": True}, None, False),
+ ({}, {"response_format": {"type": "json_object"}}, False),
+ ({"response_format": {"type": "text"}}, None, True),
+ ({"temperature": 0.2}, None, True),
+ ({"tool_choice": "required"}, {"tool_choice": "auto"}, False),
+ ({"response_format": {"type": "json_object"}}, {"response_format": {"type": "text"}}, True),
+ ({"extra_body": {"n": 2}}, None, False),
+ ],
+ ids=[
+ "deployment_response_format",
+ "deployment_n",
+ "deployment_tool_choice",
+ "deployment_merged_reasoning",
+ "router_default_response_format",
+ "deployment_plain_text_format",
+ "deployment_unrelated_param",
+ "deployment_tool_choice_over_router_default",
+ "router_default_format_over_deployment",
+ "deployment_extra_body_n",
+ ],
+)
+async def test_mid_stream_continuation_reads_constraints_configured_on_the_cut_deployment(
+ primary_params: Mapping[str, object], default_litellm_params: dict[str, object] | None, continues: bool
+) -> None:
+ assert await _continues_after_content({}, primary_params, default_litellm_params) is continues
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "extra_body,continues",
+ [({"response_format": {"type": "json_object"}}, False), ({"n": 2}, False), ({"top_k": 5}, True)],
+ ids=["json_format_in_extra_body", "n_in_extra_body", "unrelated_extra_body_key"],
+)
+async def test_mid_stream_continuation_reads_constraints_the_request_sends_in_extra_body(
+ extra_body: dict[str, object], continues: bool
+) -> None:
+ assert await _continues_after_content({"extra_body": extra_body}) is continues
+
+
+@pytest.mark.asyncio
+async def test_mid_stream_continuation_declines_when_the_cut_deployment_is_not_in_the_router() -> None:
+ assert await _continues_after_content({}, cut_model_id="per-request-credential-deployment") is False
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "routing_strategy",
+ [
+ "simple-shuffle",
+ "usage-based-routing",
+ "usage-based-routing-v2",
+ "least-busy",
+ "latency-based-routing",
+ "cost-based-routing",
+ ],
+)
+async def test_continuation_routes_only_to_deployments_that_can_continue_under_every_routing_strategy(
+ routing_strategy: str,
+) -> None:
+ from litellm.router_utils.pre_call_checks.continuation_prefill_check import (
+ MID_STREAM_CONTINUATION_KWARG,
+ MID_STREAM_CONTINUATION_MARKER,
+ )
+ from litellm.types.router import RouterRateLimitError
+
+ router: Final = litellm.Router(
+ model_list=[
+ {
+ "model_name": "backup",
+ "litellm_params": {"model": "gpt-4o", "api_key": "k1"},
+ "model_info": {"id": "no-prefill-deployment", "supports_assistant_prefill": False},
+ },
+ ],
+ routing_strategy=routing_strategy,
+ enable_mid_stream_fallback_continuation=True,
+ )
+ messages: Final = [{"role": "user", "content": "Hi"}]
+ plain_pick: Final = await router.async_get_available_deployment(
+ model="backup", messages=messages, request_kwargs={}
+ )
+ assert plain_pick["model_info"]["id"] == "no-prefill-deployment"
+ with pytest.raises(RouterRateLimitError):
+ await router.async_get_available_deployment(
+ model="backup",
+ messages=messages,
+ request_kwargs={MID_STREAM_CONTINUATION_KWARG: MID_STREAM_CONTINUATION_MARKER},
+ )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+ "flag_kwargs", [{}, {"enable_mid_stream_fallback_continuation": "false"}], ids=["default", "string_false"]
+)
+async def test_acompletion_streaming_iterator_flag_off_declines_after_content(flag_kwargs):
+ """Regression guard for the opt-in default: with the flag off, a post-content
+ break re-raises and never falls back, exactly as before this feature."""
+ from unittest.mock import AsyncMock, patch
+
+ from litellm.exceptions import MidStreamFallbackError
+
+ router = litellm.Router(
+ model_list=[
+ {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}},
+ {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}},
+ ],
+ fallbacks=[{"gpt-4": ["backup"]}],
+ **flag_kwargs,
+ )
+ error = MidStreamFallbackError(
+ message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello",
+ is_pre_first_chunk=False, emitted_disqualifying_content=False,
+ )
+ source = _make_midstream_source(error)
+
+ with patch.object(router, "async_function_with_fallbacks_common_utils", new=AsyncMock()) as mock_fallback:
+ result = await router._acompletion_streaming_iterator(
+ model_response=source, messages=[{"role": "user", "content": "Hi"}],
+ initial_kwargs={"model": "gpt-4", "stream": True},
+ )
+ with pytest.raises(MidStreamFallbackError):
+ async for _ in result:
+ pass
+
+ mock_fallback.assert_not_awaited()
+
+
+def test_build_completion_continuation_input_appends_assistant_prefill():
+ messages = [{"role": "user", "content": "hi"}]
+ built = litellm.Router._build_completion_continuation_input(messages, "partial answer")
+ assert built[:-1] == messages
+ assert built[-1] == {"role": "assistant", "content": "partial answer", "prefix": True}
+
+
+@pytest.mark.parametrize(
+ "prefill_content,second_partial,expected_content",
+ [
+ ("part one ", " part two ", "part one part two "),
+ ("part one ", "part two ", "part one part two "),
+ ("First.\n\n", "Second", "First.\n\nSecond"),
+ ([{"type": "text", "text": "part one "}], " part two ", [{"type": "text", "text": "part one part two "}]),
+ ],
+ ids=["separator_sent_again", "separator_not_sent_again", "paragraph_break_not_sent_again", "content_blocks"],
+)
+def test_build_completion_continuation_input_folds_into_existing_prefill(
+ prefill_content, second_partial, expected_content
+):
+ """A nested break must not leave two trailing assistant turns: the new partial folds into the prior prefill,
+ which keeps reading as the text the caller already received."""
+ messages = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": prefill_content, "prefix": True}]
+ built = litellm.Router._build_completion_continuation_input(messages, second_partial)
+ assert built == [messages[0], {"role": "assistant", "content": expected_content, "prefix": True}]
+
+
+@pytest.mark.parametrize(
+ "prefill_content,expected_content",
+ [
+ ("Hello", "Hello world"),
+ ([{"type": "text", "text": "Hello"}], [{"type": "text", "text": "Hello world"}]),
+ ],
+ ids=["string_prefill", "content_block_prefill"],
+)
+def test_build_completion_continuation_input_does_not_double_an_echoed_prefill(prefill_content, expected_content):
+ messages = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": prefill_content, "prefix": True}]
+ built = litellm.Router._build_completion_continuation_input(messages, "Hello world")
+ assert built is not None
+ assert built[-1] == {"role": "assistant", "content": expected_content, "prefix": True}
+
+
+@pytest.mark.parametrize(
+ "cut_partial,expected_content",
+ [("Hel", None), ("Hello", "Hello"), ("Hello wor", "Hello wor")],
+ ids=["cut_mid_echo", "cut_after_a_full_echo", "cut_after_new_text"],
+)
+def test_build_completion_continuation_input_declines_a_cut_in_the_middle_of_a_prefill_echo(
+ cut_partial, expected_content
+):
+ messages = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "Hello", "prefix": True}]
+ built = litellm.Router._build_completion_continuation_input(messages, cut_partial)
+ expected = None if expected_content is None else [messages[0], {**messages[1], "content": expected_content}]
+ assert built == expected
+
+
+@pytest.mark.parametrize(
+ "assistant_content", ["Earlier answer.", [{"type": "text", "text": "Earlier answer."}]], ids=["string", "blocks"]
+)
+def test_build_completion_continuation_input_declines_a_trailing_plain_assistant_turn(assistant_content):
+ """The caller's own assistant turn is history, not a prefill, so the partial is never fused into it."""
+ messages = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": assistant_content}]
+ assert litellm.Router._build_completion_continuation_input(messages, "The capital of France is") is None
+
+
+@pytest.mark.asyncio
+async def test_acompletion_streaming_iterator_declines_when_request_ends_in_plain_assistant_turn():
+ from unittest.mock import AsyncMock, patch
+
+ from litellm.exceptions import MidStreamFallbackError
+
+ router = litellm.Router(
+ model_list=[
+ {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}},
+ {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}},
+ ],
+ fallbacks=[{"gpt-4": ["backup"]}],
+ enable_mid_stream_fallback_continuation=True,
+ )
+ error = MidStreamFallbackError(
+ message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello",
+ is_pre_first_chunk=False, emitted_disqualifying_content=False,
+ )
+
+ with patch.object(router, "async_function_with_fallbacks_common_utils", new=AsyncMock()) as mock_fallback:
+ result = await router._acompletion_streaming_iterator(
+ model_response=_make_midstream_source(error),
+ messages=[{"role": "user", "content": "Hi"}, {"role": "assistant", "content": "Earlier answer."}],
+ initial_kwargs={"model": "gpt-4", "stream": True},
+ )
+ with pytest.raises(MidStreamFallbackError):
+ async for _ in result:
+ pass
+
+ mock_fallback.assert_not_awaited()
+
+
+def test_build_completion_continuation_input_keeps_structured_assistant_content():
+ """Content blocks on a trailing assistant turn stay blocks: the partial continues the last text
+ block in place (a block boundary would let a provider drop the space between "Sure, " and "here"),
+ lands as a new text block only after a non-text block, and a content shape that cannot be extended declines."""
+ from litellm.router import _extend_trailing_text_block
+
+ messages = [
+ {"role": "user", "content": "hi"},
+ {"role": "assistant", "content": [{"type": "text", "text": "Sure, "}], "prefix": True},
+ ]
+ built = litellm.Router._build_completion_continuation_input(messages, "here it is")
+ assert built is not None
+ assert [m["role"] for m in built] == ["user", "assistant"]
+ assert built[-1] == {"role": "assistant", "content": [{"type": "text", "text": "Sure, here it is"}], "prefix": True}
+ image_block = {"type": "image_url", "image_url": {"url": "https://example.test/a.png"}}
+ assert _extend_trailing_text_block([image_block], "here it is") == [image_block, {"type": "text", "text": "here it is"}]
+ assert _extend_trailing_text_block([], "here it is") == [{"type": "text", "text": "here it is"}]
+ assert (
+ litellm.Router._build_completion_continuation_input([{"role": "assistant", "content": 42, "prefix": True}], "x")
+ is None
+ )
+
+
+def test_mid_stream_continuation_eligible_allows_text_response_format():
+ """response_format={"type": "text"} is the unconstrained default and must stay
+ eligible, unlike json_object / json_schema."""
+ from litellm.exceptions import MidStreamFallbackError
+
+ router = litellm.Router(
+ model_list=[{"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k"}}],
+ enable_mid_stream_fallback_continuation=True,
+ )
+ e = MidStreamFallbackError(
+ message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello",
+ is_pre_first_chunk=False, emitted_disqualifying_content=False,
+ )
+ assert router._mid_stream_continuation_eligible(e=e, request_kwargs={"response_format": {"type": "text"}}) is True
+ assert router._mid_stream_continuation_eligible(e=e, request_kwargs={"response_format": {"type": "json_object"}}) is False
+
+
+def test_raise_original_mid_stream_error_surfaces_the_provider_exception():
+ from litellm.exceptions import MidStreamFallbackError, RateLimitError
+
+ provider_error = RateLimitError(message="rate limited", llm_provider="openai", model="gpt-4")
+ wrapped = MidStreamFallbackError(
+ message="rate limited", model="gpt-4", llm_provider="openai",
+ original_exception=provider_error, generated_content="Hello",
+ )
+ with pytest.raises(RateLimitError) as raised:
+ litellm.Router._raise_original_mid_stream_error(wrapped)
+ assert raised.value is provider_error
+ assert raised.value.__cause__ is wrapped
+
+ bare = MidStreamFallbackError(message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello")
+ with pytest.raises(MidStreamFallbackError) as bare_raised:
+ litellm.Router._raise_original_mid_stream_error(bare)
+ assert bare_raised.value is bare
+
+
@pytest.mark.asyncio
async def test_acompletion_streaming_iterator_reraises_original_exception_when_available():
"""Async: when the mid-stream MidStreamFallbackError wraps a real provider
@@ -8692,6 +9905,11 @@ def test_combine_fallback_usage():
assert chunk.usage.completion_tokens == 5
assert chunk.usage.total_tokens == 15
+ Router._combine_fallback_usage(chunk, Usage(prompt_tokens=0, completion_tokens=4, total_tokens=4))
+ assert chunk.usage.prompt_tokens == 10
+ assert chunk.usage.completion_tokens == 9
+ assert chunk.usage.total_tokens == 19
+
@pytest.mark.asyncio
async def test_acompletion_streaming_iterator_does_not_log_success_on_terminal_failure():
@@ -12190,7 +13408,7 @@ def test_pre_call_checks_uses_deployment_model_when_model_info_lookup_raises(mon
model="custom-alias",
healthy_deployments=deployments,
messages=[{"role": "user", "content": "hi"}],
- request_kwargs={},
+ request_kwargs={"response_format": {"type": "json_object"}},
)
assert len(result) == 1
@@ -26764,3 +27982,471 @@ def test_completion_streams_a_wrapper_that_hands_iteration_to_another_iterator()
streamed = list(response)
assert streamed == chunks
+
+
+_PRE_CALL_CHECK_UNPARSEABLE_MESSAGES: Final = [
+ {
+ "role": "user",
+ "content": [
+ {"type": "image_url", "image_url": {"detail": "not-a-real-detail-value"}},
+ ],
+ }
+]
+
+
+def test_pre_call_checks_filters_deployment_when_token_count_exceeds_limit():
+ model_list: Final = [
+ {
+ "model_name": "fake-openai-endpoint",
+ "litellm_params": {"model": "openai/my-fake-model", "api_key": "my-fake-key"},
+ "model_info": {"id": "small", "max_input_tokens": 5},
+ },
+ {
+ "model_name": "fake-openai-endpoint",
+ "litellm_params": {"model": "openai/my-other-fake-model", "api_key": "my-fake-key"},
+ "model_info": {"id": "large", "max_input_tokens": 10000},
+ },
+ ]
+ router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0)
+
+ filtered: Final = router._pre_call_checks(
+ model="fake-openai-endpoint",
+ healthy_deployments=model_list,
+ messages=[{"role": "user", "content": " ".join(["word"] * 100)}],
+ )
+
+ assert [d["model_info"]["id"] for d in filtered] == ["large"]
+
+
+def test_pre_call_checks_rpm_check_still_enforced_when_token_counting_fails():
+ from litellm.types.router import RouterRateLimitErrorBasic
+
+ model_list: Final = [
+ {
+ "model_name": "fake-openai-endpoint",
+ "litellm_params": {"model": "openai/my-fake-model", "api_key": "my-fake-key", "rpm": 0},
+ "model_info": {"max_input_tokens": 100},
+ },
+ ]
+ router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0)
+
+ with pytest.raises(RouterRateLimitErrorBasic):
+ router._pre_call_checks(
+ model="fake-openai-endpoint",
+ healthy_deployments=model_list,
+ messages=_PRE_CALL_CHECK_UNPARSEABLE_MESSAGES,
+ )
+
+
+def test_pre_call_checks_context_window_check_is_skipped_but_not_fatal_when_token_counting_fails():
+ model_list: Final = [
+ {
+ "model_name": "fake-openai-endpoint",
+ "litellm_params": {"model": "openai/my-fake-model", "api_key": "my-fake-key"},
+ "model_info": {"max_input_tokens": 100},
+ },
+ ]
+ router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0)
+
+ filtered: Final = router._pre_call_checks(
+ model="fake-openai-endpoint",
+ healthy_deployments=model_list,
+ messages=_PRE_CALL_CHECK_UNPARSEABLE_MESSAGES,
+ )
+
+ assert len(filtered) == 1
+
+
+def test_pre_call_checks_keeps_unsupported_param_filter_when_token_counting_fails(monkeypatch: pytest.MonkeyPatch):
+ monkeypatch.setattr(litellm, "drop_params", False)
+ no_json_mode_model: Final = "gpt-3.5-turbo-16k"
+ assert "response_format" not in (
+ litellm.get_supported_openai_params(model=no_json_mode_model, custom_llm_provider="openai") or []
+ )
+ model_list: Final = [
+ {
+ "model_name": "fake-openai-endpoint",
+ "litellm_params": {"model": f"openai/{deployment_model}", "api_key": "my-fake-key"},
+ "model_info": {"id": deployment_model, "max_input_tokens": 100},
+ }
+ for deployment_model in (no_json_mode_model, "gpt-4o")
+ ]
+ router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0)
+
+ filtered: Final = router._pre_call_checks(
+ model="fake-openai-endpoint",
+ healthy_deployments=model_list,
+ messages=_PRE_CALL_CHECK_UNPARSEABLE_MESSAGES,
+ request_kwargs={"response_format": {"type": "json_object"}},
+ )
+
+ assert [d["model_info"]["id"] for d in filtered] == ["gpt-4o"]
+
+
+async def test_async_get_healthy_deployments_keeps_region_check_when_token_counting_fails():
+ model_list: Final = [
+ {
+ "model_name": "fake-openai-endpoint",
+ "litellm_params": {"model": "openai/my-fake-model", "api_key": "my-fake-key", "region_name": region},
+ "model_info": {"id": f"{region}-deployment", "max_input_tokens": 100},
+ }
+ for region in ("eu", "us")
+ ]
+ router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0)
+
+ healthy: Final = await router.async_get_healthy_deployments(
+ model="fake-openai-endpoint",
+ request_kwargs={"allowed_model_region": "eu"},
+ messages=_PRE_CALL_CHECK_UNPARSEABLE_MESSAGES,
+ input=None,
+ specific_deployment=False,
+ parent_otel_span=None,
+ )
+
+ assert [deployment["model_info"]["id"] for deployment in healthy] == ["eu-deployment"]
+
+
+@pytest.mark.parametrize(
+ "messages",
+ [_PRE_CALL_CHECK_UNPARSEABLE_MESSAGES, [{"role": "user", "content": "hi"}]],
+ ids=["uncountable", "countable"],
+)
+def test_completion_answers_400_when_no_deployment_supports_the_requested_param(
+ monkeypatch: pytest.MonkeyPatch, messages: list[dict[str, object]]
+):
+ monkeypatch.setattr(litellm, "drop_params", False)
+ router: Final = Router(
+ model_list=[
+ {
+ "model_name": "no-json-mode",
+ "litellm_params": {"model": "openai/gpt-3.5-turbo-16k", "api_key": "my-fake-key"},
+ "model_info": {"max_input_tokens": 100},
+ }
+ ],
+ enable_pre_call_checks=True,
+ num_retries=2,
+ )
+
+ with pytest.raises(litellm.UnsupportedParamsError) as error:
+ router.completion(
+ model="no-json-mode",
+ messages=messages,
+ response_format={"type": "json_object"},
+ mock_response="should not be served",
+ )
+
+ assert error.value.status_code == 400
+ assert "allowed_openai_params" in str(error.value)
+
+
+@pytest.mark.parametrize(
+ ("deployment_params", "request_kwargs"),
+ [
+ ({}, {"allowed_openai_params": ["response_format"]}),
+ ({}, {"drop_params": True}),
+ ({}, {"additional_drop_params": ["response_format"]}),
+ ({"allowed_openai_params": ["response_format"]}, {}),
+ ({"drop_params": True}, {}),
+ ],
+ ids=["request-allowlist", "request-drop", "request-additional-drop", "deployment-allowlist", "deployment-drop"],
+)
+def test_pre_call_checks_keeps_deployment_whose_unsupported_param_is_allowed_or_dropped(
+ monkeypatch: pytest.MonkeyPatch, deployment_params: Mapping[str, object], request_kwargs: Mapping[str, object]
+):
+ monkeypatch.setattr(litellm, "drop_params", False)
+ model_list: Final = [
+ {
+ "model_name": "no-json-mode",
+ "litellm_params": {"model": "openai/gpt-3.5-turbo-16k", "api_key": "my-fake-key", **deployment_params},
+ "model_info": {"id": "only", "max_input_tokens": 100},
+ }
+ ]
+ router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0)
+
+ filtered: Final = router._pre_call_checks(
+ model="no-json-mode",
+ healthy_deployments=model_list,
+ messages=[{"role": "user", "content": "hi"}],
+ request_kwargs={"response_format": {"type": "json_object"}, **request_kwargs},
+ )
+
+ assert [d["model_info"]["id"] for d in filtered] == ["only"]
+
+
+def test_pre_call_checks_keeps_retryable_429_when_a_cooling_down_deployment_supports_the_param(monkeypatch: pytest.MonkeyPatch):
+ monkeypatch.setattr(litellm, "drop_params", False)
+ model_list: Final = [
+ {
+ "model_name": "mixed-json-mode",
+ "litellm_params": {"model": f"openai/{deployment_model}", "api_key": "my-fake-key"},
+ "model_info": {"id": deployment_model, "max_input_tokens": 100},
+ }
+ for deployment_model in ("gpt-3.5-turbo-16k", "gpt-4o")
+ ]
+ router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0)
+ without_json_mode: Final = [d for d in router.model_list if d["model_info"]["id"] == "gpt-3.5-turbo-16k"]
+
+ filtered: Final = router._pre_call_checks(
+ model="mixed-json-mode",
+ healthy_deployments=without_json_mode,
+ messages=[{"role": "user", "content": "hi"}],
+ request_kwargs={"response_format": {"type": "json_object"}},
+ )
+
+ assert filtered == []
+
+
+def test_completion_uses_context_window_fallback_when_the_small_deployment_also_lacks_the_param(monkeypatch: pytest.MonkeyPatch):
+ monkeypatch.setattr(litellm, "drop_params", False)
+ router: Final = Router(
+ model_list=[
+ {
+ "model_name": "small",
+ "litellm_params": {"model": "openai/gpt-3.5-turbo-16k", "api_key": "my-fake-key"},
+ "model_info": {"id": "small-id", "max_input_tokens": 5},
+ },
+ {
+ "model_name": "big",
+ "litellm_params": {"model": "openai/gpt-4o", "api_key": "my-fake-key"},
+ "model_info": {"id": "big-id"},
+ },
+ ],
+ enable_pre_call_checks=True,
+ context_window_fallbacks=[{"small": ["big"]}],
+ num_retries=0,
+ )
+
+ response: Final = router.completion(
+ model="small",
+ messages=[{"role": "user", "content": "hi " * 50}],
+ response_format={"type": "json_object"},
+ mock_response="served by the fallback",
+ )
+
+ assert response._hidden_params["model_id"] == "big-id"
+
+
+@pytest.mark.parametrize(
+ ("deployment_params", "request_kwargs"),
+ [
+ ({}, {"allowed_openai_params": ["response_format"]}),
+ ({}, {"drop_params": True}),
+ ({}, {"additional_drop_params": ["response_format"]}),
+ ({"drop_params": True}, {}),
+ ({"allowed_openai_params": ["response_format"]}, {}),
+ ],
+ ids=["request-allowlist", "request-drop", "request-additional-drop", "deployment-drop", "deployment-allowlist"],
+)
+def test_pre_call_checks_prefers_a_deployment_that_takes_the_param_over_dropping_it(
+ monkeypatch: pytest.MonkeyPatch, deployment_params: Mapping[str, object], request_kwargs: Mapping[str, object]
+):
+ monkeypatch.setattr(litellm, "drop_params", False)
+ model_list: Final = [
+ {
+ "model_name": "mixed-json-mode",
+ "litellm_params": {"model": "openai/gpt-3.5-turbo-16k", "api_key": "my-fake-key", **deployment_params},
+ "model_info": {"id": "gpt-3.5-turbo-16k", "max_input_tokens": 100},
+ },
+ {
+ "model_name": "mixed-json-mode",
+ "litellm_params": {"model": "openai/gpt-4o", "api_key": "my-fake-key"},
+ "model_info": {"id": "gpt-4o", "max_input_tokens": 100},
+ },
+ ]
+ router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0)
+
+ filtered: Final = router._pre_call_checks(
+ model="mixed-json-mode",
+ healthy_deployments=model_list,
+ messages=[{"role": "user", "content": "hi"}],
+ request_kwargs={"response_format": {"type": "json_object"}, **request_kwargs},
+ )
+
+ assert [d["model_info"]["id"] for d in filtered] == ["gpt-4o"]
+
+
+def test_pre_call_checks_drops_the_param_when_the_capable_deployment_is_outside_the_region(monkeypatch: pytest.MonkeyPatch):
+ monkeypatch.setattr(litellm, "drop_params", False)
+ model_list: Final = [
+ {
+ "model_name": "regional-json-mode",
+ "litellm_params": {"model": "openai/gpt-3.5-turbo-16k", "api_key": "my-fake-key", "region_name": "eu"},
+ "model_info": {"id": "eu-without-json-mode", "max_input_tokens": 100},
+ },
+ {
+ "model_name": "regional-json-mode",
+ "litellm_params": {"model": "openai/gpt-4o", "api_key": "my-fake-key", "region_name": "us"},
+ "model_info": {"id": "us-with-json-mode", "max_input_tokens": 100},
+ },
+ ]
+ router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0)
+
+ filtered: Final = router._pre_call_checks(
+ model="regional-json-mode",
+ healthy_deployments=model_list,
+ messages=[{"role": "user", "content": "hi"}],
+ request_kwargs={"response_format": {"type": "json_object"}, "drop_params": True, "allowed_model_region": "eu"},
+ )
+
+ assert [d["model_info"]["id"] for d in filtered] == ["eu-without-json-mode"]
+
+
+def test_pre_call_checks_filters_an_authenticating_provider_deployment_that_lacks_the_param(
+ monkeypatch: pytest.MonkeyPatch, tmp_path: Path
+):
+ monkeypatch.setattr(litellm, "drop_params", False)
+ monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path))
+ (tmp_path / "api-key.json").write_text(json.dumps({"token": "tid=test", "expires_at": 4102444800}))
+ model_list: Final = [
+ {
+ "model_name": "mixed-json-mode",
+ "litellm_params": {"model": "github_copilot/gpt-4"},
+ "model_info": {"id": "copilot-gpt-4", "max_input_tokens": 100},
+ },
+ {
+ "model_name": "mixed-json-mode",
+ "litellm_params": {"model": "openai/gpt-4o", "api_key": "my-fake-key"},
+ "model_info": {"id": "gpt-4o", "max_input_tokens": 100},
+ },
+ ]
+ router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0)
+
+ filtered: Final = router._pre_call_checks(
+ model="mixed-json-mode",
+ healthy_deployments=model_list,
+ messages=[{"role": "user", "content": "hi"}],
+ request_kwargs={"response_format": {"type": "json_object"}},
+ )
+
+ assert [d["model_info"]["id"] for d in filtered] == ["gpt-4o"]
+
+
+def test_pre_call_checks_answers_400_when_the_group_name_lists_no_deployments(monkeypatch: pytest.MonkeyPatch):
+ monkeypatch.setattr(litellm, "drop_params", False)
+ model_list: Final = [
+ {
+ "model_name": "no-json-mode",
+ "litellm_params": {"model": "openai/gpt-3.5-turbo-16k", "api_key": "my-fake-key"},
+ "model_info": {"id": "only", "max_input_tokens": 100},
+ }
+ ]
+ router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0)
+
+ with pytest.raises(litellm.UnsupportedParamsError):
+ router._pre_call_checks(
+ model="team-public-name",
+ healthy_deployments=model_list,
+ messages=[{"role": "user", "content": "hi"}],
+ request_kwargs={"response_format": {"type": "json_object"}},
+ )
+
+
+def test_pre_call_checks_lets_a_request_setting_replace_the_deployment_setting(monkeypatch: pytest.MonkeyPatch):
+ monkeypatch.setattr(litellm, "drop_params", False)
+ model_list: Final = [
+ {
+ "model_name": "no-json-mode",
+ "litellm_params": {"model": "openai/gpt-3.5-turbo-16k", "api_key": "my-fake-key", "drop_params": True},
+ "model_info": {"id": "drops-it", "max_input_tokens": 100},
+ },
+ {
+ "model_name": "no-json-mode",
+ "litellm_params": {
+ "model": "openai/gpt-4",
+ "api_key": "my-fake-key",
+ "allowed_openai_params": ["response_format"],
+ },
+ "model_info": {"id": "allows-it", "max_input_tokens": 100},
+ },
+ ]
+ router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0)
+
+ filtered: Final = router._pre_call_checks(
+ model="no-json-mode",
+ healthy_deployments=model_list,
+ messages=[{"role": "user", "content": "hi"}],
+ request_kwargs={"response_format": {"type": "json_object"}, "drop_params": False},
+ )
+
+ assert [d["model_info"]["id"] for d in filtered] == ["allows-it"]
+
+
+async def _select_team_json_deployment(router: Router, entrypoint: str, selection: Mapping[str, object]) -> object:
+ if entrypoint == "async":
+ return await router.async_get_available_deployment(**selection)
+ if entrypoint == "sync":
+ return router.get_available_deployment(**selection)
+ return router.get_available_deployment_for_pass_through(**selection)
+
+
+@pytest.mark.parametrize("entrypoint", ["sync", "async", "pass-through"])
+@pytest.mark.parametrize(
+ ("deployment_models", "cooled_down", "expected_error"),
+ [
+ pytest.param(("gpt-3.5-turbo-16k", "gpt-4"), (), litellm.UnsupportedParamsError, id="none-takes-it"),
+ pytest.param(("gpt-3.5-turbo-16k", "gpt-4o"), ("gpt-4o",), RouterRateLimitError, id="capable-cooling-down"),
+ ],
+)
+@pytest.mark.asyncio
+async def test_team_public_model_name_answers_400_only_when_no_team_deployment_takes_the_param(
+ monkeypatch: pytest.MonkeyPatch,
+ entrypoint: str,
+ deployment_models: tuple[str, str],
+ cooled_down: tuple[str, ...],
+ expected_error: type[Exception],
+):
+ monkeypatch.setattr(litellm, "drop_params", False)
+ router: Final = Router(
+ model_list=[
+ {
+ "model_name": f"team-json-{deployment_model}",
+ "litellm_params": {
+ "model": f"openai/{deployment_model}",
+ "api_key": "my-fake-key",
+ "use_in_pass_through": True,
+ },
+ "model_info": {
+ "id": deployment_model,
+ "max_input_tokens": 100,
+ "team_id": "json-team",
+ "team_public_model_name": "team-json",
+ },
+ }
+ for deployment_model in deployment_models
+ ],
+ enable_pre_call_checks=True,
+ num_retries=0,
+ )
+ _cool_down(router, *cooled_down)
+ selection: Final = {
+ "model": "team-json",
+ "messages": [{"role": "user", "content": "hi"}],
+ "request_kwargs": {
+ "metadata": {"user_api_key_team_id": "json-team"},
+ "response_format": {"type": "json_object"},
+ },
+ }
+
+ with pytest.raises(expected_error):
+ await _select_team_json_deployment(router, entrypoint, selection)
+
+
+def test_pre_call_checks_ignores_an_image_response_format_string(monkeypatch: pytest.MonkeyPatch):
+ monkeypatch.setattr(litellm, "drop_params", False)
+ model_list: Final = [
+ {
+ "model_name": "image-model",
+ "litellm_params": {"model": "openai/gpt-3.5-turbo-16k", "api_key": "my-fake-key"},
+ "model_info": {"id": "only", "max_input_tokens": 100},
+ }
+ ]
+ router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0)
+
+ filtered: Final = router._pre_call_checks(
+ model="image-model",
+ healthy_deployments=model_list,
+ messages=[{"role": "user", "content": "prompt"}],
+ request_kwargs={"response_format": "b64_json"},
+ )
+
+ assert [d["model_info"]["id"] for d in filtered] == ["only"]
diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py
index fc1799861d1..d8625ff9da3 100644
--- a/tests/unit/test_router_model_cost_isolation.py
+++ b/tests/unit/test_router_model_cost_isolation.py
@@ -119,27 +119,33 @@ async def test_discovered_limits_survive_deployment_growth_and_removal(
async def test_discovery_discards_metadata_for_a_replaced_deployment(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost))
- router: Final = Router(model_list=[{
- "model_name": "local",
- "litellm_params": {
- "model": "hosted_vllm/local-model",
- "api_base": "https://original.test/v1",
- "api_key": "local-key",
- },
- "model_info": {"id": "replaced-deployment"},
- }])
+ router: Final = Router(
+ model_list=[
+ {
+ "model_name": "local",
+ "litellm_params": {
+ "model": "hosted_vllm/local-model",
+ "api_base": "https://original.test/v1",
+ "api_key": "local-key",
+ },
+ "model_info": {"id": "replaced-deployment"},
+ }
+ ]
+ )
def respond(request: httpx.Request) -> httpx.Response:
if request.url.host == "original.test":
- router.upsert_deployment(Deployment(
- model_name="local",
- litellm_params=LiteLLM_Params(
- model="hosted_vllm/local-model",
- api_base="https://replacement.test/v1",
- api_key="local-key",
- ),
- model_info=ModelInfo(id="replaced-deployment"),
- ))
+ router.upsert_deployment(
+ Deployment(
+ model_name="local",
+ litellm_params=LiteLLM_Params(
+ model="hosted_vllm/local-model",
+ api_base="https://replacement.test/v1",
+ api_key="local-key",
+ ),
+ model_info=ModelInfo(id="replaced-deployment"),
+ )
+ )
return httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 8192}]})
assert request.url.host == "replacement.test"
return httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 2048}]})
@@ -158,15 +164,19 @@ async def test_discovery_discards_metadata_for_a_replaced_deployment(monkeypatch
async def test_discovery_is_isolated_across_routers_and_reused_ids(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost))
first, second = tuple(
- Router(model_list=[{
- "model_name": "local",
- "litellm_params": {
- "model": "hosted_vllm/local-model",
- "api_base": f"https://{host}.test/v1",
- "api_key": "local-key",
- },
- "model_info": {"id": "shared-discovery-id"},
- }])
+ Router(
+ model_list=[
+ {
+ "model_name": "local",
+ "litellm_params": {
+ "model": "hosted_vllm/local-model",
+ "api_base": f"https://{host}.test/v1",
+ "api_key": "local-key",
+ },
+ "model_info": {"id": "shared-discovery-id"},
+ }
+ ]
+ )
for host in ("first", "second")
)
@@ -187,15 +197,17 @@ async def test_discovery_is_isolated_across_routers_and_reused_ids(monkeypatch:
assert first.get_configured_token_limits("local") == (8192, 8192)
assert second.get_configured_token_limits("local") == (2048, 2048)
assert litellm.model_cost["shared-discovery-id"].get("max_input_tokens") is None
- first.upsert_deployment(Deployment(
- model_name="local",
- litellm_params=LiteLLM_Params(
- model="hosted_vllm/local-model",
- api_base="https://unavailable.test/v1",
- api_key="local-key",
- ),
- model_info=ModelInfo(id="shared-discovery-id"),
- ))
+ first.upsert_deployment(
+ Deployment(
+ model_name="local",
+ litellm_params=LiteLLM_Params(
+ model="hosted_vllm/local-model",
+ api_base="https://unavailable.test/v1",
+ api_key="local-key",
+ ),
+ model_info=ModelInfo(id="shared-discovery-id"),
+ )
+ )
assert first.get_configured_token_limits("local") == (None, None)
await first.arefresh_model_info(client=handler)
assert first.get_configured_token_limits("local") == (None, None)
@@ -206,17 +218,19 @@ async def test_discovery_is_isolated_across_routers_and_reused_ids(monkeypatch:
async def test_discovery_refreshes_other_endpoints_while_one_is_pending(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost))
second_started: Final = asyncio.Event()
- router: Final = Router(model_list=[
- {
- "model_name": host,
- "litellm_params": {
- "model": "hosted_vllm/local-model",
- "api_base": f"https://{host}.test/v1",
- "api_key": "local-key",
- },
- }
- for host in ("first", "second", "third")
- ])
+ router: Final = Router(
+ model_list=[
+ {
+ "model_name": host,
+ "litellm_params": {
+ "model": "hosted_vllm/local-model",
+ "api_base": f"https://{host}.test/v1",
+ "api_key": "local-key",
+ },
+ }
+ for host in ("first", "second", "third")
+ ]
+ )
async def respond(request: httpx.Request) -> httpx.Response:
if request.url.host == "first.test":
@@ -240,21 +254,27 @@ async def test_discovery_refreshes_other_endpoints_while_one_is_pending(monkeypa
async def test_discovered_limits_expire_after_the_last_successful_refresh(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost))
clock: Final = Mock(return_value=0.0)
- router: Final = Router(model_list=[{
- "model_name": "local",
- "litellm_params": {
- "model": "hosted_vllm/local-model",
- "api_base": "https://expiry.test/v1",
- "api_key": "local-key",
- },
- "model_info": {"id": "expiring-discovery"},
- }])
+ router: Final = Router(
+ model_list=[
+ {
+ "model_name": "local",
+ "litellm_params": {
+ "model": "hosted_vllm/local-model",
+ "api_base": "https://expiry.test/v1",
+ "api_key": "local-key",
+ },
+ "model_info": {"id": "expiring-discovery"},
+ }
+ ]
+ )
router._discovered_model_info_cache = InMemoryCache(clock=clock, default_ttl=2 * MODEL_INFO_REFRESH_SECONDS)
- responses: Final = iter((
- httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 4096}]}),
- httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 8192}]}),
- httpx.Response(503),
- ))
+ responses: Final = iter(
+ (
+ httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 4096}]}),
+ httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 8192}]}),
+ httpx.Response(503),
+ )
+ )
handler: Final = AsyncHTTPHandler()
await handler.client.aclose()
async with httpx.AsyncClient(transport=httpx.MockTransport(lambda request: next(responses))) as client:
@@ -343,10 +363,12 @@ async def test_discovered_limits_are_isolated_overridable_and_refreshable(
async def test_discovery_preserves_input_overrides_and_survives_outages(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost))
- responses: Final = iter((
- httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 4096}]}),
- httpx.Response(503),
- ))
+ responses: Final = iter(
+ (
+ httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 4096}]}),
+ httpx.Response(503),
+ )
+ )
def respond(request: httpx.Request) -> httpx.Response:
assert request.url.host == "backend.test"
@@ -354,27 +376,29 @@ async def test_discovery_preserves_input_overrides_and_survives_outages(monkeypa
assert request.headers["x-tenant"] == "tenant"
return next(responses)
- router: Final = Router(model_list=[
- {
- "model_name": "configured",
- "litellm_params": {
- "model": "hosted_vllm/local-model",
- "api_base": "https://backend.test/v1",
- "api_key": "unused-key",
- "extra_headers": {"authorization": "Bearer local-key", "X-Tenant": "tenant"},
+ router: Final = Router(
+ model_list=[
+ {
+ "model_name": "configured",
+ "litellm_params": {
+ "model": "hosted_vllm/local-model",
+ "api_base": "https://backend.test/v1",
+ "api_key": "unused-key",
+ "extra_headers": {"authorization": "Bearer local-key", "X-Tenant": "tenant"},
+ },
+ "model_info": {"id": "configured", "max_input_tokens": 1024},
},
- "model_info": {"id": "configured", "max_input_tokens": 1024},
- },
- {
- "model_name": "byok",
- "litellm_params": {
- "model": "openai/local-model",
- "api_base": "https://caller.test/v1",
- "use_clientside_credentials": True,
+ {
+ "model_name": "byok",
+ "litellm_params": {
+ "model": "openai/local-model",
+ "api_base": "https://caller.test/v1",
+ "use_clientside_credentials": True,
+ },
},
- },
- {"model_name": "default-openai", "litellm_params": {"model": "openai/local-model", "api_key": "unused"}},
- ])
+ {"model_name": "default-openai", "litellm_params": {"model": "openai/local-model", "api_key": "unused"}},
+ ]
+ )
handler: Final = AsyncHTTPHandler()
await handler.client.aclose()
responder: Final = Mock(side_effect=respond)
@@ -520,13 +544,14 @@ def test_should_not_pollute_shared_key_with_custom_nonzero_pricing():
)
-def test_regex_lookaround_flag_stays_on_the_deployment_that_set_it() -> None:
- """A deployment's ``supports_regex_lookaround`` override must not land on the shared
+@pytest.mark.parametrize("flag", ["supports_regex_lookaround", "supports_assistant_prefill"])
+def test_capability_flag_stays_on_the_deployment_that_set_it(flag: str) -> None:
+ """A deployment's capability override must not land on the shared
``{provider}/{model}`` key, or every sibling deployment of that model would inherit it."""
backend_model = "bedrock/us.xai.grok-4.6"
deploy_id = "grok-deploy-keep-regex"
- builtin_flag = litellm.get_model_info(model=backend_model).get("supports_regex_lookaround")
+ builtin_flag = litellm.get_model_info(model=backend_model).get(flag)
model_keys = {
deploy_id: litellm.model_cost.get(deploy_id),
backend_model: copy.deepcopy(litellm.model_cost.get(backend_model)),
@@ -537,13 +562,13 @@ def test_regex_lookaround_flag_stays_on_the_deployment_that_set_it() -> None:
{
"model_name": "grok-keep-regex",
"litellm_params": {"model": backend_model},
- "model_info": {"id": deploy_id, "supports_regex_lookaround": not builtin_flag},
+ "model_info": {"id": deploy_id, flag: not builtin_flag},
}
],
)
- assert litellm.model_cost[deploy_id]["supports_regex_lookaround"] is (not builtin_flag)
- assert litellm.get_model_info(model=backend_model).get("supports_regex_lookaround") is builtin_flag
+ assert litellm.model_cost[deploy_id][flag] is (not builtin_flag)
+ assert litellm.get_model_info(model=backend_model).get(flag) is builtin_flag
finally:
_restore_model_cost_entries(model_keys)
@@ -944,8 +969,7 @@ def _register_azure_tier_backend() -> None:
def test_inherit_builtin_service_tier_pricing_fills_only_missing_fields() -> None:
model_cost_entries: Final = {
- key: copy.deepcopy(litellm.model_cost.get(key))
- for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL)
+ key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL)
}
try:
_register_tier_backend()
@@ -974,9 +998,7 @@ def test_inherit_builtin_service_tier_pricing_fills_only_missing_fields() -> Non
"output_cost_per_token_flex": _TIER_BACKEND_ENTRY["output_cost_per_token_flex"],
"input_cost_per_token_balanced": _TIER_BACKEND_ENTRY["input_cost_per_token_balanced"],
"output_cost_per_token_balanced": _TIER_BACKEND_ENTRY["output_cost_per_token_balanced"],
- "cache_read_input_token_cost_ultrafast": _TIER_BACKEND_ENTRY[
- "cache_read_input_token_cost_ultrafast"
- ],
+ "cache_read_input_token_cost_ultrafast": _TIER_BACKEND_ENTRY["cache_read_input_token_cost_ultrafast"],
"input_cost_per_token_above_272k_tokens_ultrafast": _TIER_BACKEND_ENTRY[
"input_cost_per_token_above_272k_tokens_ultrafast"
],
@@ -991,8 +1013,7 @@ def test_inherit_builtin_service_tier_pricing_fills_only_missing_fields() -> Non
def test_inherit_builtin_service_tier_pricing_noop_without_base_rate_or_backend() -> None:
model_cost_entries: Final = {
- key: copy.deepcopy(litellm.model_cost.get(key))
- for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL)
+ key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL)
}
try:
_register_tier_backend()
@@ -1029,8 +1050,7 @@ def test_inherit_builtin_service_tier_pricing_noop_without_base_rate_or_backend(
def test_router_completion_uses_custom_standard_and_backend_ultrafast_pricing() -> None:
model_id: Final = "tier-priced-deployment"
model_cost_entries: Final = {
- key: copy.deepcopy(litellm.model_cost.get(key))
- for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL, model_id)
+ key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL, model_id)
}
try:
_register_tier_backend()
@@ -1090,8 +1110,7 @@ def test_router_completion_uses_custom_standard_and_backend_ultrafast_pricing()
def test_router_completion_uses_backend_ultrafast_long_context_rates() -> None:
model_id: Final = "tier-priced-long-context-deployment"
model_cost_entries: Final = {
- key: copy.deepcopy(litellm.model_cost.get(key))
- for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL, model_id)
+ key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL, model_id)
}
try:
_register_tier_backend()
@@ -1142,8 +1161,7 @@ def test_ptu_service_tier_pricing_is_disabled_only_when_attribution_is_enabled(
) -> None:
model_id: Final = f"ptu-tier-deployment-{ptu_enabled}"
model_cost_entries: Final = {
- key: copy.deepcopy(litellm.model_cost.get(key))
- for key in (_TIER_BACKEND_KEY, model_id)
+ key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_TIER_BACKEND_KEY, model_id)
}
try:
_register_tier_backend()
@@ -1200,8 +1218,7 @@ def test_azure_base_model_inherits_service_tier_pricing_for_registration_and_pay
model_id: Final = "azure-tier-priced-alias"
payload_id: Final = "azure-tier-priced-payload"
model_cost_entries: Final = {
- key: copy.deepcopy(litellm.model_cost.get(key))
- for key in (_AZURE_TIER_BACKEND_KEY, model_id, payload_id)
+ key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_AZURE_TIER_BACKEND_KEY, model_id, payload_id)
}
try:
_register_azure_tier_backend()
@@ -1257,12 +1274,8 @@ def test_azure_base_model_inherits_service_tier_pricing_for_registration_and_pay
)
)
- assert payload["input_cost_per_token_priority"] == _AZURE_TIER_BACKEND_ENTRY[
- "input_cost_per_token_priority"
- ]
- assert payload["output_cost_per_token_priority"] == _AZURE_TIER_BACKEND_ENTRY[
- "output_cost_per_token_priority"
- ]
+ assert payload["input_cost_per_token_priority"] == _AZURE_TIER_BACKEND_ENTRY["input_cost_per_token_priority"]
+ assert payload["output_cost_per_token_priority"] == _AZURE_TIER_BACKEND_ENTRY["output_cost_per_token_priority"]
finally:
_restore_model_cost_entries(model_cost_entries)
litellm.get_model_info.cache_clear()
@@ -1364,12 +1377,10 @@ def test_inherit_builtin_base_rates_for_off_peak_carries_threshold_rates():
assert model_info["input_cost_per_token"] == builtin_info["input_cost_per_token"]
assert (
- model_info["input_cost_per_token_above_200k_tokens"]
- == builtin_info["input_cost_per_token_above_200k_tokens"]
+ model_info["input_cost_per_token_above_200k_tokens"] == builtin_info["input_cost_per_token_above_200k_tokens"]
)
assert (
- model_info["output_cost_per_token_above_200k_tokens"]
- == builtin_info["output_cost_per_token_above_200k_tokens"]
+ model_info["output_cost_per_token_above_200k_tokens"] == builtin_info["output_cost_per_token_above_200k_tokens"]
)
@@ -2260,8 +2271,7 @@ def test_deployment_model_cost_payload_folds_in_litellm_params_pricing():
def test_deployment_model_cost_payload_includes_builtin_service_tier_pricing() -> None:
model_id: Final = "tier-priced-payload"
model_cost_entries: Final = {
- key: copy.deepcopy(litellm.model_cost.get(key))
- for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL, model_id)
+ key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL, model_id)
}
try:
_register_tier_backend()
@@ -2278,12 +2288,8 @@ def test_deployment_model_cost_payload_includes_builtin_service_tier_pricing() -
)
)
- assert (
- payload["input_cost_per_token_ultrafast"] == _TIER_BACKEND_ENTRY["input_cost_per_token_ultrafast"]
- )
- assert (
- payload["output_cost_per_token_ultrafast"] == _TIER_BACKEND_ENTRY["output_cost_per_token_ultrafast"]
- )
+ assert payload["input_cost_per_token_ultrafast"] == _TIER_BACKEND_ENTRY["input_cost_per_token_ultrafast"]
+ assert payload["output_cost_per_token_ultrafast"] == _TIER_BACKEND_ENTRY["output_cost_per_token_ultrafast"]
assert payload["input_cost_per_token_balanced"] == _TIER_BACKEND_ENTRY["input_cost_per_token_balanced"]
assert payload["input_cost_per_token"] == _CUSTOM_STANDARD_INPUT_RATE
assert payload["output_cost_per_token"] == _CUSTOM_STANDARD_OUTPUT_RATE
@@ -3282,9 +3288,7 @@ def test_deployment_id_colliding_with_another_providers_catalog_key_keeps_its_ow
)
assert isinstance(response, litellm.ModelResponse)
- assert response._hidden_params["response_cost"] == pytest.approx(
- 1000 * 0.00000096 + 500 * 0.00000302
- )
+ assert response._hidden_params["response_cost"] == pytest.approx(1000 * 0.00000096 + 500 * 0.00000302)
assert litellm.model_cost[_COLLIDING_BUILTIN_KEY] == builtin_row_before
finally:
_restore_model_cost_entries(model_cost_entries)
@@ -3297,8 +3301,7 @@ def test_deployment_id_matching_its_own_providers_catalog_key_still_merges() ->
matches against."""
model_id: Final = "baseten/zai-org/glm-5.2"
model_cost_entries: Final = {
- key: copy.deepcopy(litellm.model_cost.get(key))
- for key in (model_id, _COLLIDING_BUILTIN_KEY)
+ key: copy.deepcopy(litellm.model_cost.get(key)) for key in (model_id, _COLLIDING_BUILTIN_KEY)
}
try:
router: Final = Router(model_list=[_colliding_id_deployment(model_id, "baseten")])
@@ -3313,9 +3316,244 @@ def test_deployment_id_matching_its_own_providers_catalog_key_still_merges() ->
)
assert isinstance(response, litellm.ModelResponse)
- assert response._hidden_params["response_cost"] == pytest.approx(
- 1000 * 0.00000096 + 500 * 0.00000302
- )
+ assert response._hidden_params["response_cost"] == pytest.approx(1000 * 0.00000096 + 500 * 0.00000302)
finally:
_restore_model_cost_entries(model_cost_entries)
litellm.get_model_info.cache_clear()
+
+
+_CAPACITY_TEST_DEPLOYMENTS: Final = (
+ (
+ "bedrock",
+ "bedrock/converse/anthropic.claude-router-capacity-test",
+ "bedrock/us-east-1/anthropic.claude-router-capacity-test",
+ ),
+ (
+ "bedrock_mantle",
+ "bedrock_mantle/openai.router-capacity-test",
+ "bedrock_mantle/us-east-1/openai.router-capacity-test",
+ ),
+ ("openai", "openai/router-capacity-test", "openai/us/router-capacity-test"),
+)
+
+
+def _capacity_test_metadata(provider: str, max_input_tokens: int) -> dict[str, object]:
+ return {
+ "litellm_provider": provider,
+ "mode": "chat",
+ "max_input_tokens": max_input_tokens,
+ "max_output_tokens": 512,
+ "max_tokens": 512,
+ "input_cost_per_token": 0,
+ "output_cost_per_token": 0,
+ }
+
+
+def _capacity_test_router(route: str, base: str, base_location: str | None, explicit_limit: int | None) -> Router:
+ return Router(
+ model_list=[
+ {
+ "model_name": "capacity-test",
+ "litellm_params": {
+ "model": route,
+ **({"base_model": base} if base_location == "litellm_params" else {}),
+ },
+ "model_info": {
+ "id": "capacity-test",
+ **({"base_model": base} if base_location == "model_info" else {}),
+ **({"max_input_tokens": explicit_limit} if explicit_limit is not None else {}),
+ },
+ }
+ ],
+ enable_pre_call_checks=True,
+ )
+
+
+def _deployments_allowed_for_2048_input_tokens(router: Router) -> int:
+ try:
+ return len(
+ router._pre_call_checks(
+ model="capacity-test",
+ healthy_deployments=router.model_list,
+ input="synthetic input",
+ input_token_count=2048,
+ )
+ )
+ except litellm.ContextWindowExceededError:
+ return 0
+
+
+def _effective_max_input_tokens(router: Router) -> object:
+ return router.get_router_model_info(id="capacity-test", deployment=None, received_model_name="capacity-test")[
+ "max_input_tokens"
+ ]
+
+
+@pytest.mark.parametrize(("provider", "route", "base"), _CAPACITY_TEST_DEPLOYMENTS)
+@pytest.mark.parametrize("base_location", ("model_info", "litellm_params", None))
+@pytest.mark.parametrize("explicit_limit", (None, 4096))
+def test_router_capacity_widens_to_a_larger_base_model_window(
+ monkeypatch: pytest.MonkeyPatch,
+ provider: str,
+ route: str,
+ base: str,
+ base_location: str | None,
+ explicit_limit: int | None,
+) -> None:
+ monkeypatch.setattr(
+ litellm,
+ "model_cost",
+ {
+ **litellm.model_cost,
+ route: _capacity_test_metadata(provider, 1024),
+ base: _capacity_test_metadata(provider, 8192),
+ },
+ )
+ _invalidate_model_cost_lowercase_map()
+ router: Final = _capacity_test_router(route, base, base_location, explicit_limit)
+ expected: Final = (
+ explicit_limit if explicit_limit is not None else (1024 if base_location is None else 8192)
+ )
+ assert _effective_max_input_tokens(router) == expected
+ assert _deployments_allowed_for_2048_input_tokens(router) == (1 if expected >= 2048 else 0)
+ _invalidate_model_cost_lowercase_map()
+
+
+@pytest.mark.parametrize(("provider", "route", "base"), _CAPACITY_TEST_DEPLOYMENTS)
+@pytest.mark.parametrize("base_location", ("model_info", "litellm_params"))
+def test_router_capacity_never_narrows_to_a_smaller_base_model_window(
+ monkeypatch: pytest.MonkeyPatch, provider: str, route: str, base: str, base_location: str
+) -> None:
+ monkeypatch.setattr(
+ litellm,
+ "model_cost",
+ {
+ **litellm.model_cost,
+ route: _capacity_test_metadata(provider, 8192),
+ base: _capacity_test_metadata(provider, 1024),
+ },
+ )
+ _invalidate_model_cost_lowercase_map()
+ router: Final = _capacity_test_router(route, base, base_location, None)
+ assert _effective_max_input_tokens(router) == 8192
+ assert _deployments_allowed_for_2048_input_tokens(router) == 1
+ _invalidate_model_cost_lowercase_map()
+
+
+@pytest.mark.parametrize(("provider", "route", "base"), _CAPACITY_TEST_DEPLOYMENTS)
+def test_router_capacity_never_takes_a_window_from_base_model_when_the_route_is_unmapped(
+ monkeypatch: pytest.MonkeyPatch, provider: str, route: str, base: str
+) -> None:
+ monkeypatch.setattr(litellm, "model_cost", {**litellm.model_cost, base: _capacity_test_metadata(provider, 1024)})
+ _invalidate_model_cost_lowercase_map()
+ router: Final = _capacity_test_router(f"{route}-not-in-cost-map", base, "model_info", None)
+ assert _deployments_allowed_for_2048_input_tokens(router) == 1
+ _invalidate_model_cost_lowercase_map()
+
+
+@pytest.mark.parametrize(("provider", "route", "base"), _CAPACITY_TEST_DEPLOYMENTS)
+def test_router_capacity_never_takes_a_window_from_base_model_when_the_route_declares_none(
+ monkeypatch: pytest.MonkeyPatch, provider: str, route: str, base: str
+) -> None:
+ route_without_window: Final = {
+ key: value
+ for key, value in _capacity_test_metadata(provider, 0).items()
+ if key not in ("max_input_tokens", "max_tokens")
+ }
+ monkeypatch.setattr(
+ litellm,
+ "model_cost",
+ {**litellm.model_cost, route: route_without_window, base: _capacity_test_metadata(provider, 1024)},
+ )
+ _invalidate_model_cost_lowercase_map()
+ router: Final = _capacity_test_router(route, base, "model_info", None)
+ assert _effective_max_input_tokens(router) is None
+ assert _deployments_allowed_for_2048_input_tokens(router) == 1
+ _invalidate_model_cost_lowercase_map()
+
+
+_SAME_MODEL_ROUTE: Final = "bedrock/converse/anthropic.claude-routercapacitytest-9-9"
+_SAME_MODEL_GLOBAL_ROUTE: Final = "bedrock/converse/global.anthropic.claude-routercapacitytest-9-9"
+_SAME_MODEL_QUALIFIED_ROUTE: Final = "bedrock/converse/anthropic.claude-routercapacitytest-9-9@first"
+
+
+@pytest.mark.parametrize(
+ ("route", "base_model", "expected"),
+ (
+ (_SAME_MODEL_ROUTE, "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9", 8192),
+ (_SAME_MODEL_ROUTE, "us-east-1/anthropic.claude-routercapacitytest-9-9", 1024),
+ (_SAME_MODEL_ROUTE, "anthropic.claude-routercapacitytest-9-9", 1024),
+ (_SAME_MODEL_GLOBAL_ROUTE, "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9", 8192),
+ (_SAME_MODEL_ROUTE, "bedrock/anthropic.claude-routercapacitytest-9-9@global:us-east-1", 8192),
+ (_SAME_MODEL_GLOBAL_ROUTE, "bedrock/anthropic.claude-routercapacitytest-9-9@global:us-east-1", 8192),
+ (_SAME_MODEL_QUALIFIED_ROUTE, "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9@first", 8192),
+ (_SAME_MODEL_QUALIFIED_ROUTE, "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9@second", 1024),
+ (_SAME_MODEL_ROUTE, "bedrock/us-east-1/anthropic.claude-routercapacitytest-*", 1024),
+ (_SAME_MODEL_ROUTE, "bedrock/us-east-1/anthropic.claude-routercapacitystandin-9-9", 1024),
+ (_SAME_MODEL_ROUTE, "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9-extended", 1024),
+ (_SAME_MODEL_ROUTE, "bedrock/eu-west-1/anthropic.claude-routercapacitytest-9-9", 1024),
+ ),
+)
+def test_router_capacity_widens_only_for_an_exact_base_model_entry_of_the_same_model(
+ monkeypatch: pytest.MonkeyPatch, route: str, base_model: str, expected: int
+) -> None:
+ larger_window: Final = _capacity_test_metadata("bedrock", 8192)
+ monkeypatch.setattr(
+ litellm,
+ "model_cost",
+ {
+ **litellm.model_cost,
+ _SAME_MODEL_ROUTE: _capacity_test_metadata("bedrock", 1024),
+ _SAME_MODEL_GLOBAL_ROUTE: _capacity_test_metadata("bedrock", 1024),
+ _SAME_MODEL_QUALIFIED_ROUTE: _capacity_test_metadata("bedrock", 1024),
+ "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9@first": larger_window,
+ "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9@second": larger_window,
+ "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9": larger_window,
+ "anthropic.claude-routercapacitytest-9-9": _capacity_test_metadata("bedrock", 1024),
+ "bedrock/anthropic.claude-routercapacitytest-9-9": larger_window,
+ "bedrock/anthropic.claude-routercapacitytest-9-9@global:us-east-1": larger_window,
+ "bedrock/us-east-1/anthropic.claude-routercapacitytest-*": larger_window,
+ "bedrock/us-east-1/anthropic.claude-routercapacitystandin-9-9": larger_window,
+ "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9-extended": larger_window,
+ },
+ )
+ _invalidate_model_cost_lowercase_map()
+ router: Final = _capacity_test_router(route, base_model, "model_info", None)
+ assert _effective_max_input_tokens(router) == expected
+ assert _deployments_allowed_for_2048_input_tokens(router) == (1 if expected >= 2048 else 0)
+ _invalidate_model_cost_lowercase_map()
+
+
+def test_router_capacity_reads_a_model_info_base_model_before_a_litellm_params_one(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ monkeypatch.setattr(
+ litellm,
+ "model_cost",
+ {
+ **litellm.model_cost,
+ _SAME_MODEL_ROUTE: _capacity_test_metadata("bedrock", 1024),
+ "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9": _capacity_test_metadata("bedrock", 8192),
+ "bedrock/us-east-1/anthropic.claude-routercapacitystandin-9-9": _capacity_test_metadata("bedrock", 8192),
+ },
+ )
+ _invalidate_model_cost_lowercase_map()
+ router: Final = Router(
+ model_list=[
+ {
+ "model_name": "capacity-test",
+ "litellm_params": {
+ "model": _SAME_MODEL_ROUTE,
+ "base_model": "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9",
+ },
+ "model_info": {
+ "id": "capacity-test",
+ "base_model": "bedrock/us-east-1/anthropic.claude-routercapacitystandin-9-9",
+ },
+ }
+ ],
+ enable_pre_call_checks=True,
+ )
+ assert _effective_max_input_tokens(router) == 1024
+ assert _deployments_allowed_for_2048_input_tokens(router) == 0
+ _invalidate_model_cost_lowercase_map()
diff --git a/tests/unit/types/test_litellm_params.py b/tests/unit/types/test_litellm_params.py
index d27dc733c78..122d811e491 100644
--- a/tests/unit/types/test_litellm_params.py
+++ b/tests/unit/types/test_litellm_params.py
@@ -199,6 +199,7 @@ OPTION_NAMES: Final = (
"prompt_label",
"litellm_system_prompt",
"custom_prompt_dict",
+ "initial_prompt_value",
"roles",
"final_prompt_value",
"bos_token",
diff --git a/ui/litellm-dashboard/src/components/common_components/MemberTable.test.tsx b/ui/litellm-dashboard/src/components/common_components/MemberTable.test.tsx
index 418f27c8188..3264c3a99ea 100644
--- a/ui/litellm-dashboard/src/components/common_components/MemberTable.test.tsx
+++ b/ui/litellm-dashboard/src/components/common_components/MemberTable.test.tsx
@@ -1,4 +1,4 @@
-import { fireEvent, screen, waitFor, within } from "@testing-library/react";
+import { act, fireEvent, screen, waitFor, within } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import React from "react";
import { describe, expect, it, vi } from "vitest";
@@ -160,6 +160,66 @@ describe("MemberTable sorting", () => {
});
});
+describe("MemberTable paging", () => {
+ const roster = (): Member[] =>
+ Array.from({ length: 30 }, (_, index) => ({
+ user_id: `u-${index}`,
+ user_email: `member${index}@example.com`,
+ user_alias: `Member ${String(index).padStart(2, "0")}`,
+ role: index === 0 ? "admin" : "user",
+ }));
+
+ const renderRoster = () => {
+ const props = { members: roster(), canEdit: true, onEdit: vi.fn(), onDelete: vi.fn() };
+ const view = renderWithProviders();
+ return { props, view };
+ };
+
+ const goToSecondPage = async (user: ReturnType) => {
+ await user.click(screen.getByTestId("pagination-next"));
+ expect(screen.getByText("Showing 26-30 of 30")).toBeInTheDocument();
+ };
+
+ it("stays on the current page when the roster is refetched", async () => {
+ const user = userEvent.setup();
+ const { props, view } = renderRoster();
+ await goToSecondPage(user);
+
+ view.rerender();
+ await act(() => new Promise((resolve) => setTimeout(resolve, 20)));
+
+ expect(screen.getByText("Showing 26-30 of 30")).toBeInTheDocument();
+ });
+
+ it("returns to the first page when the sort order changes", async () => {
+ const user = userEvent.setup();
+ renderRoster();
+ await goToSecondPage(user);
+
+ await user.click(screen.getByTestId("sort-header-user_email"));
+
+ expect(await screen.findByText("Showing 1-25 of 30")).toBeInTheDocument();
+ });
+
+ it("returns to the first page when the search or the role filter changes", async () => {
+ const user = userEvent.setup();
+ renderRoster();
+ await goToSecondPage(user);
+
+ search("example.com");
+ expect(await screen.findByText("Showing 1-25 of 30")).toBeInTheDocument();
+
+ search("");
+ await goToSecondPage(user);
+ await user.click(screen.getByTestId("datatable-filters-trigger"));
+ await user.click(screen.getByTestId("filter-role"));
+ await user.click(await screen.findByRole("option", { name: "user" }));
+ await user.click(screen.getByTestId("filter-drawer-apply"));
+
+ expect(await screen.findByText("Showing 1-25 of 29")).toBeInTheDocument();
+ });
+});
+
describe("MemberTable role filter", () => {
it("shows only members with the chosen role and clears on reset", async () => {
const user = userEvent.setup();
diff --git a/ui/litellm-dashboard/src/components/common_components/MemberTable.tsx b/ui/litellm-dashboard/src/components/common_components/MemberTable.tsx
index 24dd86c9e2c..0acba9ad21d 100644
--- a/ui/litellm-dashboard/src/components/common_components/MemberTable.tsx
+++ b/ui/litellm-dashboard/src/components/common_components/MemberTable.tsx
@@ -1,4 +1,4 @@
-import type { ColumnDef, ColumnFiltersState } from "@tanstack/react-table";
+import type { ColumnDef, ColumnFiltersState, OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table";
import { Crown, Info, User, UserPlus } from "lucide-react";
import React, { useState } from "react";
@@ -9,6 +9,7 @@ import {
DataTableFilterField,
DataTableSortHeader,
DataTableToolbar,
+ DEFAULT_PAGE_SIZE_OPTIONS,
} from "@/components/shared/DataTable";
import { StatusBadge } from "@/components/shared/table_cells";
import { Button } from "@/components/ui/button";
@@ -44,6 +45,8 @@ export interface MemberTableProps {
const ALL_ROLES = "all";
+const DEFAULT_SORTING: SortingState = [{ id: "user_alias", desc: false }];
+
export const memberRowId = (member: Member): string => member.user_id ?? member.user_email ?? JSON.stringify(member);
export const memberRoleOptions = (members: readonly Member[]): string[] =>
@@ -230,8 +233,21 @@ export default function MemberTable({
}: MemberTableProps) {
const [globalFilter, setGlobalFilter] = useState("");
const [columnFilters, setColumnFilters] = useState([]);
+ const [sorting, setSorting] = useState(DEFAULT_SORTING);
+ const [pagination, setPagination] = useState({
+ pageIndex: 0,
+ pageSize: DEFAULT_PAGE_SIZE_OPTIONS[0],
+ });
const [filtersOpen, setFiltersOpen] = useState(false);
+ const backToFirstPage =
+ (onChange: OnChangeFn): OnChangeFn =>
+ (updaterOrValue) => {
+ onChange(updaterOrValue);
+ setPagination((current) => ({ ...current, pageIndex: 0 }));
+ };
+ const changeSearch = backToFirstPage(setGlobalFilter);
+
const columnDeps: MemberColumnDeps = {
canEdit,
onEdit,
@@ -260,13 +276,17 @@ export default function MemberTable({
data={members}
columns={columns}
getRowId={memberRowId}
+ paginationMode="client"
+ pagination={pagination}
+ onPaginationChange={setPagination}
sortingMode="client"
- defaultSorting={[{ id: "user_alias", desc: false }]}
+ sorting={sorting}
+ onSortingChange={backToFirstPage(setSorting)}
filterMode="client"
columnFilters={columnFilters}
- onColumnFiltersChange={setColumnFilters}
+ onColumnFiltersChange={backToFirstPage(setColumnFilters)}
globalFilter={globalFilter}
- onGlobalFilterChange={setGlobalFilter}
+ onGlobalFilterChange={changeSearch}
noDataMessage={
{isNarrowed ? "No members match your search or filters" : emptyText ?? "No data"}
@@ -277,7 +297,7 @@ export default function MemberTable({
setFiltersOpen(true)}
showViewOptions={false}
diff --git a/ui/litellm-dashboard/src/components/organization/organization_view.test.tsx b/ui/litellm-dashboard/src/components/organization/organization_view.test.tsx
index ca792084f2a..88d98fac60e 100644
--- a/ui/litellm-dashboard/src/components/organization/organization_view.test.tsx
+++ b/ui/litellm-dashboard/src/components/organization/organization_view.test.tsx
@@ -149,6 +149,39 @@ test("should display empty state when organization has no members", async () =>
});
});
+test("keeps the members pager on its page when an admin opens a member's edit dialog", async () => {
+ const members = Array.from({ length: 30 }, (_, index) => ({
+ user_id: `user${index + 1}`,
+ user_email: `user${index + 1}@example.com`,
+ user_role: "internal_user",
+ spend: 0,
+ user: { user_alias: `alias${String(index + 1).padStart(2, "0")}` },
+ }));
+ mockUseOrganization.mockReturnValue({ data: { ...mockOrg, members }, isLoading: false } as unknown as ReturnType<
+ typeof useOrganization
+ >);
+
+ const user = userEvent.setup();
+ renderWithProviders(
+ {}}
+ accessToken="test-token"
+ is_org_admin={true}
+ is_proxy_admin={false}
+ userModels={[]}
+ />,
+ );
+
+ await user.click(screen.getByRole("tab", { name: "Members" }));
+ await user.click(screen.getByTestId("pagination-next"));
+ expect(screen.getByText("Showing 26-30 of 30")).toBeInTheDocument();
+
+ await user.click(within(screen.getByRole("row", { name: /alias27/ })).getByTestId("edit-member"));
+
+ expect(screen.getByText("Showing 26-30 of 30")).toBeInTheDocument();
+});
+
test("should display team aliases when teams are available", async () => {
const orgWithTeams = {
...mockOrg,
diff --git a/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx b/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx
index b8f292222ed..cd51bcedb7b 100644
--- a/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx
+++ b/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx
@@ -316,6 +316,91 @@ describe("TeamMembersComponent", () => {
expect(screen.getByText("Jul 15, 2026")).toBeInTheDocument();
});
+ it("shows each member the values from its own membership row", () => {
+ const baseMembership = createMockTeamData().team_memberships[0];
+ const teamData = createMockTeamData({
+ team_memberships: [
+ {
+ ...baseMembership,
+ spend: 1.25,
+ total_spend: 10,
+ litellm_budget_table: { ...baseMembership.litellm_budget_table, max_budget: 250 },
+ },
+ {
+ ...baseMembership,
+ user_id: "user2@test.com",
+ spend: 7.5,
+ total_spend: 20,
+ litellm_budget_table: {
+ ...baseMembership.litellm_budget_table,
+ max_budget: 2500,
+ rpm_limit: 7,
+ tpm_limit: 70,
+ },
+ },
+ ],
+ });
+
+ renderWithProviders(
+ ,
+ );
+
+ const user1Row = screen.getByRole("row", { name: /user1@test\.com/ });
+ expect(within(user1Row).getByText("$1.25")).toBeInTheDocument();
+ expect(within(user1Row).getByText("$10.00")).toBeInTheDocument();
+ expect(within(user1Row).getByText("$250.00")).toBeInTheDocument();
+ expect(within(user1Row).getByText("100 RPM / 10000 TPM")).toBeInTheDocument();
+
+ const user2Row = screen.getByRole("row", { name: /user2@test\.com/ });
+ expect(within(user2Row).getByText("$7.50")).toBeInTheDocument();
+ expect(within(user2Row).getByText("$20.00")).toBeInTheDocument();
+ expect(within(user2Row).getByText("$2,500.00")).toBeInTheDocument();
+ expect(within(user2Row).getByText("7 RPM / 70 TPM")).toBeInTheDocument();
+ });
+
+ it("shows a member's new spend and budget once the team is refetched", () => {
+ const baseMembership = createMockTeamData().team_memberships[0];
+ const teamWith = (spend: number, maxBudget: number) =>
+ createMockTeamData({
+ team_memberships: [
+ {
+ ...baseMembership,
+ spend,
+ litellm_budget_table: { ...baseMembership.litellm_budget_table, max_budget: maxBudget },
+ },
+ ],
+ });
+ const renderTab = (teamData: TeamData) => (
+
+ );
+
+ const { rerender } = renderWithProviders(renderTab(teamWith(1.25, 250)));
+ rerender(renderTab(teamWith(3.5, 400)));
+
+ const row = screen.getByRole("row", { name: /user1@test\.com/ });
+ expect(within(row).getByText("$3.50")).toBeInTheDocument();
+ expect(within(row).getByText("$400.00")).toBeInTheDocument();
+ expect(within(row).queryByText("$1.25")).not.toBeInTheDocument();
+ });
+
it("should display formatted budget and Unlimited for member with no budget", () => {
renderWithProviders(
, string> = {
@@ -74,6 +74,14 @@ export default function TeamMemberTab({
const { mutate: resetMemberBudget, isPending: isResettingBudget } = useResetTeamMemberBudget();
const teamDefaultBudget = teamData.team_info.team_member_budget_table?.max_budget ?? null;
+ const membershipsByUserId = useMemo(
+ () => new Map(teamData.team_memberships.map((membership) => [membership.user_id, membership])),
+ [teamData.team_memberships],
+ );
+
+ const findMembership = (userId: string | null): TeamMembership | undefined =>
+ userId ? membershipsByUserId.get(userId) : undefined;
+
const formatNumber = (value: number | null): string => {
if (value === null || value === undefined) return "0";
@@ -95,32 +103,28 @@ export default function TeamMemberTab({
const getUserCurrentCycleSpend = (userId: string | null): number => {
if (!userId) return 0;
- const membership = teamData.team_memberships.find((tm) => tm.user_id === userId);
- return membership?.spend ?? 0;
+ return findMembership(userId)?.spend ?? 0;
};
const getUserTotalSpend = (userId: string | null): number => {
if (!userId) return 0;
- const membership = teamData.team_memberships.find((tm) => tm.user_id === userId);
- return membership?.total_spend ?? 0;
+ return findMembership(userId)?.total_spend ?? 0;
};
const getUserBudgetSource = (userId: string | null): TeamMemberBudgetSource => {
if (!userId) return "none";
- const membership = teamData.team_memberships.find((tm) => tm.user_id === userId);
- return membership?.budget_source ?? "none";
+ return findMembership(userId)?.budget_source ?? "none";
};
const getUserBudget = (userId: string | null): number | null => {
if (!userId) return null;
- const membership = teamData.team_memberships.find((tm) => tm.user_id === userId);
- return membership?.litellm_budget_table?.max_budget ?? teamDefaultBudget;
+ return findMembership(userId)?.litellm_budget_table?.max_budget ?? teamDefaultBudget;
};
// Helper function to get rate limits for a user
const getUserRateLimits = (userId: string | null): string => {
if (!userId) return "No Limits";
- const membership = teamData.team_memberships.find((tm) => tm.user_id === userId);
+ const membership = findMembership(userId);
const rpmLimit = membership?.litellm_budget_table?.rpm_limit;
const tpmLimit = membership?.litellm_budget_table?.tpm_limit;
@@ -139,15 +143,13 @@ export default function TeamMemberTab({
const getUserAllowedModels = (userId: string | null): string[] | null => {
if (!userId) return null;
- const membership = teamData.team_memberships.find((tm) => tm.user_id === userId);
- const models = membership?.litellm_budget_table?.allowed_models;
+ const models = findMembership(userId)?.litellm_budget_table?.allowed_models;
return models && models.length > 0 ? models : null;
};
const getUserBudgetReset = (userId: string | null): string | null => {
if (!userId) return null;
- const membership = teamData.team_memberships.find((tm) => tm.user_id === userId);
- return membership?.litellm_budget_table?.budget_reset_at ?? null;
+ return findMembership(userId)?.litellm_budget_table?.budget_reset_at ?? null;
};
const extraColumns: NonNullable["extraColumns"]> = [
@@ -312,7 +314,7 @@ export default function TeamMemberTab({
members={teamData.team_info.members_with_roles}
canEdit={canEditTeam}
onEdit={(record) => {
- const membership = teamData.team_memberships.find((tm) => tm.user_id === record.user_id);
+ const membership = findMembership(record.user_id);
setSelectedEditMember(
seedMemberBudgetFields(
record,