mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-23 00:41:40 +00:00
feat(passthrough): deepgram streaming /v1/listen WebSocket passthrough with duration-based cost tracking
Adds authenticated /deepgram/v1/listen and /deepgram/listen WebSocket routes that resolve the Deepgram credential through the pass-through router, inject Authorization: Token upstream, default the model to nova-3 when the client passes none, and relay audio and transcript frames unchanged. The shared WebSocket relay no longer assumes the first upstream frame is JSON and forwards every frame as received, keeping the Vertex AI Live setup handling on Vertex routes only. A Deepgram logging handler bills the call on Metadata.duration, falling back to the furthest Results start + duration, at the deepgram/<model> per-second rate from the model cost map Resolves LIT-7937 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin
parent
351a54e849
commit
6e1b4959d1
11 changed files with 974 additions and 83 deletions
|
|
@ -317,6 +317,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,27 @@
|
|||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.constants import DEEPGRAM_DEFAULT_API_BASE, DEEPGRAM_LISTEN_DEFAULT_MODEL
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
|
||||
_WEBSOCKET_SCHEMES: Final = MappingProxyType({"https": "wss", "http": "ws", "wss": "wss", "ws": "ws"})
|
||||
|
||||
|
||||
class DeepgramException(BaseLLMException):
|
||||
pass
|
||||
|
||||
|
||||
def deepgram_listen_websocket_target(api_base: str | None, query_string: str) -> str:
|
||||
"""
|
||||
The upstream ``/listen`` socket for a streaming transcription, keeping the client's query string as sent
|
||||
and adding the default model only when the client named none
|
||||
"""
|
||||
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 = httpx.QueryParams(query_string)
|
||||
query: Final = (
|
||||
query_string if params.get("model") else str(params.remove("model").add("model", DEEPGRAM_LISTEN_DEFAULT_MODEL))
|
||||
)
|
||||
return f"{websocket_url}?{query}"
|
||||
|
|
|
|||
|
|
@ -200,6 +200,7 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = (
|
|||
"/cohere/",
|
||||
"/comprehendmedical",
|
||||
"/cursor/",
|
||||
"/deepgram/",
|
||||
"/eu.assemblyai/",
|
||||
"/gemini/",
|
||||
"/gigachat/",
|
||||
|
|
|
|||
|
|
@ -70,6 +70,7 @@ from litellm.types.utils import (
|
|||
StandardLoggingVectorStoreRequest,
|
||||
StandardPassThroughResponseObject,
|
||||
TextCompletionResponse,
|
||||
TranscriptionResponse,
|
||||
)
|
||||
from litellm.types.videos.main import VideoObject
|
||||
|
||||
|
|
@ -487,6 +488,7 @@ class LiteLLMRoutes(enum.Enum):
|
|||
"/gigachat",
|
||||
"/watsonx",
|
||||
"/nvidia_nim",
|
||||
"/deepgram",
|
||||
]
|
||||
|
||||
#########################################################
|
||||
|
|
@ -4694,6 +4696,7 @@ PassThroughEndpointLoggingResultValues = (
|
|||
| VideoObject
|
||||
| StandardPassThroughResponseObject
|
||||
| ResponsesAPIResponse
|
||||
| TranscriptionResponse
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -36,6 +36,7 @@ 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_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
|
||||
|
|
@ -2573,7 +2574,7 @@ async def _openai_websocket_refusal(
|
|||
return None
|
||||
|
||||
|
||||
class _OpenAIWebsocketRelay(Protocol):
|
||||
class _WebsocketRelay(Protocol):
|
||||
async def __call__(
|
||||
self,
|
||||
*,
|
||||
|
|
@ -2593,7 +2594,7 @@ def _proxy_general_settings() -> Mapping[str, object]:
|
|||
return general_settings
|
||||
|
||||
|
||||
def _openai_websocket_relay() -> _OpenAIWebsocketRelay:
|
||||
def _websocket_relay() -> _WebsocketRelay:
|
||||
return websocket_passthrough_request
|
||||
|
||||
|
||||
|
|
@ -2611,6 +2612,19 @@ def _proxy_model_allowlists() -> _OpenAIWebsocketModelAllowlists:
|
|||
return resolve
|
||||
|
||||
|
||||
def _negotiated_websocket_subprotocol(websocket: WebSocket) -> str | None:
|
||||
"""
|
||||
The first subprotocol the client offered, echoed back so browsers that carry the LiteLLM key in
|
||||
``Sec-WebSocket-Protocol`` complete the handshake
|
||||
"""
|
||||
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(
|
||||
|
|
@ -2618,16 +2632,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:
|
||||
|
|
@ -2686,6 +2695,47 @@ 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."
|
||||
)
|
||||
|
||||
|
||||
@router.websocket("/deepgram/v1/listen")
|
||||
@router.websocket("/deepgram/listen")
|
||||
async def deepgram_listen_websocket_route(
|
||||
websocket: WebSocket,
|
||||
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth_websocket)],
|
||||
relay: Annotated[_WebsocketRelay, Depends(_websocket_relay)],
|
||||
) -> None:
|
||||
"""
|
||||
Streaming speech to text through Deepgram's ``/v1/listen`` socket. Audio frames and transcript frames are
|
||||
relayed unchanged; the call is billed on the audio duration Deepgram reports when the socket closes
|
||||
"""
|
||||
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))
|
||||
await relay(
|
||||
websocket=websocket,
|
||||
target=deepgram_listen_websocket_target(
|
||||
api_base=get_secret_str("DEEPGRAM_API_BASE"),
|
||||
query_string=websocket.url.query,
|
||||
),
|
||||
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,132 @@
|
|||
"""
|
||||
Cost tracking for Deepgram's streaming ``/v1/listen`` WebSocket. Deepgram bills the audio it processed, which it
|
||||
reports as ``duration`` on the closing ``Metadata`` frame; a stream that ends without one is billed on the furthest
|
||||
``start + duration`` across its ``Results`` frames
|
||||
"""
|
||||
|
||||
import math
|
||||
from collections.abc import Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from urllib.parse import parse_qs, urlparse
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import DEEPGRAM_LISTEN_DEFAULT_MODEL
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.proxy._types import PassThroughEndpointLoggingTypedDict
|
||||
from litellm.types.utils import TranscriptionResponse
|
||||
|
||||
DEEPGRAM_LISTEN_ROUTE_SUFFIX: Final = "/listen"
|
||||
|
||||
|
||||
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
|
||||
)
|
||||
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
|
||||
)
|
||||
|
||||
|
||||
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 _audio_cost(response: TranscriptionResponse, model: str) -> float | None:
|
||||
try:
|
||||
return litellm.completion_cost(
|
||||
completion_response=response,
|
||||
model=model,
|
||||
custom_llm_provider=litellm.LlmProviders.DEEPGRAM.value,
|
||||
call_type="transcription",
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # an unpriced model must not lose the spend row, only its cost
|
||||
verbose_proxy_logger.warning("Deepgram listen passthrough: no pricing for model '%s': %s", model, e)
|
||||
return None
|
||||
|
||||
|
||||
class DeepgramListenPassthroughLoggingHandler:
|
||||
@staticmethod
|
||||
def is_deepgram_listen_route(url_route: str) -> bool:
|
||||
path: Final = urlparse(url_route).path
|
||||
return path.startswith("/deepgram/") 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)
|
||||
response: Final = TranscriptionResponse(text=deepgram_listen_transcript(websocket_messages))
|
||||
response._hidden_params["audio_transcription_duration"] = audio_seconds # pyright: ignore[reportPrivateUsage] # the cost calculator reads the billed duration off the response's hidden params
|
||||
response_cost: Final = _audio_cost(response, model)
|
||||
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, cost %s",
|
||||
model,
|
||||
audio_seconds,
|
||||
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 typing import TYPE_CHECKING, Any, Final, TypedDict, cast
|
||||
from urllib.parse import urlencode, urlparse
|
||||
|
||||
|
|
@ -2120,6 +2120,17 @@ 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:
|
||||
"""
|
||||
The frame as a JSON object when it is one, for cost tracking; audio and non-object frames yield 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
|
||||
|
|
@ -2401,70 +2412,46 @@ 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:
|
||||
"""
|
||||
Send the frame to the client exactly as received, then keep it for cost tracking when it is a JSON
|
||||
object; the Vertex AI Live setup acknowledgement only names the model, so it is read instead of kept
|
||||
"""
|
||||
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"""
|
||||
"""Relay upstream frames to the client until the upstream closes, returning its 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
|
||||
|
|
|
|||
|
|
@ -24,6 +24,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,
|
||||
)
|
||||
|
|
@ -278,6 +281,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
|
||||
|
|
|
|||
|
|
@ -0,0 +1,254 @@
|
|||
"""Deepgram ``/v1/listen`` WebSocket passthrough: duration extraction and duration based cost tracking."""
|
||||
|
||||
import math
|
||||
from collections.abc import Mapping, Sequence
|
||||
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,
|
||||
deepgram_listen_audio_seconds,
|
||||
deepgram_listen_model,
|
||||
deepgram_listen_transcript,
|
||||
)
|
||||
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"
|
||||
|
||||
|
||||
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) -> dict[str, object]:
|
||||
return {"type": "Metadata", "request_id": "req-1", "duration": duration, "channels": 1}
|
||||
|
||||
|
||||
@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((_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
|
||||
|
||||
|
||||
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(
|
||||
("url_route", "expected"),
|
||||
[
|
||||
("/deepgram/v1/listen", True),
|
||||
("/deepgram/listen", True),
|
||||
("/deepgram/v1/listen?model=nova-3", 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(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/{model}"]["input_cost_per_second"]
|
||||
assert per_second > 0
|
||||
return per_second * seconds
|
||||
|
||||
|
||||
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("nova-3", 12.5))
|
||||
assert handler_result["kwargs"]["response_cost"] == pytest.approx(_registry_cost("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("nova-3", 12.5))
|
||||
|
||||
|
||||
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("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_keeps_the_spend_row_but_no_cost_for_an_unpriced_model():
|
||||
handler_result = DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler(
|
||||
websocket_messages=(_metadata(12.5),),
|
||||
logging_obj=_logging_obj(),
|
||||
upstream_url="wss://api.deepgram.com/v1/listen?model=nova-99-not-in-registry",
|
||||
)
|
||||
|
||||
assert handler_result["kwargs"]["model"] == "nova-99-not-in-registry"
|
||||
assert handler_result["kwargs"]["response_cost"] is None
|
||||
assert handler_result["result"]._hidden_params["audio_transcription_duration"] == 12.5
|
||||
|
||||
|
||||
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("nova-3", 20.0))
|
||||
assert payload["metadata"]["user_api_key_team_id"] == "team-stt"
|
||||
assert payload["id"] == "call-dg-e2e"
|
||||
|
|
@ -0,0 +1,280 @@
|
|||
"""Deepgram ``/v1/listen`` passthrough WebSocket route: registration, auth, credential injection, target URL."""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType, SimpleNamespace
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
from starlette.routing import WebSocketRoute
|
||||
from starlette.websockets import WebSocketDisconnect
|
||||
|
||||
from litellm.proxy._lazy_features import LAZY_FEATURES
|
||||
from litellm.proxy._types import LiteLLMRoutes, UserAPIKeyAuth
|
||||
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
|
||||
_websocket_relay,
|
||||
deepgram_listen_websocket_route,
|
||||
router,
|
||||
)
|
||||
|
||||
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")
|
||||
|
||||
|
||||
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)
|
||||
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 == []
|
||||
|
||||
|
||||
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_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,
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
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]
|
||||
|
|
@ -4844,18 +4844,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:
|
||||
|
|
@ -4876,7 +4879,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()
|
||||
|
|
@ -4930,7 +4933,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
|
||||
|
|
@ -5359,6 +5362,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: Optional[dict] = None,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue