Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/management_endpoints/budget_management_endpoints.py: 82%

1""" 

2BUDGET MANAGEMENT 

3 

4All /budget management endpoints 

5 

6/budget/new 

7/budget/info 

8/budget/update 

9/budget/delete 

10/budget/settings 

11/budget/list 

12""" 

13 

14#### BUDGET TABLE MANAGEMENT #### 

15import math 

16from collections.abc import Mapping 

17from types import MappingProxyType 

18from typing import Final 

19 

20from fastapi import APIRouter, Depends, HTTPException 

21 

22from litellm.proxy._types import * 

23from litellm.proxy.auth.user_api_key_auth import user_api_key_auth 

24from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time 

25from litellm.proxy.management_endpoints.common_utils import ( 

26 _user_has_admin_view, 

27 validate_budget_duration, 

28) 

29from litellm.proxy.utils import jsonify_object 

30from litellm.repositories.budget_repository import BudgetRepository 

31 

32router: Final = APIRouter() 

33 

34 

35@router.post( 

36 "/budget/new", 

37 tags=["budget management"], 

38 dependencies=[Depends(user_api_key_auth)], 

39) 

40async def new_budget( 

41 budget_obj: BudgetNewRequest, 

42 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), 

43): 

44 """ 

45 Create a new budget object. Can apply this to teams, orgs, end-users, keys. 

46 

47 Parameters: 

48 - budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.) 

49 - budget_id: Optional[str] - The id of the budget. If not provided, a new id will be generated. 

50 - max_budget: Optional[float] - The max budget for the budget. 

51 - soft_budget: Optional[float] - The soft budget for the budget. 

52 - max_parallel_requests: Optional[int] - The max number of parallel requests for the budget. 

53 - tpm_limit: Optional[int] - The tokens per minute limit for the budget. 

54 - rpm_limit: Optional[int] - The requests per minute limit for the budget. 

55 - tpd_limit: Optional[int] - The tokens per day limit for the budget. Charged by batch submissions instead of tpm_limit/rpm_limit. 

56 - 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}} 

57 - budget_reset_at: Optional[datetime] - Datetime when the initial budget is reset. Default is now. 

58 """ 

59 from prisma.errors import UniqueViolationError 

60 

61 from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client 

62 

63 if prisma_client is None: 63 ↛ 64line 63 didn't jump to line 64 because the condition on line 63 was never true

64 raise HTTPException( 

65 status_code=500, 

66 detail={"error": CommonProxyErrors.db_not_connected_error.value}, 

67 ) 

68 

69 # Validate budget values are not negative 

70 if budget_obj.max_budget is not None and (not math.isfinite(budget_obj.max_budget) or budget_obj.max_budget < 0): 

71 raise HTTPException( 

72 status_code=400, 

73 detail={"error": f"max_budget must be a non-negative finite number. Received: {budget_obj.max_budget}"}, 

74 ) 

75 if budget_obj.soft_budget is not None and (not math.isfinite(budget_obj.soft_budget) or budget_obj.soft_budget < 0): 

76 raise HTTPException( 

77 status_code=400, 

78 detail={"error": f"soft_budget must be a non-negative finite number. Received: {budget_obj.soft_budget}"}, 

79 ) 

80 

81 validate_budget_duration(budget_obj.budget_duration) 

82 

83 # Validate model_max_budget if present 

84 if budget_obj.model_max_budget is not None and len(budget_obj.model_max_budget) > 0: 

85 from litellm.proxy.management_endpoints.key_management_endpoints import ( 

86 validate_model_max_budget, 

87 ) 

88 

89 try: 

90 validate_model_max_budget(budget_obj.model_max_budget) 

91 except ValueError as e: 

92 raise HTTPException(status_code=400, detail={"error": str(e)}) 

93 

94 # if no budget_reset_at date is set, but a budget_duration is given, then set budget_reset_at initially to the first completed duration interval in future 

95 if budget_obj.budget_reset_at is None and budget_obj.budget_duration is not None: 95 ↛ 96line 95 didn't jump to line 96 because the condition on line 95 was never true

96 budget_obj.budget_reset_at = get_budget_reset_time(budget_duration=budget_obj.budget_duration) 

97 

98 budget_obj_json: Final = budget_obj.model_dump(exclude_none=True) 

99 budget_obj_jsonified: Final[dict[str, object]] = jsonify_object(budget_obj_json) # mutable-ok: prisma create input 

100 try: 

101 response: Final = await BudgetRepository(prisma_client).table.create( 

102 data={ 

103 **budget_obj_jsonified, 

104 "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, 

105 "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, 

106 } 

107 ) 

108 except Exception as e: 

109 if not isinstance(e, UniqueViolationError): 109 ↛ 110line 109 didn't jump to line 110 because the condition on line 109 was never true

110 raise 

111 raise HTTPException( 

112 status_code=400, 

113 detail={"error": f"Budget with id '{budget_obj.budget_id}' already exists."}, 

114 ) 

115 

116 return response 

117 

118 

119@router.post( 

120 "/budget/update", 

121 tags=["budget management"], 

122 dependencies=[Depends(user_api_key_auth)], 

123) 

124async def update_budget( 

125 budget_obj: BudgetNewRequest, 

126 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), 

127): 

128 """ 

129 Update an existing budget object. 

130 

131 Parameters: 

132 - budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.) 

133 - budget_id: Optional[str] - The id of the budget. If not provided, a new id will be generated. 

134 - max_budget: Optional[float] - The max budget for the budget. 

135 - soft_budget: Optional[float] - The soft budget for the budget. 

136 - max_parallel_requests: Optional[int] - The max number of parallel requests for the budget. 

137 - tpm_limit: Optional[int] - The tokens per minute limit for the budget. 

138 - rpm_limit: Optional[int] - The requests per minute limit for the budget. 

139 - tpd_limit: Optional[int] - The tokens per day limit for the budget. Charged by batch submissions instead of tpm_limit/rpm_limit. 

140 - 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}} 

141 - budget_reset_at: Optional[datetime] - Update the Datetime when the budget was last reset. 

142 """ 

143 from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client 

144 

145 if prisma_client is None: 145 ↛ 146line 145 didn't jump to line 146 because the condition on line 145 was never true

146 raise HTTPException( 

147 status_code=500, 

148 detail={"error": CommonProxyErrors.db_not_connected_error.value}, 

149 ) 

150 if budget_obj.budget_id is None: 

151 raise HTTPException(status_code=400, detail={"error": "budget_id is required"}) 

152 

153 # Validate budget values are not negative 

154 if budget_obj.max_budget is not None and (not math.isfinite(budget_obj.max_budget) or budget_obj.max_budget < 0): 

155 raise HTTPException( 

156 status_code=400, 

157 detail={"error": f"max_budget must be a non-negative finite number. Received: {budget_obj.max_budget}"}, 

158 ) 

159 if budget_obj.soft_budget is not None and (not math.isfinite(budget_obj.soft_budget) or budget_obj.soft_budget < 0): 

160 raise HTTPException( 

161 status_code=400, 

162 detail={"error": f"soft_budget must be a non-negative finite number. Received: {budget_obj.soft_budget}"}, 

163 ) 

164 

165 validate_budget_duration(budget_obj.budget_duration) 

166 

167 # Validate model_max_budget if present in update 

168 if budget_obj.model_max_budget is not None and len(budget_obj.model_max_budget) > 0: 168 ↛ 169line 168 didn't jump to line 169 because the condition on line 168 was never true

169 from litellm.proxy.management_endpoints.key_management_endpoints import ( 

170 validate_model_max_budget, 

171 ) 

172 

173 try: 

174 validate_model_max_budget(budget_obj.model_max_budget) 

175 except ValueError as e: 

176 raise HTTPException(status_code=400, detail={"error": str(e)}) 

177 

178 # recompute budget_reset_at when the duration changes, unless the caller pinned a reset time explicitly 

179 recomputed_reset_at: Final = ( 

180 {"budget_reset_at": get_budget_reset_time(budget_duration=budget_obj.budget_duration)} 

181 if budget_obj.budget_duration is not None and "budget_reset_at" not in budget_obj.model_fields_set 

182 else MappingProxyType({"budget_reset_at": None}) 

183 if "budget_duration" in budget_obj.model_fields_set 

184 and budget_obj.budget_duration is None 

185 and "budget_reset_at" not in budget_obj.model_fields_set 

186 else {} 

187 ) 

188 

189 budget_obj_jsonified: Final[Mapping[str, object]] = jsonify_object( 

190 { 

191 **budget_obj.model_dump(exclude_unset=True), 

192 **recomputed_reset_at, 

193 "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, 

194 } 

195 ) 

196 

197 response: Final = await BudgetRepository(prisma_client).table.update( 

198 where={"budget_id": budget_obj.budget_id}, 

199 data=budget_obj_jsonified, 

200 ) 

201 

202 return response 

203 

204 

205@router.post( 

206 "/budget/info", 

207 tags=["budget management"], 

208 dependencies=[Depends(user_api_key_auth)], 

209) 

210async def info_budget(data: BudgetRequest): 

211 """ 

212 Get the budget id specific information 

213 

214 Parameters: 

215 - budgets: List[str] - The list of budget ids to get information for 

216 """ 

217 from litellm.proxy.proxy_server import prisma_client 

218 

219 if prisma_client is None: 219 ↛ 220line 219 didn't jump to line 220 because the condition on line 219 was never true

220 raise HTTPException(status_code=500, detail={"error": "No db connected"}) 

221 

222 if len(data.budgets) == 0: 

223 raise HTTPException( 

224 status_code=400, 

225 detail={"error": f"Specify list of budget id's to query. Passed in={data.budgets}"}, 

226 ) 

227 response: Final = await BudgetRepository(prisma_client).table.find_many( 

228 where={"budget_id": {"in": data.budgets}}, 

229 ) 

230 

231 return response 

232 

233 

234@router.get( 

235 "/budget/settings", 

236 tags=["budget management"], 

237 dependencies=[Depends(user_api_key_auth)], 

238) 

239async def budget_settings( 

240 budget_id: str, 

241 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), 

242): 

243 """ 

244 Get list of configurable params + current value for a budget item + description of each field 

245 

246 Used on Admin UI. 

247 

248 Query Parameters: 

249 - budget_id: str - The budget id to get information for 

250 """ 

251 from litellm.proxy.proxy_server import prisma_client 

252 

253 if prisma_client is None: 253 ↛ 254line 253 didn't jump to line 254 because the condition on line 253 was never true

254 raise HTTPException( 

255 status_code=400, 

256 detail={"error": CommonProxyErrors.db_not_connected_error.value}, 

257 ) 

258 

259 if not _user_has_admin_view(user_api_key_dict): 259 ↛ 260line 259 didn't jump to line 260 because the condition on line 259 was never true

260 raise HTTPException( 

261 status_code=400, 

262 detail={"error": f"{CommonProxyErrors.not_allowed_access.value}, your role={user_api_key_dict.user_role}"}, 

263 ) 

264 

265 ## get budget item from db 

266 db_budget_row: Final = await BudgetRepository(prisma_client).table.find_first(where={"budget_id": budget_id}) 

267 

268 if db_budget_row is not None: 

269 db_budget_row_dict = db_budget_row.model_dump(exclude_none=True) 

270 else: 

271 db_budget_row_dict = {} 

272 

273 allowed_args: Final = { 

274 "max_parallel_requests": {"type": "Integer"}, 

275 "tpm_limit": {"type": "Integer"}, 

276 "rpm_limit": {"type": "Integer"}, 

277 "tpd_limit": {"type": "Integer"}, 

278 "budget_duration": {"type": "String"}, 

279 "max_budget": {"type": "Float"}, 

280 "soft_budget": {"type": "Float"}, 

281 "model_max_budget": {"type": "Object"}, 

282 } 

283 

284 return_val: Final = [] 

285 

286 for field_name, field_info in BudgetNewRequest.model_fields.items(): 

287 if field_name in allowed_args: 

288 _stored_in_db = True 

289 

290 _response_obj = ConfigList( 

291 field_name=field_name, 

292 field_type=allowed_args[field_name]["type"], 

293 field_description=field_info.description or "", 

294 field_value=db_budget_row_dict.get(field_name, None), 

295 stored_in_db=_stored_in_db, 

296 field_default_value=field_info.default, 

297 ) 

298 return_val.append(_response_obj) 

299 

300 return return_val 

301 

302 

303@router.get( 

304 "/budget/list", 

305 tags=["budget management"], 

306 dependencies=[Depends(user_api_key_auth)], 

307) 

308async def list_budget( 

309 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), 

310): 

311 """List all the created budgets in proxy db. Used on Admin UI.""" 

312 from litellm.proxy.proxy_server import prisma_client 

313 

314 if prisma_client is None: 314 ↛ 315line 314 didn't jump to line 315 because the condition on line 314 was never true

315 raise HTTPException( 

316 status_code=400, 

317 detail={"error": CommonProxyErrors.db_not_connected_error.value}, 

318 ) 

319 

320 if not _user_has_admin_view(user_api_key_dict): 320 ↛ 321line 320 didn't jump to line 321 because the condition on line 320 was never true

321 raise HTTPException( 

322 status_code=400, 

323 detail={"error": f"{CommonProxyErrors.not_allowed_access.value}, your role={user_api_key_dict.user_role}"}, 

324 ) 

325 

326 response: Final = await BudgetRepository(prisma_client).table.find_many() 

327 

328 return response 

329 

330 

331@router.post( 

332 "/budget/delete", 

333 tags=["budget management"], 

334 dependencies=[Depends(user_api_key_auth)], 

335) 

336async def delete_budget( 

337 data: BudgetDeleteRequest, 

338 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), 

339): 

340 """ 

341 Delete budget 

342 

343 Parameters: 

344 - id: str - The budget id to delete 

345 """ 

346 from litellm.proxy.proxy_server import prisma_client 

347 

348 if prisma_client is None: 348 ↛ 349line 348 didn't jump to line 349 because the condition on line 348 was never true

349 raise HTTPException( 

350 status_code=500, 

351 detail={"error": CommonProxyErrors.db_not_connected_error.value}, 

352 ) 

353 

354 if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: 354 ↛ 355line 354 didn't jump to line 355 because the condition on line 354 was never true

355 raise HTTPException( 

356 status_code=400, 

357 detail={"error": f"{CommonProxyErrors.not_allowed_access.value}, your role={user_api_key_dict.user_role}"}, 

358 ) 

359 

360 response: Final = await BudgetRepository(prisma_client).table.delete(where={"budget_id": data.id}) 

361 

362 return response