feat: redesign reasoning trace — paired tool cards, color-coded steps, collapsible results

dashboard/reasoning.py

@@ -17,21 +17,27 @@ from typing import Any
 
 import streamlit as st
 
-from dashboard.theme import BG, BG_MUTED, BORDER, GREEN, RED, TEXT, TEXT_MUTED, TEXT_FAINT, AMBER, BLUE
+from dashboard.theme import (
+    BG, BG_MUTED, BORDER, BORDER_STRONG,
+    GREEN, RED, AMBER, BLUE,
+    TEXT, TEXT_MUTED,
+    INFO_BG, INFO_BORDER,
+    WARN_BG, WARN_BORDER,
+)
+from dashboard.components import source_badge
 
 
 SNIPPET_LEN = 200
 
 
-# ---------------------------------------------------------------------------
+# ──────────────────────────────────────────────────────────────────────────────
 # Trace builder helpers — called from app.py during graph.stream
-# ---------------------------------------------------------------------------
+# ──────────────────────────────────────────────────────────────────────────────
 
 def _stringify(content: Any) -> str:
     if isinstance(content, str):
         return content
     if isinstance(content, list):
-        # Anthropic content blocks: [{"type": "text", "text": "..."}]
         parts = []
         for block in content:
             if isinstance(block, dict) and block.get("type") == "text":
@@ -74,7 +80,6 @@ def absorb_tool_messages(trace: list[dict], messages: list) -> None:
             "snippet": snippet,
             "full": full,
         })
-        # Mark the matching pending call as done
         for step in trace:
             if step["kind"] == "tool_call" and step["id"] == tcid:
                 step["status"] = "done"
@@ -90,99 +95,204 @@ def mark_synthesis(trace: list[dict], status: str, error: str | None = None) ->
     trace.append({"kind": "synthesis", "status": status, "error": error})
 
 
-# ---------------------------------------------------------------------------
-# Renderers
-# ---------------------------------------------------------------------------
+# ──────────────────────────────────────────────────────────────────────────────
+# Pre-rendering: pair tool_call + tool_result into unified "tool_round" steps
+# ──────────────────────────────────────────────────────────────────────────────
 
-def _fmt_args(args: dict) -> str:
+def _pair_steps(trace: list[dict]) -> list[dict]:
+    """Walk the flat trace and merge each tool_call with its matching tool_result."""
+    results_by_id: dict[str, dict] = {}
+    for step in trace:
+        if step["kind"] == "tool_result":
+            results_by_id[step["tool_call_id"]] = step
+
+    paired: list[dict] = []
+    for step in trace:
+        kind = step["kind"]
+        if kind == "assistant_text":
+            paired.append(step)
+        elif kind == "tool_call":
+            paired.append({
+                "kind": "tool_round",
+                "call": step,
+                "result": results_by_id.get(step["id"]),
+            })
+        elif kind == "tool_result":
+            pass  # consumed above
+        elif kind == "synthesis":
+            paired.append(step)
+    return paired
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Helpers
+# ──────────────────────────────────────────────────────────────────────────────
+
+_TOOL_CATEGORIES: dict[str, str] = {
+    "search_filing": "filing",
+    "get_financial_metrics": "filing",
+    "get_analyst": "filing",
+    "search_transcript": "transcript",
+    "get_news": "news",
+    "search_news": "news",
+}
+
+
+def _tool_category(name: str) -> str:
+    for key, cat in _TOOL_CATEGORIES.items():
+        if key in name:
+            return cat
+    return "tool"
+
+
+def _fmt_args(args: dict, max_len: int = 120) -> str:
     if not args:
         return ""
     try:
-        return json.dumps(args, ensure_ascii=False, default=str)
+        s = json.dumps(args, ensure_ascii=False, separators=(", ", ": "), default=str).strip("{}")
     except Exception:
-        return str(args)
+        s = str(args)
+    return s if len(s) <= max_len else s[: max_len - 1] + "…"
 
 
-def _render_step(step: dict, idx: int) -> None:
-    kind = step["kind"]
+def _details_block(full: str) -> str:
+    """HTML <details> collapsible for full tool result text. Empty if content is short."""
+    if not full or len(full.strip()) <= SNIPPET_LEN:
+        return ""
+    escaped = full.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
+    return (
+        f'<details style="margin-top:7px;">'
+        f'<summary style="cursor:pointer;font-size:0.72rem;color:{TEXT_MUTED};font-weight:500;'
+        f'list-style:none;display:inline-flex;align-items:center;gap:4px;user-select:none;">'
+        f'<span style="font-size:0.58rem;">▶</span> Show full result'
+        f'</summary>'
+        f'<div style="margin-top:6px;padding:8px 10px;background:{BG_MUTED};'
+        f'border:1px solid {BORDER};border-radius:6px;'
+        f'font-family:ui-monospace,SFMono-Regular,Menlo,monospace;'
+        f'font-size:0.72rem;line-height:1.5;color:{TEXT_MUTED};'
+        f'white-space:pre-wrap;word-break:break-all;'
+        f'max-height:300px;overflow-y:auto;">{escaped}</div>'
+        f'</details>'
+    )
 
-    if kind == "assistant_text":
-        st.markdown(
-            f'<div style="border-left:3px solid {BLUE};background:{BG_MUTED};'
-            f'padding:8px 12px;margin:4px 0 8px;border-radius:0 6px 6px 0;">'
-            f'<div style="font-size:0.6rem;font-weight:700;text-transform:uppercase;'
-            f'letter-spacing:0.08em;color:{TEXT_FAINT};margin-bottom:4px;">💭 Thinking</div>'
-            f'<div style="font-size:0.82rem;line-height:1.5;color:{TEXT};white-space:pre-wrap;">'
-            f'{step["text"]}</div></div>',
-            unsafe_allow_html=True,
-        )
 
-    elif kind == "tool_call":
-        icon = "✅" if step["status"] == "done" else "🔄"
-        args_str = _fmt_args(step["args"])
-        args_html = (
-            f'<div style="font-size:0.72rem;color:{TEXT_MUTED};margin-top:3px;'
-            f'font-family:ui-monospace,SFMono-Regular,Menlo,monospace;word-break:break-all;">'
-            f'{args_str}</div>'
-            if args_str else ""
-        )
-        st.markdown(
-            f'<div style="border:1px solid {BORDER};background:{BG};'
-            f'padding:8px 12px;margin:4px 0;border-radius:6px;">'
-            f'<div style="font-size:0.82rem;color:{TEXT};">'
-            f'{icon} <strong>tool</strong> · '
-            f'<code style="background:{BG_MUTED};padding:1px 6px;border-radius:4px;'
-            f'font-size:0.78rem;">{step["name"]}</code></div>'
-            f'{args_html}</div>',
-            unsafe_allow_html=True,
-        )
+# ──────────────────────────────────────────────────────────────────────────────
+# Step renderers
+# ──────────────────────────────────────────────────────────────────────────────
+
+def _render_assistant_text(step: dict) -> None:
+    st.markdown(
+        f'<div style="background:{INFO_BG};border:1px solid {INFO_BORDER};'
+        f'border-left:4px solid {BLUE};border-radius:0 10px 10px 0;'
+        f'padding:10px 14px;margin:6px 0;">'
+        f'<div style="font-size:0.6rem;font-weight:700;text-transform:uppercase;'
+        f'letter-spacing:0.09em;color:{BLUE};margin-bottom:5px;">💭 Thinking</div>'
+        f'<div style="font-size:0.82rem;line-height:1.55;color:{TEXT};white-space:pre-wrap;">'
+        f'{step["text"]}</div>'
+        f'</div>',
+        unsafe_allow_html=True,
+    )
+
+
+def _render_tool_round(step: dict) -> None:
+    call = step["call"]
+    result = step["result"]
+    pending = result is None
+
+    if pending:
+        bg, border_color, accent, icon = WARN_BG, WARN_BORDER, AMBER, "🔄"
+    else:
+        bg, border_color, accent, icon = BG_MUTED, BORDER, GREEN, "✅"
 
-    elif kind == "tool_result":
-        snippet = step["snippet"] or "(empty)"
-        st.markdown(
-            f'<div style="border-left:3px solid {GREEN};background:{BG};'
-            f'padding:6px 12px;margin:0 0 8px 16px;border-radius:0 6px 6px 0;">'
+    name = call["name"]
+    badge_html = source_badge(_tool_category(name))
+    args_str = _fmt_args(call.get("args") or {})
+
+    args_html = (
+        f'<div style="margin-top:6px;padding:4px 9px;background:{BG};'
+        f'border:1px solid {border_color};border-radius:5px;'
+        f'font-family:ui-monospace,SFMono-Regular,Menlo,monospace;'
+        f'font-size:0.71rem;color:{TEXT_MUTED};word-break:break-all;">'
+        f'{args_str}</div>'
+    ) if args_str else ""
+
+    if result:
+        snippet = result.get("snippet") or "(empty)"
+        details = _details_block(result.get("full") or "")
+        result_html = (
+            f'<div style="margin-top:9px;padding-top:8px;border-top:1px solid {border_color};">'
             f'<div style="font-size:0.6rem;font-weight:700;text-transform:uppercase;'
-            f'letter-spacing:0.08em;color:{TEXT_FAINT};margin-bottom:3px;">↩ result</div>'
-            f'<div style="font-size:0.78rem;color:{TEXT_MUTED};line-height:1.45;">'
-            f'{snippet}</div></div>',
-            unsafe_allow_html=True,
+            f'letter-spacing:0.08em;color:{GREEN};margin-bottom:4px;">↩ Result</div>'
+            f'<div style="font-size:0.78rem;color:{TEXT_MUTED};line-height:1.45;">{snippet}</div>'
+            f'{details}'
+            f'</div>'
         )
-        full = step.get("full") or ""
-        if full and len(full) > len(step["snippet"]):
-            st.code(full, language="text")
-
-    elif kind == "synthesis":
-        if step.get("error"):
-            icon, label, color = "❌", f"Synthesis failed: {step['error']}", RED
-        elif step["status"] == "in_progress":
-            icon, label, color = "🔄", "Synthesizing brief…", AMBER
-        else:
-            icon, label, color = "✅", "Brief synthesized", GREEN
-        st.markdown(
-            f'<div style="border:1px solid {BORDER};background:{BG_MUTED};'
-            f'padding:8px 12px;margin:8px 0;border-radius:6px;'
-            f'border-left:3px solid {color};">'
-            f'<div style="font-size:0.85rem;font-weight:600;color:{TEXT};">'
-            f'{icon} {label}</div></div>',
-            unsafe_allow_html=True,
+    else:
+        result_html = (
+            f'<div style="margin-top:8px;padding-top:8px;border-top:1px solid {border_color};'
+            f'font-size:0.75rem;color:{AMBER};font-style:italic;">Waiting for result…</div>'
         )
 
+    st.markdown(
+        f'<div style="background:{bg};border:1px solid {border_color};'
+        f'border-left:4px solid {accent};border-radius:0 10px 10px 0;'
+        f'padding:10px 14px;margin:6px 0;">'
+        f'<div style="display:flex;align-items:center;justify-content:space-between;gap:8px;">'
+        f'<div style="font-size:0.82rem;color:{TEXT};display:flex;align-items:center;gap:7px;">'
+        f'{icon}'
+        f'<code style="background:{BG};border:1px solid {BORDER_STRONG};'
+        f'padding:1px 7px;border-radius:5px;font-size:0.77rem;">{name}</code>'
+        f'</div>'
+        f'{badge_html}'
+        f'</div>'
+        f'{args_html}'
+        f'{result_html}'
+        f'</div>',
+        unsafe_allow_html=True,
+    )
+
+
+def _render_synthesis(step: dict) -> None:
+    if step.get("error"):
+        icon, label, color = "❌", f"Synthesis failed: {step['error']}", RED
+    elif step["status"] == "in_progress":
+        icon, label, color = "🔄", "Synthesizing brief…", AMBER
+    else:
+        icon, label, color = "✅", "Brief synthesized", GREEN
+    st.markdown(
+        f'<div style="background:{BG_MUTED};border:1px solid {BORDER};'
+        f'border-left:4px solid {color};border-radius:0 10px 10px 0;'
+        f'padding:10px 14px;margin:6px 0;">'
+        f'<div style="font-size:0.85rem;font-weight:600;color:{TEXT};">'
+        f'{icon} {label}</div></div>',
+        unsafe_allow_html=True,
+    )
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Public API
+# ──────────────────────────────────────────────────────────────────────────────
 
 def render_trace_body(trace: list[dict]) -> None:
     """Render every step of the trace, top to bottom."""
     if not trace:
         st.caption("No reasoning steps yet.")
         return
-    for i, step in enumerate(trace):
-        _render_step(step, i)
+    for step in _pair_steps(trace):
+        kind = step["kind"]
+        if kind == "assistant_text":
+            _render_assistant_text(step)
+        elif kind == "tool_round":
+            _render_tool_round(step)
+        elif kind == "synthesis":
+            _render_synthesis(step)
 
 
 def render(trace: list[dict]) -> None:
     """Render the reasoning expander on the Verdict tab (collapsed by default)."""
     if not trace:
         return
-    n_calls = sum(1 for s in trace if s["kind"] == "tool_call")
-    label = f"🧠 Reasoning trace — {n_calls} tool call{'s' if n_calls != 1 else ''}, {len(trace)} steps"
+    n_rounds = sum(1 for s in trace if s["kind"] == "tool_call")
+    label = f"🧠 Reasoning trace — {n_rounds} tool round{'s' if n_rounds != 1 else ''}, {len(trace)} steps"
     with st.expander(label, expanded=False):
         render_trace_body(trace)