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
d8a47ebb69
commit
b479aa0014
3 changed files with 26 additions and 7 deletions
|
|
@ -667,6 +667,9 @@ async def _common_key_generation_helper( # noqa: PLR0915
|
|||
team_pool=team_pool,
|
||||
)
|
||||
|
||||
# When effective_models is non-empty, enforce per-member restrictions.
|
||||
# When empty (team.models=[] = "allow all" and no overrides configured),
|
||||
# skip validation — any model is permitted, matching runtime auth behavior.
|
||||
if effective_models:
|
||||
# if 'all-team-models' was requested, restrict it to the effective models
|
||||
if "all-team-models" in (data.models or []):
|
||||
|
|
|
|||
|
|
@ -2541,6 +2541,7 @@ async def team_member_update( # noqa: PLR0915
|
|||
# (it was just written to the DB by _upsert_budget_and_membership).
|
||||
# Only fetch from DB when models weren't provided (e.g., role-only update).
|
||||
stored_models = data.models if data.models is not None else []
|
||||
_tm_row = None
|
||||
if data.models is None:
|
||||
from litellm.proxy.auth.auth_checks import get_team_membership
|
||||
from litellm.proxy.proxy_server import user_api_key_cache
|
||||
|
|
@ -2557,11 +2558,24 @@ async def team_member_update( # noqa: PLR0915
|
|||
)
|
||||
stored_models = (_tm_row.models or []) if _tm_row is not None else []
|
||||
else:
|
||||
# Models were explicitly changed — invalidate cache so key-gen sees fresh data
|
||||
# Models were explicitly changed — invalidate cache and fetch fresh membership
|
||||
from litellm.proxy.auth.auth_checks import get_team_membership
|
||||
from litellm.proxy.proxy_server import user_api_key_cache
|
||||
|
||||
_cache_key = f"team_membership:{received_user_id}:{data.team_id}"
|
||||
await user_api_key_cache.async_delete_cache(key=_cache_key)
|
||||
_tm_row = await get_team_membership(
|
||||
user_id=received_user_id,
|
||||
team_id=data.team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
|
||||
# Resolve tpm/rpm: prefer request values, then budget table (authoritative),
|
||||
# then existing JSON blob (may be stale for pre-deployment members).
|
||||
_budget = _tm_row.litellm_budget_table if (_tm_row and hasattr(_tm_row, "litellm_budget_table") and _tm_row.litellm_budget_table) else None
|
||||
resolved_tpm = data.tpm_limit if data.tpm_limit is not None else (getattr(_budget, "tpm_limit", None) if _budget else None)
|
||||
resolved_rpm = data.rpm_limit if data.rpm_limit is not None else (getattr(_budget, "rpm_limit", None) if _budget else None)
|
||||
|
||||
if data.role is not None or data.models is not None:
|
||||
team_members: List[Member] = []
|
||||
|
|
@ -2573,8 +2587,8 @@ async def team_member_update( # noqa: PLR0915
|
|||
role=data.role or member.role,
|
||||
user_email=data.user_email or member.user_email,
|
||||
models=stored_models,
|
||||
tpm_limit=data.tpm_limit if data.tpm_limit is not None else getattr(member, "tpm_limit", None),
|
||||
rpm_limit=data.rpm_limit if data.rpm_limit is not None else getattr(member, "rpm_limit", None),
|
||||
tpm_limit=resolved_tpm,
|
||||
rpm_limit=resolved_rpm,
|
||||
)
|
||||
)
|
||||
else:
|
||||
|
|
|
|||
|
|
@ -234,8 +234,10 @@ async def add_new_member( # noqa: PLR0915
|
|||
else:
|
||||
_budget_id = default_team_budget_id
|
||||
|
||||
# Create membership row when budget, or models is explicitly provided (even []).
|
||||
# models=None means "not specified" (skip), models=[] means "explicitly no overrides".
|
||||
if (
|
||||
(_budget_id or new_member.models)
|
||||
(_budget_id or new_member.models is not None)
|
||||
and returned_user is not None
|
||||
and returned_user.user_id is not None
|
||||
):
|
||||
|
|
@ -245,9 +247,9 @@ async def add_new_member( # noqa: PLR0915
|
|||
}
|
||||
if _budget_id:
|
||||
membership_create_data["budget_id"] = _budget_id
|
||||
if new_member.models:
|
||||
# Defense-in-depth: validate member models are within team models
|
||||
if team_models:
|
||||
if new_member.models is not None:
|
||||
# Defense-in-depth: validate non-empty member models are within team models
|
||||
if new_member.models and team_models:
|
||||
disallowed = set(new_member.models) - set(team_models)
|
||||
if disallowed:
|
||||
raise HTTPException(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue