litellm/tests/router_unit_tests/test_router_index_management.py

import sys
import os
import pytest
import ast
import ast

sys.path.insert(
    0, os.path.abspath("../..")
)  # Adds the parent directory to the system path
from litellm import Router


class TestRouterIndexManagement:
    """Test cases for router index management functions"""

    @pytest.fixture
    def router(self):
        """Create a router instance for testing"""
        return Router(model_list=[])

    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_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, "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 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"""
        model_list = [
            {
                "model_name": "gpt-3.5-turbo",
                "litellm_params": {"model": "gpt-3.5-turbo"},
                "model_info": {"id": "model-1"},
            },
            {
                "model_name": "gpt-4",
                "litellm_params": {"model": "gpt-4"},
                "model_info": {"id": "model-2"},
            },
        ]

        # Test: Build index from model list
        router._build_model_id_to_deployment_index_map(model_list)

        # Verify: model_list is populated
        assert len(router.model_list) == 2
        # Verify: model_id_to_deployment_index_map is correctly built
        assert router.model_id_to_deployment_index_map["model-1"] == 0
        assert router.model_id_to_deployment_index_map["model-2"] == 1

    def test_add_model_to_list_and_index_map_from_model_info(self, router):
        """Test _add_model_to_list_and_index_map extracting model_id from model_info"""
        # Setup: Empty router
        router.model_list = []
        router.model_id_to_deployment_index_map = {}
        
        # Test: Add model without explicit model_id
        model = {"model": "test-model", "model_info": {"id": "model-info-id"}}
        router._add_model_to_list_and_index_map(model=model)
        
        # Verify: Model added to list
        assert len(router.model_list) == 1
        assert router.model_list[0] == model
        
        # Verify: Index map uses model_info.id
        assert router.model_id_to_deployment_index_map["model-info-id"] == 0

    def test_add_model_to_list_and_index_map_multiple_models(self, router):
        """Test _add_model_to_list_and_index_map with multiple models to verify indexing"""
        # Setup: Empty router
        router.model_list = []
        router.model_id_to_deployment_index_map = {}
        
        # Test: Add multiple models
        model1 = {"model": "model1", "model_info": {"id": "id-1"}}
        model2 = {"model": "model2", "model_info": {"id": "id-2"}}
        model3 = {"model": "model3", "model_info": {"id": "id-3"}}
        
        router._add_model_to_list_and_index_map(model=model1, model_id="id-1")
        router._add_model_to_list_and_index_map(model=model2, model_id="id-2")
        router._add_model_to_list_and_index_map(model=model3, model_id="id-3")
        
        # Verify: All models added to list
        assert len(router.model_list) == 3
        assert router.model_list[0] == model1
        assert router.model_list[1] == model2
        assert router.model_list[2] == model3
        
        # Verify: Correct indices in map
        assert router.model_id_to_deployment_index_map["id-1"] == 0
        assert router.model_id_to_deployment_index_map["id-2"] == 1
        assert router.model_id_to_deployment_index_map["id-3"] == 2

    def test_has_model_id(self, router):
        """Test has_model_id function for O(1) membership check"""
        # Setup: Add models to router
        router.model_list = [
            {"model": "test1", "model_info": {"id": "model-1"}}, 
            {"model": "test2", "model_info": {"id": "model-2"}}, 
            {"model": "test3", "model_info": {"id": "model-3"}}
        ]
        router.model_id_to_deployment_index_map = {"model-1": 0, "model-2": 1, "model-3": 2}

        # Test: Check existing model IDs
        assert router.has_model_id("model-1") == True
        assert router.has_model_id("model-2") == True
        assert router.has_model_id("model-3") == True

        # Test: Check non-existing model IDs
        assert router.has_model_id("non-existent") == False
        assert router.has_model_id("") == False
        assert router.has_model_id("model-4") == False

        # Test: Empty router
        empty_router = Router(model_list=[])
        assert empty_router.has_model_id("any-id") == False

    def test_build_model_name_index(self, router):
        """Test _build_model_name_index function"""
        model_list = [
            {
                "model_name": "gpt-3.5-turbo",
                "litellm_params": {"model": "gpt-3.5-turbo"},
                "model_info": {"id": "model-1"},
            },
            {
                "model_name": "gpt-4",
                "litellm_params": {"model": "gpt-4"},
                "model_info": {"id": "model-2"},
            },
            {
                "model_name": "gpt-4",  # Duplicate model_name, different deployment
                "litellm_params": {"model": "gpt-4"},
                "model_info": {"id": "model-3"},
            },
        ]

        # Test: Build index from model list
        router._build_model_name_index(model_list)

        # Verify: model_name_to_deployment_indices is correctly built
        assert "gpt-3.5-turbo" in router.model_name_to_deployment_indices
        assert "gpt-4" in router.model_name_to_deployment_indices
        
        # Verify: gpt-3.5-turbo has single deployment
        assert router.model_name_to_deployment_indices["gpt-3.5-turbo"] == [0]
        
        # Verify: gpt-4 has multiple deployments
        assert router.model_name_to_deployment_indices["gpt-4"] == [1, 2]
        
        # Test: Rebuild index (should clear and rebuild)
        new_model_list = [
            {
                "model_name": "claude-3",
                "litellm_params": {"model": "claude-3"},
                "model_info": {"id": "model-4"},
            },
        ]
        router._build_model_name_index(new_model_list)
        
        # Verify: Old entries are cleared
        assert "gpt-3.5-turbo" not in router.model_name_to_deployment_indices
        assert "gpt-4" not in router.model_name_to_deployment_indices
        
        # Verify: New entry is added
        assert "claude-3" in router.model_name_to_deployment_indices
        assert router.model_name_to_deployment_indices["claude-3"] == [0]

    def test_no_linear_scans_in_router(self):
        """
        Static analysis test to ensure Router doesn't use O(n) linear scans.
        
        Scans router.py for 'in self.model_list' pattern which indicates
        inefficient O(n) iteration instead of using index-based O(1) lookups.
        
        Methods should use:
        - model_id_to_deployment_index_map for O(1) model_id lookups
        - model_name_to_deployment_indices for O(1) + O(k) model_name lookups
        """
        # Methods that are allowed to iterate through self.model_list
        ALLOWED_METHODS = [
            "_get_deployment_by_litellm_model",  # Edge case: lookup by litellm_params.model (not indexed)
        ]
        
        # Get path to router.py
        router_file = os.path.join(
            os.path.dirname(os.path.dirname(os.path.dirname(__file__))),
            "litellm",
            "router.py"
        )
        
        # Read the file
        with open(router_file, 'r') as f:
            content = f.read()
        
        # Parse with AST
        tree = ast.parse(content)
        
        # Find violations
        violations = []
        ignore_methods = set(ALLOWED_METHODS)
        
        for node in ast.walk(tree):
            if isinstance(node, ast.FunctionDef):
                method_name = node.name
                
                # Skip ignored methods
                if method_name in ignore_methods:
                    continue
                
                # Get source for this method
                try:
                    method_source = ast.get_source_segment(content, node)
                    if not method_source:
                        continue
                    
                    # Check for the anti-pattern: "in self.model_list"
                    # This catches: for x in self.model_list, if x in self.model_list, etc.
                    if "in self.model_list" in method_source:
                        # Extract the specific line for better error reporting
                        lines = method_source.split('\n')
                        pattern_line = None
                        for line in lines:
                            if "in self.model_list" in line:
                                pattern_line = line.strip()
                                break
                        
                        violations.append({
                            "method": method_name,
                            "line": node.lineno,
                            "pattern": pattern_line or "in self.model_list"
                        })
                except Exception:
                    # Skip if we can't get source segment
                    pass
        
        # Assert no violations
        if violations:
            error_msg = "\n".join([
                f"  - {v['method']}() at line {v['line']}: {v['pattern']}"
                for v in violations
            ])
            
            pytest.fail(
                f"\n{'='*70}\n"
                f"Found O(n) linear scan pattern in router.py:\n\n"
                f"{error_msg}\n\n"
                f"These methods should use index maps instead:\n"
                f"  - model_id_to_deployment_index_map (for model_id lookups)\n"
                f"  - model_name_to_deployment_indices (for model_name lookups)\n\n"
                f"If a method legitimately needs O(n) iteration, add it to\n"
                f"ALLOWED_METHODS in this test method.\n"
                f"{'='*70}\n"
            )
    def test_model_names_is_set(self):
        """Verify that model_names uses a set for O(1) lookups, not a list (O(n))"""
        router = Router(model_list=[])
        
        assert isinstance(router.model_names, set), (
            f"model_names should be a set for O(1) lookups, but got {type(router.model_names)}"
        )