Skip to main content

aft/db/
compression_events.rs

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}