Add cumulative statistics to routing charts

- Show total tokens/cost with breakdown by model
- Annotations: Total, Base Model, Model 1, Model 2...
- Fix model_name mismatch in load_all_trajectories_calculated
- Use step_model from steps instead of reading model_name separately

app.py

@@ -265,25 +265,6 @@ def load_all_trajectories_calculated(folder: str) -> pd.DataFrame:
 
     trajectory_steps = load_all_trajectory_steps(folder)
 
-    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"))
-
-    model_name = ""
-    if traj_files:
-        try:
-            with open(traj_files[0], "r") as f:
-                first_data = json.load(f)
-                config = first_data.get("info", {}).get("config", {}).get("model", {})
-                model_name = config.get("cost_calc_model_override", config.get("model_name", ""))
-        except Exception:
-            pass
-
     rows = []
     for instance_id, steps in trajectory_steps.items():
         if not steps:
@@ -291,7 +272,8 @@ def load_all_trajectories_calculated(folder: str) -> pd.DataFrame:
 
         try:
             model_totals = calculate_routing_tokens(steps)
-            totals = model_totals.get(model_name, {})
+            step_model = steps[0].get("model", "") if steps else ""
+            totals = model_totals.get(step_model, {})
 
             cache_read = totals.get("cache_read", 0)
             uncached_input = totals.get("uncached_input", 0)
@@ -302,7 +284,7 @@ def load_all_trajectories_calculated(folder: str) -> pd.DataFrame:
 
             rows.append({
                 "instance_id": instance_id,
-                "model_name": model_name,
+                "model_name": step_model,
                 "api_calls": len(steps),
                 "instance_cost": 0,
                 "prompt_tokens": prompt_tokens,
@@ -1099,6 +1081,7 @@ def create_routed_token_chart(base_tokens: dict, additional_models: list):
 
     fig = go.Figure()
 
+    base_total = sum(base_tokens.get(k, 0) for k in ["uncached_input", "cache_read", "cache_creation", "completion"])
     base_values = [
         base_tokens.get("uncached_input", 0) / 1e6,
         base_tokens.get("cache_read", 0) / 1e6,
@@ -1107,7 +1090,11 @@ def create_routed_token_chart(base_tokens: dict, additional_models: list):
     ]
     fig.add_trace(go.Bar(name="Base Model", x=categories, y=base_values, marker_color=colors[0]))
 
+    model_totals = [("Base Model", base_total)]
+
     for i, (model_name, tokens) in enumerate(additional_models):
+        model_total = sum(tokens.get(k, 0) for k in ["uncached_input", "cache_read", "cache_creation", "completion"])
+        model_totals.append((model_name or f"Model {i+1}", model_total))
         values = [
             tokens.get("uncached_input", 0) / 1e6,
             tokens.get("cache_read", 0) / 1e6,
@@ -1117,13 +1104,28 @@ def create_routed_token_chart(base_tokens: dict, additional_models: list):
         color = colors[(i + 1) % len(colors)]
         fig.add_trace(go.Bar(name=model_name or f"Model {i+1}", x=categories, y=values, marker_color=color))
 
+    grand_total = sum(t for _, t in model_totals)
+    annotation_lines = [f"<b>Total: {grand_total/1e6:.2f}M</b>"]
+    for name, total in model_totals:
+        annotation_lines.append(f"{name}: {total/1e6:.2f}M")
+
     fig.update_layout(
         title="Tokens by Type (per Model)",
         yaxis_title="Tokens (M)",
         barmode="group",
-        margin=dict(l=40, r=40, t=60, b=40),
+        margin=dict(l=40, r=40, t=80, b=40),
         legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1),
     )
+    fig.add_annotation(
+        text="<br>".join(annotation_lines),
+        xref="paper", yref="paper",
+        x=0.02, y=0.98, showarrow=False,
+        font=dict(size=11),
+        align="left",
+        bgcolor="rgba(255,255,255,0.8)",
+        bordercolor="gray",
+        borderwidth=1,
+    )
     return fig
 
 
@@ -1142,6 +1144,7 @@ def create_routed_cost_chart(base_costs: dict, additional_models: list):
 
     fig = go.Figure()
 
+    base_total = sum(base_costs.get(k, 0) for k in ["uncached_input", "cache_read", "cache_creation", "completion"])
     base_values = [
         base_costs.get("uncached_input", 0),
         base_costs.get("cache_read", 0),
@@ -1150,7 +1153,11 @@ def create_routed_cost_chart(base_costs: dict, additional_models: list):
     ]
     fig.add_trace(go.Bar(name="Base Model", x=categories, y=base_values, marker_color=colors[0]))
 
+    model_totals = [("Base Model", base_total)]
+
     for i, (model_name, costs) in enumerate(additional_models):
+        model_total = sum(costs.get(k, 0) for k in ["uncached_input", "cache_read", "cache_creation", "completion"])
+        model_totals.append((model_name or f"Model {i+1}", model_total))
         values = [
             costs.get("uncached_input", 0),
             costs.get("cache_read", 0),
@@ -1160,13 +1167,28 @@ def create_routed_cost_chart(base_costs: dict, additional_models: list):
         color = colors[(i + 1) % len(colors)]
         fig.add_trace(go.Bar(name=model_name or f"Model {i+1}", x=categories, y=values, marker_color=color))
 
+    grand_total = sum(t for _, t in model_totals)
+    annotation_lines = [f"<b>Total: ${grand_total:.2f}</b>"]
+    for name, total in model_totals:
+        annotation_lines.append(f"{name}: ${total:.2f}")
+
     fig.update_layout(
         title="Cost by Type (per Model) ($)",
         yaxis_title="Cost ($)",
         barmode="group",
-        margin=dict(l=40, r=40, t=60, b=40),
+        margin=dict(l=40, r=40, t=80, b=40),
         legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1),
     )
+    fig.add_annotation(
+        text="<br>".join(annotation_lines),
+        xref="paper", yref="paper",
+        x=0.02, y=0.98, showarrow=False,
+        font=dict(size=11),
+        align="left",
+        bgcolor="rgba(255,255,255,0.8)",
+        bordercolor="gray",
+        borderwidth=1,
+    )
     return fig