Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/management_endpoints/management_v1/budgets.py: 87%

1"""`GET /management/v1/budgets`.""" 

2 

3from collections.abc import Mapping, Sequence 

4from dataclasses import dataclass 

5from datetime import datetime 

6from types import MappingProxyType 

7from typing import Annotated, Final 

8 

9from fastapi import APIRouter, Depends, Request 

10from pydantic import BaseModel, TypeAdapter 

11 

12from litellm._logging import verbose_proxy_logger 

13from litellm.proxy._types import ( 

14 CommonProxyErrors, 

15 UserAPIKeyAuth, 

16 user_api_key_has_admin_view, 

17) 

18from litellm.proxy.auth.user_api_key_auth import user_api_key_auth 

19from litellm.proxy.list_api.common import ( 

20 PROBLEM_TYPE_BASE, 

21 ManagementProblem, 

22) 

23from litellm.proxy.list_api.list_framework import ( 

24 FilterSpec, 

25 ListSpec, 

26 Predicate, 

27 QueryPlan, 

28 Scope, 

29 ScopeAll, 

30 ScopeDenied, 

31 SortKey, 

32 handle_list, 

33 order_by_sql, 

34 where_sql, 

35) 

36from litellm.proxy.management_endpoints.management_v1.common import MANAGEMENT_V1_PREFIX 

37from litellm.proxy.utils import PrismaClient 

38from litellm.types.proxy.management_endpoints.management_v1 import ( 

39 ListResponse, 

40 ProblemDetail, 

41) 

42 

43router: Final = APIRouter(prefix=MANAGEMENT_V1_PREFIX) 

44 

45BUDGET_TABLE: Final = '"LiteLLM_BudgetTable"' 

46 

47 

48class BudgetListItem(BaseModel): 

49 """One budget as the Budgets page reads it, and as it comes back off the table. 

50 

51 Validating the raw row through here is what makes `tpm_limit` / `rpm_limit` 

52 numbers: they are `BigInt?` in the schema, which the query engine hands back as 

53 decimal strings, and a quoted "60000" breaks arithmetic in the dashboard. 

54 """ 

55 

56 budget_id: str 

57 max_budget: float | None = None 

58 soft_budget: float | None = None 

59 tpm_limit: int | None = None 

60 rpm_limit: int | None = None 

61 tpd_limit: int | None = None 

62 budget_duration: str | None = None 

63 budget_reset_at: datetime | None = None 

64 created_at: datetime 

65 updated_at: datetime 

66 

67 

68class _RowCount(BaseModel): 

69 count: int 

70 

71 

72_BUDGET_ROWS: Final = TypeAdapter(tuple[BudgetListItem, ...]) 

73_ROW_COUNTS: Final = TypeAdapter(tuple[_RowCount, ...]) 

74 

75SELECTED_COLUMNS: Final = ", ".join(f'"{name}"' for name in BudgetListItem.model_fields) 

76 

77 

78@dataclass(frozen=True, slots=True) 

79class PrismaBudgetListExecutor: 

80 """The database half of the budgets list. Every caller-supplied value is bound to a 

81 placeholder by `where_sql`; only the spec's own column names reach the SQL text.""" 

82 

83 prisma_client: PrismaClient 

84 

85 async def count(self, where: tuple[Predicate, ...]) -> int: 

86 clauses, params = where_sql(where) 

87 sql: Final = f"SELECT COUNT(*) AS count FROM {BUDGET_TABLE}" + (f" WHERE {clauses}" if clauses else "") 

88 rows: Final = await self.prisma_client.db.query_raw(sql, *params) 

89 counted: Final = _ROW_COUNTS.validate_python(rows) 

90 return counted[0].count if counted else 0 

91 

92 async def find_many(self, plan: QueryPlan) -> Sequence[BudgetListItem]: 

93 clauses, params = where_sql(plan.where) 

94 sql: Final = ( 

95 f"SELECT {SELECTED_COLUMNS} FROM {BUDGET_TABLE}" 

96 + (f" WHERE {clauses}" if clauses else "") 

97 + f" ORDER BY {order_by_sql(plan.order)}" 

98 + f" LIMIT ${len(params) + 1} OFFSET ${len(params) + 2}" 

99 ) 

100 rows: Final = await self.prisma_client.db.query_raw(sql, *params, plan.take, plan.skip) 

101 return _BUDGET_ROWS.validate_python(rows) 

102 

103 

104def _serialize(row: BudgetListItem) -> BudgetListItem: 

105 """The row shape is the wire shape: the query selects exactly the columns served.""" 

106 return row 

107 

108 

109def _scope(caller: UserAPIKeyAuth) -> Scope: 

110 if user_api_key_has_admin_view(caller): 110 ↛ 112line 110 didn't jump to line 112 because the condition on line 110 was always true

111 return ScopeAll() 

112 return ScopeDenied(reason=f"Only proxy admins can list budgets, your role={caller.user_role}") 

113 

114 

115# budget_duration is deliberately absent from `sortable`: the column holds strings 

116# like "7d" and "30d", so a lexicographic ORDER BY puts "30d" ahead of "7d". 

117BUDGET_FILTERS: Final[Mapping[str, FilterSpec]] = MappingProxyType( 

118 { 

119 "budget_duration": FilterSpec(type=str, ops=frozenset(("in", "is_null"))), 

120 "max_budget": FilterSpec(type=float, ops=frozenset(("gte", "lte", "is_null"))), 

121 "created_at": FilterSpec(type=datetime, ops=frozenset(("gte", "lte"))), 

122 } 

123) 

124 

125BUDGETS_LIST_SPEC: Final[ListSpec[BudgetListItem, BudgetListItem]] = ListSpec( 

126 resource="budgets", 

127 sortable=frozenset(("budget_id", "max_budget", "tpm_limit", "rpm_limit", "tpd_limit", "created_at")), 

128 searchable=frozenset(("budget_id",)), 

129 filters=BUDGET_FILTERS, 

130 default_sort=(SortKey(field="created_at", descending=True),), 

131 default_page_size=50, 

132 max_page_size=100, 

133 scope=_scope, 

134 serialize=_serialize, 

135 tiebreaker="budget_id", 

136) 

137 

138 

139@router.get( 

140 "/budgets", 

141 tags=("budget management",), 

142 dependencies=(Depends(user_api_key_auth),), 

143 response_model=ListResponse[BudgetListItem], 

144) 

145async def list_budgets( 

146 request: Request, 

147 user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], 

148) -> ListResponse[BudgetListItem]: 

149 """ 

150 The budgets defined on this proxy, paged, sortable and filterable, for the 

151 Budgets page. 

152 

153 Readable by a proxy admin or an admin viewer; anyone else is refused 403. The 

154 older `/budget/list` answers with the whole table as a bare array and has no 

155 way to page, sort or filter it. 

156 

157 `sort` takes a comma-separated list of `budget_id`, `max_budget`, `tpm_limit`, 

158 `rpm_limit`, `tpd_limit` or `created_at`, each optionally prefixed with `-` for descending, 

159 and defaults to `-created_at`. `budget_id` is appended to every sort as the 

160 tiebreaker. `q` is a case-insensitive substring match on `budget_id`. 

161 `page_size` defaults to 50 and is capped at 100. Filters are 

162 `filter[budget_duration][in|is_null]`, `filter[max_budget][gte|lte|is_null]` 

163 and `filter[created_at][gte|lte]`. 

164 

165 Example curl: 

166 ``` 

167 curl --location --globoff 'http://0.0.0.0:4000/management/v1/budgets?sort=-max_budget&filter[budget_duration][in]=7d,30d&page_size=25' \ 

168 --header 'Authorization: Bearer sk-1234' 

169 ``` 

170 """ 

171 try: 

172 from litellm.proxy.proxy_server import prisma_client 

173 

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

175 raise ManagementProblem( 

176 ProblemDetail( 

177 type=f"{PROBLEM_TYPE_BASE}database-not-connected", 

178 title="Database not connected", 

179 status=503, 

180 detail=CommonProxyErrors.db_not_connected_error.value, 

181 ) 

182 ) 

183 

184 return await handle_list( 

185 spec=BUDGETS_LIST_SPEC, 

186 executor=PrismaBudgetListExecutor(prisma_client=prisma_client), 

187 request=request, 

188 caller=user_api_key_dict, 

189 ) 

190 

191 except ManagementProblem: 

192 raise 

193 except Exception as e: # noqa: BLE001 # a driver error answers as a problem document, not the OpenAI error shape 

194 verbose_proxy_logger.exception( 

195 "litellm.proxy.management_endpoints.management_v1.budgets.list_budgets(): Exception occured - %s", e 

196 ) 

197 raise ManagementProblem( 

198 ProblemDetail( 

199 type=f"{PROBLEM_TYPE_BASE}internal-server-error", 

200 title="Internal server error", 

201 status=500, 

202 detail="Failed to list budgets.", 

203 ) 

204 )