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(): """