perf(routing): optimize team model checks and improve test coverage

- Use O(1) team index lookup instead of map_team_model in alias guard
- Fix MockPrismaClient to validate where clause filters
- Add comment explaining DB query trade-off for team deployments

Made-with: Cursor
This commit is contained in:
Sameer Kankute 2026-03-23 17:56:03 +05:30
parent c5a6104eef
commit 2f38bc6c34
3 changed files with 32 additions and 6 deletions

View file

@ -1310,12 +1310,11 @@ def _update_model_if_team_alias_exists(
# Skip alias rewrite if this model resolves to team-specific deployments
# (team models use team_public_model_name, not model_aliases)
if (
llm_router
and user_api_key_dict.team_id
and llm_router.map_team_model(_model, user_api_key_dict.team_id) is not None
):
return
# Use O(1) index lookup instead of map_team_model to avoid O(n) scan
if llm_router and user_api_key_dict.team_id:
key = (user_api_key_dict.team_id, _model)
if key in llm_router.team_model_to_deployment_indices:
return
data["model"] = user_api_key_dict.team_model_aliases[_model]
return

View file

@ -479,6 +479,10 @@ async def _update_existing_team_model_assignment(
"prisma_client not initialized; skipping old public name cleanup to preserve sibling deployments"
)
else:
# Query DB for all deployments in this team, then filter by public name.
# Note: Prisma's JSON filtering doesn't support compound AND conditions
# across multiple JSON paths, so we filter team_public_model_name in Python.
# For most teams (typically <100 deployments), this is acceptable.
response = await prisma_client.db.litellm_proxymodeltable.find_many(
where={
"model_info": {

View file

@ -53,6 +53,29 @@ class MockPrismaClient:
return None
async def find_many(self, where):
# Filter sibling deployments by team_id if where clause specifies it
if not self.sibling_deployments:
return []
# Extract team_id from where clause if present
team_id_filter = None
if where and "model_info" in where:
model_info_filter = where["model_info"]
if isinstance(model_info_filter, dict) and "path" in model_info_filter:
if (
model_info_filter["path"] == ["team_id"]
and "equals" in model_info_filter
):
team_id_filter = model_info_filter["equals"]
# Filter deployments by team_id if specified
if team_id_filter:
return [
d
for d in self.sibling_deployments
if d.model_info.get("team_id") == team_id_filter
]
return self.sibling_deployments
@property