mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-20 00:11:50 +00:00
Merge pull request #41542 from BerriAI/litellm_bedrock_realtime_sdk_0_11
fix(bedrock): support aws-sdk-bedrock-runtime 0.10/0.11 in Bedrock Realtime
This commit is contained in:
commit
bc9f4fec5b
7 changed files with 344 additions and 82 deletions
|
|
@ -320,6 +320,8 @@ WEBSOCKET_CLOSE_REASON_MAX_BYTES: Final = 123
|
|||
BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY: Final = "litellm.bedrock_realtime.pending_session_update"
|
||||
BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY: Final = "litellm.bedrock_realtime.session_committed"
|
||||
BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY: Final = "litellm.bedrock_realtime.committed_failure"
|
||||
BEDROCK_REALTIME_SDK_DISTRIBUTION: Final = "aws-sdk-bedrock-runtime"
|
||||
BEDROCK_REALTIME_SDK_SUPPORTED_RANGE: Final = ">=0.10.0,<0.12.0"
|
||||
CLIENT_REQUESTED_MODEL_SCOPE_KEY: Final = "litellm.client_requested_model"
|
||||
MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY: Final = "litellm.model_group_alias_resolved"
|
||||
REALTIME_SESSION_SUCCESS_LOGGED_KEY: Final = "realtime_session_success_logged"
|
||||
|
|
|
|||
|
|
@ -6,11 +6,12 @@ This uses aws_sdk_bedrock_runtime for bidirectional streaming with Nova Sonic.
|
|||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import importlib.metadata
|
||||
import json
|
||||
from collections.abc import AsyncIterator, Mapping, MutableMapping
|
||||
from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, MutableMapping
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import Final, NoReturn, Protocol
|
||||
from typing import Final, NoReturn, Protocol, runtime_checkable
|
||||
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
|
|
@ -19,6 +20,8 @@ from litellm._logging import _redact_string, verbose_proxy_logger
|
|||
from litellm.constants import (
|
||||
BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY,
|
||||
BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY,
|
||||
BEDROCK_REALTIME_SDK_DISTRIBUTION,
|
||||
BEDROCK_REALTIME_SDK_SUPPORTED_RANGE,
|
||||
BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY,
|
||||
REALTIME_SESSION_SUCCESS_LOGGED_KEY,
|
||||
)
|
||||
|
|
@ -121,6 +124,38 @@ class BedrockBidirectionalStream(Protocol):
|
|||
async def await_output(self) -> tuple[object, BedrockOutputStream]: ...
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class ClosableBedrockRuntimeClient(Protocol):
|
||||
async def close(self) -> None: ...
|
||||
|
||||
|
||||
def _installed_sdk_version() -> str | None:
|
||||
try:
|
||||
return importlib.metadata.version(BEDROCK_REALTIME_SDK_DISTRIBUTION)
|
||||
except importlib.metadata.PackageNotFoundError:
|
||||
return None
|
||||
|
||||
|
||||
def _sdk_import_error(installed_version: str | None, cause: ImportError) -> ImportError:
|
||||
install_hint: Final = "pip install 'litellm[bedrock-realtime]'"
|
||||
requirement: Final = f"{BEDROCK_REALTIME_SDK_DISTRIBUTION}[awscrt]{BEDROCK_REALTIME_SDK_SUPPORTED_RANGE}"
|
||||
verbose_proxy_logger.error("Bedrock Realtime: SDK import failed (installed=%s): %s", installed_version, cause)
|
||||
if installed_version is None:
|
||||
return ImportError(f"Missing aws_sdk_bedrock_runtime: {install_hint} ({requirement})")
|
||||
return ImportError(
|
||||
f"{BEDROCK_REALTIME_SDK_DISTRIBUTION} {installed_version} is installed but Bedrock realtime needs "
|
||||
f"[awscrt]{BEDROCK_REALTIME_SDK_SUPPORTED_RANGE}: {install_hint}"
|
||||
)
|
||||
|
||||
|
||||
async def _close_bedrock_client(bedrock_client: object) -> None:
|
||||
if not isinstance(bedrock_client, ClosableBedrockRuntimeClient):
|
||||
return
|
||||
with contextlib.suppress(Exception):
|
||||
await bedrock_client.close()
|
||||
verbose_proxy_logger.debug("Bedrock Realtime: closed SDK client")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _BridgeOutcome:
|
||||
logged_events: tuple[OpenAIRealtimeEvents, ...]
|
||||
|
|
@ -199,8 +234,9 @@ async def _ack_session_update(
|
|||
class BedrockRealtime(BaseAWSLLM):
|
||||
"""Handler for Bedrock Nova Sonic realtime speech-to-speech API."""
|
||||
|
||||
def __init__(self):
|
||||
def __init__(self, sdk_version_lookup: Callable[[], str | None] = _installed_sdk_version):
|
||||
super().__init__()
|
||||
self._sdk_version_lookup: Final = sdk_version_lookup
|
||||
|
||||
async def async_realtime(
|
||||
self,
|
||||
|
|
@ -234,14 +270,13 @@ class BedrockRealtime(BaseAWSLLM):
|
|||
Various AWS authentication parameters
|
||||
"""
|
||||
try:
|
||||
from aws_sdk_bedrock_runtime.client import (
|
||||
BedrockRuntimeClient,
|
||||
InvokeModelWithBidirectionalStreamOperationInput,
|
||||
)
|
||||
from aws_sdk_bedrock_runtime.config import Config
|
||||
from smithy_aws_core.identity import StaticCredentialsResolver
|
||||
except ImportError:
|
||||
raise ImportError("Missing aws_sdk_bedrock_runtime. Install with: pip install aws-sdk-bedrock-runtime")
|
||||
from aws_sdk_bedrock_runtime.client import AsyncBedrockRuntimeClient
|
||||
from aws_sdk_bedrock_runtime.config import AsyncBedrockRuntimeConfig
|
||||
from aws_sdk_bedrock_runtime.models import InvokeModelWithBidirectionalStreamOperationInput
|
||||
from smithy_aws_core.identity import AWSCredentialsIdentity, StaticCredentialsResolver
|
||||
from smithy_http.aio.crt import AWSCRTHTTPClient
|
||||
except ImportError as e:
|
||||
raise _sdk_import_error(self._sdk_version_lookup(), e) from e
|
||||
|
||||
pending_session_update: Final = _pending_session_update(websocket.scope)
|
||||
|
||||
|
|
@ -285,22 +320,37 @@ class BedrockRealtime(BaseAWSLLM):
|
|||
)
|
||||
frozen_credentials: Final = await run_aws_signing(credentials.get_frozen_credentials)
|
||||
|
||||
# Initialize Bedrock client with aws_sdk_bedrock_runtime
|
||||
config: Final = Config(
|
||||
credentials_identity: Final = AWSCredentialsIdentity(
|
||||
access_key_id=frozen_credentials.access_key,
|
||||
secret_access_key=frozen_credentials.secret_key,
|
||||
session_token=frozen_credentials.token,
|
||||
)
|
||||
config: Final = await AsyncBedrockRuntimeConfig.resolve(
|
||||
endpoint_uri=endpoint_uri,
|
||||
region=aws_region_name,
|
||||
aws_access_key_id=frozen_credentials.access_key,
|
||||
aws_secret_access_key=frozen_credentials.secret_key,
|
||||
aws_session_token=frozen_credentials.token,
|
||||
aws_credentials_identity_resolver=StaticCredentialsResolver(),
|
||||
aws_credentials_identity_resolver=StaticCredentialsResolver(identity=credentials_identity),
|
||||
transport=AWSCRTHTTPClient(),
|
||||
)
|
||||
bedrock_client: Final = BedrockRuntimeClient(config=config)
|
||||
bedrock_client: Final = AsyncBedrockRuntimeClient(config=config)
|
||||
|
||||
async def open_bidirectional_stream() -> BedrockBidirectionalStream:
|
||||
return await bedrock_client.invoke_model_with_bidirectional_stream(
|
||||
InvokeModelWithBidirectionalStreamOperationInput(model_id=model)
|
||||
)
|
||||
|
||||
try:
|
||||
await self._run_session(websocket, open_bidirectional_stream, model, logging_obj, pending_session_update)
|
||||
finally:
|
||||
await _close_bedrock_client(bedrock_client)
|
||||
|
||||
async def _run_session(
|
||||
self,
|
||||
websocket: RealtimeClientWebSocket,
|
||||
open_bidirectional_stream: Callable[[], Awaitable[BedrockBidirectionalStream]],
|
||||
model: str,
|
||||
logging_obj: LiteLLMLogging,
|
||||
pending_session_update: str | None,
|
||||
) -> None:
|
||||
transformation_config: Final = BedrockRealtimeConfig()
|
||||
|
||||
bedrock_stream: Final = await open_bidirectional_stream()
|
||||
|
|
|
|||
|
|
@ -143,8 +143,9 @@ bedrock-realtime = [
|
|||
# InvokeModelWithBidirectionalStream API, which boto3 cannot do. This
|
||||
# experimental AWS SDK (with its smithy-* deps, pulled transitively)
|
||||
# provides the bidirectional stream; imported lazily in the realtime
|
||||
# handler so litellm core stays usable without it.
|
||||
"aws-sdk-bedrock-runtime>=0.7.0,<0.8.0; python_version >= '3.12'",
|
||||
# handler so litellm core stays usable without it. The awscrt extra is
|
||||
# required: the SDK's default aiohttp transport has no duplex streaming.
|
||||
"aws-sdk-bedrock-runtime[awscrt]>=0.10.0,<0.12.0; python_version >= '3.12'",
|
||||
]
|
||||
proxy-runtime = [
|
||||
# Historically bundled in the proxy Docker images via requirements.txt.
|
||||
|
|
|
|||
|
|
@ -20,7 +20,9 @@ import pytest
|
|||
|
||||
IMAGE: Final = os.getenv("LITELLM_IMAGE")
|
||||
NON_ROOT_UID: Final = "12345:0"
|
||||
IMPORT_PROBE: Final = "import aws_sdk_bedrock_runtime, smithy_aws_core; print('bedrock-realtime ok')"
|
||||
IMPORT_PROBE: Final = (
|
||||
"import aws_sdk_bedrock_runtime, smithy_aws_core, smithy_http.aio.crt; print('bedrock-realtime ok')"
|
||||
)
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.skipif(IMAGE is None, reason="requires a built image (set LITELLM_IMAGE)"),
|
||||
|
|
@ -52,7 +54,7 @@ def test_image_imports_bedrock_realtime_sdk():
|
|||
)
|
||||
|
||||
assert probe.returncode == 0 and "bedrock-realtime ok" in probe.stdout, (
|
||||
f"{IMAGE} cannot import aws_sdk_bedrock_runtime as uid {NON_ROOT_UID}, so Bedrock Nova Sonic "
|
||||
"/v1/realtime sessions fail with 'Missing aws_sdk_bedrock_runtime'. Is `--extra bedrock-realtime` "
|
||||
f"{IMAGE} cannot import aws_sdk_bedrock_runtime with its awscrt transport as uid {NON_ROOT_UID}, so "
|
||||
"Bedrock Nova Sonic /v1/realtime sessions fail at SDK import. Is `--extra bedrock-realtime` "
|
||||
f"passed to every `uv sync` in its Dockerfile?\nstdout:\n{probe.stdout}\nstderr:\n{probe.stderr}"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -8,7 +8,11 @@ from unittest.mock import MagicMock
|
|||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.constants import REALTIME_SESSION_SUCCESS_LOGGED_KEY
|
||||
from litellm.constants import (
|
||||
BEDROCK_REALTIME_SDK_SUPPORTED_RANGE,
|
||||
REALTIME_SESSION_SUCCESS_LOGGED_KEY,
|
||||
WEBSOCKET_CLOSE_REASON_MAX_BYTES,
|
||||
)
|
||||
from litellm.llms.bedrock.common_utils import BedrockError
|
||||
from litellm.llms.bedrock.realtime.handler import BedrockRealtime
|
||||
from litellm.llms.bedrock.realtime.transformation import BedrockRealtimeConfig
|
||||
|
|
@ -207,7 +211,19 @@ class ScriptedBedrockStream:
|
|||
return (None, self._receiver)
|
||||
|
||||
|
||||
class FakeAWSCredentialsIdentity:
|
||||
def __init__(self, access_key_id, secret_access_key, session_token=None):
|
||||
self.access_key_id = access_key_id
|
||||
self.secret_access_key = secret_access_key
|
||||
self.session_token = session_token
|
||||
|
||||
|
||||
class FakeStaticCredentialsResolver:
|
||||
def __init__(self, identity=None):
|
||||
self.identity = identity
|
||||
|
||||
|
||||
class FakeAWSCRTHTTPClient:
|
||||
pass
|
||||
|
||||
|
||||
|
|
@ -227,48 +243,32 @@ class StubCredentialsBedrockRealtime(BedrockRealtime):
|
|||
return SimpleNamespace(get_frozen_credentials=lambda: self.frozen_credentials)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def stub_aws_sdk_client(monkeypatch):
|
||||
captured = {}
|
||||
class FakeOperationInput:
|
||||
def __init__(self, model_id):
|
||||
self.model_id = model_id
|
||||
|
||||
class CapturingConfig:
|
||||
def __init__(self, **kwargs):
|
||||
captured["config_kwargs"] = kwargs
|
||||
self.kwargs = kwargs
|
||||
|
||||
class FakeOperationInput:
|
||||
def __init__(self, model_id):
|
||||
self.model_id = model_id
|
||||
|
||||
class FakeBedrockRuntimeClient:
|
||||
def __init__(self, config):
|
||||
captured["client_config"] = config
|
||||
|
||||
async def invoke_model_with_bidirectional_stream(self, operation_input):
|
||||
captured["operation_input"] = operation_input
|
||||
if captured.get("streams"):
|
||||
stream = captured["streams"].pop(0)
|
||||
if isinstance(stream, Exception):
|
||||
raise stream
|
||||
return stream
|
||||
return ScriptedBedrockStream(captured.get("scripted_payloads", []))
|
||||
|
||||
def _install_fake_sdk_modules(monkeypatch, client_module, config_module):
|
||||
"""Wire fake aws_sdk_bedrock_runtime / smithy packages into sys.modules for the handler's lazy imports."""
|
||||
package = types.ModuleType("aws_sdk_bedrock_runtime")
|
||||
client_module = types.ModuleType("aws_sdk_bedrock_runtime.client")
|
||||
client_module.BedrockRuntimeClient = FakeBedrockRuntimeClient
|
||||
client_module.InvokeModelWithBidirectionalStreamOperationInput = FakeOperationInput
|
||||
config_module = types.ModuleType("aws_sdk_bedrock_runtime.config")
|
||||
config_module.Config = CapturingConfig
|
||||
models_module = types.ModuleType("aws_sdk_bedrock_runtime.models")
|
||||
models_module.BidirectionalInputPayloadPart = FakePayloadPart
|
||||
models_module.InvokeModelWithBidirectionalStreamInputChunk = FakeInputChunk
|
||||
models_module.InvokeModelWithBidirectionalStreamOperationInput = FakeOperationInput
|
||||
package.client = client_module
|
||||
package.config = config_module
|
||||
package.models = models_module
|
||||
smithy_package = types.ModuleType("smithy_aws_core")
|
||||
identity_module = types.ModuleType("smithy_aws_core.identity")
|
||||
identity_module.AWSCredentialsIdentity = FakeAWSCredentialsIdentity
|
||||
identity_module.StaticCredentialsResolver = FakeStaticCredentialsResolver
|
||||
smithy_package.identity = identity_module
|
||||
smithy_http_package = types.ModuleType("smithy_http")
|
||||
smithy_http_aio = types.ModuleType("smithy_http.aio")
|
||||
crt_module = types.ModuleType("smithy_http.aio.crt")
|
||||
crt_module.AWSCRTHTTPClient = FakeAWSCRTHTTPClient
|
||||
smithy_http_aio.crt = crt_module
|
||||
smithy_http_package.aio = smithy_http_aio
|
||||
|
||||
stubbed_modules = {
|
||||
"aws_sdk_bedrock_runtime": package,
|
||||
|
|
@ -277,10 +277,56 @@ def stub_aws_sdk_client(monkeypatch):
|
|||
"aws_sdk_bedrock_runtime.models": models_module,
|
||||
"smithy_aws_core": smithy_package,
|
||||
"smithy_aws_core.identity": identity_module,
|
||||
"smithy_http": smithy_http_package,
|
||||
"smithy_http.aio": smithy_http_aio,
|
||||
"smithy_http.aio.crt": crt_module,
|
||||
}
|
||||
for module_name, module in stubbed_modules.items():
|
||||
monkeypatch.setitem(sys.modules, module_name, module)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def stub_aws_sdk_client(monkeypatch):
|
||||
"""Fake of the aws-sdk-bedrock-runtime 0.10/0.11 surface: async config resolve, async client with close()"""
|
||||
captured = {}
|
||||
|
||||
class FakeAsyncBedrockRuntimeConfig:
|
||||
def __init__(self, kwargs):
|
||||
self.kwargs = kwargs
|
||||
|
||||
@classmethod
|
||||
async def resolve(cls, **kwargs):
|
||||
captured["config_kwargs"] = kwargs
|
||||
return cls(kwargs)
|
||||
|
||||
class FakeAsyncBedrockRuntimeClient:
|
||||
def __init__(self, config):
|
||||
captured["client_config"] = config
|
||||
captured["client_closed"] = False
|
||||
|
||||
async def invoke_model_with_bidirectional_stream(self, operation_input):
|
||||
captured["operation_input"] = operation_input
|
||||
if captured.get("streams"):
|
||||
stream = captured["streams"].pop(0)
|
||||
if isinstance(stream, Exception):
|
||||
raise stream
|
||||
captured["open_stream"] = stream
|
||||
return stream
|
||||
stream = ScriptedBedrockStream(captured.get("scripted_payloads", []))
|
||||
captured["open_stream"] = stream
|
||||
return stream
|
||||
|
||||
async def close(self):
|
||||
open_stream = captured.get("open_stream")
|
||||
captured["input_closed_before_client_close"] = open_stream is None or open_stream.input_stream.closed
|
||||
captured["client_closed"] = True
|
||||
|
||||
client_module = types.ModuleType("aws_sdk_bedrock_runtime.client")
|
||||
client_module.AsyncBedrockRuntimeClient = FakeAsyncBedrockRuntimeClient
|
||||
config_module = types.ModuleType("aws_sdk_bedrock_runtime.config")
|
||||
config_module.AsyncBedrockRuntimeConfig = FakeAsyncBedrockRuntimeConfig
|
||||
_install_fake_sdk_modules(monkeypatch, client_module, config_module)
|
||||
|
||||
for env_var in (
|
||||
"AWS_ACCESS_KEY_ID",
|
||||
"AWS_SECRET_ACCESS_KEY",
|
||||
|
|
@ -764,15 +810,33 @@ class TestBedrockRealtimeAwsAuth:
|
|||
)
|
||||
|
||||
config_kwargs = stub_aws_sdk_client["config_kwargs"]
|
||||
assert config_kwargs["aws_access_key_id"] == "litellm-params-access-key"
|
||||
assert config_kwargs["aws_secret_access_key"] == "litellm-params-secret-key"
|
||||
assert config_kwargs["aws_session_token"] == "litellm-params-session-token"
|
||||
assert isinstance(config_kwargs["aws_credentials_identity_resolver"], FakeStaticCredentialsResolver)
|
||||
resolver = config_kwargs["aws_credentials_identity_resolver"]
|
||||
assert isinstance(resolver, FakeStaticCredentialsResolver)
|
||||
assert resolver.identity.access_key_id == "litellm-params-access-key"
|
||||
assert resolver.identity.secret_access_key == "litellm-params-secret-key"
|
||||
assert resolver.identity.session_token == "litellm-params-session-token"
|
||||
assert config_kwargs["region"] == "us-east-1"
|
||||
assert config_kwargs["endpoint_uri"] == "https://bedrock-runtime.us-east-1.amazonaws.com"
|
||||
assert isinstance(config_kwargs["transport"], FakeAWSCRTHTTPClient)
|
||||
assert stub_aws_sdk_client["client_config"].kwargs is config_kwargs
|
||||
assert stub_aws_sdk_client["operation_input"].model_id == "amazon.nova-sonic-v1:0"
|
||||
assert websocket.closed
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_api_base_overrides_default_endpoint(self, stub_aws_sdk_client):
|
||||
await BedrockRealtime().async_realtime(
|
||||
model="amazon.nova-sonic-v1:0",
|
||||
websocket=RealtimeClientWS(),
|
||||
logging_obj=FakeLogging(),
|
||||
aws_region_name="us-east-1",
|
||||
aws_access_key_id="k",
|
||||
aws_secret_access_key="s",
|
||||
api_base="https://vpce-bedrock.example.internal",
|
||||
aws_bedrock_runtime_endpoint="https://ignored.example.internal",
|
||||
)
|
||||
|
||||
assert stub_aws_sdk_client["config_kwargs"]["endpoint_uri"] == "https://vpce-bedrock.example.internal"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_role_assumption_params_forwarded_to_get_credentials(self, stub_aws_sdk_client):
|
||||
handler = StubCredentialsBedrockRealtime(
|
||||
|
|
@ -805,11 +869,11 @@ class TestBedrockRealtimeAwsAuth:
|
|||
"aws_sts_endpoint": None,
|
||||
"aws_external_id": "realtime-external-id",
|
||||
}
|
||||
config_kwargs = stub_aws_sdk_client["config_kwargs"]
|
||||
assert config_kwargs["aws_access_key_id"] == "assumed-access-key"
|
||||
assert config_kwargs["aws_secret_access_key"] == "assumed-secret-key"
|
||||
assert config_kwargs["aws_session_token"] == "assumed-session-token"
|
||||
assert isinstance(config_kwargs["aws_credentials_identity_resolver"], FakeStaticCredentialsResolver)
|
||||
resolver = stub_aws_sdk_client["config_kwargs"]["aws_credentials_identity_resolver"]
|
||||
assert isinstance(resolver, FakeStaticCredentialsResolver)
|
||||
assert resolver.identity.access_key_id == "assumed-access-key"
|
||||
assert resolver.identity.secret_access_key == "assumed-secret-key"
|
||||
assert resolver.identity.session_token == "assumed-session-token"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unresolvable_credentials_raise_clear_auth_error(self, stub_aws_sdk_client):
|
||||
|
|
@ -826,5 +890,118 @@ class TestBedrockRealtimeAwsAuth:
|
|||
assert "config_kwargs" not in stub_aws_sdk_client
|
||||
|
||||
|
||||
class TestBedrockRealtimeSdkLifecycle:
|
||||
"""aws-sdk-bedrock-runtime 0.10/0.11: async config, async client, CRT transport, close() (LIT-7938 regression)"""
|
||||
|
||||
AWS_ARGS = {
|
||||
"model": "amazon.nova-sonic-v1:0",
|
||||
"aws_region_name": "us-east-1",
|
||||
"aws_access_key_id": "k",
|
||||
"aws_secret_access_key": "s",
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_closed_after_input_stream_on_normal_completion(self, stub_aws_sdk_client):
|
||||
await BedrockRealtime().async_realtime(websocket=RealtimeClientWS(), logging_obj=FakeLogging(), **self.AWS_ARGS)
|
||||
|
||||
assert stub_aws_sdk_client["client_closed"]
|
||||
assert stub_aws_sdk_client["input_closed_before_client_close"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_closed_when_stream_open_fails(self, stub_aws_sdk_client):
|
||||
stub_aws_sdk_client["streams"] = [ServiceUnavailableException("bedrock unavailable")]
|
||||
|
||||
with pytest.raises(ServiceUnavailableException):
|
||||
await BedrockRealtime().async_realtime(
|
||||
websocket=RealtimeClientWS(), logging_obj=FakeLogging(), **self.AWS_ARGS
|
||||
)
|
||||
|
||||
assert stub_aws_sdk_client["client_closed"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_closed_when_provider_stream_fails_mid_session(self, stub_aws_sdk_client):
|
||||
stub_aws_sdk_client["streams"] = [ScriptedBedrockStream([], receiver_type=BreakingBedrockReceiver)]
|
||||
|
||||
with pytest.raises(BedrockError):
|
||||
await BedrockRealtime().async_realtime(
|
||||
websocket=ConnectedClientWS([]), logging_obj=FakeLogging(), **self.AWS_ARGS
|
||||
)
|
||||
|
||||
assert stub_aws_sdk_client["client_closed"]
|
||||
assert stub_aws_sdk_client["input_closed_before_client_close"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_without_close_completes_session(self, monkeypatch):
|
||||
class ClientWithoutClose:
|
||||
def __init__(self, config):
|
||||
pass
|
||||
|
||||
async def invoke_model_with_bidirectional_stream(self, operation_input):
|
||||
return ScriptedBedrockStream([])
|
||||
|
||||
class ConfigWithoutCapture:
|
||||
@classmethod
|
||||
async def resolve(cls, **kwargs):
|
||||
return cls()
|
||||
|
||||
client_module = types.ModuleType("aws_sdk_bedrock_runtime.client")
|
||||
client_module.AsyncBedrockRuntimeClient = ClientWithoutClose
|
||||
config_module = types.ModuleType("aws_sdk_bedrock_runtime.config")
|
||||
config_module.AsyncBedrockRuntimeConfig = ConfigWithoutCapture
|
||||
_install_fake_sdk_modules(monkeypatch, client_module, config_module)
|
||||
websocket = RealtimeClientWS()
|
||||
|
||||
await BedrockRealtime().async_realtime(websocket=websocket, logging_obj=FakeLogging(), **self.AWS_ARGS)
|
||||
|
||||
assert websocket.closed
|
||||
|
||||
|
||||
class TestBedrockRealtimeSdkImportErrors:
|
||||
"""Init errors must tell 'SDK not installed' apart from 'SDK installed but unsupported version' (LIT-7938)"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_absent_sdk_names_install_extra(self, monkeypatch):
|
||||
monkeypatch.setitem(sys.modules, "aws_sdk_bedrock_runtime", None)
|
||||
handler = BedrockRealtime(sdk_version_lookup=lambda: None)
|
||||
|
||||
with pytest.raises(ImportError) as exc_info:
|
||||
await handler.async_realtime(
|
||||
model="amazon.nova-sonic-v1:0", websocket=RealtimeClientWS(), logging_obj=FakeLogging()
|
||||
)
|
||||
|
||||
message = str(exc_info.value)
|
||||
assert message.startswith("Missing aws_sdk_bedrock_runtime")
|
||||
assert "litellm[bedrock-realtime]" in message
|
||||
assert "is installed but" not in message
|
||||
close_reason = message.encode()[:WEBSOCKET_CLOSE_REASON_MAX_BYTES].decode()
|
||||
assert BEDROCK_REALTIME_SDK_SUPPORTED_RANGE in close_reason
|
||||
assert "pip install 'litellm[bedrock-realtime]'" in close_reason
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_incompatible_sdk_names_installed_version_and_supported_range(self, monkeypatch):
|
||||
legacy_client_module = types.ModuleType("aws_sdk_bedrock_runtime.client")
|
||||
legacy_client_module.BedrockRuntimeClient = object
|
||||
legacy_config_module = types.ModuleType("aws_sdk_bedrock_runtime.config")
|
||||
legacy_config_module.Config = object
|
||||
_install_fake_sdk_modules(monkeypatch, legacy_client_module, legacy_config_module)
|
||||
handler = BedrockRealtime(sdk_version_lookup=lambda: "0.7.0")
|
||||
|
||||
with pytest.raises(ImportError) as exc_info:
|
||||
await handler.async_realtime(
|
||||
model="amazon.nova-sonic-v1:0", websocket=RealtimeClientWS(), logging_obj=FakeLogging()
|
||||
)
|
||||
|
||||
message = str(exc_info.value)
|
||||
assert "aws-sdk-bedrock-runtime 0.7.0 is installed but" in message
|
||||
assert ">=0.10.0,<0.12.0" in message
|
||||
assert not message.startswith("Missing aws_sdk_bedrock_runtime")
|
||||
assert isinstance(exc_info.value.__cause__, ImportError)
|
||||
assert str(exc_info.value.__cause__) not in message
|
||||
assert "cannot import name" not in message
|
||||
close_reason = message.encode()[:WEBSOCKET_CLOSE_REASON_MAX_BYTES].decode()
|
||||
assert "0.7.0 is installed" in close_reason
|
||||
assert BEDROCK_REALTIME_SDK_SUPPORTED_RANGE in close_reason
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
|
|
|
|||
|
|
@ -4,15 +4,23 @@ Static checks that every proxy Docker image installs the `bedrock-realtime` extr
|
|||
Bedrock Nova Sonic speech-to-speech (`/v1/realtime`) needs `aws-sdk-bedrock-runtime`,
|
||||
which only ships in the `bedrock-realtime` extra. An image whose `uv sync` stages
|
||||
omit the extra fails every Nova Sonic realtime session with
|
||||
"Missing aws_sdk_bedrock_runtime. Install with: pip install aws-sdk-bedrock-runtime".
|
||||
"Missing aws_sdk_bedrock_runtime: pip install 'litellm[bedrock-realtime]' ...".
|
||||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.constants import BEDROCK_REALTIME_SDK_DISTRIBUTION, BEDROCK_REALTIME_SDK_SUPPORTED_RANGE
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
import tomllib
|
||||
else:
|
||||
import tomli as tomllib
|
||||
|
||||
REPO_ROOT: Final = os.path.join(os.path.dirname(__file__), "..", "..")
|
||||
|
||||
PROXY_DOCKERFILES: Final = (
|
||||
|
|
@ -54,3 +62,16 @@ def test_every_uv_sync_installs_bedrock_realtime_extra(relative_path: str):
|
|||
"`--extra bedrock-realtime`, so aws-sdk-bedrock-runtime is absent and Bedrock Nova Sonic "
|
||||
"/v1/realtime sessions fail with 'Missing aws_sdk_bedrock_runtime'"
|
||||
)
|
||||
|
||||
|
||||
def test_bedrock_realtime_extra_pins_the_range_named_in_the_runtime_error():
|
||||
with open(os.path.join(REPO_ROOT, "pyproject.toml"), "rb") as f:
|
||||
extra_specs: Final = tomllib.load(f)["project"]["optional-dependencies"]["bedrock-realtime"]
|
||||
|
||||
sdk_specs: Final = tuple(spec for spec in extra_specs if spec.startswith(BEDROCK_REALTIME_SDK_DISTRIBUTION))
|
||||
assert len(sdk_specs) == 1, f"expected exactly one {BEDROCK_REALTIME_SDK_DISTRIBUTION} spec, got {extra_specs}"
|
||||
requirement: Final = sdk_specs[0].split(";")[0].strip()
|
||||
assert requirement == f"{BEDROCK_REALTIME_SDK_DISTRIBUTION}[awscrt]{BEDROCK_REALTIME_SDK_SUPPORTED_RANGE}", (
|
||||
f"pyproject pins {requirement!r} but the handler's install hint names "
|
||||
f"{BEDROCK_REALTIME_SDK_SUPPORTED_RANGE!r} with the awscrt extra; keep them in sync"
|
||||
)
|
||||
|
|
|
|||
47
uv.lock
generated
47
uv.lock
generated
|
|
@ -10,7 +10,7 @@ resolution-markers = [
|
|||
]
|
||||
|
||||
[options]
|
||||
exclude-newer = "2026-09-14T20:08:36.384435Z"
|
||||
exclude-newer = "2026-09-14T20:32:38.482736111Z"
|
||||
exclude-newer-span = "P3D"
|
||||
|
||||
[manifest]
|
||||
|
|
@ -535,16 +535,21 @@ wheels = [
|
|||
|
||||
[[package]]
|
||||
name = "aws-sdk-bedrock-runtime"
|
||||
version = "0.7.0"
|
||||
version = "0.11.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "smithy-aws-core", extra = ["eventstream", "json"], marker = "python_full_version >= '3.12'" },
|
||||
{ name = "smithy-core", marker = "python_full_version >= '3.12'" },
|
||||
{ name = "smithy-http", extra = ["awscrt"], marker = "python_full_version >= '3.12'" },
|
||||
{ name = "smithy-http", extra = ["aiohttp"], marker = "python_full_version >= '3.12'" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/67/8a/ed3fd98775273b0b7f6006b4970aa876d506668b7fe29145f54fcb941c3b/aws_sdk_bedrock_runtime-0.7.0.tar.gz", hash = "sha256:0cb172cbc03ff060e5c1d6f9cfa9a8ac5e71d9e0d58d3117006ebf614cbb4677", size = 170304, upload-time = "2026-06-23T04:04:52.382Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/8e/b3/9c225cbfe9f17ea2e3d75a0fdd0b325ef79839b9c09a376bda63a7bf3bb3/aws_sdk_bedrock_runtime-0.11.0.tar.gz", hash = "sha256:f2c45d34625bf6a7b56375e29a53a16b376880bda771e4bbf7d84491622eb193", size = 173854, upload-time = "2026-08-24T21:17:16.304Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/9e/e1/f86d50f0ad9c8200645f315c524d285e86b30b94bb65118e1108597714e6/aws_sdk_bedrock_runtime-0.7.0-py3-none-any.whl", hash = "sha256:de67ede6f441bbb77ef61c237945d559513843fc827abe1af12535c2519650c5", size = 94948, upload-time = "2026-06-23T04:04:51.281Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/29/0c/9512304ed017ce49992df6661eac2b914550247e13bccb55be6ca594170d/aws_sdk_bedrock_runtime-0.11.0-py3-none-any.whl", hash = "sha256:ef01c26ddfd83a5d3e438ab72ebb3c13b41fc0ef11d81095b22c8016f97e9795", size = 97112, upload-time = "2026-08-24T21:17:17.396Z" },
|
||||
]
|
||||
|
||||
[package.optional-dependencies]
|
||||
awscrt = [
|
||||
{ name = "smithy-http", extra = ["awscrt"], marker = "python_full_version >= '3.12'" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -4483,7 +4488,7 @@ dependencies = [
|
|||
|
||||
[package.optional-dependencies]
|
||||
bedrock-realtime = [
|
||||
{ name = "aws-sdk-bedrock-runtime", marker = "python_full_version >= '3.12'" },
|
||||
{ name = "aws-sdk-bedrock-runtime", extra = ["awscrt"], marker = "python_full_version >= '3.12'" },
|
||||
]
|
||||
caching = [
|
||||
{ name = "diskcache" },
|
||||
|
|
@ -4693,7 +4698,7 @@ requires-dist = [
|
|||
{ name = "apscheduler", marker = "extra == 'proxy'", specifier = ">=3.11.2,<4.0" },
|
||||
{ name = "audioread", marker = "extra == 'stt-nvidia-riva'", specifier = ">=3.0.1" },
|
||||
{ name = "aurelio-sdk", marker = "python_full_version < '3.14' and extra == 'semantic-router'", specifier = ">=0.0.19,<1.0" },
|
||||
{ name = "aws-sdk-bedrock-runtime", marker = "python_full_version >= '3.12' and extra == 'bedrock-realtime'", specifier = ">=0.7.0,<0.8.0" },
|
||||
{ name = "aws-sdk-bedrock-runtime", extras = ["awscrt"], marker = "python_full_version >= '3.12' and extra == 'bedrock-realtime'", specifier = ">=0.10.0,<0.12.0" },
|
||||
{ name = "azure-ai-contentsafety", marker = "extra == 'proxy-runtime'", specifier = ">=1.0.0,<2.0" },
|
||||
{ name = "azure-identity", marker = "extra == 'extra-proxy'", specifier = ">=1.25.2,<2.0" },
|
||||
{ name = "azure-identity", marker = "extra == 'proxy'", specifier = ">=1.25.2,<2.0" },
|
||||
|
|
@ -9126,16 +9131,16 @@ wheels = [
|
|||
|
||||
[[package]]
|
||||
name = "smithy-aws-core"
|
||||
version = "0.7.0"
|
||||
version = "0.11.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "aws-sdk-signers", marker = "python_full_version >= '3.12'" },
|
||||
{ name = "smithy-core", marker = "python_full_version >= '3.12'" },
|
||||
{ name = "smithy-http", marker = "python_full_version >= '3.12'" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/fc/a8/37bfde59519f45d2047d0033b791aca6574d867aaf57bb56a6de42ab5c26/smithy_aws_core-0.7.0.tar.gz", hash = "sha256:34e82d09fc808acd5ffc80f03828d0609c6a211f49f0884dc6ee7ca095a1b6af", size = 15670, upload-time = "2026-06-23T04:04:50.365Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/7d/d3/501c0023548173416109ac42298ca33b708469dc922005770811a597949f/smithy_aws_core-0.11.0.tar.gz", hash = "sha256:29ee89976a520a87e3db557e03e115fdc21a0a60b81161e95174395a1b064da1", size = 38791, upload-time = "2026-08-24T21:16:59.631Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/fd/54/2d06dd9a3972a380d71bb8c3312e317aa8f1ea68dd28cffc06955ccf0220/smithy_aws_core-0.7.0-py3-none-any.whl", hash = "sha256:6c60c8fbb9431c60e80ea7f2d37e7ae48409cc1541f587fe073f202eca067e92", size = 24894, upload-time = "2026-06-23T04:04:49.349Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/e4/f6/fefda9aab809fa1a62bf7073bd6d8ab427bd9989f39b13c0d6e29d4d1045/smithy_aws_core-0.11.0-py3-none-any.whl", hash = "sha256:77cf130c22deac14a8cbeb8ccc4bcfe5a91798f4b38cb53a987080ec58c89f23", size = 58855, upload-time = "2026-08-24T21:16:58.657Z" },
|
||||
]
|
||||
|
||||
[package.optional-dependencies]
|
||||
|
|
@ -9160,41 +9165,45 @@ wheels = [
|
|||
|
||||
[[package]]
|
||||
name = "smithy-core"
|
||||
version = "0.6.0"
|
||||
version = "0.8.1"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/e9/45/688d52c61cd4d843bb230694259e91d4c7d6954eeecbadf452a168001d45/smithy_core-0.6.0.tar.gz", hash = "sha256:ba2e5d860d716aff75004a23f53e09dfaca3e2b94f8a00c1f76dcb355b769ce0", size = 52095, upload-time = "2026-06-23T04:04:44.687Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/7c/c6/93e9eea3c6163228dfe972c3e989e0553047858805ab7aa4a59f074ba129/smithy_core-0.8.1.tar.gz", hash = "sha256:3d2f8fca5960d74bd7ef380f70901c7bcdebe53f929d2d3d2fa6cb790b3f5214", size = 54259, upload-time = "2026-08-20T17:55:30.354Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/3a/b6/06795faa9844b9667ae492e6293370393e19e7f0c2df8da1b4bf7e5f6ed9/smithy_core-0.6.0-py3-none-any.whl", hash = "sha256:51e347ed309d60ab9d36b783dbf88de614c460d51bec79d39cd403956b00f063", size = 66879, upload-time = "2026-06-23T04:04:43.596Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/0c/23/c6430bbf406477fc7d16254b9908a723b299a4a21a94c99db9d12c84a8bf/smithy_core-0.8.1-py3-none-any.whl", hash = "sha256:44bd9bdf702f76919af58e44a6a1bb3dc136a745b2f955281743022ce767e347", size = 68805, upload-time = "2026-08-20T17:55:29.366Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "smithy-http"
|
||||
version = "0.4.2"
|
||||
version = "0.5.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "smithy-core", marker = "python_full_version >= '3.12'" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/66/58/5a772d212e066d6fc1398946c4aae19bcdaa75209879d776f641b6a06b5b/smithy_http-0.4.2.tar.gz", hash = "sha256:50d11b6a55e42448450a01e3d0f605ccee65a72abf52d02eed82862a15be5937", size = 29616, upload-time = "2026-06-23T04:04:45.687Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/98/78/b5f3113d6c8f0bc1f9777a7f5ca84b892d29efac05850e14f7d4f7e645b5/smithy_http-0.5.0.tar.gz", hash = "sha256:bb4a19672f7c7eeb872a308f777eb505281a5bafb1ee3d1ea9c760c06c352510", size = 31122, upload-time = "2026-08-24T21:16:56.488Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/57/3e/7b2464d40893bec0b5d1f479d25116d4aa09f9f66536b4c4b3126202215d/smithy_http-0.4.2-py3-none-any.whl", hash = "sha256:a158f107e9fab925289d20772c2e38b0bba94e55c05d0edc9290310f22a60454", size = 41025, upload-time = "2026-06-23T04:04:46.764Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/27/27/e414082643028846b73afa52a1a8f934548196ee12b187a06803f02a3e66/smithy_http-0.5.0-py3-none-any.whl", hash = "sha256:af273d5f42e7733ce7a6e9bd6fdd6a59ef1b61f6cd1f4a89dd53dfce99da7bef", size = 42198, upload-time = "2026-08-24T21:16:57.52Z" },
|
||||
]
|
||||
|
||||
[package.optional-dependencies]
|
||||
aiohttp = [
|
||||
{ name = "aiohttp", marker = "python_full_version >= '3.12'" },
|
||||
{ name = "yarl", marker = "python_full_version >= '3.12'" },
|
||||
]
|
||||
awscrt = [
|
||||
{ name = "awscrt", marker = "python_full_version >= '3.12'" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "smithy-json"
|
||||
version = "0.2.3"
|
||||
version = "0.3.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "ijson", marker = "python_full_version >= '3.12'" },
|
||||
{ name = "smithy-core", marker = "python_full_version >= '3.12'" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/b9/6c/418b5687d8933b7a135d5e1a98c61fe814b98f72517dbae0e666860cb876/smithy_json-0.2.3.tar.gz", hash = "sha256:686e9b55a36dacb08e472732b358573ef78009055e05e9fce2e806d61490b2b3", size = 7805, upload-time = "2026-06-23T04:04:47.71Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/c7/ac/04164eefb3da7479f52f6535b4b39cc8384c292cb2bb74279f2acc4f4b4d/smithy_json-0.3.0.tar.gz", hash = "sha256:c81c7034587e01bc64767cbbecb05a7d65ca9070612fd94e8a03e80540290a22", size = 7956, upload-time = "2026-08-20T17:55:32.177Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/3c/14/eabb26b355415bcd9feef27fb5b18f1dad3fabd4208cfcbaf152025fa9ae/smithy_json-0.2.3-py3-none-any.whl", hash = "sha256:594e1bbe3d480963237f8fd0fc648dbd4e988b4503fea90157b5f07706796327", size = 10252, upload-time = "2026-06-23T04:04:48.46Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/9d/cf/0104c40a0e18fa307ea3da4310eba949f474a5bc1df3cc2b5851a72e8486/smithy_json-0.3.0-py3-none-any.whl", hash = "sha256:ffb73d2e60cf5e616e5d0a1019e7b9f518edba076cb423f10981457725dcddc4", size = 10252, upload-time = "2026-08-20T17:55:31.204Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue