fix (#238) non api key warning on default model providers

Author: nabinchha

src/data_designer/config/default_model_settings.py

@@ -71,12 +71,20 @@ def get_default_model_configs() -> list[ModelConfig]:
     return []
 
 
-def get_default_model_providers_missing_api_keys() -> list[str]:
-    missing_api_keys = []
-    for predefined_provider in PREDEFINED_PROVIDERS:
-        if os.environ.get(predefined_provider["api_key"]) is None:
-            missing_api_keys.append(predefined_provider["api_key"])
-    return missing_api_keys
+def get_providers_with_missing_api_keys(providers: list[ModelProvider]) -> list[ModelProvider]:
+    providers_with_missing_keys = []
+
+    for provider in providers:
+        if provider.api_key is None:
+            # No API key specified at all
+            providers_with_missing_keys.append(provider)
+        elif provider.api_key.isupper() and "_" in provider.api_key:
+            # Looks like an environment variable name, check if it's set
+            if os.environ.get(provider.api_key) is None:
+                providers_with_missing_keys.append(provider)
+        # else: It's an actual API key value (not an env var), so it's valid
+
+    return providers_with_missing_keys
 
 
 def get_default_providers() -> list[ModelProvider]:

src/data_designer/interface/data_designer.py

@@ -12,9 +12,9 @@
 from data_designer.config.data_designer_config import DataDesignerConfig
 from data_designer.config.default_model_settings import (
     get_default_model_configs,
-    get_default_model_providers_missing_api_keys,
     get_default_provider_name,
     get_default_providers,
+    get_providers_with_missing_api_keys,
 )
 from data_designer.config.interface import DataDesignerInterface
 from data_designer.config.models import (
@@ -28,7 +28,6 @@
     MANAGED_ASSETS_PATH,
     MODEL_CONFIGS_FILE_PATH,
     MODEL_PROVIDERS_FILE_PATH,
-    PREDEFINED_PROVIDERS,
 )
 from data_designer.config.utils.info import InfoType, InterfaceInfo
 from data_designer.engine.analysis.dataset_profiler import DataDesignerDatasetProfiler, DatasetProfilerConfig
@@ -334,8 +333,11 @@ def set_run_config(self, run_config: RunConfig) -> None:
     def _resolve_model_providers(self, model_providers: list[ModelProvider] | None) -> list[ModelProvider]:
         if model_providers is None:
             model_providers = get_default_providers()
-            missing_api_keys = get_default_model_providers_missing_api_keys()
-            if len(missing_api_keys) == len(PREDEFINED_PROVIDERS):
+            # Check which providers have missing API keys (from YAML file or env vars)
+            providers_with_missing_keys = get_providers_with_missing_api_keys(model_providers)
+
+            if len(providers_with_missing_keys) == len(model_providers):
+                # All providers have missing API keys
                 logger.warning(
                     "🚨 You are trying to use a default model provider but your API keys are missing."
                     "\n\t\t\tSet the API key for the default providers you intend to use and re-initialize the Data Designer object."

tests/config/test_default_model_settings.py

@@ -13,12 +13,12 @@
     get_builtin_model_providers,
     get_default_inference_parameters,
     get_default_model_configs,
-    get_default_model_providers_missing_api_keys,
     get_default_provider_name,
     get_default_providers,
+    get_providers_with_missing_api_keys,
     resolve_seed_default_model_settings,
 )
-from data_designer.config.models import ChatCompletionInferenceParams, EmbeddingInferenceParams
+from data_designer.config.models import ChatCompletionInferenceParams, EmbeddingInferenceParams, ModelProvider
 from data_designer.config.utils.visualization import get_nvidia_api_key, get_openai_api_key
 
 
@@ -190,7 +190,77 @@ def test_resolve_seed_default_model_settings(tmp_path: Path):
             assert providers_data == {"providers": [p.model_dump() for p in get_builtin_model_providers()]}
 
 
-@patch("data_designer.config.default_model_settings.os.environ.get")
-def test_get_default_model_providers_missing_api_keys(mock_environ_get):
-    mock_environ_get.return_value = None
-    assert get_default_model_providers_missing_api_keys() == ["NVIDIA_API_KEY", "OPENAI_API_KEY", "OPENROUTER_API_KEY"]
+def test_get_providers_with_missing_api_keys():
+    """Test detection of providers with missing API keys."""
+    # Test providers with various API key configurations
+    providers = [
+        ModelProvider(name="provider1", endpoint="http://test1.com", api_key="NVIDIA_API_KEY"),  # env var
+        ModelProvider(name="provider2", endpoint="http://test2.com", api_key="sk-actual-key-12345"),  # actual key
+        ModelProvider(name="provider3", endpoint="http://test3.com", api_key=None),  # no key
+    ]
+
+    with patch("data_designer.config.default_model_settings.os.environ.get") as mock_env:
+        # Mock env to have NVIDIA_API_KEY set but not MISSING_VAR
+        def mock_get(key: str) -> str | None:
+            return "test-key" if key == "NVIDIA_API_KEY" else None
+
+        mock_env.side_effect = mock_get
+
+        missing = get_providers_with_missing_api_keys(providers)
+
+        # provider1 has env var set -> OK
+        # provider2 has actual API key -> OK
+        # provider3 has no API key -> MISSING
+        assert len(missing) == 1
+        assert missing[0].name == "provider3"
+
+
+def test_get_providers_with_missing_api_keys_env_var_not_set():
+    """Test detection when environment variable is not set."""
+    providers = [
+        ModelProvider(name="provider1", endpoint="http://test1.com", api_key="MISSING_ENV_VAR"),
+    ]
+
+    with patch("data_designer.config.default_model_settings.os.environ.get", return_value=None):
+        missing = get_providers_with_missing_api_keys(providers)
+        assert len(missing) == 1
+        assert missing[0].name == "provider1"
+
+
+def test_get_providers_with_missing_api_keys_all_valid():
+    """Test when all providers have valid API keys."""
+    providers = [
+        ModelProvider(name="provider1", endpoint="http://test1.com", api_key="sk-actual-key-1"),
+        ModelProvider(name="provider2", endpoint="http://test2.com", api_key="sk-actual-key-2"),
+    ]
+
+    missing = get_providers_with_missing_api_keys(providers)
+    assert len(missing) == 0
+
+
+def test_get_providers_with_missing_api_keys_all_missing():
+    """Test when all providers have missing API keys."""
+    providers = [
+        ModelProvider(name="provider1", endpoint="http://test1.com", api_key="MISSING_VAR_1"),
+        ModelProvider(name="provider2", endpoint="http://test2.com", api_key=None),
+    ]
+
+    with patch("data_designer.config.default_model_settings.os.environ.get", return_value=None):
+        missing = get_providers_with_missing_api_keys(providers)
+        assert len(missing) == 2
+        assert {p.name for p in missing} == {"provider1", "provider2"}
+
+
+def test_get_providers_with_missing_api_keys_mixed_case():
+    """Test that lowercase API keys are treated as actual keys, not env vars."""
+    providers = [
+        ModelProvider(name="provider1", endpoint="http://test1.com", api_key="lowercase_key"),
+        ModelProvider(name="provider2", endpoint="http://test2.com", api_key="UPPERCASE_KEY"),
+    ]
+
+    with patch("data_designer.config.default_model_settings.os.environ.get", return_value=None):
+        missing = get_providers_with_missing_api_keys(providers)
+        # provider1 has lowercase key (treated as actual key) -> OK
+        # provider2 has uppercase key but env var not set -> MISSING
+        assert len(missing) == 1
+        assert missing[0].name == "provider2"