fix(bedrock): per-response realtime usage deltas, spend-log event filter, single transcript completed

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)

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(

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
 

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"])