diff --git a/docs/my-website/docs/secret.md b/docs/my-website/docs/secret.md index 9f0ff7059c..6e6e2ad506 100644 --- a/docs/my-website/docs/secret.md +++ b/docs/my-website/docs/secret.md @@ -71,6 +71,10 @@ general_settings: store_virtual_keys: true # OPTIONAL. Defaults to False, when True will store virtual keys in secret manager prefix_for_stored_virtual_keys: "litellm/" # OPTIONAL. If set, this prefix will be used for stored virtual keys in the secret manager access_mode: "write_only" # Literal["read_only", "write_only", "read_and_write"] + description: "litellm virtual key" # OPTIONAL, if set will set this as the description for all virtual keys + tags: # OPTIONAL, if set will set this as the tags for all virtual keys + Environment: "Prod" + Owner: "AI Platform team" ``` diff --git a/litellm/proxy/hooks/key_management_event_hooks.py b/litellm/proxy/hooks/key_management_event_hooks.py index 7ad1aaffc1..49072fe841 100644 --- a/litellm/proxy/hooks/key_management_event_hooks.py +++ b/litellm/proxy/hooks/key_management_event_hooks.py @@ -234,11 +234,21 @@ class KeyManagementEventHooks: # store the key in the secret manager if isinstance(litellm.secret_manager_client, BaseSecretManager): + tags = getattr(litellm._key_management_settings, "tags", None) + description = getattr( + litellm._key_management_settings, "description", None + ) + verbose_proxy_logger.debug( + f"Creating secret with {secret_name} and tags={tags} and description={description}" + ) + await litellm.secret_manager_client.async_write_secret( secret_name=KeyManagementEventHooks._get_secret_name( secret_name ), + description=description, secret_value=secret_token, + tags=tags ) @staticmethod diff --git a/litellm/secret_managers/aws_secret_manager_v2.py b/litellm/secret_managers/aws_secret_manager_v2.py index ff9921ebcb..5f8fa795f7 100644 --- a/litellm/secret_managers/aws_secret_manager_v2.py +++ b/litellm/secret_managers/aws_secret_manager_v2.py @@ -198,12 +198,13 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): return primary_secret_kv_pairs.get(secret_name) async def async_write_secret( - self, - secret_name: str, - secret_value: str, - description: Optional[str] = None, - optional_params: Optional[dict] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, + self, + secret_name: str, + secret_value: str, + description: Optional[str] = None, + optional_params: Optional[dict] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + tags: Optional[Union[dict, list]] = None ) -> dict: """ Async function to write a secret to AWS Secrets Manager @@ -214,22 +215,37 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): description: Optional description for the secret optional_params: Additional AWS parameters timeout: Request timeout + tags: Optional dict or list of tags to apply, e.g. + {"Environment": "Prod", "Owner": "AI-Platform"} or + [{"Key": "Environment", "Value": "Prod"}] """ from litellm._uuid import uuid - # Prepare the request data - data = {"Name": secret_name, "SecretString": secret_value} + data = { + "Name": secret_name, + "SecretString": secret_value, + "ClientRequestToken": str(uuid.uuid4()), + } + if description: data["Description"] = description - data["ClientRequestToken"] = str(uuid.uuid4()) + # ✅ Normalize tags to AWS format + if tags: + if isinstance(tags, dict): + tags_list = [{"Key": k, "Value": str(v)} for k, v in tags.items()] + elif isinstance(tags, list): + tags_list = tags + else: + raise ValueError("Tags must be a dict or list of {Key, Value} pairs") + data["Tags"] = tags_list endpoint_url, headers, body = self._prepare_request( action="CreateSecret", secret_name=secret_name, secret_value=secret_value, optional_params=optional_params, - request_data=data, # Pass the complete request data + request_data=data, ) async_client = get_async_httpx_client( @@ -354,7 +370,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): # if __name__ == "__main__": # print("loading aws secret manager v2") # aws_secret_manager_v2 = AWSSecretsManagerV2() - +# import asyncio # print("writing secret to aws secret manager v2") # asyncio.run(aws_secret_manager_v2.async_write_secret(secret_name="test_secret_3", secret_value="test_value_2")) # print("reading secret from aws secret manager v2") diff --git a/litellm/secret_managers/base_secret_manager.py b/litellm/secret_managers/base_secret_manager.py index 0b5bf74db7..af77c5f45b 100644 --- a/litellm/secret_managers/base_secret_manager.py +++ b/litellm/secret_managers/base_secret_manager.py @@ -59,6 +59,7 @@ class BaseSecretManager(ABC): description: Optional[str] = None, optional_params: Optional[dict] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, + tags: Optional[Union[dict, list]] = None ) -> Dict[str, Any]: """ Asynchronously write a secret to the secret manager. @@ -69,6 +70,9 @@ class BaseSecretManager(ABC): description (Optional[str]): Description of the secret. Some secret managers allow storing a description with the secret. optional_params (Optional[dict]): Additional parameters specific to the secret manager timeout (Optional[Union[float, httpx.Timeout]]): Request timeout + tags: Optional dict or list of tags to apply, e.g. + {"Environment": "Prod", "Owner": "AI-Platform"} or + [{"Key": "Environment", "Value": "Prod"}] Returns: Dict[str, Any]: Response from the secret manager containing write operation details """ diff --git a/litellm/secret_managers/hashicorp_secret_manager.py b/litellm/secret_managers/hashicorp_secret_manager.py index c462d0b464..d2fb9ff70f 100644 --- a/litellm/secret_managers/hashicorp_secret_manager.py +++ b/litellm/secret_managers/hashicorp_secret_manager.py @@ -202,6 +202,7 @@ class HashicorpSecretManager(BaseSecretManager): description: Optional[str] = None, optional_params: Optional[dict] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, + tags: Optional[Union[dict, list]] = None ) -> Dict[str, Any]: """ Writes a secret to Vault KV v2 using an async HTTPX client. diff --git a/litellm/types/secret_managers/main.py b/litellm/types/secret_managers/main.py index e582df47df..fd7d8d4f01 100644 --- a/litellm/types/secret_managers/main.py +++ b/litellm/types/secret_managers/main.py @@ -1,5 +1,5 @@ import enum -from typing import List, Literal, Optional +from typing import List, Literal, Optional, Dict from litellm.types.llms.base import LiteLLMPydanticObjectBase @@ -35,4 +35,11 @@ class KeyManagementSettings(LiteLLMPydanticObjectBase): 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 - """ \ No newline at end of file + """ + + description: Optional[str] = None + """Optional description attached when creating secrets (visible in AWS console).""" + + + tags: Optional[Dict[str, str]] = None + """Optional tags to attach when creating secrets (e.g. {"Environment": "Prod", "Owner": "AI-Platform"}).""" \ No newline at end of file diff --git a/tests/litellm_utils_tests/test_aws_secret_manager.py b/tests/litellm_utils_tests/test_aws_secret_manager.py index 5aa608acb8..b6ccd41768 100644 --- a/tests/litellm_utils_tests/test_aws_secret_manager.py +++ b/tests/litellm_utils_tests/test_aws_secret_manager.py @@ -193,3 +193,63 @@ async def test_primary_secret_functionality(): ) print("Delete Response:", delete_response) assert delete_response is not None + +@pytest.mark.asyncio +async def test_write_secret_with_description_and_tags(): + """Test writing a secret with description and tags""" + check_aws_credentials() + + secret_manager = AWSSecretsManagerV2() + test_secret_name = f"litellm_test_{uuid.uuid4().hex[:8]}_tags" + test_secret_value = "test_value_with_tags" + + test_description = "LiteLLM Secret with Description and Tags" + test_tags = { + "Environment": "Test", + "Owner": "IntelligenceLayer", + "Purpose": "UnitTest", + } + + try: + # Write secret with tags and description + write_response = await secret_manager.async_write_secret( + secret_name=test_secret_name, + secret_value=test_secret_value, + description=test_description, + tags=test_tags, + ) + + print("Write Response:", write_response) + assert write_response is not None + assert "ARN" in write_response + assert "Name" in write_response + assert write_response["Name"] == test_secret_name + + # --- Validate the secret metadata via AWS CLI / boto3 --- + import boto3 + + client = boto3.client("secretsmanager", region_name=os.getenv("AWS_REGION_NAME")) + describe_resp = client.describe_secret(SecretId=test_secret_name) + print("Describe Response:", describe_resp) + + # Validate description + assert describe_resp.get("Description") == test_description + + # Validate tags (as list of dicts in AWS) + if "Tags" in describe_resp: + tag_dict = {t["Key"]: t["Value"] for t in describe_resp["Tags"]} + for k, v in test_tags.items(): + assert tag_dict.get(k) == v, f"Expected tag {k}={v}, got {tag_dict.get(k)}" + else: + pytest.fail("No tags found in describe_secret response") + + # --- Validate secret value --- + read_value = await secret_manager.async_read_secret(secret_name=test_secret_name) + print("Read Value:", read_value) + assert read_value == test_secret_value + + finally: + # Cleanup: Delete the secret + delete_response = await secret_manager.async_delete_secret(secret_name=test_secret_name) + print("Delete Response:", delete_response) + assert delete_response is not None diff --git a/tests/litellm_utils_tests/test_secret_manager.py b/tests/litellm_utils_tests/test_secret_manager.py index 8c604fd2e6..7099f6e13d 100644 --- a/tests/litellm_utils_tests/test_secret_manager.py +++ b/tests/litellm_utils_tests/test_secret_manager.py @@ -24,6 +24,7 @@ from litellm.secret_managers.main import ( get_secret, _should_read_secret_from_secret_manager, ) +from unittest.mock import AsyncMock def load_vertex_ai_credentials(): @@ -358,3 +359,110 @@ def test_get_secret_with_access_mode(): litellm.secret_manager_client = None litellm._key_management_settings = KeyManagementSettings() del os.environ[test_secret_name] + +def test_key_management_settings_defaults(): + """ + Test that KeyManagementSettings initializes with correct default values. + """ + from litellm.types.secret_managers.main import KeyManagementSettings + + settings = KeyManagementSettings() + + assert settings.store_virtual_keys is False + assert settings.prefix_for_stored_virtual_keys == "litellm/" + assert settings.access_mode == "read_only" + assert settings.description is None + assert settings.tags is None + assert settings.primary_secret_name is None + + +def test_key_management_settings_custom_values(): + """ + Test that KeyManagementSettings correctly stores custom description and tags. + """ + from litellm.types.secret_managers.main import KeyManagementSettings + + custom_tags = {"Environment": "Dev", "Team": "Intelligence"} + custom_description = "LiteLLM-managed API key for development" + + settings = KeyManagementSettings( + store_virtual_keys=True, + prefix_for_stored_virtual_keys="litellm/custom/", + access_mode="read_and_write", + primary_secret_name="primary/litellm/keys", + description=custom_description, + tags=custom_tags, + ) + + assert settings.store_virtual_keys is True + assert settings.prefix_for_stored_virtual_keys == "litellm/custom/" + assert settings.access_mode == "read_and_write" + assert settings.primary_secret_name == "primary/litellm/keys" + assert settings.description == custom_description + assert settings.tags == custom_tags + + +@pytest.mark.asyncio +async def test_async_write_secret_receives_description_and_tags(monkeypatch): + """ + Test that AWSSecretsManagerV2.async_write_secret receives description and tags when KeyManagementSettings is set. + """ + from litellm import litellm + from litellm.secret_managers.aws_secret_manager_v2 import AWSSecretsManagerV2 + from litellm.types.secret_managers.main import KeyManagementSettings + + # Mock out AWS network calls + mock_async_write = AsyncMock(return_value={"Name": "litellm/test_secret"}) + monkeypatch.setattr(AWSSecretsManagerV2, "async_write_secret", mock_async_write) + + # Setup settings + litellm._key_management_settings = KeyManagementSettings( + store_virtual_keys=True, + description="LiteLLM Unit Test Secret", + tags={"Owner": "UnitTest", "Purpose": "Validation"}, + ) + + # Instantiate fake client + litellm.secret_manager_client = AWSSecretsManagerV2() + + # Call the helper method that stores a virtual key + from litellm.proxy.hooks.key_management_event_hooks import ( + KeyManagementEventHooks, + ) + + await KeyManagementEventHooks._store_virtual_key_in_secret_manager( + secret_name="test_secret", secret_token="test_value" + ) + + # Verify async_write_secret was called with correct metadata + mock_async_write.assert_called_once() + args, kwargs = mock_async_write.call_args + + assert kwargs["secret_name"].endswith("test_secret") + assert kwargs["secret_value"] == "test_value" + assert kwargs["description"] == "LiteLLM Unit Test Secret" + assert kwargs["tags"] == {"Owner": "UnitTest", "Purpose": "Validation"} + + +def test_key_management_settings_serialization_roundtrip(): + """ + Test that KeyManagementSettings serializes and deserializes consistently (Pydantic behavior). + """ + from litellm.types.secret_managers.main import KeyManagementSettings + + original = KeyManagementSettings( + store_virtual_keys=True, + prefix_for_stored_virtual_keys="litellm/dev/", + access_mode="read_and_write", + description="Roundtrip test", + tags={"Env": "QA"}, + ) + + as_dict = original.model_dump() + reloaded = KeyManagementSettings(**as_dict) + + assert reloaded.store_virtual_keys is True + assert reloaded.prefix_for_stored_virtual_keys == "litellm/dev/" + assert reloaded.access_mode == "read_and_write" + assert reloaded.description == "Roundtrip test" + assert reloaded.tags == {"Env": "QA"}