Refactor routing strategy UI

- Replace per-model strategy with global Router Strategy
- Add three strategies: Random weights, Every k-th step, Replace part of trajectory
- Random weights: each step randomly assigned based on weights (must sum to 1.0)
- Every k-th step: first model has priority on overlaps
- Replace part: non-overlapping ranges for each model
- Simplify UI: remove containers, use individual visible components
- Fix strategy parameter visibility on strategy change

app.py

@@ -1,5 +1,6 @@
 import json
 import os
+import random
 import re
 import subprocess
 from pathlib import Path
@@ -29,53 +30,6 @@ _calculated_tokens_cache = {}
 _trajectory_steps_cache = {}
 
 
-def parse_start_end(start: float, end: float, total_steps: int, mode: str) -> tuple[int, int]:
-    """
-    Parse start and end values based on mode.
-
-    Args:
-        start: start value
-        end: end value
-        total_steps: total number of steps in trajectory
-        mode: "Indexes" or "Percentages"
-
-    Returns: (start_idx, end_idx) - both 0-based
-    """
-    if mode == "Indexes":
-        return int(start), int(end)
-    else:
-        return int(start * total_steps / 100), int(end * total_steps / 100)
-
-
-def get_routed_steps(total_steps: int, strategy: str, params: dict) -> set:
-    """
-    Determine which steps should be routed to alternative model.
-
-    Returns set of step indices (0-based) that should use the routing model.
-    """
-    import random
-
-    routed = set()
-
-    if strategy == "Replace on random steps":
-        pct = params.get("percentage", 50) / 100.0
-        num_to_route = int(total_steps * pct)
-        if num_to_route > 0:
-            routed = set(random.sample(range(total_steps), min(num_to_route, total_steps)))
-
-    elif strategy == "Replace every step k":
-        k = int(params.get("k", 2))
-        if k > 0:
-            routed = set(range(0, total_steps, k))
-
-    elif strategy == "Replace part of trajectory":
-        mode = params.get("mode", "Percentages")
-        start, end = parse_start_end(params.get("start", 0), params.get("end", 30), total_steps, mode)
-        routed = set(range(start, min(end, total_steps)))
-
-    return routed
-
-
 def calculate_routing_tokens(steps: list[dict]) -> dict:
     """
     Calculate token breakdown per model with proper caching simulation.
@@ -1275,12 +1229,6 @@ def build_app():
                 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")
@@ -1295,26 +1243,6 @@ def build_app():
                                 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.Column(visible=False) as part_params_1:
-                                part_mode_1 = gr.Radio(
-                                    choices=["Indexes", "Percentages"],
-                                    value="Percentages",
-                                    label="Mode",
-                                    interactive=True,
-                                )
-                                with gr.Row():
-                                    start_step_1 = gr.Number(label="Start", value=0, minimum=0, precision=0, interactive=True)
-                                    end_step_1 = gr.Number(label="End", value=30, minimum=0, precision=0, interactive=True)
 
                         add_model_2_btn = gr.Button("+ Add another model", size="sm", visible=False)
 
@@ -1332,26 +1260,6 @@ def build_app():
                                     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.Column(visible=False) as part_params_2:
-                                    part_mode_2 = gr.Radio(
-                                        choices=["Indexes", "Percentages"],
-                                        value="Percentages",
-                                        label="Mode",
-                                        interactive=True,
-                                    )
-                                    with gr.Row():
-                                        start_step_2 = gr.Number(label="Start", value=0, minimum=0, precision=0, interactive=True)
-                                        end_step_2 = gr.Number(label="End", value=30, minimum=0, precision=0, interactive=True)
 
                             add_model_3_btn = gr.Button("+ Add another model", size="sm", visible=False)
 
@@ -1369,39 +1277,48 @@ def build_app():
                                     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.Column(visible=False) as part_params_3:
-                                    part_mode_3 = gr.Radio(
-                                        choices=["Indexes", "Percentages"],
-                                        value="Percentages",
-                                        label="Mode",
-                                        interactive=True,
-                                    )
-                                    with gr.Row():
-                                        start_step_3 = gr.Number(label="Start", value=0, minimum=0, precision=0, interactive=True)
-                                        end_step_3 = gr.Number(label="End", value=30, minimum=0, precision=0, interactive=True)
+
+                    gr.Markdown("---")
+                    gr.Markdown("### 🎯 Router Strategy")
+
+                    selected_strategy = gr.Radio(
+                        choices=["Random weights", "Every k-th step", "Replace part of trajectory"],
+                        value="Random weights",
+                        label="Strategy",
+                        interactive=True,
+                    )
+
+                    random_hint = gr.Markdown("*Weights must sum to 1.0*", visible=True)
+                    weight_base = gr.Number(label="Base weight", value=0.5, minimum=0, maximum=1, precision=2, interactive=True, visible=True)
+                    weight_model_1 = gr.Number(label="Model 1 weight", value=0.5, minimum=0, maximum=1, precision=2, interactive=True, visible=True)
+                    weight_model_2 = gr.Number(label="Model 2 weight", value=0, minimum=0, maximum=1, precision=2, interactive=True, visible=False)
+                    weight_model_3 = gr.Number(label="Model 3 weight", value=0, minimum=0, maximum=1, precision=2, interactive=True, visible=False)
+
+                    every_k_hint = gr.Markdown("*First model has priority on overlaps*", visible=False)
+                    k_model_1 = gr.Number(label="k₁ (Model 1)", value=2, minimum=1, precision=0, interactive=True, visible=False)
+                    k_model_2 = gr.Number(label="k₂ (Model 2)", value=3, minimum=1, precision=0, interactive=True, visible=False)
+                    k_model_3 = gr.Number(label="k₃ (Model 3)", value=5, minimum=1, precision=0, interactive=True, visible=False)
+
+                    part_hint = gr.Markdown("*Ranges must not overlap*", visible=False)
+                    part_mode = gr.Radio(
+                        choices=["Indexes", "Percentages"],
+                        value="Percentages",
+                        label="Mode",
+                        interactive=True,
+                        visible=False,
+                    )
+                    start_1 = gr.Number(label="M1 Start", value=0, minimum=0, precision=0, interactive=True, visible=False)
+                    end_1 = gr.Number(label="M1 End", value=30, minimum=0, precision=0, interactive=True, visible=False)
+                    start_2 = gr.Number(label="M2 Start", value=30, minimum=0, precision=0, interactive=True, visible=False)
+                    end_2 = gr.Number(label="M2 End", value=60, minimum=0, precision=0, interactive=True, visible=False)
+                    start_3 = gr.Number(label="M3 Start", value=60, minimum=0, precision=0, interactive=True, visible=False)
+                    end_3 = gr.Number(label="M3 End", value=100, minimum=0, precision=0, interactive=True, visible=False)
 
                     gr.Markdown("---")
                     route_btn = gr.Button("🚀 Let's ROUTE!!", variant="primary", size="lg", interactive=False)
                     routing_result = gr.Markdown(visible=False)
 
 
-        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)
 
@@ -1410,22 +1327,31 @@ def build_app():
             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],
-        )
+        def on_strategy_change(strategy):
+            is_random = strategy == "Random weights"
+            is_every_k = strategy == "Every k-th step"
+            is_part = strategy == "Replace part of trajectory"
+            print(f"DEBUG on_strategy_change: strategy={strategy}")
+            return (
+                gr.update(visible=is_random),
+                gr.update(visible=is_random),
+                gr.update(visible=is_random),
+                gr.update(visible=is_every_k),
+                gr.update(visible=is_every_k),
+                gr.update(visible=is_part),
+                gr.update(visible=is_part),
+                gr.update(visible=is_part),
+                gr.update(visible=is_part),
+            )
 
-        strategy_3.change(
+        selected_strategy.change(
             fn=on_strategy_change,
-            inputs=[strategy_3],
-            outputs=[random_params_3, every_k_params_3, part_params_3],
+            inputs=[selected_strategy],
+            outputs=[
+                random_hint, weight_base, weight_model_1,
+                every_k_hint, k_model_1,
+                part_hint, part_mode, start_1, end_1,
+            ],
         )
 
         def filter_models(query):
@@ -1498,9 +1424,23 @@ def build_app():
             outputs=[routing_price_1_input, routing_price_1_cache_read, routing_price_1_cache_creation, routing_price_1_completion, add_model_2_btn, route_btn],
         )
 
+        def show_model_2(strategy):
+            is_random = strategy == "Random weights"
+            is_every_k = strategy == "Every k-th step"
+            is_part = strategy == "Replace part of trajectory"
+            return (
+                gr.update(visible=True),
+                gr.update(visible=False),
+                gr.update(visible=is_random),
+                gr.update(visible=is_every_k),
+                gr.update(visible=is_part),
+                gr.update(visible=is_part),
+            )
+
         add_model_2_btn.click(
-            fn=lambda: (gr.update(visible=True), gr.update(visible=False)),
-            outputs=[routing_block_2, add_model_2_btn],
+            fn=show_model_2,
+            inputs=[selected_strategy],
+            outputs=[routing_block_2, add_model_2_btn, weight_model_2, k_model_2, start_2, end_2],
         )
 
         routing_model_2.change(
@@ -1509,9 +1449,23 @@ def build_app():
             outputs=[routing_price_2_input, routing_price_2_cache_read, routing_price_2_cache_creation, routing_price_2_completion, add_model_3_btn],
         )
 
+        def show_model_3(strategy):
+            is_random = strategy == "Random weights"
+            is_every_k = strategy == "Every k-th step"
+            is_part = strategy == "Replace part of trajectory"
+            return (
+                gr.update(visible=True),
+                gr.update(visible=False),
+                gr.update(visible=is_random),
+                gr.update(visible=is_every_k),
+                gr.update(visible=is_part),
+                gr.update(visible=is_part),
+            )
+
         add_model_3_btn.click(
-            fn=lambda: (gr.update(visible=True), gr.update(visible=False)),
-            outputs=[routing_block_3, add_model_3_btn],
+            fn=show_model_3,
+            inputs=[selected_strategy],
+            outputs=[routing_block_3, add_model_3_btn, weight_model_3, k_model_3, start_3, end_3],
         )
 
         routing_model_3.change(
@@ -1524,11 +1478,12 @@ def build_app():
             state_data,
             base_input, base_cache_read, base_cache_creation, base_completion,
             routing_model_1_val, r1_input, r1_cache_read, r1_cache_creation, r1_completion,
-            strategy_1_val, random_pct_1_val, step_k_1_val, part_mode_1_val, start_1_val, end_1_val,
             routing_model_2_val, r2_input, r2_cache_read, r2_cache_creation, r2_completion,
-            strategy_2_val, random_pct_2_val, step_k_2_val, part_mode_2_val, start_2_val, end_2_val,
             routing_model_3_val, r3_input, r3_cache_read, r3_cache_creation, r3_completion,
-            strategy_3_val, random_pct_3_val, step_k_3_val, part_mode_3_val, start_3_val, end_3_val,
+            strategy_val,
+            weight_base_val, weight_1_val, weight_2_val, weight_3_val,
+            k_1_val, k_2_val, k_3_val,
+            part_mode_val, start_1_val, end_1_val, start_2_val, end_2_val, start_3_val, end_3_val,
             source, overhead, with_cache
         ):
             if state_data is None:
@@ -1556,6 +1511,7 @@ def build_app():
                 )
                 return
 
+
             df_calc = state_data.get("calculated")
             if df_calc is not None and not df_calc.empty:
                 df_for_cost = apply_thinking_overhead(df_calc.copy(), overhead)
@@ -1579,50 +1535,56 @@ def build_app():
                 "completion": base_completion,
             }
 
-            def build_strategy_params(strategy, random_pct, step_k, part_mode, start_val, end_val):
-                params = {}
-                if strategy == "Replace on random steps":
-                    params["percentage"] = random_pct
-                elif strategy == "Replace every step k":
-                    params["k"] = step_k
-                elif strategy == "Replace part of trajectory":
-                    params["mode"] = part_mode
-                    params["start"] = start_val
-                    params["end"] = end_val
-                return params
-
             routing_models = []
             if routing_model_1_val:
-                if strategy_1_val == "Replace part of trajectory" and start_1_val >= end_1_val:
-                    yield (gr.update(visible=True, value="❌ Model 1: Start must be less than End"), gr.update(visible=False), None, None)
-                    return
                 routing_models.append({
                     "name": routing_model_1_val,
                     "prices": {"input": r1_input, "cache_read": r1_cache_read, "cache_creation": r1_cache_creation, "completion": r1_completion},
-                    "strategy": strategy_1_val,
-                    "params": build_strategy_params(strategy_1_val, random_pct_1_val, step_k_1_val, part_mode_1_val, start_1_val, end_1_val),
                 })
             if routing_model_2_val:
-                if strategy_2_val == "Replace part of trajectory" and start_2_val >= end_2_val:
-                    yield (gr.update(visible=True, value="❌ Model 2: Start must be less than End"), gr.update(visible=False), None, None)
-                    return
                 routing_models.append({
                     "name": routing_model_2_val,
                     "prices": {"input": r2_input, "cache_read": r2_cache_read, "cache_creation": r2_cache_creation, "completion": r2_completion},
-                    "strategy": strategy_2_val,
-                    "params": build_strategy_params(strategy_2_val, random_pct_2_val, step_k_2_val, part_mode_2_val, start_2_val, end_2_val),
                 })
             if routing_model_3_val:
-                if strategy_3_val == "Replace part of trajectory" and start_3_val >= end_3_val:
-                    yield (gr.update(visible=True, value="❌ Model 3: Start must be less than End"), gr.update(visible=False), None, None)
-                    return
                 routing_models.append({
                     "name": routing_model_3_val,
                     "prices": {"input": r3_input, "cache_read": r3_cache_read, "cache_creation": r3_cache_creation, "completion": r3_completion},
-                    "strategy": strategy_3_val,
-                    "params": build_strategy_params(strategy_3_val, random_pct_3_val, step_k_3_val, part_mode_3_val, start_3_val, end_3_val),
                 })
 
+            if strategy_val == "Replace part of trajectory":
+                ranges = [(start_1_val, end_1_val)]
+                if len(routing_models) > 1:
+                    ranges.append((start_2_val, end_2_val))
+                if len(routing_models) > 2:
+                    ranges.append((start_3_val, end_3_val))
+                for i, (s, e) in enumerate(ranges):
+                    if s >= e:
+                        yield (gr.update(visible=True, value=f"❌ Model {i+1}: Start must be less than End"), gr.update(visible=False), None, None)
+                        return
+                for i in range(len(ranges)):
+                    for j in range(i+1, len(ranges)):
+                        s1, e1 = ranges[i]
+                        s2, e2 = ranges[j]
+                        if not (e1 <= s2 or e2 <= s1):
+                            yield (gr.update(visible=True, value=f"❌ Model {i+1} and Model {j+1} ranges overlap"), gr.update(visible=False), None, None)
+                            return
+
+            weights = None
+            if strategy_val == "Random weights":
+                weights = [weight_base_val, weight_1_val]
+                if len(routing_models) > 1:
+                    weights.append(weight_2_val)
+                if len(routing_models) > 2:
+                    weights.append(weight_3_val)
+                total_weight = sum(weights)
+                if abs(total_weight - 1.0) > 0.01:
+                    yield (gr.update(visible=True, value=f"❌ Weights must sum to 1.0 (current: {total_weight:.2f})"), gr.update(visible=False), None, None)
+                    return
+
+            k_values = [k_1_val, k_2_val, k_3_val][:len(routing_models)]
+            part_ranges = [(start_1_val, end_1_val), (start_2_val, end_2_val), (start_3_val, end_3_val)][:len(routing_models)]
+
             BASE_MODEL = "__base__"
             model_keys = [BASE_MODEL] + [f"__routing_{i}__" for i in range(len(routing_models))]
 
@@ -1635,17 +1597,35 @@ def build_app():
 
                 total_steps = len(steps)
 
-                routed_sets = []
-                for rm in routing_models:
-                    routed_sets.append(get_routed_steps(total_steps, rm["strategy"], rm["params"]))
+                step_to_model = {}
+
+                if strategy_val == "Random weights":
+                    model_choices = [BASE_MODEL] + [f"__routing_{j}__" for j in range(len(routing_models))]
+                    for i in range(total_steps):
+                        step_to_model[i] = random.choices(model_choices, weights=weights)[0]
+
+                elif strategy_val == "Every k-th step":
+                    for j, k_val in enumerate(k_values):
+                        if k_val and k_val > 0:
+                            for i in range(total_steps):
+                                if (i + 1) % int(k_val) == 0:
+                                    if i not in step_to_model:
+                                        step_to_model[i] = f"__routing_{j}__"
+
+                elif strategy_val == "Replace part of trajectory":
+                    for j, (start_val, end_val) in enumerate(part_ranges):
+                        if part_mode_val == "Percentages":
+                            start_idx = int(total_steps * start_val / 100)
+                            end_idx = int(total_steps * end_val / 100)
+                        else:
+                            start_idx = int(start_val)
+                            end_idx = min(int(end_val), total_steps)
+                        for i in range(start_idx, end_idx):
+                            step_to_model[i] = f"__routing_{j}__"
 
                 modified_steps = []
                 for i, step in enumerate(steps):
-                    model = BASE_MODEL
-                    for j, routed_set in enumerate(routed_sets):
-                        if i in routed_set:
-                            model = f"__routing_{j}__"
-                            break
+                    model = step_to_model.get(i, BASE_MODEL)
                     modified_steps.append({
                         "model": model,
                         "system_user": step.get("system_user", 0),
@@ -1752,11 +1732,12 @@ def build_app():
                 trajectories_state,
                 price_input, price_cache_read, price_cache_creation, price_completion,
                 routing_model_1, routing_price_1_input, routing_price_1_cache_read, routing_price_1_cache_creation, routing_price_1_completion,
-                strategy_1, random_pct_1, step_k_1, part_mode_1, start_step_1, end_step_1,
                 routing_model_2, routing_price_2_input, routing_price_2_cache_read, routing_price_2_cache_creation, routing_price_2_completion,
-                strategy_2, random_pct_2, step_k_2, part_mode_2, start_step_2, end_step_2,
                 routing_model_3, routing_price_3_input, routing_price_3_cache_read, routing_price_3_cache_creation, routing_price_3_completion,
-                strategy_3, random_pct_3, step_k_3, part_mode_3, start_step_3, end_step_3,
+                selected_strategy,
+                weight_base, weight_model_1, weight_model_2, weight_model_3,
+                k_model_1, k_model_2, k_model_3,
+                part_mode, start_1, end_1, start_2, end_2, start_3, end_3,
                 token_source, thinking_overhead, use_cache,
             ],
             outputs=[routing_result, routing_plots_row, routing_tokens_plot, routing_cost_plot],