clawhdf5/crates/clawhdf5-migrate/src/store_writer.rs

//! Write migrated SQLite data into a clawhdf5-agent store.
//!
//! Everything goes through `clawhdf5-agent`'s own API — `HDF5Memory::create`
//! (or `open` for `--incremental`), `save_batch`, `delete_batch`, the session
//! cache and the knowledge graph — so the result is an ordinary agent store
//! that `HDF5Memory::open` accepts, not a second hand-built copy of its schema.

use std::collections::{HashMap, HashSet};
use std::path::Path;

use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
use clawhdf5_format::float16::round_to_f16;

use crate::sqlite_reader::{MemoryChunk, SqliteData};

type BoxErr = Box<dyn std::error::Error>;

/// SQLite timestamps are Unix seconds; the agent's session and relation
/// timestamps are Unix microseconds (memory records stay in seconds).
pub const US_PER_SEC: f64 = 1_000_000.0;

/// Options controlling the output store.
#[derive(Debug, Clone)]
pub struct WriteOptions {
    pub agent_id: String,
    pub embedder: String,
    pub compression: bool,
    pub compression_level: u32,
    /// Store full-precision `f32` embeddings instead of the library default
    /// (half precision). Only applies to a newly created store: an existing
    /// store keeps the precision it was created with.
    pub f32: bool,
    /// Add to the store at the output path if there is one, instead of
    /// replacing it.
    pub incremental: bool,
    /// Leave out deleted source rows that are not in the store. (A deleted
    /// row that matches an active store record still tombstones it, so pass
    /// deleted rows in `data` for an incremental run.)
    pub skip_deleted: bool,
}

/// What the migration wrote, and where each source row went, so validation
/// can compare the store with the source row by row.
#[derive(Debug, Default)]
pub struct Migration {
    /// Whether the output store existed and was added to (`--incremental`).
    pub appended_to_existing: bool,
    /// The store's embedding precision.
    pub float16: bool,
    pub embedding_dim: usize,
    /// Records in the store after the migration (including tombstones).
    pub store_count: usize,
    /// `(store index, source chunk index)` of every record written.
    pub records: Vec<(usize, usize)>,
    /// Source chunks already in the store (incremental), not written again.
    pub chunks_present: usize,
    /// `(store index, source chunk index)` of records that were active in
    /// the store but whose source row is now deleted (incremental): they were
    /// tombstoned by this run.
    pub deleted_in_store: Vec<(usize, usize)>,
    /// Source rows that were deleted in the store but are active in the
    /// source (incremental): the agent has no un-delete, so each was written
    /// again as a new record (counted in `records` too).
    pub restored: usize,
    /// Deleted source rows left out because of `skip_deleted`.
    pub deleted_skipped: usize,
    /// `(store session index, source session index)` of each session written.
    pub sessions: Vec<(usize, usize)>,
    pub sessions_present: usize,
    /// `(store entity id, source entity index)` of each entity written.
    pub entities: Vec<(u64, usize)>,
    pub entities_present: usize,
    /// SQLite entity id -> store entity id, for every source entity.
    pub entity_ids: HashMap<i64, u64>,
    /// `(store relation index, source relation index)` of each relation written.
    pub relations: Vec<(usize, usize)>,
    pub relations_present: usize,
    /// Source relations naming an entity id that is not in the entities
    /// table; the knowledge graph cannot hold them, so they are skipped.
    pub dangling_relations: Vec<usize>,
    /// Messages of the write-anomaly alerts the agent raised while importing
    /// (informational; they never block a save — a bulk import typically
    /// trips the write-rate check).
    pub anomaly_alerts: Vec<String>,
}

/// Identity of a memory record for incremental de-duplication: every field
/// the agent stores except the embedding (whose stored form depends on the
/// store's precision).
type RecordKey = (String, String, String, String, u64);

fn record_key(
    chunk: &str,
    source_channel: &str,
    session_id: &str,
    tags: &str,
    ts: f64,
) -> RecordKey {
    (
        chunk.to_owned(),
        source_channel.to_owned(),
        session_id.to_owned(),
        tags.to_owned(),
        ts.to_bits(),
    )
}

/// Reject rows the agent would otherwise store differently from the source,
/// or not at all: an embedding of a different length from the store's
/// dimension (the agent pads/truncates silently), an empty embedding, or, in
/// a float16 store, a value beyond the half-precision range.
///
/// Every source row is checked, including ones that end up not being written
/// (already in the store, or deleted and skipped): the source must be
/// consistent as a whole, and the check runs before the store is touched.
fn check_chunks(chunks: &[MemoryChunk], dim: usize, float16: bool) -> Result<(), BoxErr> {
    for c in chunks {
        if c.embedding.is_empty() {
            return Err(format!(
                "chunk id {}: the embedding is empty; an agent store needs an embedding \
                 for every record",
                c.id
            )
            .into());
        }
        if c.embedding.len() != dim {
            return Err(format!(
                "chunk id {}: embedding has {} values, expected {dim}; every row must have \
                 the store's dimension (detected from the first row unless --embedding-dim \
                 is given), and rows are never truncated or padded to fit",
                c.id,
                c.embedding.len()
            )
            .into());
        }
        if float16
            && let Some((k, v)) = c
                .embedding
                .iter()
                .enumerate()
                .find(|&(_, &v)| v.is_finite() && round_to_f16(v).is_infinite())
        {
            return Err(format!(
                "chunk id {}: embedding[{k}] = {v} is outside the half-precision range \
                 (±65504) of a float16 store; migrate with --f32",
                c.id
            )
            .into());
        }
    }
    Ok(())
}

/// Migrate `data` into the agent store at `path`.
///
/// Without `opts.incremental` (or when nothing exists at `path`) a new store
/// is created, replacing any file there — but only once every source row has
/// passed [`check_chunks`], so a source that cannot be migrated leaves an
/// existing store untouched. With it, the existing store is opened and only
/// source rows it does not already hold are added: memory records are
/// matched on their content, sessions on their id, entities on name and
/// type, relations on (source, target, relation). A matched record then
/// takes the source row's deleted flag: see [`Migration::deleted_in_store`]
/// and [`Migration::restored`].
pub fn write_store(
    path: &Path,
    data: &SqliteData,
    opts: &WriteOptions,
) -> Result<Migration, BoxErr> {
    let existing = opts.incremental && path.exists();
    let mut mem = if existing {
        // `open` does not modify the store beyond what the agent itself does
        // on open; the checks below run before anything is written.
        let mem = HDF5Memory::open(path)?;
        let dim = mem.config().embedding_dim;
        // `data.embedding_dim` is 0 only for a source with no records and no
        // --embedding-dim, which has no dimension to disagree with.
        if data.embedding_dim != 0 && dim != data.embedding_dim {
            let hint = if dim == 0 {
                " (a store created from a source with no memory records; re-create it \
                 with --embedding-dim)"
            } else {
                ""
            };
            return Err(format!(
                "the store at {} has embedding_dim {dim}{hint}, the source {}; \
                 embeddings of a different dimension cannot be added to it",
                path.display(),
                data.embedding_dim
            )
            .into());
        }
        check_chunks(&data.chunks, dim, mem.config().float16)?;
        mem
    } else {
        // (With records, a dimension of 0 means an empty first embedding,
        // which `check_chunks` reports more precisely.)
        if data.embedding_dim == 0 && data.chunks.is_empty() {
            return Err(
                "the source has no memory records to detect the embedding dimension \
                        from; pass --embedding-dim (the dimension of the agent's embedder), \
                        or the store could never hold a record"
                    .into(),
            );
        }
        let mut config = MemoryConfig::new(path.to_path_buf(), &opts.agent_id, data.embedding_dim);
        config.embedder = opts.embedder.clone();
        config.compression = opts.compression;
        config.compression_level = opts.compression_level;
        // Only ever switch the library default off (as `clawhdf5-cli create`).
        if opts.f32 {
            config.float16 = false;
        }
        // Before `create`, which replaces whatever is at `path`.
        check_chunks(&data.chunks, config.embedding_dim, config.float16)?;
        HDF5Memory::create(config)?
    };
    let float16 = mem.config().float16;
    let dim = mem.config().embedding_dim;

    let mut m = Migration {
        appended_to_existing: existing,
        float16,
        embedding_dim: dim,
        ..Migration::default()
    };

    // ---- Memory records --------------------------------------------------
    // Store indices of every record the store already holds, by content, so
    // a source row that appears twice is only treated as present as often
    // as the store has it.
    let mut present: HashMap<RecordKey, Vec<usize>> = HashMap::new();
    if existing {
        let c = &mem.cache;
        for i in 0..c.len() {
            let key = record_key(
                &c.chunks[i],
                &c.source_channels[i],
                &c.session_ids[i],
                &c.tags[i],
                c.timestamps[i],
            );
            present.entry(key).or_default().push(i);
        }
    }
    let key_of = |c: &MemoryChunk| {
        record_key(
            &c.chunk,
            &c.source_channel,
            &c.session_id,
            &c.tags,
            c.timestamp,
        )
    };
    let tombstoned = |idx: usize| mem.cache.tombstones[idx] != 0;
    // Pass 1: a store record in the same deleted state as the source row.
    let mut unmatched: Vec<usize> = Vec::new();
    for (i, c) in data.chunks.iter().enumerate() {
        let src_deleted = c.deleted != 0;
        let hit = present.get_mut(&key_of(c)).and_then(|idxs| {
            let at = idxs.iter().position(|&x| tombstoned(x) == src_deleted)?;
            Some(idxs.remove(at))
        });
        match hit {
            Some(_) => m.chunks_present += 1,
            None => unmatched.push(i),
        }
    }
    // Pass 2: a store record whose deleted state differs — the source row
    // was deleted or restored since the last migration. The source wins.
    let mut new_chunks: Vec<usize> = Vec::with_capacity(unmatched.len());
    let mut delete_in_store: Vec<usize> = Vec::new();
    for i in unmatched {
        let c = &data.chunks[i];
        let hit = present
            .get_mut(&key_of(c))
            .and_then(|idxs| (!idxs.is_empty()).then(|| idxs.remove(0)));
        match hit {
            // Active in the store, deleted in the source: tombstone it.
            Some(idx) if c.deleted != 0 => {
                m.deleted_in_store.push((idx, i));
                delete_in_store.push(idx);
            }
            // Deleted in the store, active in the source. The agent has no
            // un-delete, so the row is written again as a new active record
            // (the tombstone stays until the store is compacted).
            Some(_) => {
                m.restored += 1;
                new_chunks.push(i);
            }
            None if c.deleted != 0 && opts.skip_deleted => m.deleted_skipped += 1,
            None => new_chunks.push(i),
        }
    }
    new_chunks.sort_unstable();
    let to_write: Vec<&MemoryChunk> = new_chunks.iter().map(|&i| &data.chunks[i]).collect();

    // ---- Sessions (in the cache; persisted by the save_batch checkpoint) ---
    let known_sessions: HashSet<String> = mem
        .sessions()
        .entries
        .iter()
        .map(|e| e.id.clone())
        .collect();
    for (i, s) in data.sessions.iter().enumerate() {
        if known_sessions.contains(&s.id) {
            m.sessions_present += 1;
            continue;
        }
        let sessions = mem.sessions_mut();
        let at = sessions.len();
        sessions.add_at(
            &s.id,
            s.start_idx.max(0) as usize,
            s.end_idx.max(0) as usize,
            &s.channel,
            &s.summary,
            s.timestamp * US_PER_SEC,
        );
        m.sessions.push((at, i));
    }

    // ---- Knowledge graph -------------------------------------------------
    let kg = mem.knowledge_mut();
    // Matched only against what the store held before this run: the source
    // itself is copied as it is, duplicates included.
    let by_name_type: HashMap<(String, String), u64> = kg
        .entities
        .iter()
        .map(|e| ((e.name.clone(), e.entity_type.clone()), e.id))
        .collect();
    for (i, e) in data.entities.iter().enumerate() {
        let key = (e.name.clone(), e.entity_type.clone());
        let id = match by_name_type.get(&key) {
            Some(&id) => {
                m.entities_present += 1;
                id
            }
            None => {
                let id = kg.add_entity(&e.name, &e.entity_type, e.embedding_idx);
                m.entities.push((id, i));
                id
            }
        };
        m.entity_ids.insert(e.id, id);
    }
    let known_relations: HashSet<(u64, u64, String)> = kg
        .relations
        .iter()
        .map(|r| (r.src, r.tgt, r.relation.clone()))
        .collect();
    for (i, r) in data.relations.iter().enumerate() {
        let (Some(&src), Some(&tgt)) = (m.entity_ids.get(&r.src), m.entity_ids.get(&r.tgt)) else {
            m.dangling_relations.push(i);
            continue;
        };
        if known_relations.contains(&(src, tgt, r.relation.clone())) {
            m.relations_present += 1;
            continue;
        }
        let at = kg.relations.len();
        kg.add_relation(src, tgt, &r.relation, r.weight as f32);
        kg.relations[at].ts = r.timestamp * US_PER_SEC;
        m.relations.push((at, i));
    }

    // ---- Write: one checkpoint for records, sessions and graph -----------
    let entries: Vec<MemoryEntry> = to_write
        .iter()
        .map(|c| MemoryEntry {
            chunk: c.chunk.clone(),
            embedding: c.embedding.clone(),
            source_channel: c.source_channel.clone(),
            timestamp: c.timestamp,
            session_id: c.session_id.clone(),
            tags: c.tags.clone(),
        })
        .collect();
    let indices = mem.save_batch(entries)?;
    m.records = indices
        .iter()
        .copied()
        .zip(new_chunks.iter().copied())
        .collect();

    // Rows deleted in the source stay deleted: tombstones, as the agent's own
    // `delete` leaves them (not compacted away).
    // Records matched in the store whose source row has since been deleted
    // are tombstoned too.
    let tombstones: Vec<usize> = m
        .records
        .iter()
        .filter(|&&(_, src)| data.chunks[src].deleted != 0)
        .map(|&(idx, _)| idx)
        .chain(delete_in_store)
        .collect();
    mem.delete_batch(&tombstones)?;

    m.anomaly_alerts = mem
        .take_anomaly_alerts()
        .into_iter()
        .map(|a| a.message)
        .collect();
    m.store_count = mem.count();
    drop(mem); // release the single-writer lock before anyone re-opens it
    Ok(m)
}