From 38d18084e09f54ba0f94d34e376e858a3d0e24af Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 15 Sep 2026 16:41:21 +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 | 41 ++++++- .../redis_update_buffer.py | 9 ++ .../spend_update_queue.py | 4 + .../test_redis_update_buffer.py | 45 ++++++++ .../proxy/db/test_db_spend_update_writer.py | 104 ++++++++++++++++++ 6 files changed, 203 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d40f518d4b8..b3a18e6b051 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" @@ -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] diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index eaa03c5d7f7..2cbef168666 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -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( 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..71b304e6707 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,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"), ],