mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
fix(proxy): bill pipelined websocket turns to their own frame and close the review gaps
This commit is contained in:
mateo-berri
parent
a78b998f1d
commit
2d860f05d4
10 changed files with 530 additions and 94 deletions
|
|
@ -3,7 +3,7 @@
|
|||
import copy
|
||||
import logging
|
||||
import re
|
||||
from collections.abc import Collection, Iterable, Iterator, Mapping
|
||||
from collections.abc import Collection, Iterable, Iterator, Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, cast, runtime_checkable
|
||||
|
||||
|
|
@ -59,6 +59,10 @@ def as_str_mapping(value: object) -> Mapping[str, object] | None:
|
|||
return value if isinstance(value, Mapping) else None
|
||||
|
||||
|
||||
def as_object_list(value: object) -> Sequence[object] | None:
|
||||
return value if isinstance(value, list) else None
|
||||
|
||||
|
||||
def decoded_json_object(value: object) -> Mapping[str, object]:
|
||||
mapping: Final = as_str_mapping(value)
|
||||
if mapping is not None:
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import json
|
||||
from functools import lru_cache
|
||||
from typing import Final
|
||||
|
||||
|
|
@ -10,7 +11,7 @@ def realtime_request_body(model: str | None) -> bytes:
|
|||
Generate the realtime websocket request body. Cached with LRU semantics to avoid repeated
|
||||
string formatting work while keeping memory usage bounded.
|
||||
"""
|
||||
return f'{{"model": "{model or ""}"}}'.encode()
|
||||
return json.dumps({"model": model or ""}).encode()
|
||||
|
||||
|
||||
_realtime_request_body: Final = realtime_request_body
|
||||
|
|
|
|||
|
|
@ -15,7 +15,7 @@ import os
|
|||
import posixpath
|
||||
import re
|
||||
import sys
|
||||
from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping
|
||||
from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from functools import partial
|
||||
from types import MappingProxyType
|
||||
|
|
@ -50,7 +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.litellm_core_utils.core_helpers import as_object_list, as_str_mapping, attribute_or_none
|
||||
from litellm.llms.anthropic.common_utils import AnthropicModelInfo, merge_anthropic_beta_headers
|
||||
from litellm.llms.azure.passthrough.transformation import (
|
||||
foreign_azure_deployment,
|
||||
|
|
@ -73,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 can_agent_call_model, can_key_call_resolved_model
|
||||
from litellm.proxy.auth.auth_checks import (
|
||||
can_agent_call_model,
|
||||
can_key_call_resolved_model,
|
||||
key_model_aliases_for_auth_check,
|
||||
team_model_aliases_for_auth_check,
|
||||
)
|
||||
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
|
||||
|
|
@ -102,6 +107,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg
|
|||
safe_get_request_headers,
|
||||
safe_set_request_parsed_body,
|
||||
)
|
||||
from litellm.proxy.common_utils.model_listing_utils import alias_map
|
||||
from litellm.proxy.common_utils.openai_error_payload import LITELLM_CALL_ID_HEADER
|
||||
from litellm.proxy.common_utils.resource_ownership import is_proxy_admin
|
||||
from litellm.proxy.common_utils.sse_keepalive import (
|
||||
|
|
@ -121,6 +127,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket
|
|||
openai_websocket_frame_kind,
|
||||
)
|
||||
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
|
||||
STORED_PROMPT_PROBE_MODEL,
|
||||
HttpPassThroughEndpointHelpers,
|
||||
WebsocketFrameModelAllowed,
|
||||
WebsocketFrameModelGate,
|
||||
|
|
@ -132,7 +139,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
|
|||
websocket_passthrough_request,
|
||||
)
|
||||
from litellm.proxy.utils import ProxyLogging as ProxyLoggingType
|
||||
from litellm.proxy.utils import is_known_model
|
||||
from litellm.proxy.utils import is_known_model, normalize_route_for_root_path
|
||||
from litellm.proxy.vector_store_endpoints.utils import (
|
||||
assert_proxy_admin_for_vector_store_index_management,
|
||||
assert_user_can_access_vector_store,
|
||||
|
|
@ -3304,6 +3311,11 @@ _OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL: Final = (
|
|||
"and ask a proxy admin to check the gateway logs if it keeps happening."
|
||||
)
|
||||
|
||||
_OPENAI_WS_ALIASED_MODEL_REFUSAL: Final = (
|
||||
"This key cannot use a model alias over OpenAI websocket passthrough. The passthrough sends each frame to "
|
||||
"OpenAI unchanged, so name the model this key may use directly."
|
||||
)
|
||||
|
||||
_OPENAI_WS_DISABLED_REFUSAL: Final = _OpenAIWebsocketRefusal(
|
||||
close_reason="OpenAI websocket passthrough is disabled",
|
||||
message=(
|
||||
|
|
@ -3357,36 +3369,65 @@ def _frame_model_authorizer() -> FrameModelAuthorizer:
|
|||
return user_api_key_auth_accepted_websocket_for_model
|
||||
|
||||
|
||||
def _proxy_model_list() -> Sequence[object] | None:
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
return as_object_list(attribute_or_none(proxy_server, "llm_model_list"))
|
||||
|
||||
|
||||
def _gateway_remaps_model(model: str, frame_auth: UserAPIKeyAuth, llm_router: Router | None) -> bool:
|
||||
alias_maps: Final = (
|
||||
key_model_aliases_for_auth_check(frame_auth) or {},
|
||||
team_model_aliases_for_auth_check(frame_auth) or {},
|
||||
alias_map(litellm.model_alias_map),
|
||||
)
|
||||
targets: Final = (
|
||||
*(aliases.get(model) for aliases in alias_maps),
|
||||
llm_router.get_model_from_alias(model) if llm_router is not None else None,
|
||||
)
|
||||
return any(target not in (None, model) for target in targets)
|
||||
|
||||
|
||||
def _aliased_model_verdict(
|
||||
allowed: WebsocketFrameModelAllowed, probe: WebsocketFrameModelVerdict
|
||||
) -> WebsocketFrameModelVerdict:
|
||||
match probe:
|
||||
case WebsocketFrameModelRefused(error_type=error_type) if error_type.endswith("_model_access_denied"):
|
||||
return WebsocketFrameModelRefused(message=_OPENAI_WS_ALIASED_MODEL_REFUSAL, error_type=error_type)
|
||||
case WebsocketFrameModelRefused():
|
||||
return probe
|
||||
case WebsocketFrameModelAllowed():
|
||||
return allowed
|
||||
|
||||
|
||||
def _proxy_frame_model_gate(
|
||||
authorize_frame_model: Annotated[FrameModelAuthorizer, Depends(_frame_model_authorizer)],
|
||||
) -> WebsocketFrameModelGate:
|
||||
from litellm.proxy.proxy_server import (
|
||||
llm_model_list,
|
||||
llm_router,
|
||||
prisma_client,
|
||||
proxy_logging_obj,
|
||||
user_api_key_cache,
|
||||
)
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.spend_tracking.budget_reservation import release_or_invalidate_budget_reservation
|
||||
|
||||
async def authorize(websocket: WebSocket, model: str) -> WebsocketFrameModelAllowed:
|
||||
async def authorize(websocket: WebSocket, model: str) -> WebsocketFrameModelVerdict:
|
||||
frame_auth: Final = await authorize_frame_model(websocket, model)
|
||||
access_groups: Final = tuple(frame_auth.matched_model_access_groups or ())
|
||||
llm_router: Final = proxy_server.llm_router
|
||||
try:
|
||||
await can_key_call_resolved_model(
|
||||
model=model, llm_model_list=llm_model_list, valid_token=frame_auth, llm_router=llm_router
|
||||
model=model, llm_model_list=_proxy_model_list(), valid_token=frame_auth, llm_router=llm_router
|
||||
)
|
||||
await can_agent_call_model(
|
||||
model=model,
|
||||
valid_token=frame_auth,
|
||||
llm_router=llm_router,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
prisma_client=proxy_server.prisma_client,
|
||||
user_api_key_cache=proxy_server.user_api_key_cache,
|
||||
proxy_logging_obj=proxy_server.proxy_logging_obj,
|
||||
)
|
||||
finally:
|
||||
await release_or_invalidate_budget_reservation(budget_reservation=frame_auth.budget_reservation)
|
||||
return WebsocketFrameModelAllowed(access_groups=access_groups)
|
||||
allowed: Final = WebsocketFrameModelAllowed(access_groups=access_groups)
|
||||
if not _gateway_remaps_model(model, frame_auth, llm_router):
|
||||
return allowed
|
||||
return _aliased_model_verdict(allowed, await gate(websocket, STORED_PROMPT_PROBE_MODEL))
|
||||
|
||||
async def gate(websocket: WebSocket, model: str, /) -> WebsocketFrameModelVerdict:
|
||||
try:
|
||||
|
|
@ -3478,7 +3519,7 @@ async def openai_websocket_proxy_route(
|
|||
custom_headers=custom_headers,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
forward_headers=False,
|
||||
endpoint=websocket.url.path,
|
||||
endpoint=normalize_route_for_root_path(websocket.url.path) or websocket.url.path,
|
||||
accept_websocket=False,
|
||||
client_frame_model_gate=frame_model_gate,
|
||||
upstream_frame_kind=openai_websocket_frame_kind,
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ 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.core_helpers import as_str_mapping
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.proxy._types import PassThroughEndpointLoggingTypedDict
|
||||
from litellm.responses.utils import ResponseAPILoggingUtils
|
||||
|
|
@ -139,15 +140,18 @@ def _responses_turn(event: Mapping[str, object], session_model: str) -> _Complet
|
|||
)
|
||||
|
||||
|
||||
def _responses_billing(messages: WebsocketMessages, model: str) -> _Billing | None:
|
||||
def _responses_billing(messages: WebsocketMessages, model: str, requested_model: str | None) -> _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))
|
||||
return _Billing(
|
||||
usage=usage,
|
||||
cost=sum(_responses_cost(turn.usage, (turn.model, requested_model)) for turn in turns),
|
||||
)
|
||||
|
||||
|
||||
def _responses_cost(usage: Usage, model: str) -> float:
|
||||
def _model_cost(usage: Usage, model: str) -> float | None:
|
||||
try:
|
||||
return litellm.completion_cost(
|
||||
completion_response=ModelResponse(model=model, usage=usage),
|
||||
|
|
@ -155,10 +159,25 @@ def _responses_cost(usage: Usage, model: str) -> float:
|
|||
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
|
||||
return None
|
||||
|
||||
|
||||
def _responses_cost(usage: Usage, price_as: Sequence[str | None]) -> float:
|
||||
costs: Final = (_model_cost(usage, model) for model in dict.fromkeys(price_as) if model is not None)
|
||||
cost: Final = next((cost for cost in costs if cost), None)
|
||||
if cost is None:
|
||||
verbose_proxy_logger.warning(
|
||||
"OpenAI websocket passthrough: could not price a Responses turn on %s, logging it at zero", model
|
||||
"OpenAI websocket passthrough: could not price a Responses turn on %s, logging it at zero", price_as
|
||||
)
|
||||
return 0.0
|
||||
return cost
|
||||
|
||||
|
||||
def _requested_model(kwargs: Mapping[str, object]) -> str | None:
|
||||
litellm_params: Final = as_str_mapping(kwargs.get("litellm_params"))
|
||||
metadata: Final = None if litellm_params is None else as_str_mapping(litellm_params.get("metadata"))
|
||||
model_group: Final = None if metadata is None else metadata.get("model_group")
|
||||
return model_group if isinstance(model_group, str) else None
|
||||
|
||||
|
||||
def _priced_result(model: str, billing: _Billing, start_time: datetime) -> ModelResponse:
|
||||
|
|
@ -195,7 +214,7 @@ class OpenAIWebsocketPassthroughLoggingHandler:
|
|||
)
|
||||
|
||||
billing: Final = _realtime_billing(websocket_messages, model, logging_obj) or _responses_billing(
|
||||
websocket_messages, model
|
||||
websocket_messages, model, _requested_model(kwargs)
|
||||
)
|
||||
if billing is None:
|
||||
verbose_proxy_logger.debug("OpenAI websocket passthrough (%s): no usage reported, skipping cost", url_route)
|
||||
|
|
|
|||
|
|
@ -13,9 +13,10 @@ from functools import reduce
|
|||
from itertools import chain, count, groupby
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias, TypedDict, cast
|
||||
from urllib.parse import urlencode, urlparse
|
||||
from urllib.parse import parse_qs, urlencode, urlparse
|
||||
|
||||
import httpx
|
||||
import orjson
|
||||
from fastapi import (
|
||||
APIRouter,
|
||||
Depends,
|
||||
|
|
@ -60,6 +61,7 @@ from litellm.integrations.custom_guardrail import (
|
|||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.integrations.otel.model.semconv import PASSTHROUGH_OPERATION_KEY, GenAIOperation
|
||||
from litellm.litellm_core_utils.core_helpers import (
|
||||
as_object_list,
|
||||
as_str_mapping,
|
||||
bind_budget_reservation_to_callbacks,
|
||||
get_metadata_variable_name_from_kwargs,
|
||||
|
|
@ -2441,11 +2443,16 @@ def _resolved_vertex_live_setup(
|
|||
return {**setup_data, "model": setup_model_rewriter(setup_model)}
|
||||
|
||||
|
||||
def _decoded_frame(frame: str | bytes) -> object:
|
||||
try:
|
||||
return _DECODED_JSON.validate_python(orjson.loads(frame))
|
||||
except orjson.JSONDecodeError:
|
||||
return _DECODED_JSON.validate_python(json.loads(frame if isinstance(frame, str) else frame.decode("utf-8")))
|
||||
|
||||
|
||||
def _json_object_frame(frame: str | bytes) -> dict[str, object] | None:
|
||||
try:
|
||||
decoded: Final[object] = _DECODED_JSON.validate_python(
|
||||
json.loads(frame if isinstance(frame, str) else frame.decode("utf-8"))
|
||||
)
|
||||
decoded: Final = _decoded_frame(frame)
|
||||
except (json.JSONDecodeError, UnicodeDecodeError):
|
||||
return None
|
||||
return decoded if isinstance(decoded, dict) else None
|
||||
|
|
@ -2525,9 +2532,16 @@ class WebsocketFrameModelGate(Protocol):
|
|||
WebsocketUpstreamFrameClassifier: TypeAlias = Callable[[Mapping[str, object]], WebsocketUpstreamFrameKind]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _TurnAuthorization:
|
||||
requested_model: str | None
|
||||
access_groups: tuple[str, ...]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ForwardClientFrame:
|
||||
turn_access_groups: tuple[str, ...] | None
|
||||
starts_turn: bool
|
||||
turn: _TurnAuthorization | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -2537,24 +2551,35 @@ class _RefuseClientFrame:
|
|||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ClientFrameRequests:
|
||||
starts_turn: bool
|
||||
turn_model: str | None
|
||||
models: tuple[str, ...]
|
||||
names_stored_prompt: bool
|
||||
unverifiable_model_refusal: str | None
|
||||
|
||||
|
||||
_NO_CLIENT_FRAME_REQUESTS: Final = _ClientFrameRequests(turn_model=None, models=(), names_stored_prompt=False)
|
||||
_FRAME_TEXT_WORTH_PARSING: Final = ('"model"', '"prompt"', "\\")
|
||||
_NO_CLIENT_FRAME_REQUESTS: Final = _ClientFrameRequests(
|
||||
starts_turn=False, turn_model=None, models=(), unverifiable_model_refusal=None
|
||||
)
|
||||
_FORWARD_WITHOUT_TURN: Final = _ForwardClientFrame(starts_turn=False, turn=None)
|
||||
_TURN_STARTING_FRAME: Final = "response.create"
|
||||
_FRAME_TEXT_WORTH_PARSING: Final = ('"model"', '"prompt"', '"tools"', "\\")
|
||||
_TURN_MODEL_PLACEMENTS: Final = (("model",), ("response", "model"), ("session", "model"))
|
||||
_TRANSCRIPTION_MODEL_PLACEMENTS: Final = (
|
||||
("session", "input_audio_transcription", "model"),
|
||||
("session", "audio", "input", "transcription", "model"),
|
||||
)
|
||||
_STORED_PROMPT_PLACEMENTS: Final = (("prompt",), ("response", "prompt"), ("session", "prompt"))
|
||||
_TOOL_LIST_PLACEMENTS: Final = (("tools",), ("response", "tools"), ("session", "tools"))
|
||||
_TOOL_THAT_RUNS_A_MODEL: Final = "image_generation"
|
||||
STORED_PROMPT_PROBE_MODEL: Final = "litellm-unverifiable-stored-prompt-model"
|
||||
_STORED_PROMPT_REFUSAL: Final = (
|
||||
"This key is limited to specific models, so it cannot run a stored prompt. "
|
||||
"The gateway cannot check which model a stored prompt runs."
|
||||
)
|
||||
_UNNAMED_TOOL_MODEL_REFUSAL: Final = (
|
||||
"This key is limited to specific models, so an image_generation tool has to name its model. "
|
||||
"The gateway cannot check which model an unnamed one runs."
|
||||
)
|
||||
_FIRST_SENTENCE: Final = re.compile(r".*?[.!?](?=\s|$)")
|
||||
|
||||
|
||||
|
|
@ -2577,6 +2602,21 @@ def _placed_models(frame: Mapping[str, object], placements: Iterable[Sequence[st
|
|||
return tuple(value for value in placed if isinstance(value, str))
|
||||
|
||||
|
||||
def _placed_tools(frame: Mapping[str, object]) -> tuple[Mapping[str, object], ...]:
|
||||
tool_lists: Final = (_placed_value(frame, placement) for placement in _TOOL_LIST_PLACEMENTS)
|
||||
tools: Final = chain.from_iterable(as_object_list(tool_list) or () for tool_list in tool_lists)
|
||||
return tuple(tool for tool in map(as_str_mapping, tools) if tool is not None)
|
||||
|
||||
|
||||
def _unverifiable_model_refusal(frame: Mapping[str, object], tools: Sequence[Mapping[str, object]]) -> str | None:
|
||||
prompts: Final = (_placed_value(frame, placement) for placement in _STORED_PROMPT_PLACEMENTS)
|
||||
if any(isinstance(prompt, Mapping) for prompt in prompts):
|
||||
return _STORED_PROMPT_REFUSAL
|
||||
if any(tool.get("type") == _TOOL_THAT_RUNS_A_MODEL and not isinstance(tool.get("model"), str) for tool in tools):
|
||||
return _UNNAMED_TOOL_MODEL_REFUSAL
|
||||
return None
|
||||
|
||||
|
||||
def _client_frame_requests(frame_data: str | bytes) -> _ClientFrameRequests:
|
||||
if isinstance(frame_data, str) and not _may_name_a_model(frame_data):
|
||||
return _NO_CLIENT_FRAME_REQUESTS
|
||||
|
|
@ -2584,17 +2624,21 @@ def _client_frame_requests(frame_data: str | bytes) -> _ClientFrameRequests:
|
|||
if frame is None:
|
||||
return _NO_CLIENT_FRAME_REQUESTS
|
||||
turn_models: Final = _placed_models(frame, _TURN_MODEL_PLACEMENTS)
|
||||
prompts: Final = (_placed_value(frame, placement) for placement in _STORED_PROMPT_PLACEMENTS)
|
||||
tools: Final = _placed_tools(frame)
|
||||
tool_models: Final = (model for model in (tool.get("model") for tool in tools) if isinstance(model, str))
|
||||
return _ClientFrameRequests(
|
||||
starts_turn=frame.get("type") == _TURN_STARTING_FRAME,
|
||||
turn_model=next(iter(turn_models), None),
|
||||
models=tuple(dict.fromkeys((*turn_models, *_placed_models(frame, _TRANSCRIPTION_MODEL_PLACEMENTS)))),
|
||||
names_stored_prompt=any(isinstance(prompt, Mapping) for prompt in prompts),
|
||||
models=tuple(
|
||||
dict.fromkeys((*turn_models, *_placed_models(frame, _TRANSCRIPTION_MODEL_PLACEMENTS), *tool_models))
|
||||
),
|
||||
unverifiable_model_refusal=_unverifiable_model_refusal(frame, tools),
|
||||
)
|
||||
|
||||
|
||||
def _stored_prompt_refusal(probe: WebsocketFrameModelRefused) -> WebsocketFrameModelRefused:
|
||||
def _unverifiable_model_probe_refusal(probe: WebsocketFrameModelRefused, message: str) -> WebsocketFrameModelRefused:
|
||||
if probe.error_type.endswith("_model_access_denied"):
|
||||
return WebsocketFrameModelRefused(message=_STORED_PROMPT_REFUSAL, error_type=probe.error_type)
|
||||
return WebsocketFrameModelRefused(message=message, error_type=probe.error_type)
|
||||
return probe
|
||||
|
||||
|
||||
|
|
@ -2602,20 +2646,25 @@ async def _client_frame_verdict(
|
|||
websocket: WebSocket, frame_data: str | bytes, gate: WebsocketFrameModelGate | None
|
||||
) -> _ForwardClientFrame | _RefuseClientFrame:
|
||||
if gate is None:
|
||||
return _ForwardClientFrame(turn_access_groups=None)
|
||||
return _FORWARD_WITHOUT_TURN
|
||||
requests: Final = _client_frame_requests(frame_data)
|
||||
verdicts: Final = tuple([await gate(websocket, model) for model in requests.models])
|
||||
refusal: Final = next((verdict for verdict in verdicts if isinstance(verdict, WebsocketFrameModelRefused)), None)
|
||||
if refusal is not None:
|
||||
return _RefuseClientFrame(refusal=refusal)
|
||||
if requests.names_stored_prompt and isinstance(
|
||||
if requests.unverifiable_model_refusal is not None and isinstance(
|
||||
probe := await gate(websocket, STORED_PROMPT_PROBE_MODEL), WebsocketFrameModelRefused
|
||||
):
|
||||
return _RefuseClientFrame(refusal=_stored_prompt_refusal(probe))
|
||||
return _RefuseClientFrame(refusal=_unverifiable_model_probe_refusal(probe, requests.unverifiable_model_refusal))
|
||||
allowed: Final = tuple(verdict for verdict in verdicts if isinstance(verdict, WebsocketFrameModelAllowed))
|
||||
if requests.turn_model is None or not allowed:
|
||||
return _ForwardClientFrame(turn_access_groups=None)
|
||||
return _ForwardClientFrame(turn_access_groups=allowed[0].access_groups)
|
||||
return _ForwardClientFrame(
|
||||
starts_turn=requests.starts_turn,
|
||||
turn=(
|
||||
_TurnAuthorization(requested_model=requests.turn_model, access_groups=allowed[0].access_groups)
|
||||
if requests.turn_model is not None and allowed
|
||||
else None
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _refusal_close_reason(message: str) -> str:
|
||||
|
|
@ -2631,12 +2680,54 @@ async def _refuse_client_frame(
|
|||
await upstream_ws.close()
|
||||
|
||||
|
||||
def _stamp_billed_access_groups(kwargs: Mapping[str, object], access_groups: tuple[str, ...]) -> None:
|
||||
"""Charge a turn to the access groups that authorized its model, and to none when no group did."""
|
||||
def _stamp_billed_turn(kwargs: Mapping[str, object], turn: _TurnAuthorization) -> None:
|
||||
"""Charge a turn to the access groups that authorized its model, and its model budget to the model it asked for,
|
||||
since the provider may answer under a dated snapshot name the budget was never set on."""
|
||||
litellm_params: Final = as_str_mapping(kwargs.get("litellm_params"))
|
||||
metadata: Final = None if litellm_params is None else litellm_params.get("metadata")
|
||||
if isinstance(metadata, MutableMapping):
|
||||
metadata[MODEL_ACCESS_GROUP_METADATA_KEY] = list(access_groups)
|
||||
if not isinstance(metadata, MutableMapping):
|
||||
return
|
||||
metadata[MODEL_ACCESS_GROUP_METADATA_KEY] = list(turn.access_groups)
|
||||
if turn.requested_model is not None:
|
||||
metadata["model_group"] = turn.requested_model
|
||||
|
||||
|
||||
def _opening_model(target: str) -> str | None:
|
||||
return next(iter(parse_qs(urlparse(target).query).get("model", ())), None)
|
||||
|
||||
|
||||
class _WebsocketTurnLedger:
|
||||
"""The groups that authorized each turn not yet logged, when the next row starts, and how many rows were logged.
|
||||
|
||||
The provider runs a socket's response.create frames one at a time in the order they were sent, so each row is
|
||||
charged to the groups that allowed the oldest turn still unlogged. A realtime session keeps the model it opened on,
|
||||
so its rows stay charged to the groups that opened it.
|
||||
"""
|
||||
|
||||
def __init__(self, opening: _TurnAuthorization, started: datetime) -> None:
|
||||
self.opening = opening
|
||||
self.queued_turns: tuple[_TurnAuthorization | None, ...] = ()
|
||||
self.session_opened = False
|
||||
self.started = started
|
||||
self.logged_turns = 0
|
||||
|
||||
@property
|
||||
def current(self) -> _TurnAuthorization:
|
||||
oldest: Final = self.queued_turns[0] if self.queued_turns else None
|
||||
return self.opening if oldest is None else oldest
|
||||
|
||||
def queue_turn(self, turn: _TurnAuthorization | None) -> None:
|
||||
if not self.session_opened:
|
||||
self.queued_turns = (*self.queued_turns, turn)
|
||||
|
||||
def open_session(self) -> None:
|
||||
self.session_opened = True
|
||||
self.queued_turns = ()
|
||||
|
||||
def close_turn(self, end_time: datetime) -> None:
|
||||
self.logged_turns += 1
|
||||
self.started = end_time
|
||||
self.queued_turns = self.queued_turns[1:]
|
||||
|
||||
|
||||
async def _relay_client_frame(
|
||||
|
|
@ -2645,12 +2736,16 @@ async def _relay_client_frame(
|
|||
checked_frame: str | bytes,
|
||||
forwarded_frame: str | bytes,
|
||||
gate: WebsocketFrameModelGate | None,
|
||||
ledger: _WebsocketTurnLedger,
|
||||
) -> _ForwardClientFrame | _RefuseClientFrame:
|
||||
"""Forward one client frame upstream, or refuse it and end the session."""
|
||||
verdict: Final = await _client_frame_verdict(websocket, checked_frame, gate)
|
||||
match verdict:
|
||||
case _RefuseClientFrame(refusal=refusal):
|
||||
await _refuse_client_frame(websocket, upstream_ws, refusal)
|
||||
case _ForwardClientFrame(starts_turn=True, turn=turn):
|
||||
ledger.queue_turn(turn)
|
||||
await upstream_ws.send(forwarded_frame)
|
||||
case _ForwardClientFrame():
|
||||
await upstream_ws.send(forwarded_frame)
|
||||
return verdict
|
||||
|
|
@ -2659,23 +2754,6 @@ async def _relay_client_frame(
|
|||
_EXPLICIT_SUCCESS_KWARGS: Final = frozenset({"logging_obj", "passthrough_logging_payload", "custom_llm_provider"})
|
||||
|
||||
|
||||
class _WebsocketTurnLedger:
|
||||
"""The groups that authorized the turn in flight, when it started, and how many turns the socket has logged."""
|
||||
|
||||
def __init__(self, access_groups: tuple[str, ...], started: datetime) -> None:
|
||||
self.access_groups = access_groups
|
||||
self.started = started
|
||||
self.logged_turns = 0
|
||||
|
||||
def authorize(self, access_groups: tuple[str, ...] | None) -> None:
|
||||
if access_groups is not None:
|
||||
self.access_groups = access_groups
|
||||
|
||||
def close_turn(self, end_time: datetime) -> None:
|
||||
self.logged_turns += 1
|
||||
self.started = end_time
|
||||
|
||||
|
||||
async def websocket_passthrough_request(
|
||||
websocket: WebSocket,
|
||||
target: str,
|
||||
|
|
@ -2835,7 +2913,11 @@ async def websocket_passthrough_request(
|
|||
self.request = MockWebSocketRequest(target_url)
|
||||
|
||||
ledger: Final = _WebsocketTurnLedger(
|
||||
access_groups=tuple(user_api_key_dict.matched_model_access_groups or ()), started=start_time
|
||||
opening=_TurnAuthorization(
|
||||
requested_model=_opening_model(target),
|
||||
access_groups=tuple(user_api_key_dict.matched_model_access_groups or ()),
|
||||
),
|
||||
started=start_time,
|
||||
)
|
||||
|
||||
def log_frames(
|
||||
|
|
@ -2856,7 +2938,7 @@ async def websocket_passthrough_request(
|
|||
"end_time": end_time,
|
||||
}
|
||||
if client_frame_model_gate is not None:
|
||||
_stamp_billed_access_groups(call_kwargs, ledger.access_groups)
|
||||
_stamp_billed_turn(call_kwargs, ledger.current)
|
||||
success_kwargs: Final = {
|
||||
key: value for key, value in call_kwargs.items() if key not in _EXPLICIT_SUCCESS_KWARGS
|
||||
}
|
||||
|
|
@ -2988,26 +3070,25 @@ async def websocket_passthrough_request(
|
|||
# Not a JSON message or doesn't contain setup data
|
||||
|
||||
if isinstance(
|
||||
relayed_text := await _relay_client_frame(
|
||||
await _relay_client_frame(
|
||||
websocket,
|
||||
upstream_ws,
|
||||
text_data,
|
||||
_rewrite_vertex_live_setup_model(text_data, setup_model_rewriter),
|
||||
client_frame_model_gate,
|
||||
ledger,
|
||||
),
|
||||
_RefuseClientFrame,
|
||||
):
|
||||
break
|
||||
ledger.authorize(relayed_text.turn_access_groups)
|
||||
elif bytes_data is not None:
|
||||
if isinstance(
|
||||
relayed_bytes := await _relay_client_frame(
|
||||
websocket, upstream_ws, bytes_data, bytes_data, client_frame_model_gate
|
||||
await _relay_client_frame(
|
||||
websocket, upstream_ws, bytes_data, bytes_data, client_frame_model_gate, ledger
|
||||
),
|
||||
_RefuseClientFrame,
|
||||
):
|
||||
break
|
||||
ledger.authorize(relayed_bytes.turn_access_groups)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception:
|
||||
|
|
@ -3042,8 +3123,13 @@ async def websocket_passthrough_request(
|
|||
_extract_vertex_live_model_from_setup_response(message_data)
|
||||
return
|
||||
websocket_messages.append(message_data)
|
||||
if upstream_frame_kind is not None and upstream_frame_kind(message_data) == "ends_turn":
|
||||
log_turn()
|
||||
match None if upstream_frame_kind is None else upstream_frame_kind(message_data):
|
||||
case "opens_session":
|
||||
ledger.open_session()
|
||||
case "ends_turn":
|
||||
log_turn()
|
||||
case _:
|
||||
pass
|
||||
|
||||
async def relay_upstream_frame(upstream_message: str | bytes) -> None:
|
||||
record_upstream_frame(upstream_message)
|
||||
|
|
@ -3105,7 +3191,9 @@ async def websocket_passthrough_request(
|
|||
reason=_truncated_close_reason(upstream_close.reason),
|
||||
)
|
||||
|
||||
if ledger.logged_turns == 0:
|
||||
if upstream_frame_kind is None or any(
|
||||
upstream_frame_kind(frame) != "opens_session" for frame in websocket_messages
|
||||
):
|
||||
log_turn()
|
||||
|
||||
# Call the proxy logging success hook
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import json
|
||||
from typing import Any, cast
|
||||
|
||||
import pytest
|
||||
|
|
@ -40,6 +41,12 @@ def test_realtime_request_body_caches_each_model_separately():
|
|||
assert gpt4o_body_first is not gpt4o_mini_body
|
||||
|
||||
|
||||
def test_realtime_request_body_keeps_a_model_name_with_quotes_inside_its_own_field():
|
||||
body = realtime_request_body('gpt-4o", "user": "someone-else')
|
||||
|
||||
assert json.loads(body) == {"model": 'gpt-4o", "user": "someone-else'}
|
||||
|
||||
|
||||
def test_realtime_query_params_template_caches_each_pair_separately():
|
||||
params_with_intent_first = _realtime_query_params_template("gpt-4o", "intent-a")
|
||||
params_with_intent_second = _realtime_query_params_template("gpt-4o", "intent-a")
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
|
@ -22,11 +23,11 @@ def _logging_obj() -> MagicMock:
|
|||
return logging_obj
|
||||
|
||||
|
||||
def _session_created(model: str) -> dict:
|
||||
def _session_created(model: str) -> dict[str, object]:
|
||||
return {"type": "session.created", "session": {"id": "sess_fake", "model": model}}
|
||||
|
||||
|
||||
def _realtime_response_done(text_in: int, audio_in: int, text_out: int, audio_out: int) -> dict:
|
||||
def _realtime_response_done(text_in: int, audio_in: int, text_out: int, audio_out: int) -> dict[str, object]:
|
||||
return {
|
||||
"type": "response.done",
|
||||
"response": {
|
||||
|
|
@ -45,7 +46,7 @@ def _realtime_response_done(text_in: int, audio_in: int, text_out: int, audio_ou
|
|||
|
||||
def _responses_completed(
|
||||
model: str, input_tokens: int, output_tokens: int, event: str = "response.completed"
|
||||
) -> dict:
|
||||
) -> dict[str, object]:
|
||||
return {
|
||||
"type": event,
|
||||
"response": {
|
||||
|
|
@ -63,7 +64,12 @@ def _responses_completed(
|
|||
}
|
||||
|
||||
|
||||
def _handle(messages, kwargs=MappingProxyType({}), logging_obj=None, url_route="/openai/v1/realtime"):
|
||||
def _handle(
|
||||
messages: Sequence[Mapping[str, object]],
|
||||
kwargs: Mapping[str, object] = MappingProxyType({}),
|
||||
logging_obj: MagicMock | None = None,
|
||||
url_route: str = "/openai/v1/realtime",
|
||||
):
|
||||
return OpenAIWebsocketPassthroughLoggingHandler().openai_websocket_passthrough_handler(
|
||||
websocket_messages=messages,
|
||||
logging_obj=logging_obj if logging_obj is not None else _logging_obj(),
|
||||
|
|
@ -320,6 +326,37 @@ def test_a_failed_responses_turn_that_reported_no_usage_is_not_billed():
|
|||
assert handled["kwargs"] == {"model": RESPONSES_MODEL, "custom_llm_provider": "openai"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("turn_model", "requested_model", "priced_as"),
|
||||
[
|
||||
("gpt-unreleased-qa-snapshot", RESPONSES_MODEL, RESPONSES_MODEL),
|
||||
("gpt-5.4-mini-qa-preview", RESPONSES_MODEL, RESPONSES_MODEL),
|
||||
(RESPONSES_MODEL, REASONING_MODEL, RESPONSES_MODEL),
|
||||
],
|
||||
ids=["unmapped_model", "model_the_cost_map_prices_at_zero", "priced_turn_model"],
|
||||
)
|
||||
def test_a_responses_turn_falls_back_to_the_requested_model_only_when_its_own_has_no_price(
|
||||
turn_model: str, requested_model: str, priced_as: str
|
||||
):
|
||||
"""
|
||||
The provider can answer under a model name the cost map does not price
|
||||
yet. That turn still ran the model the key was allowed to ask for, so it is
|
||||
billed at that model's price instead of $0, while a turn whose own model
|
||||
has a price keeps it.
|
||||
"""
|
||||
handled = _handle(
|
||||
[_responses_completed(turn_model, input_tokens=13, output_tokens=5)],
|
||||
kwargs=MappingProxyType({"litellm_params": {"metadata": {"model_group": requested_model}}}),
|
||||
url_route="/openai_passthrough/v1/responses",
|
||||
)
|
||||
|
||||
prices = litellm.model_cost[priced_as]
|
||||
expected_cost = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"]
|
||||
assert expected_cost > 0
|
||||
assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost)
|
||||
assert handled["kwargs"]["model"] == turn_model
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("frame_type", "kind"),
|
||||
[
|
||||
|
|
|
|||
|
|
@ -47,6 +47,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
|
|||
InitPassThroughEndpointHelpers,
|
||||
SafeRouteAdder,
|
||||
WebsocketFrameModelAllowed,
|
||||
WebsocketFrameModelGate,
|
||||
WebsocketFrameModelRefused,
|
||||
WebsocketFrameModelVerdict,
|
||||
_registered_pass_through_routes,
|
||||
|
|
@ -7651,13 +7652,17 @@ class RecordingFrameModelGate:
|
|||
return WebsocketFrameModelAllowed(access_groups=self.groups.get(model, ()))
|
||||
|
||||
|
||||
def _received_frame(frame):
|
||||
def _received_frame(frame: str | bytes) -> Mapping[str, str | bytes]:
|
||||
if isinstance(frame, bytes):
|
||||
return {"type": "websocket.receive", "bytes": frame}
|
||||
return {"type": "websocket.receive", "text": frame}
|
||||
|
||||
|
||||
async def _relay_client_frames(frames, gate, user_api_key_dict=None):
|
||||
async def _relay_client_frames(
|
||||
frames: Sequence[str | bytes],
|
||||
gate: WebsocketFrameModelGate | None,
|
||||
user_api_key_dict: UserAPIKeyAuth | None = None,
|
||||
) -> SimpleNamespace:
|
||||
upstream_ws = RecordingUpstreamWebSocket()
|
||||
websocket = _client_websocket(
|
||||
AsyncMock(side_effect=[_received_frame(frame) for frame in frames] + [{"type": "websocket.disconnect"}])
|
||||
|
|
@ -7907,6 +7912,53 @@ async def test_websocket_passthrough_keeps_the_budget_reason_when_a_stored_promp
|
|||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_checks_the_model_a_hosted_tool_names():
|
||||
"""A hosted tool such as image generation runs its own model, so that model gets the same check as the turn's"""
|
||||
gate = RecordingFrameModelGate(MappingProxyType({"gpt-image-2": KEY_MODEL_REFUSAL}))
|
||||
frame = {
|
||||
"type": "response.create",
|
||||
"model": "gpt-5.4-mini",
|
||||
"tools": [{"type": "web_search"}, {"type": "image_generation", "model": "gpt-image-2"}],
|
||||
}
|
||||
|
||||
relayed = await _relay_client_frames([json.dumps(frame)], gate)
|
||||
|
||||
assert gate.checked == ["gpt-5.4-mini", "gpt-image-2"]
|
||||
assert relayed.upstream.send.await_args_list == []
|
||||
assert relayed.client.close.await_args.kwargs["code"] == 1008
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"frame",
|
||||
[
|
||||
{"type": "response.create", "model": "gpt-5.4-mini", "tools": [{"type": "image_generation"}]},
|
||||
{"type": "response.create", "response": {"tools": [{"type": "image_generation", "model": None}]}},
|
||||
],
|
||||
)
|
||||
async def test_websocket_passthrough_refuses_an_image_tool_without_a_model_for_a_key_limited_to_some_models(frame):
|
||||
gate = RecordingFrameModelGate(MappingProxyType({STORED_PROMPT_PROBE_MODEL: KEY_MODEL_REFUSAL}))
|
||||
|
||||
relayed = await _relay_client_frames([json.dumps(frame)], gate)
|
||||
|
||||
assert gate.checked[-1] == STORED_PROMPT_PROBE_MODEL
|
||||
assert relayed.upstream.send.await_args_list == []
|
||||
assert "image_generation tool has to name its model" in relayed.client.close.await_args.kwargs["reason"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_checks_a_frame_only_the_standard_json_parser_reads():
|
||||
"""A lone surrogate escape is valid JSON syntax the fast parser refuses, so the frame is still checked"""
|
||||
gate = RecordingFrameModelGate(MappingProxyType({"gpt-5.4": KEY_MODEL_REFUSAL}))
|
||||
frame = '{"type": "response.create", "model": "gpt-5.4", "metadata": {"note": "\\ud800"}}'
|
||||
|
||||
relayed = await _relay_client_frames([frame], gate)
|
||||
|
||||
assert gate.checked == ["gpt-5.4"]
|
||||
assert relayed.upstream.send.await_args_list == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_forwards_binary_audio_untouched():
|
||||
gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": "denied"}))
|
||||
|
|
@ -8015,14 +8067,16 @@ class ScriptedUpstreamWebSocket:
|
|||
return frame
|
||||
|
||||
|
||||
def _turn_taking_client(frames: Sequence[str], answered_by: str) -> MagicMock:
|
||||
"""A client that sends its next frame only once the answer to the last one reached it, and that Starlette
|
||||
refuses to send to once it is closed"""
|
||||
def _turn_taking_client(
|
||||
frames: Sequence[str], answered_by: str, answers_due: Sequence[int] | None = None
|
||||
) -> MagicMock:
|
||||
"""A client that sends each frame, and then leaves, once the answers due by then reached it (by default the
|
||||
answer to every earlier frame), and that Starlette refuses to send to once it is closed"""
|
||||
from starlette.websockets import WebSocketState
|
||||
|
||||
delivered: list[str] = []
|
||||
script = iter(frames)
|
||||
receives = iter(range(len(frames) + 1))
|
||||
receives = iter(range(len(frames) + 1) if answers_due is None else answers_due)
|
||||
|
||||
async def receive():
|
||||
answers_due = next(receives)
|
||||
|
|
@ -8042,7 +8096,14 @@ def _turn_taking_client(frames: Sequence[str], answered_by: str) -> MagicMock:
|
|||
return websocket
|
||||
|
||||
|
||||
async def _relay_turns(client, upstream_ws, gate, endpoint="/openai_passthrough/v1/responses"):
|
||||
async def _relay_turns(
|
||||
client: MagicMock,
|
||||
upstream_ws: ScriptedUpstreamWebSocket,
|
||||
gate: WebsocketFrameModelGate,
|
||||
endpoint: str = "/openai_passthrough/v1/responses",
|
||||
target: str = "wss://api.openai.com/v1/responses",
|
||||
opening_access_groups: Sequence[str] = (),
|
||||
) -> SimpleNamespace:
|
||||
with (
|
||||
_patched_websocket_passthrough_environment(upstream_ws) as proxy_logging,
|
||||
patch(SUCCESS_HANDLER, new_callable=AsyncMock) as success_handler,
|
||||
|
|
@ -8050,9 +8111,11 @@ async def _relay_turns(client, upstream_ws, gate, endpoint="/openai_passthrough/
|
|||
await asyncio.wait_for(
|
||||
websocket_passthrough_request(
|
||||
websocket=client,
|
||||
target="wss://api.openai.com/v1/responses",
|
||||
target=target,
|
||||
custom_headers={"Authorization": "Bearer sk-provider"},
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
api_key="hashed-key", matched_model_access_groups=list(opening_access_groups)
|
||||
),
|
||||
forward_headers=False,
|
||||
endpoint=endpoint,
|
||||
accept_websocket=False,
|
||||
|
|
@ -8065,15 +8128,23 @@ async def _relay_turns(client, upstream_ws, gate, endpoint="/openai_passthrough/
|
|||
|
||||
|
||||
def _responses_turn(model: str) -> list[str]:
|
||||
return [
|
||||
json.dumps({"type": "response.created", "response": {"model": model}}),
|
||||
return [_responses_turn_started(model), _responses_turn_ended(model)]
|
||||
|
||||
|
||||
def _responses_turn_started(model: str) -> str:
|
||||
return json.dumps({"type": "response.created", "response": {"model": model}})
|
||||
|
||||
|
||||
def _responses_turn_ended(model: str) -> str:
|
||||
return (
|
||||
|
||||
json.dumps(
|
||||
{
|
||||
"type": "response.completed",
|
||||
"response": {"model": model, "usage": {"input_tokens": 9, "output_tokens": 3, "total_tokens": 12}},
|
||||
}
|
||||
),
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -8113,6 +8184,119 @@ async def test_websocket_passthrough_logs_each_turn_on_its_own_row_charged_to_th
|
|||
relayed.proxy_logging.post_call_failure_hook.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_charges_pipelined_turns_to_the_frame_that_asked_for_each():
|
||||
"""
|
||||
A client can send its next response.create before the provider finished the
|
||||
last one, and the provider then runs them one after the other. Each turn
|
||||
still has to reach spend under the groups that let its own frame through,
|
||||
and under the model that frame asked for, even when the provider answers
|
||||
under a dated snapshot name a per-model budget was never set on.
|
||||
"""
|
||||
gate = RecordingFrameModelGate(
|
||||
groups=MappingProxyType({"gpt-5.4-mini": ("tier-a",), "gpt-5.4": ("tier-b",)}),
|
||||
)
|
||||
client = _turn_taking_client(
|
||||
[
|
||||
json.dumps({"type": "response.create", "model": "gpt-5.4-mini"}),
|
||||
json.dumps({"type": "response.create", "model": "gpt-5.4"}),
|
||||
],
|
||||
answered_by='"response.completed"',
|
||||
answers_due=(0, 0, 2),
|
||||
)
|
||||
upstream_ws = ScriptedUpstreamWebSocket(
|
||||
replies=[_responses_turn("gpt-5.4-mini-2099-01-01"), _responses_turn("gpt-5.4-2099-01-01")]
|
||||
)
|
||||
|
||||
relayed = await _relay_turns(client, upstream_ws, gate)
|
||||
|
||||
rows = [logged.kwargs for logged in relayed.success_handler.call_args_list]
|
||||
assert [[frame["response"]["model"] for frame in row["response_body"]] for row in rows] == [
|
||||
["gpt-5.4-mini-2099-01-01", "gpt-5.4-mini-2099-01-01"],
|
||||
["gpt-5.4-2099-01-01", "gpt-5.4-2099-01-01"],
|
||||
]
|
||||
assert [row["litellm_params"]["metadata"][MODEL_ACCESS_GROUP_METADATA_KEY] for row in rows] == [
|
||||
["tier-a"],
|
||||
["tier-b"],
|
||||
]
|
||||
assert [row["litellm_params"]["metadata"]["model_group"] for row in rows] == ["gpt-5.4-mini", "gpt-5.4"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_charges_realtime_turns_to_the_session_the_key_opened():
|
||||
"""
|
||||
A realtime session runs the model it was opened on, so its turns stay
|
||||
charged to the groups and the model the key opened it with, whatever model
|
||||
a later session.update or response.create names.
|
||||
"""
|
||||
gate = RecordingFrameModelGate(groups=MappingProxyType({"gpt-realtime-2.1-mini": ("tier-b",)}))
|
||||
client = _turn_taking_client(
|
||||
[
|
||||
json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1-mini"}}),
|
||||
json.dumps({"type": "response.create", "response": {"model": "gpt-realtime-2.1-mini"}}),
|
||||
],
|
||||
answered_by='"type"',
|
||||
answers_due=(1, 2, 4),
|
||||
)
|
||||
upstream_ws = ScriptedUpstreamWebSocket(
|
||||
REALTIME_SESSION_CREATED,
|
||||
replies=[
|
||||
[json.dumps({"type": "session.updated"})],
|
||||
[json.dumps({"type": "response.created", "response": {"status": "in_progress"}}), REALTIME_RESPONSE_DONE],
|
||||
],
|
||||
)
|
||||
|
||||
relayed = await _relay_turns(
|
||||
client,
|
||||
upstream_ws,
|
||||
gate,
|
||||
"/openai/v1/realtime",
|
||||
target="wss://api.openai.com/v1/realtime?model=gpt-realtime-2.1",
|
||||
opening_access_groups=("tier-a",),
|
||||
)
|
||||
|
||||
assert gate.checked == ["gpt-realtime-2.1-mini", "gpt-realtime-2.1-mini"]
|
||||
rows = [logged.kwargs for logged in relayed.success_handler.call_args_list]
|
||||
assert [row["litellm_params"]["metadata"][MODEL_ACCESS_GROUP_METADATA_KEY] for row in rows] == [["tier-a"]]
|
||||
assert [row["litellm_params"]["metadata"]["model_group"] for row in rows] == ["gpt-realtime-2.1"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_logs_no_spend_row_for_a_session_refused_before_any_answer():
|
||||
"""A session refused at its first frame ran nothing, so its only record is the refusal, not an extra $0 row"""
|
||||
gate = RecordingFrameModelGate(MappingProxyType({"gpt-5.4": KEY_MODEL_REFUSAL}))
|
||||
client = _turn_taking_client([json.dumps({"type": "response.create", "model": "gpt-5.4"})], answered_by='"error"')
|
||||
|
||||
relayed = await _relay_turns(client, ScriptedUpstreamWebSocket(), gate)
|
||||
|
||||
assert relayed.success_handler.call_args_list == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_logs_the_unfinished_turn_a_client_walked_out_of():
|
||||
"""
|
||||
A client can leave while a later turn is still running. That turn never
|
||||
ends, but it was started and may be billed, so it still gets its own row.
|
||||
"""
|
||||
client = _turn_taking_client(
|
||||
[
|
||||
json.dumps({"type": "response.create", "model": "gpt-5.4-mini"}),
|
||||
json.dumps({"type": "response.create", "model": "gpt-5.4"}),
|
||||
],
|
||||
answered_by='"response.created"',
|
||||
)
|
||||
upstream_ws = ScriptedUpstreamWebSocket(
|
||||
replies=[_responses_turn("gpt-5.4-mini"), [_responses_turn_started("gpt-5.4")]]
|
||||
)
|
||||
|
||||
relayed = await _relay_turns(client, upstream_ws, RecordingFrameModelGate())
|
||||
|
||||
assert [
|
||||
[json.loads(frame)["type"] if isinstance(frame, str) else frame["type"] for frame in logged.kwargs["response_body"]]
|
||||
for logged in relayed.success_handler.call_args_list
|
||||
] == [["response.created", "response.completed"], ["response.created"]]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_carries_the_session_the_provider_opened_into_every_later_turn():
|
||||
"""
|
||||
|
|
@ -9095,7 +9279,7 @@ def test_passthrough_client_cannot_forge_session_id_omission(client_metadata_key
|
|||
)
|
||||
|
||||
|
||||
async def _billed_metadata(gate) -> dict:
|
||||
async def _billed_metadata(gate: WebsocketFrameModelGate) -> Mapping[str, object]:
|
||||
served = await _relay_client_frames(
|
||||
[json.dumps({"type": "response.create", "model": "gpt-5.4-mini"})],
|
||||
gate,
|
||||
|
|
|
|||
|
|
@ -667,7 +667,7 @@ REALTIME_FRAMES: Final = [
|
|||
]
|
||||
|
||||
|
||||
def _normalize(url_route: str, response_body):
|
||||
def _normalize(url_route: str, response_body: dict[str, object] | list[dict[str, object]]) -> dict:
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.model_call_details = {}
|
||||
return PassThroughEndpointLogging().normalize_llm_passthrough_logging_payload(
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ import litellm
|
|||
from litellm.llms.openai.workload_identity import _workload_identity_auth
|
||||
from litellm.proxy._types import Litellm_EntityType, ProxyErrorTypes, ProxyException, UserAPIKeyAuth
|
||||
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
|
||||
_OPENAI_WS_ALIASED_MODEL_REFUSAL,
|
||||
_OPENAI_WS_DISABLED_REFUSAL,
|
||||
_OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL,
|
||||
_openai_websocket_refusal,
|
||||
|
|
@ -173,6 +174,18 @@ async def test_openai_websocket_forwards_query_and_keeps_provider_auth(prefix, m
|
|||
assert websocket.closed is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_websocket_reports_its_route_without_the_proxy_root_path(monkeypatch):
|
||||
"""Route allowlists and spend rows key on the route the proxy serves, which a root path prefix is not part of"""
|
||||
monkeypatch.setenv("SERVER_ROOT_PATH", "/litellm")
|
||||
websocket = _FakeWebSocket("/litellm/openai/v1/responses", "")
|
||||
|
||||
with patch(GET_CREDENTIALS, return_value="sk-provider"):
|
||||
served = await _serve(websocket, "v1/responses", UserAPIKeyAuth(), ENABLED)
|
||||
|
||||
assert [call.endpoint for call in served.relay.calls] == ["/openai/v1/responses"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_websocket_accepts_first_client_subprotocol():
|
||||
websocket = _FakeWebSocket(
|
||||
|
|
@ -390,6 +403,48 @@ async def test_proxy_frame_model_gate_refuses_a_model_the_agent_was_not_granted(
|
|||
assert refusal_names in refusal.message
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"token",
|
||||
[
|
||||
UserAPIKeyAuth(models=["gpt-5.4-mini"], aliases={"gpt-5.4": "gpt-5.4-mini"}),
|
||||
UserAPIKeyAuth(models=["gpt-5.4-mini"], team_model_aliases={"gpt-5.4": "gpt-5.4-mini"}),
|
||||
],
|
||||
ids=["key_alias", "team_alias"],
|
||||
)
|
||||
async def test_proxy_frame_model_gate_refuses_an_alias_that_would_reach_openai_unchanged(token):
|
||||
"""
|
||||
The passthrough sends each frame to OpenAI as written, so OpenAI runs the
|
||||
alias name itself. An alias pointing at an allowed model must not let a key
|
||||
limited to that model run the model the alias is named after.
|
||||
"""
|
||||
refusal: Final = _refusal(await _gate_verdict(_FrameAuthorizer(token), "gpt-5.4"))
|
||||
|
||||
assert refusal == WebsocketFrameModelRefused(
|
||||
message=_OPENAI_WS_ALIASED_MODEL_REFUSAL, error_type=ProxyErrorTypes.key_model_access_denied.value
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_frame_model_gate_refuses_a_proxy_wide_alias_for_a_key_limited_to_its_target(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "model_alias_map", {"gpt-5.4": "gpt-5.4-mini"})
|
||||
|
||||
refusal: Final = _refusal(
|
||||
await _gate_verdict(_FrameAuthorizer(UserAPIKeyAuth(models=["gpt-5.4-mini"])), "gpt-5.4")
|
||||
)
|
||||
|
||||
assert refusal.message == _OPENAI_WS_ALIASED_MODEL_REFUSAL
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_frame_model_gate_lets_a_key_that_may_use_any_model_send_an_alias():
|
||||
token: Final = UserAPIKeyAuth(models=[], aliases={"fast": "gpt-5.4-mini"})
|
||||
|
||||
verdict: Final = await _gate_verdict(_FrameAuthorizer(token), "fast")
|
||||
|
||||
assert verdict == WebsocketFrameModelAllowed(access_groups=())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("unrestricted_models", [[], ["all-proxy-models"], ["*"]])
|
||||
async def test_proxy_frame_model_gate_allows_every_model_for_unrestricted_keys(unrestricted_models):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue