Update app.py

app.py

@@ -1,109 +1,158 @@
-import os, json
-from typing import Dict, Any
-from fastapi import FastAPI, Body, Request
-from fastapi.middleware.cors import CORSMiddleware
+---
+
+# app.py
+```python
+# Brain_v2 — lightweight router/control-plane (no audio proxy)
+import os, json, time, asyncio
+from typing import Optional, Dict, Any, List, Set
+
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request, Header
 from fastapi.responses import JSONResponse, PlainTextResponse
+from fastapi.middleware.cors import CORSMiddleware
 import httpx
 
-# Stable AsyncClient with explicit 4-part timeout
-_client = httpx.AsyncClient(
-    timeout=httpx.Timeout(
-        connect=15.0,   # max time to connect
-        read=30.0,      # max time to wait for response data
-        write=15.0,     # max time to send request body
-        pool=15.0       # max time to get a connection from the pool
-    )
-)
+APP_START_TS = time.time()
+
+# -----------------------
+# Config
+# -----------------------
+BRAIN_SECRET = (os.getenv("BRAIN_SHARED_SECRET") or "").strip()
+ALLOW_ORIGINS = [o.strip() for o in (os.getenv("ALLOW_ORIGINS") or "*").split(",") if o.strip()]
+LOG_EVENTS = (os.getenv("LOG_EVENTS") or "0").strip() == "1"
 
-# ========= Config (set via Secrets) =========
-TTS_BASE = os.getenv("TTS_BASE", "https://Percy3822-ActualTTS.hf.space")
-PY_AI_URL = os.getenv("PY_AI_URL", "https://Percy3822-Python_ai.hf.space")      # optional/future
-GODOT_AI_URL = os.getenv("GODOT_AI_URL", "https://Percy3822-Godot_ai.hf.space") # optional/future
-SECRET = os.getenv("BRAIN_SHARED_SECRET")  # same header "x-auth" for internal calls
+# Optional hint included in control-plane messages to clients.
+# (We do NOT proxy audio; models/clients connect directly.)
+FORWARD_TTS_URL = (os.getenv("FORWARD_TTS_URL") or "").strip()
 
-ROUTES = {
-    "python": PY_AI_URL,
-    "godot": GODOT_AI_URL,
-    "chat":  None,  # handled client-side (or future Chat module)
-}
+# httpx client (used only if you add out-calls later)
+http = httpx.AsyncClient(
+    timeout=httpx.Timeout(connect=10.0, read=20.0, write=10.0, pool=10.0)
+)
 
-app = FastAPI(title="Brain_v2 (Router-Only)", version="2.0.0")
+# -----------------------
+# App & CORS
+# -----------------------
+app = FastAPI(title="Brain_v2")
 
-# CORS for your local app; restrict later if desired
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],
+    allow_origins=ALLOW_ORIGINS if ALLOW_ORIGINS else ["*"],
+    allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-_client = httpx.AsyncClient(timeout=httpx.Timeout(connect=15, read=None, write=15))
+# -----------------------
+# Broadcast Hub
+# -----------------------
+class Hub:
+    def _init_(self):
+        self._peers: Set[WebSocket] = set()
+        self._lock = asyncio.Lock()
+
+    async def join(self, ws: WebSocket):
+        async with self._lock:
+            self._peers.add(ws)
 
-def _authorized(req: Request) -> bool:
-    return not SECRET or req.headers.get("x-auth") == SECRET
+    async def leave(self, ws: WebSocket):
+        async with self._lock:
+            self._peers.discard(ws)
 
+    async def broadcast(self, payload: Dict[str, Any]):
+        data = json.dumps(payload, ensure_ascii=False)
+        async with self._lock:
+            peers = list(self._peers)
+        if LOG_EVENTS:
+            print(f"[BRAIN] broadcast → {len(peers)} peer(s): {data[:160]}{'...' if len(data)>160 else ''}", flush=True)
+        for ws in peers:
+            try:
+                await ws.send_text(data)
+            except Exception:
+                # drop dead peers silently
+                await self.leave(ws)
+
+hub = Hub()
+
+# -----------------------
+# Health
+# -----------------------
 @app.get("/")
 def root():
-    return PlainTextResponse("Brain_v2 router — POST /process to get stream URLs")
+    return PlainTextResponse("Brain_v2 — POST /process (x-auth), WS /ws/brain, GET /health")
 
 @app.get("/health")
-async def health():
+def health():
     return {
         "ok": True,
-        "routes": {k: (v or "client") for k, v in ROUTES.items()},
-        "tts_ws": f"{TTS_BASE}/ws/tts",
-        "version": "2.0.0"
+        "uptime_sec": round(time.time() - APP_START_TS, 3),
+        "bus_peers": len(hub._peers),
+        "config": {
+            "allow_origins": ALLOW_ORIGINS or ["*"],
+            "log_events": LOG_EVENTS,
+            "forward_tts_url": bool(FORWARD_TTS_URL),
+        },
+        "tip": "Subscribe via WS /ws/brain; producers POST to /process with x-auth",
     }
 
+# -----------------------
+# Process (Notify in)
+# -----------------------
 @app.post("/process")
-async def process(payload: Dict[str, Any] = Body(...), request: Request = None):
-    """
-    Input: { "text": "...", "module":"auto|python|godot|chat" }
-    Output: { ok, route, stream_url, tts_url }
-    - Client connects directly to stream_url for live TEXT tokens
-    - Client connects directly to tts_url for AUDIO (WS) when needed
-    """
-    if request and not _authorized(request):
-        return JSONResponse({"error": "unauthorized"}, status_code=401)
-
-    text = (payload.get("text") or "").strip()
-    module = (payload.get("module") or "auto").lower()
-    if not text:
-        return JSONResponse({"error": "empty text"}, status_code=400)
+async def process_event(request: Request, x_auth: Optional[str] = Header(None)):
+    # Secure the router
+    if not BRAIN_SECRET:
+        return JSONResponse({"ok": False, "error": "server not configured (missing BRAIN_SHARED_SECRET)"}, status_code=500)
+    if (x_auth or "").strip() != BRAIN_SECRET:
+        return JSONResponse({"ok": False, "error": "unauthorized"}, status_code=401)
 
-    # Lightweight routing
-    low = text.lower()
-    route = "chat"
-    if module in ("python", "godot", "chat"):
-        route = module
-    elif low.startswith(("python:", "py:")):
-        route = "python"
-    elif low.startswith(("godot:", "gd:")):
-        route = "godot"
+    try:
+        body = await request.json()
+    except Exception:
+        return JSONResponse({"ok": False, "error": "invalid json"}, status_code=400)
 
-    target = ROUTES.get(route)
-    stream_url = f"{target}/ws/stream" if target else None
+    text = (body.get("text") or "").strip()
+    meta = body.get("meta") or {}
+    if not text:
+        return JSONResponse({"ok": False, "error": "empty text"}, status_code=400)
 
-    return {
-        "ok": True,
-        "route": route,
-        "stream_url": stream_url,         # direct model → client (text)
-        "tts_url": f"{TTS_BASE}/ws/tts",  # direct TTS → client (audio)
-        "note": "Connect directly to these URLs for streaming. Brain does not proxy."
+    payload = {
+        "event": "final",
+        "text": text,
+        "meta": meta,
+        "ts": time.time(),
     }
+    if FORWARD_TTS_URL:
+        payload["tts_hint"] = {"ws": FORWARD_TTS_URL}
+
+    await hub.broadcast(payload)
+    return {"ok": True, "delivered_to": len(hub._peers)}
 
-# Optional debug: say text via TTS to test wiring (non-streaming HTTP)
-@app.post("/say")
-async def say(payload: Dict[str, Any] = Body(...), request: Request = None):
-    if request and not _authorized(request):
-        return JSONResponse({"error": "unauthorized"}, status_code=401)
-    txt = (payload.get("text") or "").strip()
-    if not txt:
-        return JSONResponse({"error": "empty text"}, status_code=400)
+# -----------------------
+# Brain bus (WS)
+# -----------------------
+@app.websocket("/ws/brain")
+async def ws_brain(ws: WebSocket):
+    await ws.accept()
+    await hub.join(ws)
     try:
-        async with _client.stream("POST", f"{TTS_BASE}/speak.wav", json={"text": txt}) as r:
-            r.raise_for_status()
-            # we don’t proxy bytes here; caller just checks 200 OK for warmup
-            return {"ok": True}
-    except Exception as e:
-        return JSONResponse({"error": str(e)}, status_code=502)
+        # passive channel: we don’t require client messages; keepalive pings are fine
+        while True:
+            try:
+                # wait for client pings/messages but time out to stay responsive
+                msg = await asyncio.wait_for(ws.receive_text(), timeout=60.0)
+                if msg and LOG_EVENTS:
+                    print(f"[BRAIN] client msg: {msg[:160]}{'...' if len(msg)>160 else ''}", flush=True)
+            except asyncio.TimeoutError:
+                # send keepalive ping
+                await ws.send_text(json.dumps({"event":"ping","ts": time.time()}))
+    except WebSocketDisconnect:
+        pass
+    except Exception:
+        # swallow unexpected during client disconnects
+        pass
+    finally:
+        await hub.leave(ws)
+        try:
+            await ws.close()
+        except Exception:
+            pass