Update app.py

app.py

@@ -1,30 +1,31 @@
 import os
 import asyncio
-import time
 from typing import List, Dict, Any, Optional, Literal
 
 import httpx
-from fastapi import FastAPI, BackgroundTasks, HTTPException
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import FileResponse
+from fastapi.staticfiles import StaticFiles
 from pydantic import BaseModel, Field, ValidationError
 
 # =========================
 # Config (from env vars)
 # =========================
-PYTHON_AI_URL = os.getenv("PYTHON_AI_URL", "")          # e.g. https://<you>-python_ai_space.hf.space
-TTS_URL        = os.getenv("TTS_URL", "")               # e.g. https://<you>-tts_space.hf.space
-STT_URL        = os.getenv("STT_URL", "")               # (optional) future WS/HTTP
-VISION_URL     = os.getenv("VISION_URL", "")            # (optional) OCR/thumbnail summaries
-MEMORY_URL     = os.getenv("MEMORY_URL", "")            # (optional) external memory/RAG svc
+PYTHON_AI_URL = os.getenv("PYTHON_AI_URL", "")
+TTS_URL       = os.getenv("TTS_URL", "")
+STT_URL       = os.getenv("STT_URL", "")
+VISION_URL    = os.getenv("VISION_URL", "")
+MEMORY_URL    = os.getenv("MEMORY_URL", "")
 
 PROMPT_BUDGET_BYTES = int(os.getenv("PROMPT_BUDGET_BYTES", "12000"))
 MEMORY_BUDGET_BYTES = int(os.getenv("MEMORY_BUDGET_BYTES", "6000"))
-VIEWPORT_MAX_LINES  = int(os.getenv("VIEWPORT_MAX_LINES", "60"))  # visible ±30 default
+VIEWPORT_MAX_LINES  = int(os.getenv("VIEWPORT_MAX_LINES", "60"))
 
 REQUEST_TIMEOUT_S   = float(os.getenv("REQUEST_TIMEOUT_S", "60"))
 CONNECT_TIMEOUT_S   = float(os.getenv("CONNECT_TIMEOUT_S", "10"))
 
 # =========================
-# Models / Schemas
+# Schemas
 # =========================
 class Cursor(BaseModel):
     l: int = Field(..., description="line")
@@ -80,20 +81,45 @@ class CodeHelpOut(BaseModel):
     notes: Dict[str, Any] = {}
 
 # =========================
-# App & HTTP client
+# App + Static UI
 # =========================
 app = FastAPI(title="Brain (Router)", version="1.0")
 
-client = httpx.AsyncClient(
-    timeout=httpx.Timeout(REQUEST_TIMEOUT_S, connect=CONNECT_TIMEOUT_S),
-    headers={"User-Agent": "BrainRouter/1.0"}
-)
+# Serve the static UI
+app.mount("/static", StaticFiles(directory="static"), name="static")
+
+@app.get("/")
+async def root_ui():
+    return FileResponse("static/ui.html")
+
+# =========================
+# HTTP client (async)
+# =========================
+client: Optional[httpx.AsyncClient] = None
+
+@app.on_event("startup")
+async def _startup():
+    global client
+    client = httpx.AsyncClient(
+        timeout=httpx.Timeout(REQUEST_TIMEOUT_S, connect=CONNECT_TIMEOUT_S),
+        headers={"User-Agent": "BrainRouter/1.0"}
+    )
+    asyncio.create_task(worker_loop())
+    asyncio.create_task(worker_loop())
+
+@app.on_event("shutdown")
+async def _shutdown():
+    global client
+    try:
+        if client is not None:
+            await client.aclose()
+    except Exception:
+        pass
 
 # =========================
 # Small Utilities
 # =========================
 def _truncate_bytes(s: str, budget: int) -> str:
-    """Truncate a string to a byte budget (UTF-8 safe)."""
     b = s.encode("utf-8")
     if len(b) <= budget:
         return s
@@ -104,12 +130,13 @@ def _shrink_lines_to_max(window: Viewport, max_lines: int) -> Viewport:
     if len(lines) <= max_lines:
         return window
     keep = max_lines
-    start_line = max(window.start, window.end - keep + 1)
     slice_start = max(0, len(lines) - keep)
     new_text = "\n".join(lines[slice_start:])
-    return Viewport(start=start_line, end=window.end, text=new_text)
+    return Viewport(start=window.end - keep + 1, end=window.end, text=new_text)
 
 async def _safe_post_json(url: str, payload: Dict[str, Any]) -> Dict[str, Any]:
+    if client is None:
+        raise HTTPException(status_code=500, detail="HTTP client not initialized")
     try:
         r = await client.post(url, json=payload)
         r.raise_for_status()
@@ -118,11 +145,8 @@ async def _safe_post_json(url: str, payload: Dict[str, Any]) -> Dict[str, Any]:
         raise HTTPException(status_code=502, detail=f"POST {url} failed: {e}")
 
 # =========================
-# Priority Queue (P0/P1/P2)
+# Priority Queue (reserved)
 # =========================
-# P0: speech/telemetry (not used yet in this minimal Brain, reserved)
-# P1: code model + TTS (interactive)
-# P2: thumbnails / heavy analysis (future)
 TASK_Q: "asyncio.PriorityQueue[tuple[int,float,dict]]" = asyncio.PriorityQueue()
 
 async def worker_loop():
@@ -133,24 +157,10 @@ async def worker_loop():
             if handler:
                 await handler(**task.get("args", {}))
         except Exception:
-            # We keep the worker resilient; detailed logs would go here.
             pass
         finally:
             TASK_Q.task_done()
 
-@app.on_event("startup")
-async def _startup():
-    # Start a couple of workers
-    asyncio.create_task(worker_loop())
-    asyncio.create_task(worker_loop())
-
-@app.on_event("shutdown")
-async def _shutdown():
-    try:
-        await client.aclose()
-    except Exception:
-        pass
-
 # =========================
 # Health & Warmup
 # =========================
@@ -167,11 +177,9 @@ async def health():
 
 @app.post("/warmup")
 async def warmup():
-    """Optionally ping downstream services to avoid cold starts."""
-    notes = {}
+    notes: Dict[str, Any] = {}
     if PYTHON_AI_URL:
         try:
-            # If your Python AI exposes /health, use it. Otherwise skip.
             res = await _safe_post_json(f"{PYTHON_AI_URL}/code_help", {
                 "intent":"ping","file":"_warmup_.py","lang":"python",
                 "cursor":{"l":1,"c":1},
@@ -179,52 +187,35 @@ async def warmup():
                 "diag": [], "term":"", "mem":{"short":[],"sess":[],"proj":[]}
             })
             notes["python_ai"] = "ok" if res else "no-response"
-        except Exception as e:
-            notes["python_ai"] = f"err: {e}"
+        except HTTPException as e:
+            notes["python_ai"] = f"err: {e.detail}"
     if TTS_URL:
         try:
             res = await _safe_post_json(f"{TTS_URL}/speak", {"text":"warming up"})
             notes["tts"] = "ok" if "audio_path" in res else "no-audio"
-        except Exception as e:
-            notes["tts"] = f"err: {e}"
+        except HTTPException as e:
+            notes["tts"] = f"err: {e.detail}"
     return {"ok": True, "notes": notes}
 
 # =========================
 # Core: Code Help endpoint
 # =========================
 def _enforce_budgets(t: Telemetry, m: Memory) -> tuple[Telemetry, Memory, int, int]:
-    # shrink viewport to max lines
     t2 = Telemetry(
         file=t.file, lang=t.lang, cursor=t.cursor,
         viewport=_shrink_lines_to_max(t.viewport, VIEWPORT_MAX_LINES),
-        diag=t.diag[:5],  # cap diagnostics
+        diag=t.diag[:5],
         term=t.term
     )
-    # compress memory bullets and apply byte budget
     mem_text = " | ".join(m.short + m.sess + m.proj)
     mem_text = _truncate_bytes(mem_text, MEMORY_BUDGET_BYTES)
-    # reconstruct memory by naive split (keeps one string bucketed in 'sess')
     m2 = Memory(short=[], sess=[mem_text] if mem_text else [], proj=[])
-
-    # count budgets (approx: sum key strings + text fields)
     used_mem = len(mem_text.encode("utf-8"))
     prompt_bytes = (
         len(t2.file) + len(t2.lang) +
         len(t2.viewport.text) + sum(len(d.msg) for d in t2.diag) +
         len(t2.term) + used_mem
     )
-    if prompt_bytes > PROMPT_BUDGET_BYTES:
-        # tighten viewport again by half if still too large
-        t2 = Telemetry(
-            file=t2.file, lang=t2.lang, cursor=t2.cursor,
-            viewport=_shrink_lines_to_max(t2.viewport, max(20, VIEWPORT_MAX_LINES//2)),
-            diag=t2.diag, term=_truncate_bytes(t2.term, 1024)
-        )
-        prompt_bytes = (
-            len(t2.file) + len(t2.lang) +
-            len(t2.viewport.text) + sum(len(d.msg) for d in t2.diag) +
-            len(t2.term) + used_mem
-        )
     return t2, m2, used_mem, prompt_bytes
 
 async def _route_python_ai(payload: Dict[str, Any]) -> PythonAIOutput:
@@ -234,7 +225,6 @@ async def _route_python_ai(payload: Dict[str, Any]) -> PythonAIOutput:
     try:
         return PythonAIOutput(**res)
     except ValidationError as ve:
-        # If downstream returns bad JSON, surface as 502 with details.
         raise HTTPException(status_code=502, detail=f"Bad AI JSON schema: {ve}")
 
 async def _send_tts(text: str) -> Optional[str]:
@@ -242,23 +232,18 @@ async def _send_tts(text: str) -> Optional[str]:
         return None
     try:
         res = await _safe_post_json(f"{TTS_URL}/speak", {"text": text})
-        # HF Space fastapi static path helper
         audio_path = res.get("audio_path")
         if not audio_path:
             return None
-        # Convert to absolute URL the browser can open
         base = TTS_URL.rstrip("/")
         name = audio_path.split("/")[-1]
         return f"{base}/file/{name}"
-    except Exception:
+    except HTTPException:
         return None
 
 @app.post("/code_help", response_model=CodeHelpOut)
 async def code_help(x: CodeHelpIn):
-    # 1) enforce budgets / shrink context
     t2, m2, used_mem, used_prompt = _enforce_budgets(x.telemetry, x.memory)
-
-    # 2) build compact contract for Python AI (as agreed)
     py_in = {
         "intent": x.utterance,
         "file": t2.file,
@@ -269,14 +254,8 @@ async def code_help(x: CodeHelpIn):
         "term": t2.term,
         "mem": {"short": m2.short, "sess": m2.sess, "proj": m2.proj}
     }
-
-    # 3) call Python AI (async)
     ai_out = await _route_python_ai(py_in)
-
-    # 4) send short voice summary in parallel (explanation only)
     tts_url = await _send_tts(ai_out.explanation)
-
-    # 5) respond
     return CodeHelpOut(
         ai=ai_out,
         tts_audio_url=tts_url,