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

fix(router): count TPM/RPM usage before building rate-limit headers

Browse source 
Router.make_call now increments the deployment TPM/RPM counter before set_response_headers reads remaining usage, so the headers carry post-increment values directly and the in-flight subtraction workaround from LIT-2719 is removed. deployment_callback_on_success adds only the tokens not yet counted (streams) and never a second request. The counter key uses the resolved deployment name so wildcard routes are read back correctly, and the proxy strips the router-owned counted-tokens marker from client metadata

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin 2026-09-16 18:45:01 +00:00
parent 78848d01a3
commit 6e5cdd8f12
7 changed files with 321 additions and 182 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/constants.py
View file

@ -1499,6 +1499,7 @@ OUTPUT_TOKEN_CEILING_PARAMS: Final = frozenset({"max_tokens", "max_completion_to
CLIENT_OUTPUT_CEILING_METADATA_KEY: Final = "_client_output_ceiling"
CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags"
ROUTING_REQUEST_TAGS_METADATA_KEY: Final = "_routing_request_tags"
ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: Final = "_litellm_router_usage_counted_tokens"
INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin"
SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated"
SESSION_ID_OMITTED_METADATA_KEY: Final = "litellm_session_id_omitted"

9
litellm/proxy/litellm_pre_call_utils.py
View file

@ -25,6 +25,7 @@ from litellm.constants import (
LITELLM_PROXY_MASTER_KEY_ALIAS,
OTEL_SERVICE_NAME_METADATA_KEYS,
PRE_CALL_EXECUTED_GUARDRAILS_KEY,
ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY,
ROUTING_REQUEST_TAGS_METADATA_KEY,
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
SESSION_ID_GENERATED_METADATA_KEY,
@ -336,7 +337,13 @@ _CLIENT_PRICING_METADATA_FIELDS: Final = frozenset({"model_info", "standard_logg
# 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(
{"attempted_fallbacks", "original_model_group", "request_retry_count", CLIENT_OUTPUT_CEILING_METADATA_KEY}
{
"attempted_fallbacks",
"original_model_group",
"request_retry_count",
CLIENT_OUTPUT_CEILING_METADATA_KEY,
ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY,
}
)
_ALLOW_CLIENT_PRICING_OVERRIDE_METADATA_KEY: Final = "allow_client_pricing_override"

236
litellm/router.py
View file

@ -63,6 +63,7 @@ from litellm.constants import (
DEFAULT_MAX_LRU_CACHE_SIZE,
INTERNAL_CALL_ORIGIN_METADATA_KEY,
OUTPUT_TOKEN_CEILING_PARAMS,
ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY,
ROUTING_REQUEST_TAGS_METADATA_KEY,
RUNTIME_UPDATABLE_ROUTER_SETTINGS,
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
@ -132,7 +133,7 @@ from litellm.router_utils.add_retry_fallback_headers import (
get_hidden_params_dict,
prepare_response_for_header_attachment,
replace_complexity_router_headers,
response_in_flight_token_count,
response_total_token_count,
)
from litellm.router_utils.auto_router_model_naming import (
AUTO_ROUTER_MODEL_PREFIX,
@ -215,6 +216,8 @@ from litellm.router_utils.reasoning_effort_capability import (
resolve_supported_reasoning_efforts,
)
from litellm.router_utils.router_callbacks.track_deployment_metrics import (
find_deployment_metadata,
get_counted_usage_tokens,
increment_deployment_failures_for_current_minute,
increment_deployment_successes_for_current_minute,
)
@ -7937,6 +7940,7 @@ class Router:
response = original_function(*args, **kwargs)
if coroutine_checker.is_async_callable(response) or inspect.isawaitable(response):
response = await response
await self.increment_deployment_usage_for_response(response=response, request_kwargs=kwargs)
## PROCESS RESPONSE HEADERS
response = await self.set_response_headers(response=response, model_group=model_group, request_kwargs=kwargs)
@ -8153,8 +8157,6 @@ class Router:
"""
Track remaining tpm/rpm quota for model in model_list
"""
from litellm.types.caching import RedisPipelineIncrementOperation
try:
# WS session wrappers fire with result=None; per-turn costs tracked by inner calls.
if kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime"):
@ -8162,114 +8164,132 @@ class Router:
standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None)
if standard_logging_object is None:
raise ValueError("standard_logging_object is None")
if kwargs["litellm_params"].get("metadata") is None:
pass
else:
deployment_name: Final = kwargs["litellm_params"]["metadata"].get(
"deployment", None
) # stable name - works for wildcard routes as well
# Get model_group and id from kwargs like the sync version does
model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None)
model_info: Final = kwargs["litellm_params"].get("model_info", {}) or {}
id = model_info.get("id", None)
if model_group is None or id is None:
return
elif isinstance(id, int):
id = str(id)
litellm_params: Final = kwargs["litellm_params"]
metadata: Final = litellm_params.get("metadata")
if metadata is None:
return
model_group: Final = metadata.get("model_group", None)
model_info: Final = litellm_params.get("model_info", {}) or {}
deployment_id: Final = model_info.get("id", None)
if model_group is None or deployment_id is None or self.get_deployment(model_id=str(deployment_id)) is None:
return
## get deployment info
deployment_info: Final = self.get_deployment(model_id=id)
# Always track deployment successes for cooldown logic, regardless of TPM/RPM limits
increment_deployment_successes_for_current_minute(
litellm_router_instance=self,
deployment_id=str(deployment_id),
)
if deployment_info is None:
return
else:
deployment_model_info: Final = self.get_router_model_info(
deployment=deployment_info,
received_model_name=model_group,
)
# get tpm/rpm from deployment info
tpm: Final = deployment_info.get("tpm", None)
rpm: Final = deployment_info.get("rpm", None)
## check tpm/rpm in litellm_params
tpm_litellm_params: Final = deployment_info.litellm_params.tpm
rpm_litellm_params: Final = deployment_info.litellm_params.rpm
## check tpm/rpm in model_info
tpm_model_info: Final = deployment_model_info.get("tpm", None)
rpm_model_info: Final = deployment_model_info.get("rpm", None)
# Always track deployment successes for cooldown logic, regardless of TPM/RPM limits
increment_deployment_successes_for_current_minute(
litellm_router_instance=self,
deployment_id=id,
)
deployment_dict = deployment_info if isinstance(deployment_info, dict) else deployment_info.model_dump()
has_io_token_limits: Final = deployment_has_io_token_limits(deployment_dict)
## Nothing to track only when neither tpm/rpm nor itpm/otpm limits are
## set. IO deployments still record TPM/RPM usage here so TPM-aware
## routing strategies see their real load in mixed model groups; their
## itpm/otpm enforcement runs separately in ModelRateLimitingCheck.
if (
tpm is None
and rpm is None
and tpm_litellm_params is None
and rpm_litellm_params is None
and tpm_model_info is None
and rpm_model_info is None
and not has_io_token_limits
):
return
parent_otel_span: Final = _get_parent_otel_span_from_kwargs(kwargs)
total_tokens: Final[float] = standard_logging_object.get("total_tokens", 0)
# ------------
# Setup values
# ------------
dt: Final = get_utc_datetime()
current_minute: Final = dt.strftime("%H-%M") # use the same timezone regardless of system clock
tpm_key = RouterCacheEnum.TPM.value.format(id=id, current_minute=current_minute, model=deployment_name)
# ------------
# Update usage
# ------------
# update cache
pipeline_operations: Final[list[RedisPipelineIncrementOperation]] = []
## TPM
pipeline_operations.append(
RedisPipelineIncrementOperation(
key=tpm_key,
increment_value=total_tokens,
ttl=RoutingArgs.ttl.value,
)
)
## RPM
rpm_key = RouterCacheEnum.RPM.value.format(id=id, current_minute=current_minute, model=deployment_name)
pipeline_operations.append(
RedisPipelineIncrementOperation(
key=rpm_key,
increment_value=1,
ttl=RoutingArgs.ttl.value,
)
)
await self.cache.async_increment_cache_pipeline(
increment_list=pipeline_operations,
parent_otel_span=parent_otel_span,
)
return tpm_key
total_tokens: Final[float] = standard_logging_object.get("total_tokens", 0)
counted_tokens: Final = get_counted_usage_tokens(litellm_params)
deployment_name: Final = metadata.get("deployment", None)
return 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 if counted_tokens is None else max(0, total_tokens - counted_tokens),
rpm_increment=1 if counted_tokens is None else 0,
parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs),
)
except Exception as e:
verbose_router_logger.debug(
"litellm.router.Router::deployment_callback_on_success(): Exception occured - %s", e
)
async def increment_deployment_usage_for_response(
self,
response: object,
request_kwargs: dict[str, object],
) -> None:
"""Count the request before the headers are read; the success callback adds only what is still missing"""
if response is None:
return
try:
deployment_metadata: Final = find_deployment_metadata(request_kwargs)
model_group: Final = request_kwargs.get("model")
if deployment_metadata is None or not isinstance(model_group, str):
return
model_info: Final = deployment_metadata["model_info"]
deployment_id: Final = model_info.get("id") if isinstance(model_info, dict) else None
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,
rpm_increment=1,
parent_otel_span=_get_parent_otel_span_from_kwargs(request_kwargs),
)
except Exception:
deployment_metadata.pop(ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, None)
raise
except Exception as e:
verbose_router_logger.debug(
"litellm.router.Router::increment_deployment_usage_for_response(): Exception occured - %s", e
)
async def _increment_deployment_usage(
self,
*,
deployment_id: str,
deployment_name: str | None,
model_group: str,
total_tokens: float,
rpm_increment: int,
parent_otel_span: Span | None,
) -> str | None:
from litellm.types.caching import RedisPipelineIncrementOperation
deployment_info: Final = self.get_deployment(model_id=deployment_id)
if deployment_info is None:
return None
deployment_model_info: Final = self.get_router_model_info(
deployment=deployment_info,
received_model_name=model_group,
)
configured_limits: Final = (
deployment_info.get("tpm", None),
deployment_info.get("rpm", None),
deployment_info.litellm_params.tpm,
deployment_info.litellm_params.rpm,
deployment_model_info.get("tpm", None),
deployment_model_info.get("rpm", None),
)
## Nothing to track only when neither tpm/rpm nor itpm/otpm limits are
## set. IO deployments still record TPM/RPM usage here so TPM-aware
## routing strategies see their real load in mixed model groups; their
## itpm/otpm enforcement runs separately in ModelRateLimitingCheck.
if all(limit is None for limit in configured_limits) and not deployment_has_io_token_limits(
deployment_info.model_dump()
):
return None
if total_tokens <= 0 and rpm_increment <= 0:
return None
current_minute: Final = get_utc_datetime().strftime("%H-%M") # use the same timezone regardless of system clock
tpm_key: Final = RouterCacheEnum.TPM.value.format(
id=deployment_id, current_minute=current_minute, model=deployment_name
)
rpm_key: Final = RouterCacheEnum.RPM.value.format(
id=deployment_id, current_minute=current_minute, model=deployment_name
)
pipeline_operations: Final[list[RedisPipelineIncrementOperation]] = [
RedisPipelineIncrementOperation(key=key, increment_value=increment_value, ttl=RoutingArgs.ttl.value)
for key, increment_value in ((tpm_key, total_tokens), (rpm_key, rpm_increment))
if increment_value > 0
]
await self.cache.async_increment_cache_pipeline(
increment_list=pipeline_operations,
parent_otel_span=parent_otel_span,
)
return tpm_key
def sync_deployment_callback_on_success(
self,
kwargs, # kwargs to completion
@ -11205,15 +11225,7 @@ class Router:
if model_group is not None:
remaining_usage: Final = await self.get_remaining_model_group_usage(model_group)
# get_remaining_model_group_usage reads the router's TPM/RPM counter,
# which is incremented post-response by deployment_callback_on_success.
# Replay the in-flight increment for TPM/RPM only (LIT-2719); ITPM/OTPM
# counters are incremented at reservation time and must not be adjusted.
apply_remaining_usage_headers(
additional_headers,
remaining_usage,
response_in_flight_token_count(response),
)
apply_remaining_usage_headers(additional_headers, remaining_usage)
return response
def _build_model_name_index(self, model_list: list) -> None:

9
litellm/router_utils/add_retry_fallback_headers.py
View file

@ -151,7 +151,7 @@ def apply_quality_router_decision_headers(
additional_headers[header] = str(decision[field])
def response_in_flight_token_count(response: object) -> int:
def response_total_token_count(response: object) -> int:
usage: Final = response.get("usage") if isinstance(response, dict) else getattr(response, "usage", None)
if usage is None:
return 0
@ -166,15 +166,10 @@ def response_in_flight_token_count(response: object) -> int:
def apply_remaining_usage_headers(
additional_headers: dict[str, object],
remaining_usage: dict[str, int],
in_flight_tokens: int,
) -> None:
in_flight_delta: Final = {
"x-ratelimit-remaining-tokens": in_flight_tokens,
"x-ratelimit-remaining-requests": 1,
}
for header, value in remaining_usage.items():
if value is not None and header not in additional_headers:
additional_headers[header] = value - in_flight_delta.get(header, 0)
additional_headers[header] = value
def _normalize_hidden_params(hidden_params: object) -> dict[str, object]:

23
litellm/router_utils/router_callbacks/track_deployment_metrics.py
View file

@ -9,8 +9,11 @@ get_deployment_failures_for_current_minute
get_deployment_successes_for_current_minute
"""
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final
from litellm.constants import ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY
if TYPE_CHECKING:
from litellm.router import Router as _Router
@ -18,6 +21,26 @@ if TYPE_CHECKING:
else:
LitellmRouter = Any
_METADATA_CHANNELS: Final = ("litellm_metadata", "metadata")
def find_deployment_metadata(kwargs: Mapping[str, object]) -> dict[str, object] | None:
buckets: Final = (kwargs.get(channel) for channel in _METADATA_CHANNELS)
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:
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
),
None,
)
return counted if isinstance(counted, int) and not isinstance(counted, bool) else None
def increment_deployment_successes_for_current_minute(
litellm_router_instance: LitellmRouter,

217
tests/router_unit_tests/test_router_helper_utils.py
View file

@ -1,3 +1,4 @@
import asyncio
import json
import os
import traceback
@ -12,7 +13,7 @@ from unittest.mock import patch, MagicMock, AsyncMock
from create_mock_standard_logging_payload import create_standard_logging_payload
from litellm.types.utils import StandardLoggingPayload
from litellm.types.router import Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo
from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS
from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY
@pytest.fixture
@ -928,17 +929,10 @@ async def test_set_response_headers(model_list):
@pytest.mark.asyncio
async def test_set_response_headers_subtracts_in_flight_delta(model_list):
async def test_set_response_headers_passes_through_post_increment_counters(model_list):
"""
LIT-2719: router-derived `x-ratelimit-remaining-*` headers must be
post-decrement (match OpenAI/Anthropic vendor semantics) so the proxy's
HTTP response headers and the prometheus gauges that read them stay
comparable across providers.
Router's TPM/RPM counter is incremented post-response by
`deployment_callback_on_success`, so `get_remaining_model_group_usage`
sees pre-decrement values. `set_response_headers` must replay the
in-flight increment before writing the headers.
LIT-3058: `make_call` increments the router's TPM/RPM counter before the headers
are built, so `set_response_headers` writes the remaining values it reads as-is.
"""
from pydantic import BaseModel
@ -952,49 +946,10 @@ async def test_set_response_headers_subtracts_in_flight_delta(model_list):
router = Router(model_list=model_list)
router.get_remaining_model_group_usage = AsyncMock(
return_value={
"x-ratelimit-remaining-tokens": 1000,
"x-ratelimit-remaining-tokens": 958,
"x-ratelimit-limit-tokens": 1000,
"x-ratelimit-remaining-requests": 100,
"x-ratelimit-remaining-requests": 99,
"x-ratelimit-limit-requests": 100,
}
)
resp = _Resp()
resp._hidden_params = {}
await router.set_response_headers(response=resp, model_group="gpt-3.5-turbo")
headers = resp._hidden_params["additional_headers"]
assert headers["x-ratelimit-remaining-tokens"] == 958
assert headers["x-ratelimit-remaining-requests"] == 99
# Limit headers pass through unmodified.
assert headers["x-ratelimit-limit-tokens"] == 1000
assert headers["x-ratelimit-limit-requests"] == 100
@pytest.mark.asyncio
async def test_set_response_headers_in_flight_delta_only_adjusts_tpm_rpm(model_list):
"""
The in-flight replay applies only to the post-incremented TPM/RPM counters
(`x-ratelimit-remaining-tokens` / `-requests`). The ITPM/OTPM counters are
incremented at reservation time (pre-call), so the input/output token
headers already reflect this request and must pass through untouched.
"""
from pydantic import BaseModel
class _Usage(BaseModel):
total_tokens: int = 30
prompt_tokens: int = 20
completion_tokens: int = 10
class _Resp(BaseModel):
usage: _Usage = _Usage()
_hidden_params: dict = {}
router = Router(model_list=model_list)
router.get_remaining_model_group_usage = AsyncMock(
return_value={
"x-ratelimit-remaining-tokens": 1000,
"x-ratelimit-remaining-requests": 100,
"x-ratelimit-remaining-input-tokens": 1000,
"x-ratelimit-remaining-output-tokens": 500,
}
@ -1005,14 +960,155 @@ async def test_set_response_headers_in_flight_delta_only_adjusts_tpm_rpm(model_l
await router.set_response_headers(response=resp, model_group="gpt-3.5-turbo")
headers = resp._hidden_params["additional_headers"]
# TPM/RPM headers replay the in-flight increment...
assert headers["x-ratelimit-remaining-tokens"] == 970
assert headers["x-ratelimit-remaining-tokens"] == 958
assert headers["x-ratelimit-remaining-requests"] == 99
# ...but the reservation-based input/output headers pass through unchanged.
assert headers["x-ratelimit-limit-tokens"] == 1000
assert headers["x-ratelimit-limit-requests"] == 100
assert headers["x-ratelimit-remaining-input-tokens"] == 1000
assert headers["x-ratelimit-remaining-output-tokens"] == 500
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},
}
]
)
def _ratelimit_headers(response) -> dict:
return {k: v for k, v in response._hidden_params["additional_headers"].items() if k.startswith("x-ratelimit-")}
@pytest.mark.asyncio
async def test_acompletion_headers_read_post_increment_counter_and_count_once():
"""
LIT-3058 regression: the remaining-* headers on the response must already include
this request, and the request must land in the router counter exactly once even
though `deployment_callback_on_success` still runs after the response returns.
"""
router = _rpm_tpm_router("lit-3058-async")
response = await router.acompletion(
model="gpt-5-mini", messages=[{"role": "user", "content": "hi"}], mock_response="pong"
)
total_tokens = response.usage.total_tokens
assert total_tokens > 0
headers = _ratelimit_headers(response)
assert headers["x-ratelimit-remaining-tokens"] == 1000 - total_tokens
assert headers["x-ratelimit-remaining-requests"] == 99
assert await router.get_model_group_usage("gpt-5-mini") == (total_tokens, 1)
await asyncio.sleep(0.5)
assert await router.get_model_group_usage("gpt-5-mini") == (total_tokens, 1)
@pytest.mark.asyncio
async def test_acompletion_wildcard_route_headers_and_counter_use_resolved_deployment_name():
"""The counter key is written under the resolved model name, which is what the usage reader looks up."""
router = Router(
model_list=[
{
"model_name": "openai/*",
"litellm_params": {"model": "openai/*", "api_key": "sk-fake", "tpm": 1000, "rpm": 100},
"model_info": {"id": "lit-3058-wildcard"},
}
]
)
response = await router.acompletion(
model="openai/gpt-5-mini", messages=[{"role": "user", "content": "hi"}], mock_response="pong"
)
total_tokens = response.usage.total_tokens
headers = _ratelimit_headers(response)
assert headers["x-ratelimit-remaining-tokens"] == 1000 - total_tokens
assert headers["x-ratelimit-remaining-requests"] == 99
assert await router.get_model_group_usage("openai/gpt-5-mini") == (total_tokens, 1)
@pytest.mark.asyncio
async def test_acompletion_stream_counts_request_before_headers_and_tokens_once_on_completion():
"""
A stream has no usage when the headers are built: the request is counted before the
headers and the final token usage is added once when the stream completes.
"""
router = _rpm_tpm_router("lit-3058-stream")
stream = await router.acompletion(
model="gpt-5-mini",
messages=[{"role": "user", "content": "hi"}],
mock_response="pong pong pong",
stream=True,
stream_options={"include_usage": True},
)
headers = _ratelimit_headers(stream)
assert headers["x-ratelimit-remaining-tokens"] == 1000
assert headers["x-ratelimit-remaining-requests"] == 99
assert await router.get_model_group_usage("gpt-5-mini") == (None, 1)
chunks = [chunk async for chunk in stream]
total_tokens = chunks[-1].usage.total_tokens
assert total_tokens > 0
await asyncio.sleep(0.5)
assert await router.get_model_group_usage("gpt-5-mini") == (total_tokens, 1)
@pytest.mark.asyncio
async def test_deployment_callback_on_success_adds_only_uncounted_tokens():
"""
When the request was already counted before the headers, the success callback adds
only the tokens that were not known at that point and never a second request.
"""
import time
router = _rpm_tpm_router("lit-3058-callback")
standard_logging_payload = create_standard_logging_payload()
standard_logging_payload["total_tokens"] = 100
kwargs = {
"litellm_params": {
"metadata": {
"deployment": "gpt-5-mini",
"model_group": "gpt-5-mini",
ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: 60,
},
"model_info": {"id": "lit-3058-callback"},
},
"standard_logging_object": standard_logging_payload,
}
tpm_key = await router.deployment_callback_on_success(
kwargs=kwargs,
completion_response=litellm.ModelResponse(model="gpt-5-mini", usage={"total_tokens": 100}),
start_time=time.time(),
end_time=time.time(),
)
assert tpm_key is not None
assert await router.get_model_group_usage("gpt-5-mini") == (40, None)
@pytest.mark.asyncio
async def test_increment_deployment_usage_for_response_skips_session_wrappers():
"""WebSocket and realtime session wrappers return None and are not counted as a request."""
router = _rpm_tpm_router("lit-3058-ws")
request_kwargs = {
"model": "gpt-5-mini",
"litellm_metadata": {"model_group": "gpt-5-mini", "model_info": {"id": "lit-3058-ws"}},
}
await router.increment_deployment_usage_for_response(response=None, request_kwargs=request_kwargs)
assert await router.get_model_group_usage("gpt-5-mini") == (None, None)
assert ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY not in request_kwargs["litellm_metadata"]
@pytest.mark.asyncio
async def test_get_model_group_io_token_usage_sums_across_deployments():
"""
@ -1154,8 +1250,8 @@ async def test_set_response_headers_native_input_token_header_does_not_suppress_
await router.set_response_headers(response=resp, model_group="gpt-3.5-turbo")
headers = resp._hidden_params["additional_headers"]
assert headers["x-ratelimit-remaining-tokens"] == 958
assert headers["x-ratelimit-remaining-requests"] == 99
assert headers["x-ratelimit-remaining-tokens"] == 1000
assert headers["x-ratelimit-remaining-requests"] == 100
# the provider's native header is left untouched
assert headers["x-ratelimit-remaining-input-tokens"] == 5
@ -1187,7 +1283,7 @@ async def test_set_response_headers_native_token_header_does_not_suppress_io_hea
headers = resp._hidden_params["additional_headers"]
assert headers["x-ratelimit-remaining-tokens"] == 5
assert headers["x-ratelimit-remaining-requests"] == 99
assert headers["x-ratelimit-remaining-requests"] == 100
assert headers["x-ratelimit-remaining-input-tokens"] == 900
assert headers["x-ratelimit-remaining-output-tokens"] == 450
@ -1196,8 +1292,7 @@ async def test_set_response_headers_native_token_header_does_not_suppress_io_hea
async def test_set_response_headers_handles_missing_usage(model_list):
"""
Streaming chunks and some response shapes may lack a `usage` attribute or
populated `total_tokens`. The in-flight subtraction must default to 0
tokens (still subtract 1 from requests) and never raise.
populated `total_tokens`. Header composition must not depend on usage and never raise.
"""
from pydantic import BaseModel
@ -1218,7 +1313,7 @@ async def test_set_response_headers_handles_missing_usage(model_list):
headers = resp._hidden_params["additional_headers"]
assert headers["x-ratelimit-remaining-tokens"] == 1000
assert headers["x-ratelimit-remaining-requests"] == 99
assert headers["x-ratelimit-remaining-requests"] == 100
@pytest.mark.asyncio

8
tests/test_litellm/proxy/test_litellm_pre_call_utils.py
View file

@ -43,7 +43,11 @@ from litellm.litellm_core_utils.get_provider_specific_headers import (
from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
TRUSTED_CALLBACK_VARS_FIELD,
)
from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY, SESSION_ID_OMITTED_METADATA_KEY
from litellm.constants import (
ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY,
SESSION_ID_GENERATED_METADATA_KEY,
SESSION_ID_OMITTED_METADATA_KEY,
)
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.fireworks_ai.common_utils import get_fireworks_session_id
from litellm.types.utils import CredentialItem
@ -7354,6 +7358,7 @@ _PLANTED_STAMPS = {
"original_model_group": "spoofed-group",
"request_retry_count": -100,
"_client_output_ceiling": {"api_base": "https://attacker.example"},
ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: 10**9,
"client_key": "client_value",
}
@ -7386,6 +7391,7 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_from_bo
assert "original_model_group" not in updated["metadata"]
assert "_client_output_ceiling" not in updated["metadata"]
assert "request_retry_count" not in updated["metadata"]
assert ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY not in updated["metadata"]
assert updated["metadata"]["client_key"] == "client_value"