Add routing UI: model selection, strategy params, support .traj files

- Add routing section with up to 3 models
- Each model has: dropdown with search, price fields, strategy selection
- Strategy options: random steps, every k, part of trajectory
- Support both .traj.json and .traj file formats
- Load & Analyze auto-downloads trajectories if needed
- Add Routing button appears only after successful load
- Fix error handling for models without trajectories on S3
- Compact token prices UI

app.py

@@ -28,6 +28,21 @@ _trajectories_cache = {}
 _calculated_tokens_cache = {}
 
 
+def parse_step_or_ratio(value: float, total_steps: int) -> int:
+    """
+    Parse a value as either step number or ratio.
+
+    If value is integer (e.g., 3.0, 5.0) -> treat as step number
+    If value is float with decimal (e.g., 0.5, 0.25) -> treat as ratio of total_steps
+
+    Returns: step index (0-based)
+    """
+    if value == int(value) and value >= 1:
+        return int(value)
+    else:
+        return int(value * total_steps)
+
+
 def get_default_overhead(model_name: str) -> float:
     """Get default tokenizer overhead for model provider"""
     model_lower = model_name.lower() if model_name else ""
@@ -172,8 +187,12 @@ def load_all_trajectories_calculated(folder: str) -> pd.DataFrame:
     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"))
     
@@ -212,6 +231,12 @@ def load_all_trajectories_calculated(folder: str) -> pd.DataFrame:
     return df
 
 
+def get_litellm_model_list() -> list[str]:
+    """Get list of model names from litellm prices"""
+    prices = get_litellm_prices()
+    return sorted(prices.keys())
+
+
 def get_litellm_prices() -> dict:
     global _litellm_prices_cache
     if _litellm_prices_cache is not None:
@@ -357,6 +382,8 @@ def download_trajectories_from_s3(folder: str, progress=gr.Progress()):
     output_dir = TRAJS_DIR / folder
     if output_dir.exists() and any(output_dir.iterdir()):
         file_count = len(list(output_dir.glob("*/*.traj.json")))
+        if file_count == 0:
+            file_count = len(list(output_dir.glob("*/*.traj")))
         if file_count == 0:
             file_count = len(list(output_dir.glob("*.json")))
         return f"✅ Already downloaded: {output_dir}\n\n{file_count} trajectory files", gr.update(visible=True)
@@ -378,9 +405,14 @@ def download_trajectories_from_s3(folder: str, progress=gr.Progress()):
             return f"❌ S3 download failed:\n{result.stderr}", gr.update(visible=False)
 
         file_count = len(list(output_dir.glob("*/*.traj.json")))
+        if file_count == 0:
+            file_count = len(list(output_dir.glob("*/*.traj")))
         if file_count == 0:
             file_count = len(list(output_dir.glob("*.json")))
 
+        if file_count == 0:
+            return f"❌ No trajectory files found on S3 for {folder}", gr.update(visible=False)
+
         per_instance = model.get("per_instance_details", {})
         resolved_count = sum(1 for v in per_instance.values() if v.get("resolved"))
         total_count = len(per_instance)
@@ -452,8 +484,12 @@ def load_all_trajectories(folder: str) -> pd.DataFrame:
     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"))
 
@@ -873,12 +909,11 @@ def on_row_select(evt: gr.SelectData, df: pd.DataFrame):
     if evt.index is None:
         return (
             "", "",
-            gr.update(interactive=False),
             gr.update(visible=False),
-            gr.update(value=0, label="💲 Input"),
-            gr.update(value=0, label="💲 Cache Read"),
-            gr.update(value=0, label="💲 Cache Creation"),
-            gr.update(value=0, label="💲 Completion"),
+            gr.update(value=0, label="Input"),
+            gr.update(value=0, label="Cache Read"),
+            gr.update(value=0, label="Cache Creation"),
+            gr.update(value=0, label="Completion"),
             "",
             gr.update(value=1.0),
         )
@@ -888,8 +923,6 @@ def on_row_select(evt: gr.SelectData, df: pd.DataFrame):
     folder = row["folder"]
     name = row["name"]
 
-    show_analyze = check_trajectories_downloaded(folder)
-
     prices_dict, model_hint = get_prices_for_folder(folder)
     default_overhead = get_default_overhead(model_hint)
 
@@ -898,14 +931,13 @@ def on_row_select(evt: gr.SelectData, df: pd.DataFrame):
         if price_info["found"]:
             return gr.update(value=value, label=f"✅ {name}")
         elif value > 0:
-            return gr.update(value=value, label=f"❌ {name} (estimated)")
+            return gr.update(value=value, label=f"❌ {name} (est.)")
         else:
             return gr.update(value=0, label=f"❌ {name}")
 
     return (
         folder, name,
-        gr.update(interactive=True),
-        gr.update(visible=show_analyze),
+        gr.update(visible=True),
         price_update(prices_dict["input"], "Input"),
         price_update(prices_dict["cache_read"], "Cache Read"),
         price_update(prices_dict["cache_creation"], "Cache Creation"),
@@ -953,18 +985,17 @@ def build_app():
                 gr.Markdown("### Selected Model")
                 selected_name = gr.Textbox(label="Model Name", interactive=False)
 
-                download_btn = gr.Button("📥 Download Trajectories", interactive=False)
-                download_status = gr.Textbox(label="Status", interactive=False, lines=3)
-
                 analyze_btn = gr.Button("📊 Load & Analyze", visible=False, variant="primary")
+                download_status = gr.Textbox(label="Status", interactive=False, lines=3)
 
                 gr.Markdown("---")
                 gr.Markdown("### 💰 Token Prices ($/1M) · *[litellm](https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json)*")
                 detected_model = gr.Textbox(label="Detected Model", interactive=False)
-                price_input = gr.Number(label="💲 Input", value=0, precision=2)
-                price_cache_read = gr.Number(label="💲 Cache Read", value=0, precision=2)
-                price_cache_creation = gr.Number(label="💲 Cache Creation", value=0, precision=2)
-                price_completion = gr.Number(label="💲 Completion", value=0, precision=2)
+                with gr.Row():
+                    price_input = gr.Number(label="Input", value=0, precision=2, scale=1)
+                    price_cache_read = gr.Number(label="Cache Read", value=0, precision=2, scale=1)
+                    price_cache_creation = gr.Number(label="Cache Creation", value=0, precision=2, scale=1)
+                    price_completion = gr.Number(label="Completion", value=0, precision=2, scale=1)
 
                 gr.Markdown("---")
                 gr.Markdown("### 📊 Token Count Source")
@@ -986,6 +1017,231 @@ def build_app():
                     visible=False,
                 )
 
+                gr.Markdown("---")
+                add_routing_btn = gr.Button("➕ Add Routing", variant="primary", visible=False)
+
+                with gr.Column(visible=False) as routing_section:
+                    gr.Markdown("### 🔀 Routing Models")
+
+                    STRATEGY_CHOICES = [
+                        "Replace on random steps",
+                        "Replace every step k",
+                        "Replace part of trajectory",
+                    ]
+
+                    with gr.Column():
+                        with gr.Group():
+                            gr.Markdown("#### Route to Model 1")
+                            routing_model_1 = gr.Dropdown(
+                                label="Model (type 3+ chars to search)",
+                                choices=[],
+                                allow_custom_value=True,
+                                interactive=True,
+                            )
+                            with gr.Row():
+                                routing_price_1_input = gr.Number(label="Input", precision=3, scale=1)
+                                routing_price_1_cache_read = gr.Number(label="Cache Read", precision=3, scale=1)
+                                routing_price_1_cache_creation = gr.Number(label="Cache Creation", precision=3, scale=1)
+                                routing_price_1_completion = gr.Number(label="Completion", precision=3, scale=1)
+                            strategy_1 = gr.Dropdown(
+                                label="Strategy",
+                                choices=STRATEGY_CHOICES,
+                                value="Replace on random steps",
+                                interactive=True,
+                            )
+                            with gr.Row(visible=True) as random_params_1:
+                                random_pct_1 = gr.Number(label="Percentage (%)", value=50, minimum=0, maximum=100, precision=0, interactive=True)
+                            with gr.Row(visible=False) as every_k_params_1:
+                                step_k_1 = gr.Number(label="k", value=2, minimum=1, precision=0, interactive=True)
+                            with gr.Row(visible=False) as part_params_1:
+                                start_step_1 = gr.Number(label="Start (int=step; 0,0-1,0=ratio)", value=0, minimum=0, precision=2, interactive=True)
+                                end_step_1 = gr.Number(label="End (int=step; 0,0-1,0=ratio)", value=0.5, minimum=0, precision=2, interactive=True)
+
+                        add_model_2_btn = gr.Button("+ Add another model", size="sm", visible=False)
+
+                        with gr.Column(visible=False) as routing_block_2:
+                            with gr.Group():
+                                gr.Markdown("#### Route to Model 2")
+                                routing_model_2 = gr.Dropdown(
+                                    label="Model (type 3+ chars to search)",
+                                    choices=[],
+                                    allow_custom_value=True,
+                                    interactive=True,
+                                )
+                                with gr.Row():
+                                    routing_price_2_input = gr.Number(label="Input", precision=3, scale=1)
+                                    routing_price_2_cache_read = gr.Number(label="Cache Read", precision=3, scale=1)
+                                    routing_price_2_cache_creation = gr.Number(label="Cache Creation", precision=3, scale=1)
+                                    routing_price_2_completion = gr.Number(label="Completion", precision=3, scale=1)
+                                strategy_2 = gr.Dropdown(
+                                    label="Strategy",
+                                    choices=STRATEGY_CHOICES,
+                                    value="Replace on random steps",
+                                    interactive=True,
+                                )
+                                with gr.Row(visible=True) as random_params_2:
+                                    random_pct_2 = gr.Number(label="Percentage (%)", value=50, minimum=0, maximum=100, precision=0, interactive=True)
+                                with gr.Row(visible=False) as every_k_params_2:
+                                    step_k_2 = gr.Number(label="k", value=2, minimum=1, precision=0, interactive=True)
+                                with gr.Row(visible=False) as part_params_2:
+                                    start_step_2 = gr.Number(label="Start (int=step; 0,0-1,0=ratio)", value=0, minimum=0, precision=2, interactive=True)
+                                    end_step_2 = gr.Number(label="End (int=step; 0,0-1,0=ratio)", value=0.5, minimum=0, precision=2, interactive=True)
+
+                            add_model_3_btn = gr.Button("+ Add another model", size="sm", visible=False)
+
+                        with gr.Column(visible=False) as routing_block_3:
+                            with gr.Group():
+                                gr.Markdown("#### Route to Model 3")
+                                routing_model_3 = gr.Dropdown(
+                                    label="Model (type 3+ chars to search)",
+                                    choices=[],
+                                    allow_custom_value=True,
+                                    interactive=True,
+                                )
+                                with gr.Row():
+                                    routing_price_3_input = gr.Number(label="Input", precision=3, scale=1)
+                                    routing_price_3_cache_read = gr.Number(label="Cache Read", precision=3, scale=1)
+                                    routing_price_3_cache_creation = gr.Number(label="Cache Creation", precision=3, scale=1)
+                                    routing_price_3_completion = gr.Number(label="Completion", precision=3, scale=1)
+                                strategy_3 = gr.Dropdown(
+                                    label="Strategy",
+                                    choices=STRATEGY_CHOICES,
+                                    value="Replace on random steps",
+                                    interactive=True,
+                                )
+                                with gr.Row(visible=True) as random_params_3:
+                                    random_pct_3 = gr.Number(label="Percentage (%)", value=50, minimum=0, maximum=100, precision=0, interactive=True)
+                                with gr.Row(visible=False) as every_k_params_3:
+                                    step_k_3 = gr.Number(label="k", value=2, minimum=1, precision=0, interactive=True)
+                                with gr.Row(visible=False) as part_params_3:
+                                    start_step_3 = gr.Number(label="Start (int=step; 0,0-1,0=ratio)", value=0, minimum=0, precision=2, interactive=True)
+                                    end_step_3 = gr.Number(label="End (int=step; 0,0-1,0=ratio)", value=0.5, minimum=0, precision=2, interactive=True)
+
+        def on_strategy_change(strategy):
+            return (
+                gr.update(visible=strategy == "Replace on random steps"),
+                gr.update(visible=strategy == "Replace every step k"),
+                gr.update(visible=strategy == "Replace part of trajectory"),
+            )
+
+        def toggle_routing_section():
+            return gr.update(visible=True)
+
+        add_routing_btn.click(
+            fn=toggle_routing_section,
+            outputs=[routing_section],
+        )
+
+        strategy_1.change(
+            fn=on_strategy_change,
+            inputs=[strategy_1],
+            outputs=[random_params_1, every_k_params_1, part_params_1],
+        )
+
+        strategy_2.change(
+            fn=on_strategy_change,
+            inputs=[strategy_2],
+            outputs=[random_params_2, every_k_params_2, part_params_2],
+        )
+
+        strategy_3.change(
+            fn=on_strategy_change,
+            inputs=[strategy_3],
+            outputs=[random_params_3, every_k_params_3, part_params_3],
+        )
+
+        def filter_models(query):
+            """Filter models based on search query (starts at 3 chars)"""
+            if not query or len(query) < 3:
+                return gr.update(choices=[])
+            all_models = get_litellm_model_list()
+            query_lower = query.lower()
+            filtered = [m for m in all_models if query_lower in m.lower()][:50]
+            return gr.update(choices=filtered)
+
+        routing_model_1.input(fn=filter_models, inputs=[routing_model_1], outputs=[routing_model_1])
+        routing_model_2.input(fn=filter_models, inputs=[routing_model_2], outputs=[routing_model_2])
+        routing_model_3.input(fn=filter_models, inputs=[routing_model_3], outputs=[routing_model_3])
+
+        def get_routing_prices_with_labels(model_name):
+            """Get all 4 prices for a routing model with found/estimated labels"""
+            if not model_name:
+                return (
+                    gr.update(value=0, label="Input"),
+                    gr.update(value=0, label="Cache Read"),
+                    gr.update(value=0, label="Cache Creation"),
+                    gr.update(value=0, label="Completion"),
+                )
+
+            prices = get_litellm_prices()
+            model_prices = prices.get(model_name, {})
+
+            input_price = model_prices.get("input_cost_per_token", 0) * 1e6
+            cache_read = model_prices.get("cache_read_input_token_cost", 0) * 1e6
+            cache_creation = model_prices.get("cache_creation_input_token_cost", 0) * 1e6
+            completion = model_prices.get("output_cost_per_token", 0) * 1e6
+
+            input_found = input_price > 0
+            cache_read_found = cache_read > 0
+            cache_creation_found = cache_creation > 0
+            completion_found = completion > 0
+
+            if not cache_read_found and input_price > 0:
+                cache_read = input_price * 0.1
+            if not cache_creation_found and input_price > 0:
+                cache_creation = input_price * 1.25
+
+            def label(name, found):
+                return f"✅ {name}" if found else f"❌ {name}"
+
+            return (
+                gr.update(value=input_price, label=label("Input", input_found)),
+                gr.update(value=cache_read, label=label("Cache Read", cache_read_found)),
+                gr.update(value=cache_creation, label=label("Cache Creation", cache_creation_found)),
+                gr.update(value=completion, label=label("Completion", completion_found)),
+            )
+
+        def on_routing_model_1_select(model_name):
+            prices = get_routing_prices_with_labels(model_name)
+            show_btn = bool(model_name)
+            return *prices, gr.update(visible=show_btn)
+
+        def on_routing_model_2_select(model_name):
+            prices = get_routing_prices_with_labels(model_name)
+            show_btn = bool(model_name)
+            return *prices, gr.update(visible=show_btn)
+
+        def on_routing_model_3_select(model_name):
+            return get_routing_prices_with_labels(model_name)
+
+        routing_model_1.change(
+            fn=on_routing_model_1_select,
+            inputs=[routing_model_1],
+            outputs=[routing_price_1_input, routing_price_1_cache_read, routing_price_1_cache_creation, routing_price_1_completion, add_model_2_btn],
+        )
+
+        add_model_2_btn.click(
+            fn=lambda: (gr.update(visible=True), gr.update(visible=False)),
+            outputs=[routing_block_2, add_model_2_btn],
+        )
+
+        routing_model_2.change(
+            fn=on_routing_model_2_select,
+            inputs=[routing_model_2],
+            outputs=[routing_price_2_input, routing_price_2_cache_read, routing_price_2_cache_creation, routing_price_2_completion, add_model_3_btn],
+        )
+
+        add_model_3_btn.click(
+            fn=lambda: (gr.update(visible=True), gr.update(visible=False)),
+            outputs=[routing_block_3, add_model_3_btn],
+        )
+
+        routing_model_3.change(
+            fn=on_routing_model_3_select,
+            inputs=[routing_model_3],
+            outputs=[routing_price_3_input, routing_price_3_cache_read, routing_price_3_cache_creation, routing_price_3_completion],
+        )
+
         def update_calculated_options_visibility(source):
             is_calc = source == "Calculated"
             return gr.update(visible=is_calc), gr.update(visible=is_calc)
@@ -999,30 +1255,47 @@ def build_app():
         leaderboard_table.select(
             fn=on_row_select,
             inputs=[leaderboard_table],
-            outputs=[selected_folder, selected_name, download_btn, analyze_btn, price_input, price_cache_read, price_cache_creation, price_completion, detected_model, thinking_overhead],
-        )
-
-        download_btn.click(
-            fn=download_trajectories_from_s3,
-            inputs=[selected_folder],
-            outputs=[download_status, analyze_btn],
+            outputs=[selected_folder, selected_name, analyze_btn, price_input, price_cache_read, price_cache_creation, price_completion, detected_model, thinking_overhead],
         )
 
-        def load_and_analyze(folder, input_price, cache_read_price, cache_creation_price, completion_price, source, overhead, with_cache):
+        def load_and_analyze(folder, input_price, cache_read_price, cache_creation_price, completion_price, source, overhead, with_cache, progress=gr.Progress()):
             empty_result = (
+                "",
                 gr.update(visible=False),
                 None, None, None, None, None, None,
                 None,
+                gr.update(visible=False),
             )
 
             if not folder:
                 yield empty_result
                 return
 
+            if not check_trajectories_downloaded(folder):
+                yield (
+                    "⏳ Downloading trajectories...",
+                    gr.update(visible=False),
+                    None, None, None, None, None, None,
+                    None,
+                    gr.update(visible=False),
+                )
+                status, _ = download_trajectories_from_s3(folder)
+                if "❌" in status:
+                    yield (
+                        status,
+                        gr.update(visible=False),
+                        None, None, None, None, None, None,
+                        None,
+                        gr.update(visible=False),
+                    )
+                    return
+
             yield (
+                "⏳ Loading trajectories...",
                 gr.update(visible=True),
                 None, None, None, None, None, None,
                 None,
+                gr.update(visible=False),
             )
 
             df_meta = load_all_trajectories(folder)
@@ -1040,7 +1313,13 @@ def build_app():
                     df = apply_no_cache(df)
 
             if df.empty:
-                yield empty_result
+                yield (
+                    "❌ No trajectories found",
+                    gr.update(visible=False),
+                    None, None, None, None, None, None,
+                    None,
+                    gr.update(visible=False),
+                )
                 return
 
             fig_steps, fig_cost, fig_tokens, fig_tokens_cost, fig_stacked = create_basic_histograms(
@@ -1049,18 +1328,22 @@ def build_app():
             fig_cost_breakdown = create_cost_breakdown(df, input_price, cache_read_price, cache_creation_price, completion_price)
 
             yield (
+                f"✅ Loaded {len(df)} trajectories",
                 gr.update(visible=True),
                 fig_steps, fig_cost, fig_tokens, fig_tokens_cost, fig_stacked, fig_cost_breakdown,
                 state_data,
+                gr.update(visible=True),
             )
 
         analyze_btn.click(
             fn=load_and_analyze,
             inputs=[selected_folder, price_input, price_cache_read, price_cache_creation, price_completion, token_source, thinking_overhead, use_cache],
             outputs=[
+                download_status,
                 analysis_section,
                 plot_steps, plot_cost, plot_tokens, plot_tokens_cost, plot_stacked, plot_cost_breakdown,
                 trajectories_state,
+                add_routing_btn,
             ],
         )