diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 8887d86570a..828b5b514b7 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -5291,79 +5291,6 @@ async def _granted_model_lists( ) -async def _user_object_or_none( - valid_token: UserAPIKeyAuth, - prisma_client: PrismaClient, - user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: ProxyLogging, -) -> LiteLLM_UserTable | None: - try: - return await get_user_object( - user_id=valid_token.user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - user_id_upsert=False, - proxy_logging_obj=proxy_logging_obj, - ) - except UserNotFoundError: - return None - - -async def enforced_model_allowlists( - valid_token: UserAPIKeyAuth, - prisma_client: PrismaClient | None, - user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: ProxyLogging, -) -> tuple[Sequence[str], ...]: - """One model allowlist per level that ``common_checks`` enforces on a request from this identity.""" - key_models: Final = _resolve_key_models_for_auth_check(valid_token=valid_token) - if prisma_client is None: - return (key_models, tuple(valid_token.team_models or ())) - team_object: Final = ( - None - if valid_token.team_id is None - else await get_team_object( - team_id=valid_token.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - ) - user_object: Final = ( - None - if team_object is not None - else await _user_object_or_none( - valid_token=valid_token, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - ) - project_object: Final = ( - None - if valid_token.project_id is None - else await get_project_object( - project_id=valid_token.project_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - ) - return ( - key_models, - team_object.models if team_object is not None else (), - await _team_member_granted_models( - valid_token=valid_token, - team_object=team_object, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ), - user_object.models if user_object is not None else (), - project_object.models if project_object is not None else (), - ) - - async def collect_matched_model_access_groups( model: str | Sequence[str] | None, valid_token: UserAPIKeyAuth | None, @@ -5677,6 +5604,115 @@ async def can_key_call_resolved_model( ) +async def can_personal_user_call_model( + model: str, + valid_token: UserAPIKeyAuth, + llm_router: Router | None, +) -> None: + """Apply the personal user's model allowlist, the check ``common_checks`` runs for a key with no team.""" + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + if valid_token.team_id is not None or valid_token.user_id is None or prisma_client is None: + return + try: + user_object: Final = await get_user_object( + user_id=valid_token.user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + parent_otel_span=valid_token.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except UserNotFoundError: + return + if user_object is None: + return + await can_user_call_model( + model=model, + llm_router=llm_router, + user_object=user_object, + key_model_aliases=key_model_aliases_for_auth_check(valid_token), + ) + + +async def _matching_team_object( + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> LiteLLM_TeamTableCachedObj | None: + if valid_token.team_id is None: + return None + try: + return await get_team_object( + team_id=valid_token.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=valid_token.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception as e: # noqa: BLE001 # fail-safe: a team we cannot read grants no group, it never breaks auth + verbose_proxy_logger.debug("model access group attribution could not read the team: %s", e) + return None + + +async def enforce_model_access_group_budget( + model: str, + valid_token: UserAPIKeyAuth, + llm_router: Router | None, +) -> None: + """Attribute this model to the budgeted access groups that authorize it and refuse an exhausted one. + + The pair ``common_checks`` runs for every HTTP request, for callers that reach a model only + through a group: the stamp is what the spend writer charges afterwards, and the check is what + stops a group whose budget is already gone. + + Raises: + BudgetExceededError if a group that authorized this model is over its max budget. + """ + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + if prisma_client is None: + return + matched: Final = await stamp_matched_model_access_groups( + model=model, + valid_token=valid_token, + team_object=await _matching_team_object( + valid_token=valid_token, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ), + project_object=( + None + if valid_token.project_id is None + else await get_project_object( + project_id=valid_token.project_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + ), + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + await _model_access_group_max_budget_check( + matched_model_access_groups=matched, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + + def can_org_access_model( model: str, org_object: LiteLLM_OrganizationTable | None, diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 16fbd1de376..1d1f88aaa18 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -15,7 +15,7 @@ import os import posixpath import re import sys -from collections.abc import AsyncGenerator, Callable, Mapping, Sequence +from collections.abc import AsyncGenerator, Callable, Mapping from dataclasses import dataclass from functools import partial from types import MappingProxyType @@ -50,6 +50,7 @@ from litellm.constants import ( from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException from litellm.integrations.custom_guardrail import guardrail_request_data_with_streaming, is_guardrail_intervention from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.llms.anthropic.common_utils import AnthropicModelInfo, merge_anthropic_beta_headers from litellm.llms.azure.passthrough.transformation import ( foreign_azure_deployment, @@ -72,7 +73,12 @@ from litellm.llms.oss_decision import OssDecisionProvider, oss_connection, valid from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.passthrough.main import AsyncPassthroughStreamingResponse, is_provider_error from litellm.proxy._types import * -from litellm.proxy.auth.auth_checks import enforced_model_allowlists +from litellm.proxy.auth.auth_checks import ( + can_agent_call_model, + can_key_call_resolved_model, + can_personal_user_call_model, + enforce_model_access_group_budget, +) from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy module exports @@ -110,6 +116,7 @@ from litellm.proxy.pass_through_endpoints.injected_stream_usage import ( ) from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( HttpPassThroughEndpointHelpers, + WebsocketFrameModelGate, create_pass_through_route, create_websocket_passthrough_route, websocket_passthrough_request, @@ -173,7 +180,7 @@ def is_passthrough_request_using_router_model(request_body: dict, llm_router: li """ try: model: Final = request_body.get("model") - return is_known_model(model, llm_router) + return isinstance(model, str) and is_known_model(model, llm_router) except Exception: return False @@ -203,9 +210,8 @@ def is_passthrough_request_streaming(request_body: object) -> bool: A JSON body need not be an object, so a list or scalar can reach here; it carries no streaming flag. """ - if not isinstance(request_body, dict): - return False - return bool(request_body.get("stream", False)) + body: Final = as_str_mapping(request_body) + return body is not None and bool(body.get("stream", False)) def _optional_str(value: object) -> str | None: @@ -3261,19 +3267,6 @@ def _join_url_paths(base_url: httpx.URL, path: str, custom_llm_provider: litellm return joined_path_str -_OPENAI_WS_ALL_MODEL_ACCESS: Final = frozenset( - { - SpecialModelNames.all_proxy_models.value, - SpecialModelNames.all_team_models.value, - "*", - } -) - - -def _has_model_restrictions(model_allowlists: tuple[Sequence[str], ...]) -> bool: - return any(str(model) not in _OPENAI_WS_ALL_MODEL_ACCESS for allowlist in model_allowlists for model in allowlist) - - @dataclass(frozen=True, slots=True) class _OpenAIWebsocketRefusal: close_reason: str @@ -3290,6 +3283,11 @@ class _OpenAIWebsocketErrorFrame(TypedDict): error: ReadOnly[_OpenAIWebsocketErrorDetail] +_OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL: Final = ( + "This gateway could not authorize the model named in this frame, so the session was closed. Try again, " + "and ask a proxy admin to check the gateway logs if it keeps happening." +) + _OPENAI_WS_DISABLED_REFUSAL: Final = _OpenAIWebsocketRefusal( close_reason="OpenAI websocket passthrough is disabled", message=( @@ -3298,14 +3296,6 @@ _OPENAI_WS_DISABLED_REFUSAL: Final = _OpenAIWebsocketRefusal( ), ) -_OPENAI_WS_MODEL_RESTRICTED_REFUSAL: Final = _OpenAIWebsocketRefusal( - close_reason="Keys with model restrictions cannot use OpenAI websocket passthrough", - message=( - "Keys with model restrictions cannot use OpenAI websocket passthrough, because this route " - "relays frames to the provider without reading which model they ask for." - ), -) - def _is_openai_websocket_passthrough_enabled(general_settings: Mapping[str, object]) -> bool: setting: Final = general_settings.get("enable_openai_websocket_passthrough") @@ -3314,19 +3304,9 @@ def _is_openai_websocket_passthrough_enabled(general_settings: Mapping[str, obje return setting is True -class _OpenAIWebsocketModelAllowlists(Protocol): - async def __call__(self, valid_token: UserAPIKeyAuth, /) -> tuple[Sequence[str], ...]: ... - - -async def _openai_websocket_refusal( - user_api_key_dict: UserAPIKeyAuth, - general_settings: Mapping[str, object], - model_allowlists: _OpenAIWebsocketModelAllowlists, -) -> _OpenAIWebsocketRefusal | None: +def _openai_websocket_refusal(general_settings: Mapping[str, object]) -> _OpenAIWebsocketRefusal | None: if not _is_openai_websocket_passthrough_enabled(general_settings): return _OPENAI_WS_DISABLED_REFUSAL - if _has_model_restrictions(await model_allowlists(user_api_key_dict)): - return _OPENAI_WS_MODEL_RESTRICTED_REFUSAL return None @@ -3341,6 +3321,7 @@ class _WebsocketRelay(Protocol): forward_headers: bool, endpoint: str, accept_websocket: bool, + client_frame_model_gate: WebsocketFrameModelGate | None = None, ) -> None: ... @@ -3348,18 +3329,42 @@ def _websocket_relay() -> _WebsocketRelay: return websocket_passthrough_request -def _proxy_model_allowlists() -> _OpenAIWebsocketModelAllowlists: - from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache +def _proxy_frame_model_gate() -> WebsocketFrameModelGate: + from litellm.proxy.proxy_server import ( + llm_model_list, + llm_router, + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) - async def resolve(valid_token: UserAPIKeyAuth, /) -> tuple[Sequence[str], ...]: - return await enforced_model_allowlists( - valid_token=valid_token, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) + async def gate(model: str, valid_token: UserAPIKeyAuth, /) -> str | None: + try: + await can_key_call_resolved_model( + model=model, llm_model_list=llm_model_list, valid_token=valid_token, llm_router=llm_router + ) + await can_agent_call_model( + model=model, + valid_token=valid_token, + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + await can_personal_user_call_model(model=model, valid_token=valid_token, llm_router=llm_router) + await enforce_model_access_group_budget(model=model, valid_token=valid_token, llm_router=llm_router) + except ProxyException as denial: + return denial.message + except HTTPException as denial: + return str(denial.detail) + except litellm.BudgetExceededError as exhausted: + return exhausted.message + except Exception: # noqa: BLE001 # fail closed: a gate that cannot answer refuses, and says so in band + verbose_proxy_logger.exception("OpenAI websocket passthrough: the frame model gate failed on %s", model) + return _OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL + return None - return resolve + return gate def _negotiated_websocket_subprotocol(websocket: WebSocket) -> str | None: @@ -3379,12 +3384,12 @@ async def openai_websocket_proxy_route( user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth_websocket)], general_settings: Annotated[Mapping[str, object], Depends(_proxy_general_settings)], relay: Annotated[_WebsocketRelay, Depends(_websocket_relay)], - model_allowlists: Annotated[_OpenAIWebsocketModelAllowlists, Depends(_proxy_model_allowlists)], + frame_model_gate: Annotated[WebsocketFrameModelGate, Depends(_proxy_frame_model_gate)], ) -> None: """WebSocket passthrough for OpenAI prefixes (realtime / responses.connect).""" negotiated_subprotocol: Final = _negotiated_websocket_subprotocol(websocket) - refusal: Final = await _openai_websocket_refusal(user_api_key_dict, general_settings, model_allowlists) + refusal: Final = _openai_websocket_refusal(general_settings) if refusal is not None: await websocket.accept(subprotocol=negotiated_subprotocol) error_frame: Final[_OpenAIWebsocketErrorFrame] = { @@ -3438,6 +3443,7 @@ async def openai_websocket_proxy_route( forward_headers=False, endpoint=websocket.url.path, accept_websocket=False, + client_frame_model_gate=frame_model_gate, ) @@ -3683,7 +3689,8 @@ async def _tinyfish_json_object_field_names(request: Request) -> frozenset[str] parsed: Final = _DECODED_JSON.validate_python(json.loads(raw_body)) except (json.JSONDecodeError, UnicodeDecodeError): return None - return frozenset(parsed) if isinstance(parsed, dict) else None + fields: Final = as_str_mapping(parsed) + return None if fields is None else frozenset(fields) def _tinyfish_route_timeout() -> float | None: diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py new file mode 100644 index 00000000000..300aed08923 --- /dev/null +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py @@ -0,0 +1,207 @@ +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from datetime import datetime +from types import MappingProxyType +from typing import Final, Literal, TypeAlias, cast # noqa: TID251 # cost helpers type raw frames as a TypedDict list + +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.cost_calculator import RealtimeAPITokenUsageProcessor, handle_realtime_stream_cost_calculation +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.proxy._types import PassThroughEndpointLoggingTypedDict +from litellm.responses.utils import ResponseAPILoggingUtils +from litellm.types.llms.openai import OpenAIRealtimeStreamList, ResponseAPIUsage +from litellm.types.utils import LlmProviders, ModelResponse, Usage + +WebsocketMessages: TypeAlias = Sequence[Mapping[str, object]] + +OPENAI_WEBSOCKET_PROVIDER: Final = LlmProviders.OPENAI.value + +_SESSION_CREATED: Final = "session.created" +_RESPONSE_CREATED: Final = "response.created" +_REALTIME_RESPONSE_DONE: Final = "response.done" +_RESPONSES_COMPLETED: Final = "response.completed" +_MODEL_NAMING_EVENTS: Final = frozenset( + {_SESSION_CREATED, _RESPONSE_CREATED, _REALTIME_RESPONSE_DONE, _RESPONSES_COMPLETED} +) +_RESPONSE_COST_HEADER: Final = "llm_provider-x-litellm-response-cost" + + +class _NamedModel(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + model: str | None = None + + +class _ModelNamingEvent(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + session: _NamedModel | None = None + response: _NamedModel | None = None + + +class _CompletedResponse(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + model: str | None = None + usage: ResponseAPIUsage + + +class _ResponsesCompletedEvent(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + type: Literal["response.completed"] + response: _CompletedResponse + + +_MODEL_NAMING_EVENT: Final = TypeAdapter(_ModelNamingEvent) +_RESPONSES_COMPLETED_EVENT: Final = TypeAdapter(_ResponsesCompletedEvent) + + +@dataclass(frozen=True, slots=True) +class _Billing: + usage: Usage + cost: float + + +@dataclass(frozen=True, slots=True) +class _CompletedTurn: + model: str + usage: Usage + + +def _event_model(event: Mapping[str, object]) -> str | None: + if event.get("type") not in _MODEL_NAMING_EVENTS: + return None + try: + naming: Final = _MODEL_NAMING_EVENT.validate_python(event) + except ValidationError: + return None + return next( + ( + container.model + for container in (naming.session, naming.response) + if container is not None and container.model + ), + None, + ) + + +def observed_websocket_model(messages: WebsocketMessages) -> str | None: + return next((model for model in map(_event_model, messages) if model is not None), None) + + +def _realtime_billing(messages: WebsocketMessages, model: str, logging_obj: LiteLLMLoggingObj) -> _Billing | None: + if not any(event.get("type") == _REALTIME_RESPONSE_DONE for event in messages): + return None + results: Final = cast( # cast-ok: the realtime cost helpers read raw provider frames through their TypedDict view + OpenAIRealtimeStreamList, tuple(event for event in messages if "type" in event) + ) + usage: Final = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results(results) + cost: Final = handle_realtime_stream_cost_calculation( + results=results, + combined_usage_object=usage, + custom_llm_provider=OPENAI_WEBSOCKET_PROVIDER, + litellm_model_name=model, + litellm_logging_obj=logging_obj, + ) + return _Billing(usage=usage, cost=cost) + + +def _responses_turn(event: Mapping[str, object], session_model: str) -> _CompletedTurn | None: + if event.get("type") != _RESPONSES_COMPLETED: + return None + try: + completed: Final = _RESPONSES_COMPLETED_EVENT.validate_python(event) + except ValidationError: + verbose_proxy_logger.warning( + "OpenAI websocket passthrough: a response.completed frame carried no readable usage, so it is not billed" + ) + return None + return _CompletedTurn( + model=completed.response.model or session_model, + usage=ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( # pyright: ignore[reportPrivateUsage] # the one shared Responses usage converter + completed.response.usage + ), + ) + + +def _responses_billing(messages: WebsocketMessages, model: str) -> _Billing | None: + turns: Final = tuple(turn for turn in (_responses_turn(event, model) for event in messages) if turn is not None) + if not turns: + return None + usage: Final = RealtimeAPITokenUsageProcessor.combine_usage_objects([turn.usage for turn in turns]) + return _Billing(usage=usage, cost=sum(_responses_cost(turn.usage, turn.model) for turn in turns)) + + +def _responses_cost(usage: Usage, model: str) -> float: + try: + return litellm.completion_cost( + completion_response=ModelResponse(model=model, usage=usage), + model=model, + custom_llm_provider=OPENAI_WEBSOCKET_PROVIDER, + ) + except Exception: # noqa: BLE001 # completion_cost raises a bare Exception for an unpriced model, and the turn must still log + verbose_proxy_logger.warning( + "OpenAI websocket passthrough: could not price a Responses turn on %s, logging it at zero", model + ) + return 0.0 + + +def _priced_result(model: str, billing: _Billing, start_time: datetime) -> ModelResponse: + result: Final = ModelResponse( + id=f"openai-websocket-{start_time.timestamp()}", + object="chat.completion", + created=int(start_time.timestamp()), + model=model, + usage=billing.usage, + ) + hidden: Final = {"additional_headers": MappingProxyType({_RESPONSE_COST_HEADER: billing.cost})} + result._hidden_params = hidden # pyright: ignore[reportPrivateUsage] # the one settled-cost channel + return result + + +class OpenAIWebsocketPassthroughLoggingHandler: + def openai_websocket_passthrough_handler( + self, + websocket_messages: WebsocketMessages, + logging_obj: LiteLLMLoggingObj, + url_route: str, + start_time: datetime, + kwargs: Mapping[str, object], + ) -> PassThroughEndpointLoggingTypedDict: + model: Final = observed_websocket_model(websocket_messages) + if model is None: + verbose_proxy_logger.debug("OpenAI websocket passthrough (%s): no model named, skipping cost", url_route) + unpriced: Final[PassThroughEndpointLoggingTypedDict] = {"result": None, "kwargs": {**kwargs}} + return unpriced + + logging_obj.model = model # rebind-ok: the spend logger reads model and cost off the shared logging object + logging_obj.model_call_details.update( + MappingProxyType({"model": model, "custom_llm_provider": OPENAI_WEBSOCKET_PROVIDER}) + ) + + billing: Final = _realtime_billing(websocket_messages, model, logging_obj) or _responses_billing( + websocket_messages, model + ) + if billing is None: + verbose_proxy_logger.debug("OpenAI websocket passthrough (%s): no usage reported, skipping cost", url_route) + unbilled: Final[PassThroughEndpointLoggingTypedDict] = { + "result": None, + "kwargs": {**kwargs, "model": model, "custom_llm_provider": OPENAI_WEBSOCKET_PROVIDER}, + } + return unbilled + + logging_obj.model_call_details.update(MappingProxyType({"response_cost": billing.cost})) + priced: Final[PassThroughEndpointLoggingTypedDict] = { + "result": _priced_result(model, billing, start_time), + "kwargs": { + **kwargs, + "model": model, + "custom_llm_provider": OPENAI_WEBSOCKET_PROVIDER, + "response_cost": billing.cost, + }, + } + return priced diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index e1e8864678e..895533044c3 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -5,12 +5,13 @@ import json import posixpath import traceback from base64 import b64encode -from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterable, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterable, Mapping, MutableMapping, Sequence from dataclasses import dataclass from datetime import datetime +from functools import reduce from itertools import chain, count, groupby from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, TypedDict, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict, cast from urllib.parse import urlencode, urlparse import httpx @@ -23,6 +24,7 @@ from fastapi import ( Response, UploadFile, WebSocket, + WebSocketDisconnect, status, ) from fastapi.responses import StreamingResponse @@ -30,7 +32,7 @@ from pydantic import TypeAdapter from starlette.datastructures import UploadFile as StarletteUploadFile from starlette.routing import BaseRoute, Route from starlette.websockets import WebSocketState -from websockets.asyncio.client import connect +from websockets.asyncio.client import ClientConnection, connect from websockets.exceptions import ( ConnectionClosedError, ConnectionClosedOK, @@ -55,6 +57,7 @@ from litellm.integrations.custom_guardrail import ( ) from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, bind_budget_reservation_to_callbacks, get_metadata_variable_name_from_kwargs, get_or_create_metadata_bucket, @@ -64,6 +67,7 @@ from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.litellm_logging import get_masked_values from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.litellm_core_utils.realtime_errors import realtime_error_event from litellm.litellm_core_utils.redact_messages import should_redact_message_logging from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.anthropic.common_utils import is_anthropic_messages_url @@ -2454,6 +2458,106 @@ def _with_trace_context(headers: Mapping[str, str], parent_span: object) -> dict return inject_trace_context(headers, parent_span=parent_span) +class WebsocketFrameModelGate(Protocol): + async def __call__(self, model: str, valid_token: UserAPIKeyAuth, /) -> str | None: ... + + +@dataclass(frozen=True, slots=True) +class _ForwardClientFrame: + pass + + +@dataclass(frozen=True, slots=True) +class _RefuseClientFrame: + message: str + + +_PLAIN_MODEL_KEY: Final = '"model"' +_JSON_STRING_ESCAPE: Final = "\\" +_CLIENT_FRAME_MODEL_PLACEMENTS: Final = ( + ("model",), + ("response", "model"), + ("session", "model"), + ("session", "input_audio_transcription", "model"), + ("session", "audio", "input", "transcription", "model"), +) + + +def _may_name_a_model(text_data: str) -> bool: + return _PLAIN_MODEL_KEY in text_data or _JSON_STRING_ESCAPE in text_data + + +def _nested_frame(frame: Mapping[str, object] | None, key: str) -> Mapping[str, object] | None: + return None if frame is None else as_str_mapping(frame.get(key)) + + +def _placed_model(frame: Mapping[str, object], placement: Sequence[str]) -> str | None: + *parents, leaf = placement + container: Final = reduce(_nested_frame, parents, frame) + if container is None: + return None + value: Final = container.get(leaf) + return value if isinstance(value, str) else None + + +def _client_frame_models(frame_data: str | bytes) -> tuple[str, ...]: + if isinstance(frame_data, str) and not _may_name_a_model(frame_data): + return () + frame: Final = _json_object_frame(frame_data) + if frame is None: + return () + placed: Final = (_placed_model(frame, placement) for placement in _CLIENT_FRAME_MODEL_PLACEMENTS) + return tuple(dict.fromkeys(model for model in placed if model is not None)) + + +async def _client_frame_verdict( + frame_data: str | bytes, gate: WebsocketFrameModelGate | None, valid_token: UserAPIKeyAuth +) -> _ForwardClientFrame | _RefuseClientFrame: + if gate is None: + return _ForwardClientFrame() + for model in _client_frame_models(frame_data): + if (refusal := await gate(model, valid_token)) is not None: + return _RefuseClientFrame(message=refusal) + return _ForwardClientFrame() + + +async def _refuse_client_frame(websocket: WebSocket, upstream_ws: ClientConnection, message: str) -> None: + await websocket.send_text(realtime_error_event(message, error_type="invalid_request_error")) + await websocket.close(code=CloseCode.POLICY_VIOLATION, reason=_truncated_close_reason(message)) + await upstream_ws.close() + + +def _stamp_billed_access_groups(kwargs: Mapping[str, object], valid_token: UserAPIKeyAuth) -> None: + """Refresh the session's logging metadata with the access groups its frames were authorized through. + + The gate only learns them once a frame names a model, which is long after the relay stamped the + key's groups onto its metadata, so the spend writer would otherwise charge no group for the session. + """ + groups: Final = valid_token.matched_model_access_groups + litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) + metadata: Final = None if litellm_params is None else litellm_params.get("metadata") + if groups and isinstance(metadata, MutableMapping): + metadata[MODEL_ACCESS_GROUP_METADATA_KEY] = groups + + +async def _relay_client_frame( + websocket: WebSocket, + upstream_ws: ClientConnection, + checked_frame: str | bytes, + forwarded_frame: str | bytes, + gate: WebsocketFrameModelGate | None, + valid_token: UserAPIKeyAuth, +) -> bool: + """Forward one client frame upstream, or refuse it and report that the session is over.""" + match await _client_frame_verdict(checked_frame, gate, valid_token): + case _RefuseClientFrame(message=refusal_message): + await _refuse_client_frame(websocket, upstream_ws, refusal_message) + return False + case _ForwardClientFrame(): + await upstream_ws.send(forwarded_frame) + return True + + async def websocket_passthrough_request( websocket: WebSocket, target: str, @@ -2464,6 +2568,7 @@ async def websocket_passthrough_request( cost_per_request: float | None = None, accept_websocket: bool = True, setup_model_rewriter: Callable[[str], str] | None = None, + client_frame_model_gate: WebsocketFrameModelGate | None = None, ): """ WebSocket passthrough request handler. @@ -2673,9 +2778,25 @@ 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)) + if not await _relay_client_frame( + websocket, + upstream_ws, + text_data, + _rewrite_vertex_live_setup_model(text_data, setup_model_rewriter), + client_frame_model_gate, + user_api_key_dict, + ): + break elif bytes_data is not None: - await upstream_ws.send(bytes_data) + if not await _relay_client_frame( + websocket, + upstream_ws, + bytes_data, + bytes_data, + client_frame_model_gate, + user_api_key_dict, + ): + break except asyncio.CancelledError: raise except Exception: @@ -2702,11 +2823,7 @@ 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: - if isinstance(upstream_message, bytes): - await websocket.send_bytes(upstream_message) - else: - await websocket.send_text(upstream_message) + def record_upstream_frame(upstream_message: str | bytes) -> None: message_data: Final = _json_object_frame(upstream_message) if message_data is None: return @@ -2715,6 +2832,13 @@ async def websocket_passthrough_request( return websocket_messages.append(message_data) + async def relay_upstream_frame(upstream_message: str | bytes) -> None: + record_upstream_frame(upstream_message) + if isinstance(upstream_message, bytes): + await websocket.send_bytes(upstream_message) + else: + await websocket.send_text(upstream_message) + async def forward_upstream_to_client() -> Close | None: try: while True: @@ -2722,6 +2846,11 @@ async def websocket_passthrough_request( except (ConnectionClosedOK, ConnectionClosedError) as e: verbose_proxy_logger.debug("Upstream WebSocket connection closed: %s", e) return e.rcvd + except WebSocketDisconnect: + verbose_proxy_logger.debug( + "WebSocket passthrough (%s): client left while upstream was still sending", endpoint + ) + return None except asyncio.CancelledError: verbose_proxy_logger.debug("asyncio.CancelledError in forward_upstream_to_client") raise @@ -2767,6 +2896,8 @@ async def websocket_passthrough_request( passthrough_logging_payload["response_body"] = websocket_messages passthrough_logging_payload["end_time"] = end_time + _stamp_billed_access_groups(kwargs, user_api_key_dict) + # Remove logging_obj from kwargs to avoid duplicate keyword argument success_kwargs: Final = kwargs.copy() success_kwargs.pop("logging_obj", None) diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index 9c6c3111209..433d7501958 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -1,4 +1,5 @@ import json +from collections.abc import Mapping, Sequence from datetime import datetime from types import MappingProxyType from typing import Any, Final @@ -8,7 +9,7 @@ import httpx 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, ) @@ -64,6 +65,14 @@ def _safe_response_text(httpx_response: httpx.Response) -> str: return "" +def _websocket_frames( + response_body: Mapping[str, object] | Sequence[Mapping[str, object]] | None, +) -> tuple[Mapping[str, object], ...]: + return tuple( + message for message in (response_body if isinstance(response_body, list) else ()) if isinstance(message, dict) + ) + + class PassThroughEndpointLogging: def __init__( self, @@ -114,6 +123,7 @@ class PassThroughEndpointLogging: # Vertex AI Live API WebSocket self.TRACKED_VERTEX_AI_LIVE_ROUTES = ["/vertex_ai/live"] + self.TRACKED_OPENAI_WEBSOCKET_ROUTES: Final = ("/openai/", "/openai_passthrough/") @property def _log_dispatch(self) -> PassThroughLogDispatch: @@ -385,21 +395,18 @@ 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 ( + websocket_relay_result := self._websocket_relay_logging_result( + url_route=url_route, + websocket_frames=_websocket_frames(response_body), + logging_obj=logging_obj, + httpx_response=httpx_response, + start_time=start_time, + kwargs=kwargs, ) - standard_logging_response_object = deepgram_handler_result["result"] # rebind-ok: elif-chain - kwargs = deepgram_handler_result["kwargs"] # rebind-ok: elif-chain contract + ) is not None: + standard_logging_response_object = websocket_relay_result["result"] # rebind-ok: elif-chain + kwargs = websocket_relay_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({}), @@ -415,6 +422,36 @@ class PassThroughEndpointLogging: return_dict["kwargs"] = kwargs return return_dict + def _websocket_relay_logging_result( + self, + url_route: str, + websocket_frames: tuple[Mapping[str, object], ...], + logging_obj: LiteLLMLoggingObj, + httpx_response: httpx.Response, + start_time: datetime, + kwargs: Mapping[str, object], + ) -> PassThroughEndpointLoggingTypedDict | None: + if self.is_openai_websocket_route(url_route): + from .llm_provider_handlers.openai_websocket_passthrough_logging_handler import ( + OpenAIWebsocketPassthroughLoggingHandler, + ) + + return OpenAIWebsocketPassthroughLoggingHandler().openai_websocket_passthrough_handler( + websocket_messages=websocket_frames, + logging_obj=logging_obj, + url_route=url_route, + start_time=start_time, + kwargs=kwargs, + ) + if DeepgramListenPassthroughLoggingHandler.is_deepgram_listen_route(url_route): + return DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler( + websocket_messages=websocket_frames, + logging_obj=logging_obj, + upstream_url=str(httpx_response.request.url), + kwargs=kwargs, + ) + return None + async def pass_through_async_success_handler( self, httpx_response: httpx.Response, @@ -584,6 +621,9 @@ class PassThroughEndpointLogging: return True return False + def is_openai_websocket_route(self, url_route: str) -> bool: + return url_route.startswith(self.TRACKED_OPENAI_WEBSOCKET_ROUTES) + def is_cursor_route(self, url_route: str, custom_llm_provider: str | None = None): """Check if the URL route is a Cursor Cloud Agents API route.""" if custom_llm_provider == "cursor":