From 5db942a39e013d89958d0be5941fbc6e28971afd Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Wed, 13 Nov 2024 11:36:35 -0800
Subject: [PATCH] fix KeyManagementSettings

---
 litellm/proxy/_types.py                       |  6 ++-
 .../proxy/hooks/key_management_event_hooks.py | 50 +++++++++++--------
 litellm/proxy/proxy_config.yaml               |  7 +--
 3 files changed, 38 insertions(+), 25 deletions(-)

diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index 2d869af85..f2feee8b4 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -1128,7 +1128,11 @@ class KeyManagementSystem(enum.Enum):
 
 
 class KeyManagementSettings(LiteLLMBase):
-    hosted_keys: List
+    hosted_keys: Optional[List] = None
+    store_virtual_keys: Optional[bool] = False
+    """
+    If True, virtual keys created by litellm will be stored in the secret manager
+    """
 
 
 class TeamDefaultSettings(LiteLLMBase):
diff --git a/litellm/proxy/hooks/key_management_event_hooks.py b/litellm/proxy/hooks/key_management_event_hooks.py
index 292c53b33..9ef135df8 100644
--- a/litellm/proxy/hooks/key_management_event_hooks.py
+++ b/litellm/proxy/hooks/key_management_event_hooks.py
@@ -72,28 +72,10 @@ class KeyManagementEventHooks:
                 )
             )
         # store the generated key in the secret manager
-        _key_management_settings: Optional[dict] = general_settings.get(
-            "key_management", None
+        await KeyManagementEventHooks._store_virtual_key_in_secret_manager(
+            secret_name=data.key_alias or f"virtual-key-{uuid.uuid4()}",
+            secret_token=response.get("token", ""),
         )
-        if _key_management_settings is not None and isinstance(
-            _key_management_settings, dict
-        ):
-            if _key_management_settings.get("store_virtual_keys", None) is True:
-                from litellm.secret_managers.aws_secret_manager_v2 import (
-                    AWSSecretsManagerV2,
-                )
-
-                # store the key in the secret manager
-                if (
-                    litellm._key_management_system
-                    == KeyManagementSystem.AWS_SECRET_MANAGER
-                    and isinstance(litellm.secret_manager_client, AWSSecretsManagerV2)
-                ):
-                    _secret_name = data.key_alias or f"virtual-key-{uuid.uuid4()}"
-                    await litellm.secret_manager_client.async_write_secret(
-                        secret_name=_secret_name,
-                        secret_value=response.get("token_value", ""),
-                    )
 
     @staticmethod
     async def async_key_updated_hook(
@@ -197,6 +179,32 @@ class KeyManagementEventHooks:
                 )
         pass
 
+    @staticmethod
+    async def _store_virtual_key_in_secret_manager(secret_name: str, secret_token: str):
+        """
+        Store a virtual key in the secret manager
+
+        Args:
+            secret_name: Name of the virtual key
+            secret_token: Value of the virtual key (example: sk-1234)
+        """
+        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,
+                )
+
+                # store the key in the secret manager
+                if (
+                    litellm._key_management_system
+                    == KeyManagementSystem.AWS_SECRET_MANAGER
+                    and isinstance(litellm.secret_manager_client, AWSSecretsManagerV2)
+                ):
+                    await litellm.secret_manager_client.async_write_secret(
+                        secret_name=secret_name,
+                        secret_value=secret_token,
+                    )
+
     @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/proxy_config.yaml b/litellm/proxy/proxy_config.yaml
index 29d14c910..d36ef5df0 100644
--- a/litellm/proxy/proxy_config.yaml
+++ b/litellm/proxy/proxy_config.yaml
@@ -7,6 +7,7 @@ model_list:
 
 
 
-litellm_settings:
-  callbacks: ["gcs_bucket"]
-
+general_settings:
+  key_management_system: "aws_secret_manager" 
+  key_management_settings: 
+    store_virtual_keys: true
\ No newline at end of file