klieo-memory-sqlite 0.38.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;

/// 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 load(&self, thread: ThreadId, max_tokens: usize) -> Result<Vec<Message>, MemoryError> {
        let 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 ASC",
                )?;
                let iter = stmt.query_map(rusqlite::params![&thread.0], |row| {
                    Ok((
                        row.get::<_, String>(0)?,
                        row.get::<_, String>(1)?,
                        row.get::<_, String>(2)?,
                        row.get::<_, Option<String>>(3)?,
                    ))
                })?;
                iter.collect::<Result<Vec<_>, _>>()
            })
            .await?;

        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 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);
    }
}