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

Merge 5986a59563 into 3b8ad37296

Browse source
This commit is contained in:
Praising Harris Ratnam 2026-10-27 02:23:02 +00:00 • committed by GitHub
commit 360268fe84
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 320 additions and 145 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -5502,35 +5502,45 @@ 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)
try:
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(
access_group_ids=key_access_group_ids,
prisma_client=prisma_client,
)
allowed_models: Final = (
list(dict.fromkeys([*key_models, *models_from_groups])) if key_models else models_from_groups
)
if not allowed_models:
raise ModelAccessDeniedProxyException(
message=model_access_denied_client_message(model=model),
internal_message=(
f"key not allowed to access model. Key has access_group_ids={key_access_group_ids} "
f"which resolved to 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=key_models,
models=allowed_models,
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",
)
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(
access_group_ids=key_access_group_ids,
prisma_client=prisma_client,
)
if models_from_groups:
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",
)
raise
return can_object_call_model(
model=model,
llm_router=llm_router,
models=key_models,
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",
)
async def _personal_key_owner(
@ -5724,35 +5734,45 @@ 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
"""
try:
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:
models_from_groups: Final = await get_models_from_access_groups(
access_group_ids=team_access_group_ids,
prisma_client=prisma_client,
)
allowed_models: Final = (
list(dict.fromkeys([*team_models, *models_from_groups])) if team_models else models_from_groups
)
if not allowed_models:
raise ModelAccessDeniedProxyException(
message=model_access_denied_client_message(model=model),
internal_message=(
f"team not allowed to access model. Team has access_group_ids={team_access_group_ids} "
f"which resolved to 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=team_object.models if team_object else [],
models=allowed_models,
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",
)
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(
access_group_ids=team_access_group_ids,
prisma_client=prisma_client,
)
if models_from_groups:
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])),
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",
)
raise
return can_object_call_model(
model=model,
llm_router=llm_router,
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,
object_type="team",
)
async def get_authorized_resources_from_key_access_groups(

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

@ -1,10 +1,11 @@
# What is this?
## Tests if 'get_end_user_object' works as expected
import sys, os, asyncio, time, random, uuid
import asyncio
import hashlib
import json
import traceback
import uuid
from dotenv import load_dotenv
load_dotenv()
@ -15,39 +16,40 @@ from types import SimpleNamespace
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import pytest, litellm
import httpx
import pytest
from fastapi import Request, status
from prisma import Prisma
import litellm
import litellm.proxy.common_utils.auth_cache_invalidation_pubsub as pubsub_module
from litellm._service_logger import ServiceTypes
from litellm.proxy._types import LiteLLM_OrganizationTable, UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import get_org_object, get_user_object
from litellm.proxy.db.prisma_client import PrismaWrapper
from litellm.proxy.auth.auth_checks import get_end_user_object
from litellm.caching.caching import DualCache
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.proxy._types import (
LiteLLM_EndUserTable,
LiteLLM_BudgetTable,
LiteLLM_UserTable,
LiteLLM_TeamTable,
LiteLLM_EndUserTable,
Litellm_EntityType,
LiteLLM_OrganizationTable,
LiteLLM_TeamTable,
LiteLLM_UserTable,
LitellmUserRoles,
ProxyException,
UserAPIKeyAuth,
)
from litellm.proxy.utils import PrismaClient
import litellm.proxy.common_utils.auth_cache_invalidation_pubsub as pubsub_module
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.proxy.auth.auth_checks import (
_team_soft_budget_check,
can_team_access_model,
common_checks,
delete_cache_key_objects,
get_end_user_object,
get_org_object,
get_user_object,
is_model_cost_zero,
virtual_key_soft_budget_check,
_team_soft_budget_check,
common_checks,
)
from litellm.proxy.utils import ProxyLogging
from litellm.proxy.utils import CallInfo
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.db.prisma_client import PrismaWrapper
from litellm.proxy.utils import CallInfo, ProxyLogging
from litellm.types.passthrough_endpoints.pass_through_endpoints import LITELLM_PASS_THROUGH_ENDPOINT_MARKER
@ -160,12 +162,15 @@ async def test_team_require_trace_id_accepts_supported_trace_carriers(
@pytest.mark.parametrize("route", ["/v1/traces", "/v1/logs", "/v1/traces/query"])
@pytest.mark.asyncio
async def test_team_require_trace_id_accepts_trace_telemetry_routes(route: str) -> None:
assert await _run_team_trace_id_check(
route=route,
method="POST",
request_body={},
team_metadata={"require_trace_id": True},
) is True
assert (
await _run_team_trace_id_check(
route=route,
method="POST",
request_body={},
team_metadata={"require_trace_id": True},
)
is True
)
@pytest.mark.asyncio
@ -197,11 +202,14 @@ async def test_team_require_trace_id_rejects_overridden_body_metadata_on_chat()
async def test_team_require_trace_id_accepts_effective_metadata_on_responses_route(
request_body: dict[str, object],
) -> None:
assert await _run_team_trace_id_check(
route="/v1/responses",
request_body=request_body,
team_metadata={"require_trace_id": True},
) is True
assert (
await _run_team_trace_id_check(
route="/v1/responses",
request_body=request_body,
team_metadata={"require_trace_id": True},
)
is True
)
@pytest.mark.parametrize(
@ -240,12 +248,15 @@ async def test_team_require_trace_id_rejects_fallback_carriers_when_litellm_trac
@pytest.mark.asyncio
async def test_team_require_trace_id_accepts_header_on_mcp_with_empty_body() -> None:
assert await _run_team_trace_id_check(
route="/mcp/",
request_body={},
headers={"x-litellm-trace-id": "trace-mcp-123"},
team_metadata={"require_trace_id": True},
) is True
assert (
await _run_team_trace_id_check(
route="/mcp/",
request_body={},
headers={"x-litellm-trace-id": "trace-mcp-123"},
team_metadata={"require_trace_id": True},
)
is True
)
@pytest.mark.parametrize(
@ -278,12 +289,15 @@ async def test_team_require_trace_id_rejects_nonheader_carriers_on_pass_through_
@pytest.mark.asyncio
async def test_team_require_trace_id_accepts_header_on_pass_through_route() -> None:
assert await _run_team_trace_id_check(
route="/anthropic/v1/messages",
request_body={},
headers={"x-litellm-trace-id": "trace-provider-123"},
team_metadata={"require_trace_id": True},
) is True
assert (
await _run_team_trace_id_check(
route="/anthropic/v1/messages",
request_body={},
headers={"x-litellm-trace-id": "trace-provider-123"},
team_metadata={"require_trace_id": True},
)
is True
)
@pytest.mark.parametrize(
@ -329,7 +343,9 @@ async def test_team_require_trace_id_rejects_nonheader_carriers_on_registered_pa
"jsonrpc": "2.0",
"id": "1",
"method": "message/send",
"params": {"message": {"messageId": "msg-1", "role": "user", "parts": [{"type": "text", "text": "hi"}]}},
"params": {
"message": {"messageId": "msg-1", "role": "user", "parts": [{"type": "text", "text": "hi"}]}
},
},
),
],
@ -370,7 +386,9 @@ async def test_team_require_trace_id_rejects_traceparent_on_mcp_and_agent_routes
"jsonrpc": "2.0",
"id": "1",
"method": "message/send",
"params": {"message": {"messageId": "msg-1", "role": "user", "parts": [{"type": "text", "text": "hi"}]}},
"params": {
"message": {"messageId": "msg-1", "role": "user", "parts": [{"type": "text", "text": "hi"}]}
},
"metadata": {"trace_id": "trace-body-123"},
},
),
@ -463,7 +481,7 @@ def test_lens_management_does_not_require_an_inference_user_parameter(route: str
async def test_get_end_user_object(customer_spend, customer_budget):
"""
Scenario 1: normal - get_end_user_object returns the cached user
Scenario 2: user over budget - NOTE: budget enforcement now happens in
Scenario 2: user over budget - NOTE: budget enforcement now happens in
common_checks() via _check_end_user_budget(), not in get_end_user_object()
This test verifies that get_end_user_object correctly retrieves the end user
@ -480,7 +498,7 @@ async def test_get_end_user_object(customer_spend, customer_budget):
# UserApiKeyCache applies model_type on get/set; plain DualCache returns raw dicts
# and breaks get_end_user_object's typed async_get_cache path.
_cache = UserApiKeyCache()
_key = "end_user_id:{}".format(end_user_id)
_key = f"end_user_id:{end_user_id}"
await _cache.async_set_cache(
key=_key,
value=end_user_obj,
@ -550,8 +568,8 @@ async def test_can_key_call_model(model, expect_to_work):
"""
If wildcard model + specific model is used, choose the specific model settings
"""
from litellm.proxy.auth.auth_checks import can_key_call_model
from fastapi import HTTPException
llm_model_list = [
{
@ -591,7 +609,7 @@ async def test_can_key_call_model(model, expect_to_work):
if expect_to_work:
await can_key_call_model(**args)
else:
with pytest.raises(Exception, match='is not available for this API key') as e:
with pytest.raises(Exception, match="is not available for this API key") as e:
await can_key_call_model(**args)
print(e)
@ -603,8 +621,8 @@ async def test_can_key_call_model(model, expect_to_work):
)
@pytest.mark.asyncio
async def test_can_team_call_model(model, expect_to_work):
from litellm.proxy.auth.auth_checks import model_in_access_group
from fastapi import HTTPException
llm_model_list = [
{
@ -850,8 +868,8 @@ async def test_add_known_models_explicit_map_updates_provider_sets():
@pytest.mark.asyncio
async def test_is_valid_fallback_model():
from litellm.proxy.auth.auth_checks import is_valid_fallback_model
from litellm import Router
from litellm.proxy.auth.auth_checks import is_valid_fallback_model
router = Router(
model_list=[
@ -863,16 +881,12 @@ async def test_is_valid_fallback_model():
)
try:
await is_valid_fallback_model(
model="gpt-3.5-turbo", llm_router=router, user_model=None
)
await is_valid_fallback_model(model="gpt-3.5-turbo", llm_router=router, user_model=None)
except Exception as e:
pytest.fail(f"Expected is_valid_fallback_model to work, got exception: {e}")
with pytest.raises(Exception, match="Invalid") as exc_info:
await is_valid_fallback_model(
model="gpt-4o", llm_router=router, user_model=None
)
await is_valid_fallback_model(model="gpt-4o", llm_router=router, user_model=None)
e = exc_info.value
assert "Invalid" in str(e)
@ -886,9 +900,7 @@ async def test_is_valid_fallback_model():
],
)
@pytest.mark.asyncio
async def test_virtual_key_max_budget_check(
token_spend, max_budget, expect_budget_error
):
async def test_virtual_key_max_budget_check(token_spend, max_budget, expect_budget_error):
"""
Test if virtual key budget checks work as expected:
1. Triggers budget alert for all cases
@ -1003,13 +1015,11 @@ async def test_can_team_access_model(model, team_models, expect_to_work):
team_model_aliases=None,
)
if not expect_to_work:
pytest.fail(
f"Expected model access check to fail for model={model}, team_models={team_models}"
)
pytest.fail(f"Expected model access check to fail for model={model}, team_models={team_models}")
except Exception as e:
if expect_to_work:
pytest.fail(
f"Expected model access check to work for model={model}, team_models={team_models}. Got error: {str(e)}"
f"Expected model access check to work for model={model}, team_models={team_models}. Got error: {e!s}"
)
@ -1058,9 +1068,9 @@ async def test_virtual_key_soft_budget_check(spend, soft_budget, expect_alert):
await asyncio.sleep(0.1) # Allow time for the alert task to complete
assert (
alert_triggered == expect_alert
), f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}"
assert alert_triggered == expect_alert, (
f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}"
)
@pytest.mark.parametrize(
@ -1113,9 +1123,7 @@ async def test_virtual_key_soft_budget_check(spend, soft_budget, expect_alert):
],
)
@pytest.mark.asyncio
async def test_team_soft_budget_check(
spend, soft_budget, expect_alert, metadata, expected_alert_emails
):
async def test_team_soft_budget_check(spend, soft_budget, expect_alert, metadata, expected_alert_emails):
"""
Test cases for _team_soft_budget_check:
1. Spend over soft budget, no alert_emails configured - should NOT trigger alert (alerts only sent when alert_emails configured)
@ -1163,9 +1171,9 @@ async def test_team_soft_budget_check(
await asyncio.sleep(0.1) # Allow time for the alert task to complete
assert (
alert_triggered == expect_alert
), f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}"
assert alert_triggered == expect_alert, (
f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}"
)
if expect_alert:
assert captured_call_info is not None
@ -1177,17 +1185,14 @@ async def test_team_soft_budget_check(
if expected_alert_emails is not None:
assert captured_call_info.alert_emails == expected_alert_emails
else:
assert (
captured_call_info.alert_emails is None
or captured_call_info.alert_emails == []
)
assert captured_call_info.alert_emails is None or captured_call_info.alert_emails == []
@pytest.mark.asyncio
async def test_can_user_call_model():
from litellm.proxy.auth.auth_checks import can_user_call_model
from litellm.proxy._types import ProxyException
from litellm import Router
from litellm.proxy._types import ProxyException
from litellm.proxy.auth.auth_checks import can_user_call_model
router = Router(
model_list=[
@ -1225,10 +1230,11 @@ async def test_can_user_call_model():
@pytest.mark.asyncio
async def test_can_user_call_model_with_no_default_models():
from litellm.proxy.auth.auth_checks import can_user_call_model
from litellm.proxy._types import ProxyException, SpecialModelNames
from unittest.mock import MagicMock
from litellm.proxy._types import ProxyException, SpecialModelNames
from litellm.proxy.auth.auth_checks import can_user_call_model
args = {
"model": "anthropic-claude",
"llm_router": MagicMock(),
@ -1249,9 +1255,10 @@ async def test_can_user_call_model_with_no_default_models():
@pytest.mark.asyncio
async def test_get_fuzzy_user_object():
from litellm.proxy.auth.auth_checks import _get_fuzzy_user_object
from unittest.mock import AsyncMock, MagicMock
from litellm.proxy.auth.auth_checks import _get_fuzzy_user_object
# Setup mock Prisma client
mock_prisma = MagicMock()
mock_prisma.db = MagicMock()
@ -1311,9 +1318,7 @@ async def test_get_fuzzy_user_object():
# Test 5: Only email provided (no SSO ID)
mock_prisma.db.litellm_usertable.find_first = AsyncMock(return_value=test_user)
result = await _get_fuzzy_user_object(
prisma_client=mock_prisma, user_email="test@example.com"
)
result = await _get_fuzzy_user_object(prisma_client=mock_prisma, user_email="test@example.com")
assert result == test_user
mock_prisma.db.litellm_usertable.find_first.assert_called_with(
where={"user_email": {"equals": "test@example.com", "mode": "insensitive"}},
@ -1322,9 +1327,7 @@ async def test_get_fuzzy_user_object():
# Test 6: Only SSO ID provided (no email)
mock_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=test_user)
result = await _get_fuzzy_user_object(
prisma_client=mock_prisma, sso_user_id="sso_123"
)
result = await _get_fuzzy_user_object(prisma_client=mock_prisma, sso_user_id="sso_123")
assert result == test_user
mock_prisma.db.litellm_usertable.find_unique.assert_called_with(
where={"sso_user_id": "sso_123"}, include={"organization_memberships": True}
@ -1371,7 +1374,7 @@ async def test_can_key_call_model_with_aliases(model, alias_map, expect_to_work)
llm_router=router,
)
else:
with pytest.raises(Exception, match='is not available for this API key') as e:
with pytest.raises(Exception, match="is not available for this API key") as e:
await can_key_call_model(
model=model,
llm_model_list=llm_model_list,
@ -1388,8 +1391,8 @@ async def test_can_key_call_model_with_aliases(model, alias_map, expect_to_work)
@pytest.mark.asyncio
async def test_cache_access_object():
"""Test _cache_access_object stores access group in cache with correct key."""
from litellm.proxy.auth.auth_checks import cache_access_object
from litellm.proxy._types import LiteLLM_AccessGroupTable
from litellm.proxy.auth.auth_checks import cache_access_object
cache = DualCache()
ag_id = "ag-test-123"
@ -1416,8 +1419,8 @@ async def test_cache_access_object():
@pytest.mark.asyncio
async def test_delete_cache_access_object():
"""Test _delete_cache_access_object removes access group from in-memory cache."""
from litellm.proxy.auth.auth_checks import delete_cache_access_object
from litellm.proxy._types import LiteLLM_AccessGroupTable
from litellm.proxy.auth.auth_checks import delete_cache_access_object
cache = DualCache()
ag_id = "ag-delete-test"
@ -1457,9 +1460,7 @@ async def test_delete_cache_access_object():
],
)
@pytest.mark.asyncio
async def test_get_resources_from_access_groups(
resource_field, access_group_data, expected
):
async def test_get_resources_from_access_groups(resource_field, access_group_data, expected):
"""Test _get_resources_from_access_groups returns correct resource list from access groups."""
from unittest.mock import AsyncMock, MagicMock, patch
@ -1542,8 +1543,8 @@ async def test_can_team_access_model_access_group_ids_denied():
"""Test can_team_access_model denies when neither team models nor access_group_ids grant access."""
from unittest.mock import AsyncMock, patch
from litellm.proxy.auth.auth_checks import can_team_access_model
from litellm.proxy._types import ProxyException
from litellm.proxy.auth.auth_checks import can_team_access_model
team_object = LiteLLM_TeamTable(
team_id="test-team",
@ -1604,6 +1605,164 @@ async def test_can_key_call_model_via_access_group_ids():
)
@pytest.mark.asyncio
async def test_can_key_call_model_empty_models_empty_access_groups_denied():
"""Test can_key_call_model fails closed (denies) when key has empty models and access_group_ids resolve to empty."""
from unittest.mock import AsyncMock, patch
from litellm.proxy._types import ProxyException
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=[],
):
with pytest.raises(ProxyException):
await can_key_call_model(
model="gpt-4",
llm_model_list=[],
valid_token=user_api_key_object,
llm_router=router,
)
@pytest.mark.asyncio
async def test_can_key_call_model_union_preserves_both_key_and_group_models():
"""Test can_key_call_model preserves union semantics when key models and access_group_ids are both set."""
from unittest.mock import AsyncMock, patch
from litellm.proxy.auth.auth_checks import can_key_call_model
user_api_key_object = UserAPIKeyAuth(
token="test-token",
models=["gpt-3.5-turbo"],
access_group_ids=["ag-claude-only"],
)
router = litellm.Router(
model_list=[
{
"model_name": "claude-3",
"litellm_params": {"model": "anthropic/claude-3", "api_key": "test"},
},
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {"model": "openai/gpt-3.5-turbo", "api_key": "test"},
},
]
)
with patch(
"litellm.proxy.auth.auth_checks.get_models_from_access_groups",
new_callable=AsyncMock,
return_value=["claude-3"],
):
# Access gpt-3.5-turbo (from key_models)
assert (
await can_key_call_model(
model="gpt-3.5-turbo",
llm_model_list=[],
valid_token=user_api_key_object,
llm_router=router,
)
is True
)
# Access claude-3 (from access_group_ids)
assert (
await can_key_call_model(
model="claude-3",
llm_model_list=[],
valid_token=user_api_key_object,
llm_router=router,
)
is True
)
@pytest.mark.asyncio
async def test_can_team_access_model_empty_models_empty_access_groups_denied():
"""Test can_team_access_model fails closed (denies) when team has empty models and access_group_ids resolve to empty."""
from unittest.mock import AsyncMock, patch
from litellm.proxy._types import ProxyException
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=[],
):
with pytest.raises(ProxyException):
await can_team_access_model(
model="gpt-4",
team_object=team_object,
llm_router=None,
team_model_aliases=None,
)
@pytest.mark.asyncio
async def test_can_team_access_model_union_preserves_both_team_and_group_models():
"""Test can_team_access_model preserves union semantics when team_models and access_group_ids are both set."""
from unittest.mock import AsyncMock, patch
from litellm.proxy.auth.auth_checks import can_team_access_model
team_object = LiteLLM_TeamTable(
team_id="test-team",
models=["gpt-4"],
access_group_ids=["ag-claude"],
)
with patch(
"litellm.proxy.auth.auth_checks.get_models_from_access_groups",
new_callable=AsyncMock,
return_value=["claude-3"],
):
# Access gpt-4 (from team_models)
assert (
await can_team_access_model(
model="gpt-4",
team_object=team_object,
llm_router=None,
team_model_aliases=None,
)
is True
)
# Access claude-3 (from access_group_ids)
assert (
await can_team_access_model(
model="claude-3",
team_object=team_object,
llm_router=None,
team_model_aliases=None,
)
is True
)
# ---------------------------------------------------------------------------
# _key_access_group_grants_model (key access group overriding team restriction)
# ---------------------------------------------------------------------------
@ -2049,7 +2208,9 @@ def test_is_model_cost_zero_judges_an_alias_chain_by_the_deployment_its_entry_ro
model_group_alias={"chain-entry": "local-free", "local-free": "paid-gpt"},
)
expected: Final = {"chain-entry": True, "local-free": False, "paid-gpt": False}
order: Final = ("chain-entry", "local-free", "paid-gpt") if entry_first else ("local-free", "paid-gpt", "chain-entry")
order: Final = (
("chain-entry", "local-free", "paid-gpt") if entry_first else ("local-free", "paid-gpt", "chain-entry")
)
verdicts: Final = {name: is_model_cost_zero(model=name, llm_router=router) for name in order}
@ -2292,10 +2453,7 @@ async def test_bulk_eviction_clears_this_worker_before_waiting_on_a_wedged_redis
monkeypatch.setattr(pubsub_module, "_BULK_EVICTION_WAIT_SECONDS", 0.5)
client = _WedgedPublishRedisClient()
cache: Final = UserApiKeyCache(key_object_in_memory_cache=InMemoryCache(max_size_in_memory=2_000))
hashed_tokens: Final = tuple(
hashlib.sha256(f"token-{index}".encode()).hexdigest()
for index in range(600)
)
hashed_tokens: Final = tuple(hashlib.sha256(f"token-{index}".encode()).hexdigest() for index in range(600))
for hashed_token in hashed_tokens:
cache.in_memory_cache_for(hashed_token).set_cache(key=hashed_token, value={"token": hashed_token})
seeded: Final = tuple(
@ -2355,10 +2513,7 @@ async def test_bulk_eviction_clears_the_redis_copies_before_waiting_on_a_wedged_
monkeypatch.setattr(pubsub_module, "_pending_publishes", set())
monkeypatch.setattr(pubsub_module, "_BULK_EVICTION_WAIT_SECONDS", 0.5)
client = _WedgedPublishRedisClient()
hashed_tokens: Final = tuple(
hashlib.sha256(f"token-{index}".encode()).hexdigest()
for index in range(600)
)
hashed_tokens: Final = tuple(hashlib.sha256(f"token-{index}".encode()).hexdigest() for index in range(600))
redis = _AuthRedisCache({token: {"token": token} for token in hashed_tokens})
cache: Final = UserApiKeyCache(key_object_in_memory_cache=InMemoryCache(max_size_in_memory=2_000))
cache.attach_redis_cache(redis) # pyright: ignore[reportArgumentType] # a dict-backed fake of the Redis layer