Merge pull request #39955 from BerriAI/litellm_fix_adaptive_router_bandit_prior

fix(adaptive_router): add the persisted delta to the cold-start prior on load

litellm/router_strategy/adaptive_router/adaptive_router.py

@@ -123,7 +123,12 @@ class AdaptiveRouter:
                 self._cells[(rt, model)] = initial_cell(prefs, rt)
 
     async def load_state_from_db(self, prisma_client: Any) -> None:
-        """Override cold-start cells with persisted state. Called once at startup."""
+        """Add each row's persisted delta to a freshly computed cold-start prior.
+
+        A row holds an accumulated delta, not a full posterior, and can be one-sided
+        (e.g. beta=0) - assigning it straight into the cell would zero out a Beta shape
+        parameter and crash thompson_sample() on every later draw for that cell.
+        """
         if prisma_client is None:
             return
         try:
@@ -139,7 +144,12 @@ class AdaptiveRouter:
                     continue
                 if row.model_name not in self.config.available_models:
                     continue
-                self._cells[(rt, row.model_name)] = BanditCell(alpha=row.alpha, beta=row.beta)
+                prefs = self.model_to_prefs.get(row.model_name) or _default_prefs()
+                prior = initial_cell(prefs, rt)
+                self._cells[(rt, row.model_name)] = BanditCell(
+                    alpha=prior.alpha + row.alpha,
+                    beta=prior.beta + row.beta,
+                )
                 loaded += 1
             verbose_router_logger.info(
                 "AdaptiveRouter[%s]: loaded %d cells from DB",

tests/test_litellm/router_strategy/adaptive_router/test_adaptive_router.py

@@ -269,7 +269,9 @@ async def test_record_turn_bounds_feedback_contexts_and_evicts_least_recent_sess
 
 
 @pytest.mark.asyncio
-async def test_load_state_from_db_overrides_cold_start():
+async def test_load_state_from_db_adds_the_persisted_delta_to_the_cold_start_prior():
+    """A row holds an accumulated delta, not a full posterior; loading must add it to the
+    cold-start prior, not replace the cell outright."""
     r = _make_router()
     cold = r._cells[(RequestType.GENERAL, "fast")]
 
@@ -284,14 +286,38 @@ async def test_load_state_from_db_overrides_cold_start():
     await r.load_state_from_db(prisma)
 
     new_cell = r._cells[(RequestType.GENERAL, "fast")]
-    assert (new_cell.alpha, new_cell.beta) == (42.0, 13.0)
-    assert (new_cell.alpha, new_cell.beta) != (cold.alpha, cold.beta)
+    assert (new_cell.alpha, new_cell.beta) == (cold.alpha + 42.0, cold.beta + 13.0)
+
+
+@pytest.mark.asyncio
+async def test_load_state_from_db_keeps_a_one_sided_delta_row_sampleable():
+    """A cell whose only DB activity is one signal type persists a one-sided row (e.g.
+    beta=0.0); loading it must not zero out a Beta shape parameter and crash thompson_sample()."""
+    from litellm.router_strategy.adaptive_router.bandit import thompson_sample
+
+    r = _make_router()
+
+    one_sided_row = MagicMock()
+    one_sided_row.request_type = "general"
+    one_sided_row.model_name = "fast"
+    one_sided_row.alpha = 1.0
+    one_sided_row.beta = 0.0
+
+    prisma = MagicMock()
+    prisma.db.litellm_adaptiverouterstate.find_many = AsyncMock(return_value=[one_sided_row])
+    await r.load_state_from_db(prisma)
+
+    loaded_cell = r._cells[(RequestType.GENERAL, "fast")]
+    assert loaded_cell.alpha > 0.0
+    assert loaded_cell.beta > 0.0
+    thompson_sample(loaded_cell)  # must not raise
 
 
 @pytest.mark.asyncio
 async def test_load_state_from_db_handles_unknown_request_type():
     r = _make_router()
-    cold = r._cells[(RequestType.GENERAL, "fast")]
+    cold_general = r._cells[(RequestType.GENERAL, "fast")]
+    cold_writing = r._cells[(RequestType.WRITING, "fast")]
 
     bad_row = MagicMock()
     bad_row.request_type = "nonexistent_type_v999"
@@ -309,10 +335,11 @@ async def test_load_state_from_db_handles_unknown_request_type():
     prisma.db.litellm_adaptiverouterstate.find_many = AsyncMock(return_value=[bad_row, good_row])
     await r.load_state_from_db(prisma)
 
-    # Unknown skipped; good applied.
-    assert r._cells[(RequestType.GENERAL, "fast")].alpha == 7.0
-    # Other request types kept their cold-start values.
-    assert r._cells[(RequestType.WRITING, "fast")] == cold or True
+    # Unknown skipped; good added to the cold-start prior.
+    new_general = r._cells[(RequestType.GENERAL, "fast")]
+    assert new_general.alpha == cold_general.alpha + 7.0
+    # Other request types kept their own cold-start values.
+    assert r._cells[(RequestType.WRITING, "fast")] == cold_writing
 
 
 # ---- Session state eviction ---------------------------------------------

tests/test_litellm/router_strategy/adaptive_router/test_e2e_adaptive_router.py

@@ -186,8 +186,12 @@ async def test_failure_signal_increments_beta_after_flush():
 
 
 @pytest.mark.asyncio
-async def test_load_state_from_db_overrides_cold_start():
+async def test_load_state_from_db_adds_persisted_delta_to_cold_start():
+    """A row holds an accumulated delta, not a full posterior; loading must add it to the
+    cold-start prior, not replace the cell outright."""
     router = _make_router()
+    cold = router._cells[(RequestType.GENERAL, "gpt-4o")]
+
     fake_row = MagicMock()
     fake_row.request_type = RequestType.GENERAL.value
     fake_row.model_name = "gpt-4o"
@@ -200,8 +204,8 @@ async def test_load_state_from_db_overrides_cold_start():
     await router.load_state_from_db(prisma)
 
     cell = router._cells[(RequestType.GENERAL, "gpt-4o")]
-    assert cell.alpha == 90.0
-    assert cell.beta == 10.0
+    assert cell.alpha == cold.alpha + 90.0
+    assert cell.beta == cold.beta + 10.0
 
 
 @pytest.mark.asyncio