mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
fix(realtime): enforce model budgets and preserve translation routing
This commit is contained in:
parent
19bcd5c584
commit
8855c1385e
8 changed files with 151 additions and 29 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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 = (
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue