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

Merge 53e54b9a10 into 4c8b90945f

Browse source
This commit is contained in:
Hamza Imran 2026-10-30 16:38:53 +08:00 • committed by GitHub
commit b04d3ddc33
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 198 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

67
litellm/proxy/auth/auth_checks.py
View file

@ -5522,6 +5522,31 @@ async def can_key_call_model(
- Exception: If token not allowed to call model
"""
key_models: Final = _resolve_key_models_for_auth_check(valid_token=valid_token)
key_access_group_ids: Final = valid_token.access_group_ids or []
if key_access_group_ids and not key_models:
models_from_groups: Final = await get_models_from_access_groups(
access_group_ids=key_access_group_ids,
prisma_client=prisma_client,
)
if not models_from_groups:
raise ModelAccessDeniedProxyException(
message=model_access_denied_client_message(model=model),
internal_message=f"key access groups grant no models. Tried to access {model}",
type=ProxyErrorTypes.get_model_access_error_type_for_object(object_type="key"),
param="model",
code=status.HTTP_403_FORBIDDEN,
)
return can_object_call_model(
model=model,
llm_router=llm_router,
models=models_from_groups,
team_model_aliases=valid_token.team_model_aliases,
team_id=valid_token.team_id,
key_model_aliases=key_model_aliases_for_auth_check(valid_token),
object_type="key",
)
try:
return can_object_call_model(
model=model,
@ -5534,17 +5559,16 @@ async def can_key_call_model(
)
except ProxyException:
# Fallback: check key's access_group_ids
key_access_group_ids: Final = valid_token.access_group_ids or []
if key_access_group_ids:
models_from_groups: Final = await get_models_from_access_groups(
models_from_groups_fallback: Final = await get_models_from_access_groups(
access_group_ids=key_access_group_ids,
prisma_client=prisma_client,
)
if models_from_groups:
if models_from_groups_fallback:
return can_object_call_model(
model=model,
llm_router=llm_router,
models=models_from_groups,
models=models_from_groups_fallback,
team_model_aliases=valid_token.team_model_aliases,
team_id=valid_token.team_id,
key_model_aliases=key_model_aliases_for_auth_check(valid_token),
@ -5744,11 +5768,37 @@ async def can_team_access_model(
1. First checks native team-level model permissions (current implementation)
2. If not allowed natively, falls back to access_group_ids on the team
"""
team_models: Final = team_object.models if team_object else []
team_access_group_ids: Final = (team_object.access_group_ids or []) if team_object else []
if team_access_group_ids and not team_models:
models_from_groups: Final = await get_models_from_access_groups(
access_group_ids=team_access_group_ids,
prisma_client=prisma_client,
)
if not models_from_groups:
raise ModelAccessDeniedProxyException(
message=model_access_denied_client_message(model=model),
internal_message=f"team access groups grant no models. Tried to access {model}",
type=ProxyErrorTypes.get_model_access_error_type_for_object(object_type="team"),
param="model",
code=status.HTTP_403_FORBIDDEN,
)
return can_object_call_model(
model=model,
llm_router=llm_router,
models=models_from_groups,
team_model_aliases=team_model_aliases,
team_id=team_object.team_id if team_object else None,
key_model_aliases=key_model_aliases,
object_type="team",
)
try:
return can_object_call_model(
model=model,
llm_router=llm_router,
models=team_object.models if team_object else [],
models=team_models,
team_model_aliases=team_model_aliases,
team_id=team_object.team_id if team_object else None,
key_model_aliases=key_model_aliases,
@ -5756,17 +5806,16 @@ async def can_team_access_model(
)
except ProxyException:
# Fallback: check team's access_group_ids
team_access_group_ids: Final = (team_object.access_group_ids or []) if team_object else []
if team_access_group_ids:
models_from_groups: Final = await get_models_from_access_groups(
models_from_groups_fallback: Final = await get_models_from_access_groups(
access_group_ids=team_access_group_ids,
prisma_client=prisma_client,
)
if models_from_groups:
if models_from_groups_fallback:
return can_object_call_model(
model=model,
llm_router=llm_router,
models=list(dict.fromkeys([*(team_object.models if team_object else []), *models_from_groups])),
models=list(dict.fromkeys([*team_models, *models_from_groups_fallback])),
team_model_aliases=team_model_aliases,
team_id=team_object.team_id if team_object else None,
key_model_aliases=key_model_aliases,

140
tests/unit/proxy/auth/test_auth_checks.py
View file

@ -1610,6 +1610,66 @@ async def test_can_team_access_model_access_group_ids_denied():
)
@pytest.mark.asyncio
async def test_can_team_access_model_empty_models_access_group_denied():
"""Test can_team_access_model denies access when team has empty models and model is not in access groups."""
from unittest.mock import AsyncMock, patch
from litellm.proxy._types import ModelAccessDeniedProxyException
from litellm.proxy.auth.auth_checks import can_team_access_model
team_object = LiteLLM_TeamTable(
team_id="test-team",
models=[],
access_group_ids=["ag-with-gpt4"],
)
with (
patch(
"litellm.proxy.auth.auth_checks.get_models_from_access_groups",
new_callable=AsyncMock,
return_value=["gpt-4"],
),
pytest.raises(ModelAccessDeniedProxyException),
):
await can_team_access_model(
model="claude-3",
team_object=team_object,
llm_router=None,
team_model_aliases=None,
)
@pytest.mark.asyncio
async def test_can_team_access_model_empty_models_empty_access_group_denied():
"""Test can_team_access_model denies access when team has empty models and access groups have no models."""
from unittest.mock import AsyncMock, patch
from litellm.proxy._types import ModelAccessDeniedProxyException
from litellm.proxy.auth.auth_checks import can_team_access_model
team_object = LiteLLM_TeamTable(
team_id="test-team",
models=[],
access_group_ids=["ag-empty"],
)
with (
patch(
"litellm.proxy.auth.auth_checks.get_models_from_access_groups",
new_callable=AsyncMock,
return_value=[],
),
pytest.raises(ModelAccessDeniedProxyException),
):
await can_team_access_model(
model="gpt-4",
team_object=team_object,
llm_router=None,
team_model_aliases=None,
)
# ---------------------------------------------------------------------------
# can_key_call_model with access_group_ids fallback
# ---------------------------------------------------------------------------
@ -1649,6 +1709,86 @@ async def test_can_key_call_model_via_access_group_ids():
)
@pytest.mark.asyncio
async def test_can_key_call_model_access_group_ids_denied_for_unlisted_model():
"""Test can_key_call_model denies access when key has access_group_ids but requested model is not in access group."""
from unittest.mock import AsyncMock, patch
from litellm.proxy._types import ModelAccessDeniedProxyException
from litellm.proxy.auth.auth_checks import can_key_call_model
user_api_key_object = UserAPIKeyAuth(
token="test-token",
models=[],
access_group_ids=["ag-with-gpt4"],
)
router = litellm.Router(
model_list=[
{
"model_name": "gpt-4",
"litellm_params": {"model": "openai/gpt-4", "api_key": "test"},
},
{
"model_name": "claude-3",
"litellm_params": {"model": "anthropic/claude-3", "api_key": "test"},
},
]
)
with (
patch(
"litellm.proxy.auth.auth_checks.get_models_from_access_groups",
new_callable=AsyncMock,
return_value=["gpt-4"],
),
pytest.raises(ModelAccessDeniedProxyException),
):
await can_key_call_model(
model="claude-3",
llm_model_list=[],
valid_token=user_api_key_object,
llm_router=router,
)
@pytest.mark.asyncio
async def test_can_key_call_model_access_group_ids_empty_models_denied():
"""Test can_key_call_model denies access when key has access_group_ids but access groups contain no models."""
from unittest.mock import AsyncMock, patch
from litellm.proxy._types import ModelAccessDeniedProxyException
from litellm.proxy.auth.auth_checks import can_key_call_model
user_api_key_object = UserAPIKeyAuth(
token="test-token",
models=[],
access_group_ids=["ag-empty"],
)
router = litellm.Router(
model_list=[
{
"model_name": "gpt-4",
"litellm_params": {"model": "openai/gpt-4", "api_key": "test"},
}
]
)
with (
patch(
"litellm.proxy.auth.auth_checks.get_models_from_access_groups",
new_callable=AsyncMock,
return_value=[],
),
pytest.raises(ModelAccessDeniedProxyException),
):
await can_key_call_model(
model="gpt-4",
llm_model_list=[],
valid_token=user_api_key_object,
llm_router=router,
)
# ---------------------------------------------------------------------------
# _key_access_group_grants_model (key access group overriding team restriction)
# ---------------------------------------------------------------------------