Merge pull request #1570 from BerriAI/litellm_reset_key_budget

feat(proxy/utils.py): Support keys with budget duration (resets spend on duration end)

.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:

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):

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"]

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 {

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"
+            )

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

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 = [

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

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("")
-}
+}