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

1import asyncio 

2from typing import Final 

3 

4from litellm._logging import verbose_proxy_logger 

5from litellm.constants import LITELLM_ASYNCIO_QUEUE_MAXSIZE 

6from litellm.proxy._types import ( 

7 DBSpendUpdateTransactions, 

8 Litellm_EntityType, 

9 SpendUpdateQueueItem, 

10) 

11from litellm.proxy.db.db_transaction_queue.base_update_queue import ( 

12 BaseUpdateQueue, 

13 service_logger_obj, 

14) 

15from litellm.types.services import ServiceTypes 

16 

17 

18class SpendUpdateQueue(BaseUpdateQueue): 

19 """ 

20 In memory buffer for spend updates that should be committed to the database 

21 """ 

22 

23 def __init__(self): 

24 super().__init__() 

25 self.update_queue: asyncio.Queue[SpendUpdateQueueItem] = asyncio.Queue(maxsize=LITELLM_ASYNCIO_QUEUE_MAXSIZE) 

26 

27 async def flush_and_get_aggregated_db_spend_update_transactions( 

28 self, 

29 ) -> DBSpendUpdateTransactions: 

30 """Flush all updates from the queue and return all updates aggregated by entity type.""" 

31 updates: Final = await self.flush_all_updates_from_in_memory_queue() 

32 if len(updates) > 0: 

33 verbose_proxy_logger.info( 

34 "Spend tracking - flushed %d spend update items from in-memory queue", 

35 len(updates), 

36 ) 

37 verbose_proxy_logger.debug("Aggregating updates by entity type: %s", updates) 

38 return self.get_aggregated_db_spend_update_transactions(updates) 

39 

40 async def add_update(self, update: SpendUpdateQueueItem): 

41 """Enqueue an update to the spend update queue""" 

42 verbose_proxy_logger.debug("Adding update to queue: %s", update) 

43 await self.update_queue.put(update) 

44 

45 # if the queue is full, aggregate the updates 

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

47 verbose_proxy_logger.warning( 

48 "Spend update queue is full. Aggregating all entries in queue to concatenate entries." 

49 ) 

50 await self.aggregate_queue_updates() 

51 

52 async def aggregate_queue_updates(self): 

53 """Concatenate all updates in the queue to reduce the size of in-memory queue""" 

54 updates: Final[list[SpendUpdateQueueItem]] = await self.flush_all_updates_from_in_memory_queue() 

55 aggregated_updates: Final = self._get_aggregated_spend_update_queue_item(updates) 

56 for update in aggregated_updates: 

57 await self.update_queue.put(update) 

58 

59 def _get_aggregated_spend_update_queue_item( 

60 self, updates: list[SpendUpdateQueueItem] 

61 ) -> list[SpendUpdateQueueItem]: 

62 """ 

63 This is used to reduce the size of the in-memory queue by aggregating updates by entity type + id 

64 

65 

66 Aggregate updates by entity type + id 

67 

68 eg. 

69 

70 ``` 

71 [ 

72 { 

73 "entity_type": "user", 

74 "entity_id": "123", 

75 "response_cost": 100 

76 }, 

77 { 

78 "entity_type": "user", 

79 "entity_id": "123", 

80 "response_cost": 200 

81 } 

82 ] 

83 

84 ``` 

85 

86 becomes 

87 

88 ``` 

89 

90 [ 

91 { 

92 "entity_type": "user", 

93 "entity_id": "123", 

94 "response_cost": 300 

95 } 

96 ] 

97 

98 ``` 

99 """ 

100 verbose_proxy_logger.debug( 

101 "Aggregating spend updates, current queue size: %s", 

102 self.update_queue.qsize(), 

103 ) 

104 aggregated_spend_updates: Final[list[SpendUpdateQueueItem]] = [] 

105 

106 _in_memory_map: Final[dict[str, SpendUpdateQueueItem]] = {} 

107 """ 

108 Used for combining several updates into a single update 

109 Key=entity_type:entity_id 

110 Value=SpendUpdateQueueItem 

111 """ 

112 for update in updates: 

113 _key = f"{update.get('entity_type')}:{update.get('entity_id')}" 

114 if _key not in _in_memory_map: 

115 # avoid mutating caller-owned dicts while aggregating queue entries 

116 _in_memory_map[_key] = update.copy() 

117 else: 

118 current_cost = _in_memory_map[_key].get("response_cost", 0) or 0 

119 update_cost = update.get("response_cost", 0) or 0 

120 _in_memory_map[_key]["response_cost"] = current_cost + update_cost 

121 

122 for _key, update in _in_memory_map.items(): 

123 aggregated_spend_updates.append(update) 

124 

125 verbose_proxy_logger.debug("Aggregated spend updates: %s", aggregated_spend_updates) 

126 return aggregated_spend_updates 

127 

128 def get_aggregated_db_spend_update_transactions( 

129 self, updates: list[SpendUpdateQueueItem] 

130 ) -> DBSpendUpdateTransactions: 

131 """Aggregate updates by entity type.""" 

132 # Initialize all transaction lists as empty dicts 

133 db_spend_update_transactions: Final = DBSpendUpdateTransactions( 

134 user_list_transactions={}, 

135 end_user_list_transactions={}, 

136 key_list_transactions={}, 

137 team_list_transactions={}, 

138 team_member_list_transactions={}, 

139 org_list_transactions={}, 

140 org_member_list_transactions={}, 

141 project_list_transactions={}, 

142 tag_list_transactions={}, 

143 agent_list_transactions={}, 

144 model_access_group_list_transactions={}, 

145 ) 

146 

147 # Map entity types to their corresponding transaction dictionary keys 

148 entity_type_to_dict_key: Final = { 

149 Litellm_EntityType.USER: "user_list_transactions", 

150 Litellm_EntityType.END_USER: "end_user_list_transactions", 

151 Litellm_EntityType.KEY: "key_list_transactions", 

152 Litellm_EntityType.TEAM: "team_list_transactions", 

153 Litellm_EntityType.TEAM_MEMBER: "team_member_list_transactions", 

154 Litellm_EntityType.ORGANIZATION: "org_list_transactions", 

155 Litellm_EntityType.ORGANIZATION_MEMBER: "org_member_list_transactions", 

156 Litellm_EntityType.PROJECT: "project_list_transactions", 

157 Litellm_EntityType.TAG: "tag_list_transactions", 

158 Litellm_EntityType.AGENT: "agent_list_transactions", 

159 Litellm_EntityType.MODEL_ACCESS_GROUP: "model_access_group_list_transactions", 

160 } 

161 

162 for update in updates: 

163 entity_type = update.get("entity_type") 

164 entity_id = update.get("entity_id") or "" 

165 response_cost = update.get("response_cost") or 0 

166 

167 if entity_type is None: 167 ↛ 168line 167 didn't jump to line 168 because the condition on line 167 was never true

168 verbose_proxy_logger.debug( 

169 "Skipping update spend for update: %s, because entity_type is None", 

170 update, 

171 ) 

172 continue 

173 

174 dict_key = entity_type_to_dict_key.get(entity_type) 

175 if dict_key is None: 175 ↛ 176line 175 didn't jump to line 176 because the condition on line 175 was never true

176 verbose_proxy_logger.debug( 

177 "Skipping update spend for update: %s, because entity_type is not in entity_type_to_dict_key", 

178 update, 

179 ) 

180 continue # Skip unknown entity types 

181 

182 # Type-safe access using if/elif statements 

183 if dict_key == "user_list_transactions": 

184 transactions_dict = db_spend_update_transactions["user_list_transactions"] 

185 elif dict_key == "end_user_list_transactions": 

186 transactions_dict = db_spend_update_transactions["end_user_list_transactions"] 

187 elif dict_key == "key_list_transactions": 

188 transactions_dict = db_spend_update_transactions["key_list_transactions"] 

189 elif dict_key == "team_list_transactions": 189 ↛ 190line 189 didn't jump to line 190 because the condition on line 189 was never true

190 transactions_dict = db_spend_update_transactions["team_list_transactions"] 

191 elif dict_key == "team_member_list_transactions": 191 ↛ 192line 191 didn't jump to line 192 because the condition on line 191 was never true

192 transactions_dict = db_spend_update_transactions["team_member_list_transactions"] 

193 elif dict_key == "org_list_transactions": 193 ↛ 194line 193 didn't jump to line 194 because the condition on line 193 was never true

194 transactions_dict = db_spend_update_transactions["org_list_transactions"] 

195 elif dict_key == "org_member_list_transactions": 195 ↛ 196line 195 didn't jump to line 196 because the condition on line 195 was never true

196 transactions_dict = db_spend_update_transactions["org_member_list_transactions"] 

197 elif dict_key == "project_list_transactions": 197 ↛ 198line 197 didn't jump to line 198 because the condition on line 197 was never true

198 transactions_dict = db_spend_update_transactions["project_list_transactions"] 

199 elif dict_key == "tag_list_transactions": 

200 transactions_dict = db_spend_update_transactions["tag_list_transactions"] 

201 elif dict_key == "agent_list_transactions": 201 ↛ 203line 201 didn't jump to line 203 because the condition on line 201 was always true

202 transactions_dict = db_spend_update_transactions["agent_list_transactions"] 

203 elif dict_key == "model_access_group_list_transactions": 

204 transactions_dict = db_spend_update_transactions["model_access_group_list_transactions"] 

205 else: 

206 continue 

207 

208 if transactions_dict is None: 208 ↛ 209line 208 didn't jump to line 209 because the condition on line 208 was never true

209 transactions_dict = {} 

210 

211 # type ignore: dict_key is guaranteed to be one of "one of ("user_list_transactions", "end_user_list_transactions", "key_list_transactions", "team_list_transactions", "team_member_list_transactions", "org_list_transactions")" 

212 db_spend_update_transactions[dict_key] = transactions_dict 

213 

214 if entity_id not in transactions_dict: 

215 transactions_dict[entity_id] = 0 

216 

217 transactions_dict[entity_id] += response_cost or 0 

218 

219 return db_spend_update_transactions 

220 

221 async def _emit_new_item_added_to_queue_event( 

222 self, 

223 queue_size: int | None = None, 

224 ): 

225 asyncio.create_task( 

226 service_logger_obj.async_service_success_hook( 

227 service=ServiceTypes.IN_MEMORY_SPEND_UPDATE_QUEUE, 

228 duration=0, 

229 call_type="_emit_new_item_added_to_queue_event", 

230 event_metadata={ 

231 "gauge_labels": ServiceTypes.IN_MEMORY_SPEND_UPDATE_QUEUE, 

232 "gauge_value": queue_size, 

233 }, 

234 ) 

235 )