mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-03 07:23:33 +00:00
fix(mcp): preserve scoped dispatch and nested request limits
This commit is contained in:
commit
bd26d07780
56 changed files with 3793 additions and 328 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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":
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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")),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
"""
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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}})`"
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
},
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
):
|
||||
|
|
|
|||
|
|
@ -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."
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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[
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
21
tests/benchmarks/test_realtime_session_merge_benchmarks.py
Normal file
21
tests/benchmarks/test_realtime_session_merge_benchmarks.py
Normal 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
|
||||
|
|
@ -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.
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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=(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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"}],
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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-ü"
|
||||
|
|
|
|||
|
|
@ -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 []
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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]}}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
418
tests/unit/proxy/mcp_server/test_code_orchestration.py
Normal file
418
tests/unit/proxy/mcp_server/test_code_orchestration.py
Normal 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",
|
||||
}
|
||||
|
|
@ -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()
|
||||
|
||||
|
||||
|
|
|
|||
112
tests/unit/proxy/mcp_server/test_code_worker.py
Normal file
112
tests/unit/proxy/mcp_server/test_code_worker.py
Normal 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
|
||||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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} />);
|
||||
|
|
|
|||
|
|
@ -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(() => {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue