From fa8402ea09856c974d4cfdea8d25ee105cb95ba6 Mon Sep 17 00:00:00 2001 From: Eugene Yurtsev Date: Tue, 3 Sep 2024 11:39:54 -0400 Subject: [PATCH] core[minor]: Add support for multiple env keys for secrets_from_env (#25971) - Add support to look up secret using more than one env variable - Add overload to help mypy Needed for https://github.com/langchain-ai/langchain/pull/25491 --- libs/core/langchain_core/utils/utils.py | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/libs/core/langchain_core/utils/utils.py b/libs/core/langchain_core/utils/utils.py index 2a7a124c2ff..d4a29da9262 100644 --- a/libs/core/langchain_core/utils/utils.py +++ b/libs/core/langchain_core/utils/utils.py @@ -302,7 +302,9 @@ def from_env( @overload -def from_env(key: str, /, *, default: None) -> Callable[[], Optional[str]]: ... +def from_env( + key: Union[str, Sequence[str]], /, *, default: None +) -> Callable[[], Optional[str]]: ... def from_env( @@ -360,7 +362,7 @@ def secret_from_env(key: str, /, *, default: str) -> Callable[[], SecretStr]: .. @overload def secret_from_env( - key: str, /, *, default: None + key: Union[str, Sequence[str]], /, *, default: None ) -> Callable[[], Optional[SecretStr]]: ... @@ -369,7 +371,7 @@ def secret_from_env(key: str, /, *, error_message: str) -> Callable[[], SecretSt def secret_from_env( - key: str, + key: Union[str, Sequence[str]], /, *, default: Union[str, _NoDefaultType, None] = _NoDefault, @@ -390,9 +392,14 @@ def secret_from_env( def get_secret_from_env() -> Optional[SecretStr]: """Get a value from an environment variable.""" - if key in os.environ: - return SecretStr(os.environ[key]) - elif isinstance(default, str): + if isinstance(key, (list, tuple)): + for k in key: + if k in os.environ: + return SecretStr(os.environ[k]) + if isinstance(key, str): + if key in os.environ: + return SecretStr(os.environ[key]) + if isinstance(default, str): return SecretStr(default) elif isinstance(default, type(None)): return None