mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-06 02:48:13 +00:00
fix(proxy): count keyless JWT callers per user, require positive file caps, and take the upload slot after request validation
This commit is contained in:
parent
3b5f1bb122
commit
cc3316aa32
7 changed files with 71 additions and 18 deletions
|
|
@ -2784,15 +2784,18 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
|
|||
)
|
||||
max_batch_file_records: int | None = Field(
|
||||
None,
|
||||
description="max records (non-blank lines) per batch input file for /v1/files uploads with purpose=batch, applied per key. A key or team can carry its own value in metadata, set by a proxy admin, and the lowest applicable value wins. Unset means no limit",
|
||||
gt=0,
|
||||
description="max records (non-blank lines) per batch input file for /v1/files uploads with purpose=batch, applied per key. A key's metadata can override it and a team's metadata adds a team cap on top, both set by a proxy admin; the lower of the key's value and the team's value wins. Unset means no limit",
|
||||
)
|
||||
max_batch_file_uploads_per_day: int | None = Field(
|
||||
None,
|
||||
description="max /v1/files uploads with purpose=batch per key per UTC day. A key's metadata can override it and a team's metadata adds a shared team cap, both set by a proxy admin. Unset means no limit",
|
||||
gt=0,
|
||||
description="max /v1/files uploads with purpose=batch per key (per user for JWT callers) per UTC day. A key's metadata can override it and a team's metadata adds a shared team cap, both set by a proxy admin. Unset means no limit",
|
||||
)
|
||||
max_file_downloads_per_minute: int | None = Field(
|
||||
None,
|
||||
description="max GET /v1/files/{file_id}/content calls per key per file per minute. A key's metadata can override it and a team's metadata adds a shared team cap, both set by a proxy admin. Unset means no limit",
|
||||
gt=0,
|
||||
description="max GET /v1/files/{file_id}/content calls per key (per user for JWT callers) per file per minute. A key's metadata can override it and a team's metadata adds a shared team cap, both set by a proxy admin. Unset means no limit",
|
||||
)
|
||||
max_file_size_mb: int | None = Field(
|
||||
None,
|
||||
|
|
|
|||
|
|
@ -25,7 +25,7 @@ FileUsageSetting: TypeAlias = Literal[
|
|||
"max_file_downloads_per_minute",
|
||||
]
|
||||
LimitSource: TypeAlias = Literal["key", "team", "general_settings"]
|
||||
CounterScope: TypeAlias = Literal["key", "team"]
|
||||
CounterScope: TypeAlias = Literal["key", "user", "team"]
|
||||
|
||||
_COUNTER_PREFIX: Final = "litellm:file_usage"
|
||||
_DAY_SECONDS: Final = 24 * 60 * 60
|
||||
|
|
@ -65,9 +65,8 @@ def _read_limit(
|
|||
return FileUsageLimit(setting=setting, value=_LIMIT_ADAPTER.validate_python(raw), source=source)
|
||||
except ValidationError:
|
||||
verbose_proxy_logger.warning(
|
||||
"Ignoring invalid %s value %r in %s; expected a positive integer",
|
||||
"Ignoring invalid %s in %s; expected a positive integer",
|
||||
setting,
|
||||
raw,
|
||||
source,
|
||||
)
|
||||
return None
|
||||
|
|
@ -103,17 +102,24 @@ def batch_file_record_limit(
|
|||
return min(applicable, key=lambda limit: limit.value, default=None)
|
||||
|
||||
|
||||
def _caller_counter(user_api_key_dict: UserAPIKeyAuth, limit: FileUsageLimit | None) -> ScopedFileUsageLimit | None:
|
||||
if limit is None:
|
||||
return None
|
||||
if user_api_key_dict.api_key:
|
||||
return ScopedFileUsageLimit(scope="key", scope_id=user_api_key_dict.api_key, limit=limit)
|
||||
if user_api_key_dict.user_id:
|
||||
return ScopedFileUsageLimit(scope="user", scope_id=user_api_key_dict.user_id, limit=limit)
|
||||
return None
|
||||
|
||||
|
||||
def resolve_scoped_limits(
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
general_settings: Mapping[str, object],
|
||||
setting: FileUsageSetting,
|
||||
) -> tuple[ScopedFileUsageLimit, ...]:
|
||||
key_limit: Final = _key_limit(user_api_key_dict, general_settings, setting)
|
||||
team_limit: Final = _team_limit(user_api_key_dict, setting)
|
||||
candidates: Final = (
|
||||
ScopedFileUsageLimit(scope="key", scope_id=user_api_key_dict.api_key, limit=key_limit)
|
||||
if key_limit is not None and user_api_key_dict.api_key
|
||||
else None,
|
||||
_caller_counter(user_api_key_dict, _key_limit(user_api_key_dict, general_settings, setting)),
|
||||
ScopedFileUsageLimit(scope="team", scope_id=user_api_key_dict.team_id, limit=team_limit)
|
||||
if team_limit is not None and user_api_key_dict.team_id
|
||||
else None,
|
||||
|
|
@ -172,17 +178,19 @@ def describe_limit_source(source: LimitSource) -> str:
|
|||
case "general_settings":
|
||||
return "in general_settings"
|
||||
case _:
|
||||
assert_never(source)
|
||||
return assert_never(source)
|
||||
|
||||
|
||||
def _describe_scope(scoped: ScopedFileUsageLimit) -> str:
|
||||
match scoped.scope:
|
||||
case "key":
|
||||
return "this key"
|
||||
case "user":
|
||||
return f"user {scoped.scope_id}"
|
||||
case "team":
|
||||
return f"team {scoped.scope_id}"
|
||||
case _:
|
||||
assert_never(scoped.scope)
|
||||
return assert_never(scoped.scope)
|
||||
|
||||
|
||||
def _raise_limit_exceeded(exceeded: FileUsageLimitExceeded, what_ran_out: str, when_it_resets: str) -> NoReturn:
|
||||
|
|
|
|||
|
|
@ -681,9 +681,6 @@ async def create_file(
|
|||
)
|
||||
if batch_file_failure is not None:
|
||||
raise_batch_file_validation_failure(batch_file_failure)
|
||||
await enforce_batch_file_upload_limit(
|
||||
proxy_logging_obj.file_usage_cache, user_api_key_dict, general_settings
|
||||
)
|
||||
|
||||
data = {"passthrough": True} if passthrough else {}
|
||||
|
||||
|
|
@ -757,6 +754,11 @@ async def create_file(
|
|||
seconds=expires_after_seconds,
|
||||
)
|
||||
|
||||
if purpose == "batch":
|
||||
await enforce_batch_file_upload_limit(
|
||||
proxy_logging_obj.file_usage_cache, user_api_key_dict, general_settings
|
||||
)
|
||||
|
||||
# Include original request and headers in the data
|
||||
data = await add_litellm_data_to_request(
|
||||
data=data,
|
||||
|
|
|
|||
|
|
@ -4291,12 +4291,19 @@ def test_create_file_batch_uploads_over_daily_limit_get_429_and_only_valid_files
|
|||
|
||||
try:
|
||||
invalid = _upload_batch(b"not json\n")
|
||||
malformed_expiry = client.post(
|
||||
"/v1/files",
|
||||
files={"file": ("batch.jsonl", VALID_BATCH_LINE, "application/jsonl")},
|
||||
data={"purpose": "batch", "expires_after[anchor]": "created_at"},
|
||||
headers={"Authorization": "Bearer test-key"},
|
||||
)
|
||||
allowed = [_upload_batch(VALID_BATCH_LINE) for _ in range(2)]
|
||||
rejected = _upload_batch(VALID_BATCH_LINE)
|
||||
finally:
|
||||
_teardown_batch_upload_endpoint()
|
||||
|
||||
assert invalid.status_code == 400, invalid.text
|
||||
assert malformed_expiry.status_code == 400, malformed_expiry.text
|
||||
assert [response.status_code for response in allowed] == [200, 200]
|
||||
assert rejected.status_code == 429, rejected.text
|
||||
assert rejected.headers["retry-after"] == str(86400 - 3615)
|
||||
|
|
|
|||
|
|
@ -395,3 +395,15 @@ def test_mcp_metadata_rejects_unavailable_upstream_protocol(revision):
|
|||
for model in (NewMCPServerRequest, UpdateMCPServerRequest):
|
||||
with pytest.raises(ValidationError):
|
||||
model.model_validate(payload)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"setting", ["max_batch_file_records", "max_batch_file_uploads_per_day", "max_file_downloads_per_minute"]
|
||||
)
|
||||
def test_batch_file_caps_accept_only_positive_limits(setting):
|
||||
from litellm.proxy._types import ConfigGeneralSettings
|
||||
|
||||
for invalid in (0, -5):
|
||||
with pytest.raises(ValidationError):
|
||||
ConfigGeneralSettings.model_validate({setting: invalid})
|
||||
assert getattr(ConfigGeneralSettings.model_validate({setting: 3}), setting) == 3
|
||||
|
|
|
|||
|
|
@ -85,6 +85,12 @@ def test_record_limit_is_the_lowest_applicable_limit(key_metadata, team_metadata
|
|||
{},
|
||||
(),
|
||||
),
|
||||
(
|
||||
UserAPIKeyAuth(api_key=None, user_id="u1", team_id="t1"),
|
||||
{"max_batch_file_uploads_per_day": 5},
|
||||
(_scoped("user", "u1", 5, "general_settings"),),
|
||||
),
|
||||
(UserAPIKeyAuth(api_key=None), {"max_batch_file_uploads_per_day": 5}, ()),
|
||||
],
|
||||
)
|
||||
def test_counter_scopes_pair_each_limit_with_its_own_counter(caller, general_settings, expected):
|
||||
|
|
@ -180,3 +186,18 @@ async def test_no_configured_limit_never_rejects():
|
|||
for _ in range(50):
|
||||
await enforce_batch_file_upload_limit(cache, caller, {}, clock=lambda: MIDDAY)
|
||||
await enforce_file_download_limit(cache, caller, {}, "file-a", clock=lambda: MIDDAY)
|
||||
|
||||
|
||||
async def test_a_caller_without_a_key_is_counted_per_user():
|
||||
cache: Final = _cache()
|
||||
general_settings: Final = {"max_file_downloads_per_minute": 1}
|
||||
jwt_caller: Final = UserAPIKeyAuth(api_key=None, user_id="u1")
|
||||
|
||||
await enforce_file_download_limit(cache, jwt_caller, general_settings, "file-a", clock=lambda: MIDDAY)
|
||||
await enforce_file_download_limit(
|
||||
cache, UserAPIKeyAuth(api_key=None, user_id="u2"), general_settings, "file-a", clock=lambda: MIDDAY
|
||||
)
|
||||
with pytest.raises(ProxyException) as exc:
|
||||
await enforce_file_download_limit(cache, jwt_caller, general_settings, "file-a", clock=lambda: MIDDAY)
|
||||
|
||||
assert "max_file_downloads_per_minute is 1 for user u1 (set in general_settings)" in exc.value.message
|
||||
|
|
|
|||
6
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
6
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -28205,7 +28205,7 @@ export interface components {
|
|||
master_key?: string | null;
|
||||
/**
|
||||
* Max Batch File Records
|
||||
* @description max records (non-blank lines) per batch input file for /v1/files uploads with purpose=batch, applied per key. A key or team can carry its own value in metadata, set by a proxy admin, and the lowest applicable value wins. Unset means no limit
|
||||
* @description max records (non-blank lines) per batch input file for /v1/files uploads with purpose=batch, applied per key. A key's metadata can override it and a team's metadata adds a team cap on top, both set by a proxy admin; the lower of the key's value and the team's value wins. Unset means no limit
|
||||
*/
|
||||
max_batch_file_records?: number | null;
|
||||
/**
|
||||
|
|
@ -28215,7 +28215,7 @@ export interface components {
|
|||
max_batch_file_size_mb?: number | null;
|
||||
/**
|
||||
* Max Batch File Uploads Per Day
|
||||
* @description max /v1/files uploads with purpose=batch per key per UTC day. A key's metadata can override it and a team's metadata adds a shared team cap, both set by a proxy admin. Unset means no limit
|
||||
* @description max /v1/files uploads with purpose=batch per key (per user for JWT callers) per UTC day. A key's metadata can override it and a team's metadata adds a shared team cap, both set by a proxy admin. Unset means no limit
|
||||
*/
|
||||
max_batch_file_uploads_per_day?: number | null;
|
||||
/**
|
||||
|
|
@ -28232,7 +28232,7 @@ export interface components {
|
|||
} | null;
|
||||
/**
|
||||
* Max File Downloads Per Minute
|
||||
* @description max GET /v1/files/{file_id}/content calls per key per file per minute. A key's metadata can override it and a team's metadata adds a shared team cap, both set by a proxy admin. Unset means no limit
|
||||
* @description max GET /v1/files/{file_id}/content calls per key (per user for JWT callers) per file per minute. A key's metadata can override it and a team's metadata adds a shared team cap, both set by a proxy admin. Unset means no limit
|
||||
*/
|
||||
max_file_downloads_per_minute?: number | null;
|
||||
/**
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue