From 68e78bc53a1d8668f996f20976f4f2ae972b5b87 Mon Sep 17 00:00:00 2001 From: norkelbyte <348963847+norkelbyte@users.noreply.github.com> Date: Mon, 9 Nov 2026 20:13:51 -0800 Subject: [PATCH] fix(router): count a Responses WebSocket turn's tokens before the next turn is admitted (#50730) * fix(router): count a Responses WebSocket turn's tokens before the next turn is admitted * fix(responses): forward and log a finished websocket turn when its early count fails or is cancelled A usage payload the parser rejects no longer stops the relay before the terminal event reaches the client, and a client that leaves while the turn is being counted still gets the turn logged. A cancelled early count now clears its counted-tokens marker so the logging callback counts the turn in full. * fix(responses): finish a turn's early token count when the client leaves A client that left while a finished turn was being counted cancelled the count halfway. The router's count could be written with no marker, so the logging callback counted the turn a second time. The count now runs to the end before the turn is logged, and the failure path goes back to clearing only on errors * fix(router): count one router request per Responses WebSocket turn A socket whose finished turns each counted their own request in router usage no longer adds one more request when it closes, so N turns count N requests, as N HTTP calls do. A socket with no counted turn still counts 1 * fix(proxy): keep router usage markers out of pass-through client metadata and ClickHouse rows --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/constants.py | 1 + .../clickhouse/clickhouse_spend_logger.py | 10 +- litellm/llms/custom_httpx/llm_http_handler.py | 8 +- litellm/proxy/litellm_pre_call_utils.py | 6 +- .../pass_through_endpoints.py | 7 +- litellm/responses/main.py | 8 +- litellm/responses/streaming_iterator.py | 42 ++- litellm/router_rust/core.py | 6 +- litellm/router_rust/invoke/endpoints.py | 11 +- litellm/router_rust/invoke/websocket.py | 116 ++++++- litellm/router_rust/surface/legacy_members.py | 55 ++-- litellm/router_strategy/lowest_tpm_rpm_v2.py | 60 ++-- .../track_deployment_metrics.py | 26 +- .../test_clickhouse_spend_logger.py | 10 +- .../test_pass_through_endpoints.py | 29 ++ .../unit/proxy/test_litellm_pre_call_utils.py | 3 + .../test_responses_websocket_all_providers.py | 133 ++++++++ .../test_router/test_router_helper_utils.py | 290 +++++++++++++++++- 18 files changed, 744 insertions(+), 77 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 6cca43f8c8b..8df92c46630 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1706,6 +1706,7 @@ ROUTING_REQUEST_TAGS_METADATA_KEY: Final = "_routing_request_tags" BAGGAGE_MAPPED_TAGS_METADATA_KEY: Final = "_baggage_mapped_tags" VECTOR_STORE_CALLER_METADATA_KWARG: Final = "litellm_vector_store_caller_metadata" ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: Final = "_litellm_router_usage_counted_tokens" +USAGE_BASED_ROUTING_COUNTED_TOKENS_METADATA_KEY: Final = "_litellm_usage_based_routing_counted_tokens" INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin" RESPONSES_WEBSOCKET_CALL_ID_METADATA_KEY: Final = "responses_websocket_call_id" BILLED_RETRY_ATTEMPT_METADATA_KEY: Final = "billed_retry_attempt" diff --git a/litellm/integrations/clickhouse/clickhouse_spend_logger.py b/litellm/integrations/clickhouse/clickhouse_spend_logger.py index cd1250b17f1..ef7442143ca 100644 --- a/litellm/integrations/clickhouse/clickhouse_spend_logger.py +++ b/litellm/integrations/clickhouse/clickhouse_spend_logger.py @@ -16,6 +16,7 @@ from pydantic import JsonValue, TypeAdapter, ValidationError import litellm from litellm._logging import verbose_logger +from litellm.constants import ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, USAGE_BASED_ROUTING_COUNTED_TOKENS_METADATA_KEY from litellm.integrations.clickhouse.clickhouse_batch_logger import ClickHouseBatchLogger from litellm.integrations.clickhouse.context import is_lens_analysis from litellm.integrations.clickhouse.schema import SPEND_LOGS_TABLE @@ -37,7 +38,14 @@ TRACE_INGEST_ROUTE: Final = "/v1/traces" _METADATA_MAPPING: Final = TypeAdapter(Mapping[str, object]) _METADATA_VALUE: Final = TypeAdapter(JsonValue) _INTERNAL_METADATA_KEYS: Final = frozenset( - ("user_api_key", "user_api_key_auth", "user_api_key_budget_reservation", "proxy_server_request") + ( + "user_api_key", + "user_api_key_auth", + "user_api_key_budget_reservation", + "proxy_server_request", + ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, + USAGE_BASED_ROUTING_COUNTED_TOKENS_METADATA_KEY, + ) ) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 8c9de272c0d..72c086b74ed 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -268,7 +268,11 @@ if TYPE_CHECKING: FakeAnthropicMessagesStreamIterator, ) from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig - from litellm.responses.streaming_iterator import ResponsesWsFrameGuard, ResponsesWsTurnAdmission + from litellm.responses.streaming_iterator import ( + ResponsesWsFrameGuard, + ResponsesWsTurnAdmission, + ResponsesWsTurnUsageRecorder, + ) from litellm.types.google_genai.main import GenerateContentResponse from litellm.types.llms.openai_evals import ( CancelEvalResponse, @@ -7241,6 +7245,7 @@ class BaseLLMHTTPHandler: first_message: str | None = None, request_defaults: ResponsesWebSocketRequestDefaults | None = None, turn_admission: "ResponsesWsTurnAdmission | None" = None, + turn_usage_recorder: "ResponsesWsTurnUsageRecorder | None" = None, frame_guard: "ResponsesWsFrameGuard | None" = None, **kwargs: object, ) -> Exception | None: @@ -7407,6 +7412,7 @@ class BaseLLMHTTPHandler: custom_llm_provider=custom_llm_provider, request_defaults=request_defaults, turn_admission=_ws_turn_admission, + turn_usage_recorder=turn_usage_recorder, frame_guard=frame_guard, ) return await streaming.bidirectional_forward() diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 0bb4a803609..22399ced718 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -47,6 +47,7 @@ from litellm.constants import ( SESSION_ID_GENERATED_METADATA_KEY, SESSION_ID_OMITTED_METADATA_KEY, TYPESAFE_DANGER_ZONE_HEADER_PREFIX, + USAGE_BASED_ROUTING_COUNTED_TOKENS_METADATA_KEY, VECTOR_STORE_CALLER_METADATA_KWARG, X_LITELLM_DISABLE_CALLBACKS, ) @@ -531,7 +532,7 @@ _CLIENT_PRICING_METADATA_FIELDS: Final = frozenset({"model_info"}) # ``attempted_fallbacks`` and ``original_model_group`` are written by the router # and read by spend logs as fact; a client value has no legitimate meaning and no # key or team setting keeps it, so the strip is never gated. -_ROUTER_RESERVED_METADATA_FIELDS: Final = frozenset( +ROUTER_RESERVED_METADATA_FIELDS: Final = frozenset( { "attempted_fallbacks", "original_model_group", @@ -539,6 +540,7 @@ _ROUTER_RESERVED_METADATA_FIELDS: Final = frozenset( "request_retry_count", CLIENT_OUTPUT_CEILING_METADATA_KEY, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, + USAGE_BASED_ROUTING_COUNTED_TOKENS_METADATA_KEY, } ) _ALLOW_CLIENT_PRICING_OVERRIDE_METADATA_KEY: Final = "allow_client_pricing_override" @@ -883,7 +885,7 @@ def _strip_router_reserved_metadata( for metadata_key in ("metadata", "litellm_metadata"): if not isinstance(metadata := data.get(metadata_key), dict): continue - for field in _ROUTER_RESERVED_METADATA_FIELDS & metadata.keys(): + for field in ROUTER_RESERVED_METADATA_FIELDS & metadata.keys(): metadata.pop(field) verbose_proxy_logger.debug( "Stripped router-reserved metadata field from request body: %s.%s", metadata_key, field diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 4b84fe19559..66634688b9a 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -134,6 +134,7 @@ from litellm.proxy.common_utils.sse_keepalive import ( wrap_passthrough_sse_bytes_with_keepalive_pings, ) from litellm.proxy.litellm_pre_call_utils import ( # noqa: F401 # legacy module exports + ROUTER_RESERVED_METADATA_FIELDS, LiteLLMProxyRequestSetup, _get_dynamic_logging_metadata, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export _key_or_team_allows_client_pricing_override, # pyright: ignore[reportPrivateUsage] # reuse the proxy's pricing trust policy @@ -196,7 +197,11 @@ _PRICING_MODEL_OVERRIDE_PARAMS: Final = frozenset(("base_model",)) def _is_mergeable_client_metadata(key: str, value: object) -> bool: - return not key.startswith("user_api_key_") and (key != "model_info" or isinstance(value, Mapping)) + return ( + not key.startswith("user_api_key_") + and key not in ROUTER_RESERVED_METADATA_FIELDS + and (key != "model_info" or isinstance(value, Mapping)) + ) def _request_body_fields(parsed: object) -> Mapping[str, object]: diff --git a/litellm/responses/main.py b/litellm/responses/main.py index e085c4d7a2e..7eab3985f59 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -82,7 +82,11 @@ from litellm.types.llms.openai import ( # Handle ResponseText import with fallback if TYPE_CHECKING: - from litellm.responses.streaming_iterator import ResponsesWsFrameGuard, ResponsesWsTurnAdmission + from litellm.responses.streaming_iterator import ( + ResponsesWsFrameGuard, + ResponsesWsTurnAdmission, + ResponsesWsTurnUsageRecorder, + ) from litellm.types.llms.openai import ResponseText else: ResponseText = str # Fallback for ResponseText import @@ -2870,6 +2874,7 @@ async def _aresponses_websocket( api_key: str | None = None, timeout: float | None = None, turn_admission: "ResponsesWsTurnAdmission | None" = None, + turn_usage_recorder: "ResponsesWsTurnUsageRecorder | None" = None, responses_ws_frame_guard: "ResponsesWsFrameGuard | None" = None, **kwargs, ) -> Exception | None: @@ -2964,6 +2969,7 @@ async def _aresponses_websocket( custom_llm_provider=_custom_llm_provider, request_defaults=_build_responses_websocket_request_defaults(deployment_kwargs), turn_admission=turn_admission, + turn_usage_recorder=turn_usage_recorder, frame_guard=responses_ws_frame_guard, **remaining_kwargs, ) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 08b2c518632..28a82b500c9 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -222,6 +222,15 @@ def _running_loop() -> asyncio.AbstractEventLoop | None: return None +async def _finish_even_if_cancelled(work: Awaitable[None]) -> None: + task: Final = asyncio.ensure_future(work) + try: + await asyncio.shield(task) + except asyncio.CancelledError: + _ = await asyncio.wait((task,)) + raise + + def spawn_logging_task( running_loop: asyncio.AbstractEventLoop, coroutine: Coroutine[object, object, object], *, task_name: str ) -> asyncio.Task[object]: @@ -2344,6 +2353,7 @@ class GuardedResponsesWsFrame: ResponsesWsFrameGuard: TypeAlias = Callable[[str], Awaitable[GuardedResponsesWsFrame]] +ResponsesWsTurnUsageRecorder: TypeAlias = Callable[[Mapping[str, object], int], Awaitable[None]] class ResponsesWsFrameRejection(ValueError): @@ -2577,6 +2587,7 @@ class ResponsesWebSocketStreaming: custom_llm_provider: str | None = None, request_defaults: ResponsesWebSocketRequestDefaults | None = None, turn_admission: ResponsesWsTurnAdmission | None = None, + turn_usage_recorder: ResponsesWsTurnUsageRecorder | None = None, frame_guard: ResponsesWsFrameGuard | None = None, ): self.websocket = websocket @@ -2599,6 +2610,7 @@ class ResponsesWebSocketStreaming: self.authorized_model: str | None = authorized_model self.request_defaults: ResponsesWebSocketRequestDefaults | None = request_defaults self.turn_admission: ResponsesWsTurnAdmission | None = turn_admission + self.turn_usage_recorder: ResponsesWsTurnUsageRecorder | None = turn_usage_recorder def _start_turn( self, @@ -2681,18 +2693,36 @@ class ResponsesWebSocketStreaming: "pii_tokens": {**(known if _is_str_mapping(known) else {}), **pii_tokens}, } - def _record_backend_event(self, event: str) -> None: + async def _record_backend_event(self, event: str) -> None: event_obj: Final = _json_object_or_none(event) if event_obj is None or event_obj.get("type") not in RESPONSES_WS_LOGGED_EVENT_TYPES: return if not self._open_turns: self._start_turn(None) turn: Final = self._open_turns[0].with_event(event_obj) - if event_obj.get("type") in _RESPONSES_WS_TERMINAL_EVENT_TYPES: - self._open_turns = self._open_turns[1:] - self._log_turn(turn) - return self._open_turns = (turn, *self._open_turns[1:]) + if event_obj.get("type") not in _RESPONSES_WS_TERMINAL_EVENT_TYPES: + return + if event_obj.get("type") in _HISTORY_TERMINAL_EVENT_TYPES: + await self._record_turn_usage(turn) + self._open_turns = self._open_turns[1:] + self._log_turn(turn) + + async def _record_turn_usage(self, turn: _ResponsesWsTurn) -> None: + from litellm.cost_calculator import ResponsesWebSocketTokenUsageProcessor + + turn_litellm_params: Final = _json_object_or_none( + turn.logging_obj.litellm_params # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # Logging.litellm_params is untyped + ) + if self.turn_usage_recorder is None or turn_litellm_params is None: + return + try: + usage: Final = ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results( + turn.events + ) + await _finish_even_if_cancelled(self.turn_usage_recorder(turn_litellm_params, usage.total_tokens)) + except Exception as e: # noqa: BLE001 # a turn not counted early is still counted by its logging + verbose_logger.debug("Responses WS turn usage not counted before its terminal event: %s", e) def _log_open_turns(self) -> None: open_turns: Final = self._open_turns @@ -2778,7 +2808,7 @@ class ResponsesWebSocketStreaming: # Log the output-masked form so PII redacted by apply_to_output # guardrails does not appear in success logs. - self._record_backend_event(wrapped_str) + await self._record_backend_event(wrapped_str) await self.websocket.send_text(wrapped_str) diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index 55155a96fa3..69f8f70f793 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -124,6 +124,7 @@ from litellm.router_rust.invoke.streaming import ( shape_attempt_kwargs, ) from litellm.router_rust.invoke.walk import FallbackOutcomes, WalkMetadata, await_all, log_retry_into +from litellm.router_rust.invoke.websocket import turns_counted_their_requests from litellm.router_rust.routing_stores import RoutingStores from litellm.router_rust.steps.access import caller_fields, selection_refusal from litellm.router_rust.steps.callbacks import OptionalChecks, RustOptionalChecks @@ -5078,9 +5079,12 @@ class _Driver: def _tracks_usage(self, call: "AttemptCall", response: object, deferred: bool) -> bool: """Whether the attempt counts against its deployment's usage (``_deployment_tracks_usage``). A ``None`` response counts nothing (``increment_deployment_usage_for_response``), and an attempt - that reports no success then is a realtime session, which logs none to count it either.""" + that reports no success then is a realtime session, which logs none to count it either. A + Responses socket whose finished turns each counted their own request adds none at close.""" if deferred and response is None: return False + if turns_counted_their_requests(call.kwargs): + return False model: Final = call.kwargs.get("model") return self._core.deployment_tracks_usage( (as_str_mapping(call.kwargs.get("model_info")) or {}).get("id"), diff --git a/litellm/router_rust/invoke/endpoints.py b/litellm/router_rust/invoke/endpoints.py index 70c83bed728..0bd522ba06d 100644 --- a/litellm/router_rust/invoke/endpoints.py +++ b/litellm/router_rust/invoke/endpoints.py @@ -970,7 +970,7 @@ def _generic_attempt( ) -> AttemptCall: """`_ageneric_api_call_with_fallbacks_helper` for one deployment.""" from litellm.llms.base_llm.passthrough.transformation import replace_path_segment - from litellm.router_rust.invoke.websocket import turn_admission + from litellm.router_rust.invoke.websocket import turn_admission, turn_usage_recorder handler: Final = invoked_handler_name(core, plan) socket: Final = handler == "_aresponses_websocket" @@ -1005,7 +1005,14 @@ def _generic_attempt( **kwargs, "model": model_name, **_session_with_model(kwargs.get("session"), model_name), - **({"turn_admission": turn_admission(core, requested, deployment, kwargs)} if socket else {}), + **( + { + "turn_admission": turn_admission(core, requested, deployment, kwargs), + "turn_usage_recorder": turn_usage_recorder(core, requested, kwargs), + } + if socket + else {} + ), } if provider is not None: final["custom_llm_provider"] = provider diff --git a/litellm/router_rust/invoke/websocket.py b/litellm/router_rust/invoke/websocket.py index e7baeb92345..b38252503c7 100644 --- a/litellm/router_rust/invoke/websocket.py +++ b/litellm/router_rust/invoke/websocket.py @@ -8,15 +8,24 @@ Which strategy and selector apply is the Rust core's routing context (``routing_ selectors, limiters and checks are the live Python objects the host keeps. """ -from collections.abc import Callable, Mapping +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass, field from functools import partial from types import MappingProxyType from typing import TYPE_CHECKING, Final -from litellm.litellm_core_utils.core_helpers import without +from litellm._logging import verbose_router_logger +from litellm.constants import USAGE_BASED_ROUTING_COUNTED_TOKENS_METADATA_KEY +from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs, without from litellm.router_rust.invoke.shaping import as_dict +from litellm.router_utils.router_callbacks.track_deployment_metrics import ( + count_before_callbacks, + find_deployment_metadata, +) if TYPE_CHECKING: + from opentelemetry.trace import Span + from litellm.router_rust.invoke.endpoints import EndpointCore @@ -27,6 +36,100 @@ def turn_admission( return partial(admit_deployment_turn, core, model_group, deployment, kwargs) +@dataclass(slots=True) +class TurnUsageRecorder: + core: "EndpointCore" + model_group: str + request_kwargs: Mapping[str, object] + counted_a_turn: bool = field(default=False, init=False) + + async def __call__(self, turn_litellm_params: Mapping[str, object], total_tokens: int) -> None: + if await record_deployment_turn_usage( + self.core, self.model_group, self.request_kwargs, turn_litellm_params, total_tokens + ): + self.counted_a_turn = True + + +def turn_usage_recorder(core: "EndpointCore", model_group: str, kwargs: Mapping[str, object]) -> TurnUsageRecorder: + return TurnUsageRecorder(core, model_group, kwargs) + + +def turns_counted_their_requests(kwargs: Mapping[str, object]) -> bool: + recorder: Final = kwargs.get("turn_usage_recorder") + return isinstance(recorder, TurnUsageRecorder) and recorder.counted_a_turn + + +async def record_deployment_turn_usage( + core: "EndpointCore", + model_group: str, + request_kwargs: Mapping[str, object], + turn_litellm_params: Mapping[str, object], + total_tokens: int, +) -> bool: + """Count a finished turn early. Returns whether the router's own usage counted its request.""" + deployment_metadata: Final = find_deployment_metadata(turn_litellm_params) + if deployment_metadata is None: + return False + parent_otel_span: Final = get_parent_otel_span_from_kwargs(dict(request_kwargs)) + counted_request: Final = await _counted_early( + core.router.increment_deployment_usage_before_callbacks( + deployment_metadata, model_group, total_tokens, parent_otel_span + ) + ) + _ = await _counted_early( + _count_usage_based_routing_tokens( + core, model_group, request_kwargs, deployment_metadata, total_tokens, parent_otel_span + ) + ) + return counted_request + + +async def _counted_early(count: Awaitable[bool]) -> bool: + try: + return await count + except Exception as e: # noqa: BLE001 # a failed early count must not break the turn; its logging counts it + verbose_router_logger.debug("record_deployment_turn_usage: not counted early. Got - %s", e) + return False + + +async def _count_usage_based_routing_tokens( + core: "EndpointCore", + model_group: str, + request_kwargs: Mapping[str, object], + deployment_metadata: dict[str, object], # mutable-ok: the success callbacks read the counted tokens off it + total_tokens: int, + parent_otel_span: "Span | None", +) -> bool: + from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2 + + routing_kwargs: Final = without(request_kwargs, "litellm_logging_obj") + override: Final = routing_kwargs.get("routing_strategy") + strategy, scope, group = core.handle.routing_context(model_group, override if isinstance(override, str) else None) + selector: Final = core.selectors.selector_for(strategy, scope, group, routing_kwargs) + model_info: Final = deployment_metadata.get("model_info") + deployment_id: Final = as_dict(model_info).get("id") if isinstance(model_info, dict) else None # pyright: ignore[reportUnknownArgumentType] # the routed model_info dict + deployment_name: Final = deployment_metadata.get("deployment") + if ( + not isinstance(selector, LowestTPMLoggingHandler_v2) + or deployment_id is None + or not isinstance(deployment_name, str) + ): + return False + await count_before_callbacks( + deployment_metadata, + USAGE_BASED_ROUTING_COUNTED_TOKENS_METADATA_KEY, + total_tokens, + partial( + selector.async_increment_tpm, + deployment_id=str(deployment_id), + deployment_name=deployment_name, + tokens=total_tokens, + parent_otel_span=parent_otel_span, + ), + ) + return True + + def _deployment_as_configured_now(core: "EndpointCore", deployment: Mapping[str, object]) -> Mapping[str, object]: model_info: Final = deployment.get("model_info") model_id: Final = as_dict(model_info).get("id") if isinstance(model_info, dict) else None # pyright: ignore[reportUnknownArgumentType] # the configured model_info dict @@ -124,4 +227,11 @@ async def _admit_turn_through_routing_strategy( await selector.async_pre_call_check(dict(deployment), None) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # every selector is a CustomLogger -__all__ = ["admit_deployment_turn", "turn_admission"] +__all__ = [ + "TurnUsageRecorder", + "admit_deployment_turn", + "record_deployment_turn_usage", + "turn_admission", + "turn_usage_recorder", + "turns_counted_their_requests", +] diff --git a/litellm/router_rust/surface/legacy_members.py b/litellm/router_rust/surface/legacy_members.py index bef35673aaa..046501f38f3 100644 --- a/litellm/router_rust/surface/legacy_members.py +++ b/litellm/router_rust/surface/legacy_members.py @@ -17,6 +17,7 @@ import inspect import traceback import uuid from collections.abc import Awaitable, Callable, Iterable, Mapping +from functools import partial from types import MappingProxyType from typing import TYPE_CHECKING, Final @@ -47,7 +48,10 @@ from litellm.router_strategy.auto_router.auto_router_model_naming import ( from litellm.router_utils.cooldown_handlers import set_cooldown_deployments from litellm.router_utils.endpoint_eligibility import endpoint_selection from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import async_io_token_refund_failure -from litellm.router_utils.router_callbacks.track_deployment_metrics import find_deployment_metadata +from litellm.router_utils.router_callbacks.track_deployment_metrics import ( + count_before_callbacks, + find_deployment_metadata, +) from litellm.utils import get_utc_datetime if TYPE_CHECKING: @@ -441,28 +445,43 @@ class RustLegacyMembers(RustRegistryMembers, RustOptionalChecks): model_group: Final = request_kwargs.get("model") if deployment_metadata is None or not isinstance(model_group, str): return - deployment_id: Final = _metadata_model_id(deployment_metadata["model_info"]) - if deployment_id is None: - return - total_tokens: Final = response_total_token_count(response) - deployment_name: Final = deployment_metadata.get("deployment") - deployment_metadata[ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY] = total_tokens - try: - await self._increment_deployment_usage( - deployment_id=str(deployment_id), - deployment_name=deployment_name if isinstance(deployment_name, str) else None, - model_group=model_group, - total_tokens=total_tokens, - parent_otel_span=get_parent_otel_span_from_kwargs(dict(request_kwargs)), - ) - except Exception: - _ = deployment_metadata.pop(ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, None) - raise + _ = await self.increment_deployment_usage_before_callbacks( + deployment_metadata, + model_group, + response_total_token_count(response), + get_parent_otel_span_from_kwargs(dict(request_kwargs)), + ) except Exception as e: # noqa: BLE001 # as the Python core: usage accounting never fails a served call verbose_router_logger.debug( "litellm.router.Router::increment_deployment_usage_for_response(): Exception occured - %s", e ) + async def increment_deployment_usage_before_callbacks( + self, + deployment_metadata: dict[str, object], # mutable-ok: the success callback reads the counted tokens off it + model_group: str, + total_tokens: int, + parent_otel_span: "Span | None", + ) -> bool: + deployment_id: Final = _metadata_model_id(deployment_metadata.get("model_info")) + if deployment_id is None: + return False + deployment_name: Final = deployment_metadata.get("deployment") + await count_before_callbacks( + deployment_metadata, + ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, + total_tokens, + partial( + self._increment_deployment_usage, + deployment_id=str(deployment_id), + deployment_name=deployment_name if isinstance(deployment_name, str) else None, + model_group=model_group, + total_tokens=total_tokens, + parent_otel_span=parent_otel_span, + ), + ) + return True + @with_service_target(ROUTER_USAGE_TARGET) async def _increment_deployment_usage( self, diff --git a/litellm/router_strategy/lowest_tpm_rpm_v2.py b/litellm/router_strategy/lowest_tpm_rpm_v2.py index d34d7edf9e4..8d2f45756e4 100644 --- a/litellm/router_strategy/lowest_tpm_rpm_v2.py +++ b/litellm/router_strategy/lowest_tpm_rpm_v2.py @@ -9,17 +9,21 @@ from functools import partial from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import TypeAdapter +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._internal_context import with_service_target from litellm._logging import verbose_logger, verbose_router_logger from litellm.caching.caching import DualCache +from litellm.constants import USAGE_BASED_ROUTING_COUNTED_TOKENS_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs from litellm.litellm_core_utils.deployment_pricing import deployment_backend_models from litellm.litellm_core_utils.token_counter import embedding_aware_token_count from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET +from litellm.router_utils.router_callbacks.track_deployment_metrics import get_counted_usage_tokens from litellm.types.router import RouterErrors from litellm.types.utils import LiteLLMPydanticObjectBase, StandardLoggingPayload from litellm.utils import get_utc_datetime, print_verbose @@ -38,6 +42,13 @@ class RoutingArgs(LiteLLMPydanticObjectBase): ttl: int = 1 * 60 # 1min (RPM/TPM expire key) +class _SuccessKwargs(TypedDict, total=False): + litellm_params: ReadOnly[Mapping[str, object] | None] + + +_SUCCESS_KWARGS_ADAPTER: Final = TypeAdapter(_SuccessKwargs) + + _active_prefetched_usage: Final[ContextVar["PrefetchedUsage | None"]] = ContextVar("prefetched_usage", default=None) @@ -299,7 +310,7 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): ) @with_service_target(ROUTER_USAGE_TARGET) - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: if is_batch_retrieve_call_type(kwargs.get("call_type")): return try: @@ -311,30 +322,19 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): raise ValueError("standard_logging_object not passed in.") model_group: Final = standard_logging_object.get("model_group") model: Final = standard_logging_object["hidden_params"]["litellm_model_name"] - id = standard_logging_object.get("model_id") - if model_group is None or id is None: - return - elif isinstance(id, int): - id = str(id) + model_id: Final = standard_logging_object.get("model_id") total_tokens: Final = standard_logging_object.get("total_tokens") - # ------------ - # Setup values - # ------------ - dt: Final = get_utc_datetime() - current_minute: Final = dt.strftime("%H-%M") # use the same timezone regardless of system clock - - tpm_key: Final = f"{id}:{model}:tpm:{current_minute}" - # ------------ - # Update usage - # ------------ - # update cache - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - ## TPM - await self.router_cache.async_increment_cache_post_call( - key=tpm_key, - value=total_tokens, - ttl=self.routing_args.ttl, - parent_otel_span=parent_otel_span, + if model_group is None or not isinstance(model_id, (str, int)) or not isinstance(total_tokens, int): + return + counted_tokens: Final = get_counted_usage_tokens( + _SUCCESS_KWARGS_ADAPTER.validate_python(kwargs).get("litellm_params") or {}, + USAGE_BASED_ROUTING_COUNTED_TOKENS_METADATA_KEY, + ) + await self.async_increment_tpm( + deployment_id=str(model_id), + deployment_name=model if isinstance(model, str) else None, + tokens=total_tokens if counted_tokens is None else max(0, total_tokens - counted_tokens), + parent_otel_span=get_parent_otel_span_from_kwargs(kwargs), ) ### TESTING ### @@ -345,6 +345,18 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): "litellm.proxy.hooks.lowest_tpm_rpm_v2.py::async_log_success_event(): Exception occured - %s", e ) + @with_service_target(ROUTER_USAGE_TARGET) + async def async_increment_tpm( + self, *, deployment_id: str, deployment_name: str | None, tokens: int, parent_otel_span: Span | None + ) -> None: + current_minute: Final = get_utc_datetime().strftime("%H-%M") + await self.router_cache.async_increment_cache_post_call( + key=f"{deployment_id}:{deployment_name}:tpm:{current_minute}", + value=tokens, + ttl=self.routing_args.ttl, + parent_otel_span=parent_otel_span, + ) + def _return_potential_deployments( self, healthy_deployments: list[dict], diff --git a/litellm/router_utils/router_callbacks/track_deployment_metrics.py b/litellm/router_utils/router_callbacks/track_deployment_metrics.py index b9b638efa3f..2cf4e595296 100644 --- a/litellm/router_utils/router_callbacks/track_deployment_metrics.py +++ b/litellm/router_utils/router_callbacks/track_deployment_metrics.py @@ -9,7 +9,7 @@ get_deployment_failures_for_current_minute get_deployment_successes_for_current_minute """ -from collections.abc import Mapping +from collections.abc import Awaitable, Callable, Mapping from typing import TYPE_CHECKING, Final from litellm.constants import ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY @@ -26,19 +26,31 @@ def find_deployment_metadata(kwargs: Mapping[str, object]) -> dict[str, object] return next((bucket for bucket in buckets if isinstance(bucket, dict) and "model_info" in bucket), None) -def get_counted_usage_tokens(litellm_params: Mapping[str, object]) -> int | None: +def get_counted_usage_tokens( + litellm_params: Mapping[str, object], marker: str = ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY +) -> int | None: buckets: Final = (litellm_params.get(channel) for channel in _METADATA_CHANNELS) counted: Final = next( - ( - bucket[ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY] - for bucket in buckets - if isinstance(bucket, dict) and ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY in bucket - ), + (bucket[marker] for bucket in buckets if isinstance(bucket, dict) and marker in bucket), None, ) return counted if isinstance(counted, int) and not isinstance(counted, bool) else None +async def count_before_callbacks( + deployment_metadata: dict[str, object], # mutable-ok: the success callbacks read the marker off it + marker: str, + tokens: int, + count: Callable[[], Awaitable[object]], +) -> None: + deployment_metadata[marker] = tokens # rebind-ok: the success callbacks read the marker off this call's metadata + try: + await count() + except Exception: + _ = deployment_metadata.pop(marker, None) + raise + + @renamed_router_param("litellm_router_instance") def increment_deployment_successes_for_current_minute( ctx: "RoutingContext", diff --git a/tests/unit/integrations/clickhouse/test_clickhouse_spend_logger.py b/tests/unit/integrations/clickhouse/test_clickhouse_spend_logger.py index 42e69f95bb4..4e932b46ae1 100644 --- a/tests/unit/integrations/clickhouse/test_clickhouse_spend_logger.py +++ b/tests/unit/integrations/clickhouse/test_clickhouse_spend_logger.py @@ -20,7 +20,11 @@ from litellm.integrations.clickhouse.clickhouse_spend_logger import ( strip_cache_hit_suffix, ) from litellm.integrations.clickhouse.context import lens_analysis -from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY +from litellm.constants import ( + INTERNAL_CALL_ORIGIN_METADATA_KEY, + ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, + USAGE_BASED_ROUTING_COUNTED_TOKENS_METADATA_KEY, +) from litellm.litellm_core_utils.internal_call_metadata import BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX from litellm.litellm_core_utils import litellm_logging from litellm.litellm_core_utils.secret_redaction import REDACTED @@ -252,6 +256,8 @@ async def test_custom_request_metadata_is_redacted_before_clickhouse_logging( "user_api_key_budget_reservation": {"token": "internal-token"}, "proxy_server_request": {"headers": {"authorization": "internal-auth"}}, "parent_otel_span": object(), + ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: 35, + USAGE_BASED_ROUTING_COUNTED_TOKENS_METADATA_KEY: 35, }, }, } @@ -281,6 +287,8 @@ async def test_custom_request_metadata_is_redacted_before_clickhouse_logging( assert "user_api_key_auth" not in metadata assert "user_api_key_budget_reservation" not in metadata assert "proxy_server_request" not in metadata + assert ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY not in metadata + assert USAGE_BASED_ROUTING_COUNTED_TOKENS_METADATA_KEY not in metadata litellm_params: Final = cast(Mapping[str, object], kwargs["litellm_params"]) request_metadata: Final = cast(Mapping[str, object], litellm_params["metadata"]) assert request_metadata == { diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 31b7b499845..4f73e8935e8 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -10524,6 +10524,35 @@ def test_passthrough_client_cannot_forge_session_id_omission(client_metadata_key ) +@pytest.mark.parametrize("client_metadata_key", ["litellm_metadata", "metadata"]) +@pytest.mark.parametrize( + "counted_marker", + ["_litellm_router_usage_counted_tokens", "_litellm_usage_based_routing_counted_tokens"], +) +def test_passthrough_client_cannot_mark_its_tokens_as_already_counted( + client_metadata_key: str, counted_marker: str +) -> None: + from litellm.router_utils.router_callbacks.track_deployment_metrics import get_counted_usage_tokens + + mock_request = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.url = httpx.URL("http://0.0.0.0:4000/openai/v1/responses") + mock_request.headers = Headers({}) + mock_request.scope = {} + + kwargs = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint( + request=mock_request, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + passthrough_logging_payload=MagicMock(), + logging_obj=MagicMock(), + _parsed_body={client_metadata_key: {counted_marker: 1_000_000, "client_key": "kept"}}, + litellm_call_id="counted-marker-call-id", + ) + + assert get_counted_usage_tokens(kwargs["litellm_params"], counted_marker) is None + assert kwargs["litellm_params"]["metadata"]["client_key"] == "kept" + + async def _billed_metadata(gate: WebsocketFrameModelGate) -> Mapping[str, object]: served: Final = await _relay_client_frames( [json.dumps({"type": "response.create", "model": "gpt-5.4-mini"})], diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index e4a560adca8..28d2a4fa7ad 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -25,6 +25,7 @@ from litellm.constants import ( SERVER_STREAMING_CLASSIFICATION_MARKER, SESSION_ID_GENERATED_METADATA_KEY, SESSION_ID_OMITTED_METADATA_KEY, + USAGE_BASED_ROUTING_COUNTED_TOKENS_METADATA_KEY, ) from litellm.integrations.otel.model.config import is_otel_v2_enabled from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs @@ -9185,6 +9186,7 @@ _PLANTED_STAMPS = { "previous_models": [{"exception_type": "Timeout", "deployment_id": "spoofed", "attempted_retries": 101}], "_client_output_ceiling": {"api_base": "https://attacker.example"}, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: 10**9, + USAGE_BASED_ROUTING_COUNTED_TOKENS_METADATA_KEY: 10**9, "client_key": "client_value", } @@ -9219,6 +9221,7 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_from_bo assert "request_retry_count" not in updated["metadata"] assert "previous_models" not in updated["metadata"] assert ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY not in updated["metadata"] + assert USAGE_BASED_ROUTING_COUNTED_TOKENS_METADATA_KEY not in updated["metadata"] assert updated["metadata"]["client_key"] == "client_value" diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index 7e044b55f5a..9a73054b2e3 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -1595,6 +1595,9 @@ class TestWebSocketProjectQuotaEnforcement: quota_callback.enforce_project_io_token_quota_for_frame.assert_awaited_once() +_TURN_USAGE: Final = {"input_tokens": 20, "output_tokens": 15, "total_tokens": 35} + + def _turn_frames() -> tuple[str, str, str]: return tuple(json.dumps({"type": "response.create", "input": word}) for word in ("one", "two", "three")) @@ -1928,9 +1931,139 @@ def _errors(websocket: _ClientSocket) -> list[tuple[object, object]]: return [(frame.get("status"), frame["error"]) for frame in websocket.sent if frame.get("type") == "error"] +class _BackendSending: + """A provider socket that sends these events, then closes.""" + + def __init__(self, *events: Mapping[str, object]) -> None: + self.events: Final = list(events) + + async def recv(self, decode: bool = False) -> str: + import websockets.exceptions + + if not self.events: + raise websockets.exceptions.ConnectionClosed(None, None) + return json.dumps(self.events.pop(0)) + + +def _turn_logged_on_deployment(deployment_id: str) -> MagicMock: + logging_obj: Final = MagicMock() + logging_obj.litellm_params = {"litellm_metadata": {"model_info": {"id": deployment_id}}} + logging_obj.copy_for_new_call.return_value = logging_obj + logging_obj.dispatch_success_handlers = AsyncMock() + logging_obj.dispatch_failure_handlers = AsyncMock() + return logging_obj + + +def _logged_event_types(logging_obj: MagicMock) -> list[list[object]]: + return [[event["type"] for event in call.args[0]] for call in logging_obj.dispatch_success_handlers.call_args_list] + + class TestWebSocketTurnAdmission: """A turn after the connection's first is admitted by the deployment again before it reaches the provider.""" + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("terminal_event", "counted_before_it_reaches_the_client"), + [ + ({"type": "response.completed", "response": {"usage": _TURN_USAGE}}, [("ws-dep", 35)]), + ({"type": "response.incomplete", "response": {"usage": _TURN_USAGE}}, [("ws-dep", 35)]), + ({"type": "response.failed", "response": {"usage": _TURN_USAGE, "error": {"code": "server_error"}}}, []), + ], + ) + async def test_native_relay_counts_a_finished_turn_before_forwarding_its_terminal_event( + self, terminal_event: dict[str, object], counted_before_it_reaches_the_client: list[tuple[object, int]] + ): + counted: Final[list[tuple[object, int]]] = [] + + async def record_turn_usage(turn_litellm_params: Mapping[str, object], total_tokens: int) -> None: + await asyncio.sleep(0) + metadata: Final = turn_litellm_params["litellm_metadata"] + assert isinstance(metadata, dict) + counted.append((metadata["model_info"]["id"], total_tokens)) + + class Client: + def __init__(self) -> None: + self.seen: Final[list[tuple[object, list[tuple[object, int]]]]] = [] + + async def send_text(self, data: str) -> None: + self.seen.append((json.loads(data)["type"], list(counted))) + + client: Final = Client() + handler: Final = _make_streaming( + websocket=client, + backend_ws=_BackendSending({"type": "response.created", "response": {}}, terminal_event), + logging_obj=_turn_logged_on_deployment("ws-dep"), + turn_usage_recorder=record_turn_usage, + ) + + await handler.backend_to_client() + + assert client.seen == [ + ("response.created", []), + (terminal_event["type"], counted_before_it_reaches_the_client), + ] + + @pytest.mark.asyncio + async def test_native_relay_forwards_and_logs_a_finished_turn_whose_usage_it_cannot_read(self): + recorded: Final[list[int]] = [] + + async def record_turn_usage(turn_litellm_params: Mapping[str, object], total_tokens: int) -> None: + recorded.append(total_tokens) + + unreadable_usage: Final = {"input_tokens": None, "output_tokens": 3, "total_tokens": 3} + client: Final = _ClientSocket() + logging_obj: Final = _turn_logged_on_deployment("ws-dep") + handler: Final = _make_streaming( + websocket=client, + backend_ws=_BackendSending( + {"type": "response.created", "response": {}}, + {"type": "response.completed", "response": {"usage": unreadable_usage}}, + ), + logging_obj=logging_obj, + turn_usage_recorder=record_turn_usage, + ) + + await handler.backend_to_client() + + assert [event["type"] for event in client.sent] == ["response.created", "response.completed"] + assert _logged_event_types(logging_obj) == [["response.created", "response.completed"]] + assert recorded == [] + + @pytest.mark.asyncio + async def test_native_relay_finishes_counting_a_turn_before_logging_it_when_the_client_leaves(self): + counting: Final = asyncio.Event() + count_written: Final = asyncio.Event() + counted: Final[list[int]] = [] + + async def record_turn_usage_slowly(turn_litellm_params: Mapping[str, object], total_tokens: int) -> None: + counting.set() + await count_written.wait() + counted.append(total_tokens) + + logging_obj: Final = _turn_logged_on_deployment("ws-dep") + handler: Final = _make_streaming( + websocket=_ClientSocket(), + backend_ws=_BackendSending( + {"type": "response.created", "response": {}}, + {"type": "response.completed", "response": {"usage": _TURN_USAGE}}, + ), + logging_obj=logging_obj, + turn_usage_recorder=record_turn_usage_slowly, + ) + relay: Final = asyncio.create_task(handler.backend_to_client()) + await asyncio.wait_for(counting.wait(), timeout=5) + + relay.cancel() + _ = await asyncio.wait((relay,), timeout=0.1) + logged_while_counting: Final = _logged_event_types(logging_obj) + count_written.set() + with pytest.raises(asyncio.CancelledError): + await relay + + assert logged_while_counting == [] + assert counted == [35] + assert _logged_event_types(logging_obj) == [["response.created", "response.completed"]] + @pytest.mark.asyncio async def test_native_relay_drops_a_later_turn_the_deployment_refuses(self): from unittest.mock import AsyncMock diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index 9463424a48b..8a5d164fdff 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -6,6 +6,7 @@ import time import weakref from collections.abc import Awaitable, Callable, Iterator, Mapping from datetime import datetime, timezone +from types import MappingProxyType from typing import TYPE_CHECKING, Final from unittest.mock import AsyncMock, MagicMock, patch @@ -16,14 +17,22 @@ from pydantic import BaseModel import litellm from litellm import Router +from litellm.caching.caching import DualCache, InMemoryCache from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting +from litellm.litellm_core_utils.get_litellm_params import get_litellm_params +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.responses.streaming_iterator import ResponsesWsTurnAdmission, ResponsesWsTurnUsageRecorder +from litellm.router import router_arguments +from litellm.router_config import RouterConfig, build_router_config +from litellm.router_python_host import PythonHost from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook from litellm.types.litellm_params import RoutingStrategyName from litellm.types.router import AlertingConfig, Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo from litellm.types.utils import ( + CallTypes, ModelResponse, StandardLoggingHiddenParams, StandardLoggingMetadata, @@ -33,6 +42,8 @@ from litellm.types.utils import ( if TYPE_CHECKING: + from opentelemetry.trace import Span + from tests.unit.conftest import RouterRedis @@ -329,6 +340,265 @@ async def test_deployment_callback_on_success_counts_a_responses_websocket_turn( assert (await router.health().usage("gpt-5-mini"))[0] == 17 +def _open_websocket_turn(routed: Mapping[str, object], call_id: str) -> Logging: + connection: Final = routed["litellm_logging_obj"] + assert isinstance(connection, Logging) + connection.update_environment_variables(litellm_params=get_litellm_params(**routed), optional_params={}) + return connection.copy_for_new_call(litellm_call_id=call_id, start_time=datetime.now()) + + +async def _complete_35_token_turn(turn: Logging) -> None: + await turn.dispatch_success_handlers( + [ + { + "type": "response.completed", + "response": {"usage": {"input_tokens": 20, "output_tokens": 15, "total_tokens": 35}}, + } + ], + prefer_async_handlers=True, + ) + + +def _with_success_callbacks_of(monkeypatch: pytest.MonkeyPatch, build_router: Callable[[], Router]) -> Router: + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", []) + router: Final = build_router() + monkeypatch.setattr(litellm, "_async_success_callback", [*litellm._async_success_callback, *litellm.callbacks]) + return router + + +def _usage_based_routing_model_list(model_name: str, tpm: int) -> list[DeploymentTypedDict]: + return [ + { + "model_name": model_name, + "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "sk-fake", "tpm": tpm}, + "model_info": {"id": f"{model_name}-dep"}, + } + ] + + +def _usage_based_routing_router(monkeypatch: pytest.MonkeyPatch, model_name: str, tpm: int) -> Router: + return _with_success_callbacks_of( + monkeypatch, + lambda: Router( + model_list=_usage_based_routing_model_list(model_name, tpm), routing_strategy="usage-based-routing-v2" + ), + ) + + +async def _usage_after(router: Router, model_group: str, opened: tuple[int, int]) -> tuple[int, int]: + tokens, requests = await router.health().usage(model_group) + return (tokens or 0) - opened[0], (requests or 0) - opened[1] + + +async def _usage_now(router: Router, model_group: str) -> tuple[int, int]: + return await _usage_after(router, model_group, (0, 0)) + + +def _websocket_connection(model_name: str) -> Logging: + return Logging( + model=model_name, + messages=[], + stream=False, + call_type=CallTypes.aresponses_websocket.value, + start_time=time.time(), + litellm_call_id=f"{model_name}-connection", + function_id=model_name, + ) + + +async def _websocket_turn_hooks( + monkeypatch: pytest.MonkeyPatch, router: Router, model_name: str +) -> tuple[dict[str, object], ResponsesWsTurnAdmission, ResponsesWsTurnUsageRecorder]: + routed: Final[dict[str, object]] = {} + + async def _aresponses_websocket(**kwargs: object) -> None: + routed.update(kwargs) + + monkeypatch.setattr(litellm, "_aresponses_websocket", _aresponses_websocket) + await router._aresponses_websocket(model=model_name, litellm_logging_obj=_websocket_connection(model_name)) + admit_turn: Final = routed["turn_admission"] + record_turn: Final = routed["turn_usage_recorder"] + assert callable(admit_turn) + assert callable(record_turn) + return routed, admit_turn, record_turn + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("router_minute_pinned") +async def test_usage_based_routing_refuses_the_next_websocket_turn_before_the_finished_turn_s_logging_runs( + monkeypatch: pytest.MonkeyPatch, +): + router: Final = _usage_based_routing_router(monkeypatch, "ws-tpm-v2-early", tpm=30) + routed, admit_turn, record_turn = await _websocket_turn_hooks(monkeypatch, router, "ws-tpm-v2-early") + await admit_turn() + + await record_turn(_open_websocket_turn(routed, "ws-tpm-v2-early-turn").litellm_params, 35) + + with pytest.raises(litellm.RateLimitError, match="ws-tpm-v2-early"): + await admit_turn() + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("router_minute_pinned") +async def test_websocket_turn_counted_before_its_logging_is_not_counted_again_when_the_logging_runs( + monkeypatch: pytest.MonkeyPatch, +): + router: Final = _usage_based_routing_router(monkeypatch, "ws-tpm-v2-once", tpm=69) + routed, admit_turn, record_turn = await _websocket_turn_hooks(monkeypatch, router, "ws-tpm-v2-once") + opened: Final = await _usage_now(router, "ws-tpm-v2-once") + await admit_turn() + first_turn: Final = _open_websocket_turn(routed, "ws-tpm-v2-once-first") + + await record_turn(first_turn.litellm_params, 35) + assert await _usage_after(router, "ws-tpm-v2-once", opened) == (35, 1) + await _complete_35_token_turn(first_turn) + + assert await _usage_after(router, "ws-tpm-v2-once", opened) == (35, 1) + await admit_turn() + await record_turn(_open_websocket_turn(routed, "ws-tpm-v2-once-second").litellm_params, 35) + with pytest.raises(litellm.RateLimitError, match="ws-tpm-v2-once"): + await admit_turn() + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("router_minute_pinned") +async def test_usage_based_routing_counts_an_http_call_once(monkeypatch: pytest.MonkeyPatch): + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + router: Final = _usage_based_routing_router(monkeypatch, "http-tpm-v2", tpm=1000) + + response: Final = await router.acompletion( + model="http-tpm-v2", messages=[{"role": "user", "content": "hi"}], mock_response="pong" + ) + await asyncio.sleep(0) + await GLOBAL_LOGGING_WORKER.flush() + + tpm_keys, _ = router.lowesttpm_logger_v2.usage_counter_keys(router.get_model_list(model_name="http-tpm-v2")) + assert await router.cache.async_get_cache(key=tpm_keys[0]) == response.usage.total_tokens + + +class _FirstTpmCountFails(DualCache): + failed: bool = False + + async def async_increment_cache_post_call( + self, key: str, value: float, ttl: int | None, parent_otel_span: "Span | None" = None + ) -> None: + if self.failed or ":tpm:" not in key: + return await super().async_increment_cache_post_call(key, value, ttl, parent_otel_span) + self.failed = True + raise RuntimeError("routing cache unavailable") + + +class _HostWhoseFirstTpmCountFails(PythonHost): + def routing_cache(self, config: RouterConfig) -> DualCache: + return _FirstTpmCountFails(in_memory_cache=InMemoryCache()) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("router_minute_pinned") +async def test_usage_based_routing_counts_a_websocket_turn_whose_early_count_fails_when_its_logging_runs( + monkeypatch: pytest.MonkeyPatch, +): + router: Final = _with_success_callbacks_of( + monkeypatch, + lambda: Router.from_config( + build_router_config( + router_arguments( + model_list=_usage_based_routing_model_list("ws-tpm-v2-count-fails", tpm=30), + routing_strategy="usage-based-routing-v2", + ) + ), + _HostWhoseFirstTpmCountFails(), + ), + ) + routed, admit_turn, record_turn = await _websocket_turn_hooks(monkeypatch, router, "ws-tpm-v2-count-fails") + await admit_turn() + turn: Final = _open_websocket_turn(routed, "ws-tpm-v2-count-fails-turn") + + await record_turn(turn.litellm_params, 35) + await admit_turn() + await _complete_35_token_turn(turn) + + with pytest.raises(litellm.RateLimitError, match="ws-tpm-v2-count-fails"): + await admit_turn() + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("router_minute_pinned") +async def test_websocket_turn_counts_its_tokens_against_its_deployment_once( + monkeypatch: pytest.MonkeyPatch, +): + router: Final = _with_success_callbacks_of(monkeypatch, lambda: _rpm_tpm_router("ws-default-routing")) + routed, _, record_turn = await _websocket_turn_hooks(monkeypatch, router, "gpt-5-mini") + opened: Final = await _usage_now(router, "gpt-5-mini") + turn: Final = _open_websocket_turn(routed, "ws-default-routing-turn") + + await record_turn(turn.litellm_params, 35) + assert await _usage_after(router, "gpt-5-mini", opened) == (35, 1) + await _complete_35_token_turn(turn) + + assert await _usage_after(router, "gpt-5-mini", opened) == (35, 1) + + +_TURN_WITHOUT_DEPLOYMENT: Final = MappingProxyType({"litellm_metadata": {"model_group": "gpt-5-mini"}}) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("router_minute_pinned") +@pytest.mark.parametrize( + ("turns", "socket_usage"), + [ + ((), (0, 1)), + (({},), (35, 1)), + (({}, {}, {}), (105, 3)), + ((_TURN_WITHOUT_DEPLOYMENT,), (0, 1)), + ], + ids=["no_turns", "one_turn", "three_turns", "turn_without_deployment"], +) +async def test_websocket_counts_one_router_request_per_counted_turn_or_one_for_the_socket( + monkeypatch: pytest.MonkeyPatch, turns: tuple[Mapping[str, object], ...], socket_usage: tuple[int, int] +): + router: Final = _with_success_callbacks_of(monkeypatch, lambda: _rpm_tpm_router("ws-requests")) + + async def _aresponses_websocket(*, turn_usage_recorder: ResponsesWsTurnUsageRecorder, **kwargs: object) -> None: + for index, turn in enumerate(turns): + await turn_usage_recorder( + _open_websocket_turn({**kwargs, **turn}, f"ws-requests-turn-{index}").litellm_params, 35 + ) + + monkeypatch.setattr(litellm, "_aresponses_websocket", _aresponses_websocket) + await router._aresponses_websocket(model="gpt-5-mini", litellm_logging_obj=_websocket_connection("gpt-5-mini")) + + assert await _usage_now(router, "gpt-5-mini") == socket_usage + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("router_minute_pinned") +@pytest.mark.parametrize( + "deployment_metadata", + [{}, {"model_info": {"id": None}}], + ids=["no_deployment", "deployment_without_id"], +) +async def test_websocket_turn_not_attributed_to_a_deployment_does_not_throttle_the_next_turn( + monkeypatch: pytest.MonkeyPatch, deployment_metadata: dict[str, object] +): + router: Final = _usage_based_routing_router(monkeypatch, "ws-tpm-v2-unrouted", tpm=30) + routed, admit_turn, record_turn = await _websocket_turn_hooks(monkeypatch, router, "ws-tpm-v2-unrouted") + opened: Final = await _usage_now(router, "ws-tpm-v2-unrouted") + await admit_turn() + turn: Final = _open_websocket_turn( + {**routed, "litellm_metadata": {"model_group": "ws-tpm-v2-unrouted", **deployment_metadata}}, + "ws-tpm-v2-unrouted-turn", + ) + + await record_turn(turn.litellm_params, 35) + + await admit_turn() + assert (await _usage_after(router, "ws-tpm-v2-unrouted", opened))[0] == 0 + + def _io_limited_router(kind: str, **litellm_params: object) -> Router: return Router( model_list=[ @@ -788,16 +1058,18 @@ async def test_set_response_headers(model_list): assert resp is None +def _rpm_tpm_router_model_list(model_id: str) -> list[DeploymentTypedDict]: + return [ + { + "model_name": "gpt-5-mini", + "litellm_params": {"model": "gpt-5-mini", "api_key": "sk-fake", "tpm": 1000, "rpm": 100}, + "model_info": {"id": model_id}, + } + ] + + def _rpm_tpm_router(model_id: str) -> Router: - return Router( - model_list=[ - { - "model_name": "gpt-5-mini", - "litellm_params": {"model": "gpt-5-mini", "api_key": "sk-fake", "tpm": 1000, "rpm": 100}, - "model_info": {"id": model_id}, - } - ] - ) + return Router(model_list=_rpm_tpm_router_model_list(model_id)) def _pin_router_clock(monkeypatch: pytest.MonkeyPatch, pinned: datetime) -> None: