diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 8c980f33b01..94871ff072d 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -6,6 +6,7 @@ from fastapi import HTTPException, status import litellm from litellm.proxy._types import UserAPIKeyAuth +from litellm.router_utils.common_utils import _is_proxy_admin_request # Router-internal mock_testing_* flag names — kept in sync with # ``litellm.types.router.MockRouterTestingParams`` by the test @@ -363,6 +364,7 @@ async def route_request( team_id = get_team_id_from_data(data) router_model_names = llm_router.model_names if llm_router is not None else [] + is_proxy_admin_without_team = team_id is None and _is_proxy_admin_request(data) # Preprocess Google GenAI generate content requests if route_type in ["agenerate_content", "agenerate_content_stream"]: @@ -517,6 +519,13 @@ async def route_request( data["model"] = team_model_name return getattr(llm_router, f"{route_type}")(**data) + elif ( + is_proxy_admin_without_team + and data["model"] not in router_model_names + and data["model"] in llm_router.team_public_model_names + ): + return getattr(llm_router, f"{route_type}")(**data) + elif data["model"] in router_model_names or llm_router.has_model_id(data["model"]): return getattr(llm_router, f"{route_type}")(**data) diff --git a/litellm/router.py b/litellm/router.py index 5ffe60c2da0..8560aa2453a 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -26,6 +26,7 @@ from typing import ( AsyncGenerator, Callable, Dict, + FrozenSet, Generator, List, Literal, @@ -108,6 +109,7 @@ from litellm.router_utils.clientside_credential_handler import ( is_clientside_credential, ) from litellm.router_utils.common_utils import ( + _is_proxy_admin_request, filter_team_based_models, filter_web_search_deployments, ) @@ -490,6 +492,7 @@ class Router: 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]] = {} + self.team_public_model_names: FrozenSet[str] = frozenset() # Initialize cache attributes that ``_invalidate_model_group_info_cache`` # touches *before* the first ``set_model_list`` below (which calls @@ -7777,6 +7780,7 @@ class Router: 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.team_public_model_names = frozenset() # Reset per-strategy router registries so hot-reload doesn't leave # stale routers pointing at the old model_list. self.quality_routers = {} @@ -8128,6 +8132,9 @@ class Router: self.team_model_to_deployment_indices[key] = updated_indices else: del self.team_model_to_deployment_indices[key] + self.team_public_model_names = frozenset( + public_model_name for _, public_model_name in self.team_model_to_deployment_indices + ) def _update_team_model_index(self, model: dict, idx: int) -> None: """ @@ -8141,6 +8148,7 @@ class Router: 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) + 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] = [] if idx not in self.team_model_to_deployment_indices[key]: @@ -9095,6 +9103,7 @@ class Router: """ self.model_name_to_deployment_indices.clear() self.team_model_to_deployment_indices.clear() + self.team_public_model_names = frozenset() for idx, model in enumerate(model_list): model_name = model.get("model_name") @@ -10003,7 +10012,10 @@ class Router: return [m for m in self.model_list if m["litellm_params"]["model"] == model] def _try_early_resolve_deployments_for_model_not_in_names( - self, model: str, request_team_id: Optional[str] + self, + model: str, + request_team_id: Optional[str], + include_team_models: bool = False, ) -> Optional[Tuple[str, Union[List, Dict]]]: """ When ``model`` is not in ``self.model_names``, try team routes, pattern routes, @@ -10018,6 +10030,30 @@ class Router: team_deployments = self._get_all_deployments(model_name=model, team_id=request_team_id) if team_deployments: return model, team_deployments + elif include_team_models: + team_deployments = [ + self.model_list[index] + for (_, public_model_name), indices in self.team_model_to_deployment_indices.items() + if public_model_name == model + for index in indices + ] + team_ids = { + team_id + for deployment in team_deployments + for team_id in [(deployment.get("model_info") or {}).get("team_id")] + if team_id is not None + } + if len(team_ids) > 1: + raise litellm.BadRequestError( + message=( + f"Model name '{model}' matches deployments from multiple teams. " + "Specify the deployment ID directly to disambiguate." + ), + model=model, + llm_provider="", + ) + if team_deployments: + return model, team_deployments pattern_deployments = self.pattern_router.get_deployments_by_pattern( model=model, @@ -10082,7 +10118,11 @@ class Router: if _model_from_alias is not None: model = _model_from_alias - early = self._try_early_resolve_deployments_for_model_not_in_names(model=model, request_team_id=request_team_id) + early = self._try_early_resolve_deployments_for_model_not_in_names( + model=model, + request_team_id=request_team_id, + include_team_models=_is_proxy_admin_request(request_kwargs), + ) if early is not None: return early diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index 18bce2348f6..5cfea5e3bf2 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -1,14 +1,27 @@ import hashlib import json +from collections.abc import Mapping from typing import TYPE_CHECKING, Dict, List, Optional, Union if TYPE_CHECKING: from litellm.types.llms.openai import OpenAIFileObject +from litellm.exceptions import BadRequestError from litellm.types.router import CredentialLiteLLMParams from litellm._logging import verbose_logger +def _is_proxy_admin_request(request_kwargs: Optional[Mapping[str, object]]) -> bool: + if request_kwargs is None: + return False + metadata_value = request_kwargs.get("metadata") + litellm_metadata_value = request_kwargs.get("litellm_metadata") + metadata = metadata_value if isinstance(metadata_value, Mapping) else {} + litellm_metadata = litellm_metadata_value if isinstance(litellm_metadata_value, Mapping) else {} + user_api_key_auth = metadata.get("user_api_key_auth") or litellm_metadata.get("user_api_key_auth") + return getattr(user_api_key_auth, "user_role", None) == "proxy_admin" + + def get_litellm_params_sensitive_credential_hash(litellm_params: dict) -> str: """ Hash of the credential params, used for mapping the file id to the right model @@ -59,6 +72,40 @@ def filter_team_based_models( metadata = request_kwargs.get("metadata") or {} litellm_metadata = request_kwargs.get("litellm_metadata") or {} request_team_id = metadata.get("user_api_key_team_id") or litellm_metadata.get("user_api_key_team_id") + if request_team_id is None and _is_proxy_admin_request(request_kwargs) and isinstance(healthy_deployments, list): + requested_model = ( + request_kwargs.get("model") or metadata.get("model_group") or litellm_metadata.get("model_group") + ) + candidate_deployments = tuple( + (deployment.get("model_name"), deployment.get("model_info") or {}) for deployment in healthy_deployments + ) + team_ids = frozenset( + team_id + for _, model_info in candidate_deployments + for team_id in [model_info.get("team_id")] + if team_id is not None + ) + matches_requested_model = ( + isinstance(requested_model, str) + and bool(candidate_deployments) + and all( + model_info.get("team_id") is not None + and (model_name == requested_model or model_info.get("team_public_model_name") == requested_model) + for model_name, model_info in candidate_deployments + ) + ) + if matches_requested_model and len(team_ids) > 1: + raise BadRequestError( + message=( + f"Model name '{requested_model}' matches deployments from multiple teams. " + "Specify the deployment ID directly to disambiguate." + ), + model=requested_model, + llm_provider="", + ) + if matches_requested_model: + return healthy_deployments + ids_to_remove = set() if isinstance(healthy_deployments, dict): return healthy_deployments diff --git a/tests/test_litellm/proxy/test_route_llm_request.py b/tests/test_litellm/proxy/test_route_llm_request.py index 74a0efba43d..303871e3981 100644 --- a/tests/test_litellm/proxy/test_route_llm_request.py +++ b/tests/test_litellm/proxy/test_route_llm_request.py @@ -8,7 +8,7 @@ sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to from unittest.mock import MagicMock -from litellm.proxy.route_llm_request import route_request +from litellm.proxy.route_llm_request import ProxyModelNotFoundError, route_request @pytest.mark.parametrize( @@ -42,6 +42,200 @@ async def test_route_request_dynamic_credentials(route_type): getattr(llm_router, route_type).assert_called_once_with(**data) +@pytest.mark.asyncio +async def test_route_request_proxy_admin_can_call_all_team_scoped_deployments_without_team_id(): + import litellm + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + + router = litellm.Router( + model_list=[ + { + "model_name": "internal-team-azure-east", + "litellm_params": { + "model": "azure/gpt-4o", + "api_key": "fake", + "api_base": "https://east.example.openai.azure.com", + "api_version": "2024-02-15-preview", + "mock_response": "east", + }, + "model_info": { + "id": "team-azure-east", + "team_id": "team-a", + "team_public_model_name": "team-azure", + }, + }, + { + "model_name": "internal-team-azure-west", + "litellm_params": { + "model": "azure/gpt-4o", + "api_key": "fake", + "api_base": "https://west.example.openai.azure.com", + "api_version": "2024-02-15-preview", + "mock_response": "west", + }, + "model_info": { + "id": "team-azure-west", + "team_id": "team-a", + "team_public_model_name": "team-azure", + }, + }, + ] + ) + admin_auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) + data = { + "model": "team-azure", + "messages": [{"role": "user", "content": "Hello"}], + "metadata": {"user_api_key_auth": admin_auth}, + } + + llm_call = await route_request( + data=data, + llm_router=router, + user_model=None, + route_type="acompletion", + user_api_key_dict=admin_auth, + ) + response = await llm_call + deployments = await router.async_get_healthy_deployments( + model="team-azure", + request_kwargs=data, + ) + + assert response.choices[0].message.content in {"east", "west"} + assert {deployment["model_info"]["id"] for deployment in deployments} == { + "team-azure-east", + "team-azure-west", + } + + non_admin_auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER) + with pytest.raises(ProxyModelNotFoundError): + await route_request( + data={ + **data, + "metadata": {"user_api_key_auth": non_admin_auth}, + }, + llm_router=router, + user_model=None, + route_type="acompletion", + user_api_key_dict=non_admin_auth, + ) + + from litellm.types.router import Deployment + + router.add_deployment( + Deployment( + model_name="internal-team-only", + litellm_params={ + "model": "azure/gpt-4o", + "api_key": "fake", + "api_base": "https://internal.example.openai.azure.com", + "api_version": "2024-02-15-preview", + }, + model_info={ + "id": "internal-team-only-id", + "team_id": "team-a", + }, + ) + ) + internal_deployments = await router.async_get_healthy_deployments( + model="internal-team-only", + request_kwargs={ + **data, + "model": "internal-team-only", + }, + ) + + assert {deployment["model_info"]["id"] for deployment in internal_deployments} == {"internal-team-only-id"} + + router.add_deployment( + Deployment( + model_name="internal-other-team-azure", + litellm_params={ + "model": "azure/gpt-4o", + "api_key": "fake", + "api_base": "https://other.example.openai.azure.com", + "api_version": "2024-02-15-preview", + "mock_response": "other", + }, + model_info={ + "id": "other-team-azure", + "team_id": "team-b", + "team_public_model_name": "team-azure", + }, + ) + ) + + with pytest.raises(litellm.BadRequestError, match="multiple teams"): + ambiguous_call = await route_request( + data=data, + llm_router=router, + user_model=None, + route_type="acompletion", + user_api_key_dict=admin_auth, + ) + await ambiguous_call + + router.add_deployment( + Deployment( + model_name="team-azure", + litellm_params={ + "model": "azure/gpt-4o", + "api_key": "fake", + "api_base": "https://legacy.example.openai.azure.com", + "api_version": "2024-02-15-preview", + }, + model_info={ + "id": "legacy-team-azure", + "team_id": "team-a", + "team_public_model_name": "team-azure", + }, + ) + ) + router.add_deployment( + Deployment( + model_name="team-azure", + litellm_params={ + "model": "azure/gpt-4o", + "api_key": "fake", + "api_base": "https://other-legacy.example.openai.azure.com", + "api_version": "2024-02-15-preview", + }, + model_info={ + "id": "other-legacy-team-azure", + "team_id": "team-b", + "team_public_model_name": "team-azure", + }, + ) + ) + + with pytest.raises(litellm.BadRequestError, match="multiple teams"): + await router.async_get_healthy_deployments( + model="team-azure", + request_kwargs=data, + ) + + router.add_deployment( + Deployment( + model_name="team-azure", + litellm_params={ + "model": "azure/gpt-4o", + "api_key": "fake", + "api_base": "https://global.example.openai.azure.com", + "api_version": "2024-02-15-preview", + }, + model_info={"id": "global-team-azure"}, + ) + ) + + collision_deployments = await router.async_get_healthy_deployments( + model="team-azure", + request_kwargs=data, + ) + + assert {deployment["model_info"]["id"] for deployment in collision_deployments} == {"global-team-azure"} + + @pytest.mark.asyncio async def test_route_request_no_model_required(): """Test route types that don't require model parameter"""