cel-memory-duckdb 0.1.0

DuckDB embedded memory backend for AI agents, implementing cel-memory.
Documentation
//! Row mapping and query-filter helpers shared by the provider.

use cel_memory::{
    CallerScope, ChunkKind, ChunkSource, MemoryChunk, MemoryError, MemoryQuery, MemoryTier,
    Result as MemoryResult,
};
use chrono::{DateTime, Utc};
use duckdb::Row;

pub(crate) const EMBEDDING_DIM: usize = 384;

pub(crate) fn duck_err(e: duckdb::Error) -> MemoryError {
    MemoryError::Storage(e.to_string())
}

pub(crate) fn kind_str(k: ChunkKind) -> &'static str {
    match k {
        ChunkKind::Chat => "chat",
        ChunkKind::Action => "action",
        ChunkKind::Fire => "fire",
        ChunkKind::Observation => "observation",
        ChunkKind::Correction => "correction",
        ChunkKind::JobSummary => "job_summary",
        ChunkKind::Context => "context",
        ChunkKind::Rollup => "rollup",
    }
}

pub(crate) fn str_to_kind(s: &str) -> Option<ChunkKind> {
    match s {
        "chat" => Some(ChunkKind::Chat),
        "action" => Some(ChunkKind::Action),
        "fire" => Some(ChunkKind::Fire),
        "observation" => Some(ChunkKind::Observation),
        "correction" => Some(ChunkKind::Correction),
        "job_summary" => Some(ChunkKind::JobSummary),
        "context" => Some(ChunkKind::Context),
        "rollup" => Some(ChunkKind::Rollup),
        _ => None,
    }
}

pub(crate) fn tier_str(t: MemoryTier) -> &'static str {
    match t {
        MemoryTier::Session => "session",
        MemoryTier::LongTerm => "long_term",
    }
}

pub(crate) fn str_to_tier(s: &str) -> Option<MemoryTier> {
    match s {
        "session" => Some(MemoryTier::Session),
        "long_term" => Some(MemoryTier::LongTerm),
        _ => None,
    }
}

pub(crate) fn source_str(s: ChunkSource) -> &'static str {
    match s {
        ChunkSource::Embedded => "embedded",
        ChunkSource::Mcp => "mcp",
        ChunkSource::Gateway => "gateway",
        ChunkSource::Matcher => "matcher",
        ChunkSource::Perception => "cortex",
        ChunkSource::System => "system",
    }
}

pub(crate) fn str_to_source(s: &str) -> Option<ChunkSource> {
    match s {
        "embedded" => Some(ChunkSource::Embedded),
        "mcp" => Some(ChunkSource::Mcp),
        "gateway" => Some(ChunkSource::Gateway),
        "matcher" => Some(ChunkSource::Matcher),
        "cortex" | "perception" => Some(ChunkSource::Perception),
        "system" => Some(ChunkSource::System),
        _ => None,
    }
}

pub(crate) fn outcome_str(o: cel_memory::SessionOutcome) -> &'static str {
    use cel_memory::SessionOutcome;
    match o {
        SessionOutcome::Open => "open",
        SessionOutcome::Success => "success",
        SessionOutcome::Failure => "failure",
        SessionOutcome::Aborted => "aborted",
    }
}

pub(crate) fn str_to_outcome(s: &str) -> MemoryResult<cel_memory::SessionOutcome> {
    use cel_memory::SessionOutcome;
    match s {
        "open" => Ok(SessionOutcome::Open),
        "success" => Ok(SessionOutcome::Success),
        "failure" => Ok(SessionOutcome::Failure),
        "aborted" => Ok(SessionOutcome::Aborted),
        other => Err(MemoryError::Storage(format!("unknown outcome: {other}"))),
    }
}

pub(crate) fn row_to_chunk(row: &Row<'_>) -> duckdb::Result<MemoryChunk> {
    let metadata_raw: String = row.get(9)?;
    let metadata: serde_json::Value =
        serde_json::from_str(&metadata_raw).unwrap_or(serde_json::json!({}));
    Ok(MemoryChunk {
        id: row.get(0)?,
        created_at: row.get::<_, DateTime<Utc>>(1)?,
        kind: str_to_kind(&row.get::<_, String>(2)?).unwrap_or(ChunkKind::Chat),
        tier: str_to_tier(&row.get::<_, String>(3)?).unwrap_or(MemoryTier::Session),
        source: str_to_source(&row.get::<_, String>(4)?).unwrap_or(ChunkSource::System),
        session_id: row.get::<_, Option<String>>(5)?,
        project_root: row.get::<_, Option<String>>(6)?,
        caller_id: row.get(7)?,
        content: row.get(8)?,
        metadata,
        importance: row.get::<_, f32>(10)?,
        pinned: row.get::<_, bool>(11)?,
        shareable: row.get::<_, bool>(12)?,
        superseded_by: row.get::<_, Option<String>>(13)?,
        embedding_model: row.get(14)?,
        embedding_dim: row.get::<_, i32>(15)? as u32,
    })
}

pub(crate) fn row_to_session(row: &Row<'_>) -> duckdb::Result<cel_memory::MemorySession> {
    let metadata_raw: String = row.get(7)?;
    let metadata: serde_json::Value =
        serde_json::from_str(&metadata_raw).unwrap_or(serde_json::json!({}));
    let outcome = str_to_outcome(&row.get::<_, String>(6)?).map_err(|_| {
        duckdb::Error::InvalidColumnType(6, "outcome".into(), duckdb::types::Type::Text)
    })?;
    Ok(cel_memory::MemorySession {
        id: row.get(0)?,
        started_at: row.get::<_, DateTime<Utc>>(1)?,
        ended_at: row.get::<_, Option<DateTime<Utc>>>(2)?,
        caller_id: row.get(3)?,
        title: row.get::<_, Option<String>>(4)?,
        summary: row.get::<_, Option<String>>(5)?,
        outcome,
        metadata,
    })
}

pub(crate) fn chunk_matches_query(c: &MemoryChunk, q: &MemoryQuery) -> bool {
    if let Some(kinds) = &q.kinds {
        if !kinds.contains(&c.kind) {
            return false;
        }
    }
    if !q.include_rollups && c.kind == ChunkKind::Rollup {
        return false;
    }
    if let Some(since) = q.since {
        if c.created_at < since {
            return false;
        }
    }
    if let Some(until) = q.until {
        if c.created_at > until {
            return false;
        }
    }
    if let Some(sid) = &q.session_id {
        if c.session_id.as_deref() != Some(sid.as_str()) {
            return false;
        }
    }
    if let Some(prefix) = &q.project_root_prefix {
        match &c.project_root {
            Some(root) if root.starts_with(prefix.as_str()) => {}
            _ => return false,
        }
    }
    if let Some(min) = q.min_importance {
        if c.importance < min {
            return false;
        }
    }
    match q.caller_scope {
        CallerScope::Global => true,
        CallerScope::Own => c.caller_id == q.caller_id,
        CallerScope::OwnPlusShared => c.caller_id == q.caller_id || c.shareable,
    }
}

pub(crate) fn rrf(rank: usize, k: f32) -> f32 {
    1.0 / (k + rank as f32)
}

pub(crate) fn embedding_literal(values: &[f32]) -> String {
    format!(
        "[{}]",
        values
            .iter()
            .map(|v| v.to_string())
            .collect::<Vec<_>>()
            .join(", ")
    )
}

pub(crate) fn sql_string_literal(value: &str) -> String {
    format!("'{}'", value.replace('\'', "''"))
}

pub(crate) fn ids_in_clause(ids: &[String]) -> String {
    ids.iter()
        .map(|id| sql_string_literal(id))
        .collect::<Vec<_>>()
        .join(", ")
}

pub(crate) fn optional_row<T, F>(result: duckdb::Result<T>, f: F) -> MemoryResult<Option<T>>
where
    F: FnOnce(T) -> MemoryResult<T>,
{
    match result {
        Ok(value) => f(value).map(Some),
        Err(duckdb::Error::QueryReturnedNoRows) => Ok(None),
        Err(e) => Err(duck_err(e)),
    }
}