Update app.py

app.py

@@ -13,7 +13,7 @@ DB_PATH = "/data/keys.db" if os.path.exists("/data") else "keys.db"
 
 # --- DATABASE SETUP ---
 def init_db():
-    conn = sqlite3.connect(DB_PATH)
+    conn = sqlite3.connect(DB_PATH, check_same_thread=False)
     c = conn.cursor()
     c.execute('CREATE TABLE IF NOT EXISTS api_keys (key TEXT PRIMARY KEY)')
     conn.commit()
@@ -21,50 +21,59 @@ def init_db():
 
 init_db()
 
+# --- DATA MODELS ---
 class ChatRequest(BaseModel):
     message: str
     session_id: str = "default_user"
 
-# --- KEY GENERATION (The Faucet) ---
+# --- KEY GENERATION ---
 @app.get("/generate")
 async def generate_public_key():
     """Anyone can call this to get a new sk- key"""
     new_key = f"sk-{secrets.token_urlsafe(24)}"
-    
-    conn = sqlite3.connect(DB_PATH)
+    conn = sqlite3.connect(DB_PATH, check_same_thread=False)
     c = conn.cursor()
     c.execute('INSERT INTO api_keys (key) VALUES (?)', (new_key,))
     conn.commit()
     conn.close()
-    
     return {
-        "api_key": new_key, 
+        "api_key": new_key,
         "instructions": "Add this to your 'X-API-Key' header to use /v1/chat"
     }
 
-# --- PROXY LOGIC ---
+# --- KEY VERIFICATION ---
 async def verify_key(x_api_key: str = Header(None)):
     if not x_api_key:
         raise HTTPException(status_code=401, detail="Missing X-API-Key header")
-    
-    conn = sqlite3.connect(DB_PATH)
+    conn = sqlite3.connect(DB_PATH, check_same_thread=False)
     c = conn.cursor()
     c.execute('SELECT 1 FROM api_keys WHERE key = ?', (x_api_key,))
     exists = c.fetchone()
     conn.close()
-    
     if not exists:
         raise HTTPException(status_code=403, detail="Invalid API Key")
     return x_api_key
 
+# --- PROXY CHAT ---
 @app.post("/v1/chat")
 async def proxy_chat(payload: ChatRequest, api_key: str = Depends(verify_key)):
     try:
-        response = requests.post(TARGET_SPACE_URL, json=payload.dict(), timeout=120)
+        # Forward to your 20B backend
+        headers = {"Content-Type": "application/json"}
+        response = requests.post(
+            TARGET_SPACE_URL,
+            json=payload.dict(),
+            headers=headers,
+            timeout=120
+        )
+        response.raise_for_status()  # raise exception on HTTP errors
         return response.json()
-    except Exception as e:
-        return {"error": "The 20B engine is busy or offline."}
+    except requests.exceptions.Timeout:
+        return {"error": "The 20B engine timed out."}
+    except requests.exceptions.RequestException as e:
+        return {"error": f"The 20B engine is offline or returned an error: {str(e)}"}
 
+# --- HOME ---
 @app.get("/")
 def home():
-    return {"message": "Go to /generate to get a key, then POST to /v1/chat"}
+    return {"message": "Go to /generate to get a key, then POST to /v1/chat"}