diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index a0e404e3a18..83730324357 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -801,6 +801,7 @@ router_settings: | LITELLM_OTEL_INTEGRATION_ENABLE_EVENTS | Optionally enable semantic logs for OTEL | LITELLM_OTEL_INTEGRATION_ENABLE_METRICS | Optionally enable emantic metrics for OTEL | LITELLM_ENABLE_PYROSCOPE | If true, enables Pyroscope CPU profiling. Profiles are sent to PYROSCOPE_SERVER_ADDRESS. Off by default. See [Pyroscope profiling](/proxy/pyroscope_profiling). +| LITELLM_ENABLE_TEAM_STALE_ALIAS_BYPASS | When `true`, if a team's legacy `model_aliases` entry maps a public model name to an internal `model_name__` deployment, pre-call handling can skip that rewrite when team-scoped sibling deployments exist for the public name—so load balancing / `order` apply across siblings. Default is `false` for backwards compatibility. See [Team-scoped models and legacy aliases](./load_balancing#team-scoped-models-and-legacy-model_aliases). When stale aliases are detected and this flag is off, the proxy may log a one-time warning. | PYROSCOPE_APP_NAME | Application name reported to Pyroscope. Required when LITELLM_ENABLE_PYROSCOPE is true. No default. | PYROSCOPE_SERVER_ADDRESS | Pyroscope server URL to send profiles to. Required when LITELLM_ENABLE_PYROSCOPE is true. No default. | PYROSCOPE_SAMPLE_RATE | Optional. Sample rate for Pyroscope profiling (integer). No default; when unset, the pyroscope-io library default is used. diff --git a/docs/my-website/docs/proxy/load_balancing.md b/docs/my-website/docs/proxy/load_balancing.md index 5bf39d179f6..66857defd43 100644 --- a/docs/my-website/docs/proxy/load_balancing.md +++ b/docs/my-website/docs/proxy/load_balancing.md @@ -336,6 +336,21 @@ The `order` parameter requires `enable_pre_call_checks: true` in `router_setting If `order=1` deployment is unavailable (e.g., rate-limited), the router falls back to `order=2` deployments. +### Team-scoped models and legacy `model_aliases` {#team-scoped-models-and-legacy-model_aliases} + +Team-scoped deployments are identified by `model_info.team_id` and `model_info.team_public_model_name`. Requests should use the **public** model name; the router resolves all sibling deployments (same public name, different `api_base` / `order`, etc.) for routing, failover, and deployment `order`. + +For router internals: when a `team_id` is in scope, optimized lookups key off `(team_id, team_public_model_name)`. If code passes an internal deployment id (e.g. `model_name__`) instead of the public name, routing still works via the usual deployment-name paths, but the team-specific fast path applies only to the public name. + +**Legacy teams:** Older proxy versions could persist `model_aliases` on the team row mapping a public name to a single internal deployment id (`model_name__`). On each request, pre-call logic may still rewrite `model` to that internal name **before** routing, which collapses to one deployment and can make newer sibling deployments unreachable. + +**Migration options:** + +1. **Recommended for upgrades:** Set environment variable `LITELLM_ENABLE_TEAM_STALE_ALIAS_BYPASS=true` so that when sibling team deployments exist for the public name, the stale alias rewrite is skipped and team-scoped routing (including `order` and failover) applies. See the [Environment variables](./config_settings) table in the proxy settings doc. +2. **Data cleanup:** Remove obsolete `model_aliases` entries for team public names from the team record in the database so only `team_public_model_name` + team model list drive access. + +If a stale alias is detected and the bypass is **not** enabled, the proxy may emit a **one-time** warning in logs explaining that sibling deployments may be unreachable until the flag is set or aliases are cleaned up. + ### When You'll See Load Balancing in Action **Immediate Effects:** diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index daf2867699e..4623ca24f1b 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1,6 +1,7 @@ import asyncio import copy import time +from collections import OrderedDict from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union from fastapi import Request @@ -26,6 +27,7 @@ _SPECIAL_HEADERS_CACHE = frozenset( v.value.lower() for v in SpecialHeaders._member_map_.values() ) from litellm.router import Router +from litellm.secret_managers.main import get_secret_bool from litellm.types.llms.anthropic import ANTHROPIC_API_HEADERS from litellm.types.services import ServiceTypes from litellm.types.utils import ( @@ -36,6 +38,9 @@ from litellm.types.utils import ( ) service_logger_obj = ServiceLogging() # used for tracking latency on OTEL +# Bounded dedup for stale-alias warnings (FIFO eviction when over cap). +_MAX_STALE_ALIAS_WARNING_KEYS = 10_000 +_STALE_TEAM_ALIAS_WARNING_KEYS: OrderedDict[str, None] = OrderedDict() if TYPE_CHECKING: @@ -1295,6 +1300,10 @@ def _update_model_if_team_alias_exists( "gpt-4o": "gpt-4o-team-1" } - requested_model = "gpt-4o-team-1" + + Note: model_aliases for team models are deprecated. This function only applies + to legacy non-team-scoped aliases. Team-scoped deployments use team_public_model_name + and are resolved via map_team_model in route_llm_request. """ _model = data.get("model") if ( @@ -1302,7 +1311,48 @@ def _update_model_if_team_alias_exists( and user_api_key_dict.team_model_aliases and _model in user_api_key_dict.team_model_aliases ): - data["model"] = user_api_key_dict.team_model_aliases[_model] + from litellm.proxy.proxy_server import llm_router + + # Skip alias rewrite if this model resolves to team-specific deployments + # (team models use team_public_model_name, not model_aliases) + aliased_target = user_api_key_dict.team_model_aliases[_model] + + # Optional bypass for stale aliases from pre-PR deployments: + # only enabled via feature flag to preserve backwards compatibility. + enable_stale_alias_bypass = get_secret_bool( + "LITELLM_ENABLE_TEAM_STALE_ALIAS_BYPASS", False + ) + # Check if the alias points to a team-scoped UUID name + # (format: "model_name_{team_id}_{uuid}") + is_stale_team_alias = aliased_target.startswith( + f"model_name_{user_api_key_dict.team_id}_" + ) + if is_stale_team_alias and llm_router: + # This is a stale alias from pre-PR deployments. + # Check if current team deployments exist for the public name. + key = (user_api_key_dict.team_id, _model) + if key in llm_router.team_model_to_deployment_indices: + if enable_stale_alias_bypass: + # Team deployments exist; skip stale alias + return + warning_key = f"{user_api_key_dict.team_id}:{_model}:{aliased_target}" + if warning_key not in _STALE_TEAM_ALIAS_WARNING_KEYS: + _STALE_TEAM_ALIAS_WARNING_KEYS[warning_key] = None + while ( + len(_STALE_TEAM_ALIAS_WARNING_KEYS) + > _MAX_STALE_ALIAS_WARNING_KEYS + ): + _STALE_TEAM_ALIAS_WARNING_KEYS.popitem(last=False) + verbose_proxy_logger.warning( + "Stale team model alias detected for model='%s', team_id='%s'. " + "New sibling deployments may be unreachable. " + "Set LITELLM_ENABLE_TEAM_STALE_ALIAS_BYPASS=true to enable " + "team-scoped sibling routing.", + _model, + user_api_key_dict.team_id, + ) + + data["model"] = aliased_target return diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 44d41097833..95c44a431b5 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -13,13 +13,13 @@ model/{model_id}/update - PATCH endpoint for model update. import asyncio import datetime import json -from litellm._uuid import uuid from typing import Dict, List, Literal, Optional, Tuple, Union, cast from fastapi import APIRouter, Depends, HTTPException, Request, status from pydantic import BaseModel, ConfigDict, Field from litellm._logging import verbose_proxy_logger +from litellm._uuid import uuid from litellm.constants import LITELLM_PROXY_ADMIN_NAME from litellm.proxy._types import ( CommonProxyErrors, @@ -32,7 +32,7 @@ from litellm.proxy._types import ( ProxyErrorTypes, ProxyException, TeamModelAddRequest, - UpdateTeamRequest, + TeamModelDeleteRequest, UserAPIKeyAuth, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -40,7 +40,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helpe from litellm.proxy.management_endpoints.common_utils import _is_user_team_admin from litellm.proxy.management_endpoints.team_endpoints import ( team_model_add, - update_team, + team_model_delete, ) from litellm.proxy.management_helpers.audit_logs import create_object_audit_log from litellm.proxy.utils import PrismaClient @@ -324,7 +324,7 @@ async def _add_team_model_to_db( - generate a unique 'model_name' for the model (e.g. 'model_name_{team_id}_{uuid}) - store the model in the db with the unique 'model_name' - - store a team model alias mapping {"model_name": "model_name_{team_id}_{uuid}"} + - add the public model name to the team's allowed models list """ _team_id = model_params.model_info.team_id if _team_id is None: @@ -344,25 +344,15 @@ async def _add_team_model_to_db( prisma_client=prisma_client, ) - ## CREATE MODEL ALIAS IN DB ## - await update_team( - data=UpdateTeamRequest( - team_id=_team_id, - model_aliases={original_model_name: unique_model_name}, - ), - user_api_key_dict=user_api_key_dict, - http_request=Request(scope={"type": "http"}), - ) - - # add model to team object - await team_model_add( - data=TeamModelAddRequest( - team_id=_team_id, - models=[original_model_name], - ), - http_request=Request(scope={"type": "http"}), - user_api_key_dict=user_api_key_dict, - ) + if original_model_name: + await team_model_add( + data=TeamModelAddRequest( + team_id=_team_id, + models=[original_model_name], + ), + http_request=Request(scope={"type": "http"}), + user_api_key_dict=user_api_key_dict, + ) return model_response @@ -428,6 +418,7 @@ async def _update_team_model_in_db( db_model=db_model, patch_data=patch_data, user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, ) return update_db_model(db_model=db_model, updated_patch=patch_data) @@ -453,19 +444,10 @@ async def _setup_new_team_model_assignment( patch_data: updateDeployment, user_api_key_dict: UserAPIKeyAuth, ) -> None: - """Set up a new team model with unique name, alias, and team membership.""" + """Set up a new team model with unique name and team membership.""" unique_model_name = f"model_name_{team_id}_{uuid.uuid4()}" patch_data.model_name = unique_model_name - await update_team( - data=UpdateTeamRequest( - team_id=team_id, - model_aliases={public_model_name: unique_model_name}, - ), - user_api_key_dict=user_api_key_dict, - http_request=Request(scope={"type": "http"}), - ) - await team_model_add( data=TeamModelAddRequest( team_id=team_id, @@ -482,24 +464,95 @@ async def _update_existing_team_model_assignment( db_model: Deployment, patch_data: updateDeployment, user_api_key_dict: UserAPIKeyAuth, + prisma_client: Optional[PrismaClient], ) -> None: - """Update an existing team model if the public name changed.""" + """Update an existing team model if the public name changed. + + Note on DB scan: Prisma's JSON filtering does not support compound AND conditions + across multiple JSON paths, so we fetch all deployments for the team and filter + team_public_model_name in Python. For teams with many deployments this scan grows + linearly; if team deployment counts become large this should be revisited. + """ + + def _get_team_public_model_name( + model_info: Optional[Union[dict, str]] + ) -> Optional[str]: + if isinstance(model_info, dict): + value = model_info.get("team_public_model_name") + return value if isinstance(value, str) else None + if isinstance(model_info, str): + try: + parsed = json.loads(model_info) + except (TypeError, ValueError): + return None + if isinstance(parsed, dict): + value = parsed.get("team_public_model_name") + return value if isinstance(value, str) else None + return None + old_public_name = ( db_model.model_info.team_public_model_name if db_model.model_info else None ) - # Update alias only if public name changed if old_public_name and public_model_name != old_public_name: - await update_team( - data=UpdateTeamRequest( + # 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 + if prisma_client is None: + verbose_proxy_logger.warning( + "prisma_client not initialized; skipping public name update entirely to avoid orphaned entries" + ) + return + + response = await prisma_client.db.litellm_proxymodeltable.find_many( + where={ + "model_info": { + "path": ["team_id"], + "equals": team_id, + } + } + ) + if not response: + other_deployments_with_old_name = [] + else: + other_deployments_with_old_name = [ + d + for d in response + if d.model_name != db_model.model_name + and _get_team_public_model_name(d.model_info) == old_public_name + ] + + # Add new name first, then delete old name to prevent access loss on partial failure + await team_model_add( + data=TeamModelAddRequest( team_id=team_id, - model_aliases={public_model_name: db_model.model_name}, + models=[public_model_name], ), - user_api_key_dict=user_api_key_dict, http_request=Request(scope={"type": "http"}), + user_api_key_dict=user_api_key_dict, + ) + + if not other_deployments_with_old_name: + await team_model_delete( + data=TeamModelDeleteRequest( + team_id=team_id, + models=[old_public_name], + ), + http_request=Request(scope={"type": "http"}), + user_api_key_dict=user_api_key_dict, + ) + elif not old_public_name and public_model_name: + # First-time assignment of public name on an existing team deployment: + # ensure the team's models list is updated so team routing can resolve it. + await team_model_add( + data=TeamModelAddRequest( + team_id=team_id, + models=[public_model_name], + ), + http_request=Request(scope={"type": "http"}), + user_api_key_dict=user_api_key_dict, ) - # Keep existing unique model_name patch_data.model_name = None diff --git a/litellm/router.py b/litellm/router.py index 0fd4ba80c74..71dfcecd56e 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -466,6 +466,8 @@ class Router: # Initialize model name to deployment indices mapping for O(1) lookups # Maps model_name -> list of indices in model_list self.model_name_to_deployment_indices: Dict[str, List[int]] = {} + # Maps (team_id, team_public_model_name) -> list of indices in model_list + self.team_model_to_deployment_indices: Dict[Tuple[str, str], List[int]] = {} if model_list is not None: # set_model_list will build indices automatically @@ -6757,6 +6759,7 @@ class Router: self.model_list = [] self.model_id_to_deployment_index_map = {} # Reset the index self.model_name_to_deployment_indices = {} # Reset the model_name index + self.team_model_to_deployment_indices = {} # Reset the team_model index self._invalidate_model_group_info_cache() self._invalidate_access_groups_cache() # we add api_base/api_key each model so load balancing between azure/gpt on api_base1 and api_base2 works @@ -7054,16 +7057,17 @@ class Router: # Update model_name_to_deployment_indices for model_name, indices in list(self.model_name_to_deployment_indices.items()): - # Remove the deleted index - if removal_idx in indices: - indices.remove(removal_idx) - - # Decrement all indices greater than removal_idx + # Build new list without mutating the original updated_indices = [] for idx in indices: - if idx > removal_idx: + if idx == removal_idx: + # Skip the removed index + continue + elif idx > removal_idx: + # Decrement indices after removal updated_indices.append(idx - 1) else: + # Keep indices before removal unchanged updated_indices.append(idx) # Update or remove the entry @@ -7072,6 +7076,45 @@ class Router: else: del self.model_name_to_deployment_indices[model_name] + # Update team_model_to_deployment_indices + for key, indices in list(self.team_model_to_deployment_indices.items()): + # Build new list without mutating the original + updated_indices = [] + for idx in indices: + if idx == removal_idx: + # Skip the removed index + continue + elif idx > removal_idx: + # Decrement indices after removal + updated_indices.append(idx - 1) + else: + # Keep indices before removal unchanged + updated_indices.append(idx) + + # Update or remove the entry + if len(updated_indices) > 0: + self.team_model_to_deployment_indices[key] = updated_indices + else: + del self.team_model_to_deployment_indices[key] + + def _update_team_model_index(self, model: dict, idx: int) -> None: + """ + Helper to update team_model_to_deployment_indices for a single deployment. + + Parameters: + - model: dict - the deployment to index + - idx: int - the index in model_list + """ + team_id = (model.get("model_info") or {}).get("team_id") + team_public_model_name = (model.get("model_info") or {}).get( + "team_public_model_name" + ) + if team_id and team_public_model_name: + key = (team_id, team_public_model_name) + if key not in self.team_model_to_deployment_indices: + self.team_model_to_deployment_indices[key] = [] + self.team_model_to_deployment_indices[key].append(idx) + def _add_model_to_list_and_index_map( self, model: dict, model_id: Optional[str] = None ) -> None: @@ -7100,6 +7143,9 @@ class Router: self.model_name_to_deployment_indices[model_name] = [] self.model_name_to_deployment_indices[model_name].append(idx) + # Update team_model index for O(1) team-scoped lookup + self._update_team_model_index(model, idx) + def upsert_deployment(self, deployment: Deployment) -> Optional[Deployment]: """ Add or update deployment @@ -7118,7 +7164,10 @@ class Router: ) if _deployment_on_router is not None: # deployment with this model_id exists on the router - if deployment.litellm_params == _deployment_on_router.litellm_params: + if ( + deployment.litellm_params == _deployment_on_router.litellm_params + and deployment.model_info == _deployment_on_router.model_info + ): # No need to update return None @@ -7930,6 +7979,7 @@ class Router: instead of O(n) linear scan through the entire model_list. """ self.model_name_to_deployment_indices.clear() + self.team_model_to_deployment_indices.clear() for idx, model in enumerate(model_list): model_name = model.get("model_name") @@ -7938,6 +7988,8 @@ class Router: self.model_name_to_deployment_indices[model_name] = [] self.model_name_to_deployment_indices[model_name].append(idx) + self._update_team_model_index(model, idx) + def _build_model_id_to_deployment_index_map(self, model_list: list): """ Build model index from model list to enable O(1) lookups immediately. @@ -8070,20 +8122,23 @@ class Router: def map_team_model(self, team_model_name: str, team_id: str) -> Optional[str]: """ - Map a team model name to a team-specific model name. + Check if team_model_name resolves to team-specific deployments. + + Returns the public model name (unchanged) so the router can find all + sibling deployments via team_id filtering, instead of collapsing to a + single internal model_name. Returns: - - deployment id: str - the deployment id of the team-specific model - - None: if no team-specific model name is found + - str: the team_model_name if team deployments exist for this team + - None: if no team-specific model is found """ models = self.get_model_list(model_name=team_model_name, team_id=team_id) if not models: return None for model in models: if model.get("model_info", {}).get("team_id") == team_id: - return model.get("model_name") + return team_model_name - ## wildcard models return None def should_include_deployment( @@ -8094,12 +8149,19 @@ class Router: """ if ( team_id is not None - and model["model_info"].get("team_id") == team_id - and model_name == model["model_info"].get("team_public_model_name") + and (model.get("model_info") or {}).get("team_id") == team_id + and model_name + == (model.get("model_info") or {}).get("team_public_model_name") ): return True elif model_name is not None and model["model_name"] == model_name: - return True + model_team_id = (model.get("model_info") or {}).get("team_id") + if ( + team_id is None + or model_team_id is None # global deployment - accessible to all teams + or model_team_id == team_id + ): + return True return False def _get_all_deployments( @@ -8116,9 +8178,36 @@ class Router: if team_id specified, only return team-specific models Optimized with O(1) index lookup instead of O(n) linear scan. + + Note: when team_id is provided, O(1) lookup in + `team_model_to_deployment_indices` only applies when `model_name` is the + team public model name. If a caller passes an internal deployment model + name (for example, `model_name__`), this method falls back + to the standard model-name index / scan path. """ returned_models: List[DeploymentTypedDict] = [] + # O(1) lookup in team_model index when team_id is provided + if team_id is not None: + key = (team_id, model_name) + if key in self.team_model_to_deployment_indices: + indices = self.team_model_to_deployment_indices[key] + # O(k) where k = team deployments for this model_name (typically 1-10) + for idx in indices: + model = self.model_list[idx] + if not self.should_include_deployment( + model_name=model_name, model=model, team_id=team_id + ): + continue + if model_alias is not None: + alias_model = model.copy() + alias_model["model_name"] = model_alias + returned_models.append(alias_model) + else: + returned_models.append(model) + if returned_models: + return returned_models + # O(1) lookup in model_name index if model_name in self.model_name_to_deployment_indices: indices = self.model_name_to_deployment_indices[model_name] @@ -8778,6 +8867,16 @@ class Router: model = _model_from_alias if model not in self.model_names: + # Check for team-specific deployments by team_public_model_name. + # This intentionally takes priority over team pattern routers below, + # so that named team deployments shadow wildcard/pattern routes. + if request_team_id is not None: + team_deployments = self._get_all_deployments( + model_name=model, team_id=request_team_id + ) + if team_deployments: + return model, team_deployments + # check if provider/ specific wildcard routing use pattern matching pattern_deployments = self.pattern_router.get_deployments_by_pattern( model=model, diff --git a/pyproject.toml b/pyproject.toml index 07004f3ae16..2a2cdad4565 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,7 +61,7 @@ boto3 = { version = "^1.40.76", optional = true } redisvl = {version = "^0.4.1", optional = true, markers = "python_version >= '3.9' and python_version < '3.14'"} mcp = {version = ">=1.25.0,<2.0.0", optional = true, python = ">=3.10"} a2a-sdk = {version = "^0.3.22", optional = true, python = ">=3.10"} -litellm-proxy-extras = {version = "^0.4.56", optional = true} +litellm-proxy-extras = {version = "^0.4.63", optional = true} rich = {version = "^13.7.1", optional = true} litellm-enterprise = {version = "^0.1.33", optional = true} diskcache = {version = "^5.6.1", optional = true} diff --git a/requirements.txt b/requirements.txt index 2bdafda612a..edb40fcba30 100644 --- a/requirements.txt +++ b/requirements.txt @@ -57,7 +57,7 @@ grpcio>=1.75.0; python_version >= "3.14" sentry_sdk==2.21.0 # for sentry error handling detect-secrets==1.5.0 # Enterprise - secret detection / masking in LLM requests tzdata==2025.1 # IANA time zone database -litellm-proxy-extras==0.4.56 # for proxy extras - e.g. prisma migrations +litellm-proxy-extras==0.4.63 # for proxy extras - e.g. prisma migrations llm-sandbox==0.3.31 # for skill execution in sandbox ### LITELLM PACKAGE DEPENDENCIES python-dotenv==1.0.1 # for env diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index 7cff4807f5a..45646fd6a62 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -2044,6 +2044,49 @@ def test_update_model_if_team_alias_exists(data, user_api_key_dict, expected_mod assert test_data.get("model") == expected_model +def test_team_alias_stale_bypass_disabled_by_default(): + from litellm.proxy.litellm_pre_call_utils import _update_model_if_team_alias_exists + + class _MockRouter: + team_model_to_deployment_indices = {("team-1", "gpt-4o"): [0]} + + test_data = {"model": "gpt-4o"} + user_api_key_dict = UserAPIKeyAuth( + api_key="test_key", + team_id="team-1", + team_model_aliases={"gpt-4o": "model_name_team-1_legacy-uuid"}, + ) + + with patch("litellm.proxy.proxy_server.llm_router", _MockRouter()): + _update_model_if_team_alias_exists( + data=test_data, user_api_key_dict=user_api_key_dict + ) + + assert test_data.get("model") == "model_name_team-1_legacy-uuid" + + +def test_team_alias_stale_bypass_enabled_by_flag(monkeypatch): + from litellm.proxy.litellm_pre_call_utils import _update_model_if_team_alias_exists + + class _MockRouter: + team_model_to_deployment_indices = {("team-1", "gpt-4o"): [0]} + + test_data = {"model": "gpt-4o"} + user_api_key_dict = UserAPIKeyAuth( + api_key="test_key", + team_id="team-1", + team_model_aliases={"gpt-4o": "model_name_team-1_legacy-uuid"}, + ) + monkeypatch.setenv("LITELLM_ENABLE_TEAM_STALE_ALIAS_BYPASS", "true") + + with patch("litellm.proxy.proxy_server.llm_router", _MockRouter()): + _update_model_if_team_alias_exists( + data=test_data, user_api_key_dict=user_api_key_dict + ) + + assert test_data.get("model") == "gpt-4o" + + @pytest.fixture def mock_prisma_client(): client = MagicMock() diff --git a/tests/router_unit_tests/test_get_model_list_alias_optimization.py b/tests/router_unit_tests/test_get_model_list_alias_optimization.py index 31d992b6646..62baf0a3d22 100644 --- a/tests/router_unit_tests/test_get_model_list_alias_optimization.py +++ b/tests/router_unit_tests/test_get_model_list_alias_optimization.py @@ -46,5 +46,5 @@ def test_map_team_model_should_not_iterate_aliases_for_non_alias_team_model_name assert ( router.map_team_model(team_model_name="team-model", team_id="team-1") - == "gpt-3.5-turbo" + == "team-model" ) diff --git a/tests/router_unit_tests/test_router_index_management.py b/tests/router_unit_tests/test_router_index_management.py index 90d98b8ab0a..2694c62827c 100644 --- a/tests/router_unit_tests/test_router_index_management.py +++ b/tests/router_unit_tests/test_router_index_management.py @@ -118,6 +118,28 @@ class TestRouterIndexManagement: 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-4o", + }, + } + router._update_team_model_index(model, 0) + assert router.team_model_to_deployment_indices[("team-abc", "gpt-4o")] == [0] + router._update_team_model_index(model, 2) + assert router.team_model_to_deployment_indices[("team-abc", "gpt-4o")] == [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-4o"): [0, 2], + } + def test_has_model_id(self, router): """Test has_model_id function for O(1) membership check""" # Setup: Add models to router 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 f3c89003105..2a85e24d780 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 @@ -1,13 +1,14 @@ import json import os import sys -from litellm._uuid import uuid from typing import Dict, Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi.testclient import TestClient +from litellm._uuid import uuid + sys.path.insert( 0, os.path.abspath("../../../..") ) # Adds the parent directory to the system path @@ -27,9 +28,15 @@ from litellm.types.router import Deployment, LiteLLM_Params, updateDeployment class MockPrismaClient: - def __init__(self, team_exists: bool = True, user_admin: bool = True): + def __init__( + self, + team_exists: bool = True, + user_admin: bool = True, + sibling_deployments: list = None, + ): self.team_exists = team_exists self.user_admin = user_admin + self.sibling_deployments = sibling_deployments or [] self.db = self async def find_unique(self, where): @@ -45,10 +52,53 @@ 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: + + def _get_team_id(model_info): + if isinstance(model_info, dict): + return model_info.get("team_id") + if isinstance(model_info, str): + try: + parsed = json.loads(model_info) + except (TypeError, ValueError): + return None + if isinstance(parsed, dict): + return parsed.get("team_id") + return None + + return [ + d + for d in self.sibling_deployments + if _get_team_id(d.model_info) == team_id_filter + ] + + return self.sibling_deployments + @property def litellm_teamtable(self): return self + @property + def litellm_proxymodeltable(self): + return self + class MockLLMRouter: def __init__(self): @@ -399,7 +449,9 @@ class TestClearCache: """ Test that clear_cache clears DB models and preserves config models. """ - from litellm.proxy.management_endpoints.model_management_endpoints import clear_cache + from litellm.proxy.management_endpoints.model_management_endpoints import ( + clear_cache, + ) # Create mock router with mixed DB and config models mock_router = MagicMock() @@ -407,18 +459,18 @@ class TestClearCache: { "model_name": "gpt-4", "model_info": {"id": "db-model-1", "db_model": True}, - "litellm_params": {"model": "gpt-4"} + "litellm_params": {"model": "gpt-4"}, }, { - "model_name": "gpt-3.5-turbo", + "model_name": "gpt-3.5-turbo", "model_info": {"id": "config-model-1", "db_model": False}, - "litellm_params": {"model": "gpt-3.5-turbo"} + "litellm_params": {"model": "gpt-3.5-turbo"}, }, { "model_name": "claude-3", "model_info": {"id": "db-model-2", "db_model": True}, - "litellm_params": {"model": "claude-3"} - } + "litellm_params": {"model": "claude-3"}, + }, ] mock_router.delete_deployment = MagicMock(return_value=True) mock_router.auto_routers = MagicMock() @@ -466,8 +518,8 @@ class TestUpdatePublicModelGroups: """ import litellm from litellm.proxy.management_endpoints.model_management_endpoints import ( - update_public_model_groups, UpdatePublicModelGroupsRequest, + update_public_model_groups, ) old_db_models = ["db-model-1", "db-model-2"] @@ -525,7 +577,10 @@ class TestUpdatePublicModelGroups: ) old_links = {"Old Doc": "https://old.example.com"} - new_links = {"New Doc": "https://new.example.com", "API Ref": "https://api.example.com"} + new_links = { + "New Doc": "https://new.example.com", + "API Ref": "https://api.example.com", + } async def mock_get_config(*args, **kwargs): litellm.public_model_groups_links = old_links @@ -558,6 +613,161 @@ class TestUpdatePublicModelGroups: litellm.public_model_groups_links = original_value +class TestTeamModelSiblingRouting: + """ + Verify that sibling team deployments (same public model name, different + api_base) are all reachable through routing — no alias overwrite, no + collapse to a single deployment. + """ + + @pytest.mark.asyncio + async def test_no_model_aliases_written_for_team_models(self): + """ + _add_team_model_to_db must NOT write model_aliases (which caused + the second sibling to overwrite the first). It should only call + team_model_add to register the public name on the team's models list. + """ + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _add_team_model_to_db, + ) + from litellm.types.router import ModelInfo + + team_id = "team_no_alias" + public_name = "gpt-4.1-mini" + + async def mock_add_model_to_db(model_params, user_api_key_dict, prisma_client): + return MagicMock(model_id=str(uuid.uuid4())) + + mock_team_model_add = AsyncMock() + + user = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + prisma_client = MockPrismaClient(team_exists=True) + + for api_base in ["https://eastus.example.com", "https://westus.example.com"]: + dep = Deployment( + model_name=public_name, + litellm_params=LiteLLM_Params( + model="azure/gpt-4o-mini", + api_key="key", + api_base=api_base, + ), + model_info=ModelInfo(team_id=team_id), + ) + with patch( + "litellm.proxy.management_endpoints.model_management_endpoints._add_model_to_db", + side_effect=mock_add_model_to_db, + ), patch( + "litellm.proxy.management_endpoints.model_management_endpoints.team_model_add", + mock_team_model_add, + ): + await _add_team_model_to_db( + model_params=dep, + user_api_key_dict=user, + prisma_client=prisma_client, + ) + + assert mock_team_model_add.call_count == 2 + + @pytest.mark.asyncio + async def test_router_finds_all_sibling_team_deployments(self): + """ + When two team deployments share team_public_model_name="gpt-4.1-mini", + the router's _common_checks_available_deployment must return BOTH as + healthy_deployments (not collapse to one). + """ + import litellm + + team_id = "teamA" + public_name = "gpt-4.1-mini" + + router = litellm.Router( + model_list=[ + { + "model_name": f"model_name_{team_id}_uuid1", + "litellm_params": { + "model": "azure/gpt-4o-mini", + "api_key": "key-1", + "api_base": "https://eastus.openai.azure.com", + }, + "model_info": { + "team_id": team_id, + "team_public_model_name": public_name, + }, + }, + { + "model_name": f"model_name_{team_id}_uuid2", + "litellm_params": { + "model": "azure/gpt-4o-mini", + "api_key": "key-2", + "api_base": "https://westus.openai.azure.com", + }, + "model_info": { + "team_id": team_id, + "team_public_model_name": public_name, + }, + }, + { + "model_name": "global-gpt-4o", + "litellm_params": { + "model": "azure/gpt-4o", + "api_key": "global-key", + "api_base": "https://global.openai.azure.com", + }, + "model_info": {}, # No team_id - global deployment + }, + ], + ) + + # map_team_model should return the public name (not an internal UUID) + result = router.map_team_model(public_name, team_id) + assert result == public_name + + # _common_checks_available_deployment should return both deployments + model, healthy = router._common_checks_available_deployment( + model=public_name, + request_kwargs={"metadata": {"user_api_key_team_id": team_id}}, + ) + assert isinstance(healthy, list) + assert len(healthy) == 2 + api_bases = {d["litellm_params"]["api_base"] for d in healthy} + assert api_bases == { + "https://eastus.openai.azure.com", + "https://westus.openai.azure.com", + } + + def test_global_deployments_accessible_to_teams(self): + """Test that global deployments (no team_id) are accessible to all teams""" + import litellm + + router = litellm.Router( + model_list=[ + { + "model_name": "global-gpt-4o", + "litellm_params": { + "model": "azure/gpt-4o", + "api_key": "global-key", + "api_base": "https://global.openai.azure.com", + }, + "model_info": {}, # No team_id - global deployment + }, + ], + ) + + # Global deployment should be accessible when team_id is provided + deployments = router._get_all_deployments( + model_name="global-gpt-4o", team_id="teamA" + ) + assert len(deployments) == 1 + assert deployments[0]["model_name"] == "global-gpt-4o" + + # should_include_deployment should return True for global deployments + assert router.should_include_deployment( + model_name="global-gpt-4o", + model={"model_name": "global-gpt-4o", "model_info": {}}, + team_id="teamA", + ) + + class TestTeamModelUpdate: """Test team model update handles team_id consistently with model creation""" @@ -591,8 +801,6 @@ class TestTeamModelUpdate: "litellm.proxy.proxy_server.premium_user", True, ), patch( - "litellm.proxy.management_endpoints.model_management_endpoints.update_team" - ) as mock_update_team, patch( "litellm.proxy.management_endpoints.model_management_endpoints.team_model_add" ) as mock_team_model_add: result = await _update_team_model_in_db( @@ -604,9 +812,200 @@ class TestTeamModelUpdate: assert result.get("model_name", "").startswith("model_name_test_team_123_") assert "team_public_model_name" in str(result.get("model_info", "")) - mock_update_team.assert_called_once() + # team_model_add must be called to add public name to team's models list mock_team_model_add.assert_called_once() + @pytest.mark.asyncio + async def test_rename_preserves_old_name_when_siblings_exist(self): + """Test that renaming a deployment preserves old public name when sibling deployments still use it""" + from unittest.mock import MagicMock + + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _update_existing_team_model_assignment, + ) + from litellm.types.router import ModelInfo + + # Create a deployment being renamed + db_model = Deployment( + model_name="model_name_team_123_uuid1", + litellm_params=LiteLLM_Params(model="azure/gpt-4o-mini"), + model_info=ModelInfo( + team_id="team_123", team_public_model_name="old-public-name" + ), + ) + + # Create a sibling deployment that still uses the old public name + sibling_deployment = MagicMock() + sibling_deployment.model_name = "model_name_team_123_uuid2" + sibling_deployment.model_info = { + "team_id": "team_123", + "team_public_model_name": "old-public-name", + } + + prisma_client = MockPrismaClient( + team_exists=True, sibling_deployments=[sibling_deployment] + ) + + patch_data = updateDeployment( + model_name="new-public-name", + model_info=ModelInfo(team_id="team_123"), + ) + + user_api_key_dict = UserAPIKeyAuth( + user_id="test_user", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + with patch( + "litellm.proxy.management_endpoints.model_management_endpoints.team_model_delete" + ) as mock_delete, patch( + "litellm.proxy.management_endpoints.model_management_endpoints.team_model_add" + ) as mock_add: + await _update_existing_team_model_assignment( + team_id="team_123", + public_model_name="new-public-name", + db_model=db_model, + patch_data=patch_data, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, # type: ignore + ) + + # team_model_delete should NOT be called because sibling exists + mock_delete.assert_not_called() + # team_model_add should be called to add new public name + mock_add.assert_called_once() + + @pytest.mark.asyncio + async def test_first_time_public_name_assignment_adds_team_model(self): + """If existing team deployment had no public name, first assignment must call team_model_add.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _update_existing_team_model_assignment, + ) + from litellm.types.router import ModelInfo + + db_model = Deployment( + model_name="model_name_team_123_uuid1", + litellm_params=LiteLLM_Params(model="azure/gpt-4o-mini"), + model_info=ModelInfo(team_id="team_123"), + ) + + patch_data = updateDeployment( + model_name="new-public-name", + model_info=ModelInfo(team_id="team_123"), + ) + + user_api_key_dict = UserAPIKeyAuth( + user_id="test_user", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + with patch( + "litellm.proxy.management_endpoints.model_management_endpoints.team_model_delete" + ) as mock_delete, patch( + "litellm.proxy.management_endpoints.model_management_endpoints.team_model_add" + ) as mock_add: + await _update_existing_team_model_assignment( + team_id="team_123", + public_model_name="new-public-name", + db_model=db_model, + patch_data=patch_data, + user_api_key_dict=user_api_key_dict, + prisma_client=None, + ) + + mock_add.assert_called_once() + mock_delete.assert_not_called() + + @pytest.mark.asyncio + async def test_rename_with_prisma_none_clears_patch_model_name(self): + """Rename path must clear patch_data.model_name even when prisma is unavailable (P1).""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _update_existing_team_model_assignment, + ) + from litellm.types.router import ModelInfo + + db_model = Deployment( + model_name="model_name_team_123_uuid1", + litellm_params=LiteLLM_Params(model="azure/gpt-4o-mini"), + model_info=ModelInfo( + team_id="team_123", team_public_model_name="old-public-name" + ), + ) + patch_data = updateDeployment( + model_name="new-public-name", + model_info=ModelInfo(team_id="team_123"), + ) + user_api_key_dict = UserAPIKeyAuth( + user_id="test_user", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + await _update_existing_team_model_assignment( + team_id="team_123", + public_model_name="new-public-name", + db_model=db_model, + patch_data=patch_data, + user_api_key_dict=user_api_key_dict, + prisma_client=None, + ) + + assert patch_data.model_name is None + + @pytest.mark.asyncio + async def test_rename_handles_legacy_string_model_info(self): + """Test rename path handles legacy string-encoded model_info rows without crashing.""" + from unittest.mock import MagicMock + + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _update_existing_team_model_assignment, + ) + from litellm.types.router import ModelInfo + + db_model = Deployment( + model_name="model_name_team_123_uuid1", + litellm_params=LiteLLM_Params(model="azure/gpt-4o-mini"), + model_info=ModelInfo( + team_id="team_123", team_public_model_name="old-public-name" + ), + ) + + sibling_deployment = MagicMock() + sibling_deployment.model_name = "model_name_team_123_uuid2" + sibling_deployment.model_info = ( + '{"team_id":"team_123","team_public_model_name":"old-public-name"}' + ) + + prisma_client = MockPrismaClient( + team_exists=True, sibling_deployments=[sibling_deployment] + ) + + patch_data = updateDeployment( + model_name="new-public-name", + model_info=ModelInfo(team_id="team_123"), + ) + + user_api_key_dict = UserAPIKeyAuth( + user_id="test_user", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + with patch( + "litellm.proxy.management_endpoints.model_management_endpoints.team_model_delete" + ) as mock_delete, patch( + "litellm.proxy.management_endpoints.model_management_endpoints.team_model_add" + ) as mock_add: + await _update_existing_team_model_assignment( + team_id="team_123", + public_model_name="new-public-name", + db_model=db_model, + patch_data=patch_data, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, # type: ignore + ) + + mock_delete.assert_not_called() + mock_add.assert_called_once() + @pytest.mark.asyncio async def test_patch_model_with_team_id_validates_permissions(self): """Test PATCH with team_id runs same validation as POST for team permissions""" @@ -657,27 +1056,37 @@ class TestModelInfoEndpoint: user_id="test_user", api_key="test_key", models=["gpt-4", "claude-3"], - team_models=["gpt-3.5-turbo"] + team_models=["gpt-3.5-turbo"], ) - with patch("litellm.proxy.proxy_server.llm_router") as mock_router, \ - patch("litellm.proxy.proxy_server.get_key_models") as mock_get_key_models, \ - patch("litellm.proxy.proxy_server.get_team_models") as mock_get_team_models, \ - patch("litellm.proxy.proxy_server.get_complete_model_list") as mock_get_complete_models, \ - patch("litellm.get_llm_provider") as mock_get_provider: - + with patch("litellm.proxy.proxy_server.llm_router") as mock_router, patch( + "litellm.proxy.proxy_server.get_key_models" + ) as mock_get_key_models, patch( + "litellm.proxy.proxy_server.get_team_models" + ) as mock_get_team_models, patch( + "litellm.proxy.proxy_server.get_complete_model_list" + ) as mock_get_complete_models, patch( + "litellm.get_llm_provider" + ) as mock_get_provider: # Setup mocks - mock_router.get_model_names.return_value = ["gpt-4", "claude-3", "gpt-3.5-turbo"] + mock_router.get_model_names.return_value = [ + "gpt-4", + "claude-3", + "gpt-3.5-turbo", + ] mock_router.get_model_access_groups.return_value = {} mock_get_key_models.return_value = ["gpt-4", "claude-3"] mock_get_team_models.return_value = ["gpt-3.5-turbo"] - mock_get_complete_models.return_value = ["gpt-4", "claude-3", "gpt-3.5-turbo"] + mock_get_complete_models.return_value = [ + "gpt-4", + "claude-3", + "gpt-3.5-turbo", + ] mock_get_provider.return_value = (None, "openai", None, None) # Test accessible model result = await model_info( - model_id="gpt-4", - user_api_key_dict=user_api_key_dict + model_id="gpt-4", user_api_key_dict=user_api_key_dict ) assert result["id"] == "gpt-4" @@ -688,22 +1097,25 @@ class TestModelInfoEndpoint: @pytest.mark.asyncio async def test_model_info_inaccessible_model_returns_404(self): """Test model_info returns 404 for inaccessible models""" - from litellm.proxy.proxy_server import model_info from fastapi import HTTPException + from litellm.proxy.proxy_server import model_info + # Mock user with limited access user_api_key_dict = UserAPIKeyAuth( user_id="test_user", api_key="test_key", models=["gpt-4"], # Only has access to gpt-4 - team_models=[] + team_models=[], ) - with patch("litellm.proxy.proxy_server.llm_router") as mock_router, \ - patch("litellm.proxy.proxy_server.get_key_models") as mock_get_key_models, \ - patch("litellm.proxy.proxy_server.get_team_models") as mock_get_team_models, \ - patch("litellm.proxy.proxy_server.get_complete_model_list") as mock_get_complete_models: - + with patch("litellm.proxy.proxy_server.llm_router") as mock_router, patch( + "litellm.proxy.proxy_server.get_key_models" + ) as mock_get_key_models, patch( + "litellm.proxy.proxy_server.get_team_models" + ) as mock_get_team_models, patch( + "litellm.proxy.proxy_server.get_complete_model_list" + ) as mock_get_complete_models: # Setup mocks - user only has access to gpt-4 mock_router.get_model_names.return_value = ["gpt-4", "claude-3"] mock_router.get_model_access_groups.return_value = {} @@ -715,32 +1127,35 @@ class TestModelInfoEndpoint: with pytest.raises(HTTPException) as exc_info: await model_info( model_id="claude-3", # Not in user's accessible models - user_api_key_dict=user_api_key_dict + user_api_key_dict=user_api_key_dict, ) - + assert exc_info.value.status_code == 404 assert "does not exist or is not accessible" in exc_info.value.detail - @pytest.mark.asyncio + @pytest.mark.asyncio async def test_model_info_team_model_access(self): """Test model_info works with team model access""" from litellm.proxy.proxy_server import model_info - + # Mock user with team access user_api_key_dict = UserAPIKeyAuth( user_id="test_user", - api_key="test_key", + api_key="test_key", team_id="test_team", models=[], # No direct key models - team_models=["team-model-1"] + team_models=["team-model-1"], ) - with patch("litellm.proxy.proxy_server.llm_router") as mock_router, \ - patch("litellm.proxy.proxy_server.get_key_models") as mock_get_key_models, \ - patch("litellm.proxy.proxy_server.get_team_models") as mock_get_team_models, \ - patch("litellm.proxy.proxy_server.get_complete_model_list") as mock_get_complete_models, \ - patch("litellm.get_llm_provider") as mock_get_provider: - + with patch("litellm.proxy.proxy_server.llm_router") as mock_router, patch( + "litellm.proxy.proxy_server.get_key_models" + ) as mock_get_key_models, patch( + "litellm.proxy.proxy_server.get_team_models" + ) as mock_get_team_models, patch( + "litellm.proxy.proxy_server.get_complete_model_list" + ) as mock_get_complete_models, patch( + "litellm.get_llm_provider" + ) as mock_get_provider: # Setup mocks mock_router.get_model_names.return_value = ["team-model-1"] mock_router.get_model_access_groups.return_value = {} @@ -751,10 +1166,9 @@ class TestModelInfoEndpoint: # Test team model access result = await model_info( - model_id="team-model-1", - user_api_key_dict=user_api_key_dict + model_id="team-model-1", user_api_key_dict=user_api_key_dict ) assert result["id"] == "team-model-1" - assert result["object"] == "model" + assert result["object"] == "model" assert result["owned_by"] == "custom"