diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 7ab5a14dfc2..e83cb743aa0 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -24,7 +24,10 @@ from litellm.constants import ( 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.litellm_core_utils.realtime_streaming import ( + REALTIME_SESSION_SUCCESS_LOGGED_KEY, + DefaultLoggedRealTimeEventTypes, +) from litellm.types.llms.openai import OpenAIRealtimeEvents from litellm.types.realtime import RealtimeResponseTransformInput @@ -346,6 +349,7 @@ class BedrockRealtime(BaseAWSLLM): prefer_async_handlers=True, ) ) + logging_obj.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] = True if outcome.provider_failure is None: return 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 2eadaee9e1a..7e418817e74 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 @@ -8,6 +8,7 @@ from unittest.mock import MagicMock import pytest import litellm +from litellm.litellm_core_utils.realtime_streaming import REALTIME_SESSION_SUCCESS_LOGGED_KEY from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.realtime.handler import BedrockRealtime from litellm.llms.bedrock.realtime.transformation import BedrockRealtimeConfig @@ -77,6 +78,7 @@ class UnavailableBedrockStream: class FakeLogging: def __init__(self, trace_id="trace-nova-sonic"): self.litellm_trace_id = trace_id + self.model_call_details = {} class DisconnectingClientWS: @@ -673,6 +675,31 @@ class TestBedrockRealtimeProviderFailurePropagation: await spend_dispatch["coro"] assert [event["type"] for event in spend_dispatch["events"]] == ["response.done"] + @pytest.mark.asyncio + async def test_success_dispatch_stamps_the_ownership_marker_only_when_spend_was_logged( + self, stub_aws_sdk_client, spend_dispatch + ): + stub_aws_sdk_client["streams"] = [ScriptedBedrockStream(self.TEXT_TURN)] + await BedrockRealtime().async_realtime( + model="amazon.nova-sonic-v1:0", + websocket=ConnectedClientWS([self.SESSION_UPDATE]), + logging_obj=spend_dispatch["logging_obj"], + **self.AWS_PARAMS, + ) + await spend_dispatch["coro"] + assert [event["type"] for event in spend_dispatch["events"]] == ["response.done"] + assert spend_dispatch["logging_obj"].model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY) is True + + idle_logging = FakeLogging() + stub_aws_sdk_client["streams"] = [ScriptedBedrockStream([])] + await BedrockRealtime().async_realtime( + model="amazon.nova-sonic-v1:0", + websocket=ConnectedClientWS([self.SESSION_UPDATE]), + logging_obj=idle_logging, + **self.AWS_PARAMS, + ) + assert REALTIME_SESSION_SUCCESS_LOGGED_KEY not in idle_logging.model_call_details + @pytest.mark.asyncio async def test_stream_failure_after_client_disconnect_is_not_a_provider_failure(self, stub_aws_sdk_client): stream = ScriptedBedrockStream([], receiver_type=BreakingBedrockReceiver)