scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-21 00:21:49 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 74

refactor(proxy): freeze the member spend arrays and budget link to stay within the type discipline budget

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
ryan 2026-09-17 00:49:48 +00:00 • committed by ryan-crabbe-berri
parent 2d61fa66b1
commit b6f4ad190e
3 changed files with 8 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
litellm/proxy/db/db_spend_update_writer.py
View file

@ -188,17 +188,16 @@ SET spend = "LiteLLM_TeamMembership".spend + EXCLUDED.spend,
async def _write_team_member_spend(transaction: _SpendTransaction, spend_by_member_key: Mapping[str, float]) -> None:
# key is "team_id::<value>::user_id::<value>"; rows are sorted by (team_id, user_id) so the teams are
# locked in the same `sorted(team_ids)` order the team endpoints use, preventing deadlocks
# key is "team_id::<value>::user_id::<value>"; locks are taken in sorted team_id order like the team endpoints
rows: Final = sorted((key.split("::")[1], key.split("::")[3], cost) for key, cost in spend_by_member_key.items())
team_ids: Final = [team_id for team_id, _user_id, _cost in rows]
team_ids: Final = tuple(team_id for team_id, _user_id, _cost in rows)
for team_id in dict.fromkeys(team_ids):
_ = await transaction.execute_raw(_TEAM_ADVISORY_LOCK_SQL, team_id)
_ = await transaction.execute_raw(
_TEAM_MEMBER_SPEND_SQL,
[user_id for _team_id, user_id, _cost in rows],
tuple(user_id for _team_id, user_id, _cost in rows),
team_ids,
[cost for _team_id, _user_id, cost in rows],
tuple(cost for _team_id, _user_id, cost in rows),
)

4
litellm/proxy/management_helpers/utils.py
View file

@ -476,7 +476,9 @@ async def add_new_member(
if returned_user is not None and returned_user.user_id is not None:
membership_table: Final[_PrismaTeamMembershipTable] = _team_membership_table(prisma_client, tx)
membership_key: Final[Mapping[str, object]] = {"user_id": returned_user.user_id, "team_id": team_id}
budget_link: Final[Mapping[str, object]] = {"budget_id": _budget_id} if _budget_id is not None else {}
budget_link: Final[Mapping[str, str]] = (
MappingProxyType({"budget_id": _budget_id}) if _budget_id is not None else MappingProxyType({})
)
_returned_team_membership: Final = await membership_table.upsert(
where={"user_id_team_id": membership_key},
data={"create": {**membership_key, **budget_link}, "update": {}},

2
tests/test_litellm/proxy/db/test_db_spend_update_writer.py
View file

@ -979,7 +979,7 @@ async def test_commit_spend_updates_to_db_writes_team_member_spend_in_one_roster
assert "pg_advisory_xact_lock(hashtext($1))" in lock_statement
statement, user_ids, team_ids, costs = spend_call.args
assert statement is _TEAM_MEMBER_SPEND_SQL
assert (user_ids, team_ids, costs) == ([user_id], [team_id], [response_cost])
assert (list(user_ids), list(team_ids), list(costs)) == ([user_id], [team_id], [response_cost])
assert 'INSERT INTO "LiteLLM_TeamMembership"' in statement
assert "members_with_roles @> jsonb_build_array(jsonb_build_object('user_id', p.user_id))" in statement
assert "ON CONFLICT (user_id, team_id) DO UPDATE" in statement