Upload app.py with huggingface_hub

app.py

@@ -12,12 +12,13 @@ from models import (
     HealthResponse, ZaiModel, SamplerType
 )
 
-# Configure z.ai as OpenAI-compatible provider
+# Configure z.ai as Anthropic-compatible provider for LiteLLM
 ZAI_API_KEY = os.environ.get("ZAI_API_KEY", "")
 ZAI_BASE_URL = "https://api.z.ai/api/anthropic"
 
-os.environ["OPENAI_API_KEY"] = ZAI_API_KEY
-os.environ["OPENAI_API_BASE"] = ZAI_BASE_URL
+# 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
 
 # Global DataDesigner instance
 data_designer = None
@@ -27,7 +28,6 @@ data_designer = None
 async def lifespan(app: FastAPI):
     global data_designer
     from data_designer.interface import DataDesigner
-    # Use temp directory for artifacts
     data_designer = DataDesigner(artifact_path=tempfile.gettempdir())
     yield
 
@@ -97,11 +97,11 @@ def build_config(request: GenerateRequest | PreviewRequest):
                 )
             )
 
-    # Add model config
+    # Add model config - use anthropic provider with custom base_url via env
     model_config = ModelConfig(
         alias="zai-model",
-        model=f"openai/{model_id}",
-        provider="openai",
+        model=f"anthropic/{model_id}",  # LiteLLM Anthropic format
+        provider="anthropic",
         inference_parameters=ChatCompletionInferenceParams(
             temperature=request.temperature,
             max_tokens=request.max_tokens,
@@ -169,14 +169,12 @@ async def generate(request: GenerateRequest):
     try:
         config_builder = build_config(request)
 
-        # Use create() method - returns DatasetCreationResults
         result = data_designer.create(
             config_builder=config_builder,
             num_records=request.num_records,
             dataset_name="api-dataset"
         )
 
-        # Get DataFrame from results
         df = result.load_dataset()
         data = df.to_dict(orient="records")
 
@@ -190,7 +188,7 @@ async def generate(request: GenerateRequest):
         import traceback
         return GenerateResponse(
             success=False,
-            error=f"{str(e)}\n{traceback.format_exc()}"
+            error=f"{str(e)}"
         )
 
 
@@ -202,13 +200,11 @@ async def preview(request: PreviewRequest):
     try:
         config_builder = build_config(request)
 
-        # Use preview() method - returns PreviewResults
         preview_result = data_designer.preview(
             config_builder=config_builder,
             num_records=1
         )
 
-        # Get sample record
         sample = preview_result.dataset.to_dict(orient="records")[0] if len(preview_result.dataset) > 0 else {}
 
         return PreviewResponse(
@@ -220,7 +216,7 @@ async def preview(request: PreviewRequest):
         import traceback
         return PreviewResponse(
             success=False,
-            error=f"{str(e)}\n{traceback.format_exc()}"
+            error=f"{str(e)}"
         )