punch_memory/
usage.rs

1use chrono::{DateTime, Utc};
2use serde::{Deserialize, Serialize};
3
4use punch_types::{FighterId, PunchError, PunchResult};
5use tracing::debug;
6
7use crate::MemorySubstrate;
8
9/// A single usage / metering event.
10#[derive(Debug, Clone, Serialize, Deserialize)]
11pub struct UsageEvent {
12    pub id: i64,
13    pub fighter_id: FighterId,
14    pub model: String,
15    pub input_tokens: u64,
16    pub output_tokens: u64,
17    pub cost_usd: f64,
18    pub created_at: String,
19}
20
21/// Aggregated usage summary for a fighter.
22#[derive(Debug, Clone, Default, Serialize, Deserialize)]
23pub struct UsageSummary {
24    pub total_input_tokens: u64,
25    pub total_output_tokens: u64,
26    pub total_cost_usd: f64,
27    pub event_count: u64,
28}
29
30/// Per-model usage breakdown row.
31#[derive(Debug, Clone, Serialize, Deserialize)]
32pub struct ModelUsageBreakdown {
33    pub model: String,
34    pub input_tokens: u64,
35    pub output_tokens: u64,
36    pub cost_usd: f64,
37    pub request_count: u64,
38}
39
40/// Per-fighter usage breakdown row.
41#[derive(Debug, Clone, Serialize, Deserialize)]
42pub struct FighterUsageBreakdown {
43    pub fighter_id: FighterId,
44    pub input_tokens: u64,
45    pub output_tokens: u64,
46    pub cost_usd: f64,
47    pub request_count: u64,
48}
49
50impl MemorySubstrate {
51    /// Record a usage event for a fighter.
52    pub async fn record_usage(
53        &self,
54        fighter_id: &FighterId,
55        model: &str,
56        input_tokens: u64,
57        output_tokens: u64,
58        cost_usd: f64,
59    ) -> PunchResult<()> {
60        let fighter_str = fighter_id.to_string();
61
62        let conn = self.conn.lock().await;
63        conn.execute(
64            "INSERT INTO usage_events (fighter_id, model, input_tokens, output_tokens, cost_usd)
65             VALUES (?1, ?2, ?3, ?4, ?5)",
66            rusqlite::params![fighter_str, model, input_tokens, output_tokens, cost_usd],
67        )
68        .map_err(|e| PunchError::Memory(format!("failed to record usage: {e}")))?;
69
70        debug!(
71            fighter_id = %fighter_id,
72            model = model,
73            input_tokens = input_tokens,
74            output_tokens = output_tokens,
75            "usage recorded"
76        );
77        Ok(())
78    }
79
80    /// Get an aggregated usage summary for a fighter since the given timestamp.
81    pub async fn get_usage_summary(
82        &self,
83        fighter_id: &FighterId,
84        since: DateTime<Utc>,
85    ) -> PunchResult<UsageSummary> {
86        let fighter_str = fighter_id.to_string();
87        let since_str = since.format("%Y-%m-%dT%H:%M:%SZ").to_string();
88
89        let conn = self.conn.lock().await;
90
91        let result = conn
92            .query_row(
93                "SELECT COALESCE(SUM(input_tokens), 0),
94                    COALESCE(SUM(output_tokens), 0),
95                    COALESCE(SUM(cost_usd), 0.0),
96                    COUNT(*)
97             FROM usage_events
98             WHERE fighter_id = ?1 AND created_at >= ?2",
99                rusqlite::params![fighter_str, since_str],
100                |row| {
101                    let total_input_tokens: u64 = row.get(0)?;
102                    let total_output_tokens: u64 = row.get(1)?;
103                    let total_cost_usd: f64 = row.get(2)?;
104                    let event_count: u64 = row.get(3)?;
105                    Ok(UsageSummary {
106                        total_input_tokens,
107                        total_output_tokens,
108                        total_cost_usd,
109                        event_count,
110                    })
111                },
112            )
113            .map_err(|e| PunchError::Memory(format!("failed to get usage summary: {e}")))?;
114
115        Ok(result)
116    }
117
118    /// Get per-model usage breakdown for a fighter since the given timestamp.
119    pub async fn get_model_breakdown(
120        &self,
121        fighter_id: &FighterId,
122        since: DateTime<Utc>,
123    ) -> PunchResult<Vec<ModelUsageBreakdown>> {
124        let fighter_str = fighter_id.to_string();
125        let since_str = since.format("%Y-%m-%dT%H:%M:%SZ").to_string();
126
127        let conn = self.conn.lock().await;
128
129        let mut stmt = conn
130            .prepare(
131                "SELECT model,
132                    COALESCE(SUM(input_tokens), 0),
133                    COALESCE(SUM(output_tokens), 0),
134                    COALESCE(SUM(cost_usd), 0.0),
135                    COUNT(*)
136                 FROM usage_events
137                 WHERE fighter_id = ?1 AND created_at >= ?2
138                 GROUP BY model
139                 ORDER BY SUM(cost_usd) DESC",
140            )
141            .map_err(|e| PunchError::Memory(format!("failed to prepare model breakdown: {e}")))?;
142
143        let rows = stmt
144            .query_map(rusqlite::params![fighter_str, since_str], |row| {
145                Ok(ModelUsageBreakdown {
146                    model: row.get(0)?,
147                    input_tokens: row.get(1)?,
148                    output_tokens: row.get(2)?,
149                    cost_usd: row.get(3)?,
150                    request_count: row.get(4)?,
151                })
152            })
153            .map_err(|e| PunchError::Memory(format!("failed to query model breakdown: {e}")))?;
154
155        let mut result = Vec::new();
156        for row in rows {
157            result.push(
158                row.map_err(|e| PunchError::Memory(format!("failed to read breakdown row: {e}")))?,
159            );
160        }
161        Ok(result)
162    }
163
164    /// Get per-model usage breakdown across ALL fighters since the given timestamp.
165    pub async fn get_total_model_breakdown(
166        &self,
167        since: DateTime<Utc>,
168    ) -> PunchResult<Vec<ModelUsageBreakdown>> {
169        let since_str = since.format("%Y-%m-%dT%H:%M:%SZ").to_string();
170
171        let conn = self.conn.lock().await;
172
173        let mut stmt = conn
174            .prepare(
175                "SELECT model,
176                    COALESCE(SUM(input_tokens), 0),
177                    COALESCE(SUM(output_tokens), 0),
178                    COALESCE(SUM(cost_usd), 0.0),
179                    COUNT(*)
180                 FROM usage_events
181                 WHERE created_at >= ?1
182                 GROUP BY model
183                 ORDER BY SUM(cost_usd) DESC",
184            )
185            .map_err(|e| PunchError::Memory(format!("failed to prepare model breakdown: {e}")))?;
186
187        let rows = stmt
188            .query_map(rusqlite::params![since_str], |row| {
189                Ok(ModelUsageBreakdown {
190                    model: row.get(0)?,
191                    input_tokens: row.get(1)?,
192                    output_tokens: row.get(2)?,
193                    cost_usd: row.get(3)?,
194                    request_count: row.get(4)?,
195                })
196            })
197            .map_err(|e| PunchError::Memory(format!("failed to query model breakdown: {e}")))?;
198
199        let mut result = Vec::new();
200        for row in rows {
201            result.push(
202                row.map_err(|e| PunchError::Memory(format!("failed to read breakdown row: {e}")))?,
203            );
204        }
205        Ok(result)
206    }
207
208    /// Get per-fighter usage breakdown across all fighters since the given timestamp.
209    pub async fn get_fighter_breakdown(
210        &self,
211        since: DateTime<Utc>,
212    ) -> PunchResult<Vec<FighterUsageBreakdown>> {
213        let since_str = since.format("%Y-%m-%dT%H:%M:%SZ").to_string();
214
215        let conn = self.conn.lock().await;
216
217        let mut stmt = conn
218            .prepare(
219                "SELECT fighter_id,
220                    COALESCE(SUM(input_tokens), 0),
221                    COALESCE(SUM(output_tokens), 0),
222                    COALESCE(SUM(cost_usd), 0.0),
223                    COUNT(*)
224                 FROM usage_events
225                 WHERE created_at >= ?1
226                 GROUP BY fighter_id
227                 ORDER BY SUM(cost_usd) DESC",
228            )
229            .map_err(|e| PunchError::Memory(format!("failed to prepare fighter breakdown: {e}")))?;
230
231        let rows = stmt
232            .query_map(rusqlite::params![since_str], |row| {
233                let id_str: String = row.get(0)?;
234                let fighter_id = id_str
235                    .parse::<uuid::Uuid>()
236                    .map(FighterId)
237                    .unwrap_or_else(|_| FighterId::new());
238                Ok(FighterUsageBreakdown {
239                    fighter_id,
240                    input_tokens: row.get(1)?,
241                    output_tokens: row.get(2)?,
242                    cost_usd: row.get(3)?,
243                    request_count: row.get(4)?,
244                })
245            })
246            .map_err(|e| PunchError::Memory(format!("failed to query fighter breakdown: {e}")))?;
247
248        let mut result = Vec::new();
249        for row in rows {
250            result.push(
251                row.map_err(|e| PunchError::Memory(format!("failed to read breakdown row: {e}")))?,
252            );
253        }
254        Ok(result)
255    }
256
257    /// Get an aggregated usage summary across ALL fighters since the given timestamp.
258    pub async fn get_total_usage_summary(&self, since: DateTime<Utc>) -> PunchResult<UsageSummary> {
259        let since_str = since.format("%Y-%m-%dT%H:%M:%SZ").to_string();
260
261        let conn = self.conn.lock().await;
262
263        let result = conn
264            .query_row(
265                "SELECT COALESCE(SUM(input_tokens), 0),
266                    COALESCE(SUM(output_tokens), 0),
267                    COALESCE(SUM(cost_usd), 0.0),
268                    COUNT(*)
269             FROM usage_events
270             WHERE created_at >= ?1",
271                rusqlite::params![since_str],
272                |row| {
273                    let total_input_tokens: u64 = row.get(0)?;
274                    let total_output_tokens: u64 = row.get(1)?;
275                    let total_cost_usd: f64 = row.get(2)?;
276                    let event_count: u64 = row.get(3)?;
277                    Ok(UsageSummary {
278                        total_input_tokens,
279                        total_output_tokens,
280                        total_cost_usd,
281                        event_count,
282                    })
283                },
284            )
285            .map_err(|e| PunchError::Memory(format!("failed to get total usage summary: {e}")))?;
286
287        Ok(result)
288    }
289}
290
291#[cfg(test)]
292mod tests {
293    use chrono::{Duration, Utc};
294    use punch_types::{FighterManifest, FighterStatus, ModelConfig, Provider, WeightClass};
295
296    use crate::MemorySubstrate;
297
298    fn test_manifest() -> FighterManifest {
299        FighterManifest {
300            name: "Usage Fighter".into(),
301            description: "usage test".into(),
302            model: ModelConfig {
303                provider: Provider::Anthropic,
304                model: "claude-sonnet-4-20250514".into(),
305                api_key_env: None,
306                base_url: None,
307                max_tokens: Some(4096),
308                temperature: Some(0.7),
309            },
310            system_prompt: "test".into(),
311            capabilities: Vec::new(),
312            weight_class: WeightClass::Featherweight,
313            tenant_id: None,
314        }
315    }
316
317    #[tokio::test]
318    async fn test_record_and_summarize_usage() {
319        let substrate = MemorySubstrate::in_memory().unwrap();
320        let fid = punch_types::FighterId::new();
321        substrate
322            .save_fighter(&fid, &test_manifest(), FighterStatus::Idle)
323            .await
324            .unwrap();
325
326        substrate
327            .record_usage(&fid, "claude-sonnet-4-20250514", 1000, 500, 0.015)
328            .await
329            .unwrap();
330        substrate
331            .record_usage(&fid, "claude-sonnet-4-20250514", 2000, 800, 0.028)
332            .await
333            .unwrap();
334
335        let since = Utc::now() - Duration::hours(1);
336        let summary = substrate.get_usage_summary(&fid, since).await.unwrap();
337
338        assert_eq!(summary.event_count, 2);
339        assert_eq!(summary.total_input_tokens, 3000);
340        assert_eq!(summary.total_output_tokens, 1300);
341        assert!((summary.total_cost_usd - 0.043).abs() < 1e-9);
342    }
343
344    #[tokio::test]
345    async fn test_model_breakdown() {
346        let substrate = MemorySubstrate::in_memory().unwrap();
347        let fid = punch_types::FighterId::new();
348        substrate
349            .save_fighter(&fid, &test_manifest(), FighterStatus::Idle)
350            .await
351            .unwrap();
352
353        substrate
354            .record_usage(&fid, "claude-sonnet-4-20250514", 1000, 500, 0.015)
355            .await
356            .unwrap();
357        substrate
358            .record_usage(&fid, "gpt-4o-mini", 2000, 800, 0.002)
359            .await
360            .unwrap();
361        substrate
362            .record_usage(&fid, "claude-sonnet-4-20250514", 3000, 1000, 0.030)
363            .await
364            .unwrap();
365
366        let since = Utc::now() - Duration::hours(1);
367        let breakdown = substrate.get_model_breakdown(&fid, since).await.unwrap();
368
369        assert_eq!(breakdown.len(), 2);
370        // Ordered by cost DESC, so sonnet ($0.045) first, then gpt-4o-mini ($0.002)
371        assert_eq!(breakdown[0].model, "claude-sonnet-4-20250514");
372        assert_eq!(breakdown[0].input_tokens, 4000);
373        assert_eq!(breakdown[0].output_tokens, 1500);
374        assert_eq!(breakdown[0].request_count, 2);
375        assert_eq!(breakdown[1].model, "gpt-4o-mini");
376        assert_eq!(breakdown[1].request_count, 1);
377    }
378
379    #[tokio::test]
380    async fn test_total_model_breakdown() {
381        let substrate = MemorySubstrate::in_memory().unwrap();
382        let fid1 = punch_types::FighterId::new();
383        let fid2 = punch_types::FighterId::new();
384        substrate
385            .save_fighter(&fid1, &test_manifest(), FighterStatus::Idle)
386            .await
387            .unwrap();
388        substrate
389            .save_fighter(&fid2, &test_manifest(), FighterStatus::Idle)
390            .await
391            .unwrap();
392
393        substrate
394            .record_usage(&fid1, "claude-sonnet-4-20250514", 1000, 500, 0.015)
395            .await
396            .unwrap();
397        substrate
398            .record_usage(&fid2, "claude-sonnet-4-20250514", 2000, 800, 0.028)
399            .await
400            .unwrap();
401
402        let since = Utc::now() - Duration::hours(1);
403        let breakdown = substrate.get_total_model_breakdown(since).await.unwrap();
404
405        assert_eq!(breakdown.len(), 1);
406        assert_eq!(breakdown[0].input_tokens, 3000);
407        assert_eq!(breakdown[0].request_count, 2);
408    }
409
410    #[tokio::test]
411    async fn test_fighter_breakdown() {
412        let substrate = MemorySubstrate::in_memory().unwrap();
413        let fid1 = punch_types::FighterId::new();
414        let fid2 = punch_types::FighterId::new();
415        substrate
416            .save_fighter(&fid1, &test_manifest(), FighterStatus::Idle)
417            .await
418            .unwrap();
419        substrate
420            .save_fighter(&fid2, &test_manifest(), FighterStatus::Idle)
421            .await
422            .unwrap();
423
424        substrate
425            .record_usage(&fid1, "claude-sonnet-4-20250514", 1000, 500, 0.015)
426            .await
427            .unwrap();
428        substrate
429            .record_usage(&fid2, "gpt-4o-mini", 5000, 2000, 0.004)
430            .await
431            .unwrap();
432
433        let since = Utc::now() - Duration::hours(1);
434        let breakdown = substrate.get_fighter_breakdown(since).await.unwrap();
435
436        assert_eq!(breakdown.len(), 2);
437        // Ordered by cost DESC: sonnet ($0.015) first
438        assert_eq!(breakdown[0].fighter_id, fid1);
439        assert_eq!(breakdown[1].fighter_id, fid2);
440    }
441
442    #[tokio::test]
443    async fn test_usage_summary_empty() {
444        let substrate = MemorySubstrate::in_memory().unwrap();
445        let fid = punch_types::FighterId::new();
446        substrate
447            .save_fighter(&fid, &test_manifest(), FighterStatus::Idle)
448            .await
449            .unwrap();
450
451        let since = Utc::now() - Duration::hours(1);
452        let summary = substrate.get_usage_summary(&fid, since).await.unwrap();
453
454        assert_eq!(summary.event_count, 0);
455        assert_eq!(summary.total_input_tokens, 0);
456    }
457}
punch_memory/usage.rs

punch_memory/
usage.rs
