Create app.py

app.py

@@ -0,0 +1,171 @@
+import os
+import time
+import random
+from typing import Dict
+
+from fastapi import FastAPI, Request, HTTPException
+from fastapi.responses import StreamingResponse, HTMLResponse, JSONResponse
+
+from ollama import Client
+
+app = FastAPI()
+
+# ── CONFIG ─────────────────────────────────────────────
+MASTER_API_KEY = os.getenv("MASTER_API_KEY", "changeme")
+BASE_URL = os.getenv("BASE_URL", "https://ollama.com")
+
+# ── LOAD KEYS ──────────────────────────────────────────
+def load_keys(prefix="OLLAMA_KEY_"):
+    pairs = []
+    for k, v in os.environ.items():
+        if k.startswith(prefix) and v:
+            try:
+                idx = int(k.replace(prefix, ""))
+            except:
+                idx = 9999
+            pairs.append((idx, v.strip()))
+    pairs.sort(key=lambda x: x[0])
+    return [v for _, v in pairs]
+
+API_KEYS = load_keys()
+print(f"[INIT] Loaded {len(API_KEYS)} keys")
+
+# ── STATE ──────────────────────────────────────────────
+key_status: Dict[str, Dict] = {
+    k: {
+        "status": "active",
+        "fail": 0,
+        "cooldown": 0,
+        "req": 0,
+        "err": 0,
+        "last": "-"
+    }
+    for k in API_KEYS
+}
+
+# ── AUTH ───────────────────────────────────────────────
+def auth(req: Request):
+    if req.headers.get("Authorization") != f"Bearer {MASTER_API_KEY}":
+        raise HTTPException(401, "Unauthorized")
+
+# ── KEY PICKER ─────────────────────────────────────────
+def pick_key():
+    now = time.time()
+    valid = []
+
+    for k, v in key_status.items():
+        if v["status"] == "dead":
+            continue
+        if v["cooldown"] > now:
+            continue
+        valid.append(k)
+
+    return random.choice(valid) if valid else None
+
+# ── STATUS UPDATE ──────────────────────────────────────
+def mark_limit(k):
+    key_status[k]["status"] = "limited"
+    key_status[k]["cooldown"] = time.time() + 60
+    key_status[k]["err"] += 1
+    key_status[k]["last"] = "rate_limit"
+
+def mark_dead(k, reason="dead"):
+    key_status[k]["status"] = "dead"
+    key_status[k]["err"] += 1
+    key_status[k]["last"] = reason
+
+def mark_ok(k):
+    key_status[k]["status"] = "active"
+    key_status[k]["fail"] = 0
+    key_status[k]["last"] = "-"
+
+# ── STREAM PROXY ───────────────────────────────────────
+@app.post("/chat")
+async def chat(req: Request):
+    auth(req)
+    body = await req.json()
+
+    model = body.get("model")
+    messages = body.get("messages")
+
+    if not model:
+        return JSONResponse({"error": "model required"}, status_code=400)
+
+    if not messages:
+        return JSONResponse({"error": "messages required"}, status_code=400)
+
+    for _ in range(len(API_KEYS)):
+        key = pick_key()
+
+        if not key:
+            return JSONResponse({"error": "all keys exhausted"}, status_code=503)
+
+        try:
+            client = Client(
+                host=BASE_URL,
+                headers={"Authorization": f"Bearer {key}"}
+            )
+
+            key_status[key]["req"] += 1
+
+            def stream():
+                try:
+                    for part in client.chat(model, messages=messages, stream=True):
+                        yield part["message"]["content"]
+                    mark_ok(key)
+
+                except Exception as e:
+                    key_status[key]["fail"] += 1
+                    key_status[key]["err"] += 1
+                    key_status[key]["last"] = str(e)
+
+                    if "429" in str(e):
+                        mark_limit(key)
+                    elif key_status[key]["fail"] >= 3:
+                        mark_dead(key, "fail")
+
+                    raise e
+
+            return StreamingResponse(stream(), media_type="text/plain")
+
+        except Exception:
+            continue
+
+    return JSONResponse({"error": "all keys failed"}, status_code=500)
+
+# ── DASHBOARD ──────────────────────────────────────────
+@app.get("/dashboard")
+async def dash(req: Request):
+    auth(req)
+
+    now = time.time()
+
+    html = """
+    <html><body style="background:#0f172a;color:white;font-family:sans-serif">
+    <h2>KEY STATUS</h2>
+    <table border="1" cellpadding="8">
+    <tr>
+    <th>Key</th><th>Status</th><th>Req</th><th>Err</th><th>Cooldown</th><th>Last</th>
+    </tr>
+    """
+
+    for k, v in key_status.items():
+        cd = max(0, int(v["cooldown"] - now))
+        html += f"""
+        <tr>
+        <td>{k[:6]}...</td>
+        <td>{v['status']}</td>
+        <td>{v['req']}</td>
+        <td>{v['err']}</td>
+        <td>{cd}s</td>
+        <td>{v['last']}</td>
+        </tr>
+        """
+
+    html += "</table></body></html>"
+    return HTMLResponse(html)
+
+# ── ROOT ───────────────────────────────────────────────
+@app.get("/")
+def root():
+    return {"status": "ok", "keys": len(API_KEYS)}