zeph_memory/store/
preferences.rs

1// SPDX-FileCopyrightText: 2026 Andrei G <bug-ops>
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4use zeph_common::text::truncate_to_bytes_ref;
5#[allow(unused_imports)]
6use zeph_db::sql;
7
8use super::SqliteStore;
9use crate::error::MemoryError;
10
11#[derive(Debug, Clone)]
12pub struct LearnedPreferenceRow {
13    pub id: i64,
14    pub preference_key: String,
15    pub preference_value: String,
16    pub confidence: f64,
17    pub evidence_count: i64,
18    pub updated_at: String,
19}
20
21// `evidence_count` is `INTEGER` (`INT4`) on Postgres, so it decodes as `i32`, not `i64` (INT8).
22type PreferenceTuple = (i64, String, String, f64, i32, String);
23
24fn row_from_tuple(t: PreferenceTuple) -> LearnedPreferenceRow {
25    LearnedPreferenceRow {
26        id: t.0,
27        preference_key: t.1,
28        preference_value: t.2,
29        confidence: t.3,
30        evidence_count: i64::from(t.4),
31        updated_at: t.5,
32    }
33}
34
35impl SqliteStore {
36    /// Insert or update a learned preference.
37    ///
38    /// When a key already exists, the value and metadata are updated and
39    /// `updated_at` is refreshed. `evidence_count` is set to the provided
40    /// value (caller is responsible for accumulation logic).
41    ///
42    /// Keys longer than 128 bytes or values longer than 256 bytes are silently
43    /// truncated at a UTF-8 character boundary before storage.
44    ///
45    /// # Errors
46    ///
47    /// Returns an error if the query fails.
48    pub async fn upsert_learned_preference(
49        &self,
50        key: &str,
51        value: &str,
52        confidence: f64,
53        evidence_count: i64,
54    ) -> Result<(), MemoryError> {
55        const MAX_KEY_BYTES: usize = 128;
56        const MAX_VALUE_BYTES: usize = 256;
57        let key_trunc = truncate_to_bytes_ref(key, MAX_KEY_BYTES);
58        let value_trunc = truncate_to_bytes_ref(value, MAX_VALUE_BYTES);
59        if key_trunc.len() < key.len() {
60            tracing::warn!(
61                original_len = key.len(),
62                "learned_preferences: key truncated to 128 bytes"
63            );
64        }
65        if value_trunc.len() < value.len() {
66            tracing::warn!(
67                original_len = value.len(),
68                "learned_preferences: value truncated to 256 bytes"
69            );
70        }
71        zeph_db::query(sql!(
72            "INSERT INTO learned_preferences \
73             (preference_key, preference_value, confidence, evidence_count, updated_at) \
74             VALUES (?, ?, ?, ?, CURRENT_TIMESTAMP) \
75             ON CONFLICT(preference_key) DO UPDATE SET \
76               preference_value = excluded.preference_value, \
77               confidence = excluded.confidence, \
78               evidence_count = excluded.evidence_count, \
79               updated_at = CURRENT_TIMESTAMP"
80        ))
81        .bind(key_trunc)
82        .bind(value_trunc)
83        .bind(confidence)
84        .bind(evidence_count)
85        .execute(&self.pool)
86        .await?;
87        Ok(())
88    }
89
90    /// Load all learned preferences, ordered by confidence descending.
91    ///
92    /// # Errors
93    ///
94    /// Returns an error if the query fails.
95    pub async fn load_learned_preferences(&self) -> Result<Vec<LearnedPreferenceRow>, MemoryError> {
96        // `updated_at` is `TIMESTAMPTZ` on Postgres (`TEXT` on SQLite); project through
97        // `Dialect::select_as_text` so it decodes into the `String` field below.
98        let updated_at_sel =
99            <zeph_db::ActiveDialect as zeph_db::dialect::Dialect>::select_as_text("updated_at");
100        let raw = format!(
101            "SELECT id, preference_key, preference_value, confidence, evidence_count, {updated_at_sel} \
102             FROM learned_preferences \
103             ORDER BY confidence DESC"
104        );
105        let query_sql = zeph_db::rewrite_placeholders(&raw);
106        let rows: Vec<PreferenceTuple> = zeph_db::query_as(sqlx::AssertSqlSafe(query_sql))
107            .fetch_all(&self.pool)
108            .await?;
109        Ok(rows.into_iter().map(row_from_tuple).collect())
110    }
111
112    /// Load corrections with `id > after_id`, ordered by id ascending.
113    ///
114    /// Used by the learning engine to process only new corrections since the
115    /// last analysis run (watermark-based incremental scan).
116    ///
117    /// # Errors
118    ///
119    /// Returns an error if the query fails.
120    pub async fn load_corrections_after(
121        &self,
122        after_id: i64,
123        limit: u32,
124    ) -> Result<Vec<super::corrections::UserCorrectionRow>, MemoryError> {
125        use super::corrections::UserCorrectionRow;
126
127        type Tuple = (
128            i64,
129            Option<i64>,
130            String,
131            String,
132            Option<String>,
133            String,
134            String,
135        );
136
137        // `created_at` is `TIMESTAMPTZ` on Postgres (`TEXT` on SQLite); project through
138        // `Dialect::select_as_text` so it decodes into the `String` field below.
139        let created_at_sel =
140            <zeph_db::ActiveDialect as zeph_db::dialect::Dialect>::select_as_text("created_at");
141        let raw = format!(
142            "SELECT id, session_id, original_output, correction_text, \
143             skill_name, correction_kind, {created_at_sel} \
144             FROM user_corrections \
145             WHERE id > ? \
146             ORDER BY id ASC LIMIT ?"
147        );
148        let query_sql = zeph_db::rewrite_placeholders(&raw);
149        let rows: Vec<Tuple> = zeph_db::query_as(sqlx::AssertSqlSafe(query_sql))
150            .bind(after_id)
151            .bind(i64::from(limit))
152            .fetch_all(&self.pool)
153            .await?;
154
155        Ok(rows
156            .into_iter()
157            .map(|t| UserCorrectionRow {
158                id: t.0,
159                session_id: t.1,
160                original_output: t.2,
161                correction_text: t.3,
162                skill_name: t.4,
163                correction_kind: t.5,
164                created_at: t.6,
165            })
166            .collect())
167    }
168}
169
170#[cfg(test)]
171mod tests {
172    use super::*;
173
174    async fn store() -> SqliteStore {
175        SqliteStore::new(":memory:").await.unwrap()
176    }
177
178    #[tokio::test]
179    async fn upsert_and_load() {
180        let s = store().await;
181        s.upsert_learned_preference("verbosity", "concise", 0.9, 5)
182            .await
183            .unwrap();
184        let rows = s.load_learned_preferences().await.unwrap();
185        assert_eq!(rows.len(), 1);
186        assert_eq!(rows[0].preference_key, "verbosity");
187        assert_eq!(rows[0].preference_value, "concise");
188        assert!((rows[0].confidence - 0.9).abs() < 1e-9);
189        assert_eq!(rows[0].evidence_count, 5);
190    }
191
192    #[tokio::test]
193    async fn upsert_updates_existing() {
194        let s = store().await;
195        s.upsert_learned_preference("verbosity", "concise", 0.8, 3)
196            .await
197            .unwrap();
198        s.upsert_learned_preference("verbosity", "verbose", 0.95, 8)
199            .await
200            .unwrap();
201        let rows = s.load_learned_preferences().await.unwrap();
202        assert_eq!(rows.len(), 1);
203        assert_eq!(rows[0].preference_value, "verbose");
204        assert!((rows[0].confidence - 0.95).abs() < 1e-9);
205        assert_eq!(rows[0].evidence_count, 8);
206    }
207
208    #[tokio::test]
209    async fn load_ordered_by_confidence() {
210        let s = store().await;
211        s.upsert_learned_preference("format_preference", "bullet points", 0.75, 3)
212            .await
213            .unwrap();
214        s.upsert_learned_preference("verbosity", "concise", 0.9, 5)
215            .await
216            .unwrap();
217        let rows = s.load_learned_preferences().await.unwrap();
218        assert_eq!(rows[0].preference_key, "verbosity");
219        assert_eq!(rows[1].preference_key, "format_preference");
220    }
221
222    #[tokio::test]
223    async fn load_corrections_after_watermark() {
224        let s = store().await;
225        // Insert two corrections
226        s.store_user_correction(None, "output", "be brief", None, "explicit_rejection")
227            .await
228            .unwrap();
229        let id2 = s
230            .store_user_correction(None, "output2", "use bullets", None, "alternative_request")
231            .await
232            .unwrap();
233        // Watermark at id2-1 => only id2 returned
234        let rows = s.load_corrections_after(id2 - 1, 10).await.unwrap();
235        assert_eq!(rows.len(), 1);
236        assert_eq!(rows[0].correction_text, "use bullets");
237    }
238}
zeph_memory/store/preferences.rs

zeph_memory/store/
preferences.rs
