From bd158c8108e568b72fefca84c68c60d80aa2ca9e Mon Sep 17 00:00:00 2001 From: narmstrong-berri Date: Mon, 19 Oct 2026 20:39:39 +0000 Subject: [PATCH] perf(spend): build the daily rollup row once per spend log instead of once per entity _batch_database_updates queued the daily user, end user, agent, team, org and tag rows by calling _common_add_spend_log_transaction_to_daily_transaction once per entity, and each call re-parsed the metadata and re-ran compute_savings_spend for the same request. The shared fields do not depend on the entity, so they are now built once and handed to every daily rollup. The adders still build their own row when called without one --- litellm/proxy/db/db_spend_update_writer.py | 57 ++++++-- .../proxy/db/test_db_spend_update_writer.py | 124 ++++++++++++++++++ 2 files changed, 172 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 166a42adfc0..94d442b0201 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -1131,10 +1131,12 @@ class DBSpendUpdateWriter: traceback.format_exc(), ) + daily_spend_base: Final = await self._shared_daily_spend_base_transaction(payload_copy, prisma_client) try: await self.add_spend_log_transaction_to_daily_user_transaction( payload=payload_copy, prisma_client=prisma_client, + daily_spend_base=daily_spend_base, ) except Exception: verbose_proxy_logger.debug( @@ -1146,6 +1148,7 @@ class DBSpendUpdateWriter: await self.add_spend_log_transaction_to_daily_end_user_transaction( payload=payload_copy, prisma_client=prisma_client, + daily_spend_base=daily_spend_base, ) except Exception: verbose_proxy_logger.debug( @@ -1157,6 +1160,7 @@ class DBSpendUpdateWriter: await self.add_spend_log_transaction_to_daily_agent_transaction( payload=payload_copy, prisma_client=prisma_client, + daily_spend_base=daily_spend_base, ) except Exception: verbose_proxy_logger.debug( @@ -1168,6 +1172,7 @@ class DBSpendUpdateWriter: await self.add_spend_log_transaction_to_daily_team_transaction( payload=payload_copy, prisma_client=prisma_client, + daily_spend_base=daily_spend_base, ) except Exception: verbose_proxy_logger.debug( @@ -1187,6 +1192,7 @@ class DBSpendUpdateWriter: payload=payload_copy, org_id=org_id, prisma_client=prisma_client, + daily_spend_base=daily_spend_base, ) except Exception: verbose_proxy_logger.debug( @@ -1198,6 +1204,7 @@ class DBSpendUpdateWriter: await self.add_spend_log_transaction_to_daily_tag_transaction( payload=payload_copy, prisma_client=prisma_client, + daily_spend_base=daily_spend_base, ) except Exception: verbose_proxy_logger.debug( @@ -2761,9 +2768,10 @@ class DBSpendUpdateWriter: async def _common_add_spend_log_transaction_to_daily_transaction( self, - payload: dict | SpendLogsPayload, + payload: SpendLogsPayload, prisma_client: PrismaClient, type: Literal["user", "team", "org", "request_tags", "end_user", "agent"] = "user", + daily_spend_base: BaseDailySpendTransaction | None = None, ) -> BaseDailySpendTransaction | None: entity: Final = "tag" if type == "request_tags" else type identity_payload: Final = ( @@ -2771,6 +2779,31 @@ class DBSpendUpdateWriter: ) if not daily_spend_entity_ids(identity_payload, entity): return None + if daily_spend_base is not None: + return daily_spend_base + return await self._daily_spend_base_transaction(payload, prisma_client) + + async def _shared_daily_spend_base_transaction( + self, + payload: SpendLogsPayload, + prisma_client: PrismaClient | None, + ) -> BaseDailySpendTransaction | None: + if prisma_client is None: + return None + try: + return await self._daily_spend_base_transaction(payload, prisma_client) + except Exception: # noqa: BLE001 # each daily rollup falls back to building its own row, as before + verbose_proxy_logger.debug( + "_batch_database_updates: _daily_spend_base_transaction failed: %s", + traceback.format_exc(), + ) + return None + + async def _daily_spend_base_transaction( + self, + payload: SpendLogsPayload, + prisma_client: PrismaClient, + ) -> BaseDailySpendTransaction | None: expected_keys: Final = ("startTime", "api_key") if not all(key in payload for key in expected_keys): verbose_proxy_logger.debug( @@ -2806,7 +2839,7 @@ class DBSpendUpdateWriter: if isinstance(payload["startTime"], datetime): start_time: Final = payload["startTime"].isoformat() date = start_time.split("T")[0] - elif isinstance(payload["startTime"], str): + elif isinstance(payload["startTime"], str): # pyright: ignore[reportUnnecessaryIsInstance] # spend log payloads are not validated against the TypedDict at runtime date = payload["startTime"].split("T")[0] else: verbose_proxy_logger.debug( @@ -2874,8 +2907,9 @@ class DBSpendUpdateWriter: async def add_spend_log_transaction_to_daily_user_transaction( self, - payload: dict | SpendLogsPayload, + payload: SpendLogsPayload, prisma_client: PrismaClient | None = None, + daily_spend_base: BaseDailySpendTransaction | None = None, ): """ Add a spend log transaction to the `daily_spend_update_queue` @@ -2889,7 +2923,7 @@ class DBSpendUpdateWriter: return base_daily_transaction: Final = await self._common_add_spend_log_transaction_to_daily_transaction( - payload, prisma_client, "user" + payload, prisma_client, "user", daily_spend_base ) if base_daily_transaction is None: return @@ -2903,13 +2937,14 @@ class DBSpendUpdateWriter: self, payload: SpendLogsPayload, prisma_client: PrismaClient | None = None, + daily_spend_base: BaseDailySpendTransaction | None = None, ) -> None: if prisma_client is None: verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.") return base_daily_transaction: Final = await self._common_add_spend_log_transaction_to_daily_transaction( - payload, prisma_client, "team" + payload, prisma_client, "team", daily_spend_base ) if base_daily_transaction is None: return @@ -2927,6 +2962,7 @@ class DBSpendUpdateWriter: payload: SpendLogsPayload, prisma_client: PrismaClient | None = None, org_id: str | None = None, + daily_spend_base: BaseDailySpendTransaction | None = None, ) -> None: if prisma_client is None: verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.") @@ -2945,7 +2981,7 @@ class DBSpendUpdateWriter: ) base_daily_transaction: Final = await self._common_add_spend_log_transaction_to_daily_transaction( - payload_with_org, prisma_client, "org" + payload_with_org, prisma_client, "org", daily_spend_base ) if base_daily_transaction is None: return @@ -2959,6 +2995,7 @@ class DBSpendUpdateWriter: self, payload: SpendLogsPayload, prisma_client: PrismaClient | None = None, + daily_spend_base: BaseDailySpendTransaction | None = None, ) -> None: if prisma_client is None: verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.") @@ -2978,7 +3015,7 @@ class DBSpendUpdateWriter: ) base_daily_transaction: Final = await self._common_add_spend_log_transaction_to_daily_transaction( - payload_with_end_user_id, prisma_client, "end_user" + payload_with_end_user_id, prisma_client, "end_user", daily_spend_base ) if base_daily_transaction is None: return @@ -2992,6 +3029,7 @@ class DBSpendUpdateWriter: self, payload: SpendLogsPayload, prisma_client: PrismaClient | None = None, + daily_spend_base: BaseDailySpendTransaction | None = None, ) -> None: if prisma_client is None: verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.") @@ -3006,7 +3044,7 @@ class DBSpendUpdateWriter: }, ) base_daily_transaction: Final = await self._common_add_spend_log_transaction_to_daily_transaction( - payload_with_agent_id, prisma_client, "agent" + payload_with_agent_id, prisma_client, "agent", daily_spend_base ) if base_daily_transaction is None: return @@ -3019,13 +3057,14 @@ class DBSpendUpdateWriter: self, payload: SpendLogsPayload, prisma_client: PrismaClient | None = None, + daily_spend_base: BaseDailySpendTransaction | None = None, ) -> None: if prisma_client is None: verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.") return base_daily_transaction: Final = await self._common_add_spend_log_transaction_to_daily_transaction( - payload, prisma_client, "request_tags" + payload, prisma_client, "request_tags", daily_spend_base ) if base_daily_transaction is None: return diff --git a/tests/unit/proxy/db/test_db_spend_update_writer.py b/tests/unit/proxy/db/test_db_spend_update_writer.py index 1ec65a4fdcd..b73a985f976 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -1,5 +1,6 @@ import asyncio import copy +import itertools import json import logging import re @@ -35,6 +36,7 @@ from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( build_window_spend_transaction, ) from litellm.proxy.db.rollup_lock_timeout import ROLLUP_LOCK_TIMEOUT_SQL +from litellm.proxy.spend_tracking.savings import compute_savings_spend from tests.unit.proxy.db.fake_prisma_engine import engine_call @@ -1440,6 +1442,128 @@ async def test_failed_project_enqueue_is_reported_and_does_not_drop_the_rest_of_ assert transactions["team_list_transactions"] == {"team-1": 0.25} +@pytest.mark.asyncio +async def test_batch_database_updates_prices_savings_once_for_every_daily_rollup(): + db_writer: Final = DBSpendUpdateWriter() + prisma_client: Final = MagicMock() + prisma_client.get_request_status = MagicMock(return_value="success") + payload: Final = { + "request_id": "req-daily", + "user": "user-1", + "team_id": "team-1", + "end_user": "end-user-1", + "agent_id": "agent-1", + "request_tags": ["tag-a", "tag-b"], + "startTime": "2024-01-01T12:00:00", + "api_key": "key-1", + "model": "claude-sonnet-5", + "custom_llm_provider": "anthropic", + "model_group": "claude-sonnet-5", + "prompt_tokens": 1000, + "completion_tokens": 10, + "spend": 0.01, + "metadata": json.dumps({"usage_object": {"prompt_tokens": 1000, "cache_read_input_tokens": 800}}), + } + + with patch("litellm.proxy.db.db_spend_update_writer.compute_savings_spend", wraps=compute_savings_spend) as savings: + await db_writer._batch_database_updates( + response_cost=0.01, + user_id="user-1", + hashed_token="key-1", + team_id="team-1", + org_id="org-1", + end_user_id="end-user-1", + prisma_client=prisma_client, + litellm_proxy_budget_name=None, + payload=payload, + ) + + assert savings.call_count == 1 + queues: Final = { + "user_id": db_writer.daily_spend_update_queue, + "team_id": db_writer.daily_team_spend_update_queue, + "organization_id": db_writer.daily_org_spend_update_queue, + "end_user_id": db_writer.daily_end_user_spend_update_queue, + "agent_id": db_writer.daily_agent_spend_update_queue, + "tag": db_writer.daily_tag_spend_update_queue, + } + rows: Final = { + entity: tuple( + sorted((await queue.flush_and_get_aggregated_daily_spend_update_transactions()).values(), key=str) + ) + for entity, queue in queues.items() + } + assert {entity: tuple(row[entity] for row in entity_rows) for entity, entity_rows in rows.items()} == { + "user_id": ("user-1",), + "team_id": ("team-1",), + "organization_id": ("org-1",), + "end_user_id": ("end-user-1",), + "agent_id": ("agent-1",), + "tag": ("tag-a", "tag-b"), + } + all_rows: Final = tuple(itertools.chain.from_iterable(rows.values())) + savings_columns: Final = { + (row["prompt_caching_savings_spend"], row["cache_read_input_tokens"], row["spend"]) for row in all_rows + } + assert len(savings_columns) == 1 + assert next(iter(savings_columns))[0] > 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "overrides, dropped, has_prisma", + [ + ({"metadata": "{not json"}, (), True), + ({}, (), False), + ({}, ("api_key",), True), + ({}, ("model",), True), + ({}, ("model_group",), True), + ], +) +async def test_batch_database_updates_skips_daily_rows_it_cannot_build_but_keeps_key_spend( + overrides: dict[str, str], dropped: tuple[str, ...], has_prisma: bool +): + db_writer: Final = DBSpendUpdateWriter() + prisma_client: Final = MagicMock() if has_prisma else None + if prisma_client is not None: + prisma_client.get_request_status = MagicMock(return_value="success") + complete_payload: Final = { + "request_id": "req-bad", + "user": "user-1", + "team_id": "team-1", + "startTime": "2024-01-01T12:00:00", + "api_key": "key-1", + "model": "gpt-4o-mini", + "custom_llm_provider": "openai", + "model_group": "gpt-4o-mini", + "prompt_tokens": 10, + "completion_tokens": 5, + "spend": 0.01, + "metadata": '{"usage_object": {}}', + **overrides, + } + payload: Final = {key: value for key, value in complete_payload.items() if key not in dropped} + + await db_writer._batch_database_updates( + response_cost=0.01, + user_id="user-1", + hashed_token="key-1", + team_id="team-1", + org_id=None, + end_user_id=None, + prisma_client=prisma_client, + litellm_proxy_budget_name=None, + payload=payload, + ) + + assert await db_writer.daily_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() == {} + assert ( + await db_writer.daily_team_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() == {} + ) + transactions: Final = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() + assert transactions["key_list_transactions"] == ({"key-1": 0.01} if has_prisma else {}) + + @pytest.mark.asyncio async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_id(): """