diff --git a/.circleci/config.yml b/.circleci/config.yml index c7417a62f28..8685f45793a 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -42,6 +42,7 @@ jobs: pip install "anyio==3.7.1" pip install "aiodynamo==23.10.1" pip install "asyncio==3.4.3" + pip install "apscheduler==3.10.4" pip install "PyGithub==1.59.1" - save_cache: paths: diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index bb56ad6bf1b..d5dc841cb26 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -135,6 +135,7 @@ class GenerateKeyRequest(LiteLLMBase): metadata: Optional[dict] = {} tpm_limit: Optional[int] = None rpm_limit: Optional[int] = None + budget_duration: Optional[str] = None class UpdateKeyRequest(LiteLLMBase): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ae761c0335a..874731f1dec 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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 @@ -1130,7 +1132,9 @@ 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 + budget_duration: Optional[str] = None, # max_budget is used to Budget Per user token: Optional[str] = None, user_id: Optional[str] = None, team_id: Optional[str] = None, @@ -1175,6 +1179,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) @@ -1210,6 +1220,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) @@ -1530,7 +1542,7 @@ async def startup_event(): duration=None, models=[], aliases={}, config={}, spend=0, token=master_key ) verbose_proxy_logger.debug( - f"custom_db_client client - Inserting master key {custom_db_client}. Master_key: {master_key}" + f"custom_db_client client {custom_db_client}. Master_key: {master_key}" ) if custom_db_client is not None and master_key is not None: # add master key to db @@ -1538,6 +1550,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( @@ -2215,6 +2232,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"] diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 931a1581258..ea3bade8cd6 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -34,6 +34,8 @@ model LiteLLM_VerificationToken { tpm_limit BigInt? rpm_limit BigInt? max_budget Float? @default(0.0) + budget_duration String? + budget_reset_at DateTime? } model LiteLLM_Config { diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 9aef0304c62..1b3581427b7 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -14,10 +14,10 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.proxy.db.base_client import CustomDB from litellm._logging import verbose_proxy_logger from fastapi import HTTPException, status -import smtplib +import smtplib, re from email.mime.text import MIMEText from email.mime.multipart import MIMEMultipart -from datetime import datetime +from datetime import datetime, timedelta def print_verbose(print_statement): @@ -364,6 +364,8 @@ class PrismaClient: request_id: Optional[str] = None, table_name: Optional[Literal["user", "key", "config", "spend"]] = None, query_type: Literal["find_unique", "find_all"] = "find_unique", + expires: Optional[datetime] = None, + reset_at: Optional[datetime] = None, ): try: print_verbose("PrismaClient: get_data") @@ -392,6 +394,24 @@ class PrismaClient: for r in response: if isinstance(r.expires, datetime): r.expires = r.expires.isoformat() + elif ( + query_type == "find_all" + and expires is not None + and reset_at is not None + ): + response = await self.db.litellm_verificationtoken.find_many( + where={ # type:ignore + "OR": [ + {"expires": None}, + {"expires": {"gt": expires}}, + ], + "budget_reset_at": {"lt": reset_at}, + } + ) + if response is not None and len(response) > 0: + for r in response: + if isinstance(r.expires, datetime): + r.expires = r.expires.isoformat() elif query_type == "find_all": response = await self.db.litellm_verificationtoken.find_many( order={"spend": "desc"}, @@ -539,7 +559,10 @@ class PrismaClient: self, token: Optional[str] = None, data: dict = {}, + data_list: Optional[List] = None, user_id: Optional[str] = None, + query_type: Literal["update", "update_many"] = "update", + table_name: Optional[Literal["user", "key", "config", "spend"]] = None, ): """ Update existing data @@ -588,6 +611,33 @@ class PrismaClient: + "\033[0m" ) return {"user_id": user_id, "data": db_data} + elif ( + table_name is not None + and table_name == "key" + and query_type == "update_many" + and data_list is not None + and isinstance(data_list, list) + ): + """ + Batch write update queries + """ + batcher = self.db.batch_() + for idx, t in enumerate(data_list): + # check if plain text or hash + if t.token.startswith("sk-"): # type: ignore + t.token = self.hash_token(token=t.token) # type: ignore + try: + data_json = self.jsonify_object(data=t.model_dump()) + except: + data_json = self.jsonify_object(data=t.dict()) + batcher.litellm_verificationtoken.update( + where={"token": t.token}, # type: ignore + data={**data_json}, # type: ignore + ) + await batcher.commit() + print_verbose( + "\033[91m" + f"DB Token Table update succeeded" + "\033[0m" + ) except Exception as e: asyncio.create_task( self.proxy_logging_obj.failure_handler(original_exception=e) @@ -913,3 +963,48 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time): payload[param] = str(payload[param]) return payload + + +def _duration_in_seconds(duration: str): + match = re.match(r"(\d+)([smhd]?)", duration) + if not match: + raise ValueError("Invalid duration format") + + value, unit = match.groups() + value = int(value) + + if unit == "s": + return value + elif unit == "m": + return value * 60 + elif unit == "h": + return value * 3600 + elif unit == "d": + return value * 86400 + else: + raise ValueError("Unsupported duration unit") + + +async def reset_budget(prisma_client: PrismaClient): + """ + Gets all the non-expired keys for a db, which need spend to be reset + + Resets their spend + + Updates db + """ + if prisma_client is not None: + now = datetime.utcnow() + keys_to_reset = await prisma_client.get_data( + table_name="key", query_type="find_all", expires=now, reset_at=now + ) + + for key in keys_to_reset: + key.spend = 0.0 + duration_s = _duration_in_seconds(duration=key.budget_duration) + key.budget_reset_at = key.budget_reset_at + timedelta(seconds=duration_s) + + if len(keys_to_reset) > 0: + await prisma_client.update_data( + query_type="update_many", data_list=keys_to_reset, table_name="key" + ) diff --git a/litellm/utils.py b/litellm/utils.py index 7a6b12a820c..03d38ff35a2 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1074,7 +1074,7 @@ class Logging: or isinstance(result, EmbeddingResponse) ) and self.stream != True - ): + ): # handle streaming separately try: self.model_call_details["response_cost"] = litellm.completion_cost( completion_response=result, @@ -1137,7 +1137,7 @@ class Logging: else: self.sync_streaming_chunks.append(result) - if complete_streaming_response: + if complete_streaming_response is not None: verbose_logger.debug( f"Logging Details LiteLLM-Success Call streaming complete" ) @@ -1436,7 +1436,7 @@ class Logging: complete_streaming_response = None else: self.streaming_chunks.append(result) - if complete_streaming_response: + if complete_streaming_response is not None: print_verbose("Async success callbacks: Got a complete streaming response") self.model_call_details[ "complete_streaming_response" @@ -2910,6 +2910,9 @@ def cost_per_token( if model in model_cost_ref: verbose_logger.debug(f"Success: model={model} in model_cost_map") + verbose_logger.debug( + f"prompt_tokens={prompt_tokens}; completion_tokens={completion_tokens}" + ) if ( model_cost_ref[model].get("input_cost_per_token", None) is not None and model_cost_ref[model].get("output_cost_per_token", None) is not None diff --git a/pyproject.toml b/pyproject.toml index 2be49d95d9f..cd21db9035c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,6 +25,7 @@ backoff = {version = "*", optional = true} pyyaml = {version = "^6.0.1", optional = true} rq = {version = "*", optional = true} orjson = {version = "^3.9.7", optional = true} +apscheduler = {version = "^3.10.4", optional = true} streamlit = {version = "^1.29.0", optional = true} [tool.poetry.extras] @@ -36,6 +37,7 @@ proxy = [ "pyyaml", "rq", "orjson", + "apscheduler" ] extra_proxy = [ diff --git a/requirements.txt b/requirements.txt index 662dafd06ae..6103091b857 100644 --- a/requirements.txt +++ b/requirements.txt @@ -16,6 +16,7 @@ async_generator==1.10.0 # for async ollama calls traceloop-sdk==0.5.3 # for open telemetry logging langfuse>=2.6.3 # for langfuse self-hosted logging orjson==3.9.7 # fast /embedding responses +apscheduler==3.10.4 # for resetting budget in background ### LITELLM PACKAGE DEPENDENCIES python-dotenv>=0.2.0 # for env tiktoken>=0.4.0 # for calculating usage diff --git a/schema.prisma b/schema.prisma index 1212b0c661a..ea3bade8cd6 100644 --- a/schema.prisma +++ b/schema.prisma @@ -34,6 +34,8 @@ model LiteLLM_VerificationToken { tpm_limit BigInt? rpm_limit BigInt? max_budget Float? @default(0.0) + budget_duration String? + budget_reset_at DateTime? } model LiteLLM_Config { @@ -43,8 +45,8 @@ model LiteLLM_Config { model LiteLLM_SpendLogs { request_id String @unique - api_key String @default ("") call_type String + api_key String @default ("") spend Float @default(0.0) startTime DateTime // Assuming start_time is a DateTime field endTime DateTime // Assuming end_time is a DateTime field @@ -56,4 +58,4 @@ model LiteLLM_SpendLogs { usage Json @default("{}") metadata Json @default("{}") cache_hit String @default("") -} +} \ No newline at end of file