diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index e0ace1ac869..10341a2097b 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -52,6 +52,7 @@ def _swap_alias_for_deployment_model( create_batch_data: LiteLLMBatchCreateRequest, alias: str, llm_router: Optional["Router"], + team_id: "str | None", ) -> None: """ Replace a proxy model-group alias on the batch request with the @@ -61,11 +62,14 @@ def _swap_alias_for_deployment_model( cannot resolve a proxy alias, so a provider transform (e.g. Bedrock's, which forwards ``model`` as the batch ``modelId``) would otherwise receive the alias and the provider would reject it. Falls back to the alias when the - router is unavailable or the alias resolves to nothing. + router is unavailable or the alias resolves to nothing. ``team_id`` keeps + this lookup on the same team-usable deployment the credential resolver + picked, so a team-owned deployment sharing the alias can't leak its model + to callers outside that team. """ if llm_router is None: return - resolved_model = llm_router.get_deployment_model_for_alias(model_id=alias) + resolved_model = llm_router.get_deployment_model_for_alias(model_id=alias, team_id=team_id) if resolved_model is not None: create_batch_data["model"] = resolved_model @@ -214,6 +218,7 @@ async def create_batch( llm_router=llm_router, model_id=model_from_file_id, operation_context="batch creation (file created with model)", + team_id=user_api_key_dict.team_id, ) original_file_id = get_original_file_id(input_file_id) @@ -226,6 +231,7 @@ async def create_batch( create_batch_data=_create_batch_data, alias=model_from_file_id, llm_router=llm_router, + team_id=user_api_key_dict.team_id, ) # Create batch using model credentials @@ -308,6 +314,7 @@ async def create_batch( llm_router=llm_router, model_id=model_param, operation_context="batch creation", + team_id=user_api_key_dict.team_id, ) prepare_data_with_credentials( @@ -318,6 +325,7 @@ async def create_batch( create_batch_data=_create_batch_data, alias=model_param, llm_router=llm_router, + team_id=user_api_key_dict.team_id, ) # Create batch using model credentials @@ -518,6 +526,7 @@ async def retrieve_batch( llm_router=llm_router, model_id=model_from_id, operation_context="batch retrieval (batch created with model)", + team_id=user_api_key_dict.team_id, ) original_batch_id = get_original_file_id(batch_id) @@ -724,6 +733,7 @@ async def list_batches( llm_router=llm_router, model_id=model_param, operation_context="batch listing", + team_id=user_api_key_dict.team_id, ) data.update(credentials) @@ -905,6 +915,7 @@ async def cancel_batch( llm_router=llm_router, model_id=model_from_id, operation_context="batch cancellation (batch created with model)", + team_id=user_api_key_dict.team_id, ) original_batch_id = get_original_file_id(batch_id) diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 2960b031cd3..e9233e1a2a1 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -259,6 +259,7 @@ def get_credentials_for_model( llm_router, # Router instance model_id: str, operation_context: str = "file operation", + team_id: "str | None" = None, ): """ Retrieve API credentials for a model from the LLM Router. @@ -267,6 +268,8 @@ def get_credentials_for_model( llm_router: LiteLLM Router instance model_id: Model name or deployment ID operation_context: Description for error messages (e.g., "file upload", "batch creation") + team_id: Caller's team id; unlocks that team's own deployments and keeps + shared model names from resolving another team's credentials Returns: Dictionary with credentials (api_key, api_base, custom_llm_provider, etc.) @@ -282,7 +285,7 @@ def get_credentials_for_model( detail={"error": "Router not initialized. Cannot use model-based routing."}, ) - credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_id) + credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_id, team_id=team_id) if credentials is None: raise HTTPException( diff --git a/litellm/router.py b/litellm/router.py index 22e9a8986d1..d95c9a9f298 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -8639,19 +8639,54 @@ class Router: raise Exception("Model Name invalid - {}".format(type(model))) return None - def _resolve_unblocked_deployment(self, model_id: str) -> "Deployment | None": + def _resolve_unblocked_deployment(self, model_id: str, team_id: "str | None" = None) -> "Deployment | None": """ Resolve a model id, model-group alias, or wildcard pattern to a single deployment, returning None when nothing matches or the match is paused via ``LiteLLM_ProxyModelTable.blocked``. + + Both the credential resolver and the alias-to-model resolver delegate + here so a mixed model group can never hand one caller the credentials + of one deployment and the model of another. Name and wildcard lookups + skip deployments owned by a team other than ``team_id``; passing + ``team_id`` also unlocks that team's own deployments (exact team public + model name and team wildcard patterns). """ deployment = self.get_deployment(model_id=model_id) if deployment is None: - deployment = self.get_deployment_by_model_group_name(model_group_name=model_id) + deployment = self._get_model_group_deployment_usable_by_team(model_group_name=model_id, team_id=team_id) + # Team-scoped deployments whose team public model name exactly matches + # model_id (wildcard team names are matched via team_pattern_routers + # below). + if deployment is None and team_id is not None: + team_indices = self.team_model_to_deployment_indices.get((team_id, model_id)) or () + team_match = next((self.model_list[idx] for idx in team_indices), None) + if isinstance(team_match, dict): + deployment = Deployment(**team_match) + elif isinstance(team_match, Deployment): + deployment = team_match + + # Wildcard pattern matches. Team wildcard matches take priority so a + # global pattern (e.g. "openai/*") doesn't shadow the team's own entry. if deployment is None: - wildcard_match = next(iter(self.pattern_router.route(model_id) or ()), None) + team_pattern_router = self.team_pattern_routers.get(team_id) if team_id is not None else None + team_wildcard_match = ( + next(iter(team_pattern_router.route(model_id) or ()), None) if team_pattern_router else None + ) + wildcard_match = ( + team_wildcard_match + if team_wildcard_match is not None + else next( + ( + wildcard_model + for wildcard_model in (self.pattern_router.route(model_id) or ()) + if self._deployment_usable_by_team(wildcard_model, team_id) + ), + None, + ) + ) if isinstance(wildcard_match, dict): deployment = Deployment(**wildcard_match) elif isinstance(wildcard_match, Deployment): @@ -8661,7 +8696,7 @@ class Router: return None return deployment - def get_deployment_model_for_alias(self, model_id: str) -> "str | None": + def get_deployment_model_for_alias(self, model_id: str, team_id: "str | None" = None) -> "str | None": """ Resolve a model-group alias (or deployment id / wildcard) to the deployment's underlying ``litellm_params.model``. @@ -8671,8 +8706,10 @@ class Router: ``get_llm_provider`` cannot resolve an alias, so passing it straight through reaches the provider as an invalid model identifier. Returns None when the alias resolves to nothing or to a paused deployment. + Pass the caller's ``team_id`` so the deployment picked here is the same + one the credential resolver picks for that caller. """ - deployment = self._resolve_unblocked_deployment(model_id=model_id) + deployment = self._resolve_unblocked_deployment(model_id=model_id, team_id=team_id) if deployment is None: return None return deployment.litellm_params.model @@ -8753,43 +8790,8 @@ class Router: credentials = router.get_deployment_credentials_with_provider("gpt-4o-litellm") # Returns: {"api_key": "sk-...", "custom_llm_provider": "openai", ...} """ - # Try to get deployment by model_id first - deployment = self.get_deployment(model_id=model_id) - - # If not found, try by model_group_name + deployment = self._resolve_unblocked_deployment(model_id=model_id, team_id=team_id) if deployment is None: - deployment = self._get_model_group_deployment_usable_by_team(model_group_name=model_id, team_id=team_id) - - # If not found, check team-scoped deployments whose team public model - # name exactly matches model_id (wildcard team names are matched via - # team_pattern_routers below). - if deployment is None and team_id is not None: - team_indices = self.team_model_to_deployment_indices.get((team_id, model_id), []) - if team_indices: - team_model = self.model_list[team_indices[0]] - deployment = Deployment(**team_model) if isinstance(team_model, dict) else team_model - - # If still not found, check for wildcard pattern matches. Team wildcard - # matches take priority so a global pattern (e.g. "openai/*") doesn't - # shadow the team's own entry. - if deployment is None: - team_pattern_router = self.team_pattern_routers.get(team_id) if team_id is not None else None - team_wildcard_models = (team_pattern_router.route(model_id) or []) if team_pattern_router else [] - global_wildcard_models = [ - wildcard_model - for wildcard_model in (self.pattern_router.route(model_id) or []) - if self._deployment_usable_by_team(wildcard_model, team_id) - ] - potential_wildcard_models = team_wildcard_models or global_wildcard_models - if potential_wildcard_models: - # Use the first matching wildcard deployment - deployment_dict = potential_wildcard_models[0] - if isinstance(deployment_dict, dict): - deployment = Deployment(**deployment_dict) - elif isinstance(deployment_dict, Deployment): - deployment = deployment_dict - - if deployment is None or self._is_deployment_blocked(deployment): return None # Get basic credentials diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index 9bfa657cfe7..aa33bfcb0a6 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -146,12 +146,12 @@ class Harness: return dict(self.router_acreate.call_args.kwargs) -def _creds_lookup(*, model_id: str) -> Dict[str, str]: +def _creds_lookup(*, model_id: str, team_id: Optional[str] = None) -> Dict[str, str]: # KeyError on an unknown/hardcoded model_id - the bug cannot hide. return dict(CREDS[model_id]) -def _alias_lookup(*, model_id: str) -> str: +def _alias_lookup(*, model_id: str, team_id: Optional[str] = None) -> str: # The endpoint swaps the request model for the deployment's real provider # model before calling the provider; mirror that with the CREDS model so a # wrong/hardcoded model_id KeyErrors instead of hiding. @@ -267,7 +267,7 @@ async def test_create__model_encoded_file_id(harness): harness.router_acreate.assert_not_called() # 2. CREDENTIALS - resolved for the model decoded FROM the file id. - harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o") + harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", team_id=None) # 3. SEAM PAYLOAD - exact, whole dict. A new forwarded key breaks this. assert harness.acreate_kwargs() == { @@ -323,7 +323,7 @@ async def test_create__model_encoded_file_id__resolver_gets_decoded_model(harnes await call_create(harness) - harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o") + harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", team_id=None) # =========================================================================== # @@ -347,7 +347,7 @@ async def test_create__model_from_body(harness): assert harness.litellm_acreate.call_count == 1 harness.router_acreate.assert_not_called() - harness.creds_resolver.assert_called_once_with(model_id="vertex-model") + harness.creds_resolver.assert_called_once_with(model_id="vertex-model", team_id=None) payload = harness.acreate_kwargs() assert payload["custom_llm_provider"] == "vertex_ai" assert payload["input_file_id"] == "file-plain" @@ -367,7 +367,7 @@ async def test_create__model_from_header(harness): await call_create(harness, headers={"x-litellm-model": "vertex-model"}) - harness.creds_resolver.assert_called_once_with(model_id="vertex-model") + harness.creds_resolver.assert_called_once_with(model_id="vertex-model", team_id=None) harness.router_acreate.assert_not_called() @@ -384,7 +384,7 @@ async def test_create__model_from_query(harness): await call_create(harness, query={"model": "vertex-model"}) - harness.creds_resolver.assert_called_once_with(model_id="vertex-model") + harness.creds_resolver.assert_called_once_with(model_id="vertex-model", team_id=None) harness.router_acreate.assert_not_called() @@ -407,7 +407,7 @@ async def test_create__body_model_beats_header_and_query(harness): query={"model": "vertex-model"}, ) - harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o") + harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", team_id=None) # =========================================================================== # @@ -715,7 +715,7 @@ async def test_create__model_encoded_beats_unified(harness): assert harness.litellm_acreate.call_count == 1 harness.router_acreate.assert_not_called() - harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o") + harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", team_id=None) # =========================================================================== # @@ -761,7 +761,7 @@ async def test_create__model_encoded_beats_loadbalancing(harness): assert harness.litellm_acreate.call_count == 1 harness.router_acreate.assert_not_called() - harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o") + harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", team_id=None) # =========================================================================== # @@ -1072,7 +1072,7 @@ async def test_retrieve__model_encoded_id(retrieve_harness): retrieve_harness.router_aretrieve.assert_not_called() # 2. CREDENTIALS - resolved for the model decoded FROM the batch id. - retrieve_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o") + retrieve_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", team_id=None) # 3. SEAM PAYLOAD - exact, whole dict forwarded to the provider call. # Note `model` is the DECODED model, not the deployment from creds: the @@ -1130,7 +1130,7 @@ async def test_retrieve__model_encoded_beats_loadbalancing(retrieve_harness): assert retrieve_harness.litellm_aretrieve.call_count == 1 retrieve_harness.router_aretrieve.assert_not_called() - retrieve_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o") + retrieve_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", team_id=None) # --------------------------------------------------------------------------- # @@ -1546,7 +1546,7 @@ async def test_list__model_from_body_routes_and_encodes(list_harness): assert list_harness.litellm_alist.call_count == 1 list_harness.router_alist.assert_not_called() - list_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o") + list_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", team_id=None) assert resp.data[0].id == encode_file_id_with_model("batch-1", "azure/gpt-4o", id_type="batch") assert resp.data[1].id == encode_file_id_with_model("batch-2", "azure/gpt-4o", id_type="batch") @@ -1836,7 +1836,7 @@ async def test_cancel__model_encoded_id(cancel_harness): cancel_harness.router_acancel.assert_not_called() # CREDENTIALS - resolved for the model decoded from the batch id. - cancel_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o") + cancel_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", team_id=None) # SEAM PAYLOAD - exact dict. NOTE current behavior: `model` is the # DEPLOYMENT name from creds, NOT the decoded model (cancel, unlike @@ -1874,7 +1874,7 @@ async def test_cancel__model_encoded_beats_unified(cancel_harness): assert cancel_harness.litellm_acancel.call_count == 1 cancel_harness.router_acancel.assert_not_called() - cancel_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o") + cancel_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", team_id=None) # --------------------------------------------------------------------------- # diff --git a/tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py b/tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py index e6227057984..13664c71d75 100644 --- a/tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py +++ b/tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py @@ -444,6 +444,7 @@ async def test_create_batch_swaps_alias_for_deployment_model_before_provider_cal mock_user_api_key_dict = MagicMock() mock_user_api_key_dict.parent_otel_span = None mock_user_api_key_dict.user_id = "test_user" + mock_user_api_key_dict.team_id = "team-caller" mock_user_api_key_dict.team_metadata = {} mock_router = MagicMock() @@ -510,7 +511,7 @@ async def test_create_batch_swaps_alias_for_deployment_model_before_provider_cal user_api_key_dict=mock_user_api_key_dict, ) - mock_router.get_deployment_model_for_alias.assert_called_once_with(model_id=alias) + mock_router.get_deployment_model_for_alias.assert_called_once_with(model_id=alias, team_id="team-caller") create_kwargs = mock_create_batch.call_args.kwargs assert create_kwargs["model"] == real_model, ( "Bedrock batch transform receives modelId from this 'model'; it must be the " diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 114bc3810fb..3cfa9fe16c4 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -5844,6 +5844,55 @@ def test_get_deployment_model_for_alias_returns_none_for_blocked_deployment(): assert router.get_deployment_model_for_alias(model_id="dep-1") == "openai/gpt-4o-1" +def test_get_deployment_model_for_alias_matches_credential_deployment_per_team(): + """ + Model and credential resolution must pick the SAME deployment for a caller. + + Regression: with a team-owned deployment listed before a shared one under + the same alias, an unscoped alias lookup returned the team deployment's + model while the team-aware credential resolver returned the shared + deployment's credentials, so an outside caller's batch reached the provider + with team A's private model id on the shared account. + """ + router = litellm.Router( + model_list=[ + { + "model_name": "bedrock-batch", + "litellm_params": { + "model": "bedrock/team-a-private-model", + "aws_region_name": "team-a-region", + }, + "model_info": { + "id": "team-a-dep", + "team_id": "team-a", + "team_public_model_name": "bedrock-batch", + }, + }, + { + "model_name": "bedrock-batch", + "litellm_params": { + "model": "bedrock/shared-model", + "aws_region_name": "shared-region", + }, + "model_info": {"id": "shared-dep"}, + }, + ] + ) + + for team_id, expected_model, expected_region in [ + (None, "bedrock/shared-model", "shared-region"), + ("team-b", "bedrock/shared-model", "shared-region"), + ("team-a", "bedrock/team-a-private-model", "team-a-region"), + ]: + resolved_model = router.get_deployment_model_for_alias(model_id="bedrock-batch", team_id=team_id) + credentials = router.get_deployment_credentials_with_provider(model_id="bedrock-batch", team_id=team_id) + assert resolved_model == expected_model, f"team_id={team_id}" + assert credentials is not None + assert credentials["aws_region_name"] == expected_region, ( + f"team_id={team_id}: credentials came from a different deployment than the model" + ) + + def test_resolve_unblocked_deployment_resolves_alias_id_and_wildcard(): """ _resolve_unblocked_deployment underpins both the credential resolver and the