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

Merge pull request #2208 from BerriAI/litellm_enforce_team_limits

Browse source 
Litellm enforce team limits
This commit is contained in:
Krish Dholakia 2024-02-26 23:10:01 -08:00 • committed by GitHub
commit 365e7ed5b9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 302 additions and 47 deletions
Download patch file Download diff file Expand all files Collapse all files

59
litellm/proxy/proxy_server.py
View file

@ -356,7 +356,7 @@ async def user_api_key_auth(
verbose_proxy_logger.debug(f"api key: {api_key}")
if prisma_client is not None:
valid_token = await prisma_client.get_data(
token=api_key,
token=api_key, table_name="combined_view"
)
elif custom_db_client is not None:
@ -381,7 +381,8 @@ async def user_api_key_auth(
# 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
request_data = await _read_request_body(
request=request
) # request data, used across all checks. Making this easily available
@ -610,6 +611,47 @@ async def user_api_key_auth(
f"ExceededModelBudget: Current spend for model: {current_model_spend}; Max Budget for Model: {current_model_budget}"
)
# Check 6. Token spend is under Team budget
if (
valid_token.spend is not None
and hasattr(valid_token, "team_max_budget")
and valid_token.team_max_budget is not None
):
asyncio.create_task(
proxy_logging_obj.budget_alerts(
user_max_budget=valid_token.team_max_budget,
user_current_spend=valid_token.spend,
type="token_budget",
user_info=valid_token,
)
)
if valid_token.spend >= valid_token.team_max_budget:
raise Exception(
f"ExceededTokenBudget: Current spend for token: {valid_token.spend}; Max Budget for Team: {valid_token.team_max_budget}"
)
# Check 7. Team spend is under Team budget
if (
hasattr(valid_token, "team_spend")
and valid_token.team_spend is not None
and hasattr(valid_token, "team_max_budget")
and valid_token.team_max_budget is not None
):
asyncio.create_task(
proxy_logging_obj.budget_alerts(
user_max_budget=valid_token.team_max_budget,
user_current_spend=valid_token.team_spend,
type="token_budget",
user_info=valid_token,
)
)
if valid_token.team_spend >= valid_token.team_max_budget:
raise Exception(
f"ExceededTokenBudget: Current Team Spend: {valid_token.team_spend}; Max Budget for Team: {valid_token.team_max_budget}"
)
# Token passed all checks
api_key = valid_token.token
@ -1870,14 +1912,6 @@ async def generate_key_helper_fn(
saved_token["expires"], datetime
):
saved_token["expires"] = saved_token["expires"].isoformat()
if key_data["token"] is not None and isinstance(key_data["token"], str):
hashed_token = hash_token(key_data["token"])
saved_token["token"] = hashed_token
user_api_key_cache.set_cache(
key=hashed_token,
value=LiteLLM_VerificationToken(**saved_token), # type: ignore
ttl=600,
)
if prisma_client is not None:
## CREATE USER (If necessary)
verbose_proxy_logger.debug(f"prisma_client: Creating User={user_data}")
@ -2263,6 +2297,11 @@ async def startup_event():
duration=None, models=[], aliases={}, config={}, spend=0, token=master_key
)
### CHECK IF VIEW EXISTS ###
if prisma_client is not None:
create_view_response = await prisma_client.check_view_exists()
print(f"create_view_response: {create_view_response}") # noqa
### START BUDGET SCHEDULER ###
if prisma_client is not None:
scheduler = AsyncIOScheduler()