1use super::*;
2
3#[derive(Debug, Clone)]
4pub struct LlmUsageSummary {
5 pub requests: i64,
6 pub input_tokens: i64,
7 pub output_tokens: i64,
8 pub total_tokens: i64,
9 pub last_request_at: Option<String>,
10}
11
12#[derive(Debug, Clone)]
13pub struct LlmModelUsageSummary {
14 pub model: String,
15 pub requests: i64,
16 pub input_tokens: i64,
17 pub output_tokens: i64,
18 pub total_tokens: i64,
19}
20
21impl Database {
22 #[allow(clippy::too_many_arguments)]
23 pub fn log_llm_usage(
24 &self,
25 chat_id: i64,
26 caller_channel: &str,
27 provider: &str,
28 model: &str,
29 input_tokens: i64,
30 output_tokens: i64,
31 request_kind: &str,
32 ) -> Result<i64, MicroClawError> {
33 let conn = self.lock_conn();
34 let now = chrono::Utc::now().to_rfc3339();
35 let total_tokens = input_tokens.saturating_add(output_tokens);
36 conn.execute(
37 "INSERT INTO llm_usage_logs
38 (chat_id, caller_channel, provider, model, input_tokens, output_tokens, total_tokens, request_kind, created_at)
39 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)",
40 params![
41 chat_id,
42 caller_channel,
43 provider,
44 model,
45 input_tokens,
46 output_tokens,
47 total_tokens,
48 request_kind,
49 now,
50 ],
51 )?;
52 Ok(conn.last_insert_rowid())
53 }
54
55 pub fn get_llm_usage_summary(
56 &self,
57 chat_id: Option<i64>,
58 ) -> Result<LlmUsageSummary, MicroClawError> {
59 self.get_llm_usage_summary_since(chat_id, None)
60 }
61
62 pub fn get_llm_usage_summary_since(
63 &self,
64 chat_id: Option<i64>,
65 since: Option<&str>,
66 ) -> Result<LlmUsageSummary, MicroClawError> {
67 let conn = self.lock_conn();
68 let (requests, input_tokens, output_tokens, total_tokens, last_request_at) =
69 match (chat_id, since) {
70 (Some(id), Some(since_ts)) => conn.query_row(
71 "SELECT
72 COUNT(*),
73 COALESCE(SUM(input_tokens), 0),
74 COALESCE(SUM(output_tokens), 0),
75 COALESCE(SUM(total_tokens), 0),
76 MAX(created_at)
77 FROM llm_usage_logs
78 WHERE chat_id = ?1 AND created_at >= ?2",
79 params![id, since_ts],
80 |row| {
81 Ok((
82 row.get::<_, i64>(0)?,
83 row.get::<_, i64>(1)?,
84 row.get::<_, i64>(2)?,
85 row.get::<_, i64>(3)?,
86 row.get::<_, Option<String>>(4)?,
87 ))
88 },
89 )?,
90 (Some(id), None) => conn.query_row(
91 "SELECT
92 COUNT(*),
93 COALESCE(SUM(input_tokens), 0),
94 COALESCE(SUM(output_tokens), 0),
95 COALESCE(SUM(total_tokens), 0),
96 MAX(created_at)
97 FROM llm_usage_logs
98 WHERE chat_id = ?1",
99 params![id],
100 |row| {
101 Ok((
102 row.get::<_, i64>(0)?,
103 row.get::<_, i64>(1)?,
104 row.get::<_, i64>(2)?,
105 row.get::<_, i64>(3)?,
106 row.get::<_, Option<String>>(4)?,
107 ))
108 },
109 )?,
110 (None, Some(since_ts)) => conn.query_row(
111 "SELECT
112 COUNT(*),
113 COALESCE(SUM(input_tokens), 0),
114 COALESCE(SUM(output_tokens), 0),
115 COALESCE(SUM(total_tokens), 0),
116 MAX(created_at)
117 FROM llm_usage_logs
118 WHERE created_at >= ?1",
119 params![since_ts],
120 |row| {
121 Ok((
122 row.get::<_, i64>(0)?,
123 row.get::<_, i64>(1)?,
124 row.get::<_, i64>(2)?,
125 row.get::<_, i64>(3)?,
126 row.get::<_, Option<String>>(4)?,
127 ))
128 },
129 )?,
130 (None, None) => conn.query_row(
131 "SELECT
132 COUNT(*),
133 COALESCE(SUM(input_tokens), 0),
134 COALESCE(SUM(output_tokens), 0),
135 COALESCE(SUM(total_tokens), 0),
136 MAX(created_at)
137 FROM llm_usage_logs",
138 [],
139 |row| {
140 Ok((
141 row.get::<_, i64>(0)?,
142 row.get::<_, i64>(1)?,
143 row.get::<_, i64>(2)?,
144 row.get::<_, i64>(3)?,
145 row.get::<_, Option<String>>(4)?,
146 ))
147 },
148 )?,
149 };
150
151 Ok(LlmUsageSummary {
152 requests,
153 input_tokens,
154 output_tokens,
155 total_tokens,
156 last_request_at,
157 })
158 }
159
160 pub fn get_llm_usage_by_model(
161 &self,
162 chat_id: Option<i64>,
163 since: Option<&str>,
164 limit: Option<usize>,
165 ) -> Result<Vec<LlmModelUsageSummary>, MicroClawError> {
166 let conn = self.lock_conn();
167 let mut query = String::from(
168 "SELECT
169 model,
170 COUNT(*) AS requests,
171 COALESCE(SUM(input_tokens), 0) AS input_tokens,
172 COALESCE(SUM(output_tokens), 0) AS output_tokens,
173 COALESCE(SUM(total_tokens), 0) AS total_tokens
174 FROM llm_usage_logs",
175 );
176
177 let mut has_where = false;
178 if chat_id.is_some() {
179 query.push_str(" WHERE chat_id = ?1");
180 has_where = true;
181 }
182 if since.is_some() {
183 if has_where {
184 if chat_id.is_some() {
185 query.push_str(" AND created_at >= ?2");
186 } else {
187 query.push_str(" AND created_at >= ?1");
188 }
189 } else {
190 query.push_str(" WHERE created_at >= ?1");
191 }
192 }
193 query.push_str(" GROUP BY model ORDER BY total_tokens DESC");
194 if limit.is_some() {
195 match (chat_id.is_some(), since.is_some()) {
196 (true, true) => query.push_str(" LIMIT ?3"),
197 (true, false) | (false, true) => query.push_str(" LIMIT ?2"),
198 (false, false) => query.push_str(" LIMIT ?1"),
199 }
200 }
201
202 let mut stmt = conn.prepare(&query)?;
203 let mapper = |row: &rusqlite::Row<'_>| {
204 Ok(LlmModelUsageSummary {
205 model: row.get(0)?,
206 requests: row.get(1)?,
207 input_tokens: row.get(2)?,
208 output_tokens: row.get(3)?,
209 total_tokens: row.get(4)?,
210 })
211 };
212
213 let rows = match (chat_id, since, limit) {
214 (Some(id), Some(since_ts), Some(limit_n)) => {
215 stmt.query_map(params![id, since_ts, limit_n as i64], mapper)?
216 }
217 (Some(id), Some(since_ts), None) => stmt.query_map(params![id, since_ts], mapper)?,
218 (Some(id), None, Some(limit_n)) => {
219 stmt.query_map(params![id, limit_n as i64], mapper)?
220 }
221 (Some(id), None, None) => stmt.query_map(params![id], mapper)?,
222 (None, Some(since_ts), Some(limit_n)) => {
223 stmt.query_map(params![since_ts, limit_n as i64], mapper)?
224 }
225 (None, Some(since_ts), None) => stmt.query_map(params![since_ts], mapper)?,
226 (None, None, Some(limit_n)) => stmt.query_map(params![limit_n as i64], mapper)?,
227 (None, None, None) => stmt.query_map([], mapper)?,
228 };
229 rows.collect::<Result<Vec<_>, _>>().map_err(Into::into)
230 }
231
232 }
234
235#[cfg(test)]
236mod tests {
237 use crate::internal::storage::db::test_support::*;
238
239 #[test]
240 fn test_log_llm_usage_and_summary() {
241 let (db, dir) = test_db();
242 db.log_llm_usage(
243 100,
244 "telegram",
245 "anthropic",
246 "claude-test",
247 10,
248 5,
249 "agent_loop",
250 )
251 .unwrap();
252 db.log_llm_usage(
253 100,
254 "telegram",
255 "anthropic",
256 "claude-test",
257 20,
258 8,
259 "agent_loop",
260 )
261 .unwrap();
262 db.log_llm_usage(200, "discord", "openai", "gpt-test", 30, 7, "agent_loop")
263 .unwrap();
264
265 let chat_100 = db.get_llm_usage_summary(Some(100)).unwrap();
266 assert_eq!(chat_100.requests, 2);
267 assert_eq!(chat_100.input_tokens, 30);
268 assert_eq!(chat_100.output_tokens, 13);
269 assert_eq!(chat_100.total_tokens, 43);
270 assert!(chat_100.last_request_at.is_some());
271
272 let all = db.get_llm_usage_summary(None).unwrap();
273 assert_eq!(all.requests, 3);
274 assert_eq!(all.input_tokens, 60);
275 assert_eq!(all.output_tokens, 20);
276 assert_eq!(all.total_tokens, 80);
277 assert!(all.last_request_at.is_some());
278
279 cleanup(&dir);
280 }
281
282 #[test]
283 fn test_delete_chat_data_cleans_llm_usage() {
284 let (db, dir) = test_db();
285 db.upsert_chat(100, Some("chat-100"), "private").unwrap();
286 db.log_llm_usage(
287 100,
288 "telegram",
289 "anthropic",
290 "claude-test",
291 11,
292 9,
293 "agent_loop",
294 )
295 .unwrap();
296 db.log_llm_usage(
297 200,
298 "telegram",
299 "anthropic",
300 "claude-test",
301 3,
302 4,
303 "agent_loop",
304 )
305 .unwrap();
306
307 assert!(db.delete_chat_data(100).unwrap());
308
309 let chat_100 = db.get_llm_usage_summary(Some(100)).unwrap();
310 assert_eq!(chat_100.requests, 0);
311 let chat_200 = db.get_llm_usage_summary(Some(200)).unwrap();
312 assert_eq!(chat_200.requests, 1);
313
314 cleanup(&dir);
315 }
316
317 #[test]
318 fn test_get_llm_usage_summary_since_and_by_model() {
319 let (db, dir) = test_db();
320 db.log_llm_usage(
321 100,
322 "telegram",
323 "anthropic",
324 "claude-a",
325 10,
326 5,
327 "agent_loop",
328 )
329 .unwrap();
330 db.log_llm_usage(
331 100,
332 "telegram",
333 "anthropic",
334 "claude-a",
335 20,
336 10,
337 "agent_loop",
338 )
339 .unwrap();
340 db.log_llm_usage(100, "telegram", "anthropic", "claude-b", 3, 7, "agent_loop")
341 .unwrap();
342
343 let all = db.get_llm_usage_summary_since(Some(100), None).unwrap();
344 assert_eq!(all.requests, 3);
345 assert_eq!(all.input_tokens, 33);
346 assert_eq!(all.output_tokens, 22);
347
348 let future = db
349 .get_llm_usage_summary_since(Some(100), Some("2100-01-01T00:00:00Z"))
350 .unwrap();
351 assert_eq!(future.requests, 0);
352
353 let by_model = db
354 .get_llm_usage_by_model(Some(100), None, Some(10))
355 .unwrap();
356 assert_eq!(by_model.len(), 2);
357 assert_eq!(by_model[0].model, "claude-a");
358 assert_eq!(by_model[0].requests, 2);
359 assert_eq!(by_model[0].total_tokens, 45);
360 assert_eq!(by_model[1].model, "claude-b");
361 assert_eq!(by_model[1].requests, 1);
362 assert_eq!(by_model[1].total_tokens, 10);
363
364 cleanup(&dir);
365 }
366}