From 41aade2cc00b49272ef57c9bc557229e6f2e6589 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 15 Nov 2024 18:07:43 -0800 Subject: [PATCH] (feat) Use `litellm/` prefix when storing virtual keys in AWS secret manager (#6765) * fix - storing AWS keys in secret manager * fix test_key_generate_with_secret_manager_call * allow using prefix_for_stored_virtual_keys * add prefix_for_stored_virtual_keys * test_key_generate_with_secret_manager_call --- docs/my-website/docs/secret.md | 10 +++++++++- litellm/proxy/_types.py | 5 +++++ litellm/proxy/hooks/key_management_event_hooks.py | 7 +++++-- tests/proxy_unit_tests/test_key_generate_prisma.py | 11 +++++++++-- 4 files changed, 28 insertions(+), 5 deletions(-) diff --git a/docs/my-website/docs/secret.md b/docs/my-website/docs/secret.md index 15480ea3d..113a11750 100644 --- a/docs/my-website/docs/secret.md +++ b/docs/my-website/docs/secret.md @@ -85,7 +85,8 @@ This will only store virtual keys in AWS Secret Manager. No keys will be read fr general_settings: key_management_system: "aws_secret_manager" # 👈 KEY CHANGE key_management_settings: - store_virtual_keys: true + store_virtual_keys: true # OPTIONAL. Defaults to False, when True will store virtual keys in secret manager + prefix_for_stored_virtual_keys: "litellm/" # OPTIONAL. If set, this prefix will be used for stored virtual keys in the secret manager access_mode: "write_only" # Literal["read_only", "write_only", "read_and_write"] ``` @@ -247,7 +248,14 @@ All settings related to secret management general_settings: key_management_system: "aws_secret_manager" # REQUIRED key_management_settings: + + # Storing Virtual Keys Settings store_virtual_keys: true # OPTIONAL. Defaults to False, when True will store virtual keys in secret manager + prefix_for_stored_virtual_keys: "litellm/" # OPTIONAL.I f set, this prefix will be used for stored virtual keys in the secret manager + + # Access Mode Settings access_mode: "write_only" # OPTIONAL. Literal["read_only", "write_only", "read_and_write"]. Defaults to "read_only" + + # Hosted Keys Settings hosted_keys: ["litellm_master_key"] # OPTIONAL. Specify which env keys you stored on AWS ``` \ No newline at end of file diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 70e5e6756..d9efa6f9a 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1132,6 +1132,11 @@ class KeyManagementSettings(LiteLLMBase): If True, virtual keys created by litellm will be stored in the secret manager """ + prefix_for_stored_virtual_keys: str = "litellm/" + """ + If set, this prefix will be used for stored virtual keys in the secret manager + """ + access_mode: Literal["read_only", "write_only", "read_and_write"] = "read_only" """ Access mode for the secret manager, when write_only will only use for writing secrets diff --git a/litellm/proxy/hooks/key_management_event_hooks.py b/litellm/proxy/hooks/key_management_event_hooks.py index 08645a468..bdecc77b0 100644 --- a/litellm/proxy/hooks/key_management_event_hooks.py +++ b/litellm/proxy/hooks/key_management_event_hooks.py @@ -23,6 +23,9 @@ from litellm.proxy._types import ( WebhookEvent, ) +# NOTE: This is the prefix for all virtual keys stored in AWS Secrets Manager +LITELLM_PREFIX_STORED_VIRTUAL_KEYS = "litellm/" + class KeyManagementEventHooks: @@ -208,7 +211,7 @@ class KeyManagementEventHooks: and isinstance(litellm.secret_manager_client, AWSSecretsManagerV2) ): await litellm.secret_manager_client.async_write_secret( - secret_name=secret_name, + secret_name=f"{litellm._key_management_settings.prefix_for_stored_virtual_keys}/{secret_name}", secret_value=secret_token, ) @@ -232,7 +235,7 @@ class KeyManagementEventHooks: 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 + secret_name=f"{litellm._key_management_settings.prefix_for_stored_virtual_keys}/{key.key_alias}" ) else: verbose_proxy_logger.warning( diff --git a/tests/proxy_unit_tests/test_key_generate_prisma.py b/tests/proxy_unit_tests/test_key_generate_prisma.py index b97ab3514..fb6e2c7f5 100644 --- a/tests/proxy_unit_tests/test_key_generate_prisma.py +++ b/tests/proxy_unit_tests/test_key_generate_prisma.py @@ -3467,6 +3467,9 @@ async def test_key_generate_with_secret_manager_call(prisma_client): """ from litellm.secret_managers.aws_secret_manager_v2 import AWSSecretsManagerV2 from litellm.proxy._types import KeyManagementSystem, KeyManagementSettings + from litellm.proxy.hooks.key_management_event_hooks import ( + LITELLM_PREFIX_STORED_VIRTUAL_KEYS, + ) litellm.set_verbose = True @@ -3512,7 +3515,9 @@ async def test_key_generate_with_secret_manager_call(prisma_client): await asyncio.sleep(2) # read from the secret manager - result = await aws_secret_manager_client.async_read_secret(secret_name=key_alias) + result = await aws_secret_manager_client.async_read_secret( + secret_name=f"{litellm._key_management_settings.prefix_for_stored_virtual_keys}/{key_alias}" + ) # Assert the correct key is stored in the secret manager print("response from AWS Secret Manager") @@ -3530,7 +3535,9 @@ async def test_key_generate_with_secret_manager_call(prisma_client): await asyncio.sleep(2) # Assert the key is deleted from the secret manager - result = await aws_secret_manager_client.async_read_secret(secret_name=key_alias) + result = await aws_secret_manager_client.async_read_secret( + secret_name=f"{litellm._key_management_settings.prefix_for_stored_virtual_keys}/{key_alias}" + ) assert result is None # cleanup