v2.0: Custom z.ai ModelProvider

app.py

@@ -2,7 +2,6 @@ import os
 from contextlib import asynccontextmanager
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
-from typing import Any
 import tempfile
 
 from models import (
@@ -11,14 +10,7 @@ from models import (
     HealthResponse, ZaiModel
 )
 
-# Configure z.ai - use OPENAI_API_BASE for custom endpoint
 ZAI_API_KEY = os.environ.get("ZAI_API_KEY", "")
-ZAI_BASE_URL = "https://api.z.ai/api/anthropic"
-
-# For OpenAI-compatible APIs, set these env vars
-os.environ["OPENAI_API_KEY"] = ZAI_API_KEY
-os.environ["OPENAI_API_BASE"] = ZAI_BASE_URL
-
 data_designer = None
 
 
@@ -26,14 +18,27 @@ data_designer = None
 async def lifespan(app: FastAPI):
     global data_designer
     from data_designer.interface import DataDesigner
-    data_designer = DataDesigner(artifact_path=tempfile.gettempdir())
+    from data_designer.config.models import ModelProvider
+    
+    # Create custom z.ai provider
+    zai_provider = ModelProvider(
+        name="zai",
+        endpoint="https://api.z.ai/api/anthropic",
+        provider_type="openai",
+        api_key="ZAI_API_KEY",
+    )
+    
+    data_designer = DataDesigner(
+        artifact_path=tempfile.gettempdir(),
+        model_providers=[zai_provider]
+    )
     yield
 
 
 app = FastAPI(
     title="NeMo DataDesigner API",
-    description="Synthetic data generation with NVIDIA NeMo DataDesigner and z.ai",
-    version="1.2.0",
+    description="Synthetic data generation with z.ai",
+    version="2.0.0",
     lifespan=lifespan
 )
 
@@ -48,7 +53,7 @@ app.add_middleware(
 
 def build_config(request):
     import data_designer.config as dd
-    from data_designer.config.models import ModelConfig, ChatCompletionInferenceParams, ModelProvider
+    from data_designer.config.models import ModelConfig, ChatCompletionInferenceParams
 
     config_builder = dd.DataDesignerConfigBuilder()
     model_id = request.model.value
@@ -74,11 +79,11 @@ def build_config(request):
                 )
             )
 
-    # Use openai provider - LiteLLM will use OPENAI_API_BASE env var
+    # Use zai provider with openai format
     model_config = ModelConfig(
         alias="zai-model",
         model=f"openai/{model_id}",
-        provider="openai",
+        provider="zai",
         inference_parameters=ChatCompletionInferenceParams(
             temperature=request.temperature,
             max_tokens=request.max_tokens,