mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-25 18:54:30 +00:00
9b7ebb6a7d
* build(pyproject.toml): add new dev dependencies - for type checking * build: reformat files to fit black * ci: reformat to fit black * ci(test-litellm.yml): make tests run clear * build(pyproject.toml): add ruff * fix: fix ruff checks * build(mypy/): fix mypy linting errors * fix(hashicorp_secret_manager.py): fix passing cert for tls auth * build(mypy/): resolve all mypy errors * test: update test * fix: fix black formatting * build(pre-commit-config.yaml): use poetry run black * fix(proxy_server.py): fix linting error * fix: fix ruff safe representation error
286 lines
9 KiB
Python
286 lines
9 KiB
Python
"""
|
|
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
|