From 8103e2b2da2f06d5f18cd72b7272d47217949e7e Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Mon, 16 Sep 2024 18:01:40 -0700
Subject: [PATCH] [Fix-Proxy] Azure Key Management - Secret Manager  (#5728)

* fix azure key mgtm error

* add test for azure kms

* add test for azure kms
---
 litellm/proxy/proxy_config.yaml  |  4 ++--
 litellm/secret_managers/main.py  | 18 ++++++++++--------
 litellm/tests/test_get_secret.py | 30 ++++++++++++++++++++++++++++++
 3 files changed, 42 insertions(+), 10 deletions(-)
 create mode 100644 litellm/tests/test_get_secret.py

diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml
index b33bc35d1fe..51e67963fcc 100644
--- a/litellm/proxy/proxy_config.yaml
+++ b/litellm/proxy/proxy_config.yaml
@@ -20,8 +20,8 @@ model_list:
       api_key: fake-key
       api_base: https://exampleopenaiendpoint-production.up.railway.app
 
-general_settings: 
- master_key: sk-1234 
+general_settings:
+  key_management_system: "azure_key_vault"
 
 litellm_settings:
   success_callback: ["prometheus"]
diff --git a/litellm/secret_managers/main.py b/litellm/secret_managers/main.py
index 5d1f72cf7fb..13a61c66596 100644
--- a/litellm/secret_managers/main.py
+++ b/litellm/secret_managers/main.py
@@ -56,6 +56,7 @@ def get_secret(
 ):
     key_management_system = litellm._key_management_system
     key_management_settings = litellm._key_management_settings
+    secret = None
 
     if secret_name.startswith("os.environ/"):
         secret_name = secret_name.replace("os.environ/", "")
@@ -121,7 +122,7 @@ def get_secret(
                 },
             )
             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)
                 return oidc_token
             else:
@@ -245,8 +246,8 @@ def get_secret(
                     print_verbose(f"secret_dict: {secret_dict}")
                     for k, v in secret_dict.items():
                         secret = v
-                    print_verbose(f"secret: {secret}")
-                if key_manager == KeyManagementSystem.GOOGLE_SECRET_MANAGER.value:
+                        print_verbose(f"secret: {secret}")
+                elif key_manager == KeyManagementSystem.GOOGLE_SECRET_MANAGER.value:
                     try:
                         secret = client.get_secret_from_google_secret_manager(
                             secret_name
@@ -269,11 +270,12 @@ def get_secret(
                 )
                 secret = os.getenv(secret_name)
             try:
-                secret_value_as_bool = ast.literal_eval(secret)
-                if isinstance(secret_value_as_bool, bool):
-                    return secret_value_as_bool
-                else:
-                    return secret
+                if isinstance(secret, str):
+                    secret_value_as_bool = ast.literal_eval(secret)
+                    if isinstance(secret_value_as_bool, bool):
+                        return secret_value_as_bool
+                    else:
+                        return secret
             except:
                 return secret
         else:
diff --git a/litellm/tests/test_get_secret.py b/litellm/tests/test_get_secret.py
new file mode 100644
index 00000000000..eec67b5d765
--- /dev/null
+++ b/litellm/tests/test_get_secret.py
@@ -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"