From 3982bcc9d7b00f0ddab470571d96a8c6490001d5 Mon Sep 17 00:00:00 2001 From: Mrinal Chanshetty Date: Sat, 24 Oct 2026 20:18:20 -0500 Subject: [PATCH] fix(proxy): keep tag and team apart in the daily tag spend queue key The tag queue key joined its fields with "_", and the new team_id sat next to the free-text tag, so tag "cost_center" from team "west" and tag "cost" from team "center_west" got the same key. The queue then summed both into one entry under the first team, before the upsert could keep them apart. The key is now a JSON list of the same fields. Also types the new tests (Final, return types, the entity literal). --- litellm/proxy/db/db_spend_update_writer.py | 12 +++- .../proxy/db/test_daily_spend_bulk_upsert.py | 22 ++++--- .../proxy/db/test_db_spend_update_writer.py | 64 +++++++++++++------ 3 files changed, 67 insertions(+), 31 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 2e11fdd9c9d..22540866a77 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -3080,7 +3080,17 @@ class DBSpendUpdateWriter: if tag is None: continue endpoint_str = base_daily_transaction.get("endpoint") or "" - daily_transaction_key = f"{tag}_{team_id}_{base_daily_transaction['date']}_{payload['api_key']}_{payload['model']}_{payload['custom_llm_provider']}_{endpoint_str}" + daily_transaction_key = json.dumps( + [ + tag, + team_id, + base_daily_transaction["date"], + payload["api_key"], + payload["model"], + payload["custom_llm_provider"], + endpoint_str, + ] + ) daily_transaction = DailyTagSpendTransaction( tag=tag, team_id=team_id, **base_daily_transaction, request_id=payload["request_id"] ) diff --git a/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py b/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py index dd4267906ed..61a7d7dac7f 100644 --- a/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py +++ b/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py @@ -3,11 +3,13 @@ import re from collections.abc import AsyncIterator from contextlib import AbstractAsyncContextManager, asynccontextmanager +from typing import Final import pytest from litellm.proxy.db.daily_spend_bulk_upsert import ( DAILY_SPEND_TABLES, + DailySpendEntity, build_bulk_upsert, conflict_key, merge_by_conflict_key, @@ -142,19 +144,19 @@ def test_non_tag_tables_carry_no_request_id_column(): assert '"user_id"' in sql -def test_one_tag_from_two_teams_stays_two_rows(): - merged = merge_by_conflict_key( +def test_one_tag_from_two_teams_stays_two_rows() -> None: + merged: Final = merge_by_conflict_key( TAG_TABLE, (tag_txn(team_id="team-b", spend=2.0), tag_txn(team_id="team-a", spend=1.0)) ) assert [(txn["team_id"], txn["spend"]) for _, txn in merged] == [("team-a", 1.0), ("team-b", 2.0)], merged -def test_an_entry_without_team_id_folds_into_the_unattributed_row(): - legacy = tag_txn(spend=0.25) - unattributed = tag_txn(team_id="", spend=0.75) +def test_an_entry_without_team_id_folds_into_the_unattributed_row() -> None: + legacy: Final = tag_txn(spend=0.25) + unattributed: Final = tag_txn(team_id="", spend=0.75) - merged = merge_by_conflict_key(TAG_TABLE, (legacy, unattributed)) + merged: Final = merge_by_conflict_key(TAG_TABLE, (legacy, unattributed)) assert len(merged) == 1, merged key, folded = merged[0] @@ -163,13 +165,13 @@ def test_an_entry_without_team_id_folds_into_the_unattributed_row(): @pytest.mark.parametrize("entity", ["user", "team", "org", "end_user", "agent"]) -def test_only_the_tag_table_is_keyed_by_team(entity): - table = DAILY_SPEND_TABLES[entity] - row = {**tag_txn(), table.entity_id_column: "id-1"} +def test_only_the_tag_table_is_keyed_by_team(entity: DailySpendEntity) -> None: + table: Final = DAILY_SPEND_TABLES[entity] + row: Final = {**tag_txn(), table.entity_id_column: "id-1"} sql, _ = build_bulk_upsert(table, merge_by_conflict_key(table, (row,))) - conflict_target = re.search(r"ON CONFLICT \(([^)]*)\)", sql) + conflict_target: Final = re.search(r"ON CONFLICT \(([^)]*)\)", sql) assert conflict_target is not None assert conflict_target.group(1) == ( f'"{table.entity_id_column}", "date", "api_key", "model", "custom_llm_provider", ' diff --git a/tests/unit/proxy/db/test_db_spend_update_writer.py b/tests/unit/proxy/db/test_db_spend_update_writer.py index ec3620a32e4..ac235902925 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -20,7 +20,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm._service_logger import ServiceTypes from litellm.constants import SPEND_ROLLUP_LOCK_TIMEOUT_MS -from litellm.proxy._types import DailyTagSpendTransaction, Litellm_EntityType, SpendUpdateQueueItem +from litellm.proxy._types import DailyTagSpendTransaction, Litellm_EntityType, SpendLogsPayload, SpendUpdateQueueItem from litellm.proxy.db.db_spend_update_writer import ( _TEAM_ADVISORY_LOCK_SQL, _TEAM_MEMBER_SPEND_SQL, @@ -1490,16 +1490,12 @@ async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_i ) -@pytest.mark.asyncio -async def test_one_tag_from_two_teams_is_queued_as_two_team_attributed_entries(): - writer = DBSpendUpdateWriter() - prisma = MagicMock() - prisma.get_request_status = MagicMock(return_value="success") - - def tagged(team_id: str | None, spend: float) -> dict: - return { - "request_id": f"req-{team_id}", - "request_tags": '["cost-center:cc-42"]', +def _tagged_spend_log(tag: str, team_id: str | None, spend: float) -> SpendLogsPayload: + return cast( + SpendLogsPayload, + { + "request_id": f"req-{tag}-{team_id}", + "request_tags": json.dumps([tag]), "team_id": team_id, "user": "user-1", "startTime": "2026-10-17T00:00:00", @@ -1511,20 +1507,48 @@ async def test_one_tag_from_two_teams_is_queued_as_two_team_attributed_entries() "completion_tokens": 5, "spend": spend, "metadata": '{"usage_object": {}}', - } + }, + ) - for team_id, spend in (("team-a", 1.0), ("team-b", 2.0), (None, 0.5)): + +@pytest.mark.parametrize( + ("requests", "expected"), + [ + ( + ( + ("cost-center:cc-42", "team-a", 1.0), + ("cost-center:cc-42", "team-b", 2.0), + ("cost-center:cc-42", None, 0.5), + ), + [ + ("cost-center:cc-42", "", 0.5), + ("cost-center:cc-42", "team-a", 1.0), + ("cost-center:cc-42", "team-b", 2.0), + ], + ), + ( + (("cost_center", "west", 1.0), ("cost", "center_west", 2.0)), + [("cost", "center_west", 2.0), ("cost_center", "west", 1.0)], + ), + ], + ids=["one-tag-two-teams", "underscore-joined-names-collide"], +) +async def test_each_tag_and_team_pair_is_queued_as_its_own_entry( + requests: tuple[tuple[str, str | None, float], ...], + expected: list[tuple[str, str, float]], +) -> None: + writer: Final = DBSpendUpdateWriter() + prisma: Final = MagicMock() + prisma.get_request_status = MagicMock(return_value="success") + + for tag, team_id, spend in requests: await writer.add_spend_log_transaction_to_daily_tag_transaction( - payload=tagged(team_id, spend), prisma_client=prisma + payload=_tagged_spend_log(tag, team_id, spend), prisma_client=prisma ) - queued = await writer.daily_tag_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + queued: Final = await writer.daily_tag_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() - assert sorted((txn["tag"], txn["team_id"], txn["spend"]) for txn in queued.values()) == [ - ("cost-center:cc-42", "", 0.5), - ("cost-center:cc-42", "team-a", 1.0), - ("cost-center:cc-42", "team-b", 2.0), - ], queued + assert sorted((txn["tag"], txn["team_id"], txn["spend"]) for txn in queued.values()) == expected, queued @pytest.mark.asyncio