diff --git a/litellm/llms/openai_like/dynamic_config.py b/litellm/llms/openai_like/dynamic_config.py index 19e29bcdcb2..71300bcc04a 100644 --- a/litellm/llms/openai_like/dynamic_config.py +++ b/litellm/llms/openai_like/dynamic_config.py @@ -17,6 +17,21 @@ from litellm.types.llms.openai import AllMessageValues from .json_loader import SimpleProviderConfig +def _resolve_api_key(provider: SimpleProviderConfig, api_base: str, api_key: str | None) -> str | None: + if api_key: + return api_key + + env_key: Final = get_secret_str(provider.api_key_env) + if not env_key or not provider.require_explicit_key_for_custom_base: + return env_key + + env_base: Final = get_secret_str(provider.api_base_env) if provider.api_base_env else None + trusted_bases: Final = {base.rstrip("/") for base in (provider.base_url, env_base) if base} + if api_base.rstrip("/") not in trusted_bases: + raise ValueError(f"api_key is required for custom api_base on provider {provider.slug}") + return env_key + + def create_config_class(provider: SimpleProviderConfig): """Generate config class dynamically from JSON configuration""" @@ -63,8 +78,7 @@ def create_config_class(provider: SimpleProviderConfig): if not resolved_base: resolved_base = provider.base_url - # Resolve API key - resolved_key: Final = api_key or get_secret_str(provider.api_key_env) + resolved_key: Final = _resolve_api_key(provider, resolved_base, api_key) return resolved_base, resolved_key @@ -200,7 +214,12 @@ def create_responses_config_class(provider: SimpleProviderConfig): litellm_params: GenericLiteLLMParams | None, ) -> dict: litellm_params = litellm_params or GenericLiteLLMParams() - api_key: Final = litellm_params.api_key or get_secret_str(provider.api_key_env) + api_base: Final = ( + litellm_params.api_base + or (get_secret_str(provider.api_base_env) if provider.api_base_env else None) + or provider.base_url + ) + api_key: Final = _resolve_api_key(provider, api_base, litellm_params.api_key) if api_key: headers["Authorization"] = f"Bearer {api_key}" return headers diff --git a/litellm/llms/openai_like/json_loader.py b/litellm/llms/openai_like/json_loader.py index 5cdaff90d24..36056b8e881 100644 --- a/litellm/llms/openai_like/json_loader.py +++ b/litellm/llms/openai_like/json_loader.py @@ -22,6 +22,7 @@ class SimpleProviderConfig: self.constraints = data.get("constraints", {}) self.special_handling = data.get("special_handling", {}) self.supported_endpoints = data.get("supported_endpoints", []) + self.require_explicit_key_for_custom_base = data.get("require_explicit_key_for_custom_base", False) class JSONProviderRegistry: diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index 44f626fd987..35d1e078374 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -205,6 +205,7 @@ "base_url": "https://api.aipowergrid.io/v1", "api_key_env": "AIPG_API_KEY", "api_base_env": "AIPG_API_BASE", + "require_explicit_key_for_custom_base": true, "param_mappings": { "max_completion_tokens": "max_tokens" }, diff --git a/tests/llm_translation/test_aipg.py b/tests/llm_translation/test_aipg.py index d229ee0a0a2..d7cba712ddf 100644 --- a/tests/llm_translation/test_aipg.py +++ b/tests/llm_translation/test_aipg.py @@ -3,6 +3,8 @@ import os from unittest import mock +import pytest + import litellm AIPG_API_BASE = "https://api.aipowergrid.io/v1" @@ -16,6 +18,7 @@ def test_aipg_json_registry(): assert config.base_url == AIPG_API_BASE assert config.api_key_env == "AIPG_API_KEY" assert config.api_base_env == "AIPG_API_BASE" + assert config.require_explicit_key_for_custom_base is True assert config.supported_endpoints == ["/v1/chat/completions", "/v1/responses"] assert JSONProviderRegistry.supports_responses_api("aipg") is True @@ -45,6 +48,24 @@ def test_aipg_get_openai_compatible_provider_info(): assert api_base == "https://explicit.example/v1" assert api_key == "explicit-key" + with ( + mock.patch.dict(os.environ, {"AIPG_API_KEY": "env-key"}, clear=True), + pytest.raises(ValueError, match="api_key is required for custom api_base"), + ): + config._get_openai_compatible_provider_info("https://attacker.example/v1", None) + + with mock.patch.dict( + os.environ, + { + "AIPG_API_KEY": "env-key", + "AIPG_API_BASE": "https://operator.example/v1", + }, + clear=True, + ): + api_base, api_key = config._get_openai_compatible_provider_info("https://operator.example/v1/", None) + assert api_base == "https://operator.example/v1/" + assert api_key == "env-key" + mapped = config.map_openai_params( non_default_params={"max_completion_tokens": 12, "temperature": 0.2}, optional_params={}, @@ -67,6 +88,37 @@ def test_get_llm_provider_aipg(): assert api_base == AIPG_API_BASE +def test_aipg_responses_rejects_env_key_with_custom_api_base(): + from litellm.llms.openai_like.dynamic_config import create_responses_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + from litellm.types.router import GenericLiteLLMParams + + provider = JSONProviderRegistry.get("aipg") + assert provider is not None + config = create_responses_config_class(provider)() + + with ( + mock.patch.dict(os.environ, {"AIPG_API_KEY": "env-key"}, clear=True), + pytest.raises(ValueError, match="api_key is required for custom api_base"), + ): + config.validate_environment( + headers={}, + model="gpt-oss-120b", + litellm_params=GenericLiteLLMParams(api_base="https://attacker.example/v1"), + ) + + with mock.patch.dict(os.environ, {"AIPG_API_KEY": "env-key"}, clear=True): + headers = config.validate_environment( + headers={}, + model="gpt-oss-120b", + litellm_params=GenericLiteLLMParams( + api_base="https://attacker.example/v1", + api_key="explicit-key", + ), + ) + assert headers["Authorization"] == "Bearer explicit-key" + + def test_aipg_model_metadata(): model_cost = litellm.get_model_cost_map(url="") expected = {