scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-15 09:12:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 311

fix(router): use deployment metadata for streamed TPM

Browse source
This commit is contained in:
Devin AI 2026-11-04 20:09:13 +00:00
parent 3492453262
commit 3ef0365316
2 changed files with 95 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

31
litellm/router_strategy/lowest_tpm_rpm_v2.py
View file

@ -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)

70
tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py
View file

@ -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