Update app.py

app.py

@@ -0,0 +1,78 @@
+# app.py
+import os
+import json
+import time
+from typing import Dict
+from fastapi import FastAPI, Request, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+import httpx
+
+HF_TOKEN = os.environ.get("HF_TOKEN")  # set this as a Space secret
+REPO_ID = "deepseek-ai/DeepSeek-V3.2-Exp-Base"  # change if needed
+INFERENCE_URL = f"https://api-inference.huggingface.co/models/{REPO_ID}"
+# Optional: require a simple API key header from frontend (set as secret API_KEY)
+API_KEY = os.environ.get("API_KEY")
+
+app = FastAPI(title="HF Space Model Proxy")
+
+# Allow CORS from your frontend OR allow all for testing. Replace '*' with your origin in production.
+origins = ["*"]  # <-- replace with ["https://your-frontend-domain.com"] in production
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=origins,
+    allow_credentials=True,
+    allow_methods=["POST", "OPTIONS"],
+    allow_headers=["*"],
+)
+
+# simple health check
+@app.get("/health")
+async def health():
+    return {"status": "ok", "time": time.time()}
+
+@app.post("/predict")
+async def predict(request: Request):
+    # optional API key check (frontend should send header "x-api-key")
+    if API_KEY:
+        key = request.headers.get("x-api-key")
+        if key != API_KEY:
+            raise HTTPException(status_code=401, detail="Invalid API key")
+
+    payload = await request.json()
+    # Expect payload like: {"messages": [{"role":"user","content":"Who are you?"}], "max_new_tokens":256}
+    messages = payload.get("messages") or payload.get("inputs") or payload.get("prompt")
+    if messages is None:
+        raise HTTPException(status_code=400, detail="Missing 'messages' or 'inputs' or 'prompt' in JSON body")
+
+    max_new_tokens = int(payload.get("max_new_tokens", 256))
+
+    headers = {"Authorization": f"Bearer {HF_TOKEN}"} if HF_TOKEN else {}
+    body = {
+        "inputs": messages,
+        "parameters": {"max_new_tokens": max_new_tokens},
+        "options": {"wait_for_model": True}
+    }
+
+    async with httpx.AsyncClient(timeout=120) as client:
+        resp = await client.post(INFERENCE_URL, headers=headers, json=body)
+        if resp.status_code != 200:
+            # bubble up error with introspection
+            detail = {"hf_status": resp.status_code, "content": resp.text}
+            raise HTTPException(status_code=502, detail=detail)
+        data = resp.json()
+
+    # Normalize output to simple string if possible
+    if isinstance(data, list) and len(data) > 0:
+        first = data[0]
+        if isinstance(first, dict) and "generated_text" in first:
+            out_text = first["generated_text"]
+        elif isinstance(first, str):
+            out_text = first
+        else:
+            out_text = json.dumps(data)
+    elif isinstance(data, dict) and "generated_text" in data:
+        out_text = data["generated_text"]
+    else:
+        out_text = json.dumps(data)
+
+    return {"output": out_text, "raw": data}