fix: reconcile runtime pre-call checks

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/router.py

@@ -355,6 +355,13 @@ _PreRoutingStrategyT = TypeVar("_PreRoutingStrategyT")
 _ALIAS_PARAMS_NEVER_FORWARDED: Final = frozenset({"model", "api_base", "api_key", "api_version"})
 _ALIAS_MARKER_FORWARDED_PARAMS_KWARG: Final = "_alias_marker_forwarded_params"
 
+_RUNTIME_TOGGLEABLE_PRE_CALL_CHECKS: Final[Mapping[str, type[CustomLogger]]] = MappingProxyType(
+    {
+        "prompt_caching": PromptCachingDeploymentCheck,
+        "enforce_model_rate_limits": ModelRateLimitingCheck,
+    }
+)
+
 
 def _stream_chunks_have_generated_content(chunks: Sequence[ModelResponseStream]) -> bool:
     for chunk in chunks:
@@ -2082,6 +2089,27 @@ class Router:
             self.optional_callbacks.append(_callback)
             litellm.logging_callback_manager.add_litellm_callback(_callback)
 
+    def set_optional_pre_call_checks(self, optional_pre_call_checks: OptionalPreCallChecks | None) -> None:
+        if optional_pre_call_checks is None:
+            return
+        requested: Final = frozenset(optional_pre_call_checks)
+        for name, callback_cls in _RUNTIME_TOGGLEABLE_PRE_CALL_CHECKS.items():
+            if name not in requested:
+                self._remove_optional_callbacks_of_type(callback_cls)
+        self.add_optional_pre_call_checks(optional_pre_call_checks)
+
+    def _remove_optional_callbacks_of_type(self, callback_cls: type[CustomLogger]) -> None:
+        if self.optional_callbacks is None:
+            return
+        removed: Final = [cb for cb in self.optional_callbacks if isinstance(cb, callback_cls)]
+        if not removed:
+            return
+        self.optional_callbacks = [cb for cb in self.optional_callbacks if not isinstance(cb, callback_cls)]
+        for cb in removed:
+            litellm.logging_callback_manager.remove_callback_from_list_by_object(
+                litellm.callbacks, cb, require_self=False
+            )
+
     def print_deployment(self, deployment: dict):
         """
         returns a copy of the deployment with the api key masked
@@ -11356,7 +11384,7 @@ class Router:
                     self._routing_groups_input = kwargs[var]
                     rebuild_routing_groups = True
                 elif var == "optional_pre_call_checks":
-                    self.add_optional_pre_call_checks(kwargs[var])
+                    self.set_optional_pre_call_checks(kwargs[var])
                 elif var == "retry_policy":
                     value = kwargs[var]
                     if isinstance(value, dict):

tests/test_litellm/test_router_retry_policy_update.py

@@ -28,6 +28,8 @@ from pydantic import ValidationError
 
 
 import litellm
+from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
+from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck
 from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import PromptCachingDeploymentCheck
 from litellm.types.router import RetryPolicy, UpdateRouterConfig
 
@@ -114,6 +116,44 @@ def test_update_settings_adds_optional_pre_call_check_once():
     assert router.num_retries == 7
 
 
+def test_update_settings_clears_omitted_toggleable_pre_call_checks():
+    router = _build_router()
+
+    router.update_settings(optional_pre_call_checks=["prompt_caching"])
+    router.update_settings(optional_pre_call_checks=[])
+
+    assert not any(isinstance(callback, PromptCachingDeploymentCheck) for callback in (router.optional_callbacks or []))
+    assert not any(isinstance(callback, PromptCachingDeploymentCheck) for callback in litellm.callbacks)
+
+
+def test_update_settings_replaces_toggleable_pre_call_checks():
+    router = _build_router()
+
+    router.update_settings(optional_pre_call_checks=["prompt_caching"])
+    router.update_settings(optional_pre_call_checks=["enforce_model_rate_limits"])
+
+    assert not any(isinstance(callback, PromptCachingDeploymentCheck) for callback in (router.optional_callbacks or []))
+    assert not any(isinstance(callback, PromptCachingDeploymentCheck) for callback in litellm.callbacks)
+    assert any(isinstance(callback, ModelRateLimitingCheck) for callback in (router.optional_callbacks or []))
+
+
+@pytest.mark.asyncio
+async def test_update_settings_preserves_router_budget_limiting_when_omitted(monkeypatch):
+    async def _disable_periodic_sync(*args, **kwargs):
+        return None
+
+    monkeypatch.setattr(
+        "litellm.router_strategy.budget_limiter.RouterBudgetLimiting.periodic_sync_in_memory_spend_with_redis",
+        _disable_periodic_sync,
+    )
+    router = _build_router()
+
+    router.add_optional_pre_call_checks(["router_budget_limiting"])
+    router.update_settings(optional_pre_call_checks=[])
+
+    assert any(isinstance(callback, RouterBudgetLimiting) for callback in (router.optional_callbacks or []))
+
+
 def test_update_settings_persists_retry_policy_dict():
     """When the proxy's ``_add_router_settings_from_db_config`` calls
     ``llm_router.update_settings(retry_policy={...})`` after reading the