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

feat(proxy/utils.py): enable background process to reset key budgets

Browse source
This commit is contained in:
Krrish Dholakia 2024-01-23 12:33:13 -08:00
parent 43ea1e589f
commit e6f05e858b
5 changed files with 138 additions and 18 deletions
Download patch file Download diff file Expand all files Collapse all files

19
litellm/proxy/proxy_server.py
View file

@ -19,6 +19,7 @@ try:
import yaml
import orjson
import logging
from apscheduler.schedulers.asyncio import AsyncIOScheduler
except ImportError as e:
raise ImportError(f"Missing dependency {e}. Run `pip install 'litellm[proxy]'`")
@ -73,6 +74,7 @@ from litellm.proxy.utils import (
_cache_user_row,
send_email,
get_logging_payload,
reset_budget,
)
from litellm.proxy.secret_managers.google_kms import load_google_kms
import pydantic
@ -1125,6 +1127,7 @@ async def generate_key_helper_fn(
config: dict,
spend: float,
key_max_budget: Optional[float] = None, # key_max_budget is used to Budget Per key
key_budget_duration: Optional[str] = None,
max_budget: Optional[float] = None, # max_budget is used to Budget Per user
token: Optional[str] = None,
user_id: Optional[str] = None,
@ -1170,6 +1173,12 @@ async def generate_key_helper_fn(
duration_s = _duration_in_seconds(duration=duration)
expires = datetime.utcnow() + timedelta(seconds=duration_s)
if key_budget_duration is None: # one-time budget
key_reset_at = None
else:
duration_s = _duration_in_seconds(duration=key_budget_duration)
key_reset_at = datetime.utcnow() + timedelta(seconds=duration_s)
aliases_json = json.dumps(aliases)
config_json = json.dumps(config)
metadata_json = json.dumps(metadata)
@ -1205,6 +1214,8 @@ async def generate_key_helper_fn(
"metadata": metadata_json,
"tpm_limit": tpm_limit,
"rpm_limit": rpm_limit,
"budget_duration": key_budget_duration,
"budget_reset_at": key_reset_at,
}
if prisma_client is not None:
## CREATE USER (If necessary)
@ -1511,6 +1522,11 @@ async def startup_event():
duration=None, models=[], aliases={}, config={}, spend=0, token=master_key
)
### START BUDGET SCHEDULER ###
scheduler = AsyncIOScheduler()
scheduler.add_job(reset_budget, "interval", seconds=10, args=[prisma_client])
scheduler.start()
#### API ENDPOINTS ####
@router.get(
@ -2186,6 +2202,9 @@ async def generate_key_fn(
if "max_budget" in data_json:
data_json["key_max_budget"] = data_json.pop("max_budget", None)
if "budget_duration" in data_json:
data_json["key_budget_duration"] = data_json.pop("budget_duration", None)
response = await generate_key_helper_fn(**data_json)
return GenerateKeyResponse(
key=response["token"], expires=response["expires"], user_id=response["user_id"]