scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-13 08:52:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 174

fix(proxy): route master key to team-scoped models

Browse source
This commit is contained in:
kunal2002 2026-07-12 00:01:42 +05:30 • committed by ryan-crabbe-berri
parent 93b5ca9612
commit 71eb0b988b
4 changed files with 222 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
litellm/proxy/route_llm_request.py
View file

@ -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,16 @@ 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 any(
public_model_name == data["model"]
for _, public_model_name in llm_router.team_model_to_deployment_indices
)
):
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)

36
litellm/router.py
View file

@ -108,6 +108,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,
)
@ -10018,7 +10019,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,
@ -10033,6 +10037,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"Team-scoped model name '{model}' is configured for multiple teams. "
"Pass a team context or deployment ID to select one."
),
model=model,
llm_provider="",
)
if team_deployments:
return model, team_deployments
pattern_deployments = self.pattern_router.get_deployments_by_pattern(
model=model,
@ -10097,7 +10125,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

29
litellm/router_utils/common_utils.py
View file

@ -1,5 +1,6 @@
import hashlib
import json
from collections.abc import Mapping
from typing import TYPE_CHECKING, Dict, List, Optional, Union
if TYPE_CHECKING:
@ -9,6 +10,17 @@ 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 +71,23 @@ 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_model_info = tuple(
model_info for deployment in healthy_deployments for model_info in [deployment.get("model_info") or {}]
)
if (
isinstance(requested_model, str)
and candidate_model_info
and all(
model_info.get("team_id") is not None and model_info.get("team_public_model_name") == requested_model
for model_info in candidate_model_info
)
):
return healthy_deployments
ids_to_remove = set()
if isinstance(healthy_deployments, dict):
return healthy_deployments

148
tests/test_litellm/proxy/test_route_llm_request.py
View file

@ -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,152 @@ 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-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://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"""