diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index 98716ae89c7..8286058bbc6 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -754,6 +754,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 = None, ) -> tuple[ list[OpenAIRealtimeEvents], str | None, @@ -764,14 +765,9 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): """ Transform Bedrock contentStart event to OpenAI response events. - Args: - event: Bedrock contentStart event - current_response_id: Current response ID - current_output_item_id: Current output item ID - current_conversation_id: Current conversation ID - - Returns: - Tuple of (events, response_id, output_item_id, conversation_id, delta_type) + Bedrock streams one content block at a time (TEXT, AUDIO, TOOL, …). Only + ASSISTANT blocks open an OpenAI response/item lifecycle. Non-assistant + blocks must not clobber in-flight assistant part state. """ content_start: Final = event["contentStart"] role: Final = content_start.get("role") @@ -782,40 +778,37 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): current_response_id, current_output_item_id, current_conversation_id, - None, + current_delta_type, ) verbose_logger.debug("Handling ASSISTANT contentStart") - # Initialize IDs if needed + is_new_response: Final = not current_response_id if not current_response_id: current_response_id = f"resp_{uuid.uuid4()}" - if not current_output_item_id: - current_output_item_id = f"item_{uuid.uuid4()}" + current_output_item_id = f"item_{uuid.uuid4()}" 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]] = [] - # Send response.created - response_created: Final = OpenAIRealtimeStreamResponseBaseObject( - type="response.created", - event_id=f"event_{uuid.uuid4()}", - response={ - "object": "realtime.response", - "id": current_response_id, - "status": "in_progress", - "output": [], - "conversation_id": current_conversation_id, - }, - ) - returned_messages.append(response_created) + if is_new_response: + response_created: Final = OpenAIRealtimeStreamResponseBaseObject( + type="response.created", + event_id=f"event_{uuid.uuid4()}", + response={ + "object": "realtime.response", + "id": current_response_id, + "status": "in_progress", + "output": [], + "conversation_id": current_conversation_id, + }, + ) + returned_messages.append(response_created) - # Send response.output_item.added output_item_added: Final = OpenAIRealtimeStreamResponseOutputItemAdded( type="response.output_item.added", response_id=current_response_id, @@ -831,7 +824,6 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): ) returned_messages.append(output_item_added) - # Send response.content_part.added content_part_added: Final = OpenAIRealtimeResponseContentPartAdded( type="response.content_part.added", content_index=0, @@ -839,7 +831,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): 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": ""} + {"type": "text", "text": ""} if next_delta_type == "text" else {"type": "audio", "transcript": ""} ), response_id=current_response_id, ) @@ -850,7 +842,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): current_response_id, current_output_item_id, current_conversation_id, - current_delta_type, + next_delta_type, ) def transform_text_output_event( @@ -955,10 +947,10 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): verbose_logger.debug("Handling contentEnd: %s", content_end) if not current_output_item_id or not current_response_id: - return [], current_delta_chunks + return [], None if content_end.get("type") == "TOOL" or current_delta_type not in ("text", "audio"): - return [], current_delta_chunks + return [], None returned_messages: Final[list[OpenAIRealtimeEvents]] = [] @@ -1273,8 +1265,11 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): current_response_id, current_output_item_id, current_conversation_id, + current_delta_type, ) returned_messages.extend(events) + if events: + current_delta_chunks = None elif "textOutput" in event: events, current_delta_chunks = self.transform_text_output_event( @@ -1299,10 +1294,11 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): current_delta_chunks, ) returned_messages.extend(events) + current_delta_chunks = None + current_delta_type = None + current_output_item_id = None if content_end.get("type") == "TOOL": - current_output_item_id = None # rebind-ok: tool block ends so next ASSISTANT mints fresh ids - current_response_id = None # rebind-ok: tool block ends so next ASSISTANT mints fresh ids - current_delta_type = None # rebind-ok: tool block ends so next ASSISTANT mints fresh ids + current_response_id = None if BedrockContentEnd.model_validate(content_end).stopReason == "END_TURN": ( done_events, @@ -1311,6 +1307,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): current_delta_type, ) = self._response_done_events(current_response_id, current_conversation_id) returned_messages.extend(done_events) + current_delta_chunks = None elif "toolUse" in event: ( @@ -1321,8 +1318,10 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): tool_response_id, ) = self.transform_tool_use_event(event, current_output_item_id, current_response_id) returned_messages.extend(events) - current_output_item_id = tool_output_item_id # rebind-ok: persist minted tool item id into session - current_response_id = tool_response_id # rebind-ok: persist minted tool response id into session + current_output_item_id = tool_output_item_id + current_response_id = tool_response_id + current_delta_chunks = None + current_delta_type = None verbose_logger.debug("Tool use event: %s (ID: %s)", tool_name, tool_call_id) elif "promptEnd" in event or "completionEnd" in event: @@ -1333,6 +1332,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): current_delta_type, ) = self.transform_prompt_end_event(event, current_response_id, current_conversation_id) returned_messages.extend(events) + current_delta_chunks = None return { "response": returned_messages, diff --git a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_transformation.py b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_transformation.py index 79b662fa841..2f02ef94e20 100644 --- a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_transformation.py +++ b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_transformation.py @@ -1090,6 +1090,162 @@ class TestBedrockRealtimeSessionEvents: assert event["session"]["modalities"] == ["text", "audio"] +class TestBedrockRealtimeContentBlockLifecycle: + """ + Bedrock streams discrete content blocks. Session state must follow block + boundaries so text/audio/tool blocks cannot leak into each other. + """ + + def _state(self, **overrides): + base = { + "session_configuration_request": json.dumps({"configured": True}), + "current_output_item_id": None, + "current_response_id": None, + "current_conversation_id": "conv_1", + "current_delta_chunks": None, + "current_item_chunks": [], + "current_delta_type": None, + } + base.update(overrides) + return base + + def _apply(self, config, logging_obj, state, message): + result = config.transform_realtime_response( + json.dumps(message), + "amazon.nova-2-sonic-v1:0", + logging_obj, + realtime_response_transform_input=state, + ) + state.update( + { + "current_output_item_id": result["current_output_item_id"], + "current_response_id": result["current_response_id"], + "current_conversation_id": result["current_conversation_id"], + "current_delta_chunks": result["current_delta_chunks"], + "current_item_chunks": result["current_item_chunks"], + "current_delta_type": result["current_delta_type"], + } + ) + return result + + def test_tool_block_does_not_leak_prior_text_into_next_assistant_turn(self): + config = BedrockRealtimeConfig() + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "trace_123" + state = self._state() + + self._apply(config, logging_obj, state, {"event": {"contentStart": {"role": "ASSISTANT", "type": "TEXT"}}}) + first_response_id = state["current_response_id"] + self._apply( + config, + logging_obj, + state, + {"event": {"textOutput": {"content": "I will check the weather."}}}, + ) + assert state["current_delta_chunks"] is not None + assert len(state["current_delta_chunks"]) == 1 + + self._apply( + config, + logging_obj, + state, + {"event": {"contentEnd": {"stopReason": "PARTIAL_TURN", "type": "TEXT"}}}, + ) + assert state["current_delta_chunks"] is None + assert state["current_delta_type"] is None + assert state["current_output_item_id"] is None + assert state["current_response_id"] == first_response_id + + self._apply(config, logging_obj, state, {"event": {"contentStart": {"role": "TOOL", "type": "TOOL"}}}) + assert state["current_delta_chunks"] is None + assert state["current_response_id"] == first_response_id + + tool_result = self._apply( + config, + logging_obj, + state, + { + "event": { + "toolUse": { + "toolUseId": "tool_1", + "toolName": "get_weather", + "content": json.dumps({"location": "Seattle"}), + } + } + }, + ) + assert tool_result["response"][0]["type"] == "response.function_call_arguments.done" + assert state["current_delta_chunks"] is None + assert state["current_delta_type"] is None + + self._apply( + config, + logging_obj, + state, + {"event": {"contentEnd": {"stopReason": "TOOL_USE", "type": "TOOL"}}}, + ) + assert state["current_response_id"] is None + assert state["current_output_item_id"] is None + assert state["current_delta_chunks"] is None + assert state["current_delta_type"] is None + + post_tool = self._apply( + config, + logging_obj, + state, + {"event": {"contentStart": {"role": "ASSISTANT", "type": "TEXT"}}}, + ) + assert state["current_response_id"] != first_response_id + assert state["current_delta_chunks"] is None + assert [msg["type"] for msg in post_tool["response"]].count("response.created") == 1 + + self._apply( + config, + logging_obj, + state, + {"event": {"textOutput": {"content": "It is sunny in Seattle."}}}, + ) + done = self._apply( + config, + logging_obj, + state, + {"event": {"contentEnd": {"stopReason": "END_TURN", "type": "TEXT"}}}, + ) + text_done = [msg for msg in done["response"] if msg["type"] == "response.text.done"][0] + assert text_done["text"] == "It is sunny in Seattle." + assert "I will check the weather." not in text_done["text"] + assert any(msg["type"] == "response.done" for msg in done["response"]) + assert state["current_response_id"] is None + assert state["current_delta_chunks"] is None + + def test_second_assistant_content_block_reuses_response_not_item(self): + config = BedrockRealtimeConfig() + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "trace_123" + state = self._state() + + first = self._apply( + config, logging_obj, state, {"event": {"contentStart": {"role": "ASSISTANT", "type": "TEXT"}}} + ) + response_id = state["current_response_id"] + first_item = state["current_output_item_id"] + assert sum(1 for msg in first["response"] if msg["type"] == "response.created") == 1 + + self._apply( + config, + logging_obj, + state, + {"event": {"contentEnd": {"stopReason": "PARTIAL_TURN", "type": "TEXT"}}}, + ) + second = self._apply( + config, logging_obj, state, {"event": {"contentStart": {"role": "ASSISTANT", "type": "AUDIO"}}} + ) + assert state["current_response_id"] == response_id + assert state["current_output_item_id"] != first_item + assert sum(1 for msg in second["response"] if msg["type"] == "response.created") == 0 + assert sum(1 for msg in second["response"] if msg["type"] == "response.output_item.added") == 1 + + class TestBedrockRealtimeUsageAccounting: def _usage_event( self,