Add metadata-based single trajectory statistics, fix calculated formula (v0.3.37)

- Add "One trajectory statistics. Metadata from .traj" section
- Fix "One trajectory statistics. Calculated" to use same formula as aggregate charts
- uncached_input = prompt_tokens - cache_read - cache_creation

app.py

@@ -174,6 +174,276 @@ def calculate_per_step_tokens(steps: list[dict]) -> list[dict]:
     return result
 
 
+def _parse_usage_from_log_line(line: str) -> dict | None:
+    """
+    Parse usage info from log line containing ModelResponse or similar format.
+    Returns dict with prompt_tokens, completion_tokens, cached_tokens, etc.
+    """
+    if "usage=" not in line:
+        return None
+
+    result = {}
+
+    for field in ["completion_tokens", "prompt_tokens", "total_tokens"]:
+        match = re.search(rf'{field}=(\d+)', line)
+        if match:
+            result[field] = int(match.group(1))
+
+    cached_match = re.search(r'cached_tokens=(\d+)', line)
+    if cached_match:
+        result["cached_tokens"] = int(cached_match.group(1))
+
+    return result if result else None
+
+
+def _parse_old_format_log(log_path: Path) -> list[dict]:
+    """
+    Parse old SWE-agent format .info.log file to extract per-step token usage.
+    """
+    result = []
+    step = 0
+
+    try:
+        with open(log_path, "r", encoding="utf-8") as f:
+            for line in f:
+                if "usage=Usage(" not in line:
+                    continue
+
+                usage = _parse_usage_from_log_line(line)
+                if not usage:
+                    continue
+
+                prompt_tokens = usage.get("prompt_tokens", 0)
+                completion_tokens = usage.get("completion_tokens", 0)
+                cached_tokens = usage.get("cached_tokens", 0)
+
+                uncached_input = max(0, prompt_tokens - cached_tokens)
+
+                result.append({
+                    "step": step,
+                    "cache_read": cached_tokens,
+                    "uncached_input": uncached_input,
+                    "completion": completion_tokens,
+                    "cache_creation": 0,
+                })
+                step += 1
+    except Exception as e:
+        logging.debug("Error parsing log file %s: %s", log_path, e)
+
+    return result
+
+
+def parse_trajectory_metadata_per_step(traj_path: Path) -> list[dict]:
+    """
+    Parse trajectory file and extract per-step metadata from usage fields.
+    Supports both new format (.traj.json with messages[].extra.response.usage)
+    and old format (.traj with separate .info.log file).
+
+    Returns list of per-step data:
+        [{step: 0, cache_read: X, uncached_input: Y, completion: Z, cache_creation: W}, ...]
+    """
+    with open(traj_path, "r", encoding="utf-8") as f:
+        data = json.load(f)
+
+    messages = data.get("messages", [])
+    result = []
+    step = 0
+
+    for msg in messages:
+        if msg.get("role") != "assistant":
+            continue
+
+        usage = None
+        if "usage" in msg:
+            usage = msg["usage"]
+        elif "extra" in msg and isinstance(msg["extra"], dict):
+            response = msg["extra"].get("response", {})
+            if isinstance(response, dict):
+                usage = response.get("usage", {})
+
+        if usage:
+            prompt_tokens = usage.get("prompt_tokens", 0) or 0
+            completion_tokens = usage.get("completion_tokens", 0) or 0
+            cache_read = usage.get("cache_read_input_tokens", 0) or 0
+            cache_creation = usage.get("cache_creation_input_tokens", 0) or 0
+
+            prompt_tokens_details = usage.get("prompt_tokens_details", {})
+            if isinstance(prompt_tokens_details, dict):
+                cached_from_details = prompt_tokens_details.get("cached_tokens", 0) or 0
+                if cached_from_details > 0 and cache_read == 0:
+                    cache_read = cached_from_details
+
+            uncached_input = max(0, prompt_tokens - cache_read - cache_creation)
+
+            result.append({
+                "step": step,
+                "cache_read": cache_read,
+                "uncached_input": uncached_input,
+                "completion": completion_tokens,
+                "cache_creation": cache_creation,
+            })
+            step += 1
+
+    if not result:
+        log_path = traj_path.with_suffix(".info.log")
+        if not log_path.exists():
+            base_name = traj_path.stem.replace(".traj", "")
+            log_path = traj_path.parent / f"{base_name}.info.log"
+
+        if log_path.exists():
+            result = _parse_old_format_log(log_path)
+
+    return result
+
+
+def load_all_trajectory_metadata_steps(folder: str) -> dict[str, list[dict]]:
+    """
+    Load per-step metadata for all trajectories.
+
+    Returns:
+        dict mapping instance_id -> list of per-step metadata
+    """
+    output_dir = TRAJS_DIR / folder
+
+    traj_files = list(output_dir.glob("*/*.traj.json"))
+    if not traj_files:
+        traj_files = list(output_dir.glob("*/*.traj"))
+    if not traj_files:
+        traj_files = list(output_dir.glob("*.traj.json"))
+    if not traj_files:
+        traj_files = list(output_dir.glob("*.traj"))
+    if not traj_files:
+        traj_files = list(output_dir.glob("*.json"))
+
+    result = {}
+    for traj_path in traj_files:
+        try:
+            instance_id = traj_path.stem.replace(".traj", "")
+            steps = parse_trajectory_metadata_per_step(traj_path)
+            if steps:
+                result[instance_id] = steps
+        except Exception as e:
+            logging.error("Error parsing metadata steps for %s: %s", traj_path, e, exc_info=True)
+
+    return result
+
+
+def create_single_trajectory_meta_chart(steps: list[dict]):
+    """Create stacked bar chart for a single trajectory showing metadata tokens per step."""
+    import plotly.graph_objects as go
+
+    if not steps:
+        return None
+
+    x_labels = [f"Step {d['step']}" for d in steps]
+    uncached = [d["uncached_input"] / 1e3 for d in steps]
+    cache_read = [d["cache_read"] / 1e3 for d in steps]
+    cache_creation = [d["cache_creation"] / 1e3 for d in steps]
+    completion = [d["completion"] / 1e3 for d in steps]
+
+    fig = go.Figure()
+
+    fig.add_trace(go.Bar(
+        name="Uncached Input",
+        x=x_labels,
+        y=uncached,
+        marker_color="#EF553B",
+        hovertemplate="Step %{x}<br>Uncached Input: %{y:.2f}K<extra></extra>",
+    ))
+
+    fig.add_trace(go.Bar(
+        name="Cache Read",
+        x=x_labels,
+        y=cache_read,
+        marker_color="#19D3F3",
+        hovertemplate="Step %{x}<br>Cache Read: %{y:.2f}K<extra></extra>",
+    ))
+
+    fig.add_trace(go.Bar(
+        name="Cache Creation",
+        x=x_labels,
+        y=cache_creation,
+        marker_color="#FFA15A",
+        hovertemplate="Step %{x}<br>Cache Creation: %{y:.2f}K<extra></extra>",
+    ))
+
+    fig.add_trace(go.Bar(
+        name="Completion",
+        x=x_labels,
+        y=completion,
+        marker_color="#AB63FA",
+        hovertemplate="Step %{x}<br>Completion: %{y:.2f}K<extra></extra>",
+    ))
+
+    fig.update_layout(
+        barmode="stack",
+        xaxis_title="Step",
+        yaxis_title="Tokens (K)",
+        legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1),
+        margin=dict(l=50, r=20, t=40, b=40),
+    )
+
+    return fig
+
+
+def create_single_trajectory_meta_cost_chart(steps: list[dict], input_price: float, cache_read_price: float, cache_creation_price: float, completion_price: float):
+    """Create stacked bar chart for a single trajectory showing metadata cost per step."""
+    import plotly.graph_objects as go
+
+    if not steps:
+        return None
+
+    x_labels = [f"Step {d['step']}" for d in steps]
+    uncached_cost = [d["uncached_input"] * input_price / 1e6 for d in steps]
+    cache_read_cost = [d["cache_read"] * cache_read_price / 1e6 for d in steps]
+    cache_creation_cost = [d["cache_creation"] * cache_creation_price / 1e6 for d in steps]
+    completion_cost = [d["completion"] * completion_price / 1e6 for d in steps]
+
+    fig = go.Figure()
+
+    fig.add_trace(go.Bar(
+        name="Uncached Input",
+        x=x_labels,
+        y=uncached_cost,
+        marker_color="#EF553B",
+        hovertemplate="Step %{x}<br>Uncached Input: $%{y:.4f}<extra></extra>",
+    ))
+
+    fig.add_trace(go.Bar(
+        name="Cache Read",
+        x=x_labels,
+        y=cache_read_cost,
+        marker_color="#19D3F3",
+        hovertemplate="Step %{x}<br>Cache Read: $%{y:.4f}<extra></extra>",
+    ))
+
+    fig.add_trace(go.Bar(
+        name="Cache Creation",
+        x=x_labels,
+        y=cache_creation_cost,
+        marker_color="#FFA15A",
+        hovertemplate="Step %{x}<br>Cache Creation: $%{y:.4f}<extra></extra>",
+    ))
+
+    fig.add_trace(go.Bar(
+        name="Completion",
+        x=x_labels,
+        y=completion_cost,
+        marker_color="#AB63FA",
+        hovertemplate="Step %{x}<br>Completion: $%{y:.4f}<extra></extra>",
+    ))
+
+    fig.update_layout(
+        barmode="stack",
+        xaxis_title="Step",
+        yaxis_title="Cost ($)",
+        legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1),
+        margin=dict(l=50, r=20, t=40, b=40),
+    )
+
+    return fig
+
+
 def create_single_trajectory_chart(steps: list[dict]):
     """Create stacked bar chart for a single trajectory showing tokens per step."""
     import plotly.graph_objects as go
@@ -184,10 +454,11 @@ def create_single_trajectory_chart(steps: list[dict]):
     per_step_data = calculate_per_step_tokens(steps)
 
     x_labels = [f"Step {d['step']}" for d in per_step_data]
-    uncached = [d["uncached_input"] / 1e3 for d in per_step_data]
     cache_read = [d["cache_read"] / 1e3 for d in per_step_data]
     cache_creation = [d["cache_creation"] / 1e3 for d in per_step_data]
     completion = [d["completion"] / 1e3 for d in per_step_data]
+    prompt_tokens = [(d["cache_read"] + d["uncached_input"]) / 1e3 for d in per_step_data]
+    uncached = [max(0, p - cr - cc) for p, cr, cc in zip(prompt_tokens, cache_read, cache_creation)]
 
     fig = go.Figure()
 
@@ -244,10 +515,16 @@ def create_single_trajectory_cost_chart(steps: list[dict], input_price: float, c
     per_step_data = calculate_per_step_tokens(steps)
 
     x_labels = [f"Step {d['step']}" for d in per_step_data]
-    uncached_cost = [d["uncached_input"] * input_price / 1e6 for d in per_step_data]
-    cache_read_cost = [d["cache_read"] * cache_read_price / 1e6 for d in per_step_data]
-    cache_creation_cost = [d["cache_creation"] * cache_creation_price / 1e6 for d in per_step_data]
-    completion_cost = [d["completion"] * completion_price / 1e6 for d in per_step_data]
+    cache_read = [d["cache_read"] for d in per_step_data]
+    cache_creation = [d["cache_creation"] for d in per_step_data]
+    completion = [d["completion"] for d in per_step_data]
+    prompt_tokens = [d["cache_read"] + d["uncached_input"] for d in per_step_data]
+    uncached = [max(0, p - cr - cc) for p, cr, cc in zip(prompt_tokens, cache_read, cache_creation)]
+
+    uncached_cost = [u * input_price / 1e6 for u in uncached]
+    cache_read_cost = [cr * cache_read_price / 1e6 for cr in cache_read]
+    cache_creation_cost = [cc * cache_creation_price / 1e6 for cc in cache_creation]
+    completion_cost = [c * completion_price / 1e6 for c in completion]
 
     fig = go.Figure()
 
@@ -842,8 +1119,34 @@ def parse_trajectory(traj_path: Path) -> dict:
             result["prompt_tokens"] += usage.get("prompt_tokens", 0) or 0
             result["completion_tokens"] += usage.get("completion_tokens", 0) or 0
             result["total_tokens"] += usage.get("total_tokens", 0) or 0
-            result["cache_read_tokens"] += usage.get("cache_read_input_tokens", 0) or 0
-            result["cache_creation_tokens"] += usage.get("cache_creation_input_tokens", 0) or 0
+
+            cache_read = usage.get("cache_read_input_tokens", 0) or 0
+            cache_creation = usage.get("cache_creation_input_tokens", 0) or 0
+
+            prompt_tokens_details = usage.get("prompt_tokens_details", {})
+            if isinstance(prompt_tokens_details, dict):
+                cached_from_details = prompt_tokens_details.get("cached_tokens", 0) or 0
+                if cached_from_details > 0 and cache_read == 0:
+                    cache_read = cached_from_details
+
+            result["cache_read_tokens"] += cache_read
+            result["cache_creation_tokens"] += cache_creation
+
+    if result["prompt_tokens"] == 0 and result["completion_tokens"] == 0:
+        log_path = traj_path.with_suffix(".info.log")
+        if not log_path.exists():
+            base_name = traj_path.stem.replace(".traj", "")
+            log_path = traj_path.parent / f"{base_name}.info.log"
+
+        if log_path.exists():
+            steps = _parse_old_format_log(log_path)
+            for step_data in steps:
+                result["prompt_tokens"] += step_data["cache_read"] + step_data["uncached_input"]
+                result["completion_tokens"] += step_data["completion"]
+                result["cache_read_tokens"] += step_data["cache_read"]
+            result["total_tokens"] = result["prompt_tokens"] + result["completion_tokens"]
+            if result["api_calls"] == 0:
+                result["api_calls"] = len(steps)
 
     return result
 
@@ -1515,7 +1818,7 @@ def build_app():
         """)
         trajectories_state = gr.State(None)
 
-        gr.Markdown("# 🧮 SWE-bench Bash-Only Leaderboard `v0.3.36`")
+        gr.Markdown("# 🧮 SWE-bench Bash-Only Leaderboard `v0.3.37`")
         gr.Markdown("## 🎯 Select a base model for cost analysis (click a row)")
 
         with gr.Row():
@@ -1558,6 +1861,14 @@ def build_app():
                         with gr.Row():
                             plot_cost_breakdown_calc = gr.Plot(label="Cost per Trajectory")
 
+                    with gr.Accordion("One trajectory statistics. Metadata from .traj", open=False, visible=False) as single_traj_meta_accordion:
+                        with gr.Row():
+                            single_traj_meta_dropdown = gr.Dropdown(label="Select Issue", choices=[], interactive=True)
+                        with gr.Row():
+                            single_traj_meta_plot = gr.Plot(label="Tokens per Step (stacked)")
+                        with gr.Row():
+                            single_traj_meta_cost_plot = gr.Plot(label="Cost per Step (stacked) ($)")
+
                     with gr.Accordion("One trajectory statistics. Calculated from .traj messages", open=False, visible=False) as single_traj_accordion:
                         with gr.Row():
                             single_traj_dropdown = gr.Dropdown(label="Select Issue", choices=[], interactive=True)
@@ -2387,6 +2698,10 @@ def build_app():
                 gr.update(),
                 gr.update(),
                 gr.update(),
+                gr.update(visible=False),
+                gr.update(),
+                gr.update(),
+                gr.update(),
             )
 
             if not folder:
@@ -2408,6 +2723,10 @@ def build_app():
                     gr.update(),
                     gr.update(),
                     gr.update(),
+                    gr.update(visible=False),
+                    gr.update(),
+                    gr.update(),
+                    gr.update(),
                 )
                 progress(0.3, desc="Downloading")
                 status, _ = download_trajectories_from_s3(folder)
@@ -2425,6 +2744,10 @@ def build_app():
                         gr.update(),
                         gr.update(),
                         gr.update(),
+                        gr.update(visible=False),
+                        gr.update(),
+                        gr.update(),
+                        gr.update(),
                     )
                     return
             progress(0.45, desc="Loading trajectories")
@@ -2441,6 +2764,10 @@ def build_app():
                 gr.update(),
                 gr.update(),
                 gr.update(),
+                gr.update(visible=False),
+                gr.update(),
+                gr.update(),
+                gr.update(),
             )
 
             progress(0.6, desc="Reading metadata")
@@ -2449,8 +2776,10 @@ def build_app():
             df_calc = ensure_token_columns(load_all_trajectories_calculated(folder))
             df_calc["api_calls"] = df_meta["api_calls"].values
             df_calc["instance_cost"] = df_meta["instance_cost"].values
-            progress(0.8, desc="Reading steps")
+            progress(0.75, desc="Reading steps")
             trajectory_steps = load_all_trajectory_steps(folder)
+            progress(0.8, desc="Reading metadata steps")
+            metadata_steps = load_all_trajectory_metadata_steps(folder)
 
             model_details, _ = get_model_details(folder)
             resolved_instances = {}
@@ -2459,7 +2788,7 @@ def build_app():
                 for inst_id, details in per_instance.items():
                     resolved_instances[inst_id] = details.get("resolved", False)
 
-            state_data = {"meta": df_meta, "calculated": df_calc, "folder": folder, "steps": trajectory_steps, "resolved": resolved_instances}
+            state_data = {"meta": df_meta, "calculated": df_calc, "folder": folder, "steps": trajectory_steps, "metadata_steps": metadata_steps, "resolved": resolved_instances}
 
             if df_meta.empty:
                 progress(1, desc="No trajectories found")
@@ -2475,6 +2804,10 @@ def build_app():
                     gr.update(),
                     gr.update(),
                     gr.update(),
+                    gr.update(visible=False),
+                    gr.update(),
+                    gr.update(),
+                    gr.update(),
                 )
                 return
 
@@ -2504,6 +2837,10 @@ def build_app():
             issue_ids = sorted(trajectory_steps.keys())
             first_issue = issue_ids[0] if issue_ids else None
 
+            meta_issue_ids = sorted(metadata_steps.keys())
+            first_meta_issue = meta_issue_ids[0] if meta_issue_ids else None
+            has_meta_steps = len(meta_issue_ids) > 0
+
             progress(1, desc="Done")
             yield (
                 f"✅ Loaded {len(df_meta)} trajectories",
@@ -2517,6 +2854,10 @@ def build_app():
                 gr.update(choices=issue_ids, value=first_issue),
                 gr.update(),
                 gr.update(),
+                gr.update(visible=has_meta_steps),
+                gr.update(choices=meta_issue_ids, value=first_meta_issue),
+                gr.update(),
+                gr.update(),
             )
 
         def on_single_traj_select(state_data, issue_id, input_price, cache_read_price, cache_creation_price, completion_price):
@@ -2530,6 +2871,17 @@ def build_app():
             cost_chart = create_single_trajectory_cost_chart(steps, input_price, cache_read_price, cache_creation_price, completion_price)
             return tokens_chart, cost_chart
 
+        def on_single_traj_meta_select(state_data, issue_id, input_price, cache_read_price, cache_creation_price, completion_price):
+            if state_data is None or not issue_id:
+                return None, None
+            metadata_steps = state_data.get("metadata_steps", {})
+            if issue_id not in metadata_steps:
+                return None, None
+            steps = metadata_steps[issue_id]
+            tokens_chart = create_single_trajectory_meta_chart(steps)
+            cost_chart = create_single_trajectory_meta_cost_chart(steps, input_price, cache_read_price, cache_creation_price, completion_price)
+            return tokens_chart, cost_chart
+
         analyze_btn.click(
             fn=load_and_analyze,
             inputs=[selected_folder, price_input, price_cache_read, price_cache_creation, price_completion, thinking_overhead, use_cache],
@@ -2545,11 +2897,19 @@ def build_app():
                 single_traj_dropdown,
                 single_traj_plot,
                 single_traj_cost_plot,
+                single_traj_meta_accordion,
+                single_traj_meta_dropdown,
+                single_traj_meta_plot,
+                single_traj_meta_cost_plot,
             ],
         ).then(
             fn=on_single_traj_select,
             inputs=[trajectories_state, single_traj_dropdown, price_input, price_cache_read, price_cache_creation, price_completion],
             outputs=[single_traj_plot, single_traj_cost_plot],
+        ).then(
+            fn=on_single_traj_meta_select,
+            inputs=[trajectories_state, single_traj_meta_dropdown, price_input, price_cache_read, price_cache_creation, price_completion],
+            outputs=[single_traj_meta_plot, single_traj_meta_cost_plot],
         )
 
         def recalculate_costs(state_data, input_price, cache_read_price, cache_creation_price, completion_price, overhead, with_cache):
@@ -2613,6 +2973,12 @@ def build_app():
             outputs=[single_traj_plot, single_traj_cost_plot],
         )
 
+        single_traj_meta_dropdown.change(
+            fn=on_single_traj_meta_select,
+            inputs=[trajectories_state, single_traj_meta_dropdown, price_input, price_cache_read, price_cache_creation, price_completion],
+            outputs=[single_traj_meta_plot, single_traj_meta_cost_plot],
+        )
+
         thinking_overhead.change(
             fn=on_calc_options_change,
             inputs=calc_options_inputs,

requirements.txt

@@ -6,3 +6,4 @@ python-dotenv>=1.0.0
 tiktoken>=0.12.0
 awscli
 
+