dashboard/agent_graph.py

"""dashboard/agent_graph.py — live SVG graph of the LangGraph agent topology.

Rendered inside _live_trace_fragment() in app.py while a brief is being generated.
The graph shows the Signals → Orchestrator → Synthesis flow on the left, with the
5 tools fanning out to the right — nodes illuminate as the agent uses them.

Live state is derived from the streaming trace list (already available in app.py):
  • active tool  → green (pending tool_call this round)
  • used tools   → pale green (already called this run)
  • synthesis    → green when in_progress, pale green when done
  • signals      → green when trace is empty (startup)

No new dependencies — pure SVG injected via st.markdown(..., unsafe_allow_html=True).
Tool names mirror agent/tools.py; update TOOL_DISPLAY here if tools change.
"""
from __future__ import annotations

import streamlit as st

from dashboard.theme import (
    BG, BG_MUTED, BORDER, BORDER_STRONG,
    TEXT, TEXT_MUTED, TEXT_FAINT,
    BRAND_GREEN, BRAND_GREEN_DARK,
    GREEN, BULL_BG, BULL_BORDER,
    BLUE,
    AI_BG, AI_COLOR,
)


# ── Tool display metadata ─────────────────────────────────────────────────────
# Mirrors the 5 tools registered in agent/tools.py — update here if tools change.
TOOL_DISPLAY: dict[str, tuple[str, str]] = {
    "get_financial_metrics":    ("📊", "Financials"),
    "search_filing":            ("📄", "Filings"),
    "search_transcript":        ("🎙", "Transcript"),
    "get_analyst_expectations": ("📈", "Analysts"),
    "search_news":              ("📰", "News"),
}
_TOOL_NAMES = list(TOOL_DISPLAY.keys())

# Phase display metadata
_PHASE_META: dict[str, tuple[str, str]] = {
    "signals":   ("⚡", "Starting up…"),
    "thinking":  ("💭", "Thinking…"),
    "tool":      ("🔧", "Using tools"),
    "synthesis": ("✍️", "Synthesizing brief…"),
    "done":      ("✅", "Brief ready"),
}


# ── HTML escape helper ────────────────────────────────────────────────────────

def _he(s: str) -> str:
    """Minimal HTML escape for user-derived text injected into markup."""
    return (
        s.replace("&", "&")
         .replace("<", "&lt;")
         .replace(">", "&gt;")
         .replace('"', "&quot;")
    )


# ── State derivation (pure — no Streamlit dependency) ─────────────────────────

def derive_live_state(trace: list[dict]) -> dict:
    """Derive display state from the current trace snapshot.

    Args:
        trace: list of step dicts from dashboard/reasoning.py absorb_* helpers.

    Returns dict with:
        used_tools    set[str]   — tool names with status "done" this run
        active_tools  list[str]  — tool names with status "pending" (current round)
        phase         str        — "signals"|"thinking"|"tool"|"synthesis"|"done"
        latest_thought str       — last assistant_text, truncated to ~140 chars
        latest_result  dict|None — {name, line} of last tool_result, or None
        active_args    str       — short first arg of active tool call (may be "")
    """
    if not trace:
        return {
            "used_tools": set(),
            "active_tools": [],
            "phase": "signals",
            "latest_thought": "",
            "latest_result": None,
            "active_args": "",
        }

    used_tools: set[str] = set()
    active_tools: list[str] = []
    active_args: str = ""
    latest_thought: str = ""
    latest_result: dict | None = None
    phase: str = "thinking"

    for step in trace:
        kind = step["kind"]

        if kind == "assistant_text":
            text = step.get("text", "").strip().replace("\n", " ")
            if text:
                latest_thought = text[:140] + ("…" if len(text) > 140 else "")

        elif kind == "tool_call":
            name = step.get("name", "")
            if step.get("status") == "done":
                used_tools.add(name)
            elif step.get("status") == "pending":
                active_tools.append(name)
                if not active_args:
                    args = step.get("args") or {}
                    for key in ("query", "ticker", "since", "period"):
                        val = args.get(key)
                        if val:
                            s = str(val).strip()
                            active_args = s[:50] + ("…" if len(s) > 50 else "")
                            break

        elif kind == "tool_result":
            snippet = step.get("snippet", "").strip().replace("\n", " ")
            if snippet:
                line = snippet[:110] + ("…" if len(snippet) > 110 else "")
                latest_result = {"name": step.get("name", ""), "line": line}

        elif kind == "synthesis":
            status = step.get("status", "")
            if status == "done":
                phase = "done"
            elif status == "in_progress" and phase != "done":
                phase = "synthesis"

    # Resolve phase from tool activity if not set by synthesis
    if phase not in ("done", "synthesis"):
        phase = "tool" if active_tools else "thinking"

    return {
        "used_tools": used_tools,
        "active_tools": active_tools,
        "phase": phase,
        "latest_thought": latest_thought,
        "latest_result": latest_result,
        "active_args": active_args,
    }


# ── SVG layout constants ──────────────────────────────────────────────────────

_VW, _VH = 510, 265          # SVG viewBox dimensions

# Orchestrator node (center of composition)
_ORCH_CX, _ORCH_CY = 155, 132
_ORCH_W, _ORCH_H   = 118, 44

# Signals / synthesis pill nodes (vertically aligned with orchestrator)
_SIG_CX, _SIG_CY = 155, 24
_SYN_CX, _SYN_CY = 155, 242
_PILL_W, _PILL_H  = 90, 26

# Tool nodes — right column
_TOOL_X  = 360               # left edge x
_TOOL_W  = 138               # width
_TOOL_H  = 28                # height
_TOOL_ICON_X = _TOOL_X + 16  # emoji center x
_TOOL_LABEL_X = _TOOL_X + 32 # text start x


def _tool_cy(i: int) -> int:
    """Vertical center of tool node i (0-indexed, 5 tools total)."""
    top = 22
    step = (_VH - top * 2) // 4   # evenly spread across height
    return top + i * step          # → 22, 77, 132, 187, 242


# ── SVG primitive helpers ─────────────────────────────────────────────────────

def _rect(x: int, y: int, w: int, h: int, rx: int,
          fill: str, stroke: str, stroke_w: float = 1.5) -> str:
    return (
        f'<rect x="{x}" y="{y}" width="{w}" height="{h}" rx="{rx}" '
        f'fill="{fill}" stroke="{stroke}" stroke-width="{stroke_w}"/>'
    )


def _label(cx: float, cy: float, text: str, fill: str,
           size: float = 11.0, weight: int = 600, dy: float = 0.0) -> str:
    return (
        f'<text x="{cx}" y="{cy + dy}" text-anchor="middle" dominant-baseline="middle" '
        f'font-family="Inter, ui-sans-serif, sans-serif" font-size="{size}" '
        f'font-weight="{weight}" fill="{fill}">{_he(text)}</text>'
    )


def _bezier_edge(x1: float, y1: float, x2: float, y2: float,
                 color: str, width: float, dashed: bool = False,
                 marker_id: str = "") -> str:
    """Cubic bezier: horizontal tangents at both ends (S-curve)."""
    cx = (x1 + x2) / 2
    d = f"M {x1} {y1} C {cx} {y1} {cx} {y2} {x2} {y2}"
    dash = 'stroke-dasharray="5,4"' if dashed else ""
    mend = f'marker-end="url(#pag_{marker_id})"' if marker_id else ""
    return f'<path d="{d}" fill="none" stroke="{color}" stroke-width="{width}" {dash} {mend}/>'


def _vline(x: float, y1: float, y2: float,
           color: str, width: float, marker_id: str = "") -> str:
    mend = f'marker-end="url(#pag_{marker_id})"' if marker_id else ""
    return f'<line x1="{x}" y1="{y1}" x2="{x}" y2="{y2}" stroke="{color}" stroke-width="{width}" {mend}/>'


# ── SVG graph renderer ────────────────────────────────────────────────────────

def _svg(state: dict) -> str:
    used   = state["used_tools"]
    active = set(state["active_tools"])
    phase  = state["phase"]

    is_signals_phase   = phase == "signals"
    is_synthesis_active = phase == "synthesis"
    is_synthesis_done   = phase == "done"

    p: list[str] = []

    # ── Defs: arrowhead markers ───────────────────────────────────────────────
    # IDs prefixed "pag_" to avoid collisions with other SVGs on the page.
    p.append(
        '<defs>'
        # green (active)
        f'<marker id="pag_ag" viewBox="0 0 10 10" refX="9" refY="5" '
        f'markerWidth="5" markerHeight="5" orient="auto-start-reverse">'
        f'<path d="M 0 0 L 10 5 L 0 10 z" fill="{BRAND_GREEN}"/></marker>'
        # gray (inactive)
        f'<marker id="pag_gr" viewBox="0 0 10 10" refX="9" refY="5" '
        f'markerWidth="5" markerHeight="5" orient="auto-start-reverse">'
        f'<path d="M 0 0 L 10 5 L 0 10 z" fill="{BORDER_STRONG}"/></marker>'
        # pale green (used)
        f'<marker id="pag_gp" viewBox="0 0 10 10" refX="9" refY="5" '
        f'markerWidth="5" markerHeight="5" orient="auto-start-reverse">'
        f'<path d="M 0 0 L 10 5 L 0 10 z" fill="{BULL_BORDER}"/></marker>'
        '</defs>'
    )

    # ── Orchestrator → tool edges (drawn first so nodes render on top) ────────
    orch_right = _ORCH_CX + _ORCH_W // 2   # 214
    for i, name in enumerate(_TOOL_NAMES):
        tcy = _tool_cy(i)
        if name in active:
            color, width, dashed, mid = BRAND_GREEN, 2.2, False, "ag"
        elif name in used:
            color, width, dashed, mid = BULL_BORDER, 1.5, False, "gp"
        else:
            color, width, dashed, mid = BORDER, 1.2, True, "gr"
        p.append(_bezier_edge(orch_right, _ORCH_CY, _TOOL_X, tcy,
                              color, width, dashed, mid))

    # ── Signals → orchestrator (vertical) ────────────────────────────────────
    sig_y1 = _SIG_CY + _PILL_H // 2        # bottom of signals pill
    sig_y2 = _ORCH_CY - _ORCH_H // 2       # top of orchestrator rect
    if is_signals_phase:
        s_color, s_w, s_mid = BRAND_GREEN, 2.0, "ag"
    else:
        s_color, s_w, s_mid = BORDER_STRONG, 1.2, "gr"
    p.append(_vline(_SIG_CX, sig_y1, sig_y2, s_color, s_w, s_mid))

    # ── Orchestrator → synthesis (vertical) ──────────────────────────────────
    syn_y1 = _ORCH_CY + _ORCH_H // 2       # bottom of orchestrator rect
    syn_y2 = _SYN_CY - _PILL_H // 2        # top of synthesis pill
    if is_synthesis_active or is_synthesis_done:
        y_color, y_w, y_mid = BRAND_GREEN, 2.0, "ag"
    else:
        y_color, y_w, y_mid = BORDER_STRONG, 1.2, "gr"
    p.append(_vline(_SYN_CX, syn_y1, syn_y2, y_color, y_w, y_mid))

    # ── Signals node ─────────────────────────────────────────────────────────
    if is_signals_phase:
        sf, ss, st_ = BRAND_GREEN, BRAND_GREEN_DARK, "#ffffff"
    else:
        sf, ss, st_ = BG_MUTED, BORDER_STRONG, TEXT_MUTED
    p.append(_rect(_SIG_CX - _PILL_W // 2, _SIG_CY - _PILL_H // 2,
                   _PILL_W, _PILL_H, rx=13, fill=sf, stroke=ss))
    p.append(_label(_SIG_CX, _SIG_CY, "Signals", st_, size=10, weight=500))

    # ── Orchestrator node ─────────────────────────────────────────────────────
    ox = _ORCH_CX - _ORCH_W // 2
    oy = _ORCH_CY - _ORCH_H // 2
    orch_active = not is_signals_phase
    if orch_active:
        of_, os_, ot_, osw = AI_BG, AI_COLOR, AI_COLOR, 2.0
    else:
        of_, os_, ot_, osw = BG_MUTED, BORDER_STRONG, TEXT_MUTED, 1.5
    p.append(_rect(ox, oy, _ORCH_W, _ORCH_H, rx=10, fill=of_, stroke=os_, stroke_w=osw))
    p.append(_label(_ORCH_CX, _ORCH_CY, "Orchestrator", ot_, size=11, weight=700, dy=-7))
    p.append(_label(_ORCH_CX, _ORCH_CY, "agent", ot_, size=9, weight=400, dy=9))

    # ── Synthesis node ────────────────────────────────────────────────────────
    if is_synthesis_active:
        yf, ys, yt_ = BRAND_GREEN, BRAND_GREEN_DARK, "#ffffff"
    elif is_synthesis_done:
        yf, ys, yt_ = BULL_BG, BULL_BORDER, GREEN
    else:
        yf, ys, yt_ = BG_MUTED, BORDER_STRONG, TEXT_MUTED
    p.append(_rect(_SYN_CX - _PILL_W // 2, _SYN_CY - _PILL_H // 2,
                   _PILL_W, _PILL_H, rx=13, fill=yf, stroke=ys))
    p.append(_label(_SYN_CX, _SYN_CY, "Synthesis", yt_, size=10, weight=500))

    # ── Tool nodes ────────────────────────────────────────────────────────────
    for i, (name, (icon, disp_label)) in enumerate(TOOL_DISPLAY.items()):
        tcy = _tool_cy(i)
        ty  = tcy - _TOOL_H // 2

        if name in active:
            tf, ts, tt_, tsw = BRAND_GREEN, BRAND_GREEN_DARK, "#ffffff", 2.0
            fw = 700
        elif name in used:
            tf, ts, tt_, tsw = BULL_BG, BULL_BORDER, GREEN, 1.5
            fw = 500
        else:
            tf, ts, tt_, tsw = BG_MUTED, BORDER, TEXT_FAINT, 1.2
            fw = 500

        p.append(_rect(_TOOL_X, ty, _TOOL_W, _TOOL_H, rx=7,
                       fill=tf, stroke=ts, stroke_w=tsw))
        # Emoji icon
        p.append(
            f'<text x="{_TOOL_ICON_X}" y="{tcy}" '
            f'text-anchor="middle" dominant-baseline="middle" '
            f'font-family="\'Segoe UI Emoji\', \'Apple Color Emoji\', \'Noto Color Emoji\', sans-serif" '
            f'font-size="13">{icon}</text>'
        )
        # Text label
        p.append(
            f'<text x="{_TOOL_LABEL_X}" y="{tcy}" '
            f'text-anchor="start" dominant-baseline="middle" '
            f'font-family="Inter, ui-sans-serif, sans-serif" '
            f'font-size="10.5" font-weight="{fw}" fill="{tt_}">{disp_label}</text>'
        )

    # ── Wrap in a card div ────────────────────────────────────────────────────
    inner = "\n".join(p)
    return (
        f'<div style="background:{BG};border:1px solid {BORDER};border-radius:12px;'
        f'padding:12px 14px;">'
        f'<svg viewBox="0 0 {_VW} {_VH}" width="100%" '
        f'xmlns="http://www.w3.org/2000/svg" style="display:block;">'
        f'{inner}'
        f'</svg>'
        f'</div>'
    )


# ── Essential reasoning card ──────────────────────────────────────────────────

def _essential_card(state: dict) -> str:
    """Build the compact reasoning summary card HTML."""
    phase       = state["phase"]
    active      = state["active_tools"]
    used        = state["used_tools"]
    thought     = state["latest_thought"]
    result      = state["latest_result"]
    active_args = state["active_args"]

    icon, phase_label = _PHASE_META.get(phase, ("🔄", "Running…"))
    rows: list[str] = []

    # ── Phase header ──────────────────────────────────────────────────────────
    rows.append(
        f'<div style="display:flex;align-items:center;gap:6px;margin-bottom:11px;'
        f'padding-bottom:9px;border-bottom:1px solid {BORDER};">'
        f'<span style="font-size:0.95rem;">{icon}</span>'
        f'<span style="font-size:0.82rem;font-weight:700;color:{TEXT};">'
        f'{_he(phase_label)}</span>'
        f'</div>'
    )

    # ── Active tool ───────────────────────────────────────────────────────────
    if active:
        name = active[0]
        t_icon, t_label = TOOL_DISPLAY.get(name, ("🔧", name))
        arg_html = (
            f'<div style="font-size:0.7rem;color:{TEXT_MUTED};margin-top:2px;'
            f'overflow:hidden;text-overflow:ellipsis;white-space:nowrap;">'
            f'{_he(active_args)}</div>'
        ) if active_args else ""
        rows.append(
            f'<div style="margin-bottom:9px;">'
            f'<div style="font-size:0.65rem;font-weight:700;text-transform:uppercase;'
            f'letter-spacing:0.08em;color:{BRAND_GREEN};margin-bottom:3px;">Active tool</div>'
            f'<div style="font-size:0.8rem;color:{TEXT};font-weight:600;'
            f'display:flex;align-items:center;gap:4px;">'
            f'<span>{t_icon}</span><span>{_he(t_label)}</span></div>'
            f'{arg_html}'
            f'</div>'
        )

    # ── Latest thought ────────────────────────────────────────────────────────
    if thought:
        rows.append(
            f'<div style="margin-bottom:9px;">'
            f'<div style="font-size:0.65rem;font-weight:700;text-transform:uppercase;'
            f'letter-spacing:0.08em;color:{BLUE};margin-bottom:3px;">Thinking</div>'
            f'<div style="font-size:0.74rem;color:{TEXT_MUTED};line-height:1.45;'
            f'font-style:italic;">{_he(thought)}</div>'
            f'</div>'
        )

    # ── Latest tool output ────────────────────────────────────────────────────
    if result:
        t_label = TOOL_DISPLAY.get(result["name"], ("", result["name"]))[1]
        rows.append(
            f'<div style="margin-bottom:9px;">'
            f'<div style="font-size:0.65rem;font-weight:700;text-transform:uppercase;'
            f'letter-spacing:0.08em;color:{TEXT_MUTED};margin-bottom:3px;">Latest output</div>'
            f'<div style="font-size:0.73rem;color:{TEXT_MUTED};line-height:1.4;">'
            f'<span style="font-weight:600;color:{TEXT};">{_he(t_label)}</span>'
            f' — {_he(result["line"])}'
            f'</div>'
            f'</div>'
        )

    # ── Footer: call count ────────────────────────────────────────────────────
    n = len(used)
    if n > 0:
        rows.append(
            f'<div style="padding-top:8px;border-top:1px solid {BORDER};">'
            f'<span style="font-size:0.68rem;color:{TEXT_FAINT};">'
            f'{n} tool call{"s" if n != 1 else ""} completed</span>'
            f'</div>'
        )

    body = "\n".join(rows)
    return (
        f'<div style="background:{BG};border:1px solid {BORDER};border-radius:12px;'
        f'padding:14px 15px;min-height:220px;">'
        f'{body}'
        f'</div>'
    )


# ── Public API ────────────────────────────────────────────────────────────────

def render(trace: list[dict]) -> None:
    """Render the live agent graph (left) + essential reasoning card (right).

    Call from _live_trace_fragment() in app.py on every 400ms poll tick.
    """
    state = derive_live_state(trace)
    col_graph, col_card = st.columns([3, 2])
    with col_graph:
        st.markdown(_svg(state), unsafe_allow_html=True)
    with col_card:
        st.markdown(_essential_card(state), unsafe_allow_html=True)