Add support for multiple default model providers and config

Author: nabinchha

src/data_designer/config/models.py

@@ -21,6 +21,8 @@
     MIN_TOP_P,
     NVIDIA_API_KEY_ENV_VAR_NAME,
     NVIDIA_PROVIDER_NAME,
+    OPENAI_API_KEY_ENV_VAR_NAME,
+    OPENAI_PROVIDER_NAME,
 )
 from .utils.io_helpers import smart_load_yaml
 
@@ -233,38 +235,98 @@ def load_model_configs(model_configs: Union[list[ModelConfig], str, Path]) -> li
     return [ModelConfig.model_validate(mc) for mc in json_config["model_configs"]]
 
 
+def get_default_text_alias_inference_parameters() -> InferenceParameters:
+    return InferenceParameters(
+        temperature=0.85,
+        top_p=0.95,
+    )
+
+
+def get_default_reasoning_alias_inference_parameters() -> InferenceParameters:
+    return InferenceParameters(
+        temperature=0.35,
+        top_p=0.95,
+    )
+
+
+def get_default_vision_alias_inference_parameters() -> InferenceParameters:
+    return InferenceParameters(
+        temperature=0.85,
+        top_p=0.95,
+    )
+
+
 def get_default_nvidia_model_configs() -> list[ModelConfig]:
     if not get_nvidia_api_key():
         logger.warning(
-            "‼️🔑 'NVIDIA_API_KEY' environment variable is not set. Please set it to your API key from 'build.nvidia.com' if you want to use the default NVIDIA model configs."
+            f"🔑 {NVIDIA_API_KEY_ENV_VAR_NAME!r} environment variable is not set. Please set it to your API key from 'https://build.nvidia.com' if you want to use the default NVIDIA model configs."
         )
+        return []
     return [
         ModelConfig(
-            alias="text",
+            alias=f"{NVIDIA_PROVIDER_NAME}-text",
             model="nvidia/nvidia-nemotron-nano-9b-v2",
             provider=NVIDIA_PROVIDER_NAME,
-            inference_parameters=InferenceParameters(
-                temperature=0.85,
-                top_p=0.95,
-            ),
+            inference_parameters=get_default_text_alias_inference_parameters(),
         ),
         ModelConfig(
-            alias="reasoning",
-            model="nvidia/llama-3.3-nemotron-super-49b-v1.5",
+            alias=f"{NVIDIA_PROVIDER_NAME}-reasoning",
+            model="openai/gpt-oss-20b",
             provider=NVIDIA_PROVIDER_NAME,
-            inference_parameters=InferenceParameters(
-                temperature=0.35,
-                top_p=0.95,
-            ),
+            inference_parameters=get_default_reasoning_alias_inference_parameters(),
         ),
         ModelConfig(
-            alias="vision",
+            alias=f"{NVIDIA_PROVIDER_NAME}-vision",
             model="nvidia/nemotron-nano-12b-v2-vl",
             provider=NVIDIA_PROVIDER_NAME,
-            inference_parameters=InferenceParameters(
-                temperature=0.85,
-                top_p=0.95,
-            ),
+            inference_parameters=get_default_vision_alias_inference_parameters(),
+        ),
+    ]
+
+
+def get_default_openai_model_configs() -> list[ModelConfig]:
+    if not get_openai_api_key():
+        logger.warning(
+            f"🔑 {OPENAI_API_KEY_ENV_VAR_NAME!r} environment variable is not set. Please set it to your API key from 'https://platform.openai.com/api-keys' if you want to use the default OpenAI model configs."
+        )
+        return []
+    return [
+        ModelConfig(
+            alias=f"{OPENAI_PROVIDER_NAME}-text",
+            model="gpt-4.1",
+            provider=OPENAI_PROVIDER_NAME,
+            inference_parameters=get_default_text_alias_inference_parameters(),
+        ),
+        ModelConfig(
+            alias=f"{OPENAI_PROVIDER_NAME}-reasoning",
+            model="gpt-5",
+            provider=OPENAI_PROVIDER_NAME,
+            inference_parameters=get_default_reasoning_alias_inference_parameters(),
+        ),
+        ModelConfig(
+            alias=f"{OPENAI_PROVIDER_NAME}-vision",
+            model="gpt-5",
+            provider=OPENAI_PROVIDER_NAME,
+            inference_parameters=get_default_vision_alias_inference_parameters(),
+        ),
+    ]
+
+
+def get_default_model_configs() -> list[ModelConfig]:
+    return get_default_nvidia_model_configs() + get_default_openai_model_configs()
+
+
+def get_default_providers() -> list[ModelProvider]:
+    return [
+        ModelProvider(
+            name=NVIDIA_PROVIDER_NAME,
+            endpoint="https://integrate.api.nvidia.com/v1",
+            api_key=NVIDIA_API_KEY_ENV_VAR_NAME,
+        ),
+        ModelProvider(
+            name=OPENAI_PROVIDER_NAME,
+            endpoint="https://api.openai.com/v1",
+            api_key=OPENAI_API_KEY_ENV_VAR_NAME,
         ),
     ]
 
@@ -273,9 +335,5 @@ def get_nvidia_api_key() -> Optional[str]:
     return os.getenv(NVIDIA_API_KEY_ENV_VAR_NAME)
 
 
-def get_default_nvidia_model_provider() -> ModelProvider:
-    return ModelProvider(
-        name=NVIDIA_PROVIDER_NAME,
-        endpoint="https://integrate.api.nvidia.com/v1",
-        api_key=NVIDIA_API_KEY_ENV_VAR_NAME,
-    )
+def get_openai_api_key() -> Optional[str]:
+    return os.getenv(OPENAI_API_KEY_ENV_VAR_NAME)

src/data_designer/config/utils/constants.py

@@ -260,3 +260,6 @@ class NordColor(Enum):
 
 NVIDIA_PROVIDER_NAME = "nvidia"
 NVIDIA_API_KEY_ENV_VAR_NAME = "NVIDIA_API_KEY"
+
+OPENAI_PROVIDER_NAME = "openai"
+OPENAI_API_KEY_ENV_VAR_NAME = "OPENAI_API_KEY"

src/data_designer/config/utils/visualization.py

@@ -22,9 +22,10 @@
 
 from ..base import ConfigBase
 from ..columns import DataDesignerColumnType
-from ..models import ModelConfig, ModelProvider
+from ..models import ModelConfig, ModelProvider, get_nvidia_api_key, get_openai_api_key
 from ..sampler_params import SamplerType
 from .code_lang import code_lang_to_syntax_lexer
+from .constants import NVIDIA_API_KEY_ENV_VAR_NAME, OPENAI_API_KEY_ENV_VAR_NAME
 from .errors import DatasetSampleDisplayError
 
 if TYPE_CHECKING:
@@ -274,7 +275,15 @@ def display_model_configs_table(model_configs: list[ModelConfig]) -> None:
             str(model_config.inference_parameters.temperature),
             str(model_config.inference_parameters.top_p),
         )
-    group = Group(Rule(title="Model Configs"), table_model_configs)
+    group_args: list = [Rule(title="Model Configs"), table_model_configs]
+    if len(model_configs) == 0:
+        subtitle = Text(
+            "‼️ No model configs found. Please provide at least one model config to the config builder",
+            style="dim",
+            justify="center",
+        )
+        group_args.insert(1, subtitle)
+    group = Group(*group_args)
     console.print(group)
 
 
@@ -284,7 +293,18 @@ def display_model_providers_table(model_providers: list[ModelProvider]) -> None:
     table_model_providers.add_column("Endpoint")
     table_model_providers.add_column("API Key")
     for model_provider in model_providers:
-        table_model_providers.add_row(model_provider.name, model_provider.endpoint, model_provider.api_key)
+        api_key = model_provider.api_key
+        if model_provider.api_key == OPENAI_API_KEY_ENV_VAR_NAME:
+            if get_openai_api_key() is not None:
+                api_key = get_openai_api_key()[:1] + "********"
+            else:
+                api_key = f"* {OPENAI_API_KEY_ENV_VAR_NAME!r} not set in environment variables * "
+        elif model_provider.api_key == NVIDIA_API_KEY_ENV_VAR_NAME:
+            if get_nvidia_api_key() is not None:
+                api_key = get_nvidia_api_key()[:1] + "********"
+            else:
+                api_key = f"* {NVIDIA_API_KEY_ENV_VAR_NAME!r} not set in environment variables *"
+        table_model_providers.add_row(model_provider.name, model_provider.endpoint, api_key)
     group = Group(Rule(title="Model Providers"), table_model_providers)
     console.print(group)
 

src/data_designer/engine/models/registry.py

@@ -71,7 +71,9 @@ def get_model_provider(self, *, model_alias: str) -> ModelProvider:
     def run_health_check(self) -> None:
         logger.info("🩺 Running health checks for models...")
         for model in self._models.values():
-            logger.info(f"  |-- 👀 Checking '{model.model_name}'...")
+            logger.info(
+                f"  |-- 👀 Checking {model.model_name!r} in provider named {model.model_provider_name!r} for model alias {model.model_alias!r}..."
+            )
             try:
                 model.generate(
                     prompt="Hello!",

src/data_designer/interface/data_designer.py

@@ -12,12 +12,16 @@
 from data_designer.config.models import (
     ModelConfig,
     ModelProvider,
-    get_default_nvidia_model_configs,
-    get_default_nvidia_model_provider,
+    get_default_model_configs,
+    get_default_providers,
 )
 from data_designer.config.preview_results import PreviewResults
 from data_designer.config.seed import LocalSeedDatasetReference
-from data_designer.config.utils.constants import DEFAULT_NUM_RECORDS
+from data_designer.config.utils.constants import (
+    DEFAULT_NUM_RECORDS,
+    NVIDIA_API_KEY_ENV_VAR_NAME,
+    OPENAI_API_KEY_ENV_VAR_NAME,
+)
 from data_designer.config.utils.info import InterfaceInfo
 from data_designer.config.utils.io_helpers import write_seed_dataset
 from data_designer.engine.analysis.dataset_profiler import (
@@ -227,10 +231,15 @@ def preview(
         )
 
     def get_default_model_configs(self) -> list[ModelConfig]:
-        return get_default_nvidia_model_configs()
+        model_configs = get_default_model_configs()
+        if len(model_configs) == 0:
+            logger.warning(
+                f"‼️ Neither {NVIDIA_API_KEY_ENV_VAR_NAME!r} nor {OPENAI_API_KEY_ENV_VAR_NAME!r} environment variables are set. Please set at least one of them if you want to use the default model configs."
+            )
+        return model_configs
 
     def get_default_model_providers(self) -> list[ModelProvider]:
-        return [get_default_nvidia_model_provider()]
+        return get_default_providers()
 
     def set_buffer_size(self, buffer_size: int) -> None:
         """Set the buffer size for dataset generation.