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

fix(proxy): track per-member organization spend so the Organizations UI shows member spend

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin 2026-09-15 16:41:21 +00:00
parent c8114ba41f
commit 38d18084e0
6 changed files with 203 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/proxy/_types.py
View file

@ -246,6 +246,7 @@ class Litellm_EntityType(enum.Enum):
TEAM = "team"
TEAM_MEMBER = "team_member"
ORGANIZATION = "organization"
ORGANIZATION_MEMBER = "organization_member"
PROJECT = "project"
TAG = "tag"
AGENT = "agent"
@ -5230,6 +5231,7 @@ class DBSpendUpdateTransactions(TypedDict):
team_list_transactions: dict[str, float] | None
team_member_list_transactions: dict[str, float] | None
org_list_transactions: dict[str, float] | None
org_member_list_transactions: ReadOnly[dict[str, float] | None]
tag_list_transactions: dict[str, float] | None
agent_list_transactions: dict[str, float] | None
model_access_group_list_transactions: ReadOnly[dict[str, float] | None]

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

@ -666,6 +666,7 @@ class DBSpendUpdateWriter:
await self._update_org_db(
response_cost=response_cost,
org_id=org_id,
user_id=user_id,
prisma_client=prisma_client,
)
except Exception:
@ -901,6 +902,7 @@ class DBSpendUpdateWriter:
response_cost: float | None,
org_id: str | None,
prisma_client: PrismaClient | None,
user_id: str | None = None,
):
try:
if org_id is None or prisma_client is None:
@ -916,6 +918,16 @@ class DBSpendUpdateWriter:
response_cost=response_cost,
)
)
if user_id is not None:
org_member_key: Final = f"organization_id::{org_id}::user_id::{user_id}"
await self.spend_update_queue.add_update(
update=SpendUpdateQueueItem(
entity_type=Litellm_EntityType.ORGANIZATION_MEMBER,
entity_id=org_member_key,
response_cost=response_cost,
)
)
except Exception as e:
spend_log_error(
"Spend tracking - failed to enqueue org spend update. org_id=%s, response_cost=%s - %s",
@ -1163,14 +1175,15 @@ class DBSpendUpdateWriter:
if db_spend_update_transactions is not None:
verbose_proxy_logger.info(
"Spend tracking - committing spend updates from Redis to DB: "
"keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, tags=%d, agents=%d, "
"model_access_groups=%d",
"keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, org_members=%d, tags=%d, "
"agents=%d, model_access_groups=%d",
len(db_spend_update_transactions.get("key_list_transactions") or {}),
len(db_spend_update_transactions.get("user_list_transactions") or {}),
len(db_spend_update_transactions.get("team_list_transactions") or {}),
len(db_spend_update_transactions.get("org_list_transactions") or {}),
len(db_spend_update_transactions.get("end_user_list_transactions") or {}),
len(db_spend_update_transactions.get("team_member_list_transactions") or {}),
len(db_spend_update_transactions.get("org_member_list_transactions") or {}),
len(db_spend_update_transactions.get("tag_list_transactions") or {}),
len(db_spend_update_transactions.get("agent_list_transactions") or {}),
len(db_spend_update_transactions.get("model_access_group_list_transactions") or {}),
@ -1708,6 +1721,30 @@ class DBSpendUpdateWriter:
proxy_logging_obj=proxy_logging_obj,
)
### UPDATE ORG Membership TABLE with spend ###
org_member_list_transactions: Final = db_spend_update_transactions.get("org_member_list_transactions")
verbose_proxy_logger.debug("Org Membership Spend transactions: %s", org_member_list_transactions)
if org_member_list_transactions is not None and len(org_member_list_transactions.keys()) > 0:
for i in range(n_retry_times + 1):
start_time = time.time()
try:
async with _spend_update_tx(prisma_client) as transaction, transaction.batch_() as batcher:
for key, response_cost in sorted(org_member_list_transactions.items()):
_, org_id, _, user_id = key.split("::", 3)
batcher.litellm_organizationmembership.update_many(
where={"organization_id": org_id, "user_id": user_id},
data={"spend": {"increment": response_cost}},
)
break
except Exception as e:
await self._handle_spend_update_failure(
e=e,
attempt=i,
n_retry_times=n_retry_times,
start_time=start_time,
proxy_logging_obj=proxy_logging_obj,
)
### UPDATE TAG TABLE ###
tag_list_transactions: Final = db_spend_update_transactions["tag_list_transactions"]
await DBSpendUpdateWriter._update_entity_spend_in_db(

9
litellm/proxy/db/db_transaction_queue/redis_update_buffer.py
View file

@ -69,6 +69,7 @@ _SpendTransactionField: TypeAlias = Literal[
"team_list_transactions",
"team_member_list_transactions",
"org_list_transactions",
"org_member_list_transactions",
"tag_list_transactions",
"agent_list_transactions",
"model_access_group_list_transactions",
@ -81,6 +82,7 @@ _SPEND_TRANSACTION_FIELDS: Final[tuple[_SpendTransactionField, ...]] = (
"team_list_transactions",
"team_member_list_transactions",
"org_list_transactions",
"org_member_list_transactions",
"tag_list_transactions",
"agent_list_transactions",
"model_access_group_list_transactions",
@ -412,6 +414,10 @@ class RedisUpdateBuffer:
Litellm_EntityType.ORGANIZATION,
db_spend_update_transactions.get("org_list_transactions"),
),
(
Litellm_EntityType.ORGANIZATION_MEMBER,
db_spend_update_transactions.get("org_member_list_transactions"),
),
(
Litellm_EntityType.TAG,
db_spend_update_transactions.get("tag_list_transactions"),
@ -876,6 +882,9 @@ class RedisUpdateBuffer:
list_of_transactions, "team_member_list_transactions"
),
org_list_transactions=_merged_entity_transactions(list_of_transactions, "org_list_transactions"),
org_member_list_transactions=_merged_entity_transactions(
list_of_transactions, "org_member_list_transactions"
),
tag_list_transactions=_merged_entity_transactions(list_of_transactions, "tag_list_transactions"),
agent_list_transactions=_merged_entity_transactions(list_of_transactions, "agent_list_transactions"),
model_access_group_list_transactions=_merged_entity_transactions(

4
litellm/proxy/db/db_transaction_queue/spend_update_queue.py
View file

@ -137,6 +137,7 @@ class SpendUpdateQueue(BaseUpdateQueue):
team_list_transactions={},
team_member_list_transactions={},
org_list_transactions={},
org_member_list_transactions={},
tag_list_transactions={},
agent_list_transactions={},
model_access_group_list_transactions={},
@ -150,6 +151,7 @@ class SpendUpdateQueue(BaseUpdateQueue):
Litellm_EntityType.TEAM: "team_list_transactions",
Litellm_EntityType.TEAM_MEMBER: "team_member_list_transactions",
Litellm_EntityType.ORGANIZATION: "org_list_transactions",
Litellm_EntityType.ORGANIZATION_MEMBER: "org_member_list_transactions",
Litellm_EntityType.TAG: "tag_list_transactions",
Litellm_EntityType.AGENT: "agent_list_transactions",
Litellm_EntityType.MODEL_ACCESS_GROUP: "model_access_group_list_transactions",
@ -188,6 +190,8 @@ class SpendUpdateQueue(BaseUpdateQueue):
transactions_dict = db_spend_update_transactions["team_member_list_transactions"]
elif dict_key == "org_list_transactions":
transactions_dict = db_spend_update_transactions["org_list_transactions"]
elif dict_key == "org_member_list_transactions":
transactions_dict = db_spend_update_transactions["org_member_list_transactions"]
elif dict_key == "tag_list_transactions":
transactions_dict = db_spend_update_transactions["tag_list_transactions"]
elif dict_key == "agent_list_transactions":

45
tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py
View file

@ -266,6 +266,51 @@ async def test_get_all_transactions_from_redis_buffer_pipeline(redis_update_buff
assert popped_keys[6] == REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY
@pytest.mark.asyncio
async def test_org_member_spend_is_summed_across_pods_and_restored_on_rpush_failure(
redis_update_buffer, mock_redis_cache
):
from litellm.proxy._types import Litellm_EntityType
from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import (
DailySpendUpdateQueue,
)
from litellm.proxy.db.db_transaction_queue.spend_update_queue import (
SpendUpdateQueue,
)
member_key = "organization_id::org-1::user_id::user-1"
pod_json = json.dumps({"org_member_list_transactions": {member_key: 0.25}})
mock_redis_cache.async_lpop_pipeline = AsyncMock(
return_value=[[pod_json, pod_json], None, None, None, None, None, None]
)
(db_spend, *_rest) = await redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline()
assert db_spend is not None
assert db_spend["org_member_list_transactions"] == {member_key: 0.5}
mock_redis_cache.async_rpush_pipeline = AsyncMock(side_effect=ConnectionError("redis went away"))
spend_queue = SpendUpdateQueue()
await spend_queue.add_update(
{
"entity_type": Litellm_EntityType.ORGANIZATION_MEMBER,
"entity_id": member_key,
"response_cost": 1.5,
}
)
await redis_update_buffer.store_in_memory_spend_updates_in_redis(
spend_update_queue=spend_queue,
daily_spend_update_queue=DailySpendUpdateQueue(),
daily_team_spend_update_queue=DailySpendUpdateQueue(),
daily_org_spend_update_queue=DailySpendUpdateQueue(),
daily_end_user_spend_update_queue=DailySpendUpdateQueue(),
daily_agent_spend_update_queue=DailySpendUpdateQueue(),
)
restored_spend = await spend_queue.flush_and_get_aggregated_db_spend_update_transactions()
assert restored_spend["org_member_list_transactions"] == {member_key: 1.5}
@pytest.mark.asyncio
async def test_get_all_transactions_from_redis_buffer_pipeline_no_redis():
"""When redis_cache is None, should return all Nones"""

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

@ -944,6 +944,109 @@ async def test_commit_spend_updates_to_db_increments_team_member_spend_and_total
}
@pytest.mark.asyncio
async def test_org_spend_increments_organization_membership_row_for_the_calling_user():
"""A request made with a user_id inside an org must increment that user's
LiteLLM_OrganizationMembership.spend, not only the org total, or the
Organizations > Members UI renders '-' for every member."""
db_writer = DBSpendUpdateWriter()
await db_writer._update_org_db(
response_cost=0.75,
org_id="org-abc",
user_id="user-xyz",
prisma_client=MagicMock(),
)
transactions = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions()
mock_batcher = MagicMock()
mock_prisma_client = MagicMock()
mock_prisma_client.db.tx = MagicMock(return_value=_good_tx(mock_batcher))
proxy_logging = MagicMock()
proxy_logging.call_details = {}
await db_writer._commit_spend_updates_to_db(
prisma_client=mock_prisma_client,
n_retry_times=0,
proxy_logging_obj=proxy_logging,
db_spend_update_transactions=transactions,
)
mock_batcher.litellm_organizationtable.update_many.assert_called_once_with(
where={"organization_id": "org-abc"},
data={"spend": {"increment": 0.75}},
)
mock_batcher.litellm_organizationmembership.update_many.assert_called_once_with(
where={"organization_id": "org-abc", "user_id": "user-xyz"},
data={"spend": {"increment": 0.75}},
)
@pytest.mark.asyncio
async def test_org_spend_without_user_id_leaves_organization_membership_untouched():
db_writer = DBSpendUpdateWriter()
await db_writer._update_org_db(
response_cost=0.75,
org_id="org-abc",
user_id=None,
prisma_client=MagicMock(),
)
transactions = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions()
mock_batcher = MagicMock()
mock_prisma_client = MagicMock()
mock_prisma_client.db.tx = MagicMock(return_value=_good_tx(mock_batcher))
proxy_logging = MagicMock()
proxy_logging.call_details = {}
await db_writer._commit_spend_updates_to_db(
prisma_client=mock_prisma_client,
n_retry_times=0,
proxy_logging_obj=proxy_logging,
db_spend_update_transactions=transactions,
)
mock_batcher.litellm_organizationtable.update_many.assert_called_once()
mock_batcher.litellm_organizationmembership.update_many.assert_not_called()
@pytest.mark.asyncio
async def test_batch_database_updates_passes_user_id_to_org_spend():
db_writer = DBSpendUpdateWriter()
db_writer._update_org_db = AsyncMock()
db_writer._update_user_db = AsyncMock()
db_writer._update_key_db = AsyncMock()
db_writer._update_team_db = AsyncMock()
db_writer._update_tag_db = AsyncMock()
db_writer._update_agent_db = AsyncMock()
db_writer._update_model_access_group_db = AsyncMock()
db_writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock()
db_writer.add_spend_log_transaction_to_daily_end_user_transaction = AsyncMock()
db_writer.add_spend_log_transaction_to_daily_agent_transaction = AsyncMock()
db_writer.add_spend_log_transaction_to_daily_team_transaction = AsyncMock()
db_writer.add_spend_log_transaction_to_daily_org_transaction = AsyncMock()
db_writer.add_spend_log_transaction_to_daily_tag_transaction = AsyncMock()
prisma_client = MagicMock()
await db_writer._batch_database_updates(
response_cost=0.1,
user_id="u1",
hashed_token="t1",
team_id=None,
org_id="org1",
end_user_id=None,
prisma_client=prisma_client,
litellm_proxy_budget_name=None,
payload={"key": "value"},
)
db_writer._update_org_db.assert_awaited_once_with(
response_cost=0.1,
org_id="org1",
user_id="u1",
prisma_client=prisma_client,
)
@pytest.mark.asyncio
async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_id():
"""
@ -2904,6 +3007,7 @@ async def test_update_daily_spend_retries_deadlock(monkeypatch):
("team_list_transactions", "team-1"),
("team_member_list_transactions", "team_id::team-1::user_id::user-1"),
("org_list_transactions", "org-1"),
("org_member_list_transactions", "organization_id::org-1::user_id::user-1"),
("tag_list_transactions", "tag-1"),
("agent_list_transactions", "agent-1"),
],