diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 175f801dac1..e4b28001d57 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -324,6 +324,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..37b28baeafa 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1869,6 +1869,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=50, + 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 +1919,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