klieo-memory-sqlite 3.2.0

//! `SqliteShortTerm` — `ShortTermMemory` over a SQLite table.

use crate::connection::DbHandle;
use async_trait::async_trait;
use chrono::Utc;
use klieo_core::error::MemoryError;
use klieo_core::ids::ThreadId;
use klieo_core::llm::{Message, Role, ToolCall};
use klieo_core::memory::ShortTermMemory;

/// Extra rows fetched beyond the token budget so the in-Rust truncation always
/// has the full kept-suffix available even when messages are tiny (~1 token).
const LOAD_ROW_CAP_BUFFER: usize = 64;

/// SQLite-backed short-term conversation memory.
pub struct SqliteShortTerm {
    db: DbHandle,
}

impl SqliteShortTerm {
    pub(crate) fn new(db: DbHandle) -> Self {
        Self { db }
    }
}

fn role_to_str(r: Role) -> &'static str {
    match r {
        Role::System => "system",
        Role::User => "user",
        Role::Assistant => "assistant",
        Role::Tool => "tool",
        _ => "user",
    }
}

fn role_from_str(s: &str) -> Result<Role, MemoryError> {
    match s {
        "system" => Ok(Role::System),
        "user" => Ok(Role::User),
        "assistant" => Ok(Role::Assistant),
        "tool" => Ok(Role::Tool),
        other => Err(MemoryError::Serialization(format!("unknown role: {other}"))),
    }
}

#[async_trait]
impl ShortTermMemory for SqliteShortTerm {
    async fn append(&self, thread: ThreadId, msg: Message) -> Result<(), MemoryError> {
        let tool_calls_json = serde_json::to_string(&msg.tool_calls)
            .map_err(|e| MemoryError::Serialization(e.to_string()))?;
        let role = role_to_str(msg.role);
        let now = Utc::now().to_rfc3339();
        self.db
            .execute(move |conn| {
                let tx = conn.transaction()?;
                let next_seq: i64 = tx
                    .query_row(
                        "SELECT COALESCE(MAX(seq), 0) + 1 FROM short_term_messages WHERE thread_id = ?1",
                        rusqlite::params![&thread.0],
                        |r| r.get(0),
                    )?;
                tx.execute(
                    "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)",
                    rusqlite::params![
                        &thread.0,
                        next_seq,
                        role,
                        &msg.content,
                        &tool_calls_json,
                        &msg.tool_call_id,
                        &now,
                    ],
                )?;
                tx.commit()?;
                Ok(())
            })
            .await
    }

    async fn append_batch(
        &self,
        thread: ThreadId,
        messages: Vec<Message>,
    ) -> Result<(), MemoryError> {
        if messages.is_empty() {
            return Ok(());
        }
        // Serialize before opening the transaction so a bad message can't abort
        // a half-written batch. The whole batch then lands in one transaction
        // with a single MAX(seq) probe.
        let mut rows: Vec<(&'static str, String, String, Option<String>)> =
            Vec::with_capacity(messages.len());
        for msg in &messages {
            let tool_calls_json = serde_json::to_string(&msg.tool_calls)
                .map_err(|e| MemoryError::Serialization(e.to_string()))?;
            rows.push((
                role_to_str(msg.role),
                msg.content.clone(),
                tool_calls_json,
                msg.tool_call_id.clone(),
            ));
        }
        let now = Utc::now().to_rfc3339();
        self.db
            .execute(move |conn| {
                let tx = conn.transaction()?;
                let start_seq: i64 = tx.query_row(
                    "SELECT COALESCE(MAX(seq), 0) + 1 FROM short_term_messages WHERE thread_id = ?1",
                    rusqlite::params![&thread.0],
                    |r| r.get(0),
                )?;
                for (seq, (role, content, tool_calls_json, tool_call_id)) in (start_seq..).zip(&rows) {
                    tx.execute(
                        "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)",
                        rusqlite::params![
                            &thread.0, seq, role, content, tool_calls_json, tool_call_id, &now,
                        ],
                    )?;
                }
                tx.commit()?;
                Ok(())
            })
            .await
    }

    async fn load(&self, thread: ThreadId, max_tokens: usize) -> Result<Vec<Message>, MemoryError> {
        // Cap the row scan so a long-lived thread doesn't read every message
        // into memory before the token-budget truncation runs. Each message
        // costs >= 1 token, so at most `max_tokens` rows can survive the budget;
        // a small buffer keeps the cap safe.
        let row_cap =
            i64::try_from(max_tokens.saturating_add(LOAD_ROW_CAP_BUFFER)).unwrap_or(i64::MAX);
        let mut rows: Vec<(String, String, String, Option<String>)> = self
            .db
            .execute(move |conn| {
                let mut stmt = conn.prepare(
                    "SELECT role, content, tool_calls, tool_call_id FROM short_term_messages WHERE thread_id = ?1 ORDER BY seq DESC LIMIT ?2",
                )?;
                let iter = stmt.query_map(rusqlite::params![&thread.0, row_cap], |row| {
                    Ok((
                        row.get::<_, String>(0)?,
                        row.get::<_, String>(1)?,
                        row.get::<_, String>(2)?,
                        row.get::<_, Option<String>>(3)?,
                    ))
                })?;
                iter.collect::<Result<Vec<_>, _>>()
            })
            .await?;
        rows.reverse();

        let mut messages: Vec<Message> = rows
            .into_iter()
            .map(|(role, content, tool_calls_json, tool_call_id)| {
                let tool_calls: Vec<ToolCall> = serde_json::from_str(&tool_calls_json)
                    .map_err(|e| MemoryError::Serialization(e.to_string()))?;
                Ok(Message {
                    role: role_from_str(&role)?,
                    content,
                    tool_calls,
                    tool_call_id,
                })
            })
            .collect::<Result<Vec<_>, MemoryError>>()?;

        // Approximate token-budget truncation: ~4 chars per token. Walk
        // from newest to oldest accumulating cost; split off everything
        // older than the kept-suffix boundary in O(n).
        let total: usize = messages.iter().map(|m| (m.content.len() / 4).max(1)).sum();
        if total > max_tokens {
            let mut kept = 0usize;
            let mut keep_from = messages.len();
            for (idx, msg) in messages.iter().enumerate().rev() {
                let cost = (msg.content.len() / 4).max(1);
                if kept + cost > max_tokens {
                    break;
                }
                kept += cost;
                keep_from = idx;
            }
            messages = messages.split_off(keep_from);
        }
        Ok(messages)
    }

    async fn clear(&self, thread: ThreadId) -> Result<(), MemoryError> {
        self.db
            .execute(move |conn| {
                conn.execute(
                    "DELETE FROM short_term_messages WHERE thread_id = ?1",
                    rusqlite::params![&thread.0],
                )?;
                Ok(())
            })
            .await
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::connection::DbHandle;

    fn user(text: &str) -> Message {
        Message {
            role: Role::User,
            content: text.into(),
            tool_calls: vec![],
            tool_call_id: None,
        }
    }

    async fn fresh() -> SqliteShortTerm {
        let db = DbHandle::open(":memory:").await.unwrap();
        SqliteShortTerm::new(db)
    }

    #[tokio::test]
    async fn append_then_load_round_trips() {
        let m = fresh().await;
        let t = ThreadId::new("t1");
        m.append(t.clone(), user("hello")).await.unwrap();
        m.append(t.clone(), user("world")).await.unwrap();
        let loaded = m.load(t, 10_000).await.unwrap();
        assert_eq!(loaded.len(), 2);
        assert_eq!(loaded[0].content, "hello");
        assert_eq!(loaded[1].content, "world");
    }

    #[tokio::test]
    async fn append_batch_preserves_order_and_continues_seq() {
        let m = fresh().await;
        let t = ThreadId::new("t1");
        // A prior single append, then a batch — the batch must continue the
        // sequence after it, in order.
        m.append(t.clone(), user("first")).await.unwrap();
        m.append_batch(t.clone(), vec![user("second"), user("third")])
            .await
            .unwrap();
        let loaded = m.load(t, 10_000).await.unwrap();
        let contents: Vec<&str> = loaded.iter().map(|msg| msg.content.as_str()).collect();
        assert_eq!(contents, vec!["first", "second", "third"]);
    }

    #[tokio::test]
    async fn append_batch_with_empty_input_is_a_noop() {
        let m = fresh().await;
        let t = ThreadId::new("t1");
        m.append_batch(t.clone(), vec![]).await.unwrap();
        assert!(m.load(t, 10_000).await.unwrap().is_empty());
    }

    #[tokio::test]
    async fn load_caps_rows_at_token_budget_plus_buffer() {
        let m = fresh().await;
        let t = ThreadId::new("t1");
        // 200 single-char messages (~1 token each); a budget of 5 must keep
        // only the newest few, never load all 200.
        let batch: Vec<Message> = (0..200).map(|i| user(&i.to_string())).collect();
        m.append_batch(t.clone(), batch).await.unwrap();
        let loaded = m.load(t, 5).await.unwrap();
        assert!(
            loaded.len() <= 5 + LOAD_ROW_CAP_BUFFER,
            "load must bound the row scan; got {}",
            loaded.len()
        );
        // Far below the 200 written — proves the LIMIT actually fired.
        assert!(loaded.len() < 100, "cap must fire; got {}", loaded.len());
        // The kept suffix is the newest messages, in order.
        assert_eq!(loaded.last().unwrap().content, "199");
    }

    #[tokio::test]
    async fn load_truncates_to_token_budget() {
        let m = fresh().await;
        let t = ThreadId::new("t1");
        // Each ~40-char message ~= 10 tokens.
        for i in 0..20 {
            m.append(
                t.clone(),
                user(&format!("msg-{i:03}-padding-padding-padding")),
            )
            .await
            .unwrap();
        }
        let loaded = m.load(t, 30).await.unwrap();
        // Should keep ~3 messages (30 tokens / 10 each), oldest dropped.
        assert!(
            loaded.len() <= 4 && !loaded.is_empty(),
            "expected ~3 messages, got {}",
            loaded.len()
        );
        // Newest still present.
        assert!(
            loaded.last().unwrap().content.contains("msg-019"),
            "newest message must survive truncation"
        );
    }

    #[tokio::test]
    async fn clear_removes_thread() {
        let m = fresh().await;
        let t = ThreadId::new("t1");
        m.append(t.clone(), user("hello")).await.unwrap();
        m.clear(t.clone()).await.unwrap();
        let loaded = m.load(t, 10_000).await.unwrap();
        assert!(loaded.is_empty());
    }

    #[tokio::test]
    async fn threads_are_isolated() {
        let m = fresh().await;
        m.append(ThreadId::new("a"), user("a-msg")).await.unwrap();
        m.append(ThreadId::new("b"), user("b-msg")).await.unwrap();
        let a = m.load(ThreadId::new("a"), 10_000).await.unwrap();
        let b = m.load(ThreadId::new("b"), 10_000).await.unwrap();
        assert_eq!(a.len(), 1);
        assert_eq!(b.len(), 1);
        assert_eq!(a[0].content, "a-msg");
        assert_eq!(b[0].content, "b-msg");
    }

    #[tokio::test]
    async fn role_round_trip_for_all_variants() {
        let m = fresh().await;
        let t = ThreadId::new("r");
        for role in [Role::System, Role::User, Role::Assistant, Role::Tool] {
            m.append(
                t.clone(),
                Message {
                    role,
                    content: "x".into(),
                    tool_calls: vec![],
                    tool_call_id: None,
                },
            )
            .await
            .unwrap();
        }
        let loaded = m.load(t, 10_000).await.unwrap();
        assert_eq!(loaded.len(), 4);
        assert_eq!(loaded[0].role, Role::System);
        assert_eq!(loaded[3].role, Role::Tool);
    }
}