From fbdcde1a54e8be1a9f8817e79be7b739875b0010 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 1 Dec 2023 19:45:09 -0800 Subject: [PATCH] fix(proxy_server.py): fix linting errors --- litellm/__init__.py | 2 +- litellm/proxy/proxy_server.py | 32 ++++++++++++-------------------- 2 files changed, 13 insertions(+), 21 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index 5fd54401ba..fce4d62743 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -59,7 +59,7 @@ fallbacks: Optional[List] = None context_window_fallbacks: Optional[List] = None allowed_fails: int = 0 ####### SECRET MANAGERS ##################### -secret_manager_client = None # list of instantiated key management clients - e.g. azure kv, infisical, etc. +secret_manager_client: Optional[Any] = None # list of instantiated key management clients - e.g. azure kv, infisical, etc. ############################################# def get_model_cost_map(url: str): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b5b7702953..8c4ade4b38 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -295,32 +295,24 @@ def load_from_azure_key_vault(use_azure_key_vault: bool = False): from azure.identity import ClientSecretCredential # Set your Azure Key Vault URI - KVUri = os.getenv("AZURE_KEY_VAULT_URI") + KVUri = os.getenv("AZURE_KEY_VAULT_URI", None) # Set your Azure AD application/client ID, client secret, and tenant ID - client_id = os.getenv("AZURE_CLIENT_ID") - client_secret = os.getenv("AZURE_CLIENT_SECRET") - tenant_id = os.getenv("AZURE_TENANT_ID") + client_id = os.getenv("AZURE_CLIENT_ID", None) + client_secret = os.getenv("AZURE_CLIENT_SECRET", None) + tenant_id = os.getenv("AZURE_TENANT_ID", None) - # Initialize the ClientSecretCredential - credential = ClientSecretCredential(client_id=client_id, client_secret=client_secret, tenant_id=tenant_id) + if KVUri is not None and client_id is not None and client_secret is not None and tenant_id is not None: + # Initialize the ClientSecretCredential + credential = ClientSecretCredential(client_id=client_id, client_secret=client_secret, tenant_id=tenant_id) - # Create the SecretClient using the credential - client = SecretClient(vault_url=KVUri, credential=credential) + # Create the SecretClient using the credential + client = SecretClient(vault_url=KVUri, credential=credential) - litellm.secret_manager_client = client - # # Retrieve all secrets - # secrets = client.get_secrets() - - # # Load secrets into environment variables - # for secret in secrets: - # secret_name = secret.name - # secret_value = client.get_secret(secret_name).value - # os.environ[secret_name] = secret_value - - print(f"test key - : {litellm.get_secret('test-3')}") + litellm.secret_manager_client = client + else: + raise Exception(f"Missing KVUri or client_id or client_secret or tenant_id from environment") except Exception as e: - print(e) print("Error when loading keys from Azure Key Vault. Ensure you run `pip install azure-identity azure-keyvault-secrets`")