diff --git a/safetytooling/utils/utils.py b/safetytooling/utils/utils.py index ff4651c..7976b18 100644 --- a/safetytooling/utils/utils.py +++ b/safetytooling/utils/utils.py @@ -64,6 +64,15 @@ def setup_environment( os.environ["ANTHROPIC_API_KEY"] = os.environ[anthropic_tag] if openrouter_tag in os.environ: os.environ["OPENROUTER_API_KEY"] = os.environ[openrouter_tag] + + # Blank values in .env are loaded as ""; treat any var ending in _API_KEY or HF_TOKEN as unset. + _other_secret_env_vars = ("HF_TOKEN",) + for key in list(os.environ): + if (key.endswith("_API_KEY") or key in _other_secret_env_vars) and ( + (os.environ[key] or "").strip() == "" + ): + del os.environ[key] + # warn if we do not have an openai api key if "OPENAI_API_KEY" not in os.environ: LOGGER.warning("OPENAI_API_KEY not found in environment, OpenAI API will not be available") diff --git a/tests/test_utils.py b/tests/test_utils.py new file mode 100644 index 0000000..83ec0fa --- /dev/null +++ b/tests/test_utils.py @@ -0,0 +1,35 @@ +import os +from unittest.mock import patch + +from safetytooling.utils import utils + + +def test_setup_environment_treats_blank_api_keys_as_unset(): + with patch("safetytooling.utils.utils.dotenv.load_dotenv", return_value=True): + os.environ["TOGETHER_API_KEY"] = "" + try: + utils.setup_environment() + # Should be removed so downstream code gets None, not "" + assert "TOGETHER_API_KEY" not in os.environ + finally: + os.environ.pop("TOGETHER_API_KEY", None) + + +def test_setup_environment_treats_whitespace_only_api_keys_as_unset(): + with patch("safetytooling.utils.utils.dotenv.load_dotenv", return_value=True): + os.environ["TOGETHER_API_KEY"] = " \t " + try: + utils.setup_environment() + assert "TOGETHER_API_KEY" not in os.environ + finally: + os.environ.pop("TOGETHER_API_KEY", None) + + +def test_setup_environment_normalizes_any_var_ending_in_api_key(): + with patch("safetytooling.utils.utils.dotenv.load_dotenv", return_value=True): + os.environ["SOME_FUTURE_API_KEY"] = "" + try: + utils.setup_environment() + assert "SOME_FUTURE_API_KEY" not in os.environ + finally: + os.environ.pop("SOME_FUTURE_API_KEY", None)