diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 0e315e97432..0fb7363cbb5 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -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" diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index e766a2be9c3..42d64eb88e4 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -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): diff --git a/litellm/router.py b/litellm/router.py index 39e5a154739..a89361cc629 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -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 diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index c1e1ebb87b9..5594cdd1384 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -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 ] diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index 9174306eba0..a3ca604281c 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -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(), diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index 1f5421aa97c..228e0ad82a6 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -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(): diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index 030e051a380..ce60cce8408 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -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(