(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
This commit is contained in:
Ishaan Jaff 2024-11-15 18:07:43 -08:00 committed by GitHub
parent 6d2ee70a9a
commit 41aade2cc0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 28 additions and 5 deletions

View file

@ -85,7 +85,8 @@ This will only store virtual keys in AWS Secret Manager. No keys will be read fr
general_settings: general_settings:
key_management_system: "aws_secret_manager" # 👈 KEY CHANGE key_management_system: "aws_secret_manager" # 👈 KEY CHANGE
key_management_settings: 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"] access_mode: "write_only" # Literal["read_only", "write_only", "read_and_write"]
``` ```
</TabItem> </TabItem>
@ -247,7 +248,14 @@ All settings related to secret management
general_settings: general_settings:
key_management_system: "aws_secret_manager" # REQUIRED key_management_system: "aws_secret_manager" # REQUIRED
key_management_settings: key_management_settings:
# Storing Virtual Keys Settings
store_virtual_keys: true # OPTIONAL. Defaults to False, when True will store virtual keys in secret manager 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" 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 hosted_keys: ["litellm_master_key"] # OPTIONAL. Specify which env keys you stored on AWS
``` ```

View file

@ -1132,6 +1132,11 @@ class KeyManagementSettings(LiteLLMBase):
If True, virtual keys created by litellm will be stored in the secret manager 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: 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 Access mode for the secret manager, when write_only will only use for writing secrets

View file

@ -23,6 +23,9 @@ from litellm.proxy._types import (
WebhookEvent, WebhookEvent,
) )
# NOTE: This is the prefix for all virtual keys stored in AWS Secrets Manager
LITELLM_PREFIX_STORED_VIRTUAL_KEYS = "litellm/"
class KeyManagementEventHooks: class KeyManagementEventHooks:
@ -208,7 +211,7 @@ class KeyManagementEventHooks:
and isinstance(litellm.secret_manager_client, AWSSecretsManagerV2) and isinstance(litellm.secret_manager_client, AWSSecretsManagerV2)
): ):
await litellm.secret_manager_client.async_write_secret( 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, secret_value=secret_token,
) )
@ -232,7 +235,7 @@ class KeyManagementEventHooks:
for key in keys_being_deleted: for key in keys_being_deleted:
if key.key_alias is not None: if key.key_alias is not None:
await litellm.secret_manager_client.async_delete_secret( 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: else:
verbose_proxy_logger.warning( verbose_proxy_logger.warning(

View file

@ -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.secret_managers.aws_secret_manager_v2 import AWSSecretsManagerV2
from litellm.proxy._types import KeyManagementSystem, KeyManagementSettings 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 litellm.set_verbose = True
@ -3512,7 +3515,9 @@ async def test_key_generate_with_secret_manager_call(prisma_client):
await asyncio.sleep(2) await asyncio.sleep(2)
# read from the secret manager # 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 # Assert the correct key is stored in the secret manager
print("response from AWS 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) await asyncio.sleep(2)
# Assert the key is deleted from the secret manager # 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 assert result is None
# cleanup # cleanup