scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

Merge acdccee0fe into f285229b51

Browse source
This commit is contained in:
Emerson Gomes 2026-09-30 17:02:03 -04:00 • committed by GitHub
commit 24f8988c67
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 210 additions and 42 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/proxy/_types.py
View file

@ -1261,6 +1261,7 @@ class AllowedVectorStoreIndexItem(LiteLLMPydanticObjectBase):
class KeyRequestBase(GenerateRequestBase):
key: str | None = None
expires: datetime | None = None
tpd_limit: int | None = None
default_estimated_output_tokens: PositiveInt | None = None
default_estimated_output_tokens_per_model: Mapping[str, PositiveInt] | None = None

99
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -419,7 +419,9 @@ def _effective_key_for_generate(data: GenerateKeyRequest, now: datetime) -> Lite
folded_metadata: Final = {**metadata, **metadata_fields} # mutable-ok: encrypt_callback_vars needs a dict
columns: Final = handle_key_type(data, {**column_fields}) # mutable-ok: handle_key_type mutates in place
expires: Final = (
now + timedelta(seconds=duration_in_seconds(duration=data.duration)) if data.duration is not None else None
now + timedelta(seconds=duration_in_seconds(duration=data.duration))
if data.duration is not None
else data.expires
)
budget_reset_at: Final = (
get_budget_reset_time(budget_duration=data.budget_duration) if data.budget_duration is not None else None
@ -1107,6 +1109,20 @@ _BUDGET_NUMERIC_KEYS = frozenset(
)
def _key_expiration_exceeds_limit(
data: GenerateKeyRequest | UpdateKeyRequest, maximum_duration: str, fill_defaults: bool
) -> bool:
if data.duration:
return False
if not fill_defaults and data.duration is None and "duration" in data.model_fields_set:
return True
if data.expires is None:
return not fill_defaults and "expires" in data.model_fields_set
expires: Final = data.expires if data.expires.tzinfo is not None else data.expires.replace(tzinfo=timezone.utc)
maximum_expiration: Final = datetime.now(timezone.utc) + timedelta(seconds=duration_in_seconds(maximum_duration))
return expires > maximum_expiration
def _enforce_upperbound_key_params(
data: GenerateKeyRequest | UpdateKeyRequest,
fill_defaults: bool = True,
@ -1133,40 +1149,41 @@ def _enforce_upperbound_key_params(
if litellm.upperbound_key_generate_params is None:
return
maximum_duration: Final = litellm.upperbound_key_generate_params.duration
if maximum_duration is not None and _key_expiration_exceeds_limit(data, maximum_duration, fill_defaults):
raise HTTPException(
status_code=400,
detail={"error": f"expires is over max limit set in config - max_duration={maximum_duration}"},
)
for elem in data:
key, value = elem
upperbound_value = getattr(litellm.upperbound_key_generate_params, key, None)
if upperbound_value is not None:
if value is None:
if fill_defaults:
setattr(data, key, upperbound_value)
else:
if key in [
"max_budget",
"max_parallel_requests",
"tpm_limit",
"rpm_limit",
]:
if value > upperbound_value:
raise HTTPException(
status_code=400,
detail={
"error": f"{key} is over max limit set in config - user_value={value}; max_value={upperbound_value}"
},
)
elif key in ["budget_duration", "duration"]:
upperbound_duration = duration_in_seconds(duration=upperbound_value)
if value == "-1":
user_duration = float("inf")
else:
user_duration = duration_in_seconds(duration=value)
if user_duration > upperbound_duration:
raise HTTPException(
status_code=400,
detail={
"error": f"{key} is over max limit set in config - user_value={value}; max_value={upperbound_value}"
},
)
if (upperbound_value := getattr(litellm.upperbound_key_generate_params, key, None)) is None:
continue
if value is None:
if fill_defaults and not (key == "duration" and data.expires is not None):
setattr(data, key, upperbound_value)
continue
if key in ["max_budget", "max_parallel_requests", "tpm_limit", "rpm_limit"]:
if value > upperbound_value:
raise HTTPException(
status_code=400,
detail={
"error": f"{key} is over max limit set in config - user_value={value}; max_value={upperbound_value}"
},
)
elif key in ["budget_duration", "duration"]:
upperbound_duration, user_duration = (
duration_in_seconds(duration=upperbound_value),
float("inf") if value == "-1" else duration_in_seconds(duration=value),
)
if user_duration > upperbound_duration:
raise HTTPException(
status_code=400,
detail={
"error": f"{key} is over max limit set in config - user_value={value}; max_value={upperbound_value}"
},
)
async def _common_key_generation_helper(
@ -1240,6 +1257,7 @@ async def _common_key_generation_helper(
if (
value is None
and (key != "budget_duration" or key not in data.model_fields_set)
and (key != "duration" or data.expires is None)
and key
in [
"max_budget",
@ -3410,6 +3428,7 @@ async def update_key_fn(
- rpm_limit_type: Optional[str] - RPM rate limit type - "best_effort_throughput", "guaranteed_throughput", or "dynamic"
- allowed_cache_controls: Optional[list] - List of allowed cache control values
- duration: Optional[str] - Key validity duration ("30d", "1h", etc.), null to never expire, or "-1" to never expire (deprecated, use null)
- expires: Optional[datetime] - Absolute expiration timestamp. An explicitly supplied duration takes precedence
- permissions: Optional[dict] - Key-specific permissions
- send_invite_email: Optional[bool] - Send invite email to user_id
- guardrails: Optional[List[str]] - List of active guardrails for the key
@ -4618,6 +4637,7 @@ async def generate_key_helper_fn(
budget_limits: list | None = None, # multiple concurrent budget windows
*,
llm_router: Router | None = None,
expires: datetime | None = None,
):
from litellm.proxy.proxy_server import premium_user, prisma_client
@ -4637,12 +4657,11 @@ async def generate_key_helper_fn(
else:
token = f"sk-{secrets.token_urlsafe(LENGTH_OF_LITELLM_GENERATED_KEY)}"
if duration is None: # allow tokens that never expire
expires = None
else:
# Add duration to current time for exact expiration (not standardized reset time)
duration_seconds: Final = duration_in_seconds(duration)
expires = datetime.now(timezone.utc) + timedelta(seconds=duration_seconds)
resolved_expires: Final = (
datetime.now(timezone.utc) + timedelta(seconds=duration_in_seconds(duration))
if duration is not None
else expires
)
if key_budget_duration is None: # one-time budget
key_reset_at = None
@ -4719,7 +4738,7 @@ async def generate_key_helper_fn(
key_data: Final = {
"token": token,
"key_alias": key_alias,
"expires": expires,
"expires": resolved_expires,
"models": models,
"aliases": aliases_json,
"config": config_json,

145
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -85,6 +85,7 @@ from litellm.types.proxy.management_endpoints.key_management_endpoints import (
BulkUpdateKeyResponse,
CustomKeyPolicyRequest,
)
from litellm.types.proxy.management_endpoints.ui_sso import LiteLLM_UpperboundKeyGenerateParams
client = TestClient(app)
@ -12722,6 +12723,87 @@ def test_enforce_upperbound_duration_over_limit(monkeypatch):
assert "duration" in str(exc_info.value.detail)
@pytest.mark.parametrize(
("request_type", "fill_defaults"),
((GenerateKeyRequest, True), (UpdateKeyRequest, False), (RegenerateKeyRequest, False)),
)
@pytest.mark.parametrize("with_timezone", (False, True))
def test_enforce_upperbound_rejects_expiration_beyond_limit(
monkeypatch: pytest.MonkeyPatch,
request_type: type[GenerateKeyRequest] | type[UpdateKeyRequest],
fill_defaults: bool,
with_timezone: bool,
) -> None:
monkeypatch.setattr(litellm, "upperbound_key_generate_params", LiteLLM_UpperboundKeyGenerateParams(duration="1d"))
future: Final = datetime.now(timezone.utc) + timedelta(days=2)
expires: Final = future if with_timezone else future.replace(tzinfo=None)
data: Final = request_type(key="sk-expiration-test", expires=expires)
with pytest.raises(HTTPException) as error:
_enforce_upperbound_key_params(data, fill_defaults=fill_defaults)
assert error.value.status_code == 400
assert "expires" in str(error.value.detail)
@pytest.mark.parametrize(
("request_type", "fill_defaults"),
((GenerateKeyRequest, True), (UpdateKeyRequest, False), (RegenerateKeyRequest, False)),
)
def test_enforce_upperbound_preserves_bounded_expiration(
monkeypatch: pytest.MonkeyPatch,
request_type: type[GenerateKeyRequest] | type[UpdateKeyRequest],
fill_defaults: bool,
) -> None:
monkeypatch.setattr(litellm, "upperbound_key_generate_params", LiteLLM_UpperboundKeyGenerateParams(duration="1d"))
expires: Final = datetime.now(timezone.utc) + timedelta(hours=1)
data: Final = request_type(key="sk-expiration-test", expires=expires)
_enforce_upperbound_key_params(data, fill_defaults=fill_defaults)
assert (data.duration, data.expires) == (None, expires)
@pytest.mark.parametrize("request_type", (UpdateKeyRequest, RegenerateKeyRequest))
@pytest.mark.parametrize("clear_duration", (False, True))
def test_enforce_upperbound_rejects_clearing_expiration(
monkeypatch: pytest.MonkeyPatch,
request_type: type[UpdateKeyRequest] | type[RegenerateKeyRequest],
clear_duration: bool,
) -> None:
monkeypatch.setattr(litellm, "upperbound_key_generate_params", LiteLLM_UpperboundKeyGenerateParams(duration="1d"))
expires: Final = datetime.now(timezone.utc) + timedelta(hours=1)
data: Final = request_type.model_validate(
{"key": "sk-expiration-test", "duration": None, "expires": expires}
if clear_duration
else {"key": "sk-expiration-test", "expires": None}
)
with pytest.raises(HTTPException) as error:
_enforce_upperbound_key_params(data, fill_defaults=False)
assert error.value.status_code == 400
assert "expires" in str(error.value.detail)
@pytest.mark.parametrize(
("request_type", "fill_defaults"),
((GenerateKeyRequest, True), (UpdateKeyRequest, False), (RegenerateKeyRequest, False)),
)
def test_enforce_upperbound_uses_explicit_duration_over_expiration(
monkeypatch: pytest.MonkeyPatch,
request_type: type[GenerateKeyRequest] | type[UpdateKeyRequest],
fill_defaults: bool,
) -> None:
monkeypatch.setattr(litellm, "upperbound_key_generate_params", LiteLLM_UpperboundKeyGenerateParams(duration="1d"))
expires: Final = datetime.now(timezone.utc) + timedelta(days=2)
data: Final = request_type(key="sk-expiration-test", duration="1h", expires=expires)
_enforce_upperbound_key_params(data, fill_defaults=fill_defaults)
assert (data.duration, data.expires) == ("1h", expires)
def test_enforce_upperbound_no_config_is_noop(monkeypatch):
"""Test that no enforcement happens when upperbound params are not configured."""
import litellm
@ -12783,7 +12865,10 @@ def _make_regenerate_existing_key():
@pytest.mark.asyncio
async def test_execute_virtual_key_regeneration_rejects_over_limit_duration(monkeypatch):
@pytest.mark.parametrize("use_absolute_expiration", (False, True))
async def test_execute_virtual_key_regeneration_rejects_over_limit_duration(
monkeypatch: pytest.MonkeyPatch, use_absolute_expiration: bool
) -> None:
"""Regenerate must reject durations exceeding upperbound_key_generate_params.duration."""
from litellm.proxy._types import RegenerateKeyRequest
from litellm.proxy.management_endpoints.key_management_endpoints import (
@ -12801,7 +12886,11 @@ async def test_execute_virtual_key_regeneration_rejects_over_limit_duration(monk
),
)
existing_key = _make_regenerate_existing_key()
data = RegenerateKeyRequest(duration="2h")
data: Final = (
RegenerateKeyRequest(expires=datetime.now(timezone.utc) + timedelta(hours=2))
if use_absolute_expiration
else RegenerateKeyRequest(duration="2h")
)
user_api_key_dict = _make_regenerate_user_api_key_dict()
mock_prisma_client = _make_regenerate_mock_prisma()
@ -13907,6 +13996,40 @@ def test_effective_key_for_generate_without_duration_never_expires():
assert effective_key.key_type == "default"
@pytest.mark.parametrize("duration", (None, "1h"))
def test_effective_key_for_generate_preserves_explicit_expiration(duration: str | None) -> None:
now: Final = datetime(2026, 1, 1, tzinfo=timezone.utc)
explicit: Final = now + timedelta(days=1)
request: Final = GenerateKeyRequest(expires=explicit, duration=duration)
effective: Final = _effective_key_for_generate(data=request, now=now)
assert effective.expires == (explicit if duration is None else now + timedelta(hours=1))
assert request.expires == explicit
@pytest.mark.asyncio
async def test_prepare_key_update_data_preserves_explicit_expiration() -> None:
explicit: Final = datetime(2026, 1, 1, tzinfo=timezone.utc)
existing: Final = LiteLLM_VerificationToken(token="test-token", metadata={})
request: Final = UpdateKeyRequest(key="test-token", expires=explicit)
update: Final = await prepare_key_update_data(data=request, existing_key_row=existing)
assert update["expires"] == explicit
@pytest.mark.asyncio
async def test_generate_key_helper_persists_explicit_expiration(monkeypatch: pytest.MonkeyPatch) -> None:
explicit: Final = datetime(2026, 1, 1, tzinfo=timezone.utc)
database: Final = AsyncMock()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", database)
await generate_key_helper_fn(request_type="key", expires=explicit, table_name="key", token="sk-expiration-test-token")
assert database.insert_data.call_args.kwargs["data"]["expires"] == explicit
def _policy_request_for_generate() -> CustomKeyPolicyRequest:
return CustomKeyPolicyRequest(
operation="generate",
@ -19276,6 +19399,24 @@ async def test_key_generate_omitted_budget_duration_still_takes_default_key_gene
assert key_row["budget_reset_at"] is not None
@pytest.mark.parametrize("with_lifetime_limit", (False, True))
async def test_key_generate_expiration_overrides_default_duration(
monkeypatch: pytest.MonkeyPatch, with_lifetime_limit: bool
) -> None:
monkeypatch.setattr(litellm, "default_key_generate_params", {"duration": "1d"})
monkeypatch.setattr(
litellm,
"upperbound_key_generate_params",
LiteLLM_UpperboundKeyGenerateParams(duration="1d") if with_lifetime_limit else None,
)
insert_data: Final = _wire_key_generation_prisma(monkeypatch)
expires: Final = datetime.now(timezone.utc) + timedelta(hours=1)
key_row: Final = await _generate_key_and_get_persisted_row(GenerateKeyRequest(expires=expires), insert_data)
assert key_row["expires"] == expires
@pytest.mark.asyncio
async def test_key_generate_explicit_null_budget_duration_cannot_bypass_upperbound(monkeypatch):
"""upperbound_key_generate_params is an admin ceiling: an explicit null must not mint an uncapped key,

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

@ -8404,6 +8404,7 @@ export interface paths {
* - rpm_limit_type: Optional[str] - RPM rate limit type - "best_effort_throughput", "guaranteed_throughput", or "dynamic"
* - allowed_cache_controls: Optional[list] - List of allowed cache control values
* - duration: Optional[str] - Key validity duration ("30d", "1h", etc.), null to never expire, or "-1" to never expire (deprecated, use null)
* - expires: Optional[datetime] - Absolute expiration timestamp. An explicitly supplied duration takes precedence
* - permissions: Optional[dict] - Key-specific permissions
* - send_invite_email: Optional[bool] - Send invite email to user_id
* - guardrails: Optional[List[str]] - List of active guardrails for the key
@ -30702,6 +30703,8 @@ export interface components {
end_user_budget_id?: string | null;
/** Enforced Params */
enforced_params?: string[] | null;
/** Expires */
expires?: string | null;
/** Guardrails */
guardrails?: string[] | null;
/** Key */
@ -39379,6 +39382,8 @@ export interface components {
end_user_budget_id?: string | null;
/** Enforced Params */
enforced_params?: string[] | null;
/** Expires */
expires?: string | null;
/** Grace Period */
grace_period?: string | null;
/** Guardrails */
@ -44698,6 +44703,8 @@ export interface components {
end_user_budget_id?: string | null;
/** Enforced Params */
enforced_params?: string[] | null;
/** Expires */
expires?: string | null;
/** Guardrails */
guardrails?: string[] | null;
/** Key */