From 34a9833b85076b8fb2618f6d668317dfb49ceb12 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Sat, 4 Jan 2025 19:41:48 -0800 Subject: [PATCH] Support deleting keys by key_alias (#7552) * feat(key_management_endpoints.py): allow deleting keys based on key alias easier for proxy admin to delete known bad key * fix(key_management_event_hooks.py): fix linting error * docs(key_management_endpoints.py): document new key_aliases param * fix(key_management_endpoints.py): return deleted keys to user fixes return when passing key aliases --- litellm/proxy/_new_secret_config.yaml | 1 - litellm/proxy/_types.py | 12 ++- .../proxy/hooks/key_management_event_hooks.py | 2 +- .../key_management_endpoints.py | 83 ++++++++++++++----- 4 files changed, 72 insertions(+), 26 deletions(-) diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index dcabaac30c..191a1bbee0 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -15,6 +15,5 @@ model_list: tpm: 1000000 prompt_id: "jokes" - # litellm_settings: # callbacks: ["otel"] \ No newline at end of file diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 3a6ad4a922..b9697202d9 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -690,7 +690,17 @@ class RegenerateKeyRequest(GenerateKeyRequest): class KeyRequest(LiteLLMPydanticObjectBase): - keys: List[str] + keys: Optional[List[str]] = None + key_aliases: Optional[List[str]] = None + + @model_validator(mode="before") + @classmethod + def validate_at_least_one(cls, values): + if not values.get("keys") and not values.get("key_aliases"): + raise ValueError( + "At least one of 'keys' or 'key_aliases' must be provided." + ) + return values class LiteLLM_ModelTable(LiteLLMPydanticObjectBase): diff --git a/litellm/proxy/hooks/key_management_event_hooks.py b/litellm/proxy/hooks/key_management_event_hooks.py index c123f6df4b..2243488768 100644 --- a/litellm/proxy/hooks/key_management_event_hooks.py +++ b/litellm/proxy/hooks/key_management_event_hooks.py @@ -142,7 +142,7 @@ class KeyManagementEventHooks: # Enterprise Feature - Audit Logging. Enable with litellm.store_audit_logs = True # we do this after the first for loop, since first for loop is for validation. we only want this inserted after validation passes - if litellm.store_audit_logs is True: + if litellm.store_audit_logs is True and data.keys is not None: # make an audit log for each team deleted for key in data.keys: key_row = await prisma_client.get_data( # type: ignore diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 6a94d5bd4f..417e50dec3 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -23,6 +23,7 @@ from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, s import litellm from litellm._logging import verbose_proxy_logger +from litellm.caching import DualCache from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import ( _cache_key_object, @@ -34,6 +35,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks from litellm.proxy.management_helpers.utils import management_endpoint_wrapper from litellm.proxy.utils import ( + PrismaClient, _hash_token_if_needed, duration_in_seconds, handle_exception_on_proxy, @@ -772,6 +774,7 @@ async def delete_key_fn( 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"]} @@ -795,15 +798,6 @@ async def delete_key_fn( if prisma_client is None: raise Exception("Not connected to DB!") - keys = data.keys - if len(keys) == 0: - raise ProxyException( - message=f"No keys provided, passed in: keys={keys}", - type=ProxyErrorTypes.auth_error, - param="keys", - code=status.HTTP_400_BAD_REQUEST, - ) - ## only allow user to delete keys they own user_id = user_api_key_dict.user_id verbose_proxy_logger.debug( @@ -815,9 +809,28 @@ async def delete_key_fn( ): user_id = None # unless they're admin - number_deleted_keys, _keys_being_deleted = await delete_verification_token( - tokens=keys, user_id=user_id - ) + num_keys_to_be_deleted = 0 + deleted_keys = [] + if data.keys: + number_deleted_keys, _keys_being_deleted = await delete_verification_token( + tokens=data.keys, + user_api_key_cache=user_api_key_cache, + user_id=user_id, + ) + 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_id=user_id, + ) + 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", @@ -830,21 +843,15 @@ async def delete_key_fn( ) try: - assert len(keys) == number_deleted_keys["deleted_keys"] + assert num_keys_to_be_deleted == number_deleted_keys["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={len(keys)}, Deleted keys ={number_deleted_keys['deleted_keys']}" + "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['deleted_keys']}" }, ) - for key in keys: - user_api_key_cache.delete_cache(key) - # remove hash token from cache - hashed_token = hash_token(key) - user_api_key_cache.delete_cache(hashed_token) - verbose_proxy_logger.debug( f"/keys/delete - cache after delete: {user_api_key_cache.in_memory_cache.cache_dict}" ) @@ -859,8 +866,13 @@ async def delete_key_fn( ) ) - return {"deleted_keys": keys} + return {"deleted_keys": deleted_keys} except Exception as e: + verbose_proxy_logger.exception( + "litellm.proxy.proxy_server.delete_key_fn(): Exception occured - {}".format( + str(e) + ) + ) raise handle_exception_on_proxy(e) @@ -1277,7 +1289,9 @@ async def generate_key_helper_fn( # noqa: PLR0915 async def delete_verification_token( - tokens: List, user_id: Optional[str] = None + tokens: List, + user_api_key_cache: DualCache, + user_id: Optional[str] = None, ) -> Tuple[Optional[Dict], List[LiteLLM_VerificationToken]]: """ Helper that deletes the list of tokens from the database @@ -1327,16 +1341,39 @@ async def delete_verification_token( else: raise Exception("DB not connected. prisma_client is None") except Exception as e: - verbose_proxy_logger.error( + verbose_proxy_logger.exception( "litellm.proxy.proxy_server.delete_verification_token(): Exception occured - {}".format( str(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(key) + user_api_key_cache.delete_cache(hashed_token) + return deleted_tokens, _keys_being_deleted +async def delete_key_aliases( + key_aliases: List[str], + user_api_key_cache: DualCache, + prisma_client: PrismaClient, + user_id: Optional[str] = None, +) -> Tuple[Optional[Dict], List[LiteLLM_VerificationToken]]: + _keys_being_deleted = await prisma_client.db.litellm_verificationtoken.find_many( + where={"key_alias": {"in": key_aliases}} + ) + + tokens = [key.token for key in _keys_being_deleted] + return await delete_verification_token( + tokens=tokens, user_api_key_cache=user_api_key_cache, user_id=user_id + ) + + @router.post( "/key/{key:path}/regenerate", tags=["key management"],