Update app.py

app.py

@@ -1,6 +1,9 @@
 import os
 import time
 import hashlib
+from fastapi import WebSocket, WebSocketDisconnect
+from collections import defaultdict, deque
+import json
 from fastapi import FastAPI, Request, HTTPException, status, Header
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import (
@@ -54,6 +57,14 @@ from helper.ratelimit import (
 
 app = FastAPI()
 
+WEBSOCKET_KEY = os.getenv("WEBSOCKET_KEY")
+
+# authentication attempt tracking
+AUTH_ATTEMPTS = defaultdict(lambda: deque())
+
+AUTH_WINDOW_SECONDS = 60
+AUTH_MAX_ATTEMPTS = 10
+
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -62,6 +73,20 @@ app.add_middleware(
 )
 app.include_router(asset_router)
 
+def check_ws_auth_rate_limit(ip: str):
+    now = time.time()
+    q = AUTH_ATTEMPTS[ip]
+
+    # purge old attempts
+    while q and now - q[0] > AUTH_WINDOW_SECONDS:
+        q.popleft()
+
+    if len(q) >= AUTH_MAX_ATTEMPTS:
+        return False
+
+    q.append(now)
+    return True
+
 
 @app.get("/")
 async def reroute_to_home():
@@ -1032,6 +1057,66 @@ async def tiers():
         content=paid_plans,
     )
 
+@app.websocket("/ws/chat")
+async def websocket_chat(ws: WebSocket):
+    ip = ws.client.host
+
+    await ws.accept()
+
+    # rate limit auth attempts
+    if not check_ws_auth_rate_limit(ip):
+        await ws.close(code=4408)
+        return
+
+    try:
+        auth_msg = await ws.receive_text()
+        auth_data = json.loads(auth_msg)
+
+        provided_key = auth_data.get("key")
+
+        if not WEBSOCKET_KEY or provided_key != WEBSOCKET_KEY:
+            await ws.close(code=4403)
+            return
+
+        # authenticated
+        await ws.send_json({"type": "auth", "status": "ok"})
+
+        while True:
+            msg = await ws.receive_text()
+            data = json.loads(msg)
+
+            body = data.get("body")
+            headers = data.get("headers", {})
+
+            if not body:
+                await ws.send_json({"error": "Missing body"})
+                continue
+
+            url = str(ws.url).replace("ws://", "http://").replace("wss://", "https://")
+            url = url.split("/ws/chat")[0] + "/gen/chat/completions"
+
+            async with httpx.AsyncClient(timeout=None) as client:
+                async with client.stream(
+                    "POST",
+                    url,
+                    json=body,
+                    headers=headers,
+                ) as r:
+
+                    async for line in r.aiter_lines():
+                        if not line:
+                            continue
+
+                        await ws.send_text(line)
+
+    except WebSocketDisconnect:
+        return
+    except Exception as e:
+        try:
+            await ws.send_json({"error": str(e)})
+        except:
+            pass
+        await ws.close()
 
 @app.get("/portal")
 @app.post("/portal")