v2.0: Direct z.ai Anthropic API integration

app.py

@@ -1,9 +1,10 @@
 import os
+import httpx
 from contextlib import asynccontextmanager
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
-from typing import Any
 import tempfile
+import random
 
 from models import (
     GenerateRequest, GenerateResponse,
@@ -11,30 +12,13 @@ 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
-
-
-@asynccontextmanager
-async def lifespan(app: FastAPI):
-    global data_designer
-    from data_designer.interface import DataDesigner
-    data_designer = DataDesigner(artifact_path=tempfile.gettempdir())
-    yield
-
+ZAI_BASE_URL = "https://api.z.ai/api/anthropic"
 
 app = FastAPI(
     title="NeMo DataDesigner API",
-    description="Synthetic data generation with NVIDIA NeMo DataDesigner and z.ai",
-    version="1.3.0",
-    lifespan=lifespan
+    description="Synthetic data generation with z.ai",
+    version="2.0.0"
 )
 
 app.add_middleware(
@@ -46,207 +30,122 @@ app.add_middleware(
 )
 
 
-def build_config(request):
-    import data_designer.config as dd
-    from data_designer.config.models import ModelConfig, ChatCompletionInferenceParams, ModelProvider
-
-    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"),
-                )
-            )
-
-    # 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=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, zai_provider
-
+async def call_zai(prompt: str, model: str, temperature: float, max_tokens: int) -> str:
+    """Call z.ai API directly with Anthropic format."""
+    async with httpx.AsyncClient(timeout=60.0) as client:
+        response = await client.post(
+            f"{ZAI_BASE_URL}/v1/messages",
+            headers={
+                "x-api-key": ZAI_API_KEY,
+                "anthropic-version": "2023-06-01",
+                "content-type": "application/json"
+            },
+            json={
+                "model": model,
+                "max_tokens": max_tokens,
+                "messages": [{"role": "user", "content": prompt}]
+            }
+        )
+        if response.status_code != 200:
+            raise Exception(f"z.ai API error: {response.status_code} - {response.text}")
+        data = response.json()
+        return data["content"][0]["text"]
+
+
+def sample_value(sampler_type: str, params: dict) -> str:
+    """Sample a value based on sampler type."""
+    if sampler_type == "CATEGORY":
+        values = params.get("values", ["A", "B", "C"])
+        return random.choice(values)
+    elif sampler_type == "UNIFORM":
+        low = params.get("low", 0)
+        high = params.get("high", 100)
+        return str(random.randint(low, high))
+    elif sampler_type == "GAUSSIAN":
+        mean = params.get("mean", 0)
+        std = params.get("std", 1)
+        return str(round(random.gauss(mean, std), 2))
+    else:
+        return "default"
 
-def get_sampler_params(sampler_type, params):
-    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"]))
-    elif type_name == "UNIFORM":
-        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))
-    else:
-        return dd.CategorySamplerParams(values=["default"])
+def render_prompt(template: str, context: dict) -> str:
+    """Render prompt template with context variables."""
+    result = template
+    for key, value in context.items():
+        result = result.replace("{{ " + key + " }}", str(value))
+        result = result.replace("{{" + key + "}}", str(value))
+    return result
 
 
 @app.get("/", response_model=HealthResponse)
 async def root():
-    return HealthResponse(status="healthy", model="data-designer", api_configured=bool(ZAI_API_KEY))
+    return HealthResponse(status="healthy", model="z.ai", api_configured=bool(ZAI_API_KEY))
 
 
 @app.get("/health", response_model=HealthResponse)
 async def health():
-    return HealthResponse(status="healthy", model="data-designer", api_configured=bool(ZAI_API_KEY))
+    return HealthResponse(status="healthy", model="z.ai", api_configured=bool(ZAI_API_KEY))
 
 
 @app.post("/generate", response_model=GenerateResponse)
 async def generate(request: GenerateRequest):
+    """Generate synthetic data using z.ai API."""
     try:
-        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"
-        )
+        model = request.model.value
+        records = []
         
-        dd_custom = DataDesigner(
-            artifact_path=tempfile.gettempdir(),
-            model_providers=[zai_provider]
-        )
+        sampler_cols = {c.name: c for c in request.columns if c.type == "sampler"}
+        llm_cols = [c for c in request.columns if c.type == "llm_text"]
         
-        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"),
-                    )
+        for _ in range(request.num_records):
+            record = {}
+            
+            # Generate sampler values first
+            for name, col in sampler_cols.items():
+                record[name] = sample_value(
+                    col.params.get("sampler_type", "CATEGORY"),
+                    col.params
                 )
-
-        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"
-        )
-        df = result.load_dataset()
-        data = df.to_dict(orient="records")
-        return GenerateResponse(success=True, data=data, record_count=len(data))
+            
+            # Generate LLM text using z.ai
+            for col in llm_cols:
+                prompt = render_prompt(col.params.get("prompt", "Generate text"), record)
+                text = await call_zai(prompt, model, request.temperature, request.max_tokens)
+                record[col.name] = text
+            
+            records.append(record)
+        
+        return GenerateResponse(success=True, data=records, record_count=len(records))
+    
     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."""
     try:
-        from data_designer.interface import DataDesigner
-        import data_designer.config as dd
-        from data_designer.config.models import ModelConfig, ChatCompletionInferenceParams, ModelProvider
+        model = request.model.value
         
-        zai_provider = ModelProvider(
-            name="zai",
-            endpoint="https://api.z.ai/api/paas/v4/",
-            api_key="ZAI_API_KEY",
-            provider_type="openai"
-        )
+        sampler_cols = {c.name: c for c in request.columns if c.type == "sampler"}
+        llm_cols = [c for c in request.columns if c.type == "llm_text"]
         
-        dd_custom = DataDesigner(
-            artifact_path=tempfile.gettempdir(),
-            model_providers=[zai_provider]
-        )
+        record = {}
         
-        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)
+        for name, col in sampler_cols.items():
+            record[name] = sample_value(
+                col.params.get("sampler_type", "CATEGORY"),
+                col.params
+            )
+        
+        for col in llm_cols:
+            prompt = render_prompt(col.params.get("prompt", "Generate text"), record)
+            text = await call_zai(prompt, model, request.temperature, request.max_tokens)
+            record[col.name] = text
+        
+        return PreviewResponse(success=True, sample=record)
+    
     except Exception as e:
         return PreviewResponse(success=False, error=str(e))
 
@@ -254,9 +153,9 @@ async def preview(request: PreviewRequest):
 @app.get("/models")
 async def list_models():
     return {"models": [
-        {"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"}
+        {"id": "glm-5", "name": "GLM-5 (Opus)", "description": "Most capable"},
+        {"id": "glm-4.7", "name": "GLM-4.7 (Sonnet)", "description": "Balanced"},
+        {"id": "glm-4.5-air", "name": "GLM-4.5-Air (Haiku)", "description": "Fast"}
     ]}