fix(realtime): enforce model budgets and preserve translation routing

This commit is contained in:
Emerson Gomes 2026-10-19 23:59:53 -05:00
parent 19bcd5c584
commit 8855c1385e
No known key found for this signature in database
GPG key ID: D3DF28AB5D1B5E17
8 changed files with 151 additions and 29 deletions

View file

@ -727,10 +727,25 @@ class RealTimeStreaming:
backend, False if the provider transformation produced no output and
the message was effectively dropped.
"""
message = await self._apply_nested_transcription_model_policy(message)
from litellm.proxy._types import ProxyException
try:
authorized_message: Final = await self._apply_nested_transcription_model_policy(message)
except ProxyException as denied:
denial_message: Final = redact_internal_details_from_client_message(denied.message)
try:
await self.websocket.send_text(realtime_error_event(denial_message, error_type="permission_error"))
finally:
await self.websocket.close(
code=1008,
reason=websocket_close_reason(
denial_message, fallback="Nested transcription model is not authorized"
),
)
return False
if self.provider_config:
transformed: Final = self.provider_config.transform_realtime_request(
message, self.model, self.session_configuration_request
authorized_message, self.model, self.session_configuration_request
)
sent = False
for msg in transformed:
@ -766,14 +781,14 @@ class RealTimeStreaming:
self._content_sent_after_setup = True
sent = True
if sent:
self._capture_translation_input_audio(message)
self._capture_translation_input_audio(authorized_message)
return sent
await self.backend_ws.send(message)
await self.backend_ws.send(authorized_message)
if self._usage_meter is not None:
sent_event: Final = self._parse_backend_event(message)
sent_event: Final = self._parse_backend_event(authorized_message)
if sent_event is not None:
self._usage_meter.observe_client_event(sent_event)
self._capture_translation_input_audio(message)
self._capture_translation_input_audio(authorized_message)
return True
async def _apply_nested_transcription_model_policy(self, message: str) -> str:

View file

@ -60,7 +60,7 @@ class AzureAudioTranscription(AzureChatCompletion):
provider_specific_entry: Final = (
as_str_mapping(model_info.get("provider_specific_entry")) if model_info is not None else None
)
requires_deployment_api: Final = model_info is None or (
requires_deployment_api: Final = (
provider_specific_entry is not None and provider_specific_entry.get("transcription_deployment_api") == 1
)
resolved_api_version: Final = (

View file

@ -144,6 +144,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg
populate_request_with_path_params,
read_raw_json_body,
read_request_body,
resolve_realtime_route_model,
rewrite_request_model,
safe_get_request_headers,
safe_get_request_query_params,
@ -1180,7 +1181,13 @@ def _apply_budget_limits_to_end_user_params(
async def user_api_key_auth_websocket(websocket: WebSocket) -> UserAPIKeyAuth:
return await user_api_key_auth_websocket_for_model(websocket, model=websocket.query_params.get("model"))
route_path: Final[object] = websocket.scope.get("path")
model: Final = resolve_realtime_route_model(
model=websocket.query_params.get("model"),
intent=websocket.query_params.get("intent"),
is_translation=isinstance(route_path, str) and route_path.endswith("/realtime/translations"),
)
return await user_api_key_auth_websocket_for_model(websocket, model=model)
async def user_api_key_auth_live_websocket(websocket: WebSocket) -> UserAPIKeyAuth:

View file

@ -289,6 +289,7 @@ async def acreate_realtime_translation_client_secret(
model_name, custom_llm_provider, dynamic_api_key, dynamic_api_base = get_llm_provider(
model=requested_model_name,
custom_llm_provider=litellm_params.custom_llm_provider,
api_base=litellm_params.api_base,
api_key=litellm_params.api_key,
)
@ -492,6 +493,7 @@ async def arealtime_translation_calls(
model_name, custom_llm_provider, dynamic_api_key, dynamic_api_base = get_llm_provider(
model=requested_model_name,
custom_llm_provider=litellm_params.custom_llm_provider,
api_base=litellm_params.api_base,
api_key=litellm_params.api_key,
)

View file

@ -879,9 +879,7 @@ async def test_usage_meter_snapshot_stored_in_model_call_details():
backend_ws.send = AsyncMock()
backend_ws.recv = AsyncMock(
side_effect=[
json.dumps(
{"type": "session.created", "session": {"turn_detection": {"type": None}}}
).encode(),
json.dumps({"type": "session.created", "session": {"turn_detection": {"type": None}}}).encode(),
json.dumps(
{"type": "response.output_audio.delta", "delta": base64.b64encode(b"x" * 24000).decode()}
).encode(),
@ -1211,10 +1209,11 @@ async def test_transcription_session_update_enforces_authorized_nested_model():
@pytest.mark.asyncio
async def test_translation_session_update_rejects_disallowed_nested_transcription_model() -> None:
backend_ws = MagicMock()
client_ws: Final = MagicMock(send_text=AsyncMock(), close=AsyncMock())
backend_ws: Final = MagicMock()
backend_ws.send = AsyncMock()
streaming = RealTimeStreaming(
MagicMock(),
streaming: Final = RealTimeStreaming(
client_ws,
backend_ws,
MagicMock(),
model="gpt-realtime-translate",
@ -1222,24 +1221,29 @@ async def test_translation_session_update_rejects_disallowed_nested_transcriptio
translation_session=True,
)
with pytest.raises(Exception, match=r"gpt-live-transcribe.*not available"):
await streaming._send_to_backend(
json.dumps(
{
"type": "session.update",
"session": {
"type": "translation",
"audio": {
"input": {
"transcription": {"model": "gpt-live-transcribe"},
}
},
sent: Final = await streaming._send_to_backend(
json.dumps(
{
"type": "session.update",
"session": {
"type": "translation",
"audio": {
"input": {
"transcription": {"model": "gpt-live-transcribe"},
}
},
}
)
},
}
)
)
assert sent is False
backend_ws.send.assert_not_awaited()
assert streaming._is_transcription_session is False
error: Final = json.loads(client_ws.send_text.await_args.args[0])
assert error["type"] == "error"
assert error["error"]["type"] == "permission_error"
assert "gpt-live-transcribe" in error["error"]["message"]
assert client_ws.close.await_args.kwargs["code"] == 1008
@pytest.mark.asyncio

View file

@ -172,7 +172,7 @@ async def test_azure_transcription_parses_response_formats(
[
("whisper-1", "/openai/v1/audio/transcriptions"),
("gpt-transcribe", "/openai/deployments/gpt-transcribe/audio/transcriptions"),
("custom-transcribe-deployment", "/openai/deployments/custom-transcribe-deployment/audio/transcriptions"),
("custom-transcribe-deployment", "/openai/v1/audio/transcriptions"),
],
)
def test_azure_transcription_alias_uses_model_route(

View file

@ -11936,6 +11936,59 @@ async def test_websocket_auth_hands_the_reservation_to_the_socket_state():
assert websocket.scope["state"]["budget_reservation"] is reservation
@pytest.mark.asyncio
@pytest.mark.parametrize(
"path", ("/realtime/translations", "/v1/realtime/translations", "/openai/v1/realtime/translations")
)
@pytest.mark.parametrize("query", (b"", b"model=gpt-realtime-translate"), ids=("default_model", "explicit_model"))
async def test_translation_websocket_auth_enforces_the_default_model_budget(
monkeypatch: pytest.MonkeyPatch, path: str, query: bytes
) -> None:
from fastapi import WebSocket
from litellm.proxy import proxy_server
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth_websocket
model: Final = "gpt-realtime-translate"
token: Final = UserAPIKeyAuth(
token="test-key",
models=[model],
model_max_budget={model: {"max_budget": 5.0, "budget_duration": "1d"}},
)
limiter, _ = await _real_model_budget_limiter(key_spend={model: 10.0})
with pytest.raises(litellm.BudgetExceededError):
await limiter.is_key_within_model_budget(user_api_key_dict=token, model=model)
async def authenticate(request: Request, api_key: str) -> UserAPIKeyAuth:
return token
attrs: Final = {
**_proxy_server_attrs_for_custom_auth(user_custom_auth=authenticate),
"prisma_client": None,
"model_max_budget_limiter": limiter,
}
for name, value in attrs.items():
monkeypatch.setattr(proxy_server, name, value)
monkeypatch.setattr(litellm, "enable_post_custom_auth_checks", True, raising=False)
websocket: Final = WebSocket(
scope={
"type": "websocket",
"path": path,
"headers": [(b"authorization", b"Bearer sk-test")],
"query_string": query,
},
receive=AsyncMock(),
send=AsyncMock(),
)
with pytest.raises(HTTPException) as denied:
await user_api_key_auth_websocket(websocket)
assert denied.value.status_code == 403
assert model in str(denied.value.detail)
assert "budget" in str(denied.value.detail).lower()
@pytest.mark.asyncio
async def test_admission_and_budget_reservation_read_the_key_spend_counter_with_one_redis_mget():
from fastapi import Request

View file

@ -1222,6 +1222,47 @@ async def test_openai_realtime_http_routes_send_the_federated_bearer_for_a_keyle
assert gateway.authorization == "Bearer federated-bearer"
@pytest.mark.asyncio
@pytest.mark.parametrize("flow", ("client_secret", "calls"))
async def test_translation_http_routes_preserve_explicit_azure_provider(flow: str) -> None:
expected_path: Final = "/openai/v1/realtime/translations/" + {
"client_secret": "client_secrets",
"calls": "calls",
}[flow]
def send_response(request: httpx.Request) -> httpx.Response:
assert request.url.host == "azure-deployment.example"
assert request.url.path == expected_path
return httpx.Response(200, json={"value": "ek-test", "routed_path": request.url.path})
params: Final[Mapping[str, object]] = {
"model": "translation-deployment",
"custom_llm_provider": "azure",
"api_base": "https://azure-deployment.example/openai/v1/",
"api_key": "test-key",
"client": AsyncHTTPHandler(transport=httpx.MockTransport(send_response)),
"litellm_logging_obj": Logging(
model="translation-deployment",
messages=[],
stream=False,
call_type="acreate_realtime_translation_client_secret",
start_time=datetime(2026, 1, 1, tzinfo=timezone.utc),
litellm_call_id="test-translation-routing",
function_id="test-translation-routing",
),
}
response: Final = await (
realtime_main.acreate_realtime_translation_client_secret.__wrapped__(**params)
if flow == "client_secret"
else realtime_main.arealtime_translation_calls.__wrapped__(
openai_ephemeral_key="ek-test", sdp_body=b"v=0", **params
)
)
assert response.status_code == 200
assert response.json() == {"value": "ek-test", "routed_path": expected_path}
@pytest.mark.asyncio
@respx.mock
async def test_openai_realtime_health_check_sends_the_federated_bearer_for_a_keyless_deployment(