scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-20 05:04:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 155

fix(router): count WebSocket turn tokens toward TPM and re-run routing strategy checks on each turn

Browse source
This commit is contained in:
mateo-berri 2026-10-16 08:39:14 -07:00
parent 61ed4b7899
commit 0e0ed1e2be
7 changed files with 184 additions and 25 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
litellm/cost_calculator.py
View file

@ -3068,14 +3068,19 @@ class ResponsesWebSocketTokenUsageProcessor(BaseTokenUsageProcessor):
return response_ids[-1] if response_ids else None
@staticmethod
def create_logging_object(results: Sequence[Mapping[str, object]]) -> LiteLLMRealtimeStreamLoggingObject:
def create_logging_object(
results: Sequence[Mapping[str, object]], litellm_model_name: str | None = None
) -> LiteLLMRealtimeStreamLoggingObject:
tier_partition: Final = ResponsesWebSocketTokenUsageProcessor.partition_results_by_service_tier(results)
return LiteLLMRealtimeStreamLoggingObject(
logging_object: Final = LiteLLMRealtimeStreamLoggingObject(
id=ResponsesWebSocketTokenUsageProcessor.response_id_from_responses_ws_results(results),
usage=ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results(results),
results=cast(OpenAIRealtimeStreamList, list(results)), # cast-ok: raw WS frames, unvalidated field
service_tier=next(iter(tier_partition)) if len(tier_partition) == 1 else None,
)
if litellm_model_name is not None:
logging_object.hidden_params = {"litellm_model_name": litellm_model_name}
return logging_object
_TRANSCRIPTION_COMPLETED_EVENT_TYPE: Final = "conversation.item.input_audio_transcription.completed"

13
litellm/litellm_core_utils/litellm_logging.py
View file

@ -968,6 +968,16 @@ class Logging(LiteLLMLoggingBaseClass):
return model_id
return None
def _routed_deployment_model(self) -> str | None:
litellm_params: Final = self._model_call_details_view.get("litellm_params")
metadata: Final = (
StandardLoggingPayloadSetup.merge_litellm_metadata(_LITELLM_PARAMS.validate_python(litellm_params))
if isinstance(litellm_params, dict)
else {}
)
deployment: Final = metadata.get("deployment")
return deployment if isinstance(deployment, str) else None
def get_deployment_model_for_cost(self) -> str | None:
"""The provider-qualified model to price against.
@ -2419,7 +2429,8 @@ class Logging(LiteLLMLoggingBaseClass):
elif self.call_type == CallTypes.aresponses_websocket.value and isinstance(result, list): # pyright: ignore[reportUnknownMemberType] # Logging.call_type is untyped
logging_result = ResponsesWebSocketTokenUsageProcessor.create_logging_object(
results=result # pyright: ignore[reportUnknownArgumentType] # raw event dicts from the WS stream
results=result, # pyright: ignore[reportUnknownArgumentType] # raw event dicts from the WS stream
litellm_model_name=self._routed_deployment_model(),
)
elif self.relays_llm_passthrough and isinstance(result, Response):

44
litellm/router.py
View file

@ -9299,9 +9299,18 @@ class Router:
)
raise e
def _deployment_as_configured_now(self, deployment: Mapping[str, object]) -> Mapping[str, object]:
model_info: Final = deployment.get("model_info")
model_id: Final = (
_MODEL_INFO_ADAPTER.validate_python(model_info).get("id") if isinstance(model_info, dict) else None
)
current: Final = self.get_deployment(model_id) if isinstance(model_id, str) else None
return deployment if current is None else current.model_dump(exclude_none=True)
async def async_admit_deployment_turn(
self, model_group: str, deployment: Mapping[str, object], request_kwargs: Mapping[str, object]
self, model_group: str, opened_on: Mapping[str, object], request_kwargs: Mapping[str, object]
) -> None:
deployment: Final = self._deployment_as_configured_now(opened_on)
budget_limiter: Final = self.get_router_deployment_budget_limiter()
if budget_limiter is not None:
try:
@ -9318,10 +9327,43 @@ class Router:
model=model_group,
category=litellm.RateLimitErrorCategory.LITELLM_RATE_LIMIT,
) from e
await self._admit_turn_through_routing_strategy(model_group, deployment, request_kwargs)
await self.async_routing_strategy_pre_call_checks( # pyright: ignore[reportUnknownMemberType] # legacy signature takes an untyped deployment dict
deployment=dict(without_io_token_limits(deployment)), parent_otel_span=None
)
async def _admit_turn_through_routing_strategy(
self, model_group: str, deployment: Mapping[str, object], request_kwargs: Mapping[str, object]
) -> None:
routing_kwargs: Final = {key: value for key, value in request_kwargs.items() if key != "litellm_logging_obj"}
strategy, selector = self._get_routing_context( # pyright: ignore[reportUnknownMemberType] # legacy signature takes an untyped request dict
model=model_group, request_kwargs=routing_kwargs
)
if strategy == "simple-shuffle" or selector is None:
return
selected: Final = await self._select_deployment_async( # pyright: ignore[reportUnknownMemberType] # legacy signature takes untyped deployment dicts
strategy=strategy,
selector=selector,
model=model_group,
healthy_deployments=[dict(deployment)],
messages=None,
input=None,
request_kwargs=routing_kwargs,
)
if selected is None:
no_deployment: Final = await async_raise_no_deployment_exception(
litellm_router_instance=self, model=model_group, parent_otel_span=None
)
raise litellm.RateLimitError(
message=str(no_deployment),
llm_provider="",
model=model_group,
category=litellm.RateLimitErrorCategory.LITELLM_RATE_LIMIT,
) from no_deployment
await self._async_override_selector_pre_call_check( # pyright: ignore[reportUnknownMemberType] # legacy signature takes an untyped deployment dict
strategy, selector, dict(deployment), None
)
@contextlib.asynccontextmanager
async def _deployment_slot(
self, deployment: dict, kwargs: Mapping[str, object], parent_otel_span: Span | None

2
tests/code_coverage_tests/router_code_coverage.py
View file

@ -120,6 +120,8 @@ ignored_function_names = [
"_names_route_model", # Tested through get_router_model_info in test_router_model_cost_isolation.py
"_declared_base_model", # Tested through get_router_model_info in test_router_model_cost_isolation.py
"async_admit_deployment_turn", # Bound as the WebSocket turn_admission partial, tested through it in test_router_helper_utils.py
"_deployment_as_configured_now", # Read by async_admit_deployment_turn, tested through turn_admission in test_router_helper_utils.py
"_admit_turn_through_routing_strategy", # Run by async_admit_deployment_turn, tested through turn_admission in test_router_helper_utils.py
]

23
tests/unit/proxy/hooks/test_proxy_track_cost_callback.py
View file

@ -4,7 +4,7 @@ import logging
from datetime import datetime, timedelta, timezone
from types import SimpleNamespace
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
from unittest.mock import AsyncMock, MagicMock, create_autospec, patch
import httpx
import pytest
@ -146,7 +146,7 @@ async def test_async_post_call_failure_hook_carries_guardrail_info_from_litellm_
],
)
async def test_async_post_call_failure_hook_logs_the_routed_model_group_on_the_failure_row(
metadata_buckets: dict, expected_model_group: str
metadata_buckets: dict[str, dict[str, str]], expected_model_group: str
):
request_data: Final = {
"model": "ws-budget-sock",
@ -155,20 +155,15 @@ async def test_async_post_call_failure_hook_logs_the_routed_model_group_on_the_f
**metadata_buckets,
}
with patch(
"litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database",
new_callable=AsyncMock,
) as mock_update_database:
await ProxyDBLogger().async_post_call_failure_hook(
request_data=request_data,
original_exception=litellm.RateLimitError(
message="crossed budget", llm_provider="", model="ws-budget-sock"
),
user_api_key_dict=UserAPIKeyAuth(api_key="test_api_key"),
)
spend_writer: Final = create_autospec(DBSpendUpdateWriter, instance=True)
await ProxyDBLogger(spend_writer=lambda: spend_writer).async_post_call_failure_hook(
request_data=request_data,
original_exception=litellm.RateLimitError(message="crossed budget", llm_provider="", model="ws-budget-sock"),
user_api_key_dict=UserAPIKeyAuth(api_key="test_api_key"),
)
row: Final = get_logging_payload(
kwargs=mock_update_database.call_args[1]["kwargs"],
kwargs=spend_writer.update_database.call_args.kwargs["kwargs"],
response_obj={},
start_time=datetime.now(),
end_time=datetime.now(),

15
tests/unit/responses/test_responses_websocket_all_providers.py
View file

@ -1186,7 +1186,7 @@ class TestWebSocketProjectQuotaEnforcement:
callback, not just the first one."""
@pytest.mark.asyncio
async def test_managed_handler_blocks_frame_rejected_by_quota_callback(self, monkeypatch):
async def test_managed_handler_blocks_frame_rejected_by_quota_callback(self, monkeypatch: pytest.MonkeyPatch):
from unittest.mock import AsyncMock, MagicMock
import litellm
@ -1232,11 +1232,11 @@ class TestWebSocketProjectQuotaEnforcement:
quota_callback.enforce_project_io_token_quota_for_frame.assert_awaited_once()
assert aresponses_called is False
mock_websocket.send_text.assert_called_once()
error_event = mock_websocket.send_text.call_args[0][0]
assert "rate_limit_exceeded" in error_event
error_event = json.loads(mock_websocket.send_text.call_args[0][0])
assert (error_event["status"], error_event["error"]["type"]) == (429, "rate_limit_exceeded")
@pytest.mark.asyncio
async def test_managed_handler_forwards_frame_allowed_by_quota_callback(self, monkeypatch):
async def test_managed_handler_forwards_frame_allowed_by_quota_callback(self, monkeypatch: pytest.MonkeyPatch):
from unittest.mock import AsyncMock, MagicMock
import litellm
@ -1317,7 +1317,8 @@ class TestWebSocketProjectQuotaEnforcement:
assert allowed is False
mock_backend_ws.send.assert_not_called()
mock_websocket.send_text.assert_called_once()
assert "rate_limit_exceeded" in mock_websocket.send_text.call_args[0][0]
error_event = json.loads(mock_websocket.send_text.call_args[0][0])
assert (error_event["status"], error_event["error"]["type"]) == (429, "rate_limit_exceeded")
@pytest.mark.asyncio
async def test_native_handler_forwards_frame_allowed_by_quota_callback(self):
@ -1420,7 +1421,7 @@ class TestWebSocketTurnAdmission:
assert "crossed budget" in str(failure)
@pytest.mark.asyncio
async def test_managed_relay_skips_a_later_turn_the_deployment_refuses(self, monkeypatch):
async def test_managed_relay_skips_a_later_turn_the_deployment_refuses(self, monkeypatch: pytest.MonkeyPatch):
from unittest.mock import AsyncMock
import litellm
@ -1430,7 +1431,7 @@ class TestWebSocketTurnAdmission:
answered: Final[list[object]] = []
async def fake_aresponses(*args, **kwargs):
async def fake_aresponses(*args: object, **kwargs: object):
answered.append(kwargs.get("input"))
async def _empty():

103
tests/unit/test_router/test_router_helper_utils.py
View file

@ -1,6 +1,7 @@
import json
import os
from datetime import datetime, timezone
from collections.abc import Mapping
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
@ -13,6 +14,7 @@ from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, ROUTER_USAGE_
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.types.router import Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo
from litellm.types.utils import (
CallTypes,
ModelResponse,
StandardLoggingHiddenParams,
StandardLoggingMetadata,
@ -537,6 +539,27 @@ async def _kwargs_routed_to(router: Router, model: str, function_name: str) -> d
return routed
async def _finish_35_token_websocket_turn(routed: Mapping[str, object], call_id: str) -> None:
import time
from litellm.litellm_core_utils.litellm_logging import Logging
turn: Final = Logging(
model="gpt-5-mini",
messages=[],
stream=False,
call_type=CallTypes.aresponses_websocket.value,
start_time=time.time(),
litellm_call_id=call_id,
function_id=call_id,
)
turn.update_environment_variables(litellm_params={"litellm_metadata": routed["litellm_metadata"]}, optional_params={})
await turn.dispatch_success_handlers(
[{"type": "response.completed", "response": {"usage": {"input_tokens": 20, "output_tokens": 15, "total_tokens": 35}}}],
prefer_async_handlers=True,
)
@pytest.mark.asyncio
async def test_responses_websocket_turn_is_refused_once_its_deployment_crosses_budget():
router = Router(
@ -583,6 +606,86 @@ async def test_responses_websocket_turn_after_the_first_counts_toward_deployment
await admit_turn()
@pytest.mark.asyncio
async def test_responses_websocket_turn_is_checked_against_limits_changed_after_the_socket_opened():
router = Router(
model_list=[
{
"model_name": "ws-rpm-lowered",
"litellm_params": {"model": "openai/gpt-5-mini", "api_key": "sk-fake", "rpm": 100},
"model_info": {"id": "ws-rpm-lowered-dep"},
}
],
optional_pre_call_checks=["enforce_model_rate_limits"],
)
admit_turn: Final = (await _kwargs_routed_to(router, "ws-rpm-lowered", "_aresponses_websocket"))["turn_admission"]
assert callable(admit_turn)
await admit_turn()
router.upsert_deployment(
Deployment(
model_name="ws-rpm-lowered",
litellm_params=LiteLLM_Params(model="openai/gpt-5-mini", api_key="sk-fake", rpm=1),
model_info=ModelInfo(id="ws-rpm-lowered-dep"),
)
)
with pytest.raises(litellm.RateLimitError, match="RPM limit=1"):
await admit_turn()
@pytest.mark.asyncio
async def test_responses_websocket_turn_counts_toward_deployment_tpm(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(litellm, "callbacks", [])
router = Router(
model_list=[
{
"model_name": "ws-tpm",
"litellm_params": {"model": "openai/gpt-5-mini", "api_key": "sk-fake", "tpm": 30},
"model_info": {"id": "ws-tpm-dep"},
}
],
optional_pre_call_checks=["enforce_model_rate_limits"],
)
monkeypatch.setattr(litellm, "_async_success_callback", list(litellm.callbacks))
routed: Final = await _kwargs_routed_to(router, "ws-tpm", "_aresponses_websocket")
admit_turn: Final = routed["turn_admission"]
assert callable(admit_turn)
await admit_turn()
await _finish_35_token_websocket_turn(routed, "ws-tpm-turn")
with pytest.raises(litellm.RateLimitError, match="TPM limit=30, current usage=35"):
await admit_turn()
@pytest.mark.asyncio
async def test_responses_websocket_turn_is_refused_by_usage_based_routing_once_deployment_tpm_is_used(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setattr(litellm, "callbacks", [])
router = Router(
model_list=[
{
"model_name": "ws-tpm-v2",
"litellm_params": {"model": "openai/gpt-5-mini", "api_key": "sk-fake", "tpm": 30},
"model_info": {"id": "ws-tpm-v2-dep"},
}
],
routing_strategy="usage-based-routing-v2",
)
monkeypatch.setattr(litellm, "_async_success_callback", list(litellm.callbacks))
routed: Final = await _kwargs_routed_to(router, "ws-tpm-v2", "_aresponses_websocket")
admit_turn: Final = routed["turn_admission"]
assert callable(admit_turn)
await admit_turn()
await _finish_35_token_websocket_turn(routed, "ws-tpm-v2-turn")
with pytest.raises(litellm.RateLimitError):
await admit_turn()
@pytest.mark.asyncio
async def test_responses_websocket_turns_do_not_reserve_deployment_otpm():
router = Router(