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 mut next_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 (role, content, tool_calls_json, tool_call_id) in &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, next_seq, role, content, tool_calls_json, tool_call_id, &now,
117                        ],
118                    )?;
119                    next_seq += 1;
120                }
121                tx.commit()?;
122                Ok(())
123            })
124            .await
125    }
126
127    async fn load(&self, thread: ThreadId, max_tokens: usize) -> Result<Vec<Message>, MemoryError> {
128        // Cap the row scan so a long-lived thread doesn't read every message
129        // into memory before the token-budget truncation runs. Each message
130        // costs >= 1 token, so at most `max_tokens` rows can survive the budget;
131        // a small buffer keeps the cap safe.
132        let row_cap =
133            i64::try_from(max_tokens.saturating_add(LOAD_ROW_CAP_BUFFER)).unwrap_or(i64::MAX);
134        let mut rows: Vec<(String, String, String, Option<String>)> = self
135            .db
136            .execute(move |conn| {
137                let mut stmt = conn.prepare(
138                    "SELECT role, content, tool_calls, tool_call_id FROM short_term_messages WHERE thread_id = ?1 ORDER BY seq DESC LIMIT ?2",
139                )?;
140                let iter = stmt.query_map(rusqlite::params![&thread.0, row_cap], |row| {
141                    Ok((
142                        row.get::<_, String>(0)?,
143                        row.get::<_, String>(1)?,
144                        row.get::<_, String>(2)?,
145                        row.get::<_, Option<String>>(3)?,
146                    ))
147                })?;
148                iter.collect::<Result<Vec<_>, _>>()
149            })
150            .await?;
151        rows.reverse();
152
153        let mut messages: Vec<Message> = rows
154            .into_iter()
155            .map(|(role, content, tool_calls_json, tool_call_id)| {
156                let tool_calls: Vec<ToolCall> = serde_json::from_str(&tool_calls_json)
157                    .map_err(|e| MemoryError::Serialization(e.to_string()))?;
158                Ok(Message {
159                    role: role_from_str(&role)?,
160                    content,
161                    tool_calls,
162                    tool_call_id,
163                })
164            })
165            .collect::<Result<Vec<_>, MemoryError>>()?;
166
167        // Approximate token-budget truncation: ~4 chars per token. Walk
168        // from newest to oldest accumulating cost; split off everything
169        // older than the kept-suffix boundary in O(n).
170        let total: usize = messages.iter().map(|m| (m.content.len() / 4).max(1)).sum();
171        if total > max_tokens {
172            let mut kept = 0usize;
173            let mut keep_from = messages.len();
174            for (idx, msg) in messages.iter().enumerate().rev() {
175                let cost = (msg.content.len() / 4).max(1);
176                if kept + cost > max_tokens {
177                    break;
178                }
179                kept += cost;
180                keep_from = idx;
181            }
182            messages = messages.split_off(keep_from);
183        }
184        Ok(messages)
185    }
186
187    async fn clear(&self, thread: ThreadId) -> Result<(), MemoryError> {
188        self.db
189            .execute(move |conn| {
190                conn.execute(
191                    "DELETE FROM short_term_messages WHERE thread_id = ?1",
192                    rusqlite::params![&thread.0],
193                )?;
194                Ok(())
195            })
196            .await
197    }
198}
199
200#[cfg(test)]
201mod tests {
202    use super::*;
203    use crate::connection::DbHandle;
204
205    fn user(text: &str) -> Message {
206        Message {
207            role: Role::User,
208            content: text.into(),
209            tool_calls: vec![],
210            tool_call_id: None,
211        }
212    }
213
214    async fn fresh() -> SqliteShortTerm {
215        let db = DbHandle::open(":memory:").await.unwrap();
216        SqliteShortTerm::new(db)
217    }
218
219    #[tokio::test]
220    async fn append_then_load_round_trips() {
221        let m = fresh().await;
222        let t = ThreadId::new("t1");
223        m.append(t.clone(), user("hello")).await.unwrap();
224        m.append(t.clone(), user("world")).await.unwrap();
225        let loaded = m.load(t, 10_000).await.unwrap();
226        assert_eq!(loaded.len(), 2);
227        assert_eq!(loaded[0].content, "hello");
228        assert_eq!(loaded[1].content, "world");
229    }
230
231    #[tokio::test]
232    async fn append_batch_preserves_order_and_continues_seq() {
233        let m = fresh().await;
234        let t = ThreadId::new("t1");
235        // A prior single append, then a batch — the batch must continue the
236        // sequence after it, in order.
237        m.append(t.clone(), user("first")).await.unwrap();
238        m.append_batch(t.clone(), vec![user("second"), user("third")])
239            .await
240            .unwrap();
241        let loaded = m.load(t, 10_000).await.unwrap();
242        let contents: Vec<&str> = loaded.iter().map(|msg| msg.content.as_str()).collect();
243        assert_eq!(contents, vec!["first", "second", "third"]);
244    }
245
246    #[tokio::test]
247    async fn append_batch_with_empty_input_is_a_noop() {
248        let m = fresh().await;
249        let t = ThreadId::new("t1");
250        m.append_batch(t.clone(), vec![]).await.unwrap();
251        assert!(m.load(t, 10_000).await.unwrap().is_empty());
252    }
253
254    #[tokio::test]
255    async fn load_caps_rows_at_token_budget_plus_buffer() {
256        let m = fresh().await;
257        let t = ThreadId::new("t1");
258        // 200 single-char messages (~1 token each); a budget of 5 must keep
259        // only the newest few, never load all 200.
260        let batch: Vec<Message> = (0..200).map(|i| user(&i.to_string())).collect();
261        m.append_batch(t.clone(), batch).await.unwrap();
262        let loaded = m.load(t, 5).await.unwrap();
263        assert!(
264            loaded.len() <= 5 + LOAD_ROW_CAP_BUFFER,
265            "load must bound the row scan; got {}",
266            loaded.len()
267        );
268        // Far below the 200 written — proves the LIMIT actually fired.
269        assert!(loaded.len() < 100, "cap must fire; got {}", loaded.len());
270        // The kept suffix is the newest messages, in order.
271        assert_eq!(loaded.last().unwrap().content, "199");
272    }
273
274    #[tokio::test]
275    async fn load_truncates_to_token_budget() {
276        let m = fresh().await;
277        let t = ThreadId::new("t1");
278        // Each ~40-char message ~= 10 tokens.
279        for i in 0..20 {
280            m.append(
281                t.clone(),
282                user(&format!("msg-{i:03}-padding-padding-padding")),
283            )
284            .await
285            .unwrap();
286        }
287        let loaded = m.load(t, 30).await.unwrap();
288        // Should keep ~3 messages (30 tokens / 10 each), oldest dropped.
289        assert!(
290            loaded.len() <= 4 && !loaded.is_empty(),
291            "expected ~3 messages, got {}",
292            loaded.len()
293        );
294        // Newest still present.
295        assert!(
296            loaded.last().unwrap().content.contains("msg-019"),
297            "newest message must survive truncation"
298        );
299    }
300
301    #[tokio::test]
302    async fn clear_removes_thread() {
303        let m = fresh().await;
304        let t = ThreadId::new("t1");
305        m.append(t.clone(), user("hello")).await.unwrap();
306        m.clear(t.clone()).await.unwrap();
307        let loaded = m.load(t, 10_000).await.unwrap();
308        assert!(loaded.is_empty());
309    }
310
311    #[tokio::test]
312    async fn threads_are_isolated() {
313        let m = fresh().await;
314        m.append(ThreadId::new("a"), user("a-msg")).await.unwrap();
315        m.append(ThreadId::new("b"), user("b-msg")).await.unwrap();
316        let a = m.load(ThreadId::new("a"), 10_000).await.unwrap();
317        let b = m.load(ThreadId::new("b"), 10_000).await.unwrap();
318        assert_eq!(a.len(), 1);
319        assert_eq!(b.len(), 1);
320        assert_eq!(a[0].content, "a-msg");
321        assert_eq!(b[0].content, "b-msg");
322    }
323
324    #[tokio::test]
325    async fn role_round_trip_for_all_variants() {
326        let m = fresh().await;
327        let t = ThreadId::new("r");
328        for role in [Role::System, Role::User, Role::Assistant, Role::Tool] {
329            m.append(
330                t.clone(),
331                Message {
332                    role,
333                    content: "x".into(),
334                    tool_calls: vec![],
335                    tool_call_id: None,
336                },
337            )
338            .await
339            .unwrap();
340        }
341        let loaded = m.load(t, 10_000).await.unwrap();
342        assert_eq!(loaded.len(), 4);
343        assert_eq!(loaded[0].role, Role::System);
344        assert_eq!(loaded[3].role, Role::Tool);
345    }
346}
klieo_memory_sqlite/short_term.rs

klieo_memory_sqlite/
short_term.rs
