[Fix] router - regression when adding/removing models

litellm/router.py

@@ -421,8 +421,7 @@ def __init__(  # noqa: PLR0915
         self.model_name_to_deployment_indices: Dict[str, List[int]] = {}
 
         if model_list is not None:
-            # Build model index immediately to enable O(1) lookups from the start
-            self._build_model_id_to_deployment_index_map(model_list)
+            # set_model_list will build indices automatically
             self.set_model_list(model_list)
             self.healthy_deployments: List = self.model_list  # type: ignore
             for m in model_list:
@@ -5143,8 +5142,8 @@ def set_model_list(self, model_list: list):
         )
         self.model_names = [m["model_name"] for m in model_list]
 
-        # Build model_name index for O(1) lookups
-        self._build_model_name_index(self.model_list)
+        # Note: model_name_to_deployment_indices is already built incrementally
+        # by _create_deployment -> _add_model_to_list_and_index_map
 
     def _add_deployment(self, deployment: Deployment) -> Deployment:
         import os
@@ -5367,6 +5366,26 @@ def _update_deployment_indices_after_removal(
         # Remove the deleted model from index
         if model_id in self.model_id_to_deployment_index_map:
             del self.model_id_to_deployment_index_map[model_id]
+
+        # Update model_name_to_deployment_indices
+        for model_name, indices in list(self.model_name_to_deployment_indices.items()):
+            # Remove the deleted index
+            if removal_idx in indices:
+                indices.remove(removal_idx)
+
+            # Decrement all indices greater than removal_idx
+            updated_indices = []
+            for idx in indices:
+                if idx > removal_idx:
+                    updated_indices.append(idx - 1)
+                else:
+                    updated_indices.append(idx)
+
+            # Update or remove the entry
+            if len(updated_indices) > 0:
+                self.model_name_to_deployment_indices[model_name] = updated_indices
+            else:
+                del self.model_name_to_deployment_indices[model_name]
 
     def _add_model_to_list_and_index_map(
         self, model: dict, model_id: Optional[str] = None

tests/router_unit_tests/test_router_index_management.py

@@ -16,25 +16,38 @@ def router(self):
         """Create a router instance for testing"""
         return Router(model_list=[])
 
-    def test_update_deployment_indices_after_removal(self, router):
-        """Test _update_deployment_indices_after_removal function"""
-        # Setup: Add models to router with proper structure
+    def test_deletion_updates_model_name_indices(self, router):
+        """Test that deleting a deployment updates model_name_to_deployment_indices correctly"""
         router.model_list = [
-            {"model": "test1", "model_info": {"id": "model-1"}}, 
-            {"model": "test2", "model_info": {"id": "model-2"}}, 
-            {"model": "test3", "model_info": {"id": "model-3"}}
+            {"model_name": "gpt-3.5", "model_info": {"id": "model-1"}},
+            {"model_name": "gpt-4", "model_info": {"id": "model-2"}},
+            {"model_name": "gpt-4", "model_info": {"id": "model-3"}},
+            {"model_name": "claude", "model_info": {"id": "model-4"}}
         ]
-        router.model_id_to_deployment_index_map = {"model-1": 0, "model-2": 1, "model-3": 2}
-
-        # Test: Remove model-2 (index 1)
+        router.model_id_to_deployment_index_map = {
+            "model-1": 0, "model-2": 1, "model-3": 2, "model-4": 3
+        }
+        router.model_name_to_deployment_indices = {
+            "gpt-3.5": [0],
+            "gpt-4": [1, 2],
+            "claude": [3]
+        }
+
+        # Remove one of the duplicate gpt-4 deployments
         router._update_deployment_indices_after_removal(model_id="model-2", removal_idx=1)
 
-        # Verify: model-2 is removed from index
-        assert "model-2" not in router.model_id_to_deployment_index_map
-        # Verify: model-3 index is updated (2 -> 1)
-        assert router.model_id_to_deployment_index_map["model-3"] == 1
-        # Verify: model-1 index remains unchanged
-        assert router.model_id_to_deployment_index_map["model-1"] == 0
+        # Verify indices are shifted correctly
+        assert router.model_name_to_deployment_indices["gpt-3.5"] == [0]
+        assert router.model_name_to_deployment_indices["gpt-4"] == [1]  # was [1,2], removed 1, shifted 2->1
+        assert router.model_name_to_deployment_indices["claude"] == [2]  # was [3], shifted to [2]
+
+        # Remove the last gpt-4 deployment
+        router._update_deployment_indices_after_removal(model_id="model-3", removal_idx=1)
+
+        # Verify gpt-4 is removed from dict when no deployments remain
+        assert "gpt-4" not in router.model_name_to_deployment_indices
+        assert router.model_name_to_deployment_indices["gpt-3.5"] == [0]
+        assert router.model_name_to_deployment_indices["claude"] == [1]
 
     def test_build_model_id_to_deployment_index_map(self, router):
         """Test _build_model_id_to_deployment_index_map function"""