scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

fix(proxy): update existing budget row instead of creating new one in /team/member_update

Browse source 
When updating a team member's budget via /team/member_update, the
existing budget row was being orphaned because _upsert_budget_and_membership
always created a new row even when existing_budget_id was provided.

Now the function updates the existing budget row in-place when
existing_budget_id is present, and only creates a new row when no
prior budget exists.

Closes #25508
This commit is contained in:
Kiyeon Jeon 2026-04-13 04:32:17 +09:00
parent 9e6d2d2069
commit 11662ef18e
2 changed files with 99 additions and 129 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

87
litellm/proxy/management_endpoints/common_utils.py
View file

@ -378,45 +378,60 @@ async def _upsert_budget_and_membership(
)
return
# create a new budget
create_data: Dict[str, Any] = {
"created_by": user_api_key_dict.user_id or "",
"updated_by": user_api_key_dict.user_id or "",
}
if max_budget is not None:
create_data["max_budget"] = max_budget
if tpm_limit is not None:
create_data["tpm_limit"] = tpm_limit
if rpm_limit is not None:
create_data["rpm_limit"] = rpm_limit
if existing_budget_id is not None:
# update the existing budget row in-place to avoid orphaning it
update_data: Dict[str, Any] = {
"updated_by": user_api_key_dict.user_id or "",
}
if max_budget is not None:
update_data["max_budget"] = max_budget
if tpm_limit is not None:
update_data["tpm_limit"] = tpm_limit
if rpm_limit is not None:
update_data["rpm_limit"] = rpm_limit
await tx.litellm_budgettable.update(
where={"budget_id": existing_budget_id},
data=update_data,
)
else:
# create a new budget and link it to the team membership
create_data: Dict[str, Any] = {
"created_by": user_api_key_dict.user_id or "",
"updated_by": user_api_key_dict.user_id or "",
}
if max_budget is not None:
create_data["max_budget"] = max_budget
if tpm_limit is not None:
create_data["tpm_limit"] = tpm_limit
if rpm_limit is not None:
create_data["rpm_limit"] = rpm_limit
new_budget = await tx.litellm_budgettable.create(
data=create_data,
include={"team_membership": True},
)
# upsert the team membership with the new/updated budget
await tx.litellm_teammembership.upsert(
where={
"user_id_team_id": {
"user_id": user_id,
"team_id": team_id,
}
},
data={
"create": {
"user_id": user_id,
"team_id": team_id,
"litellm_budget_table": {
"connect": {"budget_id": new_budget.budget_id},
new_budget = await tx.litellm_budgettable.create(
data=create_data,
include={"team_membership": True},
)
await tx.litellm_teammembership.upsert(
where={
"user_id_team_id": {
"user_id": user_id,
"team_id": team_id,
}
},
data={
"create": {
"user_id": user_id,
"team_id": team_id,
"litellm_budget_table": {
"connect": {"budget_id": new_budget.budget_id},
},
},
"update": {
"litellm_budget_table": {
"connect": {"budget_id": new_budget.budget_id},
},
},
},
"update": {
"litellm_budget_table": {
"connect": {"budget_id": new_budget.budget_id},
},
},
},
)
)
def _update_metadata_field(updated_kv: dict, field_name: str) -> None:

141
tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py
View file

@ -12,18 +12,19 @@ from litellm.proxy.management_endpoints.common_utils import (
# Fixtures: a fake Prisma transaction and a fake UserAPIKeyAuth object
# ---------------------------------------------------------------------------
@pytest.fixture
def mock_tx():
"""
Builds an object that looks just enough like the Prisma tx you use
inside _upsert_budget_and_membership.
"""
# membership “table”
# membership "table"
membership = MagicMock()
membership.update = AsyncMock()
membership.upsert = AsyncMock()
# budget “table”
# budget "table"
budget = MagicMock()
budget.update = AsyncMock()
# budget.create returns a fake row that has .budget_id
@ -42,6 +43,7 @@ def fake_user():
"""Cheap stand-in for UserAPIKeyAuth."""
return types.SimpleNamespace(user_id="tester@example.com")
# TEST: max_budget is None, disconnect only
@pytest.mark.asyncio
async def test_upsert_disconnect(mock_tx, fake_user):
@ -63,54 +65,38 @@ async def test_upsert_disconnect(mock_tx, fake_user):
mock_tx.litellm_teammembership.upsert.assert_not_called()
# TEST: existing budget id, creates new budget (current behavior)
# TEST: existing budget id -> update existing budget in-place
@pytest.mark.asyncio
async def test_upsert_with_existing_budget_id_creates_new(mock_tx, fake_user):
async def test_upsert_with_existing_budget_id_updates_in_place(mock_tx, fake_user):
"""
Test that even when existing_budget_id is provided, the function creates a new budget.
This reflects the current implementation behavior.
When existing_budget_id is provided, the function should update the
existing budget row instead of creating a new one.
"""
await _upsert_budget_and_membership(
mock_tx,
team_id="team-2",
user_id="user-2",
max_budget=42.0,
existing_budget_id="bud-999", # This parameter is currently unused
existing_budget_id="bud-999",
user_api_key_dict=fake_user,
)
# Should create a new budget, not update existing
mock_tx.litellm_budgettable.create.assert_awaited_once_with(
# Should update the existing budget row
mock_tx.litellm_budgettable.update.assert_awaited_once_with(
where={"budget_id": "bud-999"},
data={
"max_budget": 42.0,
"created_by": fake_user.user_id,
"updated_by": fake_user.user_id,
},
include={"team_membership": True},
)
# Should upsert team membership with the new budget ID
new_budget_id = mock_tx.litellm_budgettable.create.return_value.budget_id
mock_tx.litellm_teammembership.upsert.assert_awaited_once_with(
where={"user_id_team_id": {"user_id": "user-2", "team_id": "team-2"}},
data={
"create": {
"user_id": "user-2",
"team_id": "team-2",
"litellm_budget_table": {"connect": {"budget_id": new_budget_id}},
},
"update": {
"litellm_budget_table": {"connect": {"budget_id": new_budget_id}},
},
},
)
# Should NOT update existing budget
mock_tx.litellm_budgettable.update.assert_not_called()
# Should NOT create a new budget or upsert membership
mock_tx.litellm_budgettable.create.assert_not_called()
mock_tx.litellm_teammembership.upsert.assert_not_called()
mock_tx.litellm_teammembership.update.assert_not_called()
# TEST: create new budget and link membership
# TEST: create new budget and link membership (no existing budget)
@pytest.mark.asyncio
async def test_upsert_create_and_link(mock_tx, fake_user):
await _upsert_budget_and_membership(
@ -152,12 +138,12 @@ async def test_upsert_create_and_link(mock_tx, fake_user):
mock_tx.litellm_budgettable.update.assert_not_called()
# TEST: create new budget and link membership, then create another new budget
# TEST: create budget then update it via existing_budget_id
@pytest.mark.asyncio
async def test_upsert_create_then_create_another(mock_tx, fake_user):
async def test_upsert_create_then_update_existing(mock_tx, fake_user):
"""
Test that multiple calls to _upsert_budget_and_membership create separate budgets,
reflecting the current implementation behavior.
First call creates a new budget, second call updates it in-place
using the existing_budget_id.
"""
# FIRST CALL – create new budget and link membership
await _upsert_budget_and_membership(
@ -176,62 +162,42 @@ async def test_upsert_create_then_create_another(mock_tx, fake_user):
mock_tx.litellm_budgettable.create.assert_awaited_once()
mock_tx.litellm_teammembership.upsert.assert_awaited_once()
# SECOND CALL – reset call history and create another budget
# SECOND CALL – reset call history and update the existing budget
mock_tx.litellm_budgettable.create.reset_mock()
mock_tx.litellm_teammembership.upsert.reset_mock()
mock_tx.litellm_budgettable.update.reset_mock()
# Set up a new budget ID for the second create call
mock_tx.litellm_budgettable.create.return_value = types.SimpleNamespace(budget_id="new-budget-456")
await _upsert_budget_and_membership(
mock_tx,
team_id="team-42",
user_id="user-42",
max_budget=25.0, # new limit
existing_budget_id=created_bid, # this is ignored in current implementation
max_budget=25.0,
existing_budget_id=created_bid,
user_api_key_dict=fake_user,
)
# Should create another new budget (not update existing)
mock_tx.litellm_budgettable.create.assert_awaited_once_with(
# Should update the existing budget, not create a new one
mock_tx.litellm_budgettable.update.assert_awaited_once_with(
where={"budget_id": created_bid},
data={
"max_budget": 25.0,
"created_by": fake_user.user_id,
"updated_by": fake_user.user_id,
},
include={"team_membership": True},
)
# Should upsert team membership with the new budget ID
new_budget_id = mock_tx.litellm_budgettable.create.return_value.budget_id
mock_tx.litellm_teammembership.upsert.assert_awaited_once_with(
where={"user_id_team_id": {"user_id": "user-42", "team_id": "team-42"}},
data={
"create": {
"user_id": "user-42",
"team_id": "team-42",
"litellm_budget_table": {"connect": {"budget_id": new_budget_id}},
},
"update": {
"litellm_budget_table": {"connect": {"budget_id": new_budget_id}},
},
},
)
# Should NOT call update
mock_tx.litellm_budgettable.update.assert_not_called()
mock_tx.litellm_budgettable.create.assert_not_called()
mock_tx.litellm_teammembership.upsert.assert_not_called()
# TEST: update rpm_limit for member with existing budget_id
@pytest.mark.asyncio
async def test_upsert_rpm_limit_update_creates_new_budget(mock_tx, fake_user):
async def test_upsert_rpm_limit_updates_existing_budget(mock_tx, fake_user):
"""
Test that updating rpm_limit for a member with an existing budget_id
creates a new budget with the new rpm/tpm limits and assigns it to the user.
Updating rpm_limit for a member with an existing budget_id should
update the existing budget row in-place with all specified limits.
"""
existing_budget_id = "existing-budget-456"
await _upsert_budget_and_membership(
mock_tx,
team_id="team-rpm-test",
@ -240,51 +206,35 @@ async def test_upsert_rpm_limit_update_creates_new_budget(mock_tx, fake_user):
existing_budget_id=existing_budget_id,
user_api_key_dict=fake_user,
tpm_limit=1000,
rpm_limit=100, # updating rpm_limit
rpm_limit=100,
)
# Should create a new budget with all the specified limits
mock_tx.litellm_budgettable.create.assert_awaited_once_with(
# Should update the existing budget with all specified limits
mock_tx.litellm_budgettable.update.assert_awaited_once_with(
where={"budget_id": existing_budget_id},
data={
"max_budget": 50.0,
"tpm_limit": 1000,
"rpm_limit": 100,
"created_by": fake_user.user_id,
"updated_by": fake_user.user_id,
},
include={"team_membership": True},
)
# Should NOT update the existing budget
mock_tx.litellm_budgettable.update.assert_not_called()
# Should upsert team membership with the new budget ID
new_budget_id = mock_tx.litellm_budgettable.create.return_value.budget_id
mock_tx.litellm_teammembership.upsert.assert_awaited_once_with(
where={"user_id_team_id": {"user_id": "user-rpm-test", "team_id": "team-rpm-test"}},
data={
"create": {
"user_id": "user-rpm-test",
"team_id": "team-rpm-test",
"litellm_budget_table": {"connect": {"budget_id": new_budget_id}},
},
"update": {
"litellm_budget_table": {"connect": {"budget_id": new_budget_id}},
},
},
)
# Should NOT create a new budget or touch membership
mock_tx.litellm_budgettable.create.assert_not_called()
mock_tx.litellm_teammembership.upsert.assert_not_called()
# TEST: create new budget with only rpm_limit (no max_budget)
# TEST: create new budget with only rpm_limit (no max_budget, no existing budget)
@pytest.mark.asyncio
async def test_upsert_rpm_only_creates_new_budget(mock_tx, fake_user):
"""
Test that setting only rpm_limit creates a new budget with just the rpm_limit.
Setting only rpm_limit with no existing budget creates a new budget.
"""
await _upsert_budget_and_membership(
mock_tx,
team_id="team-rpm-only",
user_id="user-rpm-only",
user_id="user-rpm-only",
max_budget=None,
existing_budget_id=None,
user_api_key_dict=fake_user,
@ -304,7 +254,12 @@ async def test_upsert_rpm_only_creates_new_budget(mock_tx, fake_user):
# Should upsert team membership with the new budget ID
new_budget_id = mock_tx.litellm_budgettable.create.return_value.budget_id
mock_tx.litellm_teammembership.upsert.assert_awaited_once_with(
where={"user_id_team_id": {"user_id": "user-rpm-only", "team_id": "team-rpm-only"}},
where={
"user_id_team_id": {
"user_id": "user-rpm-only",
"team_id": "team-rpm-only",
}
},
data={
"create": {
"user_id": "user-rpm-only",