Update server.py

server.py

@@ -1,38 +1,59 @@
-from fastapi import FastAPI, Request
+from fastapi import FastAPI, Request, HTTPException, Depends
+from fastapi.middleware.cors import CORSMiddleware
 import os, requests
 
 app = FastAPI()
 
-OLLAMA_URL = os.getenv("OLLAMA_URL")
+# --- Config from env ---
+OLLAMA_URL = os.getenv("OLLAMA_URL")  # https://studio-ollama.decidrapp.com
 CF_ID = os.getenv("CF_ACCESS_CLIENT_ID")
 CF_SECRET = os.getenv("CF_ACCESS_CLIENT_SECRET")
+API_KEY = os.getenv("API_KEY")  # <-- NEW: server-side API key
+
+# --- CORS (optional; safe defaults) ---
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["GET", "POST", "OPTIONS"],
+    allow_headers=["*"],
+)
+
+def require_api_key(request: Request):
+    """Simple header-based auth for this proxy."""
+    if not API_KEY:
+        # If you forgot to set it, fail CLOSED
+        raise HTTPException(status_code=500, detail="Server missing API_KEY")
+    if request.headers.get("x-api-key") != API_KEY:
+        raise HTTPException(status_code=401, detail="Unauthorized")
 
 @app.get("/")
 def home():
     return {"status": "Decidr proxy running"}
 
-@app.post("/api/generate")
-async def generate(req: Request):
-    body = await req.json()
-    r = requests.post(
-        f"{OLLAMA_URL}/api/generate",
-        json=body,
+@app.get("/api/tags")
+def tags(_: None = Depends(require_api_key)):
+    r = requests.get(
+        f"{OLLAMA_URL}/api/tags",
         headers={
             "CF-Access-Client-Id": CF_ID,
-            "CF-Access-Client-Secret": CF_SECRET
+            "CF-Access-Client-Secret": CF_SECRET,
         },
-        timeout=60
+        timeout=30
     )
+    r.raise_for_status()
     return r.json()
 
-@app.get("/api/tags")
-def tags():
-    r = requests.get(
-        f"{OLLAMA_URL}/api/tags",
+@app.post("/api/generate")
+async def generate(req: Request, _: None = Depends(require_api_key)):
+    body = await req.json()
+    r = requests.post(
+        f"{OLLAMA_URL}/api/generate",
+        json=body,
         headers={
             "CF-Access-Client-Id": CF_ID,
-            "CF-Access-Client-Secret": CF_SECRET
+            "CF-Access-Client-Secret": CF_SECRET,
         },
-        timeout=30
+        timeout=60
     )
+    r.raise_for_status()
     return r.json()