klieo_memory_sqlite/
short_term.rs

1//! `SqliteShortTerm` — `ShortTermMemory` over a SQLite table.
2
3use crate::connection::DbHandle;
4use async_trait::async_trait;
5use chrono::Utc;
6use klieo_core::error::MemoryError;
7use klieo_core::ids::ThreadId;
8use klieo_core::llm::{Message, Role, ToolCall};
9use klieo_core::memory::ShortTermMemory;
10
11/// Extra rows fetched beyond the token budget so the in-Rust truncation always
12/// has the full kept-suffix available even when messages are tiny (~1 token).
13const LOAD_ROW_CAP_BUFFER: usize = 64;
14
15/// SQLite-backed short-term conversation memory.
16pub struct SqliteShortTerm {
17    db: DbHandle,
18}
19
20impl SqliteShortTerm {
21    pub(crate) fn new(db: DbHandle) -> Self {
22        Self { db }
23    }
24}
25
26fn role_to_str(r: Role) -> &'static str {
27    match r {
28        Role::System => "system",
29        Role::User => "user",
30        Role::Assistant => "assistant",
31        Role::Tool => "tool",
32        _ => "user",
33    }
34}
35
36fn role_from_str(s: &str) -> Result<Role, MemoryError> {
37    match s {
38        "system" => Ok(Role::System),
39        "user" => Ok(Role::User),
40        "assistant" => Ok(Role::Assistant),
41        "tool" => Ok(Role::Tool),
42        other => Err(MemoryError::Serialization(format!("unknown role: {other}"))),
43    }
44}
45
46#[async_trait]
47impl ShortTermMemory for SqliteShortTerm {
48    async fn append(&self, thread: ThreadId, msg: Message) -> Result<(), MemoryError> {
49        let tool_calls_json = serde_json::to_string(&msg.tool_calls)
50            .map_err(|e| MemoryError::Serialization(e.to_string()))?;
51        let role = role_to_str(msg.role);
52        let now = Utc::now().to_rfc3339();
53        self.db
54            .execute(move |conn| {
55                let tx = conn.transaction()?;
56                let next_seq: i64 = tx
57                    .query_row(
58                        "SELECT COALESCE(MAX(seq), 0) + 1 FROM short_term_messages WHERE thread_id = ?1",
59                        rusqlite::params![&thread.0],
60                        |r| r.get(0),
61                    )?;
62                tx.execute(
63                    "INSERT INTO short_term_messages (thread_id, seq, role, content, tool_calls, tool_call_id, created_at) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
64                    rusqlite::params![
65                        &thread.0,
66                        next_seq,
67                        role,
68                        &msg.content,
69                        &tool_calls_json,
70                        &msg.tool_call_id,
71                        &now,
72                    ],
73                )?;
74                tx.commit()?;
75                Ok(())
76            })
77            .await
78    }
79
80    async fn append_batch(
81        &self,
82        thread: ThreadId,
83        messages: Vec<Message>,
84    ) -> Result<(), MemoryError> {
85        if messages.is_empty() {
86            return Ok(());
87        }
88        // Serialize before opening the transaction so a bad message can't abort
89        // a half-written batch. The whole batch then lands in one transaction
90        // with a single MAX(seq) probe.
91        let mut rows: Vec<(&'static str, String, String, Option<String>)> =
92            Vec::with_capacity(messages.len());
93        for msg in &messages {
94            let tool_calls_json = serde_json::to_string(&msg.tool_calls)
95                .map_err(|e| MemoryError::Serialization(e.to_string()))?;
96            rows.push((
97                role_to_str(msg.role),
98                msg.content.clone(),
99                tool_calls_json,
100                msg.tool_call_id.clone(),
101            ));
102        }
103        let now = Utc::now().to_rfc3339();
104        self.db
105            .execute(move |conn| {
106                let tx = conn.transaction()?;
107                let start_seq: i64 = tx.query_row(
108                    "SELECT COALESCE(MAX(seq), 0) + 1 FROM short_term_messages WHERE thread_id = ?1",
109                    rusqlite::params![&thread.0],
110                    |r| r.get(0),
111                )?;
112                for (seq, (role, content, tool_calls_json, tool_call_id)) in (start_seq..).zip(&rows) {
113                    tx.execute(
114                        "INSERT INTO short_term_messages (thread_id, seq, role, content, tool_calls, tool_call_id, created_at) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
115                        rusqlite::params![
116                            &thread.0, seq, role, content, tool_calls_json, tool_call_id, &now,
117                        ],
118                    )?;
119                }
120                tx.commit()?;
121                Ok(())
122            })
123            .await
124    }
125
126    async fn load(&self, thread: ThreadId, max_tokens: usize) -> Result<Vec<Message>, MemoryError> {
127        // Cap the row scan so a long-lived thread doesn't read every message
128        // into memory before the token-budget truncation runs. Each message
129        // costs >= 1 token, so at most `max_tokens` rows can survive the budget;
130        // a small buffer keeps the cap safe.
131        let row_cap =
132            i64::try_from(max_tokens.saturating_add(LOAD_ROW_CAP_BUFFER)).unwrap_or(i64::MAX);
133        let mut rows: Vec<(String, String, String, Option<String>)> = self
134            .db
135            .execute(move |conn| {
136                let mut stmt = conn.prepare(
137                    "SELECT role, content, tool_calls, tool_call_id FROM short_term_messages WHERE thread_id = ?1 ORDER BY seq DESC LIMIT ?2",
138                )?;
139                let iter = stmt.query_map(rusqlite::params![&thread.0, row_cap], |row| {
140                    Ok((
141                        row.get::<_, String>(0)?,
142                        row.get::<_, String>(1)?,
143                        row.get::<_, String>(2)?,
144                        row.get::<_, Option<String>>(3)?,
145                    ))
146                })?;
147                iter.collect::<Result<Vec<_>, _>>()
148            })
149            .await?;
150        rows.reverse();
151
152        let mut messages: Vec<Message> = rows
153            .into_iter()
154            .map(|(role, content, tool_calls_json, tool_call_id)| {
155                let tool_calls: Vec<ToolCall> = serde_json::from_str(&tool_calls_json)
156                    .map_err(|e| MemoryError::Serialization(e.to_string()))?;
157                Ok(Message {
158                    role: role_from_str(&role)?,
159                    content,
160                    tool_calls,
161                    tool_call_id,
162                })
163            })
164            .collect::<Result<Vec<_>, MemoryError>>()?;
165
166        // Approximate token-budget truncation: ~4 chars per token. Walk
167        // from newest to oldest accumulating cost; split off everything
168        // older than the kept-suffix boundary in O(n).
169        let total: usize = messages.iter().map(|m| (m.content.len() / 4).max(1)).sum();
170        if total > max_tokens {
171            let mut kept = 0usize;
172            let mut keep_from = messages.len();
173            for (idx, msg) in messages.iter().enumerate().rev() {
174                let cost = (msg.content.len() / 4).max(1);
175                if kept + cost > max_tokens {
176                    break;
177                }
178                kept += cost;
179                keep_from = idx;
180            }
181            messages = messages.split_off(keep_from);
182        }
183        Ok(messages)
184    }
185
186    async fn clear(&self, thread: ThreadId) -> Result<(), MemoryError> {
187        self.db
188            .execute(move |conn| {
189                conn.execute(
190                    "DELETE FROM short_term_messages WHERE thread_id = ?1",
191                    rusqlite::params![&thread.0],
192                )?;
193                Ok(())
194            })
195            .await
196    }
197}
198
199#[cfg(test)]
200mod tests {
201    use super::*;
202    use crate::connection::DbHandle;
203
204    fn user(text: &str) -> Message {
205        Message {
206            role: Role::User,
207            content: text.into(),
208            tool_calls: vec![],
209            tool_call_id: None,
210        }
211    }
212
213    async fn fresh() -> SqliteShortTerm {
214        let db = DbHandle::open(":memory:").await.unwrap();
215        SqliteShortTerm::new(db)
216    }
217
218    #[tokio::test]
219    async fn append_then_load_round_trips() {
220        let m = fresh().await;
221        let t = ThreadId::new("t1");
222        m.append(t.clone(), user("hello")).await.unwrap();
223        m.append(t.clone(), user("world")).await.unwrap();
224        let loaded = m.load(t, 10_000).await.unwrap();
225        assert_eq!(loaded.len(), 2);
226        assert_eq!(loaded[0].content, "hello");
227        assert_eq!(loaded[1].content, "world");
228    }
229
230    #[tokio::test]
231    async fn append_batch_preserves_order_and_continues_seq() {
232        let m = fresh().await;
233        let t = ThreadId::new("t1");
234        // A prior single append, then a batch — the batch must continue the
235        // sequence after it, in order.
236        m.append(t.clone(), user("first")).await.unwrap();
237        m.append_batch(t.clone(), vec![user("second"), user("third")])
238            .await
239            .unwrap();
240        let loaded = m.load(t, 10_000).await.unwrap();
241        let contents: Vec<&str> = loaded.iter().map(|msg| msg.content.as_str()).collect();
242        assert_eq!(contents, vec!["first", "second", "third"]);
243    }
244
245    #[tokio::test]
246    async fn append_batch_with_empty_input_is_a_noop() {
247        let m = fresh().await;
248        let t = ThreadId::new("t1");
249        m.append_batch(t.clone(), vec![]).await.unwrap();
250        assert!(m.load(t, 10_000).await.unwrap().is_empty());
251    }
252
253    #[tokio::test]
254    async fn load_caps_rows_at_token_budget_plus_buffer() {
255        let m = fresh().await;
256        let t = ThreadId::new("t1");
257        // 200 single-char messages (~1 token each); a budget of 5 must keep
258        // only the newest few, never load all 200.
259        let batch: Vec<Message> = (0..200).map(|i| user(&i.to_string())).collect();
260        m.append_batch(t.clone(), batch).await.unwrap();
261        let loaded = m.load(t, 5).await.unwrap();
262        assert!(
263            loaded.len() <= 5 + LOAD_ROW_CAP_BUFFER,
264            "load must bound the row scan; got {}",
265            loaded.len()
266        );
267        // Far below the 200 written — proves the LIMIT actually fired.
268        assert!(loaded.len() < 100, "cap must fire; got {}", loaded.len());
269        // The kept suffix is the newest messages, in order.
270        assert_eq!(loaded.last().unwrap().content, "199");
271    }
272
273    #[tokio::test]
274    async fn load_truncates_to_token_budget() {
275        let m = fresh().await;
276        let t = ThreadId::new("t1");
277        // Each ~40-char message ~= 10 tokens.
278        for i in 0..20 {
279            m.append(
280                t.clone(),
281                user(&format!("msg-{i:03}-padding-padding-padding")),
282            )
283            .await
284            .unwrap();
285        }
286        let loaded = m.load(t, 30).await.unwrap();
287        // Should keep ~3 messages (30 tokens / 10 each), oldest dropped.
288        assert!(
289            loaded.len() <= 4 && !loaded.is_empty(),
290            "expected ~3 messages, got {}",
291            loaded.len()
292        );
293        // Newest still present.
294        assert!(
295            loaded.last().unwrap().content.contains("msg-019"),
296            "newest message must survive truncation"
297        );
298    }
299
300    #[tokio::test]
301    async fn clear_removes_thread() {
302        let m = fresh().await;
303        let t = ThreadId::new("t1");
304        m.append(t.clone(), user("hello")).await.unwrap();
305        m.clear(t.clone()).await.unwrap();
306        let loaded = m.load(t, 10_000).await.unwrap();
307        assert!(loaded.is_empty());
308    }
309
310    #[tokio::test]
311    async fn threads_are_isolated() {
312        let m = fresh().await;
313        m.append(ThreadId::new("a"), user("a-msg")).await.unwrap();
314        m.append(ThreadId::new("b"), user("b-msg")).await.unwrap();
315        let a = m.load(ThreadId::new("a"), 10_000).await.unwrap();
316        let b = m.load(ThreadId::new("b"), 10_000).await.unwrap();
317        assert_eq!(a.len(), 1);
318        assert_eq!(b.len(), 1);
319        assert_eq!(a[0].content, "a-msg");
320        assert_eq!(b[0].content, "b-msg");
321    }
322
323    #[tokio::test]
324    async fn role_round_trip_for_all_variants() {
325        let m = fresh().await;
326        let t = ThreadId::new("r");
327        for role in [Role::System, Role::User, Role::Assistant, Role::Tool] {
328            m.append(
329                t.clone(),
330                Message {
331                    role,
332                    content: "x".into(),
333                    tool_calls: vec![],
334                    tool_call_id: None,
335                },
336            )
337            .await
338            .unwrap();
339        }
340        let loaded = m.load(t, 10_000).await.unwrap();
341        assert_eq!(loaded.len(), 4);
342        assert_eq!(loaded[0].role, Role::System);
343        assert_eq!(loaded[3].role, Role::Tool);
344    }
345}
klieo_memory_sqlite/short_term.rs

klieo_memory_sqlite/
short_term.rs
