[Bug Fix] fix parsing environment_variables from config.yaml (arize logger integration fix) (#12482)

* fix parsing environment_variables from config.yaml

* add tests
This commit is contained in:
Ishaan Jaff 2025-07-10 09:15:07 -07:00 • committed by GitHub
parent 01c76229f2
commit a0d6900223
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 183 additions and 11 deletions

View file

@ -4950,6 +4950,48 @@
"supports_tool_choice": true,
"supports_web_search": true
},
"xai/grok-4": {
"max_tokens": 256000,
"max_input_tokens": 256000,
"max_output_tokens": 256000,
"input_cost_per_token": 3e-06,
"output_cost_per_token": 1.5e-05,
"litellm_provider": "xai",
"mode": "chat",
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_reasoning": true,
"source": "https://docs.x.ai/docs/models",
"supports_web_search": true
},
"xai/grok-4-0709": {
"max_tokens": 256000,
"max_input_tokens": 256000,
"max_output_tokens": 256000,
"input_cost_per_token": 3e-06,
"output_cost_per_token": 1.5e-05,
"litellm_provider": "xai",
"mode": "chat",
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_reasoning": true,
"source": "https://docs.x.ai/docs/models",
"supports_web_search": true
},
"xai/grok-4-latest": {
"max_tokens": 256000,
"max_input_tokens": 256000,
"max_output_tokens": 256000,
"input_cost_per_token": 3e-06,
"output_cost_per_token": 1.5e-05,
"litellm_provider": "xai",
"mode": "chat",
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_reasoning": true,
"source": "https://docs.x.ai/docs/models",
"supports_web_search": true
},
"deepseek/deepseek-coder": {
"max_tokens": 4096,
"max_input_tokens": 128000,
@ -15796,4 +15838,4 @@
"notes": "ElevenLabs Scribe v1 experimental - enhanced version of the main Scribe model"
}
}
}
}

View file

@ -1640,6 +1640,39 @@ class ProxyConfig:
if credential_list_dict:
credential_list = [CredentialItem(**cred) for cred in credential_list_dict]
return credential_list
def _load_environment_variables(self, config: dict):
## ENVIRONMENT VARIABLES
global premium_user
environment_variables = config.get("environment_variables", None)
if environment_variables:
for key, value in environment_variables.items():
#########################################################
# handles this scenario:
# ```yaml
# environment_variables:
# ARIZE_ENDPOINT: os.environ/ARIZE_ENDPOINT
# ```
#########################################################
if isinstance(value, str) and value.startswith("os.environ/"):
resolved_secret_string: Optional[str] = get_secret_str(secret_name=value)
if resolved_secret_string is not None:
os.environ[key] = resolved_secret_string
else:
#########################################################
# handles this scenario:
# ```yaml
# environment_variables:
# ARIZE_ENDPOINT: https://otlp.arize.com/v1
# ```
#########################################################
os.environ[key] = str(value)
# check if litellm_license in general_settings
if "LITELLM_LICENSE" in environment_variables:
_license_check.license_str = os.getenv("LITELLM_LICENSE", None)
premium_user = _license_check.is_premium()
return
async def load_config( # noqa: PLR0915
self, router: Optional[litellm.Router], config_file_path: str
@ -1651,16 +1684,7 @@ class ProxyConfig:
config: dict = await self.get_config(config_file_path=config_file_path)
## ENVIRONMENT VARIABLES
environment_variables = config.get("environment_variables", None)
if environment_variables:
for key, value in environment_variables.items():
os.environ[key] = str(get_secret(secret_name=key, default_value=value))
# check if litellm_license in general_settings
if "LITELLM_LICENSE" in environment_variables:
_license_check.license_str = os.getenv("LITELLM_LICENSE", None)
premium_user = _license_check.is_premium()
self._load_environment_variables(config=config)
## Callback settings
callback_settings = config.get("callback_settings", None)

View file

@ -715,3 +715,109 @@ async def test_custom_ui_sso_sign_in_handler_config_loading():
# Clean up temporary file
import os
os.unlink(config_file_path)
@pytest.mark.asyncio
async def test_load_environment_variables_direct_and_os_environ():
"""
Test _load_environment_variables method with direct values and os.environ/ prefixed values
"""
from unittest.mock import patch
from litellm.proxy.proxy_server import ProxyConfig
proxy_config = ProxyConfig()
# Test config with both direct values and os.environ/ prefixed values
test_config = {
"environment_variables": {
"DIRECT_VAR": "direct_value",
"NUMERIC_VAR": 12345,
"BOOL_VAR": True,
"SECRET_VAR": "os.environ/ACTUAL_SECRET_VAR"
}
}
# Mock get_secret_str to return a resolved value
mock_secret_value = "resolved_secret_value"
with patch("litellm.proxy.proxy_server.get_secret_str", return_value=mock_secret_value) as mock_get_secret:
with patch.dict(os.environ, {}, clear=False): # Don't clear existing env vars, just track changes
# Call the method under test
proxy_config._load_environment_variables(test_config)
# Verify direct environment variables were set correctly
assert os.environ["DIRECT_VAR"] == "direct_value"
assert os.environ["NUMERIC_VAR"] == "12345" # Should be converted to string
assert os.environ["BOOL_VAR"] == "True" # Should be converted to string
# Verify os.environ/ prefixed variable was resolved and set
assert os.environ["SECRET_VAR"] == mock_secret_value
# Verify get_secret_str was called with the correct value
mock_get_secret.assert_called_once_with(secret_name="os.environ/ACTUAL_SECRET_VAR")
@pytest.mark.asyncio
async def test_load_environment_variables_litellm_license_and_edge_cases():
"""
Test _load_environment_variables method with LITELLM_LICENSE special handling and edge cases
"""
from unittest.mock import MagicMock, patch
from litellm.proxy.proxy_server import ProxyConfig
proxy_config = ProxyConfig()
# Test Case 1: LITELLM_LICENSE in environment_variables
test_config_with_license = {
"environment_variables": {
"LITELLM_LICENSE": "test_license_key",
"OTHER_VAR": "other_value"
}
}
# Mock _license_check
mock_license_check = MagicMock()
mock_license_check.is_premium.return_value = True
with patch("litellm.proxy.proxy_server._license_check", mock_license_check):
with patch.dict(os.environ, {}, clear=False):
# Call the method under test
proxy_config._load_environment_variables(test_config_with_license)
# Verify LITELLM_LICENSE was set in environment
assert os.environ["LITELLM_LICENSE"] == "test_license_key"
# Verify license check was updated
assert mock_license_check.license_str == "test_license_key"
mock_license_check.is_premium.assert_called_once()
# Test Case 2: No environment_variables in config
test_config_no_env_vars = {}
# This should not raise any errors and should return without doing anything
result = proxy_config._load_environment_variables(test_config_no_env_vars)
assert result is None # Method returns None
# Test Case 3: environment_variables is None
test_config_none_env_vars = {"environment_variables": None}
# This should not raise any errors and should return without doing anything
result = proxy_config._load_environment_variables(test_config_none_env_vars)
assert result is None # Method returns None
# Test Case 4: os.environ/ prefix but get_secret_str returns None
test_config_secret_none = {
"environment_variables": {
"FAILED_SECRET": "os.environ/NONEXISTENT_SECRET"
}
}
with patch("litellm.proxy.proxy_server.get_secret_str", return_value=None):
with patch.dict(os.environ, {}, clear=False):
# Call the method under test
proxy_config._load_environment_variables(test_config_secret_none)
# Verify that the environment variable was not set when secret resolution fails
assert "FAILED_SECRET" not in os.environ