mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-07 07:58:09 +00:00
fix(proxy): stop /tag/update 500 when the body carries the tag's budget_id (#50324)
Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
5d69250cd0
commit
ddd162439f
3 changed files with 50 additions and 1 deletions
|
|
@ -217,7 +217,7 @@ async def handle_budget_for_entity(
|
|||
)
|
||||
|
||||
# Get all budget field names
|
||||
budget_params: Final = LiteLLM_BudgetTable.model_fields.keys()
|
||||
budget_params: Final = LiteLLM_BudgetTable.model_fields.keys() - {"budget_id"}
|
||||
|
||||
# Extract budget fields from data
|
||||
_json_data: Final = _BUDGET_DATA_MAPPING.validate_python(
|
||||
|
|
|
|||
20
tests/integration/management/test_tag_budget_update.py
Normal file
20
tests/integration/management/test_tag_budget_update.py
Normal file
|
|
@ -0,0 +1,20 @@
|
|||
from typing import Final
|
||||
from uuid import uuid4
|
||||
|
||||
from tests.integration._support.client import Gateway, object_value, string_value
|
||||
|
||||
|
||||
def test_tag_budget_can_be_updated_with_its_existing_budget_id(gateway: Gateway) -> None:
|
||||
name: Final = f"integration-tag-{uuid4().hex}"
|
||||
gateway.post("/tag/new", {"name": name, "max_budget": 10})
|
||||
with gateway.scenario() as scenario:
|
||||
scenario.cleanups.callback(gateway.post, "/tag/delete", {"name": name})
|
||||
before: Final = object_value(gateway.post("/tag/info", {"names": [name]})[name])
|
||||
budget_id: Final = string_value(object_value(before["litellm_budget_table"])["budget_id"])
|
||||
scenario.cleanups.callback(gateway.post, "/budget/delete", {"id": budget_id})
|
||||
response: Final = gateway.request(
|
||||
"POST", "/tag/update", {"name": name, "budget_id": budget_id, "max_budget": 20}
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
after: Final = object_value(gateway.post("/tag/info", {"names": [name]})[name])
|
||||
assert object_value(after["litellm_budget_table"])["max_budget"] == 20, after
|
||||
|
|
@ -249,6 +249,35 @@ async def test_add_new_member_links_default_team_budget_id():
|
|||
assert create_data["budget_id"] == test_default_budget_id
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_budget_for_entity_updates_with_existing_budget_id_in_request() -> None:
|
||||
from litellm.proxy.management_helpers.utils import handle_budget_for_entity
|
||||
from litellm.types.tag_management import TagUpdateRequest
|
||||
|
||||
budget_update: Final = AsyncMock()
|
||||
prisma_client: Final = cast(
|
||||
PrismaClient,
|
||||
SimpleNamespace(db=SimpleNamespace(litellm_budgettable=SimpleNamespace(update=budget_update))),
|
||||
)
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", prisma_client),
|
||||
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
|
||||
):
|
||||
budget_id: Final = await handle_budget_for_entity(
|
||||
data=TagUpdateRequest(name="budget-tag", budget_id="budget-1", max_budget=20),
|
||||
existing_budget_id="budget-1",
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id="admin"),
|
||||
prisma_client=prisma_client,
|
||||
litellm_proxy_admin_name="admin",
|
||||
)
|
||||
|
||||
assert budget_id == "budget-1"
|
||||
budget_update.assert_awaited_once_with(
|
||||
where={"budget_id": "budget-1"},
|
||||
data={"budget_id": "budget-1", "max_budget": 20.0, "updated_by": "admin"},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("request_fields", "cleared_budget_fields", "should_update", "expected_max_budget"),
|
||||
cast(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue