diff options
Diffstat (limited to '.venv/lib/python3.12/site-packages/litellm/proxy/management_endpoints/budget_management_endpoints.py')
-rw-r--r-- | .venv/lib/python3.12/site-packages/litellm/proxy/management_endpoints/budget_management_endpoints.py | 287 |
1 files changed, 287 insertions, 0 deletions
diff --git a/.venv/lib/python3.12/site-packages/litellm/proxy/management_endpoints/budget_management_endpoints.py b/.venv/lib/python3.12/site-packages/litellm/proxy/management_endpoints/budget_management_endpoints.py new file mode 100644 index 00000000..20aa1c6b --- /dev/null +++ b/.venv/lib/python3.12/site-packages/litellm/proxy/management_endpoints/budget_management_endpoints.py @@ -0,0 +1,287 @@ +""" +BUDGET MANAGEMENT + +All /budget management endpoints + +/budget/new +/budget/info +/budget/update +/budget/delete +/budget/settings +/budget/list +""" + +#### BUDGET TABLE MANAGEMENT #### +from fastapi import APIRouter, Depends, HTTPException + +from litellm.proxy._types import * +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.utils import jsonify_object + +router = APIRouter() + + +@router.post( + "/budget/new", + tags=["budget management"], + dependencies=[Depends(user_api_key_auth)], +) +async def new_budget( + budget_obj: BudgetNewRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Create a new budget object. Can apply this to teams, orgs, end-users, keys. + + Parameters: + - budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.) + - budget_id: Optional[str] - The id of the budget. If not provided, a new id will be generated. + - max_budget: Optional[float] - The max budget for the budget. + - soft_budget: Optional[float] - The soft budget for the budget. + - max_parallel_requests: Optional[int] - The max number of parallel requests for the budget. + - tpm_limit: Optional[int] - The tokens per minute limit for the budget. + - rpm_limit: Optional[int] - The requests per minute limit for the budget. + - model_max_budget: Optional[dict] - Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d", "tpm_limit": 100000, "rpm_limit": 100000}} + """ + from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + + budget_obj_json = budget_obj.model_dump(exclude_none=True) + budget_obj_jsonified = jsonify_object(budget_obj_json) # json dump any dictionaries + response = await prisma_client.db.litellm_budgettable.create( + data={ + **budget_obj_jsonified, # type: ignore + "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, + "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, + } # type: ignore + ) + + return response + + +@router.post( + "/budget/update", + tags=["budget management"], + dependencies=[Depends(user_api_key_auth)], +) +async def update_budget( + budget_obj: BudgetNewRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Update an existing budget object. + + Parameters: + - budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.) + - budget_id: Optional[str] - The id of the budget. If not provided, a new id will be generated. + - max_budget: Optional[float] - The max budget for the budget. + - soft_budget: Optional[float] - The soft budget for the budget. + - max_parallel_requests: Optional[int] - The max number of parallel requests for the budget. + - tpm_limit: Optional[int] - The tokens per minute limit for the budget. + - rpm_limit: Optional[int] - The requests per minute limit for the budget. + - model_max_budget: Optional[dict] - Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d", "tpm_limit": 100000, "rpm_limit": 100000}} + """ + from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + if budget_obj.budget_id is None: + raise HTTPException(status_code=400, detail={"error": "budget_id is required"}) + + response = await prisma_client.db.litellm_budgettable.update( + where={"budget_id": budget_obj.budget_id}, + data={ + **budget_obj.model_dump(exclude_none=True), # type: ignore + "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, + }, # type: ignore + ) + + return response + + +@router.post( + "/budget/info", + tags=["budget management"], + dependencies=[Depends(user_api_key_auth)], +) +async def info_budget(data: BudgetRequest): + """ + Get the budget id specific information + + Parameters: + - budgets: List[str] - The list of budget ids to get information for + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException(status_code=500, detail={"error": "No db connected"}) + + if len(data.budgets) == 0: + raise HTTPException( + status_code=400, + detail={ + "error": f"Specify list of budget id's to query. Passed in={data.budgets}" + }, + ) + response = await prisma_client.db.litellm_budgettable.find_many( + where={"budget_id": {"in": data.budgets}}, + ) + + return response + + +@router.get( + "/budget/settings", + tags=["budget management"], + dependencies=[Depends(user_api_key_auth)], +) +async def budget_settings( + budget_id: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Get list of configurable params + current value for a budget item + description of each field + + Used on Admin UI. + + Query Parameters: + - budget_id: str - The budget id to get information for + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=400, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=400, + detail={ + "error": "{}, your role={}".format( + CommonProxyErrors.not_allowed_access.value, + user_api_key_dict.user_role, + ) + }, + ) + + ## get budget item from db + db_budget_row = await prisma_client.db.litellm_budgettable.find_first( + where={"budget_id": budget_id} + ) + + if db_budget_row is not None: + db_budget_row_dict = db_budget_row.model_dump(exclude_none=True) + else: + db_budget_row_dict = {} + + allowed_args = { + "max_parallel_requests": {"type": "Integer"}, + "tpm_limit": {"type": "Integer"}, + "rpm_limit": {"type": "Integer"}, + "budget_duration": {"type": "String"}, + "max_budget": {"type": "Float"}, + "soft_budget": {"type": "Float"}, + } + + return_val = [] + + for field_name, field_info in BudgetNewRequest.model_fields.items(): + if field_name in allowed_args: + + _stored_in_db = True + + _response_obj = ConfigList( + field_name=field_name, + field_type=allowed_args[field_name]["type"], + field_description=field_info.description or "", + field_value=db_budget_row_dict.get(field_name, None), + stored_in_db=_stored_in_db, + field_default_value=field_info.default, + ) + return_val.append(_response_obj) + + return return_val + + +@router.get( + "/budget/list", + tags=["budget management"], + dependencies=[Depends(user_api_key_auth)], +) +async def list_budget( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """List all the created budgets in proxy db. Used on Admin UI.""" + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=400, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=400, + detail={ + "error": "{}, your role={}".format( + CommonProxyErrors.not_allowed_access.value, + user_api_key_dict.user_role, + ) + }, + ) + + response = await prisma_client.db.litellm_budgettable.find_many() + + return response + + +@router.post( + "/budget/delete", + tags=["budget management"], + dependencies=[Depends(user_api_key_auth)], +) +async def delete_budget( + data: BudgetDeleteRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Delete budget + + Parameters: + - id: str - The budget id to delete + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=400, + detail={ + "error": "{}, your role={}".format( + CommonProxyErrors.not_allowed_access.value, + user_api_key_dict.user_role, + ) + }, + ) + + response = await prisma_client.db.litellm_budgettable.delete( + where={"budget_id": data.id} + ) + + return response |