mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
fix(router): never resolve another team's deployment credentials for shared model names
This commit is contained in:
parent
17ce2c4e92
commit
6d607ca3c2
4 changed files with 243 additions and 4 deletions
|
|
@ -319,7 +319,7 @@ def get_team_provider_credentials(
|
|||
return None
|
||||
|
||||
def _provider_credentials(model_id: str) -> Optional[dict]:
|
||||
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 not None and credentials.get("custom_llm_provider") == custom_llm_provider:
|
||||
return credentials
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -30,6 +30,7 @@ from typing import (
|
|||
Generator,
|
||||
List,
|
||||
Literal,
|
||||
Mapping,
|
||||
Optional,
|
||||
Set,
|
||||
Tuple,
|
||||
|
|
@ -8630,6 +8631,33 @@ class Router:
|
|||
raise Exception("Model Name invalid - {}".format(type(model)))
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _deployment_usable_by_team(model: Union[Mapping, Deployment], team_id: str | None) -> bool:
|
||||
"""
|
||||
A team-scoped deployment (``model_info.team_id`` set) is only usable by
|
||||
callers from that same team; deployments without a team owner are shared.
|
||||
"""
|
||||
model_info = model.get("model_info") if isinstance(model, dict) else model.model_info
|
||||
owner_team_id = model_info.get("team_id") if model_info is not None else None
|
||||
return owner_team_id is None or owner_team_id == team_id
|
||||
|
||||
def _get_model_group_deployment_usable_by_team(
|
||||
self, model_group_name: str, team_id: str | None
|
||||
) -> Deployment | None:
|
||||
"""
|
||||
Like ``get_deployment_by_model_group_name``, but skips deployments owned
|
||||
by other teams so a shared model name never resolves another team's
|
||||
credentials.
|
||||
"""
|
||||
indices = self.model_name_to_deployment_indices.get(model_group_name) or ()
|
||||
usable = (
|
||||
self.model_list[idx] for idx in indices if self._deployment_usable_by_team(self.model_list[idx], team_id)
|
||||
)
|
||||
first_usable = next(usable, None)
|
||||
if first_usable is None:
|
||||
return None
|
||||
return Deployment(**first_usable) if isinstance(first_usable, dict) else first_usable
|
||||
|
||||
def get_configured_token_limits(self, model_name: str) -> "tuple[int | None, int | None]":
|
||||
"""
|
||||
Return (max_input_tokens, max_output_tokens) explicitly configured in a concrete
|
||||
|
|
@ -8664,7 +8692,10 @@ class Router:
|
|||
model_id: Model ID or model name from model_list (e.g., "gpt-4o-litellm")
|
||||
team_id: Optional team id of the caller. When set, team-scoped
|
||||
deployments (indexed by team public model name, including team
|
||||
wildcard models like "openai/*") are also considered.
|
||||
wildcard models like "openai/*") are also considered. Name and
|
||||
wildcard lookups never resolve a deployment owned by a
|
||||
different team, so shared model names can't leak another
|
||||
team's credentials.
|
||||
|
||||
Returns:
|
||||
Dictionary containing api_key, api_base, custom_llm_provider, etc.
|
||||
|
|
@ -8681,7 +8712,7 @@ class Router:
|
|||
|
||||
# If not found, try by model_group_name
|
||||
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)
|
||||
|
||||
# If not found, check team-scoped deployments whose team public model
|
||||
# name exactly matches model_id (wildcard team names are matched via
|
||||
|
|
@ -8698,7 +8729,12 @@ class Router:
|
|||
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 []
|
||||
potential_wildcard_models = team_wildcard_models or self.pattern_router.route(model_id) or []
|
||||
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]
|
||||
|
|
|
|||
|
|
@ -2867,3 +2867,87 @@ def test_delete_file_provider_only_resolves_named_vertex_credentials(
|
|||
assert captured_kwargs.get("file_id") == "file-abc123"
|
||||
_assert_vertex_named_credentials_attached(captured_kwargs)
|
||||
proxy_logging_obj.post_call_failure_hook.assert_not_called()
|
||||
|
||||
|
||||
def test_create_file_provider_only_skips_other_team_vertex_deployment(
|
||||
mocker: MockerFixture, monkeypatch
|
||||
):
|
||||
"""
|
||||
Regression: with a team-scoped vertex deployment indexed before a global
|
||||
one under the same model name, a provider-only upload from a different
|
||||
team must use the global deployment's credentials, never the other
|
||||
team's.
|
||||
"""
|
||||
import litellm.proxy.proxy_server as ps
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gemini-2.5-pro",
|
||||
"litellm_params": {
|
||||
"model": "vertex_ai/gemini-2.5-pro",
|
||||
"vertex_project": "team-b-project",
|
||||
},
|
||||
"model_info": {
|
||||
"id": "team-b-vertex",
|
||||
"team_id": "team-b",
|
||||
"team_public_model_name": "gemini-2.5-pro",
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "gemini-2.5-pro",
|
||||
"litellm_params": {
|
||||
"model": "vertex_ai/gemini-2.5-pro",
|
||||
"vertex_project": "shared-project",
|
||||
},
|
||||
},
|
||||
]
|
||||
)
|
||||
proxy_logging_obj = setup_proxy_logging_object(monkeypatch, router)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router)
|
||||
proxy_logging_obj.update_request_status = mocker.AsyncMock()
|
||||
proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock()
|
||||
|
||||
captured_kwargs: dict = {}
|
||||
|
||||
async def _mock_acreate_file(**kwargs):
|
||||
captured_kwargs.update(kwargs)
|
||||
return OpenAIFileObject(
|
||||
id="file-vertex-456",
|
||||
object="file",
|
||||
bytes=2,
|
||||
created_at=1234567890,
|
||||
filename="batch.jsonl",
|
||||
purpose="batch",
|
||||
status="uploaded",
|
||||
)
|
||||
|
||||
monkeypatch.setattr(litellm, "acreate_file", _mock_acreate_file)
|
||||
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
api_key="test-key",
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
user_id="test-user",
|
||||
team_id="team-a",
|
||||
team_models=["gemini-2.5-pro"],
|
||||
)
|
||||
|
||||
try:
|
||||
response = client.post(
|
||||
"/v1/files",
|
||||
files={"file": ("batch.jsonl", b"{}", "application/jsonl")},
|
||||
data={"purpose": "batch"},
|
||||
headers={
|
||||
"Authorization": "Bearer test-key",
|
||||
"custom-llm-provider": "vertex_ai",
|
||||
},
|
||||
)
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert captured_kwargs.get("vertex_project") == "shared-project"
|
||||
proxy_logging_obj.post_call_failure_hook.assert_not_called()
|
||||
|
|
|
|||
|
|
@ -3755,6 +3755,125 @@ def test_get_deployment_credentials_with_provider_team_wildcard_priority():
|
|||
assert global_credentials["api_key"] == "global-key"
|
||||
|
||||
|
||||
def test_get_deployment_credentials_with_provider_skips_other_team_deployment():
|
||||
"""
|
||||
Regression: a team-scoped deployment sharing a model_name with a global
|
||||
deployment must never resolve for another team's (or an unscoped) caller,
|
||||
even when it is indexed first; the shared global deployment wins instead.
|
||||
"""
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gemini-2.5-pro",
|
||||
"litellm_params": {
|
||||
"model": "vertex_ai/gemini-2.5-pro",
|
||||
"vertex_project": "team-b-project",
|
||||
},
|
||||
"model_info": {
|
||||
"id": "team-b-vertex",
|
||||
"team_id": "team-b",
|
||||
"team_public_model_name": "gemini-2.5-pro",
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "gemini-2.5-pro",
|
||||
"litellm_params": {
|
||||
"model": "vertex_ai/gemini-2.5-pro",
|
||||
"vertex_project": "shared-project",
|
||||
},
|
||||
},
|
||||
],
|
||||
)
|
||||
|
||||
other_team_credentials = router.get_deployment_credentials_with_provider(
|
||||
model_id="gemini-2.5-pro", team_id="team-a"
|
||||
)
|
||||
assert other_team_credentials is not None
|
||||
assert other_team_credentials["vertex_project"] == "shared-project"
|
||||
|
||||
unscoped_credentials = router.get_deployment_credentials_with_provider(
|
||||
model_id="gemini-2.5-pro"
|
||||
)
|
||||
assert unscoped_credentials is not None
|
||||
assert unscoped_credentials["vertex_project"] == "shared-project"
|
||||
|
||||
owner_credentials = router.get_deployment_credentials_with_provider(
|
||||
model_id="gemini-2.5-pro", team_id="team-b"
|
||||
)
|
||||
assert owner_credentials is not None
|
||||
assert owner_credentials["vertex_project"] == "team-b-project"
|
||||
|
||||
|
||||
def test_get_deployment_credentials_with_provider_no_fallback_to_other_team_only_name():
|
||||
"""
|
||||
When the only deployments under a model name belong to another team, other
|
||||
callers must get None (env fallback) instead of that team's credentials.
|
||||
"""
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gemini-2.5-pro",
|
||||
"litellm_params": {
|
||||
"model": "vertex_ai/gemini-2.5-pro",
|
||||
"vertex_project": "team-b-project",
|
||||
},
|
||||
"model_info": {
|
||||
"id": "team-b-vertex",
|
||||
"team_id": "team-b",
|
||||
"team_public_model_name": "gemini-2.5-pro",
|
||||
},
|
||||
},
|
||||
],
|
||||
)
|
||||
|
||||
assert (
|
||||
router.get_deployment_credentials_with_provider(
|
||||
model_id="gemini-2.5-pro", team_id="team-a"
|
||||
)
|
||||
is None
|
||||
)
|
||||
assert (
|
||||
router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro")
|
||||
is None
|
||||
)
|
||||
|
||||
|
||||
def test_get_deployment_credentials_with_provider_skips_other_team_wildcard():
|
||||
"""
|
||||
Global wildcard resolution must skip a team-scoped wildcard deployment for
|
||||
callers outside that team, falling through to the shared wildcard entry.
|
||||
"""
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "openai/*",
|
||||
"litellm_params": {"model": "openai/*", "api_key": "team-b-key"},
|
||||
"model_info": {
|
||||
"id": "team-b-wildcard",
|
||||
"team_id": "team-b",
|
||||
"team_public_model_name": "openai/*",
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "openai/*",
|
||||
"litellm_params": {"model": "openai/*", "api_key": "global-key"},
|
||||
},
|
||||
],
|
||||
)
|
||||
|
||||
other_team_credentials = router.get_deployment_credentials_with_provider(
|
||||
model_id="openai/gpt-5.2", team_id="team-a"
|
||||
)
|
||||
assert other_team_credentials is not None
|
||||
assert other_team_credentials["api_key"] == "global-key"
|
||||
|
||||
owner_credentials = router.get_deployment_credentials_with_provider(
|
||||
model_id="openai/gpt-5.2", team_id="team-b"
|
||||
)
|
||||
assert owner_credentials is not None
|
||||
assert owner_credentials["api_key"] == "team-b-key"
|
||||
|
||||
|
||||
def test_team_wildcard_credentials_not_usable_after_delete_deployment():
|
||||
"""
|
||||
Regression: team_pattern_routers retained deleted deployments, so a team
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue