Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py: 92%
65 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1import asyncio
2from collections.abc import Coroutine
3from copy import deepcopy
4from typing import Final
6from litellm._logging import verbose_proxy_logger
7from litellm.constants import LITELLM_ASYNCIO_QUEUE_MAXSIZE
8from litellm.proxy._types import BaseDailySpendTransaction
9from litellm.proxy.db.db_transaction_queue.base_update_queue import (
10 BaseUpdateQueue,
11 service_logger_obj,
12)
13from litellm.types.services import ServiceTypes
16class DailySpendUpdateQueue(BaseUpdateQueue):
17 """
18 In memory buffer for daily spend updates that should be committed to the database
20 To add a new daily spend update transaction, use the following format:
21 daily_spend_update_queue.add_update({
22 "user1_date_api_key_model_custom_llm_provider": {
23 "spend": 10,
24 "prompt_tokens": 100,
25 "completion_tokens": 100,
26 }
27 })
29 Queue contains a list of daily spend update transactions
31 eg
32 queue = [
33 {
34 "user1_date_api_key_model_custom_llm_provider": {
35 "spend": 10,
36 "prompt_tokens": 100,
37 "completion_tokens": 100,
38 "api_requests": 100,
39 "successful_requests": 100,
40 "failed_requests": 100,
41 }
42 },
43 {
44 "user2_date_api_key_model_custom_llm_provider": {
45 "spend": 10,
46 "prompt_tokens": 100,
47 "completion_tokens": 100,
48 "api_requests": 100,
49 "successful_requests": 100,
50 "failed_requests": 100,
51 }
52 }
53 ]
54 """
56 def __init__(self):
57 super().__init__()
58 self.update_queue: asyncio.Queue[dict[str, BaseDailySpendTransaction]] = asyncio.Queue(
59 maxsize=LITELLM_ASYNCIO_QUEUE_MAXSIZE
60 )
61 self.interrupted_commits: set[asyncio.Task[None]] = (
62 set()
63 ) # mutable-ok: registry of in-flight commit outcomes, entries leave via their done callback
65 def track_interrupted_commit(self, settle: Coroutine[object, object, None]) -> None:
66 task: Final = asyncio.ensure_future(settle)
67 self.interrupted_commits.add(task)
68 task.add_done_callback(self.interrupted_commits.discard)
70 async def settle_interrupted_commits(self) -> None:
71 while self.interrupted_commits: 71 ↛ 72line 71 didn't jump to line 72 because the condition on line 71 was never true
72 await asyncio.wait(tuple(self.interrupted_commits))
74 async def add_update(self, update: dict[str, BaseDailySpendTransaction]):
75 """Enqueue an update."""
76 verbose_proxy_logger.debug("Adding update to queue: %s", update)
77 await self.update_queue.put(update)
78 if self.update_queue.qsize() >= self.MAX_SIZE_IN_MEMORY_QUEUE:
79 verbose_proxy_logger.warning(
80 "Spend update queue is full. Aggregating all entries in queue to concatenate entries."
81 )
82 await self.aggregate_queue_updates()
84 async def aggregate_queue_updates(self):
85 """
86 Combine all updates in the queue into a single update.
87 This is used to reduce the size of the in-memory queue.
88 """
89 updates: Final[list[dict[str, BaseDailySpendTransaction]]] = await self.flush_all_updates_from_in_memory_queue()
90 aggregated_updates: Final = self.get_aggregated_daily_spend_update_transactions(updates)
91 await self.update_queue.put(aggregated_updates)
93 async def flush_and_get_aggregated_daily_spend_update_transactions(
94 self,
95 ) -> dict[str, BaseDailySpendTransaction]:
96 """Get all updates from the queue and return all updates aggregated by daily_transaction_key. Works for both user and team spend updates."""
97 await self.settle_interrupted_commits()
98 updates: Final = await self.flush_all_updates_from_in_memory_queue()
99 if len(updates) > 0:
100 verbose_proxy_logger.info(
101 "Spend tracking - flushed %d daily spend update items from in-memory queue",
102 len(updates),
103 )
104 aggregated_daily_spend_update_transactions: Final = (
105 DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions(updates)
106 )
107 verbose_proxy_logger.debug(
108 "Aggregated daily spend update transactions: %s",
109 aggregated_daily_spend_update_transactions,
110 )
111 return aggregated_daily_spend_update_transactions
113 @staticmethod
114 def get_aggregated_daily_spend_update_transactions(
115 updates: list[dict[str, BaseDailySpendTransaction]],
116 ) -> dict[str, BaseDailySpendTransaction]:
117 """Aggregate updates by daily_transaction_key."""
118 aggregated_daily_spend_update_transactions: Final[dict[str, BaseDailySpendTransaction]] = {}
119 for _update in updates:
120 for _key, payload in _update.items():
121 if _key in aggregated_daily_spend_update_transactions:
122 daily_transaction = aggregated_daily_spend_update_transactions[_key]
123 daily_transaction["spend"] += payload["spend"]
124 daily_transaction["prompt_tokens"] += payload["prompt_tokens"]
125 daily_transaction["completion_tokens"] += payload["completion_tokens"]
126 daily_transaction["api_requests"] += payload["api_requests"]
127 daily_transaction["successful_requests"] += payload["successful_requests"]
128 daily_transaction["failed_requests"] += payload["failed_requests"]
130 # Add optional metrics cache_read_input_tokens and cache_creation_input_tokens
131 daily_transaction["cache_read_input_tokens"] = (
132 payload.get("cache_read_input_tokens", 0) or 0
133 ) + daily_transaction.get("cache_read_input_tokens", 0)
135 daily_transaction["cache_creation_input_tokens"] = (
136 payload.get("cache_creation_input_tokens", 0) or 0
137 ) + daily_transaction.get("cache_creation_input_tokens", 0)
139 daily_transaction["compression_saved_tokens"] = (
140 payload.get("compression_saved_tokens", 0) or 0
141 ) + daily_transaction.get("compression_saved_tokens", 0)
143 daily_transaction["compression_savings_spend"] = (
144 payload.get("compression_savings_spend", 0) or 0
145 ) + daily_transaction.get("compression_savings_spend", 0)
147 daily_transaction["prompt_caching_savings_spend"] = (
148 payload.get("prompt_caching_savings_spend", 0) or 0
149 ) + daily_transaction.get("prompt_caching_savings_spend", 0)
151 daily_transaction["gateway_injected_caching_savings_spend"] = (
152 payload.get("gateway_injected_caching_savings_spend", 0) or 0
153 ) + daily_transaction.get("gateway_injected_caching_savings_spend", 0)
155 daily_transaction["autorouter_savings_spend"] = (
156 payload.get("autorouter_savings_spend", 0) or 0
157 ) + daily_transaction.get("autorouter_savings_spend", 0)
159 daily_transaction["total_response_time_ms"] = (
160 payload.get("total_response_time_ms", 0) or 0
161 ) + daily_transaction.get("total_response_time_ms", 0)
163 daily_transaction["timed_requests"] = (
164 payload.get("timed_requests", 0) or 0
165 ) + daily_transaction.get("timed_requests", 0)
167 else:
168 aggregated_daily_spend_update_transactions[_key] = deepcopy(payload)
169 return aggregated_daily_spend_update_transactions
171 async def _emit_new_item_added_to_queue_event(
172 self,
173 queue_size: int | None = None,
174 ):
175 asyncio.create_task(
176 service_logger_obj.async_service_success_hook(
177 service=ServiceTypes.IN_MEMORY_DAILY_SPEND_UPDATE_QUEUE,
178 duration=0,
179 call_type="_emit_new_item_added_to_queue_event",
180 event_metadata={
181 "gauge_labels": ServiceTypes.IN_MEMORY_DAILY_SPEND_UPDATE_QUEUE,
182 "gauge_value": queue_size,
183 },
184 )
185 )