""" 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'' ) svg = ( f'' + "".join(bars) + "" ) legend = ( '
' '■ user   ' '■ assistant   ' '■ tool   ' '■ system' "
" ) return f'
{svg}{legend}
' 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'### [{role.upper()}] 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"
💭 Reasoning\n\n{rc}\n\n
") # 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()