app.py

"""
app.py — OCT-B agent as a Hugging Face *Docker* Space (Chainlit chat UI)
=======================================================================
This is the Chainlit variant of the Docker-Space frontend. It replaces the
Gradio single-shot form with a multi-turn chat: the user attaches an OCT B-scan
to a message and asks for a read; follow-up questions reuse the same loaded scan
and any segmentation already produced.

Dashboard behaviour
--------------------
When a turn produces visualizations, each one becomes a **selector button** on
that turn's reply, and the most recent is rendered in a **right-hand panel**
(Chainlit's ``ElementSidebar``). Clicking any selector — including ones from
earlier turns — swaps the panel to that visualization. This keeps the Gradio
side-by-side "read on the left, picture on the right" feel inside the chat.

The agent itself is unchanged. As before it is a model-driven ``SandboxAgent``
with shell + filesystem access: after segmenting it can author its own
visualization scripts inside a ``UnixLocalSandboxClient`` workspace and run them.

Per chat session it: opens the MIRAGE + LO-VLM MCP connections, provisions one
sandbox session (staging the oct-overlay skill), and keeps both alive for the
whole session. Each user turn runs one agent turn with that session wired via
``RunConfig``, collects whatever HTML the agent wrote to ``output/``, and exposes
it via the selector buttons + right-hand panel (plus a downloadable file). If the
sandbox can't be provisioned it falls back to the base agent for a text-only read.

SLO upload has been removed — the UI only accepts a B-scan.

Repo layout expected (this file at the Space root):
    Dockerfile
    app.py                <- this file
    requirements.txt
    README.md             <- carries the HF Docker Space header (sdk: docker)
    chainlit.md           <- Chainlit splash/readme (optional)
    .chainlit/config.toml <- Chainlit config (enables image attachments)
    public/elements/OverlayFrame.jsx
    oct_b_agent/          <- the unchanged project (agent/, tools/, skills/, ...)

Secrets (Space -> Settings -> Secrets):
    OPENAI_API_KEY     required
    MIRAGE_MCP_URL     deployed MIRAGE FastMCP URL (else config default)
    LO_VLM_MCP_URL     deployed LO-VLM FastMCP URL (else config default)
    AGENTOPS_API_KEY   optional (observability; no-op when unset)
    OCT_MODEL          optional (default gpt-5.4-mini)
    REASONING_EFFORT   optional none|low|medium|high (default low)
"""

from __future__ import annotations

import asyncio
import json
import os
import sys
import uuid
from pathlib import Path

# Flat-layout note: putting oct_b_agent/ on sys.path lets the project's
# `from agent... / tools / run` imports resolve. There is no top-level `agents/`
# package, so `from agents import ...` still binds to the OpenAI Agents SDK.
ROOT = Path(__file__).resolve().parent / "oct_b_agent"
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

import chainlit as cl  # noqa: E402
from chainlit.input_widget import Switch  # noqa: E402

from agents import Runner  # noqa: E402
from agents.run import RunConfig  # noqa: E402
from agents.sandbox import SandboxRunConfig  # noqa: E402
from agents.sandbox.sandboxes.unix_local import UnixLocalSandboxClient  # noqa: E402

from agent.oct_b_agent import build_oct_b_agent, build_oct_b_sandbox_agent  # noqa: E402
from tools import (  # noqa: E402
    OCTContext,
    OCTModelClients,
    SandboxManager,
    build_overlay_manifest,
    collect_overlay_outputs,
    init_observability,
)
from run import load_server_specs, load_infra_cost_model  # noqa: E402

# ── Config (resolved once at startup) ─────────────────────────────────────────
CONFIG = ROOT / "config" / "mcp_servers.json"
OVERLAY_SKILL_DIR = ROOT / "sandbox" / "oct-overlay"
MODEL = os.environ.get("OCT_MODEL", "gpt-5.4-mini")
_re = os.environ.get("REASONING_EFFORT", "low")
REASONING = None if _re == "none" else _re

SPECS = load_server_specs(CONFIG)
INFRA = load_infra_cost_model(CONFIG)

# Persistent per-session output dir so produced HTML can be read back + served.
OUT_ROOT = Path(os.environ.get("OCT_OUT_DIR", "/tmp/octb_out"))
OUT_ROOT.mkdir(parents=True, exist_ok=True)

# Name of the action that selects a visualization into the right-hand panel.
VIEW_VIZ_ACTION = "view_viz"

# AgentOps auto-instruments the SDK; init before the first agent build. No-op
# without AGENTOPS_API_KEY.
init_observability(tags=["oct-b", "hf-space", "docker", "chainlit", MODEL])

# Agents are cached across sessions: the skills registry is independent of the
# live MCP clients AND of the live sandbox session — both are injected per
# request (the clients via OCTContext, the session via RunConfig). Keep the
# (agent, registry) tuples referenced so the registries stay alive.
_SANDBOX_AGENT_CACHE: tuple | None = None
_BASE_AGENT_CACHE: tuple | None = None
_AGENT_LOCK = asyncio.Lock()


async def _get_sandbox_agent():
    global _SANDBOX_AGENT_CACHE
    async with _AGENT_LOCK:
        if _SANDBOX_AGENT_CACHE is None:
            _SANDBOX_AGENT_CACHE = await build_oct_b_sandbox_agent(
                ROOT / "skills", OVERLAY_SKILL_DIR,
                model=MODEL, reasoning_effort=REASONING,
            )
    return _SANDBOX_AGENT_CACHE[0]


async def _get_base_agent():
    global _BASE_AGENT_CACHE
    async with _AGENT_LOCK:
        if _BASE_AGENT_CACHE is None:
            _BASE_AGENT_CACHE = await build_oct_b_agent(
                ROOT / "skills", model=MODEL, reasoning_effort=REASONING,
            )
    return _BASE_AGENT_CACHE[0]


INTRO = (
    "### OCT-B — retinal OCT B-scan interpretation\n\n"
    "Attach an **OCT B-scan** image to your message and ask for a read "
    "(e.g. *“Interpret this B-scan”* or *“show me the layers as an overlay”*). "
    "Follow-up questions reuse the same scan and any segmentation already "
    "produced — no need to re-attach.\n\n"
    "Any visualization the agent produces shows up as a button on its reply and "
    "opens in the panel on the right; use the buttons to switch between them.\n\n"
    "_Decision support only — not a diagnostic device. Always confirm with a "
    "qualified ophthalmologist._"
)


# ── Right-hand panel helpers ──────────────────────────────────────────────────

async def _show_in_sidebar(viz: dict, vid: str) -> None:
    """Render one visualization in the right-hand ElementSidebar panel."""
    await cl.ElementSidebar.set_title(viz["title"])
    await cl.ElementSidebar.set_elements(
        [cl.CustomElement(
            name="OverlayFrame",
            props={"html": viz["html"], "title": viz["title"]},
            display="inline",
        )],
        key=f"viz-{vid}",
    )


@cl.action_callback(VIEW_VIZ_ACTION)
async def on_view_viz(action: cl.Action):
    """A selector button was clicked — swap the right-hand panel to that viz."""
    registry: dict = cl.user_session.get("viz_registry") or {}
    vid = (action.payload or {}).get("viz_id")
    viz = registry.get(vid)
    if not viz:
        await cl.ElementSidebar.set_title("Visualization unavailable")
        await cl.ElementSidebar.set_elements(
            [cl.Text(content="This visualization is no longer available in this session.")]
        )
        return
    await _show_in_sidebar(viz, vid)


# ── Per-session setup ─────────────────────────────────────────────────────────

@cl.on_chat_start
async def on_chat_start():
    # A visualization toggle, mirroring the old "Produce a visualization"
    # checkbox. The agent always *can* render on its own judgement; this just
    # nudges it for the whole session.
    settings = await cl.ChatSettings([
        Switch(
            id="want_visual",
            label="Produce a visualization (overlay / chart)",
            initial=False,
        ),
    ]).send()
    cl.user_session.set("want_visual", bool(settings.get("want_visual", False)))

    out_dir = OUT_ROOT / uuid.uuid4().hex
    out_dir.mkdir(parents=True, exist_ok=True)

    notes: list[str] = []
    async with cl.Step(name="Connecting to MIRAGE / LO-VLM…", type="tool"):
        # Keep the MCP connections alive for the whole session: enter the async
        # context manually here, exit it in on_chat_end. (A first connect after
        # idle may ride out a cold start via the bridge's built-in retries.)
        clients = OCTModelClients.from_specs(SPECS, infra_cost=INFRA)
        await clients.__aenter__()
        notes += [
            f"⚠️ MCP server `{k}` unavailable — {e}"
            for k, e in clients.unavailable.items()
        ]

        octx = OCTContext(clients=clients, output_dir=str(out_dir))

        # Provision one sandbox session for the session; fall back to a base
        # (text-only / lazy-overlay) agent if it can't be created.
        sandbox_session = None
        run_config = None
        try:
            sb_client = UnixLocalSandboxClient()
            sandbox_session = await sb_client.create(
                manifest=build_overlay_manifest(OVERLAY_SKILL_DIR))
            await sandbox_session.apply_manifest()   # stage skills/oct-overlay + inputs/output
            octx.sandbox_session = sandbox_session
            run_config = RunConfig(sandbox=SandboxRunConfig(session=sandbox_session))
            agent = await _get_sandbox_agent()
        except Exception as e:  # noqa: BLE001 — degrade to text-only rather than failing the chat
            notes.append(f"⚠️ Sandbox unavailable — running text-only ({e}).")
            octx.sandbox = SandboxManager(OVERLAY_SKILL_DIR)   # on-demand fallback path
            agent = await _get_base_agent()

    cl.user_session.set("clients", clients)
    cl.user_session.set("octx", octx)
    cl.user_session.set("agent", agent)
    cl.user_session.set("run_config", run_config)
    cl.user_session.set("sandbox_session", sandbox_session)
    cl.user_session.set("out_dir", out_dir)
    cl.user_session.set("conversation", None)     # None until the first turn runs
    cl.user_session.set("seen_overlays", set())   # filenames already collected
    cl.user_session.set("viz_registry", {})       # viz_id -> {name,title,html,path}

    content = INTRO + (("\n\n" + "\n".join(notes)) if notes else "")
    await cl.Message(content=content).send()


@cl.on_settings_update
async def on_settings_update(settings):
    cl.user_session.set("want_visual", bool(settings.get("want_visual", False)))


# ── Live "thinking" + LO-VLM caption helpers ──────────────────────────────────

# Friendly labels for the agent's tools, shown as live steps while it works.
# Anything not listed falls back to the raw tool name.
FRIENDLY_TOOL = {
    "load_oct_bscan":      "Loading B-scan",
    "load_slo":            "Loading SLO image",
    "list_loaded_images":  "Listing loaded images",
    "caption_oct":         "LO-VLM captioning",
    "segment_layers":      "MIRAGE layer segmentation",
    "extract_features":    "MIRAGE feature extraction",
    "reconstruct_oct":     "MIRAGE reconstruction",
    "mcp_health":          "Checking model health",
    "render_layer_overlay":"Rendering layer overlay",
    "stage_overlay":       "Staging overlay inputs",
    "save_artifact":       "Saving artifact",
    "get_skill_body":      "Reading skill instructions",
    "get_skill_metadata":  "Reading skill metadata",
    "get_skill_reference": "Reading skill reference",
    "get_skill_asset":     "Reading skill asset",
    "get_skill_script":    "Reading skill script",
    "run_skill_script":    "Running skill script",
}


def _compact(value, limit: int = 800) -> str:
    """Stringify + trim a tool arg/output blob for display inside a Step."""
    if value is None:
        return ""
    s = value if isinstance(value, str) else json.dumps(value, default=str)
    s = s.strip()
    return s if len(s) <= limit else s[:limit] + " …"


def _extract_caption(output) -> str | None:
    """Pull the full caption string out of a caption_oct tool result."""
    try:
        data = json.loads(output) if isinstance(output, str) else output
    except Exception:  # noqa: BLE001
        return None
    if isinstance(data, dict):
        cap = data.get("caption")
        if cap:
            return str(cap).strip()
    return None


def _caption_block(caption: str) -> str:
    """Render the verbatim LO-VLM caption as a labelled blockquote section."""
    cap = caption.strip()
    quoted = "\n".join(
        (f"> {ln}" if ln.strip() else ">") for ln in cap.splitlines()
    ) or f"> {cap}"
    return (
        "### Narrative (LO-VLM) — full caption\n\n"
        f"{quoted}\n\n"
        "_Verbatim model output above; the structured read and summary follow._"
    )


# ── One chat turn ─────────────────────────────────────────────────────────────


def _first_bscan_path(message: cl.Message):
    """Return the local path of the first image attached to the message, if any."""
    for el in (message.elements or []):
        mime = (getattr(el, "mime", "") or "")
        if mime.startswith("image") and getattr(el, "path", None):
            return el.path
    return None


def _build_user_text(prompt: str, bscan_path, want_visual: bool) -> str:
    parts = [prompt or "Interpret this OCT B-scan and produce a structured read."]
    if bscan_path:
        # The agent loads images from a path and never routes bytes through the
        # model context, so we pass the source path, not the image itself.
        parts.append(f"\n\nB-scan source: {bscan_path}")
    if want_visual:
        parts.append(
            "\n\nThe user would also like a visualization (the standard layer "
            "overlay, or a more apt chart if the findings call for one)."
        )
    return "".join(parts)


async def _collect_new_visualizations(octx, sandbox_session, out_dir, seen) -> list[str]:
    """Names of HTML visualizations produced this turn (not collected before)."""
    if sandbox_session is not None:
        produced = await collect_overlay_outputs(sandbox_session, out_dir)
    else:
        produced = list(octx.overlays)
    new = [n for n in produced if n not in seen]
    seen.update(produced)
    return new


def _register_visualizations(new_files: list[str], out_dir: Path, registry: dict):
    """Register each new viz, returning (selector actions, download files, last_vid).

    Each visualization becomes a clickable selector (Action) that opens it in the
    right-hand panel, plus a downloadable File. The registry persists for the whole
    session so buttons from earlier turns keep working.
    """
    actions: list = []
    files: list = []
    last_vid: str | None = None
    for name in new_files:
        fpath = out_dir / name
        if not fpath.exists():
            continue
        vid = uuid.uuid4().hex[:10]
        registry[vid] = {
            "name": name,
            "title": name,
            "html": fpath.read_text(errors="replace"),
            "path": str(fpath),
        }
        actions.append(cl.Action(
            name=VIEW_VIZ_ACTION,
            payload={"viz_id": vid},
            label=name,
            icon="image",
            tooltip="Show this visualization in the panel on the right",
        ))
        # Set mime explicitly: Chainlit infers mime via `filetype`, which only
        # detects binary signatures and returns None for .html — and the File
        # renderer calls .startsWith on the mime, crashing the UI on null.
        files.append(cl.File(
            name=name, path=str(fpath), mime="text/html", display="inline"))
        last_vid = vid
    return actions, files, last_vid


@cl.on_message
async def on_message(message: cl.Message):
    octx = cl.user_session.get("octx")
    agent = cl.user_session.get("agent")
    run_config = cl.user_session.get("run_config")
    sandbox_session = cl.user_session.get("sandbox_session")
    out_dir = cl.user_session.get("out_dir")
    conversation = cl.user_session.get("conversation")
    seen = cl.user_session.get("seen_overlays")
    registry = cl.user_session.get("viz_registry") or {}
    want_visual = cl.user_session.get("want_visual", False)

    if octx is None:
        await cl.Message(content="Session not initialised — please refresh to start a new chat.").send()
        return

    bscan_path = _first_bscan_path(message)

    # The first turn needs a B-scan to load. Later turns reuse it via history.
    if conversation is None and not bscan_path:
        await cl.Message(
            content="Please attach an OCT B-scan image to your message to begin."
        ).send()
        return

    user_text = _build_user_text((message.content or "").strip(), bscan_path, want_visual)
    run_input = user_text if conversation is None else conversation + [
        {"role": "user", "content": user_text}
    ]

    # Stream one agent turn. run_streamed() returns synchronously; the run is not
    # complete until stream_events() is fully consumed. We surface each tool call
    # as a live Step (the "thinking" view), stream the final answer tokens, and
    # intercept the raw caption_oct output so we can show the full LO-VLM caption.
    answer = cl.Message(content="")
    await answer.send()

    open_steps: list[list] = []     # [call_id, cl.Step, tool_name] awaiting output
    caption_text: str | None = None

    def _match_output(call_id):
        # Prefer matching by call_id; fall back to FIFO (calls resolve in order).
        for i, entry in enumerate(open_steps):
            if call_id is not None and entry[0] == call_id:
                return open_steps.pop(i)
        return open_steps.pop(0) if open_steps else None

    try:
        result = Runner.run_streamed(
            agent, run_input, context=octx, max_turns=20, run_config=run_config)

        async for ev in result.stream_events():
            # Token-level deltas → stream the final answer as it's generated.
            if ev.type == "raw_response_event":
                data = getattr(ev, "data", None)
                if getattr(data, "type", "") == "response.output_text.delta":
                    await answer.stream_token(getattr(data, "delta", "") or "")
                continue

            if ev.type != "run_item_stream_event":
                continue

            # A tool was invoked → open a live "thinking" step for it.
            if ev.name == "tool_called":
                raw = getattr(ev.item, "raw_item", None)
                tname = getattr(raw, "name", None) or "tool"
                cid = getattr(raw, "call_id", None) or getattr(raw, "id", None)
                step = cl.Step(name=FRIENDLY_TOOL.get(tname, tname), type="tool")
                step.input = _compact(getattr(raw, "arguments", None))
                await step.send()
                open_steps.append([cid, step, tname])

            # That tool returned → close its step and (if it was the captioner)
            # capture the full caption verbatim.
            elif ev.name == "tool_output":
                raw = getattr(ev.item, "raw_item", None)
                cid = (raw.get("call_id") if isinstance(raw, dict)
                       else getattr(raw, "call_id", None))
                output = getattr(ev.item, "output", None)
                if output is None and isinstance(raw, dict):
                    output = raw.get("output")

                match = _match_output(cid)
                if match is not None:
                    _, step, tname = match
                    step.output = _compact(output)
                    await step.update()
                    if tname == "caption_oct":
                        cap = _extract_caption(output)
                        if cap:
                            caption_text = cap
    except Exception as e:  # noqa: BLE001 — keep the session alive on a failed turn
        await answer.stream_token(f"\n\n❌ Run failed: {e}")
        await answer.update()
        return

    # Collect any visualizations the agent produced this turn.
    new_files = await _collect_new_visualizations(octx, sandbox_session, out_dir, seen)
    cl.user_session.set("seen_overlays", seen)
    actions, files, last_vid = _register_visualizations(new_files, out_dir, registry)
    cl.user_session.set("viz_registry", registry)

    # Compose the final reply: the verbatim LO-VLM caption block (if any) on top,
    # then the agent's structured read, then the visualization footer. We rewrite
    # the streamed message once so the caption block lands above the report.
    sections: list[str] = []
    if caption_text:
        sections.append(_caption_block(caption_text))
        sections.append("---")
    sections.append(str(result.final_output))

    content = "\n\n".join(sections)
    if actions:
        n = len(actions)
        content += (
            f"\n\n---\n*{n} visualization{'s' if n > 1 else ''} ready — "
            f"use the button{'s' if n > 1 else ''} below to view "
            f"{'them' if n > 1 else 'it'} in the panel on the right.*"
        )

    answer.content = content
    answer.elements = files
    answer.actions = actions
    await answer.update()

    # Auto-open the right-hand panel on the most recent visualization so the
    # dashboard view appears without a click; the buttons switch between them.
    if last_vid:
        await _show_in_sidebar(registry[last_vid], last_vid)

    # Persist history so follow-ups reuse the loaded scan + produced artifacts.
    cl.user_session.set("conversation", result.to_input_list())


# ── Teardown ──────────────────────────────────────────────────────────────────

@cl.on_chat_end
async def on_chat_end():
    sandbox_session = cl.user_session.get("sandbox_session")
    octx = cl.user_session.get("octx")
    clients = cl.user_session.get("clients")

    if sandbox_session is not None:
        try:
            await sandbox_session.aclose()
        except Exception:  # noqa: BLE001
            pass
    if octx is not None and getattr(octx, "sandbox", None) is not None:
        try:
            await octx.sandbox.aclose()
        except Exception:  # noqa: BLE001
            pass
    if clients is not None:
        try:
            await clients.__aexit__(None, None, None)
        except Exception:  # noqa: BLE001
            pass