1use rusqlite::{params, Connection};
2
3pub struct CompressionEventRow<'a> {
4 pub harness: &'a str,
5 pub session_id: Option<&'a str>,
6 pub project_key: &'a str,
7 pub tool: &'a str,
8 pub task_id: Option<&'a str>,
9 pub command: Option<&'a str>,
10 pub compressor: &'a str,
11 pub original_bytes: i64,
12 pub compressed_bytes: i64,
13 pub original_tokens: u32,
14 pub compressed_tokens: u32,
15 pub created_at: i64,
16}
17
18#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, serde::Serialize)]
19pub struct CompressionAggregate {
20 pub events: u64,
21 pub original_tokens: u64,
22 pub compressed_tokens: u64,
23}
24
25impl CompressionAggregate {
26 pub fn savings_tokens(&self) -> u64 {
27 self.original_tokens.saturating_sub(self.compressed_tokens)
28 }
29}
30
31pub fn insert_compression_event(
32 conn: &Connection,
33 row: &CompressionEventRow<'_>,
34) -> rusqlite::Result<()> {
35 conn.execute(
36 r#"
37 INSERT INTO compression_events (
38 harness, session_id, project_key, tool, task_id, command, compressor,
39 original_bytes, compressed_bytes, original_tokens, compressed_tokens, created_at
40 )
41 SELECT ?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12
42 WHERE NOT EXISTS (
43 SELECT 1
44 FROM compression_events
45 WHERE harness = ?1
46 AND session_id IS ?2
47 AND task_id IS ?5
48 AND tool = ?4
49 LIMIT 1
50 )
51 "#,
52 params![
53 row.harness,
54 row.session_id,
55 row.project_key,
56 row.tool,
57 row.task_id,
58 row.command,
59 row.compressor,
60 row.original_bytes,
61 row.compressed_bytes,
62 row.original_tokens,
63 row.compressed_tokens,
64 row.created_at,
65 ],
66 )?;
67 Ok(())
68}
69
70pub fn aggregate_for_project(
71 conn: &Connection,
72 harness: &str,
73 project_key: &str,
74) -> rusqlite::Result<CompressionAggregate> {
75 conn.query_row(
76 r#"
77 SELECT
78 COUNT(*) AS events,
79 COALESCE(SUM(original_tokens), 0) AS original,
80 COALESCE(SUM(compressed_tokens), 0) AS compressed
81 FROM compression_events
82 WHERE harness = ?1 AND project_key = ?2
83 "#,
84 params![harness, project_key],
85 |row| {
86 Ok(CompressionAggregate {
87 events: row.get::<_, i64>(0)? as u64,
88 original_tokens: row.get::<_, i64>(1)? as u64,
89 compressed_tokens: row.get::<_, i64>(2)? as u64,
90 })
91 },
92 )
93}
94
95pub fn aggregate_for_session(
96 conn: &Connection,
97 harness: &str,
98 project_key: &str,
99 session_id: &str,
100) -> rusqlite::Result<CompressionAggregate> {
101 conn.query_row(
102 r#"
103 SELECT
104 COUNT(*) AS events,
105 COALESCE(SUM(original_tokens), 0) AS original,
106 COALESCE(SUM(compressed_tokens), 0) AS compressed
107 FROM compression_events
108 WHERE harness = ?1 AND project_key = ?2 AND session_id = ?3
109 "#,
110 params![harness, project_key, session_id],
111 |row| {
112 Ok(CompressionAggregate {
113 events: row.get::<_, i64>(0)? as u64,
114 original_tokens: row.get::<_, i64>(1)? as u64,
115 compressed_tokens: row.get::<_, i64>(2)? as u64,
116 })
117 },
118 )
119}