From 67478a90746fd30ca577a0b0ae439743c8de264c Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 11 Nov 2025 18:53:48 -0800 Subject: [PATCH] [Fix] Litellm tags usage add request_id (#16111) * Add request_id into tag spend * Linting --- .../litellm_proxy_extras/schema.prisma | 1 + litellm/proxy/_types.py | 1 + litellm/proxy/db/db_spend_update_writer.py | 5 ++- litellm/proxy/schema.prisma | 1 + schema.prisma | 1 + .../proxy/db/test_db_spend_update_writer.py | 45 +++++++++++++++++++ 6 files changed, 53 insertions(+), 1 deletion(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 1ab193bba7..8890456112 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/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 diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index a5d9078921..a212eab076 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3541,6 +3541,7 @@ class DailyUserSpendTransaction(BaseDailySpendTransaction): class DailyTagSpendTransaction(BaseDailySpendTransaction): + request_id: Optional[str] tag: str diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 819c7daec1..06b5301424 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/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( diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 025a1a0e3c..51e6ea9454 100644 --- a/litellm/proxy/schema.prisma +++ b/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 diff --git a/schema.prisma b/schema.prisma index 025a1a0e3c..51e6ea9454 100644 --- a/schema.prisma +++ b/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 diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 09ac3c70c3..6dbbbdd744 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/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')}" \ No newline at end of file