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

1""" 

2Handles buffering database `UPDATE` transactions in Redis before committing them to the database 

3 

4This is to prevent deadlocks and improve reliability 

5""" 

6 

7import asyncio 

8import json 

9from collections.abc import Mapping, Sequence 

10from datetime import datetime 

11from functools import reduce 

12from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypeVar, cast 

13 

14from redis.exceptions import RedisError 

15 

16from litellm._logging import verbose_proxy_logger 

17from litellm.caching import RedisCache 

18from litellm.constants import ( 

19 MAX_REDIS_BUFFER_DEQUEUE_COUNT, 

20 REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, 

21 REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, 

22 REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, 

23 REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, 

24 REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, 

25 REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, 

26 REDIS_SPEND_LOGS_BUFFER_KEY, 

27 REDIS_SPEND_LOGS_BUFFER_MAX_ROWS, 

28 REDIS_UPDATE_BUFFER_KEY, 

29 REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, 

30) 

31from litellm.litellm_core_utils.safe_json_dumps import safe_dumps 

32from litellm.proxy._types import ( 

33 BaseDailySpendTransaction, 

34 DailyAgentSpendTransaction, 

35 DailyEndUserSpendTransaction, 

36 DailyOrganizationSpendTransaction, 

37 DailyTagSpendTransaction, 

38 DailyTeamSpendTransaction, 

39 DailyUserSpendTransaction, 

40 DBSpendUpdateTransactions, 

41 Litellm_EntityType, 

42 SpendUpdateQueueItem, 

43) 

44from litellm.proxy.db.db_transaction_queue.base_update_queue import service_logger_obj 

45from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( 

46 DailySpendUpdateQueue, 

47) 

48from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue 

49from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( 

50 WindowSpendTransaction, 

51 WindowSpendUpdateQueue, 

52 to_wire_payload, 

53) 

54from litellm.proxy.db.spend_log_batching import SpendLogRow 

55from litellm.secret_managers.main import str_to_bool 

56from litellm.types.caching import ( 

57 RedisPipelineLpopOperation, 

58 RedisPipelineRpushOperation, 

59) 

60from litellm.types.services import ServiceTypes 

61 

62if TYPE_CHECKING: 62 ↛ 63line 62 didn't jump to line 63 because the condition on line 62 was never true

63 from litellm.proxy.utils import PrismaClient 

64else: 

65 PrismaClient = Any 

66 

67BufferedSpendTransactions: TypeAlias = DBSpendUpdateTransactions | Mapping[str, BaseDailySpendTransaction] 

68 

69_SpendTransactionField: TypeAlias = Literal[ 

70 "user_list_transactions", 

71 "end_user_list_transactions", 

72 "key_list_transactions", 

73 "team_list_transactions", 

74 "team_member_list_transactions", 

75 "org_list_transactions", 

76 "org_member_list_transactions", 

77 "project_list_transactions", 

78 "tag_list_transactions", 

79 "agent_list_transactions", 

80 "model_access_group_list_transactions", 

81] 

82 

83_SPEND_TRANSACTION_FIELDS: Final[tuple[_SpendTransactionField, ...]] = ( 

84 "user_list_transactions", 

85 "end_user_list_transactions", 

86 "key_list_transactions", 

87 "team_list_transactions", 

88 "team_member_list_transactions", 

89 "org_list_transactions", 

90 "org_member_list_transactions", 

91 "project_list_transactions", 

92 "tag_list_transactions", 

93 "agent_list_transactions", 

94 "model_access_group_list_transactions", 

95) 

96 

97_ValueT = TypeVar("_ValueT") 

98 

99 

100def _spend_log_json_default(value: object) -> str: 

101 return value.isoformat() if isinstance(value, datetime) else str(value) 

102 

103 

104def _encode_spend_log_row(row: SpendLogRow) -> str: 

105 return json.dumps(row, default=_spend_log_json_default) 

106 

107 

108def _decode_spend_log_row(encoded: str) -> dict[str, object] | None: 

109 decoded: Final = json.loads(encoded) 

110 return decoded if isinstance(decoded, dict) else None 

111 

112 

113def _accumulated_spend(totals: Mapping[str, float], entities: Mapping[str, float]) -> dict[str, float]: 

114 return {**totals, **{entity_id: totals.get(entity_id, 0) + amount for entity_id, amount in entities.items()}} 

115 

116 

117def _entity_transactions(transaction: DBSpendUpdateTransactions, field: _SpendTransactionField) -> dict[str, float]: 

118 entities: Final[dict[str, float] | None] = transaction.get(field) 

119 return entities if isinstance(entities, dict) else {} 

120 

121 

122def _merged_entity_transactions( 

123 list_of_transactions: Sequence[DBSpendUpdateTransactions], 

124 field: _SpendTransactionField, 

125) -> dict[str, float]: 

126 return reduce( 

127 _accumulated_spend, 

128 (_entity_transactions(transaction, field) for transaction in list_of_transactions), 

129 {}, 

130 ) 

131 

132 

133class RedisUpdateBuffer: 

134 """ 

135 Handles buffering database `UPDATE` transactions in Redis before committing them to the database 

136 

137 This is to prevent deadlocks and improve reliability 

138 """ 

139 

140 def __init__( 

141 self, 

142 redis_cache: RedisCache | None = None, 

143 ): 

144 self.redis_cache = redis_cache 

145 

146 @staticmethod 

147 def _should_commit_spend_updates_to_redis() -> bool: 

148 """ 

149 Checks if the Pod should commit spend updates to Redis 

150 

151 This setting enables buffering database transactions in Redis 

152 to improve reliability and reduce database contention 

153 """ 

154 from litellm.proxy.proxy_server import general_settings 

155 

156 _use_redis_transaction_buffer: bool | str | None = general_settings.get("use_redis_transaction_buffer", False) 

157 if isinstance(_use_redis_transaction_buffer, str): 157 ↛ 158line 157 didn't jump to line 158 because the condition on line 157 was never true

158 _use_redis_transaction_buffer = str_to_bool(_use_redis_transaction_buffer) 

159 if _use_redis_transaction_buffer is None: 159 ↛ 160line 159 didn't jump to line 160 because the condition on line 159 was never true

160 return False 

161 return _use_redis_transaction_buffer 

162 

163 async def _store_transactions_in_redis( 

164 self, 

165 transactions: Mapping[str, BaseDailySpendTransaction] | None, 

166 redis_key: str, 

167 service_type: ServiceTypes, 

168 ) -> None: 

169 """ 

170 Helper method to store transactions in Redis and emit an event 

171 

172 Args: 

173 transactions: The transactions to store 

174 redis_key: The Redis key to store under 

175 service_type: The service type for event emission 

176 """ 

177 if transactions is None or len(transactions) == 0: 

178 return 

179 

180 list_of_transactions: Final = [safe_dumps(transactions)] 

181 if self.redis_cache is None: 

182 return 

183 try: 

184 current_redis_buffer_size: Final = await self.redis_cache.async_rpush( 

185 key=redis_key, 

186 values=list_of_transactions, 

187 ) 

188 verbose_proxy_logger.debug( 

189 "Spend tracking - pushed spend updates to Redis buffer. redis_key=%s, buffer_size=%s", 

190 redis_key, 

191 current_redis_buffer_size, 

192 ) 

193 await self._emit_new_item_added_to_redis_buffer_event( 

194 queue_size=current_redis_buffer_size, 

195 service=service_type, 

196 ) 

197 except Exception as e: 

198 verbose_proxy_logger.error( 

199 "Spend tracking - failed to push spend updates to Redis (redis_key=%s). Error: %s", 

200 redis_key, 

201 str(e), 

202 ) 

203 

204 async def store_in_memory_spend_updates_in_redis( 

205 self, 

206 spend_update_queue: SpendUpdateQueue, 

207 daily_spend_update_queue: DailySpendUpdateQueue, 

208 daily_team_spend_update_queue: DailySpendUpdateQueue, 

209 daily_org_spend_update_queue: DailySpendUpdateQueue, 

210 daily_end_user_spend_update_queue: DailySpendUpdateQueue, 

211 daily_agent_spend_update_queue: DailySpendUpdateQueue, 

212 window_spend_update_queue: WindowSpendUpdateQueue | None = None, 

213 ): 

214 """ 

215 Stores the in-memory spend updates to Redis 

216 

217 Stores the following in memory data structures in Redis: 

218 - SpendUpdateQueue - Key, User, Team, TeamMember, Org, EndUser Spend updates 

219 - DailySpendUpdateQueue - Daily Spend updates Aggregate view 

220 

221 For SpendUpdateQueue: 

222 Each transaction is a dict stored as following: 

223 - key is the entity id 

224 - value is the spend amount 

225 

226 ``` 

227 Redis List: 

228 key_list_transactions: 

229 [ 

230 "0929880201": 1.2, 

231 "0929880202": 0.01, 

232 "0929880203": 0.001, 

233 ] 

234 ``` 

235 

236 For DailySpendUpdateQueue: 

237 Each transaction is a Dict[str, DailyUserSpendTransaction] stored as following: 

238 - key is the daily_transaction_key 

239 - value is the DailyUserSpendTransaction 

240 

241 ``` 

242 Redis List: 

243 daily_spend_update_transactions: 

244 [ 

245 { 

246 "user_keyhash_1_model_1": { 

247 "spend": 1.2, 

248 "prompt_tokens": 1000, 

249 "completion_tokens": 1000, 

250 "api_requests": 1000, 

251 "successful_requests": 1000, 

252 }, 

253 

254 } 

255 ] 

256 ``` 

257 """ 

258 if self.redis_cache is None: 

259 verbose_proxy_logger.debug("redis_cache is None, skipping store_in_memory_spend_updates_in_redis") 

260 return 

261 

262 # Get all transactions 

263 db_spend_update_transactions: Final = ( 

264 await spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() 

265 ) 

266 daily_spend_update_transactions: Final = ( 

267 await daily_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() 

268 ) 

269 daily_team_spend_update_transactions: Final = ( 

270 await daily_team_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() 

271 ) 

272 daily_org_spend_update_transactions: Final = ( 

273 await daily_org_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() 

274 ) 

275 daily_end_user_spend_update_transactions: Final = ( 

276 await daily_end_user_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() 

277 ) 

278 daily_agent_spend_update_transactions: Final = ( 

279 await daily_agent_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() 

280 ) 

281 window_spend_update_transactions: Final = ( 

282 await window_spend_update_queue.flush_and_get_aggregated_window_spend_transactions() 

283 if window_spend_update_queue is not None 

284 else () 

285 ) 

286 

287 verbose_proxy_logger.debug("ALL DB SPEND UPDATE TRANSACTIONS: %s", db_spend_update_transactions) 

288 verbose_proxy_logger.debug("ALL DAILY SPEND UPDATE TRANSACTIONS: %s", daily_spend_update_transactions) 

289 

290 # Build a list of rpush operations, skipping empty/None transaction sets 

291 _queue_configs: Final[list[tuple[BufferedSpendTransactions | None, str, ServiceTypes]]] = [ 

292 ( 

293 db_spend_update_transactions, 

294 REDIS_UPDATE_BUFFER_KEY, 

295 ServiceTypes.REDIS_SPEND_UPDATE_QUEUE, 

296 ), 

297 ( 

298 daily_spend_update_transactions, 

299 REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, 

300 ServiceTypes.REDIS_DAILY_SPEND_UPDATE_QUEUE, 

301 ), 

302 ( 

303 daily_team_spend_update_transactions, 

304 REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, 

305 ServiceTypes.REDIS_DAILY_TEAM_SPEND_UPDATE_QUEUE, 

306 ), 

307 ( 

308 daily_org_spend_update_transactions, 

309 REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, 

310 ServiceTypes.REDIS_DAILY_ORG_SPEND_UPDATE_QUEUE, 

311 ), 

312 ( 

313 daily_end_user_spend_update_transactions, 

314 REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, 

315 ServiceTypes.REDIS_DAILY_END_USER_SPEND_UPDATE_QUEUE, 

316 ), 

317 ( 

318 daily_agent_spend_update_transactions, 

319 REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, 

320 ServiceTypes.REDIS_DAILY_AGENT_SPEND_UPDATE_QUEUE, 

321 ), 

322 ( 

323 tuple(map(to_wire_payload, window_spend_update_transactions)), 

324 REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, 

325 ServiceTypes.REDIS_WINDOW_SPEND_UPDATE_QUEUE, 

326 ), 

327 ] 

328 

329 rpush_list: Final[list[RedisPipelineRpushOperation]] = [] 

330 service_types: Final[list[ServiceTypes]] = [] 

331 for transactions, redis_key, service_type in _queue_configs: 

332 if transactions is None or len(transactions) == 0: 

333 continue 

334 rpush_list.append( 

335 RedisPipelineRpushOperation( 

336 key=redis_key, 

337 values=[safe_dumps(transactions)], 

338 ) 

339 ) 

340 service_types.append(service_type) 

341 

342 if len(rpush_list) == 0: 

343 return 

344 

345 try: 

346 result_lengths: Final = await self.redis_cache.async_rpush_pipeline( 

347 rpush_list=rpush_list, 

348 ) 

349 except Exception as e: 

350 # The in-memory queues were already drained above. If we let the 

351 # exception propagate without restoring, the aggregated spend is 

352 # permanently lost. Re-enqueue so the next scheduler tick retries. 

353 verbose_proxy_logger.error( 

354 "Spend tracking - failed to push aggregated spend updates to Redis. " 

355 "Restoring %d transaction sets to in-memory queues for retry on next tick. " 

356 "Error: %s", 

357 len(rpush_list), 

358 str(e), 

359 ) 

360 await self._restore_spend_updates_to_in_memory_queues( 

361 db_spend_update_transactions=db_spend_update_transactions, 

362 daily_spend_update_transactions=daily_spend_update_transactions, 

363 daily_team_spend_update_transactions=daily_team_spend_update_transactions, 

364 daily_org_spend_update_transactions=daily_org_spend_update_transactions, 

365 daily_end_user_spend_update_transactions=daily_end_user_spend_update_transactions, 

366 daily_agent_spend_update_transactions=daily_agent_spend_update_transactions, 

367 window_spend_update_transactions=window_spend_update_transactions, 

368 spend_update_queue=spend_update_queue, 

369 daily_spend_update_queue=daily_spend_update_queue, 

370 daily_team_spend_update_queue=daily_team_spend_update_queue, 

371 daily_org_spend_update_queue=daily_org_spend_update_queue, 

372 daily_end_user_spend_update_queue=daily_end_user_spend_update_queue, 

373 daily_agent_spend_update_queue=daily_agent_spend_update_queue, 

374 window_spend_update_queue=window_spend_update_queue, 

375 ) 

376 return 

377 

378 # Emit gauge events for each queue 

379 for i, queue_size in enumerate(result_lengths): 

380 if i < len(service_types): 

381 await self._emit_new_item_added_to_redis_buffer_event( 

382 queue_size=queue_size, 

383 service=service_types[i], 

384 ) 

385 

386 @staticmethod 

387 async def _restore_spend_updates_to_in_memory_queues( 

388 db_spend_update_transactions: DBSpendUpdateTransactions | None, 

389 daily_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None, 

390 daily_team_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None, 

391 daily_org_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None, 

392 daily_end_user_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None, 

393 daily_agent_spend_update_transactions: dict[str, BaseDailySpendTransaction] | None, 

394 window_spend_update_transactions: tuple[WindowSpendTransaction, ...] | None, 

395 spend_update_queue: SpendUpdateQueue, 

396 daily_spend_update_queue: DailySpendUpdateQueue, 

397 daily_team_spend_update_queue: DailySpendUpdateQueue, 

398 daily_org_spend_update_queue: DailySpendUpdateQueue, 

399 daily_end_user_spend_update_queue: DailySpendUpdateQueue, 

400 daily_agent_spend_update_queue: DailySpendUpdateQueue, 

401 window_spend_update_queue: WindowSpendUpdateQueue | None, 

402 ) -> None: 

403 """ 

404 Put drained-but-unpushed transactions back into in-memory queues. 

405 

406 Called when the Redis rpush pipeline raises. Without this, all spend 

407 data aggregated during the current scheduler tick is permanently lost 

408 because the source queues were already drained before the rpush. 

409 """ 

410 if db_spend_update_transactions is not None: 

411 entity_entries: Final[list[tuple[Litellm_EntityType, dict[str, float] | None]]] = [ 

412 ( 

413 Litellm_EntityType.USER, 

414 db_spend_update_transactions.get("user_list_transactions"), 

415 ), 

416 ( 

417 Litellm_EntityType.END_USER, 

418 db_spend_update_transactions.get("end_user_list_transactions"), 

419 ), 

420 ( 

421 Litellm_EntityType.KEY, 

422 db_spend_update_transactions.get("key_list_transactions"), 

423 ), 

424 ( 

425 Litellm_EntityType.TEAM, 

426 db_spend_update_transactions.get("team_list_transactions"), 

427 ), 

428 ( 

429 Litellm_EntityType.TEAM_MEMBER, 

430 db_spend_update_transactions.get("team_member_list_transactions"), 

431 ), 

432 ( 

433 Litellm_EntityType.ORGANIZATION, 

434 db_spend_update_transactions.get("org_list_transactions"), 

435 ), 

436 ( 

437 Litellm_EntityType.ORGANIZATION_MEMBER, 

438 db_spend_update_transactions.get("org_member_list_transactions"), 

439 ), 

440 ( 

441 Litellm_EntityType.PROJECT, 

442 db_spend_update_transactions.get("project_list_transactions"), 

443 ), 

444 ( 

445 Litellm_EntityType.TAG, 

446 db_spend_update_transactions.get("tag_list_transactions"), 

447 ), 

448 ( 

449 Litellm_EntityType.AGENT, 

450 db_spend_update_transactions.get("agent_list_transactions"), 

451 ), 

452 ( 

453 Litellm_EntityType.MODEL_ACCESS_GROUP, 

454 db_spend_update_transactions.get("model_access_group_list_transactions"), 

455 ), 

456 ] 

457 for entity_type, entities in entity_entries: 

458 if not entities: 

459 continue 

460 for entity_id, cost in entities.items(): 

461 await spend_update_queue.add_update( 

462 SpendUpdateQueueItem( 

463 entity_type=entity_type, 

464 entity_id=entity_id, 

465 response_cost=cost, 

466 ) 

467 ) 

468 

469 daily_pairs: Final[list[tuple[dict[str, BaseDailySpendTransaction] | None, DailySpendUpdateQueue]]] = [ 

470 (daily_spend_update_transactions, daily_spend_update_queue), 

471 (daily_team_spend_update_transactions, daily_team_spend_update_queue), 

472 (daily_org_spend_update_transactions, daily_org_spend_update_queue), 

473 ( 

474 daily_end_user_spend_update_transactions, 

475 daily_end_user_spend_update_queue, 

476 ), 

477 (daily_agent_spend_update_transactions, daily_agent_spend_update_queue), 

478 ] 

479 for daily_txns, daily_queue in daily_pairs: 

480 if daily_txns: 

481 await daily_queue.update_queue.put(daily_txns) 

482 

483 if window_spend_update_transactions and window_spend_update_queue is not None: 

484 await window_spend_update_queue.update_queue.put(window_spend_update_transactions) 

485 

486 async def restore_transactions_to_redis( 

487 self, 

488 db_spend_update_transactions: DBSpendUpdateTransactions | None = None, 

489 daily_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None, 

490 daily_team_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None, 

491 daily_org_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None, 

492 daily_end_user_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None, 

493 daily_agent_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None, 

494 daily_tag_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None, 

495 window_spend_update_transactions: Sequence[WindowSpendTransaction] | None = None, 

496 ) -> None: 

497 """ 

498 Re-push transactions that were popped from Redis but not committed to the DB. 

499 

500 The leader drains the buffers with a destructive ``lpop`` before committing to 

501 the database. When a commit fails after its retries are exhausted, the popped 

502 transactions must be pushed back so a later scheduler tick can retry them; 

503 otherwise the aggregated spend is lost permanently. The re-pushed payloads use 

504 the same JSON encoding as the store path, so the next drain parses them normally. 

505 """ 

506 if self.redis_cache is None: 

507 return 

508 

509 restore_configs: Final = ( 

510 (db_spend_update_transactions, REDIS_UPDATE_BUFFER_KEY), 

511 (daily_spend_update_transactions, REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY), 

512 (daily_team_spend_update_transactions, REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY), 

513 (daily_org_spend_update_transactions, REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY), 

514 (daily_end_user_spend_update_transactions, REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY), 

515 (daily_agent_spend_update_transactions, REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY), 

516 (daily_tag_spend_update_transactions, REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY), 

517 ( 

518 None 

519 if window_spend_update_transactions is None 

520 else tuple(map(to_wire_payload, window_spend_update_transactions)), 

521 REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, 

522 ), 

523 ) 

524 

525 rpush_list: Final = tuple( 

526 RedisPipelineRpushOperation(key=redis_key, values=(safe_dumps(transactions),)) 

527 for transactions, redis_key in restore_configs 

528 if transactions 

529 ) 

530 if len(rpush_list) == 0: 

531 return 

532 

533 try: 

534 await self.redis_cache.async_rpush_pipeline(rpush_list=rpush_list) 

535 verbose_proxy_logger.info( 

536 "Spend tracking - restored %d uncommitted transaction set(s) to Redis for retry on next tick.", 

537 len(rpush_list), 

538 ) 

539 except RedisError as e: 

540 verbose_proxy_logger.error( 

541 "Spend tracking - failed to restore uncommitted transactions to Redis. " 

542 "These spend updates are lost. Error: %s", 

543 str(e), 

544 ) 

545 

546 async def store_spend_logs_in_redis( 

547 self, 

548 rows: Sequence[SpendLogRow], 

549 max_rows: int = REDIS_SPEND_LOGS_BUFFER_MAX_ROWS, 

550 ) -> bool: 

551 """Park spend-log rows in Redis so they outlive this pod, dropping the oldest past ``max_rows``.""" 

552 if self.redis_cache is None or len(rows) == 0 or not self._should_commit_spend_updates_to_redis(): 

553 return False 

554 try: 

555 buffer_size: Final = await self.redis_cache.async_rpush_and_trim( 

556 key=REDIS_SPEND_LOGS_BUFFER_KEY, 

557 values=tuple(_encode_spend_log_row(row) for row in rows), 

558 max_len=max_rows, 

559 ) 

560 overflow: Final = buffer_size - max_rows 

561 if overflow > 0: 

562 verbose_proxy_logger.error( 

563 "Spend tracking - Redis spend log buffer is at its %d row cap; dropped the %d oldest spend logs", 

564 max_rows, 

565 overflow, 

566 ) 

567 except Exception as e: # noqa: BLE001 # the caller falls back to the in-memory queue on any Redis fault 

568 verbose_proxy_logger.error( 

569 "Spend tracking - failed to park %d spend log rows in Redis. Error: %s", len(rows), str(e) 

570 ) 

571 return False 

572 verbose_proxy_logger.info("Spend tracking - parked %d spend log rows in Redis for a later flush", len(rows)) 

573 return True 

574 

575 async def get_spend_logs_from_redis_buffer(self, limit: int) -> tuple[dict[str, object], ...]: 

576 """Atomically take up to ``limit`` parked spend-log rows out of Redis.""" 

577 if self.redis_cache is None or not self._should_commit_spend_updates_to_redis(): 577 ↛ 579line 577 didn't jump to line 579 because the condition on line 577 was always true

578 return () 

579 popped: Final[str | list[str] | None] = await self.redis_cache.async_lpop( 

580 key=REDIS_SPEND_LOGS_BUFFER_KEY, 

581 count=limit, 

582 ) 

583 if popped is None: 

584 return () 

585 encoded_rows: Final = tuple(popped) if isinstance(popped, list) else (popped,) 

586 decoded_rows: Final = (_decode_spend_log_row(encoded) for encoded in encoded_rows) 

587 return tuple(row for row in decoded_rows if row is not None) 

588 

589 @staticmethod 

590 def _number_of_transactions_to_store_in_redis( 

591 db_spend_update_transactions: DBSpendUpdateTransactions, 

592 ) -> int: 

593 """ 

594 Gets the number of transactions to store in Redis 

595 """ 

596 return sum( 

597 len(_entity_transactions(db_spend_update_transactions, field)) for field in _SPEND_TRANSACTION_FIELDS 

598 ) 

599 

600 @staticmethod 

601 def _remove_prefix_from_keys(data: Mapping[str, _ValueT], prefix: str) -> dict[str, _ValueT]: 

602 """ 

603 Removes the specified prefix from the keys of a dictionary. 

604 """ 

605 return {key.replace(prefix, "", 1): value for key, value in data.items()} 

606 

607 async def get_all_update_transactions_from_redis_buffer( 

608 self, 

609 ) -> DBSpendUpdateTransactions | None: 

610 """ 

611 Gets all the update transactions from Redis 

612 

613 On Redis we store a list of transactions as a JSON string 

614 

615 eg. 

616 [ 

617 DBSpendUpdateTransactions( 

618 user_list_transactions={ 

619 "user_id_1": 1.2, 

620 "user_id_2": 0.01, 

621 }, 

622 end_user_list_transactions={}, 

623 key_list_transactions={ 

624 "0929880201": 1.2, 

625 "0929880202": 0.01, 

626 }, 

627 team_list_transactions={}, 

628 team_member_list_transactions={}, 

629 org_list_transactions={}, 

630 ), 

631 DBSpendUpdateTransactions( 

632 user_list_transactions={ 

633 "user_id_3": 1.2, 

634 "user_id_4": 0.01, 

635 }, 

636 end_user_list_transactions={}, 

637 key_list_transactions={ 

638 "key_id_1": 1.2, 

639 "key_id_2": 0.01, 

640 }, 

641 team_list_transactions={}, 

642 team_member_list_transactions={}, 

643 org_list_transactions={}, 

644 ] 

645 """ 

646 if self.redis_cache is None: 

647 return None 

648 list_of_transactions: Final[str | list[str] | None] = await self.redis_cache.async_lpop( 

649 key=REDIS_UPDATE_BUFFER_KEY, 

650 count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, 

651 ) 

652 if list_of_transactions is None: 

653 return None 

654 

655 verbose_proxy_logger.info( 

656 "Spend tracking - popped %d spend update batches from Redis buffer (key=%s). " 

657 "These items are now removed from Redis and must be committed to DB.", 

658 len(list_of_transactions) if isinstance(list_of_transactions, list) else 1, 

659 REDIS_UPDATE_BUFFER_KEY, 

660 ) 

661 

662 # Parse the list of transactions from JSON strings 

663 parsed_transactions: Final = self._parse_list_of_transactions(list_of_transactions) 

664 

665 # If there are no transactions, return None 

666 if len(parsed_transactions) == 0: 

667 return None 

668 

669 # Combine all transactions into a single transaction 

670 combined_transaction: Final = self._combine_list_of_transactions(parsed_transactions) 

671 

672 return combined_transaction 

673 

674 async def get_all_transactions_from_redis_buffer_pipeline( 

675 self, 

676 ) -> tuple[ 

677 DBSpendUpdateTransactions | None, 

678 dict[str, DailyUserSpendTransaction] | None, 

679 dict[str, DailyTeamSpendTransaction] | None, 

680 dict[str, DailyOrganizationSpendTransaction] | None, 

681 dict[str, DailyEndUserSpendTransaction] | None, 

682 dict[str, DailyAgentSpendTransaction] | None, 

683 tuple[WindowSpendTransaction, ...] | None, 

684 ]: 

685 """ 

686 Drains the main 7 Redis buffer queues in a single pipeline round-trip. 

687 

688 Returns a 7-tuple of parsed results in this order: 

689 0: DBSpendUpdateTransactions 

690 1: daily user spend 

691 2: daily team spend 

692 3: daily org spend 

693 4: daily end-user spend 

694 5: daily agent spend 

695 6: budget window spend 

696 """ 

697 if self.redis_cache is None: 

698 return None, None, None, None, None, None, None 

699 

700 lpop_list: Final[list[RedisPipelineLpopOperation]] = [ 

701 RedisPipelineLpopOperation(key=REDIS_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT), 

702 RedisPipelineLpopOperation( 

703 key=REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, 

704 count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, 

705 ), 

706 RedisPipelineLpopOperation( 

707 key=REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, 

708 count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, 

709 ), 

710 RedisPipelineLpopOperation( 

711 key=REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, 

712 count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, 

713 ), 

714 RedisPipelineLpopOperation( 

715 key=REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, 

716 count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, 

717 ), 

718 RedisPipelineLpopOperation( 

719 key=REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, 

720 count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, 

721 ), 

722 RedisPipelineLpopOperation( 

723 key=REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, 

724 count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, 

725 ), 

726 ] 

727 

728 raw_results: Final = await self.redis_cache.async_lpop_pipeline(lpop_list=lpop_list) 

729 

730 # Pad with None if pipeline returned fewer results than expected 

731 while len(raw_results) < 7: 

732 raw_results.append(None) 

733 

734 # Slot 0: DBSpendUpdateTransactions 

735 db_spend: DBSpendUpdateTransactions | None = None 

736 if raw_results[0] is not None: 

737 parsed: Final = self._parse_list_of_transactions(raw_results[0]) 

738 if len(parsed) > 0: 

739 db_spend = self._combine_list_of_transactions(parsed) 

740 

741 # Slots 1-5: daily spend categories 

742 daily_results: Final[list[dict[str, BaseDailySpendTransaction] | None]] = [] 

743 for slot in range(1, 6): 

744 slot_result = raw_results[slot] 

745 if slot_result is None: 

746 daily_results.append(None) 

747 else: 

748 list_of_daily = [json.loads(t) for t in slot_result] 

749 aggregated = DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions(list_of_daily) 

750 daily_results.append(aggregated) 

751 

752 window_spend: Final = ( 

753 WindowSpendUpdateQueue.get_aggregated_window_spend_transactions( 

754 tuple(json.loads(transaction) for transaction in raw_results[6]) 

755 ) 

756 if raw_results[6] is not None 

757 else None 

758 ) 

759 

760 return ( 

761 db_spend, 

762 cast(dict[str, DailyUserSpendTransaction] | None, daily_results[0]), 

763 cast(dict[str, DailyTeamSpendTransaction] | None, daily_results[1]), 

764 cast(dict[str, DailyOrganizationSpendTransaction] | None, daily_results[2]), 

765 cast(dict[str, DailyEndUserSpendTransaction] | None, daily_results[3]), 

766 cast(dict[str, DailyAgentSpendTransaction] | None, daily_results[4]), 

767 window_spend, 

768 ) 

769 

770 async def store_in_memory_daily_tag_spend_updates_in_redis( 

771 self, 

772 daily_tag_spend_update_queue: DailySpendUpdateQueue, 

773 ) -> None: 

774 """ 

775 Flush in-memory daily tag spend updates and append them to Redis. 

776 """ 

777 daily_tag_spend_update_transactions: Final = ( 

778 await daily_tag_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() 

779 ) 

780 await self._store_transactions_in_redis( 

781 transactions=daily_tag_spend_update_transactions, 

782 redis_key=REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, 

783 service_type=ServiceTypes.REDIS_DAILY_TAG_SPEND_UPDATE_QUEUE, 

784 ) 

785 

786 async def _lpop_daily_spend_transactions( 

787 self, 

788 redis_key: str, 

789 ) -> list[dict[str, BaseDailySpendTransaction]] | None: 

790 """ 

791 Drains a daily spend buffer key and parses each popped item as JSON. 

792 """ 

793 if self.redis_cache is None: 

794 return None 

795 list_of_transactions: Final[list[str] | None] = await self.redis_cache.async_lpop( 

796 key=redis_key, 

797 count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, 

798 ) 

799 if list_of_transactions is None: 

800 return None 

801 return [json.loads(transaction) for transaction in list_of_transactions] 

802 

803 async def get_all_daily_spend_update_transactions_from_redis_buffer( 

804 self, 

805 ) -> dict[str, DailyUserSpendTransaction] | None: 

806 """ 

807 Gets all the daily spend update transactions from Redis 

808 """ 

809 if self.redis_cache is None: 

810 return None 

811 list_of_daily_spend_update_transactions: Final = await self._lpop_daily_spend_transactions( 

812 REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY 

813 ) 

814 if list_of_daily_spend_update_transactions is None: 

815 return None 

816 return cast( 

817 dict[str, DailyUserSpendTransaction], 

818 DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( 

819 list_of_daily_spend_update_transactions 

820 ), 

821 ) 

822 

823 async def get_all_daily_team_spend_update_transactions_from_redis_buffer( 

824 self, 

825 ) -> dict[str, DailyTeamSpendTransaction] | None: 

826 """ 

827 Gets all the daily team spend update transactions from Redis 

828 """ 

829 if self.redis_cache is None: 

830 return None 

831 list_of_daily_spend_update_transactions: Final = await self._lpop_daily_spend_transactions( 

832 REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY 

833 ) 

834 if list_of_daily_spend_update_transactions is None: 

835 return None 

836 return cast( 

837 dict[str, DailyTeamSpendTransaction], 

838 DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( 

839 list_of_daily_spend_update_transactions 

840 ), 

841 ) 

842 

843 async def get_all_daily_org_spend_update_transactions_from_redis_buffer( 

844 self, 

845 ) -> dict[str, DailyOrganizationSpendTransaction] | None: 

846 """ 

847 Gets all the daily organization spend update transactions from Redis 

848 """ 

849 if self.redis_cache is None: 

850 return None 

851 list_of_daily_spend_update_transactions: Final = await self._lpop_daily_spend_transactions( 

852 REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY 

853 ) 

854 if list_of_daily_spend_update_transactions is None: 

855 return None 

856 return cast( 

857 dict[str, DailyOrganizationSpendTransaction], 

858 DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( 

859 list_of_daily_spend_update_transactions 

860 ), 

861 ) 

862 

863 async def get_all_daily_end_user_spend_update_transactions_from_redis_buffer( 

864 self, 

865 ) -> dict[str, DailyEndUserSpendTransaction] | None: 

866 """ 

867 Gets all the daily end-user spend update transactions from Redis 

868 """ 

869 if self.redis_cache is None: 

870 return None 

871 list_of_daily_spend_update_transactions: Final = await self._lpop_daily_spend_transactions( 

872 REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY 

873 ) 

874 if list_of_daily_spend_update_transactions is None: 

875 return None 

876 return cast( 

877 dict[str, DailyEndUserSpendTransaction], 

878 DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( 

879 list_of_daily_spend_update_transactions 

880 ), 

881 ) 

882 

883 async def get_all_daily_agent_spend_update_transactions_from_redis_buffer( 

884 self, 

885 ) -> dict[str, DailyAgentSpendTransaction] | None: 

886 """ 

887 Gets all the daily agent spend update transactions from Redis 

888 """ 

889 if self.redis_cache is None: 

890 return None 

891 list_of_daily_spend_update_transactions: Final = await self._lpop_daily_spend_transactions( 

892 REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY 

893 ) 

894 if list_of_daily_spend_update_transactions is None: 

895 return None 

896 return cast( 

897 dict[str, DailyAgentSpendTransaction], 

898 DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( 

899 list_of_daily_spend_update_transactions 

900 ), 

901 ) 

902 

903 async def get_all_daily_tag_spend_update_transactions_from_redis_buffer(