mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
9dabd72f2d
Every repository handed its `.table` back untyped, so a dozen modules had each grown a private `_PrismaTableActions` Protocol to paper over it. They had drifted: some declared `update` as returning the row, others the row or None, and none agreed on whether `find_many` was covariant Replace all of them with a single `TableActions[RowT_co]` in `litellm/repositories/prisma_protocols.py`, keyed to the prisma row each repository is bound to. Query inputs stay `Mapping[str, object]` so callers keep passing plain dicts, and `find_many` returns `Sequence` so the row type stays covariant Typing the nullable returns honestly surfaced paths that were already crashing. A team admin could never edit or delete a memory entry owned by their team: the write-auth check fed a raw prisma row to a helper that expects the domain model, so `members_with_roles` arrived as plain dicts and the request died as a 500 instead of applying the edit. Non-admin members hit the same 500 in place of the 403 they were owed, so refusal and breakage were indistinguishable. `/v2/model/info?user_models_only=true` dereferenced a missing user row rather than returning the 400 the route already had, three team routes dereferenced a team deleted between the read and the write, and the agent registry dereferenced a missing agent instead of naming it basedpyright drops 2,132 errors, 1,454 of them reportAny and 73 reportExplicitAny. The dashboard's generated types pick up `string[]` where they had `unknown[]` for a team's members, admins and models
6756 lines
276 KiB
Python
6756 lines
276 KiB
Python
"""
|
||
KEY MANAGEMENT
|
||
|
||
All /key management endpoints
|
||
|
||
/key/generate
|
||
/key/info
|
||
/key/update
|
||
/key/delete
|
||
"""
|
||
|
||
import asyncio
|
||
import copy
|
||
import inspect
|
||
import json
|
||
import math
|
||
import os
|
||
import re
|
||
import secrets
|
||
import traceback
|
||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||
from datetime import datetime, timedelta, timezone
|
||
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypeVar, cast
|
||
|
||
import fastapi
|
||
import yaml
|
||
from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, status
|
||
|
||
import litellm
|
||
from litellm._logging import verbose_proxy_logger
|
||
from litellm._uuid import uuid
|
||
from litellm.caching.dual_cache import DualCache
|
||
from litellm.constants import (
|
||
LENGTH_OF_LITELLM_GENERATED_KEY,
|
||
LITELLM_PROXY_ADMIN_NAME,
|
||
MINIMUM_CUSTOM_KEY_LENGTH,
|
||
UI_SESSION_TOKEN_TEAM_ID,
|
||
)
|
||
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
|
||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||
from litellm.models.credentials import CredentialItem
|
||
from litellm.proxy._experimental.mcp_server.db import (
|
||
rotate_mcp_server_credentials_master_key,
|
||
rotate_mcp_user_credentials_master_key,
|
||
rotate_mcp_user_env_vars_master_key,
|
||
)
|
||
from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import (
|
||
rotate_sso_identity_assertions_master_key,
|
||
)
|
||
from litellm.proxy._types import *
|
||
from litellm.proxy._types import Litellm_EntityType, LiteLLM_VerificationToken, hash_token
|
||
from litellm.proxy.auth.auth_checks import (
|
||
_delete_cache_key_object,
|
||
can_team_access_model,
|
||
get_org_object,
|
||
get_project_object,
|
||
get_team_object,
|
||
)
|
||
from litellm.proxy.auth.auth_utils import (
|
||
abbreviate_api_key,
|
||
enforce_batch_enqueued_token_limit_is_admin_only,
|
||
enforce_output_token_estimates_are_admin_only,
|
||
)
|
||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||
from litellm.proxy.common_utils.callback_config_validation import logging_metadata_config_error
|
||
from litellm.proxy.common_utils.callback_utils import (
|
||
decrypt_callback_vars,
|
||
encrypt_callback_vars,
|
||
)
|
||
from litellm.proxy.common_utils.config_sync_pubsub import (
|
||
coordination_redis_cache,
|
||
publish_config_change,
|
||
)
|
||
from litellm.proxy.common_utils.rbac_utils import check_org_admin_can_generate_keys
|
||
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
|
||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||
from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks
|
||
from litellm.proxy.hooks.model_max_budget_limiter import build_model_max_budget_usage
|
||
from litellm.proxy.management_endpoints.common_utils import (
|
||
_check_passthrough_routes_caller_permission,
|
||
_is_user_org_admin_for_team,
|
||
_is_user_team_admin,
|
||
_set_object_metadata_field,
|
||
_team_member_has_permission,
|
||
_user_has_admin_view,
|
||
validate_budget_duration,
|
||
validate_finite_spend,
|
||
)
|
||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||
_add_model_to_db,
|
||
)
|
||
from litellm.proxy.management_helpers.access_group_key_sync import (
|
||
sync_key_access_group_membership,
|
||
sync_key_regeneration_access_group_membership,
|
||
sync_key_update_access_group_membership,
|
||
)
|
||
from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at
|
||
from litellm.proxy.management_helpers.object_permission_utils import (
|
||
_set_object_permission,
|
||
attach_object_permission_to_dict,
|
||
handle_update_object_permission_common,
|
||
validate_key_mcp_servers_against_team,
|
||
validate_key_search_tools_against_team,
|
||
validate_key_vector_stores_against_team,
|
||
)
|
||
from litellm.proxy.management_helpers.team_member_permission_checks import (
|
||
TeamMemberPermissionChecks,
|
||
)
|
||
from litellm.proxy.management_helpers.utils import management_endpoint_wrapper
|
||
from litellm.proxy.spend_tracking.spend_tracking_utils import _is_master_key
|
||
from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
|
||
get_ui_settings_cached,
|
||
)
|
||
from litellm.proxy.utils import (
|
||
PrismaClient,
|
||
ProxyLogging,
|
||
_hash_token_if_needed,
|
||
handle_exception_on_proxy,
|
||
is_valid_api_key,
|
||
)
|
||
from litellm.repositories.base_repository import BaseRepository
|
||
from litellm.repositories.budget_repository import BudgetRepository
|
||
from litellm.repositories.config_repository import ConfigParam, ConfigRepository
|
||
from litellm.repositories.credentials_repository import CredentialsRepository
|
||
from litellm.repositories.model_repository import ModelRepository
|
||
from litellm.repositories.prisma_protocols import TableActions
|
||
from litellm.repositories.table_repositories import (
|
||
DeletedVerificationTokenRepository,
|
||
DeprecatedVerificationTokenRepository,
|
||
)
|
||
from litellm.repositories.team_repository import TeamRepository
|
||
from litellm.repositories.user_repository import UserRepository
|
||
from litellm.repositories.verification_token_repository import (
|
||
VerificationTokenRepository,
|
||
)
|
||
from litellm.router import Router
|
||
from litellm.secret_managers.base_secret_manager import raise_if_unsafe_secret_name
|
||
from litellm.secret_managers.main import get_secret
|
||
from litellm.types.proxy.management_endpoints.key_management_endpoints import (
|
||
BulkUpdateKeyRequest,
|
||
BulkUpdateKeyResponse,
|
||
BulkUpdateTeamKeysRequest,
|
||
FailedKeyUpdate,
|
||
SuccessfulKeyUpdate,
|
||
)
|
||
from litellm.types.router import Deployment
|
||
from litellm.types.utils import (
|
||
BudgetConfig,
|
||
PersonalUIKeyGenerationConfig,
|
||
TeamUIKeyGenerationConfig,
|
||
)
|
||
|
||
if TYPE_CHECKING:
|
||
from prisma import Prisma
|
||
from prisma import models as prisma_models
|
||
|
||
_RepositoryModelT = TypeVar("_RepositoryModelT", bound=BaseModel)
|
||
|
||
|
||
class _UserRowLike(Protocol):
|
||
"""Read-only view of the user columns ``/key/list`` expands keys with."""
|
||
|
||
@property
|
||
def user_id(self) -> str | None: ...
|
||
|
||
@property
|
||
def user_email(self) -> str | None: ...
|
||
|
||
@property
|
||
def user_alias(self) -> str | None: ...
|
||
|
||
def model_dump(self) -> Mapping[str, object]: ...
|
||
|
||
def dict(self) -> Mapping[str, object]: ...
|
||
|
||
|
||
class _TxTables(Protocol):
|
||
litellm_proxymodeltable: TableActions[object]
|
||
|
||
|
||
class _ConfigTableActions(Protocol):
|
||
"""Config table surface this module needs; the shared repository seam exposes no ``update``."""
|
||
|
||
async def find_many(self) -> Sequence[ConfigParam]: ...
|
||
|
||
async def update(
|
||
self,
|
||
*,
|
||
where: Mapping[str, object],
|
||
data: Mapping[str, object],
|
||
) -> ConfigParam | None: ...
|
||
|
||
|
||
def _prisma_table(
|
||
repository: BaseRepository[_RepositoryModelT],
|
||
) -> TableActions[_RepositoryModelT]:
|
||
return cast( # cast-ok: callers read only the field names the prisma row and repository model share
|
||
"TableActions[_RepositoryModelT]", repository.table
|
||
)
|
||
|
||
|
||
def _deleted_verification_token_table(
|
||
prisma_client: PrismaClient,
|
||
) -> "TableActions[prisma_models.LiteLLM_DeletedVerificationToken]":
|
||
return DeletedVerificationTokenRepository(prisma_client).table
|
||
|
||
|
||
def _deprecated_verification_token_table(
|
||
prisma_client: PrismaClient,
|
||
) -> "TableActions[prisma_models.LiteLLM_DeprecatedVerificationToken]":
|
||
return DeprecatedVerificationTokenRepository(prisma_client).table
|
||
|
||
|
||
def _user_table(prisma_client: PrismaClient) -> TableActions[_UserRowLike]:
|
||
return UserRepository(prisma_client).table
|
||
|
||
|
||
def _credentials_table(prisma_client: PrismaClient) -> TableActions[CredentialItem]:
|
||
return cast( # cast-ok: the rotation loop reads and rewrites these rows through CredentialItem names only
|
||
"TableActions[CredentialItem]", CredentialsRepository(prisma_client).table
|
||
)
|
||
|
||
|
||
def _config_table(prisma_client: PrismaClient) -> _ConfigTableActions:
|
||
return cast( # cast-ok: ConfigRepository.table hides the write actions this module needs on that same object
|
||
"_ConfigTableActions", ConfigRepository(prisma_client).table
|
||
)
|
||
|
||
|
||
async def _check_custom_key_allowed(custom_key_value: str | None) -> None:
|
||
"""Raise 403 if custom API keys are disabled and a custom key was provided."""
|
||
if custom_key_value is None:
|
||
return
|
||
|
||
ui_settings: Final = await get_ui_settings_cached()
|
||
if ui_settings.get("disable_custom_api_keys", False) is True:
|
||
verbose_proxy_logger.warning("Custom API key rejected: disable_custom_api_keys is enabled")
|
||
raise HTTPException(
|
||
status_code=403,
|
||
detail={"error": "Custom API key values are disabled by your administrator. Keys must be auto-generated."},
|
||
)
|
||
|
||
|
||
def _is_team_key(data: GenerateKeyRequest | LiteLLM_VerificationToken):
|
||
return data.team_id is not None
|
||
|
||
|
||
def _get_user_in_team(team_table: LiteLLM_TeamTableCachedObj, user_id: str | None) -> Member | None:
|
||
if user_id is None:
|
||
return None
|
||
for member in team_table.members_with_roles:
|
||
if member.user_id is not None and member.user_id == user_id:
|
||
return member
|
||
|
||
return None
|
||
|
||
|
||
def _calculate_key_rotation_time(rotation_interval: str) -> datetime:
|
||
"""
|
||
Helper function to calculate the next rotation time for a key based on the rotation interval.
|
||
|
||
Args:
|
||
rotation_interval: String representing the rotation interval (e.g., '30d', '90d', '1h')
|
||
|
||
Returns:
|
||
datetime: The calculated next rotation time in UTC
|
||
"""
|
||
now: Final = datetime.now(timezone.utc)
|
||
interval_seconds: Final = duration_in_seconds(rotation_interval)
|
||
return now + timedelta(seconds=interval_seconds)
|
||
|
||
|
||
def _set_key_rotation_fields(
|
||
data: dict,
|
||
auto_rotate: bool,
|
||
rotation_interval: str | None,
|
||
existing_key_alias: str | None = None,
|
||
) -> None:
|
||
"""
|
||
Helper function to set rotation fields in key data if auto_rotate is enabled.
|
||
|
||
Args:
|
||
data: Dictionary to update with rotation fields
|
||
auto_rotate: Whether auto rotation is enabled
|
||
rotation_interval: The rotation interval string (required if auto_rotate is True)
|
||
existing_key_alias: The existing key alias from the database (if any)
|
||
"""
|
||
if auto_rotate and rotation_interval:
|
||
if (
|
||
litellm._key_management_settings is not None
|
||
and litellm._key_management_settings.store_virtual_keys is True
|
||
and data.get("key_alias") is None
|
||
and existing_key_alias is None
|
||
):
|
||
raise ProxyException(
|
||
message="key_alias is required when auto_rotate=True and store_virtual_keys is enabled. This ensures stable secret naming during rotation.",
|
||
type=ProxyErrorTypes.bad_request_error,
|
||
param="key_alias",
|
||
code=400,
|
||
)
|
||
data.update(
|
||
{
|
||
"auto_rotate": auto_rotate,
|
||
"rotation_interval": rotation_interval,
|
||
"key_rotation_at": _calculate_key_rotation_time(rotation_interval),
|
||
}
|
||
)
|
||
|
||
|
||
def _is_allowed_to_make_key_request(
|
||
user_api_key_dict: UserAPIKeyAuth,
|
||
user_id: str | None,
|
||
team_id: str | None,
|
||
) -> bool:
|
||
"""
|
||
Assert user only creates/updates keys for themselves
|
||
|
||
Relevant issue: https://github.com/BerriAI/litellm/issues/7336
|
||
"""
|
||
## BASE CASE - PROXY ADMIN
|
||
if user_api_key_dict.user_role is not None and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
|
||
return True
|
||
|
||
if user_id is not None:
|
||
assert user_id == user_api_key_dict.user_id, (
|
||
f"User can only create keys for themselves. Got user_id={user_id}, Your ID={user_api_key_dict.user_id}"
|
||
)
|
||
|
||
if team_id is not None:
|
||
if user_api_key_dict.team_id is not None and user_api_key_dict.team_id == UI_TEAM_ID:
|
||
return True # handle https://github.com/BerriAI/litellm/issues/7482
|
||
|
||
return True
|
||
|
||
|
||
def _team_key_operation_team_member_check(
|
||
assigned_user_id: str | None,
|
||
team_table: LiteLLM_TeamTableCachedObj,
|
||
user_api_key_dict: UserAPIKeyAuth,
|
||
team_key_generation: TeamUIKeyGenerationConfig,
|
||
route: KeyManagementRoutes,
|
||
):
|
||
if assigned_user_id is not None:
|
||
key_assigned_user_in_team: Final = _get_user_in_team(team_table=team_table, user_id=assigned_user_id)
|
||
|
||
if key_assigned_user_in_team is None:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail=f"User={assigned_user_id} not assigned to team={team_table.team_id}",
|
||
)
|
||
|
||
team_member_object: Final = _get_user_in_team(team_table=team_table, user_id=user_api_key_dict.user_id)
|
||
|
||
is_admin: Final = (
|
||
user_api_key_dict.user_role is not None and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
|
||
)
|
||
|
||
if is_admin:
|
||
return True
|
||
elif team_member_object is None:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail=f"User={user_api_key_dict.user_id} not assigned to team={team_table.team_id}",
|
||
)
|
||
elif (
|
||
"allowed_team_member_roles" in team_key_generation
|
||
and team_member_object.role not in team_key_generation["allowed_team_member_roles"]
|
||
):
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail=f"Team member role {team_member_object.role} not in allowed_team_member_roles={team_key_generation['allowed_team_member_roles']}",
|
||
)
|
||
|
||
TeamMemberPermissionChecks.does_team_member_have_permissions_for_endpoint(
|
||
team_member_object=team_member_object,
|
||
team_table=team_table,
|
||
route=route,
|
||
)
|
||
return True
|
||
|
||
|
||
def _key_generation_required_param_check(data: GenerateKeyRequest, required_params: list[str] | None):
|
||
if required_params is None:
|
||
return True
|
||
|
||
data_dict: Final = data.model_dump(exclude_unset=True)
|
||
for param in required_params:
|
||
if param not in data_dict:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail=f"Required param {param} not in data",
|
||
)
|
||
return True
|
||
|
||
|
||
def _team_key_generation_check(
|
||
team_table: LiteLLM_TeamTableCachedObj,
|
||
user_api_key_dict: UserAPIKeyAuth,
|
||
data: GenerateKeyRequest,
|
||
route: KeyManagementRoutes,
|
||
):
|
||
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
|
||
return True
|
||
if litellm.key_generation_settings is not None and "team_key_generation" in litellm.key_generation_settings:
|
||
_team_key_generation = litellm.key_generation_settings["team_key_generation"]
|
||
else:
|
||
_team_key_generation = TeamUIKeyGenerationConfig(
|
||
allowed_team_member_roles=["admin", "user"],
|
||
)
|
||
|
||
_team_key_operation_team_member_check(
|
||
assigned_user_id=data.user_id,
|
||
team_table=team_table,
|
||
user_api_key_dict=user_api_key_dict,
|
||
team_key_generation=_team_key_generation,
|
||
route=route,
|
||
)
|
||
_key_generation_required_param_check(
|
||
data,
|
||
_team_key_generation.get("required_params"),
|
||
)
|
||
|
||
# Field-level opt-in: non-admin members may only assign access groups when
|
||
# the team has enabled KEY_ACCESS_GROUP_ASSIGNMENT.
|
||
TeamMemberPermissionChecks.enforce_member_can_assign_access_groups(
|
||
user_api_key_dict=user_api_key_dict,
|
||
team_table=team_table,
|
||
access_group_ids=data.access_group_ids,
|
||
)
|
||
|
||
return True
|
||
|
||
|
||
def _personal_key_membership_check(
|
||
user_api_key_dict: UserAPIKeyAuth,
|
||
personal_key_generation: PersonalUIKeyGenerationConfig | None,
|
||
):
|
||
if personal_key_generation is None or "allowed_user_roles" not in personal_key_generation:
|
||
return True
|
||
|
||
if user_api_key_dict.user_role not in personal_key_generation["allowed_user_roles"]:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail=f"Personal key creation has been restricted by admin. Allowed roles={personal_key_generation['allowed_user_roles']}. Your role={user_api_key_dict.user_role}",
|
||
)
|
||
|
||
return True
|
||
|
||
|
||
def _object_permission_to_dict(
|
||
object_permission: LiteLLM_ObjectPermissionBase | None,
|
||
) -> ObjectPermissionDict | None:
|
||
if object_permission is None:
|
||
return None
|
||
return cast(ObjectPermissionDict, object_permission.model_dump(exclude_unset=True))
|
||
|
||
|
||
def _personal_key_generation_check(user_api_key_dict: UserAPIKeyAuth, data: GenerateKeyRequest):
|
||
TeamMemberPermissionChecks.enforce_member_can_assign_access_groups(
|
||
user_api_key_dict=user_api_key_dict,
|
||
team_table=None,
|
||
access_group_ids=data.access_group_ids,
|
||
)
|
||
|
||
if (
|
||
litellm.key_generation_settings is None
|
||
or litellm.key_generation_settings.get("personal_key_generation") is None
|
||
):
|
||
return True
|
||
|
||
_personal_key_generation: Final = litellm.key_generation_settings["personal_key_generation"]
|
||
|
||
_personal_key_membership_check(
|
||
user_api_key_dict,
|
||
personal_key_generation=_personal_key_generation,
|
||
)
|
||
|
||
_key_generation_required_param_check(
|
||
data,
|
||
_personal_key_generation.get("required_params"),
|
||
)
|
||
|
||
return True
|
||
|
||
|
||
def key_generation_check(
|
||
team_table: LiteLLM_TeamTableCachedObj | None,
|
||
user_api_key_dict: UserAPIKeyAuth,
|
||
data: GenerateKeyRequest,
|
||
route: KeyManagementRoutes,
|
||
) -> bool:
|
||
"""
|
||
Check if admin has restricted key creation to certain roles for teams or individuals
|
||
"""
|
||
|
||
## check if key is for team or individual
|
||
is_team_key: Final = _is_team_key(data=data)
|
||
_is_admin: Final = (
|
||
user_api_key_dict.user_role is not None and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
|
||
)
|
||
if is_team_key:
|
||
if team_table is None and litellm.key_generation_settings is not None:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail=f"Unable to find team object in database. Team ID: {data.team_id}",
|
||
)
|
||
elif team_table is None:
|
||
if _is_admin:
|
||
return True # admins can assign team_id without team table
|
||
# Non-admin callers must have a valid team (LIT-1884)
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail=f"Unable to find team object in database. Team ID: {data.team_id}",
|
||
)
|
||
return _team_key_generation_check(
|
||
team_table=team_table,
|
||
user_api_key_dict=user_api_key_dict,
|
||
data=data,
|
||
route=route,
|
||
)
|
||
else:
|
||
return _personal_key_generation_check(user_api_key_dict=user_api_key_dict, data=data)
|
||
|
||
|
||
def raise_on_invalid_key_logging_config(metadata: Mapping[str, object] | None) -> None:
|
||
"""Key-level logging writes go through key metadata, not /team/callback.
|
||
|
||
Without this the same New Relic config the team endpoint rejects would be
|
||
accepted here and then silently ignored or misrouted at request time.
|
||
"""
|
||
error: Final = logging_metadata_config_error(metadata)
|
||
if error is not None:
|
||
raise HTTPException(status_code=400, detail={"error": error}) # mutable-ok: FastAPI detail contract
|
||
|
||
|
||
def common_key_access_checks(
|
||
user_api_key_dict: UserAPIKeyAuth,
|
||
data: GenerateKeyRequest | UpdateKeyRequest,
|
||
llm_router: Router | None,
|
||
premium_user: bool,
|
||
user_id: str | None = None,
|
||
) -> Literal[True]:
|
||
"""
|
||
Check if user is allowed to make a key request, for this key
|
||
"""
|
||
try:
|
||
_is_allowed_to_make_key_request(
|
||
user_api_key_dict=user_api_key_dict,
|
||
user_id=user_id or data.user_id,
|
||
team_id=data.team_id,
|
||
)
|
||
except AssertionError as e:
|
||
raise HTTPException(
|
||
status_code=403,
|
||
detail=str(e),
|
||
)
|
||
except Exception as e:
|
||
raise HTTPException(
|
||
status_code=500,
|
||
detail=str(e),
|
||
)
|
||
|
||
_check_model_access_group(
|
||
models=data.models,
|
||
llm_router=llm_router,
|
||
premium_user=premium_user,
|
||
)
|
||
return True
|
||
|
||
|
||
router: Final = APIRouter()
|
||
|
||
|
||
def handle_key_type(data: GenerateKeyRequest, data_json: dict) -> dict:
|
||
"""
|
||
Handle the key type.
|
||
"""
|
||
key_type: Final = data.key_type
|
||
if key_type is None:
|
||
data_json.pop("key_type", None)
|
||
return data_json
|
||
data_json["key_type"] = key_type.value
|
||
if key_type == LiteLLMKeyType.LLM_API:
|
||
data_json["allowed_routes"] = ["llm_api_routes"]
|
||
elif key_type == LiteLLMKeyType.MANAGEMENT:
|
||
data_json["allowed_routes"] = ["management_routes"]
|
||
elif key_type == LiteLLMKeyType.READ_ONLY:
|
||
data_json["allowed_routes"] = ["info_routes"]
|
||
return data_json
|
||
|
||
|
||
_NON_ADMIN_SAFE_ALLOWED_ROUTES_PRESETS: Final = frozenset({"llm_api_routes", "info_routes"})
|
||
|
||
|
||
def _validate_caller_can_change_key_ownership(
|
||
data: BaseModel | None,
|
||
existing_key_row: LiteLLM_VerificationToken,
|
||
user_api_key_dict: UserAPIKeyAuth,
|
||
) -> None:
|
||
"""
|
||
Non-admin callers must not rebind a key's ``user_id`` to a different
|
||
user. The ``user_id`` on a verification token is what
|
||
``_return_user_api_key_auth_obj`` resolves against ``litellm_usertable``
|
||
to derive the request's role; a non-admin rebinding their own key's
|
||
``user_id`` to a ``PROXY_ADMIN`` row promotes themselves.
|
||
|
||
``/key/update`` already enforces this inline; ``/key/regenerate`` did
|
||
not. Sharing the check keeps both endpoints — and any future
|
||
regenerate-style endpoint — consistent.
|
||
"""
|
||
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
|
||
return
|
||
if data is None:
|
||
return
|
||
# Distinguish "user_id omitted" from "user_id explicitly set to None".
|
||
# Both leave ``getattr(data, 'user_id', None)`` at None, but only the
|
||
# explicit-null variant survives ``model_dump(exclude_unset=True)`` in
|
||
# ``prepare_key_update_data`` and writes NULL to the token row —
|
||
# detaching the key from its user and bypassing the user-row
|
||
# role check on subsequent requests.
|
||
fields_set: Final = getattr(data, "model_fields_set", None) or set()
|
||
if "user_id" not in fields_set:
|
||
return
|
||
incoming_user_id: Final = getattr(data, "user_id", None)
|
||
if incoming_user_id is None or incoming_user_id == "":
|
||
raise HTTPException(
|
||
status_code=403,
|
||
detail="Non-admin users cannot remove the user_id from a key.",
|
||
)
|
||
existing_user_id: Final = getattr(existing_key_row, "user_id", None)
|
||
if incoming_user_id != existing_user_id:
|
||
raise HTTPException(
|
||
status_code=403,
|
||
detail=(
|
||
f"Non-admin caller is not allowed to rebind the key from "
|
||
f"user={existing_user_id} to user={incoming_user_id}"
|
||
),
|
||
)
|
||
|
||
|
||
def _check_allowed_routes_caller_permission(
|
||
allowed_routes: list | None,
|
||
user_api_key_dict: UserAPIKeyAuth,
|
||
*,
|
||
allowed_routes_was_provided: bool = False,
|
||
allow_safe_presets: bool = False,
|
||
) -> None:
|
||
"""
|
||
Require PROXY_ADMIN when `allowed_routes` is present in the request body,
|
||
unless the caller went through the `key_type` preset flow.
|
||
|
||
Raw-body call sites pass
|
||
`allowed_routes_was_provided="allowed_routes" in data.model_fields_set` so a
|
||
caller that omits the field (model default flows through) is distinct from
|
||
one that sends any explicit value.
|
||
|
||
Post-`handle_key_type` call sites pass `allow_safe_presets=True` with the
|
||
values derived by `handle_key_type`; those values are not from the request
|
||
body, so `allowed_routes_was_provided` stays False and the safe-preset
|
||
carve-out below accepts any list of tokens in
|
||
`_NON_ADMIN_SAFE_ALLOWED_ROUTES_PRESETS`.
|
||
"""
|
||
if not allowed_routes_was_provided and not allowed_routes:
|
||
return
|
||
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
|
||
return
|
||
if (
|
||
allow_safe_presets
|
||
and allowed_routes
|
||
and all(r in _NON_ADMIN_SAFE_ALLOWED_ROUTES_PRESETS for r in allowed_routes)
|
||
):
|
||
return
|
||
raise HTTPException(
|
||
status_code=403,
|
||
detail={
|
||
"error": (
|
||
"Only proxy admins can set `allowed_routes` on a key. "
|
||
"Use `key_type` to pick a preset route bucket instead."
|
||
)
|
||
},
|
||
)
|
||
|
||
|
||
def _check_permissions_caller_permission(
|
||
data: GenerateRequestBase,
|
||
user_api_key_dict: UserAPIKeyAuth,
|
||
) -> None:
|
||
"""
|
||
Require PROXY_ADMIN when `permissions` is present in the request body.
|
||
|
||
Presence is detected via `data.model_fields_set` so a caller that
|
||
omits the field (default flows through) is distinct from one that
|
||
sends any explicit value.
|
||
"""
|
||
permissions_in_request: Final = "permissions" in data.model_fields_set
|
||
if not permissions_in_request and not data.permissions:
|
||
return
|
||
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
|
||
return
|
||
raise HTTPException(
|
||
status_code=403,
|
||
detail={"error": "Only proxy admins can set `permissions`."},
|
||
)
|
||
|
||
|
||
def _check_budget_limits_delegation_ceiling(
|
||
budget_limits: list[BudgetLimitEntry] | None,
|
||
delegation_ceiling: float | None,
|
||
user_api_key_dict: UserAPIKeyAuth,
|
||
is_ui_session_team_key: bool,
|
||
team_table: LiteLLM_TeamTableCachedObj | None,
|
||
) -> None:
|
||
"""
|
||
Enforce three invariants on `budget_limits`:
|
||
|
||
- Every `budget_limits[*].max_budget` must be a finite number; applies
|
||
to every caller including proxy admin.
|
||
- A CLI session token caller may not set `budget_limits` on a personal
|
||
key (one with no `team_id`); mirrors the scalar `max_budget` guard in
|
||
`_common_key_generation_helper`.
|
||
- Non-admin callers may not set a window above their delegation ceiling.
|
||
"""
|
||
if not budget_limits:
|
||
return
|
||
non_finite: Final = next((w for w in budget_limits if not math.isfinite(w.max_budget)), None)
|
||
if non_finite is not None:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail={"error": (f"budget_limits entry max_budget ({non_finite.max_budget}) must be a finite number.")},
|
||
)
|
||
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
|
||
return
|
||
if is_ui_session_team_key:
|
||
return
|
||
if user_api_key_dict.is_session_token and team_table is None:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail={
|
||
"error": ("budget_limits cannot be set without specifying team_id when using a CLI session token.")
|
||
},
|
||
)
|
||
if delegation_ceiling is None:
|
||
return
|
||
over_ceiling: Final = next((w for w in budget_limits if w.max_budget > delegation_ceiling), None)
|
||
if over_ceiling is not None:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail={
|
||
"error": (
|
||
f"budget_limits entry max_budget ({over_ceiling.max_budget}) "
|
||
f"cannot exceed the caller's own max_budget ({delegation_ceiling})."
|
||
)
|
||
},
|
||
)
|
||
|
||
|
||
async def validate_team_id_used_in_service_account_request(
|
||
team_id: str | None,
|
||
prisma_client: PrismaClient | None,
|
||
):
|
||
"""
|
||
Validate team_id is used in the request body for generating a service account key
|
||
"""
|
||
if team_id is None:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail="team_id is required for service account keys. Please specify `team_id` in the request body.",
|
||
)
|
||
|
||
if prisma_client is None:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail="prisma_client is required for service account keys. Please specify `prisma_client` in the request body.",
|
||
)
|
||
|
||
# check if team_id exists in the database
|
||
team: Final = await _prisma_table(TeamRepository(prisma_client)).find_unique(
|
||
where={"team_id": team_id},
|
||
)
|
||
if team is None:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail="team_id does not exist in the database. Please specify a valid `team_id` in the request body.",
|
||
)
|
||
return True
|
||
|
||
|
||
_BUDGET_NUMERIC_KEYS = frozenset(["max_budget", "soft_budget", "max_parallel_requests", "tpm_limit", "rpm_limit"])
|
||
|
||
|
||
def _enforce_upperbound_key_params(
|
||
data: GenerateKeyRequest | UpdateKeyRequest,
|
||
fill_defaults: bool = True,
|
||
) -> None:
|
||
"""
|
||
Enforce upperbound limits on key parameters.
|
||
|
||
For key generation (fill_defaults=True): fills None values with upperbound defaults.
|
||
For key update (fill_defaults=False): only validates explicitly provided values.
|
||
"""
|
||
# Always reject NaN / Inf regardless of whether an upperbound config is set
|
||
# (GHSA-2rv4-xv66-fpjg): float('nan') passes every `< 0` check because
|
||
# nan < 0 is False, and spend >= nan is always False, permanently disabling
|
||
# budget enforcement for any key that carries it.
|
||
for elem in data:
|
||
key, value = elem
|
||
if key in _BUDGET_NUMERIC_KEYS and value is not None:
|
||
if not math.isfinite(value):
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail={"error": f"{key} must be a finite number. Received: {value}"},
|
||
)
|
||
|
||
if litellm.upperbound_key_generate_params is None:
|
||
return
|
||
|
||
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}"
|
||
},
|
||
)
|
||
|
||
|
||
async def _common_key_generation_helper(
|
||
data: GenerateKeyRequest,
|
||
user_api_key_dict: UserAPIKeyAuth,
|
||
litellm_changed_by: str | None,
|
||
team_table: LiteLLM_TeamTableCachedObj | None,
|
||
) -> GenerateKeyResponse:
|
||
from litellm.proxy.proxy_server import (
|
||
litellm_proxy_admin_name,
|
||
llm_router,
|
||
premium_user,
|
||
prisma_client,
|
||
)
|
||
|
||
common_key_access_checks(
|
||
user_api_key_dict=user_api_key_dict,
|
||
data=data,
|
||
llm_router=llm_router,
|
||
premium_user=premium_user,
|
||
)
|
||
|
||
validate_budget_duration(data.budget_duration)
|
||
raise_on_invalid_key_logging_config(data.metadata)
|
||
|
||
if data.throttle_on_budget_exceeded is True and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value:
|
||
raise HTTPException(
|
||
status_code=403,
|
||
detail={"error": "Only proxy admins can enable throttle_on_budget_exceeded on a key."},
|
||
)
|
||
|
||
enforce_output_token_estimates_are_admin_only(
|
||
data=data,
|
||
existing_metadata=None,
|
||
user_api_key_dict=user_api_key_dict,
|
||
entity="key",
|
||
)
|
||
enforce_batch_enqueued_token_limit_is_admin_only(
|
||
data=data,
|
||
existing_metadata=None,
|
||
user_api_key_dict=user_api_key_dict,
|
||
entity="key",
|
||
)
|
||
|
||
if data.metadata is not None and data.metadata.get("service_account_id") is not None and data.team_id is None:
|
||
await validate_team_id_used_in_service_account_request(
|
||
team_id=data.team_id,
|
||
prisma_client=prisma_client,
|
||
)
|
||
|
||
# Capture caller-supplied max_budget and team_id before any defaults or
|
||
# upperbound params can fill them, so the ceiling check and its team-key
|
||
# exemption key off what the caller explicitly requested, not a value that
|
||
# default_key_generate_params injected.
|
||
_requested_max_budget: Final = data.max_budget
|
||
_requested_team_id: Final = data.team_id
|
||
|
||
# check if user set default key/generate params on config.yaml
|
||
if litellm.default_key_generate_params is not None:
|
||
for elem in data:
|
||
key, value = elem
|
||
if (
|
||
value is None
|
||
and (key != "budget_duration" or key not in data.model_fields_set)
|
||
and key
|
||
in [
|
||
"max_budget",
|
||
"user_id",
|
||
"team_id",
|
||
"max_parallel_requests",
|
||
"tpm_limit",
|
||
"rpm_limit",
|
||
"budget_duration",
|
||
"duration",
|
||
]
|
||
):
|
||
default_value = litellm.default_key_generate_params.get(key)
|
||
if default_value is not None:
|
||
setattr(data, key, default_value)
|
||
elif key == "models" and value == []:
|
||
setattr(data, key, litellm.default_key_generate_params.get(key, []))
|
||
elif key == "metadata" and value == {}:
|
||
setattr(data, key, litellm.default_key_generate_params.get(key, {}))
|
||
|
||
# check if user set upperbound key/generate params on config.yaml
|
||
_enforce_upperbound_key_params(data, fill_defaults=True)
|
||
|
||
# Delegated-authority ceiling (GHSA-q775-qw9r-2r4g): a non-admin caller
|
||
# cannot grant a key a higher budget than their own authority.
|
||
is_ui_session_team_key = user_api_key_dict.team_id == UI_SESSION_TOKEN_TEAM_ID and _requested_team_id is not None
|
||
# Session tokens (lite login) carry max_budget=None to avoid a per-session
|
||
# LLM spend cap, but that None must not be read as "unlimited delegation
|
||
# authority". A personal key (no team) has no team-budget enforcement at
|
||
# request time, so a session token cannot delegate any budget for one.
|
||
if (
|
||
user_api_key_dict.is_session_token
|
||
and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value
|
||
and not is_ui_session_team_key
|
||
and _requested_max_budget is not None
|
||
and team_table is None
|
||
):
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail={
|
||
"error": (
|
||
f"max_budget ({_requested_max_budget}) cannot be set without "
|
||
"specifying team_id when using a CLI session token."
|
||
)
|
||
},
|
||
)
|
||
delegation_ceiling: Final = (
|
||
user_api_key_dict.max_budget
|
||
if user_api_key_dict.max_budget is not None
|
||
else (team_table.max_budget if user_api_key_dict.is_session_token and team_table is not None else None)
|
||
)
|
||
if (
|
||
user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value
|
||
and not is_ui_session_team_key
|
||
and _requested_max_budget is not None
|
||
and delegation_ceiling is not None
|
||
and _requested_max_budget > delegation_ceiling
|
||
):
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail={
|
||
"error": (
|
||
f"max_budget ({_requested_max_budget}) cannot exceed the caller's "
|
||
f"own max_budget ({delegation_ceiling})."
|
||
)
|
||
},
|
||
)
|
||
|
||
_check_budget_limits_delegation_ceiling(
|
||
budget_limits=data.budget_limits,
|
||
delegation_ceiling=delegation_ceiling,
|
||
user_api_key_dict=user_api_key_dict,
|
||
is_ui_session_team_key=is_ui_session_team_key,
|
||
team_table=team_table,
|
||
)
|
||
_check_permissions_caller_permission(
|
||
data=data,
|
||
user_api_key_dict=user_api_key_dict,
|
||
)
|
||
|
||
# APPLY ENTERPRISE KEY MANAGEMENT PARAMS
|
||
try:
|
||
from litellm_enterprise.proxy.management_endpoints.key_management_endpoints import (
|
||
apply_enterprise_key_management_params,
|
||
)
|
||
|
||
data = apply_enterprise_key_management_params(data, team_table)
|
||
except Exception as e:
|
||
verbose_proxy_logger.debug(
|
||
"litellm.proxy.proxy_server.generate_key_fn(): Enterprise key management params not applied - %s", e
|
||
)
|
||
|
||
# TODO: @ishaan-jaff: Migrate all budget tracking to use LiteLLM_BudgetTable
|
||
_budget_id = data.budget_id
|
||
if prisma_client is not None and data.soft_budget is not None:
|
||
# create the Budget Row for the LiteLLM Verification Token
|
||
budget_row: Final = LiteLLM_BudgetTable(
|
||
soft_budget=data.soft_budget,
|
||
model_max_budget=data.model_max_budget or {},
|
||
)
|
||
new_budget: Final = prisma_client.jsonify_object(budget_row.json(exclude_none=True))
|
||
|
||
_budget: Final[prisma_models.LiteLLM_BudgetTable] = await BudgetRepository(prisma_client).table.create(
|
||
data={
|
||
**new_budget,
|
||
"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,
|
||
}
|
||
)
|
||
_budget_id = getattr(_budget, "budget_id", None)
|
||
|
||
# ADD METADATA FIELDS
|
||
# Set Management Endpoint Metadata Fields
|
||
for field in LiteLLM_ManagementEndpoint_MetadataFields_Premium:
|
||
if getattr(data, field, None) is not None:
|
||
_set_object_metadata_field(
|
||
object_data=data,
|
||
field_name=field,
|
||
value=getattr(data, field),
|
||
)
|
||
delattr(data, field)
|
||
|
||
for field in LiteLLM_ManagementEndpoint_MetadataFields:
|
||
if getattr(data, field, None) is not None:
|
||
_set_object_metadata_field(
|
||
object_data=data,
|
||
field_name=field,
|
||
value=getattr(data, field),
|
||
)
|
||
delattr(data, field)
|
||
|
||
data_json = data.model_dump(exclude_unset=True, exclude_none=True)
|
||
|
||
data_json = handle_key_type(data, data_json)
|
||
|
||
# Re-check allowed_routes after handle_key_type, since key_type can derive
|
||
# an elevated bucket (e.g. ["management_routes"]) that wasn't present in
|
||
# the original request body. The safe presets produced by handle_key_type
|
||
# for non-elevated buckets are accepted here; the raw-body pre-checks at
|
||
# the entry of each handler keep their default strictness.
|
||
_check_allowed_routes_caller_permission(
|
||
allowed_routes=data_json.get("allowed_routes"),
|
||
user_api_key_dict=user_api_key_dict,
|
||
allow_safe_presets=True,
|
||
)
|
||
|
||
# if we get max_budget passed to /key/generate, then use it as key_max_budget. Since generate_key_helper_fn is used to make new users
|
||
if "max_budget" in data_json:
|
||
data_json["key_max_budget"] = data_json.pop("max_budget", None)
|
||
if _budget_id is not None:
|
||
data_json["budget_id"] = _budget_id
|
||
|
||
# Only set budget_duration on key when explicitly provided. Keys with budget_id
|
||
# but no explicit budget_duration follow their linked budget tier's schedule;
|
||
# reset_budget_for_litellm_budget_table() resets them when the tier resets.
|
||
# This avoids duplicating budget_duration on keys so tier updates apply automatically.
|
||
if "budget_duration" in data_json:
|
||
data_json["key_budget_duration"] = data_json.pop("budget_duration", None)
|
||
|
||
if user_api_key_dict.user_id is not None:
|
||
data_json["created_by"] = user_api_key_dict.user_id
|
||
data_json["updated_by"] = user_api_key_dict.user_id
|
||
|
||
# Set tags on the new key
|
||
if "tags" in data_json:
|
||
from litellm.proxy.proxy_server import premium_user
|
||
|
||
if premium_user is not True and data_json["tags"] is not None:
|
||
raise ValueError(f"Only premium users can add tags to keys. {CommonProxyErrors.not_premium_user.value}")
|
||
|
||
_metadata: Final = data_json.get("metadata")
|
||
if not _metadata:
|
||
data_json["metadata"] = {"tags": data_json["tags"]}
|
||
else:
|
||
data_json["metadata"]["tags"] = data_json["tags"]
|
||
|
||
data_json.pop("tags")
|
||
|
||
# Validate MCP servers in object_permission are within team scope
|
||
_is_proxy_admin_caller: Final = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
|
||
normalized_object_permission: Final = await validate_key_mcp_servers_against_team(
|
||
object_permission=data_json.get("object_permission"),
|
||
team_obj=team_table,
|
||
prisma_client=prisma_client,
|
||
is_proxy_admin=_is_proxy_admin_caller,
|
||
)
|
||
if normalized_object_permission is not None:
|
||
data_json["object_permission"] = normalized_object_permission
|
||
await validate_key_search_tools_against_team(
|
||
object_permission=data_json.get("object_permission"),
|
||
team_obj=team_table,
|
||
is_proxy_admin=_is_proxy_admin_caller,
|
||
)
|
||
await validate_key_vector_stores_against_team(
|
||
object_permission=data_json.get("object_permission"),
|
||
team_obj=team_table,
|
||
is_proxy_admin=_is_proxy_admin_caller,
|
||
)
|
||
|
||
# Merge default_key_generate_params.object_permission in *after* the team-scope
|
||
# checks above, so an admin-configured default (e.g. vector_stores, search_tools)
|
||
# is never mistaken for a caller-requested permission and rejected by those
|
||
# non-admin/no-team checks. Only fields the caller left unset are filled in.
|
||
_default_object_permission: Final = (
|
||
litellm.default_key_generate_params.get("object_permission")
|
||
if litellm.default_key_generate_params is not None
|
||
else None
|
||
)
|
||
if isinstance(_default_object_permission, dict):
|
||
_caller_object_permission: Final = data_json.get("object_permission")
|
||
if _caller_object_permission is None:
|
||
data_json["object_permission"] = dict(_default_object_permission)
|
||
elif isinstance(_caller_object_permission, dict):
|
||
for _op_field, _op_default_value in _default_object_permission.items():
|
||
_caller_object_permission.setdefault(_op_field, _op_default_value)
|
||
|
||
data_json = await _set_object_permission(
|
||
data_json=data_json,
|
||
prisma_client=prisma_client,
|
||
)
|
||
|
||
_validate_key_alias_format(key_alias=data_json.get("key_alias", None))
|
||
|
||
await _enforce_unique_key_alias(
|
||
key_alias=data_json.get("key_alias", None),
|
||
prisma_client=prisma_client,
|
||
)
|
||
|
||
# Reject custom key values if disabled by admin
|
||
await _check_custom_key_allowed(data.key)
|
||
|
||
# Validate user-provided key format
|
||
if data.key is not None and not data.key.startswith("sk-"):
|
||
_masked: Final = f"{data.key[:4]}****{data.key[-4:]}" if len(data.key) > 8 else "****"
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail={"error": f"Invalid key format. LiteLLM Virtual Key must start with 'sk-'. Received: {_masked}"},
|
||
)
|
||
|
||
if data.key is not None and len(data.key) < MINIMUM_CUSTOM_KEY_LENGTH:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail={
|
||
"error": f"Invalid key format. LiteLLM Virtual Key must be at least {MINIMUM_CUSTOM_KEY_LENGTH} characters long."
|
||
},
|
||
)
|
||
|
||
# check org key limits - done here to handle inheriting org id from team
|
||
if data.organization_id is not None:
|
||
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
|
||
|
||
if prisma_client:
|
||
# Mirror the membership rule applied to /key/update: when the
|
||
# caller specifies an organization_id, require that they are a
|
||
# member of (or proxy admin over) the target organization.
|
||
_is_proxy_admin: Final = (
|
||
user_api_key_dict.user_role is not None
|
||
and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
|
||
)
|
||
_org_inherited_from_team: Final = (
|
||
team_table is not None
|
||
and team_table.organization_id is not None
|
||
and data.organization_id == team_table.organization_id
|
||
)
|
||
if not _is_proxy_admin and not _org_inherited_from_team:
|
||
await _validate_caller_can_assign_key_org(
|
||
user_api_key_dict=user_api_key_dict,
|
||
organization_id=data.organization_id,
|
||
prisma_client=prisma_client,
|
||
)
|
||
|
||
org_table: Final = await get_org_object(
|
||
org_id=data.organization_id,
|
||
user_api_key_cache=user_api_key_cache,
|
||
prisma_client=prisma_client,
|
||
)
|
||
if org_table is None:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail=f"Organization not found for organization_id={data.organization_id}",
|
||
)
|
||
await _check_org_key_limits(
|
||
org_table=org_table,
|
||
data=data,
|
||
prisma_client=prisma_client,
|
||
)
|
||
|
||
response = await generate_key_helper_fn(request_type="key", **data_json, table_name="key")
|
||
|
||
response["soft_budget"] = data.soft_budget # include the user-input soft budget in the response
|
||
|
||
response = GenerateKeyResponse.model_validate(response)
|
||
|
||
response.token = response.token_id # remap token to use the hash, and leave the key in the `key` field [TODO]: clean up generate_key_helper_fn to do this
|
||
|
||
asyncio.create_task(
|
||
KeyManagementEventHooks.async_key_generated_hook(
|
||
data=data,
|
||
response=response,
|
||
user_api_key_dict=user_api_key_dict,
|
||
litellm_changed_by=litellm_changed_by,
|
||
)
|
||
)
|
||
|
||
return response
|
||
|
||
|
||
def _check_key_model_specific_limits(
|
||
keys: Sequence[LiteLLM_VerificationToken],
|
||
data: GenerateKeyRequest | UpdateKeyRequest,
|
||
entity_rpm_limit: int | None,
|
||
entity_tpm_limit: int | None,
|
||
entity_model_rpm_limit_dict: dict[str, int],
|
||
entity_model_tpm_limit_dict: dict[str, int],
|
||
entity_type: str, # "team" or "organization"
|
||
) -> None:
|
||
"""
|
||
Generic function to check if a key is allocating model specific limits.
|
||
Raises an error if we're overallocating.
|
||
"""
|
||
model_rpm_limit: Final = getattr(data, "model_rpm_limit", None) or (
|
||
data.metadata.get("model_rpm_limit", None) if data.metadata else None
|
||
)
|
||
model_tpm_limit: Final = getattr(data, "model_tpm_limit", None) or (
|
||
data.metadata.get("model_tpm_limit", None) if data.metadata else None
|
||
)
|
||
if model_rpm_limit is None and model_tpm_limit is None:
|
||
return
|
||
|
||
# get total model specific tpm/rpm limit
|
||
model_specific_rpm_limit: Final[dict[str, int]] = {}
|
||
model_specific_tpm_limit: Final[dict[str, int]] = {}
|
||
|
||
for key in keys:
|
||
if key.metadata.get("model_rpm_limit", None) is not None:
|
||
for model, rpm_limit in key.metadata.get("model_rpm_limit", {}).items():
|
||
model_specific_rpm_limit[model] = model_specific_rpm_limit.get(model, 0) + rpm_limit
|
||
if key.metadata.get("model_tpm_limit", None) is not None:
|
||
for model, tpm_limit in key.metadata.get("model_tpm_limit", {}).items():
|
||
model_specific_tpm_limit[model] = model_specific_tpm_limit.get(model, 0) + tpm_limit
|
||
|
||
if model_rpm_limit is not None:
|
||
for model, rpm_limit in model_rpm_limit.items():
|
||
if entity_rpm_limit is not None and model_specific_rpm_limit.get(model, 0) + rpm_limit > entity_rpm_limit:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail=f"Allocated RPM limit={model_specific_rpm_limit.get(model, 0)} + Key RPM limit={rpm_limit} is greater than {entity_type} RPM limit={entity_rpm_limit}",
|
||
)
|
||
elif entity_model_rpm_limit_dict:
|
||
entity_model_specific_rpm_limit = entity_model_rpm_limit_dict.get(model)
|
||
if (
|
||
entity_model_specific_rpm_limit
|
||
and model_specific_rpm_limit.get(model, 0) + rpm_limit > entity_model_specific_rpm_limit
|
||
):
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail=f"Allocated RPM limit={model_specific_rpm_limit.get(model, 0)} + Key RPM limit={rpm_limit} is greater than {entity_type} RPM limit={entity_model_specific_rpm_limit}",
|
||
)
|
||
|
||
if model_tpm_limit is not None:
|
||
for model, tpm_limit in model_tpm_limit.items():
|
||
if entity_tpm_limit is not None and model_specific_tpm_limit.get(model, 0) + tpm_limit > entity_tpm_limit:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail=f"Allocated TPM limit={model_specific_tpm_limit.get(model, 0)} + Key TPM limit={tpm_limit} is greater than {entity_type} TPM limit={entity_tpm_limit}",
|
||
)
|
||
elif entity_model_tpm_limit_dict:
|
||
entity_model_specific_tpm_limit = entity_model_tpm_limit_dict.get(model)
|
||
if (
|
||
entity_model_specific_tpm_limit
|
||
and model_specific_tpm_limit.get(model, 0) + tpm_limit > entity_model_specific_tpm_limit
|
||
):
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail=f"Allocated TPM limit={model_specific_tpm_limit.get(model, 0)} + Key TPM limit={tpm_limit} is greater than {entity_type} TPM limit={entity_model_specific_tpm_limit}",
|
||
)
|
||
|
||
|
||
def _check_key_rpm_tpm_limits(
|
||
keys: Sequence[LiteLLM_VerificationToken],
|
||
data: GenerateKeyRequest | UpdateKeyRequest,
|
||
entity_rpm_limit: int | None,
|
||
entity_tpm_limit: int | None,
|
||
entity_type: str, # "team" or "organization"
|
||
) -> None:
|
||
"""
|
||
Generic function to check if a key is allocating rpm/tpm limits.
|
||
Raises an error if we're overallocating.
|
||
"""
|
||
if keys is not None and len(keys) > 0:
|
||
allocated_tpm = sum(key.tpm_limit for key in keys if key.tpm_limit is not None)
|
||
allocated_rpm = sum(key.rpm_limit for key in keys if key.rpm_limit is not None)
|
||
else:
|
||
allocated_tpm = 0
|
||
allocated_rpm = 0
|
||
|
||
if (
|
||
data.tpm_limit is not None
|
||
and entity_tpm_limit is not None
|
||
and data.tpm_limit + allocated_tpm > entity_tpm_limit
|
||
):
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail=f"Allocated TPM limit={allocated_tpm} + Key TPM limit={data.tpm_limit} is greater than {entity_type} TPM limit={entity_tpm_limit}",
|
||
)
|
||
if (
|
||
data.rpm_limit is not None
|
||
and entity_rpm_limit is not None
|
||
and data.rpm_limit + allocated_rpm > entity_rpm_limit
|
||
):
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail=f"Allocated RPM limit={allocated_rpm} + Key RPM limit={data.rpm_limit} is greater than {entity_type} RPM limit={entity_rpm_limit}",
|
||
)
|
||
|
||
|
||
def check_team_key_model_specific_limits(
|
||
keys: Sequence[LiteLLM_VerificationToken],
|
||
team_table: LiteLLM_TeamTableCachedObj,
|
||
data: GenerateKeyRequest | UpdateKeyRequest,
|
||
) -> None:
|
||
"""
|
||
Check if the team key is allocating model specific limits. If so, raise an error if we're overallocating.
|
||
"""
|
||
entity_model_rpm_limit_dict = {}
|
||
entity_model_tpm_limit_dict = {}
|
||
if team_table.metadata:
|
||
entity_model_rpm_limit_dict = team_table.metadata.get("model_rpm_limit", {})
|
||
entity_model_tpm_limit_dict = team_table.metadata.get("model_tpm_limit", {})
|
||
|
||
_check_key_model_specific_limits(
|
||
keys=keys,
|
||
data=data,
|
||
entity_rpm_limit=team_table.rpm_limit,
|
||
entity_tpm_limit=team_table.tpm_limit,
|
||
entity_model_rpm_limit_dict=entity_model_rpm_limit_dict,
|
||
entity_model_tpm_limit_dict=entity_model_tpm_limit_dict,
|
||
entity_type="team",
|
||
)
|
||
|
||
|
||
def check_team_key_rpm_tpm_limits(
|
||
keys: Sequence[LiteLLM_VerificationToken],
|
||
team_table: LiteLLM_TeamTableCachedObj,
|
||
data: GenerateKeyRequest | UpdateKeyRequest,
|
||
) -> None:
|
||
"""
|
||
Check if the team key is allocating rpm/tpm limits. If so, raise an error if we're overallocating.
|
||
"""
|
||
_check_key_rpm_tpm_limits(
|
||
keys=keys,
|
||
data=data,
|
||
entity_rpm_limit=team_table.rpm_limit,
|
||
entity_tpm_limit=team_table.tpm_limit,
|
||
entity_type="team",
|
||
)
|
||
|
||
|
||
async def _check_team_key_limits(
|
||
team_table: LiteLLM_TeamTableCachedObj,
|
||
data: GenerateKeyRequest | UpdateKeyRequest,
|
||
prisma_client: PrismaClient,
|
||
) -> None:
|
||
"""
|
||
Check if the team key is allocating guaranteed throughput limits. If so, raise an error if we're overallocating.
|
||
|
||
Only runs check if tpm_limit_type or rpm_limit_type is "guaranteed_throughput"
|
||
"""
|
||
if data.tpm_limit_type != "guaranteed_throughput" and data.rpm_limit_type != "guaranteed_throughput":
|
||
return
|
||
# get all team keys
|
||
# calculate allocated tpm/rpm limit
|
||
# check if specified tpm/rpm limit is greater than allocated tpm/rpm limit
|
||
|
||
keys = await _prisma_table(VerificationTokenRepository(prisma_client)).find_many(
|
||
where={"team_id": team_table.team_id},
|
||
)
|
||
# Exclude the key being updated to avoid double-counting its limits.
|
||
# data.key may be a raw key (sk-...) or a pre-hashed token_id.
|
||
if isinstance(data, UpdateKeyRequest) and data.key is not None:
|
||
hashed_key: Final = _hash_token_if_needed(data.key)
|
||
keys = [key for key in keys if key.token != hashed_key]
|
||
check_team_key_model_specific_limits(
|
||
keys=keys,
|
||
team_table=team_table,
|
||
data=data,
|
||
)
|
||
check_team_key_rpm_tpm_limits(
|
||
keys=keys,
|
||
team_table=team_table,
|
||
data=data,
|
||
)
|
||
|
||
|
||
_INHERITED_MODEL_SENTINELS: Final = frozenset(
|
||
{SpecialModelNames.all_team_models.value, SpecialModelNames.all_proxy_models.value}
|
||
)
|
||
|
||
|
||
async def _check_project_key_limits(
|
||
project_id: str,
|
||
data: GenerateKeyRequest | UpdateKeyRequest,
|
||
prisma_client: PrismaClient,
|
||
user_api_key_cache: UserApiKeyCache,
|
||
) -> None:
|
||
"""
|
||
Validate that key's models and budget respect its project's limits.
|
||
|
||
- Key models must be a subset of project models, except the all-team-models / all-proxy-models
|
||
sentinels, which inherit a parent scope and are narrowed by the project at request time
|
||
- Key max_budget must be <= project max_budget
|
||
"""
|
||
project_obj: Final = await get_project_object(
|
||
project_id=project_id,
|
||
prisma_client=prisma_client,
|
||
user_api_key_cache=user_api_key_cache,
|
||
)
|
||
|
||
if project_obj is None:
|
||
raise HTTPException(
|
||
status_code=404,
|
||
detail={"error": f"Project not found, project_id={project_id}"},
|
||
)
|
||
|
||
# Validate key models are a subset of project models
|
||
if data.models and len(project_obj.models) > 0:
|
||
for m in data.models:
|
||
if m not in project_obj.models and m not in _INHERITED_MODEL_SENTINELS:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail={
|
||
"error": f"Model '{m}' not in project's allowed models. Project allowed models={project_obj.models}. Project: {project_id}"
|
||
},
|
||
)
|
||
|
||
# Validate key max_budget <= project max_budget
|
||
project_max_budget = None
|
||
if project_obj.litellm_budget_table is not None:
|
||
project_max_budget = getattr(project_obj.litellm_budget_table, "max_budget", None)
|
||
|
||
if data.max_budget is not None and project_max_budget is not None and data.max_budget > project_max_budget:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail={
|
||
"error": f"Key max_budget ({data.max_budget}) exceeds project's max_budget ({project_max_budget}). Project: {project_id}"
|
||
},
|
||
)
|
||
|
||
|
||
def check_org_key_model_specific_limits(
|
||
keys: Sequence[LiteLLM_VerificationToken],
|
||
org_table: LiteLLM_OrganizationTable,
|
||
data: GenerateKeyRequest | UpdateKeyRequest,
|
||
) -> None:
|
||
"""
|
||
Check if the organization key is allocating model specific limits. If so, raise an error if we're overallocating.
|
||
"""
|
||
# Get org limits from budget table if available
|
||
entity_rpm_limit = None
|
||
entity_tpm_limit = None
|
||
entity_model_rpm_limit_dict = {}
|
||
entity_model_tpm_limit_dict = {}
|
||
|
||
if org_table.litellm_budget_table is not None:
|
||
entity_rpm_limit = org_table.litellm_budget_table.rpm_limit
|
||
entity_tpm_limit = org_table.litellm_budget_table.tpm_limit
|
||
|
||
if org_table.metadata:
|
||
entity_model_rpm_limit_dict = org_table.metadata.get("model_rpm_limit", {})
|
||
entity_model_tpm_limit_dict = org_table.metadata.get("model_tpm_limit", {})
|
||
|
||
_check_key_model_specific_limits(
|
||
keys=keys,
|
||
data=data,
|
||
entity_rpm_limit=entity_rpm_limit,
|
||
entity_tpm_limit=entity_tpm_limit,
|
||
entity_model_rpm_limit_dict=entity_model_rpm_limit_dict,
|
||
entity_model_tpm_limit_dict=entity_model_tpm_limit_dict,
|
||
entity_type="organization",
|
||
)
|
||
|
||
|
||
def check_org_key_rpm_tpm_limits(
|
||
keys: Sequence[LiteLLM_VerificationToken],
|
||
org_table: LiteLLM_OrganizationTable,
|
||
data: GenerateKeyRequest | UpdateKeyRequest,
|
||
) -> None:
|
||
"""
|
||
Check if the organization key is allocating rpm/tpm limits. If so, raise an error if we're overallocating.
|
||
"""
|
||
# Get org limits from budget table if available
|
||
entity_rpm_limit = None
|
||
entity_tpm_limit = None
|
||
|
||
if org_table.litellm_budget_table is not None:
|
||
entity_rpm_limit = org_table.litellm_budget_table.rpm_limit
|
||
entity_tpm_limit = org_table.litellm_budget_table.tpm_limit
|
||
|
||
_check_key_rpm_tpm_limits(
|
||
keys=keys,
|
||
data=data,
|
||
entity_rpm_limit=entity_rpm_limit,
|
||
entity_tpm_limit=entity_tpm_limit,
|
||
entity_type="organization",
|
||
)
|
||
|
||
|
||
async def _validate_caller_can_assign_key_org(
|
||
user_api_key_dict: UserAPIKeyAuth,
|
||
organization_id: str,
|
||
prisma_client: PrismaClient,
|
||
) -> None:
|
||
"""Reject ``/key/update`` requests that point a key at an organization
|
||
the caller does not belong to.
|
||
|
||
Mirrors the org-membership rule already enforced on ``/key/list`` in
|
||
``validate_key_list_check``. Proxy admins are checked at the call site.
|
||
"""
|
||
if user_api_key_dict.user_id is None:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_403_FORBIDDEN,
|
||
detail="Cannot assign a key to an organization without a user_id on the caller's token",
|
||
)
|
||
|
||
user_row: Final = await _prisma_table(UserRepository(prisma_client)).find_unique(
|
||
where={"user_id": user_api_key_dict.user_id},
|
||
include={"organization_memberships": True},
|
||
)
|
||
memberships: Final = getattr(user_row, "organization_memberships", None) if user_row else None
|
||
member_org_ids: Final = {
|
||
membership.organization_id for membership in (memberships or []) if membership.organization_id is not None
|
||
}
|
||
if organization_id not in member_org_ids:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_403_FORBIDDEN,
|
||
detail=f"Caller is not a member of organization_id={organization_id}",
|
||
)
|
||
|
||
|
||
async def _check_org_key_limits(
|
||
org_table: LiteLLM_OrganizationTable,
|
||
data: GenerateKeyRequest | UpdateKeyRequest,
|
||
prisma_client: PrismaClient,
|
||
) -> None:
|
||
"""
|
||
Check if the organization key is allocating guaranteed throughput limits. If so, raise an error if we're overallocating.
|
||
|
||
Only runs check if tpm_limit_type or rpm_limit_type is "guaranteed_throughput"
|
||
"""
|
||
|
||
rpm_limit_type: Final = getattr(data, "rpm_limit_type", None) or (
|
||
data.metadata.get("rpm_limit_type", None) if data.metadata else None
|
||
)
|
||
tpm_limit_type: Final = getattr(data, "tpm_limit_type", None) or (
|
||
data.metadata.get("tpm_limit_type", None) if data.metadata else None
|
||
)
|
||
|
||
if tpm_limit_type != "guaranteed_throughput" and rpm_limit_type != "guaranteed_throughput":
|
||
return
|
||
# get all organization keys
|
||
# calculate allocated tpm/rpm limit
|
||
# check if specified tpm/rpm limit is greater than allocated tpm/rpm limit
|
||
keys = await _prisma_table(VerificationTokenRepository(prisma_client)).find_many(
|
||
where={"organization_id": org_table.organization_id},
|
||
)
|
||
# Exclude the key being updated to avoid double-counting its limits.
|
||
# data.key may be a raw key (sk-...) or a pre-hashed token_id.
|
||
if isinstance(data, UpdateKeyRequest) and data.key is not None:
|
||
hashed_key: Final = _hash_token_if_needed(data.key)
|
||
keys = [key for key in keys if key.token != hashed_key]
|
||
check_org_key_model_specific_limits(
|
||
keys=keys,
|
||
org_table=org_table,
|
||
data=data,
|
||
)
|
||
check_org_key_rpm_tpm_limits(
|
||
keys=keys,
|
||
org_table=org_table,
|
||
data=data,
|
||
)
|
||
|
||
|
||
@router.post(
|
||
"/key/generate",
|
||
tags=["key management"],
|
||
dependencies=[Depends(user_api_key_auth)],
|
||
response_model=GenerateKeyResponse,
|
||
)
|
||
@management_endpoint_wrapper
|
||
async def generate_key_fn(
|
||
data: GenerateKeyRequest,
|
||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||
litellm_changed_by: str | None = Header(
|
||
None,
|
||
description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability",
|
||
),
|
||
):
|
||
"""
|
||
Generate an API key based on the provided data.
|
||
|
||
Docs: https://docs.litellm.ai/docs/proxy/virtual_keys
|
||
|
||
Parameters:
|
||
- duration: Optional[str] - Specify the length of time the token is valid for. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d").
|
||
- key_alias: Optional[str] - User defined key alias
|
||
- key: Optional[str] - User defined key value. Must start with 'sk-' and be at least 16 characters long. If not set, a 16-digit unique sk-key is created for you.
|
||
- team_id: Optional[str] - The team id of the key
|
||
- user_id: Optional[str] - The user id of the key
|
||
- agent_id: Optional[str] - The agent id associated with the key.
|
||
- organization_id: Optional[str] - The organization id of the key. If not set, and team_id is set, the organization id will be the same as the team id. If conflict, an error will be raised.
|
||
- project_id: Optional[str] - The project id of the key. When set, models and max_budget are validated against the project's limits.
|
||
- budget_id: Optional[str] - The budget id associated with the key. Created by calling `/budget/new`.
|
||
- models: Optional[list] - Model_name's a user is allowed to call. (if empty, key is allowed to call all models)
|
||
- aliases: Optional[dict] - Any alias mappings, on top of anything in the config.yaml model list. - https://docs.litellm.ai/docs/proxy/virtual_keys#managing-auth---upgradedowngrade-models
|
||
- config: Optional[dict] - any key-specific configs, overrides config in config.yaml
|
||
- spend: Optional[int] - Amount spent by key. Default is 0. Will be updated by proxy whenever key is used. https://docs.litellm.ai/docs/proxy/virtual_keys#managing-auth---tracking-spend
|
||
- send_invite_email: Optional[bool] - Whether to send an invite email to the user_id, with the generate key
|
||
- max_budget: Optional[float] - Specify max budget for a given key.
|
||
- budget_duration: Optional[str] - Budget is reset at the end of specified duration. If not set, budget is never reset. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d").
|
||
- max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x.
|
||
- metadata: Optional[dict] - Metadata for key, store information for key. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" }
|
||
- guardrails: Optional[List[str]] - List of active guardrails for the key
|
||
- policies: Optional[List[str]] - List of policy names to apply to the key. Policies define guardrails, conditions, and inheritance rules.
|
||
- disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key.
|
||
- throttle_on_budget_exceeded: Optional[bool] - When the key exceeds its max_budget, throttle its tpm/rpm to the global budget_exceeded_throttle_percentage instead of blocking the key entirely.
|
||
- enable_prompt_caching: Optional[bool] - Auto-inject prompt caching breakpoints (Anthropic cache_control markers) on requests made with this key. Anthropic and Bedrock Claude models only.
|
||
- permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false}
|
||
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget.
|
||
- budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||
- model_rpm_limit: Optional[dict] - key-specific model rpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific rpm limit.
|
||
- model_tpm_limit: Optional[dict] - key-specific model tpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific tpm limit.
|
||
- default_estimated_output_tokens: Optional[int] - Proxy admin only. Expected output tokens reserved for TPM limiting when a request omits max_tokens. Positive integer. Falls back to the team setting, then to the built-in estimate.
|
||
- default_estimated_output_tokens_per_model: Optional[dict] - Proxy admin only. Per-model override of the above. Example - {"gpt-4": 4096, "gpt-3.5-turbo": 1024}. Takes precedence over the key-wide value.
|
||
- mcp_rpm_limit: Optional[dict] - key-specific per-MCP-server rpm limit, keyed by MCP server name (alias if set, else the configured name). Example - {"github": 100, "slack": 200}. IF null or {} then no MCP-specific rpm limit.
|
||
- tag_rpm_limit: Optional[dict] - key-specific per-request-tag rpm limit, keyed by request tag. Example - {"cell-1": 1000, "cell-2": 500}. Each tag gets an independent counter; requests whose tag is absent fall back to the key-level rpm limit.
|
||
- tpm_limit_type: Optional[str] - Type of tpm limit. Options: "best_effort_throughput" (no error if we're overallocating tpm), "guaranteed_throughput" (raise an error if we're overallocating tpm), "dynamic" (dynamically exceed limit when no 429 errors). Defaults to "best_effort_throughput".
|
||
- rpm_limit_type: Optional[str] - Type of rpm limit. Options: "best_effort_throughput" (no error if we're overallocating rpm), "guaranteed_throughput" (raise an error if we're overallocating rpm), "dynamic" (dynamically exceed limit when no 429 errors). Defaults to "best_effort_throughput".
|
||
- allowed_cache_controls: Optional[list] - List of allowed cache control values. Example - ["no-cache", "no-store"]. See all values - https://docs.litellm.ai/docs/proxy/caching#turn-on--off-caching-per-request
|
||
- blocked: Optional[bool] - Whether the key is blocked.
|
||
- rpm_limit: Optional[int] - Specify rpm limit for a given key (Requests per minute)
|
||
- tpm_limit: Optional[int] - Specify tpm limit for a given key (Tokens per minute)
|
||
- soft_budget: Optional[float] - Specify soft budget for a given key. Will trigger a slack alert when this soft budget is reached.
|
||
- tags: Optional[List[str]] - Tags for [tracking spend](https://litellm.vercel.app/docs/proxy/enterprise#tracking-spend-for-custom-tags) and/or doing [tag-based routing](https://litellm.vercel.app/docs/proxy/tag_routing).
|
||
- prompts: Optional[List[str]] - List of prompts that the key is allowed to use.
|
||
- enforced_params: Optional[List[str]] - List of enforced params for the key (Enterprise only). [Docs](https://docs.litellm.ai/docs/proxy/enterprise#enforce-required-params-for-llm-requests)
|
||
- prompts: Optional[List[str]] - List of prompts that the key is allowed to use.
|
||
- allowed_routes: Optional[list] - List of allowed routes for the key. Store the actual route or store a wildcard pattern for a set of routes. Example - ["/chat/completions", "/embeddings", "/keys/*"]
|
||
- allowed_passthrough_routes: Optional[list] - List of allowed pass through endpoints for the key. Store the actual endpoint or store a wildcard pattern for a set of endpoints. Example - ["/my-custom-endpoint"]. Use this instead of allowed_routes, if you just want to specify which pass through endpoints the key can access, without specifying the routes. If allowed_routes is specified, allowed_pass_through_endpoints is ignored.
|
||
- object_permission: Optional[LiteLLM_ObjectPermissionBase] - key-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission.
|
||
- key_type: Optional[str] - Type of key that determines default allowed routes. Options: "llm_api" (can call LLM API routes), "management" (can call management routes), "read_only" (can only call info/read routes), "default" (uses default allowed routes). Defaults to "default".
|
||
- prompts: Optional[List[str]] - List of allowed prompts for the key. If specified, the key will only be able to use these specific prompts.
|
||
- auto_rotate: Optional[bool] - Whether this key should be automatically rotated (regenerated)
|
||
- rotation_interval: Optional[str] - How often to auto-rotate this key (e.g., '30s', '30m', '30h', '30d'). Required if auto_rotate=True.
|
||
- allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint.
|
||
- router_settings: Optional[UpdateRouterConfig] - key-specific router settings. Example - {"model_group_retry_policy": {"gpt-4": {"RateLimitErrorRetries": 5}}}. IF null or {} then no router settings.
|
||
- access_group_ids: Optional[List[str]] - List of access group IDs to associate with the key. Access groups define which models a key can access. Example - ["access_group_1", "access_group_2"].
|
||
- budget_limits: Optional[list] - List of concurrent budget windows for the key. Each window specifies a budget_limit, time_period, and optional budget_duration. Example - [{"budget_limit": 10.0, "time_period": "1d"}, {"budget_limit": 50.0, "time_period": "7d"}].
|
||
|
||
Examples:
|
||
|
||
1. Allow users to turn on/off pii masking
|
||
|
||
```bash
|
||
curl --location 'http://0.0.0.0:4000/key/generate' \
|
||
--header 'Authorization: Bearer sk-1234' \
|
||
--header 'Content-Type: application/json' \
|
||
--data '{
|
||
"permissions": {"allow_pii_controls": true}
|
||
}'
|
||
```
|
||
|
||
Returns:
|
||
- key: (str) The generated api key
|
||
- expires: (datetime) Datetime object for when key expires.
|
||
- user_id: (str) Unique user id - used for tracking spend across multiple keys for same user id.
|
||
"""
|
||
try:
|
||
from litellm.proxy._types import CommonProxyErrors
|
||
from litellm.proxy.proxy_server import (
|
||
prisma_client,
|
||
user_api_key_cache,
|
||
user_custom_key_generate,
|
||
)
|
||
|
||
if prisma_client is None:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||
detail={"error": CommonProxyErrors.db_not_connected_error.value},
|
||
)
|
||
|
||
verbose_proxy_logger.debug("entered /key/generate")
|
||
|
||
await check_org_admin_can_generate_keys(user_api_key_dict=user_api_key_dict)
|
||
|
||
# Validate budget values are not negative and are finite numbers
|
||
# (GHSA-2rv4-xv66-fpjg): float('nan') passes `< 0` because nan < 0 is False.
|
||
if data.max_budget is not None and (not math.isfinite(data.max_budget) or data.max_budget < 0):
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail={"error": f"max_budget must be a non-negative finite number. Received: {data.max_budget}"},
|
||
)
|
||
if data.soft_budget is not None and (not math.isfinite(data.soft_budget) or data.soft_budget < 0):
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"},
|
||
)
|
||
|
||
custom_key_generate_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = (
|
||
user_custom_key_generate
|
||
)
|
||
if custom_key_generate_hook is not None:
|
||
if inspect.iscoroutinefunction(custom_key_generate_hook):
|
||
result: Final = await custom_key_generate_hook(data)
|
||
else:
|
||
raise ValueError("user_custom_key_generate must be a coroutine")
|
||
decision: Final = result.get("decision", True)
|
||
message: Final = result.get("message", "Authentication Failed - Custom Auth Rule")
|
||
if not decision:
|
||
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=message)
|
||
|
||
_check_allowed_routes_caller_permission(
|
||
allowed_routes=data.allowed_routes,
|
||
user_api_key_dict=user_api_key_dict,
|
||
allowed_routes_was_provided="allowed_routes" in data.model_fields_set,
|
||
)
|
||
_check_passthrough_routes_caller_permission(
|
||
data=data,
|
||
user_api_key_dict=user_api_key_dict,
|
||
)
|
||
|
||
# For non-admin internal users: auto-assign caller's user_id if not provided
|
||
# This prevents creating unbound keys with no user association (LIT-1884)
|
||
_is_proxy_admin: Final = (
|
||
user_api_key_dict.user_role is not None
|
||
and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
|
||
)
|
||
if not _is_proxy_admin and data.user_id is None:
|
||
data.user_id = user_api_key_dict.user_id
|
||
verbose_proxy_logger.warning(
|
||
"key/generate: auto-assigning user_id=%s for non-admin caller",
|
||
user_api_key_dict.user_id,
|
||
)
|
||
|
||
team_table: LiteLLM_TeamTableCachedObj | None = None
|
||
if data.team_id is not None:
|
||
try:
|
||
team_table = await get_team_object(
|
||
team_id=data.team_id,
|
||
prisma_client=prisma_client,
|
||
user_api_key_cache=user_api_key_cache,
|
||
parent_otel_span=user_api_key_dict.parent_otel_span,
|
||
check_db_only=True,
|
||
)
|
||
except Exception as e:
|
||
verbose_proxy_logger.debug("Error getting team object in `/key/generate`: %s", e)
|
||
# For non-admin callers, team must exist (LIT-1884)
|
||
if not _is_proxy_admin:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail=f"Team not found for team_id={data.team_id}. Non-admin users cannot create keys for non-existent teams.",
|
||
)
|
||
|
||
key_generation_check(
|
||
team_table=team_table,
|
||
user_api_key_dict=user_api_key_dict,
|
||
data=data,
|
||
route=KeyManagementRoutes.KEY_GENERATE,
|
||
)
|
||
|
||
if team_table is not None:
|
||
await _check_team_key_limits(
|
||
team_table=team_table,
|
||
data=data,
|
||
prisma_client=prisma_client,
|
||
)
|
||
|
||
# Validate key against project limits if project_id is set
|
||
if data.project_id is not None:
|
||
await _check_project_key_limits(
|
||
project_id=data.project_id,
|
||
data=data,
|
||
prisma_client=prisma_client,
|
||
user_api_key_cache=user_api_key_cache,
|
||
)
|
||
|
||
return await _common_key_generation_helper(
|
||
data=data,
|
||
user_api_key_dict=user_api_key_dict,
|
||
litellm_changed_by=litellm_changed_by,
|
||
team_table=team_table,
|
||
)
|
||
|
||
except Exception as e:
|
||
verbose_proxy_logger.exception("litellm.proxy.proxy_server.generate_key_fn(): Exception occured - %s", e)
|
||
raise handle_exception_on_proxy(e)
|
||
|
||
|
||
@router.post(
|
||
"/key/service-account/generate",
|
||
tags=["key management"],
|
||
dependencies=[Depends(user_api_key_auth)],
|
||
)
|
||
@management_endpoint_wrapper
|
||
async def generate_service_account_key_fn(
|
||
data: GenerateKeyRequest,
|
||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||
litellm_changed_by: str | None = Header(
|
||
None,
|
||
description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability",
|
||
),
|
||
):
|
||
"""
|
||
Generate a Service Account API key based on the provided data. This key does not belong to any user. It belongs to the team.
|
||
|
||
Why use a service account key?
|
||
- Prevent key from being deleted when user is deleted.
|
||
- Apply team limits, not team member limits to key.
|
||
|
||
Docs: https://docs.litellm.ai/docs/proxy/virtual_keys
|
||
|
||
Parameters:
|
||
- duration: Optional[str] - Specify the length of time the token is valid for. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d").
|
||
- key_alias: Optional[str] - User defined key alias
|
||
- key: Optional[str] - User defined key value. Must start with 'sk-' and be at least 16 characters long. If not set, a 16-digit unique sk-key is created for you.
|
||
- team_id: Optional[str] - The team id of the key
|
||
- user_id: Optional[str] - [NON-FUNCTIONAL] THIS WILL BE IGNORED. The user id of the key
|
||
- budget_id: Optional[str] - The budget id associated with the key. Created by calling `/budget/new`.
|
||
- models: Optional[list] - Model_name's a user is allowed to call. (if empty, key is allowed to call all models)
|
||
- aliases: Optional[dict] - Any alias mappings, on top of anything in the config.yaml model list. - https://docs.litellm.ai/docs/proxy/virtual_keys#managing-auth---upgradedowngrade-models
|
||
- config: Optional[dict] - any key-specific configs, overrides config in config.yaml
|
||
- spend: Optional[int] - Amount spent by key. Default is 0. Will be updated by proxy whenever key is used. https://docs.litellm.ai/docs/proxy/virtual_keys#managing-auth---tracking-spend
|
||
- send_invite_email: Optional[bool] - Whether to send an invite email to the user_id, with the generate key
|
||
- max_budget: Optional[float] - Specify max budget for a given key.
|
||
- budget_duration: Optional[str] - Budget is reset at the end of specified duration. If not set, budget is never reset. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d").
|
||
- max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x.
|
||
- metadata: Optional[dict] - Metadata for key, store information for key. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" }
|
||
- guardrails: Optional[List[str]] - List of active guardrails for the key
|
||
- permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false}
|
||
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget.
|
||
- budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||
- model_rpm_limit: Optional[dict] - key-specific model rpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific rpm limit.
|
||
- model_tpm_limit: Optional[dict] - key-specific model tpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific tpm limit.
|
||
- default_estimated_output_tokens: Optional[int] - Proxy admin only. Expected output tokens reserved for TPM limiting when a request omits max_tokens. Positive integer. Falls back to the team setting, then to the built-in estimate.
|
||
- default_estimated_output_tokens_per_model: Optional[dict] - Proxy admin only. Per-model override of the above. Example - {"gpt-4": 4096, "gpt-3.5-turbo": 1024}. Takes precedence over the key-wide value.
|
||
- mcp_rpm_limit: Optional[dict] - key-specific per-MCP-server rpm limit, keyed by MCP server name (alias if set, else the configured name). Example - {"github": 100, "slack": 200}. IF null or {} then no MCP-specific rpm limit.
|
||
- tpm_limit_type: Optional[str] - TPM rate limit type - "best_effort_throughput", "guaranteed_throughput", or "dynamic"
|
||
- 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. Example - ["no-cache", "no-store"]. See all values - https://docs.litellm.ai/docs/proxy/caching#turn-on--off-caching-per-request
|
||
- blocked: Optional[bool] - Whether the key is blocked.
|
||
- rpm_limit: Optional[int] - Specify rpm limit for a given key (Requests per minute)
|
||
- tpm_limit: Optional[int] - Specify tpm limit for a given key (Tokens per minute)
|
||
- soft_budget: Optional[float] - Specify soft budget for a given key. Will trigger a slack alert when this soft budget is reached.
|
||
- tags: Optional[List[str]] - Tags for [tracking spend](https://litellm.vercel.app/docs/proxy/enterprise#tracking-spend-for-custom-tags) and/or doing [tag-based routing](https://litellm.vercel.app/docs/proxy/tag_routing).
|
||
- enforced_params: Optional[List[str]] - List of enforced params for the key (Enterprise only). [Docs](https://docs.litellm.ai/docs/proxy/enterprise#enforce-required-params-for-llm-requests)
|
||
- allowed_routes: Optional[list] - List of allowed routes for the key. Store the actual route or store a wildcard pattern for a set of routes. Example - ["/chat/completions", "/embeddings", "/keys/*"]
|
||
- object_permission: Optional[LiteLLM_ObjectPermissionBase] - key-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission.
|
||
Examples:
|
||
- allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint.
|
||
|
||
|
||
1. Allow users to turn on/off pii masking
|
||
|
||
```bash
|
||
curl --location 'http://0.0.0.0:4000/key/generate' \
|
||
--header 'Authorization: Bearer sk-1234' \
|
||
--header 'Content-Type: application/json' \
|
||
--data '{
|
||
"permissions": {"allow_pii_controls": true}
|
||
}'
|
||
```
|
||
|
||
Returns:
|
||
- key: (str) The generated api key
|
||
- expires: (datetime) Datetime object for when key expires.
|
||
- user_id: (str) Unique user id - used for tracking spend across multiple keys for same user id.
|
||
|
||
"""
|
||
from litellm.proxy._types import CommonProxyErrors
|
||
from litellm.proxy.proxy_server import (
|
||
prisma_client,
|
||
user_api_key_cache,
|
||
user_custom_key_generate,
|
||
)
|
||
|
||
if prisma_client is None:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||
detail={"error": CommonProxyErrors.db_not_connected_error.value},
|
||
)
|
||
|
||
await check_org_admin_can_generate_keys(user_api_key_dict=user_api_key_dict)
|
||
|
||
_check_allowed_routes_caller_permission(
|
||
allowed_routes=data.allowed_routes,
|
||
user_api_key_dict=user_api_key_dict,
|
||
allowed_routes_was_provided="allowed_routes" in data.model_fields_set,
|
||
)
|
||
_check_passthrough_routes_caller_permission(
|
||
data=data,
|
||
user_api_key_dict=user_api_key_dict,
|
||
)
|
||
|
||
await validate_team_id_used_in_service_account_request(
|
||
team_id=data.team_id,
|
||
prisma_client=prisma_client,
|
||
)
|
||
|
||
verbose_proxy_logger.debug("entered /key/generate")
|
||
|
||
custom_key_generate_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = user_custom_key_generate
|
||
if custom_key_generate_hook is not None:
|
||
if inspect.iscoroutinefunction(custom_key_generate_hook):
|
||
result: Final = await custom_key_generate_hook(data)
|
||
else:
|
||
raise ValueError("user_custom_key_generate must be a coroutine")
|
||
decision: Final = result.get("decision", True)
|
||
message: Final = result.get("message", "Authentication Failed - Custom Auth Rule")
|
||
if not decision:
|
||
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=message)
|
||
team_table: LiteLLM_TeamTableCachedObj | None = None
|
||
if data.team_id is not None:
|
||
try:
|
||
team_table = await get_team_object(
|
||
team_id=data.team_id,
|
||
prisma_client=prisma_client,
|
||
user_api_key_cache=user_api_key_cache,
|
||
parent_otel_span=user_api_key_dict.parent_otel_span,
|
||
check_db_only=True,
|
||
)
|
||
except Exception as e:
|
||
verbose_proxy_logger.debug("Error getting team object in `/key/generate`: %s", e)
|
||
team_table = None
|
||
|
||
if team_table is not None:
|
||
await _check_team_key_limits(
|
||
team_table=team_table,
|
||
data=data,
|
||
prisma_client=prisma_client,
|
||
)
|
||
|
||
key_generation_check(
|
||
team_table=team_table,
|
||
user_api_key_dict=user_api_key_dict,
|
||
data=data,
|
||
route=KeyManagementRoutes.KEY_GENERATE_SERVICE_ACCOUNT,
|
||
)
|
||
|
||
data.user_id = None # do not allow user_id to be set for service account keys
|
||
|
||
return await _common_key_generation_helper(
|
||
data=data,
|
||
user_api_key_dict=user_api_key_dict,
|
||
litellm_changed_by=litellm_changed_by,
|
||
team_table=team_table,
|
||
)
|
||
|
||
|
||
def prepare_metadata_fields(data: BaseModel, non_default_values: dict, existing_metadata: dict) -> dict:
|
||
"""
|
||
Check LiteLLM_ManagementEndpoint_MetadataFields (proxy/_types.py) for fields that are allowed to be updated
|
||
"""
|
||
raise_on_invalid_key_logging_config(non_default_values.get("metadata"))
|
||
|
||
if "metadata" not in non_default_values: # allow user to set metadata to none
|
||
non_default_values["metadata"] = existing_metadata.copy()
|
||
|
||
casted_metadata: Final = cast(dict, non_default_values["metadata"])
|
||
|
||
# Reserved metadata fields are immutable once set. Preserve the existing value
|
||
# when omitted, reject any explicit attempt to change it (including null).
|
||
for reserved_field in LiteLLM_Reserved_Metadata_Fields:
|
||
existing_value = existing_metadata.get(reserved_field)
|
||
if existing_value is None:
|
||
continue
|
||
if casted_metadata is None or (
|
||
reserved_field in casted_metadata and casted_metadata[reserved_field] != existing_value
|
||
):
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail=f"{reserved_field} is immutable once set and cannot be changed via update.",
|
||
)
|
||
casted_metadata[reserved_field] = existing_value
|
||
|
||
data_json: Final[Mapping[str, object]] = data.model_dump(exclude_unset=True, exclude_none=True)
|
||
|
||
try:
|
||
for k, v in data_json.items():
|
||
if k in LiteLLM_ManagementEndpoint_MetadataFields:
|
||
if isinstance(v, datetime):
|
||
casted_metadata[k] = v.isoformat()
|
||
else:
|
||
casted_metadata[k] = v
|
||
if k in LiteLLM_ManagementEndpoint_MetadataFields_Premium:
|
||
from litellm.proxy.utils import _premium_user_check
|
||
|
||
if v:
|
||
_premium_user_check(k)
|
||
casted_metadata[k] = v
|
||
|
||
except Exception as e:
|
||
verbose_proxy_logger.exception(
|
||
"litellm.proxy.proxy_server.prepare_metadata_fields(): Exception occured - %s", e
|
||
)
|
||
|
||
non_default_values["metadata"] = encrypt_callback_vars(casted_metadata)
|
||
return non_default_values
|
||
|
||
|
||
async def prepare_key_update_data(
|
||
data: UpdateKeyRequest | RegenerateKeyRequest,
|
||
existing_key_row: LiteLLM_VerificationToken,
|
||
):
|
||
data_json: Final[dict] = data.model_dump(exclude_unset=True)
|
||
data_json.pop("key", None)
|
||
data_json.pop("new_key", None)
|
||
data_json.pop("grace_period", None) # Request-only param, not a DB column
|
||
if (
|
||
data.metadata is not None
|
||
and data.metadata.get("service_account_id") is not None
|
||
and (data.team_id or existing_key_row.team_id) is None
|
||
):
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail="team_id is required for service account keys. Please specify `team_id` in the request body.",
|
||
)
|
||
non_default_values = {}
|
||
# ADD METADATA FIELDS
|
||
# Set Management Endpoint Metadata Fields
|
||
for field in LiteLLM_ManagementEndpoint_MetadataFields_Premium:
|
||
if getattr(data, field, None) is not None:
|
||
_set_object_metadata_field(
|
||
object_data=data,
|
||
field_name=field,
|
||
value=getattr(data, field),
|
||
)
|
||
for k, v in data_json.items():
|
||
if k in LiteLLM_ManagementEndpoint_MetadataFields or k in LiteLLM_ManagementEndpoint_MetadataFields_Premium:
|
||
continue
|
||
non_default_values[k] = v
|
||
|
||
if "duration" in non_default_values:
|
||
duration: Final = non_default_values.pop("duration")
|
||
if duration is None or duration == "-1":
|
||
# Set expires to None to indicate the key never expires
|
||
non_default_values["expires"] = None
|
||
elif duration and (isinstance(duration, str)) and len(duration) > 0:
|
||
duration_s: Final = duration_in_seconds(duration=duration)
|
||
expires: Final = datetime.now(timezone.utc) + timedelta(seconds=duration_s)
|
||
non_default_values["expires"] = expires
|
||
|
||
if "budget_duration" in non_default_values:
|
||
budget_duration: Final = non_default_values.pop("budget_duration")
|
||
if budget_duration is None:
|
||
non_default_values["budget_duration"] = None
|
||
non_default_values["budget_reset_at"] = None
|
||
elif isinstance(budget_duration, str) and len(budget_duration) > 0:
|
||
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
|
||
|
||
key_reset_at: Final = get_budget_reset_time(budget_duration=budget_duration)
|
||
non_default_values["budget_reset_at"] = key_reset_at
|
||
non_default_values["budget_duration"] = budget_duration
|
||
|
||
if "budget_limits" in non_default_values:
|
||
raw_windows: Final = non_default_values["budget_limits"]
|
||
if raw_windows:
|
||
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
|
||
|
||
initialized_windows: Final = []
|
||
for window in raw_windows:
|
||
w = window if isinstance(window, dict) else window.model_dump()
|
||
w["reset_at"] = get_budget_reset_time(budget_duration=w["budget_duration"]).isoformat()
|
||
initialized_windows.append(w)
|
||
non_default_values["budget_limits"] = json.dumps(initialized_windows)
|
||
else:
|
||
# [] / None clears the field; prisma-client-py has no DbNull
|
||
# sentinel for Json? columns, so store the JSON literal null
|
||
non_default_values["budget_limits"] = json.dumps(None)
|
||
|
||
if "object_permission" in non_default_values:
|
||
non_default_values = await _handle_update_object_permission(
|
||
data_json=non_default_values,
|
||
existing_key_row=existing_key_row,
|
||
)
|
||
|
||
_metadata: Final = existing_key_row.metadata or {}
|
||
|
||
# validate model_max_budget
|
||
if "model_max_budget" in non_default_values:
|
||
validate_model_max_budget(non_default_values["model_max_budget"])
|
||
|
||
# Serialize router_settings to JSON if present
|
||
if "router_settings" in non_default_values and non_default_values["router_settings"] is not None:
|
||
non_default_values["router_settings"] = safe_dumps(non_default_values["router_settings"])
|
||
|
||
non_default_values = prepare_metadata_fields(
|
||
data=data, non_default_values=non_default_values, existing_metadata=_metadata
|
||
)
|
||
|
||
return non_default_values
|
||
|
||
|
||
async def _handle_update_object_permission(
|
||
data_json: dict,
|
||
existing_key_row: LiteLLM_VerificationToken,
|
||
) -> dict:
|
||
"""
|
||
Handle the update of object permission.
|
||
"""
|
||
from litellm.proxy.proxy_server import prisma_client
|
||
|
||
# Use the common helper to handle the object permission update
|
||
object_permission_id: Final = await handle_update_object_permission_common(
|
||
data_json=data_json,
|
||
existing_object_permission_id=existing_key_row.object_permission_id,
|
||
prisma_client=prisma_client,
|
||
)
|
||
|
||
# Add the object_permission_id to data_json if one was created/updated
|
||
if object_permission_id is not None:
|
||
data_json["object_permission_id"] = object_permission_id
|
||
verbose_proxy_logger.debug("updated object_permission_id: %s", object_permission_id)
|
||
|
||
return data_json
|
||
|
||
|
||
def is_different_team(data: UpdateKeyRequest, existing_key_row: LiteLLM_VerificationToken) -> bool:
|
||
if data.team_id is None:
|
||
return False
|
||
if existing_key_row.team_id is None:
|
||
return True
|
||
return data.team_id != existing_key_row.team_id
|
||
|
||
|
||
def _validate_max_budget(max_budget: float | None) -> None:
|
||
"""
|
||
Validate that max_budget is not negative.
|
||
|
||
Args:
|
||
max_budget: The max_budget value to validate
|
||
|
||
Raises:
|
||
HTTPException: If max_budget is negative
|
||
"""
|
||
if max_budget is not None and (not math.isfinite(max_budget) or max_budget < 0):
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail={"error": f"max_budget must be a non-negative finite number. Received: {max_budget}"},
|
||
)
|
||
|
||
|
||
async def _get_and_validate_existing_key(
|
||
token: str | None, prisma_client: PrismaClient | None, key_alias: str | None = None
|
||
) -> LiteLLM_VerificationToken:
|
||
"""
|
||
Get existing key from database and validate it exists.
|
||
|
||
Args:
|
||
token: The key token to look up
|
||
prisma_client: Prisma client instance
|
||
key_alias: Alias to look the key up by when token is not provided
|
||
|
||
Returns:
|
||
LiteLLM_VerificationToken: The existing key row
|
||
|
||
Raises:
|
||
ProxyException: 404 if key is not found, 400 if the alias matches multiple keys
|
||
"""
|
||
if prisma_client is None:
|
||
raise HTTPException(
|
||
status_code=500,
|
||
detail={"error": "Database not connected"},
|
||
)
|
||
|
||
if token is not None:
|
||
hashed_token: Final = _hash_token_if_needed(token=token)
|
||
|
||
existing_key_row: Final[LiteLLM_VerificationToken | None] = await _prisma_table(
|
||
VerificationTokenRepository(prisma_client)
|
||
).find_unique(where={"token": hashed_token})
|
||
|
||
if existing_key_row is None:
|
||
raise ProxyException(
|
||
message="Key not found.",
|
||
type=ProxyErrorTypes.not_found_error,
|
||
param="key",
|
||
code=status.HTTP_404_NOT_FOUND,
|
||
)
|
||
|
||
return existing_key_row
|
||
|
||
if key_alias is None:
|
||
raise ProxyException(
|
||
message="either key or key_alias must be provided",
|
||
type=ProxyErrorTypes.bad_request_error,
|
||
param="key",
|
||
code=status.HTTP_400_BAD_REQUEST,
|
||
)
|
||
|
||
rows: Sequence[LiteLLM_VerificationToken] = await _prisma_table(
|
||
VerificationTokenRepository(prisma_client)
|
||
).find_many(where={"key_alias": key_alias}, take=2)
|
||
|
||
if len(rows) == 0:
|
||
raise ProxyException(
|
||
message=f"Key not found. No key with key_alias='{key_alias}'.",
|
||
type=ProxyErrorTypes.not_found_error,
|
||
param="key_alias",
|
||
code=status.HTTP_404_NOT_FOUND,
|
||
)
|
||
|
||
if len(rows) > 1:
|
||
raise ProxyException(
|
||
message=f"Multiple keys share key_alias='{key_alias}', so it cannot be used as an identifier.",
|
||
type=ProxyErrorTypes.bad_request_error,
|
||
param="key_alias",
|
||
code=status.HTTP_400_BAD_REQUEST,
|
||
)
|
||
|
||
return rows[0]
|
||
|
||
|
||
def _resolve_token_to_update(data: UpdateKeyRequest, existing_key_row: LiteLLM_VerificationToken) -> str:
|
||
if data.key is not None:
|
||
return data.key
|
||
if existing_key_row.token is None:
|
||
raise ProxyException(
|
||
message="Key not found.",
|
||
type=ProxyErrorTypes.not_found_error,
|
||
param="key",
|
||
code=status.HTTP_404_NOT_FOUND,
|
||
)
|
||
return existing_key_row.token
|
||
|
||
|
||
async def _process_single_key_update(
|
||
update_key_request: UpdateKeyRequest,
|
||
user_api_key_dict: UserAPIKeyAuth,
|
||
litellm_changed_by: str | None,
|
||
prisma_client: PrismaClient | None,
|
||
user_api_key_cache: UserApiKeyCache,
|
||
proxy_logging_obj: ProxyLogging,
|
||
llm_router: Router | None,
|
||
user_custom_key_update: Callable | None = None,
|
||
existing_key_row: LiteLLM_VerificationToken | None = None,
|
||
) -> dict[str, object]:
|
||
"""
|
||
Process a single key update with all validations and checks.
|
||
|
||
This function encapsulates all the logic for updating a single key,
|
||
including validation, permission checks, team checks, and database updates.
|
||
|
||
Args:
|
||
update_key_request: Fully-constructed UpdateKeyRequest for the target key
|
||
user_api_key_dict: The authenticated user's API key info
|
||
litellm_changed_by: Optional header for tracking who made the change
|
||
prisma_client: Prisma client instance
|
||
user_api_key_cache: User API key cache
|
||
proxy_logging_obj: Proxy logging object
|
||
llm_router: LLM router instance
|
||
existing_key_row: Optional pre-fetched key row to avoid redundant lookups
|
||
|
||
Returns:
|
||
Dict containing the updated key information
|
||
|
||
Raises:
|
||
HTTPException: For various validation and permission errors
|
||
"""
|
||
# Validate max_budget
|
||
_validate_max_budget(update_key_request.max_budget)
|
||
|
||
_check_permissions_caller_permission(
|
||
data=update_key_request,
|
||
user_api_key_dict=user_api_key_dict,
|
||
)
|
||
|
||
# Get and validate existing key
|
||
if existing_key_row is None:
|
||
existing_key_row = await _get_and_validate_existing_key(
|
||
token=update_key_request.key,
|
||
prisma_client=prisma_client,
|
||
)
|
||
|
||
_existing_row_metadata: Final = getattr(existing_key_row, "metadata", None)
|
||
enforce_batch_enqueued_token_limit_is_admin_only(
|
||
data=update_key_request,
|
||
existing_metadata=_existing_row_metadata if isinstance(_existing_row_metadata, dict) else None,
|
||
user_api_key_dict=user_api_key_dict,
|
||
entity="key",
|
||
)
|
||
|
||
# Check team member permissions
|
||
if prisma_client is not None:
|
||
await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint(
|
||
user_api_key_dict=user_api_key_dict,
|
||
route=KeyManagementRoutes.KEY_UPDATE,
|
||
prisma_client=prisma_client,
|
||
existing_key_row=existing_key_row,
|
||
user_api_key_cache=user_api_key_cache,
|
||
)
|
||
|
||
# Custom key update hook
|
||
if user_custom_key_update is not None:
|
||
if inspect.iscoroutinefunction(user_custom_key_update):
|
||
result: Final = await user_custom_key_update(update_key_request)
|
||
else:
|
||
raise ValueError("user_custom_key_update must be a coroutine")
|
||
decision: Final = result.get("decision", True)
|
||
message: Final = result.get("message", "Authentication Failed - Custom Auth Rule")
|
||
if not decision:
|
||
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=message)
|
||
|
||
# Enforce upperbound key params on update (don't fill defaults)
|
||
_enforce_upperbound_key_params(update_key_request, fill_defaults=False)
|
||
|
||
# Get team object and check team limits if team_id is provided
|
||
team_obj: LiteLLM_TeamTableCachedObj | None = None
|
||
if update_key_request.team_id is not None:
|
||
team_obj = await get_team_object(
|
||
team_id=update_key_request.team_id,
|
||
prisma_client=prisma_client,
|
||
user_api_key_cache=user_api_key_cache,
|
||
check_db_only=True,
|
||
)
|
||
|
||
if team_obj is not None and prisma_client is not None:
|
||
await _check_team_key_limits(
|
||
team_table=team_obj,
|
||
data=update_key_request,
|
||
prisma_client=prisma_client,
|
||
)
|
||
|
||
# Validate team change if team is being changed
|
||
if is_different_team(data=update_key_request, existing_key_row=existing_key_row):
|
||
if llm_router is None:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail={
|
||
"error": "LLM router not found. Please set it up by passing in a valid config.yaml or adding models via the UI."
|
||
},
|
||
)
|
||
if team_obj is None:
|
||
raise HTTPException(
|
||
status_code=500,
|
||
detail={"error": "Team object not found for team change validation"},
|
||
)
|
||
await validate_key_team_change(
|
||
key=existing_key_row,
|
||
team=team_obj,
|
||
change_initiated_by=user_api_key_dict,
|
||
llm_router=llm_router,
|
||
)
|
||
|
||
# Prepare update data
|
||
non_default_values = await prepare_key_update_data(data=update_key_request, existing_key_row=existing_key_row)
|
||
|
||
# Update key in database
|
||
if prisma_client is None:
|
||
raise HTTPException(
|
||
status_code=500,
|
||
detail={"error": "Database not connected"},
|
||
)
|
||
|
||
_data: Final = {**non_default_values, "token": update_key_request.key}
|
||
response: Final[Mapping[str, object] | None] = cast( # cast-ok: every update_data branch returns a str-keyed dict
|
||
"Mapping[str, object] | None",
|
||
await prisma_client.update_data(token=update_key_request.key, data=_data),
|
||
)
|
||
|
||
# Delete cache
|
||
await _delete_cache_key_object(
|
||
hashed_token=_hash_token_if_needed(update_key_request.key),
|
||
user_api_key_cache=user_api_key_cache,
|
||
proxy_logging_obj=proxy_logging_obj,
|
||
)
|
||
|
||
# After the key's own cache entry is dropped, so a failure here cannot leave the key
|
||
# authenticating against the access groups it just lost.
|
||
await sync_key_update_access_group_membership(
|
||
prisma_client=prisma_client,
|
||
key_token=_hash_token_if_needed(
|
||
_resolve_token_to_update(data=update_key_request, existing_key_row=existing_key_row)
|
||
),
|
||
data=update_key_request,
|
||
existing_key_row=existing_key_row,
|
||
)
|
||
|
||
# Trigger async hook
|
||
asyncio.create_task(
|
||
KeyManagementEventHooks.async_key_updated_hook(
|
||
data=update_key_request,
|
||
existing_key_row=existing_key_row,
|
||
response=response,
|
||
user_api_key_dict=user_api_key_dict,
|
||
litellm_changed_by=litellm_changed_by,
|
||
)
|
||
)
|
||
|
||
if response is None:
|
||
raise ValueError("Failed to update key got response = None")
|
||
|
||
# Extract and format updated key info
|
||
updated_key_info = response.get("data", {})
|
||
if hasattr(updated_key_info, "model_dump"):
|
||
updated_key_info = updated_key_info.model_dump()
|
||
elif hasattr(updated_key_info, "dict"):
|
||
updated_key_info = updated_key_info.dict()
|
||
|
||
updated_key_info.pop("token", None)
|
||
|
||
return updated_key_info
|
||
|
||
|
||
async def _validate_mcp_servers_for_key_update(
|
||
data: "UpdateKeyRequest",
|
||
team_obj: Optional["LiteLLM_TeamTableCachedObj"],
|
||
existing_key_row: LiteLLM_VerificationToken,
|
||
prisma_client: PrismaClient | None,
|
||
user_api_key_cache: UserApiKeyCache,
|
||
is_proxy_admin: bool,
|
||
) -> ObjectPermissionDict | None:
|
||
"""Validate MCP servers in object_permission against the effective team."""
|
||
effective_team_obj = team_obj
|
||
# If team_id isn't being changed, resolve the existing key's team
|
||
if effective_team_obj is None and existing_key_row.team_id:
|
||
effective_team_obj = await get_team_object(
|
||
team_id=existing_key_row.team_id,
|
||
prisma_client=prisma_client,
|
||
user_api_key_cache=user_api_key_cache,
|
||
check_db_only=True,
|
||
)
|
||
object_permission_dict: Final = _object_permission_to_dict(data.object_permission)
|
||
normalized_object_permission: Final = await validate_key_mcp_servers_against_team(
|
||
object_permission=object_permission_dict,
|
||
team_obj=effective_team_obj,
|
||
prisma_client=prisma_client,
|
||
is_proxy_admin=is_proxy_admin,
|
||
)
|
||
await validate_key_search_tools_against_team(
|
||
object_permission=object_permission_dict,
|
||
team_obj=effective_team_obj,
|
||
is_proxy_admin=is_proxy_admin,
|
||
)
|
||
await validate_key_vector_stores_against_team(
|
||
object_permission=object_permission_dict,
|
||
team_obj=effective_team_obj,
|
||
is_proxy_admin=is_proxy_admin,
|
||
)
|
||
return normalized_object_permission
|
||
|
||
|
||
async def _validate_update_key_data(
|
||
data: UpdateKeyRequest,
|
||
existing_key_row: LiteLLM_VerificationToken,
|
||
user_api_key_dict: UserAPIKeyAuth,
|
||
llm_router: Router | None,
|
||
premium_user: bool,
|
||
prisma_client: Any,
|
||
user_api_key_cache: UserApiKeyCache,
|
||
) -> None:
|
||
"""Validate permissions and constraints for key update."""
|
||
# Reject NaN/±inf spend before it can reach the DB / spend counter.
|
||
validate_finite_spend(data.spend)
|
||
validate_budget_duration(data.budget_duration)
|
||
|
||
_is_proxy_admin: Final = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
|
||
|
||
_check_allowed_routes_caller_permission(
|
||
allowed_routes=data.allowed_routes,
|
||
user_api_key_dict=user_api_key_dict,
|
||
allowed_routes_was_provided="allowed_routes" in data.model_fields_set,
|
||
)
|
||
_check_passthrough_routes_caller_permission(
|
||
data=data,
|
||
user_api_key_dict=user_api_key_dict,
|
||
)
|
||
_check_permissions_caller_permission(
|
||
data=data,
|
||
user_api_key_dict=user_api_key_dict,
|
||
)
|
||
|
||
_validate_caller_can_change_key_ownership(
|
||
data=data,
|
||
existing_key_row=existing_key_row,
|
||
user_api_key_dict=user_api_key_dict,
|
||
)
|
||
|
||
common_key_access_checks(
|
||
user_api_key_dict=user_api_key_dict,
|
||
data=data,
|
||
user_id=existing_key_row.user_id,
|
||
llm_router=llm_router,
|
||
premium_user=premium_user,
|
||
)
|
||
|
||
await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint(
|
||
user_api_key_dict=user_api_key_dict,
|
||
route=KeyManagementRoutes.KEY_UPDATE,
|
||
prisma_client=prisma_client,
|
||
existing_key_row=existing_key_row,
|
||
user_api_key_cache=user_api_key_cache,
|
||
)
|
||
|
||
# Cross-key authorization. Previously only gated on max_budget/spend
|
||
# changes, which let a non-admin blanket-rewrite any OTHER field on
|
||
# any key (models, alias, metadata, tpm_limit, rpm_limit,
|
||
# allowed_routes, guardrails, blocked, duration, permissions, …) as
|
||
# long as they avoided budget/spend.
|
||
#
|
||
# Policy:
|
||
# - Key owner (same user_id): may update non-budget fields on their
|
||
# own key without the admin check.
|
||
# - Team member with /key/update grant (on a team key): may update
|
||
# non-budget fields. Team membership + permission is already
|
||
# enforced by can_team_member_execute_key_management_endpoint
|
||
# above, which raises 401 for non-members or members without the
|
||
# grant — so reaching this point on a team key means the caller
|
||
# was authorized via member_permissions. This preserves the
|
||
# documented member_permissions feature while still blocking the
|
||
# cross-org attack (an outside org admin is not a member of the
|
||
# victim team and gets rejected at the earlier check).
|
||
# - Anyone else (non-PROXY_ADMIN, not the owner, not a team member
|
||
# on a team key): must pass _check_key_admin_access (PROXY_ADMIN
|
||
# / key-owner / team-admin / org-admin of the key).
|
||
# - max_budget / spend / budget_limits: always require the admin
|
||
# check, even for the key owner or a team member (matches the
|
||
# existing admin-only budget semantics). budget_limits uses
|
||
# model_fields_set because an explicit null/[] clears the field
|
||
# and must gate the same as setting or changing it.
|
||
# - spend gates on presence alone (not a value diff): the DB spend
|
||
# lags the live cross-pod counter, so letting an "unchanged" spend
|
||
# through the non-admin path would let a key owner / team member
|
||
# overwrite the live counter below real usage and silently weaken
|
||
# enforcement.
|
||
_is_budget_change: Final = (
|
||
(data.max_budget is not None and data.max_budget != existing_key_row.max_budget)
|
||
or data.spend is not None
|
||
or "budget_limits" in data.model_fields_set
|
||
)
|
||
|
||
_existing_metadata: Final = getattr(existing_key_row, "metadata", None)
|
||
_existing_throttle: Final = (
|
||
_existing_metadata.get("throttle_on_budget_exceeded") if isinstance(_existing_metadata, dict) else None
|
||
)
|
||
if data.throttle_on_budget_exceeded is True and _existing_throttle is not True and not _is_proxy_admin:
|
||
raise HTTPException(
|
||
status_code=403,
|
||
detail={"error": "Only proxy admins can enable throttle_on_budget_exceeded on a key."},
|
||
)
|
||
|
||
enforce_output_token_estimates_are_admin_only(
|
||
data=data,
|
||
existing_metadata=_existing_metadata if isinstance(_existing_metadata, dict) else None,
|
||
user_api_key_dict=user_api_key_dict,
|
||
entity="key",
|
||
)
|
||
enforce_batch_enqueued_token_limit_is_admin_only(
|
||
data=data,
|
||
existing_metadata=_existing_metadata if isinstance(_existing_metadata, dict) else None,
|
||
user_api_key_dict=user_api_key_dict,
|
||
entity="key",
|
||
)
|
||
|
||
# Personal-key bypass: the caller both created the key AND still owns it
|
||
# (user_id == caller). Checking only created_by would let a demoted admin
|
||
# who originally created a key for another user continue editing it without
|
||
# admin authorization after the key was reassigned.
|
||
caller_is_creator: Final = (
|
||
user_api_key_dict.user_id is not None
|
||
and getattr(existing_key_row, "created_by", None) == user_api_key_dict.user_id
|
||
and getattr(existing_key_row, "user_id", None) == user_api_key_dict.user_id
|
||
)
|
||
# Team keys: can_team_member_execute_key_management_endpoint (called above)
|
||
# already validated team membership + /key/update permission and would have
|
||
# raised if the caller lacked it. Reaching this point on a team key for a
|
||
# non-budget change means the caller was authorized — skip the redundant
|
||
# _check_key_admin_access that would otherwise require team/org admin status.
|
||
_key_is_team_key: Final = getattr(existing_key_row, "team_id", None) is not None
|
||
can_skip_admin_check: Final = (caller_is_creator or _key_is_team_key) and not _is_budget_change
|
||
if (not _is_proxy_admin) and prisma_client is not None and not can_skip_admin_check:
|
||
hashed_key: Final = existing_key_row.token
|
||
await _check_key_admin_access(
|
||
user_api_key_dict=user_api_key_dict,
|
||
hashed_token=hashed_key,
|
||
prisma_client=prisma_client,
|
||
user_api_key_cache=user_api_key_cache,
|
||
route=("/key/update (max_budget/spend)" if _is_budget_change else "/key/update"),
|
||
)
|
||
|
||
# Check team limits if key has a team_id (from request or existing key)
|
||
team_obj: LiteLLM_TeamTableCachedObj | None = None
|
||
_team_id_to_check: Final = data.team_id or getattr(existing_key_row, "team_id", None)
|
||
if _team_id_to_check is not None:
|
||
team_obj = await get_team_object(
|
||
team_id=_team_id_to_check,
|
||
prisma_client=prisma_client,
|
||
user_api_key_cache=user_api_key_cache,
|
||
check_db_only=True,
|
||
)
|
||
|
||
# Validate team exists when non-admin sets a new team_id (LIT-1884)
|
||
if team_obj is None and data.team_id is not None and not _is_proxy_admin:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail=f"Team not found for team_id={data.team_id}. Non-admin users cannot set keys to non-existent teams.",
|
||
)
|
||
|
||
if team_obj is not None:
|
||
await _check_team_key_limits(
|
||
team_table=team_obj,
|
||
data=data,
|
||
prisma_client=prisma_client,
|
||
)
|
||
|
||
TeamMemberPermissionChecks.enforce_member_can_assign_access_groups(
|
||
user_api_key_dict=user_api_key_dict,
|
||
team_table=team_obj,
|
||
access_group_ids=data.access_group_ids,
|
||
)
|
||
|
||
# Validate key against project limits if project_id is being set
|
||
_project_id_to_check: Final = getattr(data, "project_id", None) or getattr(existing_key_row, "project_id", None)
|
||
if _project_id_to_check is not None and (data.models is not None or data.max_budget is not None):
|
||
await _check_project_key_limits(
|
||
project_id=_project_id_to_check,
|
||
data=data,
|
||
prisma_client=prisma_client,
|
||
user_api_key_cache=user_api_key_cache,
|
||
)
|
||
|
||
# When the caller asks to change the key's organization_id, require that
|
||
# they are a member of (or a proxy admin over) the target organization.
|
||
# Without this gate, any caller could assign their key to an arbitrary
|
||
# organization_id by passing it in the request body — VERIA-55 secondary
|
||
# IDOR. The check mirrors the membership rule already used on the
|
||
# `/key/list` filter path in `validate_key_list_check`.
|
||
_existing_org_id: Final = getattr(existing_key_row, "organization_id", None)
|
||
if data.organization_id is not None and data.organization_id != _existing_org_id and not _is_proxy_admin:
|
||
await _validate_caller_can_assign_key_org(
|
||
user_api_key_dict=user_api_key_dict,
|
||
organization_id=data.organization_id,
|
||
prisma_client=prisma_client,
|
||
)
|
||
|
||
# Check org key limits only when throughput-related fields or organization_id change
|
||
_org_id_to_check: Final = data.organization_id or _existing_org_id
|
||
_throughput_fields_changed: Final = (
|
||
data.organization_id is not None
|
||
or data.tpm_limit is not None
|
||
or data.rpm_limit is not None
|
||
or data.tpm_limit_type is not None
|
||
or data.rpm_limit_type is not None
|
||
)
|
||
if _org_id_to_check is not None and _throughput_fields_changed:
|
||
org_table: Final = await get_org_object(
|
||
org_id=_org_id_to_check,
|
||
user_api_key_cache=user_api_key_cache,
|
||
prisma_client=prisma_client,
|
||
)
|
||
if org_table is None:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail=f"Organization not found for organization_id={_org_id_to_check}",
|
||
)
|
||
await _check_org_key_limits(
|
||
org_table=org_table,
|
||
data=data,
|
||
prisma_client=prisma_client,
|
||
)
|
||
|
||
# if team change - check if this is possible
|
||
if is_different_team(data=data, existing_key_row=existing_key_row):
|
||
if llm_router is None:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail={
|
||
"error": "LLM router not found. Please set it up by passing in a valid config.yaml or adding models via the UI."
|
||
},
|
||
)
|
||
if team_obj is None:
|
||
raise HTTPException(
|
||
status_code=500,
|
||
detail={"error": "Team object not found for team change validation"},
|
||
)
|
||
await validate_key_team_change(
|
||
key=existing_key_row,
|
||
team=team_obj,
|
||
change_initiated_by=user_api_key_dict,
|
||
llm_router=llm_router,
|
||
)
|
||
|
||
# Validate MCP servers in object_permission against the effective team
|
||
if data.object_permission is not None:
|
||
normalized_object_permission: Final = await _validate_mcp_servers_for_key_update(
|
||
data=data,
|
||
team_obj=team_obj,
|
||
existing_key_row=existing_key_row,
|
||
prisma_client=prisma_client,
|
||
user_api_key_cache=user_api_key_cache,
|
||
is_proxy_admin=_is_proxy_admin,
|
||
)
|
||
if normalized_object_permission is not None:
|
||
data.object_permission = LiteLLM_ObjectPermissionBase(**normalized_object_permission)
|
||
|
||
|
||
@router.post("/key/update", tags=["key management"], dependencies=[Depends(user_api_key_auth)])
|
||
@management_endpoint_wrapper
|
||
async def update_key_fn(
|
||
request: Request,
|
||
data: UpdateKeyRequest,
|
||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||
litellm_changed_by: str | None = Header(
|
||
None,
|
||
description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability",
|
||
),
|
||
):
|
||
"""
|
||
Update an existing API key's parameters.
|
||
|
||
Parameters:
|
||
- key: Optional[str] - The key to update. Either key or key_alias must be provided.
|
||
- key_alias: Optional[str] - User-friendly key alias. If key is omitted, also identifies the key to update (must match exactly one key, same as /key/delete's key_aliases)
|
||
- user_id: Optional[str] - User ID associated with key
|
||
- team_id: Optional[str] - Team ID associated with key
|
||
- agent_id: Optional[str] - The agent id associated with the key.
|
||
- organization_id: Optional[str] - The organization id of the key.
|
||
- budget_id: Optional[str] - The budget id associated with the key. Created by calling `/budget/new`.
|
||
- models: Optional[list] - Model_name's a user is allowed to call
|
||
- tags: Optional[List[str]] - Tags for organizing keys (Enterprise only)
|
||
- prompts: Optional[List[str]] - List of prompts that the key is allowed to use.
|
||
- enforced_params: Optional[List[str]] - List of enforced params for the key (Enterprise only). [Docs](https://docs.litellm.ai/docs/proxy/enterprise#enforce-required-params-for-llm-requests)
|
||
- spend: Optional[float] - Amount spent by key
|
||
- max_budget: Optional[float] - Max budget for key
|
||
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}
|
||
- budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||
- budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.)
|
||
- soft_budget: Optional[float] - [TODO] Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached.
|
||
- max_parallel_requests: Optional[int] - Rate limit for parallel requests
|
||
- metadata: Optional[dict] - Metadata for key. Example {"team": "core-infra", "app": "app2"}
|
||
- tpm_limit: Optional[int] - Tokens per minute limit
|
||
- rpm_limit: Optional[int] - Requests per minute limit
|
||
- model_rpm_limit: Optional[dict] - Model-specific RPM limits {"gpt-4": 100, "claude-v1": 200}
|
||
- mcp_rpm_limit: Optional[dict] - Per-MCP-server RPM limits, keyed by MCP server name {"github": 100, "slack": 200}
|
||
- tag_rpm_limit: Optional[dict] - Per-request-tag RPM limits, keyed by request tag {"cell-1": 1000, "cell-2": 500}. Each tag gets an independent counter; absent tags fall back to the key-level rpm limit.
|
||
- model_tpm_limit: Optional[dict] - Model-specific TPM limits {"gpt-4": 100000, "claude-v1": 200000}
|
||
- default_estimated_output_tokens: Optional[int] - Proxy admin only. Expected output tokens reserved for TPM limiting when a request omits max_tokens. Positive integer.
|
||
- default_estimated_output_tokens_per_model: Optional[dict] - Proxy admin only. Per-model override of the above {"gpt-4": 4096, "gpt-3.5-turbo": 1024}
|
||
- tpm_limit_type: Optional[str] - TPM rate limit type - "best_effort_throughput", "guaranteed_throughput", or "dynamic"
|
||
- 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)
|
||
- 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
|
||
- policies: Optional[List[str]] - List of policy names to apply to the key. Policies define guardrails, conditions, and inheritance rules.
|
||
- disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key.
|
||
- throttle_on_budget_exceeded: Optional[bool] - When the key exceeds its max_budget, throttle its tpm/rpm to the global budget_exceeded_throttle_percentage instead of blocking the key entirely.
|
||
- enable_prompt_caching: Optional[bool] - Auto-inject prompt caching breakpoints (Anthropic cache_control markers) on requests made with this key. Anthropic and Bedrock Claude models only.
|
||
- prompts: Optional[List[str]] - List of prompts that the key is allowed to use.
|
||
- blocked: Optional[bool] - Whether the key is blocked
|
||
- aliases: Optional[dict] - Model aliases for the key - [Docs](https://litellm.vercel.app/docs/proxy/virtual_keys#model-aliases)
|
||
- config: Optional[dict] - [DEPRECATED PARAM] Key-specific config.
|
||
- temp_budget_increase: Optional[float] - Temporary budget increase for the key (Enterprise only).
|
||
- temp_budget_expiry: Optional[str] - Expiry time for the temporary budget increase (Enterprise only).
|
||
- allowed_routes: Optional[list] - List of allowed routes for the key. Store the actual route or store a wildcard pattern for a set of routes. Example - ["/chat/completions", "/embeddings", "/keys/*"]
|
||
- allowed_passthrough_routes: Optional[list] - List of allowed pass through routes for the key. Store the actual route or store a wildcard pattern for a set of routes. Example - ["/my-custom-endpoint"]. Use this instead of allowed_routes, if you just want to specify which pass through routes the key can access, without specifying the routes. If allowed_routes is specified, allowed_passthrough_routes is ignored.
|
||
- prompts: Optional[List[str]] - List of allowed prompts for the key. If specified, the key will only be able to use these specific prompts.
|
||
- object_permission: Optional[LiteLLM_ObjectPermissionBase] - key-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission.
|
||
- auto_rotate: Optional[bool] - Whether this key should be automatically rotated
|
||
- rotation_interval: Optional[str] - How often to rotate this key (e.g., '30d', '90d'). Required if auto_rotate=True
|
||
- allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint.
|
||
- router_settings: Optional[UpdateRouterConfig] - key-specific router settings. Example - {"model_group_retry_policy": {"gpt-4": {"RateLimitErrorRetries": 5}}}. IF null or {} then no router settings.
|
||
- access_group_ids: Optional[List[str]] - List of access group IDs to associate with the key. Access groups define which models a key can access. Example - ["access_group_1", "access_group_2"].
|
||
- budget_limits: Optional[list] - List of concurrent budget windows for the key. Each window specifies a budget_limit, time_period, and optional budget_duration. Example - [{"budget_limit": 10.0, "time_period": "1d"}, {"budget_limit": 50.0, "time_period": "7d"}].
|
||
|
||
Example:
|
||
```bash
|
||
curl --location 'http://0.0.0.0:4000/key/update' \
|
||
--header 'Authorization: Bearer sk-1234' \
|
||
--header 'Content-Type: application/json' \
|
||
--data '{
|
||
"key": "sk-1234",
|
||
"key_alias": "my-key",
|
||
"user_id": "user-1234",
|
||
"team_id": "team-1234",
|
||
"max_budget": 100,
|
||
"metadata": {"any_key": "any-val"},
|
||
}'
|
||
```
|
||
"""
|
||
from litellm.proxy.proxy_server import (
|
||
llm_router,
|
||
premium_user,
|
||
prisma_client,
|
||
proxy_logging_obj,
|
||
user_api_key_cache,
|
||
user_custom_key_update,
|
||
)
|
||
|
||
try:
|
||
# Validate budget values are not negative and are finite numbers
|
||
if data.max_budget is not None and (not math.isfinite(data.max_budget) or data.max_budget < 0):
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail={"error": f"max_budget must be a non-negative finite number. Received: {data.max_budget}"},
|
||
)
|
||
|
||
# get the row from db
|
||
existing_key_row: Final = await _get_and_validate_existing_key(
|
||
token=data.key,
|
||
prisma_client=prisma_client,
|
||
key_alias=data.key_alias,
|
||
)
|
||
key: Final = _resolve_token_to_update(data=data, existing_key_row=existing_key_row)
|
||
data.key = key
|
||
|
||
await _validate_update_key_data(
|
||
data=data,
|
||
existing_key_row=existing_key_row,
|
||
user_api_key_dict=user_api_key_dict,
|
||
llm_router=llm_router,
|
||
premium_user=premium_user,
|
||
prisma_client=prisma_client,
|
||
user_api_key_cache=user_api_key_cache,
|
||
)
|
||
|
||
# Custom key update hook
|
||
custom_key_update_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = user_custom_key_update
|
||
if custom_key_update_hook is not None:
|
||
if inspect.iscoroutinefunction(custom_key_update_hook):
|
||
result: Final = await custom_key_update_hook(data)
|
||
else:
|
||
raise ValueError("user_custom_key_update must be a coroutine")
|
||
decision: Final = result.get("decision", True)
|
||
message: Final = result.get("message", "Authentication Failed - Custom Auth Rule")
|
||
if not decision:
|
||
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=message)
|
||
|
||
# Enforce upperbound key params on update (don't fill defaults)
|
||
_enforce_upperbound_key_params(data, fill_defaults=False)
|
||
non_default_values: Final = await prepare_key_update_data(data=data, existing_key_row=existing_key_row)
|
||
|
||
# Only validate key_alias format if it's actually being changed
|
||
new_key_alias: Final = non_default_values.get("key_alias", None)
|
||
if new_key_alias != existing_key_row.key_alias:
|
||
_validate_key_alias_format(key_alias=new_key_alias)
|
||
|
||
await _enforce_unique_key_alias(
|
||
key_alias=non_default_values.get("key_alias", None),
|
||
prisma_client=prisma_client,
|
||
existing_key_token=existing_key_row.token,
|
||
)
|
||
|
||
# Handle rotation fields if auto_rotate is being enabled
|
||
_set_key_rotation_fields(
|
||
non_default_values,
|
||
non_default_values.get("auto_rotate", False),
|
||
non_default_values.get("rotation_interval"),
|
||
existing_key_alias=existing_key_row.key_alias,
|
||
)
|
||
|
||
_data: Final = {**non_default_values, "token": key}
|
||
if prisma_client is None:
|
||
raise Exception("Not connected to DB!")
|
||
response: Final = await prisma_client.update_data(token=key, data=_data)
|
||
|
||
# Delete - key from cache, since it's been updated!
|
||
# key updated - a new model could have been added to this key. it should not block requests after this is done
|
||
await _delete_cache_key_object(
|
||
hashed_token=_hash_token_if_needed(key),
|
||
user_api_key_cache=user_api_key_cache,
|
||
proxy_logging_obj=proxy_logging_obj,
|
||
)
|
||
|
||
# After the key's own cache entry is dropped, so a failure here cannot leave the key
|
||
# authenticating against the access groups it just lost.
|
||
await sync_key_update_access_group_membership(
|
||
prisma_client=prisma_client,
|
||
key_token=_hash_token_if_needed(key),
|
||
data=data,
|
||
existing_key_row=existing_key_row,
|
||
)
|
||
|
||
if data.spend is not None:
|
||
from litellm.proxy.proxy_server import spend_counter_cache
|
||
|
||
counter_key: Final = f"spend:key:{_hash_token_if_needed(key)}"
|
||
spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=data.spend, ttl=60)
|
||
if spend_counter_cache.redis_cache is not None:
|
||
try:
|
||
await spend_counter_cache.redis_cache.async_set_cache(key=counter_key, value=data.spend, ttl=60)
|
||
except Exception as redis_err:
|
||
verbose_proxy_logger.warning(
|
||
"Failed to update spend counter %s in Redis after key spend update: %s. "
|
||
"Budget checks may use stale value until counter expires.",
|
||
counter_key,
|
||
redis_err,
|
||
)
|
||
|
||
asyncio.create_task(
|
||
KeyManagementEventHooks.async_key_updated_hook(
|
||
data=data,
|
||
existing_key_row=existing_key_row,
|
||
response=response,
|
||
user_api_key_dict=user_api_key_dict,
|
||
litellm_changed_by=litellm_changed_by,
|
||
)
|
||
)
|
||
|
||
if response is None:
|
||
raise ValueError("Failed to update key got response = None")
|
||
|
||
return {"key": key, **response["data"]}
|
||
# update based on remaining passed in values
|
||
except Exception as e:
|
||
verbose_proxy_logger.exception("litellm.proxy.proxy_server.update_key_fn(): Exception occured - %s", e)
|
||
if isinstance(e, HTTPException):
|
||
raise ProxyException(
|
||
message=getattr(e, "detail", f"Authentication Error({e})"),
|
||
type=ProxyErrorTypes.auth_error,
|
||
param=getattr(e, "param", "None"),
|
||
code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST),
|
||
)
|
||
elif isinstance(e, ProxyException):
|
||
raise e
|
||
raise ProxyException(
|
||
message="Authentication Error, " + str(e),
|
||
type=ProxyErrorTypes.auth_error,
|
||
param=getattr(e, "param", "None"),
|
||
code=status.HTTP_400_BAD_REQUEST,
|
||
)
|
||
|
||
|
||
@router.post(
|
||
"/key/bulk_update",
|
||
tags=["key management"],
|
||
dependencies=[Depends(user_api_key_auth)],
|
||
response_model=BulkUpdateKeyResponse,
|
||
)
|
||
@management_endpoint_wrapper
|
||
async def bulk_update_keys(
|
||
data: BulkUpdateKeyRequest,
|
||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||
litellm_changed_by: str | None = Header(
|
||
None,
|
||
description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability",
|
||
),
|
||
):
|
||
"""
|
||
Bulk update multiple keys at once.
|
||
|
||
This endpoint allows updating multiple keys in a single request. Each key update
|
||
is processed independently - if some updates fail, others will still succeed.
|
||
|
||
Parameters:
|
||
- keys: List[BulkUpdateKeyRequestItem] - List of key update requests, each containing:
|
||
- key: str - The key identifier (token) to update
|
||
- budget_id: Optional[str] - Budget ID associated with the key
|
||
- max_budget: Optional[float] - Max budget for key
|
||
- team_id: Optional[str] - Team ID associated with key
|
||
- tags: Optional[List[str]] - Tags for organizing keys
|
||
|
||
Returns:
|
||
- total_requested: int - Total number of keys requested for update
|
||
- successful_updates: List[SuccessfulKeyUpdate] - List of successfully updated keys with their updated info
|
||
- failed_updates: List[FailedKeyUpdate] - List of failed updates with key_info and failed_reason
|
||
|
||
Example request:
|
||
```bash
|
||
curl --location 'http://0.0.0.0:4000/key/bulk_update' \
|
||
--header 'Authorization: Bearer sk-1234' \
|
||
--header 'Content-Type: application/json' \
|
||
--data '{
|
||
"keys": [
|
||
{
|
||
"key": "sk-1234",
|
||
"max_budget": 100.0,
|
||
"team_id": "team-123",
|
||
"tags": ["production", "api"]
|
||
},
|
||
{
|
||
"key": "sk-5678",
|
||
"budget_id": "budget-456",
|
||
"tags": ["staging"]
|
||
}
|
||
]
|
||
}'
|
||
```
|
||
"""
|
||
from litellm.proxy.proxy_server import (
|
||
llm_router,
|
||
prisma_client,
|
||
proxy_logging_obj,
|
||
user_api_key_cache,
|
||
user_custom_key_update,
|
||
)
|
||
|
||
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value:
|
||
raise HTTPException(
|
||
status_code=403,
|
||
detail={"error": "Only proxy admins can perform bulk key updates"},
|
||
)
|
||
|
||
if prisma_client is None:
|
||
raise HTTPException(
|
||
status_code=500,
|
||
detail={"error": "Database not connected"},
|
||
)
|
||
|
||
if not data.keys:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail={"error": "No keys provided for update"},
|
||
)
|
||
|
||
MAX_BATCH_SIZE: Final = 500
|
||
if len(data.keys) > MAX_BATCH_SIZE:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail={"error": f"Maximum {MAX_BATCH_SIZE} keys can be updated at once. Found {len(data.keys)} keys."},
|
||
)
|
||
|
||
successful_updates: Final[list[SuccessfulKeyUpdate]] = []
|
||
failed_updates: Final[list[FailedKeyUpdate]] = []
|
||
|
||
for key_update_item in data.keys:
|
||
try:
|
||
update_key_request = UpdateKeyRequest(
|
||
key=key_update_item.key,
|
||
budget_id=key_update_item.budget_id,
|
||
max_budget=key_update_item.max_budget,
|
||
team_id=key_update_item.team_id,
|
||
tags=key_update_item.tags,
|
||
)
|
||
updated_key_info = await _process_single_key_update(
|
||
update_key_request=update_key_request,
|
||
user_api_key_dict=user_api_key_dict,
|
||
litellm_changed_by=litellm_changed_by,
|
||
prisma_client=prisma_client,
|
||
user_api_key_cache=user_api_key_cache,
|
||
proxy_logging_obj=proxy_logging_obj,
|
||
llm_router=llm_router,
|
||
user_custom_key_update=user_custom_key_update,
|
||
)
|
||
|
||
successful_updates.append(
|
||
SuccessfulKeyUpdate(
|
||
key=key_update_item.key,
|
||
key_info=updated_key_info,
|
||
)
|
||
)
|
||
|
||
except Exception as e:
|
||
verbose_proxy_logger.exception("Failed to update key %s: %s", key_update_item.key, e)
|
||
|
||
if isinstance(e, HTTPException):
|
||
error_detail = e.detail
|
||
if isinstance(error_detail, dict):
|
||
error_message = error_detail.get("error", str(e))
|
||
else:
|
||
error_message = str(error_detail)
|
||
elif isinstance(e, ProxyException):
|
||
error_message = e.message
|
||
else:
|
||
error_message = str(e)
|
||
|
||
key_info = None
|
||
try:
|
||
existing_key_row = await prisma_client.get_data(
|
||
token=key_update_item.key,
|
||
table_name="key",
|
||
query_type="find_unique",
|
||
)
|
||
if existing_key_row is not None:
|
||
if hasattr(existing_key_row, "model_dump"):
|
||
key_info = existing_key_row.model_dump()
|
||
elif hasattr(existing_key_row, "dict"):
|
||
key_info = existing_key_row.dict()
|
||
if key_info:
|
||
key_info.pop("token", None)
|
||
except Exception:
|
||
pass
|
||
|
||
failed_updates.append(
|
||
FailedKeyUpdate(
|
||
key=key_update_item.key,
|
||
key_info=key_info,
|
||
failed_reason=error_message,
|
||
)
|
||
)
|
||
|
||
return BulkUpdateKeyResponse(
|
||
total_requested=len(data.keys),
|
||
successful_updates=successful_updates,
|
||
failed_updates=failed_updates,
|
||
)
|
||
|
||
|
||
def _build_failed_team_key_update(
|
||
token: str,
|
||
exception: Exception,
|
||
existing_key_row: LiteLLM_VerificationToken | None,
|
||
) -> FailedKeyUpdate:
|
||
"""Normalize an exception from the per-key update loop into a FailedKeyUpdate."""
|
||
if isinstance(exception, HTTPException):
|
||
detail: Final = exception.detail
|
||
if isinstance(detail, dict):
|
||
error_message = detail.get("error", str(exception))
|
||
else:
|
||
error_message = str(detail)
|
||
elif isinstance(exception, ProxyException):
|
||
error_message = exception.message
|
||
else:
|
||
error_message = str(exception)
|
||
|
||
key_info: dict[str, object] | None = None
|
||
if existing_key_row is not None:
|
||
if hasattr(existing_key_row, "model_dump"):
|
||
key_info = existing_key_row.model_dump()
|
||
elif hasattr(existing_key_row, "dict"):
|
||
key_info = existing_key_row.dict()
|
||
if key_info:
|
||
key_info.pop("token", None)
|
||
|
||
return FailedKeyUpdate(key=token, key_info=key_info, failed_reason=error_message)
|
||
|
||
|
||
@router.post(
|
||
"/team/key/bulk_update",
|
||
tags=["key management"],
|
||
dependencies=[Depends(user_api_key_auth)],
|
||
response_model=BulkUpdateKeyResponse,
|
||
)
|
||
@management_endpoint_wrapper
|
||
async def bulk_update_team_keys(
|
||
data: BulkUpdateTeamKeysRequest,
|
||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||
litellm_changed_by: str | None = Header(
|
||
None,
|
||
description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability",
|
||
),
|
||
):
|
||
"""
|
||
Apply one update payload to many keys inside a single team.
|
||
|
||
Pass `team_id` plus either `key_ids` or `all_keys_in_team=True`. The
|
||
`update_fields` payload is broadcast to every selected key. Per-key
|
||
failures are returned in `failed_updates` rather than aborting the batch.
|
||
|
||
Callable by proxy admins, or by team admins with `KEY_UPDATE` permission.
|
||
"""
|
||
from litellm.proxy.proxy_server import (
|
||
llm_router,
|
||
prisma_client,
|
||
proxy_logging_obj,
|
||
user_api_key_cache,
|
||
user_custom_key_update,
|
||
)
|
||
|
||
if prisma_client is None:
|
||
raise HTTPException(
|
||
status_code=500,
|
||
detail={"error": "Database not connected"},
|
||
)
|
||
|
||
if not data.team_id:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail={"error": "team_id is required"},
|
||
)
|
||
|
||
MAX_BATCH_SIZE: Final = 500
|
||
if data.key_ids is not None and len(data.key_ids) > MAX_BATCH_SIZE:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail={
|
||
"error": f"Maximum {MAX_BATCH_SIZE} keys can be updated at once. Found {len(data.key_ids)} key_ids."
|
||
},
|
||
)
|
||
|
||
if data.all_keys_in_team:
|
||
# "all" excludes blocked/expired — bulk refresh shouldn't revive a key an admin disabled.
|
||
# `blocked` is Boolean? with no default; `/key/generate` writes NULL. Prisma's `NOT`
|
||
# excludes NULLs, so explicitly OR `false` with `null` to include them.
|
||
now: Final = datetime.now(timezone.utc)
|
||
existing_keys = await _prisma_table(VerificationTokenRepository(prisma_client)).find_many(
|
||
where={
|
||
"team_id": data.team_id,
|
||
"AND": [
|
||
{"OR": [{"blocked": False}, {"blocked": None}]},
|
||
{"OR": [{"expires": None}, {"expires": {"gt": now}}]},
|
||
],
|
||
},
|
||
order={"token": "asc"},
|
||
take=MAX_BATCH_SIZE + 1,
|
||
)
|
||
if len(existing_keys) > MAX_BATCH_SIZE:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail={
|
||
"error": f"Team {data.team_id} has more than {MAX_BATCH_SIZE} keys. Use `key_ids` to update in batches of {MAX_BATCH_SIZE}."
|
||
},
|
||
)
|
||
requested_tokens = cast( # cast-ok: token is the table's primary key, so a row read back always carries one
|
||
"list[str]", [row.token for row in existing_keys]
|
||
)
|
||
else:
|
||
if data.key_ids is None or len(data.key_ids) == 0:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail={"error": "key_ids must be provided when all_keys_in_team is False"},
|
||
)
|
||
# Dedupe by hashed form — duplicates collapse to one update.
|
||
requested_tokens = []
|
||
hashed_key_ids: Final = []
|
||
seen_hashes: Final = set()
|
||
for k in data.key_ids:
|
||
h = _hash_token_if_needed(k)
|
||
if h in seen_hashes:
|
||
continue
|
||
seen_hashes.add(h)
|
||
requested_tokens.append(k)
|
||
hashed_key_ids.append(h)
|
||
existing_keys = await _prisma_table(VerificationTokenRepository(prisma_client)).find_many(
|
||
where={"team_id": data.team_id, "token": {"in": hashed_key_ids}}
|
||
)
|
||
|
||
# Anchor membership check on data.team_id (not existing_keys[0]); empty result must still gate non-admins.
|
||
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value:
|
||
auth_anchor: Final = (
|
||
existing_keys[0]
|
||
if existing_keys
|
||
else LiteLLM_VerificationToken(
|
||
token="__team_scope_auth_check__",
|
||
team_id=data.team_id,
|
||
models=[],
|
||
)
|
||
)
|
||
await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint(
|
||
user_api_key_dict=user_api_key_dict,
|
||
route=KeyManagementRoutes.KEY_UPDATE,
|
||
prisma_client=prisma_client,
|
||
existing_key_row=auth_anchor,
|
||
user_api_key_cache=user_api_key_cache,
|
||
)
|
||
|
||
# Block metadata.allowed_passthrough_routes for non-admins — the runtime
|
||
# route checker reads it from key/team metadata to grant passthrough.
|
||
_check_passthrough_routes_caller_permission(data=data.update_fields, user_api_key_dict=user_api_key_dict)
|
||
|
||
if not requested_tokens:
|
||
raise HTTPException(
|
||
status_code=404,
|
||
detail={"error": f"No keys found for team {data.team_id}"},
|
||
)
|
||
|
||
existing_by_token: Final = {row.token: row for row in existing_keys}
|
||
update_field_dict: Final = data.update_fields.model_dump(exclude_unset=True)
|
||
|
||
successful_updates: Final[list[SuccessfulKeyUpdate]] = []
|
||
failed_updates: Final[list[FailedKeyUpdate]] = []
|
||
|
||
for token in requested_tokens:
|
||
db_token = _hash_token_if_needed(token)
|
||
try:
|
||
if db_token not in existing_by_token:
|
||
raise HTTPException(
|
||
status_code=404,
|
||
detail={"error": f"Key not found in team {data.team_id}"},
|
||
)
|
||
|
||
# team_id from validated scope, never user payload — drives _check_team_key_limits.
|
||
update_key_request = UpdateKeyRequest.model_validate(
|
||
{
|
||
"key": token,
|
||
"team_id": data.team_id,
|
||
**update_field_dict,
|
||
}
|
||
)
|
||
updated_key_info = await _process_single_key_update(
|
||
update_key_request=update_key_request,
|
||
user_api_key_dict=user_api_key_dict,
|
||
litellm_changed_by=litellm_changed_by,
|
||
prisma_client=prisma_client,
|
||
user_api_key_cache=user_api_key_cache,
|
||
proxy_logging_obj=proxy_logging_obj,
|
||
llm_router=llm_router,
|
||
user_custom_key_update=user_custom_key_update,
|
||
existing_key_row=existing_by_token[db_token],
|
||
)
|
||
|
||
successful_updates.append(SuccessfulKeyUpdate(key=token, key_info=updated_key_info))
|
||
|
||
except Exception as e:
|
||
# Log the hashed prefix — `token` may be a raw sk-... and ERROR logs persist.
|
||
verbose_proxy_logger.exception("Failed to update key %s... in team %s: %s", db_token[:12], data.team_id, e)
|
||
failed_updates.append(
|
||
_build_failed_team_key_update(
|
||
token=token,
|
||
exception=e,
|
||
existing_key_row=existing_by_token.get(db_token),
|
||
)
|
||
)
|
||
|
||
return BulkUpdateKeyResponse(
|
||
total_requested=len(requested_tokens),
|
||
successful_updates=successful_updates,
|
||
failed_updates=failed_updates,
|
||
)
|
||
|
||
|
||
async def validate_key_team_change(
|
||
key: LiteLLM_VerificationToken,
|
||
team: LiteLLM_TeamTable,
|
||
change_initiated_by: UserAPIKeyAuth,
|
||
llm_router: Router,
|
||
):
|
||
"""
|
||
Validate that a key can be moved to a new team.
|
||
|
||
- The team must have access to the key's models
|
||
- The key's user_id must be a member of the team
|
||
- The key's tpm/rpm limit must be less than the team's tpm/rpm limit
|
||
- The person initiating the change must be either Proxy Admin or Team Admin
|
||
"""
|
||
# Check if the team has access to the key's models
|
||
if len(key.models) > 0:
|
||
for model in key.models:
|
||
# Skip special sentinel values — "all-team-models" means
|
||
# "use whatever the team allows", so it's always valid.
|
||
if model == SpecialModelNames.all_team_models.value:
|
||
continue
|
||
await can_team_access_model(
|
||
model=model,
|
||
team_object=team,
|
||
llm_router=llm_router,
|
||
)
|
||
|
||
# Check if the key's tpm/rpm limit is less than the team's tpm/rpm limit
|
||
if key.tpm_limit is not None:
|
||
if team.tpm_limit and key.tpm_limit > team.tpm_limit:
|
||
raise HTTPException(
|
||
status_code=403,
|
||
detail=f"Key={key.token} has a tpm_limit={key.tpm_limit} which is greater than the team's tpm_limit={team.tpm_limit}.",
|
||
)
|
||
if team.rpm_limit and key.rpm_limit and key.rpm_limit > team.rpm_limit:
|
||
raise HTTPException(
|
||
status_code=403,
|
||
detail=f"Key={key.token} has a rpm_limit={key.rpm_limit} which is greater than the team's rpm_limit={team.rpm_limit}.",
|
||
)
|
||
|
||
# Check if the key's user_id is a member of the team
|
||
member_object: Final = _get_user_in_team(team_table=cast(LiteLLM_TeamTableCachedObj, team), user_id=key.user_id)
|
||
if key.user_id is not None:
|
||
if not member_object:
|
||
raise HTTPException(
|
||
status_code=403,
|
||
detail=f"User={key.user_id} is not a member of the team={team.team_id}. Check team members via `/team/info`.",
|
||
)
|
||
|
||
# Check if the person initiating the change is a Proxy Admin or Team Admin
|
||
if (
|
||
change_initiated_by.user_role == LitellmUserRoles.PROXY_ADMIN.value
|
||
or _is_user_team_admin(
|
||
user_api_key_dict=change_initiated_by,
|
||
team_obj=team,
|
||
)
|
||
or TeamMemberPermissionChecks.does_team_member_have_permissions_for_endpoint(
|
||
team_member_object=member_object,
|
||
team_table=cast(LiteLLM_TeamTableCachedObj, team),
|
||
route=KeyManagementRoutes.KEY_UPDATE.value,
|
||
)
|
||
):
|
||
return
|
||
else:
|
||
raise HTTPException(
|
||
status_code=403,
|
||
detail=f"User={change_initiated_by.user_id} is not a Proxy Admin or Team Admin for team={team.team_id}. Please ask your Proxy Admin to allow this action under 'Member Permissions' for this team.",
|
||
)
|
||
|
||
|
||
@router.post("/key/delete", tags=["key management"], dependencies=[Depends(user_api_key_auth)])
|
||
@management_endpoint_wrapper
|
||
async def delete_key_fn(
|
||
data: KeyRequest,
|
||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||
litellm_changed_by: str | None = Header(
|
||
None,
|
||
description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability",
|
||
),
|
||
):
|
||
"""
|
||
Delete a key from the key management system.
|
||
|
||
Parameters::
|
||
- keys (List[str]): A list of keys or hashed keys to delete. Example {"keys": ["sk-QWrxEynunsNpV1zT48HIrw", "837e17519f44683334df5291321d97b8bf1098cd490e49e215f6fea935aa28be"]}
|
||
- key_aliases (List[str]): A list of key aliases to delete. Can be passed instead of `keys`.Example {"key_aliases": ["alias1", "alias2"]}
|
||
|
||
Returns:
|
||
- deleted_keys (List[str]): A list of deleted keys. Example {"deleted_keys": ["sk-QWrxEynunsNpV1zT48HIrw", "837e17519f44683334df5291321d97b8bf1098cd490e49e215f6fea935aa28be"]}
|
||
|
||
Example:
|
||
```bash
|
||
curl --location 'http://0.0.0.0:4000/key/delete' \
|
||
--header 'Authorization: Bearer sk-1234' \
|
||
--header 'Content-Type: application/json' \
|
||
--data '{
|
||
"keys": ["sk-QWrxEynunsNpV1zT48HIrw"]
|
||
}'
|
||
```
|
||
|
||
Raises:
|
||
HTTPException: If an error occurs during key deletion.
|
||
"""
|
||
try:
|
||
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
|
||
|
||
if prisma_client is None:
|
||
raise Exception("Not connected to DB!")
|
||
|
||
# Normalize litellm_changed_by: if it's a Header object or not a string, convert to None
|
||
if litellm_changed_by is not None and not isinstance(litellm_changed_by, str):
|
||
litellm_changed_by = None
|
||
|
||
## only allow user to delete keys they own
|
||
verbose_proxy_logger.debug("user_api_key_dict.user_role: %s", user_api_key_dict.user_role)
|
||
|
||
num_keys_to_be_deleted = 0
|
||
deleted_keys = []
|
||
if data.keys:
|
||
number_deleted_keys, _keys_being_deleted = await delete_verification_tokens(
|
||
tokens=data.keys,
|
||
user_api_key_cache=user_api_key_cache,
|
||
user_api_key_dict=user_api_key_dict,
|
||
litellm_changed_by=litellm_changed_by,
|
||
)
|
||
num_keys_to_be_deleted = len(data.keys)
|
||
deleted_keys = data.keys
|
||
elif data.key_aliases:
|
||
number_deleted_keys, _keys_being_deleted = await delete_key_aliases(
|
||
key_aliases=data.key_aliases,
|
||
prisma_client=prisma_client,
|
||
user_api_key_cache=user_api_key_cache,
|
||
user_api_key_dict=user_api_key_dict,
|
||
litellm_changed_by=litellm_changed_by,
|
||
)
|
||
num_keys_to_be_deleted = len(data.key_aliases)
|
||
deleted_keys = data.key_aliases
|
||
else:
|
||
raise ValueError("Invalid request type")
|
||
|
||
if number_deleted_keys is None:
|
||
raise ProxyException(
|
||
message="Failed to delete keys got None response from delete_verification_token",
|
||
type=ProxyErrorTypes.internal_server_error,
|
||
param="keys",
|
||
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||
)
|
||
verbose_proxy_logger.debug("/key/delete - deleted_keys=%s", number_deleted_keys)
|
||
|
||
try:
|
||
assert num_keys_to_be_deleted == len(deleted_keys)
|
||
except Exception:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail={
|
||
"error": f"Not all keys passed in were deleted. This probably means you don't have access to delete all the keys passed in. Keys passed in={num_keys_to_be_deleted}, Deleted keys ={number_deleted_keys}"
|
||
},
|
||
)
|
||
|
||
verbose_proxy_logger.debug(
|
||
"/keys/delete - cache after delete: %s", user_api_key_cache.in_memory_cache.cache_dict
|
||
)
|
||
|
||
asyncio.create_task(
|
||
KeyManagementEventHooks.async_key_deleted_hook(
|
||
data=data,
|
||
keys_being_deleted=_keys_being_deleted,
|
||
user_api_key_dict=user_api_key_dict,
|
||
litellm_changed_by=litellm_changed_by,
|
||
response=number_deleted_keys,
|
||
)
|
||
)
|
||
|
||
return {"deleted_keys": deleted_keys}
|
||
except Exception as e:
|
||
verbose_proxy_logger.exception("litellm.proxy.proxy_server.delete_key_fn(): Exception occured - %s", e)
|
||
raise handle_exception_on_proxy(e)
|
||
|
||
|
||
async def _build_model_max_budget_usage(
|
||
api_key_hash: str,
|
||
model_max_budget: Mapping[str, Mapping[str, object]],
|
||
user_api_key_cache: DualCache | None,
|
||
) -> dict[str, dict[str, object]]:
|
||
return await build_model_max_budget_usage(
|
||
entity_type=Litellm_EntityType.KEY,
|
||
entity_id=api_key_hash,
|
||
model_max_budget=model_max_budget,
|
||
cache=user_api_key_cache,
|
||
)
|
||
|
||
|
||
@router.post(
|
||
"/v2/key/info",
|
||
tags=["key management"],
|
||
dependencies=[Depends(user_api_key_auth)],
|
||
include_in_schema=False,
|
||
)
|
||
async def info_key_fn_v2(
|
||
data: KeyRequest | None = None,
|
||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||
):
|
||
"""
|
||
Retrieve information about a list of keys.
|
||
|
||
**New endpoint**. Currently admin only.
|
||
Parameters:
|
||
keys: Optional[list] = body parameter representing the key(s) in the request
|
||
user_api_key_dict: UserAPIKeyAuth = Dependency representing the user's API key
|
||
Returns:
|
||
Dict containing the key and its associated information
|
||
|
||
Example Curl:
|
||
```
|
||
curl -X GET "http://0.0.0.0:4000/key/info" \
|
||
-H "Authorization: Bearer sk-1234" \
|
||
-d {"keys": ["sk-1", "sk-2", "sk-3"]}
|
||
```
|
||
"""
|
||
from litellm.proxy.proxy_server import (
|
||
model_max_budget_limiter,
|
||
prisma_client,
|
||
)
|
||
|
||
try:
|
||
if prisma_client is None:
|
||
raise Exception(
|
||
"Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys"
|
||
)
|
||
if data is None:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
|
||
detail={"message": "Malformed request. No keys passed in."},
|
||
)
|
||
|
||
# Resolve key_aliases to tokens so we never pass token=None (unbounded query)
|
||
tokens_to_query: Final = list(data.keys) if data.keys else []
|
||
if data.key_aliases:
|
||
alias_rows: Final = await _prisma_table(VerificationTokenRepository(prisma_client)).find_many(
|
||
where={"key_alias": {"in": data.key_aliases}},
|
||
include={"litellm_budget_table": True},
|
||
)
|
||
alias_tokens: Final = [row.token for row in alias_rows if row.token]
|
||
tokens_to_query.extend(alias_tokens)
|
||
|
||
if not tokens_to_query:
|
||
return {"key": data.keys, "info": []}
|
||
|
||
key_info: Final = await prisma_client.get_data(token=tokens_to_query, table_name="key", query_type="find_all")
|
||
if not key_info:
|
||
return {"key": data.keys, "info": []}
|
||
|
||
filtered_key_info: Final = []
|
||
for k in key_info:
|
||
if not await _can_user_query_key_info(
|
||
user_api_key_dict=user_api_key_dict,
|
||
key=k.token,
|
||
key_info=k,
|
||
):
|
||
continue
|
||
try:
|
||
k_dict = k.model_dump()
|
||
except Exception:
|
||
k_dict = k.dict()
|
||
k_token_hash = k_dict.pop("token", None)
|
||
|
||
model_max_budget = k_dict.get("model_max_budget") or {}
|
||
budget_table = k_dict.get("litellm_budget_table") or {}
|
||
if not model_max_budget and isinstance(budget_table, dict):
|
||
model_max_budget = budget_table.get("model_max_budget") or {}
|
||
if model_max_budget and k_token_hash:
|
||
k_dict["model_max_budget_usage"] = await _build_model_max_budget_usage(
|
||
api_key_hash=k_token_hash,
|
||
model_max_budget=model_max_budget,
|
||
user_api_key_cache=model_max_budget_limiter.dual_cache,
|
||
)
|
||
|
||
filtered_key_info.append(k_dict)
|
||
return {"key": data.keys, "info": filtered_key_info}
|
||
|
||
except Exception as e:
|
||
raise handle_exception_on_proxy(e)
|
||
|
||
|
||
@router.get("/key/info", tags=["key management"], dependencies=[Depends(user_api_key_auth)])
|
||
@management_endpoint_wrapper
|
||
async def info_key_fn(
|
||
key: str | None = fastapi.Query(default=None, description="Key in the request parameters"),
|
||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||
):
|
||
"""
|
||
Retrieve information about a key.
|
||
|
||
Parameters:
|
||
- key: str | None (query parameter) - The key to look up. Accepts the plaintext key or its hash.
|
||
Defaults to the key in the Authorization header.
|
||
|
||
Returns:
|
||
- key: str - The key that was looked up, echoed back as it was passed in
|
||
- info: dict - The key's row, minus the hashed token
|
||
- key_alias: str | None - User-friendly key alias
|
||
- spend: float - Amount spent by the key. When budget_duration is set this covers only the
|
||
current budget window, not the key's lifetime
|
||
- max_budget: float | None - Max budget for the key, enforced against spend
|
||
- budget_duration: str | None - Budget reset period ("30d", "1h", etc.)
|
||
- budget_reset_at: datetime | None - When the current budget window ends and spend is next
|
||
reset to 0, not when it was last reset. Reset times snap to standard boundaries in the
|
||
configured timezone (30d and 1mo land on the 1st of the month, 7d on Monday, 1h on the
|
||
hour), so subtracting budget_duration from it does not give the window's start
|
||
- model_max_budget: dict - Per-model budgets, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}
|
||
- model_max_budget_usage: dict | None - Current-window spend per model, present only when
|
||
the key has per-model budgets
|
||
- models: list - Model_name's the key is allowed to call
|
||
- tpm_limit / rpm_limit: int | None - Tokens and requests per minute limits
|
||
- metadata: dict - Metadata for the key, e.g. {"team": "core-infra"}
|
||
- blocked: bool | None - Whether the key is blocked
|
||
- expires: datetime | None - When the key stops authenticating requests
|
||
- last_active: datetime | None - When the key was last used
|
||
- object_permission: dict | None - Resolved vector store / MCP permissions when the key has
|
||
an object_permission_id
|
||
|
||
Example Curl:
|
||
```
|
||
curl -X GET "http://0.0.0.0:4000/key/info?key=sk-test-example-key-123" \
|
||
-H "Authorization: Bearer sk-1234"
|
||
```
|
||
|
||
Example Curl - if no key is passed, it will use the Key Passed in Authorization Header
|
||
```
|
||
curl -X GET "http://0.0.0.0:4000/key/info" \
|
||
-H "Authorization: Bearer sk-test-example-key-123"
|
||
```
|
||
"""
|
||
from litellm.proxy.proxy_server import (
|
||
model_max_budget_limiter,
|
||
prisma_client,
|
||
)
|
||
|
||
try:
|
||
if prisma_client is None:
|
||
raise Exception(
|
||
"Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys"
|
||
)
|
||
|
||
# default to using Auth token if no key is passed in
|
||
key = key or user_api_key_dict.api_key
|
||
hashed_key: str | None = key
|
||
if key is not None:
|
||
hashed_key = _hash_token_if_needed(token=key)
|
||
key_info = await _prisma_table(VerificationTokenRepository(prisma_client)).find_unique(
|
||
where={"token": hashed_key},
|
||
include={"litellm_budget_table": True},
|
||
)
|
||
if key_info is None:
|
||
raise ProxyException(
|
||
message="Key not found in database",
|
||
type=ProxyErrorTypes.not_found_error,
|
||
param="key",
|
||
code=status.HTTP_404_NOT_FOUND,
|
||
)
|
||
|
||
if (
|
||
await _can_user_query_key_info(
|
||
user_api_key_dict=user_api_key_dict,
|
||
key=key,
|
||
key_info=key_info,
|
||
)
|
||
is not True
|
||
):
|
||
raise HTTPException(
|
||
status_code=status.HTTP_403_FORBIDDEN,
|
||
detail=f"You are not allowed to access this key's info. Your role={user_api_key_dict.user_role}",
|
||
)
|
||
## REMOVE HASHED TOKEN INFO BEFORE RETURNING ##
|
||
try:
|
||
key_info = key_info.model_dump()
|
||
except Exception:
|
||
# if using pydantic v1
|
||
key_info = key_info.dict() # pyright: ignore[reportDeprecated] # deliberate pydantic v1 fallback
|
||
key_token_hash: Final = key_info.pop("token")
|
||
|
||
model_max_budget = key_info.get("model_max_budget") or {}
|
||
budget_table: Final = key_info.get("litellm_budget_table") or {}
|
||
if not model_max_budget and isinstance(budget_table, dict):
|
||
model_max_budget = budget_table.get("model_max_budget") or {}
|
||
if model_max_budget and key_token_hash:
|
||
key_info["model_max_budget_usage"] = await _build_model_max_budget_usage(
|
||
api_key_hash=key_token_hash,
|
||
model_max_budget=model_max_budget,
|
||
user_api_key_cache=model_max_budget_limiter.dual_cache,
|
||
)
|
||
|
||
# Attach object_permission if object_permission_id is set
|
||
key_info = await attach_object_permission_to_dict(key_info, prisma_client)
|
||
|
||
return {"key": key, "info": key_info}
|
||
except Exception as e:
|
||
raise handle_exception_on_proxy(e)
|
||
|
||
|
||
def _check_model_access_group(models: list[str] | None, llm_router: Router | None, premium_user: bool) -> Literal[True]:
|
||
"""
|
||
if is_model_access_group is True + is_wildcard_route is True, check if user is a premium user
|
||
|
||
Return True if user is a premium user, False otherwise
|
||
"""
|
||
if models is None or llm_router is None:
|
||
return True
|
||
|
||
for model in models:
|
||
if llm_router._is_model_access_group_for_wildcard_route(model_access_group=model):
|
||
if not premium_user:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_403_FORBIDDEN,
|
||
detail={
|
||
"error": f"Setting a model access group on a wildcard model is only available for LiteLLM Enterprise users.{CommonProxyErrors.not_premium_user.value}"
|
||
},
|
||
)
|
||
|
||
return True
|
||
|
||
|
||
async def generate_key_helper_fn(
|
||
request_type: Literal["user", "key"], # identifies if this request is from /user/new or /key/generate
|
||
duration: str | None = None,
|
||
models: list = [],
|
||
aliases: dict = {},
|
||
config: dict = {},
|
||
spend: float = 0.0,
|
||
key_max_budget: float | None = None, # key_max_budget is used to Budget Per key
|
||
key_budget_duration: str | None = None,
|
||
budget_id: float | None = None, # budget id <-> LiteLLM_BudgetTable
|
||
soft_budget: float | None = None, # soft_budget is used to set soft Budgets Per user
|
||
max_budget: float | None = None, # max_budget is used to Budget Per user
|
||
blocked: bool | None = None,
|
||
budget_duration: str | None = None, # max_budget is used to Budget Per user
|
||
token: str | None = None,
|
||
key: str
|
||
| None = None, # dev-friendly alt param for 'token'. Exposed on `/key/generate` for setting key value yourself.
|
||
user_id: str | None = None,
|
||
user_alias: str | None = None,
|
||
team_id: str | None = None,
|
||
agent_id: str | None = None,
|
||
user_email: str | None = None,
|
||
user_role: str | None = None,
|
||
max_parallel_requests: int | None = None,
|
||
metadata: dict | None = {},
|
||
tpm_limit: int | None = None,
|
||
rpm_limit: int | None = None,
|
||
query_type: Literal["insert_data", "update_data"] = "insert_data",
|
||
update_key_values: dict | None = None,
|
||
key_alias: str | None = None,
|
||
allowed_cache_controls: list | None = [],
|
||
permissions: dict | None = {},
|
||
model_max_budget: dict | None = {},
|
||
budget_fallbacks: dict | None = None,
|
||
model_rpm_limit: dict | None = None,
|
||
model_tpm_limit: dict | None = None,
|
||
mcp_rpm_limit: dict | None = None,
|
||
tag_rpm_limit: dict | None = None,
|
||
guardrails: list | None = None,
|
||
policies: list | None = None,
|
||
prompts: list | None = None,
|
||
teams: list | None = None,
|
||
organization_id: str | None = None,
|
||
project_id: str | None = None,
|
||
table_name: Literal["key", "user"] | None = None,
|
||
send_invite_email: bool | None = None,
|
||
created_by: str | None = None,
|
||
updated_by: str | None = None,
|
||
allowed_routes: list | None = None,
|
||
key_type: str | None = None,
|
||
sso_user_id: str | None = None,
|
||
object_permission_id: str | None = None, # object_permission_id <-> LiteLLM_ObjectPermissionTable
|
||
object_permission: LiteLLM_ObjectPermissionBase | None = None,
|
||
auto_rotate: bool | None = None,
|
||
rotation_interval: str | None = None,
|
||
router_settings: dict | None = None,
|
||
access_group_ids: list[str] | None = None,
|
||
budget_limits: list | None = None, # multiple concurrent budget windows
|
||
):
|
||
from litellm.proxy.proxy_server import premium_user, prisma_client
|
||
|
||
if prisma_client is None:
|
||
raise Exception("Connect Proxy to database to generate keys - https://docs.litellm.ai/docs/proxy/virtual_keys ")
|
||
|
||
if token is None:
|
||
if key is not None:
|
||
token = key
|
||
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)
|
||
|
||
if key_budget_duration is None: # one-time budget
|
||
key_reset_at = None
|
||
else:
|
||
key_reset_at = get_budget_reset_time(budget_duration=key_budget_duration)
|
||
|
||
if budget_duration is None: # one-time budget
|
||
reset_at = None
|
||
else:
|
||
reset_at = get_budget_reset_time(budget_duration=budget_duration)
|
||
|
||
# Initialize reset_at for each budget window
|
||
budget_limits_json: str | None = None
|
||
if budget_limits:
|
||
initialized_windows: Final = []
|
||
for window in budget_limits:
|
||
w = dict(window) if not isinstance(window, dict) else {**window}
|
||
w["reset_at"] = get_budget_reset_time(budget_duration=w["budget_duration"]).isoformat()
|
||
initialized_windows.append(w)
|
||
budget_limits_json = json.dumps(initialized_windows)
|
||
|
||
aliases_json: Final = json.dumps(aliases)
|
||
config_json: Final = json.dumps(config)
|
||
permissions_json: Final = json.dumps(permissions)
|
||
router_settings_json: Final = safe_dumps(router_settings) if router_settings is not None else safe_dumps({})
|
||
|
||
# Add model_rpm_limit and model_tpm_limit to metadata
|
||
if model_rpm_limit is not None:
|
||
metadata = metadata or {}
|
||
metadata["model_rpm_limit"] = model_rpm_limit
|
||
if model_tpm_limit is not None:
|
||
metadata = metadata or {}
|
||
metadata["model_tpm_limit"] = model_tpm_limit
|
||
if mcp_rpm_limit is not None:
|
||
metadata = metadata or {}
|
||
metadata["mcp_rpm_limit"] = mcp_rpm_limit
|
||
if tag_rpm_limit is not None:
|
||
metadata = metadata or {}
|
||
metadata["tag_rpm_limit"] = tag_rpm_limit
|
||
if guardrails is not None:
|
||
metadata = metadata or {}
|
||
metadata["guardrails"] = guardrails
|
||
if policies is not None:
|
||
metadata = metadata or {}
|
||
metadata["policies"] = policies
|
||
if prompts is not None:
|
||
metadata = metadata or {}
|
||
metadata["prompts"] = prompts
|
||
|
||
metadata = encrypt_callback_vars(metadata)
|
||
metadata_json: Final = json.dumps(metadata)
|
||
validate_model_max_budget(model_max_budget)
|
||
model_max_budget_json: Final = json.dumps(model_max_budget)
|
||
budget_fallbacks_json: Final = json.dumps(budget_fallbacks or {})
|
||
user_role = user_role
|
||
tpm_limit = tpm_limit
|
||
rpm_limit = rpm_limit
|
||
allowed_cache_controls = allowed_cache_controls
|
||
|
||
try:
|
||
# Create a new verification token (you may want to enhance this logic based on your needs)
|
||
|
||
user_data: Final = {
|
||
"max_budget": max_budget,
|
||
"user_email": user_email,
|
||
"user_id": user_id,
|
||
"user_alias": user_alias,
|
||
"team_id": team_id,
|
||
"organization_id": organization_id,
|
||
"user_role": user_role,
|
||
"spend": spend,
|
||
"models": models,
|
||
"metadata": metadata_json,
|
||
"max_parallel_requests": max_parallel_requests,
|
||
"tpm_limit": tpm_limit,
|
||
"rpm_limit": rpm_limit,
|
||
"budget_duration": budget_duration,
|
||
"budget_reset_at": reset_at,
|
||
"allowed_cache_controls": allowed_cache_controls,
|
||
"sso_user_id": sso_user_id,
|
||
"object_permission_id": object_permission_id,
|
||
}
|
||
if teams is not None:
|
||
user_data["teams"] = teams
|
||
if model_max_budget:
|
||
# Only when supplied: the SSO and default-key callers reach this with the
|
||
# empty default, and writing that would clear an existing user's budgets.
|
||
user_data["model_max_budget"] = model_max_budget_json
|
||
key_data: Final = {
|
||
"token": token,
|
||
"key_alias": key_alias,
|
||
"expires": expires,
|
||
"models": models,
|
||
"aliases": aliases_json,
|
||
"config": config_json,
|
||
"spend": spend,
|
||
"max_budget": key_max_budget,
|
||
"user_id": user_id,
|
||
"team_id": team_id,
|
||
"agent_id": agent_id,
|
||
"project_id": project_id,
|
||
"max_parallel_requests": max_parallel_requests,
|
||
"metadata": metadata_json,
|
||
"tpm_limit": tpm_limit,
|
||
"rpm_limit": rpm_limit,
|
||
"budget_duration": key_budget_duration,
|
||
"budget_reset_at": key_reset_at,
|
||
"allowed_cache_controls": allowed_cache_controls,
|
||
"permissions": permissions_json,
|
||
"model_max_budget": model_max_budget_json,
|
||
"budget_fallbacks": budget_fallbacks_json,
|
||
"organization_id": organization_id,
|
||
"budget_id": budget_id,
|
||
"blocked": blocked,
|
||
"budget_limits": budget_limits_json,
|
||
"created_by": created_by,
|
||
"updated_by": updated_by,
|
||
"allowed_routes": allowed_routes or [],
|
||
"key_type": key_type,
|
||
"object_permission_id": object_permission_id,
|
||
"router_settings": router_settings_json,
|
||
"access_group_ids": access_group_ids or [],
|
||
}
|
||
|
||
# Add rotation fields if auto_rotate is enabled
|
||
_set_key_rotation_fields(
|
||
data=key_data,
|
||
auto_rotate=auto_rotate or False,
|
||
rotation_interval=rotation_interval,
|
||
)
|
||
|
||
if (
|
||
get_secret("DISABLE_KEY_NAME", False) is True
|
||
): # allow user to disable storing abbreviated key name (shown in UI, to help figure out which key spent how much)
|
||
pass
|
||
else:
|
||
key_data["key_name"] = abbreviate_api_key(api_key=token)
|
||
saved_token: Final = copy.deepcopy(key_data)
|
||
if isinstance(saved_token["aliases"], str):
|
||
saved_token["aliases"] = json.loads(saved_token["aliases"])
|
||
if isinstance(saved_token["config"], str):
|
||
saved_token["config"] = json.loads(saved_token["config"])
|
||
if isinstance(saved_token["metadata"], str):
|
||
saved_token["metadata"] = json.loads(saved_token["metadata"])
|
||
if isinstance(saved_token["permissions"], str):
|
||
if "get_spend_routes" in saved_token["permissions"] and premium_user is not True:
|
||
raise ValueError("get_spend_routes permission is only available for LiteLLM Enterprise users")
|
||
|
||
saved_token["permissions"] = json.loads(saved_token["permissions"])
|
||
if isinstance(saved_token["model_max_budget"], str):
|
||
saved_token["model_max_budget"] = json.loads(saved_token["model_max_budget"])
|
||
router_settings = cast(dict | None, saved_token.get("router_settings"))
|
||
if router_settings is not None and isinstance(router_settings, str):
|
||
try:
|
||
saved_token["router_settings"] = yaml.safe_load(router_settings)
|
||
except yaml.YAMLError:
|
||
# If it's not valid JSON/YAML, keep as is or set to empty dict
|
||
saved_token["router_settings"] = {}
|
||
|
||
if saved_token.get("expires", None) is not None and isinstance(saved_token["expires"], datetime):
|
||
saved_token["expires"] = saved_token["expires"].isoformat()
|
||
if prisma_client is not None:
|
||
if table_name is None or table_name == "user": # do not auto-create users for `/key/generate`
|
||
## CREATE USER (If necessary)
|
||
if query_type == "insert_data":
|
||
user_row = cast( # cast-ok: table_name="user" is the insert_data branch returning the user row
|
||
"prisma_models.LiteLLM_UserTable | None",
|
||
await prisma_client.insert_data(data=user_data, table_name="user"),
|
||
)
|
||
|
||
if user_row is None:
|
||
raise Exception("Failed to create user")
|
||
## use default user model list if no key-specific model list provided
|
||
if len(user_row.models) > 0 and len(key_data["models"]) == 0:
|
||
key_data["models"] = user_row.models
|
||
elif query_type == "update_data":
|
||
user_row = await prisma_client.update_data(
|
||
data=user_data,
|
||
table_name="user",
|
||
update_key_values=update_key_values,
|
||
)
|
||
if table_name is not None and table_name == "user":
|
||
# do not create a key if table name is set to just 'user'
|
||
# we only need to ensure this exists in the user table
|
||
# the LiteLLM_VerificationToken table will increase in size if we don't do this check
|
||
return user_data
|
||
|
||
## CREATE KEY
|
||
verbose_proxy_logger.debug(
|
||
"prisma_client: Creating Key= %s",
|
||
{**key_data, "token": hash_token(token=token)},
|
||
)
|
||
create_key_response: Final = await prisma_client.insert_data(data=key_data, table_name="key")
|
||
|
||
key_data["token_id"] = getattr(create_key_response, "token", None)
|
||
created_token_hash: Final = getattr(create_key_response, "token", None)
|
||
if isinstance(created_token_hash, str):
|
||
await sync_key_access_group_membership(
|
||
prisma_client=prisma_client,
|
||
key_token=created_token_hash,
|
||
previous_access_group_ids=None,
|
||
updated_access_group_ids=access_group_ids,
|
||
)
|
||
key_data["litellm_budget_table"] = getattr(create_key_response, "litellm_budget_table", None)
|
||
key_data["created_at"] = getattr(create_key_response, "created_at", None)
|
||
key_data["updated_at"] = getattr(create_key_response, "updated_at", None)
|
||
|
||
# Deserialize router_settings from JSON string to dict for response
|
||
router_settings_value: Final = key_data.get("router_settings")
|
||
if router_settings_value is not None and isinstance(router_settings_value, str):
|
||
try:
|
||
key_data["router_settings"] = yaml.safe_load(router_settings_value)
|
||
except yaml.YAMLError:
|
||
# If it's not valid JSON/YAML, keep as is or set to empty dict
|
||
key_data["router_settings"] = {}
|
||
except Exception as e:
|
||
verbose_proxy_logger.error("litellm.proxy.proxy_server.generate_key_helper_fn(): Exception occured - %s", e)
|
||
verbose_proxy_logger.debug(traceback.format_exc())
|
||
if isinstance(e, HTTPException):
|
||
raise e
|
||
raise HTTPException(
|
||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||
detail={"error": "Internal Server Error."},
|
||
)
|
||
|
||
# Add budget related info in key_data - this ensures it's returned
|
||
key_data["budget_id"] = budget_id
|
||
|
||
if request_type == "user":
|
||
# if this is a /user/new request update the key_date with user_data fields
|
||
key_data.update(user_data)
|
||
|
||
return key_data
|
||
|
||
|
||
async def _team_key_deletion_check(
|
||
user_api_key_dict: UserAPIKeyAuth,
|
||
key_info: LiteLLM_VerificationToken,
|
||
prisma_client: PrismaClient,
|
||
user_api_key_cache: UserApiKeyCache,
|
||
):
|
||
is_team_key: Final = _is_team_key(data=key_info)
|
||
|
||
if is_team_key and key_info.team_id is not None:
|
||
team_table: Final = await get_team_object(
|
||
team_id=key_info.team_id,
|
||
prisma_client=prisma_client,
|
||
user_api_key_cache=user_api_key_cache,
|
||
check_db_only=True,
|
||
)
|
||
if litellm.key_generation_settings is not None and "team_key_generation" in litellm.key_generation_settings:
|
||
_team_key_generation = litellm.key_generation_settings["team_key_generation"]
|
||
else:
|
||
_team_key_generation = TeamUIKeyGenerationConfig(
|
||
allowed_team_member_roles=["admin", "user"],
|
||
)
|
||
# check if user is team admin
|
||
if team_table is not None:
|
||
return _team_key_operation_team_member_check(
|
||
assigned_user_id=user_api_key_dict.user_id,
|
||
team_table=team_table,
|
||
user_api_key_dict=user_api_key_dict,
|
||
team_key_generation=_team_key_generation,
|
||
route=KeyManagementRoutes.KEY_DELETE,
|
||
)
|
||
else:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_404_NOT_FOUND,
|
||
detail={"error": f"Team not found in db, and user not proxy admin. Team id = {key_info.team_id}"},
|
||
)
|
||
return False
|
||
|
||
|
||
async def can_modify_verification_token(
|
||
key_info: LiteLLM_VerificationToken,
|
||
user_api_key_cache: UserApiKeyCache,
|
||
user_api_key_dict: UserAPIKeyAuth,
|
||
prisma_client: PrismaClient,
|
||
) -> bool:
|
||
"""
|
||
Check if user has permission to modify (delete/regenerate) a verification token.
|
||
|
||
Rules:
|
||
- Proxy admin can modify any key
|
||
- Internal jobs service account can modify any key (for auto-rotation)
|
||
- For team keys: only team admin or key owner can modify
|
||
- For personal keys: only key owner can modify
|
||
|
||
Args:
|
||
key_info: The verification token to check
|
||
user_api_key_cache: Cache for user API keys
|
||
user_api_key_dict: The user making the request
|
||
prisma_client: Prisma client for database access
|
||
|
||
Returns:
|
||
True if user can modify the key, False otherwise
|
||
"""
|
||
from litellm.constants import LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME
|
||
|
||
is_team_key: Final = _is_team_key(data=key_info)
|
||
|
||
# 1. Proxy admin can modify any key
|
||
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
|
||
return True
|
||
|
||
# 2. Internal jobs service account can modify any key (for auto-rotation)
|
||
if user_api_key_dict.api_key == LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME:
|
||
return True
|
||
|
||
# 3. For team keys: only team admin or key owner can modify
|
||
if is_team_key and key_info.team_id is not None:
|
||
# Get team object to check if user is team admin
|
||
team_table: Final = await get_team_object(
|
||
team_id=key_info.team_id,
|
||
prisma_client=prisma_client,
|
||
user_api_key_cache=user_api_key_cache,
|
||
check_db_only=True,
|
||
)
|
||
|
||
if team_table is None:
|
||
return False
|
||
|
||
# Check if user is team admin
|
||
if _is_user_team_admin(
|
||
user_api_key_dict=user_api_key_dict,
|
||
team_obj=team_table,
|
||
):
|
||
return True
|
||
|
||
# Check if the key belongs to the user (they own it)
|
||
if key_info.user_id is not None and key_info.user_id == user_api_key_dict.user_id:
|
||
return True
|
||
|
||
# Not team admin and doesn't own the key
|
||
return False
|
||
|
||
# 4. For personal keys: only key owner can modify
|
||
if key_info.user_id is not None and key_info.user_id == user_api_key_dict.user_id:
|
||
return True
|
||
|
||
# Default: deny
|
||
return False
|
||
|
||
|
||
async def delete_verification_tokens(
|
||
tokens: list,
|
||
user_api_key_cache: UserApiKeyCache,
|
||
user_api_key_dict: UserAPIKeyAuth,
|
||
litellm_changed_by: str | None = None,
|
||
) -> tuple[dict | None, list[LiteLLM_VerificationToken]]:
|
||
"""
|
||
Helper that deletes the list of tokens from the database
|
||
|
||
- check if user is proxy admin
|
||
- check if user is team admin and key is a team key
|
||
|
||
Args:
|
||
tokens: List of tokens to delete
|
||
user_id: Optional user_id to filter by
|
||
|
||
Returns:
|
||
Tuple[Optional[Dict], List[LiteLLM_VerificationToken]]:
|
||
Optional[Dict]:
|
||
- Number of deleted tokens
|
||
List[LiteLLM_VerificationToken]:
|
||
- List of keys being deleted, this contains information about the key_alias, token, and user_id being deleted,
|
||
this is passed down to the KeyManagementEventHooks to delete the keys from the secret manager and handle audit logs
|
||
"""
|
||
from litellm.proxy.proxy_server import prisma_client
|
||
|
||
failed_tokens: list = []
|
||
try:
|
||
if prisma_client:
|
||
hashed_tokens: Final[list[str]] = [_hash_token_if_needed(token=key) for key in tokens]
|
||
tokens = hashed_tokens
|
||
_keys_being_deleted: Final[list[LiteLLM_VerificationToken]] = cast( # cast-ok: find_many returns a list
|
||
"list[LiteLLM_VerificationToken]",
|
||
await _prisma_table(VerificationTokenRepository(prisma_client)).find_many(
|
||
where={"token": {"in": hashed_tokens}}
|
||
),
|
||
)
|
||
|
||
if len(_keys_being_deleted) == 0:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_404_NOT_FOUND,
|
||
detail={"error": "No keys found"},
|
||
)
|
||
|
||
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
|
||
authorized_keys = _keys_being_deleted
|
||
else:
|
||
authorized_keys = []
|
||
for key in _keys_being_deleted:
|
||
if await can_modify_verification_token(
|
||
key_info=key,
|
||
user_api_key_cache=user_api_key_cache,
|
||
user_api_key_dict=user_api_key_dict,
|
||
prisma_client=prisma_client,
|
||
):
|
||
authorized_keys.append(key)
|
||
else:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_403_FORBIDDEN,
|
||
detail={"error": "You are not authorized to delete this key"},
|
||
)
|
||
await _persist_deleted_verification_tokens(
|
||
keys=authorized_keys,
|
||
prisma_client=prisma_client,
|
||
user_api_key_dict=user_api_key_dict,
|
||
litellm_changed_by=litellm_changed_by,
|
||
)
|
||
|
||
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
|
||
deleted_tokens = await prisma_client.delete_data(tokens=tokens)
|
||
if deleted_tokens is not None and len(deleted_tokens) != len(tokens):
|
||
failed_tokens = [token for token in tokens if token not in deleted_tokens]
|
||
else:
|
||
deletion_tasks: Final = [prisma_client.delete_data(tokens=[key.token]) for key in authorized_keys]
|
||
await asyncio.gather(*deletion_tasks)
|
||
|
||
deleted_tokens = [key.token for key in authorized_keys]
|
||
if len(deleted_tokens) != len(tokens):
|
||
failed_tokens = [token for token in tokens if token not in deleted_tokens]
|
||
|
||
else:
|
||
raise Exception("DB not connected. prisma_client is None")
|
||
except Exception as e:
|
||
verbose_proxy_logger.exception(
|
||
"litellm.proxy.proxy_server.delete_verification_tokens(): Exception occured - %s", e
|
||
)
|
||
verbose_proxy_logger.debug(traceback.format_exc())
|
||
raise e
|
||
|
||
for key in tokens:
|
||
user_api_key_cache.delete_cache(key)
|
||
# remove hash token from cache
|
||
hashed_token = hash_token(cast(str, key))
|
||
user_api_key_cache.delete_cache(hashed_token)
|
||
|
||
# After credential invalidation, so a failure here can never keep a deleted key alive.
|
||
for deleted_key in authorized_keys:
|
||
if deleted_key.token is not None:
|
||
await sync_key_access_group_membership(
|
||
prisma_client=prisma_client,
|
||
key_token=deleted_key.token,
|
||
previous_access_group_ids=deleted_key.access_group_ids,
|
||
updated_access_group_ids=None,
|
||
)
|
||
|
||
return {
|
||
"deleted_keys": deleted_tokens,
|
||
"failed_tokens": failed_tokens,
|
||
}, _keys_being_deleted
|
||
|
||
|
||
def _transform_verification_tokens_to_deleted_records(
|
||
keys: Sequence[LiteLLM_VerificationToken],
|
||
user_api_key_dict: UserAPIKeyAuth,
|
||
litellm_changed_by: str | None = None,
|
||
) -> list[dict[str, object]]:
|
||
"""Transform verification tokens into deleted token records ready for persistence."""
|
||
if not keys:
|
||
return []
|
||
|
||
deleted_at: Final = datetime.now(timezone.utc)
|
||
records: Final = []
|
||
for key in keys:
|
||
key_payload = key.model_dump()
|
||
deleted_record = LiteLLM_DeletedVerificationToken.model_validate(
|
||
{
|
||
**key_payload,
|
||
"deleted_at": deleted_at,
|
||
"deleted_by": user_api_key_dict.user_id,
|
||
"deleted_by_api_key": user_api_key_dict.api_key,
|
||
"litellm_changed_by": litellm_changed_by,
|
||
}
|
||
)
|
||
record = deleted_record.model_dump()
|
||
|
||
# Map org_id to organization_id (model uses org_id, but schema expects organization_id)
|
||
org_id_value: object = record.pop("org_id", None)
|
||
if org_id_value is not None:
|
||
record["organization_id"] = org_id_value
|
||
|
||
for json_field in [
|
||
"aliases",
|
||
"config",
|
||
"permissions",
|
||
"metadata",
|
||
"model_spend",
|
||
"model_max_budget",
|
||
"budget_fallbacks",
|
||
"router_settings",
|
||
]:
|
||
if json_field in record and record[json_field] is not None:
|
||
record[json_field] = json.dumps(record[json_field])
|
||
|
||
for rel_key in (
|
||
"litellm_budget_table",
|
||
"litellm_organization_table",
|
||
"object_permission",
|
||
"id",
|
||
"budget_limits",
|
||
):
|
||
record.pop(rel_key, None)
|
||
|
||
records.append(record)
|
||
|
||
return records
|
||
|
||
|
||
async def _save_deleted_verification_token_records(
|
||
records: Sequence[Mapping[str, object]],
|
||
prisma_client: PrismaClient,
|
||
tx: "Prisma | None" = None,
|
||
) -> None:
|
||
"""Save deleted verification token records to the database.
|
||
|
||
``tx`` runs the write on that transaction's connection instead of a fresh
|
||
one, so a caller batching this with other writes gets one all-or-nothing
|
||
commit.
|
||
"""
|
||
if not records:
|
||
return
|
||
if tx is not None:
|
||
await tx.litellm_deletedverificationtoken.create_many(data=records)
|
||
return
|
||
await _deleted_verification_token_table(prisma_client).create_many(data=records)
|
||
|
||
|
||
async def _persist_deleted_verification_tokens(
|
||
keys: Sequence[LiteLLM_VerificationToken],
|
||
prisma_client: PrismaClient,
|
||
user_api_key_dict: UserAPIKeyAuth,
|
||
litellm_changed_by: str | None = None,
|
||
tx: "Prisma | None" = None,
|
||
) -> None:
|
||
"""Persist deleted verification token records by transforming and saving them."""
|
||
records: Final = _transform_verification_tokens_to_deleted_records(
|
||
keys=keys,
|
||
user_api_key_dict=user_api_key_dict,
|
||
litellm_changed_by=litellm_changed_by,
|
||
)
|
||
await _save_deleted_verification_token_records(
|
||
records=records,
|
||
prisma_client=prisma_client,
|
||
tx=tx,
|
||
)
|
||
|
||
|
||
async def delete_key_aliases(
|
||
key_aliases: list[str],
|
||
user_api_key_cache: UserApiKeyCache,
|
||
prisma_client: PrismaClient,
|
||
user_api_key_dict: UserAPIKeyAuth,
|
||
litellm_changed_by: str | None = None,
|
||
) -> tuple[dict | None, list[LiteLLM_VerificationToken]]:
|
||
_keys_being_deleted: Final = await _prisma_table(VerificationTokenRepository(prisma_client)).find_many(
|
||
where={"key_alias": {"in": key_aliases}}
|
||
)
|
||
|
||
tokens: Final = [key.token for key in _keys_being_deleted]
|
||
return await delete_verification_tokens(
|
||
tokens=tokens,
|
||
user_api_key_cache=user_api_key_cache,
|
||
user_api_key_dict=user_api_key_dict,
|
||
litellm_changed_by=litellm_changed_by,
|
||
)
|
||
|
||
|
||
async def _rotate_master_key(
|
||
prisma_client: PrismaClient,
|
||
user_api_key_dict: UserAPIKeyAuth,
|
||
current_master_key: str,
|
||
new_master_key: str,
|
||
) -> None:
|
||
"""
|
||
Rotate the master key
|
||
|
||
1. Get the values from the DB
|
||
- Get models from DB
|
||
- Get config from DB
|
||
2. Decrypt the values
|
||
- ModelTable
|
||
- [{"model_name": "str", "litellm_params": {}}]
|
||
- ConfigTable
|
||
3. Encrypt the values with the new master key
|
||
4. Update the values in the DB
|
||
"""
|
||
import prisma
|
||
|
||
from litellm.proxy.proxy_server import proxy_config
|
||
|
||
try:
|
||
models: list | None = cast( # cast-ok: find_many returns a real list, which TableActions widens to Sequence
|
||
"list[object]", await _prisma_table(ModelRepository(prisma_client)).find_many()
|
||
)
|
||
except Exception:
|
||
models = None
|
||
# 2. process model table
|
||
if models:
|
||
decrypted_models: Final = proxy_config.decrypt_model_list_from_db(new_models=models)
|
||
verbose_proxy_logger.debug("ABLE TO DECRYPT MODELS - len(decrypted_models): %s", len(decrypted_models))
|
||
new_models: Final[list[dict[str, object]]] = []
|
||
for model in decrypted_models:
|
||
new_model = await _add_model_to_db(
|
||
model_params=Deployment(**model),
|
||
user_api_key_dict=user_api_key_dict,
|
||
prisma_client=prisma_client,
|
||
new_encryption_key=new_master_key,
|
||
should_create_model_in_db=False,
|
||
)
|
||
if new_model:
|
||
_dumped = new_model.model_dump(exclude_none=True)
|
||
_dumped["litellm_params"] = prisma.Json(_dumped["litellm_params"])
|
||
_dumped["model_info"] = prisma.Json(_dumped["model_info"])
|
||
new_models.append(_dumped)
|
||
verbose_proxy_logger.debug("Resetting proxy model table")
|
||
async with prisma_client.db.tx() as tx_ctx:
|
||
tx: Final[_TxTables] = tx_ctx
|
||
await tx.litellm_proxymodeltable.delete_many()
|
||
verbose_proxy_logger.debug("Creating %s models", len(new_models))
|
||
await tx.litellm_proxymodeltable.create_many(
|
||
data=new_models,
|
||
)
|
||
await publish_config_change(redis_cache=coordination_redis_cache(), object_type="litellm_proxymodeltable")
|
||
# 3. process config table
|
||
try:
|
||
config = await _config_table(prisma_client).find_many()
|
||
except Exception:
|
||
config = None
|
||
|
||
if config:
|
||
"""If environment_variables is found, decrypt it and encrypt it with the new master key"""
|
||
environment_variables_dict = {}
|
||
for c in config:
|
||
if c.param_name == "environment_variables":
|
||
environment_variables_dict = c.param_value
|
||
|
||
if environment_variables_dict:
|
||
decrypted_env_vars: Final = proxy_config._decrypt_and_set_db_env_variables(
|
||
environment_variables=environment_variables_dict
|
||
)
|
||
encrypted_env_vars: Final = proxy_config._encrypt_env_variables(
|
||
environment_variables=decrypted_env_vars,
|
||
new_encryption_key=new_master_key,
|
||
)
|
||
|
||
if encrypted_env_vars:
|
||
await _config_table(prisma_client).update(
|
||
where={"param_name": "environment_variables"},
|
||
data={"param_value": prisma.Json(encrypted_env_vars)},
|
||
)
|
||
|
||
# 4. process MCP server table
|
||
try:
|
||
await rotate_mcp_server_credentials_master_key(
|
||
prisma_client=prisma_client,
|
||
touched_by=user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME,
|
||
new_master_key=new_master_key,
|
||
)
|
||
except Exception as e:
|
||
verbose_proxy_logger.warning("Failed to rotate MCP server credentials: %s", str(e))
|
||
|
||
# 4b. process MCP user-scoped credentials table (BYOK + OAuth2 tokens)
|
||
try:
|
||
await rotate_mcp_user_credentials_master_key(
|
||
prisma_client=prisma_client,
|
||
new_master_key=new_master_key,
|
||
)
|
||
except Exception as e:
|
||
verbose_proxy_logger.warning("Failed to rotate MCP user credentials: %s", str(e))
|
||
|
||
# 4c. process MCP per-user environment variables table
|
||
try:
|
||
await rotate_mcp_user_env_vars_master_key(
|
||
prisma_client=prisma_client,
|
||
new_master_key=new_master_key,
|
||
)
|
||
except Exception as e:
|
||
verbose_proxy_logger.warning("Failed to rotate MCP user env vars: %s", str(e))
|
||
|
||
# 4d. process SSO identity assertion table (EMA subject tokens)
|
||
try:
|
||
await rotate_sso_identity_assertions_master_key(
|
||
prisma_client=prisma_client,
|
||
new_master_key=new_master_key,
|
||
)
|
||
except Exception as e: # noqa: BLE001 # one store's failure must not abort the master-key rotation
|
||
verbose_proxy_logger.warning("Failed to rotate SSO identity assertions: %s", str(e))
|
||
|
||
# 5. process credentials table
|
||
try:
|
||
credentials = await _credentials_table(prisma_client).find_many()
|
||
except Exception:
|
||
credentials = None
|
||
if credentials:
|
||
from litellm.proxy.credential_endpoints.endpoints import update_db_credential
|
||
|
||
for cred in credentials:
|
||
try:
|
||
decrypted_cred = proxy_config.decrypt_credentials(cred)
|
||
encrypted_cred = update_db_credential(
|
||
db_credential=cred,
|
||
updated_patch=decrypted_cred,
|
||
new_encryption_key=new_master_key,
|
||
)
|
||
_cred_data = encrypted_cred.model_dump(exclude_none=True)
|
||
if "credential_values" in _cred_data:
|
||
_cred_data["credential_values"] = prisma.Json(_cred_data["credential_values"])
|
||
if "credential_info" in _cred_data:
|
||
_cred_data["credential_info"] = prisma.Json(_cred_data["credential_info"])
|
||
await _credentials_table(prisma_client).update(
|
||
where={"credential_name": cred.credential_name},
|
||
data={
|
||
**_cred_data,
|
||
"updated_by": user_api_key_dict.user_id,
|
||
},
|
||
)
|
||
except Exception as e:
|
||
verbose_proxy_logger.error("Failed to re-encrypt credential %s: %s", cred.credential_name, e)
|
||
# Continue with next credential instead of failing entire rotation
|
||
continue
|
||
verbose_proxy_logger.debug("Successfully re-encrypted %s credentials with new master key", len(credentials))
|
||
|
||
|
||
def _require_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> None:
|
||
from litellm.proxy._types import CommonProxyErrors
|
||
|
||
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value:
|
||
raise HTTPException(
|
||
status_code=403,
|
||
detail={"error": CommonProxyErrors.not_allowed_access.value},
|
||
)
|
||
|
||
|
||
@router.post(
|
||
"/credentials/migrate-encryption",
|
||
tags=["credential management"],
|
||
dependencies=[Depends(user_api_key_auth)],
|
||
)
|
||
async def migrate_encryption_endpoint(
|
||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||
dry_run: bool = Query(
|
||
False,
|
||
description="If true, scan and report without writing any changes.",
|
||
),
|
||
):
|
||
"""
|
||
Re-encrypt all at-rest credentials into the AES-256-GCM (``v2:gcm:``) format.
|
||
|
||
Admin only. Requires ``general_settings.encryption_algorithm: aes-256-gcm``.
|
||
Idempotent and resumable — re-running skips already-migrated values. Pass
|
||
``dry_run=true`` for a non-mutating scan (equivalent to ``--check``).
|
||
"""
|
||
from litellm.proxy._types import CommonProxyErrors
|
||
from litellm.proxy.management_endpoints.credential_migration import (
|
||
migrate_encryption,
|
||
)
|
||
from litellm.proxy.proxy_server import prisma_client
|
||
|
||
_require_proxy_admin(user_api_key_dict)
|
||
if prisma_client is None:
|
||
raise HTTPException(
|
||
status_code=500,
|
||
detail={"error": CommonProxyErrors.db_not_connected_error.value},
|
||
)
|
||
|
||
report: Final = await migrate_encryption(
|
||
prisma_client=prisma_client,
|
||
user_api_key_dict=user_api_key_dict,
|
||
dry_run=dry_run,
|
||
)
|
||
return {"status": "success", "dry_run": dry_run, "report": report.as_dict()}
|
||
|
||
|
||
@router.get(
|
||
"/credentials/migrate-encryption/check",
|
||
tags=["credential management"],
|
||
dependencies=[Depends(user_api_key_auth)],
|
||
)
|
||
async def check_encryption_endpoint(
|
||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||
):
|
||
"""
|
||
Read-only residual scan for compliance attestation. Reports how many at-rest
|
||
values are still in the legacy format. ``residual_legacy == 0`` attests no
|
||
legacy ciphertext remains. Admin only; performs no writes.
|
||
"""
|
||
from litellm.proxy._types import CommonProxyErrors
|
||
from litellm.proxy.management_endpoints.credential_migration import (
|
||
check_encryption,
|
||
)
|
||
from litellm.proxy.proxy_server import prisma_client
|
||
|
||
_require_proxy_admin(user_api_key_dict)
|
||
if prisma_client is None:
|
||
raise HTTPException(
|
||
status_code=500,
|
||
detail={"error": CommonProxyErrors.db_not_connected_error.value},
|
||
)
|
||
|
||
report: Final = await check_encryption(prisma_client=prisma_client)
|
||
return {"status": "success", "report": report.as_dict()}
|
||
|
||
|
||
async def get_new_token(data: RegenerateKeyRequest | None) -> str:
|
||
if data and data.new_key is not None:
|
||
# Reject custom key values if disabled by admin
|
||
await _check_custom_key_allowed(data.new_key)
|
||
if not data.new_key.startswith("sk-"):
|
||
raise HTTPException(
|
||
status_code=status.HTTP_400_BAD_REQUEST,
|
||
detail={
|
||
"error": "New key must start with 'sk-'. This is to distinguish a key hash (used by litellm for logging / internal logic) from the actual key."
|
||
},
|
||
)
|
||
if len(data.new_key) < MINIMUM_CUSTOM_KEY_LENGTH:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_400_BAD_REQUEST,
|
||
detail={"error": f"New key must be at least {MINIMUM_CUSTOM_KEY_LENGTH} characters long."},
|
||
)
|
||
new_token = data.new_key
|
||
else:
|
||
new_token = f"sk-{secrets.token_urlsafe(LENGTH_OF_LITELLM_GENERATED_KEY)}"
|
||
return new_token
|
||
|
||
|
||
async def _insert_deprecated_key(
|
||
prisma_client: "PrismaClient",
|
||
old_token_hash: str,
|
||
new_token_hash: str,
|
||
grace_period: str | None,
|
||
) -> None:
|
||
"""
|
||
Insert old key into deprecated table so it remains valid during grace period.
|
||
|
||
Uses upsert to handle concurrent rotations gracefully.
|
||
|
||
Parameters:
|
||
prisma_client: DB client
|
||
old_token_hash: Hash of the old key being rotated out
|
||
new_token_hash: Hash of the new replacement key
|
||
grace_period: Duration string (e.g. "24h", "2d") or None/empty for immediate revoke
|
||
"""
|
||
grace_period_value: Final = grace_period or os.getenv("LITELLM_KEY_ROTATION_GRACE_PERIOD", "")
|
||
if not grace_period_value:
|
||
return
|
||
|
||
try:
|
||
grace_seconds: Final = duration_in_seconds(grace_period_value)
|
||
except ValueError:
|
||
verbose_proxy_logger.warning(
|
||
"Invalid grace_period format: %s. Expected format like '24h', '2d'.",
|
||
grace_period_value,
|
||
)
|
||
return
|
||
|
||
if grace_seconds <= 0:
|
||
return
|
||
|
||
try:
|
||
revoke_at: Final = datetime.now(timezone.utc) + timedelta(seconds=grace_seconds)
|
||
await _deprecated_verification_token_table(prisma_client).upsert(
|
||
where={"token": old_token_hash},
|
||
data={
|
||
"create": {
|
||
"token": old_token_hash,
|
||
"active_token_id": new_token_hash,
|
||
"revoke_at": revoke_at,
|
||
},
|
||
"update": {
|
||
"active_token_id": new_token_hash,
|
||
"revoke_at": revoke_at,
|
||
},
|
||
},
|
||
)
|
||
verbose_proxy_logger.debug(
|
||
"Deprecated key retained for %s (revoke_at: %s)",
|
||
grace_period_value,
|
||
revoke_at,
|
||
)
|
||
except Exception as deprecated_err:
|
||
verbose_proxy_logger.warning(
|
||
"Failed to insert deprecated key for grace period: %s",
|
||
deprecated_err,
|
||
)
|
||
|
||
|
||
async def _execute_virtual_key_regeneration(
|
||
*,
|
||
prisma_client: PrismaClient,
|
||
key_in_db: LiteLLM_VerificationToken,
|
||
hashed_api_key: str,
|
||
key: str,
|
||
data: RegenerateKeyRequest | None,
|
||
user_api_key_dict: UserAPIKeyAuth,
|
||
litellm_changed_by: str | None,
|
||
user_api_key_cache: UserApiKeyCache,
|
||
proxy_logging_obj: ProxyLogging,
|
||
) -> GenerateKeyResponse:
|
||
"""Generate new token, update DB, invalidate cache, and return response."""
|
||
from litellm.proxy.proxy_server import hash_token
|
||
|
||
# Mirror the /key/update ownership rebind guard. See helper docstring.
|
||
_validate_caller_can_change_key_ownership(
|
||
data=data,
|
||
existing_key_row=key_in_db,
|
||
user_api_key_dict=user_api_key_dict,
|
||
)
|
||
|
||
# Apply the same membership rule used on /key/update: when the caller
|
||
# asks to point the regenerated key at a different organization_id,
|
||
# require they are a member of (or proxy admin over) the target org.
|
||
if data is not None and data.organization_id is not None:
|
||
_existing_org_id: Final = getattr(key_in_db, "organization_id", None)
|
||
_is_proxy_admin: Final = (
|
||
user_api_key_dict.user_role is not None
|
||
and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
|
||
)
|
||
if data.organization_id != _existing_org_id and not _is_proxy_admin:
|
||
await _validate_caller_can_assign_key_org(
|
||
user_api_key_dict=user_api_key_dict,
|
||
organization_id=data.organization_id,
|
||
prisma_client=prisma_client,
|
||
)
|
||
|
||
if data is not None:
|
||
_existing_key_metadata: Final = getattr(key_in_db, "metadata", None)
|
||
enforce_output_token_estimates_are_admin_only(
|
||
data=data,
|
||
existing_metadata=_existing_key_metadata if isinstance(_existing_key_metadata, dict) else None,
|
||
user_api_key_dict=user_api_key_dict,
|
||
entity="key",
|
||
)
|
||
enforce_batch_enqueued_token_limit_is_admin_only(
|
||
data=data,
|
||
existing_metadata=_existing_key_metadata if isinstance(_existing_key_metadata, dict) else None,
|
||
user_api_key_dict=user_api_key_dict,
|
||
entity="key",
|
||
)
|
||
|
||
new_token: Final = await get_new_token(data=data)
|
||
new_token_hash: Final = hash_token(new_token)
|
||
new_token_key_name: Final = abbreviate_api_key(api_key=new_token)
|
||
update_data = {"token": new_token_hash, "key_name": new_token_key_name}
|
||
|
||
non_default_values = {}
|
||
if data is not None:
|
||
# Enforce upperbound key params on regenerate (don't fill defaults)
|
||
_enforce_upperbound_key_params(data, fill_defaults=False)
|
||
non_default_values = await prepare_key_update_data(data=data, existing_key_row=key_in_db)
|
||
# Only validate key_alias format if it's actually being changed
|
||
new_key_alias: Final = non_default_values.get("key_alias")
|
||
if new_key_alias != key_in_db.key_alias:
|
||
_validate_key_alias_format(key_alias=new_key_alias)
|
||
verbose_proxy_logger.debug("non_default_values: %s", non_default_values)
|
||
update_data.update(non_default_values)
|
||
jsonified_update_data: Final[Mapping[str, object]] = prisma_client.jsonify_object(data=update_data)
|
||
|
||
# If grace period set, insert deprecated key so old key remains valid
|
||
await _insert_deprecated_key(
|
||
prisma_client=prisma_client,
|
||
old_token_hash=hashed_api_key,
|
||
new_token_hash=new_token_hash,
|
||
grace_period=data.grace_period if data else None,
|
||
)
|
||
|
||
updated_token: Final[LiteLLM_VerificationToken | None] = await _prisma_table(
|
||
VerificationTokenRepository(prisma_client)
|
||
).update(
|
||
where={"token": hashed_api_key},
|
||
data=with_settings_updated_at(jsonified_update_data),
|
||
)
|
||
updated_token_dict: Final[dict[str, object]] = dict(updated_token) if updated_token is not None else {}
|
||
updated_token_dict["key"] = new_token
|
||
updated_token_dict["token_id"] = updated_token_dict.pop("token")
|
||
|
||
if hashed_api_key or key:
|
||
await _delete_cache_key_object(
|
||
hashed_token=_hash_token_if_needed(key),
|
||
user_api_key_cache=user_api_key_cache,
|
||
proxy_logging_obj=proxy_logging_obj,
|
||
)
|
||
|
||
# After credential invalidation, so a failure here can never keep the old key alive.
|
||
await sync_key_regeneration_access_group_membership(
|
||
prisma_client=prisma_client,
|
||
previous_key_token=hashed_api_key,
|
||
new_key_token=new_token_hash,
|
||
data=data,
|
||
existing_key_row=key_in_db,
|
||
)
|
||
|
||
response: Final = GenerateKeyResponse.model_validate(updated_token_dict)
|
||
asyncio.create_task(
|
||
KeyManagementEventHooks.async_key_rotated_hook(
|
||
data=data,
|
||
existing_key_row=key_in_db,
|
||
response=response,
|
||
user_api_key_dict=user_api_key_dict,
|
||
litellm_changed_by=litellm_changed_by,
|
||
)
|
||
)
|
||
return response
|
||
|
||
|
||
@router.post(
|
||
"/key/{key:path}/regenerate",
|
||
tags=["key management"],
|
||
dependencies=[Depends(user_api_key_auth)],
|
||
)
|
||
@router.post(
|
||
"/key/regenerate",
|
||
tags=["key management"],
|
||
dependencies=[Depends(user_api_key_auth)],
|
||
)
|
||
@management_endpoint_wrapper
|
||
async def regenerate_key_fn(
|
||
key: str | None = None,
|
||
data: RegenerateKeyRequest | None = None,
|
||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||
litellm_changed_by: str | None = Header(
|
||
None,
|
||
description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability",
|
||
),
|
||
) -> GenerateKeyResponse | None:
|
||
"""
|
||
Regenerate an existing API key while optionally updating its parameters.
|
||
|
||
Parameters:
|
||
- key: str (path parameter) - The key to regenerate
|
||
- data: Optional[RegenerateKeyRequest] - Request body containing optional parameters to update
|
||
- key: Optional[str] - The key to regenerate.
|
||
- new_master_key: Optional[str] - The new master key to use, if key is the master key.
|
||
- new_key: Optional[str] - The new key to use, if key is not the master key. Must start with 'sk-' and be at least 16 characters long. If both set, new_master_key will be used.
|
||
- key_alias: Optional[str] - User-friendly key alias
|
||
- user_id: Optional[str] - User ID associated with key
|
||
- team_id: Optional[str] - Team ID associated with key
|
||
- models: Optional[list] - Model_name's a user is allowed to call
|
||
- tags: Optional[List[str]] - Tags for organizing keys (Enterprise only)
|
||
- spend: Optional[float] - Amount spent by key
|
||
- max_budget: Optional[float] - Max budget for key
|
||
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}
|
||
- budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
|
||
- budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.)
|
||
- soft_budget: Optional[float] - Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached.
|
||
- max_parallel_requests: Optional[int] - Rate limit for parallel requests
|
||
- metadata: Optional[dict] - Metadata for key. Example {"team": "core-infra", "app": "app2"}
|
||
- tpm_limit: Optional[int] - Tokens per minute limit
|
||
- rpm_limit: Optional[int] - Requests per minute limit
|
||
- model_rpm_limit: Optional[dict] - Model-specific RPM limits {"gpt-4": 100, "claude-v1": 200}
|
||
- model_tpm_limit: Optional[dict] - Model-specific TPM limits {"gpt-4": 100000, "claude-v1": 200000}
|
||
- allowed_cache_controls: Optional[list] - List of allowed cache control values
|
||
- duration: Optional[str] - Key validity duration ("30d", "1h", etc.)
|
||
- permissions: Optional[dict] - Key-specific permissions
|
||
- guardrails: Optional[List[str]] - List of active guardrails for the key
|
||
- blocked: Optional[bool] - Whether the key is blocked
|
||
- grace_period: Optional[str] - Duration to keep old key valid after rotation (e.g. "24h", "2d"). Omitted = immediate revoke. Env: LITELLM_KEY_ROTATION_GRACE_PERIOD
|
||
|
||
|
||
Returns:
|
||
- GenerateKeyResponse containing the new key and its updated parameters
|
||
|
||
Example:
|
||
```bash
|
||
curl --location --request POST 'http://localhost:4000/key/sk-1234/regenerate' \
|
||
--header 'Authorization: Bearer sk-1234' \
|
||
--header 'Content-Type: application/json' \
|
||
--data-raw '{
|
||
"max_budget": 100,
|
||
"metadata": {"team": "core-infra"},
|
||
"models": ["gpt-4", "gpt-3.5-turbo"]
|
||
}'
|
||
```
|
||
|
||
Note: This is an Enterprise feature. It requires a premium license to use.
|
||
"""
|
||
try:
|
||
from litellm.proxy.proxy_server import (
|
||
hash_token,
|
||
master_key,
|
||
premium_user,
|
||
prisma_client,
|
||
proxy_logging_obj,
|
||
user_api_key_cache,
|
||
)
|
||
|
||
if data is not None:
|
||
_check_allowed_routes_caller_permission(
|
||
allowed_routes=data.allowed_routes,
|
||
user_api_key_dict=user_api_key_dict,
|
||
allowed_routes_was_provided="allowed_routes" in data.model_fields_set,
|
||
)
|
||
_check_passthrough_routes_caller_permission(
|
||
data=data,
|
||
user_api_key_dict=user_api_key_dict,
|
||
)
|
||
_check_permissions_caller_permission(
|
||
data=data,
|
||
user_api_key_dict=user_api_key_dict,
|
||
)
|
||
# Mirror /key/generate's post-handle_key_type recheck so a
|
||
# non-admin can't elevate via a key_type preset that the
|
||
# regenerate flow would otherwise carry through unchecked.
|
||
# The empty dict is intentional — `handle_key_type` is reused
|
||
# purely as a side-effect-free lookup of the preset bucket, not
|
||
# to mutate an existing `data_json`. Do not pass a real
|
||
# `data_json` here; that path would write the derived routes
|
||
# into the DB update payload and is owned by
|
||
# `_common_key_generation_helper`.
|
||
_check_allowed_routes_caller_permission(
|
||
allowed_routes=handle_key_type(data, {}).get("allowed_routes"),
|
||
user_api_key_dict=user_api_key_dict,
|
||
allow_safe_presets=True,
|
||
)
|
||
|
||
# Premium-gate bypass for master-key rotation must verify the
|
||
# caller actually holds the master key, not just that the request
|
||
# body has a ``new_master_key`` field. A presence-only check let
|
||
# any non-premium caller skip the enterprise gate by sending any
|
||
# value in that field.
|
||
regenerate_target_key: Final = data.key if data and data.key else key
|
||
is_master_key_regeneration: Final = (
|
||
data is not None
|
||
and data.new_master_key is not None
|
||
and _is_master_key(api_key=regenerate_target_key, _master_key=master_key)
|
||
)
|
||
|
||
if (
|
||
premium_user is not True and not is_master_key_regeneration
|
||
): # allow master key regeneration for non-premium users
|
||
raise ValueError(
|
||
f"Regenerating Virtual Keys is an Enterprise feature, {CommonProxyErrors.not_premium_user.value}"
|
||
)
|
||
|
||
# Check if key exists, raise exception if key is not in the DB
|
||
key = data.key if data and data.key else key
|
||
if not key:
|
||
raise HTTPException(status_code=400, detail={"error": "No key passed in."})
|
||
### 1. Create New copy that is duplicate of existing key
|
||
######################################################################
|
||
|
||
# create duplicate of existing key
|
||
# set token = new token generated
|
||
# insert new token in DB
|
||
|
||
# create hash of token
|
||
if prisma_client is None:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||
detail={"error": "DB not connected. prisma_client is None"},
|
||
)
|
||
|
||
_is_master_key_valid: Final = _is_master_key(api_key=key, _master_key=master_key)
|
||
|
||
if master_key is not None and data and _is_master_key_valid:
|
||
if data.new_master_key is None:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_400_BAD_REQUEST,
|
||
detail={"error": "New master key is required."},
|
||
)
|
||
await _rotate_master_key(
|
||
prisma_client=prisma_client,
|
||
user_api_key_dict=user_api_key_dict,
|
||
current_master_key=master_key,
|
||
new_master_key=data.new_master_key,
|
||
)
|
||
return GenerateKeyResponse(
|
||
key=data.new_master_key,
|
||
token=data.new_master_key,
|
||
key_name=data.new_master_key,
|
||
expires=None,
|
||
)
|
||
|
||
if "sk" not in key:
|
||
hashed_api_key = key
|
||
else:
|
||
hashed_api_key = hash_token(key)
|
||
|
||
_key_in_db: Final = await _prisma_table(VerificationTokenRepository(prisma_client)).find_unique(
|
||
where={"token": hashed_api_key},
|
||
)
|
||
if _key_in_db is None:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_404_NOT_FOUND,
|
||
detail={"error": f"Key {key} not found."},
|
||
)
|
||
|
||
# check if user has permission to regenerate key
|
||
await TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint(
|
||
user_api_key_dict=user_api_key_dict,
|
||
route=KeyManagementRoutes.KEY_REGENERATE,
|
||
prisma_client=prisma_client,
|
||
existing_key_row=_key_in_db,
|
||
user_api_key_cache=user_api_key_cache,
|
||
)
|
||
|
||
# check if user has ownership permission to regenerate key
|
||
if not await can_modify_verification_token(
|
||
key_info=_key_in_db,
|
||
user_api_key_cache=user_api_key_cache,
|
||
user_api_key_dict=user_api_key_dict,
|
||
prisma_client=prisma_client,
|
||
):
|
||
raise HTTPException(
|
||
status_code=status.HTTP_403_FORBIDDEN,
|
||
detail={"error": "You are not authorized to regenerate this key"},
|
||
)
|
||
|
||
if data is not None and (data.access_group_ids or data.object_permission is not None):
|
||
regenerate_team_table: LiteLLM_TeamTableCachedObj | None = None
|
||
if _key_in_db.team_id is not None:
|
||
regenerate_team_table = await get_team_object(
|
||
team_id=_key_in_db.team_id,
|
||
prisma_client=prisma_client,
|
||
user_api_key_cache=user_api_key_cache,
|
||
check_db_only=True,
|
||
)
|
||
_regen_is_proxy_admin: Final = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
|
||
TeamMemberPermissionChecks.enforce_member_can_assign_access_groups(
|
||
user_api_key_dict=user_api_key_dict,
|
||
team_table=regenerate_team_table,
|
||
access_group_ids=data.access_group_ids,
|
||
)
|
||
_regen_object_permission_dict = _object_permission_to_dict(data.object_permission)
|
||
normalized_object_permission: Final = await validate_key_mcp_servers_against_team(
|
||
object_permission=_regen_object_permission_dict,
|
||
team_obj=regenerate_team_table,
|
||
prisma_client=prisma_client,
|
||
is_proxy_admin=_regen_is_proxy_admin,
|
||
)
|
||
if normalized_object_permission is not None:
|
||
data.object_permission = LiteLLM_ObjectPermissionBase(**normalized_object_permission)
|
||
_regen_object_permission_dict = normalized_object_permission
|
||
await validate_key_search_tools_against_team(
|
||
object_permission=_regen_object_permission_dict,
|
||
team_obj=regenerate_team_table,
|
||
is_proxy_admin=_regen_is_proxy_admin,
|
||
)
|
||
await validate_key_vector_stores_against_team(
|
||
object_permission=_regen_object_permission_dict,
|
||
team_obj=regenerate_team_table,
|
||
is_proxy_admin=_regen_is_proxy_admin,
|
||
)
|
||
|
||
verbose_proxy_logger.info(
|
||
"Key regeneration requested: key_alias=%s",
|
||
getattr(_key_in_db, "key_alias", None),
|
||
)
|
||
verbose_proxy_logger.debug("key_in_db: %s", _key_in_db)
|
||
|
||
# Normalize litellm_changed_by: if it's a Header object or not a string, convert to None
|
||
if litellm_changed_by is not None and not isinstance(litellm_changed_by, str):
|
||
litellm_changed_by = None
|
||
|
||
# Save the old key record to deleted table before regeneration.
|
||
# This preserves key_alias and team_id metadata for historical spend records.
|
||
# If this fails, abort the regeneration to avoid permanently losing the
|
||
# old hash→metadata mapping.
|
||
await _persist_deleted_verification_tokens(
|
||
keys=[_key_in_db],
|
||
prisma_client=prisma_client,
|
||
user_api_key_dict=user_api_key_dict,
|
||
litellm_changed_by=litellm_changed_by,
|
||
)
|
||
|
||
return await _execute_virtual_key_regeneration(
|
||
prisma_client=prisma_client,
|
||
key_in_db=_key_in_db,
|
||
hashed_api_key=hashed_api_key,
|
||
key=key,
|
||
data=data,
|
||
user_api_key_dict=user_api_key_dict,
|
||
litellm_changed_by=litellm_changed_by,
|
||
user_api_key_cache=user_api_key_cache,
|
||
proxy_logging_obj=proxy_logging_obj,
|
||
)
|
||
except Exception as e:
|
||
verbose_proxy_logger.exception("Error regenerating key: %s", e)
|
||
raise handle_exception_on_proxy(e)
|
||
|
||
|
||
async def _check_proxy_or_team_admin_for_key(
|
||
key_in_db: LiteLLM_VerificationToken,
|
||
user_api_key_dict: UserAPIKeyAuth,
|
||
prisma_client: PrismaClient,
|
||
user_api_key_cache: UserApiKeyCache,
|
||
) -> None:
|
||
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
|
||
return
|
||
|
||
if key_in_db.team_id is not None:
|
||
team_table: Final = await get_team_object(
|
||
team_id=key_in_db.team_id,
|
||
prisma_client=prisma_client,
|
||
user_api_key_cache=user_api_key_cache,
|
||
check_db_only=True,
|
||
)
|
||
if team_table is not None:
|
||
if _is_user_team_admin(
|
||
user_api_key_dict=user_api_key_dict,
|
||
team_obj=team_table,
|
||
):
|
||
return
|
||
|
||
raise HTTPException(
|
||
status_code=status.HTTP_403_FORBIDDEN,
|
||
detail={"error": "You must be a proxy admin or team admin to reset key spend"},
|
||
)
|
||
|
||
|
||
def _validate_reset_spend_value(reset_to: object, key_in_db: LiteLLM_VerificationToken) -> float:
|
||
if not isinstance(reset_to, (int, float)):
|
||
raise HTTPException(
|
||
status_code=status.HTTP_400_BAD_REQUEST,
|
||
detail={"error": "reset_to must be a float"},
|
||
)
|
||
|
||
reset_to = float(reset_to)
|
||
|
||
if reset_to < 0:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_400_BAD_REQUEST,
|
||
detail={"error": "reset_to must be >= 0"},
|
||
)
|
||
|
||
current_spend: Final = key_in_db.spend or 0.0
|
||
if reset_to > current_spend:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_400_BAD_REQUEST,
|
||
detail={"error": f"reset_to ({reset_to}) must be <= current spend ({current_spend})"},
|
||
)
|
||
|
||
max_budget = key_in_db.max_budget
|
||
if key_in_db.litellm_budget_table is not None:
|
||
budget_max_budget: Final = getattr(key_in_db.litellm_budget_table, "max_budget", None)
|
||
if budget_max_budget is not None:
|
||
if max_budget is None or budget_max_budget < max_budget:
|
||
max_budget = budget_max_budget
|
||
|
||
if max_budget is not None and reset_to > max_budget:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_400_BAD_REQUEST,
|
||
detail={"error": f"reset_to ({reset_to}) must be <= budget ({max_budget})"},
|
||
)
|
||
|
||
return reset_to
|
||
|
||
|
||
@router.post(
|
||
"/key/{key:path}/reset_spend",
|
||
tags=["key management"],
|
||
dependencies=[Depends(user_api_key_auth)],
|
||
)
|
||
@management_endpoint_wrapper
|
||
async def reset_key_spend_fn(
|
||
key: str,
|
||
data: ResetSpendRequest,
|
||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||
litellm_changed_by: str | None = Header(
|
||
None,
|
||
description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability",
|
||
),
|
||
) -> dict[str, Any]:
|
||
try:
|
||
from litellm.proxy.proxy_server import (
|
||
hash_token,
|
||
prisma_client,
|
||
proxy_logging_obj,
|
||
user_api_key_cache,
|
||
)
|
||
|
||
if prisma_client is None:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||
detail={"error": "DB not connected. prisma_client is None"},
|
||
)
|
||
|
||
if "sk" not in key:
|
||
hashed_api_key = key
|
||
else:
|
||
hashed_api_key = hash_token(key)
|
||
|
||
_key_in_db: Final = await _prisma_table(VerificationTokenRepository(prisma_client)).find_unique(
|
||
where={"token": hashed_api_key},
|
||
include={"litellm_budget_table": True},
|
||
)
|
||
if _key_in_db is None:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_404_NOT_FOUND,
|
||
detail={"error": f"Key {key} not found."},
|
||
)
|
||
|
||
current_spend: Final = _key_in_db.spend or 0.0
|
||
reset_to: Final = _validate_reset_spend_value(data.reset_to, _key_in_db)
|
||
|
||
await _check_proxy_or_team_admin_for_key(
|
||
key_in_db=_key_in_db,
|
||
user_api_key_dict=user_api_key_dict,
|
||
prisma_client=prisma_client,
|
||
user_api_key_cache=user_api_key_cache,
|
||
)
|
||
|
||
updated_key: Final = await _prisma_table(VerificationTokenRepository(prisma_client)).update(
|
||
where={"token": hashed_api_key},
|
||
data={"spend": reset_to},
|
||
)
|
||
|
||
if updated_key is None:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||
detail={"error": "Failed to update key spend"},
|
||
)
|
||
|
||
await _delete_cache_key_object(
|
||
hashed_token=hashed_api_key,
|
||
user_api_key_cache=user_api_key_cache,
|
||
proxy_logging_obj=proxy_logging_obj,
|
||
)
|
||
|
||
# Set Redis spend counter to the new value so get_current_spend()
|
||
# returns the correct amount immediately instead of the stale pre-reset value.
|
||
# We use reset_to (not 0.0) so partial resets are reflected correctly.
|
||
from litellm.proxy.proxy_server import spend_counter_cache
|
||
|
||
_counter_key: Final = f"spend:key:{hashed_api_key}"
|
||
spend_counter_cache.in_memory_cache.set_cache(key=_counter_key, value=reset_to, ttl=60)
|
||
if spend_counter_cache.redis_cache is not None:
|
||
try:
|
||
await spend_counter_cache.redis_cache.async_set_cache(key=_counter_key, value=reset_to, ttl=60)
|
||
except Exception as redis_err:
|
||
verbose_proxy_logger.warning(
|
||
"Failed to update spend counter %s in Redis: %s. "
|
||
"Budget checks may use stale value until counter expires.",
|
||
_counter_key,
|
||
redis_err,
|
||
)
|
||
|
||
max_budget: Final = updated_key.max_budget
|
||
budget_reset_at: Final = updated_key.budget_reset_at
|
||
|
||
return {
|
||
"key_hash": hashed_api_key,
|
||
"spend": reset_to,
|
||
"previous_spend": current_spend,
|
||
"max_budget": max_budget,
|
||
"budget_reset_at": budget_reset_at,
|
||
}
|
||
except HTTPException:
|
||
raise
|
||
except Exception as e:
|
||
verbose_proxy_logger.exception("Error resetting key spend: %s", e)
|
||
raise handle_exception_on_proxy(e)
|
||
|
||
|
||
async def validate_key_list_check(
|
||
user_api_key_dict: UserAPIKeyAuth,
|
||
user_id: str | None,
|
||
team_id: str | None,
|
||
organization_id: str | None,
|
||
key_alias: str | None,
|
||
key_hash: str | None,
|
||
prisma_client: PrismaClient,
|
||
) -> LiteLLM_UserTable | None:
|
||
if _user_has_admin_view(user_api_key_dict):
|
||
return None
|
||
|
||
if user_api_key_dict.user_id is None:
|
||
raise ProxyException(
|
||
message="You are not authorized to access this endpoint. No 'user_id' is associated with your API key.",
|
||
type=ProxyErrorTypes.bad_request_error,
|
||
param="user_id",
|
||
code=status.HTTP_403_FORBIDDEN,
|
||
)
|
||
complete_user_info_db_obj: Final[BaseModel | None] = await _prisma_table(UserRepository(prisma_client)).find_unique(
|
||
where={"user_id": user_api_key_dict.user_id},
|
||
include={"organization_memberships": True},
|
||
)
|
||
|
||
if complete_user_info_db_obj is None:
|
||
raise ProxyException(
|
||
message="You are not authorized to access this endpoint. No 'user_id' is associated with your API key.",
|
||
type=ProxyErrorTypes.bad_request_error,
|
||
param="user_id",
|
||
code=status.HTTP_403_FORBIDDEN,
|
||
)
|
||
|
||
complete_user_info: Final = LiteLLM_UserTable.model_validate(complete_user_info_db_obj.model_dump())
|
||
|
||
# internal user can only see their own keys
|
||
if user_id:
|
||
if complete_user_info.user_id != user_id:
|
||
raise ProxyException(
|
||
message="You are not authorized to check another user's keys",
|
||
type=ProxyErrorTypes.bad_request_error,
|
||
param="user_id",
|
||
code=status.HTTP_403_FORBIDDEN,
|
||
)
|
||
|
||
if team_id:
|
||
if team_id not in complete_user_info.teams:
|
||
raise ProxyException(
|
||
message="You are not authorized to check this team's keys",
|
||
type=ProxyErrorTypes.bad_request_error,
|
||
param="team_id",
|
||
code=status.HTTP_403_FORBIDDEN,
|
||
)
|
||
|
||
if organization_id:
|
||
if complete_user_info.organization_memberships is None or organization_id not in [
|
||
membership.organization_id for membership in complete_user_info.organization_memberships
|
||
]:
|
||
raise ProxyException(
|
||
message="You are not authorized to check this organization's keys",
|
||
type=ProxyErrorTypes.bad_request_error,
|
||
param="organization_id",
|
||
code=status.HTTP_403_FORBIDDEN,
|
||
)
|
||
|
||
if key_hash:
|
||
try:
|
||
key_info: Final[LiteLLM_VerificationToken | None] = await _prisma_table(
|
||
VerificationTokenRepository(prisma_client)
|
||
).find_unique(
|
||
where={"token": key_hash},
|
||
)
|
||
except Exception:
|
||
raise ProxyException(
|
||
message="Key Hash not found.",
|
||
type=ProxyErrorTypes.bad_request_error,
|
||
param="key_hash",
|
||
code=status.HTTP_403_FORBIDDEN,
|
||
)
|
||
if key_info is None:
|
||
raise ProxyException(
|
||
message="Key Hash not found.",
|
||
type=ProxyErrorTypes.bad_request_error,
|
||
param="key_hash",
|
||
code=status.HTTP_403_FORBIDDEN,
|
||
)
|
||
can_user_query_key_info: Final = await _can_user_query_key_info(
|
||
user_api_key_dict=user_api_key_dict,
|
||
key=key_hash,
|
||
key_info=key_info,
|
||
)
|
||
if not can_user_query_key_info:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_403_FORBIDDEN,
|
||
detail=f"You are not allowed to access this key's info. Your role={user_api_key_dict.user_role}",
|
||
)
|
||
return complete_user_info
|
||
|
||
|
||
async def _fetch_user_team_objects(
|
||
complete_user_info: LiteLLM_UserTable | None,
|
||
prisma_client: PrismaClient,
|
||
) -> list[LiteLLM_TeamTable]:
|
||
"""Fetch team objects for all teams a user belongs to (single DB query)."""
|
||
if complete_user_info is None or not complete_user_info.teams:
|
||
return []
|
||
|
||
teams: Final[Sequence[BaseModel] | None] = cast( # cast-ok: the None guard below predates the non-optional seam
|
||
"Sequence[BaseModel] | None",
|
||
await TeamRepository(prisma_client).table.find_many(where={"team_id": {"in": complete_user_info.teams}}),
|
||
)
|
||
if teams is None:
|
||
return []
|
||
|
||
return [LiteLLM_TeamTable.model_validate(team.model_dump()) for team in teams]
|
||
|
||
|
||
def _get_admin_team_ids_from_objects(
|
||
user_api_key_dict: UserAPIKeyAuth,
|
||
team_objects: list[LiteLLM_TeamTable],
|
||
) -> list[str]:
|
||
"""Filter team objects to those where the user is an admin."""
|
||
return [
|
||
team.team_id for team in team_objects if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team)
|
||
]
|
||
|
||
|
||
def _get_team_ids_with_key_list_permission_from_objects(
|
||
user_api_key_dict: UserAPIKeyAuth,
|
||
team_objects: list[LiteLLM_TeamTable],
|
||
) -> list[str]:
|
||
"""Filter team objects to non-admin teams where the caller has /key/list
|
||
permission via team_member_permissions. These teams should grant the
|
||
caller full key visibility (same as a team admin), so other members'
|
||
keys and service account keys (user_id=NULL) are returned."""
|
||
return [
|
||
team.team_id
|
||
for team in team_objects
|
||
if not _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team)
|
||
and _team_member_has_permission(
|
||
user_api_key_dict=user_api_key_dict,
|
||
team_obj=team,
|
||
permission=KeyManagementRoutes.KEY_LIST.value,
|
||
)
|
||
]
|
||
|
||
|
||
def _get_member_team_ids_from_objects(
|
||
user_api_key_dict: UserAPIKeyAuth,
|
||
team_objects: list[LiteLLM_TeamTable],
|
||
) -> list[str]:
|
||
"""Filter team objects to those where the user is a member (any role)."""
|
||
return [
|
||
team.team_id
|
||
for team in team_objects
|
||
if any(
|
||
member.user_id is not None and member.user_id == user_api_key_dict.user_id
|
||
for member in team.members_with_roles
|
||
)
|
||
]
|
||
|
||
|
||
async def get_admin_team_ids(
|
||
complete_user_info: LiteLLM_UserTable | None,
|
||
user_api_key_dict: UserAPIKeyAuth,
|
||
prisma_client: PrismaClient,
|
||
) -> list[str]:
|
||
"""Get all team IDs where the user is an admin."""
|
||
team_objects: Final = await _fetch_user_team_objects(complete_user_info, prisma_client)
|
||
return _get_admin_team_ids_from_objects(user_api_key_dict, team_objects)
|
||
|
||
|
||
async def get_member_team_ids(
|
||
complete_user_info: LiteLLM_UserTable | None,
|
||
user_api_key_dict: UserAPIKeyAuth,
|
||
prisma_client: PrismaClient,
|
||
) -> list[str]:
|
||
"""
|
||
Get all team IDs where the user is a member (any role, including admin).
|
||
|
||
Used to determine which teams' service accounts (keys with user_id=NULL)
|
||
a regular team member can see.
|
||
"""
|
||
team_objects: Final = await _fetch_user_team_objects(complete_user_info, prisma_client)
|
||
return _get_member_team_ids_from_objects(user_api_key_dict, team_objects)
|
||
|
||
|
||
VALID_EXPIRES_FILTER_VALUES: Final = frozenset({"active", "expired"})
|
||
|
||
|
||
@router.get(
|
||
"/key/list",
|
||
tags=["key management"],
|
||
dependencies=[Depends(user_api_key_auth)],
|
||
)
|
||
@management_endpoint_wrapper
|
||
async def list_keys(
|
||
request: Request,
|
||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||
page: int = Query(1, description="Page number", ge=1),
|
||
size: int = Query(10, description="Page size", ge=1, le=100),
|
||
user_id: str | None = Query(
|
||
None,
|
||
description="Filter keys by user ID. Exact match by default; set substring_matching=true (admin only) for case-insensitive substring matching.",
|
||
),
|
||
team_id: str | None = Query(None, description="Filter keys by team ID"),
|
||
organization_id: str | None = Query(None, description="Filter keys by organization ID"),
|
||
key_hash: str | None = Query(None, description="Filter keys by key hash"),
|
||
key_alias: str | None = Query(
|
||
None,
|
||
description="Filter keys by key alias. Exact match by default; set substring_matching=true (admin only) for case-insensitive substring matching.",
|
||
),
|
||
return_full_object: bool = Query(False, description="Return full key object"),
|
||
include_team_keys: bool = Query(False, description="Include all keys for teams that user is an admin of."),
|
||
include_created_by_keys: bool = Query(False, description="Include keys created by the user"),
|
||
sort_by: str | None = Query(
|
||
default=None,
|
||
description="Column to sort by (e.g. 'user_id', 'created_at', 'spend')",
|
||
),
|
||
sort_order: str = Query(default="desc", description="Sort order ('asc' or 'desc')"),
|
||
expand: list[str] | None = Query(None, description="Expand related objects (e.g. 'user')"),
|
||
status: str | None = Query(None, description="Filter by status (e.g. 'deleted')"),
|
||
project_id: str | None = Query(None, description="Filter keys by project ID"),
|
||
access_group_id: str | None = Query(None, description="Filter keys by access group ID"),
|
||
agent_id: str | None = Query(None, description="Filter keys by agent ID"),
|
||
substring_matching: bool = Query(
|
||
False,
|
||
description="If true (proxy admins only), match user_id/key_alias as case-insensitive substrings instead of exact values. Defaults to false: /key/list matched these exactly before substring search was added, and an exact user_id/key_alias filter must never return another user's keys.",
|
||
),
|
||
expires: str | None = Query(
|
||
None,
|
||
description="Filter keys by expiration. 'expired' returns keys whose expires is in the past; 'active' returns keys that never expire or expire in the future. Omit to return keys regardless of expiration.",
|
||
),
|
||
) -> KeyListResponseObject:
|
||
"""
|
||
List all keys for a given user / team / organization.
|
||
|
||
Parameters:
|
||
expand: Optional[List[str]] - Expand related objects (e.g. 'user' to include user information)
|
||
status: Optional[str] - Filter by status. Currently supports "deleted" to query deleted keys.
|
||
|
||
Returns:
|
||
{
|
||
"keys": List[str] or List[UserAPIKeyAuth],
|
||
"total_count": int,
|
||
"current_page": int,
|
||
"total_pages": int,
|
||
}
|
||
|
||
When expand includes "user", each key object will include a "user" field with the associated user object.
|
||
Note: When expand=user is specified, full key objects are returned regardless of the return_full_object parameter.
|
||
"""
|
||
try:
|
||
from litellm.proxy.proxy_server import prisma_client
|
||
|
||
verbose_proxy_logger.debug("Entering list_keys function")
|
||
|
||
if prisma_client is None:
|
||
verbose_proxy_logger.error("Database not connected")
|
||
raise Exception("Database not connected")
|
||
|
||
# Validate status parameter
|
||
if status is not None and status != "deleted":
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail={"error": "Invalid status value. Currently only 'deleted' is supported."},
|
||
)
|
||
|
||
if isinstance(expires, str) and expires not in VALID_EXPIRES_FILTER_VALUES:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail={"error": "Invalid expires value. Supported: 'active', 'expired'."},
|
||
)
|
||
|
||
complete_user_info: Final = await validate_key_list_check(
|
||
user_api_key_dict=user_api_key_dict,
|
||
user_id=user_id,
|
||
team_id=team_id,
|
||
organization_id=organization_id,
|
||
key_alias=key_alias,
|
||
key_hash=key_hash,
|
||
prisma_client=prisma_client,
|
||
)
|
||
|
||
# Fetch team objects once when needed for either admin or member filtering.
|
||
# This avoids duplicate DB queries for the same team data.
|
||
if include_team_keys or include_created_by_keys:
|
||
team_objects = await _fetch_user_team_objects(
|
||
complete_user_info=complete_user_info,
|
||
prisma_client=prisma_client,
|
||
)
|
||
member_team_ids = _get_member_team_ids_from_objects(
|
||
user_api_key_dict=user_api_key_dict,
|
||
team_objects=team_objects,
|
||
)
|
||
else:
|
||
team_objects = []
|
||
member_team_ids = None
|
||
|
||
if include_team_keys:
|
||
admin_team_ids = _get_admin_team_ids_from_objects(
|
||
user_api_key_dict=user_api_key_dict,
|
||
team_objects=team_objects,
|
||
)
|
||
# Non-admin members with /key/list permission get full team-key
|
||
# visibility for that team — matching the UI contract that
|
||
# granting this permission lets them see all keys within the team.
|
||
list_permission_team_ids: Final = _get_team_ids_with_key_list_permission_from_objects(
|
||
user_api_key_dict=user_api_key_dict,
|
||
team_objects=team_objects,
|
||
)
|
||
if list_permission_team_ids:
|
||
admin_team_ids = list({*admin_team_ids, *list_permission_team_ids})
|
||
else:
|
||
admin_team_ids = None
|
||
|
||
is_proxy_admin: Final = user_api_key_dict.user_role in [
|
||
LitellmUserRoles.PROXY_ADMIN.value,
|
||
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value,
|
||
]
|
||
|
||
# Substring matching is opt-in (admin-only). /key/list matched user_id and
|
||
# key_alias exactly before substring search was added; auto-applying a
|
||
# substring match to every admin call broke that contract and let a caller
|
||
# passing an exact user_id (e.g. an integration scoping to one user with an
|
||
# admin key) receive other users' keys (user_id="alice" -> "alice2"). Exact
|
||
# by default restores the prior behavior; the dashboard opts in explicitly.
|
||
use_substring_matching: Final = substring_matching and is_proxy_admin
|
||
|
||
# Admins may omit user_id to list all keys; non-admins are scoped to self.
|
||
if not user_id and not is_proxy_admin:
|
||
user_id = user_api_key_dict.user_id
|
||
|
||
response: Final = await _list_key_helper(
|
||
prisma_client=prisma_client,
|
||
page=page,
|
||
size=size,
|
||
user_id=user_id,
|
||
team_id=team_id,
|
||
key_alias=key_alias,
|
||
key_hash=key_hash,
|
||
return_full_object=return_full_object,
|
||
organization_id=organization_id,
|
||
admin_team_ids=admin_team_ids,
|
||
member_team_ids=member_team_ids,
|
||
include_created_by_keys=include_created_by_keys,
|
||
sort_by=sort_by,
|
||
sort_order=sort_order,
|
||
expand=expand,
|
||
status=status,
|
||
project_id=project_id,
|
||
access_group_id=access_group_id,
|
||
agent_id=agent_id,
|
||
use_substring_matching=use_substring_matching,
|
||
expires_filter=expires if isinstance(expires, str) else None,
|
||
)
|
||
|
||
verbose_proxy_logger.debug("Successfully prepared response")
|
||
|
||
return response
|
||
|
||
except Exception as e:
|
||
verbose_proxy_logger.exception("Error in list_keys: %s", e)
|
||
if isinstance(e, HTTPException):
|
||
raise ProxyException(
|
||
message=getattr(e, "detail", f"error({e})"),
|
||
type=ProxyErrorTypes.internal_server_error,
|
||
param=getattr(e, "param", "None"),
|
||
code=getattr(e, "status_code", fastapi.status.HTTP_500_INTERNAL_SERVER_ERROR),
|
||
)
|
||
elif isinstance(e, ProxyException):
|
||
raise e
|
||
raise ProxyException(
|
||
message="Authentication Error, " + str(e),
|
||
type=ProxyErrorTypes.internal_server_error,
|
||
param=getattr(e, "param", "None"),
|
||
code=fastapi.status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||
)
|
||
|
||
|
||
async def _apply_non_admin_alias_scope(
|
||
user_api_key_dict: UserAPIKeyAuth,
|
||
prisma_client: PrismaClient,
|
||
query_params: list[object],
|
||
where_parts: list[str],
|
||
) -> None:
|
||
"""Append SQL scope conditions so non-admin users only see aliases for
|
||
keys they own or keys belonging to teams they are members of."""
|
||
scope_conditions: Final[list[str]] = []
|
||
if user_api_key_dict.user_id:
|
||
query_params.append(user_api_key_dict.user_id)
|
||
scope_conditions.append(f"user_id = ${len(query_params)}")
|
||
|
||
# Look up the user's teams from the user table
|
||
user_teams: list[str] = []
|
||
if user_api_key_dict.user_id:
|
||
user_row: Final = await _prisma_table(UserRepository(prisma_client)).find_unique(
|
||
where={"user_id": user_api_key_dict.user_id}
|
||
)
|
||
if user_row is not None:
|
||
user_teams = getattr(user_row, "teams", []) or []
|
||
|
||
if user_teams:
|
||
team_placeholders: Final = ", ".join(f"${len(query_params) + i + 1}" for i in range(len(user_teams)))
|
||
query_params.extend(user_teams)
|
||
scope_conditions.append(f"team_id IN ({team_placeholders})")
|
||
|
||
if scope_conditions:
|
||
where_parts.append(f"({' OR '.join(scope_conditions)})")
|
||
else:
|
||
# No user_id and no teams — return nothing
|
||
where_parts.append("FALSE")
|
||
|
||
|
||
@router.get(
|
||
"/key/aliases",
|
||
tags=["key management"],
|
||
dependencies=[Depends(user_api_key_auth)],
|
||
)
|
||
@management_endpoint_wrapper
|
||
async def key_aliases(
|
||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||
page: int = Query(1, ge=1, description="Page number"),
|
||
size: int = Query(50, ge=1, le=100, description="Page size"),
|
||
search: str | None = Query(None, description="Search key aliases (case-insensitive partial match)"),
|
||
team_id: str | None = Query(None, description="Filter aliases to keys belonging to this team"),
|
||
) -> dict[str, Any]:
|
||
"""
|
||
Lists key aliases with pagination and optional search.
|
||
|
||
Non-admin users only see aliases for keys they own or keys belonging to
|
||
their teams.
|
||
|
||
Returns:
|
||
{
|
||
"aliases": List[str],
|
||
"total_count": int,
|
||
"current_page": int,
|
||
"total_pages": int,
|
||
"size": int,
|
||
}
|
||
"""
|
||
try:
|
||
from litellm.proxy.proxy_server import prisma_client
|
||
|
||
verbose_proxy_logger.debug("Entering key_aliases function")
|
||
|
||
if prisma_client is None:
|
||
verbose_proxy_logger.error("Database not connected")
|
||
raise Exception("Database not connected")
|
||
|
||
# Build a parameterized WHERE clause to avoid loading full rows into
|
||
# memory. Raw SQL is used because the Prisma client wrapper does not
|
||
# support column-level SELECT projection on find_many.
|
||
#
|
||
# $1 is always UI_SESSION_TOKEN_TEAM_ID (filters out UI session tokens).
|
||
query_params: Final[list[object]] = [UI_SESSION_TOKEN_TEAM_ID]
|
||
where_parts: Final = [
|
||
"key_alias IS NOT NULL",
|
||
"key_alias != ''",
|
||
"(team_id IS NULL OR team_id != $1)",
|
||
]
|
||
|
||
# Scope results for non-admin users: only show aliases for keys the
|
||
# user owns or keys belonging to teams they are a member of.
|
||
is_proxy_admin: Final = user_api_key_dict.user_role in [
|
||
LitellmUserRoles.PROXY_ADMIN.value,
|
||
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value,
|
||
]
|
||
if not is_proxy_admin:
|
||
await _apply_non_admin_alias_scope(user_api_key_dict, prisma_client, query_params, where_parts)
|
||
|
||
if search:
|
||
query_params.append(f"%{search}%")
|
||
where_parts.append(f"key_alias ILIKE ${len(query_params)}")
|
||
|
||
if team_id:
|
||
query_params.append(team_id)
|
||
where_parts.append(f"team_id = ${len(query_params)}")
|
||
|
||
where_sql: Final = " AND ".join(where_parts)
|
||
|
||
count_sql: Final = f'SELECT COUNT(*) AS count FROM "LiteLLM_VerificationToken" WHERE {where_sql}'
|
||
count_rows: Final[Sequence[Mapping[str, int]]] = await prisma_client.db.query_raw(count_sql, *query_params)
|
||
total_count: Final = int(count_rows[0]["count"]) if count_rows else 0
|
||
|
||
aliases_params: Final = query_params + [size, (page - 1) * size]
|
||
limit_idx: Final = len(aliases_params) - 1
|
||
offset_idx: Final = len(aliases_params)
|
||
aliases_sql: Final = (
|
||
f"SELECT key_alias"
|
||
f' FROM "LiteLLM_VerificationToken"'
|
||
f" WHERE {where_sql}"
|
||
f" ORDER BY key_alias ASC"
|
||
f" LIMIT ${limit_idx} OFFSET ${offset_idx}"
|
||
)
|
||
alias_rows: Final[Sequence[Mapping[str, str]]] = await prisma_client.db.query_raw(aliases_sql, *aliases_params)
|
||
aliases: Final[list[str]] = [row["key_alias"] for row in alias_rows if row.get("key_alias")]
|
||
|
||
total_pages: Final = -(-total_count // size) if total_count > 0 else 0
|
||
verbose_proxy_logger.debug(
|
||
"key_aliases: page=%s, size=%s, search=%r, total_count=%s, total_pages=%s",
|
||
page,
|
||
size,
|
||
search,
|
||
total_count,
|
||
total_pages,
|
||
)
|
||
|
||
return {
|
||
"aliases": aliases,
|
||
"total_count": total_count,
|
||
"current_page": page,
|
||
"total_pages": total_pages,
|
||
"size": size,
|
||
}
|
||
|
||
except Exception as e:
|
||
verbose_proxy_logger.exception("Error in key_aliases: %s", e)
|
||
if isinstance(e, HTTPException):
|
||
raise ProxyException(
|
||
message=getattr(e, "detail", f"error({e})"),
|
||
type=ProxyErrorTypes.internal_server_error,
|
||
param=getattr(e, "param", "None"),
|
||
code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR),
|
||
)
|
||
elif isinstance(e, ProxyException):
|
||
raise e
|
||
raise ProxyException(
|
||
message="Authentication Error, " + str(e),
|
||
type=ProxyErrorTypes.internal_server_error,
|
||
param=getattr(e, "param", "None"),
|
||
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||
)
|
||
|
||
|
||
def _validate_sort_params(sort_by: str | None, sort_order: str) -> dict[str, str] | None:
|
||
order_by: Final[dict[str, str]] = {}
|
||
|
||
if sort_by is None:
|
||
return None
|
||
# Validate sort_by is a valid column
|
||
valid_columns: Final = [
|
||
"spend",
|
||
"max_budget",
|
||
"created_at",
|
||
"updated_at",
|
||
"token",
|
||
"key_alias",
|
||
]
|
||
if sort_by not in valid_columns:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail={"error": f"Invalid sort column. Must be one of: {', '.join(valid_columns)}"},
|
||
)
|
||
|
||
# Validate sort_order
|
||
if sort_order.lower() not in ["asc", "desc"]:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail={"error": "Invalid sort order. Must be 'asc' or 'desc'"},
|
||
)
|
||
|
||
order_by[sort_by] = sort_order.lower()
|
||
|
||
return order_by
|
||
|
||
|
||
def _build_expires_where_clause(expires_filter: str, now: datetime) -> dict[str, object]:
|
||
if expires_filter == "expired":
|
||
return {"AND": [{"expires": {"not": None}}, {"expires": {"lt": now}}]}
|
||
return {"OR": [{"expires": None}, {"expires": {"gte": now}}]}
|
||
|
||
|
||
def _build_key_filter_conditions(
|
||
user_id: str | None,
|
||
team_id: str | None,
|
||
organization_id: str | None,
|
||
key_alias: str | None,
|
||
key_hash: str | None,
|
||
exclude_team_id: str | None,
|
||
admin_team_ids: list[str] | None,
|
||
member_team_ids: list[str] | None = None,
|
||
include_created_by_keys: bool = False,
|
||
project_id: str | None = None,
|
||
access_group_id: str | None = None,
|
||
agent_id: str | None = None,
|
||
use_substring_matching: bool = False,
|
||
expires_filter: str | None = None,
|
||
) -> Mapping[str, object]:
|
||
"""Build filter conditions for key listing.
|
||
|
||
Visibility rules:
|
||
- Users always see their own keys (user_id match)
|
||
- Team admins see ALL keys for their admin teams (via admin_team_ids)
|
||
- Regular team members see only service accounts (user_id=NULL) for their
|
||
teams (via member_team_ids). This prevents leaking other members' spend data.
|
||
- created_by visibility is scoped to teams the user currently belongs to,
|
||
so former members cannot see service accounts they created after leaving.
|
||
"""
|
||
# Prepare filter conditions
|
||
where: dict[str, object] = {}
|
||
where.update(_get_condition_to_filter_out_ui_session_tokens())
|
||
|
||
# Build the OR conditions for user's keys and admin team keys
|
||
or_conditions: Final[list[dict[str, object]]] = []
|
||
|
||
# Base conditions for user's own keys
|
||
user_condition: Final[dict[str, object]] = {}
|
||
if user_id and isinstance(user_id, str):
|
||
if use_substring_matching:
|
||
user_condition["user_id"] = {
|
||
"contains": user_id,
|
||
"mode": "insensitive",
|
||
}
|
||
else:
|
||
user_condition["user_id"] = user_id
|
||
if exclude_team_id and isinstance(exclude_team_id, str):
|
||
user_condition["team_id"] = {"not": exclude_team_id}
|
||
if organization_id and isinstance(organization_id, str):
|
||
user_condition["organization_id"] = organization_id
|
||
|
||
if user_condition:
|
||
or_conditions.append(user_condition)
|
||
|
||
# Add condition for created_by keys, scoped to user's current teams
|
||
if include_created_by_keys and user_id:
|
||
if member_team_ids is not None:
|
||
if member_team_ids:
|
||
# Scope created_by keys to teams user is still a member of,
|
||
# or keys that have no team (personal keys)
|
||
or_conditions.append(
|
||
{
|
||
"AND": [
|
||
{"created_by": user_id},
|
||
{
|
||
"OR": [
|
||
{"team_id": {"in": member_team_ids}},
|
||
{"team_id": None},
|
||
]
|
||
},
|
||
]
|
||
}
|
||
)
|
||
else:
|
||
# User is not a member of any team, only show non-team created_by keys
|
||
or_conditions.append({"AND": [{"created_by": user_id}, {"team_id": None}]})
|
||
else:
|
||
# No team membership info provided (backward compatibility for
|
||
# direct _list_key_helper callers like Prometheus)
|
||
or_conditions.append({"created_by": user_id})
|
||
|
||
# Add condition for admin team keys (admins see ALL team keys)
|
||
if admin_team_ids:
|
||
or_conditions.append({"team_id": {"in": admin_team_ids}})
|
||
|
||
# Add condition for member team service accounts (members only see keys with user_id=NULL)
|
||
if member_team_ids:
|
||
# Exclude teams where user is already admin (those are covered above with full visibility)
|
||
member_only_team_ids: Final = [tid for tid in member_team_ids if tid not in (admin_team_ids or [])]
|
||
if member_only_team_ids:
|
||
or_conditions.append(
|
||
{
|
||
"AND": [
|
||
{"team_id": {"in": member_only_team_ids}},
|
||
{"user_id": None},
|
||
]
|
||
}
|
||
)
|
||
|
||
# Combine conditions with OR if we have multiple conditions
|
||
if len(or_conditions) > 1:
|
||
where = {"AND": [where, {"OR": or_conditions}]}
|
||
elif len(or_conditions) == 1:
|
||
where.update(or_conditions[0])
|
||
|
||
# Apply team_id, project_id and access_group_id as global AND filters so they
|
||
# narrow results across all visibility conditions (own keys, team keys, etc.)
|
||
global_filters: Final[tuple[dict[str, object], ...]] = (
|
||
*(
|
||
(
|
||
{"key_alias": {"contains": key_alias, "mode": "insensitive"}}
|
||
if use_substring_matching
|
||
else {"key_alias": key_alias},
|
||
)
|
||
if key_alias and isinstance(key_alias, str)
|
||
else ()
|
||
),
|
||
*(({"token": key_hash},) if key_hash and isinstance(key_hash, str) else ()),
|
||
*(({"team_id": team_id},) if team_id and isinstance(team_id, str) else ()),
|
||
*(({"project_id": project_id},) if project_id else ()),
|
||
*(({"access_group_ids": {"hasSome": [access_group_id]}},) if access_group_id else ()),
|
||
*(({"agent_id": agent_id},) if agent_id and isinstance(agent_id, str) else ()),
|
||
*(
|
||
(_build_expires_where_clause(expires_filter, datetime.now(timezone.utc)),)
|
||
if expires_filter is not None and expires_filter in VALID_EXPIRES_FILTER_VALUES
|
||
else ()
|
||
),
|
||
)
|
||
combined_where: Final[Mapping[str, object]] = {"AND": [where, *global_filters]} if global_filters else where
|
||
verbose_proxy_logger.debug("Filter conditions: %s", combined_where)
|
||
return combined_where
|
||
|
||
|
||
async def _list_key_helper(
|
||
prisma_client: PrismaClient,
|
||
page: int,
|
||
size: int,
|
||
user_id: str | None,
|
||
team_id: str | None,
|
||
organization_id: str | None,
|
||
key_alias: str | None,
|
||
key_hash: str | None,
|
||
exclude_team_id: str | None = None,
|
||
return_full_object: bool = False,
|
||
admin_team_ids: list[str] | None = None, # New parameter for teams where user is admin
|
||
member_team_ids: list[str]
|
||
| None = None, # Team IDs where user is a member (any role) - for service account visibility
|
||
include_created_by_keys: bool = False,
|
||
sort_by: str | None = None,
|
||
sort_order: str = "desc",
|
||
expand: list[str] | None = None,
|
||
status: str | None = None,
|
||
project_id: str | None = None,
|
||
access_group_id: str | None = None,
|
||
agent_id: str | None = None,
|
||
use_substring_matching: bool = False,
|
||
expires_filter: str | None = None,
|
||
) -> KeyListResponseObject:
|
||
"""
|
||
Helper function to list keys
|
||
Args:
|
||
page: int
|
||
size: int
|
||
user_id: Optional[str]
|
||
team_id: Optional[str]
|
||
key_alias: Optional[str]
|
||
exclude_team_id: Optional[str] # exclude a specific team_id
|
||
return_full_object: bool # when true, will return UserAPIKeyAuth objects instead of just the token
|
||
admin_team_ids: Optional[List[str]] # list of team IDs where the user is an admin
|
||
member_team_ids: Optional[List[str]] # list of team IDs where user is a member (for service account visibility)
|
||
|
||
Returns:
|
||
KeyListResponseObject
|
||
{
|
||
"keys": List[str] or List[UserAPIKeyAuth], # Updated to reflect possible return types
|
||
"total_count": int,
|
||
"current_page": int,
|
||
"total_pages": int,
|
||
}
|
||
"""
|
||
where: Final = _build_key_filter_conditions(
|
||
user_id=user_id,
|
||
team_id=team_id,
|
||
organization_id=organization_id,
|
||
key_alias=key_alias,
|
||
key_hash=key_hash,
|
||
exclude_team_id=exclude_team_id,
|
||
admin_team_ids=admin_team_ids,
|
||
member_team_ids=member_team_ids,
|
||
include_created_by_keys=include_created_by_keys,
|
||
project_id=project_id,
|
||
access_group_id=access_group_id,
|
||
agent_id=agent_id,
|
||
use_substring_matching=use_substring_matching,
|
||
expires_filter=expires_filter,
|
||
)
|
||
|
||
# Calculate skip for pagination
|
||
skip: Final = (page - 1) * size
|
||
|
||
verbose_proxy_logger.debug("Pagination: skip=%s, take=%s", skip, size)
|
||
|
||
order_by: Final[dict[str, str] | None] = (
|
||
_validate_sort_params(sort_by, sort_order) if sort_by is not None and isinstance(sort_by, str) else None
|
||
)
|
||
|
||
# Determine which table to query based on status
|
||
use_deleted_table: Final = status == "deleted"
|
||
|
||
# Fetch keys with pagination
|
||
if use_deleted_table:
|
||
keys = await DeletedVerificationTokenRepository(prisma_client).table.find_many(
|
||
where=where,
|
||
skip=skip,
|
||
take=size,
|
||
order=(
|
||
order_by
|
||
if order_by
|
||
else [
|
||
{"created_at": "desc"},
|
||
{"token": "desc"}, # fallback sort
|
||
]
|
||
),
|
||
)
|
||
else:
|
||
keys = await VerificationTokenRepository(prisma_client).table.find_many(
|
||
where=where,
|
||
skip=skip,
|
||
take=size,
|
||
order=(
|
||
order_by
|
||
if order_by
|
||
else [
|
||
{"created_at": "desc"},
|
||
{"token": "desc"}, # fallback sort
|
||
]
|
||
),
|
||
include={"object_permission": True},
|
||
)
|
||
|
||
verbose_proxy_logger.debug("Fetched %s keys", len(keys))
|
||
|
||
# Get total count of keys
|
||
if use_deleted_table:
|
||
total_count = await _deleted_verification_token_table(prisma_client).count(where=where)
|
||
else:
|
||
total_count = await _prisma_table(VerificationTokenRepository(prisma_client)).count(where=where)
|
||
|
||
verbose_proxy_logger.debug("Total count of keys: %s", total_count)
|
||
|
||
# Calculate total pages
|
||
total_pages: Final = -(-total_count // size) # Ceiling division
|
||
|
||
# Fetch user information if expand includes "user"
|
||
user_map = dict[str | None, _UserRowLike]()
|
||
if expand and "user" in expand:
|
||
user_ids: Final = [key.user_id for key in keys if key.user_id]
|
||
created_by_ids: Final = [key.created_by for key in keys if key.created_by]
|
||
all_ids: Final = list(set(user_ids + created_by_ids)) # Remove duplicates
|
||
if all_ids:
|
||
users: Final[Sequence[_UserRowLike]] = await _user_table(prisma_client).find_many(
|
||
where={"user_id": {"in": all_ids}}
|
||
)
|
||
user_map = {user.user_id: user for user in users}
|
||
|
||
# Prepare response
|
||
key_list: Final[list[str | UserAPIKeyAuth | LiteLLM_DeletedVerificationToken]] = []
|
||
for key in keys:
|
||
# Convert Prisma model to dict (supports both Pydantic v1 and v2)
|
||
try:
|
||
key_dict = key.model_dump()
|
||
except Exception:
|
||
# Fallback for Pydantic v1 compatibility
|
||
key_dict = key.dict() # pyright: ignore[reportDeprecated] # deliberate pydantic v1 fallback
|
||
# Attach object_permission if object_permission_id is set (only for non-deleted keys)
|
||
if not use_deleted_table:
|
||
key_dict = await attach_object_permission_to_dict(key_dict, prisma_client)
|
||
|
||
# Include user information if expand includes "user"
|
||
if expand and "user" in expand:
|
||
if key.user_id and key.user_id in user_map:
|
||
try:
|
||
key_dict["user"] = user_map[key.user_id].model_dump()
|
||
except Exception:
|
||
key_dict["user"] = user_map[key.user_id].dict()
|
||
if key.created_by and key.created_by in user_map:
|
||
created_by_user = user_map[key.created_by]
|
||
key_dict["created_by_user"] = {
|
||
"user_id": created_by_user.user_id,
|
||
"user_email": created_by_user.user_email,
|
||
"user_alias": created_by_user.user_alias,
|
||
}
|
||
|
||
if return_full_object is True or (expand and "user" in expand):
|
||
if use_deleted_table:
|
||
# Use deleted key type to preserve deleted_at, deleted_by, etc.
|
||
key_list.append(LiteLLM_DeletedVerificationToken.model_validate(key_dict))
|
||
else:
|
||
key_list.append(
|
||
UserAPIKeyAuth(**key_dict) # pyright: ignore[reportAny] # model_dump() is dict[str, Any]
|
||
)
|
||
else:
|
||
_token = key_dict.get("token")
|
||
key_list.append(cast(str, _token)) # Return only the token
|
||
|
||
return KeyListResponseObject(
|
||
keys=key_list,
|
||
total_count=total_count,
|
||
current_page=page,
|
||
total_pages=total_pages,
|
||
)
|
||
|
||
|
||
def _get_condition_to_filter_out_ui_session_tokens() -> Mapping[str, object]:
|
||
"""
|
||
Condition to filter out UI session tokens
|
||
"""
|
||
return {
|
||
"OR": [
|
||
{"team_id": None}, # Include records where team_id is null
|
||
{"team_id": {"not": UI_SESSION_TOKEN_TEAM_ID}}, # Include records where team_id != UI_SESSION_TOKEN_TEAM_ID
|
||
]
|
||
}
|
||
|
||
|
||
async def _check_key_admin_access(
|
||
user_api_key_dict: UserAPIKeyAuth,
|
||
hashed_token: str | None,
|
||
prisma_client: PrismaClient | None,
|
||
user_api_key_cache: UserApiKeyCache,
|
||
route: str,
|
||
) -> None:
|
||
"""
|
||
Check that the caller has admin privileges for the target key.
|
||
|
||
Allowed callers:
|
||
- Proxy admin
|
||
- Team admin for the key's team
|
||
- Org admin for the key's team's organization
|
||
|
||
Raises HTTPException(403) if the caller is not authorized.
|
||
"""
|
||
|
||
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
|
||
return
|
||
|
||
# Look up the target key to find its team
|
||
target_key_row: Final = await _prisma_table(VerificationTokenRepository(prisma_client)).find_unique(
|
||
where={"token": hashed_token}
|
||
)
|
||
if target_key_row is None:
|
||
raise HTTPException(
|
||
status_code=404,
|
||
detail={"error": f"Key not found: {hashed_token}"},
|
||
)
|
||
|
||
# If the key belongs to a team, check team admin / org admin
|
||
if target_key_row.team_id:
|
||
team_obj: Final = await get_team_object(
|
||
team_id=target_key_row.team_id,
|
||
prisma_client=prisma_client,
|
||
user_api_key_cache=user_api_key_cache,
|
||
check_db_only=True,
|
||
)
|
||
if team_obj is not None:
|
||
if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_obj):
|
||
return
|
||
if await _is_user_org_admin_for_team(user_api_key_dict=user_api_key_dict, team_obj=team_obj):
|
||
return
|
||
|
||
raise HTTPException(
|
||
status_code=403,
|
||
detail={
|
||
"error": f"Only proxy admins, team admins, or org admins can call {route}. "
|
||
f"user_role={user_api_key_dict.user_role}, user_id={user_api_key_dict.user_id}"
|
||
},
|
||
)
|
||
|
||
|
||
@router.post("/key/block", tags=["key management"], dependencies=[Depends(user_api_key_auth)])
|
||
@management_endpoint_wrapper
|
||
async def block_key(
|
||
data: BlockKeyRequest,
|
||
http_request: Request,
|
||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||
litellm_changed_by: str | None = Header(
|
||
None,
|
||
description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability",
|
||
),
|
||
) -> LiteLLM_VerificationToken | None:
|
||
"""
|
||
Block an Virtual key from making any requests.
|
||
|
||
Parameters:
|
||
- key: str - The key to block. Can be either the unhashed key (sk-...) or the hashed key value
|
||
|
||
Example:
|
||
```bash
|
||
curl --location 'http://0.0.0.0:4000/key/block' \
|
||
--header 'Authorization: Bearer sk-1234' \
|
||
--header 'Content-Type: application/json' \
|
||
--data '{
|
||
"key": "sk-Fn8Ej39NxjAXrvpUGKghGw"
|
||
}'
|
||
```
|
||
|
||
Note: This is an admin-only endpoint. Only proxy admins, team admins, or org admins can block keys.
|
||
"""
|
||
from litellm.proxy.management_helpers.audit_logs import (
|
||
get_audit_log_changed_by,
|
||
is_audit_logging_enabled,
|
||
)
|
||
from litellm.proxy.proxy_server import (
|
||
create_audit_log_for_update,
|
||
hash_token,
|
||
litellm_proxy_admin_name,
|
||
prisma_client,
|
||
proxy_logging_obj,
|
||
user_api_key_cache,
|
||
)
|
||
|
||
if prisma_client is None:
|
||
raise Exception(f"{CommonProxyErrors.db_not_connected_error.value}")
|
||
|
||
if not is_valid_api_key(data.key):
|
||
raise ProxyException(
|
||
message="Invalid key format.",
|
||
type=ProxyErrorTypes.bad_request_error,
|
||
param="key",
|
||
code=status.HTTP_400_BAD_REQUEST,
|
||
)
|
||
if data.key.startswith("sk-"):
|
||
hashed_token = hash_token(token=data.key)
|
||
else:
|
||
hashed_token = data.key
|
||
|
||
# Admin-only: only proxy admins, team admins, or org admins can block keys
|
||
await _check_key_admin_access(
|
||
user_api_key_dict=user_api_key_dict,
|
||
hashed_token=hashed_token,
|
||
prisma_client=prisma_client,
|
||
user_api_key_cache=user_api_key_cache,
|
||
route="/key/block",
|
||
)
|
||
|
||
# Check if the key exists before trying to block it
|
||
existing_record: Final = await _prisma_table(VerificationTokenRepository(prisma_client)).find_unique(
|
||
where={"token": hashed_token}
|
||
)
|
||
if existing_record is None:
|
||
raise ProxyException(
|
||
message="Key not found.",
|
||
type=ProxyErrorTypes.not_found_error,
|
||
param="key",
|
||
code=status.HTTP_404_NOT_FOUND,
|
||
)
|
||
|
||
if is_audit_logging_enabled():
|
||
asyncio.create_task(
|
||
create_audit_log_for_update(
|
||
request_data=LiteLLM_AuditLogs(
|
||
id=str(uuid.uuid4()),
|
||
updated_at=datetime.now(timezone.utc),
|
||
changed_by=get_audit_log_changed_by(
|
||
litellm_changed_by=litellm_changed_by,
|
||
user_api_key_dict=user_api_key_dict,
|
||
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||
),
|
||
changed_by_api_key=user_api_key_dict.api_key,
|
||
table_name=LitellmTableNames.KEY_TABLE_NAME,
|
||
object_id=hashed_token,
|
||
action="blocked",
|
||
updated_values="{}",
|
||
before_value=existing_record.model_dump_json(),
|
||
)
|
||
)
|
||
)
|
||
|
||
record: Final = await _prisma_table(VerificationTokenRepository(prisma_client)).update(
|
||
where={"token": hashed_token},
|
||
data=with_settings_updated_at({"blocked": True}),
|
||
)
|
||
|
||
## UPDATE KEY CACHE - invalidate so next read re-fetches from DB
|
||
await _delete_cache_key_object(
|
||
hashed_token=hashed_token,
|
||
user_api_key_cache=user_api_key_cache,
|
||
proxy_logging_obj=proxy_logging_obj,
|
||
)
|
||
|
||
return record
|
||
|
||
|
||
@router.post("/key/unblock", tags=["key management"], dependencies=[Depends(user_api_key_auth)])
|
||
@management_endpoint_wrapper
|
||
async def unblock_key(
|
||
data: BlockKeyRequest,
|
||
http_request: Request,
|
||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||
litellm_changed_by: str | None = Header(
|
||
None,
|
||
description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability",
|
||
),
|
||
):
|
||
"""
|
||
Unblock a Virtual key to allow it to make requests again.
|
||
|
||
Parameters:
|
||
- key: str - The key to unblock. Can be either the unhashed key (sk-...) or the hashed key value
|
||
|
||
Example:
|
||
```bash
|
||
curl --location 'http://0.0.0.0:4000/key/unblock' \
|
||
--header 'Authorization: Bearer sk-1234' \
|
||
--header 'Content-Type: application/json' \
|
||
--data '{
|
||
"key": "sk-Fn8Ej39NxjAXrvpUGKghGw"
|
||
}'
|
||
```
|
||
|
||
Note: This is an admin-only endpoint. Only proxy admins, team admins, or org admins can unblock keys.
|
||
"""
|
||
from litellm.proxy.management_helpers.audit_logs import (
|
||
get_audit_log_changed_by,
|
||
is_audit_logging_enabled,
|
||
)
|
||
from litellm.proxy.proxy_server import (
|
||
create_audit_log_for_update,
|
||
hash_token,
|
||
litellm_proxy_admin_name,
|
||
prisma_client,
|
||
proxy_logging_obj,
|
||
user_api_key_cache,
|
||
)
|
||
|
||
if prisma_client is None:
|
||
raise Exception(f"{CommonProxyErrors.db_not_connected_error.value}")
|
||
|
||
if not is_valid_api_key(data.key):
|
||
raise ProxyException(
|
||
message="Invalid key format.",
|
||
type=ProxyErrorTypes.bad_request_error,
|
||
param="key",
|
||
code=status.HTTP_400_BAD_REQUEST,
|
||
)
|
||
if data.key.startswith("sk-"):
|
||
hashed_token = hash_token(token=data.key)
|
||
else:
|
||
hashed_token = data.key
|
||
|
||
# Admin-only: only proxy admins, team admins, or org admins can unblock keys
|
||
await _check_key_admin_access(
|
||
user_api_key_dict=user_api_key_dict,
|
||
hashed_token=hashed_token,
|
||
prisma_client=prisma_client,
|
||
user_api_key_cache=user_api_key_cache,
|
||
route="/key/unblock",
|
||
)
|
||
|
||
# Check if the key exists before trying to unblock it
|
||
existing_record: Final = await _prisma_table(VerificationTokenRepository(prisma_client)).find_unique(
|
||
where={"token": hashed_token}
|
||
)
|
||
if existing_record is None:
|
||
raise ProxyException(
|
||
message="Key not found.",
|
||
type=ProxyErrorTypes.not_found_error,
|
||
param="key",
|
||
code=status.HTTP_404_NOT_FOUND,
|
||
)
|
||
|
||
if is_audit_logging_enabled():
|
||
asyncio.create_task(
|
||
create_audit_log_for_update(
|
||
request_data=LiteLLM_AuditLogs(
|
||
id=str(uuid.uuid4()),
|
||
updated_at=datetime.now(timezone.utc),
|
||
changed_by=get_audit_log_changed_by(
|
||
litellm_changed_by=litellm_changed_by,
|
||
user_api_key_dict=user_api_key_dict,
|
||
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||
),
|
||
changed_by_api_key=user_api_key_dict.api_key,
|
||
table_name=LitellmTableNames.KEY_TABLE_NAME,
|
||
object_id=hashed_token,
|
||
action="unblocked",
|
||
updated_values="{}",
|
||
before_value=existing_record.model_dump_json(),
|
||
)
|
||
)
|
||
)
|
||
|
||
record: Final = await _prisma_table(VerificationTokenRepository(prisma_client)).update(
|
||
where={"token": hashed_token},
|
||
data=with_settings_updated_at({"blocked": False}),
|
||
)
|
||
|
||
## UPDATE KEY CACHE - invalidate so next read re-fetches from DB
|
||
await _delete_cache_key_object(
|
||
hashed_token=hashed_token,
|
||
user_api_key_cache=user_api_key_cache,
|
||
proxy_logging_obj=proxy_logging_obj,
|
||
)
|
||
|
||
return record
|
||
|
||
|
||
@router.post(
|
||
"/key/health",
|
||
tags=["key management"],
|
||
dependencies=[Depends(user_api_key_auth)],
|
||
response_model=KeyHealthResponse,
|
||
)
|
||
@management_endpoint_wrapper
|
||
async def key_health(
|
||
request: Request,
|
||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||
):
|
||
"""
|
||
Check the health of the key
|
||
|
||
Checks:
|
||
- If key based logging is configured correctly - sends a test log
|
||
|
||
Usage
|
||
|
||
Pass the key in the request header
|
||
|
||
```bash
|
||
curl -X POST "http://localhost:4000/key/health" \
|
||
-H "Authorization: Bearer sk-1234" \
|
||
-H "Content-Type: application/json"
|
||
```
|
||
|
||
Response when logging callbacks are setup correctly:
|
||
|
||
```json
|
||
{
|
||
"key": "healthy",
|
||
"logging_callbacks": {
|
||
"callbacks": [
|
||
"gcs_bucket"
|
||
],
|
||
"status": "healthy",
|
||
"details": "No logger exceptions triggered, system is healthy. Manually check if logs were sent to ['gcs_bucket']"
|
||
}
|
||
}
|
||
```
|
||
|
||
|
||
Response when logging callbacks are not setup correctly:
|
||
```json
|
||
{
|
||
"key": "unhealthy",
|
||
"logging_callbacks": {
|
||
"callbacks": [
|
||
"gcs_bucket"
|
||
],
|
||
"status": "unhealthy",
|
||
"details": "Logger exceptions triggered, system is unhealthy: Failed to load vertex credentials. Check to see if credentials containing partial/invalid information."
|
||
}
|
||
}
|
||
```
|
||
"""
|
||
try:
|
||
# Get the key's metadata
|
||
key_metadata: Final = user_api_key_dict.metadata
|
||
|
||
health_status: Final[KeyHealthResponse] = KeyHealthResponse(
|
||
key="healthy",
|
||
logging_callbacks=None,
|
||
)
|
||
|
||
# Check if logging is configured in metadata
|
||
if key_metadata and "logging" in key_metadata:
|
||
logging_statuses: Final = await test_key_logging(
|
||
user_api_key_dict=user_api_key_dict,
|
||
request=request,
|
||
key_logging=decrypt_callback_vars(key_metadata)["logging"],
|
||
)
|
||
health_status["logging_callbacks"] = logging_statuses
|
||
|
||
# Check if any logging callback is unhealthy
|
||
if logging_statuses.get("status") == "unhealthy":
|
||
health_status["key"] = "unhealthy"
|
||
|
||
return KeyHealthResponse(**health_status)
|
||
|
||
except Exception as e:
|
||
raise ProxyException(
|
||
message=f"Key health check failed: {e}",
|
||
type=ProxyErrorTypes.internal_server_error,
|
||
param=getattr(e, "param", "None"),
|
||
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||
)
|
||
|
||
|
||
async def _can_user_query_key_info(
|
||
user_api_key_dict: UserAPIKeyAuth,
|
||
key: str | None,
|
||
key_info: LiteLLM_VerificationToken,
|
||
) -> bool:
|
||
"""
|
||
Helper to check if the user has access to the key's info
|
||
"""
|
||
if (
|
||
(
|
||
user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
|
||
or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value
|
||
)
|
||
or user_api_key_dict.api_key == key
|
||
or key_info.user_id == user_api_key_dict.user_id
|
||
or await TeamMemberPermissionChecks.user_belongs_to_keys_team(
|
||
user_api_key_dict=user_api_key_dict,
|
||
existing_key_row=key_info,
|
||
)
|
||
):
|
||
return True
|
||
return False
|
||
|
||
|
||
async def test_key_logging(
|
||
user_api_key_dict: UserAPIKeyAuth,
|
||
request: Request,
|
||
key_logging: Sequence[Mapping[str, str]],
|
||
) -> LoggingCallbackStatus:
|
||
"""
|
||
Test the key-based logging
|
||
|
||
- Test that key logging is correctly formatted and all args are passed correctly
|
||
- Make a mock completion call -> user can check if it's correctly logged
|
||
- Check if any logger.exceptions were triggered -> if they were then returns it to the user client side
|
||
"""
|
||
import logging
|
||
from io import StringIO
|
||
|
||
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
|
||
from litellm.proxy.proxy_server import general_settings, proxy_config
|
||
|
||
logging_callbacks: Final[list[str]] = []
|
||
for callback in key_logging:
|
||
if callback.get("callback_name") is not None:
|
||
logging_callbacks.append(callback["callback_name"])
|
||
else:
|
||
raise ValueError("callback_name is required in key_logging")
|
||
|
||
log_capture_string: Final = StringIO()
|
||
ch: Final = logging.StreamHandler(log_capture_string)
|
||
ch.setLevel(logging.ERROR)
|
||
logger: Final = logging.getLogger()
|
||
logger.addHandler(ch)
|
||
|
||
try:
|
||
data = {
|
||
"model": "openai/litellm-key-health-test",
|
||
"messages": [
|
||
{
|
||
"role": "user",
|
||
"content": "Hello, this is a test from litellm /key/health. No LLM API call was made for this",
|
||
}
|
||
],
|
||
}
|
||
data = await add_litellm_data_to_request(
|
||
data=data,
|
||
user_api_key_dict=user_api_key_dict,
|
||
proxy_config=proxy_config,
|
||
general_settings=general_settings,
|
||
request=request,
|
||
)
|
||
data["mock_response"] = "test response"
|
||
await litellm.acompletion(**data) # make mock completion call to trigger key based callbacks
|
||
except Exception as e:
|
||
return LoggingCallbackStatus(
|
||
callbacks=logging_callbacks,
|
||
status="unhealthy",
|
||
details=f"Logging test failed: {e}",
|
||
)
|
||
|
||
await asyncio.sleep(2) # wait for callbacks to run, callbacks use batching so wait for the flush event
|
||
|
||
# Check if any logger exceptions were triggered
|
||
log_contents: Final = log_capture_string.getvalue()
|
||
logger.removeHandler(ch)
|
||
if log_contents:
|
||
return LoggingCallbackStatus(
|
||
callbacks=logging_callbacks,
|
||
status="unhealthy",
|
||
details=f"Logger exceptions triggered, system is unhealthy: {log_contents}",
|
||
)
|
||
else:
|
||
return LoggingCallbackStatus(
|
||
callbacks=logging_callbacks,
|
||
status="healthy",
|
||
details=f"No logger exceptions triggered, system is healthy. Manually check if logs were sent to {logging_callbacks} ",
|
||
)
|
||
|
||
|
||
_KEY_ALIAS_PATTERN: Final = re.compile(r"^[a-zA-Z0-9][a-zA-Z0-9_\-/\.@]{0,253}[a-zA-Z0-9]$")
|
||
|
||
|
||
def _validate_key_alias_format(key_alias: str | None) -> None:
|
||
"""
|
||
Validate the format of the key_alias.
|
||
|
||
A baseline validation always runs, regardless of
|
||
``litellm.enable_key_alias_format_validation``.
|
||
|
||
The remaining charset/length rules are gated behind
|
||
``litellm.enable_key_alias_format_validation`` (default **False**). When disabled,
|
||
only the baseline validation above is performed, so existing workflows are not
|
||
broken.
|
||
|
||
Rules (when enabled):
|
||
- None is OK (no alias).
|
||
- Otherwise must be 2–255 chars
|
||
- start/end with alphanumeric
|
||
- only allow a-zA-Z0-9_-/.@
|
||
"""
|
||
if key_alias is None:
|
||
return
|
||
|
||
try:
|
||
raise_if_unsafe_secret_name(key_alias)
|
||
except ValueError:
|
||
raise ProxyException(
|
||
message="Invalid key_alias",
|
||
type=ProxyErrorTypes.bad_request_error,
|
||
param="key_alias",
|
||
code=400,
|
||
)
|
||
|
||
if not litellm.enable_key_alias_format_validation:
|
||
return
|
||
|
||
if not _KEY_ALIAS_PATTERN.match(key_alias):
|
||
raise ProxyException(
|
||
message="Invalid key_alias format. Must be 2-255 characters, start/end with alphanumeric, and only contain a-zA-Z0-9_-/.@.",
|
||
type=ProxyErrorTypes.bad_request_error,
|
||
param="key_alias",
|
||
code=400,
|
||
)
|
||
|
||
|
||
async def _enforce_unique_key_alias(
|
||
key_alias: str | None,
|
||
prisma_client: PrismaClient | None,
|
||
existing_key_token: str | None = None,
|
||
) -> None:
|
||
"""
|
||
Helper to enforce unique key aliases across all keys.
|
||
|
||
Args:
|
||
key_alias (Optional[str]): The key alias to check
|
||
prisma_client (Any): Prisma client instance
|
||
existing_key_token (Optional[str]): ID of existing key being updated, to exclude from uniqueness check
|
||
(The Admin UI passes key_alias, in all Edit key requests. So we need to be sure that if we find a key with the same alias, it's not the same key we're updating)
|
||
|
||
Raises:
|
||
ProxyException: If key alias already exists on a different key
|
||
"""
|
||
if key_alias is not None and prisma_client is not None:
|
||
where_clause: Final[dict[str, object]] = {"key_alias": key_alias}
|
||
if existing_key_token:
|
||
# Exclude the current key from the uniqueness check
|
||
where_clause["NOT"] = {"token": existing_key_token}
|
||
|
||
existing_key = await _prisma_table(VerificationTokenRepository(prisma_client)).find_first(where=where_clause)
|
||
if existing_key is not None:
|
||
raise ProxyException(
|
||
message=f"Key with alias '{key_alias}' already exists. Unique key aliases across all keys are required.",
|
||
type=ProxyErrorTypes.bad_request_error,
|
||
param="key_alias",
|
||
code=status.HTTP_400_BAD_REQUEST,
|
||
)
|
||
|
||
|
||
def validate_model_max_budget(model_max_budget: dict | None) -> None:
|
||
"""
|
||
Validate the model_max_budget is GenericBudgetConfigType + enforce user has an enterprise license
|
||
|
||
Raises:
|
||
Exception: If model_max_budget is not a valid GenericBudgetConfigType
|
||
"""
|
||
try:
|
||
if model_max_budget is None:
|
||
return
|
||
if len(model_max_budget) == 0:
|
||
return
|
||
if model_max_budget is not None:
|
||
from litellm.proxy.proxy_server import CommonProxyErrors, premium_user
|
||
|
||
if premium_user is not True:
|
||
raise ValueError(
|
||
f"You must have an enterprise license to set model_max_budget. {CommonProxyErrors.not_premium_user.value}"
|
||
)
|
||
for _model, _budget_info in model_max_budget.items():
|
||
assert isinstance(_model, str)
|
||
|
||
# Normalize to dict (Pydantic may already parse nested values as BudgetConfig)
|
||
_info = _budget_info.model_dump() if hasattr(_budget_info, "model_dump") else dict(_budget_info)
|
||
# /CRUD endpoints can pass budget_limit as a string, so we need to convert it to a float
|
||
if "budget_limit" in _info:
|
||
_info["budget_limit"] = float(_info["budget_limit"])
|
||
BudgetConfig(**_info)
|
||
except Exception as e:
|
||
raise ValueError(
|
||
f"Invalid model_max_budget: {e}. Example of valid model_max_budget: https://docs.litellm.ai/docs/proxy/users"
|
||
)
|