From 2f38bc6c34f95b6ff8a6a680f79474af08f2cf7a Mon Sep 17 00:00:00 2001
From: Sameer Kankute <sameer@berri.ai>
Date: Mon, 23 Mar 2026 17:56:03 +0530
Subject: [PATCH] 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
---
 litellm/proxy/litellm_pre_call_utils.py       | 11 ++++-----
 .../model_management_endpoints.py             |  4 ++++
 .../test_model_management_endpoints.py        | 23 +++++++++++++++++++
 3 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py
index b0bd92806e2..5de27bb7145 100644
--- a/litellm/proxy/litellm_pre_call_utils.py
+++ b/litellm/proxy/litellm_pre_call_utils.py
@@ -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
diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py
index 4b06c4460e7..7682b65657a 100644
--- a/litellm/proxy/management_endpoints/model_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/model_management_endpoints.py
@@ -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": {
diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py
index dcfd5847bd5..09410c19d34 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py
@@ -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