mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
feat(proxy): add GET /management/v1/budgets (#35310)
* feat(proxy): add a generic list contract for management/v1 entity lists Paging, sorting, filtering and search for an entity collection, declared once as a ListSpec and served by handle_list. The route injects a ListExecutor that owns its table, so this module never imports Prisma. The caller's scope is derived from the caller alone and ANDed with whatever they filtered on, so a query parameter can only narrow what they may read. This is the shared half of the budgets list; it lands here so the endpoint has something to register against, and drops out when the framework arrives on its own branch. * feat(proxy): add GET /management/v1/budgets The Budgets page reads /budget/list, which returns the whole table as a bare array with no way to page, sort or filter it. A customer with enough budgets to fill the page has no way to find one. Registers LiteLLM_BudgetTable against the management/v1 list contract: sortable on budget_id, max_budget, tpm_limit, rpm_limit and created_at, default order newest-first with budget_id breaking ties, search on budget_id, and filters for budget_duration, max_budget and created_at. budget_duration is deliberately not sortable; the column holds "7d"/"30d" strings, so a lexicographic ORDER BY puts "30d" ahead of "7d". tpm_limit and rpm_limit are BigInt? in Prisma, so rows validate through a pydantic model on the way out and serialize as JSON numbers. A caller without admin view is refused 403 as a problem document rather than served an empty page. /budget/list is untouched. * fix(proxy): rework the budgets list onto the merged list contract PR #35308 landed a different shape than this branch was written against: `where` is a tuple of frozen predicates rather than a Prisma-shaped mapping, `ListSpec` carries both the row and the wire type, and `where_sql` / `order_by_sql` render for a raw-SQL executor. The budgets executor now queries through `query_raw` the way the spend logs facet does, selecting only the columns it serves. Also casts datetime binds in `where_sql`. They cross into the query engine as JSON, so an uncast placeholder arrives as text and Postgres refuses `timestamp >= text` outright; every `filter[created_at][gte|lte]` was answering 500. The cast reads the bind as an instant and drops it to naive UTC to match Prisma's TIMESTAMP(3) column, the same one /spend/logs/ui applies. * refactor(proxy): fold the predicate renderer instead of recursing recursive_detector flags `_render_all`, and the flag is fair: it recursed once per predicate, so the stack grew with the number of filters on the request for no reason. Walking a predicate list is a running bind index, which is a fold. `_render` still re-enters for `AnyOf`, but its clauses are plain comparisons built by `?q=`, so that nesting is one level deep and no caller can drive it deeper.
This commit is contained in:
yuneng-jiang
GitHub
parent
c943d9a952
commit
fcec1488e2
10 changed files with 1075 additions and 10 deletions
|
|
@ -841,6 +841,7 @@ class LiteLLMRoutes(enum.Enum):
|
|||
"/config/list",
|
||||
"/config/field/info",
|
||||
"/budget/list",
|
||||
"/management/v1/budgets",
|
||||
"/budget/settings",
|
||||
# Invitation viewing (admin viewer cannot create/delete; can read).
|
||||
"/invitation/info",
|
||||
|
|
|
|||
|
|
@ -2,11 +2,15 @@
|
|||
|
||||
from fastapi import APIRouter
|
||||
|
||||
from litellm.proxy.management_endpoints.management_v1.budgets import (
|
||||
router as budgets_router,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.management_v1.spend_logs import (
|
||||
router as spend_logs_router,
|
||||
)
|
||||
|
||||
router = APIRouter()
|
||||
router.include_router(budgets_router)
|
||||
router.include_router(spend_logs_router)
|
||||
|
||||
__all__ = ["router"]
|
||||
|
|
|
|||
205
litellm/proxy/management_endpoints/management_v1/budgets.py
Normal file
205
litellm/proxy/management_endpoints/management_v1/budgets.py
Normal file
|
|
@ -0,0 +1,205 @@
|
|||
"""`GET /management/v1/budgets`."""
|
||||
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from types import MappingProxyType
|
||||
from typing import Annotated
|
||||
|
||||
from fastapi import APIRouter, Depends, Request
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._types import (
|
||||
CommonProxyErrors,
|
||||
UserAPIKeyAuth,
|
||||
user_api_key_has_admin_view,
|
||||
)
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.management_endpoints.management_v1.common import (
|
||||
MANAGEMENT_V1_PREFIX,
|
||||
PROBLEM_TYPE_BASE,
|
||||
ManagementProblem,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.management_v1.list_framework import (
|
||||
FilterSpec,
|
||||
ListSpec,
|
||||
Predicate,
|
||||
QueryPlan,
|
||||
Scope,
|
||||
ScopeAll,
|
||||
ScopeDenied,
|
||||
SortKey,
|
||||
handle_list,
|
||||
order_by_sql,
|
||||
where_sql,
|
||||
)
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
from litellm.types.proxy.management_endpoints.management_v1 import (
|
||||
ListResponse,
|
||||
ProblemDetail,
|
||||
)
|
||||
|
||||
router = APIRouter(prefix=MANAGEMENT_V1_PREFIX)
|
||||
|
||||
BUDGET_TABLE = '"LiteLLM_BudgetTable"'
|
||||
|
||||
|
||||
class BudgetListItem(BaseModel):
|
||||
"""One budget as the Budgets page reads it, and as it comes back off the table.
|
||||
|
||||
Validating the raw row through here is what makes `tpm_limit` / `rpm_limit`
|
||||
numbers: they are `BigInt?` in the schema, which the query engine hands back as
|
||||
decimal strings, and a quoted "60000" breaks arithmetic in the dashboard.
|
||||
"""
|
||||
|
||||
budget_id: str
|
||||
max_budget: float | None = None
|
||||
soft_budget: float | None = None
|
||||
tpm_limit: int | None = None
|
||||
rpm_limit: int | None = None
|
||||
budget_duration: str | None = None
|
||||
budget_reset_at: datetime | None = None
|
||||
created_at: datetime
|
||||
updated_at: datetime
|
||||
|
||||
|
||||
class _RowCount(BaseModel):
|
||||
count: int
|
||||
|
||||
|
||||
_BUDGET_ROWS = TypeAdapter(tuple[BudgetListItem, ...])
|
||||
_ROW_COUNTS = TypeAdapter(tuple[_RowCount, ...])
|
||||
|
||||
SELECTED_COLUMNS = ", ".join(f'"{name}"' for name in BudgetListItem.model_fields)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PrismaBudgetListExecutor:
|
||||
"""The database half of the budgets list. Every caller-supplied value is bound to a
|
||||
placeholder by `where_sql`; only the spec's own column names reach the SQL text."""
|
||||
|
||||
prisma_client: PrismaClient
|
||||
|
||||
async def count(self, where: tuple[Predicate, ...]) -> int:
|
||||
clauses, params = where_sql(where)
|
||||
sql = f"SELECT COUNT(*) AS count FROM {BUDGET_TABLE}" + (f" WHERE {clauses}" if clauses else "")
|
||||
rows = await self.prisma_client.db.query_raw(sql, *params)
|
||||
counted = _ROW_COUNTS.validate_python(rows)
|
||||
return counted[0].count if counted else 0
|
||||
|
||||
async def find_many(self, plan: QueryPlan) -> Sequence[BudgetListItem]:
|
||||
clauses, params = where_sql(plan.where)
|
||||
sql = (
|
||||
f"SELECT {SELECTED_COLUMNS} FROM {BUDGET_TABLE}"
|
||||
+ (f" WHERE {clauses}" if clauses else "")
|
||||
+ f" ORDER BY {order_by_sql(plan.order)}"
|
||||
+ f" LIMIT ${len(params) + 1} OFFSET ${len(params) + 2}"
|
||||
)
|
||||
rows = await self.prisma_client.db.query_raw(sql, *params, plan.take, plan.skip)
|
||||
return _BUDGET_ROWS.validate_python(rows)
|
||||
|
||||
|
||||
def _serialize(row: BudgetListItem) -> BudgetListItem:
|
||||
"""The row shape is the wire shape: the query selects exactly the columns served."""
|
||||
return row
|
||||
|
||||
|
||||
def _scope(caller: UserAPIKeyAuth) -> Scope:
|
||||
if user_api_key_has_admin_view(caller):
|
||||
return ScopeAll()
|
||||
return ScopeDenied(reason="Only proxy admins can list budgets, your role={}".format(caller.user_role))
|
||||
|
||||
|
||||
# budget_duration is deliberately absent from `sortable`: the column holds strings
|
||||
# like "7d" and "30d", so a lexicographic ORDER BY puts "30d" ahead of "7d".
|
||||
BUDGET_FILTERS: Mapping[str, FilterSpec] = MappingProxyType(
|
||||
{ # mutable-ok: an immutable mapping has no literal form; MappingProxyType freezes this one and it never escapes
|
||||
"budget_duration": FilterSpec(type=str, ops=frozenset(("in", "is_null"))),
|
||||
"max_budget": FilterSpec(type=float, ops=frozenset(("gte", "lte", "is_null"))),
|
||||
"created_at": FilterSpec(type=datetime, ops=frozenset(("gte", "lte"))),
|
||||
}
|
||||
)
|
||||
|
||||
BUDGETS_LIST_SPEC: ListSpec[BudgetListItem, BudgetListItem] = ListSpec(
|
||||
resource="budgets",
|
||||
sortable=frozenset(("budget_id", "max_budget", "tpm_limit", "rpm_limit", "created_at")),
|
||||
searchable=frozenset(("budget_id",)),
|
||||
filters=BUDGET_FILTERS,
|
||||
default_sort=(SortKey(field="created_at", descending=True),),
|
||||
default_page_size=50,
|
||||
max_page_size=100,
|
||||
scope=_scope,
|
||||
serialize=_serialize,
|
||||
tiebreaker="budget_id",
|
||||
)
|
||||
|
||||
|
||||
@router.get(
|
||||
"/budgets",
|
||||
tags=("budget management",),
|
||||
dependencies=(Depends(user_api_key_auth),),
|
||||
response_model=ListResponse[BudgetListItem],
|
||||
)
|
||||
async def list_budgets(
|
||||
request: Request,
|
||||
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
|
||||
) -> ListResponse[BudgetListItem]:
|
||||
"""
|
||||
The budgets defined on this proxy, paged, sortable and filterable, for the
|
||||
Budgets page.
|
||||
|
||||
Readable by a proxy admin or an admin viewer; anyone else is refused 403. The
|
||||
older `/budget/list` answers with the whole table as a bare array and has no
|
||||
way to page, sort or filter it.
|
||||
|
||||
`sort` takes a comma-separated list of `budget_id`, `max_budget`, `tpm_limit`,
|
||||
`rpm_limit` or `created_at`, each optionally prefixed with `-` for descending,
|
||||
and defaults to `-created_at`. `budget_id` is appended to every sort as the
|
||||
tiebreaker. `q` is a case-insensitive substring match on `budget_id`.
|
||||
`page_size` defaults to 50 and is capped at 100. Filters are
|
||||
`filter[budget_duration][in|is_null]`, `filter[max_budget][gte|lte|is_null]`
|
||||
and `filter[created_at][gte|lte]`.
|
||||
|
||||
Example curl:
|
||||
```
|
||||
curl --location --globoff 'http://0.0.0.0:4000/management/v1/budgets?sort=-max_budget&filter[budget_duration][in]=7d,30d&page_size=25' \
|
||||
--header 'Authorization: Bearer sk-1234'
|
||||
```
|
||||
"""
|
||||
try:
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
if prisma_client is None:
|
||||
raise ManagementProblem(
|
||||
ProblemDetail(
|
||||
type=f"{PROBLEM_TYPE_BASE}database-not-connected",
|
||||
title="Database not connected",
|
||||
status=503,
|
||||
detail=CommonProxyErrors.db_not_connected_error.value,
|
||||
)
|
||||
)
|
||||
|
||||
return await handle_list(
|
||||
spec=BUDGETS_LIST_SPEC,
|
||||
executor=PrismaBudgetListExecutor(prisma_client=prisma_client),
|
||||
request=request,
|
||||
caller=user_api_key_dict,
|
||||
)
|
||||
|
||||
except ManagementProblem:
|
||||
raise
|
||||
except Exception as e: # noqa: BLE001 # a driver error answers as a problem document, not the OpenAI error shape
|
||||
verbose_proxy_logger.exception(
|
||||
"litellm.proxy.management_endpoints.management_v1.budgets.list_budgets(): Exception occured - {}".format(
|
||||
str(e)
|
||||
)
|
||||
)
|
||||
raise ManagementProblem(
|
||||
ProblemDetail(
|
||||
type=f"{PROBLEM_TYPE_BASE}internal-server-error",
|
||||
title="Internal server error",
|
||||
status=500,
|
||||
detail="Failed to list budgets.",
|
||||
)
|
||||
)
|
||||
|
|
@ -15,6 +15,7 @@ raw-SQL executor with every caller-supplied value bound to a placeholder.
|
|||
from collections.abc import Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from functools import partial, reduce
|
||||
from math import ceil
|
||||
from typing import Generic, Literal, Protocol, TypeVar
|
||||
|
||||
|
|
@ -213,12 +214,23 @@ def _sql_operator(op: ComparisonOp) -> str:
|
|||
assert_never(op)
|
||||
|
||||
|
||||
def _placeholder(index: int, value: FilterValue) -> str:
|
||||
"""`$n`, cast when the bind is a datetime.
|
||||
|
||||
Binds cross into the query engine as JSON, so a datetime arrives as text and
|
||||
Postgres refuses `timestamp >= text` outright. Prisma stores DateTime as a naive
|
||||
`TIMESTAMP(3)` holding UTC, so the bind is read as an instant and then dropped to
|
||||
naive UTC to match the column, the same cast `/spend/logs/ui` applies.
|
||||
"""
|
||||
return f"${index}::timestamptz AT TIME ZONE 'UTC'" if isinstance(value, datetime) else f"${index}"
|
||||
|
||||
|
||||
def _render(predicate: Predicate, index: int) -> tuple[str, tuple[object, ...]]:
|
||||
match predicate:
|
||||
case IsNull(field=field, negated=negated):
|
||||
return f'"{field}" IS {"NOT NULL" if negated else "NULL"}', ()
|
||||
case Within(field=field, values=values):
|
||||
placeholders = ", ".join(f"${index + offset}" for offset in range(len(values)))
|
||||
placeholders = ", ".join(_placeholder(index + offset, value) for offset, value in enumerate(values))
|
||||
return f'"{field}" IN ({placeholders})', values
|
||||
case AnyOf(clauses=clauses):
|
||||
rendered, params = _render_all(clauses, index)
|
||||
|
|
@ -226,17 +238,31 @@ def _render(predicate: Predicate, index: int) -> tuple[str, tuple[object, ...]]:
|
|||
case Compare(field=field, op="contains", value=value):
|
||||
return f"\"{field}\" ILIKE ${index} ESCAPE '\\'", (f"%{escape_like(str(value))}%",)
|
||||
case Compare(field=field, op=op, value=value):
|
||||
return f'"{field}" {_sql_operator(op)} ${index}', (value,)
|
||||
return f'"{field}" {_sql_operator(op)} {_placeholder(index, value)}', (value,)
|
||||
case _:
|
||||
assert_never(predicate)
|
||||
|
||||
|
||||
def _render_one(
|
||||
rendered: tuple[tuple[str, ...], tuple[object, ...]],
|
||||
predicate: Predicate,
|
||||
first_index: int,
|
||||
) -> tuple[tuple[str, ...], tuple[object, ...]]:
|
||||
"""Append one predicate, numbering it after the binds already consumed."""
|
||||
clauses, params = rendered
|
||||
clause, clause_params = _render(predicate, first_index + len(params))
|
||||
return (*clauses, clause), (*params, *clause_params)
|
||||
|
||||
|
||||
def _render_all(predicates: tuple[Predicate, ...], index: int) -> tuple[tuple[str, ...], tuple[object, ...]]:
|
||||
if not predicates:
|
||||
return (), ()
|
||||
head, head_params = _render(predicates[0], index)
|
||||
tail, tail_params = _render_all(predicates[1:], index + len(head_params))
|
||||
return (head, *tail), head_params + tail_params
|
||||
"""Render every predicate, numbering placeholders continuously across them.
|
||||
|
||||
Folded rather than self-recursive: walking a predicate list is a running index, and
|
||||
recursing per predicate grew the stack with the filter count for nothing. `_render`
|
||||
still re-enters here for `AnyOf`, whose clauses are plain `Compare`s from `?q=`, so
|
||||
that nesting is one level deep and cannot be driven deeper by a caller.
|
||||
"""
|
||||
return reduce(partial(_render_one, first_index=index), predicates, ((), ()))
|
||||
|
||||
|
||||
def where_sql(where: tuple[Predicate, ...], first_index: int = 1) -> tuple[str, tuple[object, ...]]:
|
||||
|
|
|
|||
|
|
@ -57,6 +57,8 @@
|
|||
- {id: mgmt.budget.update.persists, module: mgmt, tier: P1, surface: api, assertions: [persists], source: "budget_management_endpoints.py:155", rationale: "Limit changes apply"}
|
||||
- {id: mgmt.budget.delete.persists, module: mgmt, tier: P1, surface: api, assertions: [persists], source: "budget_management_endpoints.py:280", rationale: "Clears limits"}
|
||||
- {id: mgmt.budget.list.happy_path, module: mgmt, tier: P1, surface: api, assertions: [happy_path], source: "budget_management_endpoints.py:215", rationale: "Budget enumeration"}
|
||||
- {id: mgmt.budget.list_v1.happy_path, module: mgmt, tier: P1, surface: api, assertions: [happy_path], source: "management_v1/budgets.py:129", rationale: "Budget enumeration the Budgets page can page, sort and filter"}
|
||||
- {id: mgmt.budget.list_v1.admin_only, module: mgmt, tier: P1, surface: api, assertions: [admin_only], source: "management_v1/budgets.py:129", rationale: "A caller without admin view is refused, not served an empty page"}
|
||||
- {id: mgmt.callback.list.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "callback_management_endpoints.py", rationale: "Callback config (smoke)"}
|
||||
- {id: mgmt.cache_settings.update.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "cache_settings_endpoints.py", rationale: "Cache config (smoke). Deliberately uncovered: the previous test read the live settings and wrote them back, which proves nothing (identical values in, so a no-op POST still passes) while being able to break the deployment. /cache/settings persists what it receives and that row outranks YAML cache_params, re-applied on a timer, so a write that omits ssl or redis_startup_nodes turns a TLS cluster into a plaintext standalone node and every later Redis call hangs. That took out 60 of 72 tests on 2026-07-25. GET cannot round-trip it either: it resolves the stored row overlaid with REDIS_* env and never reads YAML, so on a fresh deploy it cannot see YAML ssl to echo back. A safe test needs an isolated proxy, or LIT-4816 fixed so a partial write cannot downgrade transport. Do not re-add a read-then-write-back test against a shared proxy."}
|
||||
- {id: mgmt.cost_tracking.estimate.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "cost_tracking_settings.py", rationale: "Cost estimate (smoke)"}
|
||||
|
|
|
|||
|
|
@ -19,10 +19,10 @@ import time
|
|||
from collections.abc import Callable
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel, RootModel
|
||||
from pydantic import BaseModel, Field, RootModel
|
||||
|
||||
from e2e_config import unique_marker
|
||||
from e2e_http import NoBody, unwrap
|
||||
from e2e_http import NoBody, Success, UnauthorizedError, UnknownApiError, unwrap
|
||||
from lifecycle import ResourceManager
|
||||
from management_client import ManagementClient
|
||||
from models import KeyGenerateBody, OrgInfoParams, OrgNewBody, UserNewBody
|
||||
|
|
@ -44,9 +44,11 @@ def _poll[T](client: ManagementClient, attempt: Callable[[], T | None], failure:
|
|||
|
||||
|
||||
class BudgetNewBody(BaseModel):
|
||||
max_budget: float
|
||||
max_budget: float | None = None
|
||||
soft_budget: float | None = None
|
||||
budget_duration: str | None = None
|
||||
budget_id: str | None = None
|
||||
tpm_limit: int | None = None
|
||||
|
||||
|
||||
class BudgetNewResponse(BaseModel):
|
||||
|
|
@ -204,6 +206,164 @@ class TestBudgetManagement:
|
|||
)
|
||||
|
||||
|
||||
# ---------- /management/v1/budgets ----------
|
||||
|
||||
_BUDGETS_V1 = "/management/v1/budgets"
|
||||
|
||||
|
||||
class BudgetPageParams(BaseModel):
|
||||
"""Query for GET /management/v1/budgets. The filter fields serialize to the
|
||||
bracketed keys the route reads them under, so nothing here is a raw dict."""
|
||||
|
||||
q: str | None = None
|
||||
sort: str | None = None
|
||||
page: int | None = None
|
||||
page_size: int | None = None
|
||||
duration_in: str | None = Field(default=None, serialization_alias="filter[budget_duration][in]")
|
||||
max_budget_is_null: bool | None = Field(default=None, serialization_alias="filter[max_budget][is_null]")
|
||||
not_a_parameter: str | None = Field(default=None, serialization_alias="filter[budget_id][eq]")
|
||||
|
||||
|
||||
class BudgetPageMeta(BaseModel):
|
||||
page: int
|
||||
page_size: int
|
||||
total_count: int
|
||||
total_pages: int
|
||||
|
||||
|
||||
class BudgetPageLinks(BaseModel):
|
||||
first: str
|
||||
prev: str | None = None
|
||||
next: str | None = None
|
||||
last: str
|
||||
|
||||
|
||||
class BudgetPageRow(BaseModel):
|
||||
budget_id: str
|
||||
max_budget: float | None = None
|
||||
tpm_limit: int | None = None
|
||||
budget_duration: str | None = None
|
||||
|
||||
|
||||
class BudgetPageResponse(BaseModel):
|
||||
data: list[BudgetPageRow]
|
||||
meta: BudgetPageMeta
|
||||
links: BudgetPageLinks
|
||||
|
||||
|
||||
def _list_budgets(client: ManagementClient, params: BudgetPageParams) -> BudgetPageResponse:
|
||||
return unwrap(
|
||||
client.proxy.transport.get(
|
||||
_BUDGETS_V1,
|
||||
headers=client.proxy.transport.master,
|
||||
params=params,
|
||||
response_type=BudgetPageResponse,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _list_budget_ids(client: ManagementClient, params: BudgetPageParams) -> tuple[str, ...]:
|
||||
return tuple(row.budget_id for row in _list_budgets(client, params).data)
|
||||
|
||||
|
||||
def _list_status(client: ManagementClient, params: BudgetPageParams, key: str | None = None) -> int:
|
||||
headers = client.proxy.transport.master if key is None else client.proxy.transport.bearer(key)
|
||||
outcome = client.proxy.transport.get(
|
||||
_BUDGETS_V1, headers=headers, params=params, response_type=BudgetPageResponse
|
||||
)
|
||||
match outcome:
|
||||
case Success(status_code=status_code):
|
||||
return status_code
|
||||
case UnauthorizedError():
|
||||
return 401
|
||||
case UnknownApiError(status_code=status_code):
|
||||
return status_code
|
||||
case _:
|
||||
raise AssertionError(outcome)
|
||||
|
||||
|
||||
class TestBudgetListV1:
|
||||
"""The paged, sorted, filtered budget list the Budgets page reads.
|
||||
|
||||
Every test tags its own budgets with a marker in the budget_id and searches on
|
||||
it, so budgets left behind by other suites cannot move the assertions.
|
||||
"""
|
||||
|
||||
@pytest.mark.covers("mgmt.budget.list_v1.happy_path")
|
||||
def test_sorts_pages_and_filters_the_budgets_it_created(
|
||||
self, client: ManagementClient, resources: ResourceManager
|
||||
) -> None:
|
||||
marker = unique_marker()
|
||||
small, medium, large = (f"{marker}-small", f"{marker}-medium", f"{marker}-large")
|
||||
for budget_id, max_budget, duration in (
|
||||
(small, 1.0, "7d"),
|
||||
(medium, 2.0, "30d"),
|
||||
(large, 3.0, "30d"),
|
||||
):
|
||||
_create_budget(
|
||||
client,
|
||||
resources,
|
||||
BudgetNewBody(
|
||||
budget_id=budget_id, max_budget=max_budget, budget_duration=duration, tpm_limit=60000
|
||||
),
|
||||
)
|
||||
|
||||
mine = BudgetPageParams(q=marker, sort="-max_budget")
|
||||
_ = _poll(
|
||||
client,
|
||||
lambda: mine if len(_list_budget_ids(client, mine)) == 3 else None,
|
||||
f"{_BUDGETS_V1} never listed all three budgets tagged {marker}",
|
||||
)
|
||||
|
||||
assert _list_budget_ids(client, mine) == (large, medium, small)
|
||||
|
||||
page_two = _list_budgets(client, BudgetPageParams(q=marker, sort="-max_budget", page=2, page_size=1))
|
||||
assert [row.budget_id for row in page_two.data] == [medium]
|
||||
assert page_two.meta.total_count == 3
|
||||
assert page_two.meta.total_pages == 3
|
||||
assert page_two.meta.page_size == 1
|
||||
assert page_two.links.prev is not None and page_two.links.next is not None
|
||||
|
||||
assert set(_list_budget_ids(client, BudgetPageParams(q=marker, duration_in="30d"))) == {medium, large}
|
||||
|
||||
limits = _list_budgets(client, BudgetPageParams(q=marker, sort="budget_id")).data
|
||||
assert [row.tpm_limit for row in limits] == [60000, 60000, 60000]
|
||||
|
||||
@pytest.mark.covers("mgmt.budget.list_v1.happy_path")
|
||||
def test_is_null_finds_the_budget_left_uncapped(
|
||||
self, client: ManagementClient, resources: ResourceManager
|
||||
) -> None:
|
||||
marker = unique_marker()
|
||||
uncapped = _create_budget(client, resources, BudgetNewBody(budget_id=f"{marker}-uncapped"))
|
||||
_ = _create_budget(client, resources, BudgetNewBody(budget_id=f"{marker}-capped", max_budget=4.0))
|
||||
|
||||
params = BudgetPageParams(q=marker, max_budget_is_null=True)
|
||||
found = _poll(
|
||||
client,
|
||||
lambda: params if _list_budget_ids(client, params) == (uncapped,) else None,
|
||||
f"{_BUDGETS_V1} never isolated the uncapped budget {uncapped}",
|
||||
)
|
||||
|
||||
assert [row.max_budget for row in _list_budgets(client, found).data] == [None]
|
||||
|
||||
@pytest.mark.covers("mgmt.budget.list_v1.happy_path")
|
||||
def test_refuses_a_sort_field_and_a_parameter_it_does_not_support(self, client: ManagementClient) -> None:
|
||||
assert _list_status(client, BudgetPageParams(sort="budget_duration")) == 400
|
||||
assert _list_status(client, BudgetPageParams(not_a_parameter="b-1")) == 400
|
||||
|
||||
@pytest.mark.covers("mgmt.budget.list_v1.admin_only")
|
||||
def test_is_refused_for_a_non_admin_key(self, client: ManagementClient, resources: ResourceManager) -> None:
|
||||
key = client.proxy.generate_key(KeyGenerateBody())
|
||||
resources.defer(lambda: client.proxy.delete_key(key))
|
||||
|
||||
status = _list_status(client, BudgetPageParams(), key=key)
|
||||
|
||||
assert status in (401, 403), (
|
||||
f"a non-admin key listing budgets must be refused 401/403, got {status}. Serving 200 with an "
|
||||
f"empty page would read as 'this proxy has no budgets'"
|
||||
)
|
||||
|
||||
|
||||
# ---------- customer / end-user ----------
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -58,10 +58,14 @@ def admin_viewer_client(monkeypatch):
|
|||
mock_config_table = MagicMock()
|
||||
mock_config_table.find_first = AsyncMock(return_value=None)
|
||||
|
||||
# /management/v1/budgets reads through query_raw: count first, then the page.
|
||||
mock_query_raw = AsyncMock(side_effect=[[{"count": 0}], []])
|
||||
|
||||
mock_prisma.db = types.SimpleNamespace(
|
||||
litellm_budgettable=mock_budget_table,
|
||||
litellm_invitationlink=mock_invitation_table,
|
||||
litellm_config=mock_config_table,
|
||||
query_raw=mock_query_raw,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(ps, "prisma_client", mock_prisma)
|
||||
|
|
@ -106,6 +110,14 @@ def test_budget_list_allows_admin_viewer(admin_viewer_client):
|
|||
assert resp.status_code == 200, resp.text
|
||||
|
||||
|
||||
def test_management_v1_budgets_allows_admin_viewer(admin_viewer_client):
|
||||
"""`/management/v1/budgets` is the paged/sortable budget list; same read tier as
|
||||
`/budget/list`, and it answers 403 rather than an empty page when it refuses."""
|
||||
resp = admin_viewer_client.get("/management/v1/budgets")
|
||||
_assert_not_role_blocked(resp)
|
||||
assert resp.status_code == 200, resp.text
|
||||
|
||||
|
||||
def test_budget_settings_allows_admin_viewer(admin_viewer_client):
|
||||
"""`/budget/settings` describes a budget's fields; read-only."""
|
||||
resp = admin_viewer_client.get("/budget/settings", params={"budget_id": "b1"})
|
||||
|
|
|
|||
|
|
@ -1960,6 +1960,7 @@ ADMIN_VIEWER_SETTINGS_ROUTES = [
|
|||
"/config/field/info",
|
||||
# Budgets page
|
||||
"/budget/list",
|
||||
"/management/v1/budgets",
|
||||
"/budget/settings",
|
||||
# Invitation viewing (admin viewer cannot create/delete; can read)
|
||||
"/invitation/info",
|
||||
|
|
|
|||
|
|
@ -0,0 +1,525 @@
|
|||
from dataclasses import replace
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
from fastapi import FastAPI, Request
|
||||
from fastapi.exceptions import RequestValidationError
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from litellm.proxy._types import LiteLLMRoutes, LitellmUserRoles
|
||||
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth
|
||||
from litellm.proxy.management_endpoints.management_v1 import router
|
||||
from litellm.proxy.management_endpoints.management_v1.budgets import (
|
||||
BUDGETS_LIST_SPEC,
|
||||
BudgetListItem,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.management_v1.common import (
|
||||
MANAGEMENT_V1_PREFIX,
|
||||
PROBLEM_TYPE_BASE,
|
||||
ManagementProblem,
|
||||
problem_response,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.management_v1.list_framework import (
|
||||
Compare,
|
||||
ScopeWhere,
|
||||
build_query_plan,
|
||||
)
|
||||
from litellm.types.proxy.management_endpoints.management_v1 import ProblemDetail
|
||||
|
||||
app = FastAPI()
|
||||
|
||||
|
||||
@app.exception_handler(ManagementProblem)
|
||||
async def management_problem_exception_handler(request: Request, exc: ManagementProblem):
|
||||
return problem_response(exc.problem)
|
||||
|
||||
|
||||
@app.exception_handler(RequestValidationError)
|
||||
async def validation_exception_handler(request: Request, exc: RequestValidationError):
|
||||
return problem_response(
|
||||
ProblemDetail(
|
||||
type=f"{PROBLEM_TYPE_BASE}invalid-query-parameter",
|
||||
title="Invalid query parameter",
|
||||
status=400,
|
||||
detail="The request query parameters are invalid.",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
app.include_router(router)
|
||||
client = TestClient(app)
|
||||
|
||||
BUDGETS_PATH = f"{MANAGEMENT_V1_PREFIX}/budgets"
|
||||
SORTABLE = ["budget_id", "created_at", "max_budget", "rpm_limit", "tpm_limit"]
|
||||
|
||||
|
||||
def _row(budget_id: str, **overrides: Any) -> dict[str, Any]:
|
||||
return {
|
||||
"budget_id": budget_id,
|
||||
"max_budget": 10.0,
|
||||
"soft_budget": None,
|
||||
"tpm_limit": None,
|
||||
"rpm_limit": None,
|
||||
"budget_duration": "30d",
|
||||
"budget_reset_at": None,
|
||||
"created_at": "2026-07-20T12:00:00+00:00",
|
||||
"updated_at": "2026-07-21T12:00:00+00:00",
|
||||
**overrides,
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def query_raw(monkeypatch):
|
||||
"""Mocks the one call the executor makes. `count` reads the first result, `find_many` the second."""
|
||||
mock = AsyncMock(side_effect=[[{"count": 0}], []])
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.query_raw = mock
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client)
|
||||
return mock
|
||||
|
||||
|
||||
def _serve(query_raw, rows: list[dict[str, Any]], total: int | None = None) -> None:
|
||||
query_raw.side_effect = [[{"count": len(rows) if total is None else total}], rows]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def as_proxy_admin():
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN
|
||||
)
|
||||
yield
|
||||
app.dependency_overrides.clear()
|
||||
|
||||
|
||||
def _as_role(role: LitellmUserRoles):
|
||||
original = app.dependency_overrides.copy()
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_id="u", user_role=role)
|
||||
return original
|
||||
|
||||
|
||||
def _get(query: str = ""):
|
||||
suffix = f"?{query}" if query else ""
|
||||
return client.get(f"{BUDGETS_PATH}{suffix}", headers={"Authorization": "Bearer k"})
|
||||
|
||||
|
||||
def _select_call(query_raw):
|
||||
"""The find_many call: (sql, *params). The count call comes first."""
|
||||
return query_raw.call_args_list[1].args
|
||||
|
||||
|
||||
def test_returns_flat_rows_in_the_control_plane_envelope(query_raw, as_proxy_admin):
|
||||
"""`{data, meta, links}` with flat rows; no JSON:API `{type, id, attributes}` wrapper."""
|
||||
_serve(query_raw, [_row("b-1")])
|
||||
|
||||
response = _get()
|
||||
|
||||
assert response.status_code == 200
|
||||
body = response.json()
|
||||
assert set(body) == {"data", "meta", "links"}
|
||||
assert body["data"][0]["budget_id"] == "b-1"
|
||||
assert "attributes" not in body["data"][0]
|
||||
|
||||
|
||||
def test_serves_the_columns_the_budgets_page_renders(query_raw, as_proxy_admin):
|
||||
_serve(query_raw, [_row("b-1", soft_budget=5.0, budget_reset_at="2026-08-01T00:00:00+00:00")])
|
||||
|
||||
row = _get().json()["data"][0]
|
||||
|
||||
assert set(row) == {
|
||||
"budget_id",
|
||||
"max_budget",
|
||||
"soft_budget",
|
||||
"tpm_limit",
|
||||
"rpm_limit",
|
||||
"budget_duration",
|
||||
"budget_reset_at",
|
||||
"created_at",
|
||||
"updated_at",
|
||||
}
|
||||
assert row["soft_budget"] == 5.0
|
||||
assert row["budget_reset_at"].startswith("2026-08-01T00:00:00")
|
||||
|
||||
|
||||
def test_selects_only_the_columns_it_serves(query_raw, as_proxy_admin):
|
||||
"""A `SELECT *` would ship created_by/updated_by and model_max_budget to the browser."""
|
||||
_serve(query_raw, [])
|
||||
|
||||
_get()
|
||||
|
||||
sql = _select_call(query_raw)[0]
|
||||
assert "SELECT *" not in sql
|
||||
assert '"budget_id"' in sql and '"max_budget"' in sql
|
||||
assert "created_by" not in sql and "model_max_budget" not in sql
|
||||
|
||||
|
||||
def test_defaults_to_newest_first_with_budget_id_breaking_ties(query_raw, as_proxy_admin):
|
||||
"""Two budgets created in the same transaction share a created_at; without the
|
||||
tiebreaker their relative order is undefined and pages can repeat or drop rows."""
|
||||
_serve(query_raw, [])
|
||||
|
||||
_get()
|
||||
|
||||
assert 'ORDER BY "created_at" DESC NULLS LAST, "budget_id" ASC NULLS LAST' in _select_call(query_raw)[0]
|
||||
|
||||
|
||||
def test_appends_the_tiebreaker_to_an_explicit_sort(query_raw, as_proxy_admin):
|
||||
_serve(query_raw, [])
|
||||
|
||||
_get("sort=-max_budget")
|
||||
|
||||
assert 'ORDER BY "max_budget" DESC NULLS LAST, "budget_id" ASC NULLS LAST' in _select_call(query_raw)[0]
|
||||
|
||||
|
||||
def test_refuses_to_sort_on_budget_duration(query_raw, as_proxy_admin):
|
||||
"""The column holds "7d"/"30d", so a lexicographic ORDER BY would put "30d"
|
||||
before "7d" and silently mis-order the page."""
|
||||
_serve(query_raw, [])
|
||||
|
||||
response = _get("sort=budget_duration")
|
||||
|
||||
assert response.status_code == 400
|
||||
assert response.headers["content-type"].startswith("application/problem+json")
|
||||
body = response.json()
|
||||
assert "budget_duration" in body["detail"]
|
||||
assert body["allowed"] == SORTABLE
|
||||
query_raw.assert_not_called()
|
||||
|
||||
|
||||
def test_the_advertised_sort_fields_are_the_ones_that_work(query_raw, as_proxy_admin):
|
||||
"""Guards the rejection above against drifting from what the spec actually accepts."""
|
||||
for field in SORTABLE:
|
||||
_serve(query_raw, [])
|
||||
assert _get(f"sort={field}").status_code == 200, field
|
||||
assert sorted(BUDGETS_LIST_SPEC.sortable) == SORTABLE
|
||||
|
||||
|
||||
def test_rejects_an_unknown_query_parameter(query_raw, as_proxy_admin):
|
||||
"""A silently ignored filter over-returns budgets, which is worse than a rejected request."""
|
||||
_serve(query_raw, [])
|
||||
|
||||
response = _get("filtre[max_budget][gte]=5")
|
||||
|
||||
assert response.status_code == 400
|
||||
assert response.headers["content-type"].startswith("application/problem+json")
|
||||
body = response.json()
|
||||
assert "filtre[max_budget][gte]" in body["detail"]
|
||||
assert "filter[max_budget][gte]" in body["allowed"]
|
||||
query_raw.assert_not_called()
|
||||
|
||||
|
||||
def test_rejects_an_operator_the_filter_does_not_declare(query_raw, as_proxy_admin):
|
||||
"""`max_budget` takes ranges, not `in`; accepting an undeclared operator is how a
|
||||
filter starts meaning something the query planner never checked."""
|
||||
for query in ("filter[max_budget][in]=5,10", "filter[created_at][is_null]=true"):
|
||||
_serve(query_raw, [])
|
||||
assert _get(query).status_code == 400, query
|
||||
|
||||
|
||||
def test_omitted_page_size_serves_fifty(query_raw, as_proxy_admin):
|
||||
_serve(query_raw, [])
|
||||
|
||||
body = _get().json()
|
||||
|
||||
assert body["meta"]["page_size"] == 50
|
||||
assert _select_call(query_raw)[-2] == 50
|
||||
|
||||
|
||||
def test_clamps_an_oversized_page_size_to_a_hundred(query_raw, as_proxy_admin):
|
||||
"""Unclamped, one request can ask the proxy to serialize the whole budget table."""
|
||||
_serve(query_raw, [])
|
||||
|
||||
body = _get("page_size=500").json()
|
||||
|
||||
assert body["meta"]["page_size"] == 100
|
||||
assert _select_call(query_raw)[-2] == 100
|
||||
|
||||
|
||||
def test_offsets_by_page(query_raw, as_proxy_admin):
|
||||
_serve(query_raw, [])
|
||||
|
||||
_get("page=3&page_size=25")
|
||||
|
||||
assert _select_call(query_raw)[-2:] == (25, 50)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"role",
|
||||
[
|
||||
LitellmUserRoles.INTERNAL_USER,
|
||||
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
|
||||
LitellmUserRoles.TEAM,
|
||||
],
|
||||
)
|
||||
def test_refuses_a_caller_without_admin_view(query_raw, role):
|
||||
"""Budgets are proxy-wide, so a caller who cannot read all of them must be told
|
||||
so. Answering 200 with an empty list would read as "there are no budgets"."""
|
||||
_serve(query_raw, [_row("b-1")])
|
||||
original = _as_role(role)
|
||||
try:
|
||||
response = _get()
|
||||
finally:
|
||||
app.dependency_overrides = original
|
||||
|
||||
assert response.status_code == 403
|
||||
assert response.headers["content-type"].startswith("application/problem+json")
|
||||
assert response.json()["status"] == 403
|
||||
query_raw.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("role", [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY])
|
||||
def test_admins_and_admin_viewers_may_read_every_budget(query_raw, role):
|
||||
_serve(query_raw, [_row("b-1")])
|
||||
original = _as_role(role)
|
||||
try:
|
||||
response = _get()
|
||||
finally:
|
||||
app.dependency_overrides = original
|
||||
|
||||
assert response.status_code == 200
|
||||
assert [row["budget_id"] for row in response.json()["data"]] == ["b-1"]
|
||||
assert "WHERE" not in _select_call(query_raw)[0]
|
||||
|
||||
|
||||
def test_a_denied_caller_stays_denied_whatever_they_filter_on(query_raw):
|
||||
"""The scope decision reads the caller, never the query string."""
|
||||
_serve(query_raw, [_row("b-1")])
|
||||
original = _as_role(LitellmUserRoles.INTERNAL_USER)
|
||||
try:
|
||||
response = _get("filter[max_budget][gte]=0&q=b-")
|
||||
finally:
|
||||
app.dependency_overrides = original
|
||||
|
||||
assert response.status_code == 403
|
||||
|
||||
|
||||
def test_a_filter_sits_behind_the_scope_predicate_instead_of_replacing_it():
|
||||
"""Same spec, planned for a row-scoped caller: the scope clause has to survive, and
|
||||
lead, whatever the caller filtered on. Replacing it would let a filter widen a read."""
|
||||
scoped = replace(
|
||||
BUDGETS_LIST_SPEC,
|
||||
scope=lambda _caller: ScopeWhere(where=(Compare(field="budget_id", op="eq", value="b-1"),)),
|
||||
)
|
||||
|
||||
plan = build_query_plan(
|
||||
spec=scoped,
|
||||
params={"filter[max_budget][gte]": "5"},
|
||||
caller=UserAPIKeyAuth(user_id="u", user_role=LitellmUserRoles.INTERNAL_USER),
|
||||
)
|
||||
|
||||
assert plan.where[0] == Compare(field="budget_id", op="eq", value="b-1")
|
||||
assert Compare(field="max_budget", op="gte", value=5.0) in plan.where
|
||||
|
||||
|
||||
def test_q_matches_budget_id_case_insensitively(query_raw, as_proxy_admin):
|
||||
"""budget_id is the only text identity on the row; matching anything else would
|
||||
return budgets whose ids do not contain what the user typed."""
|
||||
_serve(query_raw, [])
|
||||
|
||||
_get("q=Prod")
|
||||
|
||||
sql, *params = _select_call(query_raw)
|
||||
assert '"budget_id" ILIKE $1' in sql
|
||||
assert params[0] == "%Prod%"
|
||||
|
||||
|
||||
def test_q_escapes_like_metacharacters(query_raw, as_proxy_admin):
|
||||
"""Budget ids routinely contain '_'; unescaped it is a single-character wildcard."""
|
||||
_serve(query_raw, [])
|
||||
|
||||
_get("q=team_a%25")
|
||||
|
||||
assert _select_call(query_raw)[1] == r"%team\_a\%%"
|
||||
|
||||
|
||||
def test_q_does_not_search_any_other_column(query_raw, as_proxy_admin):
|
||||
_serve(query_raw, [])
|
||||
|
||||
_get("q=30d")
|
||||
|
||||
assert "budget_duration" not in _select_call(query_raw)[0].split("WHERE")[1]
|
||||
assert BUDGETS_LIST_SPEC.searchable == frozenset({"budget_id"})
|
||||
|
||||
|
||||
def test_is_null_selects_the_unlimited_budgets(query_raw, as_proxy_admin):
|
||||
""""Unlimited" is max_budget IS NULL; `max_budget = 0` would be a hard zero cap."""
|
||||
_serve(query_raw, [_row("b-unlimited", max_budget=None)])
|
||||
|
||||
body = _get("filter[max_budget][is_null]=true").json()
|
||||
|
||||
assert '"max_budget" IS NULL' in _select_call(query_raw)[0]
|
||||
assert body["data"][0]["max_budget"] is None
|
||||
|
||||
|
||||
def test_is_null_false_selects_the_capped_budgets(query_raw, as_proxy_admin):
|
||||
_serve(query_raw, [])
|
||||
|
||||
_get("filter[max_budget][is_null]=false")
|
||||
|
||||
assert '"max_budget" IS NOT NULL' in _select_call(query_raw)[0]
|
||||
|
||||
|
||||
def test_in_filter_binds_each_requested_duration(query_raw, as_proxy_admin):
|
||||
_serve(query_raw, [])
|
||||
|
||||
_get("filter[budget_duration][in]=7d,30d")
|
||||
|
||||
sql, *params = _select_call(query_raw)
|
||||
assert '"budget_duration" IN ($1, $2)' in sql
|
||||
assert params[:2] == ["7d", "30d"]
|
||||
|
||||
|
||||
def test_created_at_range_is_bound_as_a_timestamp(query_raw, as_proxy_admin):
|
||||
"""The bind crosses into the query engine as JSON, so an uncast placeholder reaches
|
||||
Postgres as text and `timestamp >= text` is a hard error, not a wrong answer."""
|
||||
_serve(query_raw, [])
|
||||
|
||||
_get("filter[created_at][gte]=2026-07-01T00:00:00Z")
|
||||
|
||||
sql, *params = _select_call(query_raw)
|
||||
assert "\"created_at\" >= $1::timestamptz AT TIME ZONE 'UTC'" in sql
|
||||
assert params[0] == datetime(2026, 7, 1, tzinfo=timezone.utc)
|
||||
|
||||
|
||||
def test_numbers_placeholders_continuously_across_predicates(query_raw, as_proxy_admin):
|
||||
"""Each predicate is numbered after the binds the ones before it consumed. Restart
|
||||
the count and `$1` gets read as the duration while the search string goes unbound."""
|
||||
_serve(query_raw, [])
|
||||
|
||||
_get("filter[budget_duration][in]=7d,30d&filter[max_budget][gte]=5&q=prod")
|
||||
|
||||
sql, *params = _select_call(query_raw)
|
||||
assert '"budget_duration" IN ($1, $2)' in sql
|
||||
assert '"max_budget" >= $3' in sql
|
||||
assert '"budget_id" ILIKE $4' in sql
|
||||
assert params[:4] == ["7d", "30d", 5.0, "%prod%"]
|
||||
|
||||
|
||||
def test_a_non_datetime_bind_is_not_cast(query_raw, as_proxy_admin):
|
||||
"""Guards the cast above from being applied to every placeholder."""
|
||||
_serve(query_raw, [])
|
||||
|
||||
_get("filter[max_budget][gte]=5")
|
||||
|
||||
sql = _select_call(query_raw)[0]
|
||||
assert '"max_budget" >= $1' in sql
|
||||
assert "timestamptz" not in sql
|
||||
|
||||
|
||||
def test_an_offsetless_created_at_bound_is_read_as_utc(query_raw, as_proxy_admin):
|
||||
"""The dashboard sends 'YYYY-MM-DDTHH:MM:SS' with no offset. Left naive, Postgres
|
||||
would compare it in the session timezone and shift the window off the rows shown."""
|
||||
_serve(query_raw, [])
|
||||
|
||||
_get("filter[created_at][gte]=2026-07-01T00:00:00")
|
||||
|
||||
bound = _select_call(query_raw)[1]
|
||||
assert bound == datetime(2026, 7, 1, tzinfo=timezone.utc)
|
||||
assert bound.tzinfo is not None
|
||||
|
||||
|
||||
def test_rejects_a_filter_value_that_is_not_of_the_declared_type(query_raw, as_proxy_admin):
|
||||
for query in ("filter[max_budget][gte]=lots", "filter[created_at][gte]=yesterday"):
|
||||
_serve(query_raw, [])
|
||||
assert _get(query).status_code == 400, query
|
||||
|
||||
|
||||
def test_reports_the_total_and_links_every_page_on_a_middle_page(query_raw, as_proxy_admin):
|
||||
"""The Budgets page renders a page count, so the total has to be the match total,
|
||||
not the length of the page it just received."""
|
||||
_serve(query_raw, [_row("b-3"), _row("b-4")], total=7)
|
||||
|
||||
body = _get("page=2&page_size=2").json()
|
||||
|
||||
assert body["meta"] == {"page": 2, "page_size": 2, "total_count": 7, "total_pages": 4}
|
||||
links = body["links"]
|
||||
assert "page=1" in links["first"] and "page_size=2" in links["first"]
|
||||
assert "page=1" in links["prev"]
|
||||
assert "page=3" in links["next"]
|
||||
assert "page=4" in links["last"]
|
||||
assert "page=2" in links["self"]
|
||||
|
||||
|
||||
def test_the_last_page_has_no_next(query_raw, as_proxy_admin):
|
||||
_serve(query_raw, [_row("b-5")], total=5)
|
||||
|
||||
links = _get("page=3&page_size=2").json()["links"]
|
||||
|
||||
assert links["next"] is None
|
||||
assert "page=2" in links["prev"]
|
||||
|
||||
|
||||
def test_the_first_page_has_no_prev(query_raw, as_proxy_admin):
|
||||
_serve(query_raw, [_row("b-1")], total=5)
|
||||
|
||||
links = _get("page_size=2").json()["links"]
|
||||
|
||||
assert links["prev"] is None
|
||||
assert "page=2" in links["next"]
|
||||
|
||||
|
||||
def test_an_empty_table_still_links_a_first_and_last_page(query_raw, as_proxy_admin):
|
||||
_serve(query_raw, [], total=0)
|
||||
|
||||
body = _get().json()
|
||||
|
||||
assert body["meta"]["total_count"] == 0
|
||||
assert body["meta"]["total_pages"] == 0
|
||||
assert "page=1" in body["links"]["first"] and "page=1" in body["links"]["last"]
|
||||
|
||||
|
||||
def test_counts_over_the_same_predicate_it_pages(query_raw, as_proxy_admin):
|
||||
"""A total counted without the caller's filter would page through rows the
|
||||
filter excluded."""
|
||||
_serve(query_raw, [], total=0)
|
||||
|
||||
_get("filter[budget_duration][in]=30d")
|
||||
|
||||
count_sql, *count_params = query_raw.call_args_list[0].args
|
||||
assert '"budget_duration" IN ($1)' in count_sql
|
||||
assert count_params == ["30d"]
|
||||
assert "COUNT(*)" in count_sql
|
||||
|
||||
|
||||
def test_bigint_limits_serialize_as_json_numbers(query_raw, as_proxy_admin):
|
||||
"""tpm_limit/rpm_limit are BigInt? in Prisma; the query engine can hand them back
|
||||
as decimal strings, and a quoted "60000" breaks arithmetic in the dashboard."""
|
||||
_serve(query_raw, [_row("b-1", tpm_limit="60000", rpm_limit=1200)])
|
||||
|
||||
response = _get()
|
||||
row = response.json()["data"][0]
|
||||
|
||||
assert row["tpm_limit"] == 60000
|
||||
assert row["rpm_limit"] == 1200
|
||||
assert isinstance(row["tpm_limit"], int) and not isinstance(row["tpm_limit"], bool)
|
||||
assert '"tpm_limit": "60000"' not in response.text
|
||||
|
||||
|
||||
def test_reports_a_missing_database_as_a_problem_document(monkeypatch, as_proxy_admin):
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
|
||||
|
||||
response = _get()
|
||||
|
||||
assert response.status_code == 503
|
||||
assert response.headers["content-type"].startswith("application/problem+json")
|
||||
|
||||
|
||||
def test_the_spec_serves_what_the_row_model_declares():
|
||||
"""SELECTED_COLUMNS is built off the model, so a field added to one cannot go
|
||||
missing from the other and produce a row the validator rejects."""
|
||||
from litellm.proxy.management_endpoints.management_v1.budgets import SELECTED_COLUMNS
|
||||
|
||||
assert SELECTED_COLUMNS == ", ".join(f'"{name}"' for name in BudgetListItem.model_fields)
|
||||
assert BUDGETS_LIST_SPEC.tiebreaker in BudgetListItem.model_fields
|
||||
assert BUDGETS_LIST_SPEC.sortable <= frozenset(BudgetListItem.model_fields)
|
||||
assert frozenset(BUDGETS_LIST_SPEC.filters) <= frozenset(BudgetListItem.model_fields)
|
||||
|
||||
|
||||
def test_is_reachable_by_the_roles_that_can_open_the_budgets_page():
|
||||
"""Route-level auth gate, which the dependency_overrides above bypass. The handler's
|
||||
admin-view check is dead code if RouteChecks rejects the role first."""
|
||||
assert BUDGETS_PATH in LiteLLMRoutes.admin_viewer_routes.value
|
||||
assert ("/budget/list" in LiteLLMRoutes.admin_viewer_routes.value) == (
|
||||
BUDGETS_PATH in LiteLLMRoutes.admin_viewer_routes.value
|
||||
)
|
||||
129
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
129
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -7169,6 +7169,44 @@ export interface paths {
|
|||
patch?: never;
|
||||
trace?: never;
|
||||
};
|
||||
"/management/v1/budgets": {
|
||||
parameters: {
|
||||
query?: never;
|
||||
header?: never;
|
||||
path?: never;
|
||||
cookie?: never;
|
||||
};
|
||||
/**
|
||||
* List Budgets
|
||||
* @description The budgets defined on this proxy, paged, sortable and filterable, for the
|
||||
* Budgets page.
|
||||
*
|
||||
* Readable by a proxy admin or an admin viewer; anyone else is refused 403. The
|
||||
* older `/budget/list` answers with the whole table as a bare array and has no
|
||||
* way to page, sort or filter it.
|
||||
*
|
||||
* `sort` takes a comma-separated list of `budget_id`, `max_budget`, `tpm_limit`,
|
||||
* `rpm_limit` or `created_at`, each optionally prefixed with `-` for descending,
|
||||
* and defaults to `-created_at`. `budget_id` is appended to every sort as the
|
||||
* tiebreaker. `q` is a case-insensitive substring match on `budget_id`.
|
||||
* `page_size` defaults to 50 and is capped at 100. Filters are
|
||||
* `filter[budget_duration][in|is_null]`, `filter[max_budget][gte|lte|is_null]`
|
||||
* and `filter[created_at][gte|lte]`.
|
||||
*
|
||||
* Example curl:
|
||||
* ```
|
||||
* curl --location --globoff 'http://0.0.0.0:4000/management/v1/budgets?sort=-max_budget&filter[budget_duration][in]=7d,30d&page_size=25' --header 'Authorization: Bearer sk-1234'
|
||||
* ```
|
||||
*/
|
||||
get: operations["list_budgets_management_v1_budgets_get"];
|
||||
put?: never;
|
||||
post?: never;
|
||||
delete?: never;
|
||||
options?: never;
|
||||
head?: never;
|
||||
patch?: never;
|
||||
trace?: never;
|
||||
};
|
||||
"/management/v1/spend_logs/end_users": {
|
||||
parameters: {
|
||||
query?: never;
|
||||
|
|
@ -21533,6 +21571,40 @@ export interface components {
|
|||
/** Reset At */
|
||||
reset_at?: string | null;
|
||||
};
|
||||
/**
|
||||
* BudgetListItem
|
||||
* @description One budget as the Budgets page reads it, and as it comes back off the table.
|
||||
*
|
||||
* Validating the raw row through here is what makes `tpm_limit` / `rpm_limit`
|
||||
* numbers: they are `BigInt?` in the schema, which the query engine hands back as
|
||||
* decimal strings, and a quoted "60000" breaks arithmetic in the dashboard.
|
||||
*/
|
||||
BudgetListItem: {
|
||||
/** Budget Duration */
|
||||
budget_duration?: string | null;
|
||||
/** Budget Id */
|
||||
budget_id: string;
|
||||
/** Budget Reset At */
|
||||
budget_reset_at?: string | null;
|
||||
/**
|
||||
* Created At
|
||||
* Format: date-time
|
||||
*/
|
||||
created_at: string;
|
||||
/** Max Budget */
|
||||
max_budget?: number | null;
|
||||
/** Rpm Limit */
|
||||
rpm_limit?: number | null;
|
||||
/** Soft Budget */
|
||||
soft_budget?: number | null;
|
||||
/** Tpm Limit */
|
||||
tpm_limit?: number | null;
|
||||
/**
|
||||
* Updated At
|
||||
* Format: date-time
|
||||
*/
|
||||
updated_at: string;
|
||||
};
|
||||
/** BudgetNewRequest */
|
||||
BudgetNewRequest: {
|
||||
/**
|
||||
|
|
@ -24928,6 +25000,36 @@ export interface components {
|
|||
/** Guardrails */
|
||||
guardrails: components["schemas"]["GuardrailInfoResponse"][];
|
||||
};
|
||||
/**
|
||||
* ListLinks
|
||||
* @description Page-mode counterpart to `PageLinks`. `first`/`last` are knowable here because the total count is.
|
||||
*/
|
||||
ListLinks: {
|
||||
/** First */
|
||||
first: string;
|
||||
/** Last */
|
||||
last: string;
|
||||
/** Next */
|
||||
next?: string | null;
|
||||
/** Prev */
|
||||
prev?: string | null;
|
||||
/** Self */
|
||||
self: string;
|
||||
};
|
||||
/**
|
||||
* ListMeta
|
||||
* @description Page-mode counterpart to `PageMeta`: an entity list pays for the COUNT(*) so the table can show a page count.
|
||||
*/
|
||||
ListMeta: {
|
||||
/** Page */
|
||||
page: number;
|
||||
/** Page Size */
|
||||
page_size: number;
|
||||
/** Total Count */
|
||||
total_count: number;
|
||||
/** Total Pages */
|
||||
total_pages: number;
|
||||
};
|
||||
/**
|
||||
* ListPluginsResponse
|
||||
* @description Response from listing plugins.
|
||||
|
|
@ -24943,6 +25045,13 @@ export interface components {
|
|||
/** Prompts */
|
||||
prompts: components["schemas"]["PromptSpec"][];
|
||||
};
|
||||
/** ListResponse[BudgetListItem] */
|
||||
ListResponse_BudgetListItem_: {
|
||||
/** Data */
|
||||
data: components["schemas"]["BudgetListItem"][];
|
||||
links: components["schemas"]["ListLinks"];
|
||||
meta: components["schemas"]["ListMeta"];
|
||||
};
|
||||
/**
|
||||
* ListRunsResponse
|
||||
* @description Response from listing runs
|
||||
|
|
@ -43479,6 +43588,26 @@ export interface operations {
|
|||
};
|
||||
};
|
||||
};
|
||||
list_budgets_management_v1_budgets_get: {
|
||||
parameters: {
|
||||
query?: never;
|
||||
header?: never;
|
||||
path?: never;
|
||||
cookie?: never;
|
||||
};
|
||||
requestBody?: never;
|
||||
responses: {
|
||||
/** @description Successful Response */
|
||||
200: {
|
||||
headers: {
|
||||
[name: string]: unknown;
|
||||
};
|
||||
content: {
|
||||
"application/json": components["schemas"]["ListResponse_BudgetListItem_"];
|
||||
};
|
||||
};
|
||||
};
|
||||
};
|
||||
list_spend_log_end_users_management_v1_spend_logs_end_users_get: {
|
||||
parameters: {
|
||||
query: {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue