mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-28 06:23:07 +00:00
feat(elevenlabs): stream text-to-speech over a WebSocket pass-through and bill it per character
Adds /elevenlabs/v1/text-to-speech/{voice_id}/stream-input and /multi-stream-input. The proxy authenticates the caller with a LiteLLM key, relays the session to ElevenLabs with the proxy's own xi-api-key, and writes one spend row per session at characters times the cost-map input_cost_per_character of elevenlabs/<model_id>. A missing model_id gets eleven_multilingual_v2, an unpriced model or a bad voice_id is refused with 1008 before ElevenLabs is contacted, and caller credentials are stripped from the query string and from every frame
The shared relay takes a frame policy. ElevenLabs parses each client frame once, forwards the re-serialized frame and records its text from that same parse, and closes with 1007 on a frame it cannot read, so the forwarded text and the billed text never disagree
The shared relay also changes for every WebSocket pass-through. A client that drops while upstream frames are still streaming now ends the session normally, so the success handler logs the spend for what was used instead of dropping the whole session on the failure path
This commit is contained in:
mateo-berri
parent
76720a665c
commit
9799a1f16a
15 changed files with 1779 additions and 27 deletions
|
|
@ -104,6 +104,7 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
|
|||
"/assemblyai/",
|
||||
"/eu.assemblyai/",
|
||||
"/deepgram/",
|
||||
"/elevenlabs/",
|
||||
"/fal_ai/",
|
||||
"/langfuse/",
|
||||
"/vllm/",
|
||||
|
|
|
|||
|
|
@ -420,6 +420,10 @@ DEEPGRAM_DEFAULT_API_BASE: Final = "https://api.deepgram.com/v1"
|
|||
NADIR_DEFAULT_API_BASE: Final = "https://api.getnadir.com/v1"
|
||||
DEEPGRAM_LISTEN_DEFAULT_MODEL: Final = "nova-3"
|
||||
|
||||
ELEVENLABS_DEFAULT_API_BASE: Final = "https://api.elevenlabs.io"
|
||||
ELEVENLABS_TTS_WEBSOCKET_DEFAULT_MODEL: Final = "eleven_multilingual_v2"
|
||||
WEBSOCKET_PASSTHROUGH_CLIENT_FRAMES_KEY: Final = "client_frames"
|
||||
|
||||
BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY: Final = "litellm.bedrock_realtime.pending_opening_message"
|
||||
BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY: Final = "litellm.bedrock_realtime.session_committed"
|
||||
BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY: Final = "litellm.bedrock_realtime.committed_failure"
|
||||
|
|
|
|||
|
|
@ -1,5 +1,130 @@
|
|||
import json
|
||||
import re
|
||||
from collections.abc import Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal, TypeAlias
|
||||
from urllib.parse import parse_qs, urlparse
|
||||
|
||||
import httpx
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm.constants import ELEVENLABS_DEFAULT_API_BASE, ELEVENLABS_TTS_WEBSOCKET_DEFAULT_MODEL
|
||||
from litellm.exceptions import ModelNotMappedError
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
|
||||
ClientTextDecision,
|
||||
ForwardedClientText,
|
||||
RefusedClientText,
|
||||
)
|
||||
from litellm.types.utils import LlmProviders
|
||||
|
||||
ElevenLabsTextToSpeechStream: TypeAlias = Literal["stream-input", "multi-stream-input"]
|
||||
|
||||
ELEVENLABS_API_KEY_HEADER: Final = "xi-api-key"
|
||||
ELEVENLABS_TTS_WEBSOCKET_ROUTE_SUFFIXES: Final = ("/stream-input", "/multi-stream-input")
|
||||
ELEVENLABS_TTS_CALLER_CREDENTIAL_PARAMS: Final = frozenset({"authorization", "single_use_token"})
|
||||
ELEVENLABS_TTS_CALLER_CREDENTIAL_FIELDS: Final = frozenset({"xi_api_key", "xi-api-key", "authorization"})
|
||||
ELEVENLABS_TTS_AUDIO_FIELDS: Final = frozenset({"audio", "alignment", "normalizedAlignment"})
|
||||
_MODEL_ID_PARAM: Final = "model_id"
|
||||
_VOICE_ID: Final = re.compile(r"[A-Za-z0-9_-]+")
|
||||
_WEBSOCKET_SCHEMES: Final = MappingProxyType({"https": "wss", "http": "ws", "wss": "wss", "ws": "ws"})
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, object])
|
||||
_UNREADABLE_CLIENT_TEXT: Final = RefusedClientText(
|
||||
code=1007, reason="Each text frame must be a JSON object whose text field, when present, is a string"
|
||||
)
|
||||
|
||||
|
||||
class ElevenLabsException(BaseLLMException):
|
||||
pass
|
||||
|
||||
|
||||
def elevenlabs_tts_requested_model(query_string: str) -> str:
|
||||
return httpx.QueryParams(query_string).get(_MODEL_ID_PARAM) or ELEVENLABS_TTS_WEBSOCKET_DEFAULT_MODEL
|
||||
|
||||
|
||||
def is_elevenlabs_voice_id(voice_id: str) -> bool:
|
||||
return _VOICE_ID.fullmatch(voice_id) is not None
|
||||
|
||||
|
||||
def _upstream_params(query_string: str) -> httpx.QueryParams:
|
||||
"""Auth and pricing read the first ``model_id``, so ElevenLabs must not see a second one, and the proxy's own
|
||||
key is the only upstream credential, so a caller's ``authorization`` or ``single_use_token`` never leaves."""
|
||||
items: Final = httpx.QueryParams(query_string).multi_items()
|
||||
first_model_index: Final = next((index for index, (key, _) in enumerate(items) if key == _MODEL_ID_PARAM), None)
|
||||
return httpx.QueryParams(
|
||||
tuple(
|
||||
(key, value)
|
||||
for index, (key, value) in enumerate(items)
|
||||
if key not in ELEVENLABS_TTS_CALLER_CREDENTIAL_PARAMS
|
||||
and (key != _MODEL_ID_PARAM or index == first_model_index)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def elevenlabs_tts_websocket_target(
|
||||
api_base: str | None, voice_id: str, stream: ElevenLabsTextToSpeechStream, query_string: str
|
||||
) -> str:
|
||||
base_url: Final = (api_base or ELEVENLABS_DEFAULT_API_BASE).rstrip("/")
|
||||
stream_url: Final = httpx.URL(f"{base_url}/v1/text-to-speech/{voice_id}/{stream}")
|
||||
websocket_url: Final = stream_url.copy_with(scheme=_WEBSOCKET_SCHEMES.get(stream_url.scheme, stream_url.scheme))
|
||||
params: Final = _upstream_params(query_string)
|
||||
query: Final = (
|
||||
params
|
||||
if params.get(_MODEL_ID_PARAM)
|
||||
else params.remove(_MODEL_ID_PARAM).add(_MODEL_ID_PARAM, ELEVENLABS_TTS_WEBSOCKET_DEFAULT_MODEL)
|
||||
)
|
||||
return f"{websocket_url}?{query}"
|
||||
|
||||
|
||||
def elevenlabs_tts_model(upstream_url: str) -> str:
|
||||
models: Final = parse_qs(urlparse(upstream_url).query).get(_MODEL_ID_PARAM)
|
||||
return models[0] if models else ELEVENLABS_TTS_WEBSOCKET_DEFAULT_MODEL
|
||||
|
||||
|
||||
def elevenlabs_tts_registry_key(upstream_url: str) -> str:
|
||||
return f"{LlmProviders.ELEVENLABS.value}/{elevenlabs_tts_model(upstream_url)}"
|
||||
|
||||
|
||||
def elevenlabs_tts_is_priced(upstream_url: str) -> bool:
|
||||
"""Only an exact registry hit with a per-character rate counts: the speech cost path bills a session by
|
||||
``input_cost_per_character`` and treats a missing or zero rate as unpriced."""
|
||||
registry_key: Final = elevenlabs_tts_registry_key(upstream_url)
|
||||
try:
|
||||
model_info: Final = litellm.get_model_info(
|
||||
model=registry_key, custom_llm_provider=LlmProviders.ELEVENLABS.value
|
||||
)
|
||||
except ModelNotMappedError:
|
||||
return False
|
||||
return model_info["key"] == registry_key and bool(model_info.get("input_cost_per_character"))
|
||||
|
||||
|
||||
def _json_object(text: str) -> Mapping[str, object] | None:
|
||||
try:
|
||||
return _JSON_OBJECT.validate_json(text)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def elevenlabs_tts_client_text(text: str) -> ClientTextDecision:
|
||||
"""The upstream gets the parsed frame re-serialized, never the raw text, so ElevenLabs reads exactly the
|
||||
``text`` that was counted and no caller credential field survives an escape or a parser disagreement."""
|
||||
frame: Final = _json_object(text)
|
||||
if frame is None:
|
||||
return _UNREADABLE_CLIENT_TEXT
|
||||
frame_text: Final = frame.get("text")
|
||||
if frame_text is not None and not isinstance(frame_text, str):
|
||||
return _UNREADABLE_CLIENT_TEXT
|
||||
forwarded: Final = json.dumps(
|
||||
{key: value for key, value in frame.items() if key not in ELEVENLABS_TTS_CALLER_CREDENTIAL_FIELDS},
|
||||
ensure_ascii=False,
|
||||
)
|
||||
return ForwardedClientText(text=forwarded, recorded={"text": frame_text} if frame_text is not None else None)
|
||||
|
||||
|
||||
def elevenlabs_tts_recorded_upstream_frame(frame: Mapping[str, object]) -> Mapping[str, object]:
|
||||
return {key: value for key, value in frame.items() if key not in ELEVENLABS_TTS_AUDIO_FIELDS}
|
||||
|
||||
|
||||
def elevenlabs_tts_characters(client_frames: Sequence[Mapping[str, object]]) -> int:
|
||||
return sum(len(text) for frame in client_frames if isinstance(text := frame.get("text"), str))
|
||||
|
|
|
|||
|
|
@ -221,6 +221,7 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = (
|
|||
"/comprehendmedical",
|
||||
"/cursor/",
|
||||
"/deepgram/",
|
||||
"/elevenlabs/",
|
||||
"/eu.assemblyai/",
|
||||
"/fal_ai/",
|
||||
"/gemini/",
|
||||
|
|
|
|||
|
|
@ -420,6 +420,7 @@ class LiteLLMRoutes(enum.Enum):
|
|||
"/nvidia_nim",
|
||||
"/deepgram",
|
||||
"/fal_ai",
|
||||
"/elevenlabs",
|
||||
]
|
||||
|
||||
#########################################################
|
||||
|
|
|
|||
|
|
@ -64,6 +64,17 @@ from litellm.llms.deepgram.common_utils import (
|
|||
deepgram_listen_requested_model,
|
||||
deepgram_listen_websocket_target,
|
||||
)
|
||||
from litellm.llms.elevenlabs.common_utils import (
|
||||
ELEVENLABS_API_KEY_HEADER,
|
||||
ElevenLabsTextToSpeechStream,
|
||||
elevenlabs_tts_client_text,
|
||||
elevenlabs_tts_is_priced,
|
||||
elevenlabs_tts_recorded_upstream_frame,
|
||||
elevenlabs_tts_registry_key,
|
||||
elevenlabs_tts_requested_model,
|
||||
elevenlabs_tts_websocket_target,
|
||||
is_elevenlabs_voice_id,
|
||||
)
|
||||
from litellm.llms.fal_ai.cost_calculator import fal_ai_passthrough_cost, fal_ai_queue_base
|
||||
from litellm.llms.nvidia_nim.passthrough.transformation import nvidia_nim_model_group_in_path
|
||||
from litellm.llms.openai.common_utils import OpenAIError as LiteLLMOpenAIError
|
||||
|
|
@ -110,6 +121,7 @@ from litellm.proxy.pass_through_endpoints.injected_stream_usage import (
|
|||
)
|
||||
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
|
||||
HttpPassThroughEndpointHelpers,
|
||||
WebsocketFramePolicy,
|
||||
create_pass_through_route,
|
||||
create_websocket_passthrough_route,
|
||||
websocket_passthrough_request,
|
||||
|
|
@ -3341,6 +3353,7 @@ class _WebsocketRelay(Protocol):
|
|||
forward_headers: bool,
|
||||
endpoint: str,
|
||||
accept_websocket: bool,
|
||||
frame_policy: WebsocketFramePolicy = ...,
|
||||
) -> None: ...
|
||||
|
||||
|
||||
|
|
@ -3502,6 +3515,92 @@ async def deepgram_listen_websocket_route(
|
|||
)
|
||||
|
||||
|
||||
_ELEVENLABS_WS_MISSING_KEY_REASON: Final = (
|
||||
"Required 'ELEVENLABS_API_KEY' in environment to make pass-through calls to ElevenLabs."
|
||||
)
|
||||
_ELEVENLABS_WS_VOICE_ID_REASON: Final = "voice_id may only contain letters, digits, '_' and '-'"
|
||||
_ELEVENLABS_WS_UNPRICED_REASON: Final = (
|
||||
"No per-character price for '{registry_key}': add it to the model cost map to enable it"
|
||||
)
|
||||
_ELEVENLABS_TTS_FRAME_POLICY: Final = WebsocketFramePolicy(
|
||||
client_text=elevenlabs_tts_client_text,
|
||||
recorded_upstream_frame=elevenlabs_tts_recorded_upstream_frame,
|
||||
)
|
||||
|
||||
|
||||
async def elevenlabs_tts_user_api_key_auth(websocket: WebSocket) -> UserAPIKeyAuth:
|
||||
return await user_api_key_auth_websocket_for_model(
|
||||
websocket,
|
||||
model=elevenlabs_tts_requested_model(websocket.url.query),
|
||||
api_key_headers=("api-key", ELEVENLABS_API_KEY_HEADER),
|
||||
)
|
||||
|
||||
|
||||
async def _elevenlabs_tts_websocket(
|
||||
websocket: WebSocket,
|
||||
voice_id: str,
|
||||
stream: ElevenLabsTextToSpeechStream,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
relay: _WebsocketRelay,
|
||||
) -> None:
|
||||
elevenlabs_api_key: Final = passthrough_endpoint_router.get_credentials(
|
||||
custom_llm_provider=litellm.LlmProviders.ELEVENLABS.value,
|
||||
region_name=None,
|
||||
)
|
||||
if elevenlabs_api_key is None:
|
||||
await websocket.close(code=1011, reason=_ELEVENLABS_WS_MISSING_KEY_REASON)
|
||||
return
|
||||
|
||||
await websocket.accept(subprotocol=_negotiated_websocket_subprotocol(websocket))
|
||||
if not is_elevenlabs_voice_id(voice_id):
|
||||
await websocket.close(code=1008, reason=_ELEVENLABS_WS_VOICE_ID_REASON)
|
||||
return
|
||||
|
||||
target: Final = elevenlabs_tts_websocket_target(
|
||||
api_base=get_secret_str("ELEVENLABS_API_BASE"),
|
||||
voice_id=voice_id,
|
||||
stream=stream,
|
||||
query_string=websocket.url.query,
|
||||
)
|
||||
if not elevenlabs_tts_is_priced(target):
|
||||
await websocket.close(
|
||||
code=1008,
|
||||
reason=_ELEVENLABS_WS_UNPRICED_REASON.format(registry_key=elevenlabs_tts_registry_key(target)),
|
||||
)
|
||||
return
|
||||
|
||||
await relay(
|
||||
websocket=websocket,
|
||||
target=target,
|
||||
custom_headers={ELEVENLABS_API_KEY_HEADER: elevenlabs_api_key},
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
forward_headers=False,
|
||||
endpoint=websocket.url.path,
|
||||
accept_websocket=False,
|
||||
frame_policy=_ELEVENLABS_TTS_FRAME_POLICY,
|
||||
)
|
||||
|
||||
|
||||
@router.websocket("/elevenlabs/v1/text-to-speech/{voice_id}/stream-input")
|
||||
async def elevenlabs_tts_stream_input_websocket_route(
|
||||
websocket: WebSocket,
|
||||
voice_id: str,
|
||||
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(elevenlabs_tts_user_api_key_auth)],
|
||||
relay: Annotated[_WebsocketRelay, Depends(_websocket_relay)],
|
||||
) -> None:
|
||||
await _elevenlabs_tts_websocket(websocket, voice_id, "stream-input", user_api_key_dict, relay)
|
||||
|
||||
|
||||
@router.websocket("/elevenlabs/v1/text-to-speech/{voice_id}/multi-stream-input")
|
||||
async def elevenlabs_tts_multi_stream_input_websocket_route(
|
||||
websocket: WebSocket,
|
||||
voice_id: str,
|
||||
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(elevenlabs_tts_user_api_key_auth)],
|
||||
relay: Annotated[_WebsocketRelay, Depends(_websocket_relay)],
|
||||
) -> None:
|
||||
await _elevenlabs_tts_websocket(websocket, voice_id, "multi-stream-input", user_api_key_dict, relay)
|
||||
|
||||
|
||||
class BaseOpenAIPassThroughHandler:
|
||||
@staticmethod
|
||||
async def base_openai_pass_through_handler(
|
||||
|
|
|
|||
|
|
@ -0,0 +1,89 @@
|
|||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import WEBSOCKET_PASSTHROUGH_CLIENT_FRAMES_KEY
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.elevenlabs.common_utils import (
|
||||
ELEVENLABS_TTS_WEBSOCKET_ROUTE_SUFFIXES,
|
||||
elevenlabs_tts_characters,
|
||||
elevenlabs_tts_is_priced,
|
||||
elevenlabs_tts_model,
|
||||
elevenlabs_tts_registry_key,
|
||||
)
|
||||
from litellm.proxy._types import PassThroughEndpointLoggingTypedDict
|
||||
from litellm.types.utils import StandardPassThroughResponseObject
|
||||
|
||||
_RECORDED_FRAMES: Final = TypeAdapter(tuple[dict[str, object], ...])
|
||||
|
||||
|
||||
def _speech_cost(upstream_url: str, characters: int) -> float | None:
|
||||
if not elevenlabs_tts_is_priced(upstream_url):
|
||||
verbose_proxy_logger.warning(
|
||||
"ElevenLabs TTS passthrough: no per-character registry entry '%s'",
|
||||
elevenlabs_tts_registry_key(upstream_url),
|
||||
)
|
||||
return None
|
||||
prompt_cost, completion_cost = litellm.cost_per_token(
|
||||
model=elevenlabs_tts_model(upstream_url),
|
||||
custom_llm_provider=litellm.LlmProviders.ELEVENLABS.value,
|
||||
prompt_characters=characters,
|
||||
call_type="speech",
|
||||
)
|
||||
return prompt_cost + completion_cost
|
||||
|
||||
|
||||
class ElevenLabsTextToSpeechPassthroughLoggingHandler:
|
||||
@staticmethod
|
||||
def is_elevenlabs_tts_websocket_route(url_route: str) -> bool:
|
||||
path: Final = urlparse(url_route).path
|
||||
return "/elevenlabs/" in path and path.endswith(ELEVENLABS_TTS_WEBSOCKET_ROUTE_SUFFIXES)
|
||||
|
||||
def elevenlabs_tts_passthrough_handler(
|
||||
self,
|
||||
request_body: Mapping[str, object],
|
||||
upstream_frame_count: int,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
upstream_url: str,
|
||||
kwargs: Mapping[str, object] = MappingProxyType({}),
|
||||
) -> PassThroughEndpointLoggingTypedDict:
|
||||
client_frames: Final = _RECORDED_FRAMES.validate_python(
|
||||
request_body.get(WEBSOCKET_PASSTHROUGH_CLIENT_FRAMES_KEY, ())
|
||||
)
|
||||
model: Final = elevenlabs_tts_model(upstream_url)
|
||||
characters: Final = elevenlabs_tts_characters(client_frames)
|
||||
response_cost: Final = _speech_cost(upstream_url, characters)
|
||||
|
||||
provider: Final = litellm.LlmProviders.ELEVENLABS.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(
|
||||
"ElevenLabs TTS passthrough cost tracking: model %s, characters %s, cost %s",
|
||||
model,
|
||||
characters,
|
||||
response_cost,
|
||||
)
|
||||
logging_result: Final[PassThroughEndpointLoggingTypedDict] = {
|
||||
"result": StandardPassThroughResponseObject(
|
||||
response={
|
||||
"model": model,
|
||||
"characters": characters,
|
||||
"client_text_frames": len(client_frames),
|
||||
"upstream_frames": upstream_frame_count,
|
||||
}
|
||||
),
|
||||
"kwargs": {
|
||||
**kwargs,
|
||||
"model": model,
|
||||
"custom_llm_provider": provider,
|
||||
"response_cost": response_cost,
|
||||
},
|
||||
}
|
||||
return logging_result
|
||||
|
|
@ -29,7 +29,7 @@ from fastapi.responses import StreamingResponse
|
|||
from pydantic import TypeAdapter
|
||||
from starlette.datastructures import UploadFile as StarletteUploadFile
|
||||
from starlette.routing import BaseRoute, Route
|
||||
from starlette.websockets import WebSocketState
|
||||
from starlette.websockets import WebSocketDisconnect, WebSocketState
|
||||
from websockets.asyncio.client import connect
|
||||
from websockets.exceptions import (
|
||||
ConnectionClosedError,
|
||||
|
|
@ -47,6 +47,7 @@ from litellm.constants import (
|
|||
REDACTED_BY_LITELLM,
|
||||
SESSION_ID_OMITTED_METADATA_KEY,
|
||||
WEBSOCKET_CLOSE_REASON_MAX_BYTES,
|
||||
WEBSOCKET_PASSTHROUGH_CLIENT_FRAMES_KEY,
|
||||
)
|
||||
from litellm.integrations.custom_guardrail import (
|
||||
CustomGuardrail,
|
||||
|
|
@ -136,8 +137,11 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import (
|
|||
LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY,
|
||||
LITELLM_PASS_THROUGH_ENDPOINT_MARKER,
|
||||
LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY,
|
||||
ClientTextDecision,
|
||||
EndpointType,
|
||||
ForwardedClientText,
|
||||
PassthroughStandardLoggingPayload,
|
||||
RefusedClientText,
|
||||
)
|
||||
from litellm.types.utils import TRUSTED_CALLBACK_VARS_FIELD, Usage
|
||||
|
||||
|
|
@ -2446,6 +2450,26 @@ def _upstream_close_to_relay(task_results: Iterable[object]) -> Close | None:
|
|||
_WEBSOCKET_FORWARDED_HEADERS: Final = frozenset(("x-goog-user-project",))
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class WebsocketFramePolicy:
|
||||
client_text: Callable[[str], ClientTextDecision]
|
||||
recorded_upstream_frame: Callable[[Mapping[str, object]], Mapping[str, object] | None]
|
||||
|
||||
|
||||
def _forward_client_text_unchanged(text: str) -> ClientTextDecision:
|
||||
return ForwardedClientText(text=text, recorded=None)
|
||||
|
||||
|
||||
def _record_upstream_frame_unchanged(frame: Mapping[str, object]) -> Mapping[str, object]:
|
||||
return frame
|
||||
|
||||
|
||||
RELAY_FRAMES_UNCHANGED: Final = WebsocketFramePolicy(
|
||||
client_text=_forward_client_text_unchanged,
|
||||
recorded_upstream_frame=_record_upstream_frame_unchanged,
|
||||
)
|
||||
|
||||
|
||||
def _with_trace_context(headers: Mapping[str, str], parent_span: object) -> dict[str, str]:
|
||||
try:
|
||||
from litellm.integrations.otel.plumbing.context import inject_trace_context
|
||||
|
|
@ -2464,6 +2488,7 @@ async def websocket_passthrough_request(
|
|||
cost_per_request: float | None = None,
|
||||
accept_websocket: bool = True,
|
||||
setup_model_rewriter: Callable[[str], str] | None = None,
|
||||
frame_policy: WebsocketFramePolicy = RELAY_FRAMES_UNCHANGED,
|
||||
):
|
||||
"""
|
||||
WebSocket passthrough request handler.
|
||||
|
|
@ -2477,6 +2502,7 @@ async def websocket_passthrough_request(
|
|||
endpoint: The endpoint path (for logging purposes)
|
||||
cost_per_request: Optional field - cost per request to the target endpoint
|
||||
setup_model_rewriter: Optional rewrite of the setup frame's model before it reaches the upstream
|
||||
frame_policy: What the route forwards of each client text frame and records of client and upstream frames
|
||||
"""
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.proxy.proxy_server import proxy_config, proxy_logging_obj
|
||||
|
|
@ -2487,6 +2513,7 @@ async def websocket_passthrough_request(
|
|||
# Initialize tracking variables
|
||||
start_time: Final = datetime.now()
|
||||
websocket_messages: Final[list[dict[str, object]]] = []
|
||||
client_frames: Final[list[Mapping[str, object]]] = [] # mutable-ok: the relay appends each frame as it arrives
|
||||
litellm_call_id: Final = str(uuid.uuid4())
|
||||
|
||||
verbose_proxy_logger.info("WebSocket passthrough (%s): Starting WebSocket connection to %s", endpoint, target)
|
||||
|
|
@ -2617,6 +2644,12 @@ async def websocket_passthrough_request(
|
|||
bytes_data: bytes | None = message.get("bytes")
|
||||
|
||||
if text_data is not None:
|
||||
if isinstance(decision := frame_policy.client_text(text_data), RefusedClientText):
|
||||
await websocket.close(code=decision.code, reason=decision.reason)
|
||||
await upstream_ws.close()
|
||||
break
|
||||
if decision.recorded is not None:
|
||||
client_frames.append(decision.recorded)
|
||||
# Try to extract model from client setup message for Vertex AI Live
|
||||
if endpoint and "/vertex_ai/live" in endpoint:
|
||||
verbose_proxy_logger.debug(
|
||||
|
|
@ -2673,7 +2706,9 @@ async def websocket_passthrough_request(
|
|||
)
|
||||
# Not a JSON message or doesn't contain setup data
|
||||
|
||||
await upstream_ws.send(_rewrite_vertex_live_setup_model(text_data, setup_model_rewriter))
|
||||
await upstream_ws.send(
|
||||
_rewrite_vertex_live_setup_model(decision.text, setup_model_rewriter)
|
||||
)
|
||||
elif bytes_data is not None:
|
||||
await upstream_ws.send(bytes_data)
|
||||
except asyncio.CancelledError:
|
||||
|
|
@ -2702,26 +2737,34 @@ async def websocket_passthrough_request(
|
|||
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:
|
||||
async def relay_upstream_frame(upstream_message: str | bytes) -> bool:
|
||||
if not _client_socket_is_open(websocket):
|
||||
return False
|
||||
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
|
||||
return True
|
||||
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)
|
||||
return True
|
||||
recorded_frame: Final = frame_policy.recorded_upstream_frame(message_data)
|
||||
if recorded_frame is not None:
|
||||
websocket_messages.append(dict(recorded_frame))
|
||||
return True
|
||||
|
||||
async def forward_upstream_to_client() -> Close | None:
|
||||
try:
|
||||
while True:
|
||||
await relay_upstream_frame(await upstream_ws.recv())
|
||||
while await relay_upstream_frame(await upstream_ws.recv()):
|
||||
pass
|
||||
except (ConnectionClosedOK, ConnectionClosedError) as e:
|
||||
verbose_proxy_logger.debug("Upstream WebSocket connection closed: %s", e)
|
||||
return e.rcvd
|
||||
except WebSocketDisconnect as e:
|
||||
verbose_proxy_logger.debug("Client WebSocket disconnected mid-relay: %s", e.code)
|
||||
return None
|
||||
except asyncio.CancelledError:
|
||||
verbose_proxy_logger.debug("asyncio.CancelledError in forward_upstream_to_client")
|
||||
raise
|
||||
|
|
@ -2731,6 +2774,7 @@ async def websocket_passthrough_request(
|
|||
"WebSocket passthrough (%s): error forwarding upstream message", endpoint
|
||||
)
|
||||
raise
|
||||
return None
|
||||
|
||||
# Create tasks for bidirectional message forwarding
|
||||
tasks: Final = [
|
||||
|
|
@ -2764,6 +2808,10 @@ async def websocket_passthrough_request(
|
|||
end_time: Final = datetime.now()
|
||||
|
||||
# Update passthrough logging payload with response data
|
||||
request_body: Final[dict[str, object]] = ( # mutable-ok: the success handler and payload take a dict
|
||||
{WEBSOCKET_PASSTHROUGH_CLIENT_FRAMES_KEY: client_frames} if client_frames else {}
|
||||
)
|
||||
passthrough_logging_payload["request_body"] = request_body
|
||||
passthrough_logging_payload["response_body"] = websocket_messages
|
||||
passthrough_logging_payload["end_time"] = end_time
|
||||
|
||||
|
|
@ -2812,7 +2860,7 @@ async def websocket_passthrough_request(
|
|||
end_time=end_time,
|
||||
logging_obj=logging_obj,
|
||||
cache_hit=False,
|
||||
request_body={},
|
||||
request_body=request_body,
|
||||
**success_kwargs,
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,14 +1,16 @@
|
|||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final
|
||||
from typing import Any, Final, Literal, TypeAlias
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import httpx
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.constants import AZURE_SPEECH_CUSTOM_LLM_PROVIDER
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.proxy._types import PassThroughEndpointLoggingResultValues
|
||||
from litellm.proxy._types import PassThroughEndpointLoggingResultValues, PassThroughEndpointLoggingTypedDict
|
||||
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
|
||||
PassthroughStandardLoggingPayload,
|
||||
)
|
||||
|
|
@ -29,6 +31,9 @@ from .llm_provider_handlers.cursor_passthrough_logging_handler import (
|
|||
from .llm_provider_handlers.deepgram_listen_passthrough_logging_handler import (
|
||||
DeepgramListenPassthroughLoggingHandler,
|
||||
)
|
||||
from .llm_provider_handlers.elevenlabs_tts_passthrough_logging_handler import (
|
||||
ElevenLabsTextToSpeechPassthroughLoggingHandler,
|
||||
)
|
||||
from .llm_provider_handlers.fal_ai_passthrough_logging_handler import (
|
||||
FalAIPassthroughLoggingHandler,
|
||||
)
|
||||
|
|
@ -64,6 +69,44 @@ def _safe_response_text(httpx_response: httpx.Response) -> str:
|
|||
return ""
|
||||
|
||||
|
||||
AudioWebsocketRoute: TypeAlias = Literal["deepgram_listen", "elevenlabs_tts"]
|
||||
_REQUEST_BODY: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _audio_websocket_route(url_route: str) -> AudioWebsocketRoute | None:
|
||||
if DeepgramListenPassthroughLoggingHandler.is_deepgram_listen_route(url_route):
|
||||
return "deepgram_listen"
|
||||
if ElevenLabsTextToSpeechPassthroughLoggingHandler.is_elevenlabs_tts_websocket_route(url_route):
|
||||
return "elevenlabs_tts"
|
||||
return None
|
||||
|
||||
|
||||
def _audio_websocket_handler_result(
|
||||
route: AudioWebsocketRoute,
|
||||
httpx_response: httpx.Response,
|
||||
request_body: Mapping[str, object],
|
||||
upstream_frames: Sequence[Mapping[str, object]],
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
kwargs: Mapping[str, object],
|
||||
) -> PassThroughEndpointLoggingTypedDict:
|
||||
match route:
|
||||
case "deepgram_listen":
|
||||
return DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler(
|
||||
websocket_messages=tuple(message for message in upstream_frames if isinstance(message, dict)),
|
||||
logging_obj=logging_obj,
|
||||
upstream_url=str(httpx_response.request.url),
|
||||
kwargs=kwargs,
|
||||
)
|
||||
case "elevenlabs_tts":
|
||||
return ElevenLabsTextToSpeechPassthroughLoggingHandler().elevenlabs_tts_passthrough_handler(
|
||||
request_body=request_body,
|
||||
upstream_frame_count=len(upstream_frames),
|
||||
logging_obj=logging_obj,
|
||||
upstream_url=str(httpx_response.request.url),
|
||||
kwargs=kwargs,
|
||||
)
|
||||
|
||||
|
||||
class PassThroughEndpointLogging:
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -385,21 +428,17 @@ 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,
|
||||
)
|
||||
elif (audio_websocket_route := _audio_websocket_route(url_route)) is not None:
|
||||
audio_websocket_result: Final = _audio_websocket_handler_result(
|
||||
route=audio_websocket_route,
|
||||
httpx_response=httpx_response,
|
||||
request_body=_REQUEST_BODY.validate_python(request_body),
|
||||
upstream_frames=response_body if isinstance(response_body, list) else (),
|
||||
logging_obj=logging_obj,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
standard_logging_response_object = deepgram_handler_result["result"] # rebind-ok: elif-chain
|
||||
kwargs = deepgram_handler_result["kwargs"] # rebind-ok: elif-chain contract
|
||||
standard_logging_response_object = audio_websocket_result["result"] # rebind-ok: elif-chain
|
||||
kwargs = audio_websocket_result["kwargs"] # rebind-ok: elif-chain contract
|
||||
elif FalAIPassthroughLoggingHandler.is_fal_ai_route(url_route, custom_llm_provider):
|
||||
fal_ai_handler_result: Final = FalAIPassthroughLoggingHandler().fal_ai_passthrough_handler(
|
||||
response_body=response_body if isinstance(response_body, dict) else MappingProxyType({}),
|
||||
|
|
|
|||
|
|
@ -1,5 +1,7 @@
|
|||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from typing import Final
|
||||
from typing import Final, TypeAlias
|
||||
|
||||
from pydantic import Field
|
||||
from typing_extensions import TypedDict
|
||||
|
|
@ -28,6 +30,21 @@ LITELLM_PASS_THROUGH_ENDPOINT_MARKER: Final = "__litellm_pass_through_endpoint__
|
|||
SETTLES_BACKGROUND_RESPONSE_KWARG: Final = "settles_background_response"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ForwardedClientText:
|
||||
text: str
|
||||
recorded: Mapping[str, object] | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RefusedClientText:
|
||||
code: int
|
||||
reason: str
|
||||
|
||||
|
||||
ClientTextDecision: TypeAlias = ForwardedClientText | RefusedClientText
|
||||
|
||||
|
||||
class BackgroundResponseSettlement(str, Enum):
|
||||
SETTLES = "settles"
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,227 @@
|
|||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import threading
|
||||
import uuid
|
||||
from collections.abc import AsyncIterator, Iterator
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from hashlib import sha256
|
||||
from pathlib import Path
|
||||
from queue import SimpleQueue
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
import websockets
|
||||
from integration._support.client import eventually, gateway_from_environment
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.ports import allocate_port
|
||||
from integration._support.process import OwnedProxy, owned_proxy_process
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
from websockets.asyncio.server import ServerConnection, serve
|
||||
from websockets.exceptions import ConnectionClosed
|
||||
|
||||
pytestmark: Final = [pytest.mark.timeout(180), pytest.mark.shares_worker_state]
|
||||
|
||||
PROVIDER_KEY: Final = "synthetic-elevenlabs-key"
|
||||
FLASH: Final = "eleven_flash_v2_5"
|
||||
DEFAULT_MODEL: Final = "eleven_multilingual_v2"
|
||||
AUDIO: Final = base64.b64encode(bytes(range(48))).decode()
|
||||
VOICE_SETTINGS: Final = {"stability": 0.5, "similarity_boost": 0.8}
|
||||
SENTENCES: Final = ("Hello from the gateway. ", "Hola, ¿qué tal? ")
|
||||
SPEND_SQL: Final = 'SELECT spend, model, custom_llm_provider FROM "LiteLLM_SpendLogs" WHERE api_key = %s'
|
||||
COST_MAP_PATH: Final = Path(__file__).resolve().parents[3] / "model_prices_and_context_window.json"
|
||||
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
COST_MAP: Final = TypeAdapter(dict[str, dict[str, object]])
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Upgrade:
|
||||
target: str
|
||||
provider_key: str | None
|
||||
frames: tuple[dict[str, JsonValue], ...]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Peer:
|
||||
url: str
|
||||
upgrades: SimpleQueue[Upgrade]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Session:
|
||||
frames: tuple[dict[str, JsonValue], ...]
|
||||
close_code: int | None
|
||||
close_reason: str
|
||||
|
||||
|
||||
def _per_character(model_id: str) -> float:
|
||||
row: Final = COST_MAP.validate_json(COST_MAP_PATH.read_bytes())[f"elevenlabs/{model_id}"]
|
||||
rate: Final = row["input_cost_per_character"]
|
||||
assert isinstance(rate, float) and rate > 0, rate
|
||||
return rate
|
||||
|
||||
|
||||
async def _client_frames(connection: ServerConnection) -> AsyncIterator[dict[str, JsonValue]]:
|
||||
async for raw in connection:
|
||||
frame: Final = JSON_OBJECT.validate_json(raw)
|
||||
yield frame
|
||||
text: Final = frame.get("text")
|
||||
if text == "":
|
||||
await connection.send(json.dumps({"isFinal": True}))
|
||||
return
|
||||
if isinstance(text, str) and text.strip():
|
||||
await connection.send(
|
||||
json.dumps(
|
||||
{
|
||||
"audio": AUDIO,
|
||||
"alignment": {"chars": list(text), "charStartTimesMs": [0], "charDurationsMs": [40]},
|
||||
"isFinal": None,
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
async def _speak(connection: ServerConnection, upgrades: SimpleQueue[Upgrade]) -> None:
|
||||
request: Final = connection.request
|
||||
assert request is not None
|
||||
frames: Final = tuple([frame async for frame in _client_frames(connection)])
|
||||
upgrades.put(Upgrade(request.path, request.headers.get("xi-api-key"), frames))
|
||||
|
||||
|
||||
async def _serve(port: int, upgrades: SimpleQueue[Upgrade], ready: SimpleQueue[bool], stop: asyncio.Event) -> None:
|
||||
async with serve(lambda connection: _speak(connection, upgrades), "127.0.0.1", port):
|
||||
ready.put(True)
|
||||
await stop.wait()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def elevenlabs_peer() -> Iterator[Peer]:
|
||||
port: Final = allocate_port()
|
||||
loop: Final = asyncio.new_event_loop()
|
||||
stop: Final = asyncio.Event()
|
||||
upgrades: Final = SimpleQueue[Upgrade]()
|
||||
ready: Final = SimpleQueue[bool]()
|
||||
serving: Final = _serve(port, upgrades, ready, stop)
|
||||
thread: Final = threading.Thread(target=loop.run_until_complete, args=(serving,), daemon=True)
|
||||
thread.start()
|
||||
try:
|
||||
assert ready.get(timeout=10)
|
||||
yield Peer(f"http://127.0.0.1:{port}", upgrades)
|
||||
finally:
|
||||
loop.call_soon_threadsafe(stop.set)
|
||||
thread.join(timeout=10)
|
||||
loop.close()
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def elevenlabs() -> Iterator[Peer]:
|
||||
with elevenlabs_peer() as peer:
|
||||
yield peer
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def owned(elevenlabs: Peer, tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedProxy]:
|
||||
directory: Final = tmp_path_factory.mktemp("elevenlabs-tts")
|
||||
overrides: Final = {"ELEVENLABS_API_KEY": PROVIDER_KEY, "ELEVENLABS_API_BASE": elevenlabs.url}
|
||||
with gateway_from_environment() as gateway:
|
||||
with owned_proxy_process(gateway, directory, overrides) as proxy:
|
||||
yield proxy
|
||||
|
||||
|
||||
def _ws_url(proxy: OwnedProxy, voice_id: str, query: str) -> str:
|
||||
base: Final = str(proxy.gateway.client.base_url).rstrip("/").replace("http://", "ws://")
|
||||
return f"{base}/elevenlabs/v1/text-to-speech/{voice_id}/stream-input?{query}"
|
||||
|
||||
|
||||
async def _received(connection: websockets.ClientConnection) -> AsyncIterator[dict[str, JsonValue]]:
|
||||
async for raw in connection:
|
||||
yield JSON_OBJECT.validate_json(raw)
|
||||
|
||||
|
||||
async def _exchange(
|
||||
connection: websockets.ClientConnection, key: str, texts: tuple[str, ...]
|
||||
) -> tuple[dict[str, JsonValue], ...]:
|
||||
try:
|
||||
await connection.send(json.dumps({"text": " ", "voice_settings": VOICE_SETTINGS, "xi_api_key": key}))
|
||||
for text in texts:
|
||||
await connection.send(json.dumps({"text": text}, ensure_ascii=False))
|
||||
await connection.send(json.dumps({"text": ""}))
|
||||
return tuple([frame async for frame in _received(connection)])
|
||||
except ConnectionClosed:
|
||||
return ()
|
||||
|
||||
|
||||
async def _stream(url: str, key: str, texts: tuple[str, ...]) -> Session:
|
||||
async with websockets.connect(url, additional_headers={"xi-api-key": key}, open_timeout=30) as connection:
|
||||
frames: Final = await _exchange(connection, key, texts)
|
||||
return Session(frames, connection.close_code, connection.close_reason or "")
|
||||
|
||||
|
||||
def _session(url: str, key: str, texts: tuple[str, ...]) -> Session:
|
||||
return asyncio.run(asyncio.wait_for(_stream(url, key, texts), 60))
|
||||
|
||||
|
||||
def _upgrade_for(peer: Peer, voice_id: str) -> Upgrade:
|
||||
upgrade: Final = peer.upgrades.get(timeout=30)
|
||||
assert f"/{voice_id}/" in upgrade.target, upgrade.target
|
||||
return upgrade
|
||||
|
||||
|
||||
def _spend_rows(key: str) -> list[dict[str, JsonValue]]:
|
||||
return eventually(
|
||||
lambda: read_rows(SPEND_SQL, (sha256(key.encode()).hexdigest(),)),
|
||||
lambda rows: len(rows) == 1,
|
||||
seconds=70,
|
||||
)
|
||||
|
||||
|
||||
def test_a_session_reaches_elevenlabs_with_the_proxy_key_and_is_billed_per_character(
|
||||
owned: OwnedProxy, elevenlabs: Peer
|
||||
) -> None:
|
||||
voice_id: Final = f"voice{uuid.uuid4().hex[:12]}"
|
||||
with owned.gateway.scenario() as scenario:
|
||||
key: Final = scenario.key()
|
||||
query: Final = f"model_id={FLASH}&output_format=pcm_16000"
|
||||
session: Final = _session(_ws_url(owned, voice_id, query), key, SENTENCES)
|
||||
|
||||
assert [frame.get("audio") for frame in session.frames] == [AUDIO, AUDIO, None]
|
||||
assert session.frames[-1] == {"isFinal": True}
|
||||
assert session.close_code == 1000
|
||||
upgrade: Final = _upgrade_for(elevenlabs, voice_id)
|
||||
assert upgrade.target == f"/v1/text-to-speech/{voice_id}/stream-input?{query}"
|
||||
assert upgrade.provider_key == PROVIDER_KEY
|
||||
assert upgrade.frames == (
|
||||
{"text": " ", "voice_settings": VOICE_SETTINGS},
|
||||
*({"text": text} for text in SENTENCES),
|
||||
{"text": ""},
|
||||
)
|
||||
rows: Final = _spend_rows(key)
|
||||
characters: Final = len(" ") + sum(len(text) for text in SENTENCES)
|
||||
assert float(rows[0]["spend"]) == pytest.approx(characters * _per_character(FLASH))
|
||||
assert rows[0]["model"] == FLASH
|
||||
assert rows[0]["custom_llm_provider"] == "elevenlabs"
|
||||
|
||||
|
||||
def test_an_unpriced_model_is_refused_before_elevenlabs_is_contacted_and_the_default_model_is_billed(
|
||||
owned: OwnedProxy, elevenlabs: Peer
|
||||
) -> None:
|
||||
refused_voice: Final = f"voice{uuid.uuid4().hex[:12]}"
|
||||
default_voice: Final = f"voice{uuid.uuid4().hex[:12]}"
|
||||
with owned.gateway.scenario() as scenario:
|
||||
key: Final = scenario.key()
|
||||
refused: Final = _session(_ws_url(owned, refused_voice, "model_id=eleven_not_in_the_cost_map"), key, SENTENCES)
|
||||
default: Final = _session(_ws_url(owned, default_voice, "output_format=pcm_16000"), key, SENTENCES[:1])
|
||||
|
||||
assert refused.frames == ()
|
||||
assert refused.close_code == 1008
|
||||
assert "elevenlabs/eleven_not_in_the_cost_map" in refused.close_reason
|
||||
assert default.frames[-1] == {"isFinal": True}
|
||||
upgrade: Final = _upgrade_for(elevenlabs, default_voice)
|
||||
assert upgrade.target == (
|
||||
f"/v1/text-to-speech/{default_voice}/stream-input?output_format=pcm_16000&model_id={DEFAULT_MODEL}"
|
||||
)
|
||||
assert elevenlabs.upgrades.empty()
|
||||
rows: Final = _spend_rows(key)
|
||||
assert float(rows[0]["spend"]) == pytest.approx((1 + len(SENTENCES[0])) * _per_character(DEFAULT_MODEL))
|
||||
assert rows[0]["model"] == DEFAULT_MODEL
|
||||
261
tests/unit/llms/elevenlabs/test_elevenlabs_common_utils.py
Normal file
261
tests/unit/llms/elevenlabs/test_elevenlabs_common_utils.py
Normal file
|
|
@ -0,0 +1,261 @@
|
|||
import json
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.llms.elevenlabs.common_utils import (
|
||||
elevenlabs_tts_characters,
|
||||
elevenlabs_tts_client_text,
|
||||
elevenlabs_tts_is_priced,
|
||||
elevenlabs_tts_model,
|
||||
elevenlabs_tts_recorded_upstream_frame,
|
||||
elevenlabs_tts_registry_key,
|
||||
elevenlabs_tts_requested_model,
|
||||
elevenlabs_tts_websocket_target,
|
||||
is_elevenlabs_voice_id,
|
||||
)
|
||||
from litellm.types.passthrough_endpoints.pass_through_endpoints import ForwardedClientText, RefusedClientText
|
||||
|
||||
pytestmark: Final = pytest.mark.usefixtures("local_model_cost_map")
|
||||
|
||||
DEFAULT_MODEL: Final = "eleven_multilingual_v2"
|
||||
STREAM_INPUT: Final = "wss://api.elevenlabs.io/v1/text-to-speech/voice1/stream-input"
|
||||
FLASH_URL: Final = f"{STREAM_INPUT}?model_id=eleven_flash_v2_5"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("api_base", "stream", "query_string", "expected"),
|
||||
[
|
||||
pytest.param(
|
||||
None,
|
||||
"stream-input",
|
||||
"model_id=eleven_flash_v2_5&output_format=pcm_16000",
|
||||
f"{STREAM_INPUT}?model_id=eleven_flash_v2_5&output_format=pcm_16000",
|
||||
id="default base",
|
||||
),
|
||||
pytest.param(
|
||||
None,
|
||||
"multi-stream-input",
|
||||
"model_id=eleven_flash_v2_5",
|
||||
"wss://api.elevenlabs.io/v1/text-to-speech/voice1/multi-stream-input?model_id=eleven_flash_v2_5",
|
||||
id="multi stream path",
|
||||
),
|
||||
pytest.param(
|
||||
None,
|
||||
"stream-input",
|
||||
"output_format=pcm_16000",
|
||||
f"{STREAM_INPUT}?output_format=pcm_16000&model_id={DEFAULT_MODEL}",
|
||||
id="model added when missing",
|
||||
),
|
||||
pytest.param(
|
||||
None,
|
||||
"stream-input",
|
||||
"model_id=&sync_alignment=true",
|
||||
f"{STREAM_INPUT}?sync_alignment=true&model_id={DEFAULT_MODEL}",
|
||||
id="empty model replaced",
|
||||
),
|
||||
pytest.param(
|
||||
None,
|
||||
"stream-input",
|
||||
"model_id=eleven_flash_v2_5&seed=1&model_id=eleven_multilingual_v2",
|
||||
f"{STREAM_INPUT}?model_id=eleven_flash_v2_5&seed=1",
|
||||
id="second model dropped",
|
||||
),
|
||||
pytest.param(
|
||||
None,
|
||||
"stream-input",
|
||||
"authorization=Bearer%20tok&model_id=eleven_flash_v2_5&single_use_token=sut&inactivity_timeout=60",
|
||||
f"{STREAM_INPUT}?model_id=eleven_flash_v2_5&inactivity_timeout=60",
|
||||
id="caller credentials dropped",
|
||||
),
|
||||
pytest.param(
|
||||
"http://localhost:9000/",
|
||||
"stream-input",
|
||||
"model_id=eleven_turbo_v2_5&pronunciation_dictionary_locators=a&pronunciation_dictionary_locators=b",
|
||||
"ws://localhost:9000/v1/text-to-speech/voice1/stream-input"
|
||||
"?model_id=eleven_turbo_v2_5&pronunciation_dictionary_locators=a&pronunciation_dictionary_locators=b",
|
||||
id="custom http base becomes ws and repeated keys stay",
|
||||
),
|
||||
pytest.param(
|
||||
"https://api.eu.residency.elevenlabs.io",
|
||||
"stream-input",
|
||||
"model_id=eleven_flash_v2_5",
|
||||
"wss://api.eu.residency.elevenlabs.io/v1/text-to-speech/voice1/stream-input?model_id=eleven_flash_v2_5",
|
||||
id="custom https base becomes wss",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_websocket_target(api_base, stream, query_string, expected):
|
||||
assert elevenlabs_tts_websocket_target(api_base, "voice1", stream, query_string) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("query_string", "expected"),
|
||||
[
|
||||
("model_id=eleven_flash_v2_5", "eleven_flash_v2_5"),
|
||||
("model_id=eleven_flash_v2_5&model_id=eleven_v3", "eleven_flash_v2_5"),
|
||||
("", DEFAULT_MODEL),
|
||||
("model_id=", DEFAULT_MODEL),
|
||||
("output_format=pcm_16000", DEFAULT_MODEL),
|
||||
],
|
||||
)
|
||||
def test_requested_model_is_the_model_the_target_carries(query_string, expected):
|
||||
target = elevenlabs_tts_websocket_target(None, "voice1", "stream-input", query_string)
|
||||
assert elevenlabs_tts_requested_model(query_string) == expected
|
||||
assert elevenlabs_tts_model(target) == expected
|
||||
assert elevenlabs_tts_registry_key(target) == f"elevenlabs/{expected}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("voice_id", ["21m00Tcm4TlvDq8ikWAM", "my_voice-2", "a"])
|
||||
def test_plain_voice_ids_are_accepted(voice_id):
|
||||
assert is_elevenlabs_voice_id(voice_id)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"voice_id",
|
||||
["", "..", "a/b", "a?model_id=eleven_v3", "a#b", "a%2Fb", "a b", "voixé", "abc\n", "a\\b"],
|
||||
)
|
||||
def test_voice_ids_that_could_change_the_upstream_url_are_refused(voice_id):
|
||||
assert not is_elevenlabs_voice_id(voice_id)
|
||||
|
||||
|
||||
def test_shipped_flash_row_is_priced():
|
||||
assert elevenlabs_tts_is_priced(FLASH_URL)
|
||||
|
||||
|
||||
def test_operator_added_per_character_row_is_priced(monkeypatch):
|
||||
monkeypatch.setitem(
|
||||
litellm.model_cost,
|
||||
"elevenlabs/eleven_operator_custom",
|
||||
{"input_cost_per_character": 1e-05, "litellm_provider": "elevenlabs", "mode": "audio_speech"},
|
||||
)
|
||||
assert elevenlabs_tts_is_priced(FLASH_URL.replace("eleven_flash_v2_5", "eleven_operator_custom"))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model_id", "row"),
|
||||
[
|
||||
pytest.param("eleven_unknown_model", None, id="no row"),
|
||||
pytest.param(
|
||||
"eleven_seconds_only",
|
||||
{"input_cost_per_second": 1e-04, "litellm_provider": "elevenlabs", "mode": "audio_transcription"},
|
||||
id="per second row",
|
||||
),
|
||||
pytest.param(
|
||||
"eleven_zero_rate",
|
||||
{"input_cost_per_character": 0, "litellm_provider": "elevenlabs", "mode": "audio_speech"},
|
||||
id="zero rate",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_models_without_a_per_character_rate_are_unpriced(monkeypatch, model_id, row):
|
||||
if row is not None:
|
||||
monkeypatch.setitem(litellm.model_cost, f"elevenlabs/{model_id}", row)
|
||||
assert not elevenlabs_tts_is_priced(FLASH_URL.replace("eleven_flash_v2_5", model_id))
|
||||
|
||||
|
||||
def test_a_row_under_another_key_does_not_price_the_elevenlabs_model(monkeypatch):
|
||||
monkeypatch.setitem(
|
||||
litellm.model_cost,
|
||||
"eleven_bare_key_only",
|
||||
{"input_cost_per_character": 1e-05, "litellm_provider": "elevenlabs", "mode": "audio_speech"},
|
||||
)
|
||||
assert not elevenlabs_tts_is_priced(FLASH_URL.replace("eleven_flash_v2_5", "eleven_bare_key_only"))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "forwarded", "recorded"),
|
||||
[
|
||||
pytest.param(
|
||||
'{"text":"Hello ", "flush":true}', {"text": "Hello ", "flush": True}, {"text": "Hello "}, id="text frame"
|
||||
),
|
||||
pytest.param(
|
||||
'{"text": "The authorization form is ready "}',
|
||||
{"text": "The authorization form is ready "},
|
||||
{"text": "The authorization form is ready "},
|
||||
id="credential word inside the text",
|
||||
),
|
||||
pytest.param('{"text": ""}', {"text": ""}, {"text": ""}, id="end of stream"),
|
||||
pytest.param(
|
||||
'{"context_id": "ctx", "close_context": true}',
|
||||
{"context_id": "ctx", "close_context": True},
|
||||
None,
|
||||
id="close context has no text",
|
||||
),
|
||||
pytest.param('{"text": "a", "text": "bbbb"}', {"text": "bbbb"}, {"text": "bbbb"}, id="last duplicate wins"),
|
||||
],
|
||||
)
|
||||
def test_client_text_forwards_the_parsed_frame_and_records_the_text_it_carries(text, forwarded, recorded):
|
||||
decision = elevenlabs_tts_client_text(text)
|
||||
|
||||
assert isinstance(decision, ForwardedClientText)
|
||||
assert json.loads(decision.text) == forwarded
|
||||
assert decision.recorded == recorded
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"text",
|
||||
[
|
||||
pytest.param('{"text": " ", "xi_api_key": "sk-litellm-virtual", "voice_settings": {"a": 1}}', id="xi_api_key"),
|
||||
pytest.param('{"text": " ", "xi-api-key": "sk-litellm-virtual", "voice_settings": {"a": 1}}', id="xi-api-key"),
|
||||
pytest.param('{"text": " ", "authorization": "sk-litellm-virtual", "voice_settings": {"a": 1}}', id="auth"),
|
||||
pytest.param('{"text": " ", "xi\\u005fapi_key": "sk-litellm-virtual", "voice_settings": {"a": 1}}', id="esc"),
|
||||
],
|
||||
)
|
||||
def test_caller_credential_fields_never_reach_elevenlabs(text):
|
||||
decision = elevenlabs_tts_client_text(text)
|
||||
|
||||
assert isinstance(decision, ForwardedClientText)
|
||||
assert "sk-litellm-virtual" not in decision.text
|
||||
assert json.loads(decision.text) == {"text": " ", "voice_settings": {"a": 1}}
|
||||
assert decision.recorded == {"text": " "}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"text",
|
||||
[
|
||||
pytest.param('{"text": "free audio", "x": "\\ud800"}', id="lone surrogate"),
|
||||
pytest.param('{"text": "free audio", "x": ' + "[" * 300 + "]" * 300 + "}", id="nested past the parser limit"),
|
||||
pytest.param('{"text": 5}', id="non string text"),
|
||||
pytest.param("not json", id="not json"),
|
||||
pytest.param('"just a string"', id="not an object"),
|
||||
pytest.param('["xi_api_key"]', id="array"),
|
||||
],
|
||||
)
|
||||
def test_client_text_the_billing_parser_cannot_read_is_refused_instead_of_forwarded(text):
|
||||
decision = elevenlabs_tts_client_text(text)
|
||||
|
||||
assert isinstance(decision, RefusedClientText)
|
||||
assert decision.code == 1007
|
||||
|
||||
|
||||
def test_recorded_upstream_frame_drops_audio_and_alignment():
|
||||
frame = {
|
||||
"audio": "SUQzBAAAAAAAI1RTU0UAAAAPAAADTGF2ZjYxLjcuMTAw",
|
||||
"alignment": {"chars": ["H"], "charStartTimesMs": [0], "charDurationsMs": [50]},
|
||||
"normalizedAlignment": {"chars": ["H"], "charStartTimesMs": [0], "charDurationsMs": [50]},
|
||||
"contextId": "ctx",
|
||||
"isFinal": None,
|
||||
}
|
||||
|
||||
assert elevenlabs_tts_recorded_upstream_frame(frame) == {"contextId": "ctx", "isFinal": None}
|
||||
assert elevenlabs_tts_recorded_upstream_frame({"isFinal": True}) == {"isFinal": True}
|
||||
assert elevenlabs_tts_recorded_upstream_frame({"message": "Invalid model", "error": "invalid_request"}) == {
|
||||
"message": "Invalid model",
|
||||
"error": "invalid_request",
|
||||
}
|
||||
|
||||
|
||||
def test_characters_count_every_text_character_including_whitespace():
|
||||
frames = (
|
||||
{"text": " "},
|
||||
{"text": "Hello world. "},
|
||||
{"text": "¿Qué tal? "},
|
||||
{"close_context": True},
|
||||
{"text": None},
|
||||
{"text": ""},
|
||||
)
|
||||
|
||||
assert elevenlabs_tts_characters(frames) == 1 + len("Hello world. ") + len("¿Qué tal? ")
|
||||
assert elevenlabs_tts_characters(()) == 0
|
||||
|
|
@ -0,0 +1,223 @@
|
|||
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.elevenlabs_tts_passthrough_logging_handler import (
|
||||
ElevenLabsTextToSpeechPassthroughLoggingHandler,
|
||||
)
|
||||
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
|
||||
|
||||
pytestmark: Final = pytest.mark.usefixtures("local_model_cost_map")
|
||||
|
||||
FLASH_URL: Final = "wss://api.elevenlabs.io/v1/text-to-speech/voice1/stream-input?model_id=eleven_flash_v2_5"
|
||||
SESSION_TEXT: Final = (" ", "Hello world. ", "¿Qué tal? ", "")
|
||||
|
||||
|
||||
def _client_frames(*texts: str) -> dict[str, object]:
|
||||
return {"client_frames": [{"text": text} for text in texts]}
|
||||
|
||||
|
||||
def _registry_cost(model_id: str, characters: int) -> float:
|
||||
per_character: Final = litellm.model_cost[f"elevenlabs/{model_id}"]["input_cost_per_character"]
|
||||
assert per_character > 0
|
||||
return per_character * characters
|
||||
|
||||
|
||||
def _logging_obj(call_id: str = "call-el") -> 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",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("url_route", "expected"),
|
||||
[
|
||||
("/elevenlabs/v1/text-to-speech/voice1/stream-input", True),
|
||||
("/elevenlabs/v1/text-to-speech/voice1/multi-stream-input", True),
|
||||
("/elevenlabs/v1/text-to-speech/voice1/stream-input?model_id=eleven_flash_v2_5", True),
|
||||
("/litellm/elevenlabs/v1/text-to-speech/voice1/stream-input", True),
|
||||
("/elevenlabs/v1/text-to-speech/voice1/stream", False),
|
||||
("/elevenlabs/v1/text-to-speech/voice1", False),
|
||||
("/elevenlabs/v1/speech-to-text", False),
|
||||
("/deepgram/v1/listen", False),
|
||||
("/openai/v1/realtime/stream-input", False),
|
||||
("", False),
|
||||
],
|
||||
)
|
||||
def test_is_elevenlabs_tts_websocket_route(url_route, expected):
|
||||
assert ElevenLabsTextToSpeechPassthroughLoggingHandler.is_elevenlabs_tts_websocket_route(url_route) is expected
|
||||
|
||||
|
||||
def test_handler_bills_every_sent_character_at_the_registry_rate_and_names_the_model():
|
||||
logging_obj = _logging_obj()
|
||||
characters = sum(len(text) for text in SESSION_TEXT)
|
||||
|
||||
handler_result = ElevenLabsTextToSpeechPassthroughLoggingHandler().elevenlabs_tts_passthrough_handler(
|
||||
request_body=_client_frames(*SESSION_TEXT),
|
||||
upstream_frame_count=4,
|
||||
logging_obj=logging_obj,
|
||||
upstream_url=FLASH_URL,
|
||||
kwargs={"litellm_params": {"metadata": {}}},
|
||||
)
|
||||
|
||||
expected_cost = _registry_cost("eleven_flash_v2_5", characters)
|
||||
assert handler_result["kwargs"]["response_cost"] == pytest.approx(expected_cost)
|
||||
assert handler_result["kwargs"]["model"] == "eleven_flash_v2_5"
|
||||
assert handler_result["kwargs"]["custom_llm_provider"] == "elevenlabs"
|
||||
assert handler_result["kwargs"]["litellm_params"] == {"metadata": {}}
|
||||
assert logging_obj.model == "eleven_flash_v2_5"
|
||||
assert logging_obj.model_call_details["model"] == "eleven_flash_v2_5"
|
||||
assert logging_obj.model_call_details["custom_llm_provider"] == "elevenlabs"
|
||||
assert logging_obj.model_call_details["response_cost"] == pytest.approx(expected_cost)
|
||||
assert handler_result["result"] == {
|
||||
"response": {
|
||||
"model": "eleven_flash_v2_5",
|
||||
"characters": characters,
|
||||
"client_text_frames": len(SESSION_TEXT),
|
||||
"upstream_frames": 4,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def test_handler_bills_the_model_the_session_ran_on():
|
||||
flash = ElevenLabsTextToSpeechPassthroughLoggingHandler().elevenlabs_tts_passthrough_handler(
|
||||
request_body=_client_frames("x" * 1000),
|
||||
upstream_frame_count=1,
|
||||
logging_obj=_logging_obj(),
|
||||
upstream_url=FLASH_URL,
|
||||
)
|
||||
multilingual = ElevenLabsTextToSpeechPassthroughLoggingHandler().elevenlabs_tts_passthrough_handler(
|
||||
request_body=_client_frames("x" * 1000),
|
||||
upstream_frame_count=1,
|
||||
logging_obj=_logging_obj(),
|
||||
upstream_url=FLASH_URL.replace("eleven_flash_v2_5", "eleven_multilingual_v2"),
|
||||
)
|
||||
|
||||
assert flash["kwargs"]["response_cost"] == pytest.approx(_registry_cost("eleven_flash_v2_5", 1000))
|
||||
assert multilingual["kwargs"]["response_cost"] == pytest.approx(_registry_cost("eleven_multilingual_v2", 1000))
|
||||
assert multilingual["kwargs"]["model"] == "eleven_multilingual_v2"
|
||||
|
||||
|
||||
def test_handler_charges_in_proportion_to_the_text_sent():
|
||||
short = ElevenLabsTextToSpeechPassthroughLoggingHandler().elevenlabs_tts_passthrough_handler(
|
||||
request_body=_client_frames("abcde"), upstream_frame_count=1, logging_obj=_logging_obj(), upstream_url=FLASH_URL
|
||||
)
|
||||
long = ElevenLabsTextToSpeechPassthroughLoggingHandler().elevenlabs_tts_passthrough_handler(
|
||||
request_body=_client_frames("abcde", "fghij", "klmno"),
|
||||
upstream_frame_count=3,
|
||||
logging_obj=_logging_obj(),
|
||||
upstream_url=FLASH_URL,
|
||||
)
|
||||
|
||||
assert short["kwargs"]["response_cost"] > 0
|
||||
assert long["kwargs"]["response_cost"] == pytest.approx(3 * short["kwargs"]["response_cost"])
|
||||
|
||||
|
||||
def test_handler_records_no_cost_rather_than_a_stand_in_rate_for_an_unpriced_model():
|
||||
unpriced_url = FLASH_URL.replace("eleven_flash_v2_5", "eleven_not_in_registry")
|
||||
|
||||
handler_result = ElevenLabsTextToSpeechPassthroughLoggingHandler().elevenlabs_tts_passthrough_handler(
|
||||
request_body=_client_frames("Hello "),
|
||||
upstream_frame_count=1,
|
||||
logging_obj=_logging_obj(),
|
||||
upstream_url=unpriced_url,
|
||||
)
|
||||
|
||||
assert handler_result["kwargs"]["response_cost"] is None
|
||||
assert handler_result["kwargs"]["model"] == "eleven_not_in_registry"
|
||||
assert handler_result["result"]["response"]["characters"] == len("Hello ")
|
||||
|
||||
|
||||
def test_handler_with_no_client_text_bills_nothing():
|
||||
handler_result = ElevenLabsTextToSpeechPassthroughLoggingHandler().elevenlabs_tts_passthrough_handler(
|
||||
request_body={}, upstream_frame_count=0, logging_obj=_logging_obj(), upstream_url=FLASH_URL
|
||||
)
|
||||
|
||||
assert handler_result["kwargs"]["response_cost"] == 0
|
||||
assert handler_result["result"]["response"]["characters"] == 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_elevenlabs_tts_and_logs_character_spend_without_audio(monkeypatch):
|
||||
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-el-e2e")
|
||||
request_body = _client_frames(*SESSION_TEXT)
|
||||
upstream_frames = [{"isFinal": None}, {"isFinal": None}, {"isFinal": True}]
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key="hashed-key", team_id="team-tts", user_id="user-1")
|
||||
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,
|
||||
}
|
||||
logging_obj.update_environment_variables(
|
||||
model="unknown",
|
||||
user="unknown",
|
||||
optional_params={},
|
||||
litellm_params={"metadata": metadata},
|
||||
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=FLASH_URL),
|
||||
),
|
||||
response_body=upstream_frames,
|
||||
logging_obj=logging_obj,
|
||||
url_route="/elevenlabs/v1/text-to-speech/voice1/stream-input",
|
||||
result="websocket_connection_successful",
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
cache_hit=False,
|
||||
request_body=request_body,
|
||||
passthrough_logging_payload=PassthroughStandardLoggingPayload(
|
||||
url=FLASH_URL, request_body=request_body, request_method="WEBSOCKET", cost_per_request=None
|
||||
),
|
||||
litellm_params={"metadata": metadata},
|
||||
)
|
||||
|
||||
assert len(capturing_logger.payloads) == 1
|
||||
payload = capturing_logger.payloads[0]
|
||||
assert payload["model"] == "eleven_flash_v2_5"
|
||||
assert payload["custom_llm_provider"] == "elevenlabs"
|
||||
assert payload["response_cost"] == pytest.approx(
|
||||
_registry_cost("eleven_flash_v2_5", sum(len(text) for text in SESSION_TEXT))
|
||||
)
|
||||
assert payload["metadata"]["user_api_key_team_id"] == "team-tts"
|
||||
assert payload["id"] == "call-el-e2e"
|
||||
assert payload["response"] == {
|
||||
"response": {
|
||||
"model": "eleven_flash_v2_5",
|
||||
"characters": sum(len(text) for text in SESSION_TEXT),
|
||||
"client_text_frames": len(SESSION_TEXT),
|
||||
"upstream_frames": len(upstream_frames),
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,436 @@
|
|||
import asyncio
|
||||
import json
|
||||
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,
|
||||
elevenlabs_tts_multi_stream_input_websocket_route,
|
||||
elevenlabs_tts_stream_input_websocket_route,
|
||||
router,
|
||||
)
|
||||
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import WebsocketFramePolicy
|
||||
from litellm.proxy.utils import hash_token
|
||||
from litellm.types.passthrough_endpoints.pass_through_endpoints import ForwardedClientText, RefusedClientText
|
||||
|
||||
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"
|
||||
STREAM_PATH: Final = "/elevenlabs/v1/text-to-speech/voice1/stream-input"
|
||||
MULTI_STREAM_PATH: Final = "/elevenlabs/v1/text-to-speech/voice1/multi-stream-input"
|
||||
UPSTREAM: Final = "wss://api.elevenlabs.io/v1/text-to-speech/voice1"
|
||||
PROVIDER_HEADERS: Final = MappingProxyType({"xi-api-key": "el-provider-key"})
|
||||
|
||||
pytestmark: Final = pytest.mark.usefixtures("local_model_cost_map")
|
||||
|
||||
|
||||
class _FakeWebSocket:
|
||||
def __init__(self, path: str, query: str) -> None:
|
||||
self.url = SimpleNamespace(path=path, query=query)
|
||||
self.headers = {"xi-api-key": "sk-litellm-virtual"}
|
||||
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
|
||||
frame_policy: WebsocketFramePolicy
|
||||
|
||||
|
||||
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,
|
||||
frame_policy: WebsocketFramePolicy,
|
||||
) -> 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,
|
||||
frame_policy=frame_policy,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
async def _serve(
|
||||
websocket: _FakeWebSocket, voice_id: str = "voice1", user_api_key_dict: UserAPIKeyAuth | None = None
|
||||
) -> _FakeRelay:
|
||||
relay = _FakeRelay()
|
||||
await elevenlabs_tts_stream_input_websocket_route(
|
||||
websocket=websocket,
|
||||
voice_id=voice_id,
|
||||
user_api_key_dict=user_api_key_dict or UserAPIKeyAuth(),
|
||||
relay=relay,
|
||||
)
|
||||
return relay
|
||||
|
||||
|
||||
def test_elevenlabs_tts_websocket_routes_registered():
|
||||
ws_paths = {route.path for route in router.routes if isinstance(route, WebSocketRoute)}
|
||||
assert {
|
||||
"/elevenlabs/v1/text-to-speech/{voice_id}/stream-input",
|
||||
"/elevenlabs/v1/text-to-speech/{voice_id}/multi-stream-input",
|
||||
} <= ws_paths
|
||||
|
||||
|
||||
@pytest.mark.parametrize("path", [STREAM_PATH, MULTI_STREAM_PATH])
|
||||
def test_elevenlabs_tts_is_a_lazily_loaded_mapped_pass_through_route(path):
|
||||
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
|
||||
async def test_stream_input_relays_to_elevenlabs_with_only_the_proxy_key(monkeypatch):
|
||||
monkeypatch.delenv("ELEVENLABS_API_BASE", raising=False)
|
||||
websocket = _FakeWebSocket(STREAM_PATH, "model_id=eleven_flash_v2_5&output_format=pcm_16000")
|
||||
caller = UserAPIKeyAuth(api_key="sk-litellm-virtual", team_id="team-tts")
|
||||
|
||||
with patch(GET_CREDENTIALS, return_value="el-provider-key") as get_credentials:
|
||||
relay = await _serve(websocket, user_api_key_dict=caller)
|
||||
|
||||
assert get_credentials.call_args.kwargs == {"custom_llm_provider": "elevenlabs", "region_name": None}
|
||||
assert len(relay.calls) == 1
|
||||
call = relay.calls[0]
|
||||
assert call.target == f"{UPSTREAM}/stream-input?model_id=eleven_flash_v2_5&output_format=pcm_16000"
|
||||
assert call.custom_headers == PROVIDER_HEADERS
|
||||
assert call.user_api_key_dict is caller
|
||||
assert call.forward_headers is False
|
||||
assert call.endpoint == STREAM_PATH
|
||||
assert call.accept_websocket is False
|
||||
decision = call.frame_policy.client_text('{"text": "Hi ", "xi_api_key": "sk-litellm-virtual", "flush": true}')
|
||||
assert isinstance(decision, ForwardedClientText)
|
||||
assert json.loads(decision.text) == {"text": "Hi ", "flush": True}
|
||||
assert decision.recorded == {"text": "Hi "}
|
||||
assert isinstance(call.frame_policy.client_text('{"text": "Hi ", "x": "\\ud800"}'), RefusedClientText)
|
||||
assert call.frame_policy.recorded_upstream_frame({"audio": "AAAA", "isFinal": None}) == {"isFinal": None}
|
||||
assert websocket.accepts == [None]
|
||||
assert websocket.closed is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multi_stream_input_relays_to_the_multi_stream_endpoint(monkeypatch):
|
||||
monkeypatch.delenv("ELEVENLABS_API_BASE", raising=False)
|
||||
websocket = _FakeWebSocket(MULTI_STREAM_PATH, "model_id=eleven_flash_v2_5")
|
||||
relay = _FakeRelay()
|
||||
|
||||
with patch(GET_CREDENTIALS, return_value="el-provider-key"):
|
||||
await elevenlabs_tts_multi_stream_input_websocket_route(
|
||||
websocket=websocket, voice_id="voice1", user_api_key_dict=UserAPIKeyAuth(), relay=relay
|
||||
)
|
||||
|
||||
assert [(call.target, call.endpoint) for call in relay.calls] == [
|
||||
(f"{UPSTREAM}/multi-stream-input?model_id=eleven_flash_v2_5", MULTI_STREAM_PATH)
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("query", "expected_query"),
|
||||
[
|
||||
pytest.param("", "model_id=eleven_multilingual_v2", id="model omitted"),
|
||||
pytest.param(
|
||||
"model_id=&output_format=mp3_44100_128",
|
||||
"output_format=mp3_44100_128&model_id=eleven_multilingual_v2",
|
||||
id="model blank",
|
||||
),
|
||||
pytest.param(
|
||||
"model_id=eleven_flash_v2_5&model_id=eleven_multilingual_v2",
|
||||
"model_id=eleven_flash_v2_5",
|
||||
id="second model stripped",
|
||||
),
|
||||
pytest.param(
|
||||
"authorization=Bearer%20el-caller-key&single_use_token=sut&model_id=eleven_flash_v2_5",
|
||||
"model_id=eleven_flash_v2_5",
|
||||
id="caller credentials stripped",
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_upstream_query_carries_exactly_one_authorized_model_and_no_caller_credentials(
|
||||
query, expected_query, monkeypatch
|
||||
):
|
||||
monkeypatch.delenv("ELEVENLABS_API_BASE", raising=False)
|
||||
websocket = _FakeWebSocket(STREAM_PATH, query)
|
||||
|
||||
with patch(GET_CREDENTIALS, return_value="el-provider-key"):
|
||||
relay = await _serve(websocket)
|
||||
|
||||
assert [call.target for call in relay.calls] == [f"{UPSTREAM}/stream-input?{expected_query}"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("api_base", "expected_target"),
|
||||
[
|
||||
(
|
||||
"https://api.eu.residency.elevenlabs.io/",
|
||||
"wss://api.eu.residency.elevenlabs.io/v1/text-to-speech/voice1/stream-input?model_id=eleven_flash_v2_5",
|
||||
),
|
||||
(
|
||||
"http://localhost:8080",
|
||||
"ws://localhost:8080/v1/text-to-speech/voice1/stream-input?model_id=eleven_flash_v2_5",
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_server_configured_api_base_is_honoured(api_base, expected_target, monkeypatch):
|
||||
monkeypatch.setenv("ELEVENLABS_API_BASE", api_base)
|
||||
websocket = _FakeWebSocket(STREAM_PATH, "model_id=eleven_flash_v2_5")
|
||||
|
||||
with patch(GET_CREDENTIALS, return_value="el-provider-key"):
|
||||
relay = await _serve(websocket)
|
||||
|
||||
assert [call.target for call in relay.calls] == [expected_target]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_provider_key_closes_before_the_handshake_completes():
|
||||
websocket = _FakeWebSocket(STREAM_PATH, "model_id=eleven_flash_v2_5")
|
||||
|
||||
with patch(GET_CREDENTIALS, return_value=None):
|
||||
relay = await _serve(websocket)
|
||||
|
||||
assert websocket.closed is not None
|
||||
assert websocket.closed[0] == 1011
|
||||
assert "ELEVENLABS_API_KEY" in websocket.closed[1]
|
||||
assert websocket.accepts == []
|
||||
assert relay.calls == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("voice_id", ["..", "a?model_id=eleven_v3", "a#b", "a b", "voixé"])
|
||||
async def test_voice_ids_that_could_rewrite_the_upstream_url_are_refused(voice_id, monkeypatch):
|
||||
monkeypatch.delenv("ELEVENLABS_API_BASE", raising=False)
|
||||
websocket = _FakeWebSocket(STREAM_PATH, "model_id=eleven_flash_v2_5")
|
||||
|
||||
with patch(GET_CREDENTIALS, return_value="el-provider-key"):
|
||||
relay = await _serve(websocket, voice_id=voice_id)
|
||||
|
||||
assert relay.calls == []
|
||||
assert websocket.closed is not None
|
||||
assert websocket.closed[0] == 1008
|
||||
assert "voice_id" in websocket.closed[1]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("query", "registry_key"),
|
||||
[
|
||||
pytest.param("model_id=eleven_not_in_registry", "elevenlabs/eleven_not_in_registry", id="unknown model"),
|
||||
pytest.param("model_id=scribe_v1", "elevenlabs/scribe_v1", id="row without a per character price"),
|
||||
],
|
||||
)
|
||||
async def test_sessions_the_proxy_cannot_price_are_refused_before_elevenlabs_is_contacted(
|
||||
query, registry_key, monkeypatch
|
||||
):
|
||||
monkeypatch.delenv("ELEVENLABS_API_BASE", raising=False)
|
||||
websocket = _FakeWebSocket(STREAM_PATH, query)
|
||||
|
||||
with patch(GET_CREDENTIALS, return_value="el-provider-key"):
|
||||
relay = await _serve(websocket)
|
||||
|
||||
assert relay.calls == []
|
||||
assert websocket.closed is not None
|
||||
assert websocket.closed[0] == 1008
|
||||
assert registry_key in websocket.closed[1]
|
||||
assert "el-provider-key" not in websocket.closed[1]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_session_relays_once_the_operator_prices_the_model(monkeypatch):
|
||||
monkeypatch.delenv("ELEVENLABS_API_BASE", raising=False)
|
||||
with patch(GET_CREDENTIALS, return_value="el-provider-key"):
|
||||
assert (await _serve(_FakeWebSocket(STREAM_PATH, "model_id=eleven_operator_custom"))).calls == []
|
||||
|
||||
monkeypatch.setitem(
|
||||
litellm.model_cost,
|
||||
"elevenlabs/eleven_operator_custom",
|
||||
dict(litellm.model_cost["elevenlabs/eleven_flash_v2_5"]),
|
||||
)
|
||||
with patch(GET_CREDENTIALS, return_value="el-provider-key"):
|
||||
relay = await _serve(_FakeWebSocket(STREAM_PATH, "model_id=eleven_operator_custom"))
|
||||
|
||||
assert [call.target for call in relay.calls] == [f"{UPSTREAM}/stream-input?model_id=eleven_operator_custom"]
|
||||
|
||||
|
||||
def _app_with_relay(relay: _FakeRelay) -> FastAPI:
|
||||
app = FastAPI()
|
||||
app.include_router(router)
|
||||
app.dependency_overrides[_websocket_relay] = lambda: relay
|
||||
return app
|
||||
|
||||
|
||||
def test_connections_without_a_litellm_key_are_refused():
|
||||
relay = _FakeRelay()
|
||||
client = TestClient(_app_with_relay(relay))
|
||||
|
||||
with patch(GET_CREDENTIALS, return_value="el-provider-key") as get_credentials:
|
||||
with pytest.raises(WebSocketDisconnect) as disconnect:
|
||||
with client.websocket_connect(f"{STREAM_PATH}?model_id=eleven_flash_v2_5"):
|
||||
pass
|
||||
|
||||
assert disconnect.value.code == 1008
|
||||
assert relay.calls == []
|
||||
get_credentials.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"headers",
|
||||
[
|
||||
pytest.param({"xi-api-key": "sk-litellm-virtual"}, id="elevenlabs sdk header"),
|
||||
pytest.param({"api-key": "sk-litellm-virtual"}, id="api-key header"),
|
||||
pytest.param({"Authorization": "Bearer sk-litellm-virtual"}, id="bearer header"),
|
||||
],
|
||||
)
|
||||
def test_the_litellm_key_is_accepted_where_elevenlabs_clients_send_theirs(headers, monkeypatch):
|
||||
monkeypatch.delenv("ELEVENLABS_API_BASE", raising=False)
|
||||
relay = _FakeRelay()
|
||||
client = TestClient(_app_with_relay(relay))
|
||||
caller = UserAPIKeyAuth(api_key="hashed-sk-litellm", team_id="team-tts")
|
||||
|
||||
with (
|
||||
patch(GET_CREDENTIALS, return_value="el-provider-key"),
|
||||
patch(USER_API_KEY_AUTH, new=AsyncMock(return_value=caller)) as auth,
|
||||
):
|
||||
with client.websocket_connect(f"{STREAM_PATH}?model_id=eleven_flash_v2_5", headers=headers):
|
||||
pass
|
||||
|
||||
assert auth.await_args.kwargs["api_key"] == "Bearer sk-litellm-virtual"
|
||||
assert [(call.custom_headers, call.forward_headers, call.user_api_key_dict) for call in relay.calls] == [
|
||||
(PROVIDER_HEADERS, False, caller)
|
||||
]
|
||||
|
||||
|
||||
def test_a_bad_voice_id_reaches_the_client_as_a_policy_close(monkeypatch):
|
||||
monkeypatch.delenv("ELEVENLABS_API_BASE", raising=False)
|
||||
relay = _FakeRelay()
|
||||
client = TestClient(_app_with_relay(relay))
|
||||
|
||||
with (
|
||||
patch(GET_CREDENTIALS, return_value="el-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(
|
||||
"/elevenlabs/v1/text-to-speech/a%3Fb/stream-input?model_id=eleven_flash_v2_5",
|
||||
headers={"xi-api-key": "sk-litellm-virtual"},
|
||||
) as connection:
|
||||
connection.receive_text()
|
||||
|
||||
assert disconnect.value.code == 1008
|
||||
assert "voice_id" in disconnect.value.reason
|
||||
assert relay.calls == []
|
||||
|
||||
|
||||
def test_the_browser_subprotocol_carrying_the_litellm_key_is_echoed(monkeypatch):
|
||||
monkeypatch.delenv("ELEVENLABS_API_BASE", raising=False)
|
||||
relay = _FakeRelay()
|
||||
client = TestClient(_app_with_relay(relay))
|
||||
|
||||
with (
|
||||
patch(GET_CREDENTIALS, return_value="el-provider-key"),
|
||||
patch(USER_API_KEY_AUTH, new=AsyncMock(return_value=UserAPIKeyAuth(api_key="hashed"))),
|
||||
):
|
||||
with client.websocket_connect(
|
||||
f"{STREAM_PATH}?model_id=eleven_flash_v2_5",
|
||||
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] == [PROVIDER_HEADERS]
|
||||
|
||||
|
||||
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", "expected_query"),
|
||||
[
|
||||
pytest.param("model_id=eleven_flash_v2_5", "model_id=eleven_flash_v2_5", id="allowed model named"),
|
||||
pytest.param(
|
||||
"model_id=eleven_flash_v2_5&model_id=eleven_multilingual_v2",
|
||||
"model_id=eleven_flash_v2_5",
|
||||
id="denied model behind an allowed one is stripped",
|
||||
),
|
||||
pytest.param("model_id=eleven_multilingual_v2", None, id="denied model named"),
|
||||
pytest.param("", None, id="model omitted, default denied"),
|
||||
pytest.param("model_id=", None, id="model blank, default denied"),
|
||||
],
|
||||
)
|
||||
def test_key_model_access_is_checked_against_the_model_sent_upstream(query, expected_query, monkeypatch):
|
||||
monkeypatch.delenv("ELEVENLABS_API_BASE", raising=False)
|
||||
monkeypatch.setattr(litellm, "max_budget", 0.0)
|
||||
cache = asyncio.run(_cache_restricted_key("sk-only-flash", ["eleven_flash_v2_5"]))
|
||||
relay = _FakeRelay()
|
||||
client = TestClient(_app_with_relay(relay))
|
||||
|
||||
with (
|
||||
patch(GET_CREDENTIALS, return_value="el-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 expected_query is not None:
|
||||
with client.websocket_connect(f"{STREAM_PATH}?{query}", headers={"xi-api-key": "sk-only-flash"}):
|
||||
pass
|
||||
assert [call.target for call in relay.calls] == [f"{UPSTREAM}/stream-input?{expected_query}"]
|
||||
return
|
||||
with pytest.raises(WebSocketDisconnect) as disconnect:
|
||||
with client.websocket_connect(f"{STREAM_PATH}?{query}", headers={"xi-api-key": "sk-only-flash"}):
|
||||
pass
|
||||
|
||||
assert disconnect.value.code == 1008
|
||||
assert relay.calls == []
|
||||
|
|
@ -36,6 +36,10 @@ from litellm.integrations.custom_logger import CustomLogger
|
|||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.proxy._lazy_features import LazyFeature, attach_lazy_features
|
||||
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
|
||||
from litellm.llms.elevenlabs.common_utils import (
|
||||
elevenlabs_tts_client_text,
|
||||
elevenlabs_tts_recorded_upstream_frame,
|
||||
)
|
||||
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
|
||||
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import router as llm_passthrough_router
|
||||
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
|
||||
|
|
@ -54,6 +58,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
|
|||
pass_through_request,
|
||||
resolve_llm_passthrough_timeout,
|
||||
resolve_pass_through_request_timeout,
|
||||
WebsocketFramePolicy,
|
||||
set_env_variables_in_header,
|
||||
websocket_passthrough_request,
|
||||
)
|
||||
|
|
@ -7601,7 +7606,7 @@ async def test_websocket_passthrough_sends_deepgram_audio_bytes_and_control_text
|
|||
audio_chunk = bytes(range(256)) * 4
|
||||
close_stream = json.dumps({"type": "CloseStream"})
|
||||
|
||||
await _relay_deepgram_listen(
|
||||
_, success_handler = await _relay_deepgram_listen(
|
||||
upstream_ws,
|
||||
AsyncMock(
|
||||
side_effect=[
|
||||
|
|
@ -7615,6 +7620,182 @@ async def test_websocket_passthrough_sends_deepgram_audio_bytes_and_control_text
|
|||
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()
|
||||
assert success_handler.call_args.kwargs["request_body"] == {}
|
||||
assert success_handler.call_args.kwargs["passthrough_logging_payload"]["request_body"] == {}
|
||||
|
||||
|
||||
ELEVENLABS_TTS_TARGET = (
|
||||
"wss://api.elevenlabs.io/v1/text-to-speech/voice1/stream-input?model_id=eleven_flash_v2_5&sync_alignment=true"
|
||||
)
|
||||
ELEVENLABS_TTS_FRAME_POLICY = WebsocketFramePolicy(
|
||||
client_text=elevenlabs_tts_client_text,
|
||||
recorded_upstream_frame=elevenlabs_tts_recorded_upstream_frame,
|
||||
)
|
||||
ELEVENLABS_AUDIO_FRAME = json.dumps(
|
||||
{
|
||||
"audio": base64.b64encode(bytes(range(64))).decode(),
|
||||
"alignment": {"chars": ["H"], "charStartTimesMs": [0], "charDurationsMs": [40]},
|
||||
"normalizedAlignment": {"chars": ["H"], "charStartTimesMs": [0], "charDurationsMs": [40]},
|
||||
"isFinal": None,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class ScriptedUpstreamWebSocket:
|
||||
"""Records what the client sent and, once the expected number of frames arrived, serves its own then closes"""
|
||||
|
||||
def __init__(self, expected_client_frames: int, *frames: str | bytes):
|
||||
self._expected_client_frames = expected_client_frames
|
||||
self._frames = iter(frames)
|
||||
self._client_frames_arrived = asyncio.Event()
|
||||
self.sent: list[str | bytes] = []
|
||||
self.close = AsyncMock()
|
||||
|
||||
async def send(self, message: str | bytes) -> None:
|
||||
self.sent.append(message)
|
||||
if len(self.sent) >= self._expected_client_frames:
|
||||
self._client_frames_arrived.set()
|
||||
|
||||
async def recv(self, decode: bool | None = None):
|
||||
from websockets.exceptions import ConnectionClosedOK
|
||||
from websockets.frames import Close
|
||||
|
||||
await self._client_frames_arrived.wait()
|
||||
frame = next(self._frames, None)
|
||||
if frame is None:
|
||||
raise ConnectionClosedOK(rcvd=Close(1000, ""), sent=Close(1000, ""), rcvd_then_sent=True)
|
||||
return frame
|
||||
|
||||
|
||||
class DrainingUpstreamWebSocket:
|
||||
"""Delivers one buffered frame while its closing handshake runs, the way a real connection drains on close"""
|
||||
|
||||
def __init__(self, buffered_frame: str):
|
||||
self._buffered_frame = buffered_frame
|
||||
self._closing = asyncio.Event()
|
||||
self._drained = asyncio.Event()
|
||||
self.sent: list[str | bytes] = []
|
||||
|
||||
async def send(self, message: str | bytes) -> None:
|
||||
self.sent.append(message)
|
||||
|
||||
async def recv(self, decode: bool | None = None):
|
||||
from websockets.exceptions import ConnectionClosedOK
|
||||
from websockets.frames import Close
|
||||
|
||||
await self._closing.wait()
|
||||
if self._drained.is_set():
|
||||
raise ConnectionClosedOK(rcvd=Close(1000, ""), sent=Close(1000, ""), rcvd_then_sent=True)
|
||||
self._drained.set()
|
||||
return self._buffered_frame
|
||||
|
||||
async def close(self) -> None:
|
||||
self._closing.set()
|
||||
await self._drained.wait()
|
||||
|
||||
|
||||
def _client_sends_then_waits(*texts: str):
|
||||
messages = iter([{"type": "websocket.receive", "text": text} for text in texts])
|
||||
|
||||
async def receive():
|
||||
message = next(messages, None)
|
||||
if message is None:
|
||||
await asyncio.Event().wait()
|
||||
return message
|
||||
|
||||
return receive
|
||||
|
||||
|
||||
async def _relay_elevenlabs_tts(upstream_ws, client_receive, client_send_text=None):
|
||||
websocket = _client_websocket(client_receive)
|
||||
if client_send_text is not None:
|
||||
websocket.send_text = client_send_text
|
||||
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=ELEVENLABS_TTS_TARGET,
|
||||
custom_headers={"xi-api-key": "el-provider-key"},
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
forward_headers=False,
|
||||
endpoint="/elevenlabs/v1/text-to-speech/voice1/stream-input",
|
||||
accept_websocket=False,
|
||||
frame_policy=ELEVENLABS_TTS_FRAME_POLICY,
|
||||
)
|
||||
return websocket, success_handler
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_elevenlabs_policy_keeps_caller_keys_upstream_side_and_audio_out_of_logs():
|
||||
initial_frame = json.dumps(
|
||||
{"text": " ", "xi_api_key": "sk-litellm-virtual", "voice_settings": {"stability": 0.5, "similarity_boost": 0.8}}
|
||||
)
|
||||
text_frame = json.dumps({"text": "Hola, ¿qué tal? ", "try_trigger_generation": True}, ensure_ascii=False)
|
||||
end_frame = json.dumps({"text": ""})
|
||||
final_frame = json.dumps({"isFinal": True})
|
||||
upstream_ws = ScriptedUpstreamWebSocket(3, ELEVENLABS_AUDIO_FRAME, final_frame)
|
||||
|
||||
websocket, success_handler = await _relay_elevenlabs_tts(
|
||||
upstream_ws, _client_sends_then_waits(initial_frame, text_frame, end_frame)
|
||||
)
|
||||
|
||||
assert [json.loads(frame) for frame in upstream_ws.sent] == [
|
||||
{"text": " ", "voice_settings": {"stability": 0.5, "similarity_boost": 0.8}},
|
||||
json.loads(text_frame),
|
||||
{"text": ""},
|
||||
]
|
||||
assert [call.args[0] for call in websocket.send_text.await_args_list] == [ELEVENLABS_AUDIO_FRAME, final_frame]
|
||||
success_call = success_handler.call_args.kwargs
|
||||
expected_request_body = {"client_frames": [{"text": " "}, {"text": "Hola, ¿qué tal? "}, {"text": ""}]}
|
||||
assert success_call["request_body"] == expected_request_body
|
||||
assert success_call["passthrough_logging_payload"]["request_body"] == expected_request_body
|
||||
assert success_call["response_body"] == [{"isFinal": None}, {"isFinal": True}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_elevenlabs_refuses_a_frame_its_counter_cannot_read_and_bills_the_rest():
|
||||
upstream_ws = DrainingUpstreamWebSocket(ELEVENLABS_AUDIO_FRAME)
|
||||
unreadable_frame = '{"text": "a long passage nobody pays for", "x": "\\ud800"}'
|
||||
send_after_close = AsyncMock(side_effect=RuntimeError('Cannot call "send" once a close message has been sent.'))
|
||||
|
||||
websocket, success_handler = await _relay_elevenlabs_tts(
|
||||
upstream_ws,
|
||||
_client_sends_then_waits(json.dumps({"text": " "}), json.dumps({"text": "Hello "}), unreadable_frame),
|
||||
send_after_close,
|
||||
)
|
||||
|
||||
assert [json.loads(frame) for frame in upstream_ws.sent] == [{"text": " "}, {"text": "Hello "}]
|
||||
websocket.close.assert_awaited_once()
|
||||
assert websocket.close.await_args.kwargs["code"] == 1007
|
||||
send_after_close.assert_not_awaited()
|
||||
assert success_handler.call_args.kwargs["request_body"] == {"client_frames": [{"text": " "}, {"text": "Hello "}]}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_bills_the_session_when_the_client_drops_while_audio_streams():
|
||||
from starlette.websockets import WebSocketDisconnect
|
||||
|
||||
upstream_ws = ScriptedUpstreamWebSocket(
|
||||
2, ELEVENLABS_AUDIO_FRAME, ELEVENLABS_AUDIO_FRAME, json.dumps({"isFinal": True})
|
||||
)
|
||||
client_send_text = AsyncMock(side_effect=[None, WebSocketDisconnect(code=1006)])
|
||||
|
||||
_, success_handler = await _relay_elevenlabs_tts(
|
||||
upstream_ws,
|
||||
_client_sends_then_waits(json.dumps({"text": " "}), json.dumps({"text": "Hello there "})),
|
||||
client_send_text,
|
||||
)
|
||||
|
||||
assert client_send_text.await_count == 2
|
||||
success_call = success_handler.call_args.kwargs
|
||||
assert success_call["request_body"] == {"client_frames": [{"text": " "}, {"text": "Hello there "}]}
|
||||
assert success_call["response_body"] == [{"isFinal": None}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue