scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 154

fix(proxy): separate daily spend by MCP tool

Browse source
This commit is contained in:
sathvikvittal 2026-08-10 11:49:14 +05:30
parent 4f99f0a505
commit cc3ffbe89c
1 changed files with 45 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

51
tests/test_litellm/proxy/db/test_db_spend_update_writer.py
View file

@ -12,7 +12,7 @@ sys.path.insert(
from collections.abc import Callable
from datetime import datetime, timezone
from unittest.mock import AsyncMock, MagicMock, call, patch
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from redis.exceptions import DataError
@ -375,6 +375,47 @@ async def test_daily_user_transaction_key_includes_model_group():
]
@pytest.mark.asyncio
async def test_daily_user_transaction_key_includes_mcp_tool_name():
writer = DBSpendUpdateWriter()
writer.daily_spend_update_queue = MagicMock()
writer.daily_spend_update_queue.add_update = AsyncMock()
mock_prisma = MagicMock()
base_payload = {
"request_id": "req-mcp-tool-key",
"user": "test-user",
"startTime": "2024-01-01T12:00:00",
"api_key": "test-key",
"model": "openai/gpt-4o-mini",
"model_group": "a",
"custom_llm_provider": "openai",
"prompt_tokens": 10,
"completion_tokens": 5,
"spend": 0.2,
"metadata": '{"usage_object": {}}',
}
await writer.add_spend_log_transaction_to_daily_user_transaction(
payload={**base_payload, "mcp_namespaced_tool_name": "server_a.tool"},
prisma_client=mock_prisma,
)
await writer.add_spend_log_transaction_to_daily_user_transaction(
payload={**base_payload, "mcp_namespaced_tool_name": "server_b.tool"},
prisma_client=mock_prisma,
)
queued_keys = [
next(iter(call_kwargs.kwargs["update"].keys()))
for call_kwargs in writer.daily_spend_update_queue.add_update.call_args_list
]
assert queued_keys == [
'["test-user","2024-01-01","test-key","openai/gpt-4o-mini","a","openai","server_a.tool",""]',
'["test-user","2024-01-01","test-key","openai/gpt-4o-mini","a","openai","server_b.tool",""]',
]
def test_daily_transaction_key_is_not_delimiter_ambiguous():
assert DBSpendUpdateWriter._daily_transaction_key(
"user", "model_a", "group"
@ -677,7 +718,7 @@ async def test_update_tag_db_with_valid_tags():
"""
Test that _update_tag_db correctly processes valid tags and adds them to the spend update queue.
"""
from litellm.proxy._types import Litellm_EntityType, SpendUpdateQueueItem
from litellm.proxy._types import Litellm_EntityType
writer = DBSpendUpdateWriter()
mock_prisma = MagicMock()
@ -1019,8 +1060,6 @@ async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_i
"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(
@ -1032,8 +1071,8 @@ async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_i
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"]
for update_call in writer.daily_tag_spend_update_queue.add_update.call_args_list:
transaction_dict = update_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 (