From 3ef036531680a6b52ca0d434175f8736b87a96be Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 4 Nov 2026 20:09:13 +0000 Subject: [PATCH] fix(router): use deployment metadata for streamed TPM --- litellm/router_strategy/lowest_tpm_rpm_v2.py | 31 ++++++-- .../router_strategy/test_lowest_tpm_rpm_v2.py | 70 +++++++++++++++++++ 2 files changed, 95 insertions(+), 6 deletions(-) diff --git a/litellm/router_strategy/lowest_tpm_rpm_v2.py b/litellm/router_strategy/lowest_tpm_rpm_v2.py index 0574c3d7f69..65444937d8d 100644 --- a/litellm/router_strategy/lowest_tpm_rpm_v2.py +++ b/litellm/router_strategy/lowest_tpm_rpm_v2.py @@ -16,7 +16,10 @@ from litellm._logging import verbose_logger, verbose_router_logger from litellm.caching.caching import DualCache 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.core_helpers import ( + get_litellm_metadata_from_kwargs, # pyright: ignore[reportUnknownVariableType] # helper has no typed return + 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.types.router import RouterErrors @@ -255,7 +258,13 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): return deployment # don't fail calls if eg. redis fails to connect @with_service_target("router_usage") - def log_success_event(self, kwargs, response_obj, start_time, end_time): + def log_success_event( + self, + kwargs, + response_obj, + start_time, + end_time, # pyright: ignore[reportUnknownParameterType] # CustomLogger callback signature is untyped + ) -> None: if is_batch_retrieve_call_type(kwargs.get("call_type")): return try: @@ -266,7 +275,9 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): if standard_logging_object is None: raise ValueError("standard_logging_object not passed in.") model_group: Final = standard_logging_object.get("model_group") - model: Final = standard_logging_object["hidden_params"].get("litellm_model_name") + deployment_model: Final[object] = get_litellm_metadata_from_kwargs(kwargs).get("deployment") # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType, reportUnknownVariableType] # helper has no typed metadata signature + hidden_model: Final[str | None] = standard_logging_object["hidden_params"].get("litellm_model_name") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # standard logging payload hidden parameters are untyped + model: Final[str | None] = (deployment_model if isinstance(deployment_model, str) else None) or hidden_model # pyright: ignore[reportUnknownVariableType] # hidden parameter value has an incomplete type id = standard_logging_object.get("model_id") if model_group is None or id is None or model is None: return @@ -298,7 +309,13 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): ) @with_service_target("router_usage") - 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, # pyright: ignore[reportUnknownParameterType] # CustomLogger callback signature is untyped + ) -> None: if is_batch_retrieve_call_type(kwargs.get("call_type")): return try: @@ -309,9 +326,11 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): if standard_logging_object is None: 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"] + deployment_model: Final[object] = get_litellm_metadata_from_kwargs(kwargs).get("deployment") # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType, reportUnknownVariableType] # helper has no typed metadata signature + hidden_model: Final[str | None] = standard_logging_object["hidden_params"].get("litellm_model_name") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # standard logging payload hidden parameters are untyped + model: Final[str | None] = (deployment_model if isinstance(deployment_model, str) else None) or hidden_model # pyright: ignore[reportUnknownVariableType] # hidden parameter value has an incomplete type id = standard_logging_object.get("model_id") - if model_group is None or id is None: + if model_group is None or id is None or model is None: return elif isinstance(id, int): id = str(id) diff --git a/tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py b/tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py index f04ade1ce2d..c45c61d772e 100644 --- a/tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py +++ b/tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py @@ -8,6 +8,7 @@ import io import os import time import traceback +from datetime import datetime, timezone from typing import Dict, Final, Iterator from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER @@ -99,6 +100,75 @@ def test_tpm_rpm_updated(): # test_tpm_rpm_updated() +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", [False, True]) +@pytest.mark.parametrize("metadata_key", ["litellm_metadata", "metadata"]) +@pytest.mark.parametrize( + ("deployment_model", "hidden_model"), + [ + ("openrouter/anthropic/claude-haiku-5.5", None), + ("openrouter/anthropic/claude-haiku-5.5", "claude-haiku-5.5"), + (None, "openrouter/anthropic/claude-haiku-5.5"), + (None, None), + ], +) +async def test_success_tokens_use_routing_deployment_key( + sync_mode: bool, metadata_key: str, deployment_model: str | None, hidden_model: str | None +) -> None: + from tests.unit.create_mock_standard_logging_payload import create_standard_logging_payload + + cache: Final = DualCache() + handler: Final = LowestTPMLoggingHandler(router_cache=cache) + model: Final = "openrouter/anthropic/claude-haiku-5.5" + deployment: Final[DeploymentTypedDict] = { + "model_name": "messages", + "litellm_params": {"model": model, "tpm": 24}, + "model_info": {"id": "streaming-deployment"}, + } + original_payload: Final = create_standard_logging_payload() + payload: Final[StandardLoggingPayload] = { + **original_payload, + "model_group": "messages", + "model_id": "streaming-deployment", + "total_tokens": 25, + "hidden_params": {**original_payload["hidden_params"], "litellm_model_name": hidden_model}, + } + kwargs: Final = { + "call_type": "anthropic_messages", + "stream": True, + "litellm_params": {metadata_key: {"deployment": deployment_model}}, + "standard_logging_object": payload, + } + now: Final = datetime(2026, 1, 1, 19, 22, tzinfo=timezone.utc) + + def cached_value(key: str) -> object | None: + return cache.get_cache(key) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # DualCache cache reads are untyped + + with patch("litellm.router_strategy.lowest_tpm_rpm_v2.get_utc_datetime", return_value=now): + if sync_mode: + handler.log_success_event(kwargs, {}, now, now) # pyright: ignore[reportUnknownMemberType] # callback method is untyped + else: + await handler.async_log_success_event(kwargs, {}, now, now) # pyright: ignore[reportUnknownMemberType] # callback method is untyped + + assert cached_value("streaming-deployment:None:tpm:19-22") is None + assert cached_value("streaming-deployment:claude-haiku-5.5:tpm:19-22") is None + if deployment_model is None and hidden_model is None: + assert cached_value(f"streaming-deployment:{model}:tpm:19-22") is None + return + + assert cached_value(f"streaming-deployment:{model}:tpm:19-22") == 25 + if sync_mode: + with pytest.raises(ValueError, match="No deployments available"): + handler.get_available_deployments( # pyright: ignore[reportUnknownMemberType] # router method is untyped + "messages", [deployment] + ) + else: + with pytest.raises(litellm.RateLimitError): + await handler.async_get_available_deployments( # pyright: ignore[reportUnknownMemberType] # router method is untyped + "messages", [deployment] + ) + + @pytest.mark.usefixtures("restore_verbose_logger_level") def test_get_available_deployments(): from tests.unit.create_mock_standard_logging_payload import create_standard_logging_payload