zeph_memory/store/
corrections.rs

1// SPDX-FileCopyrightText: 2026 Andrei G <bug-ops>
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4use super::SqliteStore;
5use crate::error::MemoryError;
6#[allow(unused_imports)]
7use zeph_db::sql;
8
9#[derive(Debug, Clone)]
10pub struct UserCorrectionRow {
11    pub id: i64,
12    pub session_id: Option<i64>,
13    pub original_output: String,
14    pub correction_text: String,
15    pub skill_name: Option<String>,
16    pub correction_kind: String,
17    pub created_at: String,
18}
19
20type CorrectionTuple = (
21    i64,
22    Option<i64>,
23    String,
24    String,
25    Option<String>,
26    String,
27    String,
28);
29
30fn row_from_tuple(t: CorrectionTuple) -> UserCorrectionRow {
31    UserCorrectionRow {
32        id: t.0,
33        session_id: t.1,
34        original_output: t.2,
35        correction_text: t.3,
36        skill_name: t.4,
37        correction_kind: t.5,
38        created_at: t.6,
39    }
40}
41
42impl SqliteStore {
43    /// Store a user correction and return the new row ID.
44    ///
45    /// # Errors
46    ///
47    /// Returns an error if the insert fails.
48    pub async fn store_user_correction(
49        &self,
50        session_id: Option<i64>,
51        original_output: &str,
52        correction_text: &str,
53        skill_name: Option<&str>,
54        correction_kind: &str,
55    ) -> Result<i64, MemoryError> {
56        let row: (i64,) = zeph_db::query_as(sql!(
57            "INSERT INTO user_corrections \
58             (session_id, original_output, correction_text, skill_name, correction_kind) \
59             VALUES (?, ?, ?, ?, ?) RETURNING id"
60        ))
61        .bind(session_id)
62        .bind(original_output)
63        .bind(correction_text)
64        .bind(skill_name)
65        .bind(correction_kind)
66        .fetch_one(&self.pool)
67        .await?;
68        Ok(row.0)
69    }
70
71    /// Load corrections for a specific skill, newest first.
72    ///
73    /// # Errors
74    ///
75    /// Returns an error if the query fails.
76    pub async fn load_corrections_for_skill(
77        &self,
78        skill_name: &str,
79        limit: u32,
80    ) -> Result<Vec<UserCorrectionRow>, MemoryError> {
81        // `created_at` is `TIMESTAMPTZ` on Postgres (`TEXT` on SQLite); project through
82        // `Dialect::select_as_text` so it decodes into the `String` field below.
83        let created_at_sel =
84            <zeph_db::ActiveDialect as zeph_db::dialect::Dialect>::select_as_text("created_at");
85        let raw = format!(
86            "SELECT id, session_id, original_output, correction_text, \
87             skill_name, correction_kind, {created_at_sel} \
88             FROM user_corrections WHERE skill_name = ? \
89             ORDER BY id DESC LIMIT ?"
90        );
91        let query_sql = zeph_db::rewrite_placeholders(&raw);
92        let rows: Vec<CorrectionTuple> = zeph_db::query_as(sqlx::AssertSqlSafe(query_sql))
93            .bind(skill_name)
94            .bind(i64::from(limit))
95            .fetch_all(&self.pool)
96            .await?;
97        Ok(rows.into_iter().map(row_from_tuple).collect())
98    }
99
100    /// Load the most recent corrections across all skills.
101    ///
102    /// # Errors
103    ///
104    /// Returns an error if the query fails.
105    pub async fn load_recent_corrections(
106        &self,
107        limit: u32,
108    ) -> Result<Vec<UserCorrectionRow>, MemoryError> {
109        // `created_at` is `TIMESTAMPTZ` on Postgres — see `load_corrections_for_skill`.
110        let created_at_sel =
111            <zeph_db::ActiveDialect as zeph_db::dialect::Dialect>::select_as_text("created_at");
112        let raw = format!(
113            "SELECT id, session_id, original_output, correction_text, \
114             skill_name, correction_kind, {created_at_sel} \
115             FROM user_corrections ORDER BY id DESC LIMIT ?"
116        );
117        let query_sql = zeph_db::rewrite_placeholders(&raw);
118        let rows: Vec<CorrectionTuple> = zeph_db::query_as(sqlx::AssertSqlSafe(query_sql))
119            .bind(i64::from(limit))
120            .fetch_all(&self.pool)
121            .await?;
122        Ok(rows.into_iter().map(row_from_tuple).collect())
123    }
124
125    /// Load a correction by ID (used by vector retrieval path).
126    ///
127    /// # Errors
128    ///
129    /// Returns an error if the query fails.
130    pub async fn load_corrections_for_id(
131        &self,
132        id: i64,
133    ) -> Result<Vec<UserCorrectionRow>, MemoryError> {
134        // `created_at` is `TIMESTAMPTZ` on Postgres — see `load_corrections_for_skill`.
135        let created_at_sel =
136            <zeph_db::ActiveDialect as zeph_db::dialect::Dialect>::select_as_text("created_at");
137        let raw = format!(
138            "SELECT id, session_id, original_output, correction_text, \
139             skill_name, correction_kind, {created_at_sel} \
140             FROM user_corrections WHERE id = ?"
141        );
142        let query_sql = zeph_db::rewrite_placeholders(&raw);
143        let rows: Vec<CorrectionTuple> = zeph_db::query_as(sqlx::AssertSqlSafe(query_sql))
144            .bind(id)
145            .fetch_all(&self.pool)
146            .await?;
147        Ok(rows.into_iter().map(row_from_tuple).collect())
148    }
149}
150
151#[cfg(test)]
152mod tests {
153    use super::*;
154
155    async fn test_store() -> SqliteStore {
156        SqliteStore::new(":memory:").await.unwrap()
157    }
158
159    #[tokio::test]
160    async fn store_and_load_correction() {
161        let store = test_store().await;
162
163        let id = store
164            .store_user_correction(
165                Some(1),
166                "original assistant output",
167                "that was wrong, try again",
168                Some("git"),
169                "explicit_rejection",
170            )
171            .await
172            .unwrap();
173        assert!(id > 0);
174
175        let rows = store.load_corrections_for_skill("git", 10).await.unwrap();
176        assert_eq!(rows.len(), 1);
177        assert_eq!(rows[0].correction_kind, "explicit_rejection");
178        assert_eq!(rows[0].skill_name.as_deref(), Some("git"));
179    }
180
181    #[tokio::test]
182    async fn load_recent_corrections_ordered() {
183        let store = test_store().await;
184
185        store
186            .store_user_correction(None, "out1", "fix1", None, "explicit_rejection")
187            .await
188            .unwrap();
189        store
190            .store_user_correction(None, "out2", "fix2", None, "alternative_request")
191            .await
192            .unwrap();
193
194        let rows = store.load_recent_corrections(10).await.unwrap();
195        assert_eq!(rows.len(), 2);
196        assert_eq!(rows[0].correction_text, "fix2");
197        assert_eq!(rows[1].correction_text, "fix1");
198    }
199
200    #[tokio::test]
201    async fn load_corrections_for_id_returns_single() {
202        let store = test_store().await;
203
204        let id = store
205            .store_user_correction(None, "out", "fix", Some("docker"), "repetition")
206            .await
207            .unwrap();
208
209        let rows = store.load_corrections_for_id(id).await.unwrap();
210        assert_eq!(rows.len(), 1);
211        assert_eq!(rows[0].id, id);
212    }
213
214    #[tokio::test]
215    async fn load_corrections_for_id_unknown_returns_empty() {
216        let store = test_store().await;
217        let rows = store.load_corrections_for_id(9999).await.unwrap();
218        assert!(rows.is_empty());
219    }
220
221    #[tokio::test]
222    async fn load_corrections_for_skill_unknown_returns_empty() {
223        let store = test_store().await;
224        let rows = store
225            .load_corrections_for_skill("nonexistent", 10)
226            .await
227            .unwrap();
228        assert!(rows.is_empty());
229    }
230
231    #[tokio::test]
232    async fn load_recent_corrections_empty_table() {
233        let store = test_store().await;
234        let rows = store.load_recent_corrections(10).await.unwrap();
235        assert!(rows.is_empty());
236    }
237
238    #[tokio::test]
239    async fn store_correction_without_skill_name() {
240        let store = test_store().await;
241
242        let id = store
243            .store_user_correction(
244                None,
245                "original output",
246                "correction text",
247                None,
248                "repetition",
249            )
250            .await
251            .unwrap();
252        assert!(id > 0);
253
254        let rows = store.load_recent_corrections(10).await.unwrap();
255        assert_eq!(rows.len(), 1);
256        assert!(rows[0].skill_name.is_none());
257        assert_eq!(rows[0].correction_kind, "repetition");
258    }
259
260    #[tokio::test]
261    async fn load_corrections_for_skill_respects_limit() {
262        let store = test_store().await;
263
264        for i in 0..5 {
265            store
266                .store_user_correction(
267                    None,
268                    &format!("out{i}"),
269                    &format!("fix{i}"),
270                    Some("git"),
271                    "explicit_rejection",
272                )
273                .await
274                .unwrap();
275        }
276
277        let rows = store.load_corrections_for_skill("git", 3).await.unwrap();
278        assert_eq!(rows.len(), 3);
279    }
280}
zeph_memory/store/corrections.rs

zeph_memory/store/
corrections.rs
