mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
6a0d03914c
* test: drop the cwd-relative sys.path.insert calls from the test suite
TQ003 stands at 1,077 across 1,058 files, and 1,015 of them are the same shape:
sys.path.insert(0, os.path.abspath("../..")) and its deeper siblings. The
argument resolves against the working directory rather than the file, so from
the repo root, where every job runs pytest, it inserts the directory two levels
above the checkout. It has never pointed at litellm. The package is installed
into the environment anyway, which is what actually makes the import work, and
what the rule's message has said all along.
Removing them leaves 1,634 imports of sys and os with no remaining reference,
and those go too, except where another test module imports the name back out of
the file. The rest of TQ003 is 62 call sites that resolve against __file__ or a
variable, which are a different question and are left alone.
Collection is identical either way: 45,871 tests and the same 51 pre-existing
collection errors before and after, and ruff reports no new undefined name.
* test: drop the duplicate imports the sys.path sweep exposed to F811
* test(pre-call-utils): restore the os import the new bedrock tests need
326 lines
13 KiB
Python
326 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", # Edge case: lookup by litellm_params.model (not indexed)
|
|
"_finalize_adaptive_router_if_configured", # Init-time prefix scan for "auto_router/adaptive_router" (no index for prefix match)
|
|
]
|
|
|
|
# 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)}"
|