mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-03 07:23:33 +00:00
Merge acdcad9769 into c905cd065d
This commit is contained in:
yucheng-berri
GitHub
commit
1adb465da2
6 changed files with 867 additions and 14 deletions
|
|
@ -472,6 +472,25 @@ end
|
|||
return results
|
||||
"""
|
||||
|
||||
OPEN_WINDOW_TOKEN_INCREMENT_SCRIPT: Final = """
|
||||
-- KEYS: (window_key, counter_key) pairs; ARGV: (increment, ttl_seconds) pairs at the same indexes.
|
||||
-- A counter is incremented only while its scope's rate limit window is open. A pair whose window key is
|
||||
-- its own counter key has no window to wait for and is always incremented.
|
||||
local applied = 0
|
||||
for i = 1, #KEYS, 2 do
|
||||
local counter_key = KEYS[i + 1]
|
||||
if KEYS[i] == counter_key or redis.call('EXISTS', KEYS[i]) == 1 then
|
||||
local ttl_seconds = tonumber(ARGV[i + 1])
|
||||
redis.call('INCRBYFLOAT', counter_key, ARGV[i])
|
||||
if ttl_seconds > 0 and redis.call('TTL', counter_key) == -1 then
|
||||
redis.call('EXPIRE', counter_key, ttl_seconds)
|
||||
end
|
||||
applied = applied + 1
|
||||
end
|
||||
end
|
||||
return applied
|
||||
"""
|
||||
|
||||
# Redis cluster slot count
|
||||
REDIS_CLUSTER_SLOTS: Final = 16384
|
||||
REDIS_NODE_HASHTAG_NAME: Final = "all_keys"
|
||||
|
|
@ -622,6 +641,12 @@ class ReservationAwareIncrementOperation(RedisPipelineIncrementOperation):
|
|||
reservation_backend: NotRequired[Literal["redis", "local"]]
|
||||
|
||||
|
||||
class OpenWindowIncrementOperation(RedisPipelineIncrementOperation):
|
||||
"""A token increment Redis applies only while ``open_window_key`` holds its scope's rate limit window."""
|
||||
|
||||
open_window_key: str
|
||||
|
||||
|
||||
class RateLimitResponseWithDescriptors(TypedDict):
|
||||
descriptors: list[RateLimitDescriptor]
|
||||
response: RateLimitResponse
|
||||
|
|
@ -830,6 +855,11 @@ def _as_window_guarded(operation: RedisPipelineIncrementOperation) -> Reservatio
|
|||
)
|
||||
|
||||
|
||||
def _open_window_key(operation: RedisPipelineIncrementOperation) -> str | None:
|
||||
open_window_key: Final = operation.get("open_window_key")
|
||||
return open_window_key if isinstance(open_window_key, str) else None
|
||||
|
||||
|
||||
def _call_id_from_callback_kwargs(kwargs: object) -> str | None:
|
||||
if not isinstance(kwargs, dict):
|
||||
return None
|
||||
|
|
@ -850,6 +880,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
batch_rate_limiter_script: _AsyncLuaScript | None
|
||||
batch_counter_read_script: _AsyncLuaScript | None
|
||||
token_increment_script: _AsyncLuaScript | None
|
||||
open_window_token_increment_script: _AsyncLuaScript | None
|
||||
check_and_increment_by_n_script: _AsyncLuaScript | None
|
||||
window_guarded_token_increment_script: _AsyncLuaScript | None
|
||||
parallel_acquire_script: _AsyncLuaScript | None
|
||||
|
|
@ -885,6 +916,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
self.token_increment_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script(
|
||||
TOKEN_INCREMENT_SCRIPT
|
||||
)
|
||||
self.open_window_token_increment_script = (
|
||||
self.internal_usage_cache.dual_cache.redis_cache.async_register_script(
|
||||
OPEN_WINDOW_TOKEN_INCREMENT_SCRIPT
|
||||
)
|
||||
)
|
||||
self.check_and_increment_by_n_script = (
|
||||
self.internal_usage_cache.dual_cache.redis_cache.async_register_script(CHECK_AND_INCREMENT_BY_N_SCRIPT)
|
||||
)
|
||||
|
|
@ -906,6 +942,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
self.batch_rate_limiter_script = None
|
||||
self.batch_counter_read_script = None
|
||||
self.token_increment_script = None
|
||||
self.open_window_token_increment_script = None
|
||||
self.check_and_increment_by_n_script = None
|
||||
self.window_guarded_token_increment_script = None
|
||||
self.parallel_acquire_script = None
|
||||
|
|
@ -4589,8 +4626,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
"""
|
||||
Execute token increment script grouped by hash tag for cluster compatibility.
|
||||
"""
|
||||
if self.token_increment_script is None:
|
||||
token_increment_script: Final = self.token_increment_script
|
||||
if token_increment_script is None:
|
||||
return
|
||||
open_window_script: Final = self.open_window_token_increment_script
|
||||
|
||||
# Group operations by hash tag for Redis cluster compatibility
|
||||
operation_keys: Final = [op["key"] for op in pipeline_operations]
|
||||
|
|
@ -4599,6 +4638,26 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
for _hash_tag, group_keys in key_groups.items():
|
||||
# Get operations for this hash tag group
|
||||
group_operations = [op for op in pipeline_operations if op["key"] in group_keys]
|
||||
window_keys = [_open_window_key(op) for op in group_operations]
|
||||
if open_window_script is not None and any(window_key is not None for window_key in window_keys):
|
||||
# One script for the whole group, so its plain and window-gated increments land together
|
||||
await self._run_token_increment_script(
|
||||
OPEN_WINDOW_TOKEN_INCREMENT_SCRIPT,
|
||||
open_window_script,
|
||||
tuple(
|
||||
itertools.chain.from_iterable(
|
||||
(window_key or op["key"], op["key"])
|
||||
for op, window_key in zip(group_operations, window_keys, strict=True)
|
||||
)
|
||||
),
|
||||
tuple(
|
||||
itertools.chain.from_iterable(
|
||||
(op["increment_value"], op["ttl"] or 0) for op in group_operations
|
||||
)
|
||||
),
|
||||
group_operations,
|
||||
)
|
||||
continue
|
||||
|
||||
keys = []
|
||||
args = []
|
||||
|
|
@ -4616,25 +4675,34 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
keys.append(op["key"])
|
||||
args.extend([op["increment_value"], ttl_value])
|
||||
|
||||
if self._defer_token_increment_script(keys, args, group_operations):
|
||||
continue
|
||||
await self.token_increment_script(
|
||||
keys=keys,
|
||||
args=args,
|
||||
await self._run_token_increment_script(
|
||||
TOKEN_INCREMENT_SCRIPT, token_increment_script, keys, args, group_operations
|
||||
)
|
||||
|
||||
async def _run_token_increment_script(
|
||||
self,
|
||||
source: str,
|
||||
script: _AsyncLuaScript,
|
||||
keys: Sequence[str],
|
||||
args: Sequence[float],
|
||||
group_operations: Sequence[RedisPipelineIncrementOperation],
|
||||
) -> None:
|
||||
if not self._defer_token_increment_script(source, script, keys, args, group_operations):
|
||||
await script(keys=keys, args=args)
|
||||
|
||||
def _defer_token_increment_script(
|
||||
self,
|
||||
keys: list[str],
|
||||
args: list[int],
|
||||
group_operations: list["RedisPipelineIncrementOperation"],
|
||||
source: str,
|
||||
script: _AsyncLuaScript,
|
||||
keys: Sequence[str],
|
||||
args: Sequence[float],
|
||||
group_operations: Sequence[RedisPipelineIncrementOperation],
|
||||
) -> bool:
|
||||
"""Declared into the request's post-call pipeline instead of its own EVALSHA round trip; a failed
|
||||
script falls back to the plain increment pipeline for its own group, as the direct path does."""
|
||||
redis_cache: Final = self.internal_usage_cache.dual_cache.redis_cache
|
||||
script: Final = self.token_increment_script
|
||||
batch: Final = None if redis_cache is None else active_post_call_redis_batch(redis_cache)
|
||||
if batch is None or script is None:
|
||||
if batch is None:
|
||||
return False
|
||||
|
||||
async def fall_back(future: asyncio.Future[object]) -> None:
|
||||
|
|
@ -4647,10 +4715,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
future.exception(),
|
||||
)
|
||||
await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline(
|
||||
increment_list=group_operations,
|
||||
increment_list=[*group_operations],
|
||||
)
|
||||
|
||||
batch.script(TOKEN_INCREMENT_SCRIPT, script, keys, args).on_settled(fall_back)
|
||||
batch.script(source, script, keys, args).on_settled(fall_back)
|
||||
return True
|
||||
|
||||
@with_service_target("rate_limits")
|
||||
|
|
@ -5110,6 +5178,24 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
)
|
||||
return ops
|
||||
|
||||
def _build_open_window_tpm_ops(
|
||||
self,
|
||||
targets: Sequence[tuple[str, str]],
|
||||
tokens: int,
|
||||
) -> list[OpenWindowIncrementOperation]:
|
||||
"""Charge ``tokens`` to each target's counter while that scope has an open rate limit window."""
|
||||
if tokens == 0:
|
||||
return []
|
||||
return [
|
||||
OpenWindowIncrementOperation(
|
||||
key=self.create_rate_limit_keys(scope_key, scope_value, "tokens"),
|
||||
increment_value=tokens,
|
||||
ttl=self.window_size,
|
||||
open_window_key=f"{{{scope_key}:{scope_value}}}:window",
|
||||
)
|
||||
for scope_key, scope_value in targets
|
||||
]
|
||||
|
||||
def _build_project_reservation_op(
|
||||
self,
|
||||
scope: tuple[str, str],
|
||||
|
|
@ -5271,12 +5357,25 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
)
|
||||
pipeline_operations.extend(
|
||||
self._build_reservation_aware_tpm_ops(
|
||||
targets=charged_targets,
|
||||
targets=[target for target in charged_targets if target in reserved_scopes],
|
||||
reserved_scopes=reserved_scopes,
|
||||
actual_tokens=total_tokens,
|
||||
reserved_tokens=reserved_tokens,
|
||||
)
|
||||
)
|
||||
unreserved_targets: Final = [target for target in charged_targets if target not in reserved_scopes]
|
||||
# A call that never passed admission, such as a proxy sub-call, opened no window, and the sub-call
|
||||
# precheck reads its counters directly, so it charges them unconditionally
|
||||
pipeline_operations.extend(
|
||||
self._build_open_window_tpm_ops(targets=unreserved_targets, tokens=total_tokens)
|
||||
if stash is not None
|
||||
else self._build_reservation_aware_tpm_ops(
|
||||
targets=unreserved_targets,
|
||||
reserved_scopes=reserved_scopes,
|
||||
actual_tokens=total_tokens,
|
||||
reserved_tokens=0,
|
||||
)
|
||||
)
|
||||
pipeline_operations.extend(
|
||||
self._build_team_ptu_tpm_ops(
|
||||
standard_logging_metadata, # pyright: ignore[reportUnknownArgumentType] # untyped logging metadata
|
||||
|
|
|
|||
|
|
@ -140,6 +140,20 @@ def test_caller_without_the_summary_model_sees_the_edit_error(compacting: Rig) -
|
|||
assert len(compacting.backend_calls(received)) == 1
|
||||
|
||||
|
||||
def test_caller_over_the_summary_models_tpm_limit_sees_the_edit_error(compacting: Rig) -> None:
|
||||
with compacting.gateway.scenario() as scenario:
|
||||
key: Final = scenario.key(model_tpm_limit={CLASSIFIER_MODEL: 1})
|
||||
within_nonce: Final = nonce()
|
||||
compacting.register(within_nonce, chat_backend(tool_id(within_nonce), stream=False), summarizing_or_classifying)
|
||||
within: Final = applied_edit(json_body(compacting.post(long_request(within_nonce), key=key)))
|
||||
assert within["type"] == "compact_20260112" and "error" not in within, within
|
||||
over_nonce: Final = nonce()
|
||||
compacting.register(over_nonce, chat_backend(tool_id(over_nonce), stream=False), summarizing_or_classifying)
|
||||
over: Final = json_body(compacting.post(long_request(over_nonce), key=key))
|
||||
assert applied_edit(over)["error"] == "summary_model_rate_limit_exceeded", over
|
||||
assert compacting.classifier_calls(compacting.received(over_nonce)) == ()
|
||||
|
||||
|
||||
def assert_summarized_and_classified_on_their_models(rig: Rig, request_nonce: str, body: dict[str, JsonValue]) -> str:
|
||||
received: Final = rig.received(request_nonce)
|
||||
targets: Final = [request.target for request in received]
|
||||
|
|
|
|||
533
tests/integration/spend/test_rate_limit_token_counters.py
Normal file
533
tests/integration/spend/test_rate_limit_token_counters.py
Normal file
|
|
@ -0,0 +1,533 @@
|
|||
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 pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import anthropic
|
||||
import httpx
|
||||
import openai
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import Gateway, Scenario, eventually, held, object_value
|
||||
from integration._support.openai_wire import chat_reply, responses_reply
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.redis_process import owned_redis
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from pydantic import JsonValue
|
||||
from redis import Redis
|
||||
|
||||
from litellm.constants import LITELLM_PROXY_ADMIN_NAME
|
||||
|
||||
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_on_a_healthy_model",
|
||||
charged="user",
|
||||
key=lambda _model: {"tpm_limit": 100000, "tpm_limit_type": "dynamic"},
|
||||
user=lambda _model: {"rpm_limit": 1000},
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@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_open_window_of_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)
|
||||
direct: Final = chat(gateway, identity, fallback, prompt + " direct")
|
||||
assert direct.status_code == 200, direct.text
|
||||
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()) == 2
|
||||
|
||||
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"},
|
||||
2 * 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,
|
||||
)
|
||||
|
||||
|
||||
def test_numeric_body_user_is_charged_to_its_rate_limited_end_user_counter(gateway: Gateway) -> None:
|
||||
customer: Final = uuid.uuid4().int % 10**15
|
||||
with wire_server(completion) as upstream, gateway.scenario() as scenario, redis_counters() as cache:
|
||||
model: Final = scenario.model(api_base=upstream.url + "/v1")
|
||||
gateway.post("/customer/new", {"user_id": str(customer), "rpm_limit": 1000})
|
||||
scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [str(customer)]})
|
||||
identity: Final = replace(
|
||||
caller(scenario, Cell("numeric_end_user", charged="end_user"), model), body={"user": customer}
|
||||
)
|
||||
response: Final = chat(gateway, identity, model, f"token counters numeric user {uuid.uuid4().hex}")
|
||||
assert response.status_code == 200, response.text
|
||||
assert [json.loads(seen.body)["user"] for seen in upstream.drain()] == [customer]
|
||||
|
||||
assert_only_charged(
|
||||
cache,
|
||||
{**identity.scopes, "end_user": counter("end_user", str(customer))},
|
||||
{"end_user"},
|
||||
USAGE["total_tokens"],
|
||||
)
|
||||
|
||||
|
||||
USER_HEADER: Final = "x-integration-openwebui-user"
|
||||
|
||||
|
||||
def proxy_admin(scenario: Scenario, tpm_limit: int) -> str:
|
||||
return scenario.user(user_id=LITELLM_PROXY_ADMIN_NAME, user_role="proxy_admin", tpm_limit=tpm_limit)
|
||||
|
||||
|
||||
def internal_user(scenario: Scenario, tpm_limit: int) -> str:
|
||||
return scenario.user(user_role="internal_user", tpm_limit=tpm_limit)
|
||||
|
||||
|
||||
def master_key_request(scenario: Scenario, _user: str, body: Mapping[str, JsonValue]) -> httpx.Response:
|
||||
return scenario.gateway.request("POST", "/v1/chat/completions", body)
|
||||
|
||||
|
||||
def header_mapped_service_key_request(scenario: Scenario, user: str, body: Mapping[str, JsonValue]) -> httpx.Response:
|
||||
return scenario.gateway.request(
|
||||
"POST", "/v1/chat/completions", body, key=scenario.key(), headers={USER_HEADER: user}
|
||||
)
|
||||
|
||||
|
||||
def config_with_user_header_mapping(directory: Path) -> Path:
|
||||
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
config["general_settings"] = {
|
||||
**object_value(config["general_settings"]),
|
||||
"user_header_mappings": [{"header_name": USER_HEADER, "litellm_user_role": "internal_user"}],
|
||||
}
|
||||
path: Final = directory / "config.yaml"
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
return path
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"limited_user, unchecked_request",
|
||||
[(proxy_admin, master_key_request), (internal_user, header_mapped_service_key_request)],
|
||||
ids=["master_key", "user_header_mapping"],
|
||||
)
|
||||
def test_unchecked_traffic_counts_toward_the_user_limit_another_request_opened(
|
||||
gateway: Gateway,
|
||||
tmp_path: Path,
|
||||
limited_user: Callable[[Scenario, int], str],
|
||||
unchecked_request: Callable[[Scenario, str, Mapping[str, JsonValue]], httpx.Response],
|
||||
) -> None:
|
||||
with (
|
||||
wire_server(completion) as upstream,
|
||||
owned_redis(tmp_path / "redis") as private_redis,
|
||||
owned_proxy(
|
||||
gateway,
|
||||
tmp_path,
|
||||
{"REDIS_HOST": private_redis.host, "REDIS_PORT": str(private_redis.port)},
|
||||
config=config_with_user_header_mapping(tmp_path),
|
||||
) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
Redis(host=private_redis.host, port=private_redis.port) as cache,
|
||||
):
|
||||
model: Final = scenario.model(api_base=upstream.url + "/v1")
|
||||
user: Final = limited_user(scenario, 100000)
|
||||
user_counter: Final = {"user": counter("user", user)}
|
||||
|
||||
checked: Final = candidate.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "messages": [{"role": "user", "content": f"checked {uuid.uuid4().hex}"}]},
|
||||
key=scenario.key(user_id=user),
|
||||
)
|
||||
assert checked.status_code == 200, checked.text
|
||||
eventually(lambda: read_counters(cache, user_counter)["user"], lambda tokens: tokens == USAGE["total_tokens"])
|
||||
|
||||
unchecked: Final = unchecked_request(
|
||||
scenario, user, {"model": model, "messages": [{"role": "user", "content": f"unchecked {uuid.uuid4().hex}"}]}
|
||||
)
|
||||
assert unchecked.status_code == 200, unchecked.text
|
||||
assert len(upstream.drain()) == 2
|
||||
|
||||
eventually(
|
||||
lambda: read_counters(cache, user_counter)["user"],
|
||||
lambda tokens: tokens == 2 * USAGE["total_tokens"],
|
||||
seconds=30,
|
||||
)
|
||||
|
|
@ -7,6 +7,7 @@ from redis import Redis
|
|||
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
from litellm.proxy.hooks.parallel_request_limiter_v3 import OpenWindowIncrementOperation
|
||||
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
|
||||
PROXY_MaxParallelRequestsHandler_v3 as _PROXY_MaxParallelRequestsHandler,
|
||||
)
|
||||
|
|
@ -63,3 +64,43 @@ async def test_async_increment_tokens_with_ttl_preservation() -> None:
|
|||
finally:
|
||||
await redis_cache.async_delete_cache(key_with_ttl)
|
||||
await redis_cache.async_delete_cache(key_without_ttl)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_open_window_token_increment_charges_only_scopes_whose_window_is_open() -> None:
|
||||
redis_host: Final = os.environ["REDIS_HOST"]
|
||||
redis_port: Final = int(os.environ["REDIS_PORT"])
|
||||
redis_cache: Final = RedisCache(host=redis_host, port=redis_port)
|
||||
handler: Final = _PROXY_MaxParallelRequestsHandler(
|
||||
internal_usage_cache=InternalUsageCache(DualCache(redis_cache=redis_cache))
|
||||
)
|
||||
suffix: Final = uuid.uuid4().hex[:8]
|
||||
limited: Final = f"{{user:limited-{suffix}}}"
|
||||
unlimited: Final = f"{{user:unlimited-{suffix}}}"
|
||||
ungated: Final = f"{{team:ungated-{suffix}}}"
|
||||
|
||||
with Redis(host=redis_host, port=redis_port) as raw:
|
||||
try:
|
||||
assert raw.set(f"{limited}:window", "1", ex=60) is True
|
||||
assert raw.set(f"{limited}:tokens", "5", ex=60) is True
|
||||
await handler.async_increment_tokens_with_ttl_preservation(
|
||||
pipeline_operations=[
|
||||
RedisPipelineIncrementOperation(key=f"{ungated}:tokens", increment_value=20, ttl=60),
|
||||
*(
|
||||
OpenWindowIncrementOperation(
|
||||
key=f"{scope}:tokens", increment_value=20, ttl=60, open_window_key=f"{scope}:window"
|
||||
)
|
||||
for scope in (limited, unlimited)
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
assert raw.mget([f"{limited}:tokens", f"{unlimited}:tokens", f"{unlimited}:window"]) == [b"25", None, None]
|
||||
assert raw.get(f"{ungated}:tokens") == b"20"
|
||||
assert 0 < raw.ttl(f"{limited}:tokens") <= 60
|
||||
assert 0 < raw.ttl(f"{ungated}:tokens") <= 60
|
||||
finally:
|
||||
raw.delete(
|
||||
*(f"{scope}:{suffix_key}" for scope in (limited, unlimited) for suffix_key in ("window", "tokens")),
|
||||
f"{ungated}:tokens",
|
||||
)
|
||||
|
|
|
|||
|
|
@ -26,8 +26,10 @@ from litellm.caching.redis_batch import (
|
|||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging
|
||||
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
|
||||
OPEN_WINDOW_TOKEN_INCREMENT_SCRIPT,
|
||||
PARALLEL_RELEASE_SCRIPT,
|
||||
TOKEN_INCREMENT_SCRIPT,
|
||||
OpenWindowIncrementOperation,
|
||||
ParallelSlotAcquisition,
|
||||
PROXY_MaxParallelRequestsHandler_v3,
|
||||
RequestRateLimiterStash,
|
||||
|
|
@ -227,6 +229,67 @@ async def test_a_failed_operation_fails_only_its_owner_and_the_owner_applies_its
|
|||
assert await limiter.internal_usage_cache.dual_cache.in_memory_cache.async_get_cache("{team:t1}:tokens") is None
|
||||
|
||||
|
||||
def _open_window_op(scope: str) -> OpenWindowIncrementOperation:
|
||||
return OpenWindowIncrementOperation(
|
||||
key=scope + ":tokens", increment_value=10, ttl=60, open_window_key=scope + ":window"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_group_with_open_window_charges_rides_the_post_call_pipeline_as_one_script():
|
||||
client = FakeClient(_ok_replies)
|
||||
redis_cache = PostCallFakeRedisCache(client)
|
||||
limiter = _limiter(redis_cache)
|
||||
|
||||
with request_redis_batch_scope():
|
||||
await limiter.async_increment_tokens_with_ttl_preservation(
|
||||
[*_token_ops("{team:t1}:tokens"), _open_window_op("{api_key:k1}")]
|
||||
)
|
||||
await limiter.async_increment_tokens_with_ttl_preservation(_token_ops("{user:u1}:tokens"))
|
||||
await flush_post_call_redis_batches()
|
||||
|
||||
assert len(client.pipelines) == 1
|
||||
assert [command for command in client.pipelines[0].commands if command[0] == "EVALSHA"] == [
|
||||
(
|
||||
"EVALSHA",
|
||||
sha_of(OPEN_WINDOW_TOKEN_INCREMENT_SCRIPT),
|
||||
4,
|
||||
"{team:t1}:tokens",
|
||||
"{team:t1}:tokens",
|
||||
"{api_key:k1}:window",
|
||||
"{api_key:k1}:tokens",
|
||||
10,
|
||||
60,
|
||||
10,
|
||||
60,
|
||||
),
|
||||
("EVALSHA", sha_of(TOKEN_INCREMENT_SCRIPT), 1, "{user:u1}:tokens", 10, 60),
|
||||
]
|
||||
assert redis_cache.alone == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_failed_open_window_script_falls_back_to_the_plain_increment():
|
||||
def replies(command: tuple[object, ...]) -> object:
|
||||
if command[0] == "EVALSHA" and command[1] == sha_of(OPEN_WINDOW_TOKEN_INCREMENT_SCRIPT):
|
||||
return Exception("ERR Lua")
|
||||
return _ok_replies(command)
|
||||
|
||||
client = FakeClient(replies)
|
||||
redis_cache = PostCallFakeRedisCache(client)
|
||||
limiter = _limiter(redis_cache)
|
||||
|
||||
with request_redis_batch_scope():
|
||||
await limiter.async_increment_tokens_with_ttl_preservation(
|
||||
[*_token_ops("{team:t1}:tokens"), _open_window_op("{user:u1}")]
|
||||
)
|
||||
await flush_post_call_redis_batches()
|
||||
|
||||
assert redis_cache.alone == [("INCRBYFLOAT", "{team:t1}:tokens", 10), ("INCRBYFLOAT", "{user:u1}:tokens", 10)]
|
||||
assert await limiter.internal_usage_cache.dual_cache.in_memory_cache.async_get_cache("{team:t1}:tokens") == 10
|
||||
assert await limiter.internal_usage_cache.dual_cache.in_memory_cache.async_get_cache("{user:u1}:tokens") == 10
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_failed_slot_release_script_releases_the_slot_in_memory():
|
||||
def replies(command: tuple[object, ...]) -> object:
|
||||
|
|
|
|||
|
|
@ -9067,3 +9067,106 @@ async def test_a_successful_extraction_counts_its_input_tokens_toward_the_key_tp
|
|||
call_type="aextraction",
|
||||
)
|
||||
assert refused.value.status_code == 429
|
||||
|
||||
|
||||
_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(**body: object) -> dict[str, object]:
|
||||
return {
|
||||
"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": "gpt-4o-mini",
|
||||
"user_api_key_metadata": {},
|
||||
"user_api_key_team_metadata": {},
|
||||
}
|
||||
},
|
||||
"model": "gpt-4o-mini",
|
||||
**body,
|
||||
}
|
||||
|
||||
|
||||
async def _admit_scoped_request(handler: _PROXY_MaxParallelRequestsHandler, auth: UserAPIKeyAuth) -> None:
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=auth,
|
||||
cache=handler.internal_usage_cache.dual_cache,
|
||||
data={"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hello"}], "max_tokens": 20},
|
||||
call_type="acompletion",
|
||||
)
|
||||
|
||||
|
||||
def _charged_windows(handler: _PROXY_MaxParallelRequestsHandler, kwargs: dict[str, object]) -> dict[str, object]:
|
||||
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 {operation["key"]: operation.get("open_window_key") for operation in operations}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"limits, reserved",
|
||||
[
|
||||
({}, frozenset()),
|
||||
({"rpm_limit": 100}, frozenset()),
|
||||
({"user_rpm_limit": 100}, frozenset()),
|
||||
({"tpm_limit": 1000}, frozenset({"{api_key:key-hash}:tokens"})),
|
||||
({"team_tpm_limit": 1000}, frozenset({"{team:t}:tokens"})),
|
||||
],
|
||||
ids=["no_limits", "key_rpm", "user_rpm", "key_tpm_reserved", "team_tpm_reserved"],
|
||||
)
|
||||
async def test_success_charges_every_unreserved_scope_only_while_its_rate_limit_window_is_open(
|
||||
limits: dict[str, object], reserved: 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))
|
||||
|
||||
assert _charged_windows(handler, _scoped_success_kwargs()) == {
|
||||
key: None if key in reserved else key.removesuffix(":tokens") + ":window" for key in _SCOPE_TOKEN_KEYS
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_success_charges_a_numeric_body_user_to_the_window_its_admission_opened() -> None:
|
||||
handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache()))
|
||||
await _admit_scoped_request(
|
||||
handler, UserAPIKeyAuth(api_key="key-hash", end_user_id="123", end_user_rpm_limit=100)
|
||||
)
|
||||
|
||||
assert _charged_windows(handler, _scoped_success_kwargs(user=123))["{end_user:123}:tokens"] == "{end_user:123}:window"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("parent_admitted", [True, False], ids=["sub_call_of_an_admitted_request", "no_admission"])
|
||||
async def test_success_of_a_call_that_skipped_admission_charges_every_scope_without_waiting_for_a_window(
|
||||
parent_admitted: bool,
|
||||
) -> None:
|
||||
handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache()))
|
||||
if parent_admitted:
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="key-hash", user_id="u", team_id="t"),
|
||||
cache=handler.internal_usage_cache.dual_cache,
|
||||
data={
|
||||
"model": "gpt-4o-mini",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"litellm_call_id": "parent",
|
||||
},
|
||||
call_type="acompletion",
|
||||
)
|
||||
|
||||
assert _charged_windows(handler, _scoped_success_kwargs(litellm_call_id="sub-call")) == dict.fromkeys(
|
||||
_SCOPE_TOKEN_KEYS
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue