Create app.py

app.py

@@ -0,0 +1,286 @@
+import os
+import asyncio
+import time
+from typing import List, Dict, Any, Optional, Literal
+
+import httpx
+from fastapi import FastAPI, BackgroundTasks, HTTPException
+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
+
+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
+
+REQUEST_TIMEOUT_S   = float(os.getenv("REQUEST_TIMEOUT_S", "60"))
+CONNECT_TIMEOUT_S   = float(os.getenv("CONNECT_TIMEOUT_S", "10"))
+
+# =========================
+# Models / Schemas
+# =========================
+class Cursor(BaseModel):
+    l: int = Field(..., description="line")
+    c: int = Field(..., description="col")
+
+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 Need(BaseModel):
+    function: bool = False
+    xrefs: List[str] = []
+    page_ids: List[str] = []
+
+class PythonAIOutput(BaseModel):
+    mode: Literal["patch","full","ask"]
+    patch: str = ""
+    full_text: str = ""
+    explanation: str = ""
+    confidence: float = 0.5
+    need: Need = Need()
+
+class CodeHelpOut(BaseModel):
+    ai: PythonAIOutput
+    tts_audio_url: Optional[str] = None
+    used_memory_bytes: int
+    used_prompt_bytes: int
+    notes: Dict[str, Any] = {}
+
+# =========================
+# App & HTTP client
+# =========================
+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"}
+)
+
+# =========================
+# 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
+    return b[:budget].decode("utf-8", errors="ignore")
+
+def _shrink_lines_to_max(window: Viewport, max_lines: int) -> Viewport:
+    lines = window.text.splitlines()
+    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)
+
+async def _safe_post_json(url: str, payload: Dict[str, Any]) -> Dict[str, Any]:
+    try:
+        r = await client.post(url, json=payload)
+        r.raise_for_status()
+        return r.json()
+    except Exception as e:
+        raise HTTPException(status_code=502, detail=f"POST {url} failed: {e}")
+
+# =========================
+# Priority Queue (P0/P1/P2)
+# =========================
+# 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():
+    while True:
+        priority, ts, task = await TASK_Q.get()
+        try:
+            handler = task.get("handler")
+            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
+# =========================
+@app.get("/health")
+async def health():
+    deps = {
+        "python_ai": bool(PYTHON_AI_URL),
+        "tts": bool(TTS_URL),
+        "stt": bool(STT_URL),
+        "vision": bool(VISION_URL),
+        "memory": bool(MEMORY_URL),
+    }
+    return {"ok": True, "deps": deps, "version": "1.0"}
+
+@app.post("/warmup")
+async def warmup():
+    """Optionally ping downstream services to avoid cold starts."""
+    notes = {}
+    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},
+                "viewport":{"start":1,"end":1,"text":"print('warmup')\n"},
+                "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}"
+    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}"
+    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
+        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:
+    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", payload)
+    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]:
+    if not TTS_URL or not text:
+        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:
+        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,
+        "lang": t2.lang,
+        "cursor": {"l": t2.cursor.l, "c": t2.cursor.c},
+        "viewport": {"start": t2.viewport.start, "end": t2.viewport.end, "text": t2.viewport.text},
+        "diag": [{"l": d.l, "sev": d.sev, "msg": d.msg} for d in t2.diag],
+        "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,
+        used_memory_bytes=used_mem,
+        used_prompt_bytes=used_prompt,
+        notes={"response_mode": x.response_mode}
+    )