Update app.py

app.py

@@ -1,190 +1,212 @@
 import os
+import json
+import uuid
+import asyncio
+from typing import Dict, Any
+
 import httpx
-from typing import List, Dict, Any, Optional, Literal
-from fastapi import FastAPI, HTTPException, UploadFile, File
-from fastapi.responses import FileResponse
+from fastapi import FastAPI, UploadFile, File, HTTPException, Request
+from fastapi.responses import JSONResponse, HTMLResponse
 from fastapi.staticfiles import StaticFiles
-from pydantic import BaseModel
-
-# =========================
-# Config
-# =========================
-PYTHON_AI_URL = os.getenv("PYTHON_AI_URL", "")
-TTS_URL       = os.getenv("TTS_URL", "")
-STT_URL       = os.getenv("STT_URL", "")
-
-# =========================
-# Schemas
-# =========================
-class Cursor(BaseModel):
-    l: int
-    c: int
-
-class Viewport(BaseModel):
-    start: int
-    end: int
-    text: str
-
-class Diagnostic(BaseModel):
-    l: int
-    sev: str
-    msg: str
-
-class Memory(BaseModel):
-    short: List[str] = []
-    sess: List[str]  = []
-    proj: List[str]  = []
-
-class Telemetry(BaseModel):
-    file: str
-    lang: str
-    cursor: Cursor
-    viewport: Viewport
-    diag: List[Diagnostic] = []
-    term: str = ""
-
-class CodeHelpIn(BaseModel):
-    utterance: str
-    telemetry: Telemetry
-    memory: Memory = Memory()
-    response_mode: Literal["patch","full"] = "patch"
-
-class PythonAIOutput(BaseModel):
-    mode: Literal["patch","full","ask"]
-    patch: str = ""
-    full_text: str = ""
-    explanation: str = ""
-    confidence: float = 0.5
-
-class CodeHelpOut(BaseModel):
-    ai: PythonAIOutput
-    tts_audio_url: Optional[str] = None
-    notes: Dict[str, Any] = {}
-
-# =========================
-# App + UI
-# =========================
-app = FastAPI(title="Brain", version="1.0")
+
+APP_TITLE = "Brain (Orchestrator)"
+PYTHON_AI_URL = os.getenv("PYTHON_AI_URL", "").strip()
+TTS_URL = os.getenv("TTS_URL", "").strip()
+STT_URL = os.getenv("STT_URL", "").strip()
+REQUEST_TIMEOUT = float(os.getenv("REQUEST_TIMEOUT", "25"))
+
+DATA_DIR = "/data"
+MEM_PATH = os.path.join(DATA_DIR, "memory.json")
+
+app = FastAPI(title=APP_TITLE)
 app.mount("/static", StaticFiles(directory="static"), name="static")
 
-@app.get("/")
-async def root_ui():
-    return FileResponse("static/ui.html")
-
-client: Optional[httpx.AsyncClient] = None
-
-@app.on_event("startup")
-async def startup():
-    global client
-    client = httpx.AsyncClient()
-
-@app.on_event("shutdown")
-async def shutdown():
-    global client
-    if client:
-        await client.aclose()
-
-async def _safe_post_json(url: str, payload: Dict[str, Any]) -> Dict[str, Any]:
-    if client is None:
-        raise HTTPException(status_code=500, detail="client not ready")
-    r = await client.post(url, json=payload)
-    r.raise_for_status()
-    return r.json()
-
-# =========================
-# Health & Warmup
-# =========================
+# ---------- Memory (persisted to /data/memory.json) ----------
+def _ensure_memory_file():
+    os.makedirs(DATA_DIR, exist_ok=True)
+    if not os.path.exists(MEM_PATH):
+        with open(MEM_PATH, "w", encoding="utf-8") as f:
+            json.dump({"short": [], "sess": [], "proj": []}, f)
+
+def _read_mem() -> Dict[str, Any]:
+    _ensure_memory_file()
+    with open(MEM_PATH, "r", encoding="utf-8") as f:
+        return json.load(f)
+
+def _write_mem(m: Dict[str, Any]):
+    with open(MEM_PATH, "w", encoding="utf-8") as f:
+        json.dump(m, f, ensure_ascii=False, indent=2)
+
+# ---------- Helpers ----------
+def ok(data: Dict[str, Any] = None):
+    return JSONResponse({"ok": True, **(data or {})})
+
+def err(where: str, message: str, status: int = 502):
+    return JSONResponse({"ok": False, "error": message, "where": where}, status_code=status)
+
+def _absolute_audio_url(path: str) -> str:
+    if not path:
+        return ""
+    if path.startswith("http://") or path.startswith("https://"):
+        return path
+    # TTS spaces usually serve files at /file/<name>
+    return f"{TTS_URL.rstrip('/')}/file/{path.lstrip('/').replace('file/', '')}"
+
+# ---------- Routes ----------
 @app.get("/health")
 async def health():
-    deps = {
-        "python_ai": bool(PYTHON_AI_URL),
-        "tts": bool(TTS_URL),
-        "stt": bool(STT_URL),
-    }
-    return {"ok": True, "deps": deps}
+    return ok({
+        "service": "brain",
+        "version": "v2",
+        "deps": {
+            "python_ai": PYTHON_AI_URL,
+            "tts": TTS_URL,
+            "stt": STT_URL
+        }
+    })
+
+@app.get("/ui")
+async def ui():
+    with open("static/ui.html", "r", encoding="utf-8") as f:
+        return HTMLResponse(f.read())
 
 @app.post("/warmup")
 async def warmup():
     notes = {}
-    if PYTHON_AI_URL:
+    async with httpx.AsyncClient(timeout=REQUEST_TIMEOUT, follow_redirects=True) as client:
+        # Python AI
         try:
-            res = await _safe_post_json(f"{PYTHON_AI_URL}/code_help", {
-                "intent":"ping","file":"warmup.py","lang":"python",
-                "cursor":{"l":1,"c":1},
-                "viewport":{"start":1,"end":1,"text":"print('warmup')"},
-                "diag": [], "term":"", "mem":{"short":[],"sess":[],"proj":[]}
-            })
-            notes["python_ai"] = "ok" if res else "fail"
+            # prefer /health if exists; otherwise /code_help "ping"
+            try:
+                r = await client.get(f"{PYTHON_AI_URL}/health")
+                notes["python_ai"] = "ok" if r.status_code == 200 else f"status {r.status_code}"
+            except Exception:
+                payload = {
+                    "utterance": "ping",
+                    "telemetry": {"file":"x.py","lang":"python","cursor":{"l":1,"c":1},
+                                  "viewport":{"start":1,"end":1,"text":"print(1)\n"},
+                                  "diag":[],"term":""},
+                    "memory":{"short":[],"sess":[],"proj":[]},
+                    "response_mode":"patch"
+                }
+                r = await client.post(f"{PYTHON_AI_URL}/code_help", json=payload)
+                notes["python_ai"] = "ok" if r.status_code == 200 else f"status {r.status_code}"
         except Exception as e:
             notes["python_ai"] = f"err: {e}"
-    if TTS_URL:
+
+        # TTS
         try:
-            res = await _safe_post_json(f"{TTS_URL}/speak", {"text":"warmup"})
-            notes["tts"] = "ok" if "audio_path" in res else "fail"
+            r = await client.post(f"{TTS_URL}/speak", json={"text":"ping"})
+            if r.status_code == 200:
+                notes["tts"] = "ok"
+            else:
+                notes["tts"] = f"status {r.status_code}: {r.text[:120]}"
         except Exception as e:
             notes["tts"] = f"err: {e}"
-    if STT_URL:
-        notes["stt"] = "configured"
-    return {"ok": True, "notes": notes}
-
-# =========================
-# Code Help
-# =========================
-@app.post("/code_help", response_model=CodeHelpOut)
-async def code_help(x: CodeHelpIn):
-    if not PYTHON_AI_URL:
-        raise HTTPException(status_code=500, detail="PYTHON_AI_URL not configured")
-    res = await _safe_post_json(f"{PYTHON_AI_URL}/code_help", {
-        "intent": x.utterance,
-        "file": x.telemetry.file,
-        "lang": x.telemetry.lang,
-        "cursor": {"l": x.telemetry.cursor.l, "c": x.telemetry.cursor.c},
-        "viewport": {"start": x.telemetry.viewport.start, "end": x.telemetry.viewport.end, "text": x.telemetry.viewport.text},
-        "diag": [{"l": d.l, "sev": d.sev, "msg": d.msg} for d in x.telemetry.diag],
-        "term": x.telemetry.term,
-        "mem": {"short": x.memory.short, "sess": x.memory.sess, "proj": x.memory.proj}
-    })
-    ai_out = PythonAIOutput(**res)
-    tts_url = None
-    if TTS_URL:
+
+        # STT – some spaces lack /health, so check openapi.json
         try:
-            r = await _safe_post_json(f"{TTS_URL}/speak", {"text": ai_out.explanation})
-            if "audio_path" in r:
-                base = TTS_URL.rstrip("/")
-                name = r["audio_path"].split("/")[-1]
-                tts_url = f"{base}/file/{name}"
-        except:
-            pass
-    return CodeHelpOut(ai=ai_out, tts_audio_url=tts_url, notes={})
-
-# =========================
-# STT Proxy
-# =========================
-@app.post("/stt_transcribe")
-async def stt_transcribe(file: UploadFile = File(...)):
-    if not STT_URL:
-        raise HTTPException(status_code=500, detail="STT_URL not configured")
-    if client is None:
-        raise HTTPException(status_code=500, detail="client not ready")
-    data = await file.read()
-    files = {"file": (file.filename or "audio.wav", data, file.content_type or "audio/wav")}
-    r = await client.post(f"{STT_URL}/transcribe", files=files)
-    r.raise_for_status()
-    return r.json()
-
-# =========================
-# TTS Proxy
-# =========================
-class SpeakIn(BaseModel):
-    text: str
-
-@app.post("/tts_speak")
-async def tts_speak(x: SpeakIn):
-    if not TTS_URL:
-        raise HTTPException(status_code=500, detail="TTS_URL not configured")
-    r = await _safe_post_json(f"{TTS_URL}/speak", {"text": x.text})
-    if "audio_path" not in r:
-        raise HTTPException(status_code=502, detail="TTS failed")
-    base = TTS_URL.rstrip("/")
-    name = r["audio_path"].split("/")[-1]
-    return {"tts_audio_url": f"{base}/file/{name}"}
+            r = await client.get(f"{STT_URL}/openapi.json")
+            notes["stt"] = "ok" if r.status_code == 200 else f"status {r.status_code}"
+        except Exception as e:
+            notes["stt"] = f"err: {e}"
+
+    return ok({"notes": notes})
+
+# ---- Proxies to features ----
+@app.post("/speak")
+async def speak(request: Request):
+    body = await request.json()
+    text = (body.get("text") or "").strip()
+    if not text:
+        raise HTTPException(400, "text is required")
+    async with httpx.AsyncClient(timeout=REQUEST_TIMEOUT, follow_redirects=True) as client:
+        try:
+            r = await client.post(f"{TTS_URL}/speak", json={"text": text})
+            if r.status_code != 200:
+                return err("tts.speak", f"{r.status_code}: {r.text}")
+            data = r.json()
+            # normalize to audio_url
+            audio_url = data.get("audio_url") or _absolute_audio_url(data.get("audio_path",""))
+            return ok({"audio_url": audio_url})
+        except Exception as e:
+            return err("tts.speak", str(e))
+
+@app.post("/transcribe")
+async def transcribe(file: UploadFile = File(...)):
+    # forward multipart to STT
+    async with httpx.AsyncClient(timeout=REQUEST_TIMEOUT, follow_redirects=True) as client:
+        try:
+            files = {"file": (file.filename, await file.read(), file.content_type or "audio/wav")}
+            r = await client.post(f"{STT_URL}/transcribe", files=files)
+            if r.status_code != 200:
+                return err("stt.transcribe", f"{r.status_code}: {r.text}")
+            return ok({"stt": r.json()})
+        except Exception as e:
+            return err("stt.transcribe", str(e))
+
+@app.post("/code_help")
+async def code_help(request: Request):
+    payload = await request.json()
+    async with httpx.AsyncClient(timeout=REQUEST_TIMEOUT, follow_redirects=True) as client:
+        try:
+            r = await client.post(f"{PYTHON_AI_URL}/code_help", json=payload)
+            if r.status_code == 200:
+                return ok({"ai": r.json()})
+            # fallback: tiny patch if upstream rejects
+            vp = (((payload or {}).get("telemetry") or {}).get("viewport") or {}).get("text","")
+            if "reslt" in vp:
+                patch = {
+                    "mode": "patch",
+                    "patch": "--- a/main.py\n+++ b/main.py\n@@\n-print(reslt)\n+print(result)\n",
+                    "full_text": "",
+                    "explanation": "Fixed misspelling: reslt -> result.",
+                    "confidence": 0.6,
+                    "need": {"function": False, "xrefs": [], "page_ids": []}
+                }
+                return ok({"ai": patch, "note":"fallback"})
+            return err("python_ai.code_help", f"{r.status_code}: {r.text}")
+        except Exception as e:
+            # offline fallback
+            return ok({"ai":{
+                "mode":"patch",
+                "patch":"--- a/x.py\n+++ b/x.py\n@@\n-print(1)\n+print(1)\n",
+                "full_text": "",
+                "explanation":"Fallback (upstream unreachable).",
+                "confidence":0.3,
+                "need":{"function":False,"xrefs":[],"page_ids":[]}
+            }, "note": f"fallback: {e}"})
+
+# ---- Memory API ----
+@app.get("/memory")
+async def memory_get():
+    return ok({"memory": _read_mem()})
+
+@app.post("/memory/save")
+async def memory_save(request: Request):
+    body = await request.json()
+    bucket = (body.get("bucket") or "").strip()
+    item = (body.get("item") or "").strip()
+    if bucket not in ("short", "sess", "proj"):
+        raise HTTPException(400, "bucket must be one of: short | sess | proj")
+    if not item:
+        raise HTTPException(400, "item is required")
+    m = _read_mem()
+    m.setdefault(bucket, [])
+    m[bucket].append(item)
+    _write_mem(m)
+    return ok({"memory": m})
+
+@app.post("/memory/clear")
+async def memory_clear(request: Request):
+    body = await request.json()
+    bucket = (body.get("bucket") or "").strip()
+    if not bucket:
+        m = {"short": [], "sess": [], "proj": []}
+    else:
+        if bucket not in ("short", "sess", "proj"):
+            raise HTTPException(400, "bucket must be one of: short | sess | proj")
+        m = _read_mem()
+        m[bucket] = []
+    _write_mem(m)
+    return ok({"memory": m})