From 04e839d846749161fe3370e48dcd06f610968283 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 6 Mar 2025 19:30:18 -0800 Subject: [PATCH] (AWS Secret Manager) - Using K/V pairs in 1 AWS Secret (#9039) * fixes for primary_secret_kv_pairs * _parse_primary_secret * Using K/V pairs in 1 AWS Secret * test_primary_secret_functionality --- docs/my-website/docs/secret.md | 30 ++++++++++ litellm/proxy/_types.py | 7 +++ litellm/proxy/proxy_config.yaml | 19 ++++++ .../secret_managers/aws_secret_manager_v2.py | 59 +++++++++++++++++-- litellm/secret_managers/main.py | 5 +- .../test_aws_secret_manager.py | 56 ++++++++++++++++++ 6 files changed, 171 insertions(+), 5 deletions(-) diff --git a/docs/my-website/docs/secret.md b/docs/my-website/docs/secret.md index a65c696f367..76761642594 100644 --- a/docs/my-website/docs/secret.md +++ b/docs/my-website/docs/secret.md @@ -96,6 +96,33 @@ litellm --config /path/to/config.yaml ``` +### Using K/V pairs in 1 AWS Secret + +You can read multiple keys from a single AWS Secret using the `primary_secret_name` parameter: + +```yaml +general_settings: + key_management_system: "aws_secret_manager" + key_management_settings: + hosted_keys: [ + "OPENAI_API_KEY_MODEL_1", + "OPENAI_API_KEY_MODEL_2", + ] + primary_secret_name: "litellm_secrets" # 👈 Read multiple keys from one JSON secret +``` + +The `primary_secret_name` allows you to read multiple keys from a single AWS Secret as a JSON object. For example, the "litellm_secrets" would contain: + +```json +{ + "OPENAI_API_KEY_MODEL_1": "sk-key1...", + "OPENAI_API_KEY_MODEL_2": "sk-key2..." +} +``` + +This reduces the number of AWS Secrets you need to manage. + + ## Hashicorp Vault @@ -353,4 +380,7 @@ general_settings: # Hosted Keys Settings hosted_keys: ["litellm_master_key"] # OPTIONAL. Specify which env keys you stored on AWS + + # K/V pairs in 1 AWS Secret Settings + primary_secret_name: "litellm_secrets" # OPTIONAL. Read multiple keys from one JSON secret on AWS Secret Manager ``` \ No newline at end of file diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 89f02f82283..7ac2b17c15e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1161,6 +1161,13 @@ class KeyManagementSettings(LiteLLMPydanticObjectBase): Access mode for the secret manager, when write_only will only use for writing secrets """ + primary_secret_name: Optional[str] = None + """ + If set, will read secrets from this primary secret in the secret manager + + eg. on AWS you can store multiple secret values as K/V pairs in a single secret + """ + class TeamDefaultSettings(LiteLLMPydanticObjectBase): team_id: str diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 258eef6307b..6341b3f74a2 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -1,4 +1,23 @@ model_list: + - model_name: model-1 + litellm_params: + model: openai/model-1 + api_key: os.environ/OPENAI_API_KEY_MODEL_1 + api_base: https://exampleopenaiendpoint-production.up.railway.app/ + - model_name: model-2 + litellm_params: + model: openai/model-2 + api_key: os.environ/OPENAI_API_KEY_MODEL_2 + api_base: https://exampleopenaiendpoint-production.up.railway.app/ + +general_settings: + key_management_system: "aws_secret_manager" + key_management_settings: + hosted_keys: [ + "OPENAI_API_KEY_MODEL_1", + "OPENAI_API_KEY_MODEL_2", + ] + primary_secret_name: "litellm_secrets" - model_name: bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0 litellm_params: diff --git a/litellm/secret_managers/aws_secret_manager_v2.py b/litellm/secret_managers/aws_secret_manager_v2.py index b72a709d86e..fd89d6c5584 100644 --- a/litellm/secret_managers/aws_secret_manager_v2.py +++ b/litellm/secret_managers/aws_secret_manager_v2.py @@ -63,6 +63,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): secret_name: str, optional_params: Optional[dict] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, + primary_secret_name: Optional[str] = None, ) -> Optional[str]: """ Async function to read a secret from AWS Secrets Manager @@ -72,6 +73,11 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): Raises: ValueError: If the secret is not found or an HTTP error occurs """ + if primary_secret_name: + return await self.async_read_secret_from_primary_secret( + secret_name=secret_name, primary_secret_name=primary_secret_name + ) + endpoint_url, headers, body = self._prepare_request( action="GetSecretValue", secret_name=secret_name, @@ -93,7 +99,9 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): raise ValueError("Timeout error occurred") except Exception as e: verbose_logger.exception( - "Error reading secret from AWS Secrets Manager: %s", str(e) + "Error reading secret='%s' from AWS Secrets Manager: %s", + secret_name, + str(e), ) return None @@ -102,13 +110,13 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): secret_name: str, optional_params: Optional[dict] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, + primary_secret_name: Optional[str] = None, ) -> Optional[str]: """ Sync function to read a secret from AWS Secrets Manager Done for backwards compatibility with existing codebase, since get_secret is a sync function """ - # self._prepare_request uses these env vars, we cannot read them from AWS Secrets Manager. If we do we'd get stuck in an infinite loop if secret_name in [ "AWS_ACCESS_KEY_ID", @@ -119,6 +127,11 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): ]: return os.getenv(secret_name) + if primary_secret_name: + return self.sync_read_secret_from_primary_secret( + secret_name=secret_name, primary_secret_name=primary_secret_name + ) + endpoint_url, headers, body = self._prepare_request( action="GetSecretValue", secret_name=secret_name, @@ -138,15 +151,53 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): raise ValueError("Timeout error occurred") except httpx.HTTPStatusError as e: verbose_logger.exception( - "Error reading secret from AWS Secrets Manager: %s", + "Error reading secret='%s' from AWS Secrets Manager: %s, %s", + secret_name, str(e.response.text), + str(e.response.status_code), ) except Exception as e: verbose_logger.exception( - "Error reading secret from AWS Secrets Manager: %s", str(e) + "Error reading secret='%s' from AWS Secrets Manager: %s", + secret_name, + str(e), ) return None + def _parse_primary_secret(self, primary_secret_json_str: Optional[str]) -> dict: + """ + Parse the primary secret JSON string into a dictionary + + Args: + primary_secret_json_str: JSON string containing key-value pairs + + Returns: + Dictionary of key-value pairs from the primary secret + """ + return json.loads(primary_secret_json_str or "{}") + + def sync_read_secret_from_primary_secret( + self, secret_name: str, primary_secret_name: str + ) -> Optional[str]: + """ + Read a secret from the primary secret + """ + primary_secret_json_str = self.sync_read_secret(secret_name=primary_secret_name) + primary_secret_kv_pairs = self._parse_primary_secret(primary_secret_json_str) + return primary_secret_kv_pairs.get(secret_name) + + async def async_read_secret_from_primary_secret( + self, secret_name: str, primary_secret_name: str + ) -> Optional[str]: + """ + Read a secret from the primary secret + """ + primary_secret_json_str = await self.async_read_secret( + secret_name=primary_secret_name + ) + primary_secret_kv_pairs = self._parse_primary_secret(primary_secret_json_str) + return primary_secret_kv_pairs.get(secret_name) + async def async_write_secret( self, secret_name: str, diff --git a/litellm/secret_managers/main.py b/litellm/secret_managers/main.py index aa13960c9f9..e505484b4b4 100644 --- a/litellm/secret_managers/main.py +++ b/litellm/secret_managers/main.py @@ -274,7 +274,10 @@ def get_secret( # noqa: PLR0915 ) if isinstance(client, AWSSecretsManagerV2): - secret = client.sync_read_secret(secret_name=secret_name) + secret = client.sync_read_secret( + secret_name=secret_name, + primary_secret_name=key_management_settings.primary_secret_name, + ) print_verbose(f"get_secret_value_response: {secret}") elif key_manager == KeyManagementSystem.GOOGLE_SECRET_MANAGER.value: try: diff --git a/tests/litellm_utils_tests/test_aws_secret_manager.py b/tests/litellm_utils_tests/test_aws_secret_manager.py index f2e2319cc66..e64c42b26c6 100644 --- a/tests/litellm_utils_tests/test_aws_secret_manager.py +++ b/tests/litellm_utils_tests/test_aws_secret_manager.py @@ -137,3 +137,59 @@ async def test_read_nonexistent_secret(): response = await secret_manager.async_read_secret(secret_name=nonexistent_secret) assert response is None + + +@pytest.mark.asyncio +async def test_primary_secret_functionality(): + """Test storing and retrieving secrets from a primary secret""" + check_aws_credentials() + + secret_manager = AWSSecretsManagerV2() + primary_secret_name = f"litellm_test_primary_{uuid.uuid4().hex[:8]}" + + # Create a primary secret with multiple key-value pairs + primary_secret_value = { + "api_key_1": "secret_value_1", + "api_key_2": "secret_value_2", + "database_url": "postgresql://user:password@localhost:5432/db", + "nested_secret": json.dumps({"key": "value", "number": 42}), + } + + try: + # Write the primary secret + write_response = await secret_manager.async_write_secret( + secret_name=primary_secret_name, + secret_value=json.dumps(primary_secret_value), + description="LiteLLM Test Primary Secret", + ) + + print("Primary Secret Write Response:", write_response) + assert write_response is not None + assert "ARN" in write_response + assert "Name" in write_response + assert write_response["Name"] == primary_secret_name + + # Test reading individual secrets from the primary secret + for key, expected_value in primary_secret_value.items(): + # Read using the primary_secret_name parameter + value = await secret_manager.async_read_secret( + secret_name=key, primary_secret_name=primary_secret_name + ) + + print(f"Read {key} from primary secret:", value) + assert value == expected_value + + # Test reading a non-existent key from the primary secret + non_existent_key = "non_existent_key" + value = await secret_manager.async_read_secret( + secret_name=non_existent_key, primary_secret_name=primary_secret_name + ) + assert value is None, f"Expected None for non-existent key, got {value}" + + finally: + # Cleanup: Delete the primary secret + delete_response = await secret_manager.async_delete_secret( + secret_name=primary_secret_name + ) + print("Delete Response:", delete_response) + assert delete_response is not None