mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-17 04:34:25 +00:00
fix(proxy): let a restricted key use the team deployment its wildcard or model covers
A key restricted to gpt-4o on a team whose own OpenAI deployment is openai/* never resolved that deployment's credentials, because the key allowlist was matched against the deployment's public name literally. The allowlist now matches a wildcard on either side the way request authorization does, through the auth layer's own helper, so /rag, /files and /batches resolve the team deployment for every key that may call a model on it.
This commit is contained in:
parent
500be10d3f
commit
641cf49a3d
3 changed files with 82 additions and 6 deletions
|
|
@ -4714,7 +4714,7 @@ def check_model_access_helper(
|
|||
if _model_in_team_aliases(model=model, team_model_aliases=team_model_aliases):
|
||||
return True
|
||||
|
||||
if _model_matches_any_wildcard_pattern_in_list(model=model, allowed_model_list=filtered_models):
|
||||
if model_matches_any_wildcard_pattern_in_list(model=model, allowed_model_list=filtered_models):
|
||||
return True
|
||||
|
||||
all_model_access: bool = False
|
||||
|
|
@ -7044,7 +7044,7 @@ def is_model_allowed_by_pattern(model: str, allowed_model_pattern: str) -> bool:
|
|||
return False
|
||||
|
||||
|
||||
def _model_matches_any_wildcard_pattern_in_list(model: str, allowed_model_list: list) -> bool:
|
||||
def model_matches_any_wildcard_pattern_in_list(model: str, allowed_model_list: Sequence[str]) -> bool:
|
||||
"""
|
||||
Returns True if a model matches any wildcard pattern in a list.
|
||||
|
||||
|
|
|
|||
|
|
@ -472,6 +472,19 @@ async def get_authorized_credentials_for_model(
|
|||
)
|
||||
|
||||
|
||||
def _key_allowlist_covers(public_model_name: str, key_model_allowlist: Sequence[str]) -> bool:
|
||||
from litellm.proxy.auth.auth_checks import model_matches_any_wildcard_pattern_in_list
|
||||
|
||||
if public_model_name in key_model_allowlist:
|
||||
return True
|
||||
if model_matches_any_wildcard_pattern_in_list(model=public_model_name, allowed_model_list=key_model_allowlist):
|
||||
return True
|
||||
return any(
|
||||
model_matches_any_wildcard_pattern_in_list(model=key_model, allowed_model_list=(public_model_name,))
|
||||
for key_model in key_model_allowlist
|
||||
)
|
||||
|
||||
|
||||
def get_team_provider_credentials(
|
||||
llm_router: Optional["Router"],
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
|
|
@ -496,7 +509,9 @@ def get_team_provider_credentials(
|
|||
model allowlist (``user_api_key_dict.models``), so neither a team nor a
|
||||
restricted key within a team can resolve a provider key for a deployment
|
||||
it isn't authorized to use. A key restricted to an explicit model list
|
||||
only narrows the team scope; sentinel-bearing keys (all-proxy-models /
|
||||
only narrows the team scope, matching a wildcard on either side the way
|
||||
request authorization does (a key allowed ``gpt-4o`` may use the team's
|
||||
``openai/*`` deployment); sentinel-bearing keys (all-proxy-models /
|
||||
all-team-models) defer to the team scope instead of widening past it.
|
||||
Returns None when the router is unavailable or no authorized deployment
|
||||
matches, so the caller can fall back to default credential resolution.
|
||||
|
|
@ -532,12 +547,11 @@ def get_team_provider_credentials(
|
|||
if key_is_restricted
|
||||
else ()
|
||||
)
|
||||
key_model_allowlist_set: Final = frozenset(key_model_allowlist)
|
||||
|
||||
def _key_may_use(public_model_name: str | None) -> bool:
|
||||
if not key_model_allowlist_set:
|
||||
if not key_model_allowlist:
|
||||
return True
|
||||
return public_model_name is not None and public_model_name in key_model_allowlist_set
|
||||
return public_model_name is not None and _key_allowlist_covers(public_model_name, key_model_allowlist)
|
||||
|
||||
def _provider_credentials(model_id: str) -> dict | None:
|
||||
credentials: Final = llm_router.get_deployment_credentials_with_provider(model_id=model_id, team_id=team_id)
|
||||
|
|
|
|||
|
|
@ -5,10 +5,13 @@ from unittest.mock import AsyncMock, MagicMock
|
|||
import pytest
|
||||
from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles
|
||||
|
||||
from litellm import Router
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import (
|
||||
apply_unified_file_ids,
|
||||
get_credentials_for_model,
|
||||
get_team_provider_credentials,
|
||||
is_litellm_executed_batch,
|
||||
map_raw_file_ids_to_unified,
|
||||
)
|
||||
|
|
@ -652,3 +655,62 @@ async def test_batch_cancel_updates_database(caplog: pytest.LogCaptureFixture):
|
|||
assert "batch_processed" not in written
|
||||
assert json.loads(written["file_object"])["cancelled_at"] == 1234567999
|
||||
assert f"Updating batch {batch_id} status to cancelled after cancel" in caplog.messages
|
||||
|
||||
|
||||
def _team_openai_router(team_public_model_name: str, deployment_model: str) -> Router:
|
||||
return Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": team_public_model_name,
|
||||
"litellm_params": {"model": deployment_model, "api_key": "team-openai-key"},
|
||||
"model_info": {
|
||||
"id": "team-a-openai",
|
||||
"team_id": "team-a",
|
||||
"team_public_model_name": team_public_model_name,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "claude-opus-4-6",
|
||||
"litellm_params": {"model": "anthropic/claude-opus-4-6", "api_key": "anthropic-key"},
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def _team_only_openai_credentials(
|
||||
team_public_model_name: str, deployment_model: str, key_models: list[str]
|
||||
) -> dict | None:
|
||||
return get_team_provider_credentials(
|
||||
llm_router=_team_openai_router(team_public_model_name, deployment_model),
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
api_key="test-key",
|
||||
user_id="test-user",
|
||||
team_id="team-a",
|
||||
team_models=[team_public_model_name, "claude-opus-4-6"],
|
||||
models=key_models,
|
||||
),
|
||||
custom_llm_provider="openai",
|
||||
include_accessible_deployments=False,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"team_public_model_name, deployment_model, key_models",
|
||||
[
|
||||
("openai/*", "openai/*", ["gpt-4o"]),
|
||||
("openai/*", "openai/*", ["openai/gpt-4o"]),
|
||||
("gpt-4o", "openai/gpt-4o", ["openai/*"]),
|
||||
],
|
||||
)
|
||||
def test_a_restricted_key_resolves_the_team_deployment_its_wildcard_or_model_covers(
|
||||
team_public_model_name: str, deployment_model: str, key_models: list[str]
|
||||
):
|
||||
credentials = _team_only_openai_credentials(team_public_model_name, deployment_model, key_models)
|
||||
|
||||
assert credentials is not None
|
||||
assert credentials["api_key"] == "team-openai-key"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("key_models", [["claude-opus-4-6"], ["anthropic/*"]])
|
||||
def test_a_restricted_key_outside_the_team_wildcard_never_resolves_its_credentials(key_models: list[str]):
|
||||
assert _team_only_openai_credentials("openai/*", "openai/*", key_models) is None
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue