From fb00567e4cd275d5cbe8333cb58ac8123c2b3787 Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 15 Sep 2026 18:38:07 +0000 Subject: [PATCH] fix(proxy): track per-member organization spend so the Organizations UI shows member spend Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 2 + litellm/proxy/db/db_spend_update_writer.py | 45 ++++++- .../redis_update_buffer.py | 9 ++ .../spend_update_queue.py | 4 + .../test_redis_update_buffer.py | 45 +++++++ .../proxy/db/test_db_spend_update_writer.py | 116 ++++++++++++++++++ 6 files changed, 219 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 9f85c3d1894..1d3992ab800 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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" @@ -5236,6 +5237,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] diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index eaa03c5d7f7..d207ba2f2c6 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -16,6 +16,7 @@ from collections.abc import Mapping, Sequence from datetime import datetime, timedelta, timezone from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, overload +from urllib.parse import quote, unquote import litellm from litellm._logging import verbose_proxy_logger @@ -85,6 +86,10 @@ else: RESPONSES_SESSION_CALL_TYPES: Final = frozenset({CallTypes.responses.value, CallTypes.aresponses.value}) +def _org_member_transaction_key(org_id: str, user_id: str) -> str: + return f"organization_id::{quote(org_id, safe='')}::user_id::{quote(user_id, safe='')}" + + def _is_batch_cost_row(payload: SpendLogsPayload) -> bool: return payload.get("call_type") == CallTypes.aretrieve_batch.value and payload.get("status") == "success" @@ -110,6 +115,7 @@ class _SpendBatch(Protocol): litellm_teamtable: BatchTable litellm_teammembership: BatchTable litellm_organizationtable: BatchTable + litellm_organizationmembership: BatchTable litellm_tagtable: BatchTable litellm_agentstable: BatchTable litellm_modelaccessgroupbudgettable: BatchTable @@ -666,6 +672,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: @@ -900,6 +907,7 @@ class DBSpendUpdateWriter: self, response_cost: float | None, org_id: str | None, + user_id: str | None, prisma_client: PrismaClient | None, ): try: @@ -916,6 +924,15 @@ class DBSpendUpdateWriter: response_cost=response_cost, ) ) + + if user_id is not None: + await self.spend_update_queue.add_update( + update=SpendUpdateQueueItem( + entity_type=Litellm_EntityType.ORGANIZATION_MEMBER, + entity_id=_org_member_transaction_key(org_id, user_id), + 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 +1180,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 +1726,29 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, ) + 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()): + _, quoted_org_id, _, quoted_user_id = key.split("::") + batcher.litellm_organizationmembership.update_many( + where={"organization_id": unquote(quoted_org_id), "user_id": unquote(quoted_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( diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index c06f2e04aca..6f49a00b763 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -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( diff --git a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py index 8c0076b10c1..bc068d10daf 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py @@ -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": diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py index 8f3508fc4e9..817c86a1bdf 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py @@ -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""" diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 5e977712a1e..64001146e05 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -944,6 +944,121 @@ 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_org_spend_keeps_member_attribution_when_ids_contain_the_key_delimiter(): + db_writer = DBSpendUpdateWriter() + await db_writer._update_org_db( + response_cost=0.75, + org_id="division::west", + user_id="user::42", + 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_organizationmembership.update_many.assert_called_once_with( + where={"organization_id": "division::west", "user_id": "user::42"}, + data={"spend": {"increment": 0.75}}, + ) + + +@pytest.mark.asyncio +async def test_batch_database_updates_queues_org_member_spend_for_the_request_user(): + db_writer = DBSpendUpdateWriter() + 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=MagicMock(), + litellm_proxy_budget_name=None, + payload={"request_id": "req-1", "model": "gpt-4o-mini", "spend": 0.1}, + ) + transactions = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() + + assert transactions["org_list_transactions"] == {"org1": 0.1} + assert transactions["org_member_list_transactions"] == {"organization_id::org1::user_id::u1": 0.1} + + @pytest.mark.asyncio async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_id(): """ @@ -2904,6 +3019,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"), ],