Use z.ai OpenAI-compatible endpoint (v1.3)

app.py

@@ -2,6 +2,7 @@ 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 (
@@ -10,7 +11,14 @@ from models import (
     HealthResponse, ZaiModel
 )
 
+# z.ai OpenAI-compatible endpoint
 ZAI_API_KEY = os.environ.get("ZAI_API_KEY", "")
+ZAI_OPENAI_BASE = "https://api.z.ai/api/paas/v4/"
+
+# Set for LiteLLM
+os.environ["OPENAI_API_KEY"] = ZAI_API_KEY
+os.environ["OPENAI_API_BASE"] = ZAI_OPENAI_BASE
+
 data_designer = None
 
 
@@ -18,27 +26,14 @@ data_designer = None
 async def lifespan(app: FastAPI):
     global data_designer
     from data_designer.interface import DataDesigner
-    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]
-    )
+    data_designer = DataDesigner(artifact_path=tempfile.gettempdir())
     yield
 
 
 app = FastAPI(
     title="NeMo DataDesigner API",
-    description="Synthetic data generation with z.ai",
-    version="2.0.0",
+    description="Synthetic data generation with NVIDIA NeMo DataDesigner and z.ai",
+    version="1.3.0",
     lifespan=lifespan
 )
 
@@ -53,7 +48,7 @@ app.add_middleware(
 
 def build_config(request):
     import data_designer.config as dd
-    from data_designer.config.models import ModelConfig, ChatCompletionInferenceParams
+    from data_designer.config.models import ModelConfig, ChatCompletionInferenceParams, ModelProvider
 
     config_builder = dd.DataDesignerConfigBuilder()
     model_id = request.model.value
@@ -79,19 +74,28 @@ def build_config(request):
                 )
             )
 
-    # Use zai provider with openai format
+    # Custom z.ai provider with OpenAI-compatible endpoint
+    zai_provider = ModelProvider(
+        name="zai",
+        endpoint=ZAI_OPENAI_BASE,
+        api_key="ZAI_API_KEY",
+        provider_type="openai"
+    )
+    
     model_config = ModelConfig(
         alias="zai-model",
-        model=f"openai/{model_id}",
+        model=model_id,  # Just the model name, no prefix
         provider="zai",
         inference_parameters=ChatCompletionInferenceParams(
             temperature=request.temperature,
             max_tokens=request.max_tokens,
         ),
     )
+    
+    # Pass custom provider to config builder
     config_builder.add_model_config(model_config)
-
-    return config_builder
+    
+    return config_builder, zai_provider
 
 
 def get_sampler_params(sampler_type, params):
@@ -121,8 +125,59 @@ async def health():
 @app.post("/generate", response_model=GenerateResponse)
 async def generate(request: GenerateRequest):
     try:
-        config_builder = build_config(request)
-        result = data_designer.create(
+        from data_designer.interface import DataDesigner
+        import data_designer.config as dd
+        from data_designer.config.models import ModelConfig, ChatCompletionInferenceParams, ModelProvider
+        
+        # Rebuild DataDesigner with custom provider
+        zai_provider = ModelProvider(
+            name="zai",
+            endpoint="https://api.z.ai/api/paas/v4/",
+            api_key="ZAI_API_KEY",
+            provider_type="openai"
+        )
+        
+        dd_custom = DataDesigner(
+            artifact_path=tempfile.gettempdir(),
+            model_providers=[zai_provider]
+        )
+        
+        config_builder = dd.DataDesignerConfigBuilder()
+        model_id = request.model.value
+
+        for col in request.columns:
+            if col.type == "sampler":
+                sampler_type_str = col.params.get("sampler_type", "CATEGORY")
+                sampler_type = getattr(dd.SamplerType, sampler_type_str, dd.SamplerType.CATEGORY)
+                params = get_sampler_params(sampler_type, col.params)
+                config_builder.add_column(
+                    dd.SamplerColumnConfig(
+                        name=col.name,
+                        sampler_type=sampler_type,
+                        params=params,
+                    )
+                )
+            elif col.type == "llm_text":
+                config_builder.add_column(
+                    dd.LLMTextColumnConfig(
+                        name=col.name,
+                        model_alias="zai-model",
+                        prompt=col.params.get("prompt", "Generate text"),
+                    )
+                )
+
+        model_config = ModelConfig(
+            alias="zai-model",
+            model=model_id,
+            provider="zai",
+            inference_parameters=ChatCompletionInferenceParams(
+                temperature=request.temperature,
+                max_tokens=request.max_tokens,
+            ),
+        )
+        config_builder.add_model_config(model_config)
+
+        result = dd_custom.create(
             config_builder=config_builder,
             num_records=request.num_records,
             dataset_name="api-dataset"
@@ -131,14 +186,65 @@ async def generate(request: GenerateRequest):
         data = df.to_dict(orient="records")
         return GenerateResponse(success=True, data=data, record_count=len(data))
     except Exception as e:
-        return GenerateResponse(success=False, error=str(e))
+        import traceback
+        return GenerateResponse(success=False, error=f"{str(e)}")
 
 
 @app.post("/preview", response_model=PreviewResponse)
 async def preview(request: PreviewRequest):
     try:
-        config_builder = build_config(request)
-        preview_result = data_designer.preview(config_builder=config_builder, num_records=1)
+        from data_designer.interface import DataDesigner
+        import data_designer.config as dd
+        from data_designer.config.models import ModelConfig, ChatCompletionInferenceParams, ModelProvider
+        
+        zai_provider = ModelProvider(
+            name="zai",
+            endpoint="https://api.z.ai/api/paas/v4/",
+            api_key="ZAI_API_KEY",
+            provider_type="openai"
+        )
+        
+        dd_custom = DataDesigner(
+            artifact_path=tempfile.gettempdir(),
+            model_providers=[zai_provider]
+        )
+        
+        config_builder = dd.DataDesignerConfigBuilder()
+        model_id = request.model.value
+
+        for col in request.columns:
+            if col.type == "sampler":
+                sampler_type_str = col.params.get("sampler_type", "CATEGORY")
+                sampler_type = getattr(dd.SamplerType, sampler_type_str, dd.SamplerType.CATEGORY)
+                params = get_sampler_params(sampler_type, col.params)
+                config_builder.add_column(
+                    dd.SamplerColumnConfig(
+                        name=col.name,
+                        sampler_type=sampler_type,
+                        params=params,
+                    )
+                )
+            elif col.type == "llm_text":
+                config_builder.add_column(
+                    dd.LLMTextColumnConfig(
+                        name=col.name,
+                        model_alias="zai-model",
+                        prompt=col.params.get("prompt", "Generate text"),
+                    )
+                )
+
+        model_config = ModelConfig(
+            alias="zai-model",
+            model=model_id,
+            provider="zai",
+            inference_parameters=ChatCompletionInferenceParams(
+                temperature=request.temperature,
+                max_tokens=request.max_tokens,
+            ),
+        )
+        config_builder.add_model_config(model_config)
+
+        preview_result = dd_custom.preview(config_builder=config_builder, num_records=1)
         sample = preview_result.dataset.to_dict(orient="records")[0] if len(preview_result.dataset) > 0 else {}
         return PreviewResponse(success=True, sample=sample)
     except Exception as e:
@@ -148,9 +254,9 @@ async def preview(request: PreviewRequest):
 @app.get("/models")
 async def list_models():
     return {"models": [
-        {"id": "glm-5", "name": "GLM-5 (Opus)", "description": "Most capable model"},
-        {"id": "glm-4.7", "name": "GLM-4.7 (Sonnet)", "description": "Balanced"},
-        {"id": "glm-4.5-air", "name": "GLM-4.5-Air (Haiku)", "description": "Fast"}
+        {"id": "glm-5", "name": "GLM-5", "description": "Most capable"},
+        {"id": "glm-4.7", "name": "GLM-4.7", "description": "Balanced"},
+        {"id": "glm-4.5-air", "name": "GLM-4.5-Air", "description": "Fast"}
     ]}