diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index ecd143dd087..42fe8941443 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -7,14 +7,17 @@ This uses aws_sdk_bedrock_runtime for bidirectional streaming with Nova Sonic. import asyncio import contextlib import json +from collections.abc import AsyncIterator, Mapping from typing import Final, Protocol from pydantic import JsonValue, TypeAdapter +import litellm from litellm._logging import _redact_string, verbose_proxy_logger from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.litellm_core_utils.realtime_streaming import DefaultLoggedRealTimeEventTypes from litellm.types.llms.openai import OpenAIRealtimeEvents from litellm.types.realtime import RealtimeResponseTransformInput @@ -34,6 +37,17 @@ def _json_str(value: JsonValue) -> str | None: return value if isinstance(value, str) else None +def _should_log_event(openai_message: Mapping[str, object]) -> bool: + logged_types: Final = ( + litellm.logged_real_time_event_types + if litellm.logged_real_time_event_types is not None + else DefaultLoggedRealTimeEventTypes + ) + if logged_types == "*": + return True + return openai_message.get("type") in logged_types + + class RealtimeClientWebSocket(Protocol): """The client-facing websocket surface the realtime bridge talks to.""" @@ -207,18 +221,22 @@ class BedrockRealtime(BaseAWSLLM): ) ) - logged_events: Final[list[OpenAIRealtimeEvents]] = [] # mutable-ok: filled across the stream loop - bedrock_to_client_task: Final = asyncio.create_task( - self._forward_bedrock_to_client( - bedrock_stream, - websocket, - transformation_config, - model, - logging_obj, - session_state, - logged_events, + async def forward_bedrock_and_collect_logged_events() -> tuple[OpenAIRealtimeEvents, ...]: + return tuple( + [ + event + async for event in self._forward_bedrock_to_client( + bedrock_stream, + websocket, + transformation_config, + model, + logging_obj, + session_state, + ) + ] ) - ) + + bedrock_to_client_task: Final = asyncio.create_task(forward_bedrock_and_collect_logged_events()) # Wait for both tasks to complete await asyncio.gather( @@ -227,9 +245,25 @@ class BedrockRealtime(BaseAWSLLM): return_exceptions=True, ) + forwarded_logged_events: Final = ( + bedrock_to_client_task.result() + if not bedrock_to_client_task.cancelled() and bedrock_to_client_task.exception() is None + else () + ) + logged_events: Final = ( + *forwarded_logged_events, + *( + leftover_event + for leftover_event in transformation_config.leftover_usage_done_events() + if _should_log_event(leftover_event) + ), + ) if logged_events: GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( - logging_obj.dispatch_success_handlers(logged_events, prefer_async_handlers=True) + logging_obj.dispatch_success_handlers( + list(logged_events), # mutable-ok: realtime spend logging requires a list result + prefer_async_handlers=True, + ) ) except Exception as e: @@ -313,9 +347,8 @@ class BedrockRealtime(BaseAWSLLM): model: str, logging_obj: LiteLLMLogging, session_state: RealtimeResponseTransformInput, - logged_events: "list[OpenAIRealtimeEvents] | None" = None, # mutable-ok: caller-owned spend log accumulator - ): - """Forward messages from Bedrock stream to client WebSocket.""" + ) -> AsyncIterator[OpenAIRealtimeEvents]: + """Forward messages from Bedrock to the client, yielding the ones to record for spend logging.""" try: while True: # Receive from Bedrock @@ -363,13 +396,14 @@ class BedrockRealtime(BaseAWSLLM): ) # Send transformed messages to client - openai_messages = transformed_response.get("response", []) + response_value = transformed_response["response"] + openai_messages = response_value if isinstance(response_value, list) else (response_value,) for openai_message in openai_messages: - if logged_events is not None and isinstance(openai_message, dict): - logged_events.append(openai_message) message_json = json.dumps(openai_message) await client_ws.send_text(message_json) verbose_proxy_logger.debug("Bedrock Realtime: Sent to client: %s", message_json[:200]) + if _should_log_event(openai_message): + yield openai_message except Exception as e: verbose_proxy_logger.debug("Bedrock to client forwarding ended: %s", e, exc_info=True) diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index ec5eff0a6fc..28c2e446d10 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -41,7 +41,6 @@ from litellm.types.realtime import ( RealtimeResponseTransformInput, RealtimeResponseTypedDict, ) -from litellm.utils import get_empty_usage class BedrockContentEnd(BaseModel): @@ -118,7 +117,9 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): self._user_transcript_active = False self._user_transcript_generation_stage: str | None = None self._user_item_id: str | None = None - self._latest_usage: OpenAIRealtimeResponseUsage | None = None + self._user_transcript_buffer = "" + self._cumulative_usage = BedrockUsageEvent() + self._reported_usage = BedrockUsageEvent() def validate_environment(self, headers: dict, model: str, api_key: str | None = None) -> dict: """Validate environment - no special validation needed for Bedrock.""" @@ -836,47 +837,79 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): return (speech_event,) def transform_usage_event(self, usage_event: BedrockUsageEvent) -> None: - """Record Bedrock usageEvent token totals for the next response.done.""" + """Record Bedrock's session-cumulative usage totals for the next response.done.""" verbose_logger.debug("Handling usageEvent") + self._cumulative_usage = usage_event + + def _take_usage_delta(self) -> OpenAIRealtimeResponseUsage: + """Usage for the response now completing: cumulative totals minus what prior response.done events reported.""" + prior: Final = self._reported_usage + latest: Final = self._cumulative_usage + self._reported_usage = latest input_details: Final[OpenAIRealtimeUsageTokenDetails] = { - "audio_tokens": usage_event.details.total.input.speechTokens, - "text_tokens": usage_event.details.total.input.textTokens, + "audio_tokens": latest.details.total.input.speechTokens - prior.details.total.input.speechTokens, + "text_tokens": latest.details.total.input.textTokens - prior.details.total.input.textTokens, "cached_tokens": 0, } output_details: Final[OpenAIRealtimeUsageTokenDetails] = { - "audio_tokens": usage_event.details.total.output.speechTokens, - "text_tokens": usage_event.details.total.output.textTokens, + "audio_tokens": latest.details.total.output.speechTokens - prior.details.total.output.speechTokens, + "text_tokens": latest.details.total.output.textTokens - prior.details.total.output.textTokens, } - latest_usage: Final[OpenAIRealtimeResponseUsage] = { - "input_tokens": usage_event.totalInputTokens, - "output_tokens": usage_event.totalOutputTokens, - "total_tokens": usage_event.totalTokens, + usage_delta: Final[OpenAIRealtimeResponseUsage] = { + "input_tokens": latest.totalInputTokens - prior.totalInputTokens, + "output_tokens": latest.totalOutputTokens - prior.totalOutputTokens, + "total_tokens": latest.totalTokens - prior.totalTokens, "input_token_details": input_details, "output_token_details": output_details, } - self._latest_usage = latest_usage + return usage_delta + + def leftover_usage_done_events(self) -> tuple[OpenAIRealtimeEvents, ...]: + """Logged-only response.done for usage Bedrock reports after the final turn's contentEnd.""" + if self._cumulative_usage == self._reported_usage: + return () + usage: Final = self._take_usage_delta() + leftover_done: Final = OpenAIRealtimeDoneEvent( + type="response.done", + event_id=f"event_{uuid.uuid4()}", + response=OpenAIRealtimeResponseDoneObject( + object="realtime.response", + id=f"resp_{uuid.uuid4()}", + status="completed", + conversation_id=f"conv_{uuid.uuid4()}", + usage=dict(usage), # mutable-ok: OpenAIRealtimeResponseDoneObject types usage as plain dict + ), + ) + return (leftover_done,) def transform_user_transcript_event(self, transcript: str) -> tuple[OpenAIRealtimeEvents, ...]: - """Transform a USER-role Bedrock textOutput (ASR transcript) to OpenAI transcription events.""" + """Transform a USER-role Bedrock textOutput (ASR transcript) to an OpenAI transcription delta.""" verbose_logger.debug("Handling USER textOutput (ASR transcript)") - item_id: Final = self._current_user_item_id() delta_event: Final[OpenAIRealtimeInputAudioTranscriptionDelta] = { "type": "conversation.item.input_audio_transcription.delta", "event_id": f"event_{uuid.uuid4()}", - "item_id": item_id, + "item_id": self._current_user_item_id(), "content_index": 0, "delta": transcript, } - if self._user_transcript_generation_stage == "SPECULATIVE": - return (delta_event,) + if self._user_transcript_generation_stage != "SPECULATIVE": + self._user_transcript_buffer += transcript + return (delta_event,) + + def user_transcript_completed_events(self) -> tuple[OpenAIRealtimeEvents, ...]: + """One completed event with the full transcript once the FINAL user content block ends.""" + transcript: Final = self._user_transcript_buffer + if not transcript: + return () + self._user_transcript_buffer = "" completed_event: Final[OpenAIRealtimeInputAudioTranscriptionCompleted] = { "type": "conversation.item.input_audio_transcription.completed", "event_id": f"event_{uuid.uuid4()}", - "item_id": item_id, + "item_id": self._current_user_item_id(), "content_index": 0, "transcript": transcript, } - return (delta_event, completed_event) + return (completed_event,) def transform_text_output_event( self, @@ -1096,14 +1129,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): if not current_response_id or not current_conversation_id: return [], None, None, None - empty_usage: Final = get_empty_usage() - zero_usage: Final[OpenAIRealtimeResponseUsage] = { - "input_tokens": empty_usage.prompt_tokens, - "output_tokens": empty_usage.completion_tokens, - "total_tokens": empty_usage.total_tokens, - } - usage: Final = self._latest_usage or zero_usage - self._latest_usage = None + usage: Final = self._take_usage_delta() response_done: Final = OpenAIRealtimeDoneEvent( type="response.done", event_id=f"event_{uuid.uuid4()}", @@ -1324,6 +1350,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): elif "contentEnd" in event and self._user_transcript_active: self._user_transcript_active = False self._user_transcript_generation_stage = None + returned_messages.extend(self.user_transcript_completed_events()) elif "contentEnd" in event: events, current_delta_chunks = self.transform_content_end_event( diff --git a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py index aa6573884e2..0ea5b7ad4a1 100644 --- a/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/tests/test_litellm/llms/bedrock/realtime/test_bedrock_realtime_handler.py @@ -6,7 +6,7 @@ from unittest.mock import MagicMock import pytest - +import litellm from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.realtime.handler import BedrockRealtime from litellm.llms.bedrock.realtime.transformation import BedrockRealtimeConfig @@ -124,14 +124,6 @@ class ScriptedBedrockStream: return (None, self._receiver) -class ImmediatelyEndingBedrockStream: - def __init__(self): - self.input_stream = FakeInputStream() - - async def await_output(self): - return (None, EndedBedrockReceiver()) - - class FakeStaticCredentialsResolver: pass @@ -171,7 +163,7 @@ def stub_aws_sdk_client(monkeypatch): async def invoke_model_with_bidirectional_stream(self, operation_input): captured["operation_input"] = operation_input - return ImmediatelyEndingBedrockStream() + return ScriptedBedrockStream(captured.get("scripted_payloads", [])) package = types.ModuleType("aws_sdk_bedrock_runtime") client_module = types.ModuleType("aws_sdk_bedrock_runtime.client") @@ -292,7 +284,40 @@ class TestBedrockRealtimeHandler: assert stream.input_stream.closed @pytest.mark.asyncio - async def test_forwarded_events_are_collected_for_spend_logging(self): + async def test_forwarded_events_are_filtered_to_logged_types_for_spend_logging(self): + handler = BedrockRealtime() + stream = ScriptedBedrockStream( + [ + json.dumps({"event": {"userSpeechStart": {}}}), + json.dumps({"event": {"contentStart": {"role": "ASSISTANT", "type": "TEXT"}}}), + json.dumps({"event": {"textOutput": {"content": "Hi"}}}), + json.dumps({"event": {"contentEnd": {"stopReason": "END_TURN"}}}), + ] + ) + client_ws = RealtimeClientWS() + + logged_events = [ + event + async for event in handler._forward_bedrock_to_client( + stream, + client_ws, + BedrockRealtimeConfig(), + "amazon.nova-sonic-v1:0", + FakeLogging(), + {}, + ) + ] + + assert [event["type"] for event in logged_events] == ["response.done"] + sent_types = [json.loads(message)["type"] for message in client_ws.sent_to_client] + assert "input_audio_buffer.speech_started" in sent_types + assert "response.text.delta" in sent_types + assert "response.done" in sent_types + assert client_ws.closed + + @pytest.mark.asyncio + async def test_logged_event_types_star_collects_every_forwarded_event(self, monkeypatch): + monkeypatch.setattr(litellm, "logged_real_time_event_types", "*") handler = BedrockRealtime() stream = ScriptedBedrockStream( [ @@ -301,37 +326,89 @@ class TestBedrockRealtimeHandler: ] ) client_ws = RealtimeClientWS() - logged_events = [] - await handler._forward_bedrock_to_client( - stream, - client_ws, - BedrockRealtimeConfig(), - "amazon.nova-sonic-v1:0", - FakeLogging(), - {}, - logged_events, - ) + logged_events = [ + event + async for event in handler._forward_bedrock_to_client( + stream, + client_ws, + BedrockRealtimeConfig(), + "amazon.nova-sonic-v1:0", + FakeLogging(), + {}, + ) + ] assert [event["type"] for event in logged_events] == [ "input_audio_buffer.speech_started", "input_audio_buffer.speech_stopped", ] - assert client_ws.closed + + @pytest.mark.asyncio + async def test_trailing_usage_after_last_done_is_dispatched_for_spend(self, stub_aws_sdk_client, monkeypatch): + import litellm.llms.bedrock.realtime.handler as handler_module + + dispatched = {} + + class RecordingLogging(FakeLogging): + async def dispatch_success_handlers(self, result=None, prefer_async_handlers=False, **kwargs): + dispatched["events"] = result + + class RecordingLoggingWorker: + def ensure_initialized_and_enqueue(self, coro): + dispatched["coro"] = coro + + monkeypatch.setattr(handler_module, "GLOBAL_LOGGING_WORKER", RecordingLoggingWorker()) + stub_aws_sdk_client["scripted_payloads"] = [ + json.dumps( + { + "event": { + "usageEvent": { + "totalInputTokens": 3, + "totalOutputTokens": 6, + "totalTokens": 9, + "details": { + "total": { + "input": {"speechTokens": 3, "textTokens": 0}, + "output": {"speechTokens": 0, "textTokens": 6}, + } + }, + } + } + } + ) + ] + + await BedrockRealtime().async_realtime( + model="amazon.nova-sonic-v1:0", + websocket=RealtimeClientWS(), + logging_obj=RecordingLogging(), + aws_region_name="us-east-1", + aws_access_key_id="k", + aws_secret_access_key="s", + ) + await dispatched["coro"] + + assert [event["type"] for event in dispatched["events"]] == ["response.done"] + usage = dispatched["events"][0]["response"]["usage"] + assert (usage["input_tokens"], usage["output_tokens"], usage["total_tokens"]) == (3, 6, 9) + assert usage["input_token_details"] == {"audio_tokens": 3, "text_tokens": 0, "cached_tokens": 0} + assert usage["output_token_details"] == {"audio_tokens": 0, "text_tokens": 6} @pytest.mark.asyncio async def test_bedrock_stream_end_closes_client_websocket(self): handler = BedrockRealtime() client_ws = ClosableClientWS() - await handler._forward_bedrock_to_client( + async for _ in handler._forward_bedrock_to_client( EndedBedrockStream(), client_ws, BedrockRealtimeConfig(), "amazon.nova-sonic-v1:0", MagicMock(), {}, - ) + ): + pass assert client_ws.closed 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 dbe84cacb29..a74f03449a1 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 @@ -1025,6 +1025,112 @@ class TestBedrockRealtimeUserEventsAndUsage: assert usage["output_tokens"] == 0 assert usage["total_tokens"] == 0 + @staticmethod + def _usage_event(total_input, total_output, in_speech, in_text, out_speech, out_text): + return { + "event": { + "usageEvent": { + "totalInputTokens": total_input, + "totalOutputTokens": total_output, + "totalTokens": total_input + total_output, + "details": { + "total": { + "input": {"speechTokens": in_speech, "textTokens": in_text}, + "output": {"speechTokens": out_speech, "textTokens": out_text}, + } + }, + } + } + } + + _ASSISTANT_TURN = ( + {"event": {"contentStart": {"role": "ASSISTANT", "type": "TEXT"}}}, + {"event": {"textOutput": {"content": "Hi"}}}, + {"event": {"contentEnd": {"stopReason": "END_TURN"}}}, + ) + + def test_multi_turn_usage_reports_per_response_deltas_not_cumulative_totals(self): + events = self._run( + BedrockRealtimeConfig(), + [ + self._usage_event(25, 40, in_speech=20, in_text=5, out_speech=30, out_text=10), + *self._ASSISTANT_TURN, + self._usage_event(40, 100, in_speech=30, in_text=10, out_speech=75, out_text=25), + *self._ASSISTANT_TURN, + ], + ) + usages = [e["response"]["usage"] for e in events if e["type"] == "response.done"] + assert len(usages) == 2 + assert (usages[0]["input_tokens"], usages[0]["output_tokens"], usages[0]["total_tokens"]) == (25, 40, 65) + assert (usages[1]["input_tokens"], usages[1]["output_tokens"], usages[1]["total_tokens"]) == (15, 60, 75) + assert usages[1]["input_token_details"] == {"audio_tokens": 10, "text_tokens": 5, "cached_tokens": 0} + assert usages[1]["output_token_details"] == {"audio_tokens": 45, "text_tokens": 15} + assert sum(u["total_tokens"] for u in usages) == 140 + + def test_usage_reported_after_last_response_done_flushes_as_logged_only_done(self): + config = BedrockRealtimeConfig() + self._run( + config, + [ + self._usage_event(25, 40, in_speech=20, in_text=5, out_speech=30, out_text=10), + *self._ASSISTANT_TURN, + ], + ) + assert config.leftover_usage_done_events() == () + + self._run(config, [self._usage_event(25, 46, in_speech=20, in_text=5, out_speech=30, out_text=16)]) + leftover = config.leftover_usage_done_events() + assert len(leftover) == 1 + assert leftover[0]["type"] == "response.done" + usage = leftover[0]["response"]["usage"] + assert (usage["input_tokens"], usage["output_tokens"], usage["total_tokens"]) == (0, 6, 6) + assert usage["output_token_details"] == {"audio_tokens": 0, "text_tokens": 6} + assert config.leftover_usage_done_events() == () + + def test_final_transcript_fragments_emit_one_completed_with_full_transcript(self): + events = self._run( + BedrockRealtimeConfig(), + [ + { + "event": { + "contentStart": { + "role": "USER", + "type": "TEXT", + "additionalModelFields": json.dumps({"generationStage": "FINAL"}), + } + } + }, + {"event": {"textOutput": {"content": "What is the "}}}, + {"event": {"textOutput": {"content": "capital of France?"}}}, + {"event": {"contentEnd": {"stopReason": "PARTIAL_TURN"}}}, + ], + ) + deltas = [e for e in events if e["type"] == "conversation.item.input_audio_transcription.delta"] + completed = [e for e in events if e["type"] == "conversation.item.input_audio_transcription.completed"] + assert [d["delta"] for d in deltas] == ["What is the ", "capital of France?"] + assert len(completed) == 1 + assert completed[0]["transcript"] == "What is the capital of France?" + assert {e["item_id"] for e in deltas + completed} == {completed[0]["item_id"]} + + def test_speculative_transcript_block_end_emits_no_completed(self): + events = self._run( + BedrockRealtimeConfig(), + [ + { + "event": { + "contentStart": { + "role": "USER", + "type": "TEXT", + "additionalModelFields": json.dumps({"generationStage": "SPECULATIVE"}), + } + } + }, + {"event": {"textOutput": {"content": "rea"}}}, + {"event": {"contentEnd": {"stopReason": "PARTIAL_TURN"}}}, + ], + ) + assert [e["type"] for e in events] == ["conversation.item.input_audio_transcription.delta"] + if __name__ == "__main__": pytest.main([__file__, "-v"])