From a2ca3887d1feac596db54856cb56f76857c241dc Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 6 Jun 2024 15:32:51 -0700 Subject: [PATCH 1/5] feat(aws_secret_manager.py): allows user to keep a hash of the proxy master key in their env --- litellm/proxy/_super_secret_config.yaml | 4 ++- litellm/proxy/_types.py | 1 + litellm/proxy/proxy_server.py | 11 ++++-- .../secret_managers/aws_secret_manager.py | 19 ++++++++++ litellm/utils.py | 35 ++++++++++++++++--- 5 files changed, 62 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/_super_secret_config.yaml b/litellm/proxy/_super_secret_config.yaml index 72479bd5d5..7fa1bbc19c 100644 --- a/litellm/proxy/_super_secret_config.yaml +++ b/litellm/proxy/_super_secret_config.yaml @@ -56,8 +56,10 @@ router_settings: litellm_settings: success_callback: ["langfuse"] - json_logs: true general_settings: alerting: ["email"] + key_management_system: "aws_kms" + key_management_settings: + hosted_keys: ["LITELLM_MASTER_KEY"] diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 8a95f4e1d8..c19fd7e69e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -946,6 +946,7 @@ class KeyManagementSystem(enum.Enum): AZURE_KEY_VAULT = "azure_key_vault" AWS_SECRET_MANAGER = "aws_secret_manager" LOCAL = "local" + AWS_KMS = "aws_kms" class KeyManagementSettings(LiteLLMBase): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 8cf2fa118d..e188841eb4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -112,7 +112,10 @@ from litellm import ( CreateFileRequest, ) from litellm.proxy.secret_managers.google_kms import load_google_kms -from litellm.proxy.secret_managers.aws_secret_manager import load_aws_secret_manager +from litellm.proxy.secret_managers.aws_secret_manager import ( + load_aws_secret_manager, + load_aws_kms, +) import pydantic from litellm.proxy._types import * from litellm.caching import DualCache, RedisCache @@ -2736,10 +2739,12 @@ class ProxyConfig: load_google_kms(use_google_kms=True) elif ( key_management_system - == KeyManagementSystem.AWS_SECRET_MANAGER.value + == KeyManagementSystem.AWS_SECRET_MANAGER.value # noqa: F405 ): ### LOAD FROM AWS SECRET MANAGER ### load_aws_secret_manager(use_aws_secret_manager=True) + elif key_management_system == KeyManagementSystem.AWS_KMS.value: + load_aws_kms(use_aws_kms=True) else: raise ValueError("Invalid Key Management System selected") key_management_settings = general_settings.get( @@ -2773,6 +2778,7 @@ class ProxyConfig: master_key = general_settings.get( "master_key", litellm.get_secret("LITELLM_MASTER_KEY", None) ) + if master_key and master_key.startswith("os.environ/"): master_key = litellm.get_secret(master_key) if not isinstance(master_key, str): @@ -4098,6 +4104,7 @@ async def chat_completion( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): global general_settings, user_debug, proxy_logging_obj, llm_model_list + data = {} try: body = await request.body() diff --git a/litellm/proxy/secret_managers/aws_secret_manager.py b/litellm/proxy/secret_managers/aws_secret_manager.py index a40b1dffa2..9e6a777c40 100644 --- a/litellm/proxy/secret_managers/aws_secret_manager.py +++ b/litellm/proxy/secret_managers/aws_secret_manager.py @@ -8,6 +8,7 @@ Requires: * `pip install boto3>=1.28.57` """ +import boto3.session import litellm, os from typing import Optional from litellm.proxy._types import KeyManagementSystem @@ -38,3 +39,21 @@ def load_aws_secret_manager(use_aws_secret_manager: Optional[bool]): except Exception as e: raise e + + +def load_aws_kms(use_aws_kms: Optional[bool]): + if use_aws_kms is None or use_aws_kms is False: + return + try: + import boto3 + + validate_environment() + + # Create a Secrets Manager client + kms_client = boto3.client("kms", region_name=os.getenv("AWS_REGION_NAME")) + + litellm.secret_manager_client = kms_client + litellm._key_management_system = KeyManagementSystem.AWS_KMS + + except Exception as e: + raise e diff --git a/litellm/utils.py b/litellm/utils.py index ba6a374674..6ce41ad84a 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7351,10 +7351,10 @@ def get_provider_fields(custom_llm_provider: str) -> List[ProviderField]: if custom_llm_provider == "databricks": return litellm.DatabricksConfig().get_required_params() - + elif custom_llm_provider == "ollama": return litellm.OllamaConfig().get_required_params() - + else: return [] @@ -10052,6 +10052,8 @@ def get_secret( ): key_management_system = litellm._key_management_system key_management_settings = litellm._key_management_settings + args = locals() + if secret_name.startswith("os.environ/"): secret_name = secret_name.replace("os.environ/", "") @@ -10139,13 +10141,13 @@ def get_secret( key_manager = "local" if ( - key_manager == KeyManagementSystem.AZURE_KEY_VAULT + key_manager == KeyManagementSystem.AZURE_KEY_VAULT.value or type(client).__module__ + "." + type(client).__name__ == "azure.keyvault.secrets._client.SecretClient" ): # support Azure Secret Client - from azure.keyvault.secrets import SecretClient secret = client.get_secret(secret_name).value elif ( - key_manager == KeyManagementSystem.GOOGLE_KMS + key_manager == KeyManagementSystem.GOOGLE_KMS.value or client.__class__.__name__ == "KeyManagementServiceClient" ): encrypted_secret: Any = os.getenv(secret_name) @@ -10173,6 +10175,25 @@ def get_secret( secret = response.plaintext.decode( "utf-8" ) # assumes the original value was encoded with utf-8 + elif key_manager == KeyManagementSystem.AWS_KMS.value: + """ + Only check the tokens which start with 'aws_kms/'. This prevents latency impact caused by checking all keys. + """ + encrypted_value = os.getenv(secret_name, None) + if encrypted_value is None: + raise Exception("encrypted value for AWS KMS cannot be None.") + # Decode the base64 encoded ciphertext + ciphertext_blob = base64.b64decode(encrypted_value) + + # Set up the parameters for the decrypt call + params = {"CiphertextBlob": ciphertext_blob} + + # Perform the decryption + response = client.decrypt(**params) + + # Extract and decode the plaintext + plaintext = response["Plaintext"] + secret = plaintext.decode("utf-8") elif key_manager == KeyManagementSystem.AWS_SECRET_MANAGER.value: try: get_secret_value_response = client.get_secret_value( @@ -10193,10 +10214,14 @@ def get_secret( for k, v in secret_dict.items(): secret = v print_verbose(f"secret: {secret}") + elif key_manager == "local": + secret = os.getenv(secret_name) else: # assume the default is infisicial client secret = client.get_secret(secret_name).secret_value except Exception as e: # check if it's in os.environ - print_verbose(f"An exception occurred - {str(e)}") + verbose_logger.error( + f"An exception occurred - {str(e)}\n\n{traceback.format_exc()}" + ) secret = os.getenv(secret_name) try: secret_value_as_bool = ast.literal_eval(secret) From 233d08123d79d4f85a595c49497fafbaa3faa62e Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 6 Jun 2024 15:38:39 -0700 Subject: [PATCH 2/5] docs(secret.md): add doc on using AWS KMS with proxy --- docs/my-website/docs/secret.md | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/docs/my-website/docs/secret.md b/docs/my-website/docs/secret.md index 2b945837ad..08c2e89d1f 100644 --- a/docs/my-website/docs/secret.md +++ b/docs/my-website/docs/secret.md @@ -1,11 +1,31 @@ # Secret Manager LiteLLM supports reading secrets from Azure Key Vault and Infisical +- AWS Key Managemenet Service +- AWS Secret Manager - [Azure Key Vault](#azure-key-vault) - Google Key Management Service - [Infisical Secret Manager](#infisical-secret-manager) - [.env Files](#env-files) +## AWS Key Management Service + +Use AWS KMS to storing a hashed copy of your Proxy Master Key in the environment. + +```bash +export LITELLM_MASTER_KEY="djZ9xjVaZ..." # 👈 ENCRYPTED KEY +export AWS_REGION_NAME="us-west-2" +``` + +```yaml +general_settings: + key_management_system: "aws_kms" + key_management_settings: + hosted_keys: ["LITELLM_MASTER_KEY"] # 👈 WHICH KEYS ARE STORED ON KMS +``` + +[**See Decryption Code**](https://github.com/BerriAI/litellm/blob/a2da2a8f168d45648b61279d4795d647d94f90c9/litellm/utils.py#L10182) + ## AWS Secret Manager Store your proxy keys in AWS Secret Manager. From 35ee506cd6a34a75a3efbabc2c79f1ed93ac3e58 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 6 Jun 2024 16:43:12 -0700 Subject: [PATCH 3/5] test(test_assistants.py): add more error logging --- litellm/tests/test_assistants.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/litellm/tests/test_assistants.py b/litellm/tests/test_assistants.py index 5f565f67cc..d5f047a092 100644 --- a/litellm/tests/test_assistants.py +++ b/litellm/tests/test_assistants.py @@ -198,7 +198,11 @@ async def test_aarun_thread_litellm(sync_mode, provider, is_streaming): ) assert isinstance(messages.data[0], Message) else: - pytest.fail("An unexpected error occurred when running the thread") + pytest.fail( + "An unexpected error occurred when running the thread, {}".format( + run + ) + ) else: added_message = await litellm.a_add_message(**data) @@ -226,4 +230,8 @@ async def test_aarun_thread_litellm(sync_mode, provider, is_streaming): ) assert isinstance(messages.data[0], Message) else: - pytest.fail("An unexpected error occurred when running the thread") + pytest.fail( + "An unexpected error occurred when running the thread, {}".format( + run + ) + ) From 99cf716527902a516870975d9a02cc47f13c6214 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 6 Jun 2024 22:17:30 -0700 Subject: [PATCH 4/5] test(test_completion.py): fix replicate test which was hanging --- litellm/tests/test_completion.py | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index 47c55ca4f3..9d1fa33b1e 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -2538,6 +2538,7 @@ def test_replicate_custom_prompt_dict(): "content": "what is yc write 1 paragraph", } ], + mock_response="Hello world", repetition_penalty=0.1, num_retries=3, ) From dbd442dc48cf6ff1b9cfe8a263b4cc8113e588f5 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 7 Jun 2024 08:01:35 -0700 Subject: [PATCH 5/5] fix(aws_secret_manager.py): fix boto import --- litellm/proxy/secret_managers/aws_secret_manager.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/secret_managers/aws_secret_manager.py b/litellm/proxy/secret_managers/aws_secret_manager.py index 9e6a777c40..8dd6772cf7 100644 --- a/litellm/proxy/secret_managers/aws_secret_manager.py +++ b/litellm/proxy/secret_managers/aws_secret_manager.py @@ -8,8 +8,8 @@ Requires: * `pip install boto3>=1.28.57` """ -import boto3.session -import litellm, os +import litellm +import os from typing import Optional from litellm.proxy._types import KeyManagementSystem