fix(mcp): preserve scoped dispatch and nested request limits

This commit is contained in:
moe-berri 2026-10-30 10:05:03 -07:00
commit bd26d07780
56 changed files with 3793 additions and 328 deletions

View file

@ -420,6 +420,8 @@ REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS: Final = get_env_float(
"REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS", 20.0
)
BEDROCK_REALTIME_CLOSE_TIMEOUT_SECONDS: Final = get_env_float("BEDROCK_REALTIME_CLOSE_TIMEOUT_SECONDS", 5.0)
REALTIME_LAZY_SETUP_DEADLINE_SECONDS: Final = 240.0
REALTIME_LAZY_SETUP_CONTENT_GRACE_SECONDS: Final = 1.0
# RFC 6455 caps the close frame payload at 125 bytes, 2 of which carry the status code
WEBSOCKET_CLOSE_REASON_MAX_BYTES: Final = 123

View file

@ -85,6 +85,7 @@ OPTIONAL_KWARGS_KEYS: Final = (
"github_copilot_user_session",
"fireworks_forward_user_id",
"responses_api_version",
"allowed_openai_params",
PROVIDER_AFFINITY_HEADER_KWARG_KEY,
}
)

View file

@ -2023,12 +2023,14 @@ def infer_content_type_from_url_and_content(
"""
from litellm.litellm_core_utils.token_counter import get_image_type
media_type: Final = current_content_type.split(";", 1)[0].strip().lower() if current_content_type else None
# If we have a valid content type that's not generic, use it
if current_content_type and current_content_type not in [
if media_type and media_type not in [
"binary/octet-stream",
"application/octet-stream",
]:
return current_content_type
return media_type
# Extension to MIME type mapping
# Supports images, documents, and other common file types

View file

@ -912,9 +912,9 @@ def create_anthropic_image_param(
image_url_input: str | dict,
format: str | None = None,
is_bedrock_invoke: bool = False,
) -> AnthropicMessagesImageParam:
) -> AnthropicMessagesImageParam | AnthropicMessagesDocumentParam:
"""
Create an AnthropicMessagesImageParam from an image URL input.
Create an Anthropic content block from an image URL input.
Supports both URL references (for HTTP/HTTPS URLs) and base64 encoding.
"""
@ -933,14 +933,7 @@ def create_anthropic_image_param(
if is_bedrock_invoke or image_url.startswith("http://"):
base64_url: Final = convert_url_to_base64(url=image_url)
image_chunk = convert_to_anthropic_image_obj(openai_image_url=base64_url, format=format)
return AnthropicMessagesImageParam(
type="image",
source=AnthropicContentParamSource(
type="base64",
media_type=image_chunk["media_type"],
data=image_chunk["data"],
),
)
return _anthropic_content_element_factory(image_chunk)
else:
# HTTPS URL - pass directly for regular Anthropic
return AnthropicMessagesImageParam(
@ -953,14 +946,7 @@ def create_anthropic_image_param(
else:
# Convert to base64 for data URIs or other formats
image_chunk = convert_to_anthropic_image_obj(openai_image_url=image_url, format=format)
return AnthropicMessagesImageParam(
type="image",
source=AnthropicContentParamSource(
type="base64",
media_type=image_chunk["media_type"],
data=image_chunk["data"],
),
)
return _anthropic_content_element_factory(image_chunk)
# The following XML functions will be deprecated once JSON schema support is available on Bedrock and Vertex
@ -1910,7 +1896,12 @@ def convert_to_anthropic_tool_result(
anthropic_content_element=_anthropic_image_param,
original_content_element=content,
)
anthropic_content_list.append(cast(AnthropicMessagesImageParam, _anthropic_image_param))
anthropic_content_list.append(
cast(
AnthropicMessagesImageParam | AnthropicMessagesDocumentParam,
_anthropic_image_param,
)
)
elif content["type"] == "tool_reference":
anthropic_content_list.append(ToolReference(type="tool_reference", tool_name=content["tool_name"]))
elif content["type"] == "file":

View file

@ -1,5 +1,6 @@
import asyncio
import json
import time
import traceback
from collections.abc import Callable, Coroutine, Mapping, MutableMapping, Sequence
from dataclasses import dataclass
@ -15,13 +16,19 @@ from litellm._logging import redact_internal_details_from_client_message, verbos
from litellm._uuid import uuid
from litellm.constants import (
BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY,
REALTIME_LAZY_SETUP_CONTENT_GRACE_SECONDS,
REALTIME_LAZY_SETUP_DEADLINE_SECONDS,
REALTIME_METERED_USAGE_KEY,
REALTIME_SESSION_FAILURE_LOGGED_KEY,
REALTIME_SESSION_SUCCESS_LOGGED_KEY,
)
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.llms.base_llm.realtime.transcription_protocol import RealtimeTranscriptionProtocolError
from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig, RealtimeBackend
from litellm.llms.base_llm.realtime.transformation import (
BaseRealtimeConfig,
RealtimeBackend,
RealtimeSessionUpdateRejection,
)
from litellm.types.llms.openai import (
OpenAIRealtimeEvents,
OpenAIRealtimeOutputItemDone,
@ -125,6 +132,65 @@ class _LoggingWorker(Protocol):
_DECODED_JSON: Final = TypeAdapter(object)
_SESSION_FIELDS: Final = TypeAdapter(dict[str, object])
_NO_SESSION_FIELDS: Final[Mapping[str, object]] = MappingProxyType({})
_SESSION_MERGE_MAX_DEPTH: Final = 8
class _SessionUpdateFrame(TypedDict):
type: ReadOnly[str]
session: ReadOnly[Mapping[str, object]]
def _merge_session_fields(
base: Mapping[str, object], update: Mapping[str, object], depth: int = 0
) -> Mapping[str, object]:
def merged_value(key: str) -> object:
if key not in update:
return base[key]
value: Final = update[key]
base_value: Final = base.get(key)
if depth < _SESSION_MERGE_MAX_DEPTH and isinstance(value, Mapping) and isinstance(base_value, Mapping):
return _merge_session_fields(
_SESSION_FIELDS.validate_python(base_value), _SESSION_FIELDS.validate_python(value), depth + 1
)
return update[key]
return MappingProxyType({key: merged_value(key) for key in dict.fromkeys((*base, *update))})
def _session_update_frame(session: Mapping[str, object]) -> str:
frame: Final[_SessionUpdateFrame] = {"type": "session.update", "session": session}
return json.dumps(frame, default=_SESSION_FIELDS.validate_python)
_HELD_SESSION_UPDATE_TOO_LARGE: Final = RealtimeSessionUpdateRejection(
code="session_update_too_large",
message=(
"The session.update frames sent before the first conversation content add up to more than 10 MB, "
"so this session.update was not applied"
),
param="session",
)
def _unscreened_audio_error_event(client_event_id: object) -> Mapping[str, object]:
return {
"type": "error",
"event_id": f"event_{uuid.uuid4()}",
"error": {
"type": "invalid_request_error",
"code": "audio_input_not_screenable",
"message": (
"A realtime_input_transcription guardrail is on, and Gemini Live only transcribes a turn after it "
"starts answering it, so audio input was not sent. Send text with conversation.item.create instead"
),
"param": None,
"event_id": client_event_id if isinstance(client_event_id, str) else None,
},
}
def _decode_json_object(payload: str) -> Mapping[str, object]:
"""Decode a realtime frame into its top-level field mapping."""
return json.loads(payload)
@ -243,8 +309,17 @@ class RealTimeStreaming:
event_normalizer: RealtimeEventNormalizer | None = None,
usage_meter: RealtimeUsageMeter | None = None,
logging_worker: _LoggingWorker = GLOBAL_LOGGING_WORKER,
lazy_setup_deadline_seconds: float = REALTIME_LAZY_SETUP_DEADLINE_SECONDS,
lazy_setup_content_grace_seconds: float = REALTIME_LAZY_SETUP_CONTENT_GRACE_SECONDS,
clock: Callable[[], float] = time.monotonic,
):
self.websocket: _ClientWebSocket = websocket
self._clock: Final = clock
self._lazy_setup_deadline: Final = clock() + lazy_setup_deadline_seconds
self._lazy_setup_content_grace_seconds: Final = lazy_setup_content_grace_seconds
self._held_client_content: tuple[str, ...] = ()
self._held_client_content_bytes: int = 0
self._held_client_content_deadline: float | None = None
self._replay_scope: Final = client_websocket_scope(websocket)
self.backend_ws = backend_ws
self.logging_obj = logging_obj
@ -300,6 +375,10 @@ class RealTimeStreaming:
# Gemini Live rejects a follow-up BidiGenerateContentSetup once any
# content (realtimeInput / clientContent / toolResponse) has been sent.
self._content_sent_after_setup: bool = False
self._held_session_update: Mapping[str, object] | None = None
self._applied_session_fields: Mapping[str, object] | None = None
self._held_session_update_bytes: int = 0
self._unscreened_audio_refused: bool = False
# Whether this is a transcription-only session (session.type == "transcription",
# e.g. gpt-realtime-whisper). Such sessions must not be sent response.create and
# their input_audio_transcription.completed usage drives duration-based cost.
@ -337,6 +416,17 @@ class RealTimeStreaming:
]
)
_CLIENT_AUDIO_BUFFER_COMMIT_TYPES = frozenset(["input_audio_buffer.commit", "input_audio_buffer.end"])
_CLIENT_CONTENT_TYPES = frozenset(
[
"conversation.item.create",
"response.create",
"input_audio_buffer.append",
"input_audio_buffer.commit",
"input_audio_buffer.end",
]
)
_HELD_CONTENT_TYPES: Final = frozenset(["conversation.item.create", "input_audio_buffer.append"])
_DEFERRED_SETUP_HELD_CONTENT_TYPES: Final = frozenset(["conversation.item.create"])
_AUDIO_FORMAT_MAP: dict[str, dict[str, str | int]] = {
"pcm16": {"type": "audio/pcm", "rate": 24000},
"g711_ulaw": {"type": "audio/G711-ulaw", "rate": 8000},
@ -377,7 +467,7 @@ class RealTimeStreaming:
message_obj: dict[str, Any] = cast(dict[str, Any], message)
else:
message_obj = cast(dict[str, Any], json.loads(message))
self._collect_tool_calls_from_response_done(cast(dict, message_obj))
self._collect_tool_calls_from_response_done(message_obj)
if not self._should_store_message(message_obj):
return
try:
@ -502,7 +592,7 @@ class RealTimeStreaming:
self.store_message(flush_event)
self._capture_transcription_usage(flush_event)
def _collect_tool_calls_from_response_done(self, event_obj: dict | OpenAIRealtimeEvents) -> None:
def _collect_tool_calls_from_response_done(self, event_obj: Mapping[str, object] | OpenAIRealtimeEvents) -> None:
"""Extract function_call items from response.done events for spend logging."""
try:
if event_obj.get("type") != "response.done":
@ -607,6 +697,167 @@ class RealTimeStreaming:
self._usage_meter.observe_client_event(sent_event)
return True
def _builds_setup_lazily(self) -> bool:
if self.provider_config is None:
return False
return self.provider_config.builds_setup_from_session_update()
def _holds_session_updates_until_first_content(self) -> bool:
return self.session_configuration_request is None and self._builds_setup_lazily()
async def _rejected_setup_session_update(self, message: str) -> bool:
provider_config: Final = self.provider_config
if provider_config is None or not provider_config.builds_setup_from_session_update():
return False
client_event: Final = _decode_json_object(message)
session: Final = client_event.get("session")
update: Final = _SESSION_FIELDS.validate_python(session) if isinstance(session, Mapping) else _NO_SESSION_FIELDS
if self._leaves_applied_setup_unchanged(update):
await self._acknowledge_session_setup(self._applied_session_fields or _NO_SESSION_FIELDS)
return True
rejection: Final = provider_config.setup_session_update_rejection(
self.model, update, setup_sent=self.session_configuration_request is not None
)
if rejection is None:
return False
await self.websocket.send_text(json.dumps(rejection.error_event(client_event.get("event_id"))))
return True
def _held_session_update_message(self) -> str:
return _session_update_frame(self._held_session_update or _NO_SESSION_FIELDS)
async def _hold_session_update(self, message: str) -> None:
client_event: Final = _decode_json_object(message)
if self._held_session_update_bytes + len(message.encode("utf-8")) > RealTimeStreaming._MAX_BUFFERED_BYTES:
await self.websocket.send_text(
json.dumps(_HELD_SESSION_UPDATE_TOO_LARGE.error_event(client_event.get("event_id")))
)
return
session: Final = client_event.get("session")
update: Final = _SESSION_FIELDS.validate_python(session) if isinstance(session, Mapping) else _NO_SESSION_FIELDS
held: Final = _merge_session_fields(self._held_session_update or _NO_SESSION_FIELDS, update)
self._held_session_update = held
self._held_session_update_bytes = len(_session_update_frame(held).encode("utf-8"))
await self._acknowledge_session_setup(held)
def _setup_preview(self, session: Mapping[str, object]) -> str | None:
if self.provider_config is None:
return None
preview: Final = self.provider_config.transform_realtime_request(
_session_update_frame(session), self.model, None
)
return next(
(frame for frame in preview if isinstance(frame, str) and "setup" in _decode_json_object(frame)), None
)
def _leaves_applied_setup_unchanged(self, update: Mapping[str, object]) -> bool:
applied: Final = self._applied_session_fields
if self.session_configuration_request is None or applied is None:
return False
applied_setup: Final = self._setup_preview(applied)
merged_setup: Final = self._setup_preview(_merge_session_fields(applied, update))
return (
applied_setup is not None
and merged_setup is not None
and _decode_json_object(applied_setup) == _decode_json_object(merged_setup)
)
async def _acknowledge_session_setup(self, session: Mapping[str, object]) -> None:
if self.provider_config is None:
return
setup: Final = self._setup_preview(session)
if setup is None:
return
acknowledged: Final = self.provider_config.transform_session_created_event(
self.model, str(self.logging_obj.litellm_trace_id), setup
)
if acknowledged is None:
return
event: Final[Mapping[str, object]] = {**acknowledged, "type": "session.updated"}
event_str: Final = json.dumps(event)
self.store_message(event_str)
await self._send_event_to_client(event, event_str)
async def _send_held_session_update_as_setup(self) -> None:
held: Final = self._held_session_update
if held is None:
return
message: Final = self._held_session_update_message()
self._held_session_update = None
self._held_session_update_bytes = 0
if await self._send_to_backend(message):
self._applied_session_fields = held
await self._send_held_client_content()
def _holds_client_content(self, msg_type: str | None, message: str) -> bool:
held_types: Final = (
RealTimeStreaming._DEFERRED_SETUP_HELD_CONTENT_TYPES
if self._uses_deferred_backend_setup()
else RealTimeStreaming._HELD_CONTENT_TYPES
)
return (
msg_type in held_types
and self._held_session_update is not None
and self._holds_session_updates_until_first_content()
and not self._pending_messages_until_setup
and len(self._held_client_content) < RealTimeStreaming._MAX_BUFFERED_MESSAGES
and self._held_client_content_bytes + len(message.encode("utf-8")) <= RealTimeStreaming._MAX_BUFFERED_BYTES
)
def _hold_client_content(self, message: str) -> None:
if self._held_client_content_deadline is None:
self._held_client_content_deadline = self._clock() + self._lazy_setup_content_grace_seconds
self._held_client_content = (*self._held_client_content, message)
self._held_client_content_bytes += len(message.encode("utf-8"))
async def _send_held_client_content(self) -> None:
held: Final = self._held_client_content
self._held_client_content = ()
self._held_client_content_bytes = 0
self._held_client_content_deadline = None
for message in held:
await self._send_to_backend(message)
def _setup_deadline(self) -> float:
content_deadline: Final = self._held_client_content_deadline
return (
self._lazy_setup_deadline if content_deadline is None else min(self._lazy_setup_deadline, content_deadline)
)
async def _send_setup_before_content(self, msg_type: str | None) -> None:
if self._holds_session_updates_until_first_content():
await self._send_held_session_update_as_setup()
await self._send_default_setup_before_first_client_message(msg_type)
async def _receive_client_message_before_setup_deadline(self) -> str | None:
if not self._holds_session_updates_until_first_content():
return await self._receive_client_message()
receive: Final = asyncio.ensure_future(self._receive_client_message())
try:
done, _ = await asyncio.wait((receive,), timeout=self._setup_deadline() - self._clock())
if not done:
await self._send_setup_before_content(None)
return await receive
finally:
receive.cancel()
async def _send_default_setup_before_first_client_message(self, msg_type: str | None) -> None:
if self.session_configuration_request is not None or self.provider_config is None:
return
if msg_type == "session.update" or (
self._uses_deferred_backend_setup()
and (msg_type in RealTimeStreaming._CLIENT_AUDIO_BUFFER_TYPES or self._pending_messages_until_setup)
):
return
default_setup: Final = self.provider_config.session_configuration_request(self.model)
if not default_setup:
return
setup: Final = self.maybe_inject_guardrail_auto_response_disable(default_setup)
await self.backend_ws.send(setup)
self._cache_session_configuration_request(setup)
self._applied_session_fields = _NO_SESSION_FIELDS
await self._send_held_client_content()
def _enforce_transcription_session_model(self, message: str) -> str:
"""Force client transcription session updates to the authorized model.
@ -851,6 +1102,20 @@ class RealTimeStreaming:
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 _send_blocked_turn_to_backend(self, pre_block_backend_message: str | None, item_id: str | None) -> None:
if pre_block_backend_message is None and self._is_transcription_session:
return
await self._send_setup_before_content("conversation.item.create")
# Deliver any caller-supplied backend message FIRST so that
# protocol contracts requiring a specific ordering (e.g.
# Gemini Live's mandatory ``toolResponse`` after a
# ``toolCall``) are honored before the guardrail's own
# clientContent / cancel messages are sent.
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._withdraw_blocked_turn(item_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:
@ -958,6 +1223,21 @@ class RealTimeStreaming:
"""Transcription-only sessions have no assistant turn to gate and reject realtime session updates."""
return not self._is_transcription_session and self._has_audio_transcription_guardrails()
def _refuses_unscreened_audio(self, msg_type: str | None) -> bool:
return (
msg_type in RealTimeStreaming._CLIENT_AUDIO_BUFFER_TYPES
and self._builds_setup_lazily()
and self._should_disable_vad_auto_response()
)
async def _refuse_unscreened_audio(self, msg_type: str | None, message: str) -> None:
if self._unscreened_audio_refused and msg_type not in RealTimeStreaming._CLIENT_AUDIO_BUFFER_COMMIT_TYPES:
return
self._unscreened_audio_refused = True
await self.websocket.send_text(
json.dumps(_unscreened_audio_error_event(_decode_json_object(message).get("event_id")))
)
async def run_realtime_guardrails(
self,
transcript: str,
@ -1030,15 +1310,7 @@ class RealTimeStreaming:
# Use realtime_violation_message if configured; fall back to guardrail error text.
error_msg = getattr(callback, "realtime_violation_message", None) or safe_msg
# Deliver any caller-supplied backend message FIRST so that
# protocol contracts requiring a specific ordering (e.g.
# Gemini Live's mandatory ``toolResponse`` after a
# ``toolCall``) are honored before the guardrail's own
# clientContent / cancel messages are sent.
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._withdraw_blocked_turn(item_id)
await self._send_blocked_turn_to_backend(pre_block_backend_message, item_id)
# Send the policy violation hint (shows as small gray status text in UI).
await self.websocket.send_text(
json.dumps(
@ -1136,6 +1408,9 @@ class RealTimeStreaming:
break
finally:
self._flushing_pending_messages_until_setup = False
if self.session_created_sent_to_client and self._builds_setup_lazily():
await self._maybe_send_guardrail_turn_detection_update()
continue
if self.session_created_sent_to_client:
# A synthetic session.created (with placeholder defaults) was
# already forwarded to the client when we connected. The
@ -1295,7 +1570,11 @@ class RealTimeStreaming:
return BackendClose(code=1011, reason="proxy failed while relaying the upstream websocket")
def _backend_refused_session(self, close: BackendClose) -> bool:
return close.code != 1000 and not self._backend_sent_frames
return (
close.code != 1000
and not self._backend_sent_frames
and not self._holds_session_updates_until_first_content()
)
async def _log_session_end(self, transport_refusal: Exception | None) -> None:
provider_refusal: Final = None if self.provider_config is None else self.provider_config.session_refusal()
@ -1488,7 +1767,7 @@ class RealTimeStreaming:
client_event: _ClientEventFrame
try:
while True:
message = await self._receive_client_message()
message = await self._receive_client_message_before_setup_deadline()
if message is None:
return ClientLoopExit.CLIENT_DISCONNECTED
@ -1600,6 +1879,7 @@ class RealTimeStreaming:
if (
msg_type == "session.update"
and self.session_configuration_request is None
and self._held_session_update is None
and not self._guardrail_turn_detection_update_sent
and self._should_disable_vad_auto_response()
and (
@ -1663,12 +1943,25 @@ class RealTimeStreaming:
except (json.JSONDecodeError, AttributeError):
pass
if self._refuses_unscreened_audio(msg_type):
await self._refuse_unscreened_audio(msg_type, message)
continue
## LOGGING
# Log after any in-place modifications (GA remap, guardrail
# turn_detection injection) so audit logs reflect what we
# actually forward to the backend.
self.store_input(message=message)
if self._holds_client_content(msg_type, message):
self._hold_client_content(message)
continue
if self._held_client_content and msg_type != "session.update":
await self._send_setup_before_content(msg_type)
if msg_type in RealTimeStreaming._CLIENT_CONTENT_TYPES and self._uses_deferred_backend_setup():
await self._send_setup_before_content(msg_type)
if self._should_buffer_client_message_until_setup(message):
self._buffer_pending_message_until_setup(message)
continue
@ -1697,6 +1990,15 @@ class RealTimeStreaming:
# ``_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()
if msg_type == "session.update" and await self._rejected_setup_session_update(message):
continue
if msg_type == "session.update" and self._holds_session_updates_until_first_content():
await self._hold_session_update(message)
if self._pending_messages_until_setup:
await self._send_held_session_update_as_setup()
continue
if msg_type in RealTimeStreaming._CLIENT_CONTENT_TYPES:
await self._send_setup_before_content(msg_type)
try:
sent = await self._send_to_backend(message)
except RealtimeTranscriptionProtocolError as e:

View file

@ -647,12 +647,14 @@ class LiteLLMMessagesToCompletionTransformationHandler:
request_data["output_config"] = extra_kwargs["output_config"]
custom_llm_provider: Final = extra_kwargs.get("custom_llm_provider")
api_base: Final = extra_kwargs.get("api_base")
(
openai_request,
tool_name_mapping,
) = ANTHROPIC_ADAPTER.translate_completion_input_params_with_tool_mapping(
request_data,
custom_llm_provider=custom_llm_provider if isinstance(custom_llm_provider, str) else None,
api_base=api_base if isinstance(api_base, str) else None,
drop_params=drops_unsupported_params(extra_kwargs.get("drop_params")),
)

View file

@ -24,6 +24,7 @@ OPENAI_MAX_TOOL_NAME_LENGTH: Final = 64
TOOL_NAME_HASH_LENGTH: Final = 8
TOOL_NAME_PREFIX_LENGTH: Final = OPENAI_MAX_TOOL_NAME_LENGTH - TOOL_NAME_HASH_LENGTH - 1 # 55
PROVIDERS_PROXYING_AN_UNKNOWN_BACKEND: Final = frozenset({"litellm_proxy"})
PROVIDERS_SERVED_BY_THE_MISTRAL_CHAT_CONFIG: Final = frozenset({"mistral", "codestral"})
_COMPACTION_BLOCK: Final = TypeAdapter(dict[str, JsonValue])
@ -255,7 +256,12 @@ class AnthropicAdapter:
return result
def translate_completion_input_params_with_tool_mapping(
self, kwargs, *, custom_llm_provider: str | None = None, drop_params: bool = True
self,
kwargs,
*,
custom_llm_provider: str | None = None,
api_base: str | None = None,
drop_params: bool = True,
) -> tuple[ChatCompletionRequest | None, dict[str, str]]:
"""
Translate Anthropic request params to OpenAI format, returning tool name mapping.
@ -289,6 +295,7 @@ class AnthropicAdapter:
) = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai(
anthropic_message_request=request_body,
custom_llm_provider=custom_llm_provider,
api_base=api_base,
drop_params=drop_params,
)
@ -1001,7 +1008,9 @@ class LiteLLMAnthropicMessagesAdapter:
)
@staticmethod
def _supports_prompt_cache_key(model: str | None, custom_llm_provider: str | None) -> bool:
def _supports_prompt_cache_key(
prompt_cache_key: str, model: str | None, custom_llm_provider: str | None, api_base: str | None
) -> bool:
if not model or not custom_llm_provider:
return False
if custom_llm_provider in PROVIDERS_PROXYING_AN_UNKNOWN_BACKEND:
@ -1009,7 +1018,17 @@ class LiteLLMAnthropicMessagesAdapter:
supported_params: Final = litellm.get_supported_openai_params(
model=model, custom_llm_provider=custom_llm_provider
)
return "prompt_cache_key" in (supported_params or ())
if "prompt_cache_key" not in (supported_params or ()):
return False
if custom_llm_provider not in PROVIDERS_SERVED_BY_THE_MISTRAL_CHAT_CONFIG:
return True
_, _, _, resolved_api_base = litellm.get_llm_provider(
model=model, custom_llm_provider=custom_llm_provider, api_base=api_base
)
mistral_config: Final = litellm.MistralConfig()
return mistral_config.takes_prompt_cache_key_value(
prompt_cache_key
) and mistral_config.accepts_prompt_cache_key_at(resolved_api_base)
@staticmethod
def _target_declares_reasoning_effort(model: str, custom_llm_provider: str | None) -> bool:
@ -1042,6 +1061,7 @@ class LiteLLMAnthropicMessagesAdapter:
new_kwargs: ChatCompletionRequest,
*,
custom_llm_provider: str | None = None,
api_base: str | None = None,
) -> None:
"""Translate metadata fields from Anthropic request to OpenAI request."""
if "metadata" in anthropic_message_request:
@ -1050,7 +1070,7 @@ class LiteLLMAnthropicMessagesAdapter:
new_kwargs["user"] = metadata["user_id"]
prompt_cache_key: Final = prompt_cache_key_from_user_id(metadata["user_id"])
if prompt_cache_key is not None and self._supports_prompt_cache_key(
anthropic_message_request.get("model"), custom_llm_provider
prompt_cache_key, anthropic_message_request.get("model"), custom_llm_provider, api_base
):
new_kwargs["prompt_cache_key"] = prompt_cache_key
@ -1244,6 +1264,7 @@ class LiteLLMAnthropicMessagesAdapter:
anthropic_message_request: AnthropicMessagesRequest,
*,
custom_llm_provider: str | None = None,
api_base: str | None = None,
preserve_midturn_system: bool = False,
drop_params: bool = True,
) -> tuple[ChatCompletionRequest, dict[str, str]]:
@ -1282,6 +1303,7 @@ class LiteLLMAnthropicMessagesAdapter:
anthropic_message_request=anthropic_message_request,
new_kwargs=new_kwargs,
custom_llm_provider=custom_llm_provider,
api_base=api_base,
)
## CONVERT TOOL CHOICE
self._translate_tool_choice_to_openai(

View file

@ -1,10 +1,12 @@
import dataclasses
from abc import ABC, abstractmethod
from collections.abc import Mapping, Sequence
from types import TracebackType
from typing import TYPE_CHECKING, Any, Protocol
from typing import TYPE_CHECKING, Any, Literal, Protocol
import httpx
from litellm._uuid import uuid
from litellm.types.llms.openai import OpenAIRealtimeStreamSessionEvents
from litellm.types.realtime import (
RealtimeInputAudioTranscriptionUsage,
@ -22,6 +24,26 @@ else:
LiteLLMLoggingObj = Any
@dataclasses.dataclass(frozen=True, slots=True)
class RealtimeSessionUpdateRejection:
code: Literal["invalid_type", "invalid_value", "session_update_not_applied", "session_update_too_large"]
message: str
param: str
def error_event(self, client_event_id: object) -> dict[str, object]:
return {
"type": "error",
"event_id": f"event_{uuid.uuid4()}",
"error": {
"type": "invalid_request_error",
"code": self.code,
"message": self.message,
"param": self.param,
"event_id": client_event_id if isinstance(client_event_id, str) else None,
},
}
class RealtimeBackend(Protocol):
async def __aenter__(self) -> object: ...
@ -93,6 +115,14 @@ class BaseRealtimeConfig(ABC):
def session_configuration_request(self, model: str) -> str | None: # message sent to setup the realtime session
return None
def builds_setup_from_session_update(self) -> bool:
return False
def setup_session_update_rejection(
self, model: str, session: Mapping[str, object], *, setup_sent: bool
) -> RealtimeSessionUpdateRejection | None:
return None
def unbilled_usage_on_session_close(self, model: str) -> RealtimeInputAudioTranscriptionUsage | None:
return None

View file

@ -17,7 +17,6 @@ from pydantic import JsonValue, TypeAdapter
import litellm
from litellm._logging import redact_string, verbose_proxy_logger
from litellm._uuid import uuid
from litellm.constants import (
BEDROCK_REALTIME_CLOSE_TIMEOUT_SECONDS,
BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY,
@ -31,13 +30,14 @@ from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.litellm_core_utils.realtime_streaming import DefaultLoggedRealTimeEventTypes
from litellm.llms.base_llm.realtime.transformation import RealtimeSessionUpdateRejection
from litellm.types.llms.bedrock import AwsAuthParams
from litellm.types.llms.openai import OpenAIRealtimeEvents
from litellm.types.realtime import RealtimeResponseTransformInput
from ..base_aws_llm import BaseAWSLLM, run_aws_signing
from ..common_utils import BedrockError
from .transformation import BedrockRealtimeConfig, BedrockSessionUpdateRejection, requested_modalities
from .transformation import BedrockRealtimeConfig, requested_modalities
_CLIENT_MESSAGE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
@ -256,23 +256,11 @@ async def _ack_session_update(
async def _reject_session_update(
client_ws: RealtimeClientWebSocket,
parsed_client_message: Mapping[str, JsonValue],
rejection: BedrockSessionUpdateRejection,
rejection: RealtimeSessionUpdateRejection,
) -> bool:
"""Send an error event for a session.update Bedrock cannot apply; False means the client is gone."""
client_event_id: Final = parsed_client_message.get("event_id")
error_event: Final = {
"type": "error",
"event_id": f"event_{uuid.uuid4()}",
"error": {
"type": "invalid_request_error",
"code": rejection.code,
"message": rejection.message,
"param": rejection.param,
"event_id": client_event_id if isinstance(client_event_id, str) else None,
},
}
try:
await client_ws.send_text(json.dumps(error_event))
await client_ws.send_text(json.dumps(rejection.error_event(parsed_client_message.get("event_id"))))
except Exception as e: # noqa: BLE001 # any send failure means the client is gone
verbose_proxy_logger.debug("Client to Bedrock forwarding ended: %s", e, exc_info=True)
return False

View file

@ -11,7 +11,7 @@ import math
import uuid as uuid_lib
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import Final, Literal
from typing import Final
import httpx
from pydantic import ConfigDict, JsonValue, TypeAdapter
@ -19,7 +19,7 @@ from pydantic import ConfigDict, JsonValue, TypeAdapter
from litellm._logging import verbose_logger
from litellm._uuid import uuid
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig
from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig, RealtimeSessionUpdateRejection
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.llms.bedrock.realtime.trigger_audio import ready_trigger_pcm
from litellm.types.llms.base import LiteLLMBaseModel
@ -196,13 +196,6 @@ def _malformed_settings(session_config: Mapping[str, JsonValue]) -> tuple[tuple[
return tuple((name, expected) for name, expected, is_malformed in malformed if is_malformed)
@dataclasses.dataclass(frozen=True, slots=True)
class BedrockSessionUpdateRejection:
code: Literal["invalid_type", "invalid_value", "session_update_not_applied"]
message: str
param: str
@dataclasses.dataclass(frozen=True, slots=True)
class BedrockLockedSessionSettings:
"""The session settings Bedrock Nova Sonic locks in once the prompt starts."""
@ -533,7 +526,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
max_tokens=max_tokens,
)
def session_update_rejection(self, json_message: Mapping[str, JsonValue]) -> BedrockSessionUpdateRejection | None:
def session_update_rejection(self, json_message: Mapping[str, JsonValue]) -> RealtimeSessionUpdateRejection | None:
"""The error for a session.update this session cannot apply, or None when it applies."""
session_config: Final = json_message.get("session")
if not isinstance(session_config, Mapping):
@ -541,7 +534,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
malformed: Final = () if self.drop_params else _malformed_settings(session_config)
if malformed:
name, expected = malformed[0]
return BedrockSessionUpdateRejection(
return RealtimeSessionUpdateRejection(
code="invalid_type",
message=(
f"Invalid type for 'session.{name}': Bedrock Nova Sonic needs {expected} there. "
@ -551,7 +544,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
)
unplayable_voice: Final = None if self.drop_params else _unplayable_voice(session_config)
if unplayable_voice is not None:
return BedrockSessionUpdateRejection(
return RealtimeSessionUpdateRejection(
code="invalid_value",
message=(
f"Invalid value for 'session.voice': '{unplayable_voice}' is not a Bedrock Nova Sonic voice. "
@ -562,7 +555,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
if not self.prompt_started:
return None
if self._apply_locked_settings(session_config, self._started_settings) != self._started_settings:
return BedrockSessionUpdateRejection(
return RealtimeSessionUpdateRejection(
code="session_update_not_applied",
message="Bedrock Nova Sonic cannot change tools, instructions, voice, output audio format, temperature or max output tokens after the session has started, so this session.update was not applied",
param="session",

View file

@ -6973,14 +6973,19 @@ class BaseLLMHTTPHandler:
),
)
# Auto-send session setup if the provider requires it (e.g.
# Gemini/Vertex AI Live needs a `setup` before any realtime_input).
# Auto-send session setup at connect unless the provider builds it from the
# client's first session.update (Gemini/Vertex AI Live accept exactly one
# setup, so RealTimeStreaming sends theirs on the first client message).
# Build the streaming handler first so a transcription guardrail's
# auto-response disable can be folded into this one setup: Gemini
# rejects a second setup, so a follow-up disable would be dropped
# and the guardrail bypassed.
_session_config: str | None = None
if provider_config.requires_session_configuration():
sends_setup_at_connect: Final = (
provider_config.requires_session_configuration()
and not provider_config.builds_setup_from_session_update()
)
if sends_setup_at_connect:
_session_config = provider_config.session_configuration_request(model)
if _session_config:
_session_config = realtime_streaming.maybe_inject_guardrail_auto_response_disable(

View file

@ -7,6 +7,7 @@ from collections import OrderedDict
from collections.abc import Mapping, Sequence
from typing import Final, cast
from pydantic import TypeAdapter
from typing_extensions import ReadOnly, Required, TypedDict
import litellm
@ -15,8 +16,9 @@ from litellm._uuid import uuid
from litellm.litellm_core_utils.core_helpers import as_str_mapping
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.url_utils import join_api_base
from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig
from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig, RealtimeSessionUpdateRejection
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
OPENAI_TO_GEMINI_TTS_VOICES,
VertexGeminiConfig,
)
from litellm.responses.litellm_completion_transformation.transformation import (
@ -76,27 +78,82 @@ MAP_GEMINI_FIELD_TO_OPENAI_EVENT: Final[dict[str, OpenAIRealtimeEventTypes | Res
# Keys the main transform loop handles; siblings like ``usageMetadata`` are skipped.
_KNOWN_GEMINI_TOP_LEVEL_KEYS: Final[set] = {map_key.split(".", 1)[0] for map_key in MAP_GEMINI_FIELD_TO_OPENAI_EVENT}
_GEMINI_LIVE_RESPONSE_MODALITIES: Final = frozenset({"TEXT", "AUDIO"})
OPENAI_STOCK_REALTIME_VOICES: Final[frozenset[str]] = frozenset(
{"alloy", "ash", "ballad", "cedar", "coral", "echo", "marin", "sage", "shimmer", "verse"}
)
# Per https://docs.cloud.google.com/vertex-ai/generative-ai/docs/live-api/configure-language-voice (2026-10-28)
GEMINI_LIVE_DOCUMENTED_VOICES: Final[frozenset[str]] = frozenset(
{
"achernar",
"achird",
"algenib",
"algieba",
"alnilam",
"aoede",
"autonoe",
"callirrhoe",
"charon",
"despina",
"enceladus",
"erinome",
"fenrir",
"gacrux",
"iapetus",
"kore",
"laomedeia",
"leda",
"orus",
"puck",
"pulcherrima",
"rasalgethi",
"sadachbia",
"sadaltager",
"schedar",
"sulafat",
"umbriel",
"vindemiatrix",
"zephyr",
"zubenelgenubi",
}
)
def _gemini_live_speech_config(voice: object) -> Mapping[str, object] | None:
def _is_undocumented_live_voice(voice: object) -> bool:
return (
isinstance(voice, str)
and voice.lower() not in OPENAI_STOCK_REALTIME_VOICES
and OPENAI_TO_GEMINI_TTS_VOICES.get(voice.lower(), voice).lower() not in GEMINI_LIVE_DOCUMENTED_VOICES
)
def _gemini_live_speech_config(voice: object, drop_undocumented: bool) -> Mapping[str, object] | None:
"""Build the Gemini Live speechConfig for a client-requested voice.
OpenAI stock voice names have no Gemini equivalent and Gemini Live closes
the session on an unknown voice, so they are dropped with a warning and
the model keeps its default voice. Every other name is forwarded verbatim.
OpenAI stock voice names have no Gemini equivalent, so they are always
dropped, and so is any name outside Google's documented Live voices when
``drop_undocumented`` is set, keeping the model's default voice. Every
other name is forwarded verbatim.
"""
if isinstance(voice, str) and voice.lower() in OPENAI_STOCK_REALTIME_VOICES:
if not isinstance(voice, str):
return VertexGeminiConfig().map_audio_params({"voice": voice})
if voice.lower() in OPENAI_STOCK_REALTIME_VOICES:
verbose_logger.warning(
"Gemini Realtime: voice %s is an OpenAI voice with no Gemini equivalent; "
"dropping it so the session keeps the model's default voice.",
voice,
)
return None
if drop_undocumented and _is_undocumented_live_voice(voice):
verbose_logger.warning(
"Gemini Realtime: voice %s is not a documented Gemini Live voice; "
"drop_params is set, so it is dropped and the session keeps the model's default voice.",
voice,
)
return None
return VertexGeminiConfig().map_audio_params({"voice": voice})
@ -142,16 +199,34 @@ def _base64_decoded_byte_count(data: str) -> int:
return max(len(data) * 3 // 4 - padding, 0)
_JSON_ARRAY: Final = TypeAdapter(tuple[object, ...])
def _chat_completion_tool(tool: object) -> dict[str, object]:
fields: Final = as_str_mapping(tool) or {}
if fields.get("type") != "function" or "function" in fields:
return {**fields}
return {"type": "function", "function": {key: value for key, value in fields.items() if key != "type"}}
def _system_instruction_text(system_instruction: HttpxContentType | None) -> str | None:
parts: Final = system_instruction.get("parts", []) if system_instruction else []
texts: Final = tuple(text for part in parts if isinstance(text := part.get("text"), str))
return "\n".join(texts) if texts else None
class GeminiRealtimeConfig(BaseRealtimeConfig):
_TOOL_CALL_ID_TO_NAME_MAX = 256 # LRU cap for call_id→name mapping
def __init__(self):
def __init__(self, drop_params: bool = False):
super().__init__()
self.drop_params = drop_params
self._tool_call_id_to_name: OrderedDict[str, str] = OrderedDict()
# Gemini Live sometimes emits usageMetadata in a standalone frame between
# turns; buffer it here so the next response.done carries the token counts.
self._pending_usage_metadata: dict | None = None
self._unbilled_input_audio_bytes: int = 0
self._manual_activity_open: bool = False
def is_setup_message(self, msg_obj: dict) -> bool:
return "setup" in msg_obj
@ -167,6 +242,10 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
"""Google AI Studio closes the socket (1007) on padding or silence tuning next to ``disabled: true``."""
return True
def _rejects_undocumented_voices(self, model: str) -> bool:
"""Google AI Studio closes the socket (1007) on an unknown voice only for its native-audio Live models."""
return bool(GeminiRealtimeConfig._model_cost_entry(model).get("gemini_native_audio"))
@staticmethod
def _usage_detail_alias(details: object, defaults: dict[str, int]) -> dict[str, object]:
detail_mapping: Final = as_str_mapping(details)
@ -272,17 +351,26 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
except (json.JSONDecodeError, TypeError, AttributeError):
return False
def _start_manual_activity(self, session_configuration_request: str | None) -> tuple[str, ...]:
if self._manual_activity_open or not self._manual_turn_detection_enabled(session_configuration_request):
return ()
self._manual_activity_open = True
activity_start: Final[BidiGenerateContentRealtimeInput] = {"activityStart": {}}
return (json.dumps({"realtimeInput": activity_start}),)
def _end_manual_activity(self) -> tuple[str, ...]:
if not self._manual_activity_open:
return ()
self._manual_activity_open = False
activity_end: Final[BidiGenerateContentRealtimeInput] = {"activityEnd": {}}
return (json.dumps({"realtimeInput": activity_end}),)
def _handle_input_audio_buffer_commit_or_end(self, session_configuration_request: str | None) -> list[str]:
"""Map OpenAI buffer commit/end to Gemini Live turn-boundary signals."""
if self._manual_turn_detection_enabled(session_configuration_request):
realtime_input_dict: BidiGenerateContentRealtimeInput = {
"activityEnd": True,
}
verbose_logger.debug("Gemini Realtime: Sending activityEnd realtimeInput to backend")
else:
realtime_input_dict = {"audioStreamEnd": True}
verbose_logger.debug("Gemini Realtime: Sending audioStreamEnd realtimeInput to backend")
return [json.dumps({"realtimeInput": realtime_input_dict})]
return list(self._end_manual_activity())
audio_stream_end: Final[BidiGenerateContentRealtimeInput] = {"audioStreamEnd": True}
return [json.dumps({"realtimeInput": audio_stream_end})]
def map_automatic_turn_detection(self, value: OpenAIRealtimeTurnDetection) -> AutomaticActivityDetection:
"""Map OpenAI ``server_vad`` to Gemini ``automaticActivityDetection``.
@ -325,7 +413,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
"voice",
]
def map_openai_params(self, optional_params: dict, non_default_params: dict) -> dict:
def map_openai_params(self, optional_params: dict, non_default_params: dict, *, model: str) -> dict:
if "generationConfig" not in optional_params:
optional_params["generationConfig"] = {}
for key, value in non_default_params.items():
@ -333,12 +421,12 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
optional_params["systemInstruction"] = HttpxContentType(role="user", parts=[{"text": value}])
elif key == "temperature":
optional_params["generationConfig"]["temperature"] = value
elif key == "max_response_output_tokens":
elif key == "max_response_output_tokens" and isinstance(value, int):
optional_params["generationConfig"]["maxOutputTokens"] = value
elif key == "modalities":
optional_params["generationConfig"]["responseModalities"] = [
modality.upper() for modality in cast(list[str], value)
]
optional_params["generationConfig"]["responseModalities"] = list(
GeminiRealtimeConfig._requested_response_modalities(cast(object, value))
)
elif key == "tools":
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
VertexGeminiConfig,
@ -347,34 +435,43 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
vertex_gemini_config = VertexGeminiConfig()
# Tools should be at the top level of setup, not inside generationConfig
optional_params["tools"] = vertex_gemini_config.map_function(
value=value, optional_params=optional_params
value=[_chat_completion_tool(tool) for tool in _JSON_ARRAY.validate_python(value)],
optional_params=optional_params,
)
elif key == "input_audio_transcription" and value is not None:
optional_params["inputAudioTranscription"] = {}
elif key == "turn_detection" and value is not None:
value_typed = cast(OpenAIRealtimeTurnDetection, value)
if (
isinstance(value_typed, dict)
and value_typed.get("type") == "semantic_vad"
and "create_response" not in value_typed
):
# Pipecat/OpenAI GA semantic VAD — skip; Gemini uses its own VAD.
# Only skip when there is no create_response override so that
# a guardrail-injected create_response:false is not dropped.
continue
transformed_audio_activity_config = self.map_automatic_turn_detection(value_typed)
if transformed_audio_activity_config:
optional_params["realtimeInputConfig"] = BidiGenerateContentRealtimeInputConfig(
automaticActivityDetection=transformed_audio_activity_config
)
elif key == "turn_detection":
realtime_input_config = self._realtime_input_config(cast(object, value))
if realtime_input_config is not None:
optional_params["realtimeInputConfig"] = realtime_input_config
elif key == "voice":
speech_config = _gemini_live_speech_config(value)
speech_config = _gemini_live_speech_config(
value, self.drop_params and self._rejects_undocumented_voices(model)
)
if speech_config:
optional_params["generationConfig"]["speechConfig"] = speech_config
if len(optional_params["generationConfig"]) == 0:
optional_params.pop("generationConfig")
return optional_params
def _realtime_input_config(self, turn_detection: object) -> BidiGenerateContentRealtimeInputConfig | None:
if turn_detection is None:
return BidiGenerateContentRealtimeInputConfig(automaticActivityDetection={"disabled": True})
value_typed: Final = cast(OpenAIRealtimeTurnDetection, turn_detection)
if (
isinstance(value_typed, dict)
and value_typed.get("type") == "semantic_vad"
and "create_response" not in value_typed
):
# Pipecat/OpenAI GA semantic VAD — skip; Gemini uses its own VAD.
# Only skip when there is no create_response override so that
# a guardrail-injected create_response:false is not dropped.
return None
transformed_audio_activity_config: Final = self.map_automatic_turn_detection(value_typed)
if not transformed_audio_activity_config:
return None
return BidiGenerateContentRealtimeInputConfig(automaticActivityDetection=transformed_audio_activity_config)
@staticmethod
def _extract_turn_detection(session: dict) -> dict | None:
"""Extract turn_detection from a session.update payload.
@ -423,6 +520,8 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
if isinstance(input_cfg, dict):
if "input_audio_transcription" not in normalized and "transcription" in input_cfg:
normalized["input_audio_transcription"] = input_cfg["transcription"]
if "turn_detection" not in normalized and "turn_detection" in input_cfg:
normalized["turn_detection"] = input_cfg["turn_detection"]
output_cfg: Final = audio.get("output")
if isinstance(output_cfg, dict) and output_cfg.get("voice"):
normalized["voice"] = output_cfg["voice"]
@ -434,7 +533,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
return normalized
@staticmethod
def _model_cost_entry(model: str) -> dict:
def _model_cost_entry(model: str) -> Mapping[str, object]:
entry = litellm.model_cost.get(model)
if entry is None:
stripped: Final = model.split("/", 1)[-1]
@ -455,18 +554,122 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
return "TEXT" if GeminiRealtimeConfig._is_text_only_live_model(model) else "AUDIO"
@staticmethod
def _coerce_response_modalities(model: str, modalities: Sequence[object]) -> tuple[str, ...]:
"""Swap responseModalities a Live model cannot produce: TEXT to AUDIO for
audio-only models, AUDIO to TEXT for text-only ones (e.g. transcribe-live)."""
normalized: Final = tuple(
modality.upper() if isinstance(modality, str) else str(modality).upper() for modality in modalities
def _requested_response_modalities(value: object) -> tuple[str, ...]:
if isinstance(value, str):
return (value.upper(),)
if not isinstance(value, (list, tuple)):
return ()
return tuple(str(modality).upper() for modality in _JSON_ARRAY.validate_python(value))
@staticmethod
def _unsupported_response_modalities(model: str, requested: tuple[str, ...]) -> tuple[str, ...]:
unknown: Final = tuple(modality for modality in requested if modality not in _GEMINI_LIVE_RESPONSE_MODALITIES)
if unknown:
return unknown
if GeminiRealtimeConfig._is_audio_only_live_model(model) and frozenset(requested) == frozenset(("TEXT",)):
return ("TEXT",)
if GeminiRealtimeConfig._is_text_only_live_model(model) and frozenset(requested) == frozenset(("AUDIO",)):
return ("AUDIO",)
return ()
def setup_session_update_rejection(
self, model: str, session: Mapping[str, object], *, setup_sent: bool
) -> RealtimeSessionUpdateRejection | None:
if self.drop_params:
return None
if setup_sent:
return RealtimeSessionUpdateRejection(
code="session_update_not_applied",
message=(
"Gemini Live accepts one session configuration per connection, and this session already sent "
"its configuration, so this session.update was not applied. Put every setting in session.update "
"before the first conversation content, or set drop_params to drop it instead"
),
param="session",
)
return self._modality_rejection(model, session) or self._voice_rejection(model, session)
@staticmethod
def _modality_rejection(model: str, session: Mapping[str, object]) -> RealtimeSessionUpdateRejection | None:
param: Final = "output_modalities" if "output_modalities" in session else "modalities"
unsupported: Final = GeminiRealtimeConfig._unsupported_response_modalities(
model, GeminiRealtimeConfig._requested_response_modalities(session.get(param))
)
if not unsupported:
return None
return RealtimeSessionUpdateRejection(
code="invalid_value",
message=(
f"Invalid value for 'session.{param}': {model} cannot produce "
f"{', '.join(modality.lower() for modality in unsupported)} output. Set drop_params to drop it instead"
),
param=f"session.{param}",
)
def _voice_rejection(self, model: str, session: Mapping[str, object]) -> RealtimeSessionUpdateRejection | None:
audio_output: Final = as_str_mapping((as_str_mapping(session.get("audio")) or {}).get("output")) or {}
output_voice: Final = audio_output.get("voice")
voice: Final = output_voice or session.get("voice")
if not (self._rejects_undocumented_voices(model) and _is_undocumented_live_voice(voice)):
return None
param: Final = "session.audio.output.voice" if output_voice else "session.voice"
return RealtimeSessionUpdateRejection(
code="invalid_value",
message=(
f"Invalid value for '{param}': {voice} is not a Gemini Live voice {model} offers. "
"Set drop_params to drop it instead"
),
param=param,
)
@staticmethod
def _coerce_response_modalities(model: str, modalities: Sequence[object]) -> tuple[str, ...]:
"""Drop modalities Gemini Live cannot produce at all, swap the ones this
model cannot (TEXT to AUDIO for audio-only models, AUDIO to TEXT for
text-only ones such as transcribe-live), then keep a single modality
because Gemini Live rejects a setup naming two."""
return GeminiRealtimeConfig._single_response_modality(
GeminiRealtimeConfig._swap_unsupported_response_modalities(
model, GeminiRealtimeConfig._known_response_modalities(model, modalities)
)
)
@staticmethod
def _known_response_modalities(model: str, modalities: Sequence[object]) -> tuple[str, ...]:
normalized: Final = tuple(str(modality).upper() for modality in modalities)
known: Final = tuple(modality for modality in normalized if modality in _GEMINI_LIVE_RESPONSE_MODALITIES)
if len(known) < len(normalized):
verbose_logger.warning(
"Gemini Live: %s only produces TEXT or AUDIO; ignoring the requested %s",
model,
", ".join(modality for modality in normalized if modality not in _GEMINI_LIVE_RESPONSE_MODALITIES),
)
return known or (GeminiRealtimeConfig._default_response_modality(model),)
@staticmethod
def _swap_unsupported_response_modalities(model: str, normalized: tuple[str, ...]) -> tuple[str, ...]:
if GeminiRealtimeConfig._is_audio_only_live_model(model) and "TEXT" in normalized:
verbose_logger.warning(
"Gemini Live: %s only produces audio; downgrading the requested TEXT response modality to AUDIO",
model,
)
return tuple(modality for modality in normalized if modality != "TEXT") or ("AUDIO",)
if GeminiRealtimeConfig._is_text_only_live_model(model) and "AUDIO" in normalized:
return tuple(modality for modality in normalized if modality != "AUDIO") or ("TEXT",)
return normalized
@staticmethod
def _single_response_modality(modalities: tuple[str, ...]) -> tuple[str, ...]:
if len(modalities) <= 1:
return modalities
chosen: Final = "AUDIO" if "AUDIO" in modalities else modalities[0]
verbose_logger.warning(
"Gemini Live: at most one response modality is allowed per session; using %s for the requested %s",
chosen,
", ".join(modalities),
)
return (chosen,)
@staticmethod
def _finalize_gemini_live_setup(model: str, setup: dict[str, object]) -> dict[str, object]:
generation_config: Final = setup.get("generationConfig")
@ -504,7 +707,9 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
# would be silently dropped because ``map_openai_params`` only
# recognises the flat OpenAI-beta key names.
session_payload = self._normalize_session_payload_for_mapping(session_payload)
new_overrides: Final = self.map_openai_params(optional_params={}, non_default_params=session_payload)
new_overrides: Final = self.map_openai_params(
optional_params={}, non_default_params=session_payload, model=model
)
if session_configuration_request is None:
generation_config: Final = new_overrides.setdefault("generationConfig", {})
@ -618,14 +823,13 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
message_str = str(message)
raise ValueError(f"Invalid JSON message: {message_str}")
messages: Final[list[str]] = []
msg_type: Final = json_message.get("type")
if msg_type == "session.update":
return self._handle_session_update(json_message, model, session_configuration_request)
if msg_type == "response.create":
return [] # Gemini responds automatically; nothing to forward
return list(self._end_manual_activity())
if msg_type == "conversation.item.create":
return self._handle_conversation_item(json_message)
@ -643,8 +847,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
gemini_msg: Final = json.dumps({"realtimeInput": realtime_input_dict})
verbose_logger.debug("Gemini Realtime: Sending audio realtimeInput to backend")
messages.append(gemini_msg)
return messages
return [*self._start_manual_activity(session_configuration_request), gemini_msg]
if msg_type in ("input_audio_buffer.commit", "input_audio_buffer.end"):
return self._handle_input_audio_buffer_commit_or_end(session_configuration_request)
@ -660,22 +863,21 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
logging_session_id: str,
session_configuration_request: str | None = None,
) -> OpenAIRealtimeStreamSessionEvents:
if session_configuration_request:
session_configuration_request_dict: BidiGenerateContentSetup = _parse_setup(session_configuration_request)
else:
session_configuration_request_dict = {}
session_configuration_request_dict: Final = _parse_setup(
session_configuration_request or self.session_configuration_request(model)
)
_model: Final = session_configuration_request_dict.get("model") or model
generation_config: Final = session_configuration_request_dict.get("generationConfig", {}) or {}
gemini_modalities: Final = generation_config.get("responseModalities", ["AUDIO"])
_modalities: Final = [modality.lower() for modality in cast(list[str], gemini_modalities)]
_system_instruction: Final = session_configuration_request_dict.get("systemInstruction")
instructions: Final = _system_instruction_text(session_configuration_request_dict.get("systemInstruction"))
session: Final = OpenAIRealtimeStreamSession(
id=logging_session_id,
modalities=_modalities,
)
if _system_instruction is not None and isinstance(_system_instruction, str):
session["instructions"] = _system_instruction
if instructions is not None:
session["instructions"] = instructions
if _model is not None and isinstance(_model, str):
# Strip Vertex/AI Studio path prefixes to expose the bare model name.
if "/models/" in _model:
@ -1662,6 +1864,9 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
# Deferred setup opt-in: litellm.gemini_live_defer_setup = True
return not litellm.gemini_live_defer_setup
def builds_setup_from_session_update(self) -> bool:
return True
def session_configuration_request(self, model: str) -> str:
"""

View file

@ -9,11 +9,16 @@ Docs - https://docs.mistral.ai/api/
from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence
from itertools import accumulate, count
from typing import TYPE_CHECKING, Final, Literal, cast, get_type_hints, overload
from urllib.parse import urlparse
import httpx
from pydantic import TypeAdapter
from typing_extensions import TypeIs # noqa: TID251 # TypeIs reaches typing only on 3.13
import litellm
from litellm._logging import verbose_logger
from litellm.exceptions import UnsupportedParamsError
from litellm.litellm_core_utils.core_helpers import as_str_mapping
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.prompt_templates.common_utils import (
convert_content_list_to_str,
@ -30,7 +35,15 @@ from litellm.router_utils.reasoning_effort_capability import (
nearest_declared_reasoning_effort,
)
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.mistral import MistralTextBlock, MistralThinkingBlock, MistralToolCallMessage
from litellm.types.llms.mistral import (
MistralCacheKeyParams,
MistralFilePart,
MistralReasoningEffortParams,
MistralRequestRouting,
MistralTextBlock,
MistralThinkingBlock,
MistralToolCallMessage,
)
from litellm.types.llms.openai import AllMessageValues, ChatCompletionAssistantMessage, ChatCompletionTextObject
from litellm.types.utils import ModelResponse, ModelResponseStream
from litellm.utils import convert_to_model_response_object, supports_reasoning
@ -57,6 +70,17 @@ def _accepted_reasoning_effort(model: str, requested: str, custom_llm_provider:
_NO_TOOL_CALL_TURN: Final[tuple[int, frozenset[str]]] = (-1, frozenset())
_MISTRAL_REQUEST_ROUTING: Final = TypeAdapter(MistralRequestRouting)
_MISTRAL_CACHE_KEY_PARAMS: Final = TypeAdapter(MistralCacheKeyParams)
_MISTRAL_REASONING_EFFORT_PARAMS: Final = TypeAdapter(MistralReasoningEffortParams)
_MISTRAL_FILE_PART: Final = TypeAdapter(MistralFilePart)
def _reasoning_level_and_summary(reasoning_effort: object) -> tuple[object, object]:
requested: Final = as_str_mapping(reasoning_effort)
if requested is None:
return reasoning_effort, None
return requested.get("effort"), requested.get("summary")
def _is_json_array(value: object) -> TypeIs[Sequence[object]]: # guard-ok: a decoded JSON array is a list
@ -88,6 +112,8 @@ class MistralConfig(OpenAIGPTConfig):
- `response_format` (object or null): An object specifying the format that the model must output. Setting to { "type": "json_object" } enables JSON mode, which guarantees the message the model generates is in JSON. When using JSON mode you MUST also instruct the model to produce JSON yourself with a system or a user message.
"""
_is_base_class = False
temperature: int | None = None
top_p: int | None = None
max_tokens: int | None = None
@ -131,6 +157,7 @@ class MistralConfig(OpenAIGPTConfig):
def get_supported_openai_params(self, model: str) -> list[str]:
is_magistral: Final = "magistral" in model.lower()
accepts_reasoning_effort: Final = self._accepts_reasoning(model)
accepts_prompt_cache_key: Final = self.custom_llm_provider == "mistral"
return [
"stream",
"temperature",
@ -145,6 +172,7 @@ class MistralConfig(OpenAIGPTConfig):
"parallel_tool_calls",
*(("thinking",) if is_magistral else ()),
*(("reasoning_effort",) if accepts_reasoning_effort else ()),
*(("prompt_cache_key",) if accepts_prompt_cache_key else ()),
]
def _map_tool_choice(self, tool_choice: str) -> str:
@ -211,15 +239,84 @@ class MistralConfig(OpenAIGPTConfig):
optional_params["extra_body"] = {"random_seed": value}
if param == "response_format":
optional_params["response_format"] = value
if param == "reasoning_effort" and "magistral" not in model.lower():
optional_params["reasoning_effort"] = _accepted_reasoning_effort(model, value, self.custom_llm_provider)
if (
param == "reasoning_effort"
and "magistral" not in model.lower()
and (
reasoning_level := self._reasoning_level_to_send(
model, _MISTRAL_REASONING_EFFORT_PARAMS.validate_python(non_default_params), drop_params
)
)
is not None
):
optional_params["reasoning_effort"] = reasoning_level
if param in ("reasoning_effort", "thinking") and "magistral" in model.lower():
# Flag that we need to add reasoning system prompt
optional_params["_add_reasoning_prompt"] = True
if param == "parallel_tool_calls":
optional_params["parallel_tool_calls"] = value
if param in ("parallel_tool_calls", "prompt_cache_key"):
optional_params[param] = value
return optional_params
def _reasoning_level_to_send(
self, model: str, requested: MistralReasoningEffortParams, drop_params: bool
) -> str | None:
level, summary = _reasoning_level_and_summary(requested.get("reasoning_effort"))
sendable: Final = (
_accepted_reasoning_effort(model, level, self.custom_llm_provider) if isinstance(level, str) else None
)
if sendable is not None and summary is None:
return sendable
if drop_params or litellm.drop_params is True:
return sendable
raise UnsupportedParamsError(
message=(
f"{self.custom_llm_provider} does not support parameters: ['reasoning_effort'], for model={model}, "
"since the Mistral API takes it only as a level such as 'high', with no reasoning summary. "
"To drop it, set `litellm.drop_params=True` or for proxy:\n\n`litellm_settings:\n drop_params: true`\n"
),
llm_provider=self.custom_llm_provider,
model=model,
)
def accepts_prompt_cache_key_at(self, api_base: str | None) -> bool:
resolved_api_base, _ = self._get_openai_compatible_provider_info(api_base, None)
hostname: Final = urlparse(resolved_api_base).hostname or ""
return hostname == "mistral.ai" or hostname.endswith(".mistral.ai")
@staticmethod
def takes_prompt_cache_key_value(prompt_cache_key: object) -> bool:
return isinstance(prompt_cache_key, str) and prompt_cache_key.isascii()
def _prompt_cache_key_rejection(self, prompt_cache_key: object, api_base: str | None) -> str | None:
if not self.accepts_prompt_cache_key_at(api_base):
return "since only the Mistral API (*.mistral.ai) takes it"
if not self.takes_prompt_cache_key_value(prompt_cache_key):
return "since the Mistral API takes it only as a string of ASCII characters"
return None
def _prompt_cache_key_withheld(self, model: str, prompt_cache_key: object, routing: MistralRequestRouting) -> bool:
requested_api_base: Final = routing.get("api_base")
api_base: Final = requested_api_base if isinstance(requested_api_base, str) else None
allowed_openai_params: Final = routing.get("allowed_openai_params")
if isinstance(allowed_openai_params, (list, tuple)) and "prompt_cache_key" in allowed_openai_params:
return False
rejection: Final = self._prompt_cache_key_rejection(prompt_cache_key, api_base)
if rejection is None:
return False
if litellm.drop_params is True or routing.get("drop_params") is True:
return True
requested_provider: Final = routing.get("custom_llm_provider")
provider: Final = requested_provider if isinstance(requested_provider, str) else self.custom_llm_provider
raise UnsupportedParamsError(
message=(
f"{provider} does not support parameters: ['prompt_cache_key'], for model={model}, {rejection}. "
"To send it anyway, pass `allowed_openai_params=['prompt_cache_key']`. To drop it, set "
"`litellm.drop_params=True` or for proxy:\n\n`litellm_settings:\n drop_params: true`\n"
),
llm_provider=provider,
model=model,
)
def _get_openai_compatible_provider_info(self, api_base: str | None, api_key: str | None) -> tuple[str, str | None]:
# mistral is openai compatible, we just need to set this to custom_openai and have the api_base be https://api.mistral.ai
api_base = (
@ -359,24 +456,31 @@ class MistralConfig(OpenAIGPTConfig):
messages = self._handle_message_with_file(messages)
return messages
@staticmethod
def _with_mistral_file_id(block: Mapping[str, object]) -> Mapping[str, object]:
file_part: Final = block.get("file") if block.get("type") == "file" else None
file_id: Final = (
_MISTRAL_FILE_PART.validate_python(file_part).get("file_id") if isinstance(file_part, Mapping) else None
)
if not file_id:
return block
return {**{key: value for key, value in block.items() if key != "file"}, "file_id": file_id}
@classmethod
def _with_mistral_file_ids(cls, message: AllMessageValues) -> AllMessageValues:
content: Final = message.get("content")
if not isinstance(content, list):
return message
blocks: Final[tuple[Mapping[str, object], ...]] = tuple(content)
return cast( # cast-ok: only file blocks change, to the file_id shape the Mistral API takes
AllMessageValues, {**message, "content": [cls._with_mistral_file_id(block) for block in blocks]}
)
def _handle_message_with_file(self, messages: list[AllMessageValues]) -> list[AllMessageValues]:
"""
Mistral API supports only 'file_id' in message content with type 'file'.
"""
for m in messages:
_content_block = m.get("content")
if _content_block and isinstance(_content_block, list):
if any(c.get("type") == "file" for c in _content_block):
# If file content is present, we get file_id from 'file' attribute of content block
# then replace 'file' with 'file_id' and assign the value of 'file_id' attribute to it.
file_contents = [c for c in _content_block if c.get("type") == "file"]
for file_content in file_contents:
file_id = file_content.get("file", {}).get("file_id")
if file_id:
# Replace 'file' with 'file_id'
file_content["file_id"] = file_id # pyright: ignore[reportGeneralTypeIssues] # legacy in-place rewrite of the block shape
file_content.pop("file", None)
return messages
return [self._with_mistral_file_ids(m) for m in messages]
def _add_reasoning_system_prompt_if_needed(
self, messages: list[AllMessageValues], optional_params: dict
@ -794,7 +898,16 @@ class MistralConfig(OpenAIGPTConfig):
if "magistral" in model.lower() and optional_params.get("_add_reasoning_prompt", False):
messages = self._add_reasoning_system_prompt_if_needed(messages, optional_params)
upstream_params: Final = {key: value for key, value in optional_params.items() if key != "client_metadata"}
cache_key_params: Final = _MISTRAL_CACHE_KEY_PARAMS.validate_python(optional_params)
withheld_params: Final = (
("client_metadata", "prompt_cache_key")
if "prompt_cache_key" in cache_key_params
and self._prompt_cache_key_withheld(
model, cache_key_params["prompt_cache_key"], _MISTRAL_REQUEST_ROUTING.validate_python(litellm_params)
)
else ("client_metadata",)
)
upstream_params: Final = {key: value for key, value in optional_params.items() if key not in withheld_params}
# Call parent transform_request which handles _transform_messages
return super().transform_request(

View file

@ -27,6 +27,7 @@ import litellm
from litellm import LlmProviders
from litellm._logging import verbose_logger
from litellm.constants import DEFAULT_MAX_RETRIES
from litellm.exceptions import UnsupportedParamsError
from litellm.files.types import FileContentStreamingResult
from litellm.litellm_core_utils.core_helpers import set_provider_response_headers_in_hidden_params
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@ -971,8 +972,8 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
litellm.remove_index_from_tool_calls(messages=messages)
else:
raise e
except OpenAIError as e:
raise e
except (OpenAIError, UnsupportedParamsError):
raise
except Exception as e:
status_code: Final = getattr(e, "status_code", 500)
error_headers = getattr(e, "headers", None)

View file

@ -28,7 +28,7 @@ from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import LlmProviders
from ...openai.chat.gpt_transformation import OpenAIGPTConfig
from ..common_utils import AWSEventStreamDecoder, SagemakerError
from ..common_utils import INFERENCE_COMPONENT_HEADER, OBJECT_KEYED_MAPPING, AWSEventStreamDecoder, SagemakerError
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
@ -59,7 +59,12 @@ class SagemakerChatConfig(OpenAIGPTConfig, BaseAWSLLM):
inference_component_name: Final = optional_params.get("model_id")
if not isinstance(inference_component_name, str):
return headers
return {**headers, "X-Amzn-SageMaker-Inference-Component": inference_component_name}
other_headers: Final = {
name: value
for name, value in OBJECT_KEYED_MAPPING.validate_python(headers).items()
if not isinstance(name, str) or name.lower() != INFERENCE_COMPONENT_HEADER.lower()
}
return {**other_headers, INFERENCE_COMPONENT_HEADER: inference_component_name}
def transform_request(
self,
@ -76,9 +81,8 @@ class SagemakerChatConfig(OpenAIGPTConfig, BaseAWSLLM):
litellm_params=litellm_params,
headers=headers,
)
container_request: Final = {
key: value for key, value in request.items() if key not in self.aws_authentication_params
}
header_only_params: Final = frozenset((*self.aws_authentication_params, "model_id"))
container_request: Final = {key: value for key, value in request.items() if key not in header_only_params}
served_model_name: Final = litellm_params.get("hf_model_name")
if not isinstance(served_model_name, str):
return container_request

View file

@ -14,9 +14,10 @@ from typing_extensions import ReadOnly, Required, TypedDict
from litellm import verbose_logger
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.types.utils import GenericStreamingChunk as GChunk
from litellm.types.utils import StreamingChatCompletionChunk
from litellm.types.utils import LlmProviders, StreamingChatCompletionChunk
STR_KEYED_MAPPING: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
OBJECT_KEYED_MAPPING: Final = TypeAdapter(Mapping[object, object], config=ConfigDict(hide_input_in_errors=True))
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
@ -167,6 +168,36 @@ def get_sagemaker_response_stream_shape():
return _load_sagemaker_response_stream_shape()
INFERENCE_COMPONENT_HEADER: Final = "X-Amzn-SageMaker-Inference-Component"
_SAGEMAKER_PROVIDERS: Final = frozenset(
provider.value for provider in (LlmProviders.SAGEMAKER, LlmProviders.SAGEMAKER_CHAT, LlmProviders.SAGEMAKER_NOVA)
)
def _deployment_provider(litellm_params: Mapping[str, object]) -> str | None:
custom_llm_provider: Final = litellm_params.get("custom_llm_provider")
if isinstance(custom_llm_provider, str) and custom_llm_provider:
return custom_llm_provider
model: Final = litellm_params.get("model")
return model.split("/", 1)[0] if isinstance(model, str) else None
def deployment_inference_component(litellm_params: Mapping[str, object]) -> str | None:
headers: Final = litellm_params.get("extra_headers")
if _deployment_provider(litellm_params) not in _SAGEMAKER_PROVIDERS or not isinstance(headers, Mapping):
return None
return next(
(
value
for name, value in OBJECT_KEYED_MAPPING.validate_python(headers).items()
if isinstance(name, str) and name.lower() == INFERENCE_COMPONENT_HEADER.lower() and isinstance(value, str)
),
None,
)
def with_inference_component_hint(error_message: str) -> str:
if "Inference Component Name header is required" in error_message:
return f"{error_message}\n pass in via `litellm.embedding(..., model_id={{InferenceComponentName}})`"

View file

@ -23,7 +23,7 @@ from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import ModelResponse, Usage
from litellm.utils import token_counter
from ..common_utils import SagemakerError
from ..common_utils import STR_KEYED_MAPPING, SagemakerError
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
@ -153,7 +153,9 @@ class SagemakerConfig(BaseConfig):
litellm_params: dict,
headers: dict,
) -> dict:
inference_params: Final = optional_params.copy()
inference_params: Final = {
key: value for key, value in STR_KEYED_MAPPING.validate_python(optional_params).items() if key != "model_id"
}
stream: Final = inference_params.pop("stream", False)
data: Final[dict] = {"parameters": inference_params}
if stream is True:

View file

@ -39,8 +39,8 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig):
(they require async I/O) and injected at construction time.
"""
def __init__(self, access_token: str, project: str, location: str) -> None:
super().__init__()
def __init__(self, access_token: str, project: str, location: str, drop_params: bool = False) -> None:
super().__init__(drop_params=drop_params)
self._access_token = access_token
self._project = project
self._location = location
@ -51,6 +51,9 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig):
def _rejects_vad_tuning_with_detection_off(self) -> bool:
return False
def _rejects_undocumented_voices(self, model: str) -> bool:
return True
# ------------------------------------------------------------------
# URL
# ------------------------------------------------------------------
@ -163,7 +166,7 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig):
# settings would be silently dropped because ``map_openai_params`` only
# recognises the flat OpenAI-beta key names.
session_params = self._normalize_session_payload_for_mapping(session_params)
setup_config: Final = self.map_openai_params(optional_params={}, non_default_params=session_params)
setup_config: Final = self.map_openai_params(optional_params={}, non_default_params=session_params, model=model)
# Use full Vertex AI model path
setup_config["model"] = self._vertex_model_path(model)
@ -177,15 +180,16 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig):
# ``silence_duration_ms``). ``map_automatic_turn_detection`` sets
# ``disabled=True`` whenever ``create_response`` is absent or
# ``False``. Force ``disabled=False`` only when the client did
# not explicitly request ``create_response: False`` — that path
# not explicitly request ``create_response: False`` or
# ``turn_detection: null`` — the first of those paths
# is how transcription guardrails suppress automatic responses,
# and overriding it here would silently bypass the guardrail.
# Vertex Live has no "VAD on, no auto-response" mode, so callers
# that need that behaviour must accept that VAD is off.
client_turn_detection: Final = session_params.get("turn_detection")
client_disabled_auto_response: Final = (
isinstance(client_turn_detection, dict) and client_turn_detection.get("create_response") is False
)
"turn_detection" in session_params and client_turn_detection is None
) or (isinstance(client_turn_detection, dict) and client_turn_detection.get("create_response") is False)
realtime_input_config: Final = setup_config.setdefault("realtimeInputConfig", {})
automatic_detection: Final = realtime_input_config.setdefault("automaticActivityDetection", {})
if not client_disabled_auto_response:
@ -261,4 +265,5 @@ def vertex_realtime_config(
access_token=access_token,
project=project,
location=VertexBase.get_vertex_region(vertex_region=location, model=model),
drop_params=drop_params,
)

View file

@ -6199,6 +6199,7 @@ def completion(
PROVIDER_AFFINITY_HEADER_KWARG_KEY,
"fireworks_forward_user_id",
"responses_api_version",
"allowed_openai_params",
)
if key in kwargs
},

View file

@ -32,7 +32,7 @@ from litellm.router_strategy.auto_router.auto_router_model_naming import (
classify_strategy_router_model,
strategy_router_dependencies,
)
from litellm.router_utils.common_utils import is_model_capability_field
from litellm.router_utils.common_utils import deployment_model_id_pin, is_model_capability_field
from litellm.types.utils import secret_bearing_wif_litellm_params, server_owned_wif_litellm_params
# Provider routing and server-owned federation or OAuth token-exchange fields. Allowed for proxy
@ -855,6 +855,8 @@ def update_litellm_params_for_health_check(
`deployment_model` (default: the incoming `litellm_params` model), so a `health_check_model` or a `model`
in `health_check_params` is judged on its own. Keys already set through `health_check_params` are kept,
and `modalities` set there without "audio" turns the audio fields off
- sends a SageMaker deployment's `X-Amzn-SageMaker-Inference-Component` extra header as `model_id` when the
deployment sets none, so legacy `sagemaker/` probes reach that component the way routed requests do
"""
incoming_model: Final = _PROBE_PARAMS_ADAPTER.validate_python(litellm_params).get("model")
base_model_owner: Final = deployment_model or (incoming_model if isinstance(incoming_model, str) else None)
@ -935,9 +937,11 @@ def update_litellm_params_for_health_check(
audio_output_params: Final = _health_check_audio_output_params(
_MODEL_INFO_ADAPTER.validate_python(model_info), probe_params, mode, base_model_owner
)
model_id_pin: Final = deployment_model_id_pin(probe_params)
model_id_params: Final = {} if model_id_pin is None else {"model_id": model_id_pin}
return {
key: value
for key, value in {**audio_output_params, **probe_params}.items()
for key, value in {**audio_output_params, **model_id_params, **probe_params}.items()
if not is_model_capability_field(key)
}

View file

@ -13,9 +13,9 @@ litellm_settings:
mcp_code_mode: true
```
Reconnect the MCP client or refresh its tool list. It should discover `litellm_execute_code` alongside its existing tools. The setting is off by default. Turning it off also rejects calls from clients that cached the code tool.
Reconnect the MCP client or refresh its tool list. It should discover `litellm_execute_code` alongside its existing tools. Turning it off also rejects calls from clients that cached the code tool.
The extra installs QuickJS-NG. Platforms without a compatible wheel need a source build and its compiler toolchain; these builds are not part of the verified platform matrix. A missing runtime produces an actionable startup error. The original `quickjs` package shares files with `quickjs-ng`; use a separate environment without that package before installing the extra. Regular proxy installations and disabled gateways do not require either runtime.
The extra installs QuickJS-NG. Platforms without a compatible wheel need a source build and its compiler toolchain; source builds on these platforms have not been verified. A missing runtime produces an actionable startup error. The original `quickjs` package shares files with `quickjs-ng`; use a separate environment without that package before installing the extra. Regular proxy installations and disabled gateways do not require either runtime.
The code tool accepts a `code` string. Discover a tool and its argument schema normally, then copy the code-mode `tool_id` from its description or `search_tools` into `tools.call_tool`. Copy that reference unchanged: clients can rename the displayed tool. Exact gateway catalog names also work. Calls return the full MCP result, including `content`, `structuredContent` when provided, and `isError`. Use `text(value)` to return a JSON value to the model. For example, after discovering a tool named `example-add` that accepts `a` and `b`:
@ -33,7 +33,9 @@ if (result.isError) {
Use `await` for every call, including calls grouped with `Promise.all`. Scripts run in fresh local worker processes. They have no persistent state, general network, filesystem, package imports, or shell API. The parent gateway performs tool calls using the original caller's permissions, upstream authentication, guardrails, and accounting. A worker process isolates normal runtime failures; it is not an operating-system sandbox against an exploit of the embedded JavaScript engine.
The script itself and each nested tool call count toward request-rate limits. The script holds its caller's concurrency slot until it finishes; nested calls reuse overlapping concurrency slots and retain their other limits. The wrapper has no tool charge, and each upstream call retains its configured cost.
REST clients discover the gateway code tool through `/mcp-rest/tools/list` without a `server_id` filter, and call it through `/mcp-rest/tools/call` without `server_id`. An explicit `server_id` keeps `litellm_execute_code` routed to that upstream. Per-server REST listings contain upstream tools only. Caller permissions and an optional `toolset_name` still restrict gateway code calls.
The script itself and each nested tool call count toward request-rate limits. Server request limits also count tool lookups, as they do on `/mcp/proxy`. The script holds its caller's concurrency slot until it finishes; nested calls reuse overlapping concurrency slots and retain their other limits. The wrapper has no tool charge, and each upstream call retains its configured cost.
## Tool compatibility

View file

@ -12,7 +12,7 @@ from uuid import uuid4
import anyio
from mcp.shared.exceptions import MCPError
from mcp.types import CallToolResult, TextContent, Tool
from pydantic import ConfigDict, JsonValue, TypeAdapter
from pydantic import ConfigDict, JsonValue, TypeAdapter, ValidationError
import litellm
from litellm._logging import verbose_proxy_logger
@ -273,7 +273,10 @@ async def _execute_script(arguments: Mapping[str, object], context: OperationCon
return error_result("MCP code execution requires an authenticated caller")
from litellm.proxy.mcp_server.code_runtime import execute_code
request: Final = CodeInput.model_validate(arguments or {})
try:
request: Final = CodeInput.model_validate(arguments or {})
except ValidationError:
return error_result("Invalid code request. Provide only a 'code' string.")
auth, header, servers, server_headers, oauth_headers, raw_headers, client_ip = context.legacy_auth()
forbidden: Final = {*MCP_PROXY_TOOL_NAMES, *VIRTUAL_TOOL_NAMES, CODE_TOOL_NAME}
@ -335,6 +338,7 @@ async def _execute_script(arguments: Mapping[str, object], context: OperationCon
raw_headers=raw_headers,
mcp_proxy_mode=True,
incoming_bearer_token=context.incoming_bearer_token,
resolved_tool=selected,
)
except asyncio.CancelledError:
await cleanup_cancelled_call(
@ -344,7 +348,7 @@ async def _execute_script(arguments: Mapping[str, object], context: OperationCon
assert result is not None
from litellm.proxy.mcp_server.result_conversion import INPUT_REQUIRED_UNSUPPORTED_MESSAGE
if any(
if result.is_error and any(
isinstance(block, TextContent) and block.text == INPUT_REQUIRED_UNSUPPORTED_MESSAGE
for block in result.content
):

View file

@ -48,18 +48,20 @@ class _Calls:
def check_runtime() -> None:
try:
metadata.distribution("quickjs")
except metadata.PackageNotFoundError:
pass
else:
raise ValueError(
"MCP code mode cannot use the conflicting quickjs package. Install the code-mode extra in a separate environment."
)
try:
metadata.distribution("quickjs-ng")
except metadata.PackageNotFoundError as exc:
raise ValueError(
"MCP code mode requires quickjs-ng. Install 'litellm[proxy,mcp-code-mode]' on a supported platform."
) from exc
try:
metadata.distribution("quickjs")
except metadata.PackageNotFoundError:
pass
else:
raise ValueError("MCP code mode requires removing the conflicting quickjs package and reinstalling quickjs-ng")
if importlib.util.find_spec("_quickjs") is None:
raise ValueError(
"MCP code mode requires quickjs-ng. Install 'litellm[proxy,mcp-code-mode]' on a supported platform."

View file

@ -433,11 +433,13 @@ async def _dispatch_virtual_mcp_tool(
raw_headers: dict[str, str] | None = None,
mcp_proxy_mode: bool = False,
incoming_bearer_token: SecretStr | None | EllipsisType = ...,
resolved_tool: MCPTool | None = None,
) -> CallToolResult | None:
"""Handle the mcp_tool_search / mcp_tool_call virtual tools.
Returns a CallToolResult when ``name`` is a virtual tool, else ``None`` so
the caller falls through to normal tool routing.
``resolved_tool`` reuses this request's authenticated lookup for nested code calls.
"""
from litellm.llms.litellm_proxy.skills.skill_search import DEFAULT_SKILL_SEARCH_TOP_K
from litellm.proxy.mcp_server.code_execution import CODE_TOOL_NAME, run_code_tool
@ -488,18 +490,22 @@ async def _dispatch_virtual_mcp_tool(
proxy_tool_id: Final = proxy_arguments.get("tool_id")
try:
proxy_target_tool: Final = (
await find_mcp_proxy_tool(
tool_id=proxy_tool_id,
user_api_key_auth=user_api_key_auth,
mcp_servers=mcp_servers,
client_ip=client_ip,
mcp_auth_header=mcp_auth_header,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
resolved_tool
if resolved_tool is not None
else (
await find_mcp_proxy_tool(
tool_id=proxy_tool_id,
user_api_key_auth=user_api_key_auth,
mcp_servers=mcp_servers,
client_ip=client_ip,
mcp_auth_header=mcp_auth_header,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
)
if name == MCP_PROXY_CALL_TOOL_NAME and isinstance(proxy_tool_id, str)
else None
)
if name == MCP_PROXY_CALL_TOOL_NAME and isinstance(proxy_tool_id, str)
else None
)
proxy_guardrail_auth: Final = (
await MCPRequestHandler.billing_auth_for_tool_call(

View file

@ -992,7 +992,7 @@ if MCP_AVAILABLE:
"message": f"Failed to get tools from server {server.name}: {e}",
}
return {
"tools": with_code_tool(list_tools_result),
"tools": with_code_tool(list_tools_result, first_page=False),
"error": None,
"message": "Successfully retrieved tools",
}
@ -1357,13 +1357,11 @@ if MCP_AVAILABLE:
toolset: Final = OPTIONAL_STRING_ADAPTER.validate_python(data.get("toolset_name"))
auth: Final = await _resolve_acting_auth(toolset, user_api_key_dict)
header, server_headers, raw_headers, oauth = _extract_mcp_headers_from_request(request, auth)
server: Final = OPTIONAL_STRING_ADAPTER.validate_python(data.get("server_id"))
return await run_code_tool(
arguments,
OperationContext(
_caller=auth,
mcp_auth_header=header,
mcp_servers=(server,) if server is not None else None,
mcp_server_auth_headers=server_headers,
raw_headers=raw_headers,
oauth2_headers=oauth,
@ -1405,7 +1403,9 @@ if MCP_AVAILABLE:
from litellm.proxy.mcp_server.code_execution import settings as code_mode_settings
from litellm.proxy.mcp_server.tool_search import VIRTUAL_TOOL_NAMES
if tool_name in VIRTUAL_TOOL_NAMES or (tool_name == CODE_TOOL_NAME and code_mode_settings().enabled):
if tool_name in VIRTUAL_TOOL_NAMES or (
tool_name == CODE_TOOL_NAME and code_mode_settings().enabled and data.get("server_id") is None
):
return await _handle_virtual_mcp_tool(request, data, tool_name, user_api_key_dict)
# Validate required parameters early

View file

@ -255,6 +255,7 @@ from litellm.router_utils.clientside_credential_handler import (
from litellm.router_utils.common_utils import (
MID_STREAM_CONTINUATION_KWARG,
as_retry_skipped_deployment_ids,
deployment_model_id_pin,
deployments_the_caller_may_use,
filter_team_based_models,
filter_web_search_deployments,
@ -5316,6 +5317,9 @@ class Router(DeprecatedRouterMembers):
):
kwargs[STREAM_USAGE_INJECTED_FIELD] = True
self._update_kwargs_with_default_litellm_params(kwargs=kwargs, metadata_variable_name=metadata_variable_name)
pinned_model_id: Final = deployment_model_id_pin(deployment_context)
if pinned_model_id is not None:
kwargs["model_id"] = pinned_model_id # rebind-ok: this method's contract is filling the caller's kwargs
hop_kwargs: Final = _FALLBACK_HOP_ADAPTER.validate_python(kwargs)
if _is_fallback_hop(hop_kwargs):
EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt(
@ -6242,6 +6246,9 @@ class Router(DeprecatedRouterMembers):
kwargs[k] = v
elif k == "metadata":
kwargs[k].update(v)
pinned_model_id: Final = deployment_model_id_pin(_LITELLM_PARAMS_ADAPTER.validate_python(data))
if pinned_model_id is not None:
kwargs["model_id"] = pinned_model_id # rebind-ok: this call's own kwargs, filled like the defaults
# call via litellm.completion()
return litellm.text_completion(

View file

@ -16,6 +16,7 @@ from litellm.exceptions import BadRequestError
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_structure
from litellm.llms.sagemaker.common_utils import deployment_inference_component
from litellm.types.router import (
CredentialLiteLLMParams,
Deployment,
@ -522,6 +523,16 @@ PROVIDER_SPECIFIC_CAPABILITY_FIELDS: Final = frozenset(ProviderSpecificModelInfo
_LITELLM_PARAMS_ADAPTER: Final = TypeAdapter(Mapping[str, object])
def deployment_model_id_pin(deployment_litellm_params: Mapping[str, object]) -> object | None:
clientside_params: Final = deployment_litellm_params.get("configurable_clientside_auth_params")
if isinstance(clientside_params, list | tuple) and "model_id" in clientside_params:
return None
own_model_id: Final = deployment_litellm_params.get("model_id")
if own_model_id is not None:
return own_model_id
return deployment_inference_component(deployment_litellm_params)
def is_model_capability_field(name: str) -> bool:
return name in PROVIDER_SPECIFIC_CAPABILITY_FIELDS or (
name.startswith("supports_") and name in GetModelCostMap.loaded_model_cost_map_fields()

View file

@ -81,6 +81,14 @@ class BidiGenerateContentServerMessage(TypedDict, total=False):
"""Output only. The setup complete message."""
class BidiGenerateContentActivityStart(TypedDict):
pass
class BidiGenerateContentActivityEnd(TypedDict):
pass
class BidiGenerateContentRealtimeInput(TypedDict, total=False):
text: str
"""The text to be sent to the model."""
@ -94,11 +102,11 @@ class BidiGenerateContentRealtimeInput(TypedDict, total=False):
audioStreamEnd: bool
"""Output only. If true, indicates that the audio stream has ended."""
activityStart: bool
"""Output only. If true, indicates that the activity has started."""
activityStart: BidiGenerateContentActivityStart
"""Marks the start of user activity when automatic activity detection is disabled."""
activityEnd: bool
"""Output only. If true, indicates that the activity has ended."""
activityEnd: BidiGenerateContentActivityEnd
"""Marks the end of user activity when automatic activity detection is disabled."""
StartOfSpeechSensitivityEnum = Literal[

View file

@ -23,3 +23,22 @@ class MistralThinkingBlock(TypedDict):
type: Literal["thinking"]
thinking: list[MistralTextBlock]
closed: NotRequired[ReadOnly[bool]]
class MistralRequestRouting(TypedDict, total=False):
api_base: ReadOnly[object]
drop_params: ReadOnly[object]
allowed_openai_params: ReadOnly[object]
custom_llm_provider: ReadOnly[object]
class MistralCacheKeyParams(TypedDict, total=False):
prompt_cache_key: ReadOnly[object]
class MistralReasoningEffortParams(TypedDict, total=False):
reasoning_effort: ReadOnly[object]
class MistralFilePart(TypedDict, total=False):
file_id: ReadOnly[object]

View file

@ -10911,7 +10911,7 @@ class ProviderConfigManager:
if LlmProviders.GEMINI == provider:
from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig
return GeminiRealtimeConfig()
return GeminiRealtimeConfig(drop_params=drop_params)
if LlmProviders.META == provider:
from litellm.llms.meta.realtime.transformation import MetaRealtimeConfig

View file

@ -0,0 +1,21 @@
from typing import Final
import pytest
from litellm.litellm_core_utils.realtime_streaming import _merge_session_fields
NESTED_LEVELS: Final = 8
LEAVES: Final = 2000
def _nested_session(depth: int) -> dict[str, object]:
return {"child": _nested_session(depth - 1)} if depth else {f"k{index}": index for index in range(LEAVES)}
@pytest.mark.benchmark
def test_merging_a_session_update_that_repeats_every_nested_level() -> None:
session: Final = _nested_session(NESTED_LEVELS)
merged: Final = _merge_session_fields(session, session)
assert merged == session

View file

@ -85,6 +85,7 @@ IGNORE_FUNCTIONS = [
"_sort_processed_sets", # bounded by the nesting depth of the log-record extra it walks (a finite JSON tree, no cycles possible).
"scrub_json_strings", # max depth set (MAX_SCRUB_DEPTH); fails closed by returning "[Filtered]" for anything nested past the cap.
"_without_requester_ip_address", # walks only the fresh tree _sanitize_request_body_for_spend_logs_payload returns (mapping cycles already cut), so it never goes deeper than that sanitizer did.
"_merge_session_fields", # max depth set (_SESSION_MERGE_MAX_DEPTH); deeper mappings are replaced wholesale by the later session.update.
]

View file

@ -47,6 +47,7 @@ class RealtimeProvider:
id: str
alias: str
litellm_params: LiteLLMParamsBody
modalities: tuple[str, ...] = ("text",)
PROVIDERS = (
@ -75,6 +76,7 @@ PROVIDERS = (
model="gemini/gemini-3.1-flash-live-preview",
api_key="os.environ/GEMINI_API_KEY",
),
modalities=("audio",),
),
RealtimeProvider(
"vertex_ai",
@ -84,6 +86,7 @@ PROVIDERS = (
vertex_location="us-central1",
vertex_credentials="os.environ/VERTEXAI_CREDENTIALS",
),
modalities=("audio",),
),
RealtimeProvider(
"bedrock",

View file

@ -124,7 +124,8 @@ def test_text_conversation(
session.send(
SessionUpdate(
session=SessionConfig(
instructions="You are a terse assistant. Reply in one short sentence."
modalities=list(provider.modalities),
instructions="You are a terse assistant. Reply in one short sentence.",
)
)
)
@ -159,6 +160,7 @@ def test_tool_call_round_trip(
session.send(
SessionUpdate(
session=SessionConfig(
modalities=list(provider.modalities),
tools=[WEATHER_TOOL],
tool_choice="auto",
instructions=(

View file

@ -1,3 +1,4 @@
import json
import uuid
from collections.abc import Iterator
from pathlib import Path
@ -6,8 +7,19 @@ from typing import Final
import pytest
import yaml
from integration._support.client import JSON_OBJECT, Gateway, gateway_from_environment, object_value
from integration._support.mcp import McpCaller, mcp_peer, official_client_outcomes, register_mcp, tool_calls
from integration._support.mcp import (
McpCaller,
echo_tool,
mcp_peer,
official_client_outcomes,
register_mcp,
scripted_peer,
tool_calls,
)
from integration._support.process import owned_proxy
from pydantic import JsonValue, TypeAdapter
SEARCH_RESULTS: Final = TypeAdapter(tuple[dict[str, JsonValue], ...])
@pytest.fixture(scope="module")
@ -82,3 +94,89 @@ def test_code_timeout_returns_error_and_releases_caller_slot(code_gateway: Gatew
recovery: Final = McpCaller(code_gateway, key, "mcp").call("litellm_execute_code", {"code": "text(42)"})
assert recovery.ok and recovery.text is not None, recovery
assert JSON_OBJECT.validate_json(recovery.text)["output"] == [42]
def test_rest_server_id_preserves_upstream_reserved_name(code_gateway: Gateway) -> None:
with scripted_peer(echo_tool("litellm_execute_code")) as peer, code_gateway.scenario() as scenario:
identity: Final = register_mcp(scenario, peer, "reserved" + uuid.uuid4().hex[:8])
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
arguments: Final = {"code": "qa_upstream_marker"}
peer.drain()
upstream: Final = McpCaller(code_gateway, key, "rest").call("litellm_execute_code", arguments, identity)
assert upstream.ok and upstream.text is not None, upstream
assert JSON_OBJECT.validate_json(upstream.text) == arguments
observed: Final = tool_calls(peer.drain())
assert len(observed) == 1
request: Final = JSON_OBJECT.validate_python(observed[0]["body"])
assert object_value(request["params"])["arguments"] == arguments
gateway_code: Final = McpCaller(code_gateway, key, "rest").call("litellm_execute_code", {"code": "text(42)"})
assert gateway_code.ok and gateway_code.text is not None, gateway_code
assert JSON_OBJECT.validate_json(gateway_code.text)["output"] == [42]
assert tool_calls(peer.drain()) == ()
def test_invalid_rest_code_returns_tool_error_and_releases_slot(code_gateway: Gateway) -> None:
with mcp_peer() as peer, code_gateway.scenario() as scenario:
identity: Final = register_mcp(scenario, peer, "invalid" + uuid.uuid4().hex[:8])
key: Final = scenario.key(object_permission={"mcp_servers": [identity]}, max_parallel_requests=1)
caller: Final = McpCaller(code_gateway, key, "rest")
invalid: Final = caller.call("litellm_execute_code", {"code": 42})
assert invalid.status == 200 and not invalid.ok, invalid
valid: Final = caller.call("litellm_execute_code", {"code": "text(42)"})
assert valid.ok and valid.text is not None, valid
assert JSON_OBJECT.validate_json(valid.text)["output"] == [42]
assert tool_calls(peer.drain()) == ()
def _remaining_server_calls(gateway: Gateway, key: str, identity: str) -> int:
responses: Final = tuple(
gateway.request(
"POST",
"/mcp-rest/tools/call",
{"server_id": identity, "name": "add", "arguments": {"a": 20, "b": 22}},
key=key,
)
for _ in range(7)
)
assert all(response.status_code in (200, 429) for response in responses)
return sum(response.status_code == 200 for response in responses)
def _server_budget_after_call(gateway: Gateway, *, code: bool) -> tuple[int, int]:
with mcp_peer() as target, mcp_peer() as other, gateway.scenario() as scenario:
alias: Final = "target" + uuid.uuid4().hex[:8]
target_id: Final = register_mcp(scenario, target, alias, rpm=6)
other_id: Final = register_mcp(scenario, other, "other" + uuid.uuid4().hex[:8], rpm=6)
key: Final = scenario.key(object_permission={"mcp_servers": [target_id, other_id]})
_, search = official_client_outcomes(gateway, key, "/mcp/proxy", "search_tools", {"query": "add", "top_k": 10})
assert search.ok and search.text is not None, search
matches: Final = SEARCH_RESULTS.validate_json(search.text)
identity: Final = next(hit["tool_id"] for hit in matches if hit["name"] == f"{alias}-add")
assert isinstance(identity, str)
target.drain()
other.drain()
arguments: Final = (
{"code": f"text(await tools.call_tool({{tool_id:{json.dumps(identity)},arguments:{{a:20,b:22}}}}))"}
if code
else {"tool_id": identity, "arguments": {"a": 20, "b": 22}}
)
_, result = official_client_outcomes(
gateway, key, "/mcp/proxy", "litellm_execute_code" if code else "call_tool", arguments
)
assert result.ok and result.text is not None, result
if code:
output: Final = JSON_OBJECT.validate_json(result.text)["output"]
assert isinstance(output, list) and len(output) == 1
assert object_value(output[0]).get("isError", False) is False
else:
assert result.text == "42"
assert len(tool_calls(target.drain())) == 1
assert tool_calls(other.drain()) == ()
return _remaining_server_calls(gateway, key, target_id), _remaining_server_calls(gateway, key, other_id)
def test_code_preserves_direct_server_rpm_for_selected_and_unrelated_servers(code_gateway: Gateway) -> None:
direct: Final = _server_budget_after_call(code_gateway, code=False)
nested: Final = _server_budget_after_call(code_gateway, code=True)
assert all(remaining > 0 for remaining in direct)
assert nested == direct

View file

@ -355,7 +355,13 @@ def _voice_scenario_created_after_the_first_frame(transcript: str) -> RealtimeRe
def _live_scenario() -> RealtimeResponse:
return RealtimeResponse(content_type="application/x-realtime", events=())
return RealtimeResponse(
content_type="application/x-realtime",
events=(
{"serverContent": {"modelTurn": {"parts": [{"text": "scripted $REQUEST_ID"}]}}},
{"serverContent": {"turnComplete": True}},
),
)
def _realtime_update(turn_detection: JsonValue) -> dict[str, JsonValue]:
@ -1348,10 +1354,14 @@ def test_gemini_ai_studio_voice_session_with_vad_tuning_connects_under_the_guard
_ws_base(_owned_url(guardrail_proxy)),
f"model={model}",
key,
(_step(_update_frame(update), until="session.updated"),),
(
_step(_update_frame(update), until="session.updated"),
_step(_user_text_item(CLEAN_TRANSCRIPT), RESPONSE_CREATE, until=RESPONSE_DONE),
),
)
observed: Final = _observed(upstream_url, handle.scenario_id)
assert session.types == ("session.created", "session.updated"), session
assert session.types[:2] == ("session.created", "session.updated"), session
assert session.types[-1] == RESPONSE_DONE, session
assert _setup_detection(observed) == {"disabled": True}, _sent(observed)
@ -1365,10 +1375,14 @@ def test_vertex_live_voice_session_keeps_vad_tuning_under_the_guardrail(guardrai
_ws_base(_owned_url(guardrail_proxy)),
f"model={model}",
key,
(_step(_update_frame(_realtime_update(TUNED_SERVER_VAD)), until="session.updated"),),
(
_step(_update_frame(_realtime_update(TUNED_SERVER_VAD)), until="session.updated"),
_step(_user_text_item(CLEAN_TRANSCRIPT), RESPONSE_CREATE, until=RESPONSE_DONE),
),
)
observed: Final = _observed(upstream_url, handle.scenario_id)
assert session.types == ("session.created", "session.updated"), session
assert session.types[:2] == ("session.created", "session.updated"), session
assert session.types[-1] == RESPONSE_DONE, session
assert _setup_detection(observed) == GEMINI_TUNED_DETECTION, _sent(observed)

View file

@ -243,7 +243,7 @@ def test_sagemaker_chat_signs_the_inference_component_header_and_sends_hf_model_
)
assert "x-amzn-sagemaker-inference-component" in signed_headers.removeprefix("SignedHeaders=").split(";")
assert JSON_OBJECT.validate_json(request.body) == {
**container_body(wire.url, PROMPT, stream=False, model_id=_INFERENCE_COMPONENT),
**container_body(wire.url, PROMPT, stream=False),
"model": SERVED_MODEL,
}, request.body

View file

@ -668,12 +668,12 @@ def test_an_extra_body_that_is_not_a_populated_object_adds_nothing_to_the_body(
assert _body(request) == {"input": [text]}, request.body
def test_a_caller_sent_model_id_overrides_the_deployments_inference_component(proxy: Gateway, wire: Wire) -> None:
def test_a_caller_sent_model_id_cannot_override_the_deployments_inference_component(proxy: Gateway, wire: Wire) -> None:
text: Final = _text()
response: Final = _embed(proxy, "sm-openai", [text], model_id="caller-chosen-component")
assert response.status_code == 200, response.text
request: Final = _only_request(wire)
assert request.headers[COMPONENT_HEADER] == "caller-chosen-component", dict(request.headers)
assert request.headers[COMPONENT_HEADER] == COMPONENT, dict(request.headers)
assert _body(request) == {"input": [text], "model": SERVED_MODEL}, request.body

View file

@ -1246,6 +1246,80 @@ def test_anthropic_messages_pt_keeps_signed_thinking_block():
assert thinking_blocks[0]["thinking"] == "genuine anthropic reasoning"
@pytest.mark.parametrize("llm_provider", ("anthropic", "vertex_ai"))
def test_anthropic_messages_pt_converts_pdf_image_url_data_uri_to_document(llm_provider: str) -> None:
pdf_data = "JVBERi0xLjQK"
result: Final = anthropic_messages_pt(
messages=[
{
"role": "user",
"content": [
{"type": "text", "text": "summarize"},
{"type": "image_url", "image_url": {"url": f"data:application/pdf;base64,{pdf_data}"}},
],
}
],
model="claude-opus-5-5",
llm_provider=llm_provider,
)
assert result[0]["content"][1] == {
"type": "document",
"source": {"type": "base64", "media_type": "application/pdf", "data": pdf_data},
}
@pytest.mark.parametrize("llm_provider", ("anthropic", "vertex_ai"))
def test_anthropic_messages_pt_uses_pdf_format_for_generic_mime_data_uri(llm_provider: str) -> None:
pdf_data = "JVBERi0xLjQK"
result: Final = anthropic_messages_pt(
messages=[
{
"role": "user",
"content": [
{"type": "text", "text": "summarize"},
{
"type": "image_url",
"image_url": {
"url": f"data:application/octet-stream;base64,{pdf_data}",
"format": "application/pdf",
},
},
],
}
],
model="claude-opus-5-5",
llm_provider=llm_provider,
)
assert result[0]["content"][1] == {
"type": "document",
"source": {"type": "base64", "media_type": "application/pdf", "data": pdf_data},
}
def test_anthropic_messages_pt_keeps_png_image_url_as_image() -> None:
image_data = "iVBORw0KGgo="
result: Final = anthropic_messages_pt(
messages=[
{
"role": "user",
"content": [
{"type": "text", "text": "describe"},
{"type": "image_url", "image_url": f"data:image/png;base64,{image_data}"},
],
}
],
model="claude-opus-5-5",
llm_provider="anthropic",
)
assert result[0]["content"][1] == {
"type": "image",
"source": {"type": "base64", "media_type": "image/png", "data": image_data},
}
EMPTY_USER_MESSAGES: Final = pytest.mark.parametrize(
"empty_user_message",
[{"role": "user", "content": None}, {"role": "user", "content": []}, {"role": "user"}],

View file

@ -288,6 +288,37 @@ def test_convert_url_to_base64_keeps_a_real_content_type(monkeypatch):
assert result.startswith("data:image/jpeg;base64,")
@pytest.mark.parametrize(
("url", "content_type", "expected"),
[
(
"https://docs.example/document.pdf",
"application/pdf; qs=0.001",
"data:application/pdf;base64,cGRm",
),
(
"https://img.example/image.png",
"image/png; charset=binary",
"data:image/png;base64,cGRm",
),
(
"https://docs.example/document.pdf",
"application/octet-stream; charset=binary",
"data:application/pdf;base64,cGRm",
),
],
)
def test_convert_url_to_base64_uses_bare_content_type(
monkeypatch: pytest.MonkeyPatch,
url: str,
content_type: str,
expected: str,
) -> None:
monkeypatch.setattr(litellm, "module_level_client", _stub_image_client(b"pdf", content_type))
assert convert_url_to_base64(url) == expected
def test_convert_url_to_base64_raises_when_no_content_type_is_determinable(monkeypatch):
monkeypatch.setattr(
litellm,

View file

@ -1,7 +1,7 @@
import asyncio
import base64
import json
from collections.abc import Coroutine, Mapping, Sequence
from collections.abc import Callable, Coroutine, Mapping, Sequence
from dataclasses import dataclass
from typing import Final, Literal
from unittest.mock import AsyncMock, MagicMock, patch
@ -14,6 +14,7 @@ from websockets.frames import Close
import litellm
from litellm.constants import (
BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY,
REALTIME_LAZY_SETUP_DEADLINE_SECONDS,
REALTIME_METERED_USAGE_KEY,
REALTIME_SESSION_FAILURE_LOGGED_KEY,
REALTIME_SESSION_SUCCESS_LOGGED_KEY,
@ -21,6 +22,7 @@ from litellm.constants import (
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.realtime_streaming import (
RealTimeStreaming,
_merge_session_fields,
client_sent_openai_beta_realtime_header,
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@ -1255,11 +1257,15 @@ async def test_client_ack_caches_setup_to_prevent_duplicate_session_update_setup
side_effect=[
json.dumps({"type": "session.update", "session": {"tools": []}}),
json.dumps({"type": "session.update", "session": {"tools": []}}),
_TEXT_ITEM_CREATE,
_RESPONSE_CREATE,
Exception("client done"),
]
)
websocket.send_text = AsyncMock()
provider_config = MagicMock()
provider_config.setup_session_update_rejection = MagicMock(return_value=None)
def _transform(message: str, model: str, session_configuration_request=None):
if session_configuration_request is None:
@ -2285,6 +2291,7 @@ async def test_guardrail_turn_detection_injected_into_first_session_update_defer
logging_obj.success_handler = MagicMock()
provider_config = MagicMock()
provider_config.builds_setup_from_session_update = MagicMock(return_value=False)
transformed_messages = []
def mock_transform(msg, model, session_config):
@ -2351,6 +2358,7 @@ async def test_guardrail_turn_detection_injection_tolerates_non_dict_value(
logging_obj.success_handler = MagicMock()
provider_config = MagicMock()
provider_config.builds_setup_from_session_update = MagicMock(return_value=False)
transformed_messages = []
def mock_transform(msg, model, session_config):
@ -2413,6 +2421,7 @@ async def test_subsequent_session_update_cannot_reenable_vad_when_guardrails_act
logging_obj.success_handler = MagicMock()
provider_config = MagicMock()
provider_config.setup_session_update_rejection = MagicMock(return_value=None)
transformed_messages = []
def mock_transform(msg, model, session_config):
@ -2465,6 +2474,7 @@ async def test_follow_up_setup_updates_cached_session_configuration_request():
logging_obj.success_handler = MagicMock()
provider_config = MagicMock()
provider_config.setup_session_update_rejection = MagicMock(return_value=None)
follow_up_setup = json.dumps(
{
"setup": {
@ -2498,6 +2508,636 @@ async def test_follow_up_setup_updates_cached_session_configuration_request():
assert streaming.session_configuration_request == follow_up_setup
def _gemini_default_mode_streaming(
monkeypatch: pytest.MonkeyPatch,
client_messages: Sequence[str],
model: str = "gemini-live-2.5-flash-native-audio",
drop_params: bool = False,
defer_setup: bool = False,
) -> tuple[RealTimeStreaming, MagicMock]:
monkeypatch.setattr(litellm, "gemini_live_defer_setup", defer_setup, raising=False)
from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig
client_ws = MagicMock()
client_ws.receive_text = AsyncMock(side_effect=[*client_messages, ConnectionClosed(None, None)])
client_ws.send_text = AsyncMock()
backend_ws = MagicMock()
backend_ws.send = AsyncMock()
streaming = RealTimeStreaming(
client_ws,
backend_ws,
MagicMock(),
provider_config=GeminiRealtimeConfig(drop_params=drop_params),
model=model,
)
return streaming, backend_ws
def _events_sent_to_client(streaming: RealTimeStreaming) -> list[dict]:
return [json.loads(call.args[0]) for call in streaming.websocket.send_text.await_args_list]
_TEXT_ITEM_CREATE = json.dumps(
{
"type": "conversation.item.create",
"item": {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]},
}
)
_RESPONSE_CREATE = json.dumps({"type": "response.create"})
@pytest.mark.asyncio
async def test_default_mode_sends_default_setup_before_first_non_session_client_message(monkeypatch):
"""Gemini Live setup is no longer pre-sent at connect, so a client that never
sends session.update still gets the provider default setup ahead of its
first content frame."""
streaming, backend_ws = _gemini_default_mode_streaming(monkeypatch, [_TEXT_ITEM_CREATE])
await streaming.client_ack_messages()
frames = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list]
assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"]
assert frames[0]["setup"]["generationConfig"]["responseModalities"] == ["AUDIO"]
assert streaming.session_configuration_request == backend_ws.send.await_args_list[0].args[0]
@pytest.mark.asyncio
async def test_default_mode_session_update_first_builds_the_only_setup(monkeypatch):
session_update = json.dumps({"type": "session.update", "session": {"modalities": ["text"]}})
streaming, backend_ws = _gemini_default_mode_streaming(
monkeypatch, [session_update, _TEXT_ITEM_CREATE, _RESPONSE_CREATE], model="gemini-live-2.5-flash"
)
await streaming.client_ack_messages()
frames = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list]
assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"]
assert frames[0]["setup"]["generationConfig"]["responseModalities"] == ["TEXT"]
@pytest.mark.asyncio
async def test_default_mode_merges_every_session_update_sent_before_the_first_content_frame(monkeypatch):
"""Clients configure a session over several session.update messages. Gemini
Live takes one setup, so every update sent before the first content frame is
merged into that setup instead of only the first one applying."""
instructions_update = json.dumps({"type": "session.update", "session": {"instructions": "Be brief."}})
modalities_update = json.dumps({"type": "session.update", "session": {"modalities": ["text"]}})
streaming, backend_ws = _gemini_default_mode_streaming(
monkeypatch, [instructions_update, modalities_update, _TEXT_ITEM_CREATE, _RESPONSE_CREATE], model="gemini-live-2.5-flash"
)
await streaming.client_ack_messages()
frames = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list]
assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"]
assert frames[0]["setup"]["generationConfig"]["responseModalities"] == ["TEXT"]
assert frames[0]["setup"]["systemInstruction"]["parts"][0]["text"] == "Be brief."
@pytest.mark.asyncio
async def test_default_mode_merges_nested_audio_settings_across_session_updates(monkeypatch):
vad_update = json.dumps(
{"type": "session.update", "session": {"audio": {"input": {"turn_detection": {"silence_duration_ms": 1500}}}}}
)
voice_update = json.dumps({"type": "session.update", "session": {"audio": {"output": {"voice": "Kore"}}}})
streaming, backend_ws = _gemini_default_mode_streaming(
monkeypatch, [vad_update, voice_update, _TEXT_ITEM_CREATE, _RESPONSE_CREATE], model="gemini-live-2.5-flash"
)
await streaming.client_ack_messages()
setup = json.loads(backend_ws.send.await_args_list[0].args[0])["setup"]
assert setup["realtimeInputConfig"]["automaticActivityDetection"]["silenceDurationMs"] == 1500
assert setup["generationConfig"]["speechConfig"]["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == "Kore"
def test_merge_session_fields_replaces_nested_mappings_wholesale_past_the_depth_cap():
def nested(levels: int, key: str) -> dict:
return {key: levels, **({"child": nested(levels - 1, key)} if levels else {})}
def level(node, depth: int):
return node if depth == 0 else level(node["child"], depth - 1)
merged = _merge_session_fields(nested(10, "keep"), nested(10, "new"))
assert all({"keep", "new"} <= set(level(merged, depth)) for depth in range(9))
assert "keep" not in level(merged, 9)
assert level(merged, 9)["new"] == 1
@pytest.mark.asyncio
async def test_default_mode_acknowledges_each_held_session_update_before_any_backend_frame(monkeypatch):
"""A client that waits for session.updated before sending content must not
hang while its updates are held, so each one is acknowledged right away
with the modalities the merged setup will carry."""
instructions_update = json.dumps({"type": "session.update", "session": {"instructions": "Be brief."}})
modalities_update = json.dumps({"type": "session.update", "session": {"modalities": ["text"]}})
streaming, backend_ws = _gemini_default_mode_streaming(
monkeypatch, [instructions_update, modalities_update], model="gemini-live-2.5-flash"
)
await streaming.client_ack_messages()
acks = [json.loads(call.args[0]) for call in streaming.websocket.send_text.await_args_list]
assert [ack["type"] for ack in acks] == ["session.updated", "session.updated"]
assert [ack["session"].get("output_modalities") or ack["session"].get("modalities") for ack in acks] == [
["audio"],
["text"],
]
assert acks[1]["session"]["instructions"] == "Be brief."
backend_ws.send.assert_not_awaited()
_LATE_INSTRUCTIONS_UPDATE = json.dumps(
{"type": "session.update", "event_id": "evt_late", "session": {"instructions": "Be brief."}}
)
@pytest.mark.asyncio
async def test_default_mode_answers_a_session_update_after_the_first_content_with_an_error(monkeypatch):
streaming, backend_ws = _gemini_default_mode_streaming(monkeypatch, [_TEXT_ITEM_CREATE, _LATE_INSTRUCTIONS_UPDATE])
await streaming.client_ack_messages()
errors: Final = [event["error"] for event in _events_sent_to_client(streaming) if event["type"] == "error"]
assert [(error["code"], error["event_id"]) for error in errors] == [("session_update_not_applied", "evt_late")]
frames: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list]
assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"]
@pytest.mark.asyncio
async def test_default_mode_drops_a_late_session_update_quietly_under_drop_params(monkeypatch):
streaming, backend_ws = _gemini_default_mode_streaming(
monkeypatch, [_TEXT_ITEM_CREATE, _LATE_INSTRUCTIONS_UPDATE], drop_params=True
)
await streaming.client_ack_messages()
assert _events_sent_to_client(streaming) == []
assert backend_ws.send.await_count == 2
@pytest.mark.asyncio
@pytest.mark.parametrize(
("model", "requested_output"),
[
("gemini-live-2.5-flash-native-audio", {"modalities": ["text"]}),
("gemini-live-2.5-flash", {"output_modalities": ["image"]}),
],
)
async def test_default_mode_rejects_an_output_modality_the_model_cannot_produce(monkeypatch, model, requested_output):
update: Final = json.dumps(
{"type": "session.update", "session": {**requested_output, "instructions": "Be brief."}}
)
streaming, backend_ws = _gemini_default_mode_streaming(monkeypatch, [update, _TEXT_ITEM_CREATE], model=model)
await streaming.client_ack_messages()
events: Final = _events_sent_to_client(streaming)
assert [(event["type"], event["error"]["code"]) for event in events] == [("error", "invalid_value")]
setup: Final = json.loads(backend_ws.send.await_args_list[0].args[0])["setup"]
assert "systemInstruction" not in setup
@pytest.mark.asyncio
@pytest.mark.parametrize(
("modalities", "drop_params"),
[(["text", "audio"], False), (["text"], True)],
ids=["openai-default-pair", "text-dropped-under-drop-params"],
)
async def test_default_mode_applies_the_update_with_audio_output_on_an_audio_only_model(
monkeypatch, modalities, drop_params
):
update: Final = json.dumps(
{"type": "session.update", "session": {"modalities": modalities, "instructions": "Be brief."}}
)
streaming, backend_ws = _gemini_default_mode_streaming(
monkeypatch, [update, _TEXT_ITEM_CREATE, _RESPONSE_CREATE], drop_params=drop_params
)
await streaming.client_ack_messages()
assert [event["type"] for event in _events_sent_to_client(streaming)] == ["session.updated"]
setup: Final = json.loads(backend_ws.send.await_args_list[0].args[0])["setup"]
assert setup["generationConfig"]["responseModalities"] == ["AUDIO"]
assert setup["systemInstruction"]["parts"][0]["text"] == "Be brief."
@pytest.mark.asyncio
async def test_default_mode_sends_no_extra_session_updated_when_gemini_confirms_the_setup(monkeypatch):
"""Each held session.update was already acknowledged, so Gemini's setupComplete
must not reach the client as one more session.updated nobody asked for."""
monkeypatch.setattr(litellm, "gemini_live_defer_setup", False, raising=False)
from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig
model: Final = "gemini-live-2.5-flash-native-audio"
config: Final = GeminiRealtimeConfig()
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({"setupComplete": {}}).encode(), ConnectionClosed(None, None)])
logging_obj: Final = _mock_logging()
logging_obj.litellm_trace_id = "trace_lazy"
logging_obj.async_success_handler = AsyncMock()
logging_obj.success_handler = MagicMock()
streaming: Final = RealTimeStreaming(client_ws, backend_ws, logging_obj, provider_config=config, model=model)
streaming.session_created_sent_to_client = True
streaming.session_configuration_request = config.session_configuration_request(model)
await streaming.backend_to_client_send_messages()
assert [json.loads(call.args[0])["type"] for call in client_ws.send_text.await_args_list] == []
@pytest.mark.asyncio
async def test_default_mode_folds_transcription_guardrail_vad_disable_into_the_merged_setup(monkeypatch):
"""A transcription guardrail gates every turn, so its auto-response disable
has to land in the one setup built from the held session.updates, injected
once rather than into every update that gets merged."""
monkeypatch.setattr(litellm, "callbacks", [_transcription_guardrail()])
debug = MagicMock()
monkeypatch.setattr("litellm.litellm_core_utils.realtime_streaming.verbose_logger.debug", debug)
instructions_update = json.dumps({"type": "session.update", "session": {"instructions": "Be brief."}})
vad_update = json.dumps(
{"type": "session.update", "session": {"turn_detection": {"type": "server_vad", "silence_duration_ms": 900}}}
)
streaming, backend_ws = _gemini_default_mode_streaming(
monkeypatch, [instructions_update, vad_update, _TEXT_ITEM_CREATE, _RESPONSE_CREATE], model="gemini-live-2.5-flash"
)
await streaming.client_ack_messages()
frames = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list]
assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"]
assert frames[0]["setup"]["realtimeInputConfig"]["automaticActivityDetection"] == {"disabled": True}
assert frames[0]["setup"]["systemInstruction"]["parts"][0]["text"] == "Be brief."
assert streaming._guardrail_turn_detection_update_sent is True
injected = [call for call in debug.call_args_list if str(call.args[0]).startswith("Injected turn_detection")]
assert len(injected) == 1
def _blocking_text_guardrail() -> CustomGuardrail:
class _BlockingTextGuardrail(CustomGuardrail):
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
if any("zebrafruit" in text for text in inputs.get("texts", [])):
raise ValueError("Content blocked")
return inputs
return _BlockingTextGuardrail(
guardrail_name="test_text_guard", event_hook=GuardrailEventHooks.pre_call, default_on=True
)
@pytest.mark.asyncio
async def test_default_mode_sends_the_setup_before_a_guardrail_blocks_the_first_message(monkeypatch):
"""Gemini Live closes the socket when content arrives before the setup, so a
guardrail that blocks the first typed message must not reach the backend first."""
monkeypatch.setattr(litellm, "callbacks", [_blocking_text_guardrail()])
instructions_update = json.dumps({"type": "session.update", "session": {"instructions": "Be brief."}})
blocked_item = json.dumps(
{
"type": "conversation.item.create",
"item": {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "zebrafruit"}]},
}
)
streaming, backend_ws = _gemini_default_mode_streaming(
monkeypatch, [instructions_update, blocked_item], model="gemini-live-2.5-flash"
)
await streaming.client_ack_messages()
frames: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list]
assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"]
assert frames[0]["setup"]["systemInstruction"]["parts"][0]["text"] == "Be brief."
assert "zebrafruit" not in json.dumps(frames[1])
errors: Final = [event["error"]["type"] for event in _events_sent_to_client(streaming) if event["type"] == "error"]
assert errors == ["guardrail_violation"]
class _SteppedClock:
def __init__(self, epoch: float) -> None:
self.now = epoch
def __call__(self) -> float:
return self.now
def _client_ws_receiving(*steps: Callable[[], Coroutine[object, object, str]]) -> MagicMock:
remaining: Final = iter(steps)
async def receive_text() -> str:
return await next(remaining)()
client_ws: Final = MagicMock()
client_ws.receive_text = receive_text
client_ws.send_text = AsyncMock()
return client_ws
@pytest.mark.asyncio
@pytest.mark.parametrize("epoch", [0.0, 1e9])
async def test_default_mode_sends_the_held_setup_when_no_content_arrives_before_the_deadline(monkeypatch, epoch):
"""Vertex AI refuses a setup sent minutes after connecting, so an idle session
sends what it holds at the deadline and answers any later session.update as late."""
monkeypatch.setattr(litellm, "gemini_live_defer_setup", False, raising=False)
from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig
clock: Final = _SteppedClock(epoch)
setup_sent: Final = asyncio.Event()
async def instructions_update_then_the_deadline_passes() -> str:
clock.now += REALTIME_LAZY_SETUP_DEADLINE_SECONDS
return json.dumps({"type": "session.update", "session": {"instructions": "Be brief."}})
async def late_update_once_the_setup_is_out() -> str:
await setup_sent.wait()
return json.dumps(
{"type": "session.update", "event_id": "evt_late", "session": {"instructions": "Answer in French."}}
)
async def disconnect() -> str:
raise ConnectionClosed(None, None)
backend_ws: Final = MagicMock()
backend_ws.send = AsyncMock(side_effect=lambda frame: setup_sent.set())
streaming: Final = RealTimeStreaming(
_client_ws_receiving(instructions_update_then_the_deadline_passes, late_update_once_the_setup_is_out, disconnect),
backend_ws,
MagicMock(),
provider_config=GeminiRealtimeConfig(),
model="gemini-live-2.5-flash",
clock=clock,
)
await asyncio.wait_for(streaming.client_ack_messages(), timeout=5)
frames: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list]
assert [next(iter(frame)) for frame in frames] == ["setup"]
assert frames[0]["setup"]["systemInstruction"]["parts"][0]["text"] == "Be brief."
events: Final = _events_sent_to_client(streaming)
assert [(event["type"], event.get("error", {}).get("code")) for event in events] == [
("session.updated", None),
("error", "session_update_not_applied"),
]
@pytest.mark.asyncio
async def test_default_mode_refuses_held_session_updates_past_the_size_cap(monkeypatch):
def instructions_update(text: str) -> str:
return json.dumps({"type": "session.update", "event_id": f"evt_{text[0]}", "session": {"instructions": text}})
half_the_cap: Final = RealTimeStreaming._MAX_BUFFERED_BYTES // 2
streaming, backend_ws = _gemini_default_mode_streaming(
monkeypatch,
[instructions_update("a" * half_the_cap), instructions_update("b" * half_the_cap), _TEXT_ITEM_CREATE, _RESPONSE_CREATE],
model="gemini-live-2.5-flash",
)
await streaming.client_ack_messages()
events: Final = _events_sent_to_client(streaming)
assert [(event["type"], event.get("error", {}).get("event_id")) for event in events] == [
("session.updated", None),
("error", "evt_b"),
]
assert events[1]["error"]["code"] == "session_update_too_large"
setup: Final = json.loads(backend_ws.send.await_args_list[0].args[0])["setup"]
assert setup["systemInstruction"]["parts"][0]["text"].startswith("a")
@pytest.mark.asyncio
@pytest.mark.parametrize("defer_setup", [False, True])
async def test_a_session_update_between_the_first_items_and_response_create_joins_the_setup(monkeypatch, defer_setup):
modalities_update: Final = json.dumps({"type": "session.update", "session": {"modalities": ["text"]}})
instructions_update: Final = json.dumps({"type": "session.update", "session": {"instructions": "Be brief."}})
streaming, backend_ws = _gemini_default_mode_streaming(
monkeypatch,
[modalities_update, _TEXT_ITEM_CREATE, instructions_update, _RESPONSE_CREATE],
model="gemini-live-2.5-flash",
defer_setup=defer_setup,
)
await streaming.client_ack_messages()
frames: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list]
assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"]
assert frames[0]["setup"]["generationConfig"]["responseModalities"] == ["TEXT"]
assert frames[0]["setup"]["systemInstruction"]["parts"][0]["text"] == "Be brief."
assert [event["type"] for event in _events_sent_to_client(streaming)] == ["session.updated", "session.updated"]
@pytest.mark.asyncio
async def test_default_mode_applies_a_session_update_that_arrives_while_audio_streams_in(monkeypatch):
modalities_update: Final = json.dumps({"type": "session.update", "session": {"modalities": ["text"]}})
append: Final = json.dumps({"type": "input_audio_buffer.append", "audio": "AAAA"})
instructions_update: Final = json.dumps({"type": "session.update", "session": {"instructions": "Be brief."}})
commit: Final = json.dumps({"type": "input_audio_buffer.commit"})
streaming, backend_ws = _gemini_default_mode_streaming(
monkeypatch, [modalities_update, append, instructions_update, commit], model="gemini-live-2.5-flash"
)
await streaming.client_ack_messages()
frames: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list]
assert [next(iter(frame)) for frame in frames][:2] == ["setup", "realtimeInput"]
assert frames[0]["setup"]["systemInstruction"]["parts"][0]["text"] == "Be brief."
assert frames[1]["realtimeInput"]["audio"]["data"] == "AAAA"
assert "error" not in [event["type"] for event in _events_sent_to_client(streaming)]
@pytest.mark.asyncio
@pytest.mark.parametrize("epoch", [0.0, 1e9])
async def test_held_items_go_out_with_the_setup_when_no_response_create_follows(monkeypatch, epoch):
monkeypatch.setattr(litellm, "gemini_live_defer_setup", False, raising=False)
from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig
item_sent: Final = asyncio.Event()
async def instructions_update() -> str:
return json.dumps({"type": "session.update", "session": {"instructions": "Be brief."}})
async def text_item() -> str:
return _TEXT_ITEM_CREATE
async def disconnect_once_the_item_is_out() -> str:
await item_sent.wait()
raise ConnectionClosed(None, None)
backend_ws: Final = MagicMock()
backend_ws.send = AsyncMock(
side_effect=lambda frame: item_sent.set() if "clientContent" in json.loads(frame) else None
)
streaming: Final = RealTimeStreaming(
_client_ws_receiving(instructions_update, text_item, disconnect_once_the_item_is_out),
backend_ws,
MagicMock(),
provider_config=GeminiRealtimeConfig(),
model="gemini-live-2.5-flash",
lazy_setup_content_grace_seconds=0,
clock=_SteppedClock(epoch),
)
await asyncio.wait_for(streaming.client_ack_messages(), timeout=5)
frames: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list]
assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"]
assert frames[0]["setup"]["systemInstruction"]["parts"][0]["text"] == "Be brief."
@pytest.mark.asyncio
async def test_default_mode_acknowledges_an_unchanged_session_update_after_content(monkeypatch):
"""Pipecat resends its unchanged session.update right after its first
conversation items and treats any error event as fatal."""
session_update: Final = json.dumps(
{"type": "session.update", "session": {"modalities": ["text"], "instructions": "Be brief."}}
)
streaming, backend_ws = _gemini_default_mode_streaming(
monkeypatch, [session_update, _TEXT_ITEM_CREATE, _RESPONSE_CREATE, session_update], model="gemini-live-2.5-flash"
)
await streaming.client_ack_messages()
frames: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list]
assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"]
events: Final = _events_sent_to_client(streaming)
assert [event["type"] for event in events] == ["session.updated", "session.updated"]
assert events[0]["session"] == events[1]["session"]
@pytest.mark.asyncio
async def test_deferred_setup_merges_session_updates_until_the_first_content(monkeypatch):
"""Deferred setup merges every session.update sent before the first content,
the same as the default mode, and answers a changed one after it as late."""
monkeypatch.setattr(litellm, "gemini_live_defer_setup", True, raising=False)
from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig
client_ws: Final = MagicMock()
client_ws.receive_text = AsyncMock(
side_effect=[
json.dumps({"type": "session.update", "session": {"modalities": ["text"]}}),
json.dumps({"type": "session.update", "session": {"instructions": "Be brief."}}),
_TEXT_ITEM_CREATE,
_RESPONSE_CREATE,
json.dumps(
{"type": "session.update", "event_id": "evt_late", "session": {"instructions": "Answer in French."}}
),
ConnectionClosed(None, None),
]
)
client_ws.send_text = AsyncMock()
backend_ws: Final = MagicMock()
backend_ws.send = AsyncMock()
streaming: Final = RealTimeStreaming(
client_ws, backend_ws, MagicMock(), provider_config=GeminiRealtimeConfig(), model="gemini-live-2.5-flash"
)
await streaming.client_ack_messages()
frames: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list]
assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"]
assert frames[0]["setup"]["generationConfig"]["responseModalities"] == ["TEXT"]
assert frames[0]["setup"]["systemInstruction"]["parts"][0]["text"] == "Be brief."
errors: Final = [event["error"] for event in _events_sent_to_client(streaming) if event["type"] == "error"]
assert [(error["code"], error["event_id"]) for error in errors] == [("session_update_not_applied", "evt_late")]
@pytest.mark.asyncio
async def test_deferred_setup_sends_the_default_setup_before_text_content_that_comes_first(monkeypatch):
"""Gemini Live closes the socket on content sent before any setup, so text that arrives
before a session.update in deferred mode gets the model's default setup first."""
monkeypatch.setattr(litellm, "gemini_live_defer_setup", True, raising=False)
from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig
client_ws: Final = MagicMock()
client_ws.receive_text = AsyncMock(
side_effect=[
_TEXT_ITEM_CREATE,
json.dumps({"type": "session.update", "event_id": "evt_late", "session": {"modalities": ["text"]}}),
ConnectionClosed(None, None),
]
)
client_ws.send_text = AsyncMock()
backend_ws: Final = MagicMock()
backend_ws.send = AsyncMock()
streaming: Final = RealTimeStreaming(
client_ws, backend_ws, MagicMock(), provider_config=GeminiRealtimeConfig(), model="gemini-live-2.5-flash"
)
await streaming.client_ack_messages()
frames: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list]
assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"]
errors: Final = [event["error"] for event in _events_sent_to_client(streaming) if event["type"] == "error"]
assert [(error["code"], error["event_id"]) for error in errors] == [("session_update_not_applied", "evt_late")]
@pytest.mark.asyncio
async def test_deferred_setup_sends_a_session_update_that_follows_buffered_audio_at_once(monkeypatch):
append: Final = json.dumps({"type": "input_audio_buffer.append", "audio": "AAAA"})
commit: Final = json.dumps({"type": "input_audio_buffer.commit"})
session_update: Final = json.dumps({"type": "session.update", "session": {"instructions": "Be a pirate."}})
streaming, backend_ws = _gemini_default_mode_streaming(
monkeypatch, [append, commit, session_update], model="gemini-live-2.5-flash", defer_setup=True
)
await streaming.client_ack_messages()
frames: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list]
assert [next(iter(frame)) for frame in frames] == ["setup"]
assert "Be a pirate." in json.dumps(frames[0]["setup"]["systemInstruction"])
@pytest.mark.parametrize("defer_setup", [False, True])
@pytest.mark.asyncio
async def test_a_transcription_guardrail_refuses_gemini_audio_it_cannot_screen(monkeypatch, defer_setup):
monkeypatch.setattr(litellm, "callbacks", [_transcription_guardrail()])
session_update: Final = json.dumps(
{"type": "session.update", "session": {"turn_detection": {"type": "server_vad", "create_response": False}}}
)
append: Final = json.dumps({"type": "input_audio_buffer.append", "audio": "AAAA"})
commit: Final = json.dumps({"type": "input_audio_buffer.commit", "event_id": "commit_1"})
streaming, backend_ws = _gemini_default_mode_streaming(
monkeypatch,
[session_update, append, append, commit, json.dumps({"type": "response.create"}), _TEXT_ITEM_CREATE],
model="gemini-live-2.5-flash",
defer_setup=defer_setup,
)
await streaming.client_ack_messages()
frames: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list]
assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"]
refusals: Final = [
event["error"]["event_id"]
for event in _events_sent_to_client(streaming)
if event["type"] == "error" and event["error"]["code"] == "audio_input_not_screenable"
]
assert refusals == [None, "commit_1"]
@pytest.mark.parametrize("defer_setup", [False, True])
@pytest.mark.parametrize("control_type", ["response.cancel", "conversation.item.delete", "input_audio_buffer.clear"])
@pytest.mark.asyncio
async def test_a_control_event_before_session_update_leaves_the_setup_open(monkeypatch, defer_setup, control_type):
control_event: Final = json.dumps({"type": control_type, "item_id": "item_x"})
session_update: Final = json.dumps({"type": "session.update", "session": {"modalities": ["text"]}})
streaming, backend_ws = _gemini_default_mode_streaming(
monkeypatch,
[control_event, session_update, _TEXT_ITEM_CREATE, _RESPONSE_CREATE],
model="gemini-live-2.5-flash",
defer_setup=defer_setup,
)
await streaming.client_ack_messages()
frames: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list]
assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"]
assert frames[0]["setup"]["generationConfig"]["responseModalities"] == ["TEXT"]
assert [event["type"] for event in _events_sent_to_client(streaming) if event["type"] == "error"] == []
@pytest.mark.asyncio
async def test_deferred_setup_buffers_audio_until_backend_setup_complete(monkeypatch):
"""Pipecat may send audio before session.update when setup is deferred."""
@ -2533,14 +3173,18 @@ async def test_deferred_setup_buffers_audio_until_backend_setup_complete(monkeyp
@pytest.mark.asyncio
async def test_deferred_setup_sends_session_update_before_buffered_audio(monkeypatch):
async def test_deferred_setup_sends_the_held_session_update_when_more_audio_arrives(monkeypatch):
monkeypatch.setattr(litellm, "gemini_live_defer_setup", True, raising=False)
from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig
client_ws = MagicMock()
client_ws.send_text = AsyncMock()
audio_msg = json.dumps({"type": "input_audio_buffer.append", "audio": "AA=="})
session_update = json.dumps({"type": "session.update", "session": {"modalities": ["audio"]}})
client_ws.receive_text = AsyncMock(side_effect=[audio_msg, session_update, ConnectionClosed(None, None)])
later_audio_msg = json.dumps({"type": "input_audio_buffer.append", "audio": "AQ=="})
client_ws.receive_text = AsyncMock(
side_effect=[audio_msg, session_update, later_audio_msg, ConnectionClosed(None, None)]
)
backend_ws = MagicMock()
backend_ws.send = AsyncMock()
logging_obj: Final = _mock_logging()
@ -2560,7 +3204,7 @@ async def test_deferred_setup_sends_session_update_before_buffered_audio(monkeyp
sent_payload = json.loads(backend_ws.send.await_args_list[0].args[0])
assert "setup" in sent_payload
assert "realtimeInput" not in sent_payload
assert streaming._pending_messages_until_setup == [audio_msg]
assert streaming._pending_messages_until_setup == [audio_msg, later_audio_msg]
@pytest.mark.asyncio
@ -3523,6 +4167,24 @@ async def test_upstream_refusal_after_a_synthetic_session_created_still_logs_a_f
assert session.logging.logged_sessions == ()
class _LazySetupRealtimeConfig(_PassthroughRealtimeConfig):
def builds_setup_from_session_update(self) -> bool:
return True
@pytest.mark.asyncio
async def test_upstream_outage_before_a_held_setup_was_sent_logs_the_session_not_a_refusal():
upstream_close: Final = ConnectionClosed(Close(1011, "upstream outage"), None)
session: Final = _relay_session(
_client_ws_that_never_sends(), _backend_ws_closing_with(upstream_close), _LazySetupRealtimeConfig()
)
await session.run()
assert session.logging.logged_failures == ()
assert len(session.logging.logged_sessions) == 1
@pytest.mark.asyncio
async def test_upstream_close_after_relayed_events_still_logs_the_session_as_success():
client_ws: Final = _client_ws_that_never_sends()

View file

@ -1,6 +1,7 @@
import json
import os
import sys
from typing import Final
import pytest
@ -70,3 +71,70 @@ def test_prepare_completion_kwargs_keeps_prompt_cache_key_through_responses_rero
)
assert completion_kwargs["model"] == "openai/responses/gpt-5.6-luna"
assert completion_kwargs["prompt_cache_key"] == "session-abc"
@pytest.mark.parametrize(
("api_base", "azure_api_base_env", "derives_key"),
[
(None, None, True),
("https://api.mistral.ai/v1", None, True),
("https://my-resource.openai.azure.com/openai/v1", None, False),
(None, "https://my-resource.services.ai.azure.com/openai", False),
],
)
def test_prepare_completion_kwargs_derives_prompt_cache_key_only_for_a_mistral_hosted_endpoint(
monkeypatch: pytest.MonkeyPatch, api_base: str | None, azure_api_base_env: str | None, derives_key: bool
) -> None:
if azure_api_base_env is None:
monkeypatch.delenv("MISTRAL_AZURE_API_BASE", raising=False)
else:
monkeypatch.setenv("MISTRAL_AZURE_API_BASE", azure_api_base_env)
completion_kwargs: Final = _prepare(
"mistral/mistral-small-latest", {"custom_llm_provider": "mistral", "api_base": api_base}
)
assert completion_kwargs["user"] == CLAUDE_CODE_USER_ID
assert ("prompt_cache_key" in completion_kwargs) is derives_key
@pytest.mark.parametrize(
("codestral_api_base_env", "derives_key"),
[(None, True), ("https://my-resource.services.ai.azure.com/openai/v1", False)],
)
def test_prepare_completion_kwargs_derives_prompt_cache_key_for_codestral_only_on_its_mistral_host(
monkeypatch: pytest.MonkeyPatch, codestral_api_base_env: str | None, derives_key: bool
) -> None:
monkeypatch.delenv("MISTRAL_AZURE_API_BASE", raising=False)
if codestral_api_base_env is None:
monkeypatch.delenv("CODESTRAL_API_BASE", raising=False)
else:
monkeypatch.setenv("CODESTRAL_API_BASE", codestral_api_base_env)
completion_kwargs: Final = _prepare("codestral/codestral-latest", {"custom_llm_provider": "codestral"})
assert ("prompt_cache_key" in completion_kwargs) is derives_key
def test_prepare_completion_kwargs_derives_no_prompt_cache_key_mistral_cannot_take_from_a_non_ascii_session() -> None:
non_ascii_user_id: Final = json.dumps({"device_id": "d" * 64, "account_uuid": "", "session_id": "sitzung-ü"})
mistral_kwargs: Final = LiteLLMMessagesToCompletionTransformationHandler._prepare_completion_kwargs(
max_tokens=1024,
messages=MESSAGES,
model="mistral/mistral-small-latest",
metadata={"user_id": non_ascii_user_id},
thinking=None,
extra_kwargs={"custom_llm_provider": "mistral"},
)[0]
openai_kwargs: Final = LiteLLMMessagesToCompletionTransformationHandler._prepare_completion_kwargs(
max_tokens=1024,
messages=MESSAGES,
model="openai/gpt-5.6-luna",
metadata={"user_id": non_ascii_user_id},
thinking=None,
extra_kwargs={"custom_llm_provider": "openai"},
)[0]
assert "prompt_cache_key" not in mistral_kwargs
assert openai_kwargs["prompt_cache_key"] == "sitzung-ü"

View file

@ -14,6 +14,7 @@ import httpcore
import httpx
import pytest
import respx
import websockets.exceptions
from botocore.credentials import RefreshableCredentials
from h2.events import ConnectionTerminated
@ -2787,6 +2788,194 @@ async def test_async_realtime_error_event_send_failure_still_closes():
assert client_ws.events == [("close", (1011, "Internal server error: vertex token refresh exploded"))]
_GEMINI_LIVE_SETUP_COMPLETE = json.dumps({"setupComplete": {}})
_GEMINI_LIVE_TEXT_TURN = json.dumps({"serverContent": {"modelTurn": {"parts": [{"text": "Blue."}]}}})
_GEMINI_LIVE_TURN_COMPLETE = json.dumps({"serverContent": {"turnComplete": True}})
class _ScriptedRealtimeBackend:
"""Backend socket that answers the first setup frame with the scripted replies
and closes once the client is done, so both bridge loops run to completion."""
def __init__(self, replies_after_setup):
self.sent = []
self.closed = asyncio.Event()
self._replies_after_setup = list(replies_after_setup)
self._inbox = asyncio.Queue()
async def __aenter__(self):
return self
async def __aexit__(self, exc_type, exc, tb):
return None
async def send(self, payload):
self.sent.append(payload)
if "setup" in json.loads(payload):
for reply in self._replies_after_setup:
self._inbox.put_nowait(reply)
def close_after_replies(self):
self._inbox.put_nowait(None)
async def recv(self, decode=True):
reply = await self._inbox.get()
if reply is None:
self.closed.set()
raise websockets.exceptions.ConnectionClosedOK(None, None)
return reply
class _ScriptedRealtimeClient:
def __init__(self, messages, backend):
self.scope = {"headers": []}
self.exceptions = websockets.exceptions
self.sent = []
self._messages = list(messages)
self._backend = backend
async def send_text(self, payload):
self.sent.append(payload)
async def receive_text(self):
if self._messages:
return self._messages.pop(0)
self._backend.close_after_replies()
await self._backend.closed.wait()
raise websockets.exceptions.ConnectionClosedOK(None, None)
async def close(self, code=None, reason=None):
return None
def _realtime_logging_obj():
logging_obj = Mock()
logging_obj.litellm_trace_id = "trace-lit-6974"
logging_obj.async_success_handler = AsyncMock()
logging_obj.dispatch_failure_handlers = AsyncMock()
return logging_obj
@pytest.mark.asyncio
@pytest.mark.parametrize(
"model, session, backend_modality, client_modality, backend_turn, drop_params",
[
(
"gemini-live-2.5-flash",
{"type": "realtime", "output_modalities": ["text"]},
"TEXT",
"text",
[_GEMINI_LIVE_TEXT_TURN],
False,
),
("gemini-live-2.5-flash", {"modalities": ["text"]}, "TEXT", "text", [_GEMINI_LIVE_TEXT_TURN], False),
(
"gemini-live-2.5-flash-native-audio",
{"type": "realtime", "output_modalities": ["text"]},
"AUDIO",
"audio",
[],
True,
),
],
)
async def test_async_realtime_vertex_first_session_update_builds_the_only_setup(
model, session, backend_modality, client_modality, backend_turn, drop_params: bool, monkeypatch
):
"""Regression for text-only Vertex Live sessions answered in audio: the
connect-time default setup (AUDIO) used to go out before the client's
session.update, and Gemini accepts only one setup per connection, so the
requested modality was dropped and billed at the audio rate. The setup now
goes out on the client's first message, built from session.update when that
comes first. Under drop_params an audio-only model still gets AUDIO, and the
client learns the downgrade from session.updated."""
from litellm.llms.vertex_ai.realtime.transformation import VertexAIRealtimeConfig
monkeypatch.setattr(litellm, "gemini_live_defer_setup", False, raising=False)
backend = _ScriptedRealtimeBackend([_GEMINI_LIVE_SETUP_COMPLETE, *backend_turn, _GEMINI_LIVE_TURN_COMPLETE])
client_ws = _ScriptedRealtimeClient(
[
json.dumps({"type": "session.update", "session": session}),
json.dumps(
{
"type": "conversation.item.create",
"item": {
"type": "message",
"role": "user",
"content": [{"type": "input_text", "text": "What color is the sky? One word."}],
},
}
),
json.dumps({"type": "response.create"}),
],
backend,
)
handler = BaseLLMHTTPHandler()
with patch.object(handler, "_open_realtime_backend_ws", AsyncMock(return_value=backend)):
await asyncio.wait_for(
handler.async_realtime(
model=model,
websocket=client_ws,
logging_obj=_realtime_logging_obj(),
provider_config=VertexAIRealtimeConfig(
access_token="fake-token", project="fake-project", location="us-central1", drop_params=drop_params
),
headers={},
),
timeout=10,
)
backend_frames = [json.loads(frame) for frame in backend.sent]
setups = [frame["setup"] for frame in backend_frames if "setup" in frame]
assert len(setups) == 1
assert "setup" in backend_frames[0]
assert setups[0]["generationConfig"]["responseModalities"] == [backend_modality]
assert any("clientContent" in frame for frame in backend_frames[1:])
client_events = [json.loads(frame) for frame in client_ws.sent]
event_types = [event["type"] for event in client_events]
assert event_types[:2] == ["session.created", "session.updated"]
assert client_events[1]["session"]["modalities"] == [client_modality]
if backend_turn:
assert "response.output_text.delta" in event_types
assert "response.output_audio_transcript.delta" not in event_types
@pytest.mark.asyncio
async def test_async_realtime_sends_setup_at_connect_when_provider_cannot_build_it_from_session_update():
"""Bedrock Nova Sonic style providers keep the connect-time setup and get no
synthetic session.created."""
setup = json.dumps({"event": {"sessionStart": {}}})
backend = _ScriptedRealtimeBackend([])
client_ws = _ScriptedRealtimeClient([], backend)
provider_config = Mock()
provider_config.get_complete_url.return_value = "wss://backend.example/live"
provider_config.validate_environment.return_value = {}
provider_config.requires_session_configuration.return_value = True
provider_config.builds_setup_from_session_update.return_value = False
provider_config.session_configuration_request.return_value = setup
provider_config.unbilled_usage_on_session_close.return_value = None
provider_config.open_backend = AsyncMock(return_value=None)
provider_config.wrap_backend.side_effect = lambda opened: opened
handler = BaseLLMHTTPHandler()
with patch.object(handler, "_open_realtime_backend_ws", AsyncMock(return_value=backend)):
await asyncio.wait_for(
handler.async_realtime(
model="amazon.nova-sonic-v1:0",
websocket=client_ws,
logging_obj=_realtime_logging_obj(),
provider_config=provider_config,
headers={},
),
timeout=10,
)
assert backend.sent == [setup]
assert client_ws.sent == []
class _JSONBodyAudioTranscriptionConfig(BaseAudioTranscriptionConfig):
def get_supported_openai_params(self, model):
return []

View file

@ -1,3 +1,4 @@
import base64
import json
from collections.abc import Mapping
from typing import Final, cast
@ -7,7 +8,7 @@ import pytest
import litellm
from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig
from litellm.llms.gemini.realtime.transformation import GEMINI_LIVE_DOCUMENTED_VOICES, GeminiRealtimeConfig
from litellm.types.llms.gemini import BidiGenerateContentServerMessage
@ -439,6 +440,51 @@ def test_gemini_realtime_session_update_with_tools():
assert "parameters" in function_decl
def _schema_nodes(schema: object):
if isinstance(schema, dict):
yield schema
for value in schema.values():
yield from _schema_nodes(value)
elif isinstance(schema, list):
for value in schema:
yield from _schema_nodes(value)
def test_gemini_session_update_converts_flat_realtime_tool_schemas_to_gemini_schemas():
session_update: Final = {
"type": "session.update",
"session": {
"tools": [
{
"type": "function",
"name": "lookup_weather",
"description": "Look up the weather in a city.",
"parameters": {
"type": "object",
"$defs": {
"City": {"type": "object", "properties": {"name": {"type": "string"}}, "required": ["name"]}
},
"properties": {"city": {"$ref": "#/$defs/City"}, "note": {"type": ["string", "null"]}},
"required": ["city"],
},
}
]
},
}
messages: Final = GeminiRealtimeConfig().transform_realtime_request(
json.dumps(session_update), "gemini-2.5-flash", session_configuration_request=None
)
declaration: Final = json.loads(messages[0])["setup"]["tools"][0]["function_declarations"][0]
assert declaration["name"] == "lookup_weather"
parameters: Final = declaration["parameters"]
assert parameters["properties"]["city"]["properties"]["name"]["type"] == "string"
nodes: Final = tuple(_schema_nodes(parameters))
assert not any("$defs" in node or "$ref" in node for node in nodes)
assert all(isinstance(node["type"], str) for node in nodes if "type" in node)
def test_gemini_session_update_defaults_to_audio_modality():
config = GeminiRealtimeConfig()
@ -613,6 +659,35 @@ def test_gemini_audio_only_live_models_drop_text_from_text_audio_combo(patch_gem
assert setup["generationConfig"]["responseModalities"] == ["AUDIO"]
@pytest.mark.parametrize("output_modalities", ["text", ["text"]])
def test_gemini_session_update_accepts_a_single_string_output_modality(output_modalities):
"""Regression: a bare string was iterated per character, so output_modalities
"text" reached Vertex as responseModalities ["T"], which it rejects with 1007."""
config = GeminiRealtimeConfig()
session_update = {"type": "session.update", "session": {"type": "realtime", "output_modalities": output_modalities}}
messages = config.transform_realtime_request(
json.dumps(session_update), "gemini-live-2.5-flash", session_configuration_request=None
)
assert json.loads(messages[0])["setup"]["generationConfig"]["responseModalities"] == ["TEXT"]
@pytest.mark.parametrize(
("output_modalities", "expected"),
[(["video", "text"], ["TEXT"]), (["video"], ["AUDIO"]), (None, ["AUDIO"])],
)
def test_gemini_session_update_drops_modalities_live_cannot_produce(output_modalities, expected):
config = GeminiRealtimeConfig()
session_update = {"type": "session.update", "session": {"type": "realtime", "output_modalities": output_modalities}}
messages = config.transform_realtime_request(
json.dumps(session_update), "gemini-live-2.5-flash", session_configuration_request=None
)
assert json.loads(messages[0])["setup"]["generationConfig"]["responseModalities"] == expected
def test_gemini_non_live_model_preserves_text_modality():
config = GeminiRealtimeConfig()
session_update = {
@ -1396,22 +1471,71 @@ def test_gemini_input_audio_buffer_clear_is_local_noop():
assert messages == []
def test_gemini_input_audio_buffer_commit_maps_to_activity_end_when_manual_vad():
_MANUAL_VAD_SETUP = json.dumps({"setup": {"realtimeInputConfig": {"automaticActivityDetection": {"disabled": True}}}})
_APPEND = json.dumps({"type": "input_audio_buffer.append", "audio": base64.b64encode(b"\x00\x01" * 80).decode()})
def _realtime_input_keys(config: GeminiRealtimeConfig, client_event: str, setup: str | None) -> list[str]:
frames = config.transform_realtime_request(client_event, "gemini-live-2.5-flash-native-audio", setup)
return [key for frame in frames for key in json.loads(frame)["realtimeInput"]]
@pytest.mark.parametrize("turn_end", ["input_audio_buffer.commit", "response.create"])
def test_gemini_manual_vad_wraps_each_audio_turn_in_activity_start_and_end(turn_end):
"""With activity detection off, Gemini Live answers only after activityEnd closes an
activityStart, and it closes the socket on any other shape of those markers."""
config = GeminiRealtimeConfig()
setup = {
"setup": {
"realtimeInputConfig": {
"automaticActivityDetection": {"disabled": True},
}
}
}
messages = config.transform_realtime_request(
json.dumps({"type": "input_audio_buffer.commit"}),
"gemini-live-2.5-flash-native-audio",
session_configuration_request=json.dumps(setup),
end = json.dumps({"type": turn_end})
turns = [
[_realtime_input_keys(config, event, _MANUAL_VAD_SETUP) for event in (_APPEND, _APPEND, end, end)]
for _ in range(2)
]
assert turns == [[["activityStart", "audio"], ["audio"], ["activityEnd"], []]] * 2
def test_gemini_manual_vad_activity_markers_are_empty_messages():
config = GeminiRealtimeConfig()
commit = json.dumps({"type": "input_audio_buffer.commit"})
first_append = config.transform_realtime_request(_APPEND, "gemini-live-2.5-flash", _MANUAL_VAD_SETUP)
turn_end = config.transform_realtime_request(commit, "gemini-live-2.5-flash", _MANUAL_VAD_SETUP)
assert json.loads(first_append[0]) == {"realtimeInput": {"activityStart": {}}}
assert [json.loads(frame) for frame in turn_end] == [{"realtimeInput": {"activityEnd": {}}}]
@pytest.mark.parametrize(
("session", "expected_keys"),
[
({"turn_detection": None}, [["activityStart", "audio"], ["activityEnd"]]),
({"type": "realtime", "audio": {"input": {"turn_detection": None}}}, [["activityStart", "audio"], ["activityEnd"]]),
({"type": "realtime", "audio": {"input": {}}}, [["audio"], ["audioStreamEnd"]]),
],
)
def test_gemini_null_turn_detection_makes_the_client_mark_each_turn(session, expected_keys):
"""OpenAI's turn_detection null turns server VAD off, so the client commits each turn itself."""
config = GeminiRealtimeConfig()
setup = config.transform_realtime_request(
json.dumps({"type": "session.update", "session": session}), "gemini-live-2.5-flash-native-audio", None
)
assert len(messages) == 1
assert json.loads(messages[0]) == {"realtimeInput": {"activityEnd": True}}
commit = json.dumps({"type": "input_audio_buffer.commit"})
keys = [_realtime_input_keys(config, event, setup[0]) for event in (_APPEND, commit)]
assert keys == expected_keys
def test_gemini_automatic_vad_sends_no_activity_markers():
config = GeminiRealtimeConfig()
commit = json.dumps({"type": "input_audio_buffer.commit"})
response_create = json.dumps({"type": "response.create"})
keys = [_realtime_input_keys(config, event, None) for event in (_APPEND, commit)]
assert keys == [["audio"], ["audioStreamEnd"]]
assert config.transform_realtime_request(response_create, "gemini-live-2.5-flash", None) == []
def test_gemini_tool_call_id_to_name_evicts_oldest_when_capped():
@ -1883,13 +2007,75 @@ def test_map_openai_params_drops_stock_voice_case_insensitively():
"""
cfg = GeminiRealtimeConfig()
dropped = cfg.map_openai_params(optional_params={}, non_default_params={"voice": "Alloy"})
dropped = cfg.map_openai_params(optional_params={}, non_default_params={"voice": "Alloy"}, model="gemini-3.8-live")
assert "speechConfig" not in dropped.get("generationConfig", {})
passthrough = cfg.map_openai_params(optional_params={}, non_default_params={"voice": "Kore"})
passthrough = cfg.map_openai_params(
optional_params={}, non_default_params={"voice": "Kore"}, model="gemini-3.8-live"
)
assert passthrough["generationConfig"]["speechConfig"]["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == "Kore"
_DOCUMENTED_VOICE: Final = sorted(GEMINI_LIVE_DOCUMENTED_VOICES)[0].capitalize()
@pytest.mark.parametrize("drop_params", [False, True])
def test_google_ai_studio_forwards_a_documented_voice(drop_params):
config: Final = GeminiRealtimeConfig(drop_params=drop_params)
mapped: Final = config.map_openai_params(
optional_params={}, non_default_params={"voice": _DOCUMENTED_VOICE}, model="gemini-3.8-live"
)
speech_config: Final = mapped["generationConfig"]["speechConfig"]
assert speech_config["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == _DOCUMENTED_VOICE
assert config.setup_session_update_rejection("gemini-3.8-live", {"voice": _DOCUMENTED_VOICE}, setup_sent=False) is None
@pytest.mark.parametrize(
("session", "param"),
[
({"voice": "Nova"}, "session.voice"),
({"type": "realtime", "audio": {"output": {"voice": "Nova"}}}, "session.audio.output.voice"),
],
)
def test_google_ai_studio_rejects_an_undocumented_voice_on_the_session_update(session, param):
rejection: Final = GeminiRealtimeConfig().setup_session_update_rejection(
"gemini-2.5-flash-native-audio-preview-12-2025", session, setup_sent=False
)
assert rejection is not None
assert (rejection.code, rejection.param) == ("invalid_value", param)
def _setup_speech_config(config: GeminiRealtimeConfig, model: str, session: dict[str, object]) -> object:
messages: Final = config.transform_realtime_request(
json.dumps({"type": "session.update", "session": session}), model, session_configuration_request=None
)
return json.loads(messages[0])["setup"].get("generationConfig", {}).get("speechConfig")
@pytest.mark.parametrize(
"session",
[{"voice": "Nova"}, {"type": "realtime", "audio": {"output": {"voice": "Nova"}}}],
)
def test_google_ai_studio_drops_an_undocumented_voice_with_drop_params(session):
config: Final = GeminiRealtimeConfig(drop_params=True)
model: Final = "gemini-2.5-flash-native-audio-preview-12-2025"
assert config.setup_session_update_rejection(model, session, setup_sent=False) is None
assert _setup_speech_config(config, model, session) is None
@pytest.mark.parametrize("drop_params", [False, True])
@pytest.mark.parametrize("model", ["gemini-3.8-live", "gemini-3.1-flash-live-preview"])
def test_google_ai_studio_forwards_any_voice_to_models_that_accept_it(model, drop_params):
config: Final = GeminiRealtimeConfig(drop_params=drop_params)
assert config.setup_session_update_rejection(model, {"voice": "Nova"}, setup_sent=False) is None
speech_config: Final = _setup_speech_config(config, model, {"voice": "Nova"})
assert speech_config == {"voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Nova"}}}
@pytest.fixture(autouse=False)
def patch_gemini_transcribe_live_cost_map_entry(monkeypatch):
"""Inject the gemini-3.5-transcribe-live registry entry locally.
@ -1954,6 +2140,55 @@ def test_gemini_transcribe_live_coerces_audio_modality_to_text(modalities, patch
assert setup["generationConfig"]["responseModalities"] == ["TEXT"]
@pytest.mark.parametrize("model", ["gemini-3.5-transcribe-live", "gemini/gemini-3.5-transcribe-live"])
def test_gemini_transcribe_live_synthetic_session_created_reports_text_output(
model, patch_gemini_transcribe_live_cost_map_entry
):
session = GeminiRealtimeConfig().transform_session_created_event(model, "sess_transcribe")["session"]
assert session["modalities"] == ["text"]
@pytest.mark.parametrize(
("model", "modalities", "producible"),
[
("gemini-3.5-transcribe-live", ["audio"], False),
("gemini-3.5-transcribe-live", ["audio", "text"], True),
("gemini-live-2.5-flash-native-audio", ["text"], False),
("gemini-live-2.5-flash-native-audio", ["text", "audio"], True),
],
)
@pytest.mark.parametrize("drop_params", [False, True])
def test_gemini_live_rejects_only_output_modalities_the_model_cannot_produce_at_all(
model,
modalities,
producible,
drop_params,
patch_gemini_transcribe_live_cost_map_entry,
patch_gemini_audio_cost_map_entries,
):
rejection = GeminiRealtimeConfig(drop_params=drop_params).setup_session_update_rejection(
model, {"output_modalities": modalities}, setup_sent=False
)
assert (rejection is None) is (drop_params or producible)
assert rejection is None or rejection.param == "session.output_modalities"
@pytest.mark.parametrize("drop_params", [False, True])
def test_gemini_provider_realtime_config_carries_drop_params(drop_params):
from litellm.types.utils import LlmProviders
from litellm.utils import ProviderConfigManager
model = "gemini-live-2.5-flash-native-audio"
config = ProviderConfigManager.get_provider_realtime_config(
model=model, provider=LlmProviders.GEMINI, drop_params=drop_params
)
rejection = config.setup_session_update_rejection(model, {"modalities": ["text"]}, setup_sent=False)
assert (rejection is None) is drop_params
def test_gemini_chat_model_with_text_output_modalities_keeps_audio_eager_setup(
patch_gemini_transcribe_live_cost_map_entry,
):
@ -2169,6 +2404,47 @@ def test_unbilled_usage_on_session_close_flushes_trailing_audio(patch_gemini_tra
assert config.unbilled_usage_on_session_close("gemini-3.5-transcribe-live") is None
def test_gemini_realtime_two_output_modalities_collapse_to_audio():
"""Gemini Live closes the socket with 1007 when the setup names two response
modalities, so a GA session.update asking for ["text", "audio"] must reach the
backend as AUDIO only, the same choice the beta shim makes."""
config = GeminiRealtimeConfig()
session_update = {
"type": "session.update",
"session": {"output_modalities": ["text", "audio"], "instructions": "Be brief."},
}
messages = config.transform_realtime_request(
json.dumps(session_update),
"gemini-live-2.5-flash",
session_configuration_request=None,
)
assert len(messages) == 1
assert json.loads(messages[0])["setup"]["generationConfig"]["responseModalities"] == ["AUDIO"]
def test_gemini_realtime_inf_max_response_output_tokens_is_not_forwarded():
"""OpenAI's default max_response_output_tokens is the string "inf", which Gemini
rejects with 1007 (TYPE_INT32), so only integer limits map to maxOutputTokens."""
config = GeminiRealtimeConfig()
def generation_config_for(limit):
session_update = {
"type": "session.update",
"session": {"modalities": ["text"], "max_response_output_tokens": limit},
}
messages = config.transform_realtime_request(
json.dumps(session_update),
"gemini-live-2.5-flash",
session_configuration_request=None,
)
return json.loads(messages[0])["setup"]["generationConfig"]
assert "maxOutputTokens" not in generation_config_for("inf")
assert generation_config_for(200)["maxOutputTokens"] == 200
def _grounded_live_frame(grounding_metadata: Mapping[str, object] | None) -> Mapping[str, object]:
"""One Live server frame. Grounding metadata and usageMetadata arrive together, as Vertex sends them."""
from typing import Final

View file

@ -1,7 +1,13 @@
import json
from collections.abc import Mapping
from typing import Final, List, cast
from unittest.mock import MagicMock, patch
from urllib.parse import urlparse
import httpx
import pytest
import respx
from openai import AsyncOpenAI
from litellm.litellm_core_utils.prompt_templates.common_utils import TOOL_RESULT_IMAGE_BOUNDARY
from litellm.llms.mistral.chat.transformation import (
@ -45,21 +51,15 @@ class TestMistralReasoningSupport:
mistral_config = MistralConfig()
# Test magistral model supports reasoning parameters
supported_params = mistral_config.get_supported_openai_params(
"mistral/magistral-medium-2506"
)
supported_params = mistral_config.get_supported_openai_params("mistral/magistral-medium-2506")
assert "reasoning_effort" in supported_params
assert "thinking" in supported_params
supported_params_reasoning = mistral_config.get_supported_openai_params(
"mistral/mistral-medium-latest"
)
supported_params_reasoning = mistral_config.get_supported_openai_params("mistral/mistral-medium-latest")
assert "reasoning_effort" in supported_params_reasoning
assert "thinking" not in supported_params_reasoning
supported_params_normal = mistral_config.get_supported_openai_params(
"mistral/mistral-large-latest"
)
supported_params_normal = mistral_config.get_supported_openai_params("mistral/mistral-large-latest")
assert "reasoning_effort" not in supported_params_normal
assert "thinking" not in supported_params_normal
@ -117,6 +117,36 @@ class TestMistralReasoningSupport:
)
assert optional_params["reasoning_effort"] == sent
@pytest.mark.parametrize(
("requested", "drop_params", "sent"),
[
({"effort": "medium"}, False, {"reasoning_effort": "high"}),
({"effort": "high", "summary": "auto"}, True, {"reasoning_effort": "high"}),
({"summary": "auto"}, True, {}),
],
)
def test_reasoning_effort_object_is_sent_as_its_level(self, local_model_cost_map, requested, drop_params, sent):
import litellm
optional_params = litellm.get_optional_params(
model="mistral-small-latest",
custom_llm_provider="mistral",
reasoning_effort=requested,
drop_params=drop_params,
)
assert {key: optional_params[key] for key in ("reasoning_effort",) if key in optional_params} == sent
@pytest.mark.parametrize("requested", [{"effort": "high", "summary": "auto"}, {"summary": "auto"}])
def test_reasoning_effort_object_mistral_cannot_take_is_rejected(self, local_model_cost_map, requested):
import litellm
with pytest.raises(litellm.UnsupportedParamsError, match=r"\['reasoning_effort'\]"):
litellm.get_optional_params(
model="devstral-latest",
custom_llm_provider="mistral",
reasoning_effort=requested,
)
def test_reasoning_effort_is_forwarded_verbatim_when_the_map_declares_no_levels(
self, local_model_cost_map, monkeypatch
):
@ -194,9 +224,7 @@ class TestMistralReasoningSupport:
messages = [{"role": "user", "content": "What is 2+2?"}]
optional_params = {"_add_reasoning_prompt": True}
result = mistral_config._add_reasoning_system_prompt_if_needed(
messages, optional_params
)
result = mistral_config._add_reasoning_system_prompt_if_needed(messages, optional_params)
# Should add a new system message at the beginning
assert len(result) == 2
@ -218,9 +246,7 @@ class TestMistralReasoningSupport:
]
optional_params = {"_add_reasoning_prompt": True}
result = mistral_config._add_reasoning_system_prompt_if_needed(
messages, optional_params
)
result = mistral_config._add_reasoning_system_prompt_if_needed(messages, optional_params)
# Should modify existing system message
assert len(result) == 2
@ -251,9 +277,7 @@ class TestMistralReasoningSupport:
]
optional_params = {"_add_reasoning_prompt": True}
result = mistral_config._add_reasoning_system_prompt_if_needed(
messages, optional_params
)
result = mistral_config._add_reasoning_system_prompt_if_needed(messages, optional_params)
# Should modify existing system message preserving list format
assert len(result) == 2
@ -266,10 +290,7 @@ class TestMistralReasoningSupport:
# Original content should be preserved
assert "You are a helpful assistant." in result[0]["content"][1]["text"]
assert (
"You always provide detailed explanations."
in result[0]["content"][2]["text"]
)
assert "You always provide detailed explanations." in result[0]["content"][2]["text"]
assert result[1]["role"] == "user"
@ -287,9 +308,7 @@ class TestMistralReasoningSupport:
]
string_params = {"_add_reasoning_prompt": True}
string_result = mistral_config._add_reasoning_system_prompt_if_needed(
string_messages, string_params
)
string_result = mistral_config._add_reasoning_system_prompt_if_needed(string_messages, string_params)
assert isinstance(string_result[0]["content"], str)
assert "<think>" in string_result[0]["content"]
assert "You are helpful." in string_result[0]["content"]
@ -304,9 +323,7 @@ class TestMistralReasoningSupport:
]
list_params = {"_add_reasoning_prompt": True}
list_result = mistral_config._add_reasoning_system_prompt_if_needed(
list_messages, list_params
)
list_result = mistral_config._add_reasoning_system_prompt_if_needed(list_messages, list_params)
assert isinstance(list_result[0]["content"], list)
assert list_result[0]["content"][0]["type"] == "text"
assert "<think>" in list_result[0]["content"][0]["text"]
@ -319,9 +336,7 @@ class TestMistralReasoningSupport:
messages = [{"role": "user", "content": "What is 2+2?"}]
optional_params = {}
result = mistral_config._add_reasoning_system_prompt_if_needed(
messages, optional_params
)
result = mistral_config._add_reasoning_system_prompt_if_needed(messages, optional_params)
# Should return messages unchanged
assert result == messages
@ -444,9 +459,7 @@ class TestMistralReasoningSupport:
def test_mistral_streaming_chunk_preserves_thinking_blocks():
"""Ensure streaming chunks keep magistral reasoning content."""
iterator = MistralChatResponseIterator(
streaming_response=iter([]), sync_stream=True, json_mode=False
)
iterator = MistralChatResponseIterator(streaming_response=iter([]), sync_stream=True, json_mode=False)
streamed_chunk = {
"id": "chunk-1",
@ -523,9 +536,7 @@ class TestMistralParallelToolCalls:
def test_get_supported_openai_params_includes_parallel_tool_calls(self):
"""Test that parallel_tool_calls is in supported parameters."""
mistral_config = MistralConfig()
supported_params = mistral_config.get_supported_openai_params(
"mistral/mistral-large-latest"
)
supported_params = mistral_config.get_supported_openai_params("mistral/mistral-large-latest")
assert "parallel_tool_calls" in supported_params
def test_transform_request_preserves_parallel_tool_calls(self):
@ -691,9 +702,7 @@ class TestMistralEmptyContentHandling:
result = MistralConfig._handle_empty_content_response(response_data)
assert (
result["choices"][0]["message"]["content"] == "Hello, how can I help you?"
)
assert result["choices"][0]["message"]["content"] == "Hello, how can I help you?"
def test_handle_empty_content_response_handles_multiple_choices(self):
"""Test that only the first choice is processed for empty content."""
@ -1431,9 +1440,7 @@ def test_mistral_transform_request_keeps_an_empty_tool_message_after_its_tool_ca
{
"role": "assistant",
"content": "",
"tool_calls": [
{"id": "call_1", "type": "function", "function": {"name": "read", "arguments": "{}"}}
],
"tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "read", "arguments": "{}"}}],
},
{"role": "tool", "tool_call_id": "call_1", "content": ""},
],
@ -1457,9 +1464,7 @@ def test_mistral_transform_request_sends_a_null_tool_reply_as_empty_text(reply):
{
"role": "assistant",
"content": "",
"tool_calls": [
{"id": "call_1", "type": "function", "function": {"name": "read", "arguments": "{}"}}
],
"tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "read", "arguments": "{}"}}],
},
{"role": "tool", "tool_call_id": "call_1", **reply},
],
@ -1514,9 +1519,7 @@ def test_mistral_transform_request_hoists_tool_message_image():
{
"role": "assistant",
"content": "",
"tool_calls": [
{"id": "call_1", "type": "function", "function": {"name": "read", "arguments": "{}"}}
],
"tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "read", "arguments": "{}"}}],
},
{
"role": "tool",
@ -1665,3 +1668,274 @@ def test_mistral_transform_request_drops_an_empty_tool_message_whose_call_id_is_
)
assert [m.get("role") for m in request["messages"]] == ["user", "assistant"]
@pytest.mark.parametrize(
("model", "custom_llm_provider"),
[
("mistral-large-latest", "mistral"),
("mistral-large-4", "mistral"),
("a-mistral-model-missing-from-the-cost-map", "mistral"),
("codestral-latest", "codestral"),
],
)
def test_prompt_cache_key_reaches_the_mistral_request_body(
local_model_cost_map: None, model: str, custom_llm_provider: str
) -> None:
import litellm
optional_params: Final = litellm.get_optional_params(
model=model, custom_llm_provider=custom_llm_provider, prompt_cache_key="session-1"
)
request: Final = MistralConfig().transform_request(
model=model,
messages=[{"role": "user", "content": "hi"}],
optional_params=optional_params,
litellm_params={},
headers={},
)
assert request["prompt_cache_key"] == "session-1"
AZURE_HOSTED_MISTRAL_API_BASE: Final = "https://my-resource.services.ai.azure.com/openai/v1"
def _transform_with_prompt_cache_key(
litellm_params: dict[str, object], prompt_cache_key: object = "session-1"
) -> Mapping[str, object]:
return cast(
Mapping[str, object],
MistralConfig().transform_request(
model="mistral-small-latest",
messages=[{"role": "user", "content": "hi"}],
optional_params={"prompt_cache_key": prompt_cache_key, "temperature": 0.2},
litellm_params=litellm_params,
headers={},
),
)
@pytest.mark.parametrize("api_base", ["https://api.mistral.ai/v1", "https://codestral.mistral.ai/v1"])
def test_prompt_cache_key_reaches_a_mistral_hosted_endpoint(api_base: str) -> None:
request: Final = _transform_with_prompt_cache_key({"api_base": api_base})
assert request["prompt_cache_key"] == "session-1"
@pytest.mark.parametrize("api_base", [AZURE_HOSTED_MISTRAL_API_BASE, "https://mistral.ai.example.com/v1"])
def test_prompt_cache_key_for_another_host_is_rejected_without_drop_params(api_base: str) -> None:
import litellm
with pytest.raises(litellm.UnsupportedParamsError, match="prompt_cache_key") as rejection:
_transform_with_prompt_cache_key({"api_base": api_base})
upstream_host: Final = urlparse(api_base).hostname or api_base
assert upstream_host not in str(rejection.value)
UNSENDABLE_PROMPT_CACHE_KEYS: Final = ("Büro-Sitzung-42", "日本", 123, ["session-1"])
@pytest.mark.parametrize("prompt_cache_key", UNSENDABLE_PROMPT_CACHE_KEYS)
def test_prompt_cache_key_the_mistral_api_cannot_take_is_rejected_without_drop_params(
prompt_cache_key: object,
) -> None:
import litellm
with pytest.raises(litellm.UnsupportedParamsError, match="ASCII"):
_transform_with_prompt_cache_key({"api_base": "https://api.mistral.ai/v1"}, prompt_cache_key)
@pytest.mark.parametrize("prompt_cache_key", UNSENDABLE_PROMPT_CACHE_KEYS)
def test_prompt_cache_key_the_mistral_api_cannot_take_is_dropped_under_drop_params(prompt_cache_key: object) -> None:
request: Final = _transform_with_prompt_cache_key(
{"api_base": "https://api.mistral.ai/v1", "drop_params": True}, prompt_cache_key
)
assert "prompt_cache_key" not in request
assert request["temperature"] == 0.2
def test_prompt_cache_key_with_ascii_punctuation_reaches_the_mistral_api() -> None:
request: Final = _transform_with_prompt_cache_key({}, "team a/session:42")
assert request["prompt_cache_key"] == "team a/session:42"
def test_prompt_cache_key_allowlisted_for_another_purpose_is_still_rejected_for_another_host() -> None:
import litellm
with pytest.raises(litellm.UnsupportedParamsError, match="allowed_openai_params"):
_transform_with_prompt_cache_key(
{"api_base": AZURE_HOSTED_MISTRAL_API_BASE, "allowed_openai_params": ["reasoning_effort"]}
)
def _chat_completion_reply() -> httpx.Response:
return httpx.Response(
200,
json={
"id": "chatcmpl-1",
"object": "chat.completion",
"created": 0,
"model": "mistral-small-latest",
"choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
},
)
def test_allowlisted_prompt_cache_key_reaches_a_gateway_in_front_of_mistral(respx_mock: respx.MockRouter) -> None:
import litellm
gateway_route: Final = respx_mock.post("https://gateway.example.com/mistral/v1/chat/completions").mock(
return_value=_chat_completion_reply()
)
litellm.completion(
model="mistral/mistral-small-latest",
messages=[{"role": "user", "content": "hi"}],
api_base="https://gateway.example.com/mistral/v1",
api_key="gateway-key",
prompt_cache_key="session-1",
allowed_openai_params=["prompt_cache_key"],
max_retries=0,
)
assert json.loads(gateway_route.calls.last.request.content)["prompt_cache_key"] == "session-1"
@pytest.mark.asyncio
async def test_prompt_cache_key_is_dropped_from_an_async_codestral_call_to_another_host(
respx_mock: respx.MockRouter,
) -> None:
import litellm
azure_route: Final = respx_mock.post(f"{AZURE_HOSTED_MISTRAL_API_BASE}/chat/completions").mock(
return_value=_chat_completion_reply()
)
await litellm.acompletion(
model="codestral/codestral-latest",
messages=[{"role": "user", "content": "hi"}],
api_base=AZURE_HOSTED_MISTRAL_API_BASE,
api_key="azure-key",
prompt_cache_key="session-1",
drop_params=True,
max_retries=0,
client=AsyncOpenAI(
api_key="azure-key", base_url=AZURE_HOSTED_MISTRAL_API_BASE, http_client=httpx.AsyncClient()
),
)
assert "prompt_cache_key" not in json.loads(azure_route.calls.last.request.content)
@pytest.mark.asyncio
async def test_prompt_cache_key_rejection_for_an_async_codestral_call_names_codestral(
respx_mock: respx.MockRouter,
) -> None:
import litellm
azure_route: Final = respx_mock.post(f"{AZURE_HOSTED_MISTRAL_API_BASE}/chat/completions").mock(
return_value=_chat_completion_reply()
)
with pytest.raises(litellm.BadRequestError, match=r"codestral does not support parameters: \['prompt_cache_key'\]"):
await litellm.acompletion(
model="codestral/codestral-latest",
messages=[{"role": "user", "content": "hi"}],
api_base=AZURE_HOSTED_MISTRAL_API_BASE,
api_key="azure-key",
prompt_cache_key="session-1",
max_retries=0,
client=AsyncOpenAI(
api_key="azure-key", base_url=AZURE_HOSTED_MISTRAL_API_BASE, http_client=httpx.AsyncClient()
),
)
assert not azure_route.called
@pytest.mark.asyncio
async def test_async_codestral_call_sends_an_uploaded_file_by_its_id(respx_mock: respx.MockRouter) -> None:
import litellm
codestral_api_base: Final = "https://codestral.mistral.ai/v1"
codestral_route: Final = respx_mock.post(f"{codestral_api_base}/chat/completions").mock(
return_value=_chat_completion_reply()
)
await litellm.acompletion(
model="codestral/codestral-latest",
messages=[
{
"role": "user",
"content": [
{"type": "text", "text": "Review this file"},
{"type": "file", "file": {"file_id": "file-12345"}},
],
}
],
api_base=codestral_api_base,
api_key="codestral-key",
max_retries=0,
client=AsyncOpenAI(api_key="codestral-key", base_url=codestral_api_base, http_client=httpx.AsyncClient()),
)
sent_content: Final = json.loads(codestral_route.calls.last.request.content)["messages"][0]["content"]
assert sent_content[1] == {"type": "file", "file_id": "file-12345"}
def test_prompt_cache_key_rejection_for_a_sync_codestral_call_is_an_unsupported_params_error(
respx_mock: respx.MockRouter,
) -> None:
import litellm
azure_route: Final = respx_mock.post(f"{AZURE_HOSTED_MISTRAL_API_BASE}/chat/completions").mock(
return_value=_chat_completion_reply()
)
with pytest.raises(
litellm.UnsupportedParamsError, match=r"^litellm.UnsupportedParamsError: codestral does not support"
):
litellm.completion(
model="codestral/codestral-latest",
messages=[{"role": "user", "content": "hi"}],
api_base=AZURE_HOSTED_MISTRAL_API_BASE,
api_key="azure-key",
prompt_cache_key="session-1",
max_retries=0,
)
assert not azure_route.called
def test_prompt_cache_key_for_another_host_is_dropped_under_drop_params() -> None:
request: Final = _transform_with_prompt_cache_key({"api_base": AZURE_HOSTED_MISTRAL_API_BASE, "drop_params": True})
assert "prompt_cache_key" not in request
assert request["temperature"] == 0.2
def test_prompt_cache_key_for_the_azure_endpoint_set_in_the_environment_is_dropped_under_drop_params(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setenv("MISTRAL_AZURE_API_BASE", "https://my-resource.services.ai.azure.com/openai")
request: Final = _transform_with_prompt_cache_key({"drop_params": True})
assert "prompt_cache_key" not in request
def test_prompt_cache_key_stays_unsupported_for_vertex_hosted_mistral(local_model_cost_map: None) -> None:
import litellm
with pytest.raises(litellm.UnsupportedParamsError):
litellm.get_optional_params(model="mistral-medium-3", custom_llm_provider="vertex_ai", prompt_cache_key="s")
dropped: Final = litellm.get_optional_params(
model="mistral-medium-3", custom_llm_provider="vertex_ai", prompt_cache_key="s", drop_params=True
)
assert "prompt_cache_key" not in dropped

View file

@ -343,6 +343,26 @@ def test_body_carries_no_aws_credentials(
assert client.request_headers["X-Amz-Security-Token"] == "test-session-token"
@pytest.mark.parametrize("model", ["sagemaker_chat/my-endpoint", "sagemaker_nova/my-endpoint"])
def test_inference_component_travels_only_in_the_header(monkeypatch: pytest.MonkeyPatch, model: str):
client = _invoke_sagemaker_chat(monkeypatch, model=model, model_id="my-inference-component")
assert client.request_headers[_INFERENCE_COMPONENT_HEADER] == "my-inference-component"
assert "model_id" not in client.request_body
def test_inference_component_header_keeps_header_keys_that_are_not_strings():
headers = SagemakerChatConfig().validate_environment(
headers={7: "other", "x-amzn-sagemaker-inference-component": "client-component"},
model="my-endpoint",
messages=[],
optional_params={"model_id": "my-inference-component"},
litellm_params={},
)
assert headers == {7: "other", _INFERENCE_COMPONENT_HEADER: "my-inference-component"}
def test_body_model_stays_the_endpoint_name_when_hf_model_name_is_unset(monkeypatch):
"""Without `hf_model_name` the body must keep the model it has today."""
client = _invoke_sagemaker_chat(monkeypatch)

View file

@ -372,6 +372,27 @@ def test_completion_sagemaker_non_stream(
assert response.choices[0].message.content == "SageMaker reply"
@pytest.mark.usefixtures("httpx_transport")
@pytest.mark.usefixtures("fake_provider_credentials")
def test_completion_sagemaker_sends_the_inference_component_only_as_a_header(respx_mock: respx.MockRouter) -> None:
url: Final = "https://runtime.sagemaker.us-west-2.amazonaws.com/endpoints/my-endpoint/invocations"
route: Final = respx_mock.post(url).mock(
return_value=httpx.Response(200, json={"generated_text": "SageMaker reply"})
)
litellm.completion(
model="sagemaker/my-endpoint",
messages=[{"role": "user", "content": "hi"}],
temperature=0.2,
model_id="my-inference-component",
aws_region_name="us-west-2",
num_retries=0,
)
request: Final = route.calls[0].request
assert request.headers["X-Amzn-SageMaker-Inference-Component"] == "my-inference-component"
assert json.loads(request.content) == {"inputs": "hi", "parameters": {"temperature": 0.2}}
@pytest.mark.usefixtures("httpx_transport")
@pytest.mark.usefixtures("fake_provider_credentials")
def test_completion_sagemaker_non_stream_with_aws_params(respx_mock: respx.MockRouter) -> None:

View file

@ -18,6 +18,7 @@ import websockets.exceptions # registers websockets.exceptions on the websocket
import litellm
from litellm.llms.vertex_ai.common_utils import get_vertex_base_url
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import OPENAI_TO_GEMINI_TTS_VOICES
from litellm.llms.vertex_ai.realtime.transformation import VertexAIRealtimeConfig
# ---------------------------------------------------------------------------
@ -179,6 +180,25 @@ def test_vertex_audio_only_live_model_coerces_text_modality_to_audio(
assert setup["generationConfig"]["responseModalities"] == ["AUDIO"]
@pytest.mark.parametrize("drop_params", [False, True])
def test_vertex_realtime_config_carries_drop_params(drop_params, patch_native_audio_cost_map_entry):
from litellm.llms.vertex_ai.realtime.transformation import vertex_realtime_config
cfg: Final = vertex_realtime_config(
_NATIVE_AUDIO_MODEL,
access_token="tok",
resolve_access_token=AsyncMock(return_value="tok"),
project="my-proj",
location="us-central1",
drop_params=drop_params,
)
rejection: Final = cfg.setup_session_update_rejection(
_NATIVE_AUDIO_MODEL, {"modalities": ["text"]}, setup_sent=False
)
assert (rejection is None) is drop_params
def test_vertex_session_update_normalizes_ga_remapped_fields(
patch_native_audio_cost_map_entry,
):
@ -253,6 +273,26 @@ def test_vertex_keeps_turn_detection_tuning_when_auto_response_is_off():
}
@pytest.mark.parametrize(
("session", "disabled"),
[
({"turn_detection": None}, True),
({"type": "realtime", "audio": {"input": {"turn_detection": None}}}, True),
({"type": "realtime", "audio": {"input": {}}}, False),
],
)
def test_vertex_turns_activity_detection_off_for_null_turn_detection(session, disabled):
"""OpenAI's turn_detection null is push-to-talk, so Vertex must wait for the client's activity markers."""
cfg: Final = VertexAIRealtimeConfig(access_token="tok", project="my-proj", location="us-central1")
messages: Final = cfg.transform_realtime_request(
json.dumps({"type": "session.update", "session": session}),
_NATIVE_AUDIO_MODEL,
session_configuration_request=None,
)
setup: Final = json.loads(messages[0])["setup"]
assert setup["realtimeInputConfig"]["automaticActivityDetection"]["disabled"] is disabled
# ---------------------------------------------------------------------------
# Round-trip test: text-in / text-out via RealTimeStreaming
# ---------------------------------------------------------------------------
@ -607,3 +647,73 @@ def test_vertex_native_audio_unmapped_voice_passes_through(patch_native_audio_co
generation_config = json.loads(messages[0])["setup"]["generationConfig"]
assert generation_config["speechConfig"]["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == "Kore"
def test_vertex_session_update_two_output_modalities_collapse_to_audio():
"""Vertex Live rejects a setup naming two response modalities, so a GA
["text", "audio"] request must be sent as AUDIO only."""
cfg = VertexAIRealtimeConfig(access_token="tok", project="my-proj", location="global")
session_update = {
"type": "session.update",
"session": {"output_modalities": ["text", "audio"], "instructions": "Be brief."},
}
messages = cfg.transform_realtime_request(
json.dumps(session_update),
"gemini-live-2.5-flash",
session_configuration_request=None,
)
assert len(messages) == 1
assert json.loads(messages[0])["setup"]["generationConfig"]["responseModalities"] == ["AUDIO"]
@pytest.mark.parametrize(
("session", "param"),
[
({"voice": "Nova"}, "session.voice"),
({"type": "realtime", "audio": {"output": {"voice": "Nova"}}}, "session.audio.output.voice"),
],
)
def test_vertex_rejects_an_undocumented_voice_on_the_session_update(session, param):
cfg: Final = VertexAIRealtimeConfig(access_token="tok", project="my-proj", location="us-central1")
rejection: Final = cfg.setup_session_update_rejection("gemini-live-2.5-flash", session, setup_sent=False)
assert rejection is not None
assert (rejection.code, rejection.param) == ("invalid_value", param)
@pytest.mark.parametrize("voice", ["Kore", "alloy"])
def test_vertex_accepts_documented_and_openai_stock_voices_on_the_session_update(voice):
cfg: Final = VertexAIRealtimeConfig(access_token="tok", project="my-proj", location="us-central1")
assert cfg.setup_session_update_rejection("gemini-live-2.5-flash", {"voice": voice}, setup_sent=False) is None
def test_vertex_drops_an_undocumented_voice_with_drop_params():
cfg: Final = VertexAIRealtimeConfig(access_token="tok", project="my-proj", location="us-central1", drop_params=True)
session_update: Final = {"type": "session.update", "session": {"voice": "Nova", "instructions": "Be brief."}}
assert cfg.setup_session_update_rejection("gemini-live-2.5-flash", session_update["session"], setup_sent=False) is None
messages: Final = cfg.transform_realtime_request(
json.dumps(session_update), "gemini-live-2.5-flash", session_configuration_request=None
)
assert "speechConfig" not in json.loads(messages[0])["setup"].get("generationConfig", {})
@pytest.mark.parametrize("drop_params", [False, True])
@pytest.mark.parametrize("voice", ["fable", "onyx"])
def test_vertex_sends_the_documented_voice_an_openai_tts_voice_maps_to(voice, drop_params):
cfg: Final = VertexAIRealtimeConfig(
access_token="tok", project="my-proj", location="us-central1", drop_params=drop_params
)
session_update: Final = {"type": "session.update", "session": {"voice": voice}}
assert cfg.setup_session_update_rejection("gemini-live-2.5-flash", session_update["session"], setup_sent=False) is None
messages: Final = cfg.transform_realtime_request(
json.dumps(session_update), "gemini-live-2.5-flash", session_configuration_request=None
)
assert json.loads(messages[0])["setup"]["generationConfig"]["speechConfig"] == {
"voiceConfig": {"prebuiltVoiceConfig": {"voiceName": OPENAI_TO_GEMINI_TTS_VOICES[voice]}}
}

View file

@ -144,3 +144,16 @@ def _mcp_request_ctx():
return ServerRequestContext(**kwargs)
return _mcp_request_ctx
@pytest.fixture
def mcp_code_runtime() -> None:
from importlib.metadata import PackageNotFoundError, version
from litellm.proxy.mcp_server.code_runtime import check_runtime
try:
version("quickjs-ng")
except PackageNotFoundError:
pytest.skip("Requires the optional litellm[mcp-code-mode] runtime")
check_runtime()

View file

@ -0,0 +1,418 @@
import asyncio
import subprocess
from collections.abc import Mapping, Sequence
from datetime import datetime
from typing import Final
import anyio
import pytest
from anyio.abc import ByteReceiveStream, ByteSendStream, Process
from mcp.types import TextContent, Tool
from pydantic import JsonValue
import litellm
from litellm.caching.caching import DualCache
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy import proxy_server
from litellm.proxy._types import LiteLLM_ObjectPermissionTable, UserAPIKeyAuth
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3, get_request_stash
from litellm.proxy.mcp_server.code_execution import CODE_TOOL_NAME, run_code_tool
from litellm.proxy.mcp_server.code_worker import Call, Done, Reply, Start, encode_message, load_json
from litellm.proxy.mcp_server.contracts import OperationContext
from litellm.proxy.mcp_server.mcp_server_manager import global_mcp_server_manager
from litellm.proxy.mcp_server.result_conversion import INPUT_REQUIRED_UNSUPPORTED_MESSAGE
from litellm.proxy.mcp_server.tool_registry import global_mcp_tool_registry
from litellm.proxy.mcp_server.tool_search import mcp_proxy_tool_id, with_mcp_proxy_identity
from litellm.proxy.utils import ProxyLogging
from litellm.types.mcp import MCPPostCallResponseObject, MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPServer
from litellm.types.utils import CallTypesLiteral
pytestmark: Final = pytest.mark.usefixtures("mcp_code_runtime")
class ScriptedProcess(Process, ByteReceiveStream, ByteSendStream):
"""Only the external process transport is replaced; each reply comes from LiteLLM."""
def __init__(self, calls: tuple[Call, ...]) -> None:
self.calls = calls
self.started: Start | None = None
self.replies: tuple[Reply, ...] = ()
self.stopped = False
self.sender, self.receiver = anyio.create_memory_object_stream[bytes](1)
async def send(self, item: bytes) -> None:
if self.started is None:
self.started = Start.model_validate_json(item)
else:
reply: Final = Reply.model_validate_json(item)
assert reply.id == self.calls[len(self.replies)].id
self.replies = (*self.replies, reply)
message: Final = (
self.calls[len(self.replies)]
if len(self.replies) < len(self.calls)
else Done(
output=tuple(
dict(reply.value) if reply.value is not None else {"error": reply.error} for reply in self.replies
)
)
)
await self.sender.send(encode_message(message, 2 * 1024 * 1024))
async def receive(self, max_bytes: int = 65536) -> bytes:
data: Final = await self.receiver.receive()
assert len(data) <= max_bytes
return data
async def aclose(self) -> None:
await self.sender.aclose()
await self.receiver.aclose()
async def wait(self) -> int:
assert self.stopped
return -9
def kill(self) -> None:
self.stopped = True
def terminate(self) -> None:
self.kill()
def send_signal(self, signal: int) -> None:
self.kill()
@property
def pid(self) -> int:
return 0
@property
def returncode(self) -> int | None:
return -9 if self.stopped else None
@property
def stdin(self) -> ByteSendStream:
return self
@property
def stdout(self) -> ByteReceiveStream:
return self
@property
def stderr(self) -> None:
return None
def install_transport(monkeypatch: pytest.MonkeyPatch, calls: tuple[Call, ...]) -> ScriptedProcess:
process: Final = ScriptedProcess(calls)
async def open_process(command: Sequence[str], *, cwd: str, env: Mapping[str, str], stderr: int) -> Process:
assert "-I" in command and "-u" in command
assert set(env) <= {"SYSTEMROOT"}
assert stderr == subprocess.DEVNULL
assert cwd
return process
monkeypatch.setattr(anyio, "open_process", open_process)
return process
class ObservedCalls(CustomLogger):
def __init__(self) -> None:
self.names: tuple[str, ...] = ()
self.traces: tuple[object, ...] = ()
self.responses: tuple[MCPPostCallResponseObject, ...] = ()
self.log_traces: tuple[object, ...] = ()
self.handler_started = asyncio.Event()
async def async_pre_call_hook(
self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict[str, object], call_type: CallTypesLiteral
) -> dict[str, object]:
name: Final = data.get("name") or data.get("mcp_tool_name")
if isinstance(name, str):
self.names = (*self.names, name)
self.traces = (*self.traces, data.get("litellm_trace_id"))
return data
async def async_post_mcp_tool_call_hook(
self,
kwargs: dict[str, object],
response_obj: MCPPostCallResponseObject,
start_time: datetime,
end_time: datetime,
) -> MCPPostCallResponseObject:
self.responses = (*self.responses, response_obj)
self.log_traces = (*self.log_traces, kwargs.get("litellm_trace_id"))
return response_obj
@pytest.fixture
def registered_gateway(monkeypatch: pytest.MonkeyPatch) -> tuple[OperationContext, ObservedCalls, list[str]]:
observed: Final = ObservedCalls()
logger: Final = ProxyLogging(user_api_key_cache=UserApiKeyCache())
limiter: Final = PROXY_MaxParallelRequestsHandler_v3(logger.internal_usage_cache)
monkeypatch.setattr(
litellm, "mcp_code_mode", {"enabled": True, "excluded_tools": {"allowed-blocked": "requires manual approval"}}
)
monkeypatch.setattr(litellm, "callbacks", [limiter, observed])
monkeypatch.setattr(proxy_server, "proxy_logging_obj", logger)
manager: Final = global_mcp_server_manager
monkeypatch.setattr(manager, "registry", {})
monkeypatch.setattr(manager, "tool_name_to_mcp_server_name_mapping", {})
monkeypatch.setattr(global_mcp_tool_registry, "tools", {})
received: Final[list[str]] = []
async def echo(message: str) -> str:
stash: Final = get_request_stash()
assert stash is not None and stash.borrowed_parallel_keys
received.append(message)
if message == "wait":
observed.handler_started.set()
await anyio.sleep_forever()
if message == "continuation documentation":
return INPUT_REQUIRED_UNSUPPORTED_MESSAGE
if message == "upstream failure":
raise ValueError("upstream-specific detail")
return message
for server in (
MCPServer(
server_id=alias,
name=alias,
alias=alias,
transport=MCPTransport.http,
url="https://example.invalid/mcp",
spec_path="registered-openapi.json",
)
for alias in ("allowed", "denied")
):
manager.registry[server.name] = server
manager.tool_name_to_mcp_server_name_mapping[server.name + "-echo"] = server.name
global_mcp_tool_registry.register_tool( # pyright: ignore[reportUnknownMemberType] # Legacy registry leaves handler Callable unparameterized.
name=server.name + "-echo",
description="Echo",
input_schema={"type": "object", "properties": {"message": {"type": "string"}}},
handler=echo,
server_id=server.server_id,
)
global_mcp_tool_registry.register_tool( # pyright: ignore[reportUnknownMemberType] # Legacy registry leaves handler Callable unparameterized.
name="allowed-blocked",
description="Blocked",
input_schema={"type": "object"},
handler=echo,
server_id="allowed",
)
manager.tool_name_to_mcp_server_name_mapping["allowed-blocked"] = "allowed"
auth: Final = UserAPIKeyAuth(
api_key="code-mode-component",
max_parallel_requests=1,
object_permission=LiteLLM_ObjectPermissionTable(
object_permission_id="code-mode-component", mcp_servers=["allowed"]
),
)
return (
OperationContext(_caller=auth, raw_headers={"x-litellm-trace-id": "code-component-trace"}),
observed,
received,
)
def error_text(reply: Reply) -> str:
assert reply.value is not None and reply.value["isError"] is True
content: Final = reply.value["content"]
assert isinstance(content, list) and isinstance(content[0], dict)
return str(content[0]["text"])
@pytest.mark.asyncio
@pytest.mark.parametrize("opaque", [False, True])
async def test_authenticated_coordinator_preserves_results_and_rejects_unauthorized_tools(
monkeypatch: pytest.MonkeyPatch,
registered_gateway: tuple[OperationContext, ObservedCalls, list[str]],
opaque: bool,
) -> None:
context, observed, received = registered_gateway
reference: Final = (
mcp_proxy_tool_id(with_mcp_proxy_identity(Tool(name="allowed-echo", input_schema={}), "allowed"))
if opaque
else "allowed-echo"
)
process: Final = install_transport(
monkeypatch,
(
Call(id=0, tool_id=reference, arguments={"message": "kept"}),
Call(id=1, tool_id="denied-echo", arguments={"message": "must not run"}),
Call(id=2, tool_id=CODE_TOOL_NAME, arguments={}),
Call(id=3, tool_id="allowed-blocked", arguments={"message": "must not run"}),
Call(id=4, tool_id=reference, arguments={"message": "upstream failure"}),
),
)
result: Final = await run_code_tool({"code": "/* protocol fixture */"}, context)
assert not result.is_error
assert process.stopped
assert process.started is not None and process.started.code == "/* protocol fixture */"
assert received == ["kept", "upstream failure"]
assert process.replies[0].value == {
"content": [{"type": "text", "text": "kept"}],
"isError": False,
"resultType": "complete",
}
for reply, reason in zip(process.replies[1:4], ("Unknown", "Nested", "requires manual approval"), strict=True):
assert reason in error_text(reply)
assert process.replies[4].value == {
"content": [{"type": "text", "text": "Error: upstream-specific detail"}],
"isError": True,
"resultType": "complete",
}
assert isinstance(result.content[0], TextContent)
body: Final = load_json(result.content[0].text)
assert isinstance(body, dict) and body["calls_started"] == 5
assert body["output"] == [
dict(reply.value) if reply.value is not None else {"error": reply.error} for reply in process.replies
]
assert observed.responses
assert observed.log_traces == ("code-component-trace",) * len(observed.responses)
assert "code-component-trace" in observed.traces
@pytest.mark.asyncio
async def test_invalid_script_releases_the_outer_admission(
monkeypatch: pytest.MonkeyPatch, registered_gateway: tuple[OperationContext, ObservedCalls, list[str]]
) -> None:
context, observed, received = registered_gateway
process: Final = install_transport(monkeypatch, ())
invalid: Final = await run_code_tool({"code": 42}, context)
assert invalid.is_error and isinstance(invalid.content[0], TextContent)
assert "Invalid code request" in invalid.content[0].text
assert process.started is None and received == []
result: Final = await run_code_tool({"code": "/* protocol fixture */"}, context)
assert not result.is_error
assert CODE_TOOL_NAME in observed.names
assert process.stopped
@pytest.mark.asyncio
@pytest.mark.parametrize("failure", ["exception", "unsafe number"])
async def test_runtime_sanitizes_callback_failures(monkeypatch: pytest.MonkeyPatch, failure: str) -> None:
from litellm.proxy.mcp_server.code_runtime import execute_code
async def upstream(tool_id: str, arguments: dict[str, JsonValue]) -> dict[str, JsonValue]:
assert tool_id == "echo" and arguments == {"n": 1}
if failure == "exception":
raise RuntimeError("upstream-specific detail")
return {"large": 2**53}
process: Final = install_transport(monkeypatch, (Call(id=0, tool_id="echo", arguments={"n": 1}),))
result: Final = await execute_code("/* protocol fixture */", upstream)
expected: Final = (
"Tool execution failed"
if failure == "exception"
else "Tool result is too large or contains unsupported JSON numbers"
)
assert process.replies == (Reply(id=0, error=expected),)
assert result.output == ({"error": expected},)
assert result.calls_started == 1 and process.stopped
@pytest.mark.asyncio
@pytest.mark.parametrize("boundary", ["call count", "output bytes"])
async def test_runtime_enforces_limits_on_worker_messages(monkeypatch: pytest.MonkeyPatch, boundary: str) -> None:
from dataclasses import replace
from litellm.proxy.mcp_server.code_runtime import CodeExecutionLimits, execute_code
limits: Final = (
replace(CodeExecutionLimits(), max_calls=1)
if boundary == "call count"
else replace(CodeExecutionLimits(), output_bytes=8)
)
first: Final = Call(id=0, tool_id="echo", arguments={"n": 1})
calls: Final = (first, Call(id=1, tool_id="echo", arguments={"n": 2})) if boundary == "call count" else (first,)
process: Final = install_transport(monkeypatch, calls)
received: Final[list[dict[str, JsonValue]]] = []
async def upstream(tool_id: str, arguments: dict[str, JsonValue]) -> dict[str, JsonValue]:
assert tool_id == "echo"
received.append(arguments)
return {"text": "value exceeds eight bytes"}
result: Final = await execute_code("/* protocol fixture */", upstream, limits)
assert received == [{"n": 1}]
expected: Final = "tool-call limit" if boundary == "call count" else "output exceeds its byte limit"
assert result.error is not None and expected in result.error
assert result.calls_started == 1 and result.output == ()
assert process.stopped
@pytest.mark.asyncio
@pytest.mark.parametrize(
"frame", [b"not JSON\n", b'{"kind":"unknown"}\n', b'{"kind":"done","output":[9007199254740993]}\n']
)
async def test_runtime_rejects_invalid_worker_protocol(monkeypatch: pytest.MonkeyPatch, frame: bytes) -> None:
from litellm.proxy.mcp_server.code_runtime import execute_code
process: Final = install_transport(monkeypatch, ())
async def receive(max_bytes: int = 65536) -> bytes:
assert len(frame) <= max_bytes
return frame
async def upstream(tool_id: str, arguments: dict[str, JsonValue]) -> dict[str, JsonValue]:
pytest.fail("Malformed worker messages must not dispatch tools")
monkeypatch.setattr(process, "receive", receive)
result: Final = await execute_code("/* protocol fixture */", upstream)
assert result.error == "Code execution worker failed; upstream actions may have completed"
assert result.calls_started == 0 and result.output == ()
assert process.stopped
@pytest.mark.asyncio
async def test_cancelling_coordinator_releases_parent_and_child_admission(
monkeypatch: pytest.MonkeyPatch, registered_gateway: tuple[OperationContext, ObservedCalls, list[str]]
) -> None:
context, observed, received = registered_gateway
process: Final = install_transport(
monkeypatch, (Call(id=0, tool_id="allowed-echo", arguments={"message": "wait"}),)
)
request: Final = asyncio.create_task(run_code_tool({"code": "/* protocol fixture */"}, context))
entered: Final = asyncio.create_task(observed.handler_started.wait())
try:
completed, _ = await asyncio.wait((request, entered), return_when=asyncio.FIRST_COMPLETED)
assert entered in completed, "Request exited before reaching the registered handler"
assert entered.result()
request.cancel()
with pytest.raises(asyncio.CancelledError):
await request
finally:
for task in (request, entered):
if not task.done():
task.cancel()
await asyncio.gather(request, entered, return_exceptions=True)
assert process.stopped and process.replies == ()
assert received == ["wait"]
replacement: Final = install_transport(
monkeypatch, (Call(id=0, tool_id="allowed-echo", arguments={"message": "after cancellation"}),)
)
result: Final = await run_code_tool({"code": "/* protocol fixture */"}, context)
assert not result.is_error
assert received == ["wait", "after cancellation"]
assert replacement.replies[0].value is not None and replacement.replies[0].value["isError"] is False
assert replacement.stopped
@pytest.mark.asyncio
async def test_successful_continuation_text_is_not_reclassified_as_an_error(
monkeypatch: pytest.MonkeyPatch, registered_gateway: tuple[OperationContext, ObservedCalls, list[str]]
) -> None:
context, _, received = registered_gateway
process: Final = install_transport(
monkeypatch, (Call(id=0, tool_id="allowed-echo", arguments={"message": "continuation documentation"}),)
)
result: Final = await run_code_tool({"code": "/* protocol fixture */"}, context)
assert not result.is_error and received == ["continuation documentation"]
assert process.replies[0].value == {
"content": [{"type": "text", "text": INPUT_REQUIRED_UNSUPPORTED_MESSAGE}],
"isError": False,
"resultType": "complete",
}

View file

@ -7,11 +7,16 @@ import pytest
from litellm.proxy.mcp_server.code_worker import Reply, encode_message, load_json
def test_code_mode_rejects_archived_quickjs_package(monkeypatch: pytest.MonkeyPatch) -> None:
@pytest.mark.parametrize("ng_installed", [False, True])
def test_code_mode_rejects_archived_quickjs_package(monkeypatch: pytest.MonkeyPatch, ng_installed: bool) -> None:
from litellm.proxy.mcp_server import code_runtime
monkeypatch.setattr(code_runtime.metadata, "distribution", lambda name: None)
with pytest.raises(ValueError, match="conflicting quickjs package"):
def installed(name: str) -> None:
if name == "quickjs-ng" and not ng_installed:
raise code_runtime.metadata.PackageNotFoundError(name)
monkeypatch.setattr(code_runtime.metadata, "distribution", installed)
with pytest.raises(ValueError, match=r"conflicting quickjs package.*separate environment"):
code_runtime.check_runtime()

View file

@ -0,0 +1,112 @@
import sys
from dataclasses import replace
from io import BytesIO, TextIOWrapper
from typing import Final
import pytest
from pydantic import JsonValue
from litellm.proxy.mcp_server.code_worker import (
WORKER_MESSAGE,
Call,
CodeExecutionLimits,
Done,
Reply,
Start,
encode_message,
main,
)
pytestmark: Final = pytest.mark.usefixtures("mcp_code_runtime")
def run_worker(monkeypatch: pytest.MonkeyPatch, request: bytes) -> list[Call | Done]:
source: Final = TextIOWrapper(BytesIO(request))
destination: Final = TextIOWrapper(BytesIO())
monkeypatch.setattr(sys, "stdin", source)
monkeypatch.setattr(sys, "stdout", destination)
main()
return [WORKER_MESSAGE.validate_json(line) for line in destination.buffer.getvalue().splitlines()]
def test_native_worker_chains_full_tool_results(monkeypatch: pytest.MonkeyPatch) -> None:
limits: Final = CodeExecutionLimits()
request: Final = Start(
code="""
const first = await tools.call_tool({tool_id: 'lookup', arguments: {id: '9007199254740993'}});
text(await tools.call_tool({tool_id: 'echo', arguments: first.structuredContent}));
""",
limits=limits,
)
envelope: Final[dict[str, JsonValue]] = {
"content": [{"type": "text", "text": "kept"}],
"structuredContent": {"id": "9007199254740993"},
"isError": True,
}
messages: Final = run_worker(
monkeypatch,
b"".join(
encode_message(message, limits.message_bytes)
for message in (request, Reply(id=0, value=envelope), Reply(id=1, value=envelope))
),
)
assert messages == [
Call(id=0, tool_id="lookup", arguments={"id": "9007199254740993"}),
Call(id=1, tool_id="echo", arguments={"id": "9007199254740993"}),
Done(output=(envelope,)),
]
@pytest.mark.parametrize(
("code", "error"),
[
('text("before"); throw new Error("script failed")', "script failed"),
("text(BigInt(1))", "json values"),
("text(9007199254740992)", "integer"),
("text(() => 1)", "json values"),
],
)
def test_native_worker_reports_script_errors(monkeypatch: pytest.MonkeyPatch, code: str, error: str) -> None:
limits: Final = CodeExecutionLimits()
messages: Final = run_worker(monkeypatch, encode_message(Start(code=code, limits=limits), limits.message_bytes))
result: Final = messages[-1]
assert isinstance(result, Done)
assert result.error is not None and error in result.error.lower()
assert result.output == (("before",) if "before" in code else ())
def test_native_worker_rejects_oversized_output(monkeypatch: pytest.MonkeyPatch) -> None:
limits: Final = replace(CodeExecutionLimits(), output_bytes=32)
result: Final = run_worker(
monkeypatch, encode_message(Start(code='text("猫".repeat(50))', limits=limits), limits.message_bytes)
)[-1]
assert isinstance(result, Done)
assert result.error is not None and "byte limit" in result.error
assert result.output == ()
@pytest.mark.parametrize(
"reply", [b"", b"not json\n", b'{"id":0,"value":{"number":NaN}}\n', b'{"id":0,"value":{"id":9007199254740993}}\n']
)
def test_native_worker_fails_closed_on_bad_parent_reply(monkeypatch: pytest.MonkeyPatch, reply: bytes) -> None:
limits: Final = CodeExecutionLimits()
request: Final = Start(code="text(await tools.call_tool({tool_id:'lookup', arguments:{}}))", limits=limits)
messages: Final = run_worker(monkeypatch, encode_message(request, limits.message_bytes) + reply)
assert messages[0] == Call(id=0, tool_id="lookup", arguments={})
assert messages[-1] == Done(error="JavaScript worker failed")
def test_native_worker_rejects_oversized_start(monkeypatch: pytest.MonkeyPatch) -> None:
assert run_worker(monkeypatch, b"x" * (CodeExecutionLimits().message_bytes + 1)) == [
Done(error="JavaScript worker failed")
]
def test_native_worker_limits_outgoing_tool_arguments(monkeypatch: pytest.MonkeyPatch) -> None:
limits: Final = replace(CodeExecutionLimits(), message_bytes=128)
request: Final = Start(
code="await tools.call_tool({tool_id:'echo', arguments:{text:'x'.repeat(256)}})", limits=limits
)
messages: Final = run_worker(monkeypatch, encode_message(request, CodeExecutionLimits().message_bytes))
assert len(messages) == 1 and isinstance(messages[0], Done)
assert messages[0].error is not None

View file

@ -1,5 +1,6 @@
import asyncio
import copy
import io
import json
import logging
from typing import Final
@ -1564,3 +1565,84 @@ async def test_a_failing_health_probe_leaves_the_deployment_serving_requests(mon
assert "error" in probe
assert reply.choices[0].message.content == "ok"
@pytest.mark.parametrize(
("component_params", "probed_component"),
[
({"extra_headers": {"X-Amzn-SageMaker-Inference-Component": "header-component"}}, "header-component"),
(
{"model_id": "own-component", "extra_headers": {"X-Amzn-SageMaker-Inference-Component": "header-component"}},
"own-component",
),
(
{"model_id": "own-component", "extra_headers": {1: "other", "X-Amzn-SageMaker-Inference-Component": "x"}},
"own-component",
),
(
{"extra_headers": {1: "other", "X-Amzn-SageMaker-Inference-Component": "header-component"}},
"header-component",
),
(
{
"configurable_clientside_auth_params": ["model_id"],
"extra_headers": {"X-Amzn-SageMaker-Inference-Component": "header-component"},
},
None,
),
],
ids=[
"header_pinned",
"own_model_id_beats_the_header",
"own_model_id_beside_a_non_string_header_key",
"header_pinned_beside_a_non_string_header_key",
"client_configurable_model_id_takes_no_header_pin",
],
)
@pytest.mark.asyncio
async def test_run_model_health_check_probes_a_legacy_sagemaker_embedding_at_the_deployments_component(
component_params: dict[str, object], probed_component: str | None
) -> None:
embedding_reply: Final = {"object": "list", "data": [{"object": "embedding", "index": 0, "embedding": [0.1]}]}
invoke_endpoint: Final = MagicMock(return_value={"Body": io.BytesIO(json.dumps(embedding_reply).encode())})
client: Final = MagicMock(invoke_endpoint=invoke_endpoint)
deployment: Final = {
"litellm_params": {
"model": "sagemaker/openai/embed-endpoint",
"aws_region_name": "us-east-1",
"aws_access_key_id": "test-key",
"aws_secret_access_key": "test-secret",
**component_params,
},
"model_info": {"mode": "embedding"},
}
with patch("boto3.Session", return_value=MagicMock(client=MagicMock(return_value=client))):
result: Final = await hc_module._run_model_health_check(deployment)
assert "error" not in result, result
assert invoke_endpoint.call_count == 1
assert invoke_endpoint.call_args.kwargs.get("InferenceComponentName") == probed_component
@pytest.mark.asyncio
async def test_run_model_health_check_reports_why_a_string_configurable_clientside_auth_params_fails() -> None:
invoke_endpoint: Final = MagicMock()
deployment: Final = {
"litellm_params": {
"model": "sagemaker/openai/embed-endpoint",
"model_id": "own-component",
"aws_region_name": "us-east-1",
"aws_access_key_id": "test-key",
"aws_secret_access_key": "test-secret",
},
"model_info": {"mode": "embedding", "health_check_params": {"configurable_clientside_auth_params": "model_id"}},
}
client: Final = MagicMock(invoke_endpoint=invoke_endpoint)
with patch("boto3.Session", return_value=MagicMock(client=MagicMock(return_value=client))):
result: Final = await hc_module._run_model_health_check(deployment)
assert "configurable_clientside_auth_params" in result["error"]
invoke_endpoint.assert_not_called()

View file

@ -4,6 +4,7 @@ import copy
import functools
import gc
import itertools
import io
import json
import linecache
import logging
@ -31523,6 +31524,224 @@ def test_router_keeps_supported_fallback_shapes(
assert router.fallbacks == expected_fallbacks
_SAGEMAKER_ROUTER_CREDENTIALS: Final = MappingProxyType(
{"aws_region_name": "us-east-1", "aws_access_key_id": "test-key", "aws_secret_access_key": "test-secret"}
)
_PINNED_COMPONENT: Final = MappingProxyType({"model_id": "deployment-component"})
_HEADER_PINNED_COMPONENT: Final = MappingProxyType(
{"extra_headers": {"x-amzn-sagemaker-inference-component": "deployment-component"}}
)
_INFERENCE_COMPONENT_PRECEDENCE: Final = pytest.mark.parametrize(
("deployment_params", "sent_component"),
[
(_PINNED_COMPONENT, "deployment-component"),
(_HEADER_PINNED_COMPONENT, "deployment-component"),
(MappingProxyType({}), "client-component"),
(
MappingProxyType({**_PINNED_COMPONENT, "configurable_clientside_auth_params": ["model_id"]}),
"client-component",
),
],
ids=[
"deployment_pins_the_component",
"deployment_pins_the_component_with_a_header",
"deployment_leaves_the_component_to_the_client",
"deployment_lets_the_client_override_its_component",
],
)
def _sagemaker_router(
model: str,
deployment_params: Mapping[str, object],
mode: str,
router_defaults: Mapping[str, str] = MappingProxyType({}),
model_name: str = "sagemaker-model",
) -> Router:
return Router(
model_list=[
{
"model_name": model_name,
"litellm_params": {"model": model, **_SAGEMAKER_ROUTER_CREDENTIALS, **deployment_params},
"model_info": {"mode": mode},
}
],
default_litellm_params=dict(router_defaults),
)
_CHAT_COMPLETION_REPLY: Final = MappingProxyType(
{
"id": "chatcmpl-1",
"object": "chat.completion",
"created": 1,
"model": "served-model",
"choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
}
)
def _sagemaker_chat_route(respx_mock: respx.MockRouter) -> respx.Route:
return respx_mock.post(
"https://runtime.sagemaker.us-east-1.amazonaws.com/endpoints/chat-endpoint/invocations"
).respond(json=dict(_CHAT_COMPLETION_REPLY))
@_INFERENCE_COMPONENT_PRECEDENCE
@pytest.mark.asyncio
async def test_a_client_model_id_cannot_move_a_sagemaker_chat_call_off_the_deployments_component(
respx_mock: respx.MockRouter,
_router_httpx_transport: None,
deployment_params: Mapping[str, object],
sent_component: str,
) -> None:
route: Final = _sagemaker_chat_route(respx_mock)
router: Final = _sagemaker_router("sagemaker_chat/chat-endpoint", deployment_params, "chat")
await router.acompletion(
model="sagemaker-model", messages=[{"role": "user", "content": "hi"}], model_id="client-component"
)
assert route.calls.last.request.headers["X-Amzn-SageMaker-Inference-Component"] == sent_component
@_INFERENCE_COMPONENT_PRECEDENCE
def test_a_client_model_id_cannot_move_a_sync_sagemaker_text_completion_off_the_deployments_component(
respx_mock: respx.MockRouter, deployment_params: Mapping[str, object], sent_component: str
) -> None:
route: Final = _sagemaker_chat_route(respx_mock)
router: Final = _sagemaker_router(
"sagemaker_chat/chat-endpoint", deployment_params, "completion", model_name="sagemaker_chat/chat-endpoint"
)
router.text_completion(model="sagemaker_chat/chat-endpoint", prompt="hi", model_id="client-component")
assert route.calls.last.request.headers["X-Amzn-SageMaker-Inference-Component"] == sent_component
@pytest.mark.parametrize(
"client_header",
[
MappingProxyType({"extra_headers": {"x-amzn-sagemaker-inference-component": "client-component"}}),
MappingProxyType({"headers": {"x-amzn-sagemaker-inference-component": "client-component"}}),
MappingProxyType(
{
"provider_specific_header": {
"custom_llm_provider": "sagemaker_chat",
"extra_headers": {"x-amzn-sagemaker-inference-component": "client-component"},
}
}
),
],
ids=["extra_headers", "headers", "provider_specific_header"],
)
@pytest.mark.parametrize(
"deployment_pin", [_PINNED_COMPONENT, _HEADER_PINNED_COMPONENT], ids=["model_id_pin", "header_pin"]
)
@pytest.mark.asyncio
async def test_a_client_inference_component_header_cannot_move_a_pinned_sagemaker_chat_call(
respx_mock: respx.MockRouter,
_router_httpx_transport: None,
client_header: Mapping[str, object],
deployment_pin: Mapping[str, object],
) -> None:
route: Final = _sagemaker_chat_route(respx_mock)
router: Final = _sagemaker_router("sagemaker_chat/chat-endpoint", deployment_pin, "chat")
await router.acompletion(model="sagemaker-model", messages=[{"role": "user", "content": "hi"}], **client_header)
assert route.calls.last.request.headers["X-Amzn-SageMaker-Inference-Component"] == "deployment-component"
@pytest.mark.parametrize(
"request_component",
[MappingProxyType({}), MappingProxyType({"model_id": "deployment-component"})],
ids=["request_without_model_id", "request_with_the_pinned_model_id"],
)
@pytest.mark.asyncio
async def test_a_router_default_model_id_cannot_move_a_sagemaker_call_off_the_deployments_component(
respx_mock: respx.MockRouter, _router_httpx_transport: None, request_component: Mapping[str, str]
) -> None:
route: Final = _sagemaker_chat_route(respx_mock)
router: Final = _sagemaker_router(
"sagemaker_chat/chat-endpoint",
_PINNED_COMPONENT,
"chat",
router_defaults=MappingProxyType({"model_id": "default-component"}),
)
await router.acompletion(model="sagemaker-model", messages=[{"role": "user", "content": "hi"}], **request_component)
assert route.calls.last.request.headers["X-Amzn-SageMaker-Inference-Component"] == "deployment-component"
@_INFERENCE_COMPONENT_PRECEDENCE
@pytest.mark.asyncio
async def test_a_client_model_id_cannot_move_a_sagemaker_embedding_call_off_the_deployments_component(
deployment_params: Mapping[str, object], sent_component: str
) -> None:
component_embeddings: Final = MappingProxyType({"deployment-component": [0.1], "client-component": [0.2]})
def invoke_endpoint(**request: str | bytes) -> Mapping[str, io.BytesIO]:
served: Final = component_embeddings[str(request["InferenceComponentName"])]
return {
"Body": io.BytesIO(
json.dumps({"object": "list", "data": [{"object": "embedding", "index": 0, "embedding": served}]}).encode()
)
}
client: Final = MagicMock(invoke_endpoint=invoke_endpoint)
router: Final = _sagemaker_router("sagemaker/openai/embed-endpoint", deployment_params, "embedding")
with patch("boto3.Session", return_value=MagicMock(client=MagicMock(return_value=client))):
response: Final = await router.aembedding(model="sagemaker-model", input=["hello"], model_id="client-component")
assert response.data[0]["embedding"] == component_embeddings[sent_component]
@pytest.mark.asyncio
async def test_a_header_pin_holds_when_the_deployment_names_sagemaker_through_custom_llm_provider(
respx_mock: respx.MockRouter, _router_httpx_transport: None
) -> None:
route: Final = _sagemaker_chat_route(respx_mock)
router: Final = _sagemaker_router(
"chat-endpoint", MappingProxyType({**_HEADER_PINNED_COMPONENT, "custom_llm_provider": "sagemaker_chat"}), "chat"
)
await router.acompletion(
model="sagemaker-model", messages=[{"role": "user", "content": "hi"}], model_id="client-component"
)
assert route.calls.last.request.headers["X-Amzn-SageMaker-Inference-Component"] == "deployment-component"
@pytest.mark.asyncio
async def test_a_component_header_on_a_non_sagemaker_deployment_leaves_the_clients_model_id_alone(
respx_mock: respx.MockRouter, _router_httpx_transport: None
) -> None:
route: Final = respx_mock.post("https://gateway.local/v1/chat/completions").respond(json=dict(_CHAT_COMPLETION_REPLY))
router: Final = Router(
model_list=[
{
"model_name": "gateway-model",
"litellm_params": {
"model": "openai/served-model",
"api_base": "https://gateway.local/v1",
"api_key": "test-key",
"extra_headers": {"X-Amzn-SageMaker-Inference-Component": "gateway-component"},
},
}
]
)
await router.acompletion(model="gateway-model", messages=[{"role": "user", "content": "hi"}], model_id="client-value")
assert json.loads(route.calls.last.request.content)["model_id"] == "client-value"
def test_router_warns_when_a_deployment_enforces_endpoints_without_a_list(caplog: pytest.LogCaptureFixture) -> None:
with caplog.at_level(logging.WARNING, logger="LiteLLM Router"):
router: Final = Router(

View file

@ -173,6 +173,24 @@ describe("RealtimePlayground", () => {
expect(update.session.type).toBe("realtime");
});
it("starts recording on the configuration the session was created with", async () => {
const user = userEvent.setup();
render(<RealtimePlayground {...props} />);
await connect(user);
await act(async () => {
latestSocket().emit({ type: "session.created" });
});
await user.click(screen.getByTitle("Start recording"));
const updates = latestSocket()
.sent.map((raw) => JSON.parse(raw))
.filter((payload) => payload.type === "session.update");
expect(updates).toHaveLength(2);
expect(updates[1]).toEqual(updates[0]);
expect(updates[0].session.turn_detection).toEqual({ type: "server_vad" });
});
it("sends what was typed and then asks for a response", async () => {
const user = userEvent.setup();
render(<RealtimePlayground {...props} />);

View file

@ -22,6 +22,20 @@ interface RealtimePlaygroundProps {
selectedGuardrails?: string[];
}
const sessionUpdate = (voice: string) =>
JSON.stringify({
type: "session.update",
session: {
type: "realtime",
modalities: ["text", "audio"],
voice,
input_audio_format: "pcm16",
output_audio_format: "pcm16",
input_audio_transcription: { model: "gpt-4o-mini-transcribe" },
turn_detection: { type: "server_vad" },
},
});
const RealtimePlayground: React.FC<RealtimePlaygroundProps> = ({
accessToken,
selectedModel,
@ -125,21 +139,7 @@ const RealtimePlayground: React.FC<RealtimePlaygroundProps> = ({
const type = data.type;
if (type === "session.created") {
// GA: session.type is required ("realtime" | "transcription")
ws.send(
JSON.stringify({
type: "session.update",
session: {
type: "realtime",
modalities: ["text", "audio"],
voice: selectedVoice,
input_audio_format: "pcm16",
output_audio_format: "pcm16",
input_audio_transcription: { model: "gpt-4o-mini-transcribe" },
turn_detection: null,
},
}),
);
ws.send(sessionUpdate(selectedVoice));
} else if (type === "session.updated") {
// session configured
} else if (
@ -231,22 +231,7 @@ const RealtimePlayground: React.FC<RealtimePlaygroundProps> = ({
const startRecording = useCallback(async () => {
if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN) return;
// Switch to server VAD mode for voice input
// GA: session.type is required
wsRef.current.send(
JSON.stringify({
type: "session.update",
session: {
type: "realtime",
modalities: ["text", "audio"],
voice: selectedVoice,
input_audio_format: "pcm16",
output_audio_format: "pcm16",
input_audio_transcription: { model: "gpt-4o-mini-transcribe" },
turn_detection: { type: "server_vad" },
},
}),
);
wsRef.current.send(sessionUpdate(selectedVoice));
try {
const stream = await navigator.mediaDevices.getUserMedia({ audio: true });
@ -301,7 +286,7 @@ const RealtimePlayground: React.FC<RealtimePlaygroundProps> = ({
} catch (err: any) {
addMessage("status", `Microphone error: ${err.message}`);
}
}, [addMessage]);
}, [addMessage, selectedVoice]);
const stopRecording = useCallback(() => {
processorRef.current?.disconnect();
@ -317,21 +302,7 @@ const RealtimePlayground: React.FC<RealtimePlaygroundProps> = ({
if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN) return;
if (configureSessionRef.current) return;
configureSessionRef.current = true;
// GA: session.type is required
wsRef.current.send(
JSON.stringify({
type: "session.update",
session: {
type: "realtime",
modalities: ["text", "audio"],
voice: selectedVoice,
input_audio_format: "pcm16",
output_audio_format: "pcm16",
input_audio_transcription: { model: "gpt-4o-mini-transcribe" },
turn_detection: null,
},
}),
);
wsRef.current.send(sessionUpdate(selectedVoice));
}, [selectedVoice]);
const sendTextMessage = useCallback(() => {