diff --git a/litellm/constants.py b/litellm/constants.py index 1f9e81443f7..f7ba55d8ce4 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -380,6 +380,7 @@ REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES: Final = get_env_int_or_none("REALTIME 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) # RFC 6455 caps the close frame payload at 125 bytes, 2 of which carry the status code WEBSOCKET_CLOSE_REASON_MAX_BYTES: Final = 123 diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 2ceac5899fd..c0a382b136c 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -17,7 +17,9 @@ 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, BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY, BEDROCK_REALTIME_SDK_DISTRIBUTION, @@ -35,7 +37,7 @@ from litellm.types.realtime import RealtimeResponseTransformInput from ..base_aws_llm import BaseAWSLLM, run_aws_signing from ..common_utils import BedrockError -from .transformation import BedrockRealtimeConfig +from .transformation import BedrockRealtimeConfig, BedrockSessionUpdateRejection _CLIENT_MODALITIES_ADAPTER: Final[TypeAdapter["list[str] | None"]] = TypeAdapter(list[str] | None) _CLIENT_MESSAGE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) @@ -232,12 +234,43 @@ async def _ack_session_update( return True +async def _reject_session_update( + client_ws: RealtimeClientWebSocket, + parsed_client_message: Mapping[str, JsonValue], + rejection: BedrockSessionUpdateRejection, +) -> 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)) + 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 + return True + + class BedrockRealtime(BaseAWSLLM): """Handler for Bedrock Nova Sonic realtime speech-to-speech API.""" - def __init__(self, sdk_version_lookup: Callable[[], str | None] = _installed_sdk_version): + def __init__( + self, + sdk_version_lookup: Callable[[], str | None] = _installed_sdk_version, + close_timeout: float = BEDROCK_REALTIME_CLOSE_TIMEOUT_SECONDS, + ): super().__init__() self._sdk_version_lookup: Final = sdk_version_lookup + self._close_timeout: Final = close_timeout async def async_realtime( self, @@ -259,6 +292,7 @@ class BedrockRealtime(BaseAWSLLM): aws_bedrock_runtime_endpoint: str | None = None, aws_external_id: str | None = None, aws_session_tags: object = None, + drop_params: bool = False, **kwargs: object, ): """ @@ -341,7 +375,9 @@ class BedrockRealtime(BaseAWSLLM): ) try: - await self._run_session(websocket, open_bidirectional_stream, model, logging_obj, pending_session_update) + await self._run_session( + websocket, open_bidirectional_stream, model, logging_obj, pending_session_update, drop_params + ) finally: await _close_bedrock_client(bedrock_client) @@ -352,8 +388,9 @@ class BedrockRealtime(BaseAWSLLM): model: str, logging_obj: LiteLLMLogging, pending_session_update: str | None, + drop_params: bool = False, ) -> None: - transformation_config: Final = BedrockRealtimeConfig() + transformation_config: Final = BedrockRealtimeConfig(drop_params=drop_params) bedrock_stream: Final = await open_bidirectional_stream() @@ -489,6 +526,11 @@ class BedrockRealtime(BaseAWSLLM): parsed_client_message = _parse_client_message(message) is_session_update = _json_str(parsed_client_message.get("type")) == "session.update" if is_session_update: + rejection = transformation_config.session_update_rejection(parsed_client_message) + if rejection is not None: + if not await _reject_session_update(client_ws, parsed_client_message, rejection): + break + continue client_ws.scope[BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY] = message transformed_messages = transformation_config.transform_realtime_request( @@ -504,11 +546,20 @@ class BedrockRealtime(BaseAWSLLM): ): break finally: - for close_message in transformation_config.session_close_messages(): + + async def close_bedrock_input() -> None: + for close_message in transformation_config.session_close_messages(): + with contextlib.suppress(Exception): + await send_to_bedrock(close_message) with contextlib.suppress(Exception): - await send_to_bedrock(close_message) - with contextlib.suppress(Exception): - await bedrock_stream.input_stream.close() + await bedrock_stream.input_stream.close() + + try: + await asyncio.wait_for(close_bedrock_input(), self._close_timeout) + except (asyncio.TimeoutError, TimeoutError): + verbose_proxy_logger.debug( + "Bedrock Realtime: timed out closing Bedrock input stream after %.1fs", self._close_timeout + ) async def _forward_bedrock_to_client( self, diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index 3038b0831f4..e8bc0482a28 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -5,13 +5,16 @@ Transforms between OpenAI Realtime API format and Bedrock Nova Sonic format. """ import base64 +import dataclasses import json +import math import uuid as uuid_lib -from collections.abc import Mapping -from typing import Final, cast +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Final, Literal import httpx -from pydantic import TypeAdapter +from pydantic import ConfigDict, JsonValue, TypeAdapter from litellm._logging import verbose_logger from litellm._uuid import uuid @@ -24,6 +27,8 @@ from litellm.types.llms.openai import ( OpenAIRealtimeContentPartDone, OpenAIRealtimeDoneEvent, OpenAIRealtimeEvents, + OpenAIRealtimeFunctionCallArgumentsDelta, + OpenAIRealtimeFunctionCallArgumentsDone, OpenAIRealtimeInputAudioBufferSpeechEvent, OpenAIRealtimeInputAudioTranscriptionCompleted, OpenAIRealtimeInputAudioTranscriptionDelta, @@ -35,7 +40,9 @@ from litellm.types.llms.openai import ( OpenAIRealtimeResponseTextDone, OpenAIRealtimeResponseUsage, OpenAIRealtimeStreamResponseBaseObject, + OpenAIRealtimeStreamResponseOutputItem, OpenAIRealtimeStreamResponseOutputItemAdded, + OpenAIRealtimeStreamResponseOutputItemContent, OpenAIRealtimeStreamSession, OpenAIRealtimeStreamSessionEvents, OpenAIRealtimeUsageTokenDetails, @@ -52,6 +59,26 @@ _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) class BedrockContentEnd(LiteLLMBaseModel): stopReason: str | None = None + type: str | None = None + + +def _tool_function(tool: JsonValue) -> Mapping[str, JsonValue] | None: + if not isinstance(tool, dict) or tool.get("type") != "function": + return None + function: Final = tool.get("function") + return function if isinstance(function, dict) and function else tool + + +class BedrockToolUse(LiteLLMBaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + toolUseId: str = "" + toolName: str = "" + content: JsonValue = "{}" + + @property + def arguments(self) -> str: + return self.content if isinstance(self.content, str) else json.dumps(self.content) class BedrockUsageTokenDetails(LiteLLMBaseModel): @@ -80,12 +107,63 @@ TRIGGER_AUDIO_BYTES_PER_SECOND: Final = TRIGGER_AUDIO_SAMPLE_RATE_HERTZ * 2 TRIGGER_LEADING_SILENCE: Final = bytes(TRIGGER_AUDIO_BYTES_PER_SECOND // 2) TRIGGER_TRAILING_SILENCE: Final = bytes(TRIGGER_AUDIO_BYTES_PER_SECOND * 3) TRIGGER_AUDIO_CHUNK_SIZE: Final = 1024 +PROMPT_STARTING_CLIENT_EVENT_TYPES: Final = frozenset( + ("input_audio_buffer.append", "input_audio_buffer.commit", "conversation.item.create", "response.create") +) + + +def _number_setting(value: JsonValue) -> float | None: + if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value): + return None + return float(value) + + +def _integer_setting(value: JsonValue) -> int | None: + number: Final = _number_setting(value) + return None if number is None else int(number) + + +def _malformed_number_settings(session_config: Mapping[str, JsonValue]) -> tuple[str, ...]: + temperature: Final = session_config.get("temperature") + max_output_tokens: Final = session_config.get("max_response_output_tokens") + output_rate_hertz: Final = session_config.get("output_sample_rate_hertz") + malformed: Final = ( + ("temperature", temperature is not None and _number_setting(temperature) is None), + ( + "max_response_output_tokens", + max_output_tokens is not None + and max_output_tokens != "inf" + and _integer_setting(max_output_tokens) is None, + ), + ("output_sample_rate_hertz", output_rate_hertz is not None and _integer_setting(output_rate_hertz) is None), + ) + return tuple(name for name, is_malformed in malformed if is_malformed) + + +@dataclasses.dataclass(frozen=True, slots=True) +class BedrockSessionUpdateRejection: + code: Literal["invalid_type", "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.""" + + tools_json: str + system_prompt: str + voice_id: str + output_sample_rate_hertz: int + temperature: float + max_tokens: int | None class BedrockRealtimeConfig(BaseRealtimeConfig): """Configuration for Bedrock Nova Sonic realtime transformations.""" - def __init__(self): + def __init__(self, drop_params: bool = False): + self.drop_params = drop_params # Track session state self.prompt_name = str(uuid_lib.uuid4()) self.content_name = str(uuid_lib.uuid4()) @@ -95,7 +173,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): # Default configuration values # Inference configuration - self.max_tokens = 1024 + self.max_tokens: int | None = 1024 self.top_p = 0.9 self.temperature = 0.7 @@ -128,6 +206,19 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): self._user_transcript_buffer = "" self._cumulative_usage = BedrockUsageEvent() self._reported_usage = BedrockUsageEvent() + self._response_done_items: Mapping[int, OpenAIRealtimeStreamResponseOutputItem] = MappingProxyType({}) + self._response_item_count = 0 + self._message_output_index = 0 + self._assistant_block_open = False + self._response_done_pending = False + self._started_settings = BedrockLockedSessionSettings( + tools_json="[]", + system_prompt="", + voice_id=self.voice_id, + output_sample_rate_hertz=self.output_sample_rate_hertz, + temperature=self.temperature, + max_tokens=self.max_tokens, + ) def get_error_class( self, @@ -149,7 +240,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): """Bedrock requires session configuration.""" return True - def session_configuration_request(self, model: str, tools: list[dict] | None = None) -> str: + def session_configuration_request(self, model: str, tools: Sequence[JsonValue] | None = None) -> str: """ Create initial session configuration for Bedrock Nova Sonic. @@ -159,19 +250,9 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): Returns JSON string with session start and prompt start events. """ - session_start: Final = { - "event": { - "sessionStart": { - "inferenceConfiguration": { - "maxTokens": self.max_tokens, - "topP": self.top_p, - "temperature": self.temperature, - } - } - } - } + session_start: Final = {"event": {"sessionStart": {"inferenceConfiguration": self._inference_configuration()}}} - prompt_start_config: Final = { + prompt_start_config: Final[Mapping[str, object]] = { "promptName": self.prompt_name, "textOutputConfiguration": {"mediaType": self.text_media_type}, "audioOutputConfiguration": { @@ -183,41 +264,34 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): "encoding": self.output_encoding, "audioType": self.output_audio_type, }, + **( + { + "toolUseOutputConfiguration": {"mediaType": "application/json"}, + "toolConfiguration": {"tools": self._transform_tools_to_bedrock_format(tools)}, + } + if tools + else {} + ), } - # Add tool configuration if tools are provided - if tools: - prompt_start_config["toolUseOutputConfiguration"] = {"mediaType": "application/json"} - prompt_start_config["toolConfiguration"] = {"tools": self._transform_tools_to_bedrock_format(tools)} - prompt_start: Final = {"event": {"promptStart": prompt_start_config}} # Return as a marker that we've sent the configuration return json.dumps({"session_start": session_start, "prompt_start": prompt_start}) - def _transform_tools_to_bedrock_format(self, tools: list[dict]) -> list[dict]: - """ - Transform OpenAI tool format to Bedrock tool format. - - Args: - tools: List of OpenAI format tools - - Returns: - List of Bedrock format tools - """ - bedrock_tools: Final = [] - for tool in tools: - if tool.get("type") == "function": - function = tool.get("function", {}) - bedrock_tool = { - "toolSpec": { - "name": function.get("name", ""), - "description": function.get("description", ""), - "inputSchema": {"json": json.dumps(function.get("parameters", {}))}, - } + def _transform_tools_to_bedrock_format(self, tools: Sequence[JsonValue]) -> tuple[Mapping[str, JsonValue], ...]: + functions: Final = (_tool_function(tool) for tool in tools) + return tuple( + { + "toolSpec": { + "name": function.get("name", ""), + "description": function.get("description", ""), + "inputSchema": {"json": json.dumps(function.get("parameters", {}))}, } - bedrock_tools.append(bedrock_tool) - return bedrock_tools + } + for function in functions + if function is not None + ) def _map_audio_format_to_sample_rate(self, audio_format: str, is_output: bool = True) -> int: """ @@ -252,47 +326,34 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): messages: Final[list[str]] = [] session_config: Final = json_message.get("session", {}) + locked_settings: Final = self._apply_locked_settings(session_config, self._started_settings) - # Update inference configuration from session if provided - if "max_response_output_tokens" in session_config: - self.max_tokens = session_config["max_response_output_tokens"] - if "temperature" in session_config: - self.temperature = session_config["temperature"] - - # Update audio output configuration from session if provided - if "voice" in session_config: - self.voice_id = session_config["voice"] - if "output_audio_format" in session_config: - output_format: Final = session_config["output_audio_format"] - self.output_sample_rate_hertz = self._map_audio_format_to_sample_rate(output_format, is_output=True) + # Settings Bedrock locks at prompt start may only change before it + if not self.prompt_started: + self.max_tokens = locked_settings.max_tokens + self.temperature = locked_settings.temperature + self.voice_id = locked_settings.voice_id + self.output_sample_rate_hertz = locked_settings.output_sample_rate_hertz # Update audio input configuration from session if provided if "input_audio_format" in session_config: input_format: Final = session_config["input_audio_format"] self.input_sample_rate_hertz = self._map_audio_format_to_sample_rate(input_format, is_output=False) - # Allow direct override of sample rates if provided (custom extension) - if "output_sample_rate_hertz" in session_config: - self.output_sample_rate_hertz = session_config["output_sample_rate_hertz"] + # Allow direct override of the input sample rate if provided (custom extension) if "input_sample_rate_hertz" in session_config: self.input_sample_rate_hertz = session_config["input_sample_rate_hertz"] + if self.prompt_started: + return messages + # Send session start - session_start: Final = { - "event": { - "sessionStart": { - "inferenceConfiguration": { - "maxTokens": self.max_tokens, - "topP": self.top_p, - "temperature": self.temperature, - } - } - } - } + session_start: Final = {"event": {"sessionStart": {"inferenceConfiguration": self._inference_configuration()}}} messages.append(json.dumps(session_start)) # Send prompt start - prompt_start_config: Final = { + tools: Final[Sequence[JsonValue] | None] = session_config.get("tools") + prompt_start_config: Final[Mapping[str, object]] = { "promptName": self.prompt_name, "textOutputConfiguration": {"mediaType": self.text_media_type}, "audioOutputConfiguration": { @@ -304,63 +365,138 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): "encoding": self.output_encoding, "audioType": self.output_audio_type, }, + **( + { + "toolUseOutputConfiguration": {"mediaType": "application/json"}, + "toolConfiguration": {"tools": self._transform_tools_to_bedrock_format(tools)}, + } + if tools + else {} + ), } - # Add tool configuration if tools are provided - tools: Final = session_config.get("tools") - if tools: - prompt_start_config["toolUseOutputConfiguration"] = {"mediaType": "application/json"} - prompt_start_config["toolConfiguration"] = {"tools": self._transform_tools_to_bedrock_format(tools)} - prompt_start: Final = {"event": {"promptStart": prompt_start_config}} messages.append(json.dumps(prompt_start)) self.prompt_started = True + self._started_settings = locked_settings - # Send system prompt if provided - instructions: Final = session_config.get("instructions") - if instructions: - text_content_name: Final = str(uuid_lib.uuid4()) + # Bedrock requires the first content block in a prompt to be SYSTEM role + instructions: Final = locked_settings.system_prompt + text_content_name: Final = str(uuid_lib.uuid4()) - # Content start - text_content_start: Final = { - "event": { - "contentStart": { - "promptName": self.prompt_name, - "contentName": text_content_name, - "type": "TEXT", - "interactive": False, - "role": "SYSTEM", - "textInputConfiguration": {"mediaType": self.text_media_type}, - } + # Content start + text_content_start: Final = { + "event": { + "contentStart": { + "promptName": self.prompt_name, + "contentName": text_content_name, + "type": "TEXT", + "interactive": False, + "role": "SYSTEM", + "textInputConfiguration": {"mediaType": self.text_media_type}, } } - messages.append(json.dumps(text_content_start)) + } + messages.append(json.dumps(text_content_start)) - # Text input - text_input: Final = { - "event": { - "textInput": { - "promptName": self.prompt_name, - "contentName": text_content_name, - "content": instructions, - } + # Text input + text_input: Final = { + "event": { + "textInput": { + "promptName": self.prompt_name, + "contentName": text_content_name, + "content": instructions, } } - messages.append(json.dumps(text_input)) + } + messages.append(json.dumps(text_input)) - # Content end - text_content_end: Final = { - "event": { - "contentEnd": { - "promptName": self.prompt_name, - "contentName": text_content_name, - } + # Content end + text_content_end: Final = { + "event": { + "contentEnd": { + "promptName": self.prompt_name, + "contentName": text_content_name, } } - messages.append(json.dumps(text_content_end)) + } + messages.append(json.dumps(text_content_end)) return messages + def _inference_configuration(self) -> Mapping[str, int | float]: + return { + "topP": self.top_p, + "temperature": self.temperature, + **({"maxTokens": self.max_tokens} if self.max_tokens is not None else {}), + } + + def _apply_locked_settings( + self, session_config: Mapping[str, JsonValue], base: BedrockLockedSessionSettings + ) -> BedrockLockedSessionSettings: + output_audio_format: Final = session_config.get("output_audio_format") + temperature: Final = session_config.get("temperature") + max_output_tokens: Final = session_config.get("max_response_output_tokens") + requested_tools: Final = session_config.get("tools") + instructions: Final = session_config.get("instructions") + voice: Final = session_config.get("voice") + tools_json: Final = ( + json.dumps(self._transform_tools_to_bedrock_format(requested_tools)) + if isinstance(requested_tools, list) + else base.tools_json + ) + system_prompt: Final = base.system_prompt if instructions is None else str(instructions) + voice_id: Final = base.voice_id if voice is None else str(voice) + format_sample_rate: Final = ( + base.output_sample_rate_hertz + if output_audio_format is None + else self._map_audio_format_to_sample_rate(str(output_audio_format), is_output=True) + ) + requested_output_rate: Final = _integer_setting(session_config.get("output_sample_rate_hertz")) + output_sample_rate_hertz: Final = ( + requested_output_rate if requested_output_rate is not None else format_sample_rate + ) + requested_temperature: Final = _number_setting(temperature) + requested_max_tokens: Final = _integer_setting(max_output_tokens) + locked_temperature: Final = requested_temperature if requested_temperature is not None else base.temperature + explicit_max_tokens: Final = requested_max_tokens if requested_max_tokens is not None else base.max_tokens + max_tokens: Final = None if max_output_tokens == "inf" else explicit_max_tokens + return dataclasses.replace( + base, + tools_json=tools_json, + system_prompt=system_prompt, + voice_id=voice_id, + output_sample_rate_hertz=output_sample_rate_hertz, + temperature=locked_temperature, + max_tokens=max_tokens, + ) + + def session_update_rejection(self, json_message: Mapping[str, JsonValue]) -> BedrockSessionUpdateRejection | 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): + return None + malformed: Final = () if self.drop_params else _malformed_number_settings(session_config) + if malformed: + return BedrockSessionUpdateRejection( + code="invalid_type", + message=( + f"Invalid type for 'session.{malformed[0]}': Bedrock Nova Sonic needs a number for temperature " + "and output_sample_rate_hertz, and a number or 'inf' for max_response_output_tokens. " + "Set drop_params to drop it instead" + ), + param=f"session.{malformed[0]}", + ) + if not self.prompt_started: + return None + if self._apply_locked_settings(session_config, self._started_settings) != self._started_settings: + return BedrockSessionUpdateRejection( + 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", + ) + return None + def transform_input_audio_buffer_append_event(self, json_message: dict) -> list[str]: """ Transform input_audio_buffer.append event to Bedrock audio input. @@ -655,18 +791,21 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): return [] message_type: Final = json_message.get("type") - - # Route to appropriate transformation method + prompt_start_messages: Final = ( + self.transform_session_update_event({"session": {}}) + if message_type in PROMPT_STARTING_CLIENT_EVENT_TYPES and not self.prompt_started + else [] + ) if message_type == "session.update": return self.transform_session_update_event(json_message) elif message_type == "input_audio_buffer.append": - return self.transform_input_audio_buffer_append_event(json_message) + return prompt_start_messages + self.transform_input_audio_buffer_append_event(json_message) elif message_type == "input_audio_buffer.commit": - return self.transform_input_audio_buffer_commit_event(json_message) + return prompt_start_messages + self.transform_input_audio_buffer_commit_event(json_message) elif message_type == "conversation.item.create": - return self.transform_conversation_item_create_event(json_message) + return prompt_start_messages + self.transform_conversation_item_create_event(json_message) elif message_type == "response.create": - return self.transform_response_create_event(json_message) + return prompt_start_messages + self.transform_response_create_event(json_message) elif message_type == "response.cancel": return self.transform_response_cancel_event(json_message) else: @@ -718,6 +857,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): current_response_id: str | None, current_output_item_id: str | None, current_conversation_id: str | None, + current_delta_type: ALL_DELTA_TYPES | None, ) -> tuple[ list[OpenAIRealtimeEvents], str | None, @@ -751,7 +891,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): current_response_id, current_output_item_id, current_conversation_id, - None, + current_delta_type, ) generation_stage: Final = self._parse_generation_stage(content_start.get("additionalModelFields")) @@ -776,12 +916,14 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): current_response_id = f"resp_{uuid.uuid4()}" if not current_output_item_id: current_output_item_id = f"item_{uuid.uuid4()}" + self._message_output_index = self._response_item_count + self._response_item_count += 1 if not current_conversation_id: current_conversation_id = f"conv_{uuid.uuid4()}" # Determine content type content_type: Final = content_start.get("type", "TEXT") - current_delta_type: Final[ALL_DELTA_TYPES] = "text" if content_type == "TEXT" else "audio" + next_delta_type: Final[ALL_DELTA_TYPES] = "text" if content_type == "TEXT" else "audio" returned_messages: Final[list[OpenAIRealtimeEvents]] = [] @@ -805,7 +947,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): output_item_added: Final = OpenAIRealtimeStreamResponseOutputItemAdded( type="response.output_item.added", response_id=current_response_id, - output_index=0, + output_index=self._message_output_index, item={ "id": current_output_item_id, "object": "realtime.item", @@ -815,18 +957,17 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): "content": [], }, ) + self._assistant_block_open = True returned_messages.append(output_item_added) # Send response.content_part.added content_part_added: Final = OpenAIRealtimeResponseContentPartAdded( type="response.content_part.added", content_index=0, - output_index=0, + output_index=self._message_output_index, event_id=f"event_{uuid.uuid4()}", item_id=current_output_item_id, - part=( - {"type": "text", "text": ""} if current_delta_type == "text" else {"type": "audio", "transcript": ""} - ), + part=({"type": "text", "text": ""} if next_delta_type == "text" else {"type": "audio", "transcript": ""}), response_id=current_response_id, ) returned_messages.append(content_part_added) @@ -836,7 +977,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): current_response_id, current_output_item_id, current_conversation_id, - current_delta_type, + next_delta_type, ) @staticmethod @@ -973,7 +1114,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): content_index=0, event_id=f"event_{uuid.uuid4()}", item_id=current_output_item_id, - output_index=0, + output_index=self._message_output_index, response_id=current_response_id, delta=text_content, ) @@ -1013,7 +1154,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): content_index=0, event_id=f"event_{uuid.uuid4()}", item_id=current_output_item_id, - output_index=0, + output_index=self._message_output_index, response_id=current_response_id, delta=audio_content, ) @@ -1049,22 +1190,21 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): return [], current_delta_chunks if not current_output_item_id or not current_response_id: return [], current_delta_chunks + if not self._assistant_block_open: + return [], current_delta_chunks + self._assistant_block_open = False returned_messages: Final[list[OpenAIRealtimeEvents]] = [] + accumulated_text: Final = "".join(chunk.get("delta", "") for chunk in current_delta_chunks or []) # Send appropriate done event based on type if current_delta_type == "text": - # Accumulate text - accumulated_text = "" - if current_delta_chunks: - accumulated_text = "".join([chunk.get("delta", "") for chunk in current_delta_chunks]) - text_done: Final = OpenAIRealtimeResponseTextDone( type="response.text.done", content_index=0, event_id=f"event_{uuid.uuid4()}", item_id=current_output_item_id, - output_index=0, + output_index=self._message_output_index, response_id=current_response_id, text=accumulated_text, ) @@ -1076,7 +1216,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): content_index=0, event_id=f"event_{uuid.uuid4()}", item_id=current_output_item_id, - output_index=0, + output_index=self._message_output_index, part={"type": "text", "text": accumulated_text}, response_id=current_response_id, ) @@ -1088,7 +1228,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): content_index=0, event_id=f"event_{uuid.uuid4()}", item_id=current_output_item_id, - output_index=0, + output_index=self._message_output_index, response_id=current_response_id, ) returned_messages.append(audio_done) @@ -1099,28 +1239,34 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): content_index=0, event_id=f"event_{uuid.uuid4()}", item_id=current_output_item_id, - output_index=0, + output_index=self._message_output_index, part={"type": "audio", "transcript": ""}, response_id=current_response_id, ) returned_messages.append(content_part_done) - # Send output_item.done + completed_parts: Final = self._completed_parts(current_delta_type, accumulated_text) + existing_item: Final = self._response_done_items.get(self._message_output_index) + existing_content: Final = existing_item.get("content", []) if existing_item is not None else [] + completed_message_item: Final[OpenAIRealtimeStreamResponseOutputItem] = { + "id": current_output_item_id, + "object": "realtime.item", + "type": "message", + "status": "completed", + "role": "assistant", + "content": [*existing_content, *completed_parts], + } output_item_done: Final = OpenAIRealtimeOutputItemDone( type="response.output_item.done", event_id=f"event_{uuid.uuid4()}", - output_index=0, + output_index=self._message_output_index, response_id=current_response_id, - item={ - "id": current_output_item_id, - "object": "realtime.item", - "type": "message", - "status": "completed", - "role": "assistant", - "content": [], - }, + item=completed_message_item, ) returned_messages.append(output_item_done) + self._response_done_items = MappingProxyType( + {**self._response_done_items, self._message_output_index: completed_message_item} + ) # Reset delta chunks return returned_messages, None @@ -1165,6 +1311,12 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): return [], None, None, None usage: Final = self._take_usage_delta() + output_items: Final = tuple(item for _, item in sorted(self._response_done_items.items())) + self._response_done_items = MappingProxyType({}) + self._response_item_count = 0 + self._message_output_index = 0 + self._assistant_block_open = False + self._response_done_pending = False response_done: Final = OpenAIRealtimeDoneEvent( type="response.done", event_id=f"event_{uuid.uuid4()}", @@ -1172,7 +1324,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): object="realtime.response", id=current_response_id, status="completed", - output=[], + output=list(output_items), conversation_id=current_conversation_id, usage=dict(usage), ), @@ -1184,56 +1336,148 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): def transform_tool_use_event( self, event: dict, - current_output_item_id: str | None, current_response_id: str | None, + current_conversation_id: str | None, ) -> tuple[list[OpenAIRealtimeEvents], str, str]: """ Transform Bedrock toolUse event to OpenAI format. Args: event: Bedrock toolUse event - current_output_item_id: Current output item ID current_response_id: Current response ID + current_conversation_id: Current conversation ID Returns: - Tuple of (events, tool_call_id, tool_name) for tracking + Tuple of (events, response_id, conversation_id) """ verbose_logger.debug("Handling toolUse") - tool_use: Final = event["toolUse"] + tool_use: Final = BedrockToolUse.model_validate(event["toolUse"]) + verbose_logger.debug("Tool use event: %s (ID: %s)", tool_use.toolName, tool_use.toolUseId) - if not current_output_item_id or not current_response_id: - return [], "", "" + response_id: Final = current_response_id or f"resp_{uuid.uuid4()}" + conversation_id: Final = current_conversation_id or f"conv_{uuid.uuid4()}" + item_id: Final = f"item_{uuid.uuid4()}" + output_index: Final = self._response_item_count + self._response_item_count += 1 - # Parse the tool input - tool_input = {} - if "input" in tool_use: - try: - tool_input = json.loads(tool_use["input"]) if isinstance(tool_use["input"], str) else tool_use["input"] - except json.JSONDecodeError: - tool_input = {} + response_created_events: Final[tuple[OpenAIRealtimeEvents, ...]] = ( + ( + OpenAIRealtimeStreamResponseBaseObject( + type="response.created", + event_id=f"event_{uuid.uuid4()}", + response={ + "object": "realtime.response", + "id": response_id, + "status": "in_progress", + "output": [], + "conversation_id": conversation_id, + }, + ), + ) + if current_response_id is None + else () + ) - tool_call_id: Final = tool_use.get("toolUseId", "") - tool_name: Final = tool_use.get("toolName", "") + output_item_added: Final = OpenAIRealtimeStreamResponseOutputItemAdded( + type="response.output_item.added", + response_id=response_id, + output_index=output_index, + item={ + "id": item_id, + "object": "realtime.item", + "type": "function_call", + "status": "in_progress", + "call_id": tool_use.toolUseId, + "name": tool_use.toolName, + "arguments": "", + }, + ) - # Create a function call arguments done event - # This is a custom event format that matches what clients expect - function_call_event: Final[dict[str, object]] = { - "type": "response.function_call_arguments.done", - "event_id": f"event_{uuid.uuid4()}", - "response_id": current_response_id, - "item_id": current_output_item_id, - "output_index": 0, - "call_id": tool_call_id, - "name": tool_name, - "arguments": json.dumps(tool_input), + arguments_delta: Final = OpenAIRealtimeFunctionCallArgumentsDelta( + type="response.function_call_arguments.delta", + event_id=f"event_{uuid.uuid4()}", + response_id=response_id, + item_id=item_id, + output_index=output_index, + call_id=tool_use.toolUseId, + delta=tool_use.arguments, + ) + + function_call_event: Final = OpenAIRealtimeFunctionCallArgumentsDone( + type="response.function_call_arguments.done", + event_id=f"event_{uuid.uuid4()}", + response_id=response_id, + item_id=item_id, + output_index=output_index, + call_id=tool_use.toolUseId, + name=tool_use.toolName, + arguments=tool_use.arguments, + ) + + completed_item: Final[OpenAIRealtimeStreamResponseOutputItem] = { + "id": item_id, + "object": "realtime.item", + "type": "function_call", + "status": "completed", + "call_id": tool_use.toolUseId, + "name": tool_use.toolName, + "arguments": tool_use.arguments, } + output_item_done: Final = OpenAIRealtimeOutputItemDone( + type="response.output_item.done", + event_id=f"event_{uuid.uuid4()}", + output_index=output_index, + response_id=response_id, + item=completed_item, + ) + + self._response_done_items = MappingProxyType({**self._response_done_items, output_index: completed_item}) return ( - [cast(OpenAIRealtimeEvents, function_call_event)], - tool_call_id, - tool_name, + [ + *response_created_events, + output_item_added, + arguments_delta, + function_call_event, + output_item_done, + ], + response_id, + conversation_id, ) + def _tool_content_end_events( + self, + current_output_item_id: str | None, + current_response_id: str | None, + current_conversation_id: str | None, + current_delta_type: ALL_DELTA_TYPES | None, + ) -> tuple[ + tuple[OpenAIRealtimeEvents, ...], + str | None, + str | None, + ALL_DELTA_TYPES | None, + ]: + if self._assistant_block_open: + self._response_done_pending = True + return (), current_output_item_id, current_response_id, current_delta_type + done_events, output_item_id, response_id, delta_type = self._response_done_events( + current_response_id, current_conversation_id + ) + return tuple(done_events), output_item_id, response_id, delta_type + + @staticmethod + def _completed_parts( + delta_type: str | None, accumulated_text: str + ) -> tuple[OpenAIRealtimeStreamResponseOutputItemContent, ...]: + if delta_type == "text": + return ({"type": "text", "text": accumulated_text},) + elif delta_type == "audio": + return ({"type": "audio", "transcript": ""},) + return () + + def _should_close_response(self, stop_reason: str | None) -> bool: + return stop_reason == "END_TURN" or self._response_done_pending + def transform_conversation_item_create_tool_result_event(self, json_message: dict) -> list[str]: """ Transform conversation.item.create with tool result to Bedrock format. @@ -1363,6 +1607,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): current_response_id, current_output_item_id, current_conversation_id, + current_delta_type, ) returned_messages.extend(events) @@ -1382,6 +1627,17 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): events = self.transform_audio_output_event(event, current_output_item_id, current_response_id) returned_messages.extend(events) + elif "contentEnd" in event and BedrockContentEnd.model_validate(event["contentEnd"]).type == "TOOL": + ( + done_events, + current_output_item_id, # rebind-ok: same state-threading pattern as the adjacent event branches + current_response_id, # rebind-ok: same state-threading pattern as the adjacent event branches + current_delta_type, # rebind-ok: same state-threading pattern as the adjacent event branches + ) = self._tool_content_end_events( + current_output_item_id, current_response_id, current_conversation_id, current_delta_type + ) + returned_messages.extend(done_events) + elif "contentEnd" in event and self._user_transcript_active: self._user_transcript_active = False self._user_transcript_generation_stage = None @@ -1396,7 +1652,8 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): current_delta_chunks, ) returned_messages.extend(events) - if BedrockContentEnd.model_validate(event["contentEnd"]).stopReason == "END_TURN": + content_end: Final = BedrockContentEnd.model_validate(event["contentEnd"]) + if self._should_close_response(content_end.stopReason): ( done_events, current_output_item_id, @@ -1412,12 +1669,10 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): self.transform_usage_event(BedrockUsageEvent.model_validate(event["usageEvent"])) elif "toolUse" in event: - events, tool_call_id, tool_name = self.transform_tool_use_event( - event, current_output_item_id, current_response_id + events, current_response_id, current_conversation_id = self.transform_tool_use_event( + event, current_response_id, current_conversation_id ) returned_messages.extend(events) - # Store tool call info for potential use - verbose_logger.debug("Tool use event: %s (ID: %s)", tool_name, tool_call_id) elif "promptEnd" in event or "completionEnd" in event: ( diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index 02b95c80b60..367d5be7bfc 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -33,6 +33,7 @@ from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import CallTypes, LlmProviders from litellm.utils import ProviderConfigManager +from ..litellm_core_utils.core_helpers import normalize_drop_params from ..litellm_core_utils.credential_accessor import CredentialAccessor from ..litellm_core_utils.get_litellm_params import get_litellm_params from ..litellm_core_utils.litellm_logging import Logging as LiteLLMLogging @@ -584,6 +585,7 @@ async def _arealtime( aws_bedrock_runtime_endpoint: Final = kwargs.get("aws_bedrock_runtime_endpoint") aws_external_id: Final = kwargs.get("aws_external_id") aws_session_tags: Final = kwargs.get("aws_session_tags") + drop_params: Final = litellm.drop_params is True or normalize_drop_params(kwargs.get("drop_params")) is True await bedrock_realtime.async_realtime( model=model, @@ -604,6 +606,7 @@ async def _arealtime( aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint, aws_external_id=aws_external_id, aws_session_tags=aws_session_tags, + drop_params=drop_params, ) elif _custom_llm_provider == "xai": api_base = ( diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 1e4f78b2204..8731bc2905a 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -2271,6 +2271,16 @@ class OpenAIRealtimeContentPartDone(TypedDict): type: Literal["response.content_part.done"] +class OpenAIRealtimeFunctionCallArgumentsDelta(TypedDict): + type: ReadOnly[Literal["response.function_call_arguments.delta"]] + event_id: ReadOnly[str] + response_id: ReadOnly[str] + item_id: ReadOnly[str] + output_index: ReadOnly[int] + call_id: ReadOnly[str] + delta: ReadOnly[str] + + class OpenAIRealtimeFunctionCallArgumentsDone(TypedDict): type: Literal["response.function_call_arguments.done"] event_id: str @@ -2438,6 +2448,7 @@ OpenAIRealtimeEvents = ( | OpenAIRealtimeResponseAudioDone | OpenAIRealtimeContentPartDone | OpenAIRealtimeOutputItemDone + | OpenAIRealtimeFunctionCallArgumentsDelta | OpenAIRealtimeFunctionCallArgumentsDone | OpenAIRealtimeDoneEvent | OpenAIRealtimeInputAudioBufferSpeechEvent diff --git a/tests/e2e/coverage_registry/llm_nonconversational.yaml b/tests/e2e/coverage_registry/llm_nonconversational.yaml index b05e9c916d2..330ddf88fc5 100644 --- a/tests/e2e/coverage_registry/llm_nonconversational.yaml +++ b/tests/e2e/coverage_registry/llm_nonconversational.yaml @@ -62,6 +62,7 @@ - {id: llm.files.vertex.content.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "GET /v1/files/{id}/content on a Vertex unified file streams the GCS object back (provider-transformed JSONL, so asserts non-empty JSON lines rather than byte equality)"} - {id: llm.files.bedrock.content.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "GET /v1/files/{id}/content on a Bedrock unified file streams the S3 object back (provider-transformed JSONL, so asserts non-empty JSON lines rather than byte equality)"} - {id: llm.realtime.bedrock_converse.basic.stream.works, module: llm, tier: P0, subject_endpoint: realtime, route: bedrock_converse, capability: basic, streaming: stream, assertions: [works], source: "test_realtime_bedrock_e2e.py", rationale: "Nova Sonic realtime session emits response.done (LIT-2239)"} +- {id: llm.realtime.bedrock_converse.tool_use.stream.works, module: llm, tier: P0, subject_endpoint: realtime, route: bedrock_converse, capability: tool_use, streaming: stream, assertions: [works], source: "test_realtime_e2e.py", rationale: "Nova Sonic realtime forwards each tool call with its arguments and completes the tool-result round trip (LIT-9465)"} - {id: llm.google_native.gemini.basic.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: google_native, route: gemini, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "LIT-4076 / proxy/google_endpoints/endpoints.py", fail_before_fix: proven, rationale: "google-native generateContent must stamp x-litellm-response-cost so SDK traffic reconciles against spend"} - {id: llm.google_native.gemini.basic.stream.works, module: llm, tier: P0, subject_endpoint: google_native, route: gemini, capability: basic, streaming: stream, assertions: [works], source: "PR #28213 / proxy/proxy_server.py async_data_generator", fail_before_fix: proven, rationale: "streamGenerateContent must relay single-prefixed SSE frames with no [DONE] sentinel; doubled data: prefixes and the OpenAI terminator both break the Vertex Java SDK"} - {id: llm.realtime.openai.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: realtime, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "vendor strategy ยง9.19 / LIT-4778", rationale: "HTTP /v1/realtime/client_secrets returns an ephemeral credential"} diff --git a/tests/e2e/llm_translation/realtime/realtime_client.py b/tests/e2e/llm_translation/realtime/realtime_client.py index c9f06e787f3..4d8336775ef 100644 --- a/tests/e2e/llm_translation/realtime/realtime_client.py +++ b/tests/e2e/llm_translation/realtime/realtime_client.py @@ -85,7 +85,16 @@ PROVIDERS = ( vertex_credentials="os.environ/VERTEXAI_CREDENTIALS", ), ), - # RealtimeProvider("bedrock", "bedrock-realtime", ...) # TODO: Enable when Bedrock is passing + RealtimeProvider( + "bedrock", + "bedrock-realtime", + LiteLLMParamsBody( + model="bedrock/amazon.nova-2-sonic-v1:0", + aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID", + aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", + aws_region_name="os.environ/AWS_REGION", + ), + ), # RealtimeProvider( # "xai", # "xai-realtime", diff --git a/tests/e2e/llm_translation/realtime/test_realtime_e2e.py b/tests/e2e/llm_translation/realtime/test_realtime_e2e.py index 622ed9d507f..f5592abb34d 100644 --- a/tests/e2e/llm_translation/realtime/test_realtime_e2e.py +++ b/tests/e2e/llm_translation/realtime/test_realtime_e2e.py @@ -29,13 +29,14 @@ from realtime_client import ( JsonSchemaProperty, RealtimeClient, RealtimeProvider, + RealtimeSession, + ReceivedEvent, ResponseCreate, ResponseDone, ServerEnvelope, SessionConfig, SessionUpdate, as_text, - function_call_item, parse_last, realtime_model, transcript, @@ -64,19 +65,26 @@ TEXT_PARAMS: Final = tuple( for p in PROVIDERS ) +BEDROCK_TOOL_COVERS: Final = pytest.mark.covers( + "llm.realtime.bedrock_converse.tool_use.stream.works", exercised_on=["realtime"] +) + TOOL_PARAMS: Final = tuple( pytest.param( p, id=p.id, - marks=meta( - Subject( - domain=Domain.LLM_TRANSLATION, - route=Route.REALTIME, - providers=(Provider(p.id),), - models=(p.litellm_params.model,), - capabilities=(Capability.FUNCTION_CALLING,), - mode=Mode.WEBSOCKET, - ) + marks=( + meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.REALTIME, + providers=(Provider(p.id),), + models=(p.litellm_params.model,), + capabilities=(Capability.FUNCTION_CALLING,), + mode=Mode.WEBSOCKET, + ) + ), + *((BEDROCK_TOOL_COVERS,) if p.id == "bedrock" else ()), ), ) for p in PROVIDERS @@ -164,31 +172,51 @@ def test_tool_call_round_trip( session.send(user_message("What's the weather in Paris right now?")) session.send(ResponseCreate()) - first = session.collect_until("response.done", timeout=60) + final_answer = _tool_call_loop(session, max_rounds=4) - args_event = parse_last( - first, "response.function_call_arguments.done", FunctionCallArgumentsDone + assert "72" in transcript(final_answer), ( + "the answer after the tool results did not use them" ) - assert args_event is not None, "model did not emit a function call" - args = WeatherArgs.model_validate_json(args_event.arguments) - item = function_call_item(first) - assert item is not None, "no completed function_call output item" - assert item.name == "get_weather" - assert item.call_id == args_event.call_id - tool_result = WeatherResult(city=args.city, temperature_f=72) - session.send( - ConversationItemCreate( - item=FunctionCallOutputItem( - call_id=args_event.call_id, output=tool_result.model_dump_json() +def _tool_call_loop(session: RealtimeSession, *, max_rounds: int) -> tuple[ReceivedEvent, ...]: + """Answer every function_call in each response.done until one arrives without any. + + Nova 2 Sonic emits two tool calls per question, each inside its own response.done, + so a one-call round trip cannot wait for "the" follow-up: the second response.done + is another call, not the answer. The loop sends one function_call_output per call + then a single response.create, and returns the events of the first call-free round.""" + for round_index in range(max_rounds): + events = session.collect_until("response.done", timeout=60) + done = parse_last(events, "response.done", ResponseDone) + assert done is not None + calls = tuple( + item for item in done.response.output or () if item.type == "function_call" + ) + if not calls: + assert round_index > 0, "model did not emit a function call" + return events + done_args = tuple( + FunctionCallArgumentsDone.model_validate_json(event.payload) + for event in events + if event.type == "response.function_call_arguments.done" + ) + for item in calls: + assert item.name == "get_weather", item + args_event = next( + event for event in done_args if event.call_id == item.call_id + ) + args = WeatherArgs.model_validate_json(args_event.arguments) + tool_result = WeatherResult(city=args.city, temperature_f=72) + session.send( + ConversationItemCreate( + item=FunctionCallOutputItem( + call_id=args_event.call_id, output=tool_result.model_dump_json() + ) ) ) - ) session.send(ResponseCreate()) - second = session.collect_until("response.done", timeout=60) - - assert "72" in transcript(second), "follow-up did not use the tool result" + raise AssertionError(f"still calling tools after {max_rounds} rounds") _REFUSED_UPSTREAMS = ( diff --git a/tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e.py b/tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e.py index 27d5abd1828..0e4dad7af1a 100644 --- a/tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e.py +++ b/tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e.py @@ -91,6 +91,7 @@ TOOL_PARAMS: Final = tuple( ), ) for p in PROVIDERS + if p.id != "bedrock" ) AUDIO_OUTPUT_PARAMS: Final = tuple( @@ -109,6 +110,7 @@ AUDIO_OUTPUT_PARAMS: Final = tuple( ), ) for p in PROVIDERS + if p.id != "bedrock" ) AUDIO_INPUT_PARAMS: Final = tuple( @@ -127,6 +129,7 @@ AUDIO_INPUT_PARAMS: Final = tuple( ), ) for p in PROVIDERS + if p.id != "bedrock" ) # PCM16 24 kHz mono WAV of "What is the weather in Paris?" (generated via macOS diff --git a/tests/e2e/llm_translation/realtime/test_realtime_pipecat_e2e.py b/tests/e2e/llm_translation/realtime/test_realtime_pipecat_e2e.py index 3c22ae1a928..bfb418b32dc 100644 --- a/tests/e2e/llm_translation/realtime/test_realtime_pipecat_e2e.py +++ b/tests/e2e/llm_translation/realtime/test_realtime_pipecat_e2e.py @@ -64,6 +64,9 @@ from pipecat_service import LiteLLMRealtimeLLMService # noqa: E402 # Vertex native-audio live is flaky through pipecat tool calling (upstream # pipecat-ai/pipecat#2544); raw-ws tool_call_round_trip[vertex_ai] is the # source of truth for that provider. Keep openai/azure/gemini here. +# Bedrock event shapes are not yet pipecat-compatible (missing GA fields and +# event names, and no tool-result round trip reaches Nova); raw-ws +# tool_call_round_trip[bedrock] is the source of truth for that provider. PROVIDER_PARAMS = [ pytest.param( p, @@ -79,7 +82,7 @@ PROVIDER_PARAMS = [ ) ), ) - for p in PROVIDERS if p.id != "vertex_ai" + for p in PROVIDERS if p.id not in ("vertex_ai", "bedrock") ] WEATHER_TOOL = ToolsSchema( diff --git a/tests/integration/_support/async_server.py b/tests/integration/_support/async_server.py index 324cbaa5da4..53dcbc8d92d 100644 --- a/tests/integration/_support/async_server.py +++ b/tests/integration/_support/async_server.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import ssl from collections.abc import Awaitable, Callable from contextlib import suppress from threading import Thread @@ -16,10 +17,11 @@ async def close_writer(writer: asyncio.StreamWriter) -> None: class LoopbackServer: - def __init__(self, port: int, handler: Handler) -> None: + def __init__(self, port: int, handler: Handler, ssl_context: ssl.SSLContext | None = None) -> None: self.loop: Final = asyncio.new_event_loop() self._port: Final = port self._handler: Final = handler + self._ssl_context: Final = ssl_context self._server: asyncio.Server | None = None self._thread: Final = Thread(target=self._run, name=f"integration-relay-{port}", daemon=True) self._connections: tuple[asyncio.Task[None], ...] = () @@ -42,7 +44,7 @@ class LoopbackServer: finally: self.loop.call_soon_threadsafe(self.loop.stop) self._thread.join(10) - assert not self._thread.is_alive(), f"Database relay thread on port {self._port} did not stop" + assert not self._thread.is_alive(), f"Loopback server thread on port {self._port} did not stop" def _run(self) -> None: asyncio.set_event_loop(self.loop) @@ -53,7 +55,7 @@ class LoopbackServer: self.loop.close() async def _bind(self) -> None: - self._server = await asyncio.start_server(self._serve, "127.0.0.1", self._port) + self._server = await asyncio.start_server(self._serve, "127.0.0.1", self._port, ssl=self._ssl_context) async def _serve(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: task: Final = asyncio.current_task() diff --git a/tests/integration/_support/bedrock_realtime_peer.py b/tests/integration/_support/bedrock_realtime_peer.py new file mode 100644 index 00000000000..d30bf1b92f5 --- /dev/null +++ b/tests/integration/_support/bedrock_realtime_peer.py @@ -0,0 +1,429 @@ +"""A scripted Bedrock Nova Sonic realtime peer. + +The handler dials ``api_base`` with ``invoke_model_with_bidirectional_stream`` through +``AWSCRTHTTPClient``, which only speaks h2 over TLS ALPN, so this peer is a loopback TLS +server offering ``h2`` on ``async_server.LoopbackServer``'s thread-owned asyncio loop. +It answers ``POST /model/{id}/invoke-with-bidirectional-stream``, decodes the client's +signed event-stream chunks into the ``{"event": {...}}`` documents Bedrock would read, +and answers each consumed ``Turn`` with event-stream ``chunk`` frames carrying +``{"bytes": }`` payloads, the reverse of what the SDK produces. + +Turn triggers are inbound events. ``user_input`` fires on the ``contentEnd`` closing the +first inbound USER-role TEXT content block, the end of a forwarded conversation.item.create +text message (the injected trigger-audio block is USER/AUDIO, so it never trips it). +``tool_result`` fires on the Nth inbound ``toolResult`` event, which the proxy sends inside +a TOOL contentStart/toolResult/contentEnd triple whose ``toolResultInputConfiguration`` +``toolUseId`` names the call being answered. +""" + +from __future__ import annotations + +import asyncio +import base64 +import json +import struct +import uuid +from collections.abc import Iterator, Mapping +from contextlib import contextmanager, suppress +from dataclasses import dataclass +from itertools import chain +from pathlib import Path +from queue import SimpleQueue +from typing import Final, Literal + +import h2.config +import h2.connection +import h2.events +import h2.exceptions +from integration._support.async_server import Handler, LoopbackServer +from integration._support.process import _free_port +from integration._support.tls import server_context +from integration._support.upstream import aws_event_stream_frame +from pydantic import JsonValue, TypeAdapter + +BedrockEvent = Mapping[str, JsonValue] + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +def _ids() -> tuple[str, str, str, str]: + return (uuid.uuid4().hex, uuid.uuid4().hex, uuid.uuid4().hex, uuid.uuid4().hex) + + +def _block(prompt: str, content: str, session: str, completion: str) -> dict[str, JsonValue]: + return {"promptName": prompt, "contentName": content, "sessionId": session, "completionId": completion} + + +def user_transcript(text: str) -> tuple[BedrockEvent, ...]: + """The USER ASR transcript block Bedrock emits once it has consumed the user turn.""" + prompt, content, session, completion = _ids() + return ( + { + "event": { + "contentStart": { + **_block(prompt, content, session, completion), + "role": "USER", + "type": "TEXT", + "interactive": False, + "textOutputConfiguration": {"mediaType": "text/plain"}, + "additionalModelFields": json.dumps({"generationStage": "FINAL"}), + } + } + }, + {"event": {"textOutput": {**_block(prompt, content, session, completion), "role": "USER", "content": text}}}, + { + "event": { + "contentEnd": { + **_block(prompt, content, session, completion), + "stopReason": "PARTIAL_TURN", + "type": "TEXT", + } + } + }, + ) + + +def tool_use(tool_use_id: str, name: str, content: str | None) -> tuple[BedrockEvent, ...]: + """The role-TOOL contentStart/toolUse/contentEnd triple Bedrock emits per tool call.""" + prompt, content_name, session, completion = _ids() + tool_use_body: Final[dict[str, JsonValue]] = { + **_block(prompt, content_name, session, completion), + "role": "TOOL", + "toolName": name, + "toolUseId": tool_use_id, + } + if content is not None: + tool_use_body["content"] = content + return ( + { + "event": { + "contentStart": { + **_block(prompt, content_name, session, completion), + "role": "TOOL", + "type": "TOOL", + "toolUseOutputConfiguration": {"mediaType": "application/json"}, + } + } + }, + {"event": {"toolUse": tool_use_body}}, + { + "event": { + "contentEnd": { + **_block(prompt, content_name, session, completion), + "stopReason": "TOOL_USE", + "type": "TOOL", + } + } + }, + ) + + +def assistant_text(text: str, stop_reason: str = "END_TURN") -> tuple[BedrockEvent, ...]: + """An ASSISTANT text block (SPECULATIVE stage) closed with `stop_reason`.""" + prompt, content, session, completion = _ids() + return ( + { + "event": { + "contentStart": { + **_block(prompt, content, session, completion), + "role": "ASSISTANT", + "type": "TEXT", + "interactive": False, + "textOutputConfiguration": {"mediaType": "text/plain"}, + "additionalModelFields": json.dumps({"generationStage": "SPECULATIVE"}), + } + } + }, + { + "event": { + "textOutput": {**_block(prompt, content, session, completion), "role": "ASSISTANT", "content": text} + } + }, + { + "event": { + "contentEnd": { + **_block(prompt, content, session, completion), + "stopReason": stop_reason, + "type": "TEXT", + } + } + }, + ) + + +def usage(input_tokens: int, output_tokens: int) -> BedrockEvent: + total: Final = { + "input": {"speechTokens": 0, "textTokens": input_tokens}, + "output": {"speechTokens": 0, "textTokens": output_tokens}, + } + return { + "event": { + "usageEvent": { + "promptName": uuid.uuid4().hex, + "sessionId": uuid.uuid4().hex, + "completionId": uuid.uuid4().hex, + "details": {"delta": total["output"], "total": total}, + "totalInputTokens": input_tokens, + "totalOutputTokens": output_tokens, + "totalTokens": input_tokens + output_tokens, + } + } + } + + +@dataclass(frozen=True, slots=True) +class Turn: + after: Literal["user_input", "tool_result"] + events: tuple[BedrockEvent, ...] + + +def _event_body(event: BedrockEvent) -> tuple[str, Mapping[str, JsonValue]]: + body: Final = event["event"] + assert isinstance(body, Mapping) and len(body) == 1, event + name: Final = next(iter(body)) + fields: Final = body[name] + assert isinstance(fields, Mapping), event + return name, fields + + +def _outbound_chunk(event: BedrockEvent) -> bytes: + wrapped: Final = json.dumps({"bytes": base64.b64encode(json.dumps(event).encode()).decode()}).encode() + return aws_event_stream_frame( + {":event-type": "chunk", ":content-type": "application/json", ":message-type": "event"}, wrapped + ) + + +def _parse_frames(buffer: bytes) -> tuple[tuple[bytes, ...], bytes]: + """Slice complete event-stream frames out of `buffer`; return (payloads, remainder).""" + payloads: Final[list[bytes]] = [] + view = memoryview(buffer) + while len(view) >= 12: + total, header_length = struct.unpack_from("!II", view, 0) + if total > len(view): + break + payloads.append(view[12 + header_length : total - 4].tobytes()) + view = view[total:] + return tuple(payloads), view.tobytes() + + +def _frame_payload(raw: bytes) -> bytes: + """The payload slice of one complete framed message (prelude + headers + payload + CRC).""" + if len(raw) < 16: + return b"" + _total, header_length = struct.unpack_from("!II", raw, 0) + return raw[12 + header_length : -4] + + +def _decode_inbound(chunk_payload: bytes) -> BedrockEvent | None: + """One signed chunk frame's payload -> the inner chunk's {"bytes": ...} event.""" + inner_payload: Final = _frame_payload(chunk_payload) + try: + wrapped: Final = _JSON_OBJECT.validate_json(inner_payload) + except ValueError: + return None + encoded: Final = wrapped.get("bytes") + if not isinstance(encoded, str): + return None + try: + return _JSON_OBJECT.validate_json(base64.b64decode(encoded)) + except ValueError: + return None + + +class _Inbound: + """Inbound event bookkeeping for turn triggers and accessors.""" + + def __init__(self) -> None: + self.events: Final = SimpleQueue[BedrockEvent]() + self.open_blocks: Final[dict[str, tuple[str, str]]] = {} + self.tool_use_ids: Final[dict[str, str]] = {} + self.tool_results_seen = 0 + self.user_text_closed = False + + def record(self, event: BedrockEvent) -> None: + self.events.put(event) + name, fields = _event_body(event) + content_name = fields.get("contentName") + if name == "contentStart" and isinstance(content_name, str): + self.open_blocks[content_name] = (str(fields.get("role")), str(fields.get("type"))) + configuration = fields.get("toolResultInputConfiguration") + if isinstance(configuration, Mapping) and isinstance(configuration.get("toolUseId"), str): + self.tool_use_ids[content_name] = configuration["toolUseId"] + elif name == "contentEnd" and isinstance(content_name, str): + role_type: Final = self.open_blocks.pop(content_name, ("", "")) + if role_type == ("USER", "TEXT"): + self.user_text_closed = True + elif name == "toolResult": + self.tool_results_seen += 1 + + +class _Script: + """Consumes turns in order as their inbound triggers arrive.""" + + def __init__(self, turns: tuple[Turn, ...]) -> None: + self._turns: Final = turns + self._next = 0 + + @property + def exhausted(self) -> bool: + return self._next >= len(self._turns) + + def ready_events(self, state: _Inbound, event: BedrockEvent) -> tuple[BedrockEvent, ...]: + if self._next >= len(self._turns): + return () + turn: Final = self._turns[self._next] + if turn.after == "user_input": + if not state.user_text_closed: + return () + else: + name, _fields = _event_body(event) + if name != "toolResult": + return () + self._next += 1 + return turn.events + + +class _PeerState: + def __init__(self, turns: tuple[Turn, ...]) -> None: + self.inbound: Final = _Inbound() + self.script: Final = _Script(turns) + + +@dataclass(frozen=True, slots=True) +class BedrockRealtimePeer: + api_base: str + _state: _PeerState + _drained: list[BedrockEvent] + + def received(self) -> tuple[BedrockEvent, ...]: + while not self._state.inbound.events.empty(): + self._drained.append(self._state.inbound.events.get()) # drained log, not a built value + return tuple(self._drained) + + def tool_specs(self) -> tuple[Mapping[str, JsonValue], ...]: + return tuple(chain.from_iterable(_tool_specs(event) for event in self.received())) + + def tool_results(self) -> tuple[tuple[str, str], ...]: + """(toolUseId, content) per inbound toolResult in arrival order.""" + return tuple(result for event in self.received() if (result := self._tool_result(event)) is not None) + + def _tool_result(self, event: BedrockEvent) -> tuple[str, str] | None: + name, fields = _event_body(event) + if name != "toolResult": + return None + content_name: Final = fields.get("contentName") + content: Final = fields.get("content") + if not isinstance(content_name, str) or not isinstance(content, str): + return None + return (self._state.inbound.tool_use_ids.get(content_name, ""), content) + + +def _tool_specs(event: BedrockEvent) -> tuple[Mapping[str, JsonValue], ...]: + name, fields = _event_body(event) + if name != "promptStart": + return () + configuration: Final = fields.get("toolConfiguration") + if not isinstance(configuration, Mapping): + return () + tools: Final = configuration.get("tools") + if not isinstance(tools, list): + return () + return tuple( + entry["toolSpec"] + for entry in tools + if isinstance(entry, Mapping) and isinstance(entry.get("toolSpec"), Mapping) + ) + + +class _H2Session: + """One bidirectional-stream connection: decode inbound, answer scripted turns.""" + + def __init__(self, state: _PeerState, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + self._state: Final = state + self._reader: Final = reader + self._writer: Final = writer + self._h2: Final = h2.connection.H2Connection(config=h2.config.H2Configuration(client_side=False)) + self._buffers: Final[dict[int, bytes]] = {} + self._ended: set[int] = set() + self._closed = False + + async def run(self) -> None: + self._h2.initiate_connection() + self._writer.write(self._h2.data_to_send()) + await self._writer.drain() + while True: + incoming: Final = await self._reader.read(65536) + if not incoming: + return + for event in self._h2.receive_data(incoming): + self._apply(event) + pending: Final = self._h2.data_to_send() + if pending: + self._writer.write(pending) + await self._writer.drain() + if self._closed: + self._writer.close() + return + + def _apply(self, event: h2.events.Event) -> None: + if isinstance(event, h2.events.RequestReceived): + self._h2.send_headers( + event.stream_id, + [(":status", "200"), ("content-type", "application/vnd.amazon.eventstream")], + ) + elif isinstance(event, h2.events.DataReceived): + self._h2.acknowledge_received_data(event.flow_controlled_length, event.stream_id) + self._buffers[event.stream_id] = self._buffers.get(event.stream_id, b"") + event.data + payloads, remainder = _parse_frames(self._buffers[event.stream_id]) + self._buffers[event.stream_id] = remainder + for payload in payloads: + decoded: Final = _decode_inbound(payload) + if decoded is None: + continue + self._state.inbound.record(decoded) + if event.stream_id not in self._ended: + for outbound in self._state.script.ready_events(self._state.inbound, decoded): + self._h2.send_data(event.stream_id, _outbound_chunk(outbound)) + elif isinstance(event, (h2.events.StreamEnded, h2.events.StreamReset)): + self._close_stream(event.stream_id) + elif isinstance(event, h2.events.ConnectionTerminated): + self._writer.close() + + def _close_stream(self, stream_id: int) -> None: + self._end_stream(stream_id) + self._closed = True + self._h2.close_connection() + + def _end_stream(self, stream_id: int) -> None: + if stream_id in self._ended: + return + self._ended.add(stream_id) + with suppress(h2.exceptions.ProtocolError): + self._h2.end_stream(stream_id) + + +def _session_handler(state: _PeerState) -> Handler: + async def serve(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + with suppress(ConnectionError, OSError, asyncio.IncompleteReadError, h2.exceptions.StreamClosedError): + await _H2Session(state, reader, writer).run() + + return serve + + +@contextmanager +def bedrock_realtime_peer(turns: tuple[Turn, ...], cert: tuple[Path, Path]) -> Iterator[BedrockRealtimePeer]: + """Serve scripted `turns` on a loopback TLS h2 endpoint; yield the peer. + + `cert` is the self-signed pair the owned proxy already trusts via SSL_CERT_FILE. + SigV4 headers are ignored; the stream closes cleanly when the client ends it. + """ + port: Final = _free_port() + context: Final = server_context(cert[0], cert[1]) + context.set_alpn_protocols(["h2"]) + state: Final = _PeerState(turns) + server: Final = LoopbackServer(port, _session_handler(state), ssl_context=context) + server.start() + try: + yield BedrockRealtimePeer(api_base=f"https://localhost:{port}", _state=state, _drained=[]) + finally: + server.stop() diff --git a/tests/integration/providers/test_bedrock_nova_sonic_realtime_tool_calls.py b/tests/integration/providers/test_bedrock_nova_sonic_realtime_tool_calls.py new file mode 100644 index 00000000000..ccef839a163 --- /dev/null +++ b/tests/integration/providers/test_bedrock_nova_sonic_realtime_tool_calls.py @@ -0,0 +1,378 @@ +"""Bedrock Nova Sonic realtime tool calls through a scripted h2 event-stream peer (LIT-9465). + +Nova Sonic returns a tool call as a role-TOOL contentStart/toolUse/contentEnd triple with the +arguments as a JSON string in ``content`` and no ASSISTANT block before it, and emits one such +triple per call including the two-call turn Nova 2 Sonic produces per question. The cells here +serve those triples from ``_support/bedrock_realtime_peer.py`` (the transport the handler really +dials: ``AWSCRTHTTPClient`` -> TLS ALPN h2 -> invoke-with-bidirectional-stream) and assert both +sides: the OpenAI realtime events the websocket client receives, and the toolSpec / toolResult +documents the proxy forwarded to the peer. + +The file boots its own proxy instead of the shared gateway so it can pass ``SSL_CERT_FILE`` +pointing at the peer's self-signed cert, which is the only way the CRT h2 client trusts it. +""" + +from __future__ import annotations + +import asyncio +import json +import time +import uuid +from collections.abc import AsyncIterator, Callable, Iterator, Mapping +from contextlib import asynccontextmanager +from pathlib import Path +from typing import Final + +import pytest +import websockets +from integration._support.bedrock_realtime_peer import ( + BedrockRealtimePeer, + Turn, + assistant_text, + bedrock_realtime_peer, + tool_use, + usage, + user_transcript, +) +from integration._support.client import ( + JSON_OBJECT, + Gateway, + Scenario, + eventually, + gateway_from_environment, + object_value, + string_value, +) +from integration._support.process import owned_proxy +from integration._support.tls import write_self_signed_cert +from pydantic import JsonValue +from websockets.asyncio.client import ClientConnection + +pytestmark: Final = pytest.mark.timeout(180) + +DEPLOYMENT: Final = "bedrock/amazon.nova-2-sonic-v1:0" +WEATHER_PARAMETERS: Final = {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]} +WEATHER_FUNCTION: Final = { + "name": "get_weather", + "description": "Get current weather for a city", + "parameters": WEATHER_PARAMETERS, +} +NESTED_TOOL: Final = {"type": "function", "function": WEATHER_FUNCTION} +FLAT_TOOL: Final = {"type": "function", **WEATHER_FUNCTION} +TOOL_RESULT_OUTPUT: Final = '{"temperature_f":72,"conditions":"sunny"}' + + +@pytest.fixture(scope="module") +def cert(tmp_path_factory: pytest.TempPathFactory) -> tuple[Path, Path]: + return write_self_signed_cert(tmp_path_factory.mktemp("bedrock-sonic-cert")) + + +@pytest.fixture(scope="module") +def candidate(tmp_path_factory: pytest.TempPathFactory, cert: tuple[Path, Path]) -> Iterator[Gateway]: + with gateway_from_environment() as base: + with owned_proxy(base, tmp_path_factory.mktemp("bedrock-sonic"), {"SSL_CERT_FILE": str(cert[0])}) as proxy: + yield proxy + + +def _deployment(gateway: Gateway, scenario: Scenario, peer: BedrockRealtimePeer) -> str: + name: Final = f"bedrock-sonic-{uuid.uuid4().hex[:12]}" + created: Final = gateway.post( + "/model/new", + { + "model_name": name, + "litellm_params": { + "model": DEPLOYMENT, + "aws_region_name": "us-east-1", + "aws_access_key_id": "integration-dummy-key", + "aws_secret_access_key": "integration-dummy-secret", + "api_base": peer.api_base, + }, + "model_info": {"mode": "realtime"}, + }, + ) + scenario.cleanups.callback(scenario.delete_model, string_value(object_value(created["model_info"])["id"])) + return name + + +def _session_update(tool: Mapping[str, JsonValue]) -> str: + return json.dumps( + { + "type": "session.update", + "session": { + "instructions": "Always use the get_weather tool for weather questions.", + "tools": [tool], + }, + } + ) + + +def _user_text() -> str: + return json.dumps( + { + "type": "conversation.item.create", + "item": { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "What is the weather in New York right now?"}], + }, + } + ) + + +_RESPONSE_CREATE: Final = json.dumps({"type": "response.create"}) + + +def _tool_result(call_id: str) -> str: + return json.dumps( + { + "type": "conversation.item.create", + "item": {"type": "function_call_output", "call_id": call_id, "output": TOOL_RESULT_OUTPUT}, + } + ) + + +Event = dict[str, JsonValue] +EventPredicate = Callable[[Event, tuple[Event, ...]], bool] + + +async def _events_until( + socket: ClientConnection, predicate: EventPredicate, *, timeout: float = 60 +) -> tuple[Event, ...]: + collected: Final[list[Event]] = [] + deadline: Final = time.monotonic() + timeout + while True: + try: + raw: Final = await asyncio.wait_for(socket.recv(), deadline - time.monotonic()) + except (TimeoutError, asyncio.TimeoutError): + raise AssertionError(f"timed out waiting; types={[event.get('type') for event in collected]}") + event: Final = JSON_OBJECT.validate_json(raw) + collected.append(event) + if predicate(event, tuple(collected)): + return tuple(collected) + + +def _is_done(events: tuple[Event, ...]) -> bool: + return any(event.get("type") == "response.done" for event in events) + + +def _function_calls(events: tuple[Event, ...]) -> tuple[Event, ...]: + return tuple(event for event in events if event.get("type") == "response.function_call_arguments.done") + + +def _response_dones(events: tuple[Event, ...]) -> tuple[Event, ...]: + return tuple(event for event in events if event.get("type") == "response.done") + + +def _done_ids(events: tuple[Event, ...]) -> tuple[str, ...]: + return tuple(string_value(object_value(event["response"])["id"]) for event in _response_dones(events)) + + +def _tail_types(events: tuple[Event, ...]) -> tuple[str, ...]: + """Event types from the first response.created onward.""" + start: Final = next(index for index, event in enumerate(events) if event.get("type") == "response.created") + return tuple(string_value(event["type"]) for event in events[start:]) + + +@asynccontextmanager +async def _realtime(gateway: Gateway, model: str, key: str) -> AsyncIterator[ClientConnection]: + proxy: Final = str(gateway.client.base_url).rstrip("/").replace("http://", "ws://") + async with websockets.connect( + f"{proxy}/v1/realtime?model={model}", additional_headers={"Authorization": f"Bearer {key}"} + ) as socket: + yield socket + + +async def _open_session(socket: ClientConnection, tool: Mapping[str, JsonValue]) -> None: + created: Final = JSON_OBJECT.validate_json(await socket.recv()) + assert created.get("type") == "session.created", created + await socket.send(_session_update(tool)) + await socket.send(_user_text()) + await socket.send(_RESPONSE_CREATE) + + +async def _one_call_session(gateway: Gateway, model: str, key: str, tool: Mapping[str, JsonValue]) -> tuple[Event, ...]: + async with _realtime(gateway, model, key) as socket: + await _open_session(socket, tool) + return await _events_until(socket, lambda _event, events: _is_done(events)) + + +def test_nested_tool_call_before_assistant_block(candidate: Gateway, cert: tuple[Path, Path]) -> None: + turn: Final = Turn( + "user_input", user_transcript("ready") + tool_use("tooluse_A", "get_weather", '{"location":"New York"}') + ) + with bedrock_realtime_peer((turn,), cert) as peer, candidate.scenario() as scenario: + model: Final = _deployment(candidate, scenario, peer) + key: Final = scenario.key() + events: Final = asyncio.run(_one_call_session(candidate, model, key, NESTED_TOOL)) + assert _tail_types(events) == ( + "response.created", + "response.output_item.added", + "response.function_call_arguments.done", + "response.output_item.done", + "response.done", + ), events + (call,) = _function_calls(events) + assert string_value(call["call_id"]) == "tooluse_A" + assert string_value(call["name"]) == "get_weather" + assert string_value(call["arguments"]) == '{"location":"New York"}' + (done,) = _response_dones(events) + output: Final = object_value(done["response"])["output"] + assert isinstance(output, list) and len(output) == 1, done + item: Final = object_value(output[0]) + assert string_value(item["type"]) == "function_call" + assert string_value(item["call_id"]) == "tooluse_A" + assert string_value(item["arguments"]) == '{"location":"New York"}' + specs: Final = eventually(lambda: peer.tool_specs(), lambda values: len(values) == 1, seconds=10) + assert string_value(specs[0]["name"]) == "get_weather" + assert json.loads(string_value(object_value(specs[0]["inputSchema"])["json"])) == WEATHER_PARAMETERS + + +def test_flat_tool_shape(candidate: Gateway, cert: tuple[Path, Path]) -> None: + turn: Final = Turn( + "user_input", user_transcript("ready") + tool_use("tooluse_B", "get_weather", '{"location":"New York"}') + ) + with bedrock_realtime_peer((turn,), cert) as peer, candidate.scenario() as scenario: + model: Final = _deployment(candidate, scenario, peer) + key: Final = scenario.key() + events: Final = asyncio.run(_one_call_session(candidate, model, key, FLAT_TOOL)) + (call,) = _function_calls(events) + assert string_value(call["call_id"]) == "tooluse_B" + assert string_value(call["name"]) == "get_weather" + specs: Final = eventually(lambda: peer.tool_specs(), lambda values: len(values) == 1, seconds=10) + assert string_value(specs[0]["name"]) == "get_weather" + assert string_value(specs[0]["description"]) == "Get current weather for a city" + assert json.loads(string_value(object_value(specs[0]["inputSchema"])["json"])) == WEATHER_PARAMETERS + + +async def _tool_result_round_trip(gateway: Gateway, model: str, key: str) -> tuple[Event, ...]: + async with _realtime(gateway, model, key) as socket: + await _open_session(socket, NESTED_TOOL) + first: Final = await _events_until(socket, lambda _event, events: _is_done(events)) + (call,) = _function_calls(first) + await socket.send(_tool_result(string_value(call["call_id"]))) + await socket.send(_RESPONSE_CREATE) + second: Final = await _events_until(socket, lambda _event, events: _is_done(events)) + return first + second + + +def test_tool_result_round_trip(candidate: Gateway, cert: tuple[Path, Path]) -> None: + turns: Final = ( + Turn("user_input", user_transcript("ready") + tool_use("tooluse_A", "get_weather", '{"location":"New York"}')), + Turn("tool_result", assistant_text("It is 72 and sunny") + (usage(30, 12),)), + ) + with bedrock_realtime_peer(turns, cert) as peer, candidate.scenario() as scenario: + model: Final = _deployment(candidate, scenario, peer) + key: Final = scenario.key() + events: Final = asyncio.run(_tool_result_round_trip(candidate, model, key)) + dones: Final = _response_dones(events) + assert len(dones) == 2, events + first_id, second_id = _done_ids(events) + assert first_id != second_id + second_created: Final = tuple(event for event in events if event.get("type") == "response.created") + assert len(second_created) == 2 and object_value(second_created[1]["response"])["id"] == second_id + text: Final = "".join( + string_value(event["delta"]) for event in events if event.get("type") == "response.text.delta" + ) + assert "72" in text + last_output: Final = object_value(dones[-1]["response"])["output"] + assert tuple(object_value(item).get("type") for item in last_output) == ("message",) + (answer_item,) = last_output + assert object_value(answer_item)["content"] == [{"type": "text", "text": "It is 72 and sunny"}] + results: Final = eventually(lambda: peer.tool_results(), lambda values: len(values) == 1, seconds=10) + assert results == (("tooluse_A", TOOL_RESULT_OUTPUT),) + + +async def _two_call_session(gateway: Gateway, model: str, key: str) -> tuple[Event, ...]: + async with _realtime(gateway, model, key) as socket: + await _open_session(socket, NESTED_TOOL) + events: Final[tuple[Event, ...]] = await _events_until( + socket, lambda _event, collected: len(_response_dones(collected)) == 2 + ) + for call in _function_calls(events): + await socket.send(_tool_result(string_value(call["call_id"]))) + after_answer: Final = await _events_until(socket, lambda _event, collected: _is_done(collected)) + return events + after_answer + + +def test_two_tool_calls_one_turn(candidate: Gateway, cert: tuple[Path, Path]) -> None: + args: Final = '{"location":"New York"}' + turns: Final = ( + Turn( + "user_input", + user_transcript("ready") + + tool_use("tooluse_A", "get_weather", args) + + tool_use("tooluse_B", "get_weather", args), + ), + Turn("tool_result", assistant_text("It is 72 and sunny") + (usage(40, 15),)), + ) + with bedrock_realtime_peer(turns, cert) as peer, candidate.scenario() as scenario: + model: Final = _deployment(candidate, scenario, peer) + key: Final = scenario.key() + events: Final = asyncio.run(_two_call_session(candidate, model, key)) + calls: Final = _function_calls(events) + assert len(calls) == 2, events + assert tuple(string_value(call["call_id"]) for call in calls) == ("tooluse_A", "tooluse_B") + assert tuple(string_value(call["arguments"]) for call in calls) == (args, args) + item_ids: Final = tuple(string_value(call["item_id"]) for call in calls) + response_ids: Final = tuple(string_value(call["response_id"]) for call in calls) + assert len(set(item_ids)) == 2 and len(set(response_ids)) == 2 + first_dones: Final = _response_dones(events)[:2] + for done, call_id in zip(first_dones, ("tooluse_A", "tooluse_B")): + output: Final = object_value(done["response"])["output"] + assert isinstance(output, list) and len(output) == 1, done + assert string_value(object_value(output[0])["call_id"]) == call_id + results: Final = eventually(lambda: peer.tool_results(), lambda values: len(values) == 2, seconds=10) + assert tuple(result[0] for result in results) == ("tooluse_A", "tooluse_B") + assert all(result[1] == TOOL_RESULT_OUTPUT for result in results) + text: Final = "".join( + string_value(event["delta"]) for event in events if event.get("type") == "response.text.delta" + ) + assert "72" in text + + +def test_tool_call_inside_active_assistant_response(candidate: Gateway, cert: tuple[Path, Path]) -> None: + start, output, _end = assistant_text("It is 72 and sunny") + turn: Final = Turn( + "user_input", + (start, output) + tool_use("tooluse_C", "get_weather", '{"location":"New York"}') + (_end,), + ) + with bedrock_realtime_peer((turn,), cert) as peer, candidate.scenario() as scenario: + model: Final = _deployment(candidate, scenario, peer) + key: Final = scenario.key() + events: Final = asyncio.run(_one_call_session(candidate, model, key, NESTED_TOOL)) + created: Final = tuple(event for event in events if event.get("type") == "response.created") + assert len(created) == 1, events + response_id: Final = string_value(object_value(created[0]["response"])["id"]) + (call,) = _function_calls(events) + assert string_value(call["call_id"]) == "tooluse_C" + assert string_value(call["response_id"]) == response_id + assert string_value(call["arguments"]) == '{"location":"New York"}' + assert call["output_index"] == 1, events + assert tuple(string_value(e["text"]) for e in events if e.get("type") == "response.text.done") == ( + "It is 72 and sunny", + ), events + message_done_indexes: Final = tuple( + i + for i, event in enumerate(events) + if event.get("type") == "response.output_item.done" and object_value(event["item"]).get("type") == "message" + ) + response_done_indexes: Final = tuple( + i for i, event in enumerate(events) if event.get("type") == "response.done" + ) + assert len(response_done_indexes) == 1, events + assert len(message_done_indexes) == 1 and response_done_indexes[0] > message_done_indexes[0], events + + +@pytest.mark.parametrize( + "content", + ["", None, "not json", '{"location":"' + "x" * 5000 + '"}'], + ids=["empty", "missing", "not-json", "five-kb"], +) +def test_malformed_tool_content(candidate: Gateway, cert: tuple[Path, Path], content: str | None) -> None: + turn: Final = Turn("user_input", tool_use("tooluse_M", "get_weather", content)) + with bedrock_realtime_peer((turn,), cert) as peer, candidate.scenario() as scenario: + model: Final = _deployment(candidate, scenario, peer) + key: Final = scenario.key() + events: Final = asyncio.run(_one_call_session(candidate, model, key, NESTED_TOOL)) + (call,) = _function_calls(events) + assert string_value(call["arguments"]) == (content if content is not None else "{}"), events diff --git a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py index b1836cb6c46..8f95a502137 100644 --- a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py @@ -3,6 +3,7 @@ import json import sys import types from types import SimpleNamespace +from typing import Final from unittest.mock import MagicMock import pytest @@ -68,6 +69,20 @@ class FailOnPromptEndStream(FakeInputStream): self.sent.append(event) +class HangingCloseInputStream(FakeInputStream): + """Models a CRT input stream on an already-ended h2 stream: close sends and close() never complete""" + + async def send(self, event): + payload = json.loads(event.value.bytes_.decode("utf-8")) + if "promptEnd" in payload.get("event", {}) or "sessionEnd" in payload.get("event", {}): + await asyncio.Future() + else: + self.sent.append(event) + + async def close(self): + await asyncio.Future() + + class FakeBedrockStream: def __init__(self, input_stream=None): self.input_stream = input_stream if input_stream is not None else FakeInputStream() @@ -417,6 +432,21 @@ class TestBedrockRealtimeHandler: assert stream.input_stream.closed + @pytest.mark.asyncio + async def test_close_flush_is_bounded_when_input_stream_never_completes(self, stub_aws_models): + handler = BedrockRealtime(close_timeout=0.05) + config = BedrockRealtimeConfig() + stream = FakeBedrockStream(input_stream=HangingCloseInputStream()) + client_ws = DisconnectingClientWS( + [json.dumps({"type": "session.update", "session": {"instructions": "You are helpful."}})] + ) + + await asyncio.wait_for( + handler._forward_client_to_bedrock(client_ws, stream, config, "amazon.nova-sonic-v1:0", {}), 2 + ) + + assert not stream.input_stream.closed + @pytest.mark.asyncio async def test_close_flush_continues_after_partial_send_failure(self, stub_aws_models): handler = BedrockRealtime() @@ -613,6 +643,78 @@ class TestBedrockRealtimeSessionLifecycle: assert client_ws.sent_to_client == [] + @pytest.mark.asyncio + async def test_conflicting_session_update_gets_error_and_sends_nothing_to_bedrock(self, stub_aws_models): + tool: Final = {"type": "function", "name": "get_weather", "parameters": {"type": "object"}} + other_tool: Final = {"type": "function", "name": "get_time", "parameters": {"type": "object"}} + + def update(tools) -> str: + return json.dumps( + { + "type": "session.update", + "event_id": "evt_update_1", + "session": {"tools": tools, "modalities": ["text"]}, + } + ) + + handler = BedrockRealtime() + config = BedrockRealtimeConfig() + stream = FakeBedrockStream() + client_ws = DisconnectingClientWS([update([tool]), update([other_tool]), update([tool])]) + + await handler._forward_client_to_bedrock(client_ws, stream, config, "amazon.nova-sonic-v1:0", {}, FakeLogging()) + + acked: Final = [json.loads(message) for message in client_ws.sent_to_client] + assert [event["type"] for event in acked] == ["session.updated", "error", "session.updated"] + error: Final = acked[1]["error"] + assert error["code"] == "session_update_not_applied" + assert error["event_id"] == "evt_update_1" + prompt_starts: Final = [ + payload + for payload in (json.loads(chunk.value.bytes_.decode("utf-8")) for chunk in stream.input_stream.sent) + if "promptStart" in payload.get("event", {}) + ] + assert len(prompt_starts) == 1 + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "drop_params,expected_errors", + [(False, [("invalid_type", "session.temperature")]), (True, [])], + ) + async def test_malformed_number_setting_is_rejected_or_dropped_and_the_session_still_answers( + self, stub_aws_sdk_client, drop_params, expected_errors + ): + stream: Final = ScriptedBedrockStream([]) + stub_aws_sdk_client["streams"] = [stream] + user_text: Final = {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]} + websocket: Final = ConnectedClientWS( + [ + json.dumps({"type": "session.update", "session": {"temperature": "0.2", "modalities": ["text"]}}), + json.dumps({"type": "conversation.item.create", "item": user_text}), + ] + ) + + await BedrockRealtime().async_realtime( + model="amazon.nova-sonic-v1:0", + websocket=websocket, + logging_obj=FakeLogging(), + aws_region_name="us-east-1", + aws_access_key_id="k", + aws_secret_access_key="s", + drop_params=drop_params, + ) + + client_events: Final = [json.loads(message) for message in websocket.sent_to_client] + errors: Final = [event["error"] for event in client_events if event["type"] == "error"] + sent_to_bedrock: Final = [json.loads(chunk.value.bytes_.decode("utf-8")) for chunk in stream.input_stream.sent] + session_starts: Final = [ + payload["event"]["sessionStart"] for payload in sent_to_bedrock if "sessionStart" in payload["event"] + ] + assert [(error["code"], error["param"]) for error in errors] == expected_errors + assert [start["inferenceConfiguration"]["temperature"] for start in session_starts] == [0.7] + assert next(iter(sent_to_bedrock[0]["event"])) == "sessionStart" + assert any(payload["event"].get("textInput", {}).get("content") == "hi" for payload in sent_to_bedrock) + class TestBedrockRealtimeProviderFailurePropagation: """Deferred Nova Sonic failures must escape async_realtime so the router can fall back / cool down (LIT-6484)""" diff --git a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_transformation.py b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_transformation.py index 9d034c9838c..1cc135a8671 100644 --- a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_transformation.py +++ b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_transformation.py @@ -15,6 +15,9 @@ from litellm.llms.bedrock.realtime.transformation import ( from litellm.llms.bedrock.realtime.trigger_audio import ready_trigger_pcm from litellm.types.llms.openai import OpenAIRealtimeEventTypes +OBJECT_PARAMETERS: Final = {"type": "object"} +CONFLICT_TOOL: Final = {"type": "function", "name": "get_weather", "parameters": OBJECT_PARAMETERS} + class TestBedrockRealtimeConfig: """Test suite for BedrockRealtimeConfig class""" @@ -110,6 +113,50 @@ class TestBedrockRealtimeConfig: assert schema["type"] == "object" assert "location" in schema["properties"] + def test_transform_tools_flat_realtime_shape(self): + """The flat OpenAI Realtime tool shape {type: function, name, parameters} + produces the same toolSpec as the nested chat-completions shape (LIT-9465)""" + config = BedrockRealtimeConfig() + + nested = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get current weather for a city", + "parameters": { + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"], + }, + }, + } + ] + flat = [ + { + "type": "function", + "name": "get_weather", + "description": "Get current weather for a city", + "parameters": { + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"], + }, + } + ] + + assert config._transform_tools_to_bedrock_format(flat) == config._transform_tools_to_bedrock_format(nested) + + flat_tools = config._transform_tools_to_bedrock_format(flat) + assert flat_tools[0]["toolSpec"]["name"] == "get_weather" + assert flat_tools[0]["toolSpec"]["description"] == "Get current weather for a city" + schema = json.loads(flat_tools[0]["toolSpec"]["inputSchema"]["json"]) + assert schema["required"] == ["location"] + + session_config = config.session_configuration_request("amazon.nova-2-sonic-v1:0", tools=flat) + prompt_start = json.loads(session_config)["prompt_start"]["event"]["promptStart"] + assert prompt_start["toolConfiguration"]["tools"][0]["toolSpec"]["name"] == "get_weather" + def test_audio_format_mapping(self): """Test audio format to sample rate mapping""" config = BedrockRealtimeConfig() @@ -149,6 +196,49 @@ class TestBedrockRealtimeConfig: session_start = json.loads(messages[0]) assert session_start["event"]["sessionStart"]["inferenceConfiguration"]["temperature"] == 0.9 + def test_session_update_without_instructions_sends_empty_system_prompt(self): + config = BedrockRealtimeConfig() + session_update = { + "type": "session.update", + "session": { + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get current weather", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + }, + } + + messages = config.transform_session_update_event(session_update) + event_names = [next(iter(json.loads(message)["event"])) for message in messages] + assert event_names[:3] == ["sessionStart", "promptStart", "contentStart"] + system_content_start = json.loads(messages[2])["event"]["contentStart"] + assert system_content_start["role"] == "SYSTEM" + text_input = json.loads(messages[3])["event"]["textInput"] + assert text_input["content"] == "" + assert json.loads(messages[4])["event"]["contentEnd"]["contentName"] == system_content_start["contentName"] + + def test_session_update_instructions_become_system_prompt(self): + config = BedrockRealtimeConfig() + session_update = {"type": "session.update", "session": {"instructions": "Answer tersely."}} + + messages = config.transform_session_update_event(session_update) + text_input = json.loads(messages[3])["event"]["textInput"] + assert text_input["content"] == "Answer tersely." + + def test_second_session_update_sends_no_duplicate_prompt_start(self): + config = BedrockRealtimeConfig() + session_update = {"type": "session.update", "session": {"instructions": "You are helpful."}} + + first = config.transform_session_update_event(session_update) + assert any("promptStart" in message for message in first) + assert config.transform_session_update_event(session_update) == [] + def test_transform_session_update_with_tools(self): """Test session.update with tools""" config = BedrockRealtimeConfig() @@ -274,12 +364,14 @@ class TestBedrockRealtimeResponseCreate: "amazon.nova-sonic-v1:0", ) - def test_response_create_before_session_update_is_noop(self): + def test_response_create_before_session_update_starts_prompt_with_defaults(self): config = BedrockRealtimeConfig() messages = config.transform_realtime_request(json.dumps({"type": "response.create"}), "amazon.nova-sonic-v1:0") - assert messages == [] + events = [next(iter(json.loads(message)["event"])) for message in messages] + assert events[:5] == ["sessionStart", "promptStart", "contentStart", "textInput", "contentEnd"] + assert "audioInput" in events[5:] def test_response_create_emits_spoken_trigger_audio(self): config = BedrockRealtimeConfig() @@ -528,7 +620,9 @@ class TestBedrockRealtimeResponseTransformation: assert audio_deltas[0]["delta"] == "base64_audio_content" def test_transform_tool_use_response(self): - """Test toolUse response transformation""" + """toolUse is forwarded with its arguments even when no assistant response is + active yet (Bedrock Nova Sonic sends the call inside a role TOOL block before + any ASSISTANT contentStart, LIT-9465)""" config = BedrockRealtimeConfig() logging_obj = MagicMock() logging_obj.litellm_trace_id = "trace_123" @@ -536,9 +630,77 @@ class TestBedrockRealtimeResponseTransformation: tool_use_message = { "event": { "toolUse": { - "toolUseId": "tool_call_123", + "completionId": "comp_1", + "content": '{"location":"New York"}', + "contentId": "content_1", + "promptName": "prompt_1", + "role": "TOOL", + "sessionId": "sess_1", "toolName": "get_weather", - "input": json.dumps({"location": "San Francisco"}), + "toolUseId": "tool_call_123", + } + } + } + + result = config.transform_realtime_response( + json.dumps(tool_use_message), + "amazon.nova-sonic-v1:0", + logging_obj, + realtime_response_transform_input={ + "session_configuration_request": json.dumps({"configured": True}), + "current_output_item_id": "item_message_1", + "current_response_id": None, + "current_conversation_id": None, + "current_delta_chunks": [], + "current_item_chunks": [], + "current_delta_type": None, + }, + ) + + assert [e["type"] for e in result["response"]] == [ + "response.created", + "response.output_item.added", + "response.function_call_arguments.delta", + "response.function_call_arguments.done", + "response.output_item.done", + ] + response_created = result["response"][0] + assert response_created["response"]["id"] == result["current_response_id"] + + arguments_delta = result["response"][2] + assert arguments_delta["delta"] == '{"location":"New York"}' + assert arguments_delta["call_id"] == "tool_call_123" + assert arguments_delta["item_id"] == result["response"][1]["item"]["id"] + assert arguments_delta["response_id"] == result["current_response_id"] + + function_call = result["response"][3] + assert function_call["call_id"] == "tool_call_123" + assert function_call["name"] == "get_weather" + assert function_call["arguments"] == '{"location":"New York"}' + assert function_call["response_id"] == result["current_response_id"] + assert function_call["item_id"] != "item_message_1" + assert function_call["item_id"] == result["response"][1]["item"]["id"] + assert result["response"][1]["item"]["type"] == "function_call" + + done_item = result["response"][4]["item"] + assert done_item["type"] == "function_call" + assert done_item["status"] == "completed" + assert done_item["call_id"] == "tool_call_123" + assert done_item["arguments"] == '{"location":"New York"}' + + def test_tool_use_reuses_active_response(self): + """A toolUse arriving while an assistant response is in flight must not emit a + second response.created and must reuse the active response id""" + config = BedrockRealtimeConfig() + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "trace_123" + + tool_use_message = { + "event": { + "toolUse": { + "content": '{"location":"New York"}', + "toolName": "get_weather", + "toolUseId": "tool_call_123", } } } @@ -558,20 +720,19 @@ class TestBedrockRealtimeResponseTransformation: }, ) - # Check for function call event - assert len(result["response"]) == 1 - function_call = result["response"][0] - assert function_call["type"] == "response.function_call_arguments.done" - assert function_call["call_id"] == "tool_call_123" - assert function_call["name"] == "get_weather" - - # Verify arguments are properly formatted - args = json.loads(function_call["arguments"]) - assert args["location"] == "San Francisco" + assert [e["type"] for e in result["response"]] == [ + "response.output_item.added", + "response.function_call_arguments.delta", + "response.function_call_arguments.done", + "response.output_item.done", + ] + assert result["current_response_id"] == "resp_123" + assert all(e.get("response_id") == "resp_123" for e in result["response"]) def test_transform_content_end_text(self): """Test contentEnd for text response""" config = BedrockRealtimeConfig() + config._assistant_block_open = True logging_obj = MagicMock() logging_obj.litellm_trace_id = "trace_123" @@ -1211,6 +1372,395 @@ class TestBedrockRealtimeUserEventsAndUsage: assert types.count("response.created") == 1 assert types.count("response.done") == 1 + def test_nova2_sonic_tool_call_sequence_forwards_each_call_with_arguments(self): + """The real Nova 2 Sonic event sequence for one question: a USER transcript + block, two role TOOL blocks each carrying a toolUse with a distinct toolUseId, + then the post-tool-result ASSISTANT answer (LIT-9465)""" + args_json: Final = '{"location":"New York"}' + events: Final = self._run( + BedrockRealtimeConfig(), + [ + { + "event": { + "contentStart": { + "role": "USER", + "type": "TEXT", + "additionalModelFields": json.dumps({"generationStage": "FINAL"}), + } + } + }, + {"event": {"textOutput": {"content": "ready"}}}, + {"event": {"contentEnd": {"stopReason": "PARTIAL_TURN", "type": "TEXT"}}}, + ] + + self._tool_block("tooluse_A", args_json) + + self._tool_block("tooluse_B", args_json) + + self._assistant_text_block( + "SPECULATIVE", "The current weather in New York is sunny and 72 degrees.", "END_TURN" + ), + ) + + function_calls: Final = [e for e in events if e["type"] == "response.function_call_arguments.done"] + assert [e["call_id"] for e in function_calls] == ["tooluse_A", "tooluse_B"] + assert all(e["name"] == "get_weather" for e in function_calls) + assert all(e["arguments"] == args_json for e in function_calls) + assert function_calls[0]["item_id"] != function_calls[1]["item_id"] + + response_dones: Final = [e for e in events if e["type"] == "response.done"] + assert len(response_dones) == 3 + tool_call_dones: Final = [ + d for d in response_dones if any(o["type"] == "function_call" for o in d["response"]["output"]) + ] + assert len(tool_call_dones) == 2 + assert tool_call_dones[0]["response"]["output"][0]["call_id"] == "tooluse_A" + assert tool_call_dones[1]["response"]["output"][0]["call_id"] == "tooluse_B" + for done, call in zip(tool_call_dones, function_calls): + assert len(done["response"]["output"]) == 1 + assert done["response"]["output"][0]["type"] == "function_call" + assert done["response"]["output"][0]["status"] == "completed" + assert done["response"]["id"] == call["response_id"] + final_output: Final = response_dones[2]["response"]["output"] + assert [o["type"] for o in final_output] == ["message"] + assert final_output[0]["status"] == "completed" + + response_ids: Final = {d["response"]["id"] for d in response_dones} + assert len(response_ids) == 3 + + @staticmethod + def _tool_block(tool_use_id: str, arguments: str = '{"city":"Paris"}') -> list[dict[str, object]]: + return [ + {"event": {"contentStart": {"role": "TOOL", "type": "TOOL", "contentId": f"c_{tool_use_id}"}}}, + { + "event": { + "toolUse": { + "content": arguments, + "contentId": f"c_{tool_use_id}", + "role": "TOOL", + "toolName": "get_weather", + "toolUseId": tool_use_id, + } + } + }, + {"event": {"contentEnd": {"stopReason": "TOOL_USE", "type": "TOOL", "contentId": f"c_{tool_use_id}"}}}, + ] + + @staticmethod + def _item_id_of(event) -> str | None: + if event["type"].startswith("response.output_item."): + return event["item"]["id"] + return event.get("item_id") + + def test_tool_call_after_assistant_message_uses_output_index_one(self): + events: Final = self._run( + BedrockRealtimeConfig(), + self._assistant_text_block("SPECULATIVE", "Let me check.", "PARTIAL_TURN") + self._tool_block("tooluse_A"), + ) + types: Final = [e["type"] for e in events] + assert types.count("response.created") == 1 + + added_items: Final = [e["item"]["id"] for e in events if e["type"] == "response.output_item.added"] + message_id, call_id = added_items + + message_events: Final = [e for e in events if self._item_id_of(e) == message_id] + call_events: Final = [e for e in events if self._item_id_of(e) == call_id] + assert message_events and all(e["output_index"] == 0 for e in message_events) + assert call_events and all(e["output_index"] == 1 for e in call_events) + + response_dones: Final = [e for e in events if e["type"] == "response.done"] + assert len(response_dones) == 1 + assert [o["id"] for o in response_dones[0]["response"]["output"]] == [message_id, call_id] + + def test_tool_call_inside_open_assistant_block_defers_response_done(self): + config: Final = BedrockRealtimeConfig() + events: Final = self._run( + config, + [ + { + "event": { + "contentStart": { + "role": "ASSISTANT", + "type": "TEXT", + "additionalModelFields": json.dumps({"generationStage": "SPECULATIVE"}), + } + } + }, + {"event": {"textOutput": {"content": "Let me check."}}}, + ] + + self._tool_block("tooluse_A") + + [ + {"event": {"textOutput": {"content": " It is 72."}}}, + {"event": {"contentEnd": {"stopReason": "PARTIAL_TURN"}}}, + ], + ) + types: Final = [e["type"] for e in events] + assert types.count("response.done") == 1 + assert types[-1] == "response.done" + assert [e["text"] for e in events if e["type"] == "response.text.done"] == ["Let me check. It is 72."] + + next_events: Final = self._run(config, self._assistant_text_block("SPECULATIVE", "Sunny.", "END_TURN")) + assert [e["text"] for e in next_events if e["type"] == "response.text.done"] == ["Sunny."] + assert all(e["output_index"] == 0 for e in next_events if "output_index" in e) + assert [e["response"]["id"] for e in next_events if e["type"] == "response.created"] != [ + e["response"]["id"] for e in events if e["type"] == "response.created" + ] + + def test_response_done_message_item_carries_the_text_part(self): + events: Final = self._run( + BedrockRealtimeConfig(), + self._assistant_text_block("SPECULATIVE", "It is 72.", "END_TURN"), + ) + (done,) = [e for e in events if e["type"] == "response.done"] + (message_item,) = [o for o in done["response"]["output"] if o["type"] == "message"] + assert message_item["content"] == [{"type": "text", "text": "It is 72."}] + + def test_response_done_message_item_merges_text_then_audio_parts_in_order(self): + events: Final = self._run( + BedrockRealtimeConfig(), + self._assistant_text_block("SPECULATIVE", "Hi.", "PARTIAL_TURN") + self._assistant_audio_block("END_TURN"), + ) + (done,) = [e for e in events if e["type"] == "response.done"] + assert [o["type"] for o in done["response"]["output"]] == ["message"] + assert done["response"]["output"][0]["content"] == [ + {"type": "text", "text": "Hi."}, + {"type": "audio", "transcript": ""}, + ] + + def test_session_update_rejection_before_prompt_start(self): + assert ( + BedrockRealtimeConfig().session_update_rejection( + {"type": "session.update", "session": {"tools": [CONFLICT_TOOL]}} + ) + is None + ) + + @pytest.mark.parametrize( + "session,expects_conflict", + [ + ({"modalities": ["text"]}, False), + ({"tools": [CONFLICT_TOOL]}, False), + ({"tools": [{"parameters": OBJECT_PARAMETERS, "name": "get_weather", "type": "function"}]}, False), + ( + {"tools": [{"type": "function", "function": {"name": "get_weather", "parameters": OBJECT_PARAMETERS}}]}, + False, + ), + ({"tools": None}, False), + ({"voice": None}, False), + ({"output_sample_rate_hertz": 24000.0}, False), + ({"output_sample_rate_hertz": 16000.0}, True), + ({"tools": [{"type": "function", "name": "get_time", "parameters": OBJECT_PARAMETERS}]}, True), + ({"tools": []}, True), + ({"instructions": "Be brief."}, True), + ({"instructions": None}, False), + ({"instructions": ""}, False), + ({"voice": "amy"}, True), + ({"output_audio_format": "g711_ulaw"}, True), + ({"output_sample_rate_hertz": 8000}, True), + ({"temperature": 0.1}, True), + ({"max_response_output_tokens": 512}, True), + ({"voice": "matthew", "temperature": 0.7}, False), + ({"input_audio_format": "g711_alaw"}, False), + ], + ) + def test_session_update_rejection_locked_settings(self, session, expects_conflict): + config: Final = BedrockRealtimeConfig() + config.transform_session_update_event({"type": "session.update", "session": {"tools": [CONFLICT_TOOL]}}) + assert config.prompt_started + + conflict: Final = config.session_update_rejection({"type": "session.update", "session": session}) + assert (conflict.code if conflict else None) == ("session_update_not_applied" if expects_conflict else None) + + def test_session_update_with_inf_max_tokens_omits_max_tokens_from_session_start(self): + config: Final = BedrockRealtimeConfig() + messages: Final = config.transform_session_update_event( + {"type": "session.update", "session": {"max_response_output_tokens": "inf", "modalities": ["text"]}} + ) + session_start: Final = json.loads(messages[0])["event"]["sessionStart"]["inferenceConfiguration"] + assert "maxTokens" not in session_start + assert session_start["topP"] == 0.9 + assert session_start["temperature"] == 0.7 + + @pytest.mark.parametrize( + "first_session,later_session,expects_conflict", + [ + ({"modalities": ["text"]}, {"max_response_output_tokens": "inf"}, True), + ({"max_response_output_tokens": "inf"}, {"max_response_output_tokens": "inf"}, False), + ({"max_response_output_tokens": "inf"}, {"max_response_output_tokens": 2048}, True), + ], + ) + def test_session_update_rejection_inf_max_tokens(self, first_session, later_session, expects_conflict): + config: Final = BedrockRealtimeConfig() + config.transform_session_update_event({"type": "session.update", "session": first_session}) + assert config.prompt_started + + conflict: Final = config.session_update_rejection({"type": "session.update", "session": later_session}) + assert (conflict is not None) == expects_conflict + + def test_session_update_after_start_leaves_locked_attributes_unchanged(self): + config: Final = BedrockRealtimeConfig() + config.transform_session_update_event({"type": "session.update", "session": {"modalities": ["text"]}}) + + config.transform_session_update_event( + { + "type": "session.update", + "session": {"voice": "amy", "output_audio_format": "g711_ulaw", "temperature": 0.1}, + } + ) + assert config.voice_id == "matthew" + assert config.output_sample_rate_hertz == 24000 + assert config.temperature == 0.7 + + # Input audio settings still apply mid-session + config.transform_session_update_event( + {"type": "session.update", "session": {"input_audio_format": "g711_alaw"}} + ) + assert config.input_sample_rate_hertz == 8000 + config.transform_session_update_event({"type": "session.update", "session": {"input_sample_rate_hertz": 16000}}) + assert config.input_sample_rate_hertz == 16000 + + def test_stray_non_assistant_content_end_emits_no_duplicate_done_events(self): + events: Final = self._run( + BedrockRealtimeConfig(), + self._assistant_text_block("SPECULATIVE", "Hi.", "PARTIAL_TURN") + + [ + {"event": {"contentStart": {"role": "USER", "type": "AUDIO"}}}, + {"event": {"contentEnd": {}}}, + ], + ) + assert len([e for e in events if e["type"] == "response.text.done"]) == 1 + message_item_dones: Final = [ + e for e in events if e["type"] == "response.output_item.done" and e["item"]["type"] == "message" + ] + assert len(message_item_dones) == 1 + + @staticmethod + def _inference_configuration(config: BedrockRealtimeConfig, session: dict[str, object]) -> dict[str, object]: + messages: Final = config.transform_session_update_event({"type": "session.update", "session": session}) + return json.loads(messages[0])["event"]["sessionStart"]["inferenceConfiguration"] + + @pytest.mark.parametrize("max_output_tokens", [4096.0, 4096.5]) + def test_float_max_tokens_reaches_session_start_as_an_integer(self, max_output_tokens): + inference: Final = self._inference_configuration( + BedrockRealtimeConfig(), {"max_response_output_tokens": max_output_tokens, "temperature": 1} + ) + assert inference["maxTokens"] == 4096 + assert isinstance(inference["maxTokens"], int) + assert inference["temperature"] == 1.0 + + @pytest.mark.parametrize( + "session,param", + [ + ({"temperature": "0.2"}, "session.temperature"), + ({"temperature": True}, "session.temperature"), + ({"max_response_output_tokens": "4096"}, "session.max_response_output_tokens"), + ({"temperature": float("nan")}, "session.temperature"), + ({"max_response_output_tokens": float("inf")}, "session.max_response_output_tokens"), + ({"max_response_output_tokens": True}, "session.max_response_output_tokens"), + ({"output_sample_rate_hertz": "16000"}, "session.output_sample_rate_hertz"), + ], + ) + def test_malformed_number_setting_is_rejected_without_drop_params(self, session, param): + rejection: Final = BedrockRealtimeConfig().session_update_rejection( + {"type": "session.update", "session": session} + ) + assert rejection is not None + assert (rejection.code, rejection.param) == ("invalid_type", param) + + @pytest.mark.parametrize( + "session", + [ + {"temperature": "0.2", "max_response_output_tokens": "4096"}, + {"temperature": True, "max_response_output_tokens": float("nan")}, + ], + ) + def test_malformed_number_setting_is_dropped_with_drop_params(self, session): + config: Final = BedrockRealtimeConfig(drop_params=True) + assert config.session_update_rejection({"type": "session.update", "session": session}) is None + inference: Final = self._inference_configuration(config, session) + assert (inference["temperature"], inference["maxTokens"]) == (0.7, 1024) + + def test_null_number_settings_count_as_unset(self): + session: Final = {"temperature": None, "max_response_output_tokens": None} + config: Final = BedrockRealtimeConfig() + assert config.session_update_rejection({"type": "session.update", "session": session}) is None + inference: Final = self._inference_configuration(config, session) + assert (inference["temperature"], inference["maxTokens"]) == (0.7, 1024) + + @staticmethod + def _audio_output_configuration(session: dict[str, object]) -> dict[str, object]: + messages: Final = BedrockRealtimeConfig().transform_session_update_event( + {"type": "session.update", "session": session} + ) + return json.loads(messages[1])["event"]["promptStart"]["audioOutputConfiguration"] + + @pytest.mark.parametrize("output_sample_rate_hertz", [16000, 16000.0]) + def test_output_sample_rate_reaches_prompt_start_as_an_integer(self, output_sample_rate_hertz): + audio_output: Final = self._audio_output_configuration({"output_sample_rate_hertz": output_sample_rate_hertz}) + assert audio_output["sampleRateHertz"] == 16000 + assert isinstance(audio_output["sampleRateHertz"], int) + + def test_null_voice_keeps_the_default_voice(self): + assert self._audio_output_configuration({"voice": None})["voiceId"] == "matthew" + + @pytest.mark.parametrize( + "client_event", + [ + {"type": "conversation.item.create", "item": {"type": "message", "role": "user", "content": []}}, + {"type": "input_audio_buffer.append", "audio": base64.b64encode(bytes(320)).decode()}, + {"type": "response.create"}, + ], + ) + def test_first_input_without_an_applied_session_update_starts_the_prompt_once(self, client_event): + config: Final = BedrockRealtimeConfig() + rejected: Final = {"type": "session.update", "session": {"temperature": "warm"}} + assert config.session_update_rejection(rejected) is not None + + first: Final = config.transform_realtime_request(json.dumps(client_event), "amazon.nova-2-sonic-v1:0") + second: Final = config.transform_realtime_request(json.dumps(client_event), "amazon.nova-2-sonic-v1:0") + + first_events: Final = [json.loads(message)["event"] for message in first] + assert [next(iter(event)) for event in first_events[:2]] == ["sessionStart", "promptStart"] + assert first_events[2]["contentStart"]["role"] == "SYSTEM" + assert first_events[0]["sessionStart"]["inferenceConfiguration"]["temperature"] == 0.7 + assert not any("promptStart" in json.loads(message)["event"] for message in second) + + @pytest.mark.parametrize( + "tool_use,name,arguments", + [ + ({"toolUseId": "t1", "content": '{"city":"Paris"}'}, "", '{"city":"Paris"}'), + ( + {"toolUseId": "t1", "toolName": "get_weather", "content": {"city": "Paris"}}, + "get_weather", + '{"city": "Paris"}', + ), + ], + ) + def test_tool_use_with_missing_name_or_object_content_is_forwarded(self, tool_use, name, arguments): + events: Final = self._run(BedrockRealtimeConfig(), [{"event": {"toolUse": tool_use}}]) + (function_call,) = [e for e in events if e["type"] == "response.function_call_arguments.done"] + assert (function_call["name"], function_call["arguments"]) == (name, arguments) + + def test_tool_use_without_tool_use_id_is_forwarded_with_empty_call_id(self): + events: Final = self._run( + BedrockRealtimeConfig(), + [{"event": {"toolUse": {"content": '{"city":"Paris"}', "toolName": "get_weather"}}}], + ) + (function_call,) = [e for e in events if e["type"] == "response.function_call_arguments.done"] + assert (function_call["call_id"], function_call["arguments"]) == ("", '{"city":"Paris"}') + + def test_tool_call_without_assistant_message_stays_output_index_zero(self): + events: Final = self._run(BedrockRealtimeConfig(), self._tool_block("tooluse_A")) + call_events: Final = [ + e + for e in events + if e["type"] + in ( + "response.output_item.added", + "response.output_item.done", + "response.function_call_arguments.delta", + "response.function_call_arguments.done", + ) + ] + assert call_events and all(e["output_index"] == 0 for e in call_events) + if __name__ == "__main__": pytest.main([__file__, "-v"])