diff --git a/litellm/router_strategy/adaptive_router/adaptive_router.py b/litellm/router_strategy/adaptive_router/adaptive_router.py index 1a33ea23bd4..12ccacbbc1d 100644 --- a/litellm/router_strategy/adaptive_router/adaptive_router.py +++ b/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", diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_adaptive_router.py b/tests/test_litellm/router_strategy/adaptive_router/test_adaptive_router.py index cbf5635a5ae..f36443db1e5 100644 --- a/tests/test_litellm/router_strategy/adaptive_router/test_adaptive_router.py +++ b/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 --------------------------------------------- diff --git a/tests/test_litellm/router_strategy/adaptive_router/test_e2e_adaptive_router.py b/tests/test_litellm/router_strategy/adaptive_router/test_e2e_adaptive_router.py index 3071f916ef1..23fc859d4a6 100644 --- a/tests/test_litellm/router_strategy/adaptive_router/test_e2e_adaptive_router.py +++ b/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