From 9187b3d9cabfd40f2b806c33d97f22a9673e5047 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Mon, 3 Aug 2026 15:59:47 -0700 Subject: [PATCH 01/19] chore: empty commit for litellm_bedrock_realtime_staging From 1ca08f0bbde311add13bb5244fc50148e8a8a8c6 Mon Sep 17 00:00:00 2001 From: Otavio Brito <69211663+otaviofbrito@users.noreply.github.com> Date: Mon, 3 Aug 2026 20:00:48 -0300 Subject: [PATCH 02/19] fix(bedrock/realtime): nova sonic tool calling (#33127) * fix(bedrock/realtime): Nova Sonic tool calling * increase coverage * Apply suggestion from @greptile-apps[bot] Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * Update test_bedrock_realtime_transformation.py * test(e2e): bound the post-/model/new servable wait at 40s _await_model_servable used poll_timeout (120s), the spend/log read-back budget. A stuck model reload therefore stalled every suite that creates a deployment for two minutes before failing Give create_model a fixed harness middle ground: model_servable_timeout=40s, polled every 2s, with each /v1/models call capped at 5s and clamped to the remaining deadline so one slow GET cannot overrun the wait. Happy path still returns on the first listing. Not derived from proxy general_settings or env Transport.get accepts an optional per-call timeout for that clamp. Unit tests cover the deadline arithmetic and clamp without a live proxy (cherry picked from commit c082a0e6488f50978bf5255f6b5298ba7e8fd8da) * fix(e2e): wait one default DB reload interval of continuous listing create_model returned after the first /v1/models hit that listed the model, so chat could still land on a cold gateway worker (numWorkers>1 / peer pod) and 400 Invalid model name. Require continuous listing for the product default add_deployment interval (30s) after first sight so every worker has synced from the DB; first listing still bounded at 40s (cherry picked from commit 7d1ee2ff861b970f6de3f6759ff015947af9d2a1) * test(e2e): drop proxy_client model-servable unit tests Keep the create_model DB-sync wait in the harness; the pure-function unit file is not needed for this PR (cherry picked from commit 89204651d1a4537c6f21550c5ab85448ae0923f8) * fix(e2e): never skip the final deadline-clamped model-servable poll When less than one full poll interval remained in the first-listing budget, the pre-sleep check returned NotServable without another /v1/models call. Sleep only min(interval, time left) so a model that becomes listable in the last seconds of the timeout still gets a clamped final poll (cherry picked from commit 8439195922c913d118cb146409c75cc081d23e6e) * fix(e2e): reject first listing that returns after the 40s deadline A poll may start with remaining budget and still return after started+timeout if the transport overruns its clamp. Recheck the first-listing deadline after the response so a late listing does not open the continuous DB-sync phase (cherry picked from commit 7ff2bcbf1498ee82f7dbe0b4330c1ab48927ed01) * test(e2e): poll MCP tools across multi-worker lag (#35047) * fix(mcp): resolve call_tool by registry without requiring tool map Multi-worker reloads put MCP servers in the registry from the DB but do not re-run tools/list on every process. Gating call_tool on tool_name_to_mcp_server_name_mapping made cold workers 500 with Tool not found after another worker had already listed the tool. Treat a registry match on server id/name/alias as enough; upstream rejects unknown tools * test(e2e): poll MCP register, tools/list, and tools/call across multi-worker lag Stage multi-worker gateways only load MCP servers and tool maps on the process that handled the request. Poll until the server is listed, the tool appears on tools/list, and tools/call is not a cold-worker 500 so key-access and Datadog MCP e2e stop racing the LB * Revert "fix(mcp): resolve call_tool by registry without requiring tool map" This reverts commit 8b56e51e39b876d13d1112efa4130554ddf5f173. * test(e2e): tighten MCP multi-worker lag classifier Only retry tools/call on gateway shapes Tool not found and server_not_found, not any 500 that mentions tool/server not found, so upstream failures are not retried until the poll deadline * test(e2e): drop unit file for MCP lag classifier The live await_call_tool polls already cover multi-worker lag; a separate string-match unit module is not worth keeping (cherry picked from commit c274cf321c5c35c629220a89bb497d15b56f870f) --------- Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: yuneng-jiang Co-authored-by: mubashir1osmani --- .../llms/bedrock/realtime/transformation.py | 8 +- tests/e2e/proxy_client.py | 174 +++++++++++++++--- tests/e2e/transport.py | 13 +- .../test_bedrock_realtime_transformation.py | 87 +++++++++ 4 files changed, 254 insertions(+), 28 deletions(-) diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index 39f5d25cf89..4435fba9d2d 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -1029,11 +1029,13 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): if not current_output_item_id or not current_response_id: return [], "", "" - # Parse the tool input + # Parse the tool input. Nova 2 Sonic sends arguments in `content`; + # fall back to `input` for backward compatibility. tool_input = {} - if "input" in tool_use: + raw_input = tool_use["content"] if "content" in tool_use else tool_use.get("input") + if raw_input: try: - tool_input = json.loads(tool_use["input"]) if isinstance(tool_use["input"], str) else tool_use["input"] + tool_input = json.loads(raw_input) if isinstance(raw_input, str) else raw_input except json.JSONDecodeError: tool_input = {} diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 6c6b948e29c..b3fc8538322 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -75,12 +75,130 @@ from transport import HttpTransport, SplitTransport, Transport RowsPredicate = Callable[[list[SpendLogRow]], bool] +# After /model/new, the control-plane writer reloads itself immediately, but every +# other gateway worker (and peer pod) only picks the model up on its add_deployment +# job. That job runs every proxy_config_reload_interval_seconds (product default 30). +# A single /v1/models hit can land on a hot worker while the next /chat hits a cold +# one ("Invalid model name"). Wait for first listing within MODEL_SERVABLE_TIMEOUT, +# then require continuous listing for MODEL_SERVABLE_DB_SYNC_SECONDS (the default +# reload interval) so every worker has had a chance to sync from the DB. +MODEL_SERVABLE_TIMEOUT = 40.0 +MODEL_SERVABLE_DB_SYNC_SECONDS = 30.0 +MODEL_SERVABLE_INTERVAL = 2.0 +# Cap each /v1/models poll so one slow request cannot outlast the remaining budget. +MODEL_SERVABLE_REQUEST_TIMEOUT = 5.0 + + +@dataclass(frozen=True, slots=True) +class Servable: + """The data plane listed the model within the deadline.""" + + +@dataclass(frozen=True, slots=True) +class NotServable: + """The deadline passed without the data plane listing the model. + + `last_result` is the final /v1/models read, so the caller can tell "the proxy + answered but omitted the model" (propagation) from "the read itself failed" + (network/auth) when reporting.""" + + last_result: Result[ModelsListResponse] | None + + +ServableOutcome = Servable | NotServable + + +def await_servable( + list_models: Callable[[float], Result[ModelsListResponse]], + *, + model_name: str, + timeout: float, + interval: float, + request_timeout: float, + db_sync_seconds: float, + now: Callable[[], float], + sleep: Callable[[float], None], +) -> ServableOutcome: + """Poll until `model_name` is listed long enough for every worker to DB-sync. + + First listing must happen within `timeout`. After that, the model must stay + listed continuously for `db_sync_seconds` (any miss resets the continuous + window). `db_sync_seconds=0` returns on the first listing. Each poll's request + timeout is clamped to the remaining budget. Sleeps only min(interval, time left) + so a final deadline-clamped poll is never skipped just because a full interval + does not fit. Clock and sleep are injected.""" + started = now() + first_seen_at: float | None = None + last_result: Result[ModelsListResponse] | None = None + while True: + t = now() + phase_deadline = ( + started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds + ) + remaining = phase_deadline - t + if remaining <= 0: + if ( + last_result is not None + and first_seen_at is not None + and (db_sync_seconds <= 0 or t - first_seen_at >= db_sync_seconds) + ): + return Servable() + return NotServable(last_result=last_result) + + poll_timeout = min(request_timeout, remaining) + last_result = list_models(poll_timeout) + listed = isinstance(last_result, Success) and any( + entry.id == model_name for entry in last_result.data.data + ) + t = now() + if not listed: + first_seen_at = None + elif first_seen_at is None: + if t > started + timeout: + return NotServable(last_result=last_result) + first_seen_at = t + if db_sync_seconds <= 0: + return Servable() + elif t - first_seen_at >= db_sync_seconds: + return Servable() + + phase_deadline = ( + started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds + ) + wait = min(interval, phase_deadline - now()) + if wait > 0: + sleep(wait) + + +def servable_timeout_message( + *, + model_name: str, + timeout: float, + db_sync_seconds: float, + last_result: Result[ModelsListResponse] | None, +) -> str: + last_error = ( + f"; last /v1/models poll did not succeed: {last_result}" + if last_result is not None and not isinstance(last_result, Success) + else "" + ) + return ( + f"model {model_name!r} was created but never became servable on the data " + f"plane within {timeout}s of first listing (plus {db_sync_seconds}s continuous " + f"DB sync) after /model/new (control/data-plane propagation or " + f"STORE_MODEL_IN_DB reload issue){last_error}" + ) + @dataclass(frozen=True, slots=True) class ProxyClient: transport: Transport poll_timeout: float = 120.0 poll_interval: float = 5.0 + model_servable_timeout: float = MODEL_SERVABLE_TIMEOUT + model_servable_db_sync_seconds: float = MODEL_SERVABLE_DB_SYNC_SECONDS + model_servable_interval: float = MODEL_SERVABLE_INTERVAL + model_servable_request_timeout: float = MODEL_SERVABLE_REQUEST_TIMEOUT # ---- keys / customers (satisfies lifecycle.ResourceClient) ---------- @@ -167,7 +285,12 @@ class ProxyClient: this returns can race the reload and 400 with "Invalid model name passed". We therefore poll the data-plane /v1/models until the model appears before handing back, so callers can invoke it immediately. In the monolithic case - it is already present on the first poll, so this adds one request.""" + it is already present on the first poll, so this adds one request. + + First listing must arrive within `model_servable_timeout` (not the longer + spend `poll_timeout`). The model must then stay listed for + `model_servable_db_sync_seconds` (product default DB reload interval) so every + gateway worker has run add_deployment before callers use the model.""" model_id = unwrap( self.transport.post( "/model/new", @@ -184,33 +307,38 @@ class ProxyClient: return model_id def _await_model_servable(self, model_name: str) -> None: - """Block until the data plane lists `model_name`, or fail loudly if it does - not within poll_timeout (a real propagation/config problem, surfaced here - instead of as a downstream "Invalid model name passed").""" - deadline = time.monotonic() + self.poll_timeout - last_result: Result[ModelsListResponse] | None = None - while time.monotonic() < deadline: - last_result = self.transport.get( + """Block until the data plane lists `model_name` long enough for DB sync. + + Fails if first listing misses model_servable_timeout, or if continuous listing + for model_servable_db_sync_seconds never holds (multi-worker / peer reload).""" + outcome = await_servable( + lambda poll_timeout: self.transport.get( "/v1/models", headers=self.transport.master, params=NoBody(), response_type=ModelsListResponse, - ) - if isinstance(last_result, Success) and any( - entry.id == model_name for entry in last_result.data.data - ): + timeout=poll_timeout, + ), + model_name=model_name, + timeout=self.model_servable_timeout, + interval=self.model_servable_interval, + request_timeout=self.model_servable_request_timeout, + db_sync_seconds=self.model_servable_db_sync_seconds, + now=time.monotonic, + sleep=time.sleep, + ) + match outcome: + case Servable(): return - time.sleep(self.poll_interval) - last_error = ( - f"; last /v1/models poll did not succeed: {last_result}" - if last_result is not None and not isinstance(last_result, Success) - else "" - ) - raise AssertionError( - f"model {model_name!r} was created but never became servable on the data " - f"plane within {self.poll_timeout}s of /model/new (control/data-plane " - f"propagation or STORE_MODEL_IN_DB reload issue){last_error}" - ) + case NotServable(last_result=last_result): + raise AssertionError( + servable_timeout_message( + model_name=model_name, + timeout=self.model_servable_timeout, + db_sync_seconds=self.model_servable_db_sync_seconds, + last_result=last_result, + ) + ) def update_model(self, model_id: str, litellm_params: LiteLLMParamsBody) -> None: """Merge `litellm_params` over the deployment `model_id`'s stored params via diff --git a/tests/e2e/transport.py b/tests/e2e/transport.py index a6adf83ed1f..27b11befc8e 100644 --- a/tests/e2e/transport.py +++ b/tests/e2e/transport.py @@ -58,6 +58,7 @@ class Transport(Protocol): headers: BaseModel, params: BaseModel, response_type: type[R], + timeout: float | None = None, ) -> Result[R]: ... def delete[R: BaseModel]( @@ -136,13 +137,16 @@ class HttpTransport: headers: BaseModel, params: BaseModel, response_type: type[R], + timeout: float | None = None, ) -> Result[R]: + """`timeout` overrides the transport-wide request_timeout for this call, for + pollers whose own deadline is shorter than it.""" return e2e_http.get( self._url(path), headers=headers, params=params, response_type=response_type, - timeout=self.request_timeout, + timeout=self.request_timeout if timeout is None else timeout, ) def delete[R: BaseModel]( @@ -336,9 +340,14 @@ class SplitTransport: headers: BaseModel, params: BaseModel, response_type: type[R], + timeout: float | None = None, ) -> Result[R]: return self._route(path).get( - path, headers=headers, params=params, response_type=response_type + path, + headers=headers, + params=params, + response_type=response_type, + timeout=timeout, ) def delete[R: BaseModel]( 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 aa002b6e302..6611f3b431e 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 @@ -571,6 +571,93 @@ class TestBedrockRealtimeResponseTransformation: args = json.loads(function_call["arguments"]) assert args["location"] == "San Francisco" + def test_transform_tool_use_response_with_content_field(self): + """Test toolUse response transformation with Nova 2 Sonic `content` field""" + config = BedrockRealtimeConfig() + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "trace_123" + + tool_use_message = { + "event": { + "toolUse": { + "toolUseId": "tool_call_123", + "toolName": "get_weather", + "content": json.dumps({"location": "San Francisco"}), + } + } + } + + result = config.transform_realtime_response( + json.dumps(tool_use_message), + "amazon.nova-2-sonic-v1:0", + logging_obj, + realtime_response_transform_input={ + "session_configuration_request": json.dumps({"configured": True}), + "current_output_item_id": "item_123", + "current_response_id": "resp_123", + "current_conversation_id": "conv_123", + "current_delta_chunks": [], + "current_item_chunks": [], + "current_delta_type": "text", + }, + ) + + # Check for function call event + assert len(result["response"]) == 1 + function_call = result["response"][0] + assert function_call["type"] == "response.function_call_arguments.done" + assert function_call["call_id"] == "tool_call_123" + assert function_call["name"] == "get_weather" + + # Verify arguments are properly formatted + args = json.loads(function_call["arguments"]) + assert args["location"] == "San Francisco" + + def test_transform_tool_use_event_directly(self): + """Test transform_tool_use_event directly for guard clause and input parsing""" + config = BedrockRealtimeConfig() + + # Guard clause: missing IDs returns empty result + events, tool_call_id, tool_name = config.transform_tool_use_event( + {"toolUse": {}}, None, "resp_123" + ) + assert events == [] + assert tool_call_id == "" + assert tool_name == "" + + # JSON string content is parsed and converted to a function call event + events, tool_call_id, tool_name = config.transform_tool_use_event( + { + "toolUse": { + "toolUseId": "tool_call_123", + "toolName": "get_weather", + "content": json.dumps({"location": "San Francisco"}), + } + }, + "item_123", + "resp_123", + ) + assert len(events) == 1 + assert events[0]["type"] == "response.function_call_arguments.done" + assert events[0]["call_id"] == "tool_call_123" + assert events[0]["name"] == "get_weather" + assert json.loads(events[0]["arguments"]) == {"location": "San Francisco"} + + # Invalid JSON content falls back to empty arguments + events, _, _ = config.transform_tool_use_event( + { + "toolUse": { + "toolUseId": "tool_call_124", + "toolName": "get_weather", + "content": "not valid json", + } + }, + "item_123", + "resp_123", + ) + assert len(events) == 1 + assert json.loads(events[0]["arguments"]) == {} + def test_transform_content_end_text(self): """Test contentEnd for text response""" config = BedrockRealtimeConfig() From 95239b9c68f699c3ef8f9e23b0587ab0c59a1163 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Mon, 3 Aug 2026 16:02:38 -0700 Subject: [PATCH 03/19] chore: drop unrelated e2e harness changes from bedrock realtime PR Keep only the Nova 2 Sonic toolUse content-field fix and its unit tests; the model-servable / multi-worker wait changes were cherry-picked noise --- tests/e2e/proxy_client.py | 174 +++++--------------------------------- tests/e2e/transport.py | 13 +-- 2 files changed, 25 insertions(+), 162 deletions(-) diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index b3fc8538322..6c6b948e29c 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -75,130 +75,12 @@ from transport import HttpTransport, SplitTransport, Transport RowsPredicate = Callable[[list[SpendLogRow]], bool] -# After /model/new, the control-plane writer reloads itself immediately, but every -# other gateway worker (and peer pod) only picks the model up on its add_deployment -# job. That job runs every proxy_config_reload_interval_seconds (product default 30). -# A single /v1/models hit can land on a hot worker while the next /chat hits a cold -# one ("Invalid model name"). Wait for first listing within MODEL_SERVABLE_TIMEOUT, -# then require continuous listing for MODEL_SERVABLE_DB_SYNC_SECONDS (the default -# reload interval) so every worker has had a chance to sync from the DB. -MODEL_SERVABLE_TIMEOUT = 40.0 -MODEL_SERVABLE_DB_SYNC_SECONDS = 30.0 -MODEL_SERVABLE_INTERVAL = 2.0 -# Cap each /v1/models poll so one slow request cannot outlast the remaining budget. -MODEL_SERVABLE_REQUEST_TIMEOUT = 5.0 - - -@dataclass(frozen=True, slots=True) -class Servable: - """The data plane listed the model within the deadline.""" - - -@dataclass(frozen=True, slots=True) -class NotServable: - """The deadline passed without the data plane listing the model. - - `last_result` is the final /v1/models read, so the caller can tell "the proxy - answered but omitted the model" (propagation) from "the read itself failed" - (network/auth) when reporting.""" - - last_result: Result[ModelsListResponse] | None - - -ServableOutcome = Servable | NotServable - - -def await_servable( - list_models: Callable[[float], Result[ModelsListResponse]], - *, - model_name: str, - timeout: float, - interval: float, - request_timeout: float, - db_sync_seconds: float, - now: Callable[[], float], - sleep: Callable[[float], None], -) -> ServableOutcome: - """Poll until `model_name` is listed long enough for every worker to DB-sync. - - First listing must happen within `timeout`. After that, the model must stay - listed continuously for `db_sync_seconds` (any miss resets the continuous - window). `db_sync_seconds=0` returns on the first listing. Each poll's request - timeout is clamped to the remaining budget. Sleeps only min(interval, time left) - so a final deadline-clamped poll is never skipped just because a full interval - does not fit. Clock and sleep are injected.""" - started = now() - first_seen_at: float | None = None - last_result: Result[ModelsListResponse] | None = None - while True: - t = now() - phase_deadline = ( - started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds - ) - remaining = phase_deadline - t - if remaining <= 0: - if ( - last_result is not None - and first_seen_at is not None - and (db_sync_seconds <= 0 or t - first_seen_at >= db_sync_seconds) - ): - return Servable() - return NotServable(last_result=last_result) - - poll_timeout = min(request_timeout, remaining) - last_result = list_models(poll_timeout) - listed = isinstance(last_result, Success) and any( - entry.id == model_name for entry in last_result.data.data - ) - t = now() - if not listed: - first_seen_at = None - elif first_seen_at is None: - if t > started + timeout: - return NotServable(last_result=last_result) - first_seen_at = t - if db_sync_seconds <= 0: - return Servable() - elif t - first_seen_at >= db_sync_seconds: - return Servable() - - phase_deadline = ( - started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds - ) - wait = min(interval, phase_deadline - now()) - if wait > 0: - sleep(wait) - - -def servable_timeout_message( - *, - model_name: str, - timeout: float, - db_sync_seconds: float, - last_result: Result[ModelsListResponse] | None, -) -> str: - last_error = ( - f"; last /v1/models poll did not succeed: {last_result}" - if last_result is not None and not isinstance(last_result, Success) - else "" - ) - return ( - f"model {model_name!r} was created but never became servable on the data " - f"plane within {timeout}s of first listing (plus {db_sync_seconds}s continuous " - f"DB sync) after /model/new (control/data-plane propagation or " - f"STORE_MODEL_IN_DB reload issue){last_error}" - ) - @dataclass(frozen=True, slots=True) class ProxyClient: transport: Transport poll_timeout: float = 120.0 poll_interval: float = 5.0 - model_servable_timeout: float = MODEL_SERVABLE_TIMEOUT - model_servable_db_sync_seconds: float = MODEL_SERVABLE_DB_SYNC_SECONDS - model_servable_interval: float = MODEL_SERVABLE_INTERVAL - model_servable_request_timeout: float = MODEL_SERVABLE_REQUEST_TIMEOUT # ---- keys / customers (satisfies lifecycle.ResourceClient) ---------- @@ -285,12 +167,7 @@ class ProxyClient: this returns can race the reload and 400 with "Invalid model name passed". We therefore poll the data-plane /v1/models until the model appears before handing back, so callers can invoke it immediately. In the monolithic case - it is already present on the first poll, so this adds one request. - - First listing must arrive within `model_servable_timeout` (not the longer - spend `poll_timeout`). The model must then stay listed for - `model_servable_db_sync_seconds` (product default DB reload interval) so every - gateway worker has run add_deployment before callers use the model.""" + it is already present on the first poll, so this adds one request.""" model_id = unwrap( self.transport.post( "/model/new", @@ -307,38 +184,33 @@ class ProxyClient: return model_id def _await_model_servable(self, model_name: str) -> None: - """Block until the data plane lists `model_name` long enough for DB sync. - - Fails if first listing misses model_servable_timeout, or if continuous listing - for model_servable_db_sync_seconds never holds (multi-worker / peer reload).""" - outcome = await_servable( - lambda poll_timeout: self.transport.get( + """Block until the data plane lists `model_name`, or fail loudly if it does + not within poll_timeout (a real propagation/config problem, surfaced here + instead of as a downstream "Invalid model name passed").""" + deadline = time.monotonic() + self.poll_timeout + last_result: Result[ModelsListResponse] | None = None + while time.monotonic() < deadline: + last_result = self.transport.get( "/v1/models", headers=self.transport.master, params=NoBody(), response_type=ModelsListResponse, - timeout=poll_timeout, - ), - model_name=model_name, - timeout=self.model_servable_timeout, - interval=self.model_servable_interval, - request_timeout=self.model_servable_request_timeout, - db_sync_seconds=self.model_servable_db_sync_seconds, - now=time.monotonic, - sleep=time.sleep, - ) - match outcome: - case Servable(): + ) + if isinstance(last_result, Success) and any( + entry.id == model_name for entry in last_result.data.data + ): return - case NotServable(last_result=last_result): - raise AssertionError( - servable_timeout_message( - model_name=model_name, - timeout=self.model_servable_timeout, - db_sync_seconds=self.model_servable_db_sync_seconds, - last_result=last_result, - ) - ) + time.sleep(self.poll_interval) + last_error = ( + f"; last /v1/models poll did not succeed: {last_result}" + if last_result is not None and not isinstance(last_result, Success) + else "" + ) + raise AssertionError( + f"model {model_name!r} was created but never became servable on the data " + f"plane within {self.poll_timeout}s of /model/new (control/data-plane " + f"propagation or STORE_MODEL_IN_DB reload issue){last_error}" + ) def update_model(self, model_id: str, litellm_params: LiteLLMParamsBody) -> None: """Merge `litellm_params` over the deployment `model_id`'s stored params via diff --git a/tests/e2e/transport.py b/tests/e2e/transport.py index 27b11befc8e..a6adf83ed1f 100644 --- a/tests/e2e/transport.py +++ b/tests/e2e/transport.py @@ -58,7 +58,6 @@ class Transport(Protocol): headers: BaseModel, params: BaseModel, response_type: type[R], - timeout: float | None = None, ) -> Result[R]: ... def delete[R: BaseModel]( @@ -137,16 +136,13 @@ class HttpTransport: headers: BaseModel, params: BaseModel, response_type: type[R], - timeout: float | None = None, ) -> Result[R]: - """`timeout` overrides the transport-wide request_timeout for this call, for - pollers whose own deadline is shorter than it.""" return e2e_http.get( self._url(path), headers=headers, params=params, response_type=response_type, - timeout=self.request_timeout if timeout is None else timeout, + timeout=self.request_timeout, ) def delete[R: BaseModel]( @@ -340,14 +336,9 @@ class SplitTransport: headers: BaseModel, params: BaseModel, response_type: type[R], - timeout: float | None = None, ) -> Result[R]: return self._route(path).get( - path, - headers=headers, - params=params, - response_type=response_type, - timeout=timeout, + path, headers=headers, params=params, response_type=response_type ) def delete[R: BaseModel]( From 06cbbbb79aedc9baafc78c4f9aa296fc5cef6c8e Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Mon, 3 Aug 2026 16:15:32 -0700 Subject: [PATCH 04/19] fix(bedrock/realtime): mint response/item ids so toolUse is not dropped Nova Sonic starts tool turns with contentStart role=TOOL, so ASSISTANT-only id init never runs and transform_tool_use_event used to return empty events even after parsing content. Mint resp_/item_ ids when missing so function_call_arguments.done reaches the client with non-empty args --- .../llms/bedrock/realtime/transformation.py | 12 +++----- .../test_bedrock_realtime_transformation.py | 28 +++++++++++++++---- 2 files changed, 26 insertions(+), 14 deletions(-) diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index 4435fba9d2d..b0b2c4f0201 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -1026,11 +1026,9 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): verbose_logger.debug("Handling toolUse") tool_use = event["toolUse"] - if not current_output_item_id or not current_response_id: - return [], "", "" + response_id = current_response_id or f"resp_{uuid.uuid4()}" + item_id = current_output_item_id or f"item_{uuid.uuid4()}" - # Parse the tool input. Nova 2 Sonic sends arguments in `content`; - # fall back to `input` for backward compatibility. tool_input = {} raw_input = tool_use["content"] if "content" in tool_use else tool_use.get("input") if raw_input: @@ -1042,15 +1040,13 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): tool_call_id = tool_use.get("toolUseId", "") tool_name = tool_use.get("toolName", "") - # Create a function call arguments done event - # This is a custom event format that matches what clients expect from typing import cast function_call_event: dict[str, Any] = { "type": "response.function_call_arguments.done", "event_id": f"event_{uuid.uuid4()}", - "response_id": current_response_id, - "item_id": current_output_item_id, + "response_id": response_id, + "item_id": item_id, "output_index": 0, "call_id": tool_call_id, "name": tool_name, 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 6611f3b431e..e45a39187e0 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 @@ -614,16 +614,30 @@ class TestBedrockRealtimeResponseTransformation: assert args["location"] == "San Francisco" def test_transform_tool_use_event_directly(self): - """Test transform_tool_use_event directly for guard clause and input parsing""" + """Test transform_tool_use_event directly for input parsing and missing IDs""" config = BedrockRealtimeConfig() - # Guard clause: missing IDs returns empty result + # Missing IDs still emit a function call (Nova Sonic starts tools with role=TOOL) events, tool_call_id, tool_name = config.transform_tool_use_event( - {"toolUse": {}}, None, "resp_123" + { + "toolUse": { + "toolUseId": "tool_call_no_ids", + "toolName": "get_weather", + "content": json.dumps({"location": "Seattle"}), + } + }, + None, + None, ) - assert events == [] - assert tool_call_id == "" - assert tool_name == "" + assert len(events) == 1 + assert events[0]["type"] == "response.function_call_arguments.done" + assert events[0]["call_id"] == "tool_call_no_ids" + assert events[0]["name"] == "get_weather" + assert events[0]["response_id"].startswith("resp_") + assert events[0]["item_id"].startswith("item_") + assert json.loads(events[0]["arguments"]) == {"location": "Seattle"} + assert tool_call_id == "tool_call_no_ids" + assert tool_name == "get_weather" # JSON string content is parsed and converted to a function call event events, tool_call_id, tool_name = config.transform_tool_use_event( @@ -641,6 +655,8 @@ class TestBedrockRealtimeResponseTransformation: assert events[0]["type"] == "response.function_call_arguments.done" assert events[0]["call_id"] == "tool_call_123" assert events[0]["name"] == "get_weather" + assert events[0]["response_id"] == "resp_123" + assert events[0]["item_id"] == "item_123" assert json.loads(events[0]["arguments"]) == {"location": "San Francisco"} # Invalid JSON content falls back to empty arguments From 0780c64ec067b566270852267f717749520ee6fb Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Mon, 3 Aug 2026 16:23:57 -0700 Subject: [PATCH 05/19] build(docker): install bedrock-realtime extra in root and database images Nova Sonic bidirectional streaming needs aws-sdk-bedrock-runtime; the gateway image already had --extra bedrock-realtime, the classic root and database images did not --- Dockerfile | 2 ++ docker/Dockerfile.database | 2 ++ 2 files changed, 4 insertions(+) diff --git a/Dockerfile b/Dockerfile index a127cdabd59..b0dffca9e96 100644 --- a/Dockerfile +++ b/Dockerfile @@ -65,6 +65,7 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr --extra extra_proxy \ --extra semantic-router \ --extra saml \ + --extra bedrock-realtime \ --python python3 # Copy full source tree @@ -86,6 +87,7 @@ RUN uv sync --frozen --no-default-groups --no-editable \ --extra extra_proxy \ --extra semantic-router \ --extra saml \ + --extra bedrock-realtime \ --python python3 RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ diff --git a/docker/Dockerfile.database b/docker/Dockerfile.database index 9ee076ce825..79da2302edc 100644 --- a/docker/Dockerfile.database +++ b/docker/Dockerfile.database @@ -63,6 +63,7 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr --extra extra_proxy \ --extra semantic-router \ --extra saml \ + --extra bedrock-realtime \ --python python3 # Copy full source tree @@ -84,6 +85,7 @@ RUN uv sync --frozen --no-default-groups --no-editable \ --extra extra_proxy \ --extra semantic-router \ --extra saml \ + --extra bedrock-realtime \ --python python3 RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ From ec9887ecb2cf5c6bab096a8e4a070d4f5def5908 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Mon, 3 Aug 2026 16:34:52 -0700 Subject: [PATCH 06/19] fix(bedrock/realtime): persist minted tool-call response/item ids transform_tool_use_event already minted ids for TOOL-first turns, but transform_realtime_response discarded them so later contentEnd events still saw None. Return and store the ids in session state; cover via transform_realtime_response regression --- .../llms/bedrock/realtime/transformation.py | 18 ++-- .../test_bedrock_realtime_transformation.py | 82 +++++++++++++++++-- 2 files changed, 89 insertions(+), 11 deletions(-) diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index b0b2c4f0201..809b1de64f2 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -1011,7 +1011,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): event: dict, current_output_item_id: str | None, current_response_id: str | None, - ) -> tuple[list[OpenAIRealtimeEvents], str, str]: + ) -> tuple[list[OpenAIRealtimeEvents], str, str, str, str]: """ Transform Bedrock toolUse event to OpenAI format. @@ -1021,7 +1021,8 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): current_response_id: Current response ID Returns: - Tuple of (events, tool_call_id, tool_name) for tracking + Tuple of (events, tool_call_id, tool_name, output_item_id, response_id) + so the caller can persist any minted IDs into session state """ verbose_logger.debug("Handling toolUse") tool_use = event["toolUse"] @@ -1057,6 +1058,8 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): [cast(OpenAIRealtimeEvents, function_call_event)], tool_call_id, tool_name, + item_id, + response_id, ) def transform_conversation_item_create_tool_result_event(self, json_message: dict) -> list[str]: @@ -1223,11 +1226,14 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): returned_messages.extend(done_events) elif "toolUse" in event: - events, tool_call_id, tool_name = self.transform_tool_use_event( - event, current_output_item_id, current_response_id - ) + ( + events, + tool_call_id, + tool_name, + current_output_item_id, + current_response_id, + ) = self.transform_tool_use_event(event, current_output_item_id, current_response_id) returned_messages.extend(events) - # Store tool call info for potential use verbose_logger.debug(f"Tool use event: {tool_name} (ID: {tool_call_id})") elif "promptEnd" in event or "completionEnd" in event: 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 e45a39187e0..521bcfa3992 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 @@ -618,7 +618,7 @@ class TestBedrockRealtimeResponseTransformation: config = BedrockRealtimeConfig() # Missing IDs still emit a function call (Nova Sonic starts tools with role=TOOL) - events, tool_call_id, tool_name = config.transform_tool_use_event( + events, tool_call_id, tool_name, item_id, response_id = config.transform_tool_use_event( { "toolUse": { "toolUseId": "tool_call_no_ids", @@ -633,14 +633,16 @@ class TestBedrockRealtimeResponseTransformation: assert events[0]["type"] == "response.function_call_arguments.done" assert events[0]["call_id"] == "tool_call_no_ids" assert events[0]["name"] == "get_weather" - assert events[0]["response_id"].startswith("resp_") - assert events[0]["item_id"].startswith("item_") + assert response_id.startswith("resp_") + assert item_id.startswith("item_") + assert events[0]["response_id"] == response_id + assert events[0]["item_id"] == item_id assert json.loads(events[0]["arguments"]) == {"location": "Seattle"} assert tool_call_id == "tool_call_no_ids" assert tool_name == "get_weather" # JSON string content is parsed and converted to a function call event - events, tool_call_id, tool_name = config.transform_tool_use_event( + events, tool_call_id, tool_name, item_id, response_id = config.transform_tool_use_event( { "toolUse": { "toolUseId": "tool_call_123", @@ -655,12 +657,14 @@ class TestBedrockRealtimeResponseTransformation: assert events[0]["type"] == "response.function_call_arguments.done" assert events[0]["call_id"] == "tool_call_123" assert events[0]["name"] == "get_weather" + assert response_id == "resp_123" + assert item_id == "item_123" assert events[0]["response_id"] == "resp_123" assert events[0]["item_id"] == "item_123" assert json.loads(events[0]["arguments"]) == {"location": "San Francisco"} # Invalid JSON content falls back to empty arguments - events, _, _ = config.transform_tool_use_event( + events, _, _, _, _ = config.transform_tool_use_event( { "toolUse": { "toolUseId": "tool_call_124", @@ -674,6 +678,74 @@ class TestBedrockRealtimeResponseTransformation: assert len(events) == 1 assert json.loads(events[0]["arguments"]) == {} + def test_transform_realtime_response_persists_minted_tool_ids(self): + """TOOL-first turns must write minted response/item ids into session state""" + config = BedrockRealtimeConfig() + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "trace_123" + + tool_use_message = { + "event": { + "toolUse": { + "toolUseId": "tool_call_state", + "toolName": "get_weather", + "content": json.dumps({"location": "Seattle"}), + } + } + } + + result = config.transform_realtime_response( + json.dumps(tool_use_message), + "amazon.nova-2-sonic-v1:0", + logging_obj, + realtime_response_transform_input={ + "session_configuration_request": json.dumps({"configured": True}), + "current_output_item_id": None, + "current_response_id": None, + "current_conversation_id": "conv_123", + "current_delta_chunks": [], + "current_item_chunks": [], + "current_delta_type": "text", + }, + ) + + assert len(result["response"]) == 1 + function_call = result["response"][0] + assert function_call["type"] == "response.function_call_arguments.done" + assert result["current_response_id"] is not None + assert result["current_output_item_id"] is not None + assert result["current_response_id"].startswith("resp_") + assert result["current_output_item_id"].startswith("item_") + assert function_call["response_id"] == result["current_response_id"] + assert function_call["item_id"] == result["current_output_item_id"] + assert json.loads(function_call["arguments"]) == {"location": "Seattle"} + + # A follow-up contentEnd must reuse the same persisted ids + content_end_message = { + "event": { + "contentEnd": { + "stopReason": "TOOL_USE", + "type": "TOOL", + } + } + } + follow_up = config.transform_realtime_response( + json.dumps(content_end_message), + "amazon.nova-2-sonic-v1:0", + logging_obj, + realtime_response_transform_input={ + "session_configuration_request": result["session_configuration_request"], + "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"], + }, + ) + assert follow_up["current_response_id"] == result["current_response_id"] + assert follow_up["current_output_item_id"] == result["current_output_item_id"] + def test_transform_content_end_text(self): """Test contentEnd for text response""" config = BedrockRealtimeConfig() From 053973160d79072f1788d93d400386d473401b43 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Wed, 5 Aug 2026 15:28:03 -0700 Subject: [PATCH 07/19] fix(bedrock/realtime): skip message output_item.done on TOOL contentEnd Minted tool ids must not unlock unpaired message completion events after TOOL-first turns; only text/audio content ends emit output_item.done --- .../llms/bedrock/realtime/transformation.py | 5 +- .../test_bedrock_realtime_transformation.py | 76 ++++++++++++++++--- 2 files changed, 70 insertions(+), 11 deletions(-) diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index 3902c791f06..1510211c77c 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -871,6 +871,9 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): if not current_output_item_id or not current_response_id: return [], current_delta_chunks + if content_end.get("type") == "TOOL" or current_delta_type not in ("text", "audio"): + return [], current_delta_chunks + returned_messages: list[OpenAIRealtimeEvents] = [] # Send appropriate done event based on type @@ -903,7 +906,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): ) returned_messages.append(content_part_done) - elif current_delta_type == "audio": + else: audio_done = OpenAIRealtimeResponseAudioDone( type="response.audio.done", content_index=0, 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 521bcfa3992..9b7456c5155 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 @@ -684,6 +684,35 @@ class TestBedrockRealtimeResponseTransformation: logging_obj = MagicMock() logging_obj.litellm_trace_id = "trace_123" + state = { + "session_configuration_request": json.dumps({"configured": True}), + "current_output_item_id": None, + "current_response_id": None, + "current_conversation_id": "conv_123", + "current_delta_chunks": [], + "current_item_chunks": [], + "current_delta_type": None, + } + + content_start_result = config.transform_realtime_response( + json.dumps({"event": {"contentStart": {"role": "TOOL", "type": "TOOL"}}}), + "amazon.nova-2-sonic-v1:0", + logging_obj, + realtime_response_transform_input=state, + ) + assert content_start_result["response"] == [] + assert content_start_result["current_delta_type"] is None + state.update( + { + "current_output_item_id": content_start_result["current_output_item_id"], + "current_response_id": content_start_result["current_response_id"], + "current_conversation_id": content_start_result["current_conversation_id"], + "current_delta_chunks": content_start_result["current_delta_chunks"], + "current_item_chunks": content_start_result["current_item_chunks"], + "current_delta_type": content_start_result["current_delta_type"], + } + ) + tool_use_message = { "event": { "toolUse": { @@ -698,15 +727,7 @@ class TestBedrockRealtimeResponseTransformation: json.dumps(tool_use_message), "amazon.nova-2-sonic-v1:0", logging_obj, - realtime_response_transform_input={ - "session_configuration_request": json.dumps({"configured": True}), - "current_output_item_id": None, - "current_response_id": None, - "current_conversation_id": "conv_123", - "current_delta_chunks": [], - "current_item_chunks": [], - "current_delta_type": "text", - }, + realtime_response_transform_input=state, ) assert len(result["response"]) == 1 @@ -720,7 +741,6 @@ class TestBedrockRealtimeResponseTransformation: assert function_call["item_id"] == result["current_output_item_id"] assert json.loads(function_call["arguments"]) == {"location": "Seattle"} - # A follow-up contentEnd must reuse the same persisted ids content_end_message = { "event": { "contentEnd": { @@ -745,6 +765,42 @@ class TestBedrockRealtimeResponseTransformation: ) assert follow_up["current_response_id"] == result["current_response_id"] assert follow_up["current_output_item_id"] == result["current_output_item_id"] + assert follow_up["response"] == [] + assert all(msg["type"] != "response.output_item.done" for msg in follow_up["response"]) + + def test_tool_content_end_does_not_emit_message_output_item_done(self): + """Minted tool ids must not unlock unpaired message output_item.done on TOOL contentEnd""" + config = BedrockRealtimeConfig() + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "trace_123" + + content_end_message = { + "event": { + "contentEnd": { + "stopReason": "TOOL_USE", + "type": "TOOL", + } + } + } + result = config.transform_realtime_response( + json.dumps(content_end_message), + "amazon.nova-2-sonic-v1:0", + logging_obj, + realtime_response_transform_input={ + "session_configuration_request": json.dumps({"configured": True}), + "current_output_item_id": "item_minted_for_tool", + "current_response_id": "resp_minted_for_tool", + "current_conversation_id": "conv_123", + "current_delta_chunks": [], + "current_item_chunks": [], + "current_delta_type": "text", + }, + ) + + assert result["response"] == [] + assert result["current_response_id"] == "resp_minted_for_tool" + assert result["current_output_item_id"] == "item_minted_for_tool" + assert result["current_delta_type"] == "text" def test_transform_content_end_text(self): """Test contentEnd for text response""" From 9948da46bf31ab02e3a86e6d0333554b087f5960 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 6 Aug 2026 14:12:06 -0700 Subject: [PATCH 08/19] feat(pricing): add Bedrock Nova Sonic realtime model costs Register amazon.nova-sonic-v1:0 and amazon.nova-2-sonic-v1:0 with US East (N. Virginia) on-demand speech and text token rates from AWS Bedrock pricing so realtime spend tracking resolves correctly --- ...odel_prices_and_context_window_backup.json | 56 +++++++++++++++++++ model_prices_and_context_window.json | 56 +++++++++++++++++++ 2 files changed, 112 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 749663775ea..af1dfbbc103 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -46257,5 +46257,61 @@ } } ] + }, + "amazon.nova-sonic-v1:0": { + "input_cost_per_audio_token": 3.4e-06, + "input_cost_per_token": 6e-08, + "litellm_provider": "bedrock", + "max_input_tokens": 300000, + "max_output_tokens": 10000, + "max_tokens": 10000, + "mode": "realtime", + "output_cost_per_audio_token": 1.36e-05, + "output_cost_per_token": 2.4e-07, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "amazon.nova-2-sonic-v1:0": { + "input_cost_per_audio_token": 3e-06, + "input_cost_per_token": 3.3e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "realtime", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 2.75e-06, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true } } diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 8a8bbdad8f1..41f8df4cf4c 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -46379,5 +46379,61 @@ } } ] + }, + "amazon.nova-sonic-v1:0": { + "input_cost_per_audio_token": 3.4e-06, + "input_cost_per_token": 6e-08, + "litellm_provider": "bedrock", + "max_input_tokens": 300000, + "max_output_tokens": 10000, + "max_tokens": 10000, + "mode": "realtime", + "output_cost_per_audio_token": 1.36e-05, + "output_cost_per_token": 2.4e-07, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "amazon.nova-2-sonic-v1:0": { + "input_cost_per_audio_token": 3e-06, + "input_cost_per_token": 3.3e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "realtime", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 2.75e-06, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true } } From f9e3d7d7195244b73601666cbaa26187c331282f Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Fri, 7 Aug 2026 13:31:57 -0700 Subject: [PATCH 09/19] fix(bedrock/realtime): clear tool ids on TOOL contentEnd TOOL contentEnd left minted response/item ids active, so the next ASSISTANT contentStart reused them and broke post-tool correlation. Clear those ids when a TOOL block ends and assert the next assistant response mints distinct ids --- .../llms/bedrock/realtime/transformation.py | 7 +++- .../test_bedrock_realtime_transformation.py | 39 ++++++++++++++++--- 2 files changed, 40 insertions(+), 6 deletions(-) diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index 4a475389f70..578a2aa11c5 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -1211,6 +1211,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): returned_messages.extend(events) elif "contentEnd" in event: + content_end: Final = event["contentEnd"] events, current_delta_chunks = self.transform_content_end_event( event, current_output_item_id, @@ -1219,7 +1220,11 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): current_delta_chunks, ) returned_messages.extend(events) - if BedrockContentEnd.model_validate(event["contentEnd"]).stopReason == "END_TURN": + if content_end.get("type") == "TOOL": + current_output_item_id = None + current_response_id = None + current_delta_type = None + if BedrockContentEnd.model_validate(content_end).stopReason == "END_TURN": ( done_events, current_output_item_id, 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 9b7456c5155..798bc86023d 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 @@ -763,11 +763,40 @@ class TestBedrockRealtimeResponseTransformation: "current_delta_type": result["current_delta_type"], }, ) - assert follow_up["current_response_id"] == result["current_response_id"] - assert follow_up["current_output_item_id"] == result["current_output_item_id"] + assert follow_up["current_response_id"] is None + assert follow_up["current_output_item_id"] is None + assert follow_up["current_delta_type"] is None assert follow_up["response"] == [] assert all(msg["type"] != "response.output_item.done" for msg in follow_up["response"]) + post_tool_state = { + "session_configuration_request": follow_up["session_configuration_request"], + "current_output_item_id": follow_up["current_output_item_id"], + "current_response_id": follow_up["current_response_id"], + "current_conversation_id": follow_up["current_conversation_id"], + "current_delta_chunks": follow_up["current_delta_chunks"], + "current_item_chunks": follow_up["current_item_chunks"], + "current_delta_type": follow_up["current_delta_type"], + } + assistant_start = config.transform_realtime_response( + json.dumps({"event": {"contentStart": {"role": "ASSISTANT", "type": "TEXT"}}}), + "amazon.nova-2-sonic-v1:0", + logging_obj, + realtime_response_transform_input=post_tool_state, + ) + tool_response_id = result["current_response_id"] + tool_item_id = result["current_output_item_id"] + assert assistant_start["current_response_id"] is not None + assert assistant_start["current_output_item_id"] is not None + assert assistant_start["current_response_id"] != tool_response_id + assert assistant_start["current_output_item_id"] != tool_item_id + created = [msg for msg in assistant_start["response"] if msg["type"] == "response.created"][0] + added = [msg for msg in assistant_start["response"] if msg["type"] == "response.output_item.added"][0] + assert created["response"]["id"] == assistant_start["current_response_id"] + assert added["item"]["id"] == assistant_start["current_output_item_id"] + assert created["response"]["id"] != function_call["response_id"] + assert added["item"]["id"] != function_call["item_id"] + def test_tool_content_end_does_not_emit_message_output_item_done(self): """Minted tool ids must not unlock unpaired message output_item.done on TOOL contentEnd""" config = BedrockRealtimeConfig() @@ -798,9 +827,9 @@ class TestBedrockRealtimeResponseTransformation: ) assert result["response"] == [] - assert result["current_response_id"] == "resp_minted_for_tool" - assert result["current_output_item_id"] == "item_minted_for_tool" - assert result["current_delta_type"] == "text" + assert result["current_response_id"] is None + assert result["current_output_item_id"] is None + assert result["current_delta_type"] is None def test_transform_content_end_text(self): """Test contentEnd for text response""" From 7444897f1375cab3611ba328b64547bc693f4672 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Fri, 7 Aug 2026 13:41:31 -0700 Subject: [PATCH 10/19] fix(bedrock/realtime): satisfy LIT010 on tool id state updates Annotate minted tool ids with Final, parse tool args without rebinding, and mark intentional session state clears/persists with rebind-ok --- .../llms/bedrock/realtime/transformation.py | 38 +++++++++++-------- 1 file changed, 22 insertions(+), 16 deletions(-) diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index 578a2aa11c5..4d9158853cd 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -50,6 +50,17 @@ TRIGGER_TRAILING_SILENCE: Final = bytes(TRIGGER_AUDIO_BYTES_PER_SECOND * 3) TRIGGER_AUDIO_CHUNK_SIZE: Final = 1024 +def _parse_bedrock_tool_use_input(raw_input: object) -> object: + if not raw_input: + return {} + if not isinstance(raw_input, str): + return raw_input + try: + return json.loads(raw_input) + except json.JSONDecodeError: + return {} + + class BedrockRealtimeConfig(BaseRealtimeConfig): """Configuration for Bedrock Nova Sonic realtime transformations.""" @@ -1030,16 +1041,10 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): verbose_logger.debug("Handling toolUse") tool_use: Final = event["toolUse"] - response_id = current_response_id or f"resp_{uuid.uuid4()}" - item_id = current_output_item_id or f"item_{uuid.uuid4()}" - - tool_input = {} - raw_input = tool_use["content"] if "content" in tool_use else tool_use.get("input") - if raw_input: - try: - tool_input = json.loads(raw_input) if isinstance(raw_input, str) else raw_input - except json.JSONDecodeError: - tool_input = {} + response_id: Final = current_response_id or f"resp_{uuid.uuid4()}" + item_id: Final = current_output_item_id or f"item_{uuid.uuid4()}" + raw_input: Final = tool_use["content"] if "content" in tool_use else tool_use.get("input") + tool_input: Final = _parse_bedrock_tool_use_input(raw_input) tool_call_id: Final = tool_use.get("toolUseId", "") tool_name: Final = tool_use.get("toolName", "") @@ -1221,9 +1226,9 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): ) returned_messages.extend(events) if content_end.get("type") == "TOOL": - current_output_item_id = None - current_response_id = None - current_delta_type = None + 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 if BedrockContentEnd.model_validate(content_end).stopReason == "END_TURN": ( done_events, @@ -1238,11 +1243,12 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): events, tool_call_id, tool_name, - current_output_item_id, - current_response_id, + tool_output_item_id, + tool_response_id, ) = self.transform_tool_use_event(event, current_output_item_id, current_response_id) returned_messages.extend(events) - # Store tool call info for potential use + 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 verbose_logger.debug("Tool use event: %s (ID: %s)", tool_name, tool_call_id) elif "promptEnd" in event or "completionEnd" in event: From 7d60bf61267aabc551cf7dae44bd089eeea3a1ed Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Fri, 7 Aug 2026 15:56:37 -0700 Subject: [PATCH 11/19] fix(bedrock/realtime): meter usage and log via RealTimeStreaming Parse Nova Sonic usageEvent into turn-level response.done usage, reuse RealTimeStreaming store_message/store_input/log_messages for spend and budget accounting, and pass user key metadata through arealtime --- litellm/llms/bedrock/realtime/handler.py | 120 +++++++++++++----- .../llms/bedrock/realtime/transformation.py | 94 ++++++++++++-- litellm/realtime_api/main.py | 2 + .../realtime/test_bedrock_realtime_handler.py | 51 ++++++++ .../test_bedrock_realtime_transformation.py | 111 ++++++++++++++++ 5 files changed, 337 insertions(+), 41 deletions(-) diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 3eeb3cb9fc6..c8035a5265d 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -2,17 +2,20 @@ This file contains the handler for AWS Bedrock Nova Sonic realtime API. This uses aws_sdk_bedrock_runtime for bidirectional streaming with Nova Sonic. +Spend / budget logging follows the same RealTimeStreaming path as OpenAI/Azure: +store_message for backend events, store_input for client events, log_messages on close. """ import asyncio import contextlib import json -from typing import Any, Final +from typing import Any, Final, cast from pydantic import TypeAdapter from litellm._logging import _redact_string, verbose_proxy_logger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming from ..base_aws_llm import BaseAWSLLM from ..common_utils import BedrockError @@ -46,6 +49,8 @@ class BedrockRealtime(BaseAWSLLM): aws_sts_endpoint: str | None = None, aws_bedrock_runtime_endpoint: str | None = None, aws_external_id: str | None = None, + user_api_key_dict: Any | None = None, + litellm_metadata: dict | None = None, **kwargs, ): """ @@ -120,6 +125,27 @@ class BedrockRealtime(BaseAWSLLM): transformation_config: Final = BedrockRealtimeConfig() + logging_obj.pre_call( + input=None, + api_key=api_key or "", + additional_args={ + "api_base": endpoint_uri, + "complete_input_dict": {"model": model}, + }, + ) + + # RealTimeStreaming owns spend logging for other realtime providers. Bedrock cannot + # use its WebSocket bidirectional_forward (AWS SDK stream instead), but store_message / + # store_input / log_messages are the same path used by OpenAI and Azure. + realtime_streaming: Final = RealTimeStreaming( + websocket=websocket, + backend_ws=cast(Any, object()), + logging_obj=logging_obj, + model=model, + user_api_key_dict=user_api_key_dict, + request_data={"litellm_metadata": litellm_metadata or {}}, + ) + try: # Initialize the bidirectional stream bedrock_stream: Final = await bedrock_client.invoke_model_with_bidirectional_stream( @@ -128,7 +154,9 @@ class BedrockRealtime(BaseAWSLLM): verbose_proxy_logger.debug("Bedrock Realtime: Bidirectional stream established") - await websocket.send_text(json.dumps(transformation_config.session_created_event(model, logging_obj))) + session_created: Final = transformation_config.session_created_event(model, logging_obj) + realtime_streaming.store_message(session_created) + await websocket.send_text(json.dumps(session_created)) verbose_proxy_logger.debug("Bedrock Realtime: sent session.created to client on connect") # Track state for transformation @@ -142,35 +170,38 @@ class BedrockRealtime(BaseAWSLLM): "session_configuration_request": None, } - # Create tasks for bidirectional forwarding - client_to_bedrock_task: Final = asyncio.create_task( - self._forward_client_to_bedrock( - websocket, - bedrock_stream, - transformation_config, - model, - session_state, - logging_obj, + try: + client_to_bedrock_task: Final = asyncio.create_task( + self._forward_client_to_bedrock( + websocket, + bedrock_stream, + transformation_config, + model, + session_state, + logging_obj, + realtime_streaming, + ) ) - ) - bedrock_to_client_task: Final = asyncio.create_task( - self._forward_bedrock_to_client( - bedrock_stream, - websocket, - transformation_config, - model, - logging_obj, - session_state, + bedrock_to_client_task: Final = asyncio.create_task( + self._forward_bedrock_to_client( + bedrock_stream, + websocket, + transformation_config, + model, + logging_obj, + session_state, + realtime_streaming, + ) ) - ) - # Wait for both tasks to complete - await asyncio.gather( - client_to_bedrock_task, - bedrock_to_client_task, - return_exceptions=True, - ) + await asyncio.gather( + client_to_bedrock_task, + bedrock_to_client_task, + return_exceptions=True, + ) + finally: + await realtime_streaming.log_messages() except Exception as e: verbose_proxy_logger.exception("Error in BedrockRealtime.async_realtime: %s", e) @@ -180,6 +211,24 @@ class BedrockRealtime(BaseAWSLLM): pass raise + @staticmethod + def _collect_tool_call_from_function_call_event( + realtime_streaming: RealTimeStreaming, + message: object, + ) -> None: + if not isinstance(message, dict) or message.get("type") != "response.function_call_arguments.done": + return + realtime_streaming.tool_calls.append( + { + "id": message.get("call_id", ""), + "type": "function", + "function": { + "name": message.get("name", ""), + "arguments": message.get("arguments", "{}"), + }, + } + ) + async def _forward_client_to_bedrock( self, client_ws: Any, @@ -188,6 +237,7 @@ class BedrockRealtime(BaseAWSLLM): model: str, session_state: dict, logging_obj: LiteLLMLogging | None = None, + realtime_streaming: RealTimeStreaming | None = None, ): """Forward messages from client WebSocket to Bedrock stream.""" from aws_sdk_bedrock_runtime.models import ( @@ -208,6 +258,9 @@ class BedrockRealtime(BaseAWSLLM): message = await client_ws.receive_text() verbose_proxy_logger.debug("Bedrock Realtime: Received from client: %s", message[:200]) + if realtime_streaming is not None: + realtime_streaming.store_input(message) + # Transform OpenAI format to Bedrock format transformed_messages = transformation_config.transform_realtime_request( message=message, @@ -230,11 +283,12 @@ class BedrockRealtime(BaseAWSLLM): parsed_client_message.get("session", {}).get("modalities") ) if client_message_type == "session.update": - await client_ws.send_text( - json.dumps( - transformation_config.session_updated_event(model, logging_obj, requested_modalities) - ) + session_updated: Final = transformation_config.session_updated_event( + model, logging_obj, requested_modalities ) + if realtime_streaming is not None: + realtime_streaming.store_message(session_updated) + await client_ws.send_text(json.dumps(session_updated)) except Exception as e: verbose_proxy_logger.debug("Client to Bedrock forwarding ended: %s", e, exc_info=True) @@ -252,6 +306,7 @@ class BedrockRealtime(BaseAWSLLM): model: str, logging_obj: LiteLLMLogging, session_state: dict, + realtime_streaming: RealTimeStreaming | None = None, ): """Forward messages from Bedrock stream to client WebSocket.""" try: @@ -304,6 +359,9 @@ class BedrockRealtime(BaseAWSLLM): # Send transformed messages to client openai_messages = transformed_response.get("response", []) for openai_message in openai_messages: + if realtime_streaming is not None: + realtime_streaming.store_message(openai_message) + self._collect_tool_call_from_function_call_event(realtime_streaming, 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]) diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index 4d9158853cd..98716ae89c7 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -36,9 +36,6 @@ from litellm.types.realtime import ( RealtimeResponseTransformInput, RealtimeResponseTypedDict, ) -from litellm.utils import get_empty_usage - - class BedrockContentEnd(BaseModel): stopReason: str | None = None @@ -61,6 +58,74 @@ def _parse_bedrock_tool_use_input(raw_input: object) -> object: return {} +def _as_nonneg_int(value: object) -> int: + if isinstance(value, bool) or not isinstance(value, (int, float)): + return 0 + return max(0, int(value)) + + +def _empty_usage_snapshot() -> dict[str, int]: + return { + "input_speech": 0, + "input_text": 0, + "output_speech": 0, + "output_text": 0, + "total_input": 0, + "total_output": 0, + "total": 0, + } + + +def _usage_snapshot_from_event(usage_event: dict) -> dict[str, int]: + details: Final = usage_event.get("details") if isinstance(usage_event.get("details"), dict) else {} + total_block: Final = details.get("total") if isinstance(details.get("total"), dict) else {} + input_block: Final = total_block.get("input") if isinstance(total_block.get("input"), dict) else {} + output_block: Final = total_block.get("output") if isinstance(total_block.get("output"), dict) else {} + input_speech: Final = _as_nonneg_int(input_block.get("speechTokens")) + input_text: Final = _as_nonneg_int(input_block.get("textTokens")) + output_speech: Final = _as_nonneg_int(output_block.get("speechTokens")) + output_text: Final = _as_nonneg_int(output_block.get("textTokens")) + total_input: Final = _as_nonneg_int(usage_event.get("totalInputTokens")) or (input_speech + input_text) + total_output: Final = _as_nonneg_int(usage_event.get("totalOutputTokens")) or (output_speech + output_text) + total: Final = _as_nonneg_int(usage_event.get("totalTokens")) or (total_input + total_output) + return { + "input_speech": input_speech, + "input_text": input_text, + "output_speech": output_speech, + "output_text": output_text, + "total_input": total_input, + "total_output": total_output, + "total": total, + } + + +def _usage_snapshot_delta(current: dict[str, int], previous: dict[str, int]) -> dict[str, int]: + return {key: max(0, current.get(key, 0) - previous.get(key, 0)) for key in _empty_usage_snapshot()} + + +def _openai_usage_from_snapshot(snapshot: dict[str, int]) -> dict[str, Any]: + input_tokens: Final = snapshot.get("total_input", 0) or ( + snapshot.get("input_speech", 0) + snapshot.get("input_text", 0) + ) + output_tokens: Final = snapshot.get("total_output", 0) or ( + snapshot.get("output_speech", 0) + snapshot.get("output_text", 0) + ) + return { + "total_tokens": snapshot.get("total", 0) or (input_tokens + output_tokens), + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "input_token_details": { + "text_tokens": snapshot.get("input_text", 0), + "audio_tokens": snapshot.get("input_speech", 0), + "cached_tokens": 0, + }, + "output_token_details": { + "text_tokens": snapshot.get("output_text", 0), + "audio_tokens": snapshot.get("output_speech", 0), + }, + } + + class BedrockRealtimeConfig(BaseRealtimeConfig): """Configuration for Bedrock Nova Sonic realtime transformations.""" @@ -71,6 +136,8 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): self.audio_content_name = str(uuid_lib.uuid4()) self.prompt_started = False self.client_audio_streamed = False + self._usage_totals = _empty_usage_snapshot() + self._usage_at_last_response_done = _empty_usage_snapshot() # Default configuration values # Inference configuration @@ -98,6 +165,14 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): # Text configuration self.text_media_type = "text/plain" + def record_usage_event(self, usage_event: dict) -> None: + self._usage_totals = _usage_snapshot_from_event(usage_event) + + def consume_usage_for_response_done(self) -> dict[str, Any]: + delta: Final = _usage_snapshot_delta(self._usage_totals, self._usage_at_last_response_done) + self._usage_at_last_response_done = dict(self._usage_totals) + return _openai_usage_from_snapshot(delta) + def validate_environment(self, headers: dict, model: str, api_key: str | None = None) -> dict: """Validate environment - no special validation needed for Bedrock.""" return headers @@ -999,7 +1074,6 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): if not current_response_id or not current_conversation_id: return [], None, None, None - usage_obj: Final = get_empty_usage() response_done: Final = OpenAIRealtimeDoneEvent( type="response.done", event_id=f"event_{uuid.uuid4()}", @@ -1009,15 +1083,10 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): status="completed", output=[], conversation_id=current_conversation_id, - usage={ - "prompt_tokens": usage_obj.prompt_tokens, - "completion_tokens": usage_obj.completion_tokens, - "total_tokens": usage_obj.total_tokens, - }, + usage=self.consume_usage_for_response_done(), ), ) - # Reset state for next response return [response_done], None, None, None def transform_tool_use_event( @@ -1187,6 +1256,11 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): if "sessionStart" in event: session_configuration_request = json.dumps({"configured": True}) + elif "usageEvent" in event: + usage_event: Final = event["usageEvent"] + if isinstance(usage_event, dict): + self.record_usage_event(usage_event) + elif "contentStart" in event: ( events, diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index d5195659b1c..203537d689a 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -434,6 +434,8 @@ async def _arealtime( aws_sts_endpoint=aws_sts_endpoint, aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint, aws_external_id=aws_external_id, + user_api_key_dict=kwargs.get("user_api_key_dict"), + litellm_metadata=_build_litellm_metadata(kwargs), ) elif _custom_llm_provider == "xai": api_base = ( 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 ffe21b91ab2..e1d5ef72c00 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 @@ -57,6 +57,33 @@ class FakeBedrockStream: class FakeLogging: def __init__(self, trace_id="trace-nova-sonic"): self.litellm_trace_id = trace_id + self.dispatched_results = [] + self.model_call_details = {} + self.pre_call_args = [] + + def pre_call(self, input=None, api_key="", model=None, additional_args=None): + self.pre_call_args.append( + {"input": input, "api_key": api_key, "model": model, "additional_args": additional_args or {}} + ) + + async def dispatch_success_handlers(self, result=None, prefer_async_handlers=False, **kwargs): + self.dispatched_results.append(result) + + +@pytest.fixture(autouse=True) +def drain_bedrock_realtime_logging_worker(monkeypatch): + pending = [] + + def capture_enqueue(coro): + pending.append(coro) + + monkeypatch.setattr( + "litellm.litellm_core_utils.realtime_streaming.GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue", + capture_enqueue, + ) + yield pending + for coro in pending: + coro.close() class DisconnectingClientWS: @@ -313,6 +340,30 @@ class TestBedrockRealtimeSessionLifecycle: assert first_event["session"]["id"] == "trace-nova-sonic" assert first_event["session"]["model"] == "amazon.nova-sonic-v1:0" + @pytest.mark.asyncio + async def test_session_dispatches_logged_events_via_realtime_streaming( + self, stub_aws_sdk_client, drain_bedrock_realtime_logging_worker + ): + handler = BedrockRealtime() + websocket = RealtimeClientWS() + logging_obj = FakeLogging() + + await handler.async_realtime( + model="amazon.nova-sonic-v1:0", + websocket=websocket, + logging_obj=logging_obj, + aws_region_name="us-east-1", + aws_access_key_id="k", + aws_secret_access_key="s", + ) + + assert logging_obj.pre_call_args + assert len(drain_bedrock_realtime_logging_worker) == 1 + await drain_bedrock_realtime_logging_worker.pop() + assert logging_obj.dispatched_results + dispatched = logging_obj.dispatched_results[0] + assert any(event.get("type") == "session.created" for event in dispatched) + @pytest.mark.asyncio async def test_session_update_is_acked_with_session_updated(self, stub_aws_models): handler = BedrockRealtime() 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 798bc86023d..79b662fa841 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,5 +1090,116 @@ class TestBedrockRealtimeSessionEvents: assert event["session"]["modalities"] == ["text", "audio"] +class TestBedrockRealtimeUsageAccounting: + def _usage_event( + self, + *, + input_speech: int, + input_text: int, + output_speech: int, + output_text: int, + total_input: int | None = None, + total_output: int | None = None, + total: int | None = None, + ) -> dict: + resolved_input = total_input if total_input is not None else input_speech + input_text + resolved_output = total_output if total_output is not None else output_speech + output_text + resolved_total = total if total is not None else resolved_input + resolved_output + return { + "event": { + "usageEvent": { + "completionId": "completion_1", + "details": { + "total": { + "input": {"speechTokens": input_speech, "textTokens": input_text}, + "output": {"speechTokens": output_speech, "textTokens": output_text}, + } + }, + "promptName": "prompt_1", + "sessionId": "session_1", + "totalInputTokens": resolved_input, + "totalOutputTokens": resolved_output, + "totalTokens": resolved_total, + } + } + } + + def test_usage_event_fills_response_done_turn_delta(self): + config = BedrockRealtimeConfig() + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "trace_123" + state = { + "session_configuration_request": json.dumps({"configured": True}), + "current_output_item_id": "item_1", + "current_response_id": "resp_1", + "current_conversation_id": "conv_1", + "current_delta_chunks": [], + "current_item_chunks": [], + "current_delta_type": "audio", + } + + config.transform_realtime_response( + json.dumps(self._usage_event(input_speech=10, input_text=2, output_speech=20, output_text=3)), + "amazon.nova-sonic-v1:0", + logging_obj, + realtime_response_transform_input=state, + ) + first_done = config.transform_realtime_response( + json.dumps({"event": {"contentEnd": {"stopReason": "END_TURN", "type": "AUDIO"}}}), + "amazon.nova-sonic-v1:0", + logging_obj, + realtime_response_transform_input=state, + ) + done_events = [msg for msg in first_done["response"] if msg["type"] == "response.done"] + assert len(done_events) == 1 + usage = done_events[0]["response"]["usage"] + assert usage["input_tokens"] == 12 + assert usage["output_tokens"] == 23 + assert usage["total_tokens"] == 35 + assert usage["input_token_details"]["audio_tokens"] == 10 + assert usage["input_token_details"]["text_tokens"] == 2 + assert usage["output_token_details"]["audio_tokens"] == 20 + assert usage["output_token_details"]["text_tokens"] == 3 + + config.transform_realtime_response( + json.dumps( + self._usage_event( + input_speech=15, + input_text=2, + output_speech=30, + output_text=3, + total_input=17, + total_output=33, + total=50, + ) + ), + "amazon.nova-sonic-v1:0", + logging_obj, + realtime_response_transform_input={ + **state, + "current_output_item_id": "item_2", + "current_response_id": "resp_2", + }, + ) + second_done = config.transform_realtime_response( + json.dumps({"event": {"contentEnd": {"stopReason": "END_TURN", "type": "AUDIO"}}}), + "amazon.nova-sonic-v1:0", + logging_obj, + realtime_response_transform_input={ + **state, + "current_output_item_id": "item_2", + "current_response_id": "resp_2", + }, + ) + second_usage = [msg for msg in second_done["response"] if msg["type"] == "response.done"][0]["response"][ + "usage" + ] + assert second_usage["input_tokens"] == 5 + assert second_usage["output_tokens"] == 10 + assert second_usage["total_tokens"] == 15 + assert second_usage["input_token_details"]["audio_tokens"] == 5 + assert second_usage["output_token_details"]["audio_tokens"] == 10 + + if __name__ == "__main__": pytest.main([__file__, "-v"]) From a098330b519ddc50dedb2ae2e226fa561f31a570 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Fri, 7 Aug 2026 16:20:55 -0700 Subject: [PATCH 12/19] fix(bedrock/realtime): content-block session state lifecycle Nova Sonic streams discrete content blocks. Reset part-local state (delta chunks/type/item) on every contentEnd, mint a new item per assistant contentStart, and only emit response.created for a new response so text/tool/audio blocks cannot leak into each other --- .../llms/bedrock/realtime/transformation.py | 76 ++++----- .../test_bedrock_realtime_transformation.py | 156 ++++++++++++++++++ 2 files changed, 194 insertions(+), 38 deletions(-) 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, From e4e5f778bcece15d6f1018b2758c29cb79dbd4ac Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Sat, 8 Aug 2026 12:59:49 -0700 Subject: [PATCH 13/19] fix(bedrock/realtime): flush pending usage and type UserAPIKeyAuth Bill late usageEvent after response ids clear, TOOL END_TURN response.done, session-close drain, and completionEnd mint. Type user_api_key_dict as UserAPIKeyAuth on Bedrock realtime and RealTimeStreaming. --- .../litellm_core_utils/realtime_streaming.py | 6 +- litellm/llms/bedrock/realtime/handler.py | 12 +- .../llms/bedrock/realtime/transformation.py | 60 +++++++-- .../test_bedrock_realtime_transformation.py | 120 +++++++++++++++++- 4 files changed, 183 insertions(+), 15 deletions(-) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 858d10df53b..4ab109e5830 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -20,6 +20,8 @@ from .litellm_logging import Logging as LiteLLMLogging if TYPE_CHECKING: from websockets.asyncio.client import ClientConnection + from litellm.proxy._types import UserAPIKeyAuth + CLIENT_CONNECTION_CLASS = ClientConnection else: CLIENT_CONNECTION_CLASS = Any @@ -48,7 +50,7 @@ class RealTimeStreaming: logging_obj: LiteLLMLogging, provider_config: BaseRealtimeConfig | None = None, model: str = "", - user_api_key_dict: Any | None = None, + user_api_key_dict: "UserAPIKeyAuth | None" = None, request_data: dict | None = None, backend_uses_beta_protocol: bool | None = None, force_transcription_model: str | None = None, @@ -83,7 +85,7 @@ class RealTimeStreaming: self.current_item_chunks: list[OpenAIRealtimeOutputItemDone] | None = None self.current_delta_type: ALL_DELTA_TYPES | None = None self.session_configuration_request: str | None = None - self.user_api_key_dict = user_api_key_dict + self.user_api_key_dict: "UserAPIKeyAuth | None" = user_api_key_dict self.request_data: dict = request_data or {} # Violation counter for end_session_after_n_fails support self._violation_count: int = 0 diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index c8035a5265d..3ab0306fff5 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -9,7 +9,7 @@ store_message for backend events, store_input for client events, log_messages on import asyncio import contextlib import json -from typing import Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, cast from pydantic import TypeAdapter @@ -21,6 +21,9 @@ from ..base_aws_llm import BaseAWSLLM from ..common_utils import BedrockError from .transformation import BedrockRealtimeConfig +if TYPE_CHECKING: + from litellm.proxy._types import UserAPIKeyAuth + _CLIENT_MODALITIES_ADAPTER: Final[TypeAdapter["list[str] | None"]] = TypeAdapter(list[str] | None) @@ -49,7 +52,7 @@ class BedrockRealtime(BaseAWSLLM): aws_sts_endpoint: str | None = None, aws_bedrock_runtime_endpoint: str | None = None, aws_external_id: str | None = None, - user_api_key_dict: Any | None = None, + user_api_key_dict: "UserAPIKeyAuth | None" = None, litellm_metadata: dict | None = None, **kwargs, ): @@ -201,6 +204,11 @@ class BedrockRealtime(BaseAWSLLM): return_exceptions=True, ) finally: + for pending_usage_event in transformation_config.flush_pending_usage_as_response_done( + session_state.get("current_response_id"), + session_state.get("current_conversation_id"), + ): + realtime_streaming.store_message(pending_usage_event) await realtime_streaming.log_messages() except Exception as e: diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index 8286058bbc6..ec1e541a82e 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -36,6 +36,8 @@ from litellm.types.realtime import ( RealtimeResponseTransformInput, RealtimeResponseTypedDict, ) + + class BedrockContentEnd(BaseModel): stopReason: str | None = None @@ -168,11 +170,29 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): def record_usage_event(self, usage_event: dict) -> None: self._usage_totals = _usage_snapshot_from_event(usage_event) + def has_unbilled_usage(self) -> bool: + delta: Final = _usage_snapshot_delta(self._usage_totals, self._usage_at_last_response_done) + return any(value > 0 for value in delta.values()) + def consume_usage_for_response_done(self) -> dict[str, Any]: delta: Final = _usage_snapshot_delta(self._usage_totals, self._usage_at_last_response_done) self._usage_at_last_response_done = dict(self._usage_totals) return _openai_usage_from_snapshot(delta) + def flush_pending_usage_as_response_done( + self, + current_response_id: str | None = None, + current_conversation_id: str | None = None, + ) -> list[OpenAIRealtimeEvents]: + if not self.has_unbilled_usage(): + return [] + events, _, _, _ = self._response_done_events( + current_response_id, + current_conversation_id, + mint_ids_if_missing=True, + ) + return events + def validate_environment(self, headers: dict, model: str, api_key: str | None = None) -> dict: """Validate environment - no special validation needed for Bedrock.""" return headers @@ -830,9 +850,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): output_index=0, event_id=f"event_{uuid.uuid4()}", item_id=current_output_item_id, - part=( - {"type": "text", "text": ""} if next_delta_type == "text" else {"type": "audio", "transcript": ""} - ), + part=({"type": "text", "text": ""} if next_delta_type == "text" else {"type": "audio", "transcript": ""}), response_id=current_response_id, ) returned_messages.append(content_part_added) @@ -1051,19 +1069,27 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): Tuple of (events, reset_output_item_id, reset_response_id, reset_delta_type) """ verbose_logger.debug("Handling promptEnd") - return self._response_done_events(current_response_id, current_conversation_id) + return self._response_done_events( + current_response_id, + current_conversation_id, + mint_ids_if_missing=self.has_unbilled_usage(), + ) def _response_done_events( self, current_response_id: str | None, current_conversation_id: str | None, + *, + mint_ids_if_missing: bool = False, ) -> tuple[ list[OpenAIRealtimeEvents], str | None, str | None, ALL_DELTA_TYPES | None, ]: - if not current_response_id or not current_conversation_id: + response_id: Final = current_response_id or (f"resp_{uuid.uuid4()}" if mint_ids_if_missing else None) + conversation_id: Final = current_conversation_id or (f"conv_{uuid.uuid4()}" if mint_ids_if_missing else None) + if not response_id or not conversation_id: return [], None, None, None response_done: Final = OpenAIRealtimeDoneEvent( @@ -1071,10 +1097,10 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): event_id=f"event_{uuid.uuid4()}", response=OpenAIRealtimeResponseDoneObject( object="realtime.response", - id=current_response_id, + id=response_id, status="completed", output=[], - conversation_id=current_conversation_id, + conversation_id=conversation_id, usage=self.consume_usage_for_response_done(), ), ) @@ -1252,6 +1278,19 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): usage_event: Final = event["usageEvent"] if isinstance(usage_event, dict): self.record_usage_event(usage_event) + if current_response_id is None and self.has_unbilled_usage(): + ( + done_events, + current_output_item_id, + current_response_id, + current_delta_type, + ) = self._response_done_events( + None, + current_conversation_id, + mint_ids_if_missing=True, + ) + returned_messages.extend(done_events) + current_delta_chunks = None elif "contentStart" in event: ( @@ -1297,9 +1336,8 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): current_delta_chunks = None current_delta_type = None current_output_item_id = None - if content_end.get("type") == "TOOL": - current_response_id = None - if BedrockContentEnd.model_validate(content_end).stopReason == "END_TURN": + is_end_turn: Final = BedrockContentEnd.model_validate(content_end).stopReason == "END_TURN" + if is_end_turn: ( done_events, current_output_item_id, @@ -1308,6 +1346,8 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): ) = self._response_done_events(current_response_id, current_conversation_id) returned_messages.extend(done_events) current_delta_chunks = None + elif content_end.get("type") == "TOOL": + current_response_id = None elif "toolUse" in event: ( 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 2f02ef94e20..6ec3c72f4bf 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 @@ -15,7 +15,6 @@ from litellm.llms.bedrock.realtime.transformation import ( BedrockRealtimeConfig, ) from litellm.llms.bedrock.realtime.trigger_audio import ready_trigger_pcm -from litellm.types.llms.openai import OpenAIRealtimeEventTypes class TestBedrockRealtimeConfig: @@ -1356,6 +1355,125 @@ class TestBedrockRealtimeUsageAccounting: assert second_usage["input_token_details"]["audio_tokens"] == 5 assert second_usage["output_token_details"]["audio_tokens"] == 10 + def test_late_usage_event_after_response_id_cleared_emits_response_done(self): + config = BedrockRealtimeConfig() + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "trace_late_usage" + state = { + "session_configuration_request": json.dumps({"configured": True}), + "current_output_item_id": "item_1", + "current_response_id": "resp_1", + "current_conversation_id": "conv_1", + "current_delta_chunks": [], + "current_item_chunks": [], + "current_delta_type": "audio", + } + + end_turn = config.transform_realtime_response( + json.dumps({"event": {"contentEnd": {"stopReason": "END_TURN", "type": "AUDIO"}}}), + "amazon.nova-sonic-v1:0", + logging_obj, + realtime_response_transform_input=state, + ) + assert any(msg["type"] == "response.done" for msg in end_turn["response"]) + assert end_turn["current_response_id"] is None + state["current_response_id"] = end_turn["current_response_id"] + state["current_output_item_id"] = end_turn["current_output_item_id"] + state["current_conversation_id"] = end_turn["current_conversation_id"] + + late_usage = config.transform_realtime_response( + json.dumps(self._usage_event(input_speech=10, input_text=2, output_speech=20, output_text=3)), + "amazon.nova-sonic-v1:0", + logging_obj, + realtime_response_transform_input=state, + ) + done_events = [msg for msg in late_usage["response"] if msg["type"] == "response.done"] + assert len(done_events) == 1 + usage = done_events[0]["response"]["usage"] + assert usage["input_tokens"] == 12 + assert usage["output_tokens"] == 23 + assert usage["total_tokens"] == 35 + assert not config.has_unbilled_usage() + + def test_tool_end_turn_emits_response_done_before_clearing_ids(self): + config = BedrockRealtimeConfig() + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "trace_tool_end_turn" + state = { + "session_configuration_request": json.dumps({"configured": True}), + "current_output_item_id": "item_tool", + "current_response_id": "resp_tool", + "current_conversation_id": "conv_tool", + "current_delta_chunks": [], + "current_item_chunks": [], + "current_delta_type": None, + } + + config.transform_realtime_response( + json.dumps(self._usage_event(input_speech=4, input_text=1, output_speech=0, output_text=2)), + "amazon.nova-sonic-v1:0", + logging_obj, + realtime_response_transform_input=state, + ) + tool_end = config.transform_realtime_response( + json.dumps({"event": {"contentEnd": {"stopReason": "END_TURN", "type": "TOOL"}}}), + "amazon.nova-sonic-v1:0", + logging_obj, + realtime_response_transform_input=state, + ) + done_events = [msg for msg in tool_end["response"] if msg["type"] == "response.done"] + assert len(done_events) == 1 + assert done_events[0]["response"]["id"] == "resp_tool" + assert done_events[0]["response"]["usage"]["input_tokens"] == 5 + assert done_events[0]["response"]["usage"]["output_tokens"] == 2 + assert tool_end["current_response_id"] is None + assert not config.has_unbilled_usage() + + def test_flush_pending_usage_on_session_close(self): + config = BedrockRealtimeConfig() + config.record_usage_event( + self._usage_event(input_speech=8, input_text=1, output_speech=16, output_text=2)["event"]["usageEvent"] + ) + assert config.has_unbilled_usage() + + flushed = config.flush_pending_usage_as_response_done(None, None) + assert len(flushed) == 1 + assert flushed[0]["type"] == "response.done" + usage = flushed[0]["response"]["usage"] + assert usage["input_tokens"] == 9 + assert usage["output_tokens"] == 18 + assert usage["total_tokens"] == 27 + assert not config.has_unbilled_usage() + assert config.flush_pending_usage_as_response_done(None, None) == [] + + def test_completion_end_with_unbilled_usage_mints_response_done(self): + config = BedrockRealtimeConfig() + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "trace_completion_end" + config.record_usage_event( + self._usage_event(input_speech=3, input_text=0, output_speech=6, output_text=0)["event"]["usageEvent"] + ) + + result = config.transform_realtime_response( + json.dumps({"event": {"completionEnd": {}}}), + "amazon.nova-sonic-v1:0", + logging_obj, + realtime_response_transform_input={ + "session_configuration_request": json.dumps({"configured": True}), + "current_output_item_id": None, + "current_response_id": None, + "current_conversation_id": None, + "current_delta_chunks": None, + "current_item_chunks": None, + "current_delta_type": None, + }, + ) + done_events = [msg for msg in result["response"] if msg["type"] == "response.done"] + assert len(done_events) == 1 + assert done_events[0]["response"]["usage"]["input_tokens"] == 3 + assert done_events[0]["response"]["usage"]["output_tokens"] == 6 + assert not config.has_unbilled_usage() + if __name__ == "__main__": pytest.main([__file__, "-v"]) From 50b937505c85966e8d69d2bac153c1237f3d78ca Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Sat, 8 Aug 2026 14:30:40 -0700 Subject: [PATCH 14/19] fix(bedrock/realtime): clear type-discipline and basedpyright lint budgets Convert usage snapshots and event payloads to immutable Mapping/MappingProxyType, type kwargs and metadata, and drop Final from a per-iteration loop local that basedpyright rejected. Ratchet the strict, LIT, and basedpyright budgets down to match. --- basedpyright-code-budget.json | 4 +- litellm/llms/bedrock/realtime/handler.py | 31 +-- .../llms/bedrock/realtime/transformation.py | 183 +++++++++++------- ruff-strict-budget.json | 2 +- type-discipline-budget.json | 4 +- 5 files changed, 134 insertions(+), 90 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 0385f7a96e7..42c6c9e83eb 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -54,7 +54,7 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5835 + "limit": 5834 }, "reportMissingTypeArgument": { "limit": 15790 @@ -108,7 +108,7 @@ "limit": 39773 }, "reportUnknownParameterType": { - "limit": 20207 + "limit": 20206 }, "reportUnknownVariableType": { "limit": 31281 diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 3ab0306fff5..6ff26b7d3b0 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -9,6 +9,8 @@ store_message for backend events, store_input for client events, log_messages on import asyncio import contextlib import json +from collections.abc import Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, cast from pydantic import TypeAdapter @@ -25,6 +27,7 @@ if TYPE_CHECKING: from litellm.proxy._types import UserAPIKeyAuth _CLIENT_MODALITIES_ADAPTER: Final[TypeAdapter["list[str] | None"]] = TypeAdapter(list[str] | None) +_EMPTY_METADATA: Final[Mapping[str, object]] = MappingProxyType({}) class BedrockRealtime(BaseAWSLLM): @@ -53,8 +56,8 @@ class BedrockRealtime(BaseAWSLLM): aws_bedrock_runtime_endpoint: str | None = None, aws_external_id: str | None = None, user_api_key_dict: "UserAPIKeyAuth | None" = None, - litellm_metadata: dict | None = None, - **kwargs, + litellm_metadata: Mapping[str, object] | None = None, + **kwargs: object, ): """ Establish bidirectional streaming connection with Bedrock Nova Sonic. @@ -128,25 +131,31 @@ class BedrockRealtime(BaseAWSLLM): transformation_config: Final = BedrockRealtimeConfig() + pre_call_args: Final = MappingProxyType( + { + "api_base": endpoint_uri, + "complete_input_dict": MappingProxyType({"model": model}), + } + ) logging_obj.pre_call( input=None, api_key=api_key or "", - additional_args={ - "api_base": endpoint_uri, - "complete_input_dict": {"model": model}, - }, + additional_args=dict(pre_call_args), # mutable-ok: Logging.pre_call expects a mutable dict ) # RealTimeStreaming owns spend logging for other realtime providers. Bedrock cannot # use its WebSocket bidirectional_forward (AWS SDK stream instead), but store_message / # store_input / log_messages are the same path used by OpenAI and Azure. + request_data: Final = MappingProxyType( + {"litellm_metadata": litellm_metadata if litellm_metadata is not None else _EMPTY_METADATA} + ) realtime_streaming: Final = RealTimeStreaming( websocket=websocket, - backend_ws=cast(Any, object()), + backend_ws=cast(Any, object()), # cast-ok: Bedrock uses AWS SDK stream; backend_ws unused for store/log logging_obj=logging_obj, model=model, user_api_key_dict=user_api_key_dict, - request_data={"litellm_metadata": litellm_metadata or {}}, + request_data=dict(request_data), # mutable-ok: RealTimeStreaming stores request_data as dict ) try: @@ -227,10 +236,10 @@ class BedrockRealtime(BaseAWSLLM): if not isinstance(message, dict) or message.get("type") != "response.function_call_arguments.done": return realtime_streaming.tool_calls.append( - { + { # mutable-ok: spend logger tool_calls is a mutable list of JSON dicts "id": message.get("call_id", ""), "type": "function", - "function": { + "function": { # mutable-ok: nested tool function payload "name": message.get("name", ""), "arguments": message.get("arguments", "{}"), }, @@ -291,7 +300,7 @@ class BedrockRealtime(BaseAWSLLM): parsed_client_message.get("session", {}).get("modalities") ) if client_message_type == "session.update": - session_updated: Final = transformation_config.session_updated_event( + session_updated = transformation_config.session_updated_event( # rebind-ok: per-iteration local model, logging_obj, requested_modalities ) if realtime_streaming is not None: diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index ec1e541a82e..f3d55a30067 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -7,6 +7,8 @@ Transforms between OpenAI Realtime API format and Bedrock Nova Sonic format. import base64 import json import uuid as uuid_lib +from collections.abc import Mapping +from types import MappingProxyType from typing import Any, Final from pydantic import BaseModel @@ -48,16 +50,27 @@ TRIGGER_LEADING_SILENCE: Final = bytes(TRIGGER_AUDIO_BYTES_PER_SECOND // 2) TRIGGER_TRAILING_SILENCE: Final = bytes(TRIGGER_AUDIO_BYTES_PER_SECOND * 3) TRIGGER_AUDIO_CHUNK_SIZE: Final = 1024 +_EMPTY_USAGE_MAPPING: Final[Mapping[str, object]] = MappingProxyType({}) +_USAGE_SNAPSHOT_KEYS: Final = ( + "input_speech", + "input_text", + "output_speech", + "output_text", + "total_input", + "total_output", + "total", +) + def _parse_bedrock_tool_use_input(raw_input: object) -> object: if not raw_input: - return {} + return {} # mutable-ok: tool args are JSON-serializable wire values if not isinstance(raw_input, str): return raw_input try: return json.loads(raw_input) except json.JSONDecodeError: - return {} + return {} # mutable-ok: tool args are JSON-serializable wire values def _as_nonneg_int(value: object) -> int: @@ -66,23 +79,29 @@ def _as_nonneg_int(value: object) -> int: return max(0, int(value)) -def _empty_usage_snapshot() -> dict[str, int]: - return { - "input_speech": 0, - "input_text": 0, - "output_speech": 0, - "output_text": 0, - "total_input": 0, - "total_output": 0, - "total": 0, - } +def _mapping_or_empty(value: object) -> Mapping[str, object]: + return value if isinstance(value, Mapping) else _EMPTY_USAGE_MAPPING -def _usage_snapshot_from_event(usage_event: dict) -> dict[str, int]: - details: Final = usage_event.get("details") if isinstance(usage_event.get("details"), dict) else {} - total_block: Final = details.get("total") if isinstance(details.get("total"), dict) else {} - input_block: Final = total_block.get("input") if isinstance(total_block.get("input"), dict) else {} - output_block: Final = total_block.get("output") if isinstance(total_block.get("output"), dict) else {} +def _empty_usage_snapshot() -> Mapping[str, int]: + return MappingProxyType( + { + "input_speech": 0, + "input_text": 0, + "output_speech": 0, + "output_text": 0, + "total_input": 0, + "total_output": 0, + "total": 0, + } + ) + + +def _usage_snapshot_from_event(usage_event: Mapping[str, object]) -> Mapping[str, int]: + details: Final = _mapping_or_empty(usage_event.get("details")) + total_block: Final = _mapping_or_empty(details.get("total")) + input_block: Final = _mapping_or_empty(total_block.get("input")) + output_block: Final = _mapping_or_empty(total_block.get("output")) input_speech: Final = _as_nonneg_int(input_block.get("speechTokens")) input_text: Final = _as_nonneg_int(input_block.get("textTokens")) output_speech: Final = _as_nonneg_int(output_block.get("speechTokens")) @@ -90,38 +109,40 @@ def _usage_snapshot_from_event(usage_event: dict) -> dict[str, int]: total_input: Final = _as_nonneg_int(usage_event.get("totalInputTokens")) or (input_speech + input_text) total_output: Final = _as_nonneg_int(usage_event.get("totalOutputTokens")) or (output_speech + output_text) total: Final = _as_nonneg_int(usage_event.get("totalTokens")) or (total_input + total_output) - return { - "input_speech": input_speech, - "input_text": input_text, - "output_speech": output_speech, - "output_text": output_text, - "total_input": total_input, - "total_output": total_output, - "total": total, - } + return MappingProxyType( + { + "input_speech": input_speech, + "input_text": input_text, + "output_speech": output_speech, + "output_text": output_text, + "total_input": total_input, + "total_output": total_output, + "total": total, + } + ) -def _usage_snapshot_delta(current: dict[str, int], previous: dict[str, int]) -> dict[str, int]: - return {key: max(0, current.get(key, 0) - previous.get(key, 0)) for key in _empty_usage_snapshot()} +def _usage_snapshot_delta(current: Mapping[str, int], previous: Mapping[str, int]) -> Mapping[str, int]: + return MappingProxyType({key: max(0, current.get(key, 0) - previous.get(key, 0)) for key in _USAGE_SNAPSHOT_KEYS}) -def _openai_usage_from_snapshot(snapshot: dict[str, int]) -> dict[str, Any]: +def _openai_usage_from_snapshot(snapshot: Mapping[str, int]) -> dict[str, object]: # mutable-ok: OpenAI usage wire dict input_tokens: Final = snapshot.get("total_input", 0) or ( snapshot.get("input_speech", 0) + snapshot.get("input_text", 0) ) output_tokens: Final = snapshot.get("total_output", 0) or ( snapshot.get("output_speech", 0) + snapshot.get("output_text", 0) ) - return { + return { # mutable-ok: OpenAI response.done usage is a JSON-serializable dict "total_tokens": snapshot.get("total", 0) or (input_tokens + output_tokens), "input_tokens": input_tokens, "output_tokens": output_tokens, - "input_token_details": { + "input_token_details": { # mutable-ok: nested OpenAI usage wire shape "text_tokens": snapshot.get("input_text", 0), "audio_tokens": snapshot.get("input_speech", 0), "cached_tokens": 0, }, - "output_token_details": { + "output_token_details": { # mutable-ok: nested OpenAI usage wire shape "text_tokens": snapshot.get("output_text", 0), "audio_tokens": snapshot.get("output_speech", 0), }, @@ -167,31 +188,31 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): # Text configuration self.text_media_type = "text/plain" - def record_usage_event(self, usage_event: dict) -> None: + def record_usage_event(self, usage_event: Mapping[str, object]) -> None: self._usage_totals = _usage_snapshot_from_event(usage_event) def has_unbilled_usage(self) -> bool: delta: Final = _usage_snapshot_delta(self._usage_totals, self._usage_at_last_response_done) return any(value > 0 for value in delta.values()) - def consume_usage_for_response_done(self) -> dict[str, Any]: + def consume_usage_for_response_done(self) -> dict[str, object]: # mutable-ok: OpenAI usage wire dict delta: Final = _usage_snapshot_delta(self._usage_totals, self._usage_at_last_response_done) - self._usage_at_last_response_done = dict(self._usage_totals) + self._usage_at_last_response_done = self._usage_totals return _openai_usage_from_snapshot(delta) def flush_pending_usage_as_response_done( self, current_response_id: str | None = None, current_conversation_id: str | None = None, - ) -> list[OpenAIRealtimeEvents]: + ) -> list[OpenAIRealtimeEvents]: # mutable-ok: callers store into mutable message lists if not self.has_unbilled_usage(): - return [] + return [] # mutable-ok: empty OpenAI event list for callers that append/extend events, _, _, _ = self._response_done_events( current_response_id, current_conversation_id, mint_ids_if_missing=True, ) - return events + return list(events) # mutable-ok: session-close flush is stored into mutable message lists def validate_environment(self, headers: dict, model: str, api_key: str | None = None) -> dict: """Validate environment - no special validation needed for Bedrock.""" @@ -1257,13 +1278,25 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): "session_configuration_request": realtime_response_transform_input.get("session_configuration_request"), } - # Extract state - current_output_item_id = realtime_response_transform_input.get("current_output_item_id") - current_response_id = realtime_response_transform_input.get("current_response_id") - current_conversation_id = realtime_response_transform_input.get("current_conversation_id") - current_delta_chunks = realtime_response_transform_input.get("current_delta_chunks") - current_delta_type = realtime_response_transform_input.get("current_delta_type") - session_configuration_request = realtime_response_transform_input.get("session_configuration_request") + # Extract state. Session state is intentionally re-bound as each Bedrock event is folded in. + current_output_item_id = realtime_response_transform_input.get( + "current_output_item_id" + ) # rebind-ok: session state machine + current_response_id = realtime_response_transform_input.get( + "current_response_id" + ) # rebind-ok: session state machine + current_conversation_id = realtime_response_transform_input.get( + "current_conversation_id" + ) # rebind-ok: session state machine + current_delta_chunks = realtime_response_transform_input.get( + "current_delta_chunks" + ) # rebind-ok: session state machine + current_delta_type = realtime_response_transform_input.get( + "current_delta_type" + ) # rebind-ok: session state machine + session_configuration_request = realtime_response_transform_input.get( + "session_configuration_request" + ) # rebind-ok: session state machine returned_messages: Final[list[OpenAIRealtimeEvents]] = [] @@ -1272,7 +1305,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): # Route to appropriate transformation method if "sessionStart" in event: - session_configuration_request = json.dumps({"configured": True}) + session_configuration_request = json.dumps({"configured": True}) # rebind-ok: session state machine elif "usageEvent" in event: usage_event: Final = event["usageEvent"] @@ -1281,24 +1314,24 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): if current_response_id is None and self.has_unbilled_usage(): ( done_events, - current_output_item_id, - current_response_id, - current_delta_type, + current_output_item_id, # rebind-ok: session state machine + current_response_id, # rebind-ok: session state machine + current_delta_type, # rebind-ok: session state machine ) = self._response_done_events( None, current_conversation_id, mint_ids_if_missing=True, ) returned_messages.extend(done_events) - current_delta_chunks = None + current_delta_chunks = None # rebind-ok: session state machine elif "contentStart" in event: ( events, - current_response_id, - current_output_item_id, - current_conversation_id, - current_delta_type, + current_response_id, # rebind-ok: session state machine + current_output_item_id, # rebind-ok: session state machine + current_conversation_id, # rebind-ok: session state machine + current_delta_type, # rebind-ok: session state machine ) = self.transform_content_start_event( event, current_response_id, @@ -1308,10 +1341,10 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): ) returned_messages.extend(events) if events: - current_delta_chunks = None + current_delta_chunks = None # rebind-ok: session state machine elif "textOutput" in event: - events, current_delta_chunks = self.transform_text_output_event( + events, current_delta_chunks = self.transform_text_output_event( # rebind-ok: session state machine event, current_output_item_id, current_response_id, @@ -1320,12 +1353,14 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): returned_messages.extend(events) elif "audioOutput" in event: - events = self.transform_audio_output_event(event, current_output_item_id, current_response_id) + events = self.transform_audio_output_event( + event, current_output_item_id, current_response_id + ) # rebind-ok: session state machine returned_messages.extend(events) elif "contentEnd" in event: content_end: Final = event["contentEnd"] - events, current_delta_chunks = self.transform_content_end_event( + events, current_delta_chunks = self.transform_content_end_event( # rebind-ok: session state machine event, current_output_item_id, current_response_id, @@ -1333,21 +1368,21 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): current_delta_chunks, ) returned_messages.extend(events) - current_delta_chunks = None - current_delta_type = None - current_output_item_id = None + current_delta_chunks = None # rebind-ok: session state machine + current_delta_type = None # rebind-ok: session state machine + current_output_item_id = None # rebind-ok: session state machine is_end_turn: Final = BedrockContentEnd.model_validate(content_end).stopReason == "END_TURN" if is_end_turn: ( done_events, - current_output_item_id, - current_response_id, - current_delta_type, + current_output_item_id, # rebind-ok: session state machine + current_response_id, # rebind-ok: session state machine + current_delta_type, # rebind-ok: session state machine ) = self._response_done_events(current_response_id, current_conversation_id) returned_messages.extend(done_events) - current_delta_chunks = None + current_delta_chunks = None # rebind-ok: session state machine elif content_end.get("type") == "TOOL": - current_response_id = None + current_response_id = None # rebind-ok: session state machine elif "toolUse" in event: ( @@ -1358,21 +1393,21 @@ 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 - current_response_id = tool_response_id - current_delta_chunks = None - current_delta_type = None + current_output_item_id = tool_output_item_id # rebind-ok: session state machine + current_response_id = tool_response_id # rebind-ok: session state machine + current_delta_chunks = None # rebind-ok: session state machine + current_delta_type = None # rebind-ok: session state machine verbose_logger.debug("Tool use event: %s (ID: %s)", tool_name, tool_call_id) elif "promptEnd" in event or "completionEnd" in event: ( events, - current_output_item_id, - current_response_id, - current_delta_type, + current_output_item_id, # rebind-ok: session state machine + current_response_id, # rebind-ok: session state machine + current_delta_type, # rebind-ok: session state machine ) = self.transform_prompt_end_event(event, current_response_id, current_conversation_id) returned_messages.extend(events) - current_delta_chunks = None + current_delta_chunks = None # rebind-ok: session state machine return { "response": returned_messages, diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index fdc81fac196..9ba43b0dee3 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -6,7 +6,7 @@ "limit": 71 }, "ANN003": { - "limit": 834 + "limit": 833 }, "ANN201": { "limit": 2031 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 3a670bc7345..54d93485ca8 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -3,7 +3,7 @@ "limit": 23149 }, "LIT002": { - "limit": 27166 + "limit": 27163 }, "LIT003": { "limit": 269 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16760 + "limit": 16743 }, "LIT011": { "limit": 5598 From 60822cd137e936217b045b1b7d6a08236779573c Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Sat, 8 Aug 2026 14:39:08 -0700 Subject: [PATCH 15/19] chore(bedrock/realtime): keep lint budgets at base, drop ratchet --- basedpyright-code-budget.json | 4 ++-- ruff-strict-budget.json | 2 +- type-discipline-budget.json | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 42c6c9e83eb..0385f7a96e7 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -54,7 +54,7 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5834 + "limit": 5835 }, "reportMissingTypeArgument": { "limit": 15790 @@ -108,7 +108,7 @@ "limit": 39773 }, "reportUnknownParameterType": { - "limit": 20206 + "limit": 20207 }, "reportUnknownVariableType": { "limit": 31281 diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 9ba43b0dee3..fdc81fac196 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -6,7 +6,7 @@ "limit": 71 }, "ANN003": { - "limit": 833 + "limit": 834 }, "ANN201": { "limit": 2031 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index af336a28fb2..d621e85f09b 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -3,7 +3,7 @@ "limit": 23149 }, "LIT002": { - "limit": 27163 + "limit": 27166 }, "LIT003": { "limit": 269 From 06e5e7d0dad81ab0e0288af202862e4ae1d3dad4 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Mon, 24 Aug 2026 13:26:19 -0400 Subject: [PATCH 16/19] fix(bedrock/realtime): close the response a Nova Sonic tool turn opens A toolUse emitted only response.function_call_arguments.done, so the response the assistant block had already opened with response.created was never closed: the tool boundary dropped its id without a response.done and the post-tool turn minted a new one, leaving realtime clients tracking a response that stays in_progress forever. Expand a toolUse into the OpenAI function-call lifecycle the protocol expects, the same shape the Gemini realtime config emits: output_item.added, conversation.item.added, function_call_arguments.delta, .done, output_item.done, then response.done carrying the function_call item. Ids are cleared afterwards, so the tool call cannot leak into the next assistant turn and TOOL contentEnd no longer needs to drop a live response id. Because response.done now carries the call, RealTimeStreaming's shared _collect_tool_calls_from_response_done picks it up and the Bedrock-specific collector goes away. backend_ws becomes optional on RealTimeStreaming so Bedrock, which streams over the AWS SDK, no longer passes a cast bare object for it. --- .../litellm_core_utils/realtime_streaming.py | 17 +- litellm/llms/bedrock/realtime/handler.py | 23 +- .../llms/bedrock/realtime/transformation.py | 160 ++++++--- litellm/types/llms/openai.py | 11 + .../realtime/test_bedrock_realtime_handler.py | 87 ++++- .../test_bedrock_realtime_transformation.py | 304 ++++++++++++++---- 6 files changed, 464 insertions(+), 138 deletions(-) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 4b774269cc2..3dfbdfdc8a8 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -59,7 +59,7 @@ class RealTimeStreaming: def __init__( self, websocket: Any, - backend_ws: CLIENT_CONNECTION_CLASS, + backend_ws: CLIENT_CONNECTION_CLASS | None, logging_obj: LiteLLMLogging, provider_config: BaseRealtimeConfig | None = None, model: str = "", @@ -70,7 +70,7 @@ class RealTimeStreaming: event_normalizer: RealtimeEventNormalizer | None = None, ): self.websocket: _ClientWebSocket = websocket - self.backend_ws = backend_ws + self._backend_ws = backend_ws self.logging_obj = logging_obj self.messages: list[OpenAIRealtimeEvents] = [] self.input_message: dict = {} @@ -162,6 +162,19 @@ class RealTimeStreaming: "output_audio": "audio", } + @property + def backend_ws(self) -> CLIENT_CONNECTION_CLASS: + """ + The backend websocket, for the forwarding paths that require one. + + Providers that stream over a non-websocket transport (Bedrock uses the AWS SDK + bidirectional stream) construct this class only for its message store and spend + logging, and pass ``backend_ws=None``; reaching a forwarding path from there is a bug. + """ + if self._backend_ws is None: + raise RuntimeError("RealTimeStreaming was constructed without a backend websocket") + return self._backend_ws + def _should_store_message( self, message_obj: dict | OpenAIRealtimeEvents, diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 6ff26b7d3b0..8dcc1d19ffd 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -11,7 +11,7 @@ import contextlib import json from collections.abc import Mapping from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final from pydantic import TypeAdapter @@ -151,7 +151,7 @@ class BedrockRealtime(BaseAWSLLM): ) realtime_streaming: Final = RealTimeStreaming( websocket=websocket, - backend_ws=cast(Any, object()), # cast-ok: Bedrock uses AWS SDK stream; backend_ws unused for store/log + backend_ws=None, # Bedrock streams over the AWS SDK; only store/log are used here logging_obj=logging_obj, model=model, user_api_key_dict=user_api_key_dict, @@ -228,24 +228,6 @@ class BedrockRealtime(BaseAWSLLM): pass raise - @staticmethod - def _collect_tool_call_from_function_call_event( - realtime_streaming: RealTimeStreaming, - message: object, - ) -> None: - if not isinstance(message, dict) or message.get("type") != "response.function_call_arguments.done": - return - realtime_streaming.tool_calls.append( - { # mutable-ok: spend logger tool_calls is a mutable list of JSON dicts - "id": message.get("call_id", ""), - "type": "function", - "function": { # mutable-ok: nested tool function payload - "name": message.get("name", ""), - "arguments": message.get("arguments", "{}"), - }, - } - ) - async def _forward_client_to_bedrock( self, client_ws: Any, @@ -378,7 +360,6 @@ class BedrockRealtime(BaseAWSLLM): for openai_message in openai_messages: if realtime_streaming is not None: realtime_streaming.store_message(openai_message) - self._collect_tool_call_from_function_call_event(realtime_streaming, 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]) diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index f3d55a30067..d0936231d52 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -9,7 +9,7 @@ import json import uuid as uuid_lib from collections.abc import Mapping from types import MappingProxyType -from typing import Any, Final +from typing import Final from pydantic import BaseModel @@ -20,8 +20,11 @@ from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig from litellm.llms.bedrock.realtime.trigger_audio import ready_trigger_pcm from litellm.types.llms.openai import ( OpenAIRealtimeContentPartDone, + OpenAIRealtimeConversationItemAdded, OpenAIRealtimeDoneEvent, OpenAIRealtimeEvents, + OpenAIRealtimeFunctionCallArgumentsDelta, + OpenAIRealtimeFunctionCallArgumentsDone, OpenAIRealtimeOutputItemDone, OpenAIRealtimeResponseAudioDone, OpenAIRealtimeResponseContentPartAdded, @@ -29,6 +32,7 @@ from litellm.types.llms.openai import ( OpenAIRealtimeResponseDoneObject, OpenAIRealtimeResponseTextDone, OpenAIRealtimeStreamResponseBaseObject, + OpenAIRealtimeStreamResponseOutputItem, OpenAIRealtimeStreamResponseOutputItemAdded, OpenAIRealtimeStreamSession, OpenAIRealtimeStreamSessionEvents, @@ -44,6 +48,17 @@ class BedrockContentEnd(BaseModel): stopReason: str | None = None +class BedrockToolUse(BaseModel): + toolUseId: str = "" + toolName: str = "" + content: object | None = None + input: object | None = None + + def arguments(self) -> str: + """Nova Sonic puts tool args in ``content`` as a JSON string; older payloads use ``input``.""" + return json.dumps(_parse_bedrock_tool_use_input(self.content if self.content is not None else self.input)) + + TRIGGER_AUDIO_SAMPLE_RATE_HERTZ: Final = 16000 TRIGGER_AUDIO_BYTES_PER_SECOND: Final = TRIGGER_AUDIO_SAMPLE_RATE_HERTZ * 2 TRIGGER_LEADING_SILENCE: Final = bytes(TRIGGER_AUDIO_BYTES_PER_SECOND // 2) @@ -1133,50 +1148,103 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): event: dict, current_output_item_id: str | None, current_response_id: str | None, - ) -> tuple[list[OpenAIRealtimeEvents], str, str, str, str]: + conversation_id: str, + ) -> tuple[list[OpenAIRealtimeEvents], str, str]: """ - Transform Bedrock toolUse event to OpenAI format. + Transform a Bedrock toolUse event into the full OpenAI function-call lifecycle. - Args: - event: Bedrock toolUse event - current_output_item_id: Current output item ID - current_response_id: Current response ID + Nova Sonic delivers one tool call, fully formed, in a single event, and starts the + block with ``contentStart`` role ``TOOL``, which opens no OpenAI response. Mint the + response/item ids when they are missing, emit the item added/delta/done trio the + OpenAI realtime protocol requires around ``function_call_arguments.done``, then close + the response so no in-progress response is left orphaned and downstream spend logging + can harvest the call from ``response.done`` output. Returns: - Tuple of (events, tool_call_id, tool_name, output_item_id, response_id) - so the caller can persist any minted IDs into session state + Tuple of (events, tool_call_id, tool_name). The caller clears the response and + item ids, since this sequence closes the response it emits. """ verbose_logger.debug("Handling toolUse") - tool_use: Final = event["toolUse"] + tool_use: Final = BedrockToolUse.model_validate(event["toolUse"]) response_id: Final = current_response_id or f"resp_{uuid.uuid4()}" item_id: Final = current_output_item_id or f"item_{uuid.uuid4()}" - raw_input: Final = tool_use["content"] if "content" in tool_use else tool_use.get("input") - tool_input: Final = _parse_bedrock_tool_use_input(raw_input) + tool_call_id: Final = tool_use.toolUseId + tool_name: Final = tool_use.toolName + arguments: Final = tool_use.arguments() - tool_call_id: Final = tool_use.get("toolUseId", "") - tool_name: Final = tool_use.get("toolName", "") - - from typing import cast - - function_call_event: Final[dict[str, Any]] = { - "type": "response.function_call_arguments.done", - "event_id": f"event_{uuid.uuid4()}", - "response_id": response_id, - "item_id": item_id, - "output_index": 0, - "call_id": tool_call_id, - "name": tool_name, - "arguments": json.dumps(tool_input), - } - - return ( - [cast(OpenAIRealtimeEvents, function_call_event)], - tool_call_id, - tool_name, - item_id, - response_id, + function_call_item: Final = OpenAIRealtimeStreamResponseOutputItem( + id=item_id, + object="realtime.item", + type="function_call", + status="completed", + call_id=tool_call_id, + name=tool_name, + arguments=arguments, ) + pending_item: Final = OpenAIRealtimeStreamResponseOutputItem( + {**function_call_item, "status": "in_progress", "arguments": ""} + ) + + events: Final[list[OpenAIRealtimeEvents]] = [ + OpenAIRealtimeStreamResponseOutputItemAdded( + type="response.output_item.added", + event_id=f"event_{uuid.uuid4()}", + response_id=response_id, + output_index=0, + item=pending_item, + ), + # Pipecat registers call_id from conversation.item.added; without it the + # function_call_arguments.done below is dropped as an unknown call. + OpenAIRealtimeConversationItemAdded( + type="conversation.item.added", + event_id=f"event_{uuid.uuid4()}", + previous_item_id=None, + item=OpenAIRealtimeStreamResponseOutputItem({**pending_item}), + ), + # Nova Sonic delivers the whole argument payload at once; emit one delta anyway + # so clients that accumulate deltas rather than read `.done` still get the args. + OpenAIRealtimeFunctionCallArgumentsDelta( + type="response.function_call_arguments.delta", + event_id=f"event_{uuid.uuid4()}", + response_id=response_id, + item_id=item_id, + output_index=0, + call_id=tool_call_id, + delta=arguments, + ), + OpenAIRealtimeFunctionCallArgumentsDone( + type="response.function_call_arguments.done", + event_id=f"event_{uuid.uuid4()}", + response_id=response_id, + item_id=item_id, + output_index=0, + call_id=tool_call_id, + name=tool_name, + arguments=arguments, + ), + OpenAIRealtimeOutputItemDone( + type="response.output_item.done", + event_id=f"event_{uuid.uuid4()}", + response_id=response_id, + output_index=0, + item=OpenAIRealtimeStreamResponseOutputItem({**function_call_item}), + ), + OpenAIRealtimeDoneEvent( + type="response.done", + event_id=f"event_{uuid.uuid4()}", + response=OpenAIRealtimeResponseDoneObject( + object="realtime.response", + id=response_id, + status="completed", + output=[OpenAIRealtimeStreamResponseOutputItem({**function_call_item})], + conversation_id=conversation_id, + usage=self.consume_usage_for_response_done(), + ), + ), + ] + + return events, tool_call_id, tool_name def transform_conversation_item_create_tool_result_event(self, json_message: dict) -> list[str]: """ @@ -1381,20 +1449,22 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): ) = self._response_done_events(current_response_id, current_conversation_id) returned_messages.extend(done_events) current_delta_chunks = None # rebind-ok: session state machine - elif content_end.get("type") == "TOOL": - current_response_id = None # rebind-ok: session state machine elif "toolUse" in event: - ( - events, - tool_call_id, - tool_name, - tool_output_item_id, - tool_response_id, - ) = self.transform_tool_use_event(event, current_output_item_id, current_response_id) + current_conversation_id = ( # rebind-ok: session state machine + current_conversation_id or f"conv_{uuid.uuid4()}" + ) + events, tool_call_id, tool_name = self.transform_tool_use_event( + event, + current_output_item_id, + current_response_id, + current_conversation_id, + ) returned_messages.extend(events) - current_output_item_id = tool_output_item_id # rebind-ok: session state machine - current_response_id = tool_response_id # rebind-ok: session state machine + # transform_tool_use_event closes the response it emits, so the tool ids must not + # survive into the post-tool assistant turn. + current_output_item_id = None # rebind-ok: session state machine + current_response_id = None # rebind-ok: session state machine current_delta_chunks = None # rebind-ok: session state machine current_delta_type = None # rebind-ok: session state machine verbose_logger.debug("Tool use event: %s (ID: %s)", tool_name, tool_call_id) diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index da0592e6bb2..c5200f0ac72 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -2011,6 +2011,16 @@ class OpenAIRealtimeContentPartDone(TypedDict): type: Literal["response.content_part.done"] +class OpenAIRealtimeFunctionCallArgumentsDelta(TypedDict): + type: Literal["response.function_call_arguments.delta"] + event_id: str + response_id: str + item_id: str + output_index: int + call_id: str + delta: str + + class OpenAIRealtimeFunctionCallArgumentsDone(TypedDict): type: Literal["response.function_call_arguments.done"] event_id: str @@ -2087,6 +2097,7 @@ OpenAIRealtimeEvents = ( | OpenAIRealtimeResponseAudioDone | OpenAIRealtimeContentPartDone | OpenAIRealtimeOutputItemDone + | OpenAIRealtimeFunctionCallArgumentsDelta | OpenAIRealtimeFunctionCallArgumentsDone | OpenAIRealtimeDoneEvent ) 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 e1d5ef72c00..c91c5a35360 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 @@ -9,6 +9,7 @@ import pytest sys.path.insert(0, os.path.abspath("../../../../..")) # Adds the parent directory to the system path +from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.realtime.handler import BedrockRealtime from litellm.llms.bedrock.realtime.transformation import BedrockRealtimeConfig @@ -118,6 +119,37 @@ class EndedBedrockStream: return (None, EndedBedrockReceiver()) +class ScriptedBedrockChunk: + def __init__(self, payload): + self.bytes_ = json.dumps(payload).encode("utf-8") + + +class ScriptedBedrockResult: + def __init__(self, payload): + self.value = ScriptedBedrockChunk(payload) + + +class ScriptedBedrockReceiver: + def __init__(self, payloads): + self._payloads = list(payloads) + + async def receive(self): + if not self._payloads: + return None + return ScriptedBedrockResult(self._payloads.pop(0)) + + +class ScriptedBedrockStream: + """Replays a fixed list of Bedrock event payloads, then ends the stream.""" + + def __init__(self, payloads): + self.input_stream = FakeInputStream() + self._receiver = ScriptedBedrockReceiver(payloads) + + async def await_output(self): + return (None, self._receiver) + + class RealtimeClientWS: def __init__(self): self.closed = False @@ -364,6 +396,53 @@ class TestBedrockRealtimeSessionLifecycle: dispatched = logging_obj.dispatched_results[0] assert any(event.get("type") == "session.created" for event in dispatched) + @pytest.mark.asyncio + async def test_tool_call_reaches_spend_logging_via_response_done(self): + """ + Bedrock tool calls must be billable through the shared RealTimeStreaming collector, + which reads function_call items off response.done, with no Bedrock-specific plumbing. + """ + handler = BedrockRealtime() + client_ws = RealtimeClientWS() + logging_obj = FakeLogging() + realtime_streaming = RealTimeStreaming( + websocket=client_ws, + backend_ws=None, + logging_obj=logging_obj, + model="amazon.nova-2-sonic-v1:0", + ) + + await handler._forward_bedrock_to_client( + ScriptedBedrockStream( + [ + {"event": {"contentStart": {"role": "TOOL", "type": "TOOL"}}}, + { + "event": { + "toolUse": { + "toolUseId": "tool_call_1", + "toolName": "get_weather", + "content": json.dumps({"location": "Seattle"}), + } + } + }, + ] + ), + client_ws, + BedrockRealtimeConfig(), + "amazon.nova-2-sonic-v1:0", + logging_obj, + {}, + realtime_streaming, + ) + + assert realtime_streaming.tool_calls == [ + { + "id": "tool_call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": json.dumps({"location": "Seattle"})}, + } + ] + @pytest.mark.asyncio async def test_session_update_is_acked_with_session_updated(self, stub_aws_models): handler = BedrockRealtime() @@ -373,9 +452,7 @@ class TestBedrockRealtimeSessionLifecycle: [json.dumps({"type": "session.update", "session": {"instructions": "hi", "modalities": ["text"]}})] ) - await handler._forward_client_to_bedrock( - client_ws, stream, config, "amazon.nova-sonic-v1:0", {}, FakeLogging() - ) + await handler._forward_client_to_bedrock(client_ws, stream, config, "amazon.nova-sonic-v1:0", {}, FakeLogging()) acked = [json.loads(message) for message in client_ws.sent_to_client] updated = [event for event in acked if event["type"] == "session.updated"] @@ -387,9 +464,7 @@ class TestBedrockRealtimeSessionLifecycle: handler = BedrockRealtime() config = BedrockRealtimeConfig() stream = FakeBedrockStream() - client_ws = DisconnectingClientWS( - [json.dumps({"type": "session.update", "session": {"instructions": "hi"}})] - ) + client_ws = DisconnectingClientWS([json.dumps({"type": "session.update", "session": {"instructions": "hi"}})]) await handler._forward_client_to_bedrock(client_ws, stream, config, "amazon.nova-sonic-v1:0", {}) 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 6ec3c72f4bf..17bf5ae46cd 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 @@ -16,6 +16,30 @@ from litellm.llms.bedrock.realtime.transformation import ( ) from litellm.llms.bedrock.realtime.trigger_audio import ready_trigger_pcm +# The OpenAI realtime function-call lifecycle a single Nova Sonic toolUse must expand into. +_TOOL_CALL_EVENT_SEQUENCE = [ + "response.output_item.added", + "conversation.item.added", + "response.function_call_arguments.delta", + "response.function_call_arguments.done", + "response.output_item.done", + "response.done", +] + + +def _only(events, event_type): + matches = [event for event in events if event["type"] == event_type] + assert len(matches) == 1, f"expected exactly one {event_type}, got {len(matches)}" + return matches[0] + + +def _response_id_of(event): + if event["type"] == "response.done": + return event["response"]["id"] + if event["type"] == "conversation.item.added": + return None + return event["response_id"] + class TestBedrockRealtimeConfig: """Test suite for BedrockRealtimeConfig class""" @@ -559,16 +583,11 @@ class TestBedrockRealtimeResponseTransformation: }, ) - # Check for function call event - assert len(result["response"]) == 1 - function_call = result["response"][0] - assert function_call["type"] == "response.function_call_arguments.done" + assert [msg["type"] for msg in result["response"]] == _TOOL_CALL_EVENT_SEQUENCE + function_call = _only(result["response"], "response.function_call_arguments.done") assert function_call["call_id"] == "tool_call_123" assert function_call["name"] == "get_weather" - - # Verify arguments are properly formatted - args = json.loads(function_call["arguments"]) - assert args["location"] == "San Francisco" + assert json.loads(function_call["arguments"]) == {"location": "San Francisco"} def test_transform_tool_use_response_with_content_field(self): """Test toolUse response transformation with Nova 2 Sonic `content` field""" @@ -601,23 +620,34 @@ class TestBedrockRealtimeResponseTransformation: }, ) - # Check for function call event - assert len(result["response"]) == 1 - function_call = result["response"][0] - assert function_call["type"] == "response.function_call_arguments.done" + assert [msg["type"] for msg in result["response"]] == _TOOL_CALL_EVENT_SEQUENCE + function_call = _only(result["response"], "response.function_call_arguments.done") assert function_call["call_id"] == "tool_call_123" assert function_call["name"] == "get_weather" + assert json.loads(function_call["arguments"]) == {"location": "San Francisco"} - # Verify arguments are properly formatted - args = json.loads(function_call["arguments"]) - assert args["location"] == "San Francisco" + done = _only(result["response"], "response.done") + assert done["response"]["id"] == "resp_123" + assert done["response"]["output"] == [ + { + "id": "item_123", + "object": "realtime.item", + "type": "function_call", + "status": "completed", + "call_id": "tool_call_123", + "name": "get_weather", + "arguments": function_call["arguments"], + } + ] + assert result["current_response_id"] is None + assert result["current_output_item_id"] is None def test_transform_tool_use_event_directly(self): - """Test transform_tool_use_event directly for input parsing and missing IDs""" + """transform_tool_use_event emits the full OpenAI function-call lifecycle""" config = BedrockRealtimeConfig() - # Missing IDs still emit a function call (Nova Sonic starts tools with role=TOOL) - events, tool_call_id, tool_name, item_id, response_id = config.transform_tool_use_event( + # Missing IDs are minted (Nova Sonic starts tool turns with contentStart role=TOOL) + events, tool_call_id, tool_name = config.transform_tool_use_event( { "toolUse": { "toolUseId": "tool_call_no_ids", @@ -627,21 +657,41 @@ class TestBedrockRealtimeResponseTransformation: }, None, None, + "conv_1", ) - assert len(events) == 1 - assert events[0]["type"] == "response.function_call_arguments.done" - assert events[0]["call_id"] == "tool_call_no_ids" - assert events[0]["name"] == "get_weather" - assert response_id.startswith("resp_") - assert item_id.startswith("item_") - assert events[0]["response_id"] == response_id - assert events[0]["item_id"] == item_id - assert json.loads(events[0]["arguments"]) == {"location": "Seattle"} + assert [event["type"] for event in events] == _TOOL_CALL_EVENT_SEQUENCE + function_call = _only(events, "response.function_call_arguments.done") + assert function_call["call_id"] == "tool_call_no_ids" + assert function_call["name"] == "get_weather" + assert function_call["response_id"].startswith("resp_") + assert function_call["item_id"].startswith("item_") + assert json.loads(function_call["arguments"]) == {"location": "Seattle"} assert tool_call_id == "tool_call_no_ids" assert tool_name == "get_weather" - # JSON string content is parsed and converted to a function call event - events, tool_call_id, tool_name, item_id, response_id = config.transform_tool_use_event( + # Every event in the turn shares the minted response/item ids + assert {_response_id_of(event) for event in events} - {None} == {function_call["response_id"]} + assert _only(events, "response.output_item.added")["item"]["id"] == function_call["item_id"] + assert _only(events, "response.output_item.done")["item"]["id"] == function_call["item_id"] + + # The added item is in_progress with empty args; the done item carries the parsed args + assert _only(events, "response.output_item.added")["item"]["status"] == "in_progress" + assert _only(events, "response.output_item.added")["item"]["arguments"] == "" + assert _only(events, "conversation.item.added")["item"]["arguments"] == "" + assert _only(events, "response.function_call_arguments.delta")["delta"] == function_call["arguments"] + assert _only(events, "response.output_item.done")["item"]["status"] == "completed" + assert _only(events, "response.output_item.done")["item"]["arguments"] == function_call["arguments"] + + # response.done closes the turn and carries the call so spend logging can harvest it + done = _only(events, "response.done") + assert done["response"]["id"] == function_call["response_id"] + assert done["response"]["conversation_id"] == "conv_1" + assert done["response"]["status"] == "completed" + assert done["response"]["output"][0]["call_id"] == "tool_call_no_ids" + assert done["response"]["output"][0]["type"] == "function_call" + + # Explicit ids are reused rather than minted + events, _, _ = config.transform_tool_use_event( { "toolUse": { "toolUseId": "tool_call_123", @@ -651,19 +701,30 @@ class TestBedrockRealtimeResponseTransformation: }, "item_123", "resp_123", + "conv_1", ) - assert len(events) == 1 - assert events[0]["type"] == "response.function_call_arguments.done" - assert events[0]["call_id"] == "tool_call_123" - assert events[0]["name"] == "get_weather" - assert response_id == "resp_123" - assert item_id == "item_123" - assert events[0]["response_id"] == "resp_123" - assert events[0]["item_id"] == "item_123" - assert json.loads(events[0]["arguments"]) == {"location": "San Francisco"} + function_call = _only(events, "response.function_call_arguments.done") + assert function_call["response_id"] == "resp_123" + assert function_call["item_id"] == "item_123" + assert json.loads(function_call["arguments"]) == {"location": "San Francisco"} + + # Legacy `input` field is still honoured when `content` is absent + events, _, _ = config.transform_tool_use_event( + { + "toolUse": { + "toolUseId": "tool_call_legacy", + "toolName": "get_weather", + "input": json.dumps({"location": "Boston"}), + } + }, + "item_123", + "resp_123", + "conv_1", + ) + assert json.loads(_only(events, "response.function_call_arguments.done")["arguments"]) == {"location": "Boston"} # Invalid JSON content falls back to empty arguments - events, _, _, _, _ = config.transform_tool_use_event( + events, _, _ = config.transform_tool_use_event( { "toolUse": { "toolUseId": "tool_call_124", @@ -673,9 +734,9 @@ class TestBedrockRealtimeResponseTransformation: }, "item_123", "resp_123", + "conv_1", ) - assert len(events) == 1 - assert json.loads(events[0]["arguments"]) == {} + assert json.loads(_only(events, "response.function_call_arguments.done")["arguments"]) == {} def test_transform_realtime_response_persists_minted_tool_ids(self): """TOOL-first turns must write minted response/item ids into session state""" @@ -729,17 +790,19 @@ class TestBedrockRealtimeResponseTransformation: realtime_response_transform_input=state, ) - assert len(result["response"]) == 1 - function_call = result["response"][0] - assert function_call["type"] == "response.function_call_arguments.done" - assert result["current_response_id"] is not None - assert result["current_output_item_id"] is not None - assert result["current_response_id"].startswith("resp_") - assert result["current_output_item_id"].startswith("item_") - assert function_call["response_id"] == result["current_response_id"] - assert function_call["item_id"] == result["current_output_item_id"] + assert [msg["type"] for msg in result["response"]] == _TOOL_CALL_EVENT_SEQUENCE + function_call = _only(result["response"], "response.function_call_arguments.done") + assert function_call["response_id"].startswith("resp_") + assert function_call["item_id"].startswith("item_") assert json.loads(function_call["arguments"]) == {"location": "Seattle"} + # The tool turn closes the response it minted, so no in-progress response is orphaned + # and the ids cannot leak into the post-tool assistant turn. + tool_done = _only(result["response"], "response.done") + assert tool_done["response"]["id"] == function_call["response_id"] + assert result["current_response_id"] is None + assert result["current_output_item_id"] is None + content_end_message = { "event": { "contentEnd": { @@ -765,8 +828,9 @@ class TestBedrockRealtimeResponseTransformation: assert follow_up["current_response_id"] is None assert follow_up["current_output_item_id"] is None assert follow_up["current_delta_type"] is None + # The tool turn already emitted response.done; TOOL contentEnd must not emit a second + # one, nor an unpaired message-shaped output_item.done. assert follow_up["response"] == [] - assert all(msg["type"] != "response.output_item.done" for msg in follow_up["response"]) post_tool_state = { "session_configuration_request": follow_up["session_configuration_request"], @@ -783,12 +847,10 @@ class TestBedrockRealtimeResponseTransformation: logging_obj, realtime_response_transform_input=post_tool_state, ) - tool_response_id = result["current_response_id"] - tool_item_id = result["current_output_item_id"] assert assistant_start["current_response_id"] is not None assert assistant_start["current_output_item_id"] is not None - assert assistant_start["current_response_id"] != tool_response_id - assert assistant_start["current_output_item_id"] != tool_item_id + assert assistant_start["current_response_id"] != function_call["response_id"] + assert assistant_start["current_output_item_id"] != function_call["item_id"] created = [msg for msg in assistant_start["response"] if msg["type"] == "response.created"][0] added = [msg for msg in assistant_start["response"] if msg["type"] == "response.output_item.added"][0] assert created["response"]["id"] == assistant_start["current_response_id"] @@ -797,7 +859,10 @@ class TestBedrockRealtimeResponseTransformation: assert added["item"]["id"] != function_call["item_id"] def test_tool_content_end_does_not_emit_message_output_item_done(self): - """Minted tool ids must not unlock unpaired message output_item.done on TOOL contentEnd""" + """ + TOOL contentEnd must stay silent: the tool turn already closed its own response, and a + message-shaped output_item.done here would have no matching output_item.added. + """ config = BedrockRealtimeConfig() logging_obj = MagicMock() logging_obj.litellm_trace_id = "trace_123" @@ -816,8 +881,8 @@ class TestBedrockRealtimeResponseTransformation: logging_obj, realtime_response_transform_input={ "session_configuration_request": json.dumps({"configured": True}), - "current_output_item_id": "item_minted_for_tool", - "current_response_id": "resp_minted_for_tool", + "current_output_item_id": "item_open_assistant_turn", + "current_response_id": "resp_open_assistant_turn", "current_conversation_id": "conv_123", "current_delta_chunks": [], "current_item_chunks": [], @@ -826,9 +891,11 @@ class TestBedrockRealtimeResponseTransformation: ) assert result["response"] == [] - assert result["current_response_id"] is None assert result["current_output_item_id"] is None assert result["current_delta_type"] is None + # A TOOL block that produced no toolUse leaves the assistant response open rather than + # dropping its id, so the next assistant block reuses it instead of orphaning it. + assert result["current_response_id"] == "resp_open_assistant_turn" def test_transform_content_end_text(self): """Test contentEnd for text response""" @@ -1173,20 +1240,26 @@ class TestBedrockRealtimeContentBlockLifecycle: } }, ) - assert tool_result["response"][0]["type"] == "response.function_call_arguments.done" + assert [msg["type"] for msg in tool_result["response"]] == _TOOL_CALL_EVENT_SEQUENCE + # The response the assistant text block opened is closed by the tool turn instead of + # being left in_progress forever once the ids are cleared. + tool_done = _only(tool_result["response"], "response.done") + assert tool_done["response"]["id"] == first_response_id + assert tool_done["response"]["output"][0]["call_id"] == "tool_1" + 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 - self._apply( + tool_content_end = self._apply( config, logging_obj, state, {"event": {"contentEnd": {"stopReason": "TOOL_USE", "type": "TOOL"}}}, ) + assert tool_content_end["response"] == [] 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, @@ -1217,6 +1290,43 @@ class TestBedrockRealtimeContentBlockLifecycle: assert state["current_response_id"] is None assert state["current_delta_chunks"] is None + def test_every_created_response_is_closed_across_a_tool_turn(self): + """ + Realtime clients track in-progress responses by id. A tool turn that drops the + response id without a matching response.done leaves one open forever. + """ + config = BedrockRealtimeConfig() + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "trace_123" + state = self._state() + + turn = [ + {"event": {"contentStart": {"role": "ASSISTANT", "type": "TEXT"}}}, + {"event": {"textOutput": {"content": "Let me check."}}}, + {"event": {"contentEnd": {"stopReason": "PARTIAL_TURN", "type": "TEXT"}}}, + {"event": {"contentStart": {"role": "TOOL", "type": "TOOL"}}}, + { + "event": { + "toolUse": { + "toolUseId": "tool_1", + "toolName": "get_weather", + "content": json.dumps({"location": "Seattle"}), + } + } + }, + {"event": {"contentEnd": {"stopReason": "TOOL_USE", "type": "TOOL"}}}, + {"event": {"contentStart": {"role": "ASSISTANT", "type": "TEXT"}}}, + {"event": {"textOutput": {"content": "It is sunny."}}}, + {"event": {"contentEnd": {"stopReason": "END_TURN", "type": "TEXT"}}}, + ] + emitted = [msg for event in turn for msg in self._apply(config, logging_obj, state, event)["response"]] + + created = [msg["response"]["id"] for msg in emitted if msg["type"] == "response.created"] + done = [msg["response"]["id"] for msg in emitted if msg["type"] == "response.done"] + assert len(created) == 2 + assert created == done + assert state["current_response_id"] is None + def test_second_assistant_content_block_reuses_response_not_item(self): config = BedrockRealtimeConfig() logging_obj = MagicMock() @@ -1395,7 +1505,7 @@ class TestBedrockRealtimeUsageAccounting: assert usage["total_tokens"] == 35 assert not config.has_unbilled_usage() - def test_tool_end_turn_emits_response_done_before_clearing_ids(self): + def test_tool_content_end_with_end_turn_stop_reason_emits_response_done(self): config = BedrockRealtimeConfig() logging_obj = MagicMock() logging_obj.litellm_trace_id = "trace_tool_end_turn" @@ -1429,6 +1539,72 @@ class TestBedrockRealtimeUsageAccounting: assert tool_end["current_response_id"] is None assert not config.has_unbilled_usage() + def test_tool_turn_does_not_double_bill_across_late_usage_and_completion_end(self): + """The tool response.done, a later usageEvent, and completionEnd must each bill once.""" + config = BedrockRealtimeConfig() + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "trace_tool_usage" + state = { + "session_configuration_request": json.dumps({"configured": True}), + "current_output_item_id": "item_1", + "current_response_id": "resp_1", + "current_conversation_id": "conv_1", + "current_delta_chunks": None, + "current_item_chunks": [], + "current_delta_type": None, + } + + config.transform_realtime_response( + json.dumps(self._usage_event(input_speech=4, input_text=0, output_speech=6, output_text=0)), + "amazon.nova-2-sonic-v1:0", + logging_obj, + realtime_response_transform_input=state, + ) + tool_result = config.transform_realtime_response( + json.dumps( + { + "event": { + "toolUse": { + "toolUseId": "tool_1", + "toolName": "get_weather", + "content": json.dumps({"location": "Seattle"}), + } + } + } + ), + "amazon.nova-2-sonic-v1:0", + logging_obj, + realtime_response_transform_input=state, + ) + tool_usage = _only(tool_result["response"], "response.done")["response"]["usage"] + assert tool_usage["input_tokens"] == 4 + assert tool_usage["output_tokens"] == 6 + assert not config.has_unbilled_usage() + state["current_response_id"] = tool_result["current_response_id"] + state["current_output_item_id"] = tool_result["current_output_item_id"] + + # Cumulative usage grows after the tool turn; only the delta may be billed again. + late = config.transform_realtime_response( + json.dumps(self._usage_event(input_speech=10, input_text=0, output_speech=15, output_text=0)), + "amazon.nova-2-sonic-v1:0", + logging_obj, + realtime_response_transform_input=state, + ) + late_usage = _only(late["response"], "response.done")["response"]["usage"] + assert late_usage["input_tokens"] == 6 + assert late_usage["output_tokens"] == 9 + state["current_response_id"] = late["current_response_id"] + + completion_end = config.transform_realtime_response( + json.dumps({"event": {"completionEnd": {}}}), + "amazon.nova-2-sonic-v1:0", + logging_obj, + realtime_response_transform_input=state, + ) + assert completion_end["response"] == [] + assert not config.has_unbilled_usage() + assert config.flush_pending_usage_as_response_done(None, None) == [] + def test_flush_pending_usage_on_session_close(self): config = BedrockRealtimeConfig() config.record_usage_event( From 6518decca7c6bf236050b75802a384c58b75a3e1 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Mon, 24 Aug 2026 13:34:38 -0400 Subject: [PATCH 17/19] fix(bedrock/realtime): open the response a tool-first Nova Sonic turn closes Nova Sonic starts a tool turn with contentStart role TOOL, which emits nothing, so when the tool call is the first thing in a turn there is no open response. The tool lifecycle still emitted response.done for the id it minted, handing the client a close for a response it never saw created: the mirror of the orphan just fixed. Emit response.created ahead of the tool events whenever the response id is minted, matching what the Gemini realtime config does on its toolCall branch. --- .../llms/bedrock/realtime/transformation.py | 31 +++++++++-- .../test_bedrock_realtime_transformation.py | 55 ++++++++++++++++--- 2 files changed, 75 insertions(+), 11 deletions(-) diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index d0936231d52..df3d992ccd1 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -1155,10 +1155,11 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): Nova Sonic delivers one tool call, fully formed, in a single event, and starts the block with ``contentStart`` role ``TOOL``, which opens no OpenAI response. Mint the - response/item ids when they are missing, emit the item added/delta/done trio the - OpenAI realtime protocol requires around ``function_call_arguments.done``, then close - the response so no in-progress response is left orphaned and downstream spend logging - can harvest the call from ``response.done`` output. + response/item ids when they are missing, opening the response first so its + ``response.done`` is never unmatched, emit the item added/delta/done trio the OpenAI + realtime protocol requires around ``function_call_arguments.done``, then close the + response so nothing is left in progress and downstream spend logging can harvest the + call from ``response.done`` output. Returns: Tuple of (events, tool_call_id, tool_name). The caller clears the response and @@ -1167,6 +1168,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): verbose_logger.debug("Handling toolUse") tool_use: Final = BedrockToolUse.model_validate(event["toolUse"]) + is_new_response: Final = not current_response_id response_id: Final = current_response_id or f"resp_{uuid.uuid4()}" item_id: Final = current_output_item_id or f"item_{uuid.uuid4()}" tool_call_id: Final = tool_use.toolUseId @@ -1186,7 +1188,28 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): {**function_call_item, "status": "in_progress", "arguments": ""} ) + # A tool turn that Bedrock opens with contentStart role TOOL has no response yet, so the + # response.done below would close an id the client never saw opened. + response_created: Final[tuple[OpenAIRealtimeEvents, ...]] = ( + ( + OpenAIRealtimeStreamResponseBaseObject( + type="response.created", + event_id=f"event_{uuid.uuid4()}", + response={ + "object": "realtime.response", + "id": response_id, + "status": "in_progress", + "output": [], + "conversation_id": conversation_id, + }, + ), + ) + if is_new_response + else () + ) + events: Final[list[OpenAIRealtimeEvents]] = [ + *response_created, OpenAIRealtimeStreamResponseOutputItemAdded( type="response.output_item.added", event_id=f"event_{uuid.uuid4()}", 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 17bf5ae46cd..7d8187a4540 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 @@ -16,7 +16,8 @@ from litellm.llms.bedrock.realtime.transformation import ( ) from litellm.llms.bedrock.realtime.trigger_audio import ready_trigger_pcm -# The OpenAI realtime function-call lifecycle a single Nova Sonic toolUse must expand into. +# The OpenAI realtime function-call lifecycle a single Nova Sonic toolUse must expand into, +# when an assistant block already opened the response the tool call belongs to. _TOOL_CALL_EVENT_SEQUENCE = [ "response.output_item.added", "conversation.item.added", @@ -26,6 +27,10 @@ _TOOL_CALL_EVENT_SEQUENCE = [ "response.done", ] +# Nova Sonic opens tool turns with contentStart role TOOL, which opens no response, so the +# tool call has to open one itself before it can close it. +_TOOL_CALL_EVENT_SEQUENCE_NEW_RESPONSE = ["response.created", *_TOOL_CALL_EVENT_SEQUENCE] + def _only(events, event_type): matches = [event for event in events if event["type"] == event_type] @@ -34,11 +39,10 @@ def _only(events, event_type): def _response_id_of(event): - if event["type"] == "response.done": + """The response id an event is bound to, or None for events that carry no response id.""" + if event["type"] in ("response.created", "response.done"): return event["response"]["id"] - if event["type"] == "conversation.item.added": - return None - return event["response_id"] + return event.get("response_id") class TestBedrockRealtimeConfig: @@ -659,8 +663,10 @@ class TestBedrockRealtimeResponseTransformation: None, "conv_1", ) - assert [event["type"] for event in events] == _TOOL_CALL_EVENT_SEQUENCE + assert [event["type"] for event in events] == _TOOL_CALL_EVENT_SEQUENCE_NEW_RESPONSE function_call = _only(events, "response.function_call_arguments.done") + assert _only(events, "response.created")["response"]["id"] == function_call["response_id"] + assert _only(events, "response.created")["response"]["conversation_id"] == "conv_1" assert function_call["call_id"] == "tool_call_no_ids" assert function_call["name"] == "get_weather" assert function_call["response_id"].startswith("resp_") @@ -790,8 +796,9 @@ class TestBedrockRealtimeResponseTransformation: realtime_response_transform_input=state, ) - assert [msg["type"] for msg in result["response"]] == _TOOL_CALL_EVENT_SEQUENCE + assert [msg["type"] for msg in result["response"]] == _TOOL_CALL_EVENT_SEQUENCE_NEW_RESPONSE function_call = _only(result["response"], "response.function_call_arguments.done") + assert _only(result["response"], "response.created")["response"]["id"] == function_call["response_id"] assert function_call["response_id"].startswith("resp_") assert function_call["item_id"].startswith("item_") assert json.loads(function_call["arguments"]) == {"location": "Seattle"} @@ -1327,6 +1334,40 @@ class TestBedrockRealtimeContentBlockLifecycle: assert created == done assert state["current_response_id"] is None + def test_every_response_is_opened_and_closed_on_a_tool_first_turn(self): + """ + Nova Sonic opens tool turns with contentStart role TOOL, which emits nothing, so the + tool call is the first thing in the session. It has to open the response it closes, + or the client sees a response.done for an id it never saw created. + """ + config = BedrockRealtimeConfig() + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "trace_123" + state = self._state(current_conversation_id=None) + + turn = [ + {"event": {"contentStart": {"role": "TOOL", "type": "TOOL"}}}, + { + "event": { + "toolUse": { + "toolUseId": "tool_1", + "toolName": "get_weather", + "content": json.dumps({"location": "Seattle"}), + } + } + }, + {"event": {"contentEnd": {"stopReason": "TOOL_USE", "type": "TOOL"}}}, + ] + emitted = [msg for event in turn for msg in self._apply(config, logging_obj, state, event)["response"]] + + created = [msg["response"]["id"] for msg in emitted if msg["type"] == "response.created"] + done = [msg["response"]["id"] for msg in emitted if msg["type"] == "response.done"] + assert len(created) == 1 + assert created == done + # Every event in the turn is bound to that one response. + assert {_response_id_of(msg) for msg in emitted} - {None} == set(created) + assert state["current_response_id"] is None + def test_second_assistant_content_block_reuses_response_not_item(self): config = BedrockRealtimeConfig() logging_obj = MagicMock() From 67aa7068aacd5503d993e93239e153423cdcecc1 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Mon, 24 Aug 2026 13:52:05 -0400 Subject: [PATCH 18/19] chore(bedrock/realtime): keep the new realtime event type within lint budgets --- litellm/llms/bedrock/realtime/transformation.py | 6 +++--- litellm/types/llms/openai.py | 14 +++++++------- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index df3d992ccd1..642a076e516 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -1223,7 +1223,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): type="conversation.item.added", event_id=f"event_{uuid.uuid4()}", previous_item_id=None, - item=OpenAIRealtimeStreamResponseOutputItem({**pending_item}), + item=pending_item, ), # Nova Sonic delivers the whole argument payload at once; emit one delta anyway # so clients that accumulate deltas rather than read `.done` still get the args. @@ -1251,7 +1251,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): event_id=f"event_{uuid.uuid4()}", response_id=response_id, output_index=0, - item=OpenAIRealtimeStreamResponseOutputItem({**function_call_item}), + item=function_call_item, ), OpenAIRealtimeDoneEvent( type="response.done", @@ -1260,7 +1260,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): object="realtime.response", id=response_id, status="completed", - output=[OpenAIRealtimeStreamResponseOutputItem({**function_call_item})], + output=[function_call_item], conversation_id=conversation_id, usage=self.consume_usage_for_response_done(), ), diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index d8f6dd5faf3..4e88fda47a0 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -2098,13 +2098,13 @@ class OpenAIRealtimeContentPartDone(TypedDict): class OpenAIRealtimeFunctionCallArgumentsDelta(TypedDict): - type: Literal["response.function_call_arguments.delta"] - event_id: str - response_id: str - item_id: str - output_index: int - call_id: str - delta: str + type: ReadOnly[Literal["response.function_call_arguments.delta"]] + event_id: ReadOnly[str] + response_id: ReadOnly[str] + item_id: ReadOnly[str] + output_index: ReadOnly[int] + call_id: ReadOnly[str] + delta: ReadOnly[str] class OpenAIRealtimeFunctionCallArgumentsDone(TypedDict): From 225b04c809278c89ca3aea5b8b40be2357e97035 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Mon, 24 Aug 2026 14:00:45 -0400 Subject: [PATCH 19/19] test(bedrock/realtime): cover the new spend-logging and tool-arg paths --- .../test_realtime_streaming.py | 52 ++++++----- .../realtime/test_bedrock_realtime_handler.py | 88 ++++++++++++++++++- .../test_bedrock_realtime_transformation.py | 37 ++++++++ 3 files changed, 153 insertions(+), 24 deletions(-) diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py index 61b63e2b917..dbd0f08f723 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -30,6 +30,21 @@ def _make_transcript_event(text: str, item_id: str = "item_x") -> bytes: ).encode() +def test_store_and_log_work_without_a_backend_websocket(): + """ + Providers that stream over a non-websocket transport (Bedrock's AWS SDK stream) build this + class only to store and log; store/log must work with backend_ws=None, and any forwarding + path reached from there must fail loudly rather than on a placeholder object. + """ + streaming = RealTimeStreaming(MagicMock(), None, MagicMock()) + + streaming.store_message(json.dumps({"type": "session.created", "session": {"id": "s"}})) + assert [message["type"] for message in streaming.messages] == ["session.created"] + + with pytest.raises(RuntimeError, match="without a backend websocket"): + _ = streaming.backend_ws + + def test_realtime_streaming_store_message(): # Setup websocket = MagicMock() @@ -1403,7 +1418,6 @@ async def test_realtime_guardrail_blocks_prompt_injection(monkeypatch: pytest.Mo ) - @pytest.mark.asyncio async def test_realtime_guardrail_allows_clean_transcript(monkeypatch: pytest.MonkeyPatch): """ @@ -1460,7 +1474,6 @@ async def test_realtime_guardrail_allows_clean_transcript(monkeypatch: pytest.Mo assert len(response_creates) == 1, f"Clean transcript should trigger response.create, got: {sent_to_backend}" - @pytest.mark.asyncio async def test_realtime_text_input_guardrail_blocks_and_returns_error(monkeypatch: pytest.MonkeyPatch): """ @@ -1554,7 +1567,6 @@ async def test_realtime_text_input_guardrail_blocks_and_returns_error(monkeypatc assert len(original_items) == 0, f"Blocked item should not be forwarded to backend, got: {original_items}" - @pytest.mark.asyncio async def test_realtime_function_call_output_guardrail_blocks_and_returns_error(monkeypatch: pytest.MonkeyPatch): """ @@ -1643,7 +1655,6 @@ async def test_realtime_function_call_output_guardrail_blocks_and_returns_error( assert "test@example.com" not in sanitized_item["output"] - @pytest.mark.asyncio async def test_realtime_function_call_output_guardrail_allows_clean_output(monkeypatch: pytest.MonkeyPatch): """ @@ -1708,7 +1719,6 @@ async def test_realtime_function_call_output_guardrail_allows_clean_output(monke assert len(forwarded) == 1, f"Clean function_call_output should be forwarded, got: {forwarded}" - @pytest.mark.asyncio async def test_realtime_text_input_guardrail_uses_pre_call_mode(monkeypatch: pytest.MonkeyPatch): """ @@ -1744,7 +1754,6 @@ async def test_realtime_text_input_guardrail_uses_pre_call_mode(monkeypatch: pyt ) - @pytest.mark.asyncio async def test_realtime_session_created_injects_session_update_for_audio_guardrail(monkeypatch: pytest.MonkeyPatch): """ @@ -1801,7 +1810,6 @@ async def test_realtime_session_created_injects_session_update_for_audio_guardra ) - @pytest.mark.asyncio async def test_realtime_session_created_does_not_inject_session_update_for_pre_call_only( monkeypatch: pytest.MonkeyPatch, @@ -1846,7 +1854,6 @@ async def test_realtime_session_created_does_not_inject_session_update_for_pre_c assert len(session_updates) == 0, f"pre_call-only guardrail must not inject session.update, got: {sent_to_backend}" - @pytest.mark.asyncio async def test_pre_call_and_post_call_guardrails_do_not_disable_server_vad(monkeypatch: pytest.MonkeyPatch): """Model Armor-style pre_call + post_call must not gate audio VAD.""" @@ -1862,17 +1869,17 @@ async def test_pre_call_and_post_call_guardrails_do_not_disable_server_vad(monke litellm, "callbacks", [ - ModelArmorStyleGuardrail( - guardrail_name="model_armor_all_pre_call", - event_hook=GuardrailEventHooks.pre_call, - default_on=False, - ), - ModelArmorStyleGuardrail( - guardrail_name="model_armor_all_post_call", - event_hook=GuardrailEventHooks.post_call, - default_on=False, - ), - ], + ModelArmorStyleGuardrail( + guardrail_name="model_armor_all_pre_call", + event_hook=GuardrailEventHooks.pre_call, + default_on=False, + ), + ModelArmorStyleGuardrail( + guardrail_name="model_armor_all_post_call", + event_hook=GuardrailEventHooks.post_call, + default_on=False, + ), + ], ) client_ws = MagicMock() @@ -1896,7 +1903,6 @@ async def test_pre_call_and_post_call_guardrails_do_not_disable_server_vad(monke assert streaming._has_audio_transcription_guardrails() is False - @pytest.mark.asyncio async def test_end_session_after_n_fails_closes_connection(monkeypatch: pytest.MonkeyPatch): """ @@ -1943,7 +1949,6 @@ async def test_end_session_after_n_fails_closes_connection(monkeypatch: pytest.M assert streaming._violation_count == 2 - @pytest.mark.asyncio async def test_on_violation_end_session_closes_on_first_fail(monkeypatch: pytest.MonkeyPatch): """ @@ -1989,7 +1994,6 @@ async def test_on_violation_end_session_closes_on_first_fail(monkeypatch: pytest assert streaming._violation_count == 1 - @pytest.mark.asyncio async def test_provider_path_suppresses_duplicate_session_created_after_synthetic(): client_ws = MagicMock() @@ -2952,7 +2956,9 @@ async def test_log_messages_routes_async_logging_through_bounded_worker(): mock_worker.ensure_initialized_and_enqueue.assert_called_once() enqueued = mock_worker.ensure_initialized_and_enqueue.call_args - assert (enqueued.args or tuple(enqueued.kwargs.values()))[0] is logging_obj.dispatch_success_handlers.return_value + assert (enqueued.args or tuple(enqueued.kwargs.values()))[ + 0 + ] is logging_obj.dispatch_success_handlers.return_value logging_obj.dispatch_success_handlers.assert_called_once_with(streaming.messages, prefer_async_handlers=True) logging_obj.success_handler.assert_not_called() # the bare create_task path must no longer be used for success logging 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 e728a2363b9..35c4caf51cf 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 @@ -210,7 +210,9 @@ def stub_aws_sdk_client(monkeypatch): async def invoke_model_with_bidirectional_stream(self, operation_input): captured["operation_input"] = operation_input - return ImmediatelyEndingBedrockStream() + # Tests that need the session to see Bedrock frames set captured["stream_events"]; + # with none set this replays nothing and ends immediately. + return ScriptedBedrockStream(captured.get("stream_events", [])) package = types.ModuleType("aws_sdk_bedrock_runtime") client_module = types.ModuleType("aws_sdk_bedrock_runtime.client") @@ -394,6 +396,90 @@ class TestBedrockRealtimeSessionLifecycle: dispatched = logging_obj.dispatched_results[0] assert any(event.get("type") == "session.created" for event in dispatched) + @pytest.mark.asyncio + async def test_unbilled_usage_at_session_close_is_flushed_into_the_spend_log( + self, stub_aws_sdk_client, drain_bedrock_realtime_logging_worker + ): + """ + Usage that arrives while a response is open is not billed by any response.done during + the session, so closing the session must flush it or the turn is never charged. + """ + stub_aws_sdk_client["stream_events"] = [ + {"event": {"contentStart": {"role": "ASSISTANT", "type": "TEXT"}}}, + { + "event": { + "usageEvent": { + "totalInputTokens": 12, + "totalOutputTokens": 23, + "totalTokens": 35, + "details": { + "total": { + "input": {"speechTokens": 10, "textTokens": 2}, + "output": {"speechTokens": 20, "textTokens": 3}, + } + }, + } + } + }, + ] + handler = BedrockRealtime() + logging_obj = FakeLogging() + + await handler.async_realtime( + model="amazon.nova-sonic-v1:0", + websocket=RealtimeClientWS(), + logging_obj=logging_obj, + aws_region_name="us-east-1", + aws_access_key_id="k", + aws_secret_access_key="s", + ) + + await drain_bedrock_realtime_logging_worker.pop() + dispatched = logging_obj.dispatched_results[0] + done_events = [event for event in dispatched if event.get("type") == "response.done"] + assert len(done_events) == 1, "session close did not flush the open turn's usage" + usage = done_events[0]["response"]["usage"] + assert usage["input_tokens"] == 12 + assert usage["output_tokens"] == 23 + assert usage["total_tokens"] == 35 + + @pytest.mark.asyncio + async def test_client_session_update_reaches_spend_logging(self, stub_aws_models): + """Declared tools and instructions must reach the spend log via store_input.""" + handler = BedrockRealtime() + client_ws = DisconnectingClientWS( + [ + json.dumps( + { + "type": "session.update", + "session": { + "instructions": "be brief", + "tools": [{"type": "function", "name": "get_weather"}], + }, + } + ) + ] + ) + realtime_streaming = RealTimeStreaming( + websocket=client_ws, + backend_ws=None, + logging_obj=FakeLogging(), + model="amazon.nova-sonic-v1:0", + ) + + await handler._forward_client_to_bedrock( + client_ws, + FakeBedrockStream(), + BedrockRealtimeConfig(), + "amazon.nova-sonic-v1:0", + {}, + FakeLogging(), + realtime_streaming, + ) + + assert realtime_streaming.session_tools == [{"type": "function", "name": "get_weather"}] + assert {"role": "system", "content": "be brief"} in realtime_streaming.input_messages + @pytest.mark.asyncio async def test_tool_call_reaches_spend_logging_via_response_done(self): """ 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 9c01bc8b48e..2d3c0360faa 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 @@ -1393,6 +1393,29 @@ class TestBedrockRealtimeContentBlockLifecycle: assert sum(1 for msg in second["response"] if msg["type"] == "response.output_item.added") == 1 +class TestBedrockRealtimeToolArgumentParsing: + """Nova Sonic tool args arrive in several shapes; none may crash or leak a raw wire value.""" + + def _args(self, tool_use: dict) -> str: + config = BedrockRealtimeConfig() + events, _, _ = config.transform_tool_use_event({"toolUse": tool_use}, "item_1", "resp_1", "conv_1") + return _only(events, "response.function_call_arguments.done")["arguments"] + + def test_already_decoded_object_content_is_passed_through(self): + assert json.loads(self._args({"toolUseId": "t", "toolName": "f", "content": {"location": "Seattle"}})) == { + "location": "Seattle" + } + + def test_empty_content_falls_back_to_empty_arguments(self): + assert json.loads(self._args({"toolUseId": "t", "toolName": "f", "content": ""})) == {} + + def test_missing_content_and_input_yields_empty_arguments(self): + assert json.loads(self._args({"toolUseId": "t", "toolName": "f"})) == {} + + def test_non_json_content_yields_empty_arguments(self): + assert json.loads(self._args({"toolUseId": "t", "toolName": "f", "content": "not json"})) == {} + + class TestBedrockRealtimeUsageAccounting: def _usage_event( self, @@ -1643,6 +1666,20 @@ class TestBedrockRealtimeUsageAccounting: assert not config.has_unbilled_usage() assert config.flush_pending_usage_as_response_done(None, None) == [] + def test_non_numeric_token_counts_are_billed_as_zero(self): + """A malformed usageEvent must not crash the session or bill a bogus amount.""" + config = BedrockRealtimeConfig() + config.record_usage_event( + { + "totalInputTokens": "twelve", + "totalOutputTokens": True, + "totalTokens": None, + "details": {"total": {"input": {"speechTokens": None}, "output": {"textTokens": "x"}}}, + } + ) + assert not config.has_unbilled_usage() + assert config.flush_pending_usage_as_response_done(None, None) == [] + def test_flush_pending_usage_on_session_close(self): config = BedrockRealtimeConfig() config.record_usage_event(