From 8b631a31e0137350155b9844e6fdc191ee06e1e6 Mon Sep 17 00:00:00 2001 From: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 00:54:47 -0700 Subject: [PATCH] perf(rate-limiter): charge v3 token counters only for scopes admission rate limited The v3 success path wrote a {scope}:tokens counter for every key, user, team, team member, end user, organization and per-model scope on the request, including scopes with no rate limit that no request ever reads. The pre-call hook now records every (scope, value) it built a rate limit descriptor for, plus the model group it checked, on the request stash. The success path keeps only those targets, keeps the per-model targets when a different model group served the request (fallback), and keeps every target when no admission was recorded for the call. Resolves LIT-9875 --- .../hooks/parallel_request_limiter_v3.py | 52 ++- .../spend/test_rate_limit_token_counters.py | 422 ++++++++++++++++++ .../hooks/test_parallel_request_limiter_v3.py | 183 +++++++- 3 files changed, 654 insertions(+), 3 deletions(-) create mode 100644 tests/integration/spend/test_rate_limit_token_counters.py diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index d858eff3348..972e944f729 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -515,6 +515,9 @@ _AUDIO_BYTES_PER_TOKEN: Final = 1600 PROJECT_ITPM_DESCRIPTOR_KEY: Final = "model_per_project_itpm" PROJECT_OTPM_DESCRIPTOR_KEY: Final = "model_per_project_otpm" PTU_TEAM_DESCRIPTOR_KEY: Final = "model_per_team_ptu" +_MODEL_SCOPED_TPM_TARGETS: Final = frozenset( + {"model_per_key", "model_per_team", "model_per_organization", "model_per_project"} +) # How long an acquired slot counts toward the in-flight total before it is # considered leaked (worker crashed without any release callback firing) and # pruned. Also the longest request duration the gauge can track: a request @@ -665,6 +668,14 @@ class CallTypeRateLimiter(Protocol): ) -> Exception | str | dict[str, object] | None: ... +@dataclass(frozen=True, slots=True) +class AdmittedRateLimitScopes: + """Every ``(scope, value)`` the pre-call hook built a rate limit descriptor for, and the model group it checked.""" + + scopes: frozenset[tuple[str, str]] + model_group: str | None + + @dataclass(slots=True) class RequestRateLimiterStash: """ @@ -712,6 +723,7 @@ class RequestRateLimiterStash: batch_tpd_refund_ops: tuple[ReservationAwareIncrementOperation, ...] = () reservation_released: bool = False tpm_limited_tags: frozenset[str] = field(default_factory=frozenset) + admitted_scopes: AdmittedRateLimitScopes | None = None @dataclass(frozen=True, slots=True) @@ -4197,6 +4209,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): for d in descriptors if d["key"] == "tag" and d["rate_limit"] is not None and d["rate_limit"].get("tokens_per_unit") is not None ) + stash.admitted_scopes = self._admitted_scopes(descriptors, requested_model) # Only check rate limits if we have descriptors with actual limits if descriptors: @@ -5017,6 +5030,36 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): targets.extend(("tag", tag) for tag in sorted(tpm_limited_tags)) return targets + def _admitted_scopes( + self, descriptors: Sequence[RateLimitDescriptor], requested_model: str | None + ) -> AdmittedRateLimitScopes: + model: Final = self._rate_limited_model(requested_model) + return AdmittedRateLimitScopes( + scopes=frozenset((d["key"], d["value"]) for d in descriptors), + model_group=model.group if model is not None else None, + ) + + @staticmethod + def _rate_limited_targets( + targets: Sequence[tuple[str, str]], + admitted: AdmittedRateLimitScopes | None, + model_group: str | None, + ) -> Sequence[tuple[str, str]]: + """ + Keep the TPM targets admission built a rate limit descriptor for, plus + the per-model targets when tokens are settled against a model group + admission did not check. Every target stays when no admission was + recorded for the call. + """ + if admitted is None: + return targets + served_admitted_group: Final = model_group == admitted.model_group + return tuple( + target + for target in targets + if target in admitted.scopes or (not served_admitted_group and target[0] in _MODEL_SCOPED_TPM_TARGETS) + ) + def _build_reservation_aware_tpm_ops( self, targets: Sequence[tuple[str, str]], @@ -5197,10 +5240,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): for scope in sorted(reserved_scopes) if scope[0] in ("agent", "agent_session") and scope not in targets ) + limited_targets: Final = self._rate_limited_targets( + targets=targets, + admitted=stash.admitted_scopes if stash is not None else None, + model_group=reconcile_model.group if reconcile_model is not None else None, + ) charged_targets: Final = ( - [target for target in targets if target[0] != "model_per_team"] + [target for target in limited_targets if target[0] != "model_per_team"] if self._key_owns_model_tpm_limit_from_request_metadata(request_metadata, reconcile_model) - else targets + else limited_targets ) if reserved_tokens > 0 and total_tokens < reserved_tokens: verbose_proxy_logger.debug( diff --git a/tests/integration/spend/test_rate_limit_token_counters.py b/tests/integration/spend/test_rate_limit_token_counters.py new file mode 100644 index 00000000000..d51dbe6b428 --- /dev/null +++ b/tests/integration/spend/test_rate_limit_token_counters.py @@ -0,0 +1,422 @@ +import asyncio +import json +import os +import uuid +from collections.abc import Callable, Iterator, Mapping, Set +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from dataclasses import dataclass, replace +from hashlib import sha256 +from typing import Final + +import anthropic +import httpx +import openai +import pytest +from integration._support.client import Gateway, Scenario, eventually, held +from integration._support.openai_wire import chat_reply, responses_reply +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue +from redis import Redis + +USAGE: Final = {"prompt_tokens": 12, "completion_tokens": 8, "total_tokens": 20} +WIRE_TOTAL_TOKENS: Final = 8 +Fields = Callable[[str], Mapping[str, JsonValue]] + + +def completion(request: Request) -> Reply: + assert request.target == "/v1/chat/completions" + body: Final = json.loads(request.body) + return Reply( + body=json.dumps( + { + "id": "chatcmpl-" + uuid.uuid4().hex, + "object": "chat.completion", + "created": 1, + "model": body["model"], + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": USAGE, + } + ).encode() + ) + + +def outage(request: Request) -> Reply: + assert request.target == "/v1/chat/completions" + return Reply(status=500, body=json.dumps({"error": {"message": "synthetic outage"}}).encode()) + + +def openai_endpoints(request: Request) -> Reply: + body: Final = json.loads(request.body) + stream: Final = body.get("stream") is True + if request.target == "/v1/chat/completions": + return chat_reply("chatcmpl-" + uuid.uuid4().hex, body["model"], "ok", stream=stream) + assert request.target == "/v1/responses", request.target + return responses_reply("resp_" + uuid.uuid4().hex, body["model"], "ok", stream=stream) + + +def counter(scope: str, value: str) -> str: + return f"{{{scope}:{value}}}:tokens" + + +def read_counters(cache: Redis, keys: Mapping[str, str]) -> Mapping[str, float | None]: + values: Final = cache.mget(list(keys.values())) + return {name: None if raw is None else float(raw) for name, raw in zip(keys, values, strict=True)} + + +def no_fields(_model: str) -> Mapping[str, JsonValue]: + return {} + + +@dataclass(frozen=True, slots=True) +class Caller: + key: str + user: str + team: str + scopes: Mapping[str, str] + body: Mapping[str, JsonValue] + + +@dataclass(frozen=True, slots=True) +class Cell: + name: str + charged: str + also_charged: frozenset[str] = frozenset() + key: Fields = no_fields + user: Fields = no_fields + team: Fields = no_fields + organization: Fields | None = None + project: Fields | None = None + customer: Fields | None = None + tag: Fields | None = None + + +def caller(scenario: Scenario, cell: Cell, model: str) -> Caller: + gateway: Final = scenario.gateway + organization: Final = None if cell.organization is None else scenario.organization(**cell.organization(model)) + team: Final = scenario.team( + **cell.team(model), **({} if organization is None else {"organization_id": organization}) + ) + user: Final = scenario.user(user_role="internal_user", **cell.user(model)) + gateway.post("/team/member_add", {"team_id": team, "member": {"role": "user", "user_id": user}}) + project: Final = None if cell.project is None else scenario.project(team, **cell.project(model)) + key: Final = scenario.key( + user_id=user, team_id=team, **({} if project is None else {"project_id": project}), **cell.key(model) + ) + hashed: Final = sha256(key.encode()).hexdigest() + scopes: Final = { + "api_key": counter("api_key", hashed), + "user": counter("user", user), + "team": counter("team", team), + "team_member": counter("team_member", f"{team}:{user}"), + "model_per_key": counter("model_per_key", f"{hashed}:{model}"), + "model_per_team": counter("model_per_team", f"{team}:{model}"), + } + body: Final[dict[str, JsonValue]] = {} + if organization is not None: + scopes["organization"] = counter("organization", organization) + scopes["model_per_organization"] = counter("model_per_organization", f"{organization}:{model}") + if project is not None: + scopes["model_per_project"] = counter("model_per_project", f"{project}:{model}") + if cell.customer is not None: + customer: Final = f"integration-customer-{uuid.uuid4().hex}" + gateway.post("/customer/new", {"user_id": customer, **cell.customer(model)}) + scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [customer]}) + scopes["end_user"] = counter("end_user", customer) + body["user"] = customer + if cell.tag is not None: + tag: Final = f"integration-tag-{uuid.uuid4().hex}" + gateway.post("/tag/new", {"name": tag, **cell.tag(model)}) + scenario.cleanups.callback(gateway.request, "POST", "/tag/delete", {"name": tag}) + scopes["tag"] = counter("tag", tag) + body["metadata"] = {"tags": [tag]} + return Caller(key, user, team, scopes, body) + + +def key_rpm_only(_model: str) -> Mapping[str, JsonValue]: + return {"rpm_limit": 1000} + + +CELLS: Final = ( + Cell("organization_rpm_only", charged="organization", organization=lambda _model: {"rpm_limit": 1000}), + Cell( + "organization_model_rpm_only", + charged="model_per_organization", + organization=lambda model: {"model_rpm_limit": {model: 1000}}, + ), + Cell("key_rpm_only", charged="api_key", key=key_rpm_only), + Cell("user_rpm_only", charged="user", user=lambda _model: {"rpm_limit": 1000}), + Cell("team_rpm_only", charged="team", team=lambda _model: {"rpm_limit": 1000}), + Cell("team_member_rpm_only", charged="team_member", team=lambda _model: {"team_member_rpm_limit": 1000}), + Cell("key_model_rpm_only", charged="model_per_key", key=lambda model: {"model_rpm_limit": {model: 1000}}), + Cell( + "team_model_rpm_only", + charged="model_per_team", + also_charged=frozenset({"model_per_key"}), + team=lambda model: {"model_rpm_limit": {model: 1000}}, + ), + Cell( + "project_model_rpm_only", charged="model_per_project", project=lambda model: {"model_rpm_limit": {model: 1000}} + ), + Cell("end_user_rpm_only", charged="end_user", customer=lambda _model: {"rpm_limit": 1000}), + Cell("tag_tpm", charged="tag", tag=lambda _model: {"tpm_limit": 100000}), + Cell("key_tpm_reserved", charged="api_key", key=lambda _model: {"tpm_limit": 100000}), + Cell( + "key_dynamic_tpm", + charged="api_key", + key=lambda _model: {"tpm_limit": 100000, "tpm_limit_type": "dynamic"}, + ), +) + + +@contextmanager +def redis_counters() -> Iterator[Redis]: + with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as cache: + yield cache + + +def assert_only_charged(cache: Redis, scopes: Mapping[str, str], charged: Set[str], tokens: float) -> None: + counters: Final = eventually( + lambda: read_counters(cache, scopes), + lambda values: all(values[name] == tokens for name in charged), + seconds=30, + ) + assert {name: value for name, value in counters.items() if name not in charged} == { + name: None for name in scopes if name not in charged + } + + +def chat(gateway: Gateway, identity: Caller, model: str, prompt: str, **extra: JsonValue) -> httpx.Response: + return gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": prompt}], **identity.body, **extra}, + key=identity.key, + ) + + +@pytest.mark.parametrize("cell", CELLS, ids=[cell.name for cell in CELLS]) +def test_success_charges_token_counters_only_for_rate_limited_scopes(gateway: Gateway, cell: Cell) -> None: + prompt: Final = f"token counters {cell.name} {uuid.uuid4().hex}" + with wire_server(completion) as upstream, gateway.scenario() as scenario, redis_counters() as cache: + model: Final = scenario.model(api_base=upstream.url + "/v1") + identity: Final = caller(scenario, cell, model) + response: Final = chat(gateway, identity, model, prompt, max_tokens=500) + assert response.status_code == 200, response.text + assert response.json()["usage"]["total_tokens"] == USAGE["total_tokens"] + assert [json.loads(seen.body)["messages"] for seen in upstream.drain()] == [ + [{"role": "user", "content": prompt}] + ] + + assert_only_charged(cache, identity.scopes, {cell.charged, *cell.also_charged}, USAGE["total_tokens"]) + + +def gateway_url(gateway: Gateway) -> str: + return str(gateway.client.base_url).rstrip("/") + + +def openai_chat_sync(gateway: Gateway, key: str, model: str, prompt: str) -> int | None: + with ( + httpx.Client(timeout=30, trust_env=False) as transport, + openai.OpenAI( + api_key=key, base_url=gateway_url(gateway) + "/v1", max_retries=0, http_client=transport + ) as client, + ): + response: Final = client.chat.completions.create(model=model, messages=[{"role": "user", "content": prompt}]) + return None if response.usage is None else response.usage.total_tokens + + +async def openai_chat_stream_async(gateway: Gateway, key: str, model: str, prompt: str) -> int | None: + async with ( + httpx.AsyncClient(timeout=30, trust_env=False) as transport, + openai.AsyncOpenAI( + api_key=key, base_url=gateway_url(gateway) + "/v1", max_retries=0, http_client=transport + ) as client, + ): + stream: Final = await client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": prompt}], + stream=True, + stream_options={"include_usage": True}, + ) + usage: Final = [chunk.usage.total_tokens async for chunk in stream if chunk.usage is not None] + return usage[-1] if usage else None + + +def anthropic_messages_sync(gateway: Gateway, key: str, model: str, prompt: str) -> int | None: + with ( + httpx.Client(timeout=30, trust_env=False) as transport, + anthropic.Anthropic(api_key=key, base_url=gateway_url(gateway), max_retries=0, http_client=transport) as client, + ): + message: Final = client.messages.create( + model=model, max_tokens=50, messages=[{"role": "user", "content": prompt}] + ) + return message.usage.input_tokens + message.usage.output_tokens + + +async def anthropic_messages_stream_async(gateway: Gateway, key: str, model: str, prompt: str) -> int | None: + async with ( + httpx.AsyncClient(timeout=30, trust_env=False) as transport, + anthropic.AsyncAnthropic( + api_key=key, base_url=gateway_url(gateway), max_retries=0, http_client=transport + ) as client, + client.messages.stream(model=model, max_tokens=50, messages=[{"role": "user", "content": prompt}]) as stream, + ): + message: Final = await stream.get_final_message() + return message.usage.input_tokens + message.usage.output_tokens + + +def openai_responses_sync(gateway: Gateway, key: str, model: str, prompt: str) -> int | None: + with ( + httpx.Client(timeout=30, trust_env=False) as transport, + openai.OpenAI( + api_key=key, base_url=gateway_url(gateway) + "/v1", max_retries=0, http_client=transport + ) as client, + ): + response: Final = client.responses.create(model=model, input=prompt) + return None if response.usage is None else response.usage.total_tokens + + +async def openai_responses_stream_async(gateway: Gateway, key: str, model: str, prompt: str) -> int | None: + async with ( + httpx.AsyncClient(timeout=30, trust_env=False) as transport, + openai.AsyncOpenAI( + api_key=key, base_url=gateway_url(gateway) + "/v1", max_retries=0, http_client=transport + ) as client, + ): + stream: Final = await client.responses.create(model=model, input=prompt, stream=True) + usage: Final = [ + event.response.usage.total_tokens + async for event in stream + if event.type == "response.completed" and event.response.usage is not None + ] + return usage[-1] if usage else None + + +ENDPOINTS: Final[Mapping[str, Callable[[Gateway, str, str, str], int | None]]] = { + "chat_sync": openai_chat_sync, + "chat_stream_async": lambda *args: asyncio.run(openai_chat_stream_async(*args)), + "messages_sync": anthropic_messages_sync, + "messages_stream_async": lambda *args: asyncio.run(anthropic_messages_stream_async(*args)), + "responses_sync": openai_responses_sync, + "responses_stream_async": lambda *args: asyncio.run(openai_responses_stream_async(*args)), +} + + +@pytest.mark.parametrize("endpoint", ENDPOINTS) +def test_every_endpoint_charges_only_the_rate_limited_scope(gateway: Gateway, endpoint: str) -> None: + prompt: Final = f"token counters {endpoint} {uuid.uuid4().hex}" + with wire_server(openai_endpoints) as upstream, gateway.scenario() as scenario, redis_counters() as cache: + model: Final = scenario.model(api_base=upstream.url + "/v1") + identity: Final = caller(scenario, Cell(endpoint, charged="api_key", key=key_rpm_only), model) + + assert ENDPOINTS[endpoint](gateway, identity.key, model, prompt) == WIRE_TOTAL_TOKENS + assert len(upstream.drain()) == 1 + + assert_only_charged(cache, identity.scopes, {"api_key"}, WIRE_TOTAL_TOKENS) + + +def test_fallback_hop_charges_the_rate_limited_model_it_served(gateway: Gateway) -> None: + prompt: Final = f"token counters fallback {uuid.uuid4().hex}" + with ( + wire_server(outage) as primary_upstream, + wire_server(completion) as fallback_upstream, + gateway.scenario() as scenario, + redis_counters() as cache, + ): + primary: Final = scenario.model(api_base=primary_upstream.url + "/v1") + fallback: Final = scenario.model(api_base=fallback_upstream.url + "/v1") + cell: Final = Cell( + "fallback", + charged="model_per_key", + key=lambda _model: {"model_tpm_limit": {fallback: 100000}}, + ) + identity: Final = caller(scenario, cell, fallback) + response: Final = chat(gateway, identity, primary, prompt, fallbacks=[fallback]) + assert response.status_code == 200, response.text + assert response.headers["x-litellm-model-group"] == fallback + assert len(primary_upstream.drain()) == 1 + assert len(fallback_upstream.drain()) == 1 + + hashed: Final = sha256(identity.key.encode()).hexdigest() + scopes: Final = { + name: identity.scopes[name] for name in ("api_key", "user", "team", "team_member", "model_per_key") + } + assert_only_charged( + cache, + {**scopes, "primary_model_per_key": counter("model_per_key", f"{hashed}:{primary}")}, + {"model_per_key"}, + USAGE["total_tokens"], + ) + + +def same_key(_scenario: Scenario, identity: Caller) -> Caller: + return identity + + +def another_key_of_the_same_user(scenario: Scenario, identity: Caller) -> Caller: + return replace(identity, key=scenario.key(user_id=identity.user, team_id=identity.team)) + + +@pytest.mark.parametrize( + "cell, next_caller", + [ + (Cell("key_tpm", charged="api_key", key=lambda _model: {"tpm_limit": 21}), same_key), + ( + Cell("user_tpm_across_keys", charged="user", user=lambda _model: {"tpm_limit": 21}), + another_key_of_the_same_user, + ), + ], + ids=["key_tpm", "user_tpm_across_keys"], +) +def test_rate_limited_scope_rejects_once_charged_tokens_reach_its_limit( + gateway: Gateway, cell: Cell, next_caller: Callable[[Scenario, Caller], Caller] +) -> None: + with wire_server(completion) as upstream, gateway.scenario() as scenario, redis_counters() as cache: + model: Final = scenario.model(api_base=upstream.url + "/v1") + identity: Final = caller(scenario, cell, model) + first: Final = chat(gateway, identity, model, "hi", max_tokens=1) + assert first.status_code == 200, first.text + eventually( + lambda: read_counters(cache, identity.scopes)[cell.charged], + lambda tokens: tokens == USAGE["total_tokens"], + seconds=30, + ) + + rejected: Final = chat(gateway, next_caller(scenario, identity), model, "hi again", max_tokens=1) + assert rejected.status_code == 429, rejected.text + assert len(upstream.drain()) == 1 + assert_only_charged(cache, identity.scopes, {cell.charged}, USAGE["total_tokens"]) + + +def test_concurrent_burst_charges_every_request_to_the_rate_limited_scope_only(gateway: Gateway) -> None: + requests: Final = 20 + with wire_server(completion) as upstream, gateway.scenario() as scenario, redis_counters() as cache: + model: Final = scenario.model(api_base=upstream.url + "/v1") + identity: Final = caller(scenario, Cell("burst", charged="api_key", key=key_rpm_only), model) + with ThreadPoolExecutor(max_workers=requests) as pool: + statuses: Final = list( + pool.map( + lambda index: chat(gateway, identity, model, f"token counters burst {index} {uuid.uuid4().hex}"), + range(requests), + ) + ) + assert [response.status_code for response in statuses] == [200] * requests + assert len(upstream.drain()) == requests + + assert_only_charged(cache, identity.scopes, {"api_key"}, requests * USAGE["total_tokens"]) + + +def test_upstream_outage_charges_no_token_counter(gateway: Gateway) -> None: + with wire_server(outage) as upstream, gateway.scenario() as scenario, redis_counters() as cache: + model: Final = scenario.model(api_base=upstream.url + "/v1") + identity: Final = caller(scenario, Cell("outage", charged="api_key", key=key_rpm_only), model) + response: Final = chat(gateway, identity, model, f"token counters outage {uuid.uuid4().hex}") + assert response.status_code == 500, response.text + assert len(upstream.drain()) == 1 + + held( + lambda: read_counters(cache, identity.scopes), + lambda values: set(values.values()) == {None}, + holding=3, + ) diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index 7af747fc26d..82b5bab8362 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -6,7 +6,7 @@ import asyncio import logging import sys import time -from collections.abc import Iterator, Sequence +from collections.abc import Awaitable, Callable, Iterator, Sequence from contextlib import contextmanager from datetime import datetime, timedelta, timezone from typing import Any, Dict, Final, List, Optional @@ -8686,3 +8686,184 @@ async def test_a2a_url_target_owns_invocation_fee_and_request_limit( await _rpm_request(limiter, cache, auth, "a2a/cheap") assert denied.value.status_code == 429 assert "expensive" in str(denied.value.detail) + + +_SCOPE_TOKEN_KEYS: Final = frozenset( + { + "{api_key:key-hash}:tokens", + "{user:u}:tokens", + "{team:t}:tokens", + "{team_member:t:u}:tokens", + "{model_per_key:key-hash:gpt-4o-mini}:tokens", + "{model_per_team:t:gpt-4o-mini}:tokens", + } +) + + +def _scoped_success_kwargs( + model_group: str, key_metadata: Dict[str, object] | None = None, call_id: str | None = None +) -> Dict[str, object]: + return { + **({} if call_id is None else {"litellm_call_id": call_id}), + "standard_logging_object": { + "metadata": {"user_api_key_hash": "key-hash", "user_api_key_user_id": "u", "user_api_key_team_id": "t"} + }, + "litellm_params": { + "metadata": { + "model_group": model_group, + "user_api_key_metadata": key_metadata or {}, + "user_api_key_team_metadata": {}, + } + }, + "model": model_group, + } + + +async def _admit_scoped_request( + handler: _PROXY_MaxParallelRequestsHandler, auth: UserAPIKeyAuth, model: str, call_id: str | None = None +) -> None: + await handler.async_pre_call_hook( + user_api_key_dict=auth, + cache=handler.internal_usage_cache.dual_cache, + data={ + "model": model, + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 20, + **({} if call_id is None else {"litellm_call_id": call_id}), + }, + call_type="acompletion", + ) + + +def _charged_token_keys(handler: _PROXY_MaxParallelRequestsHandler, kwargs: Dict[str, object]) -> frozenset[str]: + operations: Final = handler._build_success_event_pipeline_operations( + kwargs=kwargs, + response_obj=ModelResponse(usage=Usage(prompt_tokens=2, completion_tokens=1, total_tokens=3)), + rate_limit_type="total", + ) + return frozenset(operation["key"] for operation in operations) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "limits, charged", + [ + ({}, frozenset()), + ({"rpm_limit": 100}, frozenset({"{api_key:key-hash}:tokens"})), + ({"max_parallel_requests": 10}, frozenset({"{api_key:key-hash}:tokens"})), + ({"tpm_limit": 1000}, frozenset({"{api_key:key-hash}:tokens"})), + ({"tpm_limit": 1000, "metadata": {"tpm_limit_type": "dynamic"}}, frozenset({"{api_key:key-hash}:tokens"})), + ({"user_rpm_limit": 100}, frozenset({"{user:u}:tokens"})), + ({"team_tpm_limit": 1000}, frozenset({"{team:t}:tokens"})), + ({"team_member_rpm_limit": 100}, frozenset({"{team_member:t:u}:tokens"})), + ( + {"metadata": {"model_rpm_limit": {"gpt-4o-mini": 100}}}, + frozenset({"{model_per_key:key-hash:gpt-4o-mini}:tokens"}), + ), + ], + ids=[ + "no_limits", + "key_rpm", + "key_max_parallel", + "key_tpm_reserved", + "key_dynamic_tpm", + "user_rpm", + "team_tpm_reserved", + "team_member_rpm", + "key_model_rpm", + ], +) +async def test_success_charges_tokens_only_to_scopes_admission_rate_limited( + limits: Dict[str, object], charged: frozenset[str] +) -> None: + handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + await _admit_scoped_request( + handler, UserAPIKeyAuth(api_key="key-hash", user_id="u", team_id="t", **limits), "gpt-4o-mini" + ) + + assert _charged_token_keys(handler, _scoped_success_kwargs("gpt-4o-mini")) == charged + + +@pytest.mark.asyncio +async def test_fallback_success_charges_the_model_scopes_of_the_group_that_served_it() -> None: + key_metadata: Final = {"model_tpm_limit": {"fallback-model": 1000}} + handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + await _admit_scoped_request( + handler, UserAPIKeyAuth(api_key="key-hash", user_id="u", team_id="t", metadata=key_metadata), "primary-model" + ) + + assert _charged_token_keys(handler, _scoped_success_kwargs("fallback-model", key_metadata)) == { + "{model_per_key:key-hash:fallback-model}:tokens" + } + + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "success_end_user, charged", + [ + (None, frozenset({"{end_user:admitted-customer}:tokens"})), + ("admitted-customer", frozenset({"{end_user:admitted-customer}:tokens"})), + ("other-customer", frozenset()), + ], + ids=["metadata_end_user", "same_body_user", "different_body_user"], +) +async def test_success_charges_the_end_user_admission_checked_and_no_other( + success_end_user: str | None, charged: frozenset[str] +) -> None: + handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + await _admit_scoped_request( + handler, + UserAPIKeyAuth(api_key="key-hash", end_user_id="admitted-customer", end_user_rpm_limit=100), + "gpt-4o-mini", + ) + kwargs: Final = _scoped_success_kwargs("gpt-4o-mini") + kwargs["standard_logging_object"] = { + "metadata": {"user_api_key_hash": "key-hash", "user_api_key_end_user_id": "admitted-customer"} + } + if success_end_user is not None: + kwargs["user"] = success_end_user + + assert _charged_token_keys(handler, kwargs) == charged + + +@pytest.mark.asyncio +async def test_success_on_an_aliased_model_charges_the_model_limit_of_the_resolved_group() -> None: + key_metadata: Final = {"model_rpm_limit": {"target": 100}} + handler: Final = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()), model_group_resolver=_resolve_alias_to_target + ) + await _admit_scoped_request( + handler, UserAPIKeyAuth(api_key="key-hash", user_id="u", team_id="t", metadata=key_metadata), "alias" + ) + + assert _charged_token_keys(handler, _scoped_success_kwargs("target", key_metadata)) == { + "{model_per_key:key-hash:target}:tokens" + } + +async def _no_stash(handler: _PROXY_MaxParallelRequestsHandler) -> str | None: + _request_stash.set(None) + return None + + +async def _nested_call(handler: _PROXY_MaxParallelRequestsHandler) -> str | None: + await _admit_scoped_request( + handler, UserAPIKeyAuth(api_key="key-hash", user_id="u", team_id="t"), "gpt-4o-mini", "proxy-request" + ) + return "nested-call" + + +async def _stash_without_admission(handler: _PROXY_MaxParallelRequestsHandler) -> str | None: + _request_stash.set(RequestRateLimiterStash()) + return None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("arrange", [_no_stash, _nested_call, _stash_without_admission]) +async def test_success_charges_every_scope_when_the_call_was_not_admitted( + arrange: Callable[[_PROXY_MaxParallelRequestsHandler], Awaitable[str | None]], +) -> None: + handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + call_id: Final = await arrange(handler) + + assert _charged_token_keys(handler, _scoped_success_kwargs("gpt-4o-mini", call_id=call_id)) == _SCOPE_TOKEN_KEYS