Update app.py

app.py

@@ -1,60 +1,115 @@
-import os, json, asyncio, time
-from typing import Optional
-import httpx
+import os, json, time, asyncio
+from typing import Dict, Any
+import httpx, websockets
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Body
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import JSONResponse
+from fastapi.responses import JSONResponse, StreamingResponse
 
 FILES_DIR = os.environ.get("FILES_DIR", "/tmp/brain_app/files")
 TTS_WS_TARGET = os.environ.get("TTS_WS_TARGET", "")
 TTS_HTTP_TARGET = os.environ.get("TTS_HTTP_TARGET", "")
 STT_WS_TARGET = os.environ.get("STT_WS_TARGET", "")
+AGENT_WS_TARGET = os.environ.get("AGENT_WS_TARGET", "")
 
-app = FastAPI(title="Brain Proxy")
+os.makedirs(FILES_DIR, exist_ok=True)
+
+app = FastAPI(title="Brain (Proxy/Orchestrator)")
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=[""], allow_headers=[""], allow_methods=["*"], allow_credentials=True,
+    allow_origins=[""], allow_methods=[""], allow_headers=["*"], allow_credentials=True,
 )
 
+telemetry_subs = set()     # Server-Sent Events subscribers
+logs_subs = set()          # SSE for generic logs/events
+
+def _sse_pack(d: Dict[str, Any]) -> bytes:
+    return f"data: {json.dumps(d, ensure_ascii=False)}\n\n".encode("utf-8")
+
+async def _sse_broadcast(pool, payload):
+    dead = []
+    for q in list(pool):
+        try:
+            await q.put(payload)
+        except Exception:
+            dead.append(q)
+    for q in dead:
+        pool.discard(q)
+
 @app.get("/health")
 def health():
     return {
-        "ok": True,
-        "service": "brain",
-        "files_dir": FILES_DIR,
-        "tts_ws_target": TTS_WS_TARGET,
-        "tts_http_target": TTS_HTTP_TARGET,
-        "stt_ws_target": STT_WS_TARGET,
-        "time": time.time(),
+        "ok": True, "service": "brain",
+        "tts_ws_target": TTS_WS_TARGET, "tts_http_target": TTS_HTTP_TARGET,
+        "stt_ws_target": STT_WS_TARGET, "agent_ws_target": AGENT_WS_TARGET,
+        "files_dir": FILES_DIR, "time": time.time()
     }
 
+# ---------- SSE streams ----------
+@app.get("/stream/telemetry")
+async def stream_telemetry():
+    q: asyncio.Queue = asyncio.Queue()
+    telemetry_subs.add(q)
+    async def gen():
+        try:
+            while True:
+                item = await q.get()
+                yield item
+        finally:
+            telemetry_subs.discard(q)
+    return StreamingResponse(gen(), media_type="text/event-stream")
+
+@app.get("/stream/logs")
+async def stream_logs():
+    q: asyncio.Queue = asyncio.Queue()
+    logs_subs.add(q)
+    async def gen():
+        try:
+            while True:
+                item = await q.get()
+                yield item
+        finally:
+            logs_subs.discard(q)
+    return StreamingResponse(gen(), media_type="text/event-stream")
+
+# ---------- Telemetry ingest ----------
+@app.post("/telemetry")
+async def telemetry(body: dict = Body(...)):
+    body["ts"] = time.time()
+    await _sse_broadcast(telemetry_subs, _sse_pack(body))
+    return {"ok": True}
+
+# ---------- Generic log/event ----------
+@app.post("/log")
+async def log(body: dict = Body(...)):
+    body["ts"] = time.time()
+    await _sse_broadcast(logs_subs, _sse_pack(body))
+    return {"ok": True}
+
 # ---------- HTTP proxy to TTS /speak ----------
 @app.post("/speak")
 async def speak_proxy(body: dict = Body(...)):
     if not TTS_HTTP_TARGET:
         return JSONResponse({"ok": False, "error": "TTS HTTP target not configured"}, 500)
     try:
-        async with httpx.AsyncClient(timeout=60) as client:
+        async with httpx.AsyncClient(timeout=120) as client:
             r = await client.post(TTS_HTTP_TARGET, json=body)
             r.raise_for_status()
-            return JSONResponse(r.json(), r.status_code)
-    except httpx.HTTPError as e:
+            out = r.json()
+            await _sse_broadcast(logs_subs, _sse_pack({"kind":"tts_http", "req": body, "resp": out}))
+            return JSONResponse(out, r.status_code)
+    except Exception as e:
+        await _sse_broadcast(logs_subs, _sse_pack({"kind":"error", "where":"speak_proxy", "error": repr(e)}))
         return JSONResponse({"ok": False, "error": f"TTS proxy failed: {e}"}, 502)
 
-# ---------- WS proxy helpers ----------
-async def _bridge_ws(client_ws: WebSocket, upstream_url: str):
-    import websockets
-
-    # Establish upstream connection
+# ---------- bidirectional WS bridge ----------
+async def _bridge(client_ws: WebSocket, upstream_url: str):
     try:
         upstream = await websockets.connect(upstream_url, ping_interval=None, max_size=8_000_000)
     except Exception as e:
         await client_ws.send_json({"event":"error","error":f"Upstream connect failed: {repr(e)}"})
-        await client_ws.close()
-        return
+        await client_ws.close(); return
 
     async def c2u():
-        # client -> upstream (pass text or bytes)
         try:
             while True:
                 msg = await client_ws.receive()
@@ -62,19 +117,13 @@ async def _bridge_ws(client_ws: WebSocket, upstream_url: str):
                     await upstream.send(msg["bytes"])
                 elif "text" in msg and msg["text"] is not None:
                     await upstream.send(msg["text"])
-                else:
-                    # ignore pings/others
-                    pass
         except WebSocketDisconnect:
             pass
         finally:
-            try:
-                await upstream.close()
-            except Exception:
-                pass
+            try: await upstream.close()
+            except: pass
 
     async def u2c():
-        # upstream -> client (pass text or binary)
         try:
             while True:
                 msg = await upstream.recv()
@@ -85,27 +134,28 @@ async def _bridge_ws(client_ws: WebSocket, upstream_url: str):
         except Exception:
             pass
         finally:
-            try:
-                await client_ws.close()
-            except Exception:
-                pass
+            try: await client_ws.close()
+            except: pass
 
-    # run both directions
     await asyncio.gather(c2u(), u2c())
 
-# ---------- WS endpoints the client will use ----------
-@app.websocket("/ws/tts_proxy")
-async def ws_tts_proxy(ws: WebSocket):
+@app.websocket("/ws/tts")
+async def ws_tts(ws: WebSocket):
     await ws.accept()
     if not TTS_WS_TARGET:
-        await ws.send_json({"event":"error","error":"TTS WS target not configured"})
-        await ws.close(); return
-    await _bridge_ws(ws, TTS_WS_TARGET)
+        await ws.send_json({"event":"error","error":"TTS WS target not configured"}); await ws.close(); return
+    await _bridge(ws, TTS_WS_TARGET)
 
-@app.websocket("/ws/stt_proxy")
-async def ws_stt_proxy(ws: WebSocket):
+@app.websocket("/ws/stt")
+async def ws_stt(ws: WebSocket):
     await ws.accept()
     if not STT_WS_TARGET:
-        await ws.send_json({"event":"error","error":"STT WS target not configured"})
-        await ws.close(); return
-    await _bridge_ws(ws, STT_WS_TARGET)
+        await ws.send_json({"event":"error","error":"STT WS target not configured"}); await ws.close(); return
+    await _bridge(ws, STT_WS_TARGET)
+
+@app.websocket("/ws/agent")
+async def ws_agent(ws: WebSocket):
+    await ws.accept()
+    if not AGENT_WS_TARGET:
+        await ws.send_json({"event":"error","error":"Agent WS target not configured"}); await ws.close(); return
+    await _bridge(ws, AGENT_WS_TARGET)