diff --git a/litellm/proxy/hooks/key_management_event_hooks.py b/litellm/proxy/hooks/key_management_event_hooks.py index 9ef135df8..08645a468 100644 --- a/litellm/proxy/hooks/key_management_event_hooks.py +++ b/litellm/proxy/hooks/key_management_event_hooks.py @@ -3,16 +3,18 @@ import json import uuid from datetime import datetime, timezone from re import A -from typing import Any, Optional +from typing import Any, List, Optional from fastapi import status import litellm +from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( GenerateKeyRequest, KeyManagementSystem, KeyRequest, LiteLLM_AuditLogs, + LiteLLM_VerificationToken, LitellmTableNames, ProxyErrorTypes, ProxyException, @@ -125,6 +127,7 @@ class KeyManagementEventHooks: @staticmethod async def async_key_deleted_hook( data: KeyRequest, + keys_being_deleted: List[LiteLLM_VerificationToken], response: dict, user_api_key_dict: UserAPIKeyAuth, litellm_changed_by: Optional[str] = None, @@ -177,6 +180,10 @@ class KeyManagementEventHooks: ) ) ) + # delete the keys from the secret manager + await KeyManagementEventHooks._delete_virtual_keys_from_secret_manager( + keys_being_deleted=keys_being_deleted + ) pass @staticmethod @@ -205,6 +212,33 @@ class KeyManagementEventHooks: secret_value=secret_token, ) + @staticmethod + async def _delete_virtual_keys_from_secret_manager( + keys_being_deleted: List[LiteLLM_VerificationToken], + ): + """ + Deletes virtual keys from the secret manager + + Args: + keys_being_deleted: List of keys being deleted, this is passed down from the /key/delete operation + """ + if litellm._key_management_settings is not None: + if litellm._key_management_settings.store_virtual_keys is True: + from litellm.secret_managers.aws_secret_manager_v2 import ( + AWSSecretsManagerV2, + ) + + if isinstance(litellm.secret_manager_client, AWSSecretsManagerV2): + for key in keys_being_deleted: + if key.key_alias is not None: + await litellm.secret_manager_client.async_delete_secret( + secret_name=key.key_alias + ) + else: + verbose_proxy_logger.warning( + f"KeyManagementEventHooks._delete_virtual_key_from_secret_manager: Key alias not found for key {key.token}. Skipping deletion from secret manager." + ) + @staticmethod async def _send_key_created_email(response: dict): from litellm.proxy.proxy_server import general_settings, proxy_logging_obj diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 230e0330c..d88c1d642 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -17,7 +17,7 @@ import secrets import traceback import uuid from datetime import datetime, timedelta, timezone -from typing import List, Optional +from typing import List, Optional, Tuple import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, status @@ -450,6 +450,9 @@ async def delete_key_fn( user_custom_key_generate, ) + if prisma_client is None: + raise Exception("Not connected to DB!") + keys = data.keys if len(keys) == 0: raise ProxyException( @@ -469,7 +472,8 @@ async def delete_key_fn( and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN ): user_id = None # unless they're admin - number_deleted_keys = await delete_verification_token( + + number_deleted_keys, _keys_being_deleted = await delete_verification_token( tokens=keys, user_id=user_id ) if number_deleted_keys is None: @@ -506,6 +510,7 @@ async def delete_key_fn( 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, @@ -950,11 +955,35 @@ async def generate_key_helper_fn( # noqa: PLR0915 return key_data -async def delete_verification_token(tokens: List, user_id: Optional[str] = None): +async def delete_verification_token( + tokens: List, user_id: Optional[str] = None +) -> Tuple[Optional[Dict], List[LiteLLM_VerificationToken]]: + """ + Helper that deletes the list of tokens from the database + + 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 litellm_proxy_admin_name, prisma_client try: if prisma_client: + tokens = [_hash_token_if_needed(token=key) for key in tokens] + _keys_being_deleted = ( + await prisma_client.db.litellm_verificationtoken.find_many( + where={"token": {"in": tokens}} + ) + ) + # Assuming 'db' is your Prisma Client instance # check if admin making request - don't filter by user-id if user_id == litellm_proxy_admin_name: @@ -984,7 +1013,7 @@ async def delete_verification_token(tokens: List, user_id: Optional[str] = None) ) verbose_proxy_logger.debug(traceback.format_exc()) raise e - return deleted_tokens + return deleted_tokens, _keys_being_deleted @router.post(