forked from phoenix/litellm-mirror
[Fix-Proxy] Azure Key Management - Secret Manager (#5728)
* fix azure key mgtm error * add test for azure kms * add test for azure kms
This commit is contained in:
parent
ca6d99e1ad
commit
8103e2b2da
3 changed files with 42 additions and 10 deletions
|
@ -20,8 +20,8 @@ model_list:
|
||||||
api_key: fake-key
|
api_key: fake-key
|
||||||
api_base: https://exampleopenaiendpoint-production.up.railway.app
|
api_base: https://exampleopenaiendpoint-production.up.railway.app
|
||||||
|
|
||||||
general_settings:
|
general_settings:
|
||||||
master_key: sk-1234
|
key_management_system: "azure_key_vault"
|
||||||
|
|
||||||
litellm_settings:
|
litellm_settings:
|
||||||
success_callback: ["prometheus"]
|
success_callback: ["prometheus"]
|
||||||
|
|
|
@ -56,6 +56,7 @@ def get_secret(
|
||||||
):
|
):
|
||||||
key_management_system = litellm._key_management_system
|
key_management_system = litellm._key_management_system
|
||||||
key_management_settings = litellm._key_management_settings
|
key_management_settings = litellm._key_management_settings
|
||||||
|
secret = None
|
||||||
|
|
||||||
if secret_name.startswith("os.environ/"):
|
if secret_name.startswith("os.environ/"):
|
||||||
secret_name = secret_name.replace("os.environ/", "")
|
secret_name = secret_name.replace("os.environ/", "")
|
||||||
|
@ -121,7 +122,7 @@ def get_secret(
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
if response.status_code == 200:
|
if response.status_code == 200:
|
||||||
oidc_token = response.text["value"]
|
oidc_token = response.json().get("value", None)
|
||||||
oidc_cache.set_cache(key=secret_name, value=oidc_token, ttl=300 - 5)
|
oidc_cache.set_cache(key=secret_name, value=oidc_token, ttl=300 - 5)
|
||||||
return oidc_token
|
return oidc_token
|
||||||
else:
|
else:
|
||||||
|
@ -245,8 +246,8 @@ def get_secret(
|
||||||
print_verbose(f"secret_dict: {secret_dict}")
|
print_verbose(f"secret_dict: {secret_dict}")
|
||||||
for k, v in secret_dict.items():
|
for k, v in secret_dict.items():
|
||||||
secret = v
|
secret = v
|
||||||
print_verbose(f"secret: {secret}")
|
print_verbose(f"secret: {secret}")
|
||||||
if key_manager == KeyManagementSystem.GOOGLE_SECRET_MANAGER.value:
|
elif key_manager == KeyManagementSystem.GOOGLE_SECRET_MANAGER.value:
|
||||||
try:
|
try:
|
||||||
secret = client.get_secret_from_google_secret_manager(
|
secret = client.get_secret_from_google_secret_manager(
|
||||||
secret_name
|
secret_name
|
||||||
|
@ -269,11 +270,12 @@ def get_secret(
|
||||||
)
|
)
|
||||||
secret = os.getenv(secret_name)
|
secret = os.getenv(secret_name)
|
||||||
try:
|
try:
|
||||||
secret_value_as_bool = ast.literal_eval(secret)
|
if isinstance(secret, str):
|
||||||
if isinstance(secret_value_as_bool, bool):
|
secret_value_as_bool = ast.literal_eval(secret)
|
||||||
return secret_value_as_bool
|
if isinstance(secret_value_as_bool, bool):
|
||||||
else:
|
return secret_value_as_bool
|
||||||
return secret
|
else:
|
||||||
|
return secret
|
||||||
except:
|
except:
|
||||||
return secret
|
return secret
|
||||||
else:
|
else:
|
||||||
|
|
30
litellm/tests/test_get_secret.py
Normal file
30
litellm/tests/test_get_secret.py
Normal file
|
@ -0,0 +1,30 @@
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
import sys
|
||||||
|
from datetime import datetime
|
||||||
|
from unittest.mock import AsyncMock, Mock, patch
|
||||||
|
|
||||||
|
sys.path.insert(
|
||||||
|
0, os.path.abspath("../..")
|
||||||
|
) # Adds the parent directory to the system path
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
import litellm
|
||||||
|
from litellm.proxy._types import KeyManagementSystem
|
||||||
|
from litellm.secret_managers.main import get_secret
|
||||||
|
|
||||||
|
|
||||||
|
class MockSecretClient:
|
||||||
|
def get_secret(self, secret_name):
|
||||||
|
return Mock(value="mocked_secret_value")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_azure_kms():
|
||||||
|
"""
|
||||||
|
Basic asserts that the value from get secret is from Azure Key Vault when Key Management System is Azure Key Vault
|
||||||
|
"""
|
||||||
|
with patch("litellm.secret_manager_client", new=MockSecretClient()):
|
||||||
|
litellm._key_management_system = KeyManagementSystem.AZURE_KEY_VAULT
|
||||||
|
secret = get_secret(secret_name="ishaan-test-key")
|
||||||
|
assert secret == "mocked_secret_value"
|
Loading…
Add table
Add a link
Reference in a new issue