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 ast
import base64 import base64
import os import os
import re
from typing import Any, Dict, Optional from typing import Any, Dict, Optional
import litellm import litellm
@ -145,9 +146,14 @@ def decrypt_env_var() -> Dict[str, Any]:
# iterate through env - for `aws_kms/` # iterate through env - for `aws_kms/`
new_values = {} new_values = {}
for k, v in os.environ.items(): 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) decrypted_value = aws_kms.decrypt_value(secret_name=k)
# reset env var # reset env var
k = re.sub("litellm_secret_aws_kms", "", k, flags=re.IGNORECASE)
new_values[k] = decrypted_value new_values[k] = decrypted_value
return new_values return new_values