scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-17 23:51:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 67

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 18:38:07 +00:00
parent 41b5d47c71
commit fb00567e4c
6 changed files with 219 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"
@ -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]

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

@ -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(

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"""

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

@ -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"),
],