Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/db/db_transaction_queue/window_spend_update_queue.py: 64%

1""" 

2In memory buffer for per-budget-window spend increments. 

3 

4Kept separate from SpendUpdateQueue: an increment is only meaningful together 

5with the window it landed in, so two increments for the same entity must not be 

6merged when their window_start differs. 

7""" 

8 

9import asyncio 

10import math 

11from collections.abc import Sequence 

12from datetime import datetime, timezone 

13from itertools import chain, groupby 

14from typing import Final, TypedDict 

15 

16from typing_extensions import ReadOnly 

17 

18from litellm._logging import verbose_proxy_logger 

19from litellm.constants import LITELLM_ASYNCIO_QUEUE_MAXSIZE 

20from litellm.proxy.db.db_transaction_queue.base_update_queue import BaseUpdateQueue 

21 

22 

23class WindowSpendTransaction(TypedDict): 

24 """One increment for a single (entity, budget window) pair. 

25 

26 window_start is an ISO-8601 string rather than a datetime so the 

27 transaction survives the JSON round trip through the Redis buffer. 

28 

29 started_at is the earliest request start in the batch. The one-time seed for 

30 a window that has no row yet uses it to tell this batch's own 

31 LiteLLM_SpendLogs rows from everything else, because the spend log writer 

32 flushes on its own ~2s poll and will usually have persisted this batch's 

33 rows before the window queue flushes; without that split the seed and the 

34 increment would each count them. 

35 """ 

36 

37 entity_type: ReadOnly[str] 

38 entity_id: ReadOnly[str] 

39 window_duration: ReadOnly[str] 

40 window_start: ReadOnly[str] 

41 spend: ReadOnly[float] 

42 started_at: ReadOnly[str | None] 

43 

44 

45class WindowSpendWirePayload(WindowSpendTransaction): 

46 """How an increment is encoded in the shared Redis buffer. 

47 

48 request_ids is dead weight here: workers built before this field was 

49 dropped index it while merging whatever they pop, and the pop is 

50 destructive, so a leader still running one of those during a rolling deploy 

51 would raise on a payload without the key and lose those increments. It is 

52 always empty, which only makes such a leader seed without exclusions. 

53 

54 TODO: remove once no supported version reads it, i.e. one release after the 

55 field stopped being written. 

56 """ 

57 

58 request_ids: ReadOnly[Sequence[str]] 

59 

60 

61def to_wire_payload(transaction: WindowSpendTransaction) -> WindowSpendWirePayload: 

62 return WindowSpendWirePayload( 

63 entity_type=transaction["entity_type"], 

64 entity_id=transaction["entity_id"], 

65 window_duration=transaction["window_duration"], 

66 window_start=transaction["window_start"], 

67 spend=transaction["spend"], 

68 started_at=transaction.get("started_at"), 

69 request_ids=(), 

70 ) 

71 

72 

73def to_naive_utc(value: datetime) -> datetime: 

74 """LiteLLM_BudgetWindowSpend.window_start is TIMESTAMP(3), which holds naive UTC.""" 

75 if value.tzinfo is None: 

76 return value 

77 return value.astimezone(timezone.utc).replace(tzinfo=None) 

78 

79 

80def window_spend_group_key(transaction: WindowSpendTransaction) -> tuple[str, str, str, str]: 

81 """Identity of a window increment: the row's primary key plus the window it 

82 belongs to. Two increments only aggregate when all four match.""" 

83 return ( 

84 transaction["entity_type"], 

85 transaction["entity_id"], 

86 transaction["window_duration"], 

87 transaction["window_start"], 

88 ) 

89 

90 

91def build_window_spend_transaction( 

92 entity_type: str, 

93 entity_id: str, 

94 window_duration: str, 

95 window_start: datetime, 

96 spend: float, 

97 started_at: datetime | None = None, 

98) -> WindowSpendTransaction: 

99 return WindowSpendTransaction( 

100 entity_type=entity_type, 

101 entity_id=entity_id, 

102 window_duration=window_duration, 

103 window_start=to_naive_utc(window_start).isoformat(timespec="microseconds"), 

104 spend=spend, 

105 started_at=None 

106 if started_at is None 

107 else to_naive_utc(started_at.astimezone(timezone.utc)).isoformat(timespec="microseconds"), 

108 ) 

109 

110 

111def _merge_window_spend_transactions( 

112 payloads: tuple[WindowSpendTransaction, ...], 

113) -> WindowSpendTransaction: 

114 first: Final = payloads[0] 

115 started_ats: Final = tuple( 

116 started_at for payload in payloads if (started_at := payload.get("started_at")) is not None 

117 ) 

118 return WindowSpendTransaction( 

119 entity_type=first["entity_type"], 

120 entity_id=first["entity_id"], 

121 window_duration=first["window_duration"], 

122 window_start=first["window_start"], 

123 spend=math.fsum(payload["spend"] for payload in payloads), 

124 started_at=min(started_ats) if started_ats else None, 

125 ) 

126 

127 

128class WindowSpendUpdateQueue(BaseUpdateQueue): 

129 """ 

130 In memory buffer for budget-window spend increments committed to 

131 LiteLLM_BudgetWindowSpend. 

132 

133 Add an update with the payload built by build_window_spend_transaction: 

134 window_spend_update_queue.add_update( 

135 build_window_spend_transaction( 

136 entity_type="key", 

137 entity_id="<hashed token>", 

138 window_duration="30d", 

139 window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), 

140 spend=0.02, 

141 ) 

142 ) 

143 """ 

144 

145 def __init__(self) -> None: 

146 super().__init__() 

147 self.update_queue: asyncio.Queue[tuple[WindowSpendTransaction, ...]] = asyncio.Queue( 

148 maxsize=LITELLM_ASYNCIO_QUEUE_MAXSIZE 

149 ) 

150 

151 async def add_update(self, update: WindowSpendTransaction) -> None: 

152 """Enqueue an update.""" 

153 verbose_proxy_logger.debug("Adding budget window spend update to queue: %s", update) 

154 await self.update_queue.put((update,)) 

155 if self.update_queue.qsize() >= self.MAX_SIZE_IN_MEMORY_QUEUE: 

156 verbose_proxy_logger.warning( 

157 "Budget window spend update queue is full. Aggregating all entries in queue to concatenate entries." 

158 ) 

159 await self.aggregate_queue_updates() 

160 

161 async def aggregate_queue_updates(self) -> None: 

162 """Collapse everything currently queued into a single aggregated update.""" 

163 updates: Final = await self.flush_all_updates_from_in_memory_queue() 

164 await self.update_queue.put(WindowSpendUpdateQueue.get_aggregated_window_spend_transactions(updates)) 

165 

166 async def flush_and_get_aggregated_window_spend_transactions( 

167 self, 

168 ) -> tuple[WindowSpendTransaction, ...]: 

169 """Drain the queue and return the increments aggregated per window.""" 

170 updates: Final = await self.flush_all_updates_from_in_memory_queue() 

171 if len(updates) > 0: 171 ↛ 172line 171 didn't jump to line 172 because the condition on line 171 was never true

172 verbose_proxy_logger.info( 

173 "Spend tracking - flushed %d budget window spend update batches from in-memory queue", 

174 len(updates), 

175 ) 

176 return WindowSpendUpdateQueue.get_aggregated_window_spend_transactions(updates) 

177 

178 @staticmethod 

179 def get_aggregated_window_spend_transactions( 

180 updates: Sequence[Sequence[WindowSpendTransaction]], 

181 ) -> tuple[WindowSpendTransaction, ...]: 

182 """Sum spend per (entity_type, entity_id, window_duration, window_start). 

183 

184 Increments belonging to different windows stay separate even when they 

185 share a primary key, so a window boundary crossed mid-tick does not fold 

186 the new window's spend into the previous window's total. 

187 

188 The result is ordered by that same key, which is the order the flush 

189 needs: primary key first for cross-pod lock ordering, then window_start 

190 so an older window is applied before the roll that supersedes it. 

191 """ 

192 ordered: Final = tuple(sorted(chain.from_iterable(updates), key=window_spend_group_key)) 

193 return tuple( 

194 _merge_window_spend_transactions(tuple(group)) for _, group in groupby(ordered, key=window_spend_group_key) 

195 )