phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 11:14:04 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(user_api_key_auth.py): respect team budgets over user budget, if key belongs to team

Browse source 
Closes https://github.com/BerriAI/litellm/issues/5097
This commit is contained in:
Krrish Dholakia 2024-08-07 14:32:27 -07:00
parent f579aef740
commit d832327ccf
3 changed files with 92 additions and 110 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -55,11 +55,11 @@ def common_checks(
1. If team is blocked
2. If team can call model
3. If team is in budget
5. If user passed in (JWT or key.user_id) - is in budget
4. If end_user (either via JWT or 'user' passed to /chat/completions, /embeddings endpoint) is in budget
5. [OPTIONAL] If 'enforce_end_user' enabled - did developer pass in 'user' param for openai endpoints
6. [OPTIONAL] If 'litellm.max_budget' is set (>0), is proxy under budget
7. [OPTIONAL] If guardrails modified - is request allowed to change this
4. If user passed in (JWT or key.user_id) - is in budget
5. If end_user (either via JWT or 'user' passed to /chat/completions, /embeddings endpoint) is in budget
6. [OPTIONAL] If 'enforce_end_user' enabled - did developer pass in 'user' param for openai endpoints
7. [OPTIONAL] If 'litellm.max_budget' is set (>0), is proxy under budget
8. [OPTIONAL] If guardrails modified - is request allowed to change this
"""
_model = request_body.get("model", None)
if team_object is not None and team_object.blocked is True:
@ -91,12 +91,19 @@ def common_checks(
raise Exception(
f"Team={team_object.team_id} over budget. Spend={team_object.spend}, Budget={team_object.max_budget}"
)
if user_object is not None and user_object.max_budget is not None:
# 4. If user is in budget
## 4.1 check personal budget, if personal key
if (
(team_object is None or team_object.team_id is None)
and user_object is not None
and user_object.max_budget is not None
):
user_budget = user_object.max_budget
if user_budget > user_object.spend:
if user_budget < user_object.spend:
raise Exception(
f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_object.spend}, Budget={user_budget}"
)
## 4.2 check team member budget, if team key
# 5. If end_user ('user' passed to /chat/completions, /embeddings endpoint) is in budget
if end_user_object is not None and end_user_object.litellm_budget_table is not None:
end_user_budget = end_user_object.litellm_budget_table.max_budget