Merge remote-tracking branch 'origin/litellm_lit8516_tag_team_id' into litellm_lit8516_team_group_by_tag

litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000000_daily_tag_spend_add_team_id/migration.sql

@@ -1,3 +1,5 @@
+SET LOCAL lock_timeout = '10s';
+
 -- Records the team that sent each request on its daily tag spend rows, so one tag used by
 -- several teams lands in one row per team. "" means the request had no team, and every row
 -- written before this migration keeps "" (attribution is not retroactive).

litellm-proxy-extras/litellm_proxy_extras/migrations/20261017000003_daily_tag_spend_drop_tag_only_unique_index/migration.sql

@@ -6,6 +6,9 @@
 --
 -- Upgrade note: an old-version writer still targets this index in its ON CONFLICT clause and
 -- fails every flush once it is gone. Drain and stop old instances before running migrations.
+-- An old instance booting against the migrated database also logs "prisma schema out of sync"
+-- with SQL that drops team_id and both new indexes. Do not run it: it deletes team attribution
+-- and its CREATE UNIQUE INDEX fails once two teams share a key.
 --
 -- Disclaimer:
 -- - DROP INDEX CONCURRENTLY cannot run inside a transaction. This migration must stay a

litellm/proxy/_types.py

@@ -5159,9 +5159,6 @@ class DailyEndUserSpendTransaction(BaseDailySpendTransaction):
 class DailyTagSpendTransaction(BaseDailySpendTransaction):
     request_id: str | None
     tag: str
-    # "" when the request had no team. Not required: entries queued by a pod running the
-    # previous release, or replayed from the Redis buffer across an upgrade, carry no such key,
-    # and the upsert reads a missing value as "".
     team_id: ReadOnly[NotRequired[str]]
 
 

litellm/proxy/db/baseline_accounting.py

@@ -58,8 +58,6 @@ class DailyBaselineAttribution(LiteLLMBaseModel):
     model_group: str | None = None
     endpoint: str | None = None
     mcp_namespaced_tool_name: str | None = None
-    # The request's team, for the team-scoped tag rows. Defaulted so records stored before
-    # the column existed still validate; they land under "" like any pre-upgrade tag spend.
     team_id: str = ""
     targets: tuple[DailyBaselineTarget, ...] = ()
 

litellm/proxy/db/daily_spend_bulk_upsert.py

@@ -32,8 +32,6 @@ class DailySpendTable:
     name: str
     entity_id_column: str
     carries_request_id: bool = False
-    # Rows also keyed by the team that sent the request, so one entity id shared by several
-    # teams is kept as one row per team.
     team_scoped: bool = False
 
 

litellm/proxy/db/db_spend_update_writer.py

@@ -3035,15 +3035,12 @@ class DBSpendUpdateWriter:
             return
 
         request_tags: Final = daily_spend_entity_ids(payload, "tag")
-        # .get, not indexing: the caller swallows exceptions at debug level, so a payload without
-        # team_id must still record its tag spend (as unattributed) rather than drop it silently.
+        # .get, not indexing: the caller swallows exceptions, so a KeyError would silently drop tag spend
         team_id: Final = payload.get("team_id") or ""  # pyright: ignore[reportUnknownMemberType]  # TypedDict.get overloads carry Any defaults
         for tag in request_tags:
             if tag is None:
                 continue
             endpoint_str = base_daily_transaction.get("endpoint") or ""
-            # team_id is part of the key so the queue never folds two teams' spend on one tag
-            # into a single entry before the upsert can keep them apart.
             daily_transaction_key = f"{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/integration/spend/test_team_tag_spend_attribution.py

@@ -47,14 +47,15 @@ def _tag_breakdown(gateway: Gateway, team: str, *, key: str | None = None) -> di
     return totals
 
 
-def test_one_tag_sent_by_two_teams_is_recorded_once_per_team(gateway: Gateway) -> None:
+def test_one_key_moved_between_teams_records_its_tag_once_per_team(gateway: Gateway) -> None:
     with gateway.scenario() as scenario:
         model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
         tag: Final = f"cost-center:{uuid.uuid4().hex}"
         teams: Final = tuple(scenario.team(models=[model]) for _ in range(2))
         calls_per_team: Final = {teams[0]: 1, teams[1]: 2}
+        key: Final = scenario.key(team_id=teams[0], models=[model])
         for team, calls in calls_per_team.items():
-            key = scenario.key(team_id=team, models=[model])
+            gateway.post("/key/update", {"key": key, "team_id": team})
             for _ in range(calls):
                 reply = gateway.post(
                     "/v1/chat/completions",

tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py

@@ -435,7 +435,7 @@ class TestMigrationSQLIdempotency:
         for migration_name, sql in _guarded_migrations(all_migrations):
             for line_num, line in _statements(sql):
                 if re.search(r"DROP\s+INDEX\s+", line, re.IGNORECASE) and not re.search(
-                    r"DROP\s+INDEX\s+IF\s+EXISTS", line, re.IGNORECASE
+                    r"DROP\s+INDEX\s+(?:CONCURRENTLY\s+)?IF\s+EXISTS", line, re.IGNORECASE
                 ):
                     violations.append(f"  {migration_name}:{line_num}: {line.strip()}")
         assert (

tests/unit/proxy/db/test_daily_spend_bulk_upsert.py

@@ -88,8 +88,6 @@ def test_one_statement_carries_every_row_in_the_batch():
 
     assert sql.count("INSERT INTO") == 1
     assert len(re.findall(r"ON CONFLICT", sql)) == 1
-    # 26 bound columns per row (team_id included) plus the inlined updated_at, so the row
-    # count is what separates one multi-row statement from a hundred single-row ones.
     assert len(params) == 100 * 26
     assert "$2600::text" in sql
     assert sql.count("(NOW() AT TIME ZONE 'UTC')") == 100 + 1
@@ -153,8 +151,6 @@ def test_one_tag_from_two_teams_stays_two_rows():
 
 
 def test_an_entry_without_team_id_folds_into_the_unattributed_row():
-    """Entries queued by an older pod, or replayed from Redis across the upgrade, carry no
-    team_id key. They arbitrate to the same row as an explicit "" rather than a second one."""
     legacy = tag_txn(spend=0.25)
     unattributed = tag_txn(team_id="", spend=0.75)
 
@@ -168,8 +164,6 @@ 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):
-    """Only LiteLLM_DailyTagSpend has the extra team_id column; adding it to any other
-    table's conflict target names a column that table does not have."""
     table = DAILY_SPEND_TABLES[entity]
     row = {**tag_txn(), table.entity_id_column: "id-1"}
 

tests/unit/proxy/db/test_db_spend_update_writer.py

@@ -1490,8 +1490,6 @@ 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():
-    """The queue merges entries by key before the upsert ever sees them, so a key without the
-    team would fold both teams' spend on a shared tag into one entry."""
     writer = DBSpendUpdateWriter()
     prisma = MagicMock()
     prisma.get_request_status = MagicMock(return_value="success")