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

feat(budgets): add per-model tpm and rpm limit maps to budgets and customers (#50797)

Browse source 
* feat(budgets): add per-model tpm and rpm limit maps to budgets and customers

Budgets gain model_tpm_limit and model_rpm_limit (model name to whole number, 0 or more).
/budget/new, /budget/update, /budget/info and /customer/new carry them, and the auth
path loads a customer's maps onto UserAPIKeyAuth as end_user_model_tpm_limit and
end_user_model_rpm_limit so the limiter can enforce them in a follow-up.

* fix(budgets): keep org, team and project per-model limits in metadata; lock_timeout and type gate fixes

Organizations, teams and projects store model_tpm_limit / model_rpm_limit in their own metadata,
which the limiter reads, so the shared budget allowlist leaves those two fields to customers.
Migration starts with SET LOCAL lock_timeout. Mapping annotations and rebind-ok reasons satisfy
the type discipline gate. Customer info tests expect the two new budget keys.

* fix(organizations): keep stripping budget_id and flattening nested per-model maps on update

The update path pops every budget field from the organization row again, so a caller's budget_id
never changes the attached budget, and a nested litellm_budget_table still flattens the two maps
into org metadata. Drops the explanatory comment on CUSTOMER_ONLY_BUDGET_FIELDS.

* fix(types): keep basedpyright and ruff strict counts flat for the new budget fields

Reads the per-model maps off the end-user params with membership checks instead of .get on an
untyped dict, writes the inline customer budget through the client's typed jsonify_object, and
builds budget rows in handle_budget_for_entity with model_validate so the two Mapping fields never
receive object-typed kwargs.

* chore: drop unrelated reformatting from project endpoints

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Kerry Lu 2026-11-09 22:15:18 -08:00 • committed by GitHub
parent 07be67603c
commit 64d84e557e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
19 changed files with 371 additions and 43 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
litellm-proxy-extras/litellm_proxy_extras/migrations/20261110000000_add_budget_model_tpm_rpm_limit/migration.sql Normal file
View file

@ -0,0 +1,7 @@
SET LOCAL lock_timeout = '10s';
-- AlterTable
ALTER TABLE "LiteLLM_BudgetTable" ADD COLUMN IF NOT EXISTS "model_tpm_limit" JSONB;
-- AlterTable
ALTER TABLE "LiteLLM_BudgetTable" ADD COLUMN IF NOT EXISTS "model_rpm_limit" JSONB;

2
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -19,6 +19,8 @@ model LiteLLM_BudgetTable {
rpm_limit BigInt?
tpd_limit BigInt?
model_max_budget Json?
model_tpm_limit Json?
model_rpm_limit Json?
budget_duration String?
budget_reset_at DateTime?
allowed_models String[] @default([]) // per-member model scope; empty = inherit team models

2
litellm/models/__init__.py
View file

@ -5,6 +5,7 @@ Domain models for LiteLLM backend.
from litellm.models.access_group import LiteLLM_AccessGroupTable
from litellm.models.autorouter_session import LiteLLM_AutoRouterSession
from litellm.models.budget import (
CUSTOMER_ONLY_BUDGET_FIELDS,
LiteLLM_BudgetTable,
LiteLLM_BudgetTableFull,
LiteLLM_TeamMemberTable,
@ -38,6 +39,7 @@ from litellm.models.user import LiteLLM_UserTable
from litellm.models.verification_token import LiteLLM_VerificationToken
__all__ = [
"CUSTOMER_ONLY_BUDGET_FIELDS",
"CreateCredentialItem",
"CredentialBase",
"CredentialItem",

6
litellm/models/budget.py
View file

@ -5,6 +5,7 @@ Canonical definition for ``litellm_budgettable``. Re-exported from
``litellm.proxy._types`` for backwards compatibility.
"""
from collections.abc import Mapping
from datetime import datetime, timezone
from typing import Final
@ -29,6 +30,8 @@ class LiteLLM_BudgetTable(LiteLLMPydanticObjectBase):
rpm_limit: int | None = None
tpd_limit: int | None = None
model_max_budget: dict | None = None
model_tpm_limit: Mapping[str, int] | None = None
model_rpm_limit: Mapping[str, int] | None = None
budget_duration: str | None = None
allowed_models: list[str] | None = None # per-member model scope; empty = inherit team models
temp_budget_increase: float | None = None
@ -70,3 +73,6 @@ class LiteLLM_TeamMemberTable(LiteLLM_BudgetTable):
budget_id: str | None = None
model_config = ConfigDict(protected_namespaces=())
CUSTOMER_ONLY_BUDGET_FIELDS: Final = frozenset({"model_tpm_limit", "model_rpm_limit"})

9
litellm/proxy/_types.py
View file

@ -15,6 +15,7 @@ from pydantic import (
Field,
Json,
JsonValue,
NonNegativeInt,
PositiveInt,
TypeAdapter,
field_validator,
@ -1826,6 +1827,14 @@ class BudgetNewRequest(LiteLLMPydanticObjectBase):
default=None,
description="Max budget for each model (e.g. {'gpt-4o': {'max_budget': '0.0000001', 'budget_duration': '1d', 'tpm_limit': 1000, 'rpm_limit': 1000}})",
)
model_tpm_limit: Mapping[str, NonNegativeInt] | None = Field(
default=None,
description="Max tokens per minute per model for this budget id, shared by every key, team and customer on it (e.g. {'gpt-4o': 400000})",
)
model_rpm_limit: Mapping[str, NonNegativeInt] | None = Field(
default=None,
description="Max requests per minute per model for this budget id, shared by every key, team and customer on it (e.g. {'gpt-4o': 200})",
)
budget_reset_at: datetime | None = Field(
default=None,
description="Datetime when the budget is reset",

18
litellm/proxy/auth/user_api_key_auth.py
View file

@ -1313,6 +1313,12 @@ def _apply_budget_limits_to_end_user_params(
if budget_info.model_max_budget is not None:
end_user_params["end_user_model_max_budget"] = budget_info.model_max_budget
if budget_info.model_tpm_limit is not None:
end_user_params["end_user_model_tpm_limit"] = budget_info.model_tpm_limit # rebind-ok: in-place like siblings
if budget_info.model_rpm_limit is not None:
end_user_params["end_user_model_rpm_limit"] = budget_info.model_rpm_limit # rebind-ok: in-place like siblings
verbose_proxy_logger.debug("Applied budget limits to end user %s", end_user_id)
@ -1444,6 +1450,10 @@ def update_valid_token_with_end_user_params(valid_token: UserAPIKeyAuth, end_use
valid_token.allowed_model_region = end_user_params["allowed_model_region"]
if end_user_params.get("end_user_model_max_budget") is not None:
valid_token.end_user_model_max_budget = end_user_params["end_user_model_max_budget"]
if "end_user_model_tpm_limit" in end_user_params:
valid_token.end_user_model_tpm_limit = end_user_params["end_user_model_tpm_limit"] # rebind-ok: in-place
if "end_user_model_rpm_limit" in end_user_params:
valid_token.end_user_model_rpm_limit = end_user_params["end_user_model_rpm_limit"] # rebind-ok: in-place
return valid_token
@ -2985,6 +2995,10 @@ async def user_api_key_auth_builder(
valid_token.end_user_tpm_limit = end_user_params.get("end_user_tpm_limit")
valid_token.end_user_rpm_limit = end_user_params.get("end_user_rpm_limit")
valid_token.end_user_tpd_limit = end_user_params.get("end_user_tpd_limit")
if "end_user_model_tpm_limit" in end_user_params:
valid_token.end_user_model_tpm_limit = end_user_params["end_user_model_tpm_limit"]
if "end_user_model_rpm_limit" in end_user_params:
valid_token.end_user_model_rpm_limit = end_user_params["end_user_model_rpm_limit"]
valid_token.allowed_model_region = end_user_params.get("allowed_model_region")
return await validate_resolved_virtual_key(
@ -4134,6 +4148,10 @@ async def _apply_key_end_user_default_budget_to_token(
valid_token.end_user_tpd_limit = default_budget.tpd_limit
if not keep_token_limits or valid_token.end_user_model_max_budget is None:
valid_token.end_user_model_max_budget = default_budget.model_max_budget
if not keep_token_limits or valid_token.end_user_model_tpm_limit is None:
valid_token.end_user_model_tpm_limit = default_budget.model_tpm_limit # rebind-ok: in-place like siblings
if not keep_token_limits or valid_token.end_user_model_rpm_limit is None:
valid_token.end_user_model_rpm_limit = default_budget.model_rpm_limit # rebind-ok: in-place like siblings
async def _reserve_budget_after_common_checks(

6
litellm/proxy/management_endpoints/budget_management_endpoints.py
View file

@ -57,6 +57,8 @@ async def new_budget(
- rpm_limit: Optional[int] - The requests per minute limit for the budget.
- tpd_limit: Optional[int] - The tokens per day limit for the budget. Charged by batch submissions instead of tpm_limit/rpm_limit.
- model_max_budget: Optional[dict] - Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d", "tpm_limit": 100000, "rpm_limit": 100000}}
- model_tpm_limit: Optional[dict] - Max tokens per minute per model, shared by everything on this budget. Example: {"gpt-4o": 400000}
- model_rpm_limit: Optional[dict] - Max requests per minute per model, shared by everything on this budget. Example: {"gpt-4o": 200}
- budget_reset_at: Optional[datetime] - Datetime when the initial budget is reset. Default is now.
"""
from prisma.errors import UniqueViolationError
@ -141,6 +143,8 @@ async def update_budget(
- rpm_limit: Optional[int] - The requests per minute limit for the budget.
- tpd_limit: Optional[int] - The tokens per day limit for the budget. Charged by batch submissions instead of tpm_limit/rpm_limit.
- model_max_budget: Optional[dict] - Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d", "tpm_limit": 100000, "rpm_limit": 100000}}
- model_tpm_limit: Optional[dict] - Max tokens per minute per model, shared by everything on this budget. Example: {"gpt-4o": 400000}
- model_rpm_limit: Optional[dict] - Max requests per minute per model, shared by everything on this budget. Example: {"gpt-4o": 200}
- budget_reset_at: Optional[datetime] - Update the Datetime when the budget was last reset.
"""
from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client, user_api_key_cache
@ -287,6 +291,8 @@ async def budget_settings(
"max_budget": {"type": "Float"},
"soft_budget": {"type": "Float"},
"model_max_budget": {"type": "Object"},
"model_tpm_limit": {"type": "Object"},
"model_rpm_limit": {"type": "Object"},
}
return_val: Final = []

4
litellm/proxy/management_endpoints/customer_endpoints.py
View file

@ -359,6 +359,8 @@ async def new_end_user(
- rpm_limit: Optional[int] - [Not Implemented Yet] Specify rpm limit for a given customer (Requests per minute)
- tpd_limit: Optional[int] - Specify tpd limit for a given customer (Tokens per day). Batch submissions are charged against it instead of tpm_limit/rpm_limit
- model_max_budget: Optional[dict] - [Not Implemented Yet] Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d"}}
- model_tpm_limit: Optional[dict] - Max tokens per minute per model for this customer, shared across every key the customer is sent on. Stored on the customer's budget. Example: {"gpt-4o": 400000}
- model_rpm_limit: Optional[dict] - Max requests per minute per model for this customer. Stored on the customer's budget. Example: {"gpt-4o": 200}
- max_parallel_requests: Optional[int] - [Not Implemented Yet] Specify max parallel requests for a given customer.
- soft_budget: Optional[float] - [Not Implemented Yet] Get alerts when customer crosses given budget, doesn't block requests.
- spend: Optional[float] - Specify initial spend for a given customer.
@ -455,7 +457,7 @@ async def new_end_user(
try:
budget_record: Final = await _typed_table(BudgetRepository(prisma_client)).create(
data={
**_new_budget.model_dump(exclude_unset=True),
**prisma_client.jsonify_object(_new_budget.model_dump(exclude_unset=True)),
"created_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
"updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
}

7
litellm/proxy/management_endpoints/organization_endpoints.py
View file

@ -32,6 +32,7 @@ from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
from litellm.models.budget import CUSTOMER_ONLY_BUDGET_FIELDS
from litellm.proxy._types import *
from litellm.proxy.auth.auth_checks import (
can_user_call_model,
@ -319,7 +320,9 @@ _verify_org_access: Final = verify_org_access
_STR_OBJECT_DICT_ADAPTER: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True))
_BUDGET_SETTABLE_FIELDS: Final = frozenset(LiteLLM_BudgetTable.model_fields.keys()) - {"budget_id"}
_BUDGET_SETTABLE_FIELDS: Final = (
frozenset(LiteLLM_BudgetTable.model_fields.keys()) - {"budget_id"} - CUSTOMER_ONLY_BUDGET_FIELDS
)
_ORG_COLUMN_FIELDS: Final = frozenset({"organization_alias", "models"})
_ORG_METADATA_FIELDS: Final = tuple(
field for field in LiteLLM_ManagementEndpoint_MetadataFields if field not in _BUDGET_SETTABLE_FIELDS
@ -497,7 +500,7 @@ async def new_organization(
If none provided, create one based on provided values
"""
budget_params: Final = LiteLLM_BudgetTable.model_fields.keys()
budget_params: Final = LiteLLM_BudgetTable.model_fields.keys() - CUSTOMER_ONLY_BUDGET_FIELDS
# Only include Budget Params when creating an entry in litellm_budgettable
_json_data: Final = _STR_OBJECT_DICT_ADAPTER.validate_python(data.json(exclude_none=True))

4
litellm/proxy/management_helpers/utils.py
View file

@ -241,7 +241,7 @@ async def handle_budget_for_entity(
return data_budget_id
elif _budget_data:
# Create a new budget with the provided fields
budget_row: Final = LiteLLM_BudgetTable(**_budget_data)
budget_row: Final = LiteLLM_BudgetTable.model_validate(dict(_budget_data))
new_budget_data: Final = prisma_client.jsonify_object(budget_row.model_dump(exclude_none=True))
_budget: Final[_PrismaBudgetRecord] = await BudgetRepository(prisma_client).table.create(
@ -262,7 +262,7 @@ async def handle_budget_for_entity(
# If budget fields are provided, update the existing budget
if _budget_data:
await update_budget(
budget_obj=BudgetNewRequest(budget_id=existing_budget_id, **_budget_data),
budget_obj=BudgetNewRequest.model_validate({"budget_id": existing_budget_id, **_budget_data}),
user_api_key_dict=user_api_key_dict,
)

2
litellm/proxy/schema.prisma
View file

@ -19,6 +19,8 @@ model LiteLLM_BudgetTable {
rpm_limit BigInt?
tpd_limit BigInt?
model_max_budget Json?
model_tpm_limit Json?
model_rpm_limit Json?
budget_duration String?
budget_reset_at DateTime?
allowed_models String[] @default([]) // per-member model scope; empty = inherit team models

2
litellm/types/proxy/auth/user_api_key_auth.py
View file

@ -87,6 +87,8 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken):
end_user_tpd_limit: int | None = None
end_user_max_budget: float | None = None
end_user_model_max_budget: _UntypedDict | None = None
end_user_model_tpm_limit: Mapping[str, int] | None = None
end_user_model_rpm_limit: Mapping[str, int] | None = None
# Organization Params
organization_alias: str | None = None

2
schema.prisma
View file

@ -19,6 +19,8 @@ model LiteLLM_BudgetTable {
rpm_limit BigInt?
tpd_limit BigInt?
model_max_budget Json?
model_tpm_limit Json?
model_rpm_limit Json?
budget_duration String?
budget_reset_at DateTime?
allowed_models String[] @default([]) // per-member model scope; empty = inherit team models

78
tests/integration/management/test_budget_model_limits.py Normal file
View file

@ -0,0 +1,78 @@
import uuid
from typing import Final
import pytest
from pydantic import JsonValue, TypeAdapter
from tests.integration._support.client import Gateway, object_value
from tests.integration._support.database import read_rows
pytestmark: Final = pytest.mark.shares_worker_state
_JSON_OBJECTS: Final = TypeAdapter(list[dict[str, JsonValue]])
def _budget_columns(budget_id: str) -> dict[str, JsonValue]:
rows: Final = read_rows(
'SELECT model_tpm_limit, model_rpm_limit FROM "LiteLLM_BudgetTable" WHERE budget_id = %s', (budget_id,)
)
assert len(rows) == 1, rows
return rows[0]
def _budget_info(gateway: Gateway, budget_id: str) -> dict[str, JsonValue]:
response: Final = gateway.request("POST", "/budget/info", {"budgets": [budget_id]})
assert response.status_code == 200, response.text
entries: Final = _JSON_OBJECTS.validate_json(response.content)
assert len(entries) == 1, entries
return entries[0]
def test_budget_stores_and_returns_per_model_tpm_and_rpm_limits(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
budget_id: Final = scenario.budget(model_tpm_limit={"gpt-4o": 400000}, model_rpm_limit={"gpt-4o": 200})
info: Final = _budget_info(gateway, budget_id)
assert info["model_tpm_limit"] == {"gpt-4o": 400000}, info
assert info["model_rpm_limit"] == {"gpt-4o": 200}, info
row: Final = _budget_columns(budget_id)
assert row == {"model_tpm_limit": {"gpt-4o": 400000}, "model_rpm_limit": {"gpt-4o": 200}}, row
response: Final = gateway.request(
"POST", "/budget/update", {"budget_id": budget_id, "model_tpm_limit": {"gpt-4o": 100000, "gpt-4o-mini": 50}}
)
assert response.status_code == 200, response.text
updated: Final = _budget_columns(budget_id)
assert updated == {
"model_tpm_limit": {"gpt-4o": 100000, "gpt-4o-mini": 50},
"model_rpm_limit": {"gpt-4o": 200},
}, updated
def test_customer_created_with_inline_per_model_limits_gets_a_budget_that_carries_them(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
customer: Final = f"integration-customer-{uuid.uuid4().hex}"
created: Final = gateway.post(
"/customer/new",
{"user_id": customer, "model_tpm_limit": {"gpt-4o": 400000}, "model_rpm_limit": {"gpt-4o": 200}},
)
scenario.cleanups.callback(gateway.post, "/customer/delete", {"user_ids": [customer]})
budget_id: Final = created["budget_id"]
assert isinstance(budget_id, str) and budget_id, created
scenario.cleanups.callback(scenario.delete_budget, budget_id)
response: Final = gateway.request("GET", f"/customer/info?end_user_id={customer}", None)
assert response.status_code == 200, response.text
budget: Final = object_value(
TypeAdapter(dict[str, JsonValue]).validate_json(response.content)["litellm_budget_table"]
)
assert budget["model_tpm_limit"] == {"gpt-4o": 400000}, budget
assert budget["model_rpm_limit"] == {"gpt-4o": 200}, budget
@pytest.mark.parametrize("field", ["model_tpm_limit", "model_rpm_limit"])
def test_budget_rejects_a_negative_per_model_limit(gateway: Gateway, field: str) -> None:
budget_id: Final = f"integration-budget-{uuid.uuid4().hex}"
response: Final = gateway.request("POST", "/budget/new", {"budget_id": budget_id, field: {"gpt-4o": -1}})
assert response.status_code == 422, response.text
assert read_rows('SELECT budget_id FROM "LiteLLM_BudgetTable" WHERE budget_id = %s', (budget_id,)) == []

54
tests/unit/proxy/auth/test_custom_auth_end_user_budget.py
View file

@ -1,17 +1,19 @@
import pytest
from typing import Final
from unittest.mock import AsyncMock, patch
import pytest
from starlette.requests import Request
import litellm
from litellm.proxy.auth.user_api_key_auth import (
_run_post_custom_auth_checks,
update_valid_token_with_end_user_params,
)
from litellm.proxy._types import (
LiteLLM_BudgetTable,
LiteLLM_EndUserTable,
UserAPIKeyAuth,
)
from litellm.proxy.auth.user_api_key_auth import (
_run_post_custom_auth_checks,
update_valid_token_with_end_user_params,
)
@pytest.mark.asyncio
@ -407,3 +409,45 @@ def test_end_user_budget_tpd_limit_reaches_the_token():
assert result.end_user_rpm_limit == 5
assert result.end_user_tpd_limit == 750000
def test_update_valid_token_loads_per_model_tpm_and_rpm_limits_from_the_customer_budget():
valid_token: Final = UserAPIKeyAuth(token="test_token")
end_user_params: Final[dict[str, object]] = {
"end_user_id": "acme",
"end_user_model_tpm_limit": {"gpt-4o": 400000},
"end_user_model_rpm_limit": {"gpt-4o": 200},
}
result: Final = update_valid_token_with_end_user_params(valid_token, end_user_params)
assert result.end_user_model_tpm_limit == {"gpt-4o": 400000}
assert result.end_user_model_rpm_limit == {"gpt-4o": 200}
def test_update_valid_token_keeps_custom_auth_per_model_limits_when_the_budget_has_none():
valid_token: Final = UserAPIKeyAuth(
token="test_token",
end_user_id="acme",
end_user_model_tpm_limit={"gpt-4o": 1000},
end_user_model_rpm_limit={"gpt-4o": 10},
)
result: Final = update_valid_token_with_end_user_params(valid_token, {"end_user_id": "acme"})
assert result.end_user_model_tpm_limit == {"gpt-4o": 1000}
assert result.end_user_model_rpm_limit == {"gpt-4o": 10}
def test_apply_budget_limits_copies_per_model_limits_onto_end_user_params():
from litellm.proxy.auth.user_api_key_auth import _apply_budget_limits_to_end_user_params
end_user_params: Final[dict[str, object]] = {}
budget: Final = LiteLLM_BudgetTable(model_tpm_limit={"gpt-4o": 400000}, model_rpm_limit={"gpt-4o": 200})
_apply_budget_limits_to_end_user_params(end_user_params, budget, "acme")
assert end_user_params == {
"end_user_model_tpm_limit": {"gpt-4o": 400000},
"end_user_model_rpm_limit": {"gpt-4o": 200},
}

45
tests/unit/proxy/management_endpoints/test_budget_endpoints.py
View file

@ -4,16 +4,16 @@ import json
import types
from datetime import datetime, timedelta, timezone
from typing import Final
import pytest
from unittest.mock import AsyncMock, MagicMock
import pytest
from fastapi.testclient import TestClient
import litellm.proxy.proxy_server as ps
from litellm.proxy.proxy_server import app
from litellm.proxy._types import UserAPIKeyAuth, LitellmUserRoles, CommonProxyErrors
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.proxy_server import app
@pytest.fixture
@ -192,6 +192,43 @@ async def test_new_and_update_budget_persist_tpd_limit(client_and_mocks):
assert mock_table.update.await_args.kwargs["data"]["tpd_limit"] == 500000
@pytest.mark.asyncio
async def test_new_and_update_budget_persist_per_model_tpm_and_rpm_limits(client_and_mocks):
client, _, mock_table = client_and_mocks
resp = client.post(
"/budget/new",
json={"budget_id": "budget_models", "model_tpm_limit": {"gpt-4o": 400000}, "model_rpm_limit": {"gpt-4o": 200}},
)
assert resp.status_code == 200, resp.text
created: Final = mock_table.create.await_args.kwargs["data"]
assert json.loads(created["model_tpm_limit"]) == {"gpt-4o": 400000}, created
assert json.loads(created["model_rpm_limit"]) == {"gpt-4o": 200}, created
resp = client.post("/budget/update", json={"budget_id": "budget_models", "model_tpm_limit": {"gpt-4o": 0}})
assert resp.status_code == 200, resp.text
updated: Final = mock_table.update.await_args.kwargs["data"]
assert json.loads(updated["model_tpm_limit"]) == {"gpt-4o": 0}, updated
assert "model_rpm_limit" not in updated, updated
@pytest.mark.parametrize("route", ["/budget/new", "/budget/update"])
@pytest.mark.parametrize("field", ["model_tpm_limit", "model_rpm_limit"])
@pytest.mark.parametrize("value", [-1, 1.5])
@pytest.mark.asyncio
async def test_budget_rejects_a_per_model_limit_that_is_not_a_whole_number_of_zero_or_more(
client_and_mocks, route, field, value
):
client, _, mock_table = client_and_mocks
resp = client.post(route, json={"budget_id": "budget_models", field: {"gpt-4o": value}})
assert resp.status_code == 422, resp.text
assert field in resp.text
mock_table.create.assert_not_awaited()
mock_table.update.assert_not_awaited()
@pytest.mark.asyncio
async def test_update_budget_missing_id(client_and_mocks, monkeypatch):
client, mock_prisma, mock_table = client_and_mocks

49
tests/unit/proxy/management_endpoints/test_customer_endpoints.py
View file

@ -1,5 +1,6 @@
from collections.abc import Mapping, Sequence
from contextlib import contextmanager
from typing import Final, Mapping, Sequence
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -261,7 +262,9 @@ def test_delete_customer_drops_its_committed_spend_so_a_recreated_customer_start
memory={"spend:committed:end_user:cust-gone": 7.0, "spend:committed:end_user:cust-stays": 3.0},
)
monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_store", spend.store())
mock_prisma_client.db.litellm_endusertable.find_many = AsyncMock(return_value=[SimpleNamespace(user_id="cust-gone")])
mock_prisma_client.db.litellm_endusertable.find_many = AsyncMock(
return_value=[SimpleNamespace(user_id="cust-gone")]
)
mock_prisma_client.db.litellm_endusertable.delete_many = AsyncMock(return_value=1)
response: Final = client.post(
@ -827,6 +830,8 @@ _FULL_DB_ROW = {
"tpm_limit": None,
"rpm_limit": None,
"model_max_budget": None,
"model_tpm_limit": None,
"model_rpm_limit": None,
"budget_duration": "30d",
"allowed_models": [],
"budget_reset_at": "2024-02-01T00:00:00",
@ -873,6 +878,8 @@ _EXPECTED_CUSTOMER = {
"rpm_limit": None,
"tpd_limit": None,
"model_max_budget": None,
"model_tpm_limit": None,
"model_rpm_limit": None,
"budget_duration": "30d",
"allowed_models": [],
"temp_budget_increase": None,
@ -914,7 +921,9 @@ def test_char_info_body(mock_prisma_client, mock_user_api_key_auth):
@pytest.mark.parametrize("path", ["/customer/list", "/end_user/list"])
def test_char_list_body(
mock_prisma_client: MagicMock, mock_user_api_key_auth: None, path: str,
mock_prisma_client: MagicMock,
mock_user_api_key_auth: None,
path: str,
) -> None:
mock_prisma_client.db.litellm_endusertable.find_many = AsyncMock(return_value=[_row(_FULL_DB_ROW)])
response: Final = client.get(path, headers={"Authorization": "Bearer k"})
@ -944,10 +953,9 @@ class _CustomerListTable:
filtered: Final = tuple(row for row in self.rows if self._matches(row, where))
ordered: Final = (
tuple(sorted(filtered, key=customer_id, reverse=order["user_id"] == "desc"))
if order else filtered
tuple(sorted(filtered, key=customer_id, reverse=order["user_id"] == "desc")) if order else filtered
)
selected: Final = ordered[skip:] if take is None else ordered[skip:skip + take]
selected: Final = ordered[skip:] if take is None else ordered[skip : skip + take]
self.materialized_rows += (len(selected),)
return selected
@ -972,9 +980,7 @@ class _CustomerListTable:
@pytest.mark.parametrize("path", ["/customer/list", "/end_user/list"])
@pytest.mark.parametrize(
"role", [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY]
)
@pytest.mark.parametrize("role", [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY])
def test_customer_list_defaults_bound_database_reads(
mock_prisma_client: MagicMock,
mock_user_api_key_auth: None,
@ -997,7 +1003,9 @@ def test_customer_list_defaults_bound_database_reads(
@pytest.mark.parametrize("path", ["/customer/list", "/end_user/list"])
def test_customer_list_pages_cover_the_directory_once(
mock_prisma_client: MagicMock, mock_user_api_key_auth: None, path: str,
mock_prisma_client: MagicMock,
mock_user_api_key_auth: None,
path: str,
) -> None:
table: Final = _CustomerListTable()
mock_prisma_client.db.litellm_endusertable = table
@ -1021,8 +1029,12 @@ def test_customer_list_pages_cover_the_directory_once(
[("", 1, [0, 1, 2]), ("CuStOmEr-0249", 1, [249]), ("TITLE 24", 2, [242, 243, 244]), ("missing", 1, [])],
)
def test_customer_list_search_filters_before_pagination(
mock_prisma_client: MagicMock, mock_user_api_key_auth: None,
path: str, search: str, page: int, expected: list[int],
mock_prisma_client: MagicMock,
mock_user_api_key_auth: None,
path: str,
search: str,
page: int,
expected: list[int],
) -> None:
table: Final = _CustomerListTable()
mock_prisma_client.db.litellm_endusertable = table
@ -1036,7 +1048,10 @@ def test_customer_list_search_filters_before_pagination(
@pytest.mark.parametrize("path", ["/customer/list", "/end_user/list"])
@pytest.mark.parametrize("query", ["page=0", "page=-1", "page_size=0", "page_size=-1", "page_size=101", "page=bad"])
def test_customer_list_invalid_pagination_never_reads_database(
mock_prisma_client: MagicMock, mock_user_api_key_auth: None, path: str, query: str,
mock_prisma_client: MagicMock,
mock_user_api_key_auth: None,
path: str,
query: str,
) -> None:
table: Final = _CustomerListTable()
mock_prisma_client.db.litellm_endusertable = table
@ -1048,7 +1063,9 @@ def test_customer_list_invalid_pagination_never_reads_database(
@pytest.mark.parametrize("path", ["/customer/list", "/end_user/list"])
def test_customer_list_non_admin_never_reads_database(
mock_prisma_client: MagicMock, mock_user_api_key_auth: None, path: str,
mock_prisma_client: MagicMock,
mock_user_api_key_auth: None,
path: str,
) -> None:
table: Final = _CustomerListTable()
mock_prisma_client.db.litellm_endusertable = table
@ -1084,9 +1101,7 @@ def test_customer_new_forwards_models_to_db(mock_prisma_client, mock_user_api_ke
@pytest.mark.parametrize("bad_duration", ["0s", "-5m"])
def test_customer_new_rejects_a_duration_that_never_advances(
mock_prisma_client, mock_user_api_key_auth, bad_duration
):
def test_customer_new_rejects_a_duration_that_never_advances(mock_prisma_client, mock_user_api_key_auth, bad_duration):
"""A zero-length window resets to "now", leaving the customer's budget row
permanently due for the reset job to re-read every tick."""
mock_prisma_client.db.litellm_endusertable.create = AsyncMock(return_value=_row(_FULL_DB_ROW))

57
tests/unit/proxy/management_endpoints/test_organization_endpoints.py
View file

@ -1,15 +1,14 @@
import asyncio
import json
from collections.abc import Mapping
from contextlib import asynccontextmanager
from types import MappingProxyType, SimpleNamespace
from typing import Final, Mapping, Optional, cast
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import HTTPException
from fastapi.testclient import TestClient
from litellm._uuid import uuid
from tests.unit.proxy.management_endpoints.jwt_key_mapping_doubles import (
CascadingJWTMappingTable,
JWTMappingRow,
@ -27,8 +26,6 @@ async def test_organization_update_object_permissions_existing_permission(monkey
"""
from unittest.mock import AsyncMock, MagicMock
import pytest
from litellm.proxy._types import (
LiteLLM_ObjectPermissionBase,
LiteLLM_OrganizationTable,
@ -262,8 +259,6 @@ async def test_organization_update_object_permissions_no_existing_permission(
"""
from unittest.mock import AsyncMock, MagicMock
import pytest
from litellm.proxy._types import (
LiteLLM_ObjectPermissionBase,
LiteLLM_OrganizationTable,
@ -327,8 +322,6 @@ async def test_organization_update_object_permissions_missing_permission_record(
"""
from unittest.mock import AsyncMock, MagicMock
import pytest
from litellm.proxy._types import (
LiteLLM_ObjectPermissionBase,
LiteLLM_OrganizationTable,
@ -648,6 +641,7 @@ async def test_organization_member_add_budget_omission_and_null_leave_budget_uns
from litellm.proxy.management_endpoints.organization_endpoints import organization_member_add
user = LiteLLM_UserTable(user_id="user-1", user_role="internal_user")
async def create_membership(data):
return LiteLLM_OrganizationMembershipTable(
user_id="user-1",
@ -1445,6 +1439,43 @@ async def test_new_organization_temp_budget_fields_go_to_budget_row_not_metadata
assert json.loads(org_write.get("metadata", "{}")) == {}
@pytest.mark.asyncio
async def test_new_organization_keeps_per_model_limits_in_org_metadata_not_the_budget_row(monkeypatch):
"""The rate limiter reads an organization's model_tpm_limit / model_rpm_limit from the org metadata.
Budgets gained columns of the same name for customers, so /organization/new must still write the maps
to metadata and leave them off the budget row."""
from litellm.proxy._types import LitellmUserRoles, NewOrganizationRequest, UserAPIKeyAuth
from litellm.proxy.management_endpoints.organization_endpoints import new_organization
from litellm.proxy.utils import PrismaClient
prisma_client: Final = MagicMock()
prisma_client.jsonify_object = MagicMock(side_effect=lambda data: PrismaClient.jsonify_object(prisma_client, data))
prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
prisma_client.db.litellm_budgettable.create = AsyncMock(return_value=MagicMock(budget_id="budget-1"))
prisma_client.db.litellm_organizationtable.create = AsyncMock(return_value={"organization_id": "org-1"})
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", MagicMock())
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True, raising=False)
response: Final = await new_organization(
data=NewOrganizationRequest(
organization_alias="org",
model_tpm_limit={"gpt-4o": 400000},
model_rpm_limit={"gpt-4o": 200},
),
user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert response == {"organization_id": "org-1"}
budget_write: Final = prisma_client.db.litellm_budgettable.create.await_args.kwargs["data"]
assert "model_tpm_limit" not in budget_write and "model_rpm_limit" not in budget_write, budget_write
org_write: Final = prisma_client.db.litellm_organizationtable.create.await_args.kwargs["data"]
assert json.loads(org_write["metadata"]) == {
"model_tpm_limit": {"gpt-4o": 400000},
"model_rpm_limit": {"gpt-4o": 200},
}, org_write
def test_v2_update_organization_is_in_openapi_schema():
"""PATCH /v2/organization/{organization_id} is documented in the generated OpenAPI spec."""
from fastapi import FastAPI
@ -1511,7 +1542,7 @@ def _organization_test_client() -> TestClient:
@pytest.mark.parametrize(("method", "path"), _organization_route_targets())
def test_organization_routes_are_blocked_without_enterprise_license(monkeypatch, method, path):
"""Every /organization route is enterprise-only, even for a proxy admin sending a valid request."""
import litellm.proxy.proxy_server as proxy_server
from litellm.proxy import proxy_server
monkeypatch.setattr(proxy_server, "premium_user", False, raising=False)
monkeypatch.setattr(proxy_server, "prisma_client", None, raising=False)
@ -1525,7 +1556,7 @@ def test_organization_routes_are_blocked_without_enterprise_license(monkeypatch,
@pytest.mark.parametrize(("method", "path"), _organization_route_targets())
def test_organization_routes_reach_their_handler_with_enterprise_license(monkeypatch, method, path):
"""The same request a license refuses above now reaches the handler, which is the code reporting the missing database."""
import litellm.proxy.proxy_server as proxy_server
from litellm.proxy import proxy_server
from litellm.proxy._types import CommonProxyErrors
monkeypatch.setattr(proxy_server, "premium_user", True, raising=False)
@ -1689,7 +1720,9 @@ async def test_delete_organization_is_refused_while_one_of_its_teams_holds_a_ptu
@pytest.mark.asyncio
async def test_organization_member_add_reports_an_unknown_org_as_a_request_error_not_an_auth_error(monkeypatch: pytest.MonkeyPatch) -> None:
async def test_organization_member_add_reports_an_unknown_org_as_a_request_error_not_an_auth_error(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from fastapi import Request
from litellm.proxy._types import (

60
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -2131,6 +2131,8 @@ export interface paths {
* - rpm_limit: Optional[int] - The requests per minute limit for the budget.
* - tpd_limit: Optional[int] - The tokens per day limit for the budget. Charged by batch submissions instead of tpm_limit/rpm_limit.
* - model_max_budget: Optional[dict] - Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d", "tpm_limit": 100000, "rpm_limit": 100000}}
* - model_tpm_limit: Optional[dict] - Max tokens per minute per model, shared by everything on this budget. Example: {"gpt-4o": 400000}
* - model_rpm_limit: Optional[dict] - Max requests per minute per model, shared by everything on this budget. Example: {"gpt-4o": 200}
* - budget_reset_at: Optional[datetime] - Datetime when the initial budget is reset. Default is now.
*/
post: operations["new_budget_budget_new_post"];
@ -2188,6 +2190,8 @@ export interface paths {
* - rpm_limit: Optional[int] - The requests per minute limit for the budget.
* - tpd_limit: Optional[int] - The tokens per day limit for the budget. Charged by batch submissions instead of tpm_limit/rpm_limit.
* - model_max_budget: Optional[dict] - Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d", "tpm_limit": 100000, "rpm_limit": 100000}}
* - model_tpm_limit: Optional[dict] - Max tokens per minute per model, shared by everything on this budget. Example: {"gpt-4o": 400000}
* - model_rpm_limit: Optional[dict] - Max requests per minute per model, shared by everything on this budget. Example: {"gpt-4o": 200}
* - budget_reset_at: Optional[datetime] - Update the Datetime when the budget was last reset.
*/
post: operations["update_budget_budget_update_post"];
@ -4551,6 +4555,8 @@ export interface paths {
* - rpm_limit: Optional[int] - [Not Implemented Yet] Specify rpm limit for a given customer (Requests per minute)
* - tpd_limit: Optional[int] - Specify tpd limit for a given customer (Tokens per day). Batch submissions are charged against it instead of tpm_limit/rpm_limit
* - model_max_budget: Optional[dict] - [Not Implemented Yet] Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d"}}
* - model_tpm_limit: Optional[dict] - Max tokens per minute per model for this customer, shared across every key the customer is sent on. Stored on the customer's budget. Example: {"gpt-4o": 400000}
* - model_rpm_limit: Optional[dict] - Max requests per minute per model for this customer. Stored on the customer's budget. Example: {"gpt-4o": 200}
* - max_parallel_requests: Optional[int] - [Not Implemented Yet] Specify max parallel requests for a given customer.
* - soft_budget: Optional[float] - [Not Implemented Yet] Get alerts when customer crosses given budget, doesn't block requests.
* - spend: Optional[float] - Specify initial spend for a given customer.
@ -5229,6 +5235,8 @@ export interface paths {
* - rpm_limit: Optional[int] - [Not Implemented Yet] Specify rpm limit for a given customer (Requests per minute)
* - tpd_limit: Optional[int] - Specify tpd limit for a given customer (Tokens per day). Batch submissions are charged against it instead of tpm_limit/rpm_limit
* - model_max_budget: Optional[dict] - [Not Implemented Yet] Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d"}}
* - model_tpm_limit: Optional[dict] - Max tokens per minute per model for this customer, shared across every key the customer is sent on. Stored on the customer's budget. Example: {"gpt-4o": 400000}
* - model_rpm_limit: Optional[dict] - Max requests per minute per model for this customer. Stored on the customer's budget. Example: {"gpt-4o": 200}
* - max_parallel_requests: Optional[int] - [Not Implemented Yet] Specify max parallel requests for a given customer.
* - soft_budget: Optional[float] - [Not Implemented Yet] Get alerts when customer crosses given budget, doesn't block requests.
* - spend: Optional[float] - Specify initial spend for a given customer.
@ -28874,6 +28882,20 @@ export interface components {
model_max_budget?: {
[key: string]: components["schemas"]["BudgetConfig"];
} | null;
/**
* Model Rpm Limit
* @description Max requests per minute per model for this budget id, shared by every key, team and customer on it (e.g. {'gpt-4o': 200})
*/
model_rpm_limit?: {
[key: string]: number;
} | null;
/**
* Model Tpm Limit
* @description Max tokens per minute per model for this budget id, shared by every key, team and customer on it (e.g. {'gpt-4o': 400000})
*/
model_tpm_limit?: {
[key: string]: number;
} | null;
/**
* Rpm Limit
* @description Max requests per minute, allowed for this budget id.
@ -35519,6 +35541,14 @@ export interface components {
model_max_budget?: {
[key: string]: unknown;
} | null;
/** Model Rpm Limit */
model_rpm_limit?: {
[key: string]: number;
} | null;
/** Model Tpm Limit */
model_tpm_limit?: {
[key: string]: number;
} | null;
/** Rpm Limit */
rpm_limit?: number | null;
/** Soft Budget */
@ -35558,6 +35588,14 @@ export interface components {
model_max_budget?: {
[key: string]: unknown;
} | null;
/** Model Rpm Limit */
model_rpm_limit?: {
[key: string]: number;
} | null;
/** Model Tpm Limit */
model_tpm_limit?: {
[key: string]: number;
} | null;
/** Rpm Limit */
rpm_limit?: number | null;
/** Soft Budget */
@ -40180,6 +40218,20 @@ export interface components {
model_max_budget?: {
[key: string]: components["schemas"]["BudgetConfig"];
} | null;
/**
* Model Rpm Limit
* @description Max requests per minute per model for this budget id, shared by every key, team and customer on it (e.g. {'gpt-4o': 200})
*/
model_rpm_limit?: {
[key: string]: number;
} | null;
/**
* Model Tpm Limit
* @description Max tokens per minute per model for this budget id, shared by every key, team and customer on it (e.g. {'gpt-4o': 400000})
*/
model_tpm_limit?: {
[key: string]: number;
} | null;
/** Models */
models?: string[] | null;
object_permission?: components["schemas"]["LiteLLM_ObjectPermissionBase"] | null;
@ -51371,6 +51423,14 @@ export interface components {
end_user_model_max_budget?: {
[key: string]: unknown;
} | null;
/** End User Model Rpm Limit */
end_user_model_rpm_limit?: {
[key: string]: number;
} | null;
/** End User Model Tpm Limit */
end_user_model_tpm_limit?: {
[key: string]: number;
} | null;
end_user_object_permission?: components["schemas"]["LiteLLM_ObjectPermissionTable"] | null;
/** End User Rpm Limit */
end_user_rpm_limit?: number | null;