Update app.py

app.py

@@ -1,8 +1,10 @@
-import os, json, time, asyncio, tempfile
+import os, json, time, asyncio, tempfile, struct
 from typing import AsyncGenerator, Dict, Any, Optional
 from fastapi import FastAPI, Request, Query, UploadFile
 from fastapi.responses import JSONResponse, StreamingResponse, FileResponse
 
+import websockets  # NEW
+
 # === Directories ===
 BASE_DIR  = os.environ.get("BASE_DIR", "/tmp/brain_app")
 FILES_DIR = os.path.join(BASE_DIR, "files")
@@ -21,7 +23,7 @@ BASE_WPM       = int(os.environ.get("BASE_WPM", "165"))
 NOISE_SCALE    = float(os.environ.get("NOISE_SCALE", "0.33"))
 NOISE_W        = float(os.environ.get("NOISE_W", "0.92"))
 
-app = FastAPI(title="Brain Space (STT→TTS coordinator)", version="3.1.0")
+app = FastAPI(title="Brain Space (STT→TTS coordinator)", version="3.2.0")
 log_queue: "asyncio.Queue[Dict[str, Any]]" = asyncio.Queue()
 
 def write_event(event: Dict[str, Any]) -> None:
@@ -70,9 +72,7 @@ async def stream_logs() -> StreamingResponse:
     return StreamingResponse(gen(), media_type="text/event-stream",
                              headers={"Cache-Control":"no-cache","Connection":"keep-alive"})
 
-# ---------- TTS proxy streaming (/tts/say.wav) ----------
-# GET:  /tts/say.wav?text=...&voice=...&rate_wpm=165
-# POST: JSON {"text": "...", "voice": "...", "rate_wpm": 165}
+# ---------- TTS proxy (file-mode, not live) ----------
 async def _proxy_tts_wav_stream(text: str, voice: str, rate_wpm: Optional[int],
                                 noise_scale: float, noise_w: float) -> StreamingResponse:
     import httpx
@@ -88,7 +88,6 @@ async def _proxy_tts_wav_stream(text: str, voice: str, rate_wpm: Optional[int],
         async with httpx.AsyncClient(timeout=None) as client:
             async with client.stream("GET", f"{TTS_BASE}/speak.wav", params=params) as resp:
                 if resp.status_code != 200:
-                    # bubble up exact error body from TTS
                     yield (await resp.aread())
                     return
                 async for chunk in resp.aiter_bytes():
@@ -123,8 +122,97 @@ async def tts_say_wav_post(req: Request):
     write_event({"type":"tts_post","len":len(text),"voice":voice,"rate_wpm":rate_wpm})
     return await _proxy_tts_wav_stream(text, voice, rate_wpm, noise_s, noise_wgt)
 
-# ---------- (Optional) simple relay demo kept for later ----------
-# You can keep your /demo/relay.wav here if you still want the file-upload STT→TTS demo.
+# ---------- NEW: true LIVE streaming via TTS WebSocket ----------
+def _wav_header(sr: int, ch: int, bits: int = 16) -> bytes:
+    byte_rate = sr * ch * (bits // 8)
+    block_align = ch * (bits // 8)
+    # streaming data size; many players accept 0xFFFFFFFF as "unknown length"
+    data_size = 0xFFFFFFFF
+    riff_size = 36 + data_size
+    return b"".join([
+        b"RIFF", struct.pack("<I", riff_size), b"WAVE",
+        b"fmt ", struct.pack("<I", 16), struct.pack("<H", 1),
+        struct.pack("<H", ch), struct.pack("<I", sr),
+        struct.pack("<I", byte_rate), struct.pack("<H", block_align),
+        struct.pack("<H", bits),
+        b"data", struct.pack("<I", data_size),
+    ])
+
+def _tts_ws_url() -> str:
+    base = TTS_BASE.rstrip("/")
+    return base.replace("http://", "ws://").replace("https://", "wss://") + "/ws/tts"
+
+@app.get("/tts/say.stream.wav")
+async def tts_say_stream_wav(
+    text: str = Query(..., description="Text to synthesize (live)"),
+    voice: str = Query(DEFAULT_VOICE),
+    rate_wpm: Optional[int] = Query(BASE_WPM),
+    length_scale: Optional[float] = Query(None),
+    noise_scale: float = Query(NOISE_SCALE),
+    noise_w: float = Query(NOISE_W),
+):
+    """
+    LIVE streaming proxy: TTS WS (raw PCM16) -> HTTP chunked WAV.
+    Starts emitting audio as soon as the TTS starts producing frames.
+    """
+    ls = float(length_scale) if length_scale is not None else rate_to_length_scale(rate_wpm or BASE_WPM)
+    write_event({"type":"tts_stream_get","len":len(text),"voice":voice,"ls":ls})
+
+    async def gen():
+        ws_url = _tts_ws_url()
+        async with websockets.connect(ws_url, ping_interval=None, max_size=8_000_000) as ws:
+            # init
+            await ws.send(json.dumps({
+                "event": "init",
+                "voice": voice,
+                "length_scale": ls,
+                "noise_scale": noise_scale,
+                "noise_w": noise_w,
+            }))
+            sr, ch = 22050, 1
+            # wait for ready -> send header
+            while True:
+                m = await ws.recv()
+                if isinstance(m, (bytes, bytearray)):
+                    # unlikely before ready; ignore
+                    continue
+                try:
+                    evt = json.loads(m)
+                except Exception:
+                    continue
+                if evt.get("event") == "ready":
+                    sr = int(evt.get("sr", 22050))
+                    ch = int(evt.get("channels", 1))
+                    yield _wav_header(sr, ch)
+                    break
+                if evt.get("event") == "error":
+                    yield f'ERROR: {evt.get("detail","tts init error")}'.encode("utf-8")
+                    return
+
+            # speak
+            await ws.send(json.dumps({"event": "speak", "text": text}))
+
+            # pump frames
+            while True:
+                try:
+                    msg = await ws.recv()
+                except websockets.exceptions.ConnectionClosed:
+                    break
+                if isinstance(msg, (bytes, bytearray)):
+                    yield msg
+                else:
+                    try:
+                        evt = json.loads(msg)
+                    except Exception:
+                        continue
+                    if evt.get("event") in ("done", "end"):
+                        break
+                    if evt.get("event") == "error":
+                        yield f'ERROR: {evt.get("detail","tts error")}'.encode("utf-8")
+                        break
+
+    return StreamingResponse(gen(), media_type="audio/wav",
+                             headers={"Cache-Control":"no-cache","Connection":"keep-alive"})
 
 # ---------- Optional: serve saved files later ----------
 @app.get("/files/{name}")