v3.0: Use LiteLLM anthropic/ prefix with ANTHROPIC_BASE_URL

app.py

@@ -1,10 +1,8 @@
 import os
-import httpx
 from contextlib import asynccontextmanager
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 import tempfile
-import random
 
 from models import (
     GenerateRequest, GenerateResponse,
@@ -12,13 +10,30 @@ from models import (
     HealthResponse, ZaiModel
 )
 
+# z.ai als Anthropic-Endpunkt für LiteLLM konfigurieren
 ZAI_API_KEY = os.environ.get("ZAI_API_KEY", "")
 ZAI_BASE_URL = "https://api.z.ai/api/anthropic"
 
+# LiteLLM Anthropic-Konfiguration
+os.environ["ANTHROPIC_API_KEY"] = ZAI_API_KEY
+os.environ["ANTHROPIC_BASE_URL"] = ZAI_BASE_URL
+
+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
+
+
 app = FastAPI(
     title="NeMo DataDesigner API",
-    description="Synthetic data generation with z.ai",
-    version="2.0.0"
+    description="Synthetic data generation with DataDesigner + z.ai",
+    version="3.0.0",
+    lifespan=lifespan
 )
 
 app.add_middleware(
@@ -30,122 +45,96 @@ app.add_middleware(
 )
 
 
-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 build_config(request):
+    import data_designer.config as dd
+    from data_designer.config.models import ModelConfig, ChatCompletionInferenceParams
+
+    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"),
+                )
+            )
 
-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
+    # LiteLLM erkennt "anthropic/" Prefix und nutzt ANTHROPIC_BASE_URL
+    model_config = ModelConfig(
+        alias="zai-model",
+        model=f"anthropic/{model_id}",
+        provider="anthropic",
+        inference_parameters=ChatCompletionInferenceParams(
+            temperature=request.temperature,
+            max_tokens=request.max_tokens,
+        ),
+    )
+    config_builder.add_model_config(model_config)
+
+    return config_builder
+
+
+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"])
 
 
 @app.get("/", response_model=HealthResponse)
 async def root():
-    return HealthResponse(status="healthy", model="z.ai", 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():
-    return HealthResponse(status="healthy", model="z.ai", 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 z.ai API."""
     try:
-        model = request.model.value
-        records = []
-        
-        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"]
-        
-        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
-                )
-            
-            # 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))
-    
+        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))
     except Exception as e:
         return GenerateResponse(success=False, error=str(e))
 
 
 @app.post("/preview", response_model=PreviewResponse)
 async def preview(request: PreviewRequest):
-    """Preview a single record."""
     try:
-        model = request.model.value
-        
-        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"]
-        
-        record = {}
-        
-        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)
-    
+        config_builder = build_config(request)
+        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)
     except Exception as e:
         return PreviewResponse(success=False, error=str(e))