two version of R2R are here HEAD master

1 files changed, 143 insertions, 0 deletions

diff --git a/.venv/lib/python3.12/site-packages/litellm/secret_managers/aws_secret_manager.py b/.venv/lib/python3.12/site-packages/litellm/secret_managers/aws_secret_manager.py
new file mode 100644
index 00000000..fbe951e6
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/litellm/secret_managers/aws_secret_manager.py
@@ -0,0 +1,143 @@
+"""
+This is a file for the AWS Secret Manager Integration
+
+Relevant issue: https://github.com/BerriAI/litellm/issues/1883
+
+Requires:
+* `os.environ["AWS_REGION_NAME"], 
+* `pip install boto3>=1.28.57`
+"""
+
+import ast
+import base64
+import os
+import re
+from typing import Any, Dict, Optional
+
+import litellm
+from litellm.proxy._types import KeyManagementSystem
+
+
+def validate_environment():
+    if "AWS_REGION_NAME" not in os.environ:
+        raise ValueError("Missing required environment variable - AWS_REGION_NAME")
+
+
+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
+
+
+class AWSKeyManagementService_V2:
+    """
+    V2 Clean Class for decrypting keys from AWS KeyManagementService
+    """
+
+    def __init__(self) -> None:
+        self.validate_environment()
+        self.kms_client = self.load_aws_kms(use_aws_kms=True)
+
+    def validate_environment(
+        self,
+    ):
+        if "AWS_REGION_NAME" not in os.environ:
+            raise ValueError("Missing required environment variable - AWS_REGION_NAME")
+
+        ## CHECK IF LICENSE IN ENV ## - premium feature
+        is_litellm_license_in_env: bool = False
+
+        if os.getenv("LITELLM_LICENSE", None) is not None:
+            is_litellm_license_in_env = True
+        elif os.getenv("LITELLM_SECRET_AWS_KMS_LITELLM_LICENSE", None) is not None:
+            is_litellm_license_in_env = True
+        if is_litellm_license_in_env is False:
+            raise ValueError(
+                "AWSKeyManagementService V2 is an Enterprise Feature. Please add a valid LITELLM_LICENSE to your envionment."
+            )
+
+    def load_aws_kms(self, 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"))
+
+            return kms_client
+        except Exception as e:
+            raise e
+
+    def decrypt_value(self, secret_name: str) -> Any:
+        if self.kms_client is None:
+            raise ValueError("kms_client is None")
+        encrypted_value = os.getenv(secret_name, None)
+        if encrypted_value is None:
+            raise Exception(
+                "AWS KMS - Encrypted Value of Key={} is None".format(secret_name)
+            )
+        if isinstance(encrypted_value, str) and encrypted_value.startswith("aws_kms/"):
+            encrypted_value = encrypted_value.replace("aws_kms/", "")
+
+        # 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 = self.kms_client.decrypt(**params)
+
+        # Extract and decode the plaintext
+        plaintext = response["Plaintext"]
+        secret = plaintext.decode("utf-8")
+        if isinstance(secret, str):
+            secret = secret.strip()
+        try:
+            secret_value_as_bool = ast.literal_eval(secret)
+            if isinstance(secret_value_as_bool, bool):
+                return secret_value_as_bool
+        except Exception:
+            pass
+
+        return secret
+
+
+"""
+- look for all values in the env with `aws_kms/<hashed_key>` 
+- decrypt keys 
+- rewrite env var with decrypted key (). Note: this environment variable will only be available to the current process and any child processes spawned from it. Once the Python script ends, the environment variable will not persist.
+"""
+
+
+def decrypt_env_var() -> Dict[str, Any]:
+    # setup client class
+    aws_kms = AWSKeyManagementService_V2()
+    # iterate through env - for `aws_kms/`
+    new_values = {}
+    for k, v in os.environ.items():
+        if (
+            k is not None
+            and isinstance(k, str)
+            and k.lower().startswith("litellm_secret_aws_kms")
+        ) or (v is not None and isinstance(v, str) and v.startswith("aws_kms/")):
+            decrypted_value = aws_kms.decrypt_value(secret_name=k)
+            # reset env var
+            k = re.sub("litellm_secret_aws_kms_", "", k, flags=re.IGNORECASE)
+            new_values[k] = decrypted_value
+
+    return new_values