Auto-recalculate cost charts when prices change

- Add create_cost_by_type_chart() function for reuse
- Store trajectories DataFrame in gr.State
- Add .change() handlers on all 4 price fields
- Recalculate 'Total Cost by Token Type' and 'Cost Breakdown per Instance' on price change

app.py

@@ -250,6 +250,54 @@ def load_all_trajectories(folder: str) -> pd.DataFrame:
     return df
 
 
+def create_cost_by_type_chart(df: pd.DataFrame, input_price: float, cache_read_price: float, cache_creation_price: float, completion_price: float):
+    """Create Total Cost by Token Type chart (can be called separately for price updates)"""
+    if df.empty:
+        return None
+
+    total_completion = df["completion_tokens"].sum()
+    total_cache_read = df["cache_read_tokens"].sum()
+    total_cache_creation = df["cache_creation_tokens"].sum()
+    df_temp = df.copy()
+    df_temp["uncached_input"] = (df_temp["prompt_tokens"] - df_temp["cache_read_tokens"] - df_temp["cache_creation_tokens"]).clip(lower=0)
+    total_uncached_input = df_temp["uncached_input"].sum()
+
+    cost_uncached_input = total_uncached_input * input_price / 1e6
+    cost_cache_read = total_cache_read * cache_read_price / 1e6
+    cost_cache_creation = total_cache_creation * cache_creation_price / 1e6
+    cost_completion = total_completion * completion_price / 1e6
+
+    cost_data = pd.DataFrame({
+        "Token Type": ["Uncached Input", "Cache Read", "Cache Creation", "Completion"],
+        "Cost ($)": [cost_uncached_input, cost_cache_read, cost_cache_creation, cost_completion],
+    })
+
+    fig = px.bar(
+        cost_data,
+        x="Token Type",
+        y="Cost ($)",
+        title="Total Cost by Token Type ($)",
+        color="Token Type",
+        color_discrete_sequence=["#EF553B", "#19D3F3", "#FFA15A", "#AB63FA"],
+    )
+    fig.update_layout(
+        xaxis_title="Token Type",
+        yaxis_title="Cost ($)",
+        showlegend=False,
+        margin=dict(l=40, r=20, t=40, b=40),
+    )
+
+    total_cost = cost_uncached_input + cost_cache_read + cost_cache_creation + cost_completion
+    fig.add_annotation(
+        text=f"Total: ${total_cost:.2f}",
+        xref="paper", yref="paper",
+        x=0.95, y=0.95, showarrow=False,
+        font=dict(size=12),
+    )
+
+    return fig
+
+
 def create_basic_histograms(df: pd.DataFrame, input_price: float, cache_read_price: float, cache_creation_price: float, completion_price: float):
     if df.empty:
         return None, None, None, None, None
@@ -330,39 +378,8 @@ def create_basic_histograms(df: pd.DataFrame, input_price: float, cache_read_pri
         font=dict(size=12),
     )
 
-    # Cost by token type
-    cost_uncached_input = total_uncached_input * input_price / 1e6
-    cost_cache_read = total_cache_read * cache_read_price / 1e6
-    cost_cache_creation = total_cache_creation * cache_creation_price / 1e6
-    cost_completion = total_completion * completion_price / 1e6
-
-    cost_data = pd.DataFrame({
-        "Token Type": ["Uncached Input", "Cache Read", "Cache Creation", "Completion"],
-        "Cost ($)": [cost_uncached_input, cost_cache_read, cost_cache_creation, cost_completion],
-    })
-
-    fig_tokens_cost = px.bar(
-        cost_data,
-        x="Token Type",
-        y="Cost ($)",
-        title="Total Cost by Token Type ($)",
-        color="Token Type",
-        color_discrete_sequence=["#EF553B", "#19D3F3", "#FFA15A", "#AB63FA"],
-    )
-    fig_tokens_cost.update_layout(
-        xaxis_title="Token Type",
-        yaxis_title="Cost ($)",
-        showlegend=False,
-        margin=dict(l=40, r=20, t=40, b=40),
-    )
-
-    total_cost = cost_uncached_input + cost_cache_read + cost_cache_creation + cost_completion
-    fig_tokens_cost.add_annotation(
-        text=f"Total: ${total_cost:.2f}",
-        xref="paper", yref="paper",
-        x=0.95, y=0.95, showarrow=False,
-        font=dict(size=12),
-    )
+    # Cost by token type (use separate function)
+    fig_tokens_cost = create_cost_by_type_chart(df, input_price, cache_read_price, cache_creation_price, completion_price)
 
     df_sorted = df.sort_values("cache_read_tokens", ascending=False).reset_index(drop=True)
     df_sorted["instance_idx"] = range(len(df_sorted))
@@ -627,6 +644,7 @@ def build_app():
             empty_result = (
                 gr.update(visible=False),
                 None, None, None, None, None, None,
+                None,  # trajectories_state
             )
 
             if not folder:
@@ -636,6 +654,7 @@ def build_app():
             yield (
                 gr.update(visible=True),
                 None, None, None, None, None, None,
+                None,
             )
 
             df = load_all_trajectories(folder)
@@ -651,6 +670,7 @@ def build_app():
             yield (
                 gr.update(visible=True),
                 fig_steps, fig_cost, fig_tokens, fig_tokens_cost, fig_stacked, fig_cost_breakdown,
+                df,  # save to state
             )
 
         analyze_btn.click(
@@ -659,9 +679,25 @@ def build_app():
             outputs=[
                 analysis_section,
                 plot_steps, plot_cost, plot_tokens, plot_tokens_cost, plot_stacked, plot_cost_breakdown,
+                trajectories_state,
             ],
         )
 
+        def recalculate_costs(df, input_price, cache_read_price, cache_creation_price, completion_price):
+            if df is None or (isinstance(df, pd.DataFrame) and df.empty):
+                return None, None
+            fig_tokens_cost = create_cost_by_type_chart(df, input_price, cache_read_price, cache_creation_price, completion_price)
+            fig_cost_breakdown = create_cost_breakdown(df, input_price, cache_read_price, cache_creation_price, completion_price)
+            return fig_tokens_cost, fig_cost_breakdown
+
+        price_inputs = [trajectories_state, price_input, price_cache_read, price_cache_creation, price_completion]
+        price_outputs = [plot_tokens_cost, plot_cost_breakdown]
+
+        price_input.change(fn=recalculate_costs, inputs=price_inputs, outputs=price_outputs)
+        price_cache_read.change(fn=recalculate_costs, inputs=price_inputs, outputs=price_outputs)
+        price_cache_creation.change(fn=recalculate_costs, inputs=price_inputs, outputs=price_outputs)
+        price_completion.change(fn=recalculate_costs, inputs=price_inputs, outputs=price_outputs)
+
     return app