mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
fix(bedrock): stamp the realtime success ownership marker when Nova Sonic spend is logged
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
2441e8a2a9
commit
d2342f06ce
2 changed files with 32 additions and 1 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue