From 7143aeeb86f36669ea08083e669aa0ccfa163e47 Mon Sep 17 00:00:00 2001 From: shrey Date: Sat, 31 Oct 2026 12:46:40 -0700 Subject: [PATCH] feat(passthrough): add Deepgram streaming text-to-speech WebSocket route (#49282) --- litellm/constants.py | 1 + litellm/llms/deepgram/common_utils.py | 77 ++++++- .../llm_passthrough_endpoints.py | 56 +++++ .../pass_through_endpoints.py | 4 +- .../deepgram/test_deepgram_common_utils.py | 105 +++++++++ .../test_deepgram_ws_passthrough_routes.py | 200 +++++++++++++++++- 6 files changed, 438 insertions(+), 5 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index c175e3541f6..cd7c0176428 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -432,6 +432,7 @@ WEBSOCKET_CLOSE_REASON_MAX_BYTES: Final = 123 DEEPGRAM_DEFAULT_API_BASE: Final = "https://api.deepgram.com/v1" NADIR_DEFAULT_API_BASE: Final = "https://api.getnadir.com/v1" DEEPGRAM_LISTEN_DEFAULT_MODEL: Final = "nova-3" +DEEPGRAM_SPEAK_DEFAULT_MODEL: Final = "aura-asteria-en" BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY: Final = "litellm.bedrock_realtime.pending_opening_message" BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY: Final = "litellm.bedrock_realtime.session_committed" diff --git a/litellm/llms/deepgram/common_utils.py b/litellm/llms/deepgram/common_utils.py index 29998895bbc..77fe83b551a 100644 --- a/litellm/llms/deepgram/common_utils.py +++ b/litellm/llms/deepgram/common_utils.py @@ -1,16 +1,21 @@ +import json import math +import re from collections.abc import Mapping, Sequence +from dataclasses import dataclass from types import MappingProxyType from typing import Final -from urllib.parse import parse_qs, urlparse +from urllib.parse import parse_qs, urlparse, urlsplit, urlunsplit import httpx import litellm -from litellm.constants import DEEPGRAM_DEFAULT_API_BASE, DEEPGRAM_LISTEN_DEFAULT_MODEL +from litellm.constants import DEEPGRAM_DEFAULT_API_BASE, DEEPGRAM_LISTEN_DEFAULT_MODEL, DEEPGRAM_SPEAK_DEFAULT_MODEL from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.types.passthrough_endpoints.pass_through_endpoints import WebSocketSessionUsage from litellm.types.utils import LlmProviders +from litellm.utils import ModelInfoLookupMiss, lookup_model_info_for_registration _WEBSOCKET_SCHEMES: Final = MappingProxyType({"https": "wss", "http": "ws", "wss": "wss", "ws": "ws"}) DEEPGRAM_LISTEN_CALLBACK_PARAMS: Final = frozenset({"callback", "callback_method"}) @@ -199,3 +204,71 @@ def deepgram_listen_transcript(websocket_messages: Sequence[Mapping[str, object] if frame.get("type") == "Results" if (transcript := _final_transcript(frame)) is not None ) + + +_AURA_GENERATION: Final = re.compile(r"aura(?:-\d+)?(?=-[a-z])") + + +def deepgram_speak_websocket_target(api_base: str | None, query_string: str) -> str: + speak_url: Final = urlsplit(join_api_base(api_base or DEEPGRAM_DEFAULT_API_BASE, "speak")) + scheme: Final = _WEBSOCKET_SCHEMES.get(speak_url.scheme, speak_url.scheme) + params: Final = _first_occurrences(query_string) + query: Final = params if params.get("model") else params.remove("model").add("model", DEEPGRAM_SPEAK_DEFAULT_MODEL) + return urlunsplit((scheme, speak_url.netloc, speak_url.path, str(query), "")) + + +def deepgram_speak_model(upstream_url: str) -> str: + models: Final = parse_qs(urlparse(upstream_url).query).get("model") + return models[0] if models else DEEPGRAM_SPEAK_DEFAULT_MODEL + + +def deepgram_speak_requested_model(query_string: str) -> str: + """Read off the upstream target, so the voice a key is authorized for is the voice Deepgram receives""" + return deepgram_speak_model(deepgram_speak_websocket_target(None, query_string)) + + +def deepgram_speak_registry_key(upstream_url: str) -> str: + """Deepgram bills every voice of an Aura generation at one rate, so ``aura-2-thalia-en`` is priced at the + ``deepgram/aura-2`` row""" + model: Final = deepgram_speak_model(upstream_url) + generation: Final = _AURA_GENERATION.match(model) + return f"{LlmProviders.DEEPGRAM.value}/{generation.group() if generation else model}" + + +def deepgram_speak_character_rate(upstream_url: str) -> float | None: + """Only an exact registry hit counts, so a voice is never billed at a neighbouring row's rate""" + registry_key: Final = deepgram_speak_registry_key(upstream_url) + model_info: Final = lookup_model_info_for_registration(registry_key, LlmProviders.DEEPGRAM.value) + if isinstance(model_info, ModelInfoLookupMiss) or model_info["key"] != registry_key: + return None + return model_info.get("input_cost_per_character") + + +@dataclass(frozen=True, slots=True) +class DeepgramSpeakMeter: + """Holds the rate checked when the socket opened, so a cost-map reload mid-session can't change or drop the price""" + + upstream_url: str + character_rate: float + + def measure(self, client_text: str) -> int: + """Counts a Speak message's text as it is sent. Deepgram also speaks, and bills, unflushed text on its own (a + long buffer, Close, a close frame), so counting only flushed text would let a client get audio free, while + counting on send can only over-bill text Deepgram discards, such as a Clear or an oversized Speak. Never + raises, so a malformed frame can't end the session""" + try: + message: Final[object] = json.loads(client_text) + except (ValueError, RecursionError): + return 0 + match message: + case {"type": "Speak", "text": str() as text}: + return len(text) + case _: + return 0 + + def price(self, total: int) -> WebSocketSessionUsage: + return WebSocketSessionUsage( + model=deepgram_speak_model(self.upstream_url), + custom_llm_provider=LlmProviders.DEEPGRAM.value, + response_cost=total * self.character_rate, + ) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index fa09b229564..f7f0cf2ab3b 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -64,11 +64,16 @@ from litellm.llms.azure.passthrough.transformation import ( from litellm.llms.bedrock.passthrough.transformation import is_bedrock_streaming_endpoint from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.llms.deepgram.common_utils import ( + DeepgramSpeakMeter, deepgram_listen_callback_params, deepgram_listen_is_priced, deepgram_listen_registry_key, deepgram_listen_requested_model, deepgram_listen_websocket_target, + deepgram_speak_character_rate, + deepgram_speak_registry_key, + deepgram_speak_requested_model, + deepgram_speak_websocket_target, ) from litellm.llms.fal_ai.cost_calculator import fal_ai_passthrough_cost, fal_ai_queue_base from litellm.llms.nvidia_nim.passthrough.transformation import nvidia_nim_model_group_in_path @@ -129,6 +134,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( YieldsToConfiguredPassThroughRoute, create_pass_through_route, create_websocket_passthrough_route, + truncated_close_reason, websocket_passthrough_request, ) from litellm.proxy.utils import ProxyLogging as ProxyLoggingType @@ -145,6 +151,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, + WebSocketSessionMeter, ) from litellm.types.passthrough_endpoints.tinyfish import ( TINYFISH_AUTHENTICATED_RUN_FIELDS, @@ -3460,6 +3467,7 @@ class _WebsocketRelay(Protocol): forward_headers: bool, endpoint: str, accept_websocket: bool, + meter: WebSocketSessionMeter | None = None, ) -> None: ... @@ -3621,6 +3629,54 @@ async def deepgram_listen_websocket_route( ) +async def deepgram_speak_user_api_key_auth(websocket: WebSocket) -> UserAPIKeyAuth: + return await user_api_key_auth_websocket_for_model( + websocket, model=deepgram_speak_requested_model(websocket.url.query) + ) + + +@router.websocket("/deepgram/v1/speak") +@router.websocket("/deepgram/speak") +async def deepgram_speak_websocket_route( + websocket: WebSocket, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(deepgram_speak_user_api_key_auth)], + relay: Annotated[_WebsocketRelay, Depends(_websocket_relay)], +) -> None: + deepgram_api_key: Final = passthrough_endpoint_router.get_credentials( + custom_llm_provider=litellm.LlmProviders.DEEPGRAM.value, + region_name=None, + ) + if deepgram_api_key is None: + await websocket.close(code=1011, reason=_DEEPGRAM_WS_MISSING_KEY_REASON) + return + + await websocket.accept(subprotocol=_negotiated_websocket_subprotocol(websocket)) + target: Final = deepgram_speak_websocket_target( + api_base=get_secret_str("DEEPGRAM_API_BASE"), + query_string=websocket.url.query, + ) + character_rate: Final = deepgram_speak_character_rate(target) + if character_rate is None: + await websocket.close( + code=1008, + reason=truncated_close_reason( + _DEEPGRAM_WS_UNPRICED_REASON.format(registry_key=deepgram_speak_registry_key(target)) + ), + ) + return + + await relay( + websocket=websocket, + target=target, + custom_headers={"Authorization": f"Token {deepgram_api_key}"}, + user_api_key_dict=user_api_key_dict, + forward_headers=False, + endpoint=websocket.url.path, + accept_websocket=False, + meter=DeepgramSpeakMeter(target, character_rate), + ) + + _OPENAI_BETA_HEADER: Final = "openai-beta" _CLIENT_OPENAI_BETA_RESOURCES: Final = frozenset({"agents", "vaults", "chatkit"}) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index bda4f4c2ad2..970cd466337 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -2584,7 +2584,7 @@ def _json_object_frame(frame: str | bytes) -> dict[str, object] | None: return decoded if isinstance(decoded, dict) else None -def _truncated_close_reason(reason: str) -> str: +def truncated_close_reason(reason: str) -> str: """ Fit a close reason inside the byte budget a WebSocket close frame allows, without splitting a character """ @@ -2958,7 +2958,7 @@ async def websocket_passthrough_request( if upstream_close is not None and _client_socket_is_open(websocket): await websocket.close( code=upstream_close.code, - reason=_truncated_close_reason(upstream_close.reason), + reason=truncated_close_reason(upstream_close.reason), ) end_time: Final = datetime.now() diff --git a/tests/unit/llms/deepgram/test_deepgram_common_utils.py b/tests/unit/llms/deepgram/test_deepgram_common_utils.py index 530888b70c4..4c1b4603076 100644 --- a/tests/unit/llms/deepgram/test_deepgram_common_utils.py +++ b/tests/unit/llms/deepgram/test_deepgram_common_utils.py @@ -7,6 +7,7 @@ import pytest import litellm from litellm.llms.deepgram.common_utils import ( + DeepgramSpeakMeter, deepgram_listen_addon_pricing_models, deepgram_listen_audio_seconds, deepgram_listen_callback_params, @@ -18,7 +19,12 @@ from litellm.llms.deepgram.common_utils import ( deepgram_listen_requested_model, deepgram_listen_transcript, deepgram_listen_websocket_target, + deepgram_speak_character_rate, + deepgram_speak_registry_key, + deepgram_speak_requested_model, + deepgram_speak_websocket_target, ) +from litellm.types.passthrough_endpoints.pass_through_endpoints import WebSocketSessionUsage NOVA_3_URL: Final = "wss://api.deepgram.com/v1/listen?model=nova-3&encoding=linear16&sample_rate=16000" @@ -324,3 +330,102 @@ def test_deepgram_listen_is_priced( ) def test_deepgram_listen_addon_pricing_models(upstream_url: str, expected: tuple[str, ...]): assert deepgram_listen_addon_pricing_models(upstream_url) == expected + + +@pytest.mark.parametrize( + ("api_base", "query_string", "expected"), + [ + pytest.param( + None, + "encoding=linear16", + "wss://api.deepgram.com/v1/speak?encoding=linear16&model=aura-asteria-en", + id="no model: Deepgram's own default, named so billing and upstream agree", + ), + pytest.param( + None, + "model=aura-asteria-en&model=aura-2-thalia-en", + "wss://api.deepgram.com/v1/speak?model=aura-asteria-en", + id="second model dropped: Deepgram honours the last, auth and pricing read the first", + ), + pytest.param( + "http://localhost:8080/v1", + "model=aura-2-thalia-en", + "ws://localhost:8080/v1/speak?model=aura-2-thalia-en", + id="server-configured api base", + ), + ], +) +def test_deepgram_speak_websocket_target(api_base: str | None, query_string: str, expected: str): + assert deepgram_speak_websocket_target(api_base, query_string) == expected + + +@pytest.mark.parametrize( + "query_string", + [ + "", + "model=", + "model=&model=aura-2-thalia-en", + "model=aura-2-thalia-en", + "model=aura-asteria-en&model=aura-2-thalia-en", + "encoding=mulaw", + ], +) +def test_deepgram_speak_requested_voice_is_the_only_voice_the_upstream_target_carries(query_string: str): + """Authorization runs against the requested voice and the upstream URL is built separately, so they must agree or + a key could be authorized for one voice and reach another""" + target: Final = deepgram_speak_websocket_target(None, query_string) + assert parse_qs(urlparse(target).query)["model"] == [deepgram_speak_requested_model(query_string)] + + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize( + ("model", "registry_key", "priced"), + [ + ("aura-asteria-en", "deepgram/aura", True), + ("aura-2-thalia-en", "deepgram/aura-2", True), + ("aura-2-celeste-es", "deepgram/aura-2", True), + ("aura-99-x-en", "deepgram/aura-99", False), + ("aura-2", "deepgram/aura-2", True), + ("nova-3", "deepgram/nova-3", False), + ], +) +def test_deepgram_speak_prices_each_voice_at_its_generation_row_only(model: str, registry_key: str, priced: bool): + """A voice is never billed at another generation's rate, and a listen model's row is not a speak price""" + upstream_url: Final = f"wss://api.deepgram.com/v1/speak?model={model}" + + assert deepgram_speak_registry_key(upstream_url) == registry_key + assert deepgram_speak_character_rate(upstream_url) == ( + litellm.model_cost[registry_key]["input_cost_per_character"] if priced else None + ) + + +@pytest.mark.parametrize( + ("frame", "characters"), + [ + ('{"type": "Speak", "text": "Hello world."}', 12), + ('{"type":"Speak","text":"¿qué tal?"}', 9), + ('{"type": "Flush"}', 0), + ('{"type": "Speak"}', 0), + ('{"type": "Speak", "text": 12}', 0), + ('["Speak", "Hello"]', 0), + ("Hello", 0), + ("[" * 100_000, 0), + ], +) +def test_deepgram_speak_meter_measures_only_speak_text_and_never_raises(frame: str, characters: int): + meter: Final = DeepgramSpeakMeter("wss://api.deepgram.com/v1/speak?model=aura-2-thalia-en", character_rate=1.0) + + assert meter.measure(frame) == characters + + +@pytest.mark.parametrize( + ("upstream_url", "model"), + [ + ("wss://api.deepgram.com/v1/speak?model=aura-2-thalia-en", "aura-2-thalia-en"), + ("wss://api.deepgram.com/v1/speak?encoding=linear16&model=aura-asteria-en", "aura-asteria-en"), + ], +) +def test_deepgram_speak_meter_prices_the_total_at_the_rate_it_holds(upstream_url: str, model: str) -> None: + assert DeepgramSpeakMeter(upstream_url, character_rate=0.25).price(47) == WebSocketSessionUsage( + model=model, custom_llm_provider="deepgram", response_cost=47 * 0.25 + ) diff --git a/tests/unit/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py b/tests/unit/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py index 404ed0f5800..b21a5e6c337 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py +++ b/tests/unit/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py @@ -1,4 +1,5 @@ -"""Deepgram ``/v1/listen`` passthrough WebSocket route: registration, auth, credential injection, target URL.""" +"""Deepgram ``/v1/listen`` and ``/v1/speak`` passthrough WebSocket routes: registration, auth, credential injection, +target URL.""" import asyncio from collections.abc import Mapping @@ -15,21 +16,27 @@ from starlette.websockets import WebSocketDisconnect import litellm from litellm.caching.dual_cache import DualCache +from litellm.constants import WEBSOCKET_CLOSE_REASON_MAX_BYTES +from litellm.litellm_core_utils.get_model_cost_map import adopt_model_cost_map +from litellm.llms.deepgram.common_utils import DeepgramSpeakMeter from litellm.proxy._lazy_features import LAZY_FEATURES from litellm.proxy._types import LiteLLMRoutes, UserAPIKeyAuth from litellm.proxy.auth.auth_checks import cache_key_object from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( _websocket_relay, deepgram_listen_websocket_route, + deepgram_speak_websocket_route, router, ) from litellm.proxy.utils import hash_token +from litellm.types.passthrough_endpoints.pass_through_endpoints import WebSocketSessionMeter GET_CREDENTIALS: Final = ( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials" ) USER_API_KEY_AUTH: Final = "litellm.proxy.auth.user_api_key_auth.user_api_key_auth" LISTEN_PATHS: Final = ("/deepgram/v1/listen", "/deepgram/listen") +SPEAK_PATHS: Final = ("/deepgram/v1/speak", "/deepgram/speak") NOVA_2_STREAMING_KEY: Final = "deepgram/streaming/nova-2" pytestmark: Final = pytest.mark.usefixtures("local_model_cost_map") @@ -62,6 +69,7 @@ class _RelayCall: forward_headers: bool endpoint: str accept_websocket: bool + meter: WebSocketSessionMeter | None = None class _FakeRelay: @@ -78,6 +86,7 @@ class _FakeRelay: forward_headers: bool, endpoint: str, accept_websocket: bool, + meter: WebSocketSessionMeter | None = None, ) -> None: self.calls.append( _RelayCall( @@ -87,6 +96,7 @@ class _FakeRelay: forward_headers=forward_headers, endpoint=endpoint, accept_websocket=accept_websocket, + meter=meter, ) ) @@ -472,3 +482,191 @@ def test_deepgram_listen_echoes_the_browser_subprotocol_that_carries_the_litellm MappingProxyType({"Authorization": "Token dg-provider-key"}) ] assert [call.forward_headers for call in relay.calls] == [False] + + +async def _serve_speak(websocket: _FakeWebSocket, user_api_key_dict: UserAPIKeyAuth | None = None) -> _FakeRelay: + relay: Final = _FakeRelay() + await deepgram_speak_websocket_route( + websocket=websocket, + user_api_key_dict=user_api_key_dict or UserAPIKeyAuth(), + relay=relay, + ) + return relay + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", SPEAK_PATHS) +async def test_deepgram_speak_relays_with_provider_auth_and_meters_the_session( + path: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False) + monkeypatch.setenv("DEEPGRAM_API_KEY", "dg-provider-key") + websocket: Final = _FakeWebSocket(path, "encoding=linear16&sample_rate=24000") + caller: Final = UserAPIKeyAuth(api_key="sk-litellm-virtual", team_id="team-tts") + target: Final = "wss://api.deepgram.com/v1/speak?encoding=linear16&sample_rate=24000&model=aura-asteria-en" + + relay: Final = await _serve_speak(websocket, caller) + + assert relay.calls == [ + _RelayCall( + target=target, + custom_headers=MappingProxyType({"Authorization": "Token dg-provider-key"}), + user_api_key_dict=caller, + forward_headers=False, + endpoint=path, + accept_websocket=False, + meter=DeepgramSpeakMeter(target, litellm.model_cost["deepgram/aura"]["input_cost_per_character"]), + ) + ] + assert websocket.accepts == [None] + assert websocket.closed is None + + +@pytest.mark.asyncio +async def test_deepgram_speak_bills_at_the_connect_time_rate_after_a_cost_map_reload_drops_the_row( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False) + monkeypatch.setenv("DEEPGRAM_API_KEY", "dg-provider-key") + rate: Final = litellm.model_cost["deepgram/aura-2"]["input_cost_per_character"] + + relay: Final = await _serve_speak(_FakeWebSocket("/deepgram/v1/speak", "model=aura-2-thalia-en")) + meters: Final = [call.meter for call in relay.calls if isinstance(call.meter, DeepgramSpeakMeter)] + adopt_model_cost_map({key: row for key, row in litellm.model_cost.items() if key != "deepgram/aura-2"}) + + assert [meter.price(47).response_cost for meter in meters] == [47 * rate] + + +@pytest.mark.asyncio +async def test_deepgram_speak_honours_server_configured_api_base(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("DEEPGRAM_API_BASE", "http://localhost:8080/v1") + monkeypatch.setenv("DEEPGRAM_API_KEY", "dg-provider-key") + + relay: Final = await _serve_speak(_FakeWebSocket("/deepgram/v1/speak", "model=aura-2-thalia-en")) + + assert [call.target for call in relay.calls] == ["ws://localhost:8080/v1/speak?model=aura-2-thalia-en"] + + +@pytest.mark.asyncio +async def test_deepgram_speak_closes_cleanly_when_provider_credentials_missing( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("DEEPGRAM_API_KEY", raising=False) + websocket: Final = _FakeWebSocket("/deepgram/v1/speak", "model=aura-2-thalia-en") + + relay: Final = await _serve_speak(websocket) + + assert relay.calls == [] + assert websocket.accepts == [] + assert websocket.closed == ( + 1011, + "Required 'DEEPGRAM_API_KEY' in environment to make pass-through calls to Deepgram.", + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("query", "missing_key"), + [ + pytest.param("model=aura-99-x-en", "deepgram/aura-99", id="generation without a row"), + pytest.param("model=nova-3", "deepgram/nova-3", id="listen model has no per-character rate"), + ], +) +async def test_deepgram_speak_refuses_voices_it_cannot_price( + query: str, missing_key: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False) + monkeypatch.setenv("DEEPGRAM_API_KEY", "dg-provider-key") + websocket: Final = _FakeWebSocket("/deepgram/v1/speak", query) + + relay: Final = await _serve_speak(websocket) + + assert relay.calls == [] + assert websocket.closed == ( + 1008, + f"No streaming price for '{missing_key}': add it to the model cost map to enable it", + ) + + +@pytest.mark.asyncio +async def test_deepgram_speak_refusal_for_a_long_unpriced_model_fits_a_close_frame( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False) + monkeypatch.setenv("DEEPGRAM_API_KEY", "dg-provider-key") + websocket: Final = _FakeWebSocket("/deepgram/v1/speak", f"model={'x' * 200}") + + relay: Final = await _serve_speak(websocket) + + assert relay.calls == [] + assert websocket.closed is not None + assert (websocket.closed[0], len(websocket.closed[1].encode())) == (1008, WEBSOCKET_CLOSE_REASON_MAX_BYTES) + + +@pytest.mark.parametrize("path", SPEAK_PATHS) +@pytest.mark.parametrize( + "query", + [ + pytest.param("model=aura-2-thalia-en", id="allowed voice"), + pytest.param("model=aura-2-thalia-en&model=aura-asteria-en", id="second voice stripped"), + ], +) +def test_deepgram_speak_relays_only_the_voice_the_key_was_checked_against( + path: str, query: str, monkeypatch: pytest.MonkeyPatch +) -> None: + """A key allowed only one voice must not reach another by repeating ``model``: the real key auth path must see + the voice the upstream target will carry""" + monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False) + monkeypatch.setenv("DEEPGRAM_API_KEY", "dg-provider-key") + monkeypatch.setattr(litellm, "max_budget", 0.0) + cache: Final = asyncio.run(_cache_restricted_key("sk-only-thalia", ["aura-2-thalia-en"])) + relay: Final = _FakeRelay() + client: Final = TestClient(_app_with_relay(relay)) + + with patch.multiple( # test-quality-ok: the real key auth path reads these proxy_server globals and has no injection seam + "litellm.proxy.proxy_server", + master_key="sk-master", + prisma_client=MagicMock(), + user_api_key_cache=cache, + llm_model_list=None, + llm_router=None, + ): + with client.websocket_connect(f"{path}?{query}", headers={"Authorization": "Bearer sk-only-thalia"}): + pass + + assert [(call.target, call.endpoint) for call in relay.calls] == [ + ("wss://api.deepgram.com/v1/speak?model=aura-2-thalia-en", path) + ] + + +@pytest.mark.parametrize( + "query", + [ + pytest.param("model=aura-asteria-en", id="denied voice named"), + pytest.param("", id="voice omitted, default denied"), + ], +) +def test_deepgram_speak_refuses_a_voice_the_key_is_not_allowed(query: str, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False) + monkeypatch.setenv("DEEPGRAM_API_KEY", "dg-provider-key") + monkeypatch.setattr(litellm, "max_budget", 0.0) + cache: Final = asyncio.run(_cache_restricted_key("sk-only-thalia", ["aura-2-thalia-en"])) + relay: Final = _FakeRelay() + client: Final = TestClient(_app_with_relay(relay)) + + with patch.multiple( # test-quality-ok: the real key auth path reads these proxy_server globals and has no injection seam + "litellm.proxy.proxy_server", + master_key="sk-master", + prisma_client=MagicMock(), + user_api_key_cache=cache, + llm_model_list=None, + llm_router=None, + ): + with pytest.raises(WebSocketDisconnect) as disconnect: + with client.websocket_connect( + f"/deepgram/v1/speak?{query}", headers={"Authorization": "Bearer sk-only-thalia"} + ): + pass + + assert disconnect.value.code == 1008 + assert relay.calls == []