Merge branch 'litellm_lit8516_tag_team_id' into litellm_lit8516_team_group_by_tag

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

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

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