viewer.py

"""Streamlit viewer for StarShell eval runs."""

import json
import re
import xml.etree.ElementTree as ET
from pathlib import Path

import pandas as pd
import streamlit as st

from starshell.eval import compute_cost

RESULTS_DIR = Path("results")


# ── Custom styling ──────────────────────────────────────────────────────────

def _inject_css() -> None:
    st.markdown("""
    <style>
    /* Tighten top padding */
    .block-container { padding-top: 2rem; }

    /* Sidebar width */
    section[data-testid="stSidebar"] { width: 280px !important; }
    section[data-testid="stSidebar"] > div { width: 280px !important; }

    /* Sidebar styling */
    section[data-testid="stSidebar"] .block-container { padding-top: 1rem; }

    /* Metric cards */
    div[data-testid="stMetric"] {
        background: #f8f9fa;
        border: 1px solid #e9ecef;
        border-radius: 0.5rem;
        padding: 0.4rem 0.6rem;
    }
    div[data-testid="stMetric"]:has(div[data-testid="stMetricValue"]) {
        text-align: center;
    }
    div[data-testid="stMetric"] label[data-testid="stMetricLabel"] p {
        font-size: 0.75rem;
    }
    div[data-testid="stMetric"] div[data-testid="stMetricValue"] {
        font-size: 1.1rem;
    }

    /* Score badge colors */
    .score-pass { color: #198754; font-weight: 600; }
    .score-fail { color: #dc3545; font-weight: 600; }
    .score-skip { color: #6c757d; font-weight: 600; }

    /* Tool call expanders */
    details[data-testid="stExpander"] summary span[data-testid="stMarkdownContainer"] p {
        font-size: 0.9rem;
    }
    </style>
    """, unsafe_allow_html=True)


# ── Helpers ─────────────────────────────────────────────────────────────────

def load_run_dirs() -> list[Path]:
    """Return run directories sorted newest first."""
    if not RESULTS_DIR.exists():
        return []
    dirs = [d for d in RESULTS_DIR.iterdir() if d.is_dir() and (d / "summary.json").exists()]
    dirs.sort(key=lambda d: d.name, reverse=True)
    return dirs


def load_json(path: Path) -> dict | list | None:
    if not path.exists():
        return None
    with open(path) as f:
        return json.load(f)


_FC_PATTERN = re.compile(r"<function_calls>.*?</function_calls>", re.DOTALL)


def render_assistant_text(text: str) -> None:
    """Render assistant text, extracting inline <function_calls> XML as expanders."""
    parts = _FC_PATTERN.split(text)
    calls = _FC_PATTERN.findall(text)

    for idx, part in enumerate(parts):
        part = part.strip()
        if part:
            st.markdown(part)
        # Render the function_calls block that followed this text part
        if idx < len(calls):
            _render_function_calls_xml(calls[idx])


def _render_function_calls_xml(xml_str: str) -> None:
    """Parse a <function_calls> XML block and render each invoke as an expander."""
    try:
        root = ET.fromstring(xml_str)
    except ET.ParseError:
        st.code(xml_str, language="xml")
        return

    for invoke in root.findall("invoke"):
        tool_name = invoke.get("name", "tool")
        params = {}
        for param in invoke.findall("parameter"):
            key = param.get("name", "?")
            value = param.text or ""
            params[key] = value

        with st.expander(f"Tool call: **{tool_name}**"):
            # Try to pretty-print JSON values within params
            for key, value in params.items():
                try:
                    parsed = json.loads(value)
                    st.markdown(f"**{key}:**")
                    st.code(json.dumps(parsed, indent=2), language="json")
                except (json.JSONDecodeError, TypeError):
                    st.markdown(f"**{key}:** `{value}`")



# ── Page config ──────────────────────────────────────────────────────────────

st.set_page_config(page_title="StarShell Viewer", page_icon="public/favicon.png", layout="wide")
_inject_css()

import base64 as _b64
_favicon_b64 = _b64.b64encode(Path("public/favicon.png").read_bytes()).decode()
st.markdown(
    f'<h1 style="display:flex !important;align-items:center !important;">'
    f'<img src="data:image/png;base64,{_favicon_b64}" width="40" height="40" style="margin-bottom:10px;margin-right:8px;">'
    f'StarShell Viewer</h1>',
    unsafe_allow_html=True,
)

# ── Sidebar: filters + run selector + metadata ───────────────────────────────

run_dirs = load_run_dirs()
if not run_dirs:
    st.warning("No runs found in `results/`. Run an eval first.")
    st.stop()

# logo_light.png = dark logo (for light bg), logo_dark.png = light logo (for dark bg)
_sidebar_logo = "public/logo_dark.png" if st.context.theme.get("type") == "dark" else "public/logo_light.png"
st.logo(_sidebar_logo)

# Pre-load all summaries for filtering
_ALL = "All"
_summaries = {}
_agents, _models, _domains = set(), set(), set()
for d in run_dirs:
    s = load_json(d / "summary.json")
    if s:
        _summaries[d.name] = s
        if s.get("agent_type"):
            _agents.add(s["agent_type"])
        if s.get("model_id"):
            _models.add(s["model_id"])
        if s.get("domain"):
            _domains.add(s["domain"])

filter_agent = st.sidebar.selectbox("Agent", [_ALL] + sorted(_agents))
filter_model = st.sidebar.selectbox("Model", [_ALL] + sorted(_models))
filter_domain = st.sidebar.selectbox("Benchmark", [_ALL] + sorted(_domains))

# Apply filters
filtered_dirs = []
for d in run_dirs:
    s = _summaries.get(d.name, {})
    if filter_agent != _ALL and s.get("agent_type") != filter_agent:
        continue
    if filter_model != _ALL and s.get("model_id") != filter_model:
        continue
    if filter_domain != _ALL and s.get("domain") != filter_domain:
        continue
    filtered_dirs.append(d)

if not filtered_dirs:
    st.sidebar.warning("No runs match the selected filters.")
    st.stop()

st.sidebar.divider()
selected_name = st.sidebar.selectbox("Run", [d.name for d in filtered_dirs])
run_path = RESULTS_DIR / selected_name
summary = load_json(run_path / "summary.json")

if summary is None:
    st.error(f"Could not load summary.json for {selected_name}")
    st.stop()

# Sidebar metadata
st.sidebar.divider()
accuracy = summary.get("accuracy", 0)
total = summary.get("total", 0)
passed = summary.get("passed", 0)
skipped = summary.get("skipped", 0)

evaluated = [r for r in summary.get("results", []) if not r.get("skipped")]
avg_tool_calls = sum(r.get("num_tool_calls", 0) for r in evaluated) / len(evaluated) if evaluated else 0
avg_time = sum(r.get("agent_time_s", 0) for r in evaluated) / len(evaluated) if evaluated else 0

# Token usage → cost estimate via litellm pricing
# Strip "responses/" prefix so litellm can look up pricing (e.g. azure/responses/gpt-5.4 → azure/gpt-5.4)
_model_id = summary.get("model_id", "").replace("/responses/", "/")
_has_usage = any(r.get("usage") for r in evaluated)
total_cost = sum(compute_cost(r.get("usage", {}), _model_id) for r in evaluated) if _has_usage else None
avg_cost = total_cost / len(evaluated) if total_cost is not None and evaluated else None

col1, col2 = st.sidebar.columns(2)
col1.metric("Success Rate", f"{accuracy:.1%}")
col2.metric("Avg Steps", f"{avg_tool_calls:.1f}")
col3, col4 = st.sidebar.columns(2)
col3.metric("Avg Time", f"{avg_time:.0f}s")
col4.metric("Avg Cost", f"${avg_cost:.3f}" if avg_cost is not None else "—")

# Check for reasoning_effort from first available task.json
_reasoning_effort = None
for _task_dir in run_path.iterdir():
    _tj = _task_dir / "task.json"
    if _tj.is_file():
        _task_meta = load_json(_tj)
        if _task_meta and "reasoning_effort" in _task_meta:
            _reasoning_effort = _task_meta["reasoning_effort"]
        break

_caption = (
    f"**Run:** `{selected_name}`  \n"
    f"**Agent:** {summary.get('agent_type', '?')}  \n"
    f"**Model:** `{summary.get('model_id', '?')}`  \n"
    f"**Domain:** {summary.get('domain', '?')}  \n"
)
if _reasoning_effort:
    _caption += f"**Reasoning Effort:** {_reasoning_effort}  \n"
_caption += f"**Total:** {total}  |  **Passed:** {passed}  |  **Skipped:** {skipped}"
st.sidebar.caption(_caption)

# ── Main area: summary table ─────────────────────────────────────────────────

st.header("Run Summary")

results = summary.get("results", [])
if not results:
    st.info("No task results in this run.")
    st.stop()

df = pd.DataFrame(results)[["task_id", "goal", "score", "agent_time_s", "num_tool_calls"]]
df["agent_time_s"] = df["agent_time_s"].round(2)
df = df.rename(columns={
    "task_id": "Task",
    "goal": "Goal",
    "score": "Score",
    "agent_time_s": "Time (s)",
    "num_tool_calls": "Tool calls",
})


def _color_row(row: pd.Series) -> list[str]:
    if row["Score"] == 1:
        bg = "background-color: rgba(25, 135, 84, 0.1)"
    elif row["Score"] == 0:
        bg = "background-color: rgba(220, 53, 69, 0.08)"
    else:
        bg = "background-color: rgba(108, 117, 125, 0.08)"
    return [bg] * len(row)



styled = (
    df.style
    .apply(_color_row, axis=1)
    .format({"Time (s)": "{:.2f}"})
)
event = st.dataframe(
    styled,
    width="stretch",
    hide_index=True,
    height=min(len(df) * 40 + 40, 600),
    on_select="rerun",
    selection_mode="single-row",
)

# ── Trace Inspector ──────────────────────────────────────────────────────────

st.divider()
st.header("Trace Inspector")

task_ids = [r["task_id"] for r in results]

# Selection from dataframe click, default to first task
selected_rows = event.selection.rows if event.selection else []
selected_idx = selected_rows[0] if selected_rows else 0
selected_task = task_ids[selected_idx]

st.caption(f"Showing: **{selected_task}**")

# Derive the task subdirectory name from the task_id
task_dir = run_path / selected_task
if not task_dir.is_dir():
    # Fallback for old runs using hyphenated dir names
    task_dir = run_path / selected_task.replace(".", "-").replace("__", "-")
if not task_dir.is_dir():
    task_dir = None

if task_dir is None:
    st.warning(f"Could not find task directory for `{selected_task}`.")
    st.stop()

# ── Outcome panel ────────────────────────────────────────────────────────────

outcome = load_json(task_dir / "outcome.json")
trace = load_json(task_dir / "trace.json")

if outcome:
    num_tool_calls = outcome.get("num_tool_calls")
    if num_tool_calls is None:
        # Count from trace: separate function_call entries (bash) or inline XML (mcp)
        num_tool_calls = sum(1 for e in (trace or []) if e.get("type") == "function_call")
        if num_tool_calls == 0:
            # Count inline <invoke> tags in assistant messages (MCP traces)
            for e in (trace or []):
                if e.get("role") == "assistant" and isinstance(e.get("content"), list):
                    for block in e["content"]:
                        text = block.get("text", "")
                        num_tool_calls += len(re.findall(r"<invoke\s", text))

    score = outcome.get("score", "?")
    cols = st.columns(3)
    cols[0].metric("Score", score)
    cols[1].metric("Time (s)", f"{outcome.get('agent_time_s', 0):.2f}")
    cols[2].metric("Tool calls", num_tool_calls)

    # Console links
    agent_url = outcome.get("agent_console_url")
    env_url = outcome.get("env_console_url")
    if agent_url or env_url:
        links = []
        if agent_url:
            links.append(f"[Agent job]({agent_url})")
        if env_url:
            links.append(f"[Env job]({env_url})")
        st.caption("Console: " + " · ".join(links))


# ── Trace panel ──────────────────────────────────────────────────────────────
if trace is None:
    st.info("No trace.json found for this task.")
    st.stop()

st.subheader("Trace")

# Detect MAS traces by checking for _mas_phase tags
_is_mas = any(e.get("_mas_phase") for e in trace)


def _render_trace_entries(entries: list) -> None:
    """Render a list of trace entries as chat messages and tool calls."""
    i = 0
    while i < len(entries):
        entry = entries[i]
        role = entry.get("role")
        entry_type = entry.get("type")

        # User message
        if role == "user" and isinstance(entry.get("content"), str):
            with st.chat_message("user"):
                st.markdown(entry["content"])

        # Assistant message
        elif role == "assistant" and isinstance(entry.get("content"), list):
            texts = []
            for block in entry["content"]:
                if block.get("type") == "output_text" and block.get("text", "").strip():
                    texts.append(block["text"])
            if texts:
                full_text = "\n\n".join(texts)
                if "<function_calls>" in full_text:
                    with st.chat_message("assistant"):
                        render_assistant_text(full_text)
                else:
                    with st.chat_message("assistant"):
                        st.markdown(full_text)

        # Function call (tool use)
        elif entry_type == "function_call":
            tool_name = entry.get("name", "tool")
            args_raw = entry.get("arguments", "")

            # Look ahead for the matching function_call_output
            output_text = None
            if i + 1 < len(entries) and entries[i + 1].get("type") == "function_call_output":
                output_text = entries[i + 1].get("output", "")
                i += 1  # skip the output entry in the main loop

            with st.expander(f"Tool call: **{tool_name}**"):
                # Pretty-print arguments if valid JSON
                try:
                    args_obj = json.loads(args_raw)
                    st.code(json.dumps(args_obj, indent=2), language="json")
                except (json.JSONDecodeError, TypeError):
                    st.code(args_raw)

                if output_text is not None:
                    st.markdown("**Output:**")
                    truncated = len(output_text) > 2000
                    display_text = output_text[:2000] if truncated else output_text
                    st.code(display_text)
                    if truncated:
                        with st.expander("Show full output"):
                            st.code(output_text)

        # Standalone function_call_output (not preceded by a function_call)
        elif entry_type == "function_call_output":
            with st.expander("Tool output"):
                output_text = entry.get("output", "")
                truncated = len(output_text) > 2000
                display_text = output_text[:2000] if truncated else output_text
                st.code(display_text)
                if truncated:
                    with st.expander("Show full output"):
                        st.code(output_text)

        i += 1


if _is_mas:
    # Split trace into planner and executor entries (skip separator entries)
    planner_entries = [e for e in trace if e.get("_mas_phase") == "planner"]
    executor_entries = [e for e in trace if e.get("_mas_phase") == "executor"]

    usage = (outcome or {}).get("usage", {})
    planner_prompt = usage.get("planner", {}).get("system_prompt")
    executor_prompt = usage.get("executor", {}).get("system_prompt")

    # ── Phase 1: Planner ──
    st.markdown("### Phase 1: Planner")
    planner_stats = usage.get("planner", {})
    if planner_stats:
        pc = st.columns(3)
        pc[0].metric("Time (s)", f"{planner_stats.get('time_s', 0):.2f}")
        pc[1].metric("Messages", planner_stats.get("num_messages", 0))
        pc[2].metric("Tool calls", planner_stats.get("num_tool_calls", 0))
    if planner_prompt:
        with st.expander("Planner system prompt"):
            st.code(planner_prompt)
    _render_trace_entries(planner_entries)

    st.divider()

    # ── Phase 2: Executor ──
    st.markdown("### Phase 2: Executor")
    executor_stats = usage.get("executor", {})
    if executor_stats:
        ec = st.columns(3)
        ec[0].metric("Time (s)", f"{executor_stats.get('time_s', 0):.2f}")
        ec[1].metric("Messages", executor_stats.get("num_messages", 0))
        ec[2].metric("Tool calls", executor_stats.get("num_tool_calls", 0))
    if executor_prompt:
        with st.expander("Executor system prompt"):
            st.code(executor_prompt)
    _render_trace_entries(executor_entries)

else:
    # Standard single-agent rendering
    system_prompt = (outcome or {}).get("system_prompt") or summary.get("system_prompt")
    if system_prompt:
        with st.expander("System prompt"):
            st.code(system_prompt)
    _render_trace_entries(trace)

# ── Job Logs ────────────────────────────────────────────────────────────────

agent_logs_path = task_dir / "agent_logs.txt"
env_logs_path = task_dir / "env_logs.txt"

if agent_logs_path.exists() or env_logs_path.exists():
    st.divider()
    st.subheader("Job Logs")

    if agent_logs_path.exists():
        with st.expander("Agent logs"):
            st.code(agent_logs_path.read_text())

    if env_logs_path.exists():
        with st.expander("Env logs"):
            st.code(env_logs_path.read_text())