Create main.py

main.py

@@ -0,0 +1,35 @@
+from fastapi import FastAPI, Header, HTTPException, Request
+import secrets, sqlite3, requests, os
+
+app = FastAPI()
+DB_PATH = "/data/keys.db" # Use persistent storage if available
+BRAIN_URL = "https://your-username-gemma-brain.hf.space/completion"
+HF_TOKEN = os.getenv("HF_TOKEN")
+
+# Initialize Database
+conn = sqlite3.connect(DB_PATH, check_same_thread=False)
+conn.execute("CREATE TABLE IF NOT EXISTS keys (api_key TEXT PRIMARY KEY, user_id TEXT)")
+
+@app.post("/generate-key")
+async def create_key(user_id: str):
+    new_key = f"gemma_max_sk_{secrets.token_urlsafe(16)}"
+    conn.execute("INSERT INTO keys VALUES (?, ?)", (new_key, user_id))
+    conn.commit()
+    return {"api_key": new_key}
+
+@app.post("/v1/chat")
+async def proxy_to_brain(request: Request, authorization: str = Header(None)):
+    # 1. Check API Key
+    key = authorization.replace("Bearer ", "")
+    res = conn.execute("SELECT user_id FROM keys WHERE api_key=?", (key,)).fetchone()
+    if not res:
+        raise HTTPException(status_code=401, detail="Invalid API Key")
+
+    # 2. Forward to Private Brain
+    user_data = await request.json()
+    response = requests.post(
+        BRAIN_URL,
+        json={"prompt": user_data["prompt"], "n_predict": 512},
+        headers={"Authorization": f"Bearer {HF_TOKEN}"}
+    )
+    return response.json()