phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

feat - check team member spend, budget

Browse source
This commit is contained in:
Ishaan Jaff 2024-05-22 18:36:39 -07:00
parent c2d25b9a14
commit 7a8d78e6c9
3 changed files with 61 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

50
litellm/proxy/proxy_server.py
View file

@ -798,12 +798,13 @@ async def user_api_key_auth(
# Run checks for
# 1. If token can call model
# 2. If user_id for this token is in budget
# 3. If 'user' passed to /chat/completions, /embeddings endpoint is in budget
# 4. If token is expired
# 5. If token spend is under Budget for the token
# 6. If token spend per model is under budget per model
# 7. If token spend is under team budget
# 8. If team spend is under team budget
# 3. If the user spend within their own team is within budget
# 4. If 'user' passed to /chat/completions, /embeddings endpoint is in budget
# 5. If token is expired
# 6. If token spend is under Budget for the token
# 7. If token spend per model is under budget per model
# 8. If token spend is under team budget
# 9. If team spend is under team budget
# Check 1. If token can call model
_model_alias_map = {}
@ -1001,6 +1002,43 @@ async def user_api_key_auth(
raise Exception(
f"ExceededBudget: User {valid_token.user_id} has exceeded their budget. Current spend: {user_current_spend}; Max Budget: {user_max_budget}"
)
# Check 3. Check if user is in their team budget
if valid_token.team_member_spend is not None:
if prisma_client is not None:
_cache_key = f"{valid_token.team_id}_{valid_token.user_id}"
team_member_info = await user_api_key_cache.async_get_cache(
key=_cache_key
)
if team_member_info is None:
team_member_info = (
await prisma_client.db.litellm_teammembership.find_first(
where={
"user_id": valid_token.user_id,
"team_id": valid_token.team_id,
}, # type: ignore
include={"litellm_budget_table": True},
)
)
await user_api_key_cache.async_set_cache(
key=_cache_key,
value=team_member_info,
ttl=UserAPIKeyCacheTTLEnum.user_information_cache.value,
)
if (
team_member_info is not None
and team_member_info.litellm_budget_table is not None
):
team_member_budget = (
team_member_info.litellm_budget_table.max_budget
)
if team_member_budget is not None and team_member_budget > 0:
if valid_token.team_member_spend > team_member_budget:
raise Exception(
f"ExceededBudget: Crossed spend within team. UserID: {valid_token.user_id}, in team {valid_token.team_id} has exceeded their budget. Current spend: {valid_token.team_member_spend}; Max Budget: {team_member_budget}"
)
# Check 3. If token is expired
if valid_token.expires is not None: