mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
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:
parent
c5a6104eef
commit
2f38bc6c34
3 changed files with 32 additions and 6 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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": {
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue