feat(ui): sort budgets by reset duration (#48107)

* feat(ui): sort budgets by reset duration

* fix(ui): order unlimited budget limits as infinity

* fix(ui): sort localized budget reset durations

litellm/proxy/management_endpoints/management_v1/budgets.py

@@ -3,6 +3,7 @@
 from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
 from datetime import datetime
+from itertools import chain
 from types import MappingProxyType
 from typing import Annotated, Final
 
@@ -76,6 +77,32 @@ _ROW_COUNTS: Final = TypeAdapter(tuple[_RowCount, ...])
 SELECTED_COLUMNS: Final = ", ".join(f'"{name}"' for name in BudgetListItem.model_fields)
 
 
+# Unicode 17.0 Nd zeroes, each followed by its other nine digits in UnicodeData.txt.
+_DECIMAL_ZEROES: Final = "0٠۰߀०০੦૦୦௦౦೦൦෦๐໐༠၀႐០᠐᥆᧐᪀᪐᭐᮰᱀᱐꘠꣐꤀꧐꧰꩐꯰０𐒠𐴰𐵀𑁦𑃰𑄶𑇐𑋰𑑐𑓐𑙐𑛀𑛐𑛚𑜰𑣠𑥐𑯰𑱐𑵐𑶠\U00011de0𑽐𖄰𖩠𖫀𖭐𖵰𜳰𝟎𝟘𝟢𝟬𝟶𞅀𞋰𞓰𞗱𞥐🯰"
+_DECIMAL_DIGITS: Final = "".join(
+    chain.from_iterable(tuple(chr(ord(zero) + digit) for digit in range(10)) for zero in _DECIMAL_ZEROES)
+)
+_ASCII_DIGITS: Final = "0123456789" * len(_DECIMAL_ZEROES)
+
+
+_BUDGET_DURATION_SECONDS: Final = f"""(
+    SELECT EXTRACT(EPOCH FROM (
+        (CURRENT_TIMESTAMP AT TIME ZONE 'UTC') +
+        (parts[1] || ' ' || CASE parts[2]
+            WHEN 's' THEN 'seconds' WHEN 'm' THEN 'minutes' WHEN 'h' THEN 'hours'
+            WHEN 'd' THEN 'days' WHEN 'w' THEN 'weeks' WHEN 'mo' THEN 'months' END)::interval
+        - (CURRENT_TIMESTAMP AT TIME ZONE 'UTC')
+    )) FROM regexp_match(
+        CASE lower(trim(budget_duration))
+            WHEN 'hourly' THEN '1h' WHEN 'daily' THEN '24h'
+            WHEN 'weekly' THEN '7d' WHEN 'monthly' THEN '30d'
+            ELSE CASE WHEN octet_length(budget_duration) = length(budget_duration) THEN budget_duration
+                ELSE translate(budget_duration, '{_DECIMAL_DIGITS}', '{_ASCII_DIGITS}') END END,
+        '^([0-9]+)(mo|[smhdw])'
+    ) AS duration(parts)
+)"""
+
+
 @dataclass(frozen=True, slots=True)
 class PrismaBudgetListExecutor:
     """The database half of the budgets list. Every caller-supplied value is bound to a
@@ -91,10 +118,16 @@ class PrismaBudgetListExecutor:
 
     async def find_many(self, plan: QueryPlan) -> Sequence[BudgetListItem]:
         clauses, params = where_sql(plan.where)
+        ordering: Final = ", ".join(
+            order_by_sql((key,)).replace("NULLS LAST", "NULLS FIRST")
+            if key.descending and key.field in ("max_budget", "tpm_limit", "rpm_limit", "tpd_limit")
+            else order_by_sql((key,))
+            for key in plan.order
+        ).replace('"budget_duration"', _BUDGET_DURATION_SECONDS)
         sql: Final = (
             f"SELECT {SELECTED_COLUMNS} FROM {BUDGET_TABLE}"
             + (f" WHERE {clauses}" if clauses else "")
-            + f" ORDER BY {order_by_sql(plan.order)}"
+            + f" ORDER BY {ordering}"
             + f" LIMIT ${len(params) + 1} OFFSET ${len(params) + 2}"
         )
         return _BUDGET_ROWS.validate_python(await self.prisma_client.db.query_raw(sql, *params, plan.take, plan.skip))
@@ -111,8 +144,6 @@ def _scope(caller: UserAPIKeyAuth) -> Scope:
     return ScopeDenied(reason=f"Only proxy admins can list budgets, your role={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: Final[Mapping[str, FilterSpec]] = MappingProxyType(
     {
         "budget_duration": FilterSpec(type=str, ops=frozenset(("in", "is_null"))),
@@ -123,7 +154,9 @@ BUDGET_FILTERS: Final[Mapping[str, FilterSpec]] = MappingProxyType(
 
 BUDGETS_LIST_SPEC: Final[ListSpec[BudgetListItem, BudgetListItem]] = ListSpec(
     resource="budgets",
-    sortable=frozenset(("budget_id", "max_budget", "tpm_limit", "rpm_limit", "tpd_limit", "created_at")),
+    sortable=frozenset(
+        ("budget_id", "max_budget", "tpm_limit", "rpm_limit", "tpd_limit", "created_at", "budget_duration")
+    ),
     searchable=frozenset(("budget_id",)),
     filters=BUDGET_FILTERS,
     default_sort=(SortKey(field="created_at", descending=True),),
@@ -154,8 +187,9 @@ async def list_budgets(
     way to page, sort or filter it.
 
     `sort` takes a comma-separated list of `budget_id`, `max_budget`, `tpm_limit`,
-    `rpm_limit`, `tpd_limit` or `created_at`, each optionally prefixed with `-` for descending,
-    and defaults to `-created_at`. `budget_id` is appended to every sort as the
+    `rpm_limit`, `tpd_limit`, `created_at` or `budget_duration`, each optionally prefixed with `-` for descending,
+    and defaults to `-created_at`. Reset periods sort by elapsed time, with calendar months
+    measured from the current UTC date and unset periods last. `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]`

tests/integration/management/test_budget_list_sorting.py

@@ -0,0 +1,94 @@
+from datetime import datetime, timezone
+from itertools import chain
+from typing import Final
+from uuid import uuid4
+
+import pytest
+
+from litellm.litellm_core_utils.duration_parser import duration_in_seconds_from
+from tests.integration._support.client import Gateway, Scenario
+
+
+@pytest.mark.parametrize("descending", (False, True))
+def test_budget_reset_sort_orders_durations_before_pagination(
+    descending: bool,
+    gateway: Gateway,
+    scenario: Scenario,
+) -> None:
+    prefix: Final = f"sort-{uuid4().hex}"
+    durations: Final = ("30d", "7d", "1h", "60m", "1mo", "weekly", "monthly", "٥d", "۵d", "५d", "５d", "𝟝d", "5d", None)
+    ids: Final = tuple(
+        scenario.budget(budget_id=f"{prefix}-{index}", budget_duration=duration, max_budget=10)
+        for index, duration in enumerate(durations)
+    )
+    now: Final = datetime.now(timezone.utc)
+    expected: Final = tuple(
+        budget_id
+        for _, budget_id in sorted(
+            zip(durations, ids, strict=True),
+            key=lambda item: (
+                item[0] is None,
+                duration_in_seconds_from(now, item[0]) * (-1 if descending else 1) if item[0] else 0,
+                item[1],
+            ),
+        )
+    )
+    pages: Final = tuple(
+        gateway.request(
+            "GET",
+            "/management/v1/budgets",
+            params={
+                "q": prefix,
+                "sort": "-budget_duration" if descending else "budget_duration",
+                "page_size": "3",
+                "page": str(page),
+            },
+        )
+        for page in range(1, -(-len(ids) // 3) + 1)
+    )
+    assert all(page.status_code == 200 for page in pages), [page.text for page in pages]
+    assert tuple(row["budget_id"] for row in chain.from_iterable(page.json()["data"] for page in pages)) == expected
+    filtered: Final = gateway.request(
+        "GET",
+        "/management/v1/budgets",
+        params={
+            "q": prefix,
+            "sort": "budget_duration",
+            "filter[budget_duration][in]": "7d,30d",
+        },
+    )
+    assert filtered.status_code == 200, filtered.text
+    assert [row["budget_duration"] for row in filtered.json()["data"]] == ["7d", "30d"]
+
+
+@pytest.mark.parametrize("field", ("max_budget", "tpm_limit", "rpm_limit", "tpd_limit"))
+@pytest.mark.parametrize("descending", (False, True))
+def test_budget_limits_sort_unlimited_as_infinity(
+    field: str, descending: bool, gateway: Gateway, scenario: Scenario
+) -> None:
+    prefix: Final = f"sort-unlimited-{uuid4().hex}"
+    values: Final = (20, None, 0, 3)
+    ids: Final = tuple(scenario.budget(budget_id=f"{prefix}-{i}", **{field: value}) for i, value in enumerate(values))
+    expected: Final = tuple(
+        identity
+        for value, identity in sorted(
+            zip(values, ids, strict=True),
+            key=lambda item: float("inf") if item[0] is None else item[0],
+            reverse=descending,
+        )
+    )
+    pages: Final = tuple(
+        gateway.request(
+            "GET",
+            "/management/v1/budgets",
+            params={
+                "q": prefix,
+                "sort": ("-" if descending else "") + field,
+                "page_size": "2",
+                "page": str(page),
+            },
+        )
+        for page in (1, 2)
+    )
+    assert all(page.status_code == 200 for page in pages), [page.text for page in pages]
+    assert tuple(row["budget_id"] for row in chain.from_iterable(page.json()["data"] for page in pages)) == expected

tests/unit/proxy/management_endpoints/management_v1/test_budgets.py

@@ -52,7 +52,7 @@ app.include_router(router)
 client = TestClient(app)
 
 BUDGETS_PATH = f"{MANAGEMENT_V1_PREFIX}/budgets"
-SORTABLE = ["budget_id", "created_at", "max_budget", "rpm_limit", "tpd_limit", "tpm_limit"]
+SORTABLE = ["budget_duration", "budget_id", "created_at", "max_budget", "rpm_limit", "tpd_limit", "tpm_limit"]
 
 
 def _row(budget_id: str, **overrides: Any) -> dict[str, Any]:
@@ -172,22 +172,12 @@ def test_appends_the_tiebreaker_to_an_explicit_sort(query_raw, as_proxy_admin):
 
     _get("sort=-max_budget")
 
-    assert 'ORDER BY "max_budget" DESC NULLS LAST, "budget_id" ASC NULLS LAST' in _select_call(query_raw)[0]
+    assert 'ORDER BY "max_budget" DESC NULLS FIRST, "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, [])
-
+def test_accepts_duration_sorting(query_raw, as_proxy_admin):
     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()
+    assert response.status_code == 200
 
 
 def test_the_advertised_sort_fields_are_the_ones_that_work(query_raw, as_proxy_admin):
@@ -346,7 +336,7 @@ def test_q_does_not_search_any_other_column(query_raw, as_proxy_admin):
 
 
 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."""
+    """ "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()

ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.test.tsx

@@ -134,13 +134,16 @@ describe("BudgetTable", () => {
     }
   });
 
-  it("should not make the reset column sortable", async () => {
+  it("should request duration sorting from the server", async () => {
     const user = userEvent.setup();
-    renderWithProviders(<BudgetTable {...defaultProps} list={makeList()} />);
+    const onSortingChange = vi.fn();
+    renderWithProviders(<BudgetTable {...defaultProps} list={makeList({ onSortingChange })} />);
     await showColumn(user, "budget_duration");
     const headers = screen.getAllByRole("columnheader").map((header) => header.textContent);
     expect(headers).toContain("Reset");
-    expect(screen.queryByTestId("sort-header-budget_duration")).not.toBeInTheDocument();
+    await user.click(screen.getByTestId("sort-header-budget_duration"));
+    expect(onSortingChange).toHaveBeenCalled();
+    expect(onSortingChange.mock.calls[0][0]([])).toEqual([{ id: "budget_duration", desc: false }]);
   });
 
   it("should ask the list for a new sort when a sortable header is clicked", async () => {

ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTableColumns.tsx

@@ -139,8 +139,7 @@ export const getBudgetTableColumns = ({
     accessorKey: "budget_duration",
     filterFn: serverFilter,
     meta: { title: "Reset" },
-    // "7d"/"30d" sort lexicographically, not chronologically, so the route does not offer it.
-    enableSorting: false,
+    sortDescFirst: false,
     header: ({ column }) => <DataTableSortHeader column={column} title="Reset" />,
     size: 110,
     cell: ({ row }) => <BudgetDurationCell value={row.original.budget_duration} />,

ui/litellm-dashboard/src/lib/http/schema.d.ts

@@ -9280,8 +9280,9 @@ export interface paths {
          *     way to page, sort or filter it.
          *
          *     `sort` takes a comma-separated list of `budget_id`, `max_budget`, `tpm_limit`,
-         *     `rpm_limit`, `tpd_limit` or `created_at`, each optionally prefixed with `-` for descending,
-         *     and defaults to `-created_at`. `budget_id` is appended to every sort as the
+         *     `rpm_limit`, `tpd_limit`, `created_at` or `budget_duration`, each optionally prefixed with `-` for descending,
+         *     and defaults to `-created_at`. Reset periods sort by elapsed time, with calendar months
+         *     measured from the current UTC date and unset periods last. `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]`