fix(aws_secret_manager.py): accept 'aws_kms' being in the key name

This commit is contained in:
Krrish Dholakia 2024-07-01 13:46:20 -07:00
parent a389f6bff4
commit 223494a0e0

View file

@ -11,6 +11,7 @@ Requires:
import ast
import base64
import os
import re
from typing import Any, Dict, Optional
import litellm
@ -145,9 +146,14 @@ def decrypt_env_var() -> Dict[str, Any]:
# iterate through env - for `aws_kms/`
new_values = {}
for k, v in os.environ.items():
if v is not None and isinstance(v, str) and v.startswith("aws_kms/"):
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