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
30impl MemorySubstrate {
31    /// Record a usage event for a fighter.
32    pub async fn record_usage(
33        &self,
34        fighter_id: &FighterId,
35        model: &str,
36        input_tokens: u64,
37        output_tokens: u64,
38        cost_usd: f64,
39    ) -> PunchResult<()> {
40        let fighter_str = fighter_id.to_string();
41
42        let conn = self.conn.lock().await;
43        conn.execute(
44            "INSERT INTO usage_events (fighter_id, model, input_tokens, output_tokens, cost_usd)
45             VALUES (?1, ?2, ?3, ?4, ?5)",
46            rusqlite::params![fighter_str, model, input_tokens, output_tokens, cost_usd],
47        )
48        .map_err(|e| PunchError::Memory(format!("failed to record usage: {e}")))?;
49
50        debug!(
51            fighter_id = %fighter_id,
52            model = model,
53            input_tokens = input_tokens,
54            output_tokens = output_tokens,
55            "usage recorded"
56        );
57        Ok(())
58    }
59
60    /// Get an aggregated usage summary for a fighter since the given timestamp.
61    pub async fn get_usage_summary(
62        &self,
63        fighter_id: &FighterId,
64        since: DateTime<Utc>,
65    ) -> PunchResult<UsageSummary> {
66        let fighter_str = fighter_id.to_string();
67        let since_str = since.format("%Y-%m-%dT%H:%M:%SZ").to_string();
68
69        let conn = self.conn.lock().await;
70
71        let result = conn
72            .query_row(
73                "SELECT COALESCE(SUM(input_tokens), 0),
74                    COALESCE(SUM(output_tokens), 0),
75                    COALESCE(SUM(cost_usd), 0.0),
76                    COUNT(*)
77             FROM usage_events
78             WHERE fighter_id = ?1 AND created_at >= ?2",
79                rusqlite::params![fighter_str, since_str],
80                |row| {
81                    let total_input_tokens: u64 = row.get(0)?;
82                    let total_output_tokens: u64 = row.get(1)?;
83                    let total_cost_usd: f64 = row.get(2)?;
84                    let event_count: u64 = row.get(3)?;
85                    Ok(UsageSummary {
86                        total_input_tokens,
87                        total_output_tokens,
88                        total_cost_usd,
89                        event_count,
90                    })
91                },
92            )
93            .map_err(|e| PunchError::Memory(format!("failed to get usage summary: {e}")))?;
94
95        Ok(result)
96    }
97
98    /// Get an aggregated usage summary across ALL fighters since the given timestamp.
99    pub async fn get_total_usage_summary(&self, since: DateTime<Utc>) -> PunchResult<UsageSummary> {
100        let since_str = since.format("%Y-%m-%dT%H:%M:%SZ").to_string();
101
102        let conn = self.conn.lock().await;
103
104        let result = conn
105            .query_row(
106                "SELECT COALESCE(SUM(input_tokens), 0),
107                    COALESCE(SUM(output_tokens), 0),
108                    COALESCE(SUM(cost_usd), 0.0),
109                    COUNT(*)
110             FROM usage_events
111             WHERE created_at >= ?1",
112                rusqlite::params![since_str],
113                |row| {
114                    let total_input_tokens: u64 = row.get(0)?;
115                    let total_output_tokens: u64 = row.get(1)?;
116                    let total_cost_usd: f64 = row.get(2)?;
117                    let event_count: u64 = row.get(3)?;
118                    Ok(UsageSummary {
119                        total_input_tokens,
120                        total_output_tokens,
121                        total_cost_usd,
122                        event_count,
123                    })
124                },
125            )
126            .map_err(|e| PunchError::Memory(format!("failed to get total usage summary: {e}")))?;
127
128        Ok(result)
129    }
130}
131
132#[cfg(test)]
133mod tests {
134    use chrono::{Duration, Utc};
135    use punch_types::{FighterManifest, FighterStatus, ModelConfig, Provider, WeightClass};
136
137    use crate::MemorySubstrate;
138
139    fn test_manifest() -> FighterManifest {
140        FighterManifest {
141            name: "Usage Fighter".into(),
142            description: "usage test".into(),
143            model: ModelConfig {
144                provider: Provider::Anthropic,
145                model: "claude-sonnet-4-20250514".into(),
146                api_key_env: None,
147                base_url: None,
148                max_tokens: Some(4096),
149                temperature: Some(0.7),
150            },
151            system_prompt: "test".into(),
152            capabilities: Vec::new(),
153            weight_class: WeightClass::Featherweight,
154            tenant_id: None,
155        }
156    }
157
158    #[tokio::test]
159    async fn test_record_and_summarize_usage() {
160        let substrate = MemorySubstrate::in_memory().unwrap();
161        let fid = punch_types::FighterId::new();
162        substrate
163            .save_fighter(&fid, &test_manifest(), FighterStatus::Idle)
164            .await
165            .unwrap();
166
167        substrate
168            .record_usage(&fid, "claude-sonnet-4-20250514", 1000, 500, 0.015)
169            .await
170            .unwrap();
171        substrate
172            .record_usage(&fid, "claude-sonnet-4-20250514", 2000, 800, 0.028)
173            .await
174            .unwrap();
175
176        let since = Utc::now() - Duration::hours(1);
177        let summary = substrate.get_usage_summary(&fid, since).await.unwrap();
178
179        assert_eq!(summary.event_count, 2);
180        assert_eq!(summary.total_input_tokens, 3000);
181        assert_eq!(summary.total_output_tokens, 1300);
182        assert!((summary.total_cost_usd - 0.043).abs() < 1e-9);
183    }
184
185    #[tokio::test]
186    async fn test_usage_summary_empty() {
187        let substrate = MemorySubstrate::in_memory().unwrap();
188        let fid = punch_types::FighterId::new();
189        substrate
190            .save_fighter(&fid, &test_manifest(), FighterStatus::Idle)
191            .await
192            .unwrap();
193
194        let since = Utc::now() - Duration::hours(1);
195        let summary = substrate.get_usage_summary(&fid, since).await.unwrap();
196
197        assert_eq!(summary.event_count, 0);
198        assert_eq!(summary.total_input_tokens, 0);
199    }
200}
punch_memory/usage.rs

punch_memory/
usage.rs
