diff --git a/litellm/proxy/management_endpoints/budget_management_endpoints.py b/litellm/proxy/management_endpoints/budget_management_endpoints.py index b3f372327dc..48d78488894 100644 --- a/litellm/proxy/management_endpoints/budget_management_endpoints.py +++ b/litellm/proxy/management_endpoints/budget_management_endpoints.py @@ -141,6 +141,7 @@ async def update_budget( - rpm_limit: Optional[int] - The requests per minute limit for the budget. - tpd_limit: Optional[int] - The tokens per day limit for the budget. Charged by batch submissions instead of tpm_limit/rpm_limit. - 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}} + Send null or {} to clear all per-model limits. Omit this field to keep them unchanged. - budget_reset_at: Optional[datetime] - Update the Datetime when the budget was last reset. """ from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client, user_api_key_cache @@ -189,9 +190,13 @@ async def update_budget( else {} ) + budget_patch: Final[Mapping[str, object]] = budget_obj.model_dump(exclude_unset=True) budget_obj_jsonified: Final[Mapping[str, object]] = jsonify_object( { - **budget_obj.model_dump(exclude_unset=True), + **{ + field: {} if field == "model_max_budget" and value is None else value + for field, value in budget_patch.items() + }, **recomputed_reset_at, "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, } diff --git a/tests/integration/management/test_budget_updates.py b/tests/integration/management/test_budget_updates.py index 6939b447866..faf61876452 100644 --- a/tests/integration/management/test_budget_updates.py +++ b/tests/integration/management/test_budget_updates.py @@ -1,3 +1,4 @@ +import json from datetime import datetime, timedelta, timezone from pathlib import Path from typing import Final @@ -6,7 +7,7 @@ import pytest from pydantic import JsonValue, TypeAdapter from tests.integration._support.client import Gateway, string_value -from tests.integration._support.database import read_rows +from tests.integration._support.database import read_rows, write_rows from tests.integration._support.process import owned_proxy pytestmark: Final = pytest.mark.shares_worker_state @@ -213,3 +214,38 @@ def test_model_max_budget_update_is_refused_by_the_handler_without_a_license(gat served: Final = unlicensed.request("POST", "/budget/update", {"budget_id": budget_id, "max_budget": 20.0}) assert served.status_code == 200, served.text assert _budget_row(budget_id) == {**before, "max_budget": 20.0} + + +@pytest.mark.parametrize("model_max_budget", [None, {}], ids=["null", "empty-object"]) +def test_clearing_model_budgets_preserves_other_limits( + gateway: Gateway, model_max_budget: dict[str, JsonValue] | None +) -> None: + stored_model_budgets: Final = { + "qa-budget-a": {"max_budget": 1.0, "budget_duration": "1d"}, + "qa-budget-b": {"max_budget": 2.0, "budget_duration": "7d"}, + } + with gateway.scenario() as scenario: + budget_id: Final = scenario.budget( + max_budget=20.0, + soft_budget=5.0, + tpm_limit=1000, + rpm_limit=60, + budget_duration="30d", + ) + write_rows( + 'UPDATE "LiteLLM_BudgetTable" SET model_max_budget = %s::jsonb WHERE budget_id = %s', + (json.dumps(stored_model_budgets), budget_id), + ) + before_row: Final = _budget_row(budget_id) + before_info: Final = _budget_info(gateway, budget_id) + assert before_row["model_max_budget"] == stored_model_budgets + + response: Final = gateway.request( + "POST", "/budget/update", {"budget_id": budget_id, "model_max_budget": model_max_budget} + ) + + assert response.status_code == 200, response.text + assert _budget_row(budget_id) == {**before_row, "model_max_budget": {}} + info: Final = _budget_info(gateway, budget_id) + assert info == {**before_info, "model_max_budget": {}, "updated_at": info["updated_at"]} + assert response.json() == info, response.text diff --git a/tests/unit/proxy/management_endpoints/test_budget_endpoints.py b/tests/unit/proxy/management_endpoints/test_budget_endpoints.py index 24129488f21..9bb37a5e3b9 100644 --- a/tests/unit/proxy/management_endpoints/test_budget_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_budget_endpoints.py @@ -477,3 +477,33 @@ async def test_update_budget_serializes_model_max_budget_for_prisma( "gpt4o": {"max_budget": 5.0, "budget_duration": "1d"}, "glm-5.2": {"max_budget": 7.5, "budget_duration": "30d"}, } + + +@pytest.mark.parametrize( + ("model_patch", "expected_model_patch"), + [ + pytest.param({}, {}, id="omitted"), + pytest.param({"model_max_budget": None}, {"model_max_budget": "{}"}, id="null"), + pytest.param({"model_max_budget": {}}, {"model_max_budget": "{}"}, id="empty-object"), + ], +) +def test_update_budget_preserves_omitted_model_budgets_and_clears_explicit_empty_values( + client_and_mocks: tuple[TestClient, MagicMock, MagicMock], + model_patch: dict[str, object], + expected_model_patch: dict[str, object], +) -> None: + client, _, mock_table = client_and_mocks + captured: Final = _capture_update_data(mock_table) + + response: Final = client.post( + "/budget/update", + json={"budget_id": "budget_model_clear", "soft_budget": None, **model_patch}, + ) + + assert response.status_code == 200, response.text + assert captured == { + "budget_id": "budget_model_clear", + "soft_budget": None, + "updated_by": "test_user", + **expected_model_patch, + } diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index a439b12b6e9..98c7a9f6d24 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -2123,6 +2123,7 @@ export interface paths { * - rpm_limit: Optional[int] - The requests per minute limit for the budget. * - tpd_limit: Optional[int] - The tokens per day limit for the budget. Charged by batch submissions instead of tpm_limit/rpm_limit. * - 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}} + * Send null or {} to clear all per-model limits. Omit this field to keep them unchanged. * - budget_reset_at: Optional[datetime] - Update the Datetime when the budget was last reset. */ post: operations["update_budget_budget_update_post"];