tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-09 19:10:50 +00:00
Code Issues Packages Projects Releases Wiki Activity

perf(proxy): batch 11 create_task() calls into 1 in update_database()

Browse Source 
Replace 11 separate asyncio.create_task() calls per request with a
single batched task that runs all spend-update helpers sequentially.
This reduces task scheduling overhead at high RPS (11,000 -> 1,000
tasks/sec at 1K RPS) and cuts 5 copy.deepcopy(payload) calls to 1
shared copy.

Also fixes a mutation bug where the daily agent spend handler received
the raw payload without deepcopy, unlike all other daily helpers.
This commit is contained in:
Ryan Crabbe
2026-02-24 09:56:16 -08:00
parent 5e9f24f74c
commit 86ec2fbf84
2 changed files with 419 additions and 158 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/db/db_spend_update_writer.py
+195 -89
View File
@@ -124,53 +124,20 @@ class DBSpendUpdateWriter:
payload["startTime"] = payload["startTime"].isoformat()
if isinstance(payload["endTime"], datetime):
payload["endTime"] = payload["endTime"].isoformat()
if org_id is not None and org_id != "":
payload["organization_id"] = org_id
if team_id is not None and team_id != "":
payload["team_id"] = team_id
asyncio.create_task(
self._update_user_db(
response_cost=response_cost,
user_id=user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
litellm_proxy_budget_name=litellm_proxy_budget_name,
end_user_id=end_user_id,
)
)
asyncio.create_task(
self._update_key_db(
response_cost=response_cost,
hashed_token=hashed_token,
prisma_client=prisma_client,
)
)
asyncio.create_task(
self._update_team_db(
response_cost=response_cost,
team_id=team_id,
user_id=user_id,
prisma_client=prisma_client,
)
)
asyncio.create_task(
self._update_org_db(
response_cost=response_cost,
org_id=org_id,
prisma_client=prisma_client,
)
)
asyncio.create_task(
self._update_tag_db(
response_cost=response_cost,
request_tags=copy.deepcopy(payload.get("request_tags")),
prisma_client=prisma_client,
)
)
# One deepcopy shared by all 6 daily spend helpers (was 5, fixes agent bug)
payload_copy = copy.deepcopy(payload)
# Deepcopy request_tags for _update_tag_db
request_tags = copy.deepcopy(payload.get("request_tags"))
# Keep _insert_spend_log_to_db awaited inline (not a task, preserve current behavior)
if disable_spend_logs is False:
await self._insert_spend_log_to_db(
payload=copy.deepcopy(payload),
@@ -181,44 +148,20 @@ class DBSpendUpdateWriter:
"disable_spend_logs=True. Skipping writing spend logs to db. Other spend updates - Key/User/Team table will still occur."
)
# Single task replaces 11 create_task() calls
asyncio.create_task(
self.add_spend_log_transaction_to_daily_user_transaction(
payload=copy.deepcopy(payload),
prisma_client=prisma_client,
)
)
asyncio.create_task(
self.add_spend_log_transaction_to_daily_end_user_transaction(
payload=copy.deepcopy(payload),
prisma_client=prisma_client,
)
)
asyncio.create_task(
self.add_spend_log_transaction_to_daily_agent_transaction(
payload=payload,
prisma_client=prisma_client,
)
)
asyncio.create_task(
self.add_spend_log_transaction_to_daily_team_transaction(
payload=copy.deepcopy(payload),
prisma_client=prisma_client,
)
)
asyncio.create_task(
self.add_spend_log_transaction_to_daily_org_transaction(
payload=copy.deepcopy(payload),
self._batch_database_updates(
response_cost=response_cost,
user_id=user_id,
hashed_token=hashed_token,
team_id=team_id,
org_id=org_id,
end_user_id=end_user_id,
prisma_client=prisma_client,
)
)
asyncio.create_task(
self.add_spend_log_transaction_to_daily_tag_transaction(
payload=copy.deepcopy(payload),
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
litellm_proxy_budget_name=litellm_proxy_budget_name,
payload_copy=payload_copy,
request_tags=request_tags,
)
)
@@ -228,6 +171,157 @@ class DBSpendUpdateWriter:
f"Error updating Prisma database: {traceback.format_exc()}"
)
async def _batch_database_updates(
self,
*,
response_cost: Optional[float],
user_id: Optional[str],
hashed_token: Optional[str],
team_id: Optional[str],
org_id: Optional[str],
end_user_id: Optional[str],
prisma_client: Optional[PrismaClient],
user_api_key_cache: DualCache,
litellm_proxy_budget_name: Optional[str],
payload_copy: dict,
request_tags: Optional[Any],
):
"""
Runs all 11 spend-update helpers sequentially inside a single asyncio task.
Each helper is wrapped in try/except so one failure doesn't prevent the others.
"""
try:
await self._update_user_db(
response_cost=response_cost,
user_id=user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
litellm_proxy_budget_name=litellm_proxy_budget_name,
end_user_id=end_user_id,
)
except Exception:
verbose_proxy_logger.debug(
"_batch_database_updates: _update_user_db failed: %s",
traceback.format_exc(),
)
try:
await self._update_key_db(
response_cost=response_cost,
hashed_token=hashed_token,
prisma_client=prisma_client,
)
except Exception:
verbose_proxy_logger.debug(
"_batch_database_updates: _update_key_db failed: %s",
traceback.format_exc(),
)
try:
await self._update_team_db(
response_cost=response_cost,
team_id=team_id,
user_id=user_id,
prisma_client=prisma_client,
)
except Exception:
verbose_proxy_logger.debug(
"_batch_database_updates: _update_team_db failed: %s",
traceback.format_exc(),
)
try:
await self._update_org_db(
response_cost=response_cost,
org_id=org_id,
prisma_client=prisma_client,
)
except Exception:
verbose_proxy_logger.debug(
"_batch_database_updates: _update_org_db failed: %s",
traceback.format_exc(),
)
try:
await self._update_tag_db(
response_cost=response_cost,
request_tags=request_tags,
prisma_client=prisma_client,
)
except Exception:
verbose_proxy_logger.debug(
"_batch_database_updates: _update_tag_db failed: %s",
traceback.format_exc(),
)
try:
await self.add_spend_log_transaction_to_daily_user_transaction(
payload=payload_copy,
prisma_client=prisma_client,
)
except Exception:
verbose_proxy_logger.debug(
"_batch_database_updates: add_spend_log_transaction_to_daily_user_transaction failed: %s",
traceback.format_exc(),
)
try:
await self.add_spend_log_transaction_to_daily_end_user_transaction(
payload=payload_copy,
prisma_client=prisma_client,
)
except Exception:
verbose_proxy_logger.debug(
"_batch_database_updates: add_spend_log_transaction_to_daily_end_user_transaction failed: %s",
traceback.format_exc(),
)
try:
await self.add_spend_log_transaction_to_daily_agent_transaction(
payload=payload_copy,
prisma_client=prisma_client,
)
except Exception:
verbose_proxy_logger.debug(
"_batch_database_updates: add_spend_log_transaction_to_daily_agent_transaction failed: %s",
traceback.format_exc(),
)
try:
await self.add_spend_log_transaction_to_daily_team_transaction(
payload=payload_copy,
prisma_client=prisma_client,
)
except Exception:
verbose_proxy_logger.debug(
"_batch_database_updates: add_spend_log_transaction_to_daily_team_transaction failed: %s",
traceback.format_exc(),
)
try:
await self.add_spend_log_transaction_to_daily_org_transaction(
payload=payload_copy,
org_id=org_id,
prisma_client=prisma_client,
)
except Exception:
verbose_proxy_logger.debug(
"_batch_database_updates: add_spend_log_transaction_to_daily_org_transaction failed: %s",
traceback.format_exc(),
)
try:
await self.add_spend_log_transaction_to_daily_tag_transaction(
payload=payload_copy,
prisma_client=prisma_client,
)
except Exception:
verbose_proxy_logger.debug(
"_batch_database_updates: add_spend_log_transaction_to_daily_tag_transaction failed: %s",
traceback.format_exc(),
)
async def _update_key_db(
self,
response_cost: Optional[float],
@@ -880,7 +974,7 @@ class DBSpendUpdateWriter:
team_id = key.split("::")[1]
user_id = key.split("::")[3]
team_memberships_to_invalidate.append((user_id, team_id))
for i in range(n_retry_times + 1):
start_time = time.time()
try:
@@ -917,11 +1011,13 @@ class DBSpendUpdateWriter:
_raise_failed_update_spend_exception(
e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj
)
# Invalidate cache for updated team memberships
# This ensures budget checks read fresh spend data from the database
if team_memberships_to_invalidate and proxy_logging_obj is not None:
user_api_key_cache = proxy_logging_obj.call_details.get("user_api_key_cache")
user_api_key_cache = proxy_logging_obj.call_details.get(
"user_api_key_cache"
)
if user_api_key_cache is not None:
for user_id, team_id in team_memberships_to_invalidate:
cache_key = "team_membership:{}:{}".format(user_id, team_id)
@@ -1233,7 +1329,9 @@ class DBSpendUpdateWriter:
),
"endpoint": transaction.get("endpoint") or "",
"prompt_tokens": transaction["prompt_tokens"],
"completion_tokens": transaction["completion_tokens"],
"completion_tokens": transaction[
"completion_tokens"
],
"spend": transaction["spend"],
"api_requests": transaction["api_requests"],
"successful_requests": transaction[
@@ -1244,12 +1342,14 @@ class DBSpendUpdateWriter:
# Add cache-related fields if they exist
if "cache_read_input_tokens" in transaction:
common_data["cache_read_input_tokens"] = (
transaction.get("cache_read_input_tokens", 0)
)
common_data[
"cache_read_input_tokens"
] = transaction.get("cache_read_input_tokens", 0)
if "cache_creation_input_tokens" in transaction:
common_data["cache_creation_input_tokens"] = (
transaction.get("cache_creation_input_tokens", 0)
common_data[
"cache_creation_input_tokens"
] = transaction.get(
"cache_creation_input_tokens", 0
)
if entity_type == "tag" and "request_id" in transaction:
@@ -1292,10 +1392,14 @@ class DBSpendUpdateWriter:
}
if entity_type == "tag" and "request_id" in transaction:
update_data["request_id"] = transaction.get("request_id")
update_data["request_id"] = transaction.get(
"request_id"
)
# Add endpoint to update_data so existing rows get their endpoint field updated
update_data["endpoint"] = transaction.get("endpoint") or ""
update_data["endpoint"] = (
transaction.get("endpoint") or ""
)
table.upsert(
where=where_clause,
@@ -1479,7 +1583,9 @@ class DBSpendUpdateWriter:
self,
payload: Union[dict, SpendLogsPayload],
prisma_client: PrismaClient,
type: Literal["user", "team", "org", "request_tags", "end_user", "agent"] = "user",
type: Literal[
"user", "team", "org", "request_tags", "end_user", "agent"
] = "user",
) -> Optional[BaseDailySpendTransaction]:
common_expected_keys = ["startTime", "api_key"]
if type == "user":
@@ -1538,7 +1644,7 @@ class DBSpendUpdateWriter:
endpoint = None
if call_type:
endpoint = ROUTE_ENDPOINT_MAPPING.get(call_type, None)
daily_transaction = BaseDailySpendTransaction(
date=date,
api_key=payload["api_key"],
@@ -1750,7 +1856,7 @@ class DBSpendUpdateWriter:
endpoint_str = base_daily_transaction.get("endpoint") or ""
daily_transaction_key = f"{payload['agent_id']}_{base_daily_transaction['date']}_{payload_with_agent_id['api_key']}_{payload_with_agent_id['model']}_{payload_with_agent_id['custom_llm_provider']}_{endpoint_str}"
daily_transaction = DailyAgentSpendTransaction(
agent_id=payload['agent_id'], **base_daily_transaction
agent_id=payload["agent_id"], **base_daily_transaction
)
await self.daily_agent_spend_update_queue.add_update(
update={daily_transaction_key: daily_transaction}
tests/test_litellm/proxy/db/test_db_spend_update_writer.py
+224 -69
View File
@@ -1,3 +1,4 @@
import asyncio
import json
import os
import sys
@@ -51,6 +52,9 @@ async def test_daily_spend_tracking_with_disabled_spend_logs():
# Call the method
await db_writer.update_database(**test_data)
# Let the single batched task run
await asyncio.sleep(0)
# Verify that _insert_spend_log_to_db was NOT called (since disable_spend_logs is True)
db_writer._insert_spend_log_to_db.assert_not_called()
@@ -115,7 +119,9 @@ async def test_update_daily_spend_with_null_entity_id():
# Verify the where clause contains null entity_id
call_args = mock_table.upsert.call_args[1]
where_clause = call_args["where"]["user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint"]
where_clause = call_args["where"][
"user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint"
]
assert where_clause["user_id"] is None
assert where_clause["date"] == "2024-01-01"
assert where_clause["api_key"] == "test-api-key"
@@ -161,7 +167,7 @@ async def test_update_daily_spend_sorting():
upsert_calls = []
for i in range(50):
daily_spend_transactions[f"test_key_{i}"] = {
"user_id": f"user{60-i}", # user60 ... user11, reverse order
"user_id": f"user{60-i}", # user60 ... user11, reverse order
"date": "2024-01-01",
"api_key": "test-api-key",
"model": "gpt-4",
@@ -173,46 +179,48 @@ async def test_update_daily_spend_sorting():
"successful_requests": 1,
"failed_requests": 0,
}
upsert_calls.append(call(
where={
"user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint": {
"user_id": f"user{i+11}", # user11 ... user60, sorted order
"date": "2024-01-01",
"api_key": "test-api-key",
"model": "gpt-4",
"custom_llm_provider": "openai",
"mcp_namespaced_tool_name": "",
"endpoint": "",
}
},
data={
"create": {
"user_id": f"user{i+11}",
"date": "2024-01-01",
"api_key": "test-api-key",
"model": "gpt-4",
"model_group": None,
"mcp_namespaced_tool_name": "",
"custom_llm_provider": "openai",
"endpoint": "",
"prompt_tokens": 10,
"completion_tokens": 20,
"spend": 0.1,
"api_requests": 1,
"successful_requests": 1,
"failed_requests": 0,
upsert_calls.append(
call(
where={
"user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint": {
"user_id": f"user{i+11}", # user11 ... user60, sorted order
"date": "2024-01-01",
"api_key": "test-api-key",
"model": "gpt-4",
"custom_llm_provider": "openai",
"mcp_namespaced_tool_name": "",
"endpoint": "",
}
},
"update": {
"prompt_tokens": {"increment": 10},
"completion_tokens": {"increment": 20},
"spend": {"increment": 0.1},
"api_requests": {"increment": 1},
"successful_requests": {"increment": 1},
"failed_requests": {"increment": 0},
"endpoint": "",
data={
"create": {
"user_id": f"user{i+11}",
"date": "2024-01-01",
"api_key": "test-api-key",
"model": "gpt-4",
"model_group": None,
"mcp_namespaced_tool_name": "",
"custom_llm_provider": "openai",
"endpoint": "",
"prompt_tokens": 10,
"completion_tokens": 20,
"spend": 0.1,
"api_requests": 1,
"successful_requests": 1,
"failed_requests": 0,
},
"update": {
"prompt_tokens": {"increment": 10},
"completion_tokens": {"increment": 20},
"spend": {"increment": 0.1},
"api_requests": {"increment": 1},
"successful_requests": {"increment": 1},
"failed_requests": {"increment": 0},
"endpoint": "",
},
},
},
))
)
)
# Call the method
await DBSpendUpdateWriter._update_daily_spend(
@@ -275,7 +283,7 @@ async def test_update_daily_spend_tag_with_request_id():
# Verify that table.upsert was called
mock_table.upsert.assert_called_once()
# Verify request_id is in update_data
call_args = mock_table.upsert.call_args[1]
update_data = call_args["data"]["update"]
@@ -283,15 +291,13 @@ async def test_update_daily_spend_tag_with_request_id():
assert update_data["request_id"] == "test-request-id-123"
@pytest.mark.asyncio
async def test_update_daily_spend_with_none_values_in_sorting_fields():
"""
Test that _update_daily_spend handles None values in sorting fields correctly.
This test ensures that when fields like date, api_key, model, or custom_llm_provider
are None, the sorting doesn't crash with TypeError: '<' not supported between
are None, the sorting doesn't crash with TypeError: '<' not supported between
instances of 'NoneType' and 'str'.
"""
# Setup
@@ -509,6 +515,7 @@ 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():
"""
@@ -518,7 +525,7 @@ async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_i
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,
@@ -546,13 +553,15 @@ async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_i
# 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')}"
assert (
transaction["request_id"] == request_id
), f"request_id should be {request_id} but got {transaction.get('request_id')}"
@pytest.mark.asyncio
@@ -866,11 +875,11 @@ async def test_endpoint_field_is_correctly_mapped_from_call_type():
call_args = writer.daily_spend_update_queue.add_update.call_args[1]
update_dict = call_args["update"]
assert len(update_dict) == 1
for key, transaction in update_dict.items():
# Verify endpoint is included in the key
assert key == f"test-user_2024-01-01_test-key_gpt-4_openai_/chat/completions"
# Verify endpoint is set in the transaction
assert transaction["endpoint"] == "/chat/completions"
assert transaction["user_id"] == "test-user"
@@ -887,7 +896,7 @@ async def test_update_daily_spend_logs_detailed_error_on_batch_upsert_failure():
This ensures proper debugging information is available for issues like unique constraint violations.
"""
from litellm._logging import verbose_proxy_logger
# Setup
mock_prisma_client = MagicMock()
mock_batcher = MagicMock()
@@ -895,13 +904,13 @@ async def test_update_daily_spend_logs_detailed_error_on_batch_upsert_failure():
mock_batch_context = MagicMock()
mock_batch_context.__aenter__ = AsyncMock(return_value=mock_batcher)
mock_batcher.litellm_dailyuserspend = mock_table
# Make the batch context manager's exit raise an exception
# This simulates a batch commit failure (e.g., unique constraint violation)
test_exception = Exception("Unique constraint violation")
mock_batch_context.__aexit__ = AsyncMock(side_effect=test_exception)
mock_prisma_client.db.batch_.return_value = mock_batch_context
# Create a transaction
daily_spend_transactions = {
"test_key": {
@@ -918,13 +927,13 @@ async def test_update_daily_spend_logs_detailed_error_on_batch_upsert_failure():
"failed_requests": 0,
}
}
# Create a mock proxy_logging_obj with failure_handler as AsyncMock
mock_proxy_logging = MagicMock()
mock_proxy_logging.failure_handler = AsyncMock()
# Mock the logger to capture exception calls
with patch.object(verbose_proxy_logger, 'exception') as mock_exception_logger:
with patch.object(verbose_proxy_logger, "exception") as mock_exception_logger:
# Call the method and expect it to raise the exception
with pytest.raises(Exception, match="Unique constraint violation"):
await DBSpendUpdateWriter._update_daily_spend(
@@ -937,13 +946,16 @@ async def test_update_daily_spend_logs_detailed_error_on_batch_upsert_failure():
table_name="litellm_dailyuserspend",
unique_constraint_name="user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint",
)
# Verify that exception was logged with detailed information
assert mock_exception_logger.called
call_args = mock_exception_logger.call_args[0][0]
assert "Daily user spend batch upsert failed" in call_args
assert "Table: litellm_dailyuserspend" in call_args
assert "Constraint: user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint" in call_args
assert (
"Constraint: user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint"
in call_args
)
assert "Batch size: 1" in call_args
assert "Unique constraint violation" in call_args
@@ -961,7 +973,7 @@ async def test_update_daily_spend_re_raises_exception_after_logging():
mock_batch_context = MagicMock()
mock_batch_context.__aenter__ = AsyncMock(return_value=mock_batcher)
mock_batcher.litellm_dailyuserspend = mock_table
# Create a transaction
daily_spend_transactions = {
"test_key": {
@@ -978,16 +990,16 @@ async def test_update_daily_spend_re_raises_exception_after_logging():
"failed_requests": 0,
}
}
# Create a custom exception to verify it's re-raised
custom_exception = ValueError("Database connection lost")
mock_batch_context.__aexit__ = AsyncMock(side_effect=custom_exception)
mock_prisma_client.db.batch_.return_value = mock_batch_context
# Create a mock proxy_logging_obj with failure_handler as AsyncMock
mock_proxy_logging = MagicMock()
mock_proxy_logging.failure_handler = AsyncMock()
# Verify the exception is re-raised
with pytest.raises(ValueError, match="Database connection lost"):
await DBSpendUpdateWriter._update_daily_spend(
@@ -1018,10 +1030,12 @@ async def test_commit_key_spend_updates_includes_last_active():
mock_transaction = AsyncMock()
mock_transaction.__aenter__ = AsyncMock(return_value=mock_transaction)
mock_transaction.__aexit__ = AsyncMock(return_value=False)
mock_transaction.batch_ = MagicMock(return_value=AsyncMock(
__aenter__=AsyncMock(return_value=mock_batcher),
__aexit__=AsyncMock(return_value=False),
))
mock_transaction.batch_ = MagicMock(
return_value=AsyncMock(
__aenter__=AsyncMock(return_value=mock_batcher),
__aexit__=AsyncMock(return_value=False),
)
)
mock_prisma_client = MagicMock()
mock_prisma_client.db = MagicMock()
@@ -1049,9 +1063,7 @@ async def test_commit_key_spend_updates_includes_last_active():
before_call = datetime.now(timezone.utc)
with patch(
"litellm.proxy.utils._raise_failed_update_spend_exception"
):
with patch("litellm.proxy.utils._raise_failed_update_spend_exception"):
await db_writer._commit_spend_updates_to_db(
prisma_client=mock_prisma_client,
n_retry_times=0,
@@ -1076,3 +1088,146 @@ async def test_commit_key_spend_updates_includes_last_active():
last_active = call_kwargs["data"]["last_active"]
assert isinstance(last_active, datetime)
assert before_call <= last_active <= after_call
@pytest.mark.asyncio
async def test_update_database_creates_single_task():
"""
Test that update_database() fires exactly 1 asyncio.create_task() call
(the batched task) instead of the previous 11.
"""
db_writer = DBSpendUpdateWriter()
# Mock all helpers so nothing real runs
db_writer._insert_spend_log_to_db = AsyncMock()
db_writer._batch_database_updates = AsyncMock()
with patch("litellm.proxy.proxy_server.disable_spend_logs", False), patch(
"litellm.proxy.proxy_server.prisma_client", MagicMock()
), patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), patch(
"litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget"
), patch(
"litellm.proxy.db.db_spend_update_writer.asyncio.create_task"
) as mock_create_task:
await db_writer.update_database(
token="test-token",
user_id="test-user",
end_user_id="test-end-user",
start_time=datetime.now(),
end_time=datetime.now(),
team_id="test-team",
org_id="test-org",
completion_response=MagicMock(),
response_cost=0.1,
kwargs={"model": "gpt-4", "custom_llm_provider": "openai"},
)
# Exactly 1 create_task call (the batch), not 11
assert mock_create_task.call_count == 1
@pytest.mark.asyncio
async def test_batch_database_updates_isolation_on_failure():
"""
Test that if one helper inside _batch_database_updates raises,
all other helpers still execute.
"""
db_writer = DBSpendUpdateWriter()
# Make _update_key_db raise
db_writer._update_key_db = AsyncMock(side_effect=RuntimeError("key db boom"))
# All other helpers are normal mocks
db_writer._update_user_db = AsyncMock()
db_writer._update_team_db = AsyncMock()
db_writer._update_org_db = AsyncMock()
db_writer._update_tag_db = AsyncMock()
db_writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock()
db_writer.add_spend_log_transaction_to_daily_end_user_transaction = AsyncMock()
db_writer.add_spend_log_transaction_to_daily_agent_transaction = AsyncMock()
db_writer.add_spend_log_transaction_to_daily_team_transaction = AsyncMock()
db_writer.add_spend_log_transaction_to_daily_org_transaction = AsyncMock()
db_writer.add_spend_log_transaction_to_daily_tag_transaction = AsyncMock()
await db_writer._batch_database_updates(
response_cost=0.1,
user_id="u1",
hashed_token="t1",
team_id="team1",
org_id="org1",
end_user_id="eu1",
prisma_client=MagicMock(),
user_api_key_cache=MagicMock(),
litellm_proxy_budget_name="budget",
payload_copy={"key": "value"},
request_tags=None,
)
# _update_key_db raised, but all others should still have been called
db_writer._update_user_db.assert_awaited_once()
db_writer._update_key_db.assert_awaited_once()
db_writer._update_team_db.assert_awaited_once()
db_writer._update_org_db.assert_awaited_once()
db_writer._update_tag_db.assert_awaited_once()
db_writer.add_spend_log_transaction_to_daily_user_transaction.assert_awaited_once()
db_writer.add_spend_log_transaction_to_daily_end_user_transaction.assert_awaited_once()
db_writer.add_spend_log_transaction_to_daily_agent_transaction.assert_awaited_once()
db_writer.add_spend_log_transaction_to_daily_team_transaction.assert_awaited_once()
db_writer.add_spend_log_transaction_to_daily_org_transaction.assert_awaited_once()
db_writer.add_spend_log_transaction_to_daily_tag_transaction.assert_awaited_once()
@pytest.mark.asyncio
async def test_daily_agent_receives_deepcopied_payload():
"""
Test that the daily agent handler receives a deepcopied payload (not the original).
Previously, add_spend_log_transaction_to_daily_agent_transaction received the raw
payload without a deepcopy, which was a mutation bug.
"""
db_writer = DBSpendUpdateWriter()
original_payload = {"key": "value", "nested": {"a": 1}}
captured_payloads = []
async def capture_payload(**kwargs):
captured_payloads.append(kwargs.get("payload"))
# Mock all helpers
db_writer._update_user_db = AsyncMock()
db_writer._update_key_db = AsyncMock()
db_writer._update_team_db = AsyncMock()
db_writer._update_org_db = AsyncMock()
db_writer._update_tag_db = AsyncMock()
db_writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock()
db_writer.add_spend_log_transaction_to_daily_end_user_transaction = AsyncMock()
db_writer.add_spend_log_transaction_to_daily_agent_transaction = AsyncMock(
side_effect=capture_payload
)
db_writer.add_spend_log_transaction_to_daily_team_transaction = AsyncMock()
db_writer.add_spend_log_transaction_to_daily_org_transaction = AsyncMock()
db_writer.add_spend_log_transaction_to_daily_tag_transaction = AsyncMock()
import copy
payload_copy = copy.deepcopy(original_payload)
await db_writer._batch_database_updates(
response_cost=0.1,
user_id="u1",
hashed_token="t1",
team_id="team1",
org_id="org1",
end_user_id="eu1",
prisma_client=MagicMock(),
user_api_key_cache=MagicMock(),
litellm_proxy_budget_name="budget",
payload_copy=payload_copy,
request_tags=None,
)
# The payload passed to the agent handler should NOT be the original object
assert len(captured_payloads) == 1
assert captured_payloads[0] is not original_payload
# But it should have the same content
assert captured_payloads[0] == original_payload