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

fix(proxy): charge actual tokens to the per-customer key counter after the call, report customer limit edits as dirty fields

Browse source
This commit is contained in:
devin-ai-integration[bot] 2026-11-06 04:06:57 +00:00
parent b56fc011ee
commit 4d5f1dd144
6 changed files with 110 additions and 48 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

51
litellm/proxy/hooks/parallel_request_limiter_v3.py
View file

@ -761,6 +761,7 @@ class RequestRateLimiterStash:
reservation_released: bool = False
io_reservation_released: bool = False
tpm_limited_tags: frozenset[str] = field(default_factory=frozenset)
customer_per_key_tpm_scope: tuple[str, str] | None = None
socket_agent_windows: tuple[RateLimitDescriptor, ...] = ()
def clear_reserved_tokens(self) -> None:
@ -3219,13 +3220,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
)
def _add_customer_per_key_rate_limit_descriptor(
def _customer_per_key_rate_limit_descriptor(
self,
user_api_key_dict: UserAPIKeyAuth,
descriptors: list[RateLimitDescriptor],
) -> None:
) -> RateLimitDescriptor | None:
"""
Add a per-customer rpm/tpm descriptor scoped to the API key.
Per-customer rpm/tpm descriptor scoped to the API key.
The counter is ``{api_key}:{customer}``, so one customer exhausting its
allowance on this key never blocks the other customers on the key, nor
@ -3234,23 +3234,21 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
"""
customer: Final = user_api_key_dict.end_user_id
if not user_api_key_dict.api_key or not customer:
return
return None
rpm_limit: Final = get_key_customer_rpm_limit(user_api_key_dict).get(customer)
tpm_limit: Final = get_key_customer_tpm_limit(user_api_key_dict).get(customer)
if rpm_limit is None and tpm_limit is None:
return
return None
descriptors.append(
RateLimitDescriptor(
key="customer_per_key",
value=f"{user_api_key_dict.api_key}:{customer}",
rate_limit={
"requests_per_unit": rpm_limit,
"tokens_per_unit": tpm_limit,
"window_size": self.window_size,
},
)
return RateLimitDescriptor(
key="customer_per_key",
value=f"{user_api_key_dict.api_key}:{customer}",
rate_limit={
"requests_per_unit": rpm_limit,
"tokens_per_unit": tpm_limit,
"window_size": self.window_size,
},
)
def _add_mcp_per_key_rate_limit_descriptor(
@ -3671,10 +3669,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
descriptors=descriptors,
)
self._add_customer_per_key_rate_limit_descriptor(
user_api_key_dict=user_api_key_dict,
descriptors=descriptors,
)
customer_descriptor: Final = self._customer_per_key_rate_limit_descriptor(user_api_key_dict=user_api_key_dict)
if customer_descriptor is not None:
descriptors.append(customer_descriptor)
self._add_team_model_rate_limit_descriptor_from_metadata(
user_api_key_dict=user_api_key_dict,
@ -4500,6 +4497,16 @@ 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.customer_per_key_tpm_scope = next(
(
(d["key"], d["value"])
for d in descriptors
if d["key"] == "customer_per_key"
and d["rate_limit"] is not None
and d["rate_limit"].get("tokens_per_unit") is not None
),
None,
)
stash.socket_agent_windows = tuple(d for d in descriptors if d["key"] in _AGENT_DESCRIPTOR_KEYS)
# Only check rate limits if we have descriptors with actual limits
@ -5282,6 +5289,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
kwargs: object,
model_group: str | None,
tpm_limited_tags: Set[str] = frozenset(),
customer_per_key_scope: tuple[str, str] | None = None,
) -> list[tuple[str, str]]:
"""
Enumerate every (scope_key, scope_value) pair that *might* carry a
@ -5338,6 +5346,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
if session_id:
targets.append(("agent_session", f"{agent_id}:{session_id}"))
targets.extend(("tag", tag) for tag in sorted(tpm_limited_tags))
if customer_per_key_scope is not None:
targets.append(customer_per_key_scope)
return targets
def _build_reservation_aware_tpm_ops(
@ -5514,6 +5524,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
standard_logging_metadata=standard_logging_metadata,
kwargs=kwargs,
tpm_limited_tags=_tpm_limited_tags(stash, kwargs),
customer_per_key_scope=stash.customer_per_key_tpm_scope if stash is not None else None,
model_group=reconcile_model.group if reconcile_model is not None else None,
)
targets.extend(

54
tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py
View file

@ -10131,33 +10131,39 @@ async def test_per_customer_rpm_limit_on_key_is_scoped_to_that_customer_v3(monke
await call(other_key, "acme")
def test_per_customer_descriptor_carries_rpm_and_tpm_for_the_requesting_customer_v3():
"""
create_rate_limit_descriptors emits one customer_per_key descriptor keyed by
``{api_key}:{customer}`` with that customer's rpm and tpm, and none when the
requesting customer has no configured limit.
"""
api_key = hash_token("sk-per-customer-desc")
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache()))
metadata = {"customer_rpm_limit": {"acme": 5}, "customer_tpm_limit": {"acme": 900, "globex": 50}}
@pytest.mark.asyncio
async def test_per_customer_tpm_limit_on_key_charges_actual_usage_after_the_call_v3(monkeypatch):
monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60")
monkeypatch.setenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", "false")
_request_stash.set(None)
api_key = hash_token("sk-per-customer-tpm")
local_cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
monkeypatch.setattr(handler, "get_rate_limit_type", lambda: "total")
metadata = {"customer_tpm_limit": {"acme": 100}}
def descriptors_for(customer: str) -> list[RateLimitDescriptor]:
descriptors = handler.create_rate_limit_descriptors(
async def call(customer: str) -> None:
await handler.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key=api_key, end_user_id=customer, metadata=metadata),
cache=local_cache,
data={"model": "gpt-3.5-turbo", "user": customer},
rpm_limit_type=None,
tpm_limit_type=None,
model_has_failures=False,
call_type="",
)
return [d for d in descriptors if d["key"] == "customer_per_key"]
(acme,) = descriptors_for("acme")
assert acme["value"] == f"{api_key}:acme"
assert acme["rate_limit"]["requests_per_unit"] == 5
assert acme["rate_limit"]["tokens_per_unit"] == 900
await call("acme")
await handler.async_log_success_event(
kwargs={
"standard_logging_object": {"metadata": {"user_api_key_hash": api_key}},
"model": "gpt-3.5-turbo",
"user": "acme",
},
response_obj=ModelResponse(usage=Usage(prompt_tokens=60, completion_tokens=60, total_tokens=120)),
start_time=datetime.now(),
end_time=datetime.now(),
)
(globex,) = descriptors_for("globex")
assert globex["rate_limit"]["requests_per_unit"] is None
assert globex["rate_limit"]["tokens_per_unit"] == 50
assert descriptors_for("initech") == []
with pytest.raises(HTTPException) as exc_info:
await call("acme")
assert exc_info.value.status_code == 429
assert "customer_per_key" in str(exc_info.value.detail)
await call("globex")

21
ui/litellm-dashboard/src/components/key_team_helpers/CustomerRateLimitEditor.test.tsx
View file

@ -7,6 +7,7 @@ import {
CustomerRateLimitEntry,
customerLimitsToRows,
customerRowsToLimits,
customerLimitsEqual,
} from "./CustomerRateLimitEditor";
function Harness({
@ -101,3 +102,23 @@ describe("customerLimitsToRows", () => {
expect(customerLimitsToRows(undefined, null)).toStrictEqual([]);
});
});
describe("customerLimitsEqual", () => {
it("treats the same limits in a different row order as unchanged", () => {
expect(
customerLimitsEqual(
{ customer_rpm_limit: { acme: 2, globex: 5 }, customer_tpm_limit: {} },
{ customer_rpm_limit: { globex: 5, acme: 2 }, customer_tpm_limit: {} },
),
).toBe(true);
});
it("reports a changed, added or removed customer limit", () => {
const stored = { customer_rpm_limit: { acme: 2 }, customer_tpm_limit: { acme: 900 } };
expect(customerLimitsEqual(stored, { customer_rpm_limit: { acme: 3 }, customer_tpm_limit: { acme: 900 } })).toBe(
false,
);
expect(customerLimitsEqual(stored, { customer_rpm_limit: { acme: 2 }, customer_tpm_limit: {} })).toBe(false);
expect(customerLimitsEqual(stored, { ...stored, customer_rpm_limit: { acme: 2, globex: 1 } })).toBe(false);
});
});

7
ui/litellm-dashboard/src/components/key_team_helpers/CustomerRateLimitEditor.tsx
View file

@ -39,6 +39,13 @@ const toNumberMap = (raw: unknown): Record<string, number> => {
return out;
};
const sameLimitMap = (a: Record<string, number>, b: Record<string, number>): boolean =>
Object.keys(a).length === Object.keys(b).length &&
Object.entries(a).every(([customer, limit]) => b[customer] === limit);
export const customerLimitsEqual = (a: CustomerRateLimits, b: CustomerRateLimits): boolean =>
sameLimitMap(a.customer_rpm_limit, b.customer_rpm_limit) && sameLimitMap(a.customer_tpm_limit, b.customer_tpm_limit);
export const customerLimitsToRows = (
customerRpmLimit?: unknown,
customerTpmLimit?: unknown,

6
ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx
View file

@ -2592,6 +2592,8 @@ describe("KeyEditView", () => {
auto_rotate: false,
rotation_interval: undefined,
tag_rpm_limit: {},
customer_rpm_limit: {},
customer_tpm_limit: {},
};
describe("submit payload contract", () => {
@ -3171,6 +3173,7 @@ describe("KeyEditView", () => {
});
expect(onSubmitMock.mock.calls[0][0]).toHaveProperty("customer_rpm_limit", { acme: 7 });
expect(onSubmitMock.mock.calls[0][0]).toHaveProperty("customer_tpm_limit", { acme: 900 });
expect(onSubmitMock.mock.calls[0][1]).not.toContain("customer_rpm_limit");
});
it("clears a stored customer limit when its row is removed", async () => {
@ -3188,6 +3191,9 @@ describe("KeyEditView", () => {
expect(onSubmitMock).toHaveBeenCalled();
});
expect(onSubmitMock.mock.calls[0][0]).toHaveProperty("customer_rpm_limit", {});
expect(onSubmitMock.mock.calls[0][1]).toEqual(
expect.arrayContaining(["customer_rpm_limit", "customer_tpm_limit"]),
);
});
const setRpmLimit = (value: string) => {

19
ui/litellm-dashboard/src/components/templates/key_edit_view.tsx
View file

@ -68,6 +68,7 @@ import {
import {
CustomerRateLimitEditor,
CustomerRateLimitEntry,
customerLimitsEqual,
customerLimitsToRows,
customerRowsToLimits,
} from "../key_team_helpers/CustomerRateLimitEditor";
@ -303,9 +304,15 @@ function KeyEditSession({
// clears the stored limits ({} overwrites the metadata field).
const { tag_rpm_limit } = tagRowsToLimits(tagRateLimits);
values.tag_rpm_limit = tag_rpm_limit;
const { customer_rpm_limit, customer_tpm_limit } = customerRowsToLimits(customerRateLimits);
values.customer_rpm_limit = customer_rpm_limit;
values.customer_tpm_limit = customer_tpm_limit;
const customerLimits = customerRowsToLimits(customerRateLimits);
values.customer_rpm_limit = customerLimits.customer_rpm_limit;
values.customer_tpm_limit = customerLimits.customer_tpm_limit;
const customerLimitsUnchanged = customerLimitsEqual(
customerLimits,
customerRowsToLimits(
customerLimitsToRows(keyData.metadata?.customer_rpm_limit, keyData.metadata?.customer_tpm_limit),
),
);
const hadExistingFallbacks = keyData.budget_fallbacks != null && Object.keys(keyData.budget_fallbacks).length > 0;
if (Object.keys(budgetFallbacks).length > 0) {
@ -334,7 +341,11 @@ function KeyEditSession({
...values,
...(detachProject && enableProjectsUI && canDetachProject ? { project_id: null } : {}),
}),
[...Object.keys(dirtyFields), ...(budgetLimitsUnchanged ? [] : ["budget_limits"])],
[
...Object.keys(dirtyFields),
...(budgetLimitsUnchanged ? [] : ["budget_limits"]),
...(customerLimitsUnchanged ? [] : ["customer_rpm_limit", "customer_tpm_limit"]),
],
);
} finally {
setIsKeySaving(false);