mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
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>
This commit is contained in:
norkelbyte
GitHub
parent
69b0f1deac
commit
68e78bc53a
18 changed files with 744 additions and 77 deletions
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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]:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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],
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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 == {
|
||||
|
|
|
|||
|
|
@ -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"})],
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue