From 488b64d7d143fd385b66b03fbfcaf12214eaa1ac Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 31 Aug 2026 23:27:27 -0700 Subject: [PATCH] spike: prototype the team_id column + (team_id, model_name) router identity Exploratory spike used to stress-test the Team BYOK team_id plan against a live proxy with mock upstreams. Not a shippable change: 5 tests in TestDeleteTeamModels fail because they assert the prefix-scan query shape, and the v2 path regresses retries for every team model. - migration + schema: nullable indexed team_id on LiteLLM_ProxyModelTable (DDL only, no backfill - a backfill UPDATE is banned by check_migrations_no_data_rewrites) - ownership queries read the column with a legacy prefix fallback - rename guard reads the stored row instead of the caller payload - get_model_group(id) scopes a team deployment's group to its owning team - LITELLM_TEAM_MODEL_IDENTITY_V2=1: store the public name verbatim, keep team rows out of the global bare-name index and out of router.model_names --- .../migration.sql | 5 ++ .../litellm_proxy_extras/schema.prisma | 3 ++ litellm/models/model.py | 6 ++- litellm/proxy/_types.py | 1 + .../model_management_endpoints.py | 44 ++++++++++++---- litellm/proxy/schema.prisma | 3 ++ litellm/router.py | 50 ++++++++++++++++--- litellm/team_model_identity.py | 16 ++++++ schema.prisma | 3 ++ 9 files changed, 111 insertions(+), 20 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260902000000_add_proxy_model_team_id/migration.sql create mode 100644 litellm/team_model_identity.py diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260902000000_add_proxy_model_team_id/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260902000000_add_proxy_model_team_id/migration.sql new file mode 100644 index 00000000000..1d77ba02795 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260902000000_add_proxy_model_team_id/migration.sql @@ -0,0 +1,5 @@ +-- AlterTable +ALTER TABLE "LiteLLM_ProxyModelTable" ADD COLUMN IF NOT EXISTS "team_id" TEXT; + +-- CreateIndex +CREATE INDEX IF NOT EXISTS "LiteLLM_ProxyModelTable_team_id_idx" ON "LiteLLM_ProxyModelTable"("team_id"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 7604ceadf7a..104b82a8f2b 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -50,6 +50,7 @@ model LiteLLM_CredentialsTable { model LiteLLM_ProxyModelTable { model_id String @id @default(uuid()) model_name String + team_id String? litellm_params Json model_info Json? blocked Boolean @default(false) @@ -57,6 +58,8 @@ model LiteLLM_ProxyModelTable { created_by String updated_at DateTime @default(now()) @updatedAt @map("updated_at") updated_by String + + @@index([team_id]) } diff --git a/litellm/models/model.py b/litellm/models/model.py index a0c840341ab..ca24dce2713 100644 --- a/litellm/models/model.py +++ b/litellm/models/model.py @@ -16,6 +16,7 @@ from litellm.types.llms.base import LiteLLMPydanticObjectBase class LiteLLM_ProxyModelTable(LiteLLMPydanticObjectBase): model_id: str model_name: str + team_id: str | None = None litellm_params: dict model_info: dict | None = None blocked: bool = False @@ -48,7 +49,10 @@ class LiteLLM_ProxyModelTable(LiteLLMPydanticObjectBase): return self.blocked @property - def team_id(self) -> str | None: + def owner_team_id(self) -> str | None: + """Team that owns this deployment: the column, falling back to the legacy JSON blob.""" + if self.team_id is not None: + return self.team_id if self.model_info: return self.model_info.get("team_id") return None diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c549f48126e..8005184cafa 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4821,6 +4821,7 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): class PrismaCompatibleUpdateDBModel(TypedDict, total=False): model_name: str + team_id: str | None litellm_params: str model_info: str blocked: bool diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 012aec38458..984ac98c3b4 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -33,6 +33,7 @@ from litellm.litellm_core_utils.ptu_pricing import ( SEARCH_CONTEXT_SIZES, ptu_config_error, ) +from litellm.team_model_identity import team_model_identity_v2_enabled from litellm.proxy._types import ( BlockModelRequest, CommonProxyErrors, @@ -603,6 +604,7 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr prisma_compatible_model_dict: Final = PrismaCompatibleUpdateDBModel( model_name=merged_model_name, + team_id=merged_model_info.get("team_id"), litellm_params=json.dumps(merged_litellm_params), model_info=json.dumps(merged_model_info), ) @@ -962,6 +964,7 @@ async def _add_model_to_db( _data: Final[dict] = { "model_id": model_params.model_info.id, "model_name": model_params.model_name, + "team_id": model_params.model_info.team_id, "litellm_params": model_params.litellm_params.model_dump_json(exclude_none=True), "model_info": model_params.model_info.model_dump_json(exclude_none=True), "created_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, @@ -1004,8 +1007,9 @@ async def _add_team_model_to_db( # Generate and assign unique internal model_name LAST # (after team_public_model_name is safely stored) - unique_model_name: Final = f"model_name_{_team_id}_{uuid.uuid4()}" - model_params.model_name = unique_model_name + if not team_model_identity_v2_enabled(): + unique_model_name: Final = f"model_name_{_team_id}_{uuid.uuid4()}" + model_params.model_name = unique_model_name ## CREATE MODEL IN DB ## model_response: Final = await _add_model_to_db( @@ -1069,9 +1073,13 @@ async def _update_team_model_in_db( ), ) - patch_team_id: Final = patch_data.model_info.team_id if patch_data.model_info else None + # Ownership comes from the stored row, never from the caller's payload: a PATCH + # that sends only `model_name` used to take the standard-update path and rewrite + # a team deployment's routing key while `team_id` stayed behind, orphaning the row. + db_team_id_for_guard: Final = db_model.model_info.team_id if db_model.model_info else None + patch_team_id: Final = (patch_data.model_info.team_id if patch_data.model_info else None) or db_team_id_for_guard - # No team_id in patch, proceed with standard update + # No team ownership anywhere, proceed with standard update if patch_team_id is None: return update_db_model(db_model=db_model, updated_patch=patch_data) @@ -1165,8 +1173,11 @@ async def _setup_new_team_model_assignment( user_api_key_dict: UserAPIKeyAuth, ) -> None: """Set up a new team model with unique name and team membership.""" - unique_model_name: Final = f"model_name_{team_id}_{uuid.uuid4()}" - patch_data.model_name = unique_model_name + if not team_model_identity_v2_enabled(): + unique_model_name: Final = f"model_name_{team_id}_{uuid.uuid4()}" + patch_data.model_name = unique_model_name + else: + patch_data.model_name = public_model_name await team_model_add( data=TeamModelAddRequest( @@ -1196,17 +1207,25 @@ async def _get_team_deployments( """ prefix: Final = f"model_name_{team_id}_" table = table or _proxy_model_table(prisma_client) + # The team_id column is the ownership key. The legacy name prefix is still read + # so rows written before the column existed (or by a pod that predates it) are + # not orphaned; the JSON check then confirms ownership for those. response: Final = await table.find_many( where={ - "model_name": {"startswith": prefix}, + "OR": [ + {"team_id": team_id}, + {"model_name": {"startswith": prefix}}, + ] } ) if not response: return [] - # Confirm team_id in model_info (defensive check) result: Final = [] for row in response: + if getattr(row, "team_id", None) == team_id: + result.append(row) + continue model_info = model_info_as_mapping(row.model_info) if model_info is not None and model_info.get("team_id") == team_id: result.append(row) @@ -1383,7 +1402,9 @@ async def _update_existing_team_model_assignment( if old_public_name and public_model_name != old_public_name: # Clear user-supplied public name from patch before any early return so the # caller does not overwrite the internal UUID-based model_name in the DB. - patch_data.model_name = None + # Under the v2 identity the public name IS the stored name, so a rename must + # write it through instead of suppressing it. + patch_data.model_name = public_model_name if team_model_identity_v2_enabled() else None if prisma_client is None: verbose_proxy_logger.warning( "prisma_client not initialized; skipping public name update entirely to avoid orphaned entries" @@ -1432,8 +1453,9 @@ async def _update_existing_team_model_assignment( # No team_model_add/delete calls required; public name is already registered # Always clear patch_data.model_name to prevent caller from overwriting - # the internal UUID-based model_name in the DB with the user-supplied public name - patch_data.model_name = None + # the internal UUID-based model_name in the DB with the user-supplied public name. + # Under v2 there is no internal name: the public name is the stored name. + patch_data.model_name = public_model_name if team_model_identity_v2_enabled() else None class ModelManagementAuthChecks: diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 7604ceadf7a..104b82a8f2b 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -50,6 +50,7 @@ model LiteLLM_CredentialsTable { model LiteLLM_ProxyModelTable { model_id String @id @default(uuid()) model_name String + team_id String? litellm_params Json model_info Json? blocked Boolean @default(false) @@ -57,6 +58,8 @@ model LiteLLM_ProxyModelTable { created_by String updated_at DateTime @default(now()) @updatedAt @map("updated_at") updated_by String + + @@index([team_id]) } diff --git a/litellm/router.py b/litellm/router.py index edff8294c3e..c431f2a4c17 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -117,6 +117,7 @@ from litellm.router_utils.batch_utils import ( replace_model_in_jsonl, should_replace_model_in_jsonl, ) +from litellm.team_model_identity import team_model_identity_v2_enabled from litellm.router_utils.client_initalization_utils import InitalizeCachedClient from litellm.router_utils.clientside_credential_handler import ( get_dynamic_litellm_params, @@ -8798,7 +8799,11 @@ class Router: ) verbose_router_logger.debug("\nInitialized Model List %s", self.get_model_names()) - self.model_names = {m["model_name"] for m in model_list} + self.model_names = { + m["model_name"] + for m in model_list + if not (team_model_identity_v2_enabled() and (m.get("model_info") or {}).get("team_id") is not None) + } # Note: model_name_to_deployment_indices is already built incrementally # by _create_deployment -> _add_model_to_list_and_index_map @@ -9020,7 +9025,8 @@ class Router: # add to model names self._add_model_to_list_and_index_map(model=_deployment, model_id=deployment.model_info.id) - self.model_names.add(deployment.model_name) + if not (team_model_identity_v2_enabled() and deployment.model_info.team_id is not None): + self.model_names.add(deployment.model_name) self._sync_deployment_budget_config(deployment=deployment) return deployment @@ -9101,9 +9107,18 @@ class Router: - idx: int - the index in model_list """ team_id: Final = (model.get("model_info") or {}).get("team_id") - team_public_model_name: Final = (model.get("model_info") or {}).get("team_public_model_name") - if team_id and team_public_model_name: - key: Final = (team_id, team_public_model_name) + if not team_id: + return + public_names: Final = { + name + for name in ( + (model.get("model_info") or {}).get("team_public_model_name"), + model.get("model_name") if team_model_identity_v2_enabled() else None, + ) + if name + } + for team_public_model_name in public_names: + key = (team_id, team_public_model_name) self.team_public_model_names = self.team_public_model_names | frozenset({team_public_model_name}) if key not in self.team_model_to_deployment_indices: self.team_model_to_deployment_indices[key] = [] @@ -9129,9 +9144,13 @@ class Router: elif model.get("model_info", {}).get("id") is not None: self.model_id_to_deployment_index_map[model["model_info"]["id"]] = idx - # Update model_name index for O(1) lookup + # Update model_name index for O(1) lookup. + # Under the v2 identity a team-owned deployment lives only under + # (team_id, model_name), so a lookup that forgets the team finds nothing + # instead of crossing into another team. model_name: Final = model.get("model_name") - if model_name: + owns_team: Final = (model.get("model_info") or {}).get("team_id") is not None + if model_name and not (team_model_identity_v2_enabled() and owns_team): if model_name not in self.model_name_to_deployment_indices: self.model_name_to_deployment_indices[model_name] = [] self.model_name_to_deployment_indices[model_name].append(idx) @@ -9800,7 +9819,11 @@ class Router: def get_model_group(self, id: str) -> list | None: """ - Return list of all models in the same model group as that model id + Return list of all models in the same model group as that model id. + + A team-owned deployment's group is scoped to its owning team: the group is + ``(team_id, model_name)``, not the bare name, so a team model with sibling + replicas is not mistaken for a single-deployment group. """ model_info: Final = self.get_model_info(id=id) @@ -9808,6 +9831,10 @@ class Router: return None model_name: Final = model_info["model_name"] + owner_team_id: Final = (model_info.get("model_info") or {}).get("team_id") + if owner_team_id is not None: + public_name: Final = (model_info.get("model_info") or {}).get("team_public_model_name") or model_name + return self.get_model_list(model_name=public_name, team_id=owner_team_id) return self.get_model_list(model_name=model_name) def get_deployment_model_info(self, model_id: str, model_name: str) -> ModelInfo | None: @@ -10541,6 +10568,8 @@ class Router: # Fallback: check by internal model_name for non-team deployments # or deployments that haven't been migrated to team_public_model_name yet model_team_id: Final = (model.get("model_info") or {}).get("team_id") + if team_model_identity_v2_enabled(): + return model_team_id is None or model_team_id == team_id if ( team_id is None # requester has no team constraint or model_team_id is None # global deployment - accessible to all teams @@ -10640,6 +10669,11 @@ class Router: team_model_name = self._get_team_specific_model(deployment=deployment, team_id=team_id) if team_model_name: model_names.append(team_model_name) + elif team_model_identity_v2_enabled() and team_id is None: + # Admin listing: a team deployment still has a real public name. + public = (model_info or {}).get("team_public_model_name") or deployment.get("model_name") + if public: + model_names.append(public) else: model_names.append(deployment.get("model_name", "")) diff --git a/litellm/team_model_identity.py b/litellm/team_model_identity.py new file mode 100644 index 00000000000..7d9da62cf15 --- /dev/null +++ b/litellm/team_model_identity.py @@ -0,0 +1,16 @@ +"""Feature flag for the (team_id, model_name) deployment identity. + +v1 stores a team deployment under a synthetic ``model_name_{team_id}_{uuid}`` and +translates back at every read site. v2 stores the public name verbatim and makes +``(team_id or None, model_name)`` the router's key. +""" + +import os +from typing import Final + + +def team_model_identity_v2_enabled() -> bool: + return os.getenv("LITELLM_TEAM_MODEL_IDENTITY_V2", "").lower() in ("1", "true", "yes") + + +TEAM_MODEL_NAME_PREFIX: Final = "model_name_" diff --git a/schema.prisma b/schema.prisma index 7604ceadf7a..104b82a8f2b 100644 --- a/schema.prisma +++ b/schema.prisma @@ -50,6 +50,7 @@ model LiteLLM_CredentialsTable { model LiteLLM_ProxyModelTable { model_id String @id @default(uuid()) model_name String + team_id String? litellm_params Json model_info Json? blocked Boolean @default(false) @@ -57,6 +58,8 @@ model LiteLLM_ProxyModelTable { created_by String updated_at DateTime @default(now()) @updatedAt @map("updated_at") updated_by String + + @@index([team_id]) }