[Fix] Litellm tags usage add request_id (#16111)

* Add request_id into tag spend

* Linting

litellm-proxy-extras/litellm_proxy_extras/schema.prisma

@@ -451,6 +451,7 @@ model LiteLLM_DailyTeamSpend {
 // Track daily team spend metrics per model and key
 model LiteLLM_DailyTagSpend {
   id                  String   @id @default(uuid())
+  request_id          String?
   tag                 String?   
   date                String
   api_key             String   

litellm/proxy/_types.py

@@ -3541,6 +3541,7 @@ class DailyUserSpendTransaction(BaseDailySpendTransaction):
 
 
 class DailyTagSpendTransaction(BaseDailySpendTransaction):
+    request_id: Optional[str]
     tag: str
 
 

litellm/proxy/db/db_spend_update_writer.py

@@ -1070,6 +1070,9 @@ class DBSpendUpdateWriter:
                                     "cache_creation_input_tokens"
                                 ] = transaction.get("cache_creation_input_tokens", 0)
 
+                            if entity_type == "tag" and "request_id" in transaction:
+                                common_data["request_id"] = transaction.get("request_id")
+                            
                             # Create update data structure
                             update_data = {
                                 "prompt_tokens": {
@@ -1385,7 +1388,7 @@ class DBSpendUpdateWriter:
         for tag in request_tags:
             daily_transaction_key = f"{tag}_{base_daily_transaction['date']}_{payload['api_key']}_{payload['model']}_{payload['custom_llm_provider']}"
             daily_transaction = DailyTagSpendTransaction(
-                tag=tag, **base_daily_transaction
+                tag=tag, **base_daily_transaction, request_id=payload["request_id"]
             )
 
             await self.daily_tag_spend_update_queue.add_update(

litellm/proxy/schema.prisma

@@ -451,6 +451,7 @@ model LiteLLM_DailyTeamSpend {
 // Track daily team spend metrics per model and key
 model LiteLLM_DailyTagSpend {
   id                  String   @id @default(uuid())
+  request_id          String?
   tag                 String?   
   date                String
   api_key             String   

schema.prisma

@@ -451,6 +451,7 @@ model LiteLLM_DailyTeamSpend {
 // Track daily team spend metrics per model and key
 model LiteLLM_DailyTagSpend {
   id                  String   @id @default(uuid())
+  request_id          String?
   tag                 String?   
   date                String
   api_key             String   

tests/test_litellm/proxy/db/test_db_spend_update_writer.py

@@ -343,3 +343,48 @@ async def test_update_tag_db_without_prisma_client():
     )
 
     assert writer.spend_update_queue.add_update.call_count == 0
+
+@pytest.mark.asyncio
+async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_id():
+    """
+    Test that add_spend_log_transaction_to_daily_tag_transaction correctly processes request_id.
+    This tests that request_id is included in the DailyTagSpendTransaction for the LiteLLM_DailyTagSpend table.
+    """
+    writer = DBSpendUpdateWriter()
+    mock_prisma = MagicMock()
+    mock_prisma.get_request_status = MagicMock(return_value="success")
+    
+    request_id = "test-request-id-123"
+    payload = {
+        "request_id": request_id,
+        "request_tags": '["prod-tag", "test-tag"]',
+        "user": "test-user",
+        "startTime": "2024-01-01T00:00:00",
+        "api_key": "test-key",
+        "model": "gpt-4",
+        "custom_llm_provider": "openai",
+        "model_group": "gpt-4-group",
+        "prompt_tokens": 100,
+        "completion_tokens": 50,
+        "spend": 0.05,
+        "metadata": '{"usage_object": {}}',
+    }
+
+    # Mock the add_update method to capture what's being added
+    original_add_update = writer.daily_tag_spend_update_queue.add_update
+    writer.daily_tag_spend_update_queue.add_update = AsyncMock()
+
+    await writer.add_spend_log_transaction_to_daily_tag_transaction(
+        payload=payload,
+        prisma_client=mock_prisma,
+    )
+
+    # Should be called twice (once for each tag)
+    assert writer.daily_tag_spend_update_queue.add_update.call_count == 2
+    
+    # Check that request_id is included in both transactions
+    for call in writer.daily_tag_spend_update_queue.add_update.call_args_list:
+        transaction_dict = call[1]["update"]
+        # Each transaction should have one key with the format tag_date_api_key_model_provider
+        for key, transaction in transaction_dict.items():
+            assert transaction["request_id"] == request_id, f"request_id should be {request_id} but got {transaction.get('request_id')}"