Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/db/spend_log_tool_index.py: 53%

1""" 

2Tool usage tracking for the dashboard. 

3 

4At request time the spend writer builds one ToolUsageTransaction per request that 

5invoked tools (MCP namespaced tool name plus response tool_calls; declared-but-not- 

6invoked tools are excluded) and queues it on the prisma client. The spend-log flush 

7job drains the queue into LiteLLM_SpendLogToolIndex (per-request drill-down) and 

8LiteLLM_DailyToolSpend (the per-day rollup the Cost Optimization card reads). The 

9index rows are keyed on (request_id, tool_name) and written with skip_duplicates, 

10so they go out as bounded standalone statements; every rollup upsert stays in one 

11transaction, so a failed flush never leaves a partial rollup increment. 

12""" 

13 

14from __future__ import annotations 

15 

16import asyncio 

17import random 

18from collections.abc import Sequence 

19from dataclasses import dataclass 

20from datetime import datetime, timezone 

21from itertools import groupby 

22from typing import TYPE_CHECKING, Any, Final 

23 

24from litellm.constants import SPEND_LOG_WRITE_BATCH_MAX_BYTES, SPEND_LOG_WRITE_BATCH_MAX_ROWS 

25from litellm.proxy._types import DB_RETRY_SAFE_ERROR_TYPES 

26from litellm.proxy.db.spend_log_batching import spend_log_write_batches 

27from litellm.repositories.table_repositories import SpendLogToolIndexRepository 

28 

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

30 from litellm.proxy.utils import PrismaClient 

31 

32 

33@dataclass(frozen=True, slots=True) 

34class ToolUsageTransaction: 

35 request_id: str 

36 date: str 

37 start_time: datetime 

38 tool_names: tuple[str, ...] 

39 spend: float 

40 total_tokens: int 

41 

42 

43def response_tool_call_names(completion_response: object) -> tuple[str, ...]: 

44 """Tool names invoked in a completion response, in call order, for any response 

45 surface get_tool_calls_from_response understands (chat completions, Responses 

46 API output items, Anthropic Messages tool_use blocks). Reads every choice of 

47 an ``n>1`` chat response: each choice cost money and its tool calls ran.""" 

48 if completion_response is None or isinstance(completion_response, Exception): 

49 return () 

50 from litellm.litellm_core_utils.prompt_templates.factory import ( 

51 get_tool_calls_from_response, 

52 ) 

53 

54 return tuple( 

55 stripped 

56 for tool_call in get_tool_calls_from_response(completion_response, include_all_choices=True) 

57 if isinstance(name := tool_call.get("name"), str) and (stripped := name.strip()) 

58 ) 

59 

60 

61def build_tool_usage_transaction( 

62 request_id: str, 

63 start_time_iso: str, 

64 mcp_namespaced_tool_name: str | None, 

65 spend: float, 

66 total_tokens: int, 

67 completion_response: object, 

68 realtime_tool_calls: Any = None, 

69) -> ToolUsageTransaction | None: 

70 """None when the request invoked no tools. Realtime sessions carry invoked 

71 tools in kwargs["realtime_tool_calls"] (OpenAI tool_calls shape) rather than 

72 on a response object, so they are normalized through the same owner by 

73 wrapping them in the chat-completion shape. Date derivation must match the 

74 daily spend writer's ``startTime.split("T")[0]`` so rollup rows land in the 

75 same UTC day bucket as LiteLLM_DailyUserSpend.""" 

76 mcp_names: Final = ( 

77 (mcp_namespaced_tool_name.strip(),) if mcp_namespaced_tool_name and mcp_namespaced_tool_name.strip() else () 

78 ) 

79 realtime_names: Final = ( 

80 response_tool_call_names({"choices": [{"message": {"tool_calls": realtime_tool_calls}}]}) 

81 if realtime_tool_calls 

82 else () 

83 ) 

84 tool_names: Final = tuple(dict.fromkeys(mcp_names + response_tool_call_names(completion_response) + realtime_names)) 

85 if not tool_names: 85 ↛ 87line 85 didn't jump to line 87 because the condition on line 85 was always true

86 return None 

87 try: 

88 start_time: Final = datetime.fromisoformat(start_time_iso.replace("Z", "+00:00")) 

89 except ValueError: 

90 return None 

91 return ToolUsageTransaction( 

92 request_id=request_id, 

93 date=start_time_iso.split("T")[0], 

94 start_time=start_time if start_time.tzinfo else start_time.replace(tzinfo=timezone.utc), 

95 tool_names=tool_names, 

96 spend=spend, 

97 total_tokens=total_tokens, 

98 ) 

99 

100 

101async def flush_tool_usage_transactions( 

102 prisma_client: PrismaClient, 

103 transactions: Sequence[ToolUsageTransaction], 

104 n_retry_times: int = 3, 

105) -> None: 

106 """Write the index rows as bounded standalone statements, then every rollup 

107 upsert for the drained queue batch in one transaction. One flush fans out to 

108 transactions x tools index rows, so the index write is split by the spend-log 

109 statement budgets; a split inside ``batch_()`` would not help, since the 

110 batcher ships every queued statement to the query engine as one payload. 

111 Retries only ConnectError, the one failure that proves the statements never 

112 reached the database; replayed index rows are no-ops under skip_duplicates. 

113 Post-send failures (Read timeouts and errors) are ambiguous and are NOT 

114 retried: the engine can abandon the transaction open on the pooled 

115 connection, so a retry's statements stack into the same transaction and one 

116 commit applies both increment sets. Ambiguous failures drop the batch; the 

117 caller logs it at error. Callers must not add their own retry around this 

118 function.""" 

119 if not transactions: 119 ↛ 122line 119 didn't jump to line 122 because the condition on line 119 was always true

120 return 

121 

122 index_rows: Final = [ 

123 {"request_id": txn.request_id, "tool_name": tool_name, "start_time": txn.start_time} 

124 for txn in transactions 

125 for tool_name in txn.tool_names 

126 ] 

127 per_tool_day: Final = sorted( 

128 ((txn.date, tool_name, txn.spend, txn.total_tokens) for txn in transactions for tool_name in txn.tool_names), 

129 key=lambda entry: (entry[0], entry[1]), 

130 ) 

131 

132 index_table: Final = SpendLogToolIndexRepository(prisma_client).table 

133 for attempt in range(n_retry_times + 1): 

134 try: 

135 for statement_rows in spend_log_write_batches( 

136 index_rows, SPEND_LOG_WRITE_BATCH_MAX_BYTES, SPEND_LOG_WRITE_BATCH_MAX_ROWS 

137 ): 

138 await index_table.create_many(data=statement_rows, skip_duplicates=True) 

139 async with prisma_client.db.batch_() as batcher: 

140 for (date_key, tool_name), grouped in groupby(per_tool_day, key=lambda entry: (entry[0], entry[1])): 

141 entries = tuple(grouped) 

142 spend = sum(entry[2] for entry in entries) 

143 total_tokens = sum(entry[3] for entry in entries) 

144 batcher.litellm_dailytoolspend.upsert( 

145 where={"date_tool_name": {"date": date_key, "tool_name": tool_name}}, 

146 data={ 

147 "create": { 

148 "date": date_key, 

149 "tool_name": tool_name, 

150 "spend": spend, 

151 "total_tokens": total_tokens, 

152 "request_count": len(entries), 

153 }, 

154 "update": { 

155 "spend": {"increment": spend}, 

156 "total_tokens": {"increment": total_tokens}, 

157 "request_count": {"increment": len(entries)}, 

158 }, 

159 }, 

160 ) 

161 return 

162 except DB_RETRY_SAFE_ERROR_TYPES: 

163 if attempt >= n_retry_times: 

164 raise 

165 await asyncio.sleep(2.0**attempt + random.uniform(0, 1))