phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

(feat) Use litellm/ prefix when storing virtual keys in AWS secret manager (#6765)

Browse source 
* 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
Download patch file Download diff file Expand all files Collapse all files

11
tests/proxy_unit_tests/test_key_generate_prisma.py
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.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