mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
da5af0cb27
test_no_linear_scans_in_router: #39674 renamed heuristic_v2_router_limit_violation
to auto_router_capability_violation, so the allowlist entry stopped matching and the
same admin-only scan tripped the static check. Rename the entry to follow it.
tableScrolling.spec.ts: 9ba6cab889 (LIT-4738) gave the Tags and Model Hub tables
client-side pagination at 25 rows, so the 40 seeded rows no longer render on one
page. Select 50 rows per page before counting, as the Logs case already does.
328 lines
13 KiB
Python
328 lines
13 KiB
Python
import os
|
|
import pytest
|
|
import ast
|
|
|
|
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-5.5", "model_info": {"id": "model-2"}},
|
|
{"model_name": "gpt-5.5", "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-5.5": [1, 2],
|
|
"claude": [3],
|
|
}
|
|
|
|
# Remove one of the duplicate gpt-5.5 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-5.5"] == [
|
|
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-5.5 deployment
|
|
router._update_deployment_indices_after_removal(
|
|
model_id="model-3", removal_idx=1
|
|
)
|
|
|
|
# Verify gpt-5.5 is removed from dict when no deployments remain
|
|
assert "gpt-5.5" 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-5-mini",
|
|
"litellm_params": {"model": "gpt-5-mini"},
|
|
"model_info": {"id": "model-1"},
|
|
},
|
|
{
|
|
"model_name": "gpt-5.5",
|
|
"litellm_params": {"model": "gpt-5.5"},
|
|
"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_update_team_model_index(self, router):
|
|
"""Test _update_team_model_index updates team_model_to_deployment_indices."""
|
|
model = {
|
|
"model_name": "team-alias",
|
|
"model_info": {
|
|
"id": "dep-1",
|
|
"team_id": "team-abc",
|
|
"team_public_model_name": "gpt-5.5",
|
|
},
|
|
}
|
|
router._update_team_model_index(model, 0)
|
|
assert router.team_model_to_deployment_indices[("team-abc", "gpt-5.5")] == [0]
|
|
router._update_team_model_index(model, 2)
|
|
assert router.team_model_to_deployment_indices[("team-abc", "gpt-5.5")] == [0, 2]
|
|
|
|
router._update_team_model_index(
|
|
{"model_name": "x", "model_info": {"id": "dep-2"}}, 5
|
|
)
|
|
assert router.team_model_to_deployment_indices == {
|
|
("team-abc", "gpt-5.5"): [0, 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-5-mini",
|
|
"litellm_params": {"model": "gpt-5-mini"},
|
|
"model_info": {"id": "model-1"},
|
|
},
|
|
{
|
|
"model_name": "gpt-5.5",
|
|
"litellm_params": {"model": "gpt-5.5"},
|
|
"model_info": {"id": "model-2"},
|
|
},
|
|
{
|
|
"model_name": "gpt-5.5", # Duplicate model_name, different deployment
|
|
"litellm_params": {"model": "gpt-5.5"},
|
|
"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-5-mini" in router.model_name_to_deployment_indices
|
|
assert "gpt-5.5" in router.model_name_to_deployment_indices
|
|
|
|
# Verify: gpt-5-mini has single deployment
|
|
assert router.model_name_to_deployment_indices["gpt-5-mini"] == [0]
|
|
|
|
# Verify: gpt-5.5 has multiple deployments
|
|
assert router.model_name_to_deployment_indices["gpt-5.5"] == [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-5-mini" not in router.model_name_to_deployment_indices
|
|
assert "gpt-5.5" 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": "lookup by litellm_params.model, which is not indexed",
|
|
"_finalize_adaptive_router_if_configured": 'init-time prefix scan for "auto_router/adaptive_router"; no index for prefix match',
|
|
"config_deployments": "filters the whole list on model_info.db_model; admin path only (model add/upsert)",
|
|
"auto_router_capability_violation": "counts gated auto-routers across the whole list; admin path only (auto-router init/upsert)",
|
|
}
|
|
|
|
# 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)}"
|