Update app.py

app.py

@@ -1,39 +1,34 @@
+# brain_app.py — Brain Space: STT → TTS proxy streamer
 import os, json, time, asyncio, tempfile
-from typing import AsyncGenerator, Dict, Any, Optional, List
-from fastapi import FastAPI, Request, Query, UploadFile, File
+from typing import AsyncGenerator, Dict, Any, Optional
+from fastapi import FastAPI, Request, Query, UploadFile
 from fastapi.responses import JSONResponse, StreamingResponse, FileResponse
 
-# ========== Directories ==========
+# === Directories ===
 BASE_DIR  = os.environ.get("BASE_DIR", "/tmp/brain_app")
 FILES_DIR = os.path.join(BASE_DIR, "files")
 LOGS_DIR  = os.path.join(FILES_DIR, "logs")
 EVENTS_FILE = os.path.join(LOGS_DIR, "events.jsonl")
-
 for p in (BASE_DIR, FILES_DIR, LOGS_DIR):
     os.makedirs(p, exist_ok=True)
 
-# ========== TTS Config ==========
-TTS_BASE  = os.environ.get("TTS_BASE", "https://Percy3822-ActualTTS.hf.space")
-BASE_WPM  = int(os.environ.get("BASE_WPM", "180"))
-NOISE_SCALE = float(os.environ.get("NOISE_SCALE", "0.33"))
-NOISE_W     = float(os.environ.get("NOISE_W", "0.92"))
-DEFAULT_VOICE = os.environ.get("DEFAULT_VOICE", "en_US-amy-medium")
-
-# ========== STT Config ==========
-STT_MODEL = os.environ.get("STT_MODEL", "base.en")  # faster-whisper model id
-STT_DEVICE = os.environ.get("STT_DEVICE", "cpu")    # "cpu" | "cuda"
-STT_COMPUTE = os.environ.get("STT_COMPUTE", "int8") # "int8"|"int8_float16"|"float32"
-STT_MAXLEN_S = float(os.environ.get("STT_MAXLEN_S", "600"))  # refuse extremely long uploads
+# === External Spaces ===
+TTS_BASE  = os.environ.get("TTS_BASE",  "https://Percy3822-ActualTTS.hf.space")
+# TODO: set your STT Space base here (example):
+STT_BASE  = os.environ.get("STT_BASE",  "https://YOUR-STT-SPACE.hf.space")
 
-# ========== App ==========
-app = FastAPI(title="Brain Space (TTS+STT)", version="2.1.0")
+# === TTS defaults ===
+DEFAULT_VOICE  = os.environ.get("DEFAULT_VOICE", "en_US-amy-medium")
+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"))
 
-# In-memory queue to fan-out logs to /stream/logs clients
+# === App ===
+app = FastAPI(title="Brain Space (STT→TTS coordinator)", version="3.0.0")
 log_queue: "asyncio.Queue[Dict[str, Any]]" = asyncio.Queue()
 
 def write_event(event: Dict[str, Any]) -> None:
     event.setdefault("ts", time.time())
-    os.makedirs(LOGS_DIR, exist_ok=True)
     with open(EVENTS_FILE, "a", encoding="utf-8") as f:
         f.write(json.dumps(event, ensure_ascii=False) + "\n")
     try:
@@ -41,16 +36,13 @@ def write_event(event: Dict[str, Any]) -> None:
     except asyncio.QueueFull:
         pass
 
-def clamp_rate(rate_wpm: Optional[int]) -> int:
-    if not isinstance(rate_wpm, int):
-        return BASE_WPM
-    return max(80, min(320, rate_wpm))
-
 def rate_to_length_scale(rate_wpm: Optional[int]) -> float:
-    r = clamp_rate(rate_wpm)
-    return round(BASE_WPM / float(r), 3)
+    base = BASE_WPM
+    if not isinstance(rate_wpm, int):
+        return 1.0
+    r = max(80, min(320, rate_wpm))
+    return round(base / float(r), 3)
 
-# ========== Health & Basics ==========
 @app.get("/health")
 def health():
     return {
@@ -58,24 +50,12 @@ def health():
         "service": "brain-space",
         "time": time.time(),
         "files_dir": FILES_DIR,
-        "logs_dir": LOGS_DIR,
         "tts_base": TTS_BASE,
-        "stt_model": STT_MODEL,
-        "stt_device": STT_DEVICE,
-        "stt_compute": STT_COMPUTE,
+        "stt_base": STT_BASE,
+        "defaults": {"voice": DEFAULT_VOICE, "rate_wpm": BASE_WPM}
     }
 
-@app.post("/process")
-async def process(req: Request):
-    try:
-        payload = await req.json()
-    except Exception:
-        return JSONResponse({"ok": False, "error": "Invalid JSON body"}, status_code=400)
-    event = {"type": "process", "data": payload}
-    write_event(event)
-    return {"ok": True, "received": payload}
-
-# ========== SSE Logs ==========
+# ========== SSE logs (optional) ==========
 @app.get("/stream/logs")
 async def stream_logs() -> StreamingResponse:
     async def gen() -> AsyncGenerator[bytes, None]:
@@ -88,72 +68,40 @@ async def stream_logs() -> StreamingResponse:
             pass
         while True:
             event = await log_queue.get()
-            line = json.dumps(event, ensure_ascii=False)
-            yield b"data: " + line.encode("utf-8") + b"\n\n"
-    headers = {"Cache-Control": "no-cache", "Connection": "keep-alive"}
-    return StreamingResponse(gen(), media_type="text/event-stream", headers=headers)
-
-@app.post("/log_error")
-async def log_error(req: Request):
-    try:
-        payload = await req.json()
-    except Exception:
-        return JSONResponse({"ok": False, "error": "Invalid JSON body"}, status_code=400)
-    event = {"type": "error", "data": payload}
-    write_event(event)
-    return {"ok": True}
-
-# ========== TTS: JSON (file URL) ==========
-@app.post("/tts/say")
-async def tts_say_json(req: Request):
-    """
-    POST JSON -> call TTS /speak (JSON) and return audio_url and audio_url_full.
-    Body:
-      {
-        "text": "Hello",
-        "voice": "en_US-amy-medium",
-        "rate_wpm": 165,           # optional (maps to length_scale)
-        "length_scale": 1.05,      # optional (overrides rate_wpm)
-        "noise_scale": 0.33,       # optional
-        "noise_w": 0.92            # optional
-      }
-    """
-    try:
-        body = await req.json()
-    except Exception:
-        return JSONResponse({"ok": False, "error": "Invalid JSON body"}, status_code=400)
+            yield b"data: " + json.dumps(event, ensure_ascii=False).encode("utf-8") + b"\n\n"
+    return StreamingResponse(gen(), media_type="text/event-stream",
+                             headers={"Cache-Control":"no-cache","Connection":"keep-alive"})
 
-    text  = (body.get("text") or "").strip()
-    if not text:
-        return JSONResponse({"ok": False, "error": "Missing text"}, status_code=400)
-
-    voice = (body.get("voice") or DEFAULT_VOICE).strip()
-    length_scale = float(body["length_scale"]) if "length_scale" in body else rate_to_length_scale(int(body.get("rate_wpm", BASE_WPM)))
-    noise_scale  = float(body.get("noise_scale", NOISE_SCALE))
-    noise_w      = float(body.get("noise_w", NOISE_W))
+# ---------- Helpers ----------
+async def _download_to_temp(url: str) -> str:
+    import httpx, os
+    _, ext = os.path.splitext(url.split("?")[0])
+    if not ext: ext = ".wav"
+    fd, tmp_path = tempfile.mkstemp(prefix="mic_", suffix=ext)
+    os.close(fd)
+    async with httpx.AsyncClient(timeout=300) as client:
+        r = await client.get(url)
+        r.raise_for_status()
+        with open(tmp_path, "wb") as f:
+            f.write(r.content)
+    return tmp_path
 
+async def _call_stt_transcribe_file(path: str) -> Dict[str, Any]:
+    """POST multipart 'audio' to STT /stt/transcribe and return its JSON."""
     import httpx
-    payload = {"text": text, "voice": voice, "length_scale": length_scale, "noise_scale": noise_scale, "noise_w": noise_w}
-    async with httpx.AsyncClient(timeout=180) as client:
-        resp = await client.post(f"{TTS_BASE}/speak", json=payload)
-        ok = resp.status_code == 200
-        try:
-            data = resp.json()
-        except Exception:
-            data = None
-
-    write_event({"type": "tts_say_json", "data": {"text_len": len(text), "voice": voice, "ok": ok, "resp": data}})
-
-    if not ok or not data or not data.get("ok"):
-        return JSONResponse({"ok": False, "error": (data or {}).get("error", f"TTS error {resp.status_code}")}, status_code=500)
-
-    audio_url = data["audio_url"]
-    audio_url_full = audio_url if audio_url.startswith("http") else f"{TTS_BASE}{audio_url}"
-    return {"ok": True, "audio_url": audio_url, "audio_url_full": audio_url_full, "voice": voice, "length_scale": length_scale}
-
-# ========== TTS: Direct WAV Proxy ==========
-async def _proxy_tts_wav_stream(text: str, voice: str, length_scale: float, noise_scale: float, noise_w: float, save_local: bool = False) -> StreamingResponse:
+    stt_url = f"{STT_BASE}/stt/transcribe"
+    files = {"audio": (os.path.basename(path), open(path, "rb"), "audio/wav")}
+    async with httpx.AsyncClient(timeout=300) as client:
+        r = await client.post(stt_url, files=files)
+        ok = r.status_code == 200
+        data = r.json() if ok else {"ok": False, "error": f"STT {r.status_code}"}
+    return data
+
+async def _proxy_tts_wav_stream(text: str, voice: str, rate_wpm: Optional[int],
+                                noise_scale: float, noise_w: float) -> StreamingResponse:
+    """Proxy stream from TTS /speak.wav based on text."""
     import httpx
+    length_scale = rate_to_length_scale(rate_wpm) if rate_wpm is not None else rate_to_length_scale(BASE_WPM)
     params = {
         "text": text,
         "voice": voice,
@@ -161,233 +109,79 @@ async def _proxy_tts_wav_stream(text: str, voice: str, length_scale: float, nois
         "noise_scale": f"{noise_scale:.3f}",
         "noise_w": f"{noise_w:.3f}",
     }
-    ts = int(time.time() * 1000)
-    local_path = os.path.join(FILES_DIR, f"say-{ts}.wav") if save_local else None
-
     async def gen():
         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:
-                    err_body = await resp.aread()
-                    yield err_body
+                    yield (await resp.aread())
                     return
-                f = None
-                try:
-                    if local_path:
-                        f = open(local_path, "wb")
-                    async for chunk in resp.aiter_bytes():
-                        if chunk:
-                            if f: f.write(chunk)
-                            yield chunk
-                finally:
-                    if f: f.close()
-
-    headers = {"Cache-Control": "no-cache"}
-    if local_path:
-        headers["X-Local-Path"] = local_path
-    return StreamingResponse(gen(), media_type="audio/wav", headers=headers)
-
-@app.get("/tts/say.wav")
-async def tts_say_wav_get(
-    text: str = Query(..., description="Text to synthesize"),
-    voice: str = Query(DEFAULT_VOICE, description="Voice id"),
-    rate_wpm: Optional[int] = Query(None, description="Words-per-minute"),
-    length_scale: Optional[float] = Query(None, description="Override length_scale"),
-    noise_scale: float = Query(NOISE_SCALE),
-    noise_w: float = Query(NOISE_W),
-    save: bool = Query(False, description="Also save under /files"),
-):
-    ls = float(length_scale) if length_scale is not None else rate_to_length_scale(rate_wpm if rate_wpm is not None else BASE_WPM)
-    write_event({"type": "tts_say_wav_get", "data": {"len": len(text), "voice": voice, "ls": ls, "save": save}})
-    return await _proxy_tts_wav_stream(text, voice, ls, noise_scale, noise_w, save_local=save)
-
-@app.post("/tts/say.wav")
-async def tts_say_wav_post(req: Request, save: bool = Query(False, description="Also save under /files")):
-    try:
-        body = await req.json()
-    except Exception:
-        return JSONResponse({"ok": False, "error": "Invalid JSON body"}, status_code=400)
-    text  = (body.get("text") or "").strip()
-    if not text:
-        return JSONResponse({"ok": False, "error": "Missing text"}, status_code=400)
-    voice = (body.get("voice") or DEFAULT_VOICE).strip()
-    ls = float(body["length_scale"]) if "length_scale" in body else rate_to_length_scale(int(body.get("rate_wpm", BASE_WPM)))
-    ns = float(body.get("noise_scale", NOISE_SCALE))
-    nw = float(body.get("noise_w", NOISE_W))
-    write_event({"type": "tts_say_wav_post", "data": {"len": len(text), "voice": voice, "ls": ls, "save": save}})
-    return await _proxy_tts_wav_stream(text, voice, ls, ns, nw, save_local=save)
-
-# ========== Serve saved files ==========
-@app.get("/files/{name}")
-def get_saved_file(name: str):
-    path = os.path.join(FILES_DIR, name)
-    if not os.path.exists(path):
-        return JSONResponse({"ok": False, "error": "not found"}, status_code=404)
-    return FileResponse(path, media_type="audio/wav", filename=name)
-
-# ========== STT (faster-whisper) ==========
-_model = None
-def _stt_model():
-    global _model
-    if _model is None:
-        from faster_whisper import WhisperModel
-        _model = WhisperModel(STT_MODEL, device=STT_DEVICE, compute_type=STT_COMPUTE)
-    return _model
-
-async def _download_to_temp(url: str) -> str:
-    import httpx
-    _, ext = os.path.splitext(url.split("?")[0])
-    if not ext: ext = ".wav"
-    fd, tmp_path = tempfile.mkstemp(prefix="stt_", suffix=ext)
-    os.close(fd)
-    async with httpx.AsyncClient(timeout=300) as client:
-        r = await client.get(url)
-        r.raise_for_status()
-        with open(tmp_path, "wb") as f:
-            f.write(r.content)
-    return tmp_path
-
-def _transcribe_path(path: str, language: Optional[str] = None) -> Dict[str, Any]:
-    model = _stt_model()
-    segments, info = model.transcribe(
-        path,
-        language=language,             # "en" or None for auto
-        beam_size=5,
-        vad_filter=False,
-        word_timestamps=False
-    )
-    out_segments: List[Dict[str, Any]] = []
-    txt_parts: List[str] = []
-    dur = getattr(info, "duration", None)
-    for seg in segments:
-        out_segments.append({"start": seg.start, "end": seg.end, "text": seg.text})
-        txt_parts.append(seg.text)
-        # guard against absurdly long files
-        if STT_MAXLEN_S and dur and seg.end and float(seg.end) > STT_MAXLEN_S:
-            break
-    text = "".join(txt_parts).strip()
-    return {"text": text, "language": getattr(info, "language", language or "unknown"), "duration": dur, "segments": out_segments}
-
-@app.post("/stt/transcribe")
-async def stt_transcribe(
-    req: Request,
-    language: Optional[str] = Query(None, description="ISO code like 'en' (None = auto)"),
-    file_url: Optional[str] = Query(None, description="If provided via query")
-):
-    """
-    POST either:
-      - multipart/form-data with 'audio' file
-      - or JSON: { "file_url": "https://..." }
-      - or query param ?file_url=...
-    Returns: { ok, text, language, duration, segments:[...] }
-    """
-    tmp_path = None
-    try:
-        content_type = req.headers.get("content-type","").lower()
-        if "multipart/form-data" in content_type:
-            form = await req.form()
-            up: UploadFile = form.get("audio")  # key: audio
-            if not up:
-                return JSONResponse({"ok": False, "error": "Missing 'audio' file"}, status_code=400)
-            suffix = os.path.splitext(up.filename or "")[1] or ".wav"
-            fd, tmp_path = tempfile.mkstemp(prefix="stt_", suffix=suffix)
-            os.close(fd)
-            with open(tmp_path, "wb") as f:
-                f.write(await up.read())
-        else:
-            # JSON or query
-            try:
-                body = await req.json()
-            except Exception:
-                body = {}
-            url = file_url or (body.get("file_url") if isinstance(body, dict) else None)
-            if not url:
-                return JSONResponse({"ok": False, "error": "Provide file_url (JSON/query) or multipart 'audio' file"}, status_code=400)
-            tmp_path = await _download_to_temp(url)
-
-        res = _transcribe_path(tmp_path, language=language)
-        write_event({"type": "stt_transcribe", "data": {"ok": True, "language": res.get("language"), "dur": res.get("duration"), "text_len": len(res.get("text",""))}})
-        return {"ok": True, **res}
-    except Exception as e:
-        write_event({"type": "stt_transcribe", "data": {"ok": False, "error": str(e)}})
-        return JSONResponse({"ok": False, "error": str(e)}, status_code=500)
-    finally:
-        try:
-            if tmp_path and os.path.exists(tmp_path):
-                os.unlink(tmp_path)
-        except Exception:
-            pass
-
-# --- End-to-end: STT -> Brain -> TTS (streamed WAV) ---
-@app.post("/demo/echo.wav")
-async def demo_echo_wav(
+                async for chunk in resp.aiter_bytes():
+                    if chunk:
+                        yield chunk
+    return StreamingResponse(gen(), media_type="audio/wav", headers={"Cache-Control":"no-cache"})
+
+# ========== The simple end-to-end endpoint ==========
+@app.post("/demo/relay.wav")
+async def demo_relay_wav(
     req: Request,
-    voice: str = Query(DEFAULT_VOICE, description="Voice id (TTS)"),
-    rate_wpm: Optional[int] = Query(None, description="Words-per-minute -> length_scale"),
-    length_scale: Optional[float] = Query(None, description="Override prosody"),
+    voice: str = Query(DEFAULT_VOICE),
+    rate_wpm: Optional[int] = Query(BASE_WPM),
     noise_scale: float = Query(NOISE_SCALE),
     noise_w: float = Query(NOISE_W),
-    save: bool = Query(False, description="Also save output WAV under /files"),
 ):
     """
-    POST either:
-      - multipart/form-data with 'audio' file
-      - or JSON: { "file_url": "https://..." }
-    Returns: streaming audio/wav that says what it heard.
+    Accept 5s mic recording from client (multipart 'audio' or JSON {file_url}),
+    send to STT Space for transcription, then IMMEDIATELY proxy stream TTS WAV
+    that speaks back what was heard.
     """
     tmp_path = None
     try:
-        # --- Ingest audio (multipart or JSON URL) ---
-        content_type = req.headers.get("content-type", "").lower()
-        if "multipart/form-data" in content_type:
+        # Ingest audio
+        ctype = (req.headers.get("content-type") or "").lower()
+        if "multipart/form-data" in ctype:
             form = await req.form()
-            up = form.get("audio")  # UploadFile
+            up: UploadFile = form.get("audio")
             if not up:
                 return JSONResponse({"ok": False, "error": "Missing 'audio' file"}, status_code=400)
-            suffix = os.path.splitext(getattr(up, "filename", "") or "")[1] or ".wav"
-            fd, tmp_path = tempfile.mkstemp(prefix="demo_echo_", suffix=suffix)
+            import os, tempfile
+            suffix = os.path.splitext(up.filename or "")[1] or ".wav"
+            fd, tmp_path = tempfile.mkstemp(prefix="mic_", suffix=suffix)
             os.close(fd)
             with open(tmp_path, "wb") as f:
                 f.write(await up.read())
         else:
-            # JSON with file_url
+            # JSON with {file_url}
             try:
                 body = await req.json()
             except Exception:
                 body = {}
             url = (body or {}).get("file_url")
             if not url:
-                return JSONResponse({"ok": False, "error": "Provide multipart 'audio' or JSON {file_url}"}, status_code=400)
+                return JSONResponse({"ok": False, "error": "Provide multipart 'audio' or JSON {file_url}"},
+                                    status_code=400)
             tmp_path = await _download_to_temp(url)
 
-        # --- STT ---
-        stt_res = _transcribe_path(tmp_path, language=None)
-        text = (stt_res.get("text") or "").strip()
+        # STT
+        stt = await _call_stt_transcribe_file(tmp_path)
+        if not stt.get("ok"):
+            write_event({"type":"relay","ok":False,"stage":"stt","err":stt.get("error")})
+            return JSONResponse({"ok": False, "error": f"STT failed: {stt.get('error')}"}, status_code=502)
+
+        text = (stt.get("text") or "").strip()
         if not text:
-            write_event({"type": "demo_echo", "data": {"ok": False, "error": "No speech detected"}})
+            write_event({"type":"relay","ok":False,"stage":"stt","err":"empty transcript"})
             return JSONResponse({"ok": False, "error": "No speech detected"}, status_code=422)
 
-        # --- Brain reply (simple confirmation) ---
-        reply_text = f"I heard: {text}"
-        reply_text = reply_text[:800]  # safety bound
+        # Brain reply (for demo we just echo; you can replace with actual brain logic later)
+        reply_text = f"I heard: {text}"[:800]
 
-        # Prosody parameters
-        ls = float(length_scale) if length_scale is not None else rate_to_length_scale(rate_wpm if rate_wpm is not None else BASE_WPM)
+        write_event({"type":"relay","ok":True,"heard_len":len(text),"voice":voice,"rate_wpm":rate_wpm})
 
-        write_event({"type": "demo_echo", "data": {"ok": True, "heard_len": len(text), "voice": voice, "ls": ls, "save": save}})
-
-        # --- TTS (stream WAV back to the caller) ---
-        return await _proxy_tts_wav_stream(
-            text=reply_text,
-            voice=voice,
-            length_scale=ls,
-            noise_scale=noise_scale,
-            noise_w=noise_w,
-            save_local=save
-        )
+        # TTS proxy stream (immediate)
+        return await _proxy_tts_wav_stream(reply_text, voice, rate_wpm, noise_scale, noise_w)
 
     except Exception as e:
-        write_event({"type": "demo_echo", "data": {"ok": False, "error": str(e)}})
+        write_event({"type":"relay","ok":False,"err":str(e)})
         return JSONResponse({"ok": False, "error": str(e)}, status_code=500)
     finally:
         try:
@@ -396,7 +190,14 @@ async def demo_echo_wav(
         except Exception:
             pass
 
-# ========== Optional direct runner ==========
+# Optional: serve saved files if you decide to persist later
+@app.get("/files/{name}")
+def get_file(name: str):
+    path = os.path.join(FILES_DIR, name)
+    if not os.path.exists(path):
+        return JSONResponse({"ok": False, "error": "not found"}, status_code=404)
+    return FileResponse(path, media_type="application/octet-stream", filename=name)
+
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run("brain_app:app", host="0.0.0.0", port=7861, reload=False)