mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
Merge pull request #41554 from BerriAI/litellm_deepgram_listen_websocket_passthrough
feat(passthrough): deepgram streaming /v1/listen WebSocket passthrough with duration-based cost tracking
This commit is contained in:
Yassin Kortam
GitHub
commit
52d6aab421
18 changed files with 1886 additions and 91 deletions
|
|
@ -95,6 +95,7 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
|
|||
"/vertex-ai/",
|
||||
"/assemblyai/",
|
||||
"/eu.assemblyai/",
|
||||
"/deepgram/",
|
||||
"/langfuse/",
|
||||
"/vllm/",
|
||||
"/mistral/",
|
||||
|
|
|
|||
|
|
@ -320,6 +320,9 @@ REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS: Final = float(
|
|||
# RFC 6455 caps the close frame payload at 125 bytes, 2 of which carry the status code
|
||||
WEBSOCKET_CLOSE_REASON_MAX_BYTES: Final = 123
|
||||
|
||||
DEEPGRAM_DEFAULT_API_BASE: Final = "https://api.deepgram.com/v1"
|
||||
DEEPGRAM_LISTEN_DEFAULT_MODEL: Final = "nova-3"
|
||||
|
||||
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"
|
||||
|
|
|
|||
|
|
@ -1,5 +1,200 @@
|
|||
import math
|
||||
from collections.abc import Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from urllib.parse import parse_qs, urlparse
|
||||
|
||||
import httpx
|
||||
|
||||
import litellm
|
||||
from litellm.constants import DEEPGRAM_DEFAULT_API_BASE, DEEPGRAM_LISTEN_DEFAULT_MODEL
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.types.utils import LlmProviders
|
||||
|
||||
_WEBSOCKET_SCHEMES: Final = MappingProxyType({"https": "wss", "http": "ws", "wss": "wss", "ws": "ws"})
|
||||
DEEPGRAM_LISTEN_CALLBACK_PARAMS: Final = frozenset({"callback", "callback_method"})
|
||||
DEEPGRAM_LISTEN_STREAMING_PRICING_PREFIX: Final = "streaming/"
|
||||
DEEPGRAM_LISTEN_MULTILINGUAL_LANGUAGE: Final = "multi"
|
||||
DEEPGRAM_LISTEN_MULTILINGUAL_PRICING_SUFFIX: Final = "-multilingual"
|
||||
DEEPGRAM_LISTEN_ADDON_PRICING_PARAMS: Final = MappingProxyType(
|
||||
{
|
||||
"redact": "redact",
|
||||
"keyterm": "keyterm",
|
||||
"detect_entities": "detect_entities",
|
||||
"diarize": "diarize",
|
||||
"diarize_model": "diarize",
|
||||
}
|
||||
)
|
||||
_DISABLED_PARAM_VALUES: Final = frozenset({"", "false"})
|
||||
_SINGLE_VALUED_PARAMS: Final = frozenset({"model", "language"})
|
||||
|
||||
|
||||
class DeepgramException(BaseLLMException):
|
||||
pass
|
||||
|
||||
|
||||
def deepgram_listen_requested_model(query_string: str) -> str:
|
||||
return httpx.QueryParams(query_string).get("model") or DEEPGRAM_LISTEN_DEFAULT_MODEL
|
||||
|
||||
|
||||
def _first_occurrences(query_string: str) -> httpx.QueryParams:
|
||||
"""Authorization and pricing read the first ``model`` and ``language`` value; Deepgram must not see a second one."""
|
||||
items: Final = httpx.QueryParams(query_string).multi_items()
|
||||
return httpx.QueryParams(
|
||||
tuple(
|
||||
(key, value)
|
||||
for index, (key, value) in enumerate(items)
|
||||
if key not in _SINGLE_VALUED_PARAMS or all(earlier != key for earlier, _ in items[:index])
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def deepgram_listen_websocket_target(api_base: str | None, query_string: str) -> str:
|
||||
listen_url: Final = httpx.URL(f"{(api_base or DEEPGRAM_DEFAULT_API_BASE).rstrip('/')}/listen")
|
||||
websocket_url: Final = listen_url.copy_with(scheme=_WEBSOCKET_SCHEMES.get(listen_url.scheme, listen_url.scheme))
|
||||
params: Final = _first_occurrences(query_string)
|
||||
query: Final = params if params.get("model") else params.remove("model").add("model", DEEPGRAM_LISTEN_DEFAULT_MODEL)
|
||||
return f"{websocket_url}?{query}"
|
||||
|
||||
|
||||
def deepgram_listen_callback_params(query_string: str) -> tuple[str, ...]:
|
||||
return tuple(sorted(DEEPGRAM_LISTEN_CALLBACK_PARAMS.intersection(httpx.QueryParams(query_string).keys())))
|
||||
|
||||
|
||||
def deepgram_listen_model(upstream_url: str) -> str:
|
||||
models: Final = parse_qs(urlparse(upstream_url).query).get("model")
|
||||
return models[0] if models else DEEPGRAM_LISTEN_DEFAULT_MODEL
|
||||
|
||||
|
||||
def _param_enabled(values: Sequence[str]) -> bool:
|
||||
return any(value.strip().lower() not in _DISABLED_PARAM_VALUES for value in values)
|
||||
|
||||
|
||||
def deepgram_listen_pricing_model(upstream_url: str) -> str:
|
||||
"""Registry key, without the provider prefix, for the per-second base rate Deepgram bills a streaming session at:
|
||||
the multilingual streaming entry when ``language=multi``, otherwise the model's own streaming entry. Pre-recorded
|
||||
entries are never a substitute: Deepgram prices the two products differently."""
|
||||
streaming: Final = f"{DEEPGRAM_LISTEN_STREAMING_PRICING_PREFIX}{deepgram_listen_model(upstream_url)}"
|
||||
language: Final = parse_qs(urlparse(upstream_url).query).get("language", ("",))[0]
|
||||
if language.strip().lower() == DEEPGRAM_LISTEN_MULTILINGUAL_LANGUAGE:
|
||||
return f"{streaming}{DEEPGRAM_LISTEN_MULTILINGUAL_PRICING_SUFFIX}"
|
||||
return streaming
|
||||
|
||||
|
||||
def deepgram_listen_registry_key(upstream_url: str) -> str:
|
||||
return f"{LlmProviders.DEEPGRAM.value}/{deepgram_listen_pricing_model(upstream_url)}"
|
||||
|
||||
|
||||
def deepgram_listen_is_priced(upstream_url: str) -> bool:
|
||||
"""Only an exact registry hit counts: the cost calculator resolves a missing ``streaming/<model>`` row to the
|
||||
pre-recorded ``<model>`` row, which is not the rate Deepgram bills a WebSocket session at."""
|
||||
registry_key: Final = deepgram_listen_registry_key(upstream_url)
|
||||
try:
|
||||
model_info: Final = litellm.get_model_info(model=registry_key, custom_llm_provider=LlmProviders.DEEPGRAM.value)
|
||||
except Exception:
|
||||
return False
|
||||
return model_info["key"] == registry_key
|
||||
|
||||
|
||||
def deepgram_listen_addon_pricing_models(upstream_url: str) -> tuple[str, ...]:
|
||||
params: Final = parse_qs(urlparse(upstream_url).query)
|
||||
return tuple(
|
||||
sorted(
|
||||
frozenset(
|
||||
f"{DEEPGRAM_LISTEN_STREAMING_PRICING_PREFIX}{addon}"
|
||||
for param, addon in DEEPGRAM_LISTEN_ADDON_PRICING_PARAMS.items()
|
||||
if _param_enabled(params.get(param, ()))
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _channel_count(value: object) -> int | None:
|
||||
if isinstance(value, bool) or not isinstance(value, int):
|
||||
return None
|
||||
return value if value >= 1 else None
|
||||
|
||||
|
||||
def _results_channel_count(frame: Mapping[str, object]) -> int | None:
|
||||
channel_index: Final = frame.get("channel_index")
|
||||
if not isinstance(channel_index, list) or len(channel_index) != 2:
|
||||
return None
|
||||
return _channel_count(channel_index[1])
|
||||
|
||||
|
||||
def _declared_channel_count(upstream_url: str) -> int | None:
|
||||
declared: Final = parse_qs(urlparse(upstream_url).query).get("channels")
|
||||
if not declared or not declared[0].isdigit():
|
||||
return None
|
||||
return _channel_count(int(declared[0]))
|
||||
|
||||
|
||||
def deepgram_listen_channel_count(websocket_messages: Sequence[Mapping[str, object]], upstream_url: str) -> int:
|
||||
metadata_channels: Final = tuple(
|
||||
channels
|
||||
for frame in websocket_messages
|
||||
if frame.get("type") == "Metadata"
|
||||
if (channels := _channel_count(frame.get("channels"))) is not None
|
||||
)
|
||||
if metadata_channels:
|
||||
return metadata_channels[-1]
|
||||
results_channels: Final = tuple(
|
||||
channels
|
||||
for frame in websocket_messages
|
||||
if frame.get("type") == "Results"
|
||||
if (channels := _results_channel_count(frame)) is not None
|
||||
)
|
||||
if results_channels:
|
||||
return max(results_channels)
|
||||
return _declared_channel_count(upstream_url) or 1
|
||||
|
||||
|
||||
def _seconds(value: object) -> float | None:
|
||||
if isinstance(value, bool) or not isinstance(value, (int, float)):
|
||||
return None
|
||||
return float(value) if math.isfinite(value) and value >= 0 else None
|
||||
|
||||
|
||||
def _results_frame_end(frame: Mapping[str, object]) -> float | None:
|
||||
start: Final = _seconds(frame.get("start"))
|
||||
duration: Final = _seconds(frame.get("duration"))
|
||||
return None if start is None or duration is None else start + duration
|
||||
|
||||
|
||||
def _final_transcript(frame: Mapping[str, object]) -> str | None:
|
||||
if frame.get("is_final") is not True:
|
||||
return None
|
||||
channel: Final = frame.get("channel")
|
||||
alternatives: Final = channel.get("alternatives") if isinstance(channel, Mapping) else None
|
||||
first: Final = alternatives[0] if isinstance(alternatives, list) and alternatives else None
|
||||
transcript: Final = first.get("transcript") if isinstance(first, Mapping) else None
|
||||
return transcript if isinstance(transcript, str) and transcript else None
|
||||
|
||||
|
||||
def deepgram_listen_audio_seconds(websocket_messages: Sequence[Mapping[str, object]]) -> float:
|
||||
metadata_durations: Final = tuple(
|
||||
duration
|
||||
for frame in websocket_messages
|
||||
if frame.get("type") == "Metadata"
|
||||
if (duration := _seconds(frame.get("duration"))) is not None and duration > 0
|
||||
)
|
||||
if metadata_durations:
|
||||
return metadata_durations[-1]
|
||||
return max(
|
||||
(
|
||||
end
|
||||
for frame in websocket_messages
|
||||
if frame.get("type") == "Results"
|
||||
if (end := _results_frame_end(frame)) is not None
|
||||
),
|
||||
default=0.0,
|
||||
)
|
||||
|
||||
|
||||
def deepgram_listen_transcript(websocket_messages: Sequence[Mapping[str, object]]) -> str:
|
||||
return " ".join(
|
||||
transcript
|
||||
for frame in websocket_messages
|
||||
if frame.get("type") == "Results"
|
||||
if (transcript := _final_transcript(frame)) is not None
|
||||
)
|
||||
|
|
|
|||
|
|
@ -20853,6 +20853,96 @@
|
|||
"/v1/audio/transcriptions"
|
||||
]
|
||||
},
|
||||
"deepgram/streaming/nova-3": {
|
||||
"input_cost_per_second": 8e-05,
|
||||
"litellm_provider": "deepgram",
|
||||
"metadata": {
|
||||
"calculation": "$0.0048/60 seconds = $0.00008000 per second",
|
||||
"note": "Nova-3 monolingual streaming, pay as you go",
|
||||
"original_pricing_per_minute": 0.0048
|
||||
},
|
||||
"mode": "audio_transcription",
|
||||
"output_cost_per_second": 0.0,
|
||||
"source": "https://deepgram.com/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/listen"
|
||||
]
|
||||
},
|
||||
"deepgram/streaming/nova-3-multilingual": {
|
||||
"input_cost_per_second": 9.667e-05,
|
||||
"litellm_provider": "deepgram",
|
||||
"metadata": {
|
||||
"calculation": "$0.0058/60 seconds = $0.00009667 per second",
|
||||
"note": "Nova-3 multilingual (language=multi) streaming, pay as you go",
|
||||
"original_pricing_per_minute": 0.0058
|
||||
},
|
||||
"mode": "audio_transcription",
|
||||
"output_cost_per_second": 0.0,
|
||||
"source": "https://deepgram.com/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/listen"
|
||||
]
|
||||
},
|
||||
"deepgram/streaming/redact": {
|
||||
"input_cost_per_second": 3.333e-05,
|
||||
"litellm_provider": "deepgram",
|
||||
"metadata": {
|
||||
"calculation": "$0.0020/60 seconds = $0.00003333 per second",
|
||||
"note": "Redaction add-on (redact query param), streaming, pay as you go",
|
||||
"original_pricing_per_minute": 0.002
|
||||
},
|
||||
"mode": "audio_transcription",
|
||||
"output_cost_per_second": 0.0,
|
||||
"source": "https://deepgram.com/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/listen"
|
||||
]
|
||||
},
|
||||
"deepgram/streaming/keyterm": {
|
||||
"input_cost_per_second": 2.167e-05,
|
||||
"litellm_provider": "deepgram",
|
||||
"metadata": {
|
||||
"calculation": "$0.0013/60 seconds = $0.00002167 per second",
|
||||
"note": "Keyterm Prompting add-on (keyterm query param), streaming, pay as you go",
|
||||
"original_pricing_per_minute": 0.0013
|
||||
},
|
||||
"mode": "audio_transcription",
|
||||
"output_cost_per_second": 0.0,
|
||||
"source": "https://deepgram.com/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/listen"
|
||||
]
|
||||
},
|
||||
"deepgram/streaming/detect_entities": {
|
||||
"input_cost_per_second": 2.833e-05,
|
||||
"litellm_provider": "deepgram",
|
||||
"metadata": {
|
||||
"calculation": "$0.0017/60 seconds = $0.00002833 per second",
|
||||
"note": "Entity Detection add-on (detect_entities query param), streaming, pay as you go",
|
||||
"original_pricing_per_minute": 0.0017
|
||||
},
|
||||
"mode": "audio_transcription",
|
||||
"output_cost_per_second": 0.0,
|
||||
"source": "https://deepgram.com/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/listen"
|
||||
]
|
||||
},
|
||||
"deepgram/streaming/diarize": {
|
||||
"input_cost_per_second": 3.333e-05,
|
||||
"litellm_provider": "deepgram",
|
||||
"metadata": {
|
||||
"calculation": "$0.0020/60 seconds = $0.00003333 per second",
|
||||
"note": "Speaker Diarization add-on (diarize / diarize_model query params), streaming, pay as you go",
|
||||
"original_pricing_per_minute": 0.002
|
||||
},
|
||||
"mode": "audio_transcription",
|
||||
"output_cost_per_second": 0.0,
|
||||
"source": "https://deepgram.com/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/listen"
|
||||
]
|
||||
},
|
||||
"deepgram/whisper": {
|
||||
"input_cost_per_second": 0.0001,
|
||||
"litellm_provider": "deepgram",
|
||||
|
|
|
|||
|
|
@ -201,6 +201,7 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = (
|
|||
"/cohere/",
|
||||
"/comprehendmedical",
|
||||
"/cursor/",
|
||||
"/deepgram/",
|
||||
"/eu.assemblyai/",
|
||||
"/gemini/",
|
||||
"/gigachat/",
|
||||
|
|
|
|||
|
|
@ -71,6 +71,7 @@ from litellm.types.utils import (
|
|||
StandardLoggingVectorStoreRequest,
|
||||
StandardPassThroughResponseObject,
|
||||
TextCompletionResponse,
|
||||
TranscriptionResponse,
|
||||
)
|
||||
from litellm.types.videos.main import VideoObject
|
||||
|
||||
|
|
@ -491,6 +492,7 @@ class LiteLLMRoutes(enum.Enum):
|
|||
"/gigachat",
|
||||
"/watsonx",
|
||||
"/nvidia_nim",
|
||||
"/deepgram",
|
||||
]
|
||||
|
||||
#########################################################
|
||||
|
|
@ -4772,6 +4774,7 @@ PassThroughEndpointLoggingResultValues = (
|
|||
| VideoObject
|
||||
| StandardPassThroughResponseObject
|
||||
| ResponsesAPIResponse
|
||||
| TranscriptionResponse
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -632,9 +632,11 @@ def _apply_budget_limits_to_end_user_params(
|
|||
verbose_proxy_logger.debug("Applied budget limits to end user %s", end_user_id)
|
||||
|
||||
|
||||
async def user_api_key_auth_websocket(websocket: WebSocket):
|
||||
# Accept the WebSocket connection
|
||||
async def user_api_key_auth_websocket(websocket: WebSocket) -> UserAPIKeyAuth:
|
||||
return await user_api_key_auth_websocket_for_model(websocket, model=websocket.query_params.get("model"))
|
||||
|
||||
|
||||
async def user_api_key_auth_websocket_for_model(websocket: WebSocket, model: str | None) -> UserAPIKeyAuth:
|
||||
ws_scope: Final = websocket.scope or {}
|
||||
scope_headers: Final = list(ws_scope.get("headers") or [])
|
||||
# ``get_request_route`` falls back to ``request.url.path`` when
|
||||
|
|
@ -654,10 +656,6 @@ async def user_api_key_auth_websocket(websocket: WebSocket):
|
|||
|
||||
request._url = websocket.url
|
||||
|
||||
query_params: Final = websocket.query_params
|
||||
|
||||
model: Final = query_params.get("model")
|
||||
|
||||
async def return_body():
|
||||
return _realtime_request_body(model)
|
||||
|
||||
|
|
|
|||
|
|
@ -45,6 +45,13 @@ from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
|
|||
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
|
||||
from litellm.llms.azure.passthrough.transformation import foreign_azure_deployment
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.llms.deepgram.common_utils import (
|
||||
deepgram_listen_callback_params,
|
||||
deepgram_listen_is_priced,
|
||||
deepgram_listen_registry_key,
|
||||
deepgram_listen_requested_model,
|
||||
deepgram_listen_websocket_target,
|
||||
)
|
||||
from litellm.llms.nvidia_nim.passthrough.transformation import nvidia_nim_model_group_in_path
|
||||
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
|
||||
from litellm.passthrough.main import AsyncPassthroughStreamingResponse
|
||||
|
|
@ -57,6 +64,7 @@ from litellm.proxy.auth.user_api_key_auth import (
|
|||
is_no_auth_dev_mode,
|
||||
user_api_key_auth,
|
||||
user_api_key_auth_websocket,
|
||||
user_api_key_auth_websocket_for_model,
|
||||
)
|
||||
from litellm.proxy.common_request_processing import open_sse_before_first_byte
|
||||
from litellm.proxy.common_utils.http_parsing_utils import (
|
||||
|
|
@ -2925,7 +2933,7 @@ async def _openai_websocket_refusal(
|
|||
return None
|
||||
|
||||
|
||||
class _OpenAIWebsocketRelay(Protocol):
|
||||
class _WebsocketRelay(Protocol):
|
||||
async def __call__(
|
||||
self,
|
||||
*,
|
||||
|
|
@ -2939,7 +2947,7 @@ class _OpenAIWebsocketRelay(Protocol):
|
|||
) -> None: ...
|
||||
|
||||
|
||||
def _openai_websocket_relay() -> _OpenAIWebsocketRelay:
|
||||
def _websocket_relay() -> _WebsocketRelay:
|
||||
return websocket_passthrough_request
|
||||
|
||||
|
||||
|
|
@ -2957,6 +2965,15 @@ def _proxy_model_allowlists() -> _OpenAIWebsocketModelAllowlists:
|
|||
return resolve
|
||||
|
||||
|
||||
def _negotiated_websocket_subprotocol(websocket: WebSocket) -> str | None:
|
||||
requested_subprotocols: Final = tuple(
|
||||
protocol.strip()
|
||||
for protocol in (websocket.headers.get("sec-websocket-protocol") or "").split(",")
|
||||
if protocol.strip()
|
||||
)
|
||||
return requested_subprotocols[0] if requested_subprotocols else None
|
||||
|
||||
|
||||
@router.websocket("/openai_passthrough/{endpoint:path}")
|
||||
@router.websocket("/openai/{endpoint:path}")
|
||||
async def openai_websocket_proxy_route(
|
||||
|
|
@ -2964,16 +2981,11 @@ async def openai_websocket_proxy_route(
|
|||
endpoint: str,
|
||||
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth_websocket)],
|
||||
general_settings: Annotated[Mapping[str, object], Depends(_proxy_general_settings)],
|
||||
relay: Annotated[_OpenAIWebsocketRelay, Depends(_openai_websocket_relay)],
|
||||
relay: Annotated[_WebsocketRelay, Depends(_websocket_relay)],
|
||||
model_allowlists: Annotated[_OpenAIWebsocketModelAllowlists, Depends(_proxy_model_allowlists)],
|
||||
) -> None:
|
||||
"""WebSocket passthrough for OpenAI prefixes (realtime / responses.connect)."""
|
||||
requested_subprotocols: Final = tuple(
|
||||
protocol.strip()
|
||||
for protocol in (websocket.headers.get("sec-websocket-protocol") or "").split(",")
|
||||
if protocol.strip()
|
||||
)
|
||||
negotiated_subprotocol: Final = requested_subprotocols[0] if requested_subprotocols else None
|
||||
negotiated_subprotocol: Final = _negotiated_websocket_subprotocol(websocket)
|
||||
|
||||
refusal: Final = await _openai_websocket_refusal(user_api_key_dict, general_settings, model_allowlists)
|
||||
if refusal is not None:
|
||||
|
|
@ -3032,6 +3044,69 @@ async def openai_websocket_proxy_route(
|
|||
)
|
||||
|
||||
|
||||
_DEEPGRAM_WS_MISSING_KEY_REASON: Final = (
|
||||
"Required 'DEEPGRAM_API_KEY' in environment to make pass-through calls to Deepgram."
|
||||
)
|
||||
_DEEPGRAM_WS_CALLBACK_REASON: Final = "Deepgram callback delivery is not supported through the proxy: remove {params}"
|
||||
_DEEPGRAM_WS_UNPRICED_REASON: Final = (
|
||||
"No streaming price for '{registry_key}': add it to the model cost map to enable it"
|
||||
)
|
||||
|
||||
|
||||
async def deepgram_listen_user_api_key_auth(websocket: WebSocket) -> UserAPIKeyAuth:
|
||||
return await user_api_key_auth_websocket_for_model(
|
||||
websocket, model=deepgram_listen_requested_model(websocket.url.query)
|
||||
)
|
||||
|
||||
|
||||
@router.websocket("/deepgram/v1/listen")
|
||||
@router.websocket("/deepgram/listen")
|
||||
async def deepgram_listen_websocket_route(
|
||||
websocket: WebSocket,
|
||||
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(deepgram_listen_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))
|
||||
callback_params: Final = deepgram_listen_callback_params(websocket.url.query)
|
||||
if callback_params:
|
||||
await websocket.close(
|
||||
code=1008,
|
||||
reason=_DEEPGRAM_WS_CALLBACK_REASON.format(params=", ".join(callback_params)),
|
||||
)
|
||||
return
|
||||
|
||||
target: Final = deepgram_listen_websocket_target(
|
||||
api_base=get_secret_str("DEEPGRAM_API_BASE"),
|
||||
query_string=websocket.url.query,
|
||||
)
|
||||
if not deepgram_listen_is_priced(target):
|
||||
await websocket.close(
|
||||
code=1008,
|
||||
reason=_DEEPGRAM_WS_UNPRICED_REASON.format(registry_key=deepgram_listen_registry_key(target)),
|
||||
)
|
||||
return
|
||||
|
||||
await relay(
|
||||
websocket=websocket,
|
||||
target=target,
|
||||
custom_headers={ # mutable-ok: websocket_passthrough_request requires a plain dict of upstream 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,
|
||||
)
|
||||
|
||||
|
||||
class BaseOpenAIPassThroughHandler:
|
||||
@staticmethod
|
||||
async def _base_openai_pass_through_handler(
|
||||
|
|
|
|||
|
|
@ -0,0 +1,96 @@
|
|||
from collections.abc import Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.deepgram.common_utils import (
|
||||
deepgram_listen_addon_pricing_models,
|
||||
deepgram_listen_audio_seconds,
|
||||
deepgram_listen_channel_count,
|
||||
deepgram_listen_is_priced,
|
||||
deepgram_listen_model,
|
||||
deepgram_listen_pricing_model,
|
||||
deepgram_listen_registry_key,
|
||||
deepgram_listen_transcript,
|
||||
)
|
||||
from litellm.proxy._types import PassThroughEndpointLoggingTypedDict
|
||||
from litellm.types.utils import TranscriptionResponse
|
||||
|
||||
DEEPGRAM_LISTEN_ROUTE_SUFFIX: Final = "/listen"
|
||||
|
||||
|
||||
def _registry_cost(response: TranscriptionResponse, pricing_model: str) -> float | None:
|
||||
try:
|
||||
return litellm.completion_cost(
|
||||
completion_response=response,
|
||||
model=pricing_model,
|
||||
custom_llm_provider=litellm.LlmProviders.DEEPGRAM.value,
|
||||
call_type="transcription",
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # an unpriced entry must not lose the spend row, only its cost
|
||||
verbose_proxy_logger.debug("Deepgram listen passthrough: no registry price for '%s': %s", pricing_model, e)
|
||||
return None
|
||||
|
||||
|
||||
def _audio_cost(response: TranscriptionResponse, upstream_url: str) -> float | None:
|
||||
if not deepgram_listen_is_priced(upstream_url):
|
||||
verbose_proxy_logger.warning(
|
||||
"Deepgram listen passthrough: no registry entry '%s'", deepgram_listen_registry_key(upstream_url)
|
||||
)
|
||||
return None
|
||||
base_cost: Final = _registry_cost(response, deepgram_listen_pricing_model(upstream_url))
|
||||
if base_cost is None:
|
||||
return None
|
||||
addon_costs: Final = tuple(
|
||||
_registry_cost(response, pricing_model) for pricing_model in deepgram_listen_addon_pricing_models(upstream_url)
|
||||
)
|
||||
return base_cost + sum(cost for cost in addon_costs if cost is not None)
|
||||
|
||||
|
||||
class DeepgramListenPassthroughLoggingHandler:
|
||||
@staticmethod
|
||||
def is_deepgram_listen_route(url_route: str) -> bool:
|
||||
path: Final = urlparse(url_route).path
|
||||
return "/deepgram/" in path and path.endswith(DEEPGRAM_LISTEN_ROUTE_SUFFIX)
|
||||
|
||||
def deepgram_listen_passthrough_handler(
|
||||
self,
|
||||
websocket_messages: Sequence[Mapping[str, object]],
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
upstream_url: str,
|
||||
kwargs: Mapping[str, object] = MappingProxyType({}),
|
||||
) -> PassThroughEndpointLoggingTypedDict:
|
||||
model: Final = deepgram_listen_model(upstream_url)
|
||||
audio_seconds: Final = deepgram_listen_audio_seconds(websocket_messages)
|
||||
channels: Final = deepgram_listen_channel_count(websocket_messages, upstream_url)
|
||||
billed_seconds: Final = audio_seconds * channels
|
||||
response: Final = TranscriptionResponse(text=deepgram_listen_transcript(websocket_messages))
|
||||
response._hidden_params["audio_transcription_duration"] = billed_seconds # pyright: ignore[reportPrivateUsage] # the cost calculator reads the billed duration off the response's hidden params
|
||||
response_cost: Final = _audio_cost(response, upstream_url)
|
||||
response._hidden_params["response_cost"] = response_cost # pyright: ignore[reportPrivateUsage] # the logger reads a precomputed cost off the response's hidden params
|
||||
|
||||
provider: Final = litellm.LlmProviders.DEEPGRAM.value
|
||||
logging_obj.model = model # rebind-ok: the spend logger reads model and cost off the shared logging object
|
||||
logging_obj.model_call_details["model"] = model # rebind-ok: same shared logging object
|
||||
logging_obj.model_call_details["custom_llm_provider"] = provider # rebind-ok: same shared logging object
|
||||
logging_obj.model_call_details["response_cost"] = response_cost # rebind-ok: same shared logging object
|
||||
verbose_proxy_logger.debug(
|
||||
"Deepgram listen passthrough cost tracking: model %s, audio seconds %s, channels %s, cost %s",
|
||||
model,
|
||||
audio_seconds,
|
||||
channels,
|
||||
response_cost,
|
||||
)
|
||||
logging_result: Final[PassThroughEndpointLoggingTypedDict] = {
|
||||
"result": response,
|
||||
"kwargs": {
|
||||
**kwargs,
|
||||
"model": model,
|
||||
"custom_llm_provider": provider,
|
||||
"response_cost": response_cost,
|
||||
},
|
||||
}
|
||||
return logging_result
|
||||
|
|
@ -8,7 +8,7 @@ from base64 import b64encode
|
|||
from collections.abc import AsyncGenerator, Callable, Iterable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from itertools import groupby
|
||||
from itertools import count, groupby
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, TypedDict, cast
|
||||
from urllib.parse import urlencode, urlparse
|
||||
|
|
@ -2122,6 +2122,14 @@ def _resolved_vertex_live_setup(
|
|||
return {**setup_data, "model": setup_model_rewriter(setup_model)}
|
||||
|
||||
|
||||
def _json_object_frame(frame: str | bytes) -> dict[str, object] | None:
|
||||
try:
|
||||
decoded: Final = json.loads(frame if isinstance(frame, str) else frame.decode("utf-8"))
|
||||
except (json.JSONDecodeError, UnicodeDecodeError):
|
||||
return None
|
||||
return decoded if isinstance(decoded, dict) else None
|
||||
|
||||
|
||||
def _truncated_close_reason(reason: str) -> str:
|
||||
"""
|
||||
Fit a close reason inside the byte budget a WebSocket close frame allows, without splitting a character
|
||||
|
|
@ -2403,70 +2411,41 @@ async def websocket_passthrough_request(
|
|||
)
|
||||
await upstream_ws.close()
|
||||
|
||||
def _extract_vertex_live_model_from_setup_response(setup_response: Mapping[str, object]) -> None:
|
||||
extracted_model: Final = _extract_model_from_vertex_ai_setup(setup_response)
|
||||
if not extracted_model:
|
||||
verbose_proxy_logger.warning(
|
||||
"WebSocket passthrough (%s): Failed to extract model from server setup response: %s",
|
||||
endpoint,
|
||||
setup_response,
|
||||
)
|
||||
return
|
||||
kwargs["model"] = extracted_model
|
||||
kwargs["custom_llm_provider"] = "vertex_ai_language_models"
|
||||
logging_obj.model = extracted_model
|
||||
logging_obj.model_call_details["model"] = extracted_model
|
||||
logging_obj.model_call_details["custom_llm_provider"] = "vertex_ai_language_models"
|
||||
|
||||
is_vertex_live: Final = bool(endpoint and "/vertex_ai/live" in endpoint)
|
||||
json_frame_ordinal: Final = count()
|
||||
|
||||
async def relay_upstream_frame(upstream_message: str | bytes) -> None:
|
||||
if isinstance(upstream_message, bytes):
|
||||
await websocket.send_bytes(upstream_message)
|
||||
else:
|
||||
await websocket.send_text(upstream_message)
|
||||
message_data: Final = _json_object_frame(upstream_message)
|
||||
if message_data is None:
|
||||
return
|
||||
if is_vertex_live and next(json_frame_ordinal) == 0:
|
||||
_extract_vertex_live_model_from_setup_response(message_data)
|
||||
return
|
||||
websocket_messages.append(message_data)
|
||||
|
||||
async def forward_upstream_to_client() -> Close | None:
|
||||
"""Forward messages from upstream to client WebSocket, returning the upstream's close frame"""
|
||||
try:
|
||||
# Wait for the first response from upstream
|
||||
raw_response = await upstream_ws.recv(decode=False)
|
||||
# Ensure raw_response is bytes before decoding
|
||||
if isinstance(raw_response, str):
|
||||
raw_response = raw_response.encode("utf-8")
|
||||
setup_response: Final[Mapping[str, object]] = json.loads(raw_response.decode("utf-8"))
|
||||
verbose_proxy_logger.debug("Setup response: %s", setup_response)
|
||||
|
||||
# Extract model and provider from setup response for Vertex AI Live
|
||||
if endpoint and "/vertex_ai/live" in endpoint:
|
||||
verbose_proxy_logger.debug(
|
||||
"WebSocket passthrough (%s): Processing server setup response for model extraction",
|
||||
endpoint,
|
||||
)
|
||||
extracted_model: Final = _extract_model_from_vertex_ai_setup(setup_response)
|
||||
if extracted_model:
|
||||
kwargs["model"] = extracted_model
|
||||
kwargs["custom_llm_provider"] = "vertex_ai_language_models"
|
||||
# Update logging object with correct model
|
||||
logging_obj.model = extracted_model
|
||||
logging_obj.model_call_details["model"] = extracted_model
|
||||
logging_obj.model_call_details["custom_llm_provider"] = "vertex_ai_language_models"
|
||||
verbose_proxy_logger.debug(
|
||||
"WebSocket passthrough (%s): Successfully extracted model '%s' and set provider to 'vertex_ai' from server setup response",
|
||||
endpoint,
|
||||
extracted_model,
|
||||
)
|
||||
else:
|
||||
verbose_proxy_logger.warning(
|
||||
"WebSocket passthrough (%s): Failed to extract model from server setup response: %s",
|
||||
endpoint,
|
||||
setup_response,
|
||||
)
|
||||
else:
|
||||
verbose_proxy_logger.debug(
|
||||
"WebSocket passthrough (%s): Not a Vertex AI Live endpoint, skipping model extraction",
|
||||
endpoint,
|
||||
)
|
||||
|
||||
# Send the setup response to the client
|
||||
await websocket.send_text(json.dumps(setup_response))
|
||||
|
||||
# Now continuously forward messages from upstream to client
|
||||
async for upstream_message in upstream_ws:
|
||||
if isinstance(upstream_message, bytes):
|
||||
await websocket.send_bytes(upstream_message)
|
||||
# Parse and collect for cost tracking
|
||||
try:
|
||||
message_data: dict[str, object] = json.loads(upstream_message.decode())
|
||||
websocket_messages.append(message_data)
|
||||
except (json.JSONDecodeError, UnicodeDecodeError):
|
||||
pass
|
||||
else:
|
||||
await websocket.send_text(upstream_message)
|
||||
# Parse and collect for cost tracking
|
||||
try:
|
||||
message_data = json.loads(upstream_message)
|
||||
websocket_messages.append(message_data)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
while True:
|
||||
await relay_upstream_frame(await upstream_ws.recv())
|
||||
except (ConnectionClosedOK, ConnectionClosedError) as e:
|
||||
verbose_proxy_logger.debug("Upstream WebSocket connection closed: %s", e)
|
||||
return e.rcvd
|
||||
|
|
|
|||
|
|
@ -26,6 +26,9 @@ from .llm_provider_handlers.cohere_passthrough_logging_handler import (
|
|||
from .llm_provider_handlers.cursor_passthrough_logging_handler import (
|
||||
CursorPassthroughLoggingHandler,
|
||||
)
|
||||
from .llm_provider_handlers.deepgram_listen_passthrough_logging_handler import (
|
||||
DeepgramListenPassthroughLoggingHandler,
|
||||
)
|
||||
from .llm_provider_handlers.gemini_passthrough_logging_handler import (
|
||||
GeminiPassthroughLoggingHandler,
|
||||
)
|
||||
|
|
@ -349,6 +352,21 @@ class PassThroughEndpointLogging:
|
|||
|
||||
standard_logging_response_object = vertex_ai_live_handler_result["result"]
|
||||
kwargs = vertex_ai_live_handler_result["kwargs"]
|
||||
elif DeepgramListenPassthroughLoggingHandler.is_deepgram_listen_route(url_route):
|
||||
deepgram_handler_result: Final = (
|
||||
DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler(
|
||||
websocket_messages=tuple(
|
||||
message
|
||||
for message in (response_body if isinstance(response_body, list) else ())
|
||||
if isinstance(message, dict)
|
||||
),
|
||||
logging_obj=logging_obj,
|
||||
upstream_url=str(httpx_response.request.url),
|
||||
kwargs=kwargs,
|
||||
)
|
||||
)
|
||||
standard_logging_response_object = deepgram_handler_result["result"] # rebind-ok: elif-chain
|
||||
kwargs = deepgram_handler_result["kwargs"] # rebind-ok: elif-chain contract
|
||||
return_dict["standard_logging_response_object"] = standard_logging_response_object
|
||||
|
||||
return_dict["kwargs"] = kwargs
|
||||
|
|
|
|||
|
|
@ -20853,6 +20853,96 @@
|
|||
"/v1/audio/transcriptions"
|
||||
]
|
||||
},
|
||||
"deepgram/streaming/nova-3": {
|
||||
"input_cost_per_second": 8e-05,
|
||||
"litellm_provider": "deepgram",
|
||||
"metadata": {
|
||||
"calculation": "$0.0048/60 seconds = $0.00008000 per second",
|
||||
"note": "Nova-3 monolingual streaming, pay as you go",
|
||||
"original_pricing_per_minute": 0.0048
|
||||
},
|
||||
"mode": "audio_transcription",
|
||||
"output_cost_per_second": 0.0,
|
||||
"source": "https://deepgram.com/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/listen"
|
||||
]
|
||||
},
|
||||
"deepgram/streaming/nova-3-multilingual": {
|
||||
"input_cost_per_second": 9.667e-05,
|
||||
"litellm_provider": "deepgram",
|
||||
"metadata": {
|
||||
"calculation": "$0.0058/60 seconds = $0.00009667 per second",
|
||||
"note": "Nova-3 multilingual (language=multi) streaming, pay as you go",
|
||||
"original_pricing_per_minute": 0.0058
|
||||
},
|
||||
"mode": "audio_transcription",
|
||||
"output_cost_per_second": 0.0,
|
||||
"source": "https://deepgram.com/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/listen"
|
||||
]
|
||||
},
|
||||
"deepgram/streaming/redact": {
|
||||
"input_cost_per_second": 3.333e-05,
|
||||
"litellm_provider": "deepgram",
|
||||
"metadata": {
|
||||
"calculation": "$0.0020/60 seconds = $0.00003333 per second",
|
||||
"note": "Redaction add-on (redact query param), streaming, pay as you go",
|
||||
"original_pricing_per_minute": 0.002
|
||||
},
|
||||
"mode": "audio_transcription",
|
||||
"output_cost_per_second": 0.0,
|
||||
"source": "https://deepgram.com/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/listen"
|
||||
]
|
||||
},
|
||||
"deepgram/streaming/keyterm": {
|
||||
"input_cost_per_second": 2.167e-05,
|
||||
"litellm_provider": "deepgram",
|
||||
"metadata": {
|
||||
"calculation": "$0.0013/60 seconds = $0.00002167 per second",
|
||||
"note": "Keyterm Prompting add-on (keyterm query param), streaming, pay as you go",
|
||||
"original_pricing_per_minute": 0.0013
|
||||
},
|
||||
"mode": "audio_transcription",
|
||||
"output_cost_per_second": 0.0,
|
||||
"source": "https://deepgram.com/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/listen"
|
||||
]
|
||||
},
|
||||
"deepgram/streaming/detect_entities": {
|
||||
"input_cost_per_second": 2.833e-05,
|
||||
"litellm_provider": "deepgram",
|
||||
"metadata": {
|
||||
"calculation": "$0.0017/60 seconds = $0.00002833 per second",
|
||||
"note": "Entity Detection add-on (detect_entities query param), streaming, pay as you go",
|
||||
"original_pricing_per_minute": 0.0017
|
||||
},
|
||||
"mode": "audio_transcription",
|
||||
"output_cost_per_second": 0.0,
|
||||
"source": "https://deepgram.com/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/listen"
|
||||
]
|
||||
},
|
||||
"deepgram/streaming/diarize": {
|
||||
"input_cost_per_second": 3.333e-05,
|
||||
"litellm_provider": "deepgram",
|
||||
"metadata": {
|
||||
"calculation": "$0.0020/60 seconds = $0.00003333 per second",
|
||||
"note": "Speaker Diarization add-on (diarize / diarize_model query params), streaming, pay as you go",
|
||||
"original_pricing_per_minute": 0.002
|
||||
},
|
||||
"mode": "audio_transcription",
|
||||
"output_cost_per_second": 0.0,
|
||||
"source": "https://deepgram.com/pricing",
|
||||
"supported_endpoints": [
|
||||
"/v1/listen"
|
||||
]
|
||||
},
|
||||
"deepgram/whisper": {
|
||||
"input_cost_per_second": 0.0001,
|
||||
"litellm_provider": "deepgram",
|
||||
|
|
|
|||
326
tests/test_litellm/llms/deepgram/test_deepgram_common_utils.py
Normal file
326
tests/test_litellm/llms/deepgram/test_deepgram_common_utils.py
Normal file
|
|
@ -0,0 +1,326 @@
|
|||
import math
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Final
|
||||
from urllib.parse import parse_qs, urlparse
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.llms.deepgram.common_utils import (
|
||||
deepgram_listen_addon_pricing_models,
|
||||
deepgram_listen_audio_seconds,
|
||||
deepgram_listen_callback_params,
|
||||
deepgram_listen_channel_count,
|
||||
deepgram_listen_is_priced,
|
||||
deepgram_listen_model,
|
||||
deepgram_listen_pricing_model,
|
||||
deepgram_listen_registry_key,
|
||||
deepgram_listen_requested_model,
|
||||
deepgram_listen_transcript,
|
||||
deepgram_listen_websocket_target,
|
||||
)
|
||||
|
||||
NOVA_3_URL: Final = "wss://api.deepgram.com/v1/listen?model=nova-3&encoding=linear16&sample_rate=16000"
|
||||
|
||||
|
||||
def _results(
|
||||
start: object,
|
||||
duration: object,
|
||||
transcript: str = "",
|
||||
is_final: object = True,
|
||||
channel_index: object = (0, 1),
|
||||
) -> dict[str, object]:
|
||||
return {
|
||||
"type": "Results",
|
||||
"start": start,
|
||||
"duration": duration,
|
||||
"is_final": is_final,
|
||||
"channel_index": list(channel_index) if isinstance(channel_index, tuple) else channel_index,
|
||||
"channel": {"alternatives": [{"transcript": transcript, "confidence": 0.9}]},
|
||||
}
|
||||
|
||||
|
||||
def _metadata(duration: object, channels: object = 1) -> dict[str, object]:
|
||||
return {"type": "Metadata", "request_id": "req-1", "duration": duration, "channels": channels}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("api_base", "query_string", "expected"),
|
||||
[
|
||||
pytest.param(
|
||||
None,
|
||||
"model=nova-3&encoding=linear16",
|
||||
"wss://api.deepgram.com/v1/listen?model=nova-3&encoding=linear16",
|
||||
id="default",
|
||||
),
|
||||
pytest.param(
|
||||
None,
|
||||
"encoding=linear16&sample_rate=16000",
|
||||
"wss://api.deepgram.com/v1/listen?encoding=linear16&sample_rate=16000&model=nova-3",
|
||||
id="model added when missing",
|
||||
),
|
||||
pytest.param(
|
||||
None,
|
||||
"model=&encoding=linear16",
|
||||
"wss://api.deepgram.com/v1/listen?encoding=linear16&model=nova-3",
|
||||
id="empty model replaced",
|
||||
),
|
||||
pytest.param(
|
||||
"http://localhost:9000/v1/",
|
||||
"model=nova-2",
|
||||
"ws://localhost:9000/v1/listen?model=nova-2",
|
||||
id="custom base becomes ws",
|
||||
),
|
||||
pytest.param(
|
||||
"wss://dg.internal/v1",
|
||||
"model=nova-3&keywords=a&keywords=b",
|
||||
"wss://dg.internal/v1/listen?model=nova-3&keywords=a&keywords=b",
|
||||
id="repeated keys preserved",
|
||||
),
|
||||
pytest.param(
|
||||
None,
|
||||
"model=nova-2&encoding=linear16&model=nova-3",
|
||||
"wss://api.deepgram.com/v1/listen?model=nova-2&encoding=linear16",
|
||||
id="only the authorized first model reaches deepgram",
|
||||
),
|
||||
pytest.param(
|
||||
None,
|
||||
"language=en&model=nova-3&language=multi",
|
||||
"wss://api.deepgram.com/v1/listen?language=en&model=nova-3",
|
||||
id="only the priced first language reaches deepgram",
|
||||
),
|
||||
pytest.param(
|
||||
None,
|
||||
"model=&model=nova-2",
|
||||
"wss://api.deepgram.com/v1/listen?model=nova-3",
|
||||
id="blank first model is the default, later models dropped",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_deepgram_listen_websocket_target(api_base: str | None, query_string: str, expected: str):
|
||||
assert deepgram_listen_websocket_target(api_base=api_base, query_string=query_string) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("query_string", "expected"),
|
||||
[
|
||||
pytest.param("model=nova-3&encoding=linear16", (), id="no callback"),
|
||||
pytest.param("model=nova-3&callback=https%3A%2F%2Fevil.example%2Fsink", ("callback",), id="callback"),
|
||||
pytest.param(
|
||||
"callback_method=put&model=nova-3&callback=wss%3A%2F%2Fevil.example",
|
||||
("callback", "callback_method"),
|
||||
id="callback and method",
|
||||
),
|
||||
pytest.param("model=nova-3&callback_method=put", ("callback_method",), id="method alone"),
|
||||
pytest.param("model=nova-3&callbacks=x&my_callback=y", (), id="only exact names match"),
|
||||
],
|
||||
)
|
||||
def test_deepgram_listen_callback_params(query_string: str, expected: tuple[str, ...]):
|
||||
assert deepgram_listen_callback_params(query_string) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("frames", "expected_seconds"),
|
||||
[
|
||||
pytest.param((_results(0.0, 2.0), _results(2.0, 3.5), _metadata(6.25)), 6.25, id="metadata wins"),
|
||||
pytest.param((_metadata(4.0), _results(0.0, 9.0), _metadata(5.5)), 5.5, id="last metadata wins"),
|
||||
pytest.param((_results(0.0, 2.0), _results(2.0, 3.5), _results(1.0, 1.0)), 5.5, id="furthest results end"),
|
||||
pytest.param((_results(0.0, 0.0), _metadata(0.0)), 0.0, id="zero metadata is a real zero"),
|
||||
pytest.param(
|
||||
(_metadata(0.0), _results(0.0, 2.0), _results(2.0, 3.5)),
|
||||
5.5,
|
||||
id="handshake metadata zero does not hide streamed results",
|
||||
),
|
||||
pytest.param((_metadata(0.0), _results(0.0, 2.0), _metadata(0.0)), 2.0, id="only zero metadata frames"),
|
||||
pytest.param((_results(0.0, 1.5), _metadata("6.25")), 1.5, id="string metadata is ignored"),
|
||||
pytest.param((_results(0.0, 1.5), _metadata(True)), 1.5, id="boolean metadata is ignored"),
|
||||
pytest.param((_results(0.0, 1.5), _metadata(-3.0)), 1.5, id="negative metadata is ignored"),
|
||||
pytest.param((_results(0.0, 1.5), _metadata(math.nan), _metadata(math.inf)), 1.5, id="nan/inf ignored"),
|
||||
pytest.param((_results("0", 2.0), _results(0.0, None), _results(0.0, 0.75)), 0.75, id="malformed results"),
|
||||
pytest.param(({"type": "SpeechStarted", "timestamp": 3.0}, {"type": "UtteranceEnd"}), 0.0, id="no usage"),
|
||||
pytest.param((), 0.0, id="no frames"),
|
||||
],
|
||||
)
|
||||
def test_deepgram_listen_audio_seconds(frames: Sequence[Mapping[str, object]], expected_seconds: float):
|
||||
assert deepgram_listen_audio_seconds(frames) == expected_seconds
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("frames", "upstream_url", "expected_channels"),
|
||||
[
|
||||
pytest.param((_results(0.0, 2.0), _metadata(6.25)), NOVA_3_URL, 1, id="mono"),
|
||||
pytest.param((_results(0.0, 2.0, channel_index=(0, 2)), _metadata(6.25, 2)), NOVA_3_URL, 2, id="stereo"),
|
||||
pytest.param((_metadata(1.0, 3), _metadata(1.0, 5)), NOVA_3_URL, 5, id="last metadata wins"),
|
||||
pytest.param(
|
||||
(_metadata(1.0, 20), _results(0.0, 1.0, channel_index=(1, 2))),
|
||||
NOVA_3_URL,
|
||||
20,
|
||||
id="metadata beats channel_index",
|
||||
),
|
||||
pytest.param(
|
||||
(_results(0.0, 1.0, channel_index=(0, 2)), _results(0.0, 1.0, channel_index=(3, 4))),
|
||||
NOVA_3_URL,
|
||||
4,
|
||||
id="widest channel_index without metadata",
|
||||
),
|
||||
pytest.param(
|
||||
(_results(0.0, 1.0, channel_index=(0, 2)),),
|
||||
f"{NOVA_3_URL}&channels=7&multichannel=true",
|
||||
2,
|
||||
id="frames beat the declared query",
|
||||
),
|
||||
pytest.param((), f"{NOVA_3_URL}&channels=7&multichannel=true", 7, id="declared query when no frames"),
|
||||
pytest.param((), f"{NOVA_3_URL}&channels=0", 1, id="zero declared channels"),
|
||||
pytest.param((), f"{NOVA_3_URL}&channels=-2", 1, id="negative declared channels"),
|
||||
pytest.param((), f"{NOVA_3_URL}&channels=two", 1, id="non numeric declared channels"),
|
||||
pytest.param((), NOVA_3_URL, 1, id="nothing declared"),
|
||||
pytest.param((_metadata(1.0, "2"), _metadata(1.0, True), _metadata(1.0, 0)), NOVA_3_URL, 1, id="bad metadata"),
|
||||
pytest.param((_metadata(1.0, 3), _metadata(1.0, True)), NOVA_3_URL, 3, id="boolean does not shadow a count"),
|
||||
pytest.param((_metadata(1.0, 2.0), _metadata(1.0, -1)), NOVA_3_URL, 1, id="float and negative metadata"),
|
||||
pytest.param(
|
||||
(_metadata(1.0, 2), {**_results(0.0, 1.0), "channels": 9}, {"type": "UtteranceEnd", "channels": 11}),
|
||||
NOVA_3_URL,
|
||||
2,
|
||||
id="channels on non metadata frames ignored",
|
||||
),
|
||||
pytest.param(
|
||||
(
|
||||
_results(0.0, 1.0, channel_index=[0]),
|
||||
_results(0.0, 1.0, channel_index=(0, "2")),
|
||||
_results(0.0, 1.0, channel_index=(0, 0)),
|
||||
),
|
||||
NOVA_3_URL,
|
||||
1,
|
||||
id="bad channel_index",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_deepgram_listen_channel_count(
|
||||
frames: Sequence[Mapping[str, object]], upstream_url: str, expected_channels: int
|
||||
):
|
||||
assert deepgram_listen_channel_count(frames, upstream_url) == expected_channels
|
||||
|
||||
|
||||
def test_deepgram_listen_transcript_joins_final_results_only():
|
||||
frames = (
|
||||
_results(0.0, 1.0, "hello wor", is_final=False),
|
||||
_results(0.0, 1.5, "hello world"),
|
||||
_results(1.5, 0.5, "", is_final=True),
|
||||
_results(2.0, 1.0, "how are you", is_final="yes"),
|
||||
{"type": "Results", "start": 3.0, "duration": 1.0, "is_final": True, "channel": {"alternatives": []}},
|
||||
_results(4.0, 1.0, "goodbye"),
|
||||
_metadata(5.0),
|
||||
)
|
||||
assert deepgram_listen_transcript(frames) == "hello world goodbye"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("upstream_url", "expected_model"),
|
||||
[
|
||||
(NOVA_3_URL, "nova-3"),
|
||||
("wss://api.deepgram.com/v1/listen?encoding=linear16&model=nova-2-medical", "nova-2-medical"),
|
||||
("wss://api.deepgram.com/v1/listen?model=nova-3&model=nova-2", "nova-3"),
|
||||
("wss://api.deepgram.com/v1/listen?encoding=linear16", litellm.constants.DEEPGRAM_LISTEN_DEFAULT_MODEL),
|
||||
],
|
||||
)
|
||||
def test_deepgram_listen_model_comes_from_the_upstream_query(upstream_url: str, expected_model: str):
|
||||
assert deepgram_listen_model(upstream_url) == expected_model
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"query_string",
|
||||
[
|
||||
"model=nova-2&language=en",
|
||||
"language=en",
|
||||
"model=&language=en",
|
||||
"",
|
||||
"model=nova-3-medical",
|
||||
"model=nova-2&model=nova-3",
|
||||
"model=&model=nova-3-medical",
|
||||
],
|
||||
)
|
||||
def test_requested_model_is_the_only_model_the_upstream_target_carries(query_string: str):
|
||||
"""Authorization runs against ``deepgram_listen_requested_model``; the upstream URL is built separately, so the
|
||||
two must always agree or a key could be authorized for one model and reach another. Deepgram reads the last
|
||||
repeated ``model``, so the target must carry exactly one."""
|
||||
target: Final = deepgram_listen_websocket_target(None, query_string)
|
||||
assert parse_qs(urlparse(target).query)["model"] == [deepgram_listen_requested_model(query_string)]
|
||||
assert deepgram_listen_requested_model(query_string) == deepgram_listen_model(target)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("upstream_url", "expected"),
|
||||
[
|
||||
pytest.param(NOVA_3_URL, "streaming/nova-3", id="monolingual"),
|
||||
pytest.param(f"{NOVA_3_URL}&language=en", "streaming/nova-3", id="explicit language"),
|
||||
pytest.param(f"{NOVA_3_URL}&language=multi", "streaming/nova-3-multilingual", id="multilingual"),
|
||||
pytest.param(f"{NOVA_3_URL}&language=MULTI", "streaming/nova-3-multilingual", id="multilingual any case"),
|
||||
pytest.param(
|
||||
"wss://api.deepgram.com/v1/listen?model=nova-2&language=multi",
|
||||
"streaming/nova-2-multilingual",
|
||||
id="other model",
|
||||
),
|
||||
pytest.param("wss://api.deepgram.com/v1/listen?encoding=linear16", "streaming/nova-3", id="default model"),
|
||||
],
|
||||
)
|
||||
def test_deepgram_listen_pricing_model_is_the_streaming_entry_never_the_prerecorded_one(
|
||||
upstream_url: str, expected: str
|
||||
):
|
||||
assert deepgram_listen_pricing_model(upstream_url) == expected
|
||||
assert deepgram_listen_registry_key(upstream_url) == f"deepgram/{expected}"
|
||||
|
||||
|
||||
NOVA_2_URL: Final = "wss://api.deepgram.com/v1/listen?model=nova-2"
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("local_model_cost_map")
|
||||
@pytest.mark.parametrize(
|
||||
("upstream_url", "extra_rows", "expected"),
|
||||
[
|
||||
pytest.param(NOVA_3_URL, (), True, id="streaming entry present"),
|
||||
pytest.param(f"{NOVA_3_URL}&language=multi", (), True, id="multilingual entry present"),
|
||||
pytest.param(NOVA_2_URL, (), False, id="only the pre-recorded entry"),
|
||||
pytest.param(f"{NOVA_2_URL}&language=multi", ("deepgram/streaming/nova-2",), False, id="needs multilingual"),
|
||||
pytest.param("wss://api.deepgram.com/v1/listen?model=nova-99-unmapped", (), False, id="nothing priced"),
|
||||
pytest.param(NOVA_2_URL, ("deepgram/streaming/nova-2",), True, id="operator-supplied streaming entry"),
|
||||
pytest.param(NOVA_2_URL, ("streaming/nova-2",), False, id="a row under another key is not the entry"),
|
||||
],
|
||||
)
|
||||
def test_deepgram_listen_is_priced(
|
||||
monkeypatch: pytest.MonkeyPatch, upstream_url: str, extra_rows: tuple[str, ...], expected: bool
|
||||
):
|
||||
"""The bundled map prices only nova-3 for streaming; nova-2 has a pre-recorded row, which must never count."""
|
||||
monkeypatch.delitem(litellm.model_cost, "deepgram/streaming/nova-2", raising=False)
|
||||
assert "deepgram/nova-2" in litellm.model_cost
|
||||
for row in extra_rows:
|
||||
monkeypatch.setitem(litellm.model_cost, row, dict(litellm.model_cost["deepgram/streaming/nova-3"]))
|
||||
|
||||
assert deepgram_listen_is_priced(upstream_url) is expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("upstream_url", "expected"),
|
||||
[
|
||||
pytest.param(NOVA_3_URL, (), id="no add-ons"),
|
||||
pytest.param(f"{NOVA_3_URL}&redact=pci", ("streaming/redact",), id="redact"),
|
||||
pytest.param(f"{NOVA_3_URL}&redact=pci&redact=ssn", ("streaming/redact",), id="repeated redact once"),
|
||||
pytest.param(f"{NOVA_3_URL}&keyterm=a&keyterm=b", ("streaming/keyterm",), id="keyterm"),
|
||||
pytest.param(f"{NOVA_3_URL}&detect_entities=true", ("streaming/detect_entities",), id="detect_entities"),
|
||||
pytest.param(f"{NOVA_3_URL}&diarize=true", ("streaming/diarize",), id="diarize"),
|
||||
pytest.param(f"{NOVA_3_URL}&diarize_model=v1", ("streaming/diarize",), id="diarize_model"),
|
||||
pytest.param(f"{NOVA_3_URL}&diarize=true&diarize_model=latest", ("streaming/diarize",), id="diarize both once"),
|
||||
pytest.param(f"{NOVA_3_URL}&detect_entities=false&diarize=FALSE&redact=", (), id="disabled"),
|
||||
pytest.param(
|
||||
f"{NOVA_3_URL}&detect_entities=false&detect_entities=true",
|
||||
("streaming/detect_entities",),
|
||||
id="any enabling value wins",
|
||||
),
|
||||
pytest.param(
|
||||
f"{NOVA_3_URL}&diarize=true&redact=pci&keyterm=x&detect_entities=true",
|
||||
("streaming/detect_entities", "streaming/diarize", "streaming/keyterm", "streaming/redact"),
|
||||
id="all, sorted",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_deepgram_listen_addon_pricing_models(upstream_url: str, expected: tuple[str, ...]):
|
||||
assert deepgram_listen_addon_pricing_models(upstream_url) == expected
|
||||
|
|
@ -0,0 +1,302 @@
|
|||
"""Deepgram ``/v1/listen`` WebSocket passthrough: duration extraction and duration based cost tracking."""
|
||||
|
||||
from datetime import datetime
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.deepgram_listen_passthrough_logging_handler import (
|
||||
DeepgramListenPassthroughLoggingHandler,
|
||||
)
|
||||
from litellm.proxy.pass_through_endpoints.success_handler import PassThroughEndpointLogging
|
||||
from litellm.types.passthrough_endpoints.pass_through_endpoints import PassthroughStandardLoggingPayload
|
||||
from litellm.types.utils import StandardLoggingPayload, TranscriptionResponse
|
||||
|
||||
NOVA_3_URL: Final = "wss://api.deepgram.com/v1/listen?model=nova-3&encoding=linear16&sample_rate=16000"
|
||||
|
||||
pytestmark: Final = pytest.mark.usefixtures("local_model_cost_map")
|
||||
|
||||
|
||||
def _results(start: object, duration: object, transcript: str = "", is_final: object = True) -> dict[str, object]:
|
||||
return {
|
||||
"type": "Results",
|
||||
"start": start,
|
||||
"duration": duration,
|
||||
"is_final": is_final,
|
||||
"channel": {"alternatives": [{"transcript": transcript, "confidence": 0.9}]},
|
||||
}
|
||||
|
||||
|
||||
def _metadata(duration: object, channels: int = 1) -> dict[str, object]:
|
||||
return {"type": "Metadata", "request_id": "req-1", "duration": duration, "channels": channels}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("url_route", "expected"),
|
||||
[
|
||||
("/deepgram/v1/listen", True),
|
||||
("/deepgram/listen", True),
|
||||
("/deepgram/v1/listen?model=nova-3", True),
|
||||
("/litellm/deepgram/v1/listen", True),
|
||||
("/deepgram/v1/speak", False),
|
||||
("/deepgram/v1/listen/extra", False),
|
||||
("/openai/v1/realtime", False),
|
||||
("/vertex_ai/live", False),
|
||||
("", False),
|
||||
],
|
||||
)
|
||||
def test_is_deepgram_listen_route(url_route: str, expected: bool):
|
||||
assert DeepgramListenPassthroughLoggingHandler.is_deepgram_listen_route(url_route) is expected
|
||||
|
||||
|
||||
def _logging_obj(call_id: str = "call-dg") -> LiteLLMLoggingObj:
|
||||
return LiteLLMLoggingObj(
|
||||
model="unknown",
|
||||
messages=[{"role": "user", "content": "WebSocket connection"}],
|
||||
stream=True,
|
||||
call_type="pass_through_endpoint",
|
||||
start_time=datetime.now(),
|
||||
litellm_call_id=call_id,
|
||||
function_id="websocket_passthrough",
|
||||
)
|
||||
|
||||
|
||||
def _registry_cost(pricing_model: str, seconds: float) -> float:
|
||||
"""Derives the expected charge from the live cost map rather than pinning a vendor price."""
|
||||
per_second: Final = litellm.model_cost[f"deepgram/{pricing_model}"]["input_cost_per_second"]
|
||||
assert per_second > 0
|
||||
return per_second * seconds
|
||||
|
||||
|
||||
def _cost(upstream_url: str, *frames: dict[str, object]) -> float:
|
||||
handler_result = DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler(
|
||||
websocket_messages=frames, logging_obj=_logging_obj(), upstream_url=upstream_url
|
||||
)
|
||||
response_cost = handler_result["kwargs"]["response_cost"]
|
||||
assert isinstance(response_cost, float)
|
||||
return response_cost
|
||||
|
||||
|
||||
def test_handler_bills_metadata_duration_at_the_registry_rate_and_names_the_model():
|
||||
frames = (_results(0.0, 5.0, "first sentence"), _results(5.0, 7.5, "second sentence"), _metadata(12.5))
|
||||
logging_obj = _logging_obj()
|
||||
|
||||
handler_result = DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler(
|
||||
websocket_messages=frames,
|
||||
logging_obj=logging_obj,
|
||||
upstream_url=NOVA_3_URL,
|
||||
kwargs={"litellm_params": {"metadata": {}}},
|
||||
)
|
||||
|
||||
result = handler_result["result"]
|
||||
assert isinstance(result, TranscriptionResponse)
|
||||
assert result.text == "first sentence second sentence"
|
||||
assert result._hidden_params["audio_transcription_duration"] == 12.5
|
||||
assert result._hidden_params["response_cost"] == pytest.approx(_registry_cost("streaming/nova-3", 12.5))
|
||||
assert handler_result["kwargs"]["response_cost"] == pytest.approx(_registry_cost("streaming/nova-3", 12.5))
|
||||
assert handler_result["kwargs"]["model"] == "nova-3"
|
||||
assert handler_result["kwargs"]["custom_llm_provider"] == "deepgram"
|
||||
assert handler_result["kwargs"]["litellm_params"] == {"metadata": {}}
|
||||
assert logging_obj.model == "nova-3"
|
||||
assert logging_obj.model_call_details["model"] == "nova-3"
|
||||
assert logging_obj.model_call_details["custom_llm_provider"] == "deepgram"
|
||||
assert logging_obj.model_call_details["response_cost"] == pytest.approx(_registry_cost("streaming/nova-3", 12.5))
|
||||
|
||||
|
||||
def test_handler_bills_streaming_not_prerecorded_rates():
|
||||
"""Deepgram prices /v1/listen over a WebSocket separately from pre-recorded transcription, so the streaming entry
|
||||
must be the one charged; the two registry rows only need to differ for this to matter, whatever their values."""
|
||||
streaming = litellm.model_cost["deepgram/streaming/nova-3"]["input_cost_per_second"]
|
||||
prerecorded = litellm.model_cost["deepgram/nova-3"]["input_cost_per_second"]
|
||||
assert streaming != prerecorded
|
||||
|
||||
assert _cost(NOVA_3_URL, _metadata(60.0)) == pytest.approx(60.0 * streaming)
|
||||
|
||||
|
||||
def test_handler_bills_multilingual_streaming_when_language_is_multi():
|
||||
monolingual = _cost(NOVA_3_URL, _metadata(60.0))
|
||||
multilingual = _cost(f"{NOVA_3_URL}&language=multi", _metadata(60.0))
|
||||
|
||||
assert multilingual == pytest.approx(_registry_cost("streaming/nova-3-multilingual", 60.0))
|
||||
assert multilingual > monolingual
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("query", "addons"),
|
||||
[
|
||||
pytest.param("redact=pci", ("redact",), id="redaction"),
|
||||
pytest.param("redact=pci&redact=numbers", ("redact",), id="redaction counted once"),
|
||||
pytest.param("keyterm=LiteLLM&keyterm=Deepgram", ("keyterm",), id="keyterm prompting"),
|
||||
pytest.param("detect_entities=true", ("detect_entities",), id="entity detection"),
|
||||
pytest.param("diarize=true", ("diarize",), id="diarization"),
|
||||
pytest.param("diarize_model=v1", ("diarize",), id="diarization via diarize_model"),
|
||||
pytest.param("diarize=true&diarize_model=v1", ("diarize",), id="diarization counted once"),
|
||||
pytest.param(
|
||||
"redact=pci&keyterm=x&detect_entities=true&diarize=true",
|
||||
("redact", "keyterm", "detect_entities", "diarize"),
|
||||
id="every add-on",
|
||||
),
|
||||
pytest.param("detect_entities=false&diarize=False&redact=", (), id="disabled add-ons cost nothing"),
|
||||
],
|
||||
)
|
||||
def test_handler_adds_each_priced_add_on_once_on_top_of_the_base_rate(query: str, addons: tuple[str, ...]):
|
||||
base = _cost(NOVA_3_URL, _metadata(60.0))
|
||||
expected = base + sum(_registry_cost(f"streaming/{addon}", 60.0) for addon in addons)
|
||||
|
||||
assert _cost(f"{NOVA_3_URL}&{query}", _metadata(60.0)) == pytest.approx(expected)
|
||||
|
||||
|
||||
def test_handler_add_ons_scale_with_channels_like_the_base_rate():
|
||||
stereo_plain = _cost(f"{NOVA_3_URL}&channels=2", _metadata(60.0, channels=2))
|
||||
stereo_redacted = _cost(f"{NOVA_3_URL}&channels=2&redact=pci", _metadata(60.0, channels=2))
|
||||
|
||||
assert stereo_redacted - stereo_plain == pytest.approx(_registry_cost("streaming/redact", 120.0))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"upstream_url",
|
||||
[
|
||||
pytest.param("wss://api.deepgram.com/v1/listen?model=nova-2", id="only a pre-recorded entry"),
|
||||
pytest.param("wss://api.deepgram.com/v1/listen?model=nova-99-not-in-registry", id="no entry at all"),
|
||||
],
|
||||
)
|
||||
def test_handler_never_substitutes_another_rate_for_a_missing_streaming_entry(monkeypatch, upstream_url):
|
||||
"""The route refuses these sessions up front; should the registry change under a live one, the spend row
|
||||
keeps the duration and carries no cost, rather than the pre-recorded rate or any other stand-in."""
|
||||
monkeypatch.delitem(litellm.model_cost, "deepgram/streaming/nova-2", raising=False)
|
||||
assert "deepgram/nova-2" in litellm.model_cost
|
||||
|
||||
handler_result = DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler(
|
||||
websocket_messages=(_metadata(60.0),), logging_obj=_logging_obj(), upstream_url=upstream_url
|
||||
)
|
||||
|
||||
assert handler_result["kwargs"]["response_cost"] is None
|
||||
assert handler_result["result"]._hidden_params["audio_transcription_duration"] == 60.0
|
||||
|
||||
|
||||
def test_handler_falls_back_to_results_frames_when_the_stream_ends_without_metadata():
|
||||
frames = (_results(0.0, 30.0, "a"), _results(30.0, 30.0, "b"), _results(60.0, 12.5, "c"))
|
||||
|
||||
handler_result = DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler(
|
||||
websocket_messages=frames, logging_obj=_logging_obj(), upstream_url=NOVA_3_URL
|
||||
)
|
||||
|
||||
assert handler_result["result"]._hidden_params["audio_transcription_duration"] == 72.5
|
||||
assert handler_result["kwargs"]["response_cost"] == pytest.approx(_registry_cost("streaming/nova-3", 72.5))
|
||||
|
||||
|
||||
def test_handler_charges_more_for_more_audio_on_the_same_model():
|
||||
short = DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler(
|
||||
websocket_messages=(_metadata(10.0),), logging_obj=_logging_obj(), upstream_url=NOVA_3_URL
|
||||
)
|
||||
long = DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler(
|
||||
websocket_messages=(_metadata(30.0),), logging_obj=_logging_obj(), upstream_url=NOVA_3_URL
|
||||
)
|
||||
|
||||
assert long["kwargs"]["response_cost"] == pytest.approx(3 * short["kwargs"]["response_cost"])
|
||||
assert short["kwargs"]["response_cost"] > 0
|
||||
|
||||
|
||||
def test_handler_bills_every_channel_of_a_multichannel_session():
|
||||
"""Deepgram bills processed audio per channel (deepgram.com/pricing FAQ, 2026-09-17), so a stereo session must be
|
||||
charged for twice its wall-clock duration or budgets can be bypassed by requesting more channels."""
|
||||
mono = DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler(
|
||||
websocket_messages=(_metadata(30.0),), logging_obj=_logging_obj(), upstream_url=NOVA_3_URL
|
||||
)
|
||||
stereo = DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler(
|
||||
websocket_messages=(_metadata(30.0, channels=2),),
|
||||
logging_obj=_logging_obj(),
|
||||
upstream_url=f"{NOVA_3_URL}&multichannel=true&channels=2",
|
||||
)
|
||||
|
||||
assert stereo["result"]._hidden_params["audio_transcription_duration"] == 60.0
|
||||
assert stereo["kwargs"]["response_cost"] == pytest.approx(2 * mono["kwargs"]["response_cost"])
|
||||
assert stereo["kwargs"]["response_cost"] == pytest.approx(_registry_cost("streaming/nova-3", 60.0))
|
||||
|
||||
|
||||
def test_handler_bills_the_declared_channels_when_the_stream_dies_before_any_frame_reports_them():
|
||||
handler_result = DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler(
|
||||
websocket_messages=(_results(0.0, 10.0, "a"),),
|
||||
logging_obj=_logging_obj(),
|
||||
upstream_url=f"{NOVA_3_URL}&multichannel=true&channels=3",
|
||||
)
|
||||
|
||||
assert handler_result["result"]._hidden_params["audio_transcription_duration"] == 30.0
|
||||
assert handler_result["kwargs"]["response_cost"] == pytest.approx(_registry_cost("streaming/nova-3", 30.0))
|
||||
|
||||
|
||||
class _CapturingLogger(CustomLogger):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.payloads: list[StandardLoggingPayload] = []
|
||||
|
||||
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
self.payloads.append(kwargs["standard_logging_object"])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_success_handler_dispatches_deepgram_listen_and_logs_duration_based_spend(monkeypatch):
|
||||
"""Drives the shared passthrough success handler the way the WebSocket relay does at socket close and reads
|
||||
what a spend logger receives: Deepgram model and provider, the audio duration billed at the registry rate."""
|
||||
capturing_logger = _CapturingLogger()
|
||||
monkeypatch.setattr(litellm, "_async_success_callback", [capturing_logger])
|
||||
monkeypatch.setattr(litellm, "success_callback", [])
|
||||
monkeypatch.setattr(litellm, "callbacks", [])
|
||||
logging_obj = _logging_obj("call-dg-e2e")
|
||||
frames = [_results(0.0, 5.0, "hello world", is_final=False), _results(0.0, 5.0, "hello world"), _metadata(20.0)]
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key="hashed-key", team_id="team-stt", user_id="user-1")
|
||||
start_time = datetime.now()
|
||||
passthrough_logging_payload = PassthroughStandardLoggingPayload(
|
||||
url=NOVA_3_URL, request_body={}, request_method="WEBSOCKET", cost_per_request=None
|
||||
)
|
||||
logging_obj.update_environment_variables(
|
||||
model="unknown",
|
||||
user="unknown",
|
||||
optional_params={},
|
||||
litellm_params={
|
||||
"metadata": {
|
||||
"user_api_key": user_api_key_dict.api_key,
|
||||
"user_api_key_team_id": user_api_key_dict.team_id,
|
||||
"user_api_key_user_id": user_api_key_dict.user_id,
|
||||
}
|
||||
},
|
||||
call_type="pass_through_endpoint",
|
||||
)
|
||||
|
||||
await PassThroughEndpointLogging().pass_through_async_success_handler(
|
||||
httpx_response=SimpleNamespace(
|
||||
status_code=200,
|
||||
text="WebSocket connection successful",
|
||||
headers={},
|
||||
request=SimpleNamespace(method="WEBSOCKET", url=NOVA_3_URL),
|
||||
),
|
||||
response_body=frames,
|
||||
logging_obj=logging_obj,
|
||||
url_route="/deepgram/v1/listen",
|
||||
result="websocket_connection_successful",
|
||||
start_time=start_time,
|
||||
end_time=datetime.now(),
|
||||
cache_hit=False,
|
||||
request_body={},
|
||||
passthrough_logging_payload=passthrough_logging_payload,
|
||||
litellm_params={
|
||||
"metadata": {
|
||||
"user_api_key": user_api_key_dict.api_key,
|
||||
"user_api_key_team_id": user_api_key_dict.team_id,
|
||||
"user_api_key_user_id": user_api_key_dict.user_id,
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
assert len(capturing_logger.payloads) == 1
|
||||
payload = capturing_logger.payloads[0]
|
||||
assert payload["model"] == "nova-3"
|
||||
assert payload["custom_llm_provider"] == "deepgram"
|
||||
assert payload["response_cost"] == pytest.approx(_registry_cost("streaming/nova-3", 20.0))
|
||||
assert payload["metadata"]["user_api_key_team_id"] == "team-stt"
|
||||
assert payload["id"] == "call-dg-e2e"
|
||||
|
|
@ -0,0 +1,474 @@
|
|||
"""Deepgram ``/v1/listen`` passthrough WebSocket route: registration, auth, credential injection, target URL."""
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType, SimpleNamespace
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
from starlette.routing import WebSocketRoute
|
||||
from starlette.websockets import WebSocketDisconnect
|
||||
|
||||
import litellm
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
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,
|
||||
router,
|
||||
)
|
||||
from litellm.proxy.utils import hash_token
|
||||
|
||||
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")
|
||||
NOVA_2_STREAMING_KEY: Final = "deepgram/streaming/nova-2"
|
||||
|
||||
pytestmark: Final = pytest.mark.usefixtures("local_model_cost_map")
|
||||
|
||||
|
||||
def _price_nova_2_streaming(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""An operator-supplied streaming row: the bundled map prices only nova-3 for streaming."""
|
||||
monkeypatch.setitem(litellm.model_cost, NOVA_2_STREAMING_KEY, dict(litellm.model_cost["deepgram/streaming/nova-3"]))
|
||||
|
||||
|
||||
class _FakeWebSocket:
|
||||
def __init__(self, path: str, query: str) -> None:
|
||||
self.url = SimpleNamespace(path=path, query=query)
|
||||
self.headers = {"authorization": "Bearer sk-litellm-virtual", "x-api-key": "sk-caller-secret"}
|
||||
self.accepts: list[str | None] = []
|
||||
self.closed: tuple[int, str] | None = None
|
||||
|
||||
async def accept(self, subprotocol: str | None = None) -> None:
|
||||
self.accepts.append(subprotocol)
|
||||
|
||||
async def close(self, code: int = 1000, reason: str = "") -> None:
|
||||
self.closed = (code, reason)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _RelayCall:
|
||||
target: str
|
||||
custom_headers: Mapping[str, str]
|
||||
user_api_key_dict: UserAPIKeyAuth
|
||||
forward_headers: bool
|
||||
endpoint: str
|
||||
accept_websocket: bool
|
||||
|
||||
|
||||
class _FakeRelay:
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[_RelayCall] = []
|
||||
|
||||
async def __call__(
|
||||
self,
|
||||
*,
|
||||
websocket: object,
|
||||
target: str,
|
||||
custom_headers: dict[str, str],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
forward_headers: bool,
|
||||
endpoint: str,
|
||||
accept_websocket: bool,
|
||||
) -> None:
|
||||
self.calls.append(
|
||||
_RelayCall(
|
||||
target=target,
|
||||
custom_headers=MappingProxyType(dict(custom_headers)),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
forward_headers=forward_headers,
|
||||
endpoint=endpoint,
|
||||
accept_websocket=accept_websocket,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
async def _serve(websocket: _FakeWebSocket, user_api_key_dict: UserAPIKeyAuth | None = None) -> _FakeRelay:
|
||||
relay = _FakeRelay()
|
||||
await deepgram_listen_websocket_route(
|
||||
websocket=websocket,
|
||||
user_api_key_dict=user_api_key_dict or UserAPIKeyAuth(),
|
||||
relay=relay,
|
||||
)
|
||||
return relay
|
||||
|
||||
|
||||
def test_deepgram_listen_websocket_routes_registered():
|
||||
ws_paths = {route.path for route in router.routes if isinstance(route, WebSocketRoute)}
|
||||
assert set(LISTEN_PATHS) <= ws_paths
|
||||
|
||||
|
||||
@pytest.mark.parametrize("path", LISTEN_PATHS)
|
||||
def test_deepgram_listen_is_a_lazily_loaded_mapped_pass_through_route(path):
|
||||
"""The route must be reachable before the passthrough module is imported and must be authed and
|
||||
billed as a mapped pass-through route like the other provider prefixes."""
|
||||
feature = next(feature for feature in LAZY_FEATURES if feature.name == "llm_passthrough")
|
||||
assert feature.matches(path)
|
||||
assert any(path.startswith(prefix) for prefix in LiteLLMRoutes.mapped_pass_through_routes.value)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("path", LISTEN_PATHS)
|
||||
async def test_deepgram_listen_forwards_query_and_injects_only_provider_auth(path, monkeypatch):
|
||||
monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False)
|
||||
websocket = _FakeWebSocket(path, "encoding=linear16&sample_rate=16000&keywords=hi%3A2&keywords=there")
|
||||
caller = UserAPIKeyAuth(api_key="sk-litellm-virtual", team_id="team-stt")
|
||||
|
||||
with patch(GET_CREDENTIALS, return_value="dg-provider-key") as get_credentials:
|
||||
relay = await _serve(websocket, caller)
|
||||
|
||||
assert get_credentials.call_args.kwargs == {"custom_llm_provider": "deepgram", "region_name": None}
|
||||
assert relay.calls == [
|
||||
_RelayCall(
|
||||
target=(
|
||||
"wss://api.deepgram.com/v1/listen"
|
||||
"?encoding=linear16&sample_rate=16000&keywords=hi%3A2&keywords=there&model=nova-3"
|
||||
),
|
||||
custom_headers=MappingProxyType({"Authorization": "Token dg-provider-key"}),
|
||||
user_api_key_dict=caller,
|
||||
forward_headers=False,
|
||||
endpoint=path,
|
||||
accept_websocket=False,
|
||||
)
|
||||
]
|
||||
assert websocket.accepts == [None]
|
||||
assert websocket.closed is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deepgram_listen_keeps_caller_chosen_model(monkeypatch):
|
||||
monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False)
|
||||
_price_nova_2_streaming(monkeypatch)
|
||||
websocket = _FakeWebSocket("/deepgram/v1/listen", "model=nova-2&language=en")
|
||||
|
||||
with patch(GET_CREDENTIALS, return_value="dg-provider-key"):
|
||||
relay = await _serve(websocket)
|
||||
|
||||
assert [call.target for call in relay.calls] == ["wss://api.deepgram.com/v1/listen?model=nova-2&language=en"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("query", "expected_target"),
|
||||
[
|
||||
("", "wss://api.deepgram.com/v1/listen?model=nova-3"),
|
||||
("model=", "wss://api.deepgram.com/v1/listen?model=nova-3"),
|
||||
("model=&language=en", "wss://api.deepgram.com/v1/listen?language=en&model=nova-3"),
|
||||
],
|
||||
)
|
||||
async def test_deepgram_listen_defaults_to_nova_3_when_no_model_is_named(query, expected_target, monkeypatch):
|
||||
monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False)
|
||||
websocket = _FakeWebSocket("/deepgram/listen", query)
|
||||
|
||||
with patch(GET_CREDENTIALS, return_value="dg-provider-key"):
|
||||
relay = await _serve(websocket)
|
||||
|
||||
assert [call.target for call in relay.calls] == [expected_target]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("api_base", "expected_target"),
|
||||
[
|
||||
("https://api.eu.deepgram.com/v1/", "wss://api.eu.deepgram.com/v1/listen?model=nova-3"),
|
||||
("http://localhost:8080/v1", "ws://localhost:8080/v1/listen?model=nova-3"),
|
||||
("wss://deepgram.internal.example/v1", "wss://deepgram.internal.example/v1/listen?model=nova-3"),
|
||||
],
|
||||
)
|
||||
async def test_deepgram_listen_honours_server_configured_api_base(api_base, expected_target, monkeypatch):
|
||||
monkeypatch.setenv("DEEPGRAM_API_BASE", api_base)
|
||||
websocket = _FakeWebSocket("/deepgram/v1/listen", "")
|
||||
|
||||
with patch(GET_CREDENTIALS, return_value="dg-provider-key"):
|
||||
relay = await _serve(websocket)
|
||||
|
||||
assert [call.target for call in relay.calls] == [expected_target]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deepgram_listen_ignores_caller_supplied_api_base(monkeypatch):
|
||||
"""V1: the server-configured Deepgram key must only ever go to the server-configured host."""
|
||||
monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False)
|
||||
websocket = _FakeWebSocket("/deepgram/v1/listen", "api_base=wss%3A%2F%2Fattacker.example%2Fv1&model=nova-3")
|
||||
|
||||
with patch(GET_CREDENTIALS, return_value="dg-provider-key"):
|
||||
relay = await _serve(websocket)
|
||||
|
||||
assert [call.target for call in relay.calls] == [
|
||||
"wss://api.deepgram.com/v1/listen?api_base=wss%3A%2F%2Fattacker.example%2Fv1&model=nova-3"
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deepgram_listen_closes_cleanly_when_provider_credentials_missing():
|
||||
websocket = _FakeWebSocket("/deepgram/v1/listen", "model=nova-3")
|
||||
|
||||
with patch(GET_CREDENTIALS, return_value=None):
|
||||
relay = await _serve(websocket)
|
||||
|
||||
assert websocket.closed is not None
|
||||
assert websocket.closed[0] == 1011
|
||||
assert "DEEPGRAM_API_KEY" in websocket.closed[1]
|
||||
assert websocket.accepts == []
|
||||
assert relay.calls == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"query",
|
||||
[
|
||||
pytest.param("model=nova-3&callback=https%3A%2F%2Fsink.example%2Fdg", id="http callback"),
|
||||
pytest.param("callback=wss%3A%2F%2Fsink.example&callback_method=put&model=nova-3", id="ws callback"),
|
||||
],
|
||||
)
|
||||
async def test_deepgram_listen_rejects_callback_delivery_that_would_go_unbilled(query, monkeypatch):
|
||||
"""With ``callback`` set, Deepgram sends every Results and Metadata frame to the caller's URL and only a
|
||||
request id down this socket, so the proxy would meter zero seconds of audio; refuse before contacting Deepgram."""
|
||||
monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False)
|
||||
websocket = _FakeWebSocket("/deepgram/v1/listen", query)
|
||||
|
||||
with patch(GET_CREDENTIALS, return_value="dg-provider-key"):
|
||||
relay = await _serve(websocket)
|
||||
|
||||
assert relay.calls == []
|
||||
assert websocket.closed is not None
|
||||
assert websocket.closed[0] == 1008
|
||||
assert "callback" in websocket.closed[1]
|
||||
assert "dg-provider-key" not in websocket.closed[1]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("query", "missing_key"),
|
||||
[
|
||||
pytest.param("model=nova-2", "deepgram/streaming/nova-2", id="model with only a pre-recorded price"),
|
||||
pytest.param("model=nova-99", "deepgram/streaming/nova-99", id="model unknown to the registry"),
|
||||
pytest.param(
|
||||
"model=nova-3&language=multi",
|
||||
"deepgram/streaming/nova-3-multilingual",
|
||||
id="multilingual session without its own price",
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_deepgram_listen_refuses_sessions_it_cannot_price(query, missing_key, monkeypatch):
|
||||
"""A session with no streaming price would be logged at zero (or at the pre-recorded rate), letting a caller run
|
||||
up unmetered spend, so the proxy closes it before Deepgram is contacted and names the registry row to add."""
|
||||
monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False)
|
||||
monkeypatch.delitem(litellm.model_cost, missing_key, raising=False)
|
||||
assert "deepgram/nova-2" in litellm.model_cost
|
||||
websocket = _FakeWebSocket("/deepgram/v1/listen", query)
|
||||
|
||||
with patch(GET_CREDENTIALS, return_value="dg-provider-key"):
|
||||
relay = await _serve(websocket)
|
||||
|
||||
assert relay.calls == []
|
||||
assert websocket.closed is not None
|
||||
assert websocket.closed[0] == 1008
|
||||
assert missing_key in websocket.closed[1]
|
||||
assert "dg-provider-key" not in websocket.closed[1]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deepgram_listen_relays_once_the_operator_prices_the_model(monkeypatch):
|
||||
monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False)
|
||||
websocket = _FakeWebSocket("/deepgram/v1/listen", "model=nova-2")
|
||||
with patch(GET_CREDENTIALS, return_value="dg-provider-key"):
|
||||
assert (await _serve(websocket)).calls == []
|
||||
|
||||
_price_nova_2_streaming(monkeypatch)
|
||||
priced_websocket = _FakeWebSocket("/deepgram/v1/listen", "model=nova-2")
|
||||
with patch(GET_CREDENTIALS, return_value="dg-provider-key"):
|
||||
relay = await _serve(priced_websocket)
|
||||
|
||||
assert [call.target for call in relay.calls] == ["wss://api.deepgram.com/v1/listen?model=nova-2"]
|
||||
assert priced_websocket.closed is None
|
||||
|
||||
|
||||
def _app_with_relay(relay: _FakeRelay) -> FastAPI:
|
||||
app = FastAPI()
|
||||
app.include_router(router)
|
||||
app.dependency_overrides[_websocket_relay] = lambda: relay
|
||||
return app
|
||||
|
||||
|
||||
def test_deepgram_listen_rejects_connections_without_a_litellm_key():
|
||||
relay = _FakeRelay()
|
||||
client = TestClient(_app_with_relay(relay))
|
||||
|
||||
with patch(GET_CREDENTIALS, return_value="dg-provider-key") as get_credentials:
|
||||
with pytest.raises(WebSocketDisconnect) as disconnect:
|
||||
with client.websocket_connect("/deepgram/v1/listen?model=nova-3"):
|
||||
pass
|
||||
|
||||
assert disconnect.value.code == 1008
|
||||
assert relay.calls == []
|
||||
get_credentials.assert_not_called()
|
||||
|
||||
|
||||
def test_deepgram_listen_callback_rejection_reaches_the_client_as_a_policy_close(monkeypatch):
|
||||
monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False)
|
||||
relay = _FakeRelay()
|
||||
client = TestClient(_app_with_relay(relay))
|
||||
|
||||
with (
|
||||
patch(GET_CREDENTIALS, return_value="dg-provider-key"),
|
||||
patch(USER_API_KEY_AUTH, new=AsyncMock(return_value=UserAPIKeyAuth(api_key="hashed"))),
|
||||
):
|
||||
with pytest.raises(WebSocketDisconnect) as disconnect:
|
||||
with client.websocket_connect(
|
||||
"/deepgram/v1/listen?model=nova-3&callback=https%3A%2F%2Fsink.example%2Fdg",
|
||||
headers={"Authorization": "Bearer sk-litellm-virtual"},
|
||||
) as connection:
|
||||
connection.receive_text()
|
||||
|
||||
assert disconnect.value.code == 1008
|
||||
assert "callback" in disconnect.value.reason
|
||||
assert relay.calls == []
|
||||
|
||||
|
||||
def test_deepgram_listen_authenticates_the_litellm_key_and_relays_to_deepgram(monkeypatch):
|
||||
monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False)
|
||||
relay = _FakeRelay()
|
||||
client = TestClient(_app_with_relay(relay))
|
||||
caller = UserAPIKeyAuth(api_key="hashed-sk-litellm", team_id="team-stt")
|
||||
|
||||
with (
|
||||
patch(GET_CREDENTIALS, return_value="dg-provider-key"),
|
||||
patch(USER_API_KEY_AUTH, new=AsyncMock(return_value=caller)) as auth,
|
||||
):
|
||||
with client.websocket_connect(
|
||||
"/deepgram/v1/listen?model=nova-3&punctuate=true",
|
||||
headers={"Authorization": "Bearer sk-litellm-virtual"},
|
||||
):
|
||||
pass
|
||||
|
||||
assert auth.await_args.kwargs["api_key"] == "Bearer sk-litellm-virtual"
|
||||
assert relay.calls == [
|
||||
_RelayCall(
|
||||
target="wss://api.deepgram.com/v1/listen?model=nova-3&punctuate=true",
|
||||
custom_headers=MappingProxyType({"Authorization": "Token dg-provider-key"}),
|
||||
user_api_key_dict=caller,
|
||||
forward_headers=False,
|
||||
endpoint="/deepgram/v1/listen",
|
||||
accept_websocket=False,
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
async def _cache_restricted_key(virtual_key: str, models: list[str]) -> DualCache:
|
||||
cache = DualCache()
|
||||
await _cache_key_object(
|
||||
hashed_token=hash_token(virtual_key),
|
||||
user_api_key_obj=UserAPIKeyAuth(token=hash_token(virtual_key), models=models),
|
||||
user_api_key_cache=cache,
|
||||
proxy_logging_obj=None,
|
||||
)
|
||||
return cache
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("query", "expect_relay"),
|
||||
[
|
||||
pytest.param("model=nova-2", True, id="allowed model named"),
|
||||
pytest.param("model=nova-3", False, id="denied model named"),
|
||||
pytest.param("", False, id="model omitted, default denied"),
|
||||
pytest.param("model=&language=en", False, id="model blank, default denied"),
|
||||
],
|
||||
)
|
||||
def test_deepgram_listen_authorizes_the_model_it_will_actually_send_upstream(query, expect_relay, monkeypatch):
|
||||
"""A key allowed only ``nova-2`` must not reach ``nova-3`` by leaving ``model`` out and letting the proxy fill
|
||||
in its default: the real key auth path must see the same model the upstream target will carry."""
|
||||
monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False)
|
||||
monkeypatch.setattr(litellm, "max_budget", 0.0)
|
||||
_price_nova_2_streaming(monkeypatch)
|
||||
cache = asyncio.run(_cache_restricted_key("sk-only-nova-2", ["nova-2"]))
|
||||
relay = _FakeRelay()
|
||||
client = TestClient(_app_with_relay(relay))
|
||||
|
||||
with (
|
||||
patch(GET_CREDENTIALS, return_value="dg-provider-key"),
|
||||
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,
|
||||
),
|
||||
):
|
||||
if expect_relay:
|
||||
with client.websocket_connect(
|
||||
f"/deepgram/v1/listen?{query}", headers={"Authorization": "Bearer sk-only-nova-2"}
|
||||
):
|
||||
pass
|
||||
assert [call.target for call in relay.calls] == [f"wss://api.deepgram.com/v1/listen?{query}"]
|
||||
return
|
||||
with pytest.raises(WebSocketDisconnect) as disconnect:
|
||||
with client.websocket_connect(
|
||||
f"/deepgram/v1/listen?{query}", headers={"Authorization": "Bearer sk-only-nova-2"}
|
||||
):
|
||||
pass
|
||||
|
||||
assert disconnect.value.code == 1008
|
||||
assert relay.calls == []
|
||||
|
||||
|
||||
def test_deepgram_listen_strips_a_second_model_that_would_outrank_the_authorized_one(monkeypatch):
|
||||
"""Deepgram honours the last repeated ``model``; auth and pricing read the first. A key allowed only ``nova-2``
|
||||
must not smuggle ``nova-3`` past authorization behind an authorized first value."""
|
||||
monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False)
|
||||
monkeypatch.setattr(litellm, "max_budget", 0.0)
|
||||
_price_nova_2_streaming(monkeypatch)
|
||||
cache = asyncio.run(_cache_restricted_key("sk-only-nova-2", ["nova-2"]))
|
||||
relay = _FakeRelay()
|
||||
client = TestClient(_app_with_relay(relay))
|
||||
|
||||
with (
|
||||
patch(GET_CREDENTIALS, return_value="dg-provider-key"),
|
||||
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(
|
||||
"/deepgram/v1/listen?model=nova-2&language=en&model=nova-3&language=multi",
|
||||
headers={"Authorization": "Bearer sk-only-nova-2"},
|
||||
):
|
||||
pass
|
||||
|
||||
assert [call.target for call in relay.calls] == ["wss://api.deepgram.com/v1/listen?model=nova-2&language=en"]
|
||||
|
||||
|
||||
def test_deepgram_listen_echoes_the_browser_subprotocol_that_carries_the_litellm_key(monkeypatch):
|
||||
"""Browsers cannot set headers, so they send the key as a subprotocol and abort the handshake unless the
|
||||
server echoes that subprotocol back; the key itself must still stay off the upstream connection."""
|
||||
monkeypatch.delenv("DEEPGRAM_API_BASE", raising=False)
|
||||
relay = _FakeRelay()
|
||||
client = TestClient(_app_with_relay(relay))
|
||||
|
||||
with (
|
||||
patch(GET_CREDENTIALS, return_value="dg-provider-key"),
|
||||
patch(USER_API_KEY_AUTH, new=AsyncMock(return_value=UserAPIKeyAuth(api_key="hashed"))),
|
||||
):
|
||||
with client.websocket_connect(
|
||||
"/deepgram/v1/listen?model=nova-3",
|
||||
subprotocols=["openai-insecure-api-key.sk-litellm-virtual"],
|
||||
) as connection:
|
||||
assert connection.accepted_subprotocol == "openai-insecure-api-key.sk-litellm-virtual"
|
||||
|
||||
assert [call.custom_headers for call in relay.calls] == [
|
||||
MappingProxyType({"Authorization": "Token dg-provider-key"})
|
||||
]
|
||||
assert [call.forward_headers for call in relay.calls] == [False]
|
||||
|
|
@ -4942,18 +4942,21 @@ async def test_unusable_upstream_cost_records_zero_not_the_flat_estimate():
|
|||
|
||||
|
||||
class FakeUpstreamWebSocket:
|
||||
def __init__(self, first_frame: bytes):
|
||||
self._first_frame = first_frame
|
||||
"""Serves the given frames in order, then closes normally, the way a real websockets connection does"""
|
||||
|
||||
def __init__(self, *frames: str | bytes):
|
||||
self._frames = iter(frames)
|
||||
self.close = AsyncMock()
|
||||
self.send = AsyncMock()
|
||||
|
||||
async def recv(self, decode: bool = True):
|
||||
return self._first_frame
|
||||
async def recv(self, decode: bool | None = None):
|
||||
from websockets.exceptions import ConnectionClosedOK
|
||||
from websockets.frames import Close
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
raise StopAsyncIteration
|
||||
frame = next(self._frames, None)
|
||||
if frame is None:
|
||||
raise ConnectionClosedOK(rcvd=Close(1000, ""), sent=Close(1000, ""), rcvd_then_sent=True)
|
||||
return frame
|
||||
|
||||
|
||||
class FakeUpstreamConnect:
|
||||
|
|
@ -4974,7 +4977,7 @@ async def test_websocket_passthrough_forwards_non_ascii_first_frame():
|
|||
first_frame = json.dumps(
|
||||
{"type": "session.created", "session": {"instructions": "Hablas español, ¿sí?"}},
|
||||
ensure_ascii=False,
|
||||
).encode("utf-8")
|
||||
)
|
||||
upstream_ws = FakeUpstreamWebSocket(first_frame)
|
||||
|
||||
websocket = MagicMock()
|
||||
|
|
@ -5028,7 +5031,7 @@ async def test_websocket_passthrough_propagates_active_trace_context(
|
|||
from starlette.websockets import WebSocketState
|
||||
|
||||
captured: dict[str, dict[str, str]] = {}
|
||||
upstream_ws = FakeUpstreamWebSocket(b"{}")
|
||||
upstream_ws = FakeUpstreamWebSocket("{}")
|
||||
|
||||
def fake_connect(target, additional_headers):
|
||||
captured["headers"] = additional_headers
|
||||
|
|
@ -5457,6 +5460,144 @@ async def test_websocket_passthrough_does_not_close_twice_when_success_logging_f
|
|||
websocket.close.assert_awaited_once_with(code=1008, reason=upstream_reason)
|
||||
|
||||
|
||||
DEEPGRAM_LISTEN_TARGET = "wss://api.deepgram.com/v1/listen?model=nova-3&encoding=linear16&sample_rate=16000"
|
||||
DEEPGRAM_INTERIM_FRAME = json.dumps(
|
||||
{
|
||||
"type": "Results",
|
||||
"start": 0.0,
|
||||
"duration": 1.02,
|
||||
"is_final": False,
|
||||
"channel": {"alternatives": [{"transcript": "hello wor", "confidence": 0.71}]},
|
||||
}
|
||||
)
|
||||
DEEPGRAM_FINAL_FRAME = json.dumps(
|
||||
{
|
||||
"type": "Results",
|
||||
"start": 0.0,
|
||||
"duration": 2.5,
|
||||
"is_final": True,
|
||||
"speech_final": True,
|
||||
"channel": {"alternatives": [{"transcript": "hello world, ¿qué tal?", "confidence": 0.98}]},
|
||||
},
|
||||
ensure_ascii=False,
|
||||
)
|
||||
DEEPGRAM_METADATA_FRAME = json.dumps({"type": "Metadata", "request_id": "req-1", "duration": 2.5, "channels": 1})
|
||||
|
||||
|
||||
async def _relay_deepgram_listen(upstream_ws, client_receive):
|
||||
"""Runs the generic relay the way the Deepgram route does and returns (client websocket, success handler mock)"""
|
||||
websocket = _client_websocket(client_receive)
|
||||
with (
|
||||
_patched_websocket_passthrough_environment(upstream_ws),
|
||||
patch( # test-quality-ok: pass_through_endpoint_logging is a module global read inside websocket_passthrough_request; there is no injection seam
|
||||
"litellm.proxy.pass_through_endpoints.pass_through_endpoints."
|
||||
"pass_through_endpoint_logging.pass_through_async_success_handler",
|
||||
new=AsyncMock(),
|
||||
) as success_handler,
|
||||
):
|
||||
await websocket_passthrough_request(
|
||||
websocket=websocket,
|
||||
target=DEEPGRAM_LISTEN_TARGET,
|
||||
custom_headers={"Authorization": "Token dg-provider-key"},
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
forward_headers=False,
|
||||
endpoint="/deepgram/v1/listen",
|
||||
accept_websocket=False,
|
||||
)
|
||||
return websocket, success_handler
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_relays_deepgram_transcript_frames_verbatim_and_keeps_them_for_billing():
|
||||
"""Interim, final and Metadata frames reach the client byte for byte (no JSON round trip, non-ASCII intact,
|
||||
a binary frame first) and every JSON object frame is what the success handler gets to bill from."""
|
||||
upstream_ws = FakeUpstreamWebSocket(
|
||||
b"\x00\x01binary-first",
|
||||
DEEPGRAM_INTERIM_FRAME,
|
||||
"not json at all",
|
||||
DEEPGRAM_FINAL_FRAME,
|
||||
DEEPGRAM_METADATA_FRAME,
|
||||
)
|
||||
|
||||
websocket, success_handler = await _relay_deepgram_listen(upstream_ws, _pending_receive)
|
||||
|
||||
assert [call.args[0] for call in websocket.send_bytes.await_args_list] == [b"\x00\x01binary-first"]
|
||||
assert [call.args[0] for call in websocket.send_text.await_args_list] == [
|
||||
DEEPGRAM_INTERIM_FRAME,
|
||||
"not json at all",
|
||||
DEEPGRAM_FINAL_FRAME,
|
||||
DEEPGRAM_METADATA_FRAME,
|
||||
]
|
||||
success_call = success_handler.call_args.kwargs
|
||||
assert success_call["url_route"] == "/deepgram/v1/listen"
|
||||
assert success_call["response_body"] == [
|
||||
json.loads(DEEPGRAM_INTERIM_FRAME),
|
||||
json.loads(DEEPGRAM_FINAL_FRAME),
|
||||
json.loads(DEEPGRAM_METADATA_FRAME),
|
||||
]
|
||||
assert success_call["httpx_response"].request.url == DEEPGRAM_LISTEN_TARGET
|
||||
assert success_call["logging_obj"].model_call_details.get("custom_llm_provider") is None
|
||||
websocket.close.assert_awaited_once_with()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_sends_deepgram_audio_bytes_and_control_text_upstream_unchanged():
|
||||
upstream_ws = RecordingUpstreamWebSocket()
|
||||
audio_chunk = bytes(range(256)) * 4
|
||||
close_stream = json.dumps({"type": "CloseStream"})
|
||||
|
||||
await _relay_deepgram_listen(
|
||||
upstream_ws,
|
||||
AsyncMock(
|
||||
side_effect=[
|
||||
{"type": "websocket.receive", "bytes": audio_chunk},
|
||||
{"type": "websocket.receive", "text": close_stream},
|
||||
{"type": "websocket.disconnect"},
|
||||
]
|
||||
),
|
||||
)
|
||||
|
||||
assert [call.args[0] for call in upstream_ws.send.await_args_list] == [audio_chunk, close_stream]
|
||||
assert isinstance(upstream_ws.send.await_args_list[0].args[0], bytes)
|
||||
upstream_ws.close.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_vertex_live_setup_ack_names_the_model_but_is_not_billed_as_usage():
|
||||
"""Vertex Live keeps its special first frame: the setup acknowledgement is forwarded verbatim, read for the
|
||||
model, and left out of the frames the usage handler sees; later frames are kept as before."""
|
||||
setup_ack = json.dumps(
|
||||
{"setupComplete": {}, "model": "projects/p/locations/global/publishers/google/models/gemini-live-2.5-flash"}
|
||||
)
|
||||
server_content = json.dumps({"serverContent": {"turnComplete": True}, "usageMetadata": {"totalTokenCount": 12}})
|
||||
upstream_ws = FakeUpstreamWebSocket(setup_ack, server_content)
|
||||
websocket = _client_websocket(_pending_receive)
|
||||
|
||||
with (
|
||||
_patched_websocket_passthrough_environment(upstream_ws),
|
||||
patch( # test-quality-ok: pass_through_endpoint_logging is a module global read inside websocket_passthrough_request; there is no injection seam
|
||||
"litellm.proxy.pass_through_endpoints.pass_through_endpoints."
|
||||
"pass_through_endpoint_logging.pass_through_async_success_handler",
|
||||
new=AsyncMock(),
|
||||
) as success_handler,
|
||||
):
|
||||
await websocket_passthrough_request(
|
||||
websocket=websocket,
|
||||
target="wss://aiplatform.googleapis.com/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent",
|
||||
custom_headers={"Authorization": "Bearer token"},
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
forward_headers=False,
|
||||
endpoint="/vertex_ai/live",
|
||||
accept_websocket=False,
|
||||
)
|
||||
|
||||
assert [call.args[0] for call in websocket.send_text.await_args_list] == [setup_ack, server_content]
|
||||
success_call = success_handler.call_args.kwargs
|
||||
assert success_call["response_body"] == [json.loads(server_content)]
|
||||
assert success_call["logging_obj"].model == "gemini-live-2.5-flash"
|
||||
assert success_call["logging_obj"].model_call_details["custom_llm_provider"] == "vertex_ai_language_models"
|
||||
|
||||
|
||||
def _passthrough_kwargs_for_reservation(
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
parsed_body: dict | None = None,
|
||||
|
|
|
|||
|
|
@ -7607,8 +7607,9 @@ async def test_deleting_the_stored_pass_through_row_takes_the_route_out_of_servi
|
|||
|
||||
settings: Final = patch("litellm.proxy.proxy_server.general_settings", {}) # test-quality-ok: the method reads this module global; no injection seam
|
||||
yaml_endpoints: Final = patch("litellm.proxy.proxy_server.config_passthrough_endpoints", None) # test-quality-ok: module global holding the YAML endpoints; this case has none
|
||||
app_routes: Final = patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.SafeRouteAdder.add_api_route_if_not_exists") # test-quality-ok: the registry is the observable; a real route would stay on the shared FastAPI app for the rest of the xdist worker
|
||||
try:
|
||||
with settings, yaml_endpoints:
|
||||
with settings, yaml_endpoints, app_routes:
|
||||
pc = ProxyConfig()
|
||||
await pc._update_general_settings(db_general_settings={"pass_through_endpoints": [db_endpoint]})
|
||||
assert live_routes(), "the stored endpoint should be serving before the row is deleted"
|
||||
|
|
@ -7648,8 +7649,9 @@ async def test_a_stored_pass_through_row_never_disturbs_the_config_declared_rout
|
|||
|
||||
settings: Final = patch("litellm.proxy.proxy_server.general_settings", {"pass_through_endpoints": [config_endpoint]}) # test-quality-ok: the method reads this module global; no injection seam
|
||||
yaml_endpoints: Final = patch("litellm.proxy.proxy_server.config_passthrough_endpoints", [config_endpoint]) # test-quality-ok: module global holding the YAML endpoints the reload merges in
|
||||
app_routes: Final = patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.SafeRouteAdder.add_api_route_if_not_exists") # test-quality-ok: the registry is the observable; a real route would stay on the shared FastAPI app for the rest of the xdist worker
|
||||
try:
|
||||
with settings, yaml_endpoints:
|
||||
with settings, yaml_endpoints, app_routes:
|
||||
await initialize_pass_through_endpoints(pass_through_endpoints=[config_endpoint])
|
||||
assert live_paths() == {config_path}
|
||||
|
||||
|
|
|
|||
|
|
@ -903,6 +903,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid():
|
|||
"/v1/audio/speech",
|
||||
"/v1/ocr",
|
||||
"/vertex_ai/live",
|
||||
"/v1/listen",
|
||||
"/v1beta/interactions",
|
||||
],
|
||||
},
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue