memory_core/store/
metrics.rs

1use rusqlite::params;
2
3use crate::types::MemoryMetric;
4use crate::Result;
5
6use super::Store;
7
8impl Store {
9    pub fn record_injection(&self, memory_ids: &[i64], tokens_per_memory: i32) -> Result<()> {
10        let tx = self.conn().unchecked_transaction()?;
11        for id in memory_ids {
12            tx.execute(
13                "INSERT INTO metrics (memory_id, injections, tokens_injected, last_injected_at)
14                 VALUES (?1, 1, ?2, strftime('%Y-%m-%dT%H:%M:%fZ', 'now'))
15                 ON CONFLICT(memory_id) DO UPDATE SET
16                    injections = injections + 1,
17                    tokens_injected = tokens_injected + ?2,
18                    last_injected_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now')",
19                params![id, tokens_per_memory],
20            )?;
21        }
22        tx.commit()?;
23        Ok(())
24    }
25
26    pub fn record_hit(&self, memory_id: i64) -> Result<()> {
27        self.conn().execute(
28            "INSERT INTO metrics (memory_id, hits, last_hit_at)
29             VALUES (?1, 1, strftime('%Y-%m-%dT%H:%M:%fZ', 'now'))
30             ON CONFLICT(memory_id) DO UPDATE SET
31                hits = hits + 1,
32                last_hit_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now')",
33            params![memory_id],
34        )?;
35        Ok(())
36    }
37
38    pub fn record_hit_batch(&self, memory_ids: &[i64]) -> Result<()> {
39        let tx = self.conn().unchecked_transaction()?;
40        for id in memory_ids {
41            tx.execute(
42                "INSERT INTO metrics (memory_id, hits, last_hit_at)
43                 VALUES (?1, 1, strftime('%Y-%m-%dT%H:%M:%fZ', 'now'))
44                 ON CONFLICT(memory_id) DO UPDATE SET
45                    hits = hits + 1,
46                    last_hit_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now')",
47                params![id],
48            )?;
49        }
50        tx.commit()?;
51        Ok(())
52    }
53
54    /// Cumulative token stats across all time (from the metrics table).
55    pub fn cumulative_stats(&self) -> Result<crate::types::TokenStats> {
56        let (injections, hits, tokens_injected) = self.conn().query_row(
57            "SELECT COALESCE(SUM(injections), 0), COALESCE(SUM(hits), 0), COALESCE(SUM(tokens_injected), 0) FROM metrics",
58            [],
59            |row| Ok((row.get::<_, i64>(0)?, row.get::<_, i64>(1)?, row.get::<_, i64>(2)?)),
60        )?;
61        let unique = self.conn().query_row(
62            "SELECT COUNT(*) FROM metrics WHERE injections > 0",
63            [],
64            |row| row.get::<_, i64>(0),
65        )?;
66        Ok(crate::types::TokenStats {
67            injections,
68            hits,
69            unique_memories_injected: unique,
70            tokens_injected,
71        })
72    }
73
74    pub fn dedup_total(&self) -> Result<i64> {
75        Ok(self.conn().query_row(
76            "SELECT COALESCE(SUM(duplicate_count), 0) FROM memories WHERE deleted_at IS NULL",
77            [],
78            |row| row.get(0),
79        )?)
80    }
81
82    pub fn revision_total(&self) -> Result<i64> {
83        Ok(self.conn().query_row(
84            "SELECT COALESCE(SUM(revision_count), 0) FROM memories WHERE deleted_at IS NULL",
85            [],
86            |row| row.get(0),
87        )?)
88    }
89
90    pub fn low_roi_count(&self) -> Result<i64> {
91        Ok(self.conn().query_row(
92            "SELECT COUNT(*) FROM metrics WHERE injections > 10 AND CAST(hits AS REAL) / injections < 0.1",
93            [],
94            |row| row.get(0),
95        )?)
96    }
97
98    pub fn get_metrics(&self) -> Result<Vec<MemoryMetric>> {
99        let mut stmt = self.conn().prepare(
100            "SELECT m.id, m.key, m.scope,
101                    COALESCE(mt.injections, 0),
102                    COALESCE(mt.hits, 0),
103                    COALESCE(mt.tokens_injected, 0),
104                    CASE WHEN COALESCE(mt.injections, 0) > 0
105                         THEN CAST(COALESCE(mt.hits, 0) AS REAL) / mt.injections
106                         ELSE 0.0 END
107             FROM memories m
108             LEFT JOIN metrics mt ON mt.memory_id = m.id
109             WHERE m.deleted_at IS NULL
110             ORDER BY COALESCE(mt.injections, 0) DESC
111             LIMIT 100",
112        )?;
113        let results = stmt
114            .query_map([], |row| {
115                Ok(MemoryMetric {
116                    id: row.get(0)?,
117                    key: row.get(1)?,
118                    scope: row.get(2)?,
119                    injections: row.get(3)?,
120                    hits: row.get(4)?,
121                    tokens_injected: row.get(5)?,
122                    hit_rate: row.get(6)?,
123                })
124            })?
125            .collect::<std::result::Result<Vec<_>, _>>()?;
126        Ok(results)
127    }
128}
129
130#[cfg(test)]
131mod tests {
132    use crate::store::Store;
133    use crate::types::SaveParams;
134
135    fn make_memory(store: &Store, key: &str) -> i64 {
136        store
137            .save(SaveParams {
138                key: key.to_string(),
139                value: "test value".to_string(),
140                ..Default::default()
141            })
142            .unwrap()
143            .id()
144    }
145
146    #[test]
147    fn test_record_injection_increments() {
148        let store = Store::open_in_memory().unwrap();
149        let id = make_memory(&store, "test/key");
150        store.record_injection(&[id], 0).unwrap();
151        store.record_injection(&[id], 0).unwrap();
152        let metrics = store.get_metrics().unwrap();
153        let m = metrics.iter().find(|m| m.id == id).unwrap();
154        assert_eq!(m.injections, 2);
155    }
156
157    #[test]
158    fn test_record_hit_increments() {
159        let store = Store::open_in_memory().unwrap();
160        let id = make_memory(&store, "test/key");
161        store.record_hit(id).unwrap();
162        store.record_hit(id).unwrap();
163        store.record_hit(id).unwrap();
164        let metrics = store.get_metrics().unwrap();
165        let m = metrics.iter().find(|m| m.id == id).unwrap();
166        assert_eq!(m.hits, 3);
167    }
168
169    #[test]
170    fn test_hit_rate_calculation() {
171        let store = Store::open_in_memory().unwrap();
172        let id = make_memory(&store, "test/key");
173        store.record_injection(&[id], 0).unwrap();
174        store.record_injection(&[id], 0).unwrap();
175        store.record_hit(id).unwrap();
176        let metrics = store.get_metrics().unwrap();
177        let m = metrics.iter().find(|m| m.id == id).unwrap();
178        assert!((m.hit_rate - 0.5).abs() < f64::EPSILON);
179    }
180
181    #[test]
182    fn test_record_injection_accumulates_tokens() {
183        let store = Store::open_in_memory().unwrap();
184        let id = make_memory(&store, "key/one");
185        store.record_injection(&[id], 50).unwrap();
186        store.record_injection(&[id], 50).unwrap();
187        let metrics = store.get_metrics().unwrap();
188        let m = metrics.iter().find(|m| m.id == id).unwrap();
189        assert_eq!(m.tokens_injected, 100);
190    }
191}
memory_core/store/metrics.rs

memory_core/store/
metrics.rs
