mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
fix: req changes by greptile
This commit is contained in:
parent
237d8354a0
commit
0199f21119
3 changed files with 9 additions and 95 deletions
|
|
@ -2700,7 +2700,8 @@ def compute_effective_models(
|
|||
- If cap empties the list (all stale), falls back to team_pool (NOT [] which = allow-all).
|
||||
- team_pool=[] means "allow all" — cap is skipped.
|
||||
"""
|
||||
effective = list(set(team_defaults + member_models))
|
||||
# dict.fromkeys preserves insertion order while deduplicating
|
||||
effective = list(dict.fromkeys(team_defaults + member_models))
|
||||
|
||||
if not effective:
|
||||
return team_pool
|
||||
|
|
|
|||
|
|
@ -2528,6 +2528,13 @@ async def team_member_update( # noqa: PLR0915
|
|||
models=data.models,
|
||||
)
|
||||
|
||||
# Invalidate the get_team_membership cache so subsequent key-gen calls
|
||||
# see the updated models immediately (not stale cached values).
|
||||
from litellm.proxy.proxy_server import user_api_key_cache
|
||||
|
||||
_cache_key = f"team_membership:{received_user_id}:{data.team_id}"
|
||||
user_api_key_cache.delete_cache(key=_cache_key)
|
||||
|
||||
### update team member role
|
||||
# Resolve the effective models for this member (from the authoritative
|
||||
# LiteLLM_TeamMembership table via cached helper) so we can: (a) keep the
|
||||
|
|
|
|||
|
|
@ -1,94 +0,0 @@
|
|||
import sys
|
||||
import os
|
||||
import pytest
|
||||
|
||||
# Add the parent directory to the system path to import litellm
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
|
||||
|
||||
import litellm
|
||||
from litellm.proxy._types import UserAPIKeyAuth, LiteLLM_TeamTable
|
||||
from litellm.proxy.auth.auth_checks import (
|
||||
can_team_access_model,
|
||||
get_effective_team_models,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_effective_team_models():
|
||||
original_flag = litellm.team_model_overrides_enabled
|
||||
original_env = os.environ.pop("TEAM_MODEL_OVERRIDES", None)
|
||||
try:
|
||||
litellm.team_model_overrides_enabled = True
|
||||
|
||||
# Case 1: No overrides, should return team.models
|
||||
team = LiteLLM_TeamTable(team_id="t1", models=["m1"])
|
||||
assert get_effective_team_models(team) == ["m1"]
|
||||
|
||||
# Case 2: Team defaults exist (d1 must be in team.models pool)
|
||||
team = LiteLLM_TeamTable(team_id="t1", models=["m1", "d1"], default_models=["d1"])
|
||||
assert set(get_effective_team_models(team)) == {"d1"}
|
||||
|
||||
# Case 3: Team defaults + Member overrides (all in team.models pool)
|
||||
team = LiteLLM_TeamTable(
|
||||
team_id="t1", models=["m1", "d1", "mo1"], default_models=["d1"]
|
||||
)
|
||||
token = UserAPIKeyAuth(team_member_models=["mo1"])
|
||||
assert set(get_effective_team_models(team, token)) == {"d1", "mo1"}
|
||||
|
||||
# Case 4: No team object (should use token values if available)
|
||||
token.team_default_models = ["td1"]
|
||||
assert set(get_effective_team_models(None, token)) == {"td1", "mo1"}
|
||||
|
||||
# Case 5: Feature disabled — also ensure env var is cleared
|
||||
litellm.team_model_overrides_enabled = False
|
||||
os.environ.pop("TEAM_MODEL_OVERRIDES", None)
|
||||
assert get_effective_team_models(team, token) == ["m1", "d1", "mo1"]
|
||||
finally:
|
||||
litellm.team_model_overrides_enabled = original_flag
|
||||
if original_env is not None:
|
||||
os.environ["TEAM_MODEL_OVERRIDES"] = original_env
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_can_team_access_model_with_overrides():
|
||||
original_flag = litellm.team_model_overrides_enabled
|
||||
try:
|
||||
litellm.team_model_overrides_enabled = True
|
||||
|
||||
# Team pool includes m1, d1, g1. default_models=["d1"].
|
||||
team = LiteLLM_TeamTable(
|
||||
team_id="t1", models=["m1", "d1", "g1"], default_models=["d1"]
|
||||
)
|
||||
|
||||
# With only defaults, should NOT have access to m1
|
||||
with pytest.raises(Exception):
|
||||
await can_team_access_model(model="m1", team_object=team, llm_router=None)
|
||||
|
||||
# Should have access to d1 (it's a default)
|
||||
assert (
|
||||
await can_team_access_model(model="d1", team_object=team, llm_router=None)
|
||||
is True
|
||||
)
|
||||
|
||||
# Member has extra access to g1
|
||||
token = UserAPIKeyAuth(team_member_models=["g1"])
|
||||
assert (
|
||||
await can_team_access_model(
|
||||
model="g1", team_object=team, llm_router=None, valid_token=token
|
||||
)
|
||||
is True
|
||||
)
|
||||
assert (
|
||||
await can_team_access_model(
|
||||
model="d1", team_object=team, llm_router=None, valid_token=token
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
# Should NOT have access to m1
|
||||
with pytest.raises(Exception):
|
||||
await can_team_access_model(
|
||||
model="m1", team_object=team, llm_router=None, valid_token=token
|
||||
)
|
||||
finally:
|
||||
litellm.team_model_overrides_enabled = original_flag
|
||||
Loading…
Add table
Reference in a new issue