Update app.py

app.py

@@ -1,171 +1,129 @@
 import os
 import time
 import random
+import httpx
 from typing import Dict
 
 from fastapi import FastAPI, Request, HTTPException
-from fastapi.responses import StreamingResponse, HTMLResponse, JSONResponse
-
-from ollama import Client
+from fastapi.responses import StreamingResponse, Response
 
 app = FastAPI()
 
-# ── CONFIG ─────────────────────────────────────────────
+# ── CONFIG ─────────────────────────────────────
 MASTER_API_KEY = os.getenv("MASTER_API_KEY", "changeme")
 BASE_URL = os.getenv("BASE_URL", "https://ollama.com")
 
-# ── LOAD KEYS ──────────────────────────────────────────
+# ── LOAD KEYS ──────────────────────────────────
 def load_keys(prefix="OLLAMA_KEY_"):
-    pairs = []
+    keys = []
     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]
+            keys.append(v.strip())
+    return sorted(keys)
 
 API_KEYS = load_keys()
-print(f"[INIT] Loaded {len(API_KEYS)} keys")
 
-# ── STATE ──────────────────────────────────────────────
+# ── STATE ──────────────────────────────────────
 key_status: Dict[str, Dict] = {
-    k: {
-        "status": "active",
-        "fail": 0,
-        "cooldown": 0,
-        "req": 0,
-        "err": 0,
-        "last": "-"
-    }
+    k: {"fail": 0, "cooldown": 0, "status": "active"}
     for k in API_KEYS
 }
 
-# ── AUTH ───────────────────────────────────────────────
+# ── AUTH ───────────────────────────────────────
 def auth(req: Request):
     if req.headers.get("Authorization") != f"Bearer {MASTER_API_KEY}":
         raise HTTPException(401, "Unauthorized")
 
-# ── KEY PICKER ─────────────────────────────────────────
+# ── KEY PICK ───────────────────────────────────
 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)
-
+    valid = [
+        k for k, v in key_status.items()
+        if v["status"] != "dead" and v["cooldown"] < now
+    ]
     return random.choice(valid) if valid else None
 
-# ── STATUS UPDATE ──────────────────────────────────────
+# ── STATUS ─────────────────────────────────────
 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"):
+def mark_dead(k):
     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):
+# ── UNIVERSAL PROXY ────────────────────────────
+@app.api_route("/v1/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
+async def proxy(req: Request, path: str):
     auth(req)
-    body = await req.json()
 
-    model = body.get("model")
-    messages = body.get("messages")
+    # build target url dinamis
+    target_url = f"{BASE_URL}/v1/{path}"
 
-    if not model:
-        return JSONResponse({"error": "model required"}, status_code=400)
+    # ambil body (kalau ada)
+    body = await req.body()
 
-    if not messages:
-        return JSONResponse({"error": "messages required"}, status_code=400)
+    # copy headers kecuali auth (kita ganti)
+    headers = dict(req.headers)
+    headers.pop("host", None)
 
     for _ in range(len(API_KEYS)):
         key = pick_key()
-
         if not key:
-            return JSONResponse({"error": "all keys exhausted"}, status_code=503)
+            return Response("All keys exhausted", status_code=503)
 
-        try:
-            client = Client(
-                host=BASE_URL,
-                headers={"Authorization": f"Bearer {key}"}
-            )
-
-            key_status[key]["req"] += 1
+        headers["Authorization"] = f"Bearer {key}"
 
-            def stream():
-                try:
-                    for part in client.chat(model, messages=messages, stream=True):
-                        yield part["message"]["content"]
+        try:
+            async with httpx.AsyncClient(timeout=None) as client:
+
+                r = await client.request(
+                    method=req.method,
+                    url=target_url,
+                    headers=headers,
+                    content=body,
+                    params=req.query_params
+                )
+
+                # SUCCESS
+                if r.status_code < 400:
                     mark_ok(key)
 
-                except Exception as e:
+                    # STREAMING
+                    if "text/event-stream" in r.headers.get("content-type", ""):
+                        return StreamingResponse(
+                            r.aiter_raw(),
+                            media_type="text/event-stream"
+                        )
+
+                    return Response(
+                        content=r.content,
+                        status_code=r.status_code,
+                        headers=dict(r.headers)
+                    )
+
+                # RATE LIMIT
+                elif r.status_code == 429:
+                    mark_limit(key)
+
+                # SERVER ERROR
+                elif r.status_code >= 500:
                     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")
+                    if key_status[key]["fail"] >= 3:
+                        mark_dead(key)
 
         except Exception:
-            continue
-
-    return JSONResponse({"error": "all keys failed"}, status_code=500)
+            key_status[key]["fail"] += 1
+            if key_status[key]["fail"] >= 3:
+                mark_dead(key)
 
-# ── DASHBOARD ──────────────────────────────────────────
-@app.get("/dashboard")
-async def dash(req: Request):
-    auth(req)
-
-    now = time.time()
+    return Response("All keys failed", status_code=500)
 
-    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 ───────────────────────────────────────────────
+# ── ROOT ───────────────────────────────────────
 @app.get("/")
-def root():
-    return {"status": "ok", "keys": len(API_KEYS)}
+def home():
+    return {
+        "status": "universal proxy running",
+        "keys": len(API_KEYS)
+    }