app.py

"""
Moonshiner Trace Explorer
─────────────────────────
Interactive explorer for agentic coding traces published in the
Moonshiner canonical dataset format.

See MOONSHINER_SCHEMA.md for the dataset contract.
"""

import json
import difflib
import tempfile
import os
from collections import Counter, defaultdict

import gradio as gr
import plotly.graph_objects as go
from datasets import load_dataset

# ── Data loading ────────────────────────────────────────────────

DATASET_ID = os.environ.get(
    "DATASET_ID", "greghavens/fable-5-coding-and-debugging-traces"
)

# Derive app title from dataset ID
if "kimi" in DATASET_ID.lower():
    APP_TITLE = "Kimi K3 Trace Explorer"
elif "fable" in DATASET_ID.lower():
    APP_TITLE = "Fable 5 Trace Explorer"
else:
    APP_TITLE = "Trace Explorer"

print(f"Loading dataset {DATASET_ID}…")
ds = load_dataset(DATASET_ID, split="train", download_mode="force_redownload")
print(f"Loaded {len(ds)} rows")

# Group by source_trajectory_id; keep only the final row (max assistant_step).
task_index: dict[str, dict] = {}
for row in ds:
    tid = row.get("source_trajectory_id") or row.get("task", "")
    if tid not in task_index or row["assistant_step"] > task_index[tid]["assistant_step"]:
        task_index[tid] = row

# Normalize language labels
LANG_NORMALIZE = {
    "en": "English",
    "ts": "typescript",
    "js": "javascript",
    "py": "python",
}
for tid in task_index:
    raw = task_index[tid].get("lang") or ""
    task_index[tid]["lang"] = LANG_NORMALIZE.get(raw, raw)

all_tasks_sorted = sorted(task_index.keys())
all_langs = sorted(set(r["lang"] for r in task_index.values() if r.get("lang")))
all_categories = sorted(set(r["category"] for r in task_index.values() if r.get("category")))

lang_tasks: dict[str, list[str]] = defaultdict(list)
cat_tasks: dict[str, list[str]] = defaultdict(list)
for tid, r in task_index.items():
    if r.get("lang"):
        lang_tasks[r["lang"]].append(tid)
    if r.get("category"):
        cat_tasks[r["category"]].append(tid)

# Pre-build search index: trajectory_id → concatenated lowercase text
search_corpus: dict[str, str] = {}
for tid, r in task_index.items():
    parts = []
    for m in r["messages"]:
        content = m.get("content", "")
        if content:
            parts.append(content)
        rc = m.get("reasoning_content", "")
        if rc:
            parts.append(rc)
        for tc in m.get("tool_calls") or []:
            fn = tc.get("function", {}) if isinstance(tc, dict) else {}
            if isinstance(fn, dict):
                parts.append(fn.get("name", ""))
                args_raw = fn.get("arguments", "")
                if args_raw:
                    parts.append(args_raw)
    search_corpus[tid] = "\n".join(parts).lower()

print(f"Indexed {len(task_index)} trajectories, {len(all_langs)} langs, {len(all_categories)} categories")

# ── Language mapping for syntax fences ──────────────────────────

LANG_MAP = {
    "python": "python", "typescript": "typescript", "javascript": "javascript",
    "go": "go", "rust": "rust", "java": "java", "ruby": "ruby",
    "csharp": "csharp", "c": "c", "cpp": "cpp", "bash": "bash",
    "zsh": "bash", "asm": "nasm",
}

# ── Precompute stats & charts ───────────────────────────────────

total_trajectories = len(task_index)
total_rows = len(ds)
total_languages = len(all_langs)
total_categories = len(all_categories)

avg_messages = sum(r["n_messages"] for r in task_index.values()) / max(total_trajectories, 1)
avg_agent_steps = sum(r["assistant_steps"] for r in task_index.values()) / max(total_trajectories, 1)

total_tool_calls = 0
for r in task_index.values():
    for m in r["messages"]:
        if m.get("role") == "tool":
            total_tool_calls += 1
        for tc in m.get("tool_calls") or []:
            total_tool_calls += 1

verification_pass_rate = "100%"

CHART_FONT = dict(color="#e2e8f0")
CHART_LAYOUT = dict(
    paper_bgcolor="rgba(0,0,0,0)",
    plot_bgcolor="rgba(0,0,0,0)",
    font=CHART_FONT,
)

# Chart 1: Trajectories by language
lang_counts = Counter(r["lang"] for r in task_index.values() if r.get("lang"))
lang_sorted = sorted(lang_counts.items(), key=lambda x: -x[1])
fig_lang = go.Figure(go.Bar(
    x=[l for l, _ in lang_sorted],
    y=[c for _, c in lang_sorted],
    marker_color="#10b981",
))
fig_lang.update_layout(title="Trajectories by Language", xaxis_title="Language", yaxis_title="Count", **CHART_LAYOUT)

# Chart 2: Top 20 task categories (horizontal)
cat_counts = Counter(r["category"] for r in task_index.values() if r.get("category"))
cat_sorted = sorted(cat_counts.items(), key=lambda x: x[1])[-20:]
fig_cat = go.Figure(go.Bar(
    y=[c for c, _ in cat_sorted],
    x=[n for _, n in cat_sorted],
    orientation="h",
    marker_color="#3b82f6",
))
fig_cat.update_layout(title="Top 20 Task Categories", xaxis_title="Count", yaxis_title="", height=600, **CHART_LAYOUT)

# Chart 3: Trajectory length distribution
msg_counts = [r["n_messages"] for r in task_index.values()]
fig_msgdist = go.Figure(go.Histogram(x=msg_counts, marker_color="#8b5cf6"))
fig_msgdist.update_layout(title="Trajectory Length Distribution (messages)", xaxis_title="Message Count", yaxis_title="Trajectories", **CHART_LAYOUT)

# Chart 4: Agent decision points per trajectory
step_counts = [r["assistant_steps"] for r in task_index.values()]
fig_steps = go.Figure(go.Histogram(x=step_counts, marker_color="#f59e0b"))
fig_steps.update_layout(title="Agent Decision Points per Trajectory", xaxis_title="Assistant Steps", yaxis_title="Trajectories", **CHART_LAYOUT)

# Chart 5: Avg tool calls per trajectory by language
lang_tool_calls: dict[str, list[int]] = defaultdict(list)
for r in task_index.values():
    tc = 0
    for m in r["messages"]:
        if m.get("role") == "tool":
            tc += 1
        for _ in m.get("tool_calls") or []:
            tc += 1
    lang = r.get("lang")
    if lang:
        lang_tool_calls[lang].append(tc)

avg_tc_by_lang = {l: sum(v) / len(v) for l, v in lang_tool_calls.items()}
avg_tc_sorted = sorted(avg_tc_by_lang.items(), key=lambda x: -x[1])
fig_toolcalls = go.Figure(go.Bar(
    x=[l for l, _ in avg_tc_sorted],
    y=[v for _, v in avg_tc_sorted],
    marker_color="#ef4444",
))
fig_toolcalls.update_layout(title="Avg Tool Calls per Trajectory by Language", xaxis_title="Language", yaxis_title="Avg Tool Calls", **CHART_LAYOUT)

# ── Rendering helpers ───────────────────────────────────────────

TRUNCATE_LIMIT = 80_000

ROLE_COLORS = {
    "user": "#3b82f6",
    "assistant": "#10b981",
    "tool": "#f59e0b",
    "system": "#6b7280",
}


def render_minimap(messages, current_idx):
    n = len(messages)
    bar_w = max(2, min(6, 500 // max(n, 1)))
    total_w = n * (bar_w + 1)
    bars = []
    for i, m in enumerate(messages):
        role = m.get("role", "system")
        color = ROLE_COLORS.get(role, "#6b7280")
        stroke = ' stroke="white" stroke-width="1.5"' if i == current_idx else ""
        onclick = (
            f"onclick=\"(function(){{const s=document.querySelector('#msg-slider input[type=range]');"
            f"if(s){{s.value={i+1};s.dispatchEvent(new Event('input',{{bubbles:true}}));"
            f"s.dispatchEvent(new Event('change',{{bubbles:true}}));}}}})();\""
        )
        bars.append(
            f'<rect x="{i*(bar_w+1)}" y="0" width="{bar_w}" height="32" '
            f'fill="{color}"{stroke} rx="1" style="cursor:pointer" {onclick}/>'
        )
    svg = (
        f'<svg width="{total_w}" height="32" xmlns="http://www.w3.org/2000/svg" '
        f'style="display:block;max-width:100%;overflow-x:auto;">'
        + "".join(bars)
        + "</svg>"
    )
    legend = (
        '<div style="margin-top:6px;font-size:12px;color:#94a3b8;">'
        '<span style="color:#3b82f6;">&#9632;</span> user &nbsp; '
        '<span style="color:#10b981;">&#9632;</span> assistant &nbsp; '
        '<span style="color:#f59e0b;">&#9632;</span> tool &nbsp; '
        '<span style="color:#6b7280;">&#9632;</span> system'
        "</div>"
    )
    return f'<div class="minimap-container">{svg}{legend}</div>'


def render_message(messages, idx, lang="python"):
    if idx < 0 or idx >= len(messages):
        return "*No message at this index.*"
    m = messages[idx]
    role = m.get("role", "unknown")
    color = ROLE_COLORS.get(role, "#6b7280")
    header = f'### <span style="color:{color};">[{role.upper()}]</span> Message {idx+1} / {len(messages)}\n\n---\n\n'

    content = m.get("content", "")
    fence_lang = LANG_MAP.get(lang, "text")

    # ── Tool result messages ────────────────────────────────────
    if role == "tool":
        tool_call_id = m.get("tool_call_id", "")
        tool_name = m.get("name", "")
        # Fall back to looking up name from preceding assistant tool_calls
        if not tool_name and tool_call_id:
            for prev in messages[:idx]:
                if prev.get("role") == "assistant":
                    for tc in prev.get("tool_calls") or []:
                        if isinstance(tc, dict) and tc.get("id") == tool_call_id:
                            fn = tc.get("function", {})
                            tool_name = fn.get("name", "") if isinstance(fn, dict) else ""
                            break
                    if tool_name:
                        break
        name_line = f"**📋 Result from `{tool_name}`**\n\n" if tool_name else ""
        body = _truncate(content)
        return header + name_line + f"```\n{body}\n```"

    # ── Assistant & other roles ─────────────────────────────────
    parts = []

    # Reasoning (canonical field)
    reasoning = m.get("reasoning_content", "")
    if reasoning:
        rc = _truncate(reasoning)
        parts.append(f"<details><summary><b>💭 Reasoning</b></summary>\n\n{rc}\n\n</details>")

    # Visible content
    if content:
        parts.append(_truncate(content))

    # Tool calls
    tool_calls = m.get("tool_calls") or []
    for tc in tool_calls:
        if not isinstance(tc, dict):
            continue
        fn = tc.get("function", {})
        if not isinstance(fn, dict):
            continue
        name = fn.get("name", "unknown")
        args_raw = fn.get("arguments", "{}")
        try:
            args = json.loads(args_raw) if isinstance(args_raw, str) else args_raw
        except (json.JSONDecodeError, TypeError):
            args = {}
        parts.append(_render_tool_call(name, args, fence_lang))

    if not parts:
        return header + "*No content*"
    return header + "\n\n".join(parts)


def _truncate(text):
    if len(text) > TRUNCATE_LIMIT:
        return text[:TRUNCATE_LIMIT] + f"\n\n✂️ *truncated — {len(text):,} chars total*"
    return text


def _render_tool_call(name, args, fence_lang):
    """Render a single tool call with parsed arguments."""
    tc_header = f"**🔧 Tool Call: `{name}`**\n\n"

    if not isinstance(args, dict) or not args:
        return tc_header.rstrip("\n")

    # Edit/str_replace with old_str and new_str → diff
    if "old_str" in args and "new_str" in args:
        old = args["old_str"]
        new = args["new_str"]
        fpath = args.get("file_path", args.get("path", ""))
        path_line = f"**File:** `{fpath}`\n\n" if fpath else ""
        if old == new:
            return tc_header + path_line + "*No changes detected*"
        diff_lines = list(difflib.unified_diff(
            old.splitlines(keepends=True),
            new.splitlines(keepends=True),
            fromfile="before", tofile="after",
        ))
        diff_text = "".join(diff_lines) if diff_lines else "*No changes detected*"
        return tc_header + path_line + f"```diff\n{diff_text}\n```"

    # Bash/shell command
    if "command" in args:
        return tc_header + f"```bash\n{args['command']}\n```"

    # File read
    if "file_path" in args and len(args) <= 2 and "command" not in args and "old_str" not in args:
        return tc_header + f"`{args['file_path']}`"

    # Write/create with file content
    for key in ("file_text", "content", "code"):
        if key in args and isinstance(args[key], str) and len(args[key]) > 50:
            val = _truncate(args[key])
            fpath = args.get("file_path", args.get("path", ""))
            path_line = f"**File:** `{fpath}`\n\n" if fpath else ""
            return tc_header + path_line + f"```{fence_lang}\n{val}\n```"

    # Default: JSON
    args_str = json.dumps(args, indent=2)
    return tc_header + f"```json\n{_truncate(args_str)}\n```"


# ── Filtering logic ────────────────────────────────────────────

def filter_tasks(langs, cats, search_text):
    tasks = set(all_tasks_sorted)

    if langs and "All" not in langs:
        tasks &= {t for l in langs for t in lang_tasks.get(l, [])}
    if cats and "All" not in cats:
        tasks &= {t for c in cats for t in cat_tasks.get(c, [])}
    if search_text and search_text.strip():
        query = search_text.strip().lower()
        tasks = {t for t in tasks if query in search_corpus.get(t, "")}

    return sorted(tasks)


def get_task_stats(task_id):
    if task_id not in task_index:
        return "*No trajectory selected.*"
    r = task_index[task_id]
    msgs = r["messages"]
    role_counts = Counter(m.get("role", "?") for m in msgs)
    lines = [
        f"**Task:** `{r.get('task', task_id)}`",
    ]
    model = r.get("teacher_model")
    if model:
        lines.append(f"**Model:** {model}")
    provider = r.get("provider")
    if provider:
        lines.append(f"**Provider:** {provider}")
    lang = r.get("lang")
    if lang:
        lines.append(f"**Language:** {lang}")
    cat = r.get("category")
    if cat:
        lines.append(f"**Category:** {cat}")
    lines.append(f"**Messages:** {r['n_messages']}")
    for role in ["user", "assistant", "tool", "system"]:
        if role_counts.get(role, 0):
            lines.append(f"  - {role}: {role_counts[role]}")
    lines.append(f"**Agent steps:** {r['assistant_steps']}")
    tools_used = r.get("tools_used")
    if tools_used:
        lines.append(f"**Tools:** {', '.join(tools_used)}")
    return "\n\n".join(lines)


# ── Export logic ────────────────────────────────────────────────

def export_filtered(langs, cats, search_text):
    matching = set(filter_tasks(langs, cats, search_text))
    if not matching:
        return gr.update(visible=False)

    # Map trajectory IDs back to source_trajectory_ids for matching rows
    matching_tids = set()
    for tid in matching:
        r = task_index[tid]
        matching_tids.add(r.get("source_trajectory_id") or r.get("task", tid))

    tmp = tempfile.NamedTemporaryFile(
        mode="w", suffix=".jsonl", delete=False, prefix="trace_export_"
    )
    for row in ds:
        row_tid = row.get("source_trajectory_id") or row.get("task", "")
        if row_tid in matching_tids:
            row_out = {k: v for k, v in row.items() if k != "messages"}
            row_out["messages"] = json.dumps(row["messages"])
            if "tools" in row and row["tools"]:
                row_out["tools"] = row["tools"] if isinstance(row["tools"], str) else json.dumps(row["tools"])
            tmp.write(json.dumps(row_out) + "\n")
    tmp.close()
    return gr.update(value=tmp.name, visible=True)


# ── UI ──────────────────────────────────────────────────────────

CSS = """
.message-viewer { max-height: 70vh; overflow-y: auto; border: 1px solid var(--border-color-primary); border-radius: 8px; padding: 20px; }
.minimap-container { border: 1px solid var(--border-color-primary); border-radius: 8px; padding: 12px 8px; }
footer { display: none !important; }
"""

SUMMARY_MD = f"""
| Metric | Value |
|--------|-------|
| Total trajectories | **{total_trajectories}** |
| Total rows | **{total_rows:,}** |
| Languages | **{total_languages}** |
| Categories | **{total_categories}** |
| Avg messages / trajectory | **{avg_messages:.1f}** |
| Avg agent steps / trajectory | **{avg_agent_steps:.1f}** |
| Total tool calls | **{total_tool_calls:,}** |
| Verification pass rate | **{verification_pass_rate}** |
"""

THEME = gr.themes.Soft(primary_hue="emerald", secondary_hue="blue", neutral_hue="slate")

with gr.Blocks(theme=THEME, css=CSS, title=APP_TITLE) as demo:
    gr.Markdown(f"# 🔍 {APP_TITLE}")

    with gr.Tabs():
        # ── Tab 1: Dataset Overview ─────────────────────────────
        with gr.Tab("📊 Dataset Overview"):
            gr.Markdown(SUMMARY_MD)
            with gr.Row():
                gr.Plot(fig_lang)
                gr.Plot(fig_cat)
            with gr.Row():
                gr.Plot(fig_msgdist)
                gr.Plot(fig_steps)
            with gr.Row():
                gr.Plot(fig_toolcalls)

        # ── Tab 2: Explore Traces ───────────────────────────────
        with gr.Tab("🔍 Explore Traces"):
            with gr.Row():
                # ── Left sidebar ────────────────────────────────
                with gr.Column(scale=1, min_width=240):
                    lang_dd = gr.Dropdown(
                        label="Language",
                        choices=["All"] + all_langs,
                        value=["All"],
                        multiselect=True,
                    )
                    cat_dd = gr.Dropdown(
                        label="Category",
                        choices=["All"] + all_categories,
                        value=["All"],
                        multiselect=True,
                    )
                    search_box = gr.Textbox(
                        label="Search",
                        placeholder="e.g. segfault, backtrack, deadlock…",
                    )
                    match_count = gr.Markdown(f"**{total_trajectories}** trajectories match")
                    task_dd = gr.Dropdown(
                        label="Trajectory",
                        choices=all_tasks_sorted,
                        value=all_tasks_sorted[0] if all_tasks_sorted else None,
                    )
                    stats_md = gr.Markdown(
                        get_task_stats(all_tasks_sorted[0]) if all_tasks_sorted else ""
                    )
                    export_btn = gr.Button("📦 Export filtered subset (.jsonl)")
                    download_file = gr.File(visible=False)

                # ── Main area ───────────────────────────────────
                with gr.Column(scale=3):
                    minimap = gr.HTML(
                        render_minimap(
                            task_index[all_tasks_sorted[0]]["messages"], 0
                        ) if all_tasks_sorted else ""
                    )
                    with gr.Row():
                        prev_agent_btn = gr.Button("⏮ Prev Agent", variant="secondary")
                        prev_btn = gr.Button("◀ Prev")
                        msg_slider = gr.Slider(
                            minimum=1,
                            maximum=max(
                                task_index[all_tasks_sorted[0]]["n_messages"] if all_tasks_sorted else 1, 1
                            ),
                            value=1,
                            step=1,
                            label="Message",
                            scale=4,
                            elem_id="msg-slider",
                        )
                        next_btn = gr.Button("Next ▶")
                        next_agent_btn = gr.Button("Next Agent ⏭", variant="secondary")
                    message_display = gr.Markdown(
                        render_message(
                            task_index[all_tasks_sorted[0]]["messages"], 0,
                            task_index[all_tasks_sorted[0]].get("lang", ""),
                        ) if all_tasks_sorted else "",
                        elem_classes=["message-viewer"],
                    )

    # ── Event handlers ──────────────────────────────────────────

    def on_filter(langs, cats, search_text):
        tasks = filter_tasks(langs, cats, search_text)
        count_text = f"**{len(tasks)}** trajectories match"
        if not tasks:
            return (
                gr.update(choices=[], value=None),
                count_text,
                "*No matching trajectories.*",
                "",
                gr.update(minimum=1, maximum=1, value=1),
                "",
            )
        first = tasks[0]
        r = task_index[first]
        msgs = r["messages"]
        return (
            gr.update(choices=tasks, value=first),
            count_text,
            get_task_stats(first),
            render_minimap(msgs, 0),
            gr.update(minimum=1, maximum=len(msgs), value=1),
            render_message(msgs, 0, r.get("lang", "")),
        )

    def on_task_select(task_id):
        if not task_id or task_id not in task_index:
            return "", gr.update(minimum=1, maximum=1, value=1), "", get_task_stats(task_id or "")
        r = task_index[task_id]
        msgs = r["messages"]
        return (
            get_task_stats(task_id),
            render_minimap(msgs, 0),
            gr.update(minimum=1, maximum=len(msgs), value=1),
            render_message(msgs, 0, r.get("lang", "")),
        )

    def on_slider(step, task_id):
        if not task_id or task_id not in task_index:
            return "", ""
        r = task_index[task_id]
        msgs = r["messages"]
        idx = int(step) - 1
        return render_minimap(msgs, idx), render_message(msgs, idx, r.get("lang", ""))

    def on_prev(step, task_id):
        new_step = max(1, int(step) - 1)
        mm, md = on_slider(new_step, task_id)
        return mm, gr.update(value=new_step), md

    def on_next(step, task_id):
        if not task_id or task_id not in task_index:
            return "", gr.update(), ""
        max_step = len(task_index[task_id]["messages"])
        new_step = min(max_step, int(step) + 1)
        mm, md = on_slider(new_step, task_id)
        return mm, gr.update(value=new_step), md

    def on_prev_agent(step, task_id):
        if not task_id or task_id not in task_index:
            return "", gr.update(), ""
        msgs = task_index[task_id]["messages"]
        idx = int(step) - 2
        while idx >= 0:
            if msgs[idx].get("role") == "assistant":
                break
            idx -= 1
        if idx < 0:
            idx = 0
        new_step = idx + 1
        mm, md = on_slider(new_step, task_id)
        return mm, gr.update(value=new_step), md

    def on_next_agent(step, task_id):
        if not task_id or task_id not in task_index:
            return "", gr.update(), ""
        msgs = task_index[task_id]["messages"]
        idx = int(step)
        while idx < len(msgs):
            if msgs[idx].get("role") == "assistant":
                break
            idx += 1
        if idx >= len(msgs):
            idx = len(msgs) - 1
        new_step = idx + 1
        mm, md = on_slider(new_step, task_id)
        return mm, gr.update(value=new_step), md

    # Wire filter events
    filter_outputs = [task_dd, match_count, stats_md, minimap, msg_slider, message_display]
    lang_dd.change(on_filter, [lang_dd, cat_dd, search_box], filter_outputs)
    cat_dd.change(on_filter, [lang_dd, cat_dd, search_box], filter_outputs)
    search_box.submit(on_filter, [lang_dd, cat_dd, search_box], filter_outputs)

    # Wire task selection
    task_dd.change(on_task_select, [task_dd], [stats_md, minimap, msg_slider, message_display])

    # Wire slider
    msg_slider.change(on_slider, [msg_slider, task_dd], [minimap, message_display])

    # Wire nav buttons
    step_outputs = [minimap, msg_slider, message_display]
    prev_btn.click(on_prev, [msg_slider, task_dd], step_outputs)
    next_btn.click(on_next, [msg_slider, task_dd], step_outputs)
    prev_agent_btn.click(on_prev_agent, [msg_slider, task_dd], step_outputs)
    next_agent_btn.click(on_next_agent, [msg_slider, task_dd], step_outputs)

    # Wire export
    export_btn.click(export_filtered, [lang_dd, cat_dd, search_box], [download_file])

if __name__ == "__main__":
    demo.launch()