12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485 |
- import contextlib
- import os
- import logging
- logger = logging.getLogger("env_config")
- # List of environment variables prefixes that are allowed to be used for configuration.
- env_prefix_whitelist = [
- 'OPENAI',
- 'AZURE_OPENAI'
- ]
- def _obfuscate(secret):
- l = len(secret)
- return '.' * (l - 4) + secret[-4:]
- def _log_env(use_prefix: str, env: dict):
- """
- Logs each name value pair of the given environment. If the name indicates that it might store a secret such as an API key, then obfuscate the value.
- """
- log_prefix = f"'{use_prefix}'" if use_prefix else "no"
- logger.info(f"Resolved OpenAI env vars with {log_prefix} prefix:")
- for key, value in env.items():
- if any(prefix in key for prefix in ['KEY', 'SECRET', 'TOKEN']):
- value = _obfuscate(value)
- logger.info(f" - {key}={value}")
- def read_env_config(use_prefix: str, env: dict = os.environ) -> str:
- """
- Read whitelisted environment variables and return them in a dictionary.
- Overrides the whitelisted environment variable with ones prefixed with the given use_prefix if available.
- """
- config = {}
- for prefix in [None, use_prefix]:
- read_env_config_prefixed(prefix, config, env)
- _log_env(use_prefix, config)
- return config
- def read_env_config_prefixed(use_prefix: str, config: dict, env: dict = os.environ) -> str:
- """
- Read whitelisted environment variables prefixed with use_prefix and adds them to the dictionary
- with use_prefix stripped.
- """
- use_prefix = format_prefix(use_prefix)
- for key in env:
- for env_prefix in env_prefix_whitelist:
- key_prefix = f"{use_prefix}{format_prefix(env_prefix)}"
- if key.startswith(key_prefix):
- striped_key = key.removeprefix(use_prefix)
- config[striped_key] = env[key]
- def format_prefix(prefix: str) -> str:
- """
- Format the prefix to be used in the environment variable.
- """
- if prefix and len(prefix) > 0 and not prefix.endswith("_"):
- prefix = f"{prefix}_"
- if not prefix:
- prefix = ""
- return prefix
- @contextlib.contextmanager
- def set_env(**environ: dict[str, str]):
- """
- Temporarily set the process environment variables.
- Warning, this is not thread safe as the environment is updated for the whole process.
- >>> with set_env(PLUGINS_DIR='test/plugins'):
- ... "PLUGINS_DIR" in os.environ
- True
- >>> "PLUGINS_DIR" in os.environ
- False
- :type environ: dict[str, unicode]
- :param environ: Environment variables to set
- """
- old_environ = dict(os.environ)
- os.environ.update(environ)
- try:
- yield
- finally:
- os.environ.clear()
- os.environ.update(old_environ)
|