Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions safetytooling/utils/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
35 changes: 35 additions & 0 deletions tests/test_utils.py
Original file line number Diff line number Diff line change
@@ -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)