Fix: Use Anthropic provider for z.ai API

app.py

@@ -4,23 +4,21 @@ from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from typing import Any
 import tempfile
-from pathlib import Path
 
 from models import (
     GenerateRequest, GenerateResponse,
     PreviewRequest, PreviewResponse,
-    HealthResponse, ZaiModel, SamplerType
+    HealthResponse, ZaiModel
 )
 
-# Configure z.ai as Anthropic-compatible provider for LiteLLM
+# Configure z.ai as OpenAI-compatible provider
 ZAI_API_KEY = os.environ.get("ZAI_API_KEY", "")
-ZAI_BASE_URL = "https://api.z.ai/api/anthropic"
+ZAI_BASE_URL = "https://api.z.ai/api/openai"
 
-# LiteLLM uses these env vars for custom Anthropic endpoint
-os.environ["ANTHROPIC_API_KEY"] = ZAI_API_KEY
-os.environ["ANTHROPIC_BASE_URL"] = ZAI_BASE_URL
+# LiteLLM OpenAI-compatible config
+os.environ["OPENAI_API_KEY"] = ZAI_API_KEY
+os.environ["OPENAI_API_BASE"] = ZAI_BASE_URL
 
-# Global DataDesigner instance
 data_designer = None
 
 
@@ -54,10 +52,8 @@ def build_config(request: GenerateRequest | PreviewRequest):
     from data_designer.config.models import ModelConfig, ChatCompletionInferenceParams
 
     config_builder = dd.DataDesignerConfigBuilder()
-
     model_id = request.model.value
 
-    # Process columns
     for col in request.columns:
         if col.type == "sampler":
             sampler_type_str = col.params.get("sampler_type", "CATEGORY")
@@ -97,11 +93,10 @@ def build_config(request: GenerateRequest | PreviewRequest):
                 )
             )
 
-    # Add model config - use anthropic provider with custom base_url via env
     model_config = ModelConfig(
         alias="zai-model",
-        model=f"anthropic/{model_id}",  # LiteLLM Anthropic format
-        provider="anthropic",
+        model=f"openai/{model_id}",
+        provider="openai",
         inference_parameters=ChatCompletionInferenceParams(
             temperature=request.temperature,
             max_tokens=request.max_tokens,
@@ -113,116 +108,60 @@ def build_config(request: GenerateRequest | PreviewRequest):
 
 
 def get_sampler_params(sampler_type, params: dict) -> Any:
-    """Get appropriate sampler params based on type."""
     import data_designer.config as dd
-
     type_name = sampler_type.name if hasattr(sampler_type, 'name') else str(sampler_type)
 
     if type_name == "CATEGORY":
-        return dd.CategorySamplerParams(
-            values=params.get("values", ["A", "B", "C"])
-        )
+        return dd.CategorySamplerParams(values=params.get("values", ["A", "B", "C"]))
     elif type_name == "UNIFORM":
-        return dd.UniformSamplerParams(
-            low=params.get("low", 0),
-            high=params.get("high", 100)
-        )
+        return dd.UniformSamplerParams(low=params.get("low", 0), high=params.get("high", 100))
     elif type_name == "GAUSSIAN":
-        return dd.GaussianSamplerParams(
-            mean=params.get("mean", 0),
-            std=params.get("std", 1)
-        )
+        return dd.GaussianSamplerParams(mean=params.get("mean", 0), std=params.get("std", 1))
     elif type_name == "DATETIME":
-        return dd.DateTimeSamplerParams(
-            start_date=params.get("start_date", "2020-01-01"),
-            end_date=params.get("end_date", "2025-12-31")
-        )
+        return dd.DateTimeSamplerParams(start_date=params.get("start_date", "2020-01-01"), end_date=params.get("end_date", "2025-12-31"))
     else:
         return dd.CategorySamplerParams(values=["default"])
 
 
 @app.get("/", response_model=HealthResponse)
 async def root():
-    """Health check endpoint."""
-    return HealthResponse(
-        status="healthy",
-        model="data-designer",
-        api_configured=bool(ZAI_API_KEY)
-    )
+    return HealthResponse(status="healthy", model="data-designer", api_configured=bool(ZAI_API_KEY))
 
 
 @app.get("/health", response_model=HealthResponse)
 async def health():
-    """Health check endpoint."""
-    return HealthResponse(
-        status="healthy",
-        model="data-designer",
-        api_configured=bool(ZAI_API_KEY)
-    )
+    return HealthResponse(status="healthy", model="data-designer", api_configured=bool(ZAI_API_KEY))
 
 
 @app.post("/generate", response_model=GenerateResponse)
 async def generate(request: GenerateRequest):
-    """
-    Generate synthetic data using DataDesigner.create().
-    """
     try:
         config_builder = build_config(request)
-
         result = data_designer.create(
             config_builder=config_builder,
             num_records=request.num_records,
             dataset_name="api-dataset"
         )
-
         df = result.load_dataset()
         data = df.to_dict(orient="records")
-
-        return GenerateResponse(
-            success=True,
-            data=data,
-            record_count=len(data)
-        )
-
+        return GenerateResponse(success=True, data=data, record_count=len(data))
     except Exception as e:
-        import traceback
-        return GenerateResponse(
-            success=False,
-            error=f"{str(e)}"
-        )
+        return GenerateResponse(success=False, error=str(e))
 
 
 @app.post("/preview", response_model=PreviewResponse)
 async def preview(request: PreviewRequest):
-    """
-    Preview a single record using DataDesigner.preview().
-    """
     try:
         config_builder = build_config(request)
-
-        preview_result = data_designer.preview(
-            config_builder=config_builder,
-            num_records=1
-        )
-
+        preview_result = data_designer.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
-        )
-
+        return PreviewResponse(success=True, sample=sample)
     except Exception as e:
-        import traceback
-        return PreviewResponse(
-            success=False,
-            error=f"{str(e)}"
-        )
+        return PreviewResponse(success=False, error=str(e))
 
 
 @app.get("/models")
 async def list_models():
-    """List available z.ai models."""
     return {
         "models": [
             {"id": "glm-5", "name": "GLM-5 (Opus)", "description": "Most capable model"},
@@ -234,14 +173,12 @@ async def list_models():
 
 @app.get("/sampler-types")
 async def list_sampler_types():
-    """List available sampler types."""
     return {
         "sampler_types": [
             {"id": "CATEGORY", "params": ["values"]},
             {"id": "UNIFORM", "params": ["low", "high"]},
             {"id": "GAUSSIAN", "params": ["mean", "std"]},
             {"id": "UUID", "params": []},
-            {"id": "DATETIME", "params": ["start_date", "end_date"]},
-            {"id": "PERSON", "params": ["locale", "include_attributes"]}
+            {"id": "DATETIME", "params": ["start_date", "end_date"]}
         ]
     }