diff --git a/docs/my-website/docs/proxy/virtual_keys.md b/docs/my-website/docs/proxy/virtual_keys.md index e350ce9d57b..70fd6e6a8d9 100644 --- a/docs/my-website/docs/proxy/virtual_keys.md +++ b/docs/my-website/docs/proxy/virtual_keys.md @@ -79,6 +79,7 @@ curl 'http://0.0.0.0:8000/key/generate' \ "metadata": {"user": "ishaan@berri.ai"}, "team_id": "core-infra", "max_budget": 10, + "soft_budget": 5, }' ``` @@ -93,6 +94,7 @@ Request Params: - `config`: *Optional[dict]* - any key-specific configs, overrides config in config.yaml - `spend`: *Optional[int]* - Amount spent by key. Default is 0. Will be updated by proxy whenever key is used. https://docs.litellm.ai/docs/proxy/virtual_keys#managing-auth---tracking-spend - `max_budget`: *Optional[float]* - Specify max budget for a given key. +- `soft_budget`: *Optional[float]* - Specify soft limit budget for a given key. Get Alerts when key hits its soft budget - `model_max_budget`: *Optional[dict[str, float]]* - Specify max budget for each model, `model_max_budget={"gpt4": 0.5, "gpt-5": 0.01}` - `max_parallel_requests`: *Optional[int]* - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - `metadata`: *Optional[dict]* - Metadata for key, store information for key. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } diff --git a/litellm/__init__.py b/litellm/__init__.py index cd639ddb9b7..f218fe036e9 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -79,6 +79,9 @@ max_budget: float = 0.0 # set the max budget across all providers budget_duration: Optional[str] = ( None # proxy only - resets budget after fixed duration. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). ) +default_soft_budget: float = ( + 50.0 # by default all litellm proxy keys have a soft budget of 50.0 +) _openai_finish_reasons = ["stop", "length", "function_call", "content_filter", "null"] _openai_completion_params = [ "functions", diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 175f801dac1..6196f18a2ee 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -151,6 +151,7 @@ class GenerateRequestBase(LiteLLMBase): rpm_limit: Optional[int] = None budget_duration: Optional[str] = None allowed_cache_controls: Optional[list] = [] + soft_budget: Optional[float] = None class GenerateKeyRequest(GenerateRequestBase): @@ -324,6 +325,21 @@ class TeamRequest(LiteLLMBase): teams: List[str] +class LiteLLM_BudgetTable(LiteLLMBase): + """Represents user-controllable params for a LiteLLM_BudgetTable record""" + + max_budget: Optional[float] = None + soft_budget: Optional[float] = None + max_parallel_requests: Optional[int] = None + tpm_limit: Optional[int] = None + rpm_limit: Optional[int] = None + model_max_budget: dict + budget_duration: Optional[str] = None + budget_reset_at: Optional[datetime] = None + created_by: str + updated_by: str + + class KeyManagementSystem(enum.Enum): GOOGLE_KMS = "google_kms" AZURE_KEY_VAULT = "azure_key_vault" diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 17db8c3ab4f..482397b86e5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1810,6 +1810,9 @@ async def generate_key_helper_fn( spend: float, key_max_budget: Optional[float] = None, # key_max_budget is used to Budget Per key key_budget_duration: Optional[str] = None, + key_soft_budget: Optional[ + float + ] = None, # key_soft_budget is used to Budget Per key 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, @@ -1869,6 +1872,19 @@ async def generate_key_helper_fn( rpm_limit = rpm_limit allowed_cache_controls = allowed_cache_controls + # TODO: @ishaan-jaff: Migrate all budget tracking to use LiteLLM_BudgetTable + if prisma_client is not None: + # create the Budget Row for the LiteLLM Verification Token + budget_row = LiteLLM_BudgetTable( + soft_budget=key_soft_budget or litellm.default_soft_budget, + model_max_budget=model_max_budget or {}, + created_by=user_id, + updated_by=user_id, + ) + new_budget = prisma_client.jsonify_object(budget_row.json(exclude_none=True)) + _budget = await prisma_client.db.litellm_budgettable.create(data={**new_budget}) # type: ignore + _budget_id = getattr(_budget, "id", None) + try: # Create a new verification token (you may want to enhance this logic based on your needs) user_data = { @@ -1906,6 +1922,7 @@ async def generate_key_helper_fn( "allowed_cache_controls": allowed_cache_controls, "permissions": permissions_json, "model_max_budget": model_max_budget_json, + "budget_id": _budget_id, } if ( general_settings.get("allow_user_auth", False) == True @@ -1978,6 +1995,9 @@ async def generate_key_helper_fn( except Exception as e: traceback.print_exc() raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR) + + # Add budget related info in key_data - this ensures it's returned + key_data["soft_budget"] = key_soft_budget return key_data @@ -3333,6 +3353,8 @@ async def generate_key_fn( # if we get max_budget passed to /key/generate, then use it as key_max_budget. Since generate_key_helper_fn is used to make new users if "max_budget" in data_json: data_json["key_max_budget"] = data_json.pop("max_budget", None) + if "soft_budget" in data_json: + data_json["key_soft_budget"] = data_json.pop("soft_budget", None) if "budget_duration" in data_json: data_json["key_budget_duration"] = data_json.pop("budget_duration", None) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 7eb59ee4839..1fe55f24e34 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -11,6 +11,7 @@ generator client { model LiteLLM_BudgetTable { budget_id String @id @default(uuid()) max_budget Float? + soft_budget Float? max_parallel_requests Int? tpm_limit BigInt? rpm_limit BigInt? @@ -107,6 +108,7 @@ model LiteLLM_VerificationToken { allowed_cache_controls String[] @default([]) model_spend Json @default("{}") model_max_budget Json @default("{}") + budget_id String? } // store proxy config.yaml diff --git a/schema.prisma b/schema.prisma index 7eb59ee4839..1fe55f24e34 100644 --- a/schema.prisma +++ b/schema.prisma @@ -11,6 +11,7 @@ generator client { model LiteLLM_BudgetTable { budget_id String @id @default(uuid()) max_budget Float? + soft_budget Float? max_parallel_requests Int? tpm_limit BigInt? rpm_limit BigInt? @@ -107,6 +108,7 @@ model LiteLLM_VerificationToken { allowed_cache_controls String[] @default([]) model_spend Json @default("{}") model_max_budget Json @default("{}") + budget_id String? } // store proxy config.yaml