Omar Sobh
b08df7b628
Mission: 01a00c41-bac0-7eb3-a8c8-8b7044f3086d Phase: 01a00c41-bac2-71e3-a58b-c473421200ee Committed by the ClawMates delivery pipeline from the agents' working tree. Authored by agents, not by the named committer.
962 lines
34 KiB
Rust
962 lines
34 KiB
Rust
mod hdf5_reader;
|
|
mod hdf5_writer;
|
|
mod sqlite_reader;
|
|
mod validate;
|
|
|
|
use clap::Parser;
|
|
|
|
use sqlite_reader::SchemaConfig;
|
|
|
|
/// Migrate ZeroClaw agent memory from SQLite to HDF5 format.
|
|
#[derive(Parser, Debug)]
|
|
#[command(name = "clawhdf5-migrate", version, about)]
|
|
struct Cli {
|
|
/// Source SQLite database path
|
|
#[arg(long)]
|
|
sqlite: String,
|
|
|
|
/// Destination HDF5 file path
|
|
#[arg(long)]
|
|
hdf5: String,
|
|
|
|
/// Agent ID for metadata
|
|
#[arg(long, default_value = "migrated")]
|
|
agent_id: String,
|
|
|
|
/// Embedder name for metadata
|
|
#[arg(long, default_value = "unknown")]
|
|
embedder: String,
|
|
|
|
/// Embedding dimension (auto-detect from first row if not specified)
|
|
#[arg(long)]
|
|
embedding_dim: Option<usize>,
|
|
|
|
/// Skip deleted/tombstoned entries
|
|
#[arg(long)]
|
|
skip_deleted: bool,
|
|
|
|
/// Enable deflate compression on embeddings
|
|
#[arg(long)]
|
|
compression: bool,
|
|
|
|
/// Compression level 1-9
|
|
#[arg(long, default_value_t = 4)]
|
|
compression_level: u32,
|
|
|
|
/// Store embeddings as float16 (halves storage)
|
|
#[arg(long)]
|
|
float16: bool,
|
|
|
|
/// Validate without writing
|
|
#[arg(long)]
|
|
dry_run: bool,
|
|
|
|
/// Content-check every migrated row (default: a representative sample)
|
|
#[arg(long)]
|
|
validate_full: bool,
|
|
|
|
/// Append only rows newer than the existing output (by chunk id), merging
|
|
/// into the file at --hdf5 if it exists
|
|
#[arg(long)]
|
|
incremental: bool,
|
|
|
|
/// Override the SQLite table name for memory chunks
|
|
#[arg(long)]
|
|
chunks_table: Option<String>,
|
|
|
|
/// Override the SQLite table name for sessions
|
|
#[arg(long)]
|
|
sessions_table: Option<String>,
|
|
|
|
/// Override the SQLite table name for entities
|
|
#[arg(long)]
|
|
entities_table: Option<String>,
|
|
|
|
/// Override the SQLite table name for relations
|
|
#[arg(long)]
|
|
relations_table: Option<String>,
|
|
|
|
/// Print progress
|
|
#[arg(long)]
|
|
verbose: bool,
|
|
}
|
|
|
|
/// Build the schema config from CLI table-name overrides (defaults otherwise).
|
|
fn schema_from_cli(cli: &Cli) -> SchemaConfig {
|
|
let mut c = SchemaConfig::default();
|
|
if let Some(t) = &cli.chunks_table {
|
|
c.chunks.table = t.clone();
|
|
}
|
|
if let Some(t) = &cli.sessions_table {
|
|
c.sessions.table = t.clone();
|
|
}
|
|
if let Some(t) = &cli.entities_table {
|
|
c.entities.table = t.clone();
|
|
}
|
|
if let Some(t) = &cli.relations_table {
|
|
c.relations.table = t.clone();
|
|
}
|
|
c
|
|
}
|
|
|
|
fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|
let cli = Cli::parse();
|
|
let schema = schema_from_cli(&cli);
|
|
|
|
// Dry run: a fast count-only pass that does not buffer the database.
|
|
if cli.dry_run {
|
|
let counts = sqlite_reader::read_counts(&cli.sqlite, cli.skip_deleted, &schema)?;
|
|
eprintln!("Dry run — no output file written.");
|
|
eprintln!(
|
|
"Would migrate: {} chunks, {} sessions, {} entities, {} relations",
|
|
counts.chunks, counts.sessions, counts.entities, counts.relations
|
|
);
|
|
return Ok(());
|
|
}
|
|
|
|
if cli.verbose {
|
|
eprintln!("Reading SQLite database: {}", cli.sqlite);
|
|
}
|
|
|
|
// Incremental: merge new rows into the existing output (if present).
|
|
let incremental_base = if cli.incremental && std::path::Path::new(&cli.hdf5).exists() {
|
|
Some(hdf5_reader::read_hdf5(&cli.hdf5)?)
|
|
} else {
|
|
None
|
|
};
|
|
let min_chunk_id = incremental_base
|
|
.as_ref()
|
|
.map(|d| d.chunks.iter().map(|c| c.id).max().unwrap_or(0))
|
|
.unwrap_or(0);
|
|
let dim_hint = cli
|
|
.embedding_dim
|
|
.or_else(|| incremental_base.as_ref().map(|d| d.embedding_dim));
|
|
|
|
let source = if min_chunk_id > 0 {
|
|
sqlite_reader::read_sqlite_filtered(
|
|
&cli.sqlite,
|
|
cli.skip_deleted,
|
|
dim_hint,
|
|
&schema,
|
|
min_chunk_id,
|
|
)?
|
|
} else {
|
|
sqlite_reader::read_sqlite(&cli.sqlite, cli.skip_deleted, dim_hint, &schema)?
|
|
};
|
|
|
|
// Build the dataset to write: either the source alone, or the existing
|
|
// output plus the newly-read rows (metadata groups refreshed from source).
|
|
let data = match incremental_base {
|
|
Some(mut base) => {
|
|
let added = source.chunks.len();
|
|
base.chunks.extend(source.chunks);
|
|
base.sessions = source.sessions;
|
|
base.entities = source.entities;
|
|
base.relations = source.relations;
|
|
base.embedding_dim = source.embedding_dim.max(base.embedding_dim);
|
|
// Carry the current run's real SQLite source forward for
|
|
// provenance — `base` (re-read from the prior HDF5 output) has
|
|
// no meaningful source_path of its own.
|
|
base.source_path = source.source_path;
|
|
if cli.verbose {
|
|
eprintln!("Incremental: appended {added} new chunks (id > {min_chunk_id})");
|
|
}
|
|
base
|
|
}
|
|
None => source,
|
|
};
|
|
|
|
if cli.verbose {
|
|
eprintln!(
|
|
"Migrating {} chunks, {} sessions, {} entities, {} relations (dim={})",
|
|
data.chunks.len(),
|
|
data.sessions.len(),
|
|
data.entities.len(),
|
|
data.relations.len(),
|
|
data.embedding_dim
|
|
);
|
|
eprintln!("Writing HDF5 file: {}", cli.hdf5);
|
|
}
|
|
|
|
let opts = hdf5_writer::WriteOptions {
|
|
agent_id: cli.agent_id,
|
|
embedder: cli.embedder,
|
|
compression: cli.compression,
|
|
compression_level: cli.compression_level.clamp(1, 9),
|
|
float16: cli.float16,
|
|
};
|
|
|
|
hdf5_writer::write_hdf5(&cli.hdf5, &data, &opts)?;
|
|
|
|
if cli.verbose {
|
|
eprintln!("Validating output (content check)...");
|
|
}
|
|
|
|
let summary = validate::validate_hdf5(&cli.hdf5, &data, cli.validate_full, cli.float16)?;
|
|
|
|
eprintln!(
|
|
"Migration complete: {} chunks, {} sessions, {} entities, {} relations (dim={}); {} rows content-verified",
|
|
summary.chunks,
|
|
summary.sessions,
|
|
summary.entities,
|
|
summary.relations,
|
|
summary.embedding_dim,
|
|
summary.rows_checked,
|
|
);
|
|
if summary.provenance_verified {
|
|
eprintln!("Provenance: chunks/text and chunks/embeddings SHA-256 hashes verified.");
|
|
} else if cli.verbose {
|
|
eprintln!("Provenance: no provenance hash found to verify (older output format?).");
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use rusqlite::Connection;
|
|
use tempfile::TempDir;
|
|
|
|
/// Create a test SQLite database with the ZeroClaw schema.
|
|
fn create_test_db(dir: &TempDir) -> String {
|
|
let db_path = dir.path().join("test.db");
|
|
let path_str = db_path.to_str().unwrap().to_string();
|
|
let conn = Connection::open(&path_str).unwrap();
|
|
conn.execute_batch(
|
|
"CREATE TABLE memory_chunks (
|
|
id INTEGER PRIMARY KEY,
|
|
chunk TEXT NOT NULL,
|
|
embedding BLOB NOT NULL,
|
|
source_channel TEXT DEFAULT 'api',
|
|
timestamp REAL NOT NULL,
|
|
session_id TEXT,
|
|
tags TEXT DEFAULT '',
|
|
deleted INTEGER DEFAULT 0
|
|
);
|
|
CREATE TABLE sessions (
|
|
id TEXT PRIMARY KEY,
|
|
start_idx INTEGER,
|
|
end_idx INTEGER,
|
|
channel TEXT,
|
|
timestamp REAL,
|
|
summary TEXT
|
|
);
|
|
CREATE TABLE entities (
|
|
id INTEGER PRIMARY KEY,
|
|
name TEXT NOT NULL,
|
|
type TEXT NOT NULL,
|
|
embedding_idx INTEGER DEFAULT -1
|
|
);
|
|
CREATE TABLE relations (
|
|
src INTEGER NOT NULL,
|
|
tgt INTEGER NOT NULL,
|
|
relation TEXT NOT NULL,
|
|
weight REAL DEFAULT 1.0,
|
|
timestamp REAL,
|
|
FOREIGN KEY (src) REFERENCES entities(id),
|
|
FOREIGN KEY (tgt) REFERENCES entities(id)
|
|
);",
|
|
)
|
|
.unwrap();
|
|
path_str
|
|
}
|
|
|
|
/// Insert a memory chunk with a known embedding.
|
|
fn insert_chunk(conn: &Connection, id: i64, text: &str, embedding: &[f32], deleted: i32) {
|
|
let blob: Vec<u8> = embedding.iter().flat_map(|v| v.to_le_bytes()).collect();
|
|
conn.execute(
|
|
"INSERT INTO memory_chunks (id, chunk, embedding, source_channel, timestamp, session_id, tags, deleted)
|
|
VALUES (?1, ?2, ?3, 'api', 1700000000.0, 'sess-1', 'tag1,tag2', ?4)",
|
|
rusqlite::params![id, text, blob, deleted],
|
|
)
|
|
.unwrap();
|
|
}
|
|
|
|
fn insert_session(conn: &Connection, id: &str, start: i64, end: i64) {
|
|
conn.execute(
|
|
"INSERT INTO sessions (id, start_idx, end_idx, channel, timestamp, summary)
|
|
VALUES (?1, ?2, ?3, 'discord', 1700000000.0, 'test summary')",
|
|
rusqlite::params![id, start, end],
|
|
)
|
|
.unwrap();
|
|
}
|
|
|
|
fn insert_entity(conn: &Connection, id: i64, name: &str, etype: &str) {
|
|
conn.execute(
|
|
"INSERT INTO entities (id, name, type, embedding_idx) VALUES (?1, ?2, ?3, -1)",
|
|
rusqlite::params![id, name, etype],
|
|
)
|
|
.unwrap();
|
|
}
|
|
|
|
fn insert_relation(conn: &Connection, src: i64, tgt: i64, rel: &str) {
|
|
conn.execute(
|
|
"INSERT INTO relations (src, tgt, relation, weight, timestamp)
|
|
VALUES (?1, ?2, ?3, 1.0, 1700000000.0)",
|
|
rusqlite::params![src, tgt, rel],
|
|
)
|
|
.unwrap();
|
|
}
|
|
|
|
fn make_embedding(dim: usize, seed: f32) -> Vec<f32> {
|
|
(0..dim).map(|i| seed + i as f32 * 0.1).collect()
|
|
}
|
|
|
|
// ---------- Test 1: Basic end-to-end migration ----------
|
|
#[test]
|
|
fn test_basic_migration() {
|
|
let dir = TempDir::new().unwrap();
|
|
let db_path = create_test_db(&dir);
|
|
let h5_path = dir.path().join("out.h5");
|
|
|
|
let conn = Connection::open(&db_path).unwrap();
|
|
insert_chunk(&conn, 1, "hello world", &make_embedding(8, 1.0), 0);
|
|
insert_chunk(&conn, 2, "goodbye world", &make_embedding(8, 2.0), 0);
|
|
insert_session(&conn, "s1", 0, 1);
|
|
insert_entity(&conn, 1, "Alice", "person");
|
|
insert_relation(&conn, 1, 1, "self");
|
|
drop(conn);
|
|
|
|
let data =
|
|
sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
|
let opts = hdf5_writer::WriteOptions {
|
|
agent_id: "test-agent".into(),
|
|
embedder: "test-embed".into(),
|
|
compression: false,
|
|
compression_level: 4,
|
|
float16: false,
|
|
};
|
|
hdf5_writer::write_hdf5(h5_path.to_str().unwrap(), &data, &opts).unwrap();
|
|
|
|
let summary =
|
|
validate::validate_hdf5(h5_path.to_str().unwrap(), &data, false, false).unwrap();
|
|
assert_eq!(summary.chunks, 2);
|
|
assert_eq!(summary.sessions, 1);
|
|
assert_eq!(summary.entities, 1);
|
|
assert_eq!(summary.relations, 1);
|
|
assert_eq!(summary.embedding_dim, 8);
|
|
}
|
|
|
|
// ---------- Test 2: Skip deleted rows ----------
|
|
#[test]
|
|
fn test_skip_deleted() {
|
|
let dir = TempDir::new().unwrap();
|
|
let db_path = create_test_db(&dir);
|
|
let h5_path = dir.path().join("out.h5");
|
|
|
|
let conn = Connection::open(&db_path).unwrap();
|
|
insert_chunk(&conn, 1, "active", &make_embedding(4, 1.0), 0);
|
|
insert_chunk(&conn, 2, "deleted", &make_embedding(4, 2.0), 1);
|
|
insert_chunk(&conn, 3, "also active", &make_embedding(4, 3.0), 0);
|
|
drop(conn);
|
|
|
|
let data =
|
|
sqlite_reader::read_sqlite(&db_path, true, None, &SchemaConfig::default()).unwrap();
|
|
assert_eq!(data.chunks.len(), 2);
|
|
|
|
let opts = hdf5_writer::WriteOptions {
|
|
agent_id: "t".into(),
|
|
embedder: "t".into(),
|
|
compression: false,
|
|
compression_level: 4,
|
|
float16: false,
|
|
};
|
|
hdf5_writer::write_hdf5(h5_path.to_str().unwrap(), &data, &opts).unwrap();
|
|
|
|
let summary =
|
|
validate::validate_hdf5(h5_path.to_str().unwrap(), &data, false, false).unwrap();
|
|
assert_eq!(summary.chunks, 2);
|
|
}
|
|
|
|
// ---------- Test 3: Include deleted rows ----------
|
|
#[test]
|
|
fn test_include_deleted() {
|
|
let dir = TempDir::new().unwrap();
|
|
let db_path = create_test_db(&dir);
|
|
|
|
let conn = Connection::open(&db_path).unwrap();
|
|
insert_chunk(&conn, 1, "active", &make_embedding(4, 1.0), 0);
|
|
insert_chunk(&conn, 2, "deleted", &make_embedding(4, 2.0), 1);
|
|
drop(conn);
|
|
|
|
let data =
|
|
sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
|
assert_eq!(data.chunks.len(), 2);
|
|
}
|
|
|
|
// ---------- Test 4: Auto-detect embedding dimension ----------
|
|
#[test]
|
|
fn test_auto_detect_dim() {
|
|
let dir = TempDir::new().unwrap();
|
|
let db_path = create_test_db(&dir);
|
|
|
|
let conn = Connection::open(&db_path).unwrap();
|
|
insert_chunk(&conn, 1, "test", &make_embedding(16, 0.5), 0);
|
|
drop(conn);
|
|
|
|
let data =
|
|
sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
|
assert_eq!(data.embedding_dim, 16);
|
|
}
|
|
|
|
// ---------- Test 5: Manual embedding dimension ----------
|
|
#[test]
|
|
fn test_manual_dim() {
|
|
let dir = TempDir::new().unwrap();
|
|
let db_path = create_test_db(&dir);
|
|
|
|
let conn = Connection::open(&db_path).unwrap();
|
|
insert_chunk(&conn, 1, "test", &make_embedding(16, 0.5), 0);
|
|
drop(conn);
|
|
|
|
let data =
|
|
sqlite_reader::read_sqlite(&db_path, false, Some(8), &SchemaConfig::default()).unwrap();
|
|
assert_eq!(data.embedding_dim, 8);
|
|
// Embedding truncated to dim 8
|
|
assert_eq!(data.chunks[0].embedding.len(), 8);
|
|
}
|
|
|
|
// ---------- Test 6: Float16 conversion ----------
|
|
#[test]
|
|
fn test_float16_conversion() {
|
|
let dir = TempDir::new().unwrap();
|
|
let db_path = create_test_db(&dir);
|
|
let h5_path = dir.path().join("out.h5");
|
|
|
|
let conn = Connection::open(&db_path).unwrap();
|
|
let emb = vec![1.0f32, 2.5, -0.5, 3.125];
|
|
insert_chunk(&conn, 1, "test", &emb, 0);
|
|
drop(conn);
|
|
|
|
let data =
|
|
sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
|
let opts = hdf5_writer::WriteOptions {
|
|
agent_id: "t".into(),
|
|
embedder: "t".into(),
|
|
compression: false,
|
|
compression_level: 4,
|
|
float16: true,
|
|
};
|
|
hdf5_writer::write_hdf5(h5_path.to_str().unwrap(), &data, &opts).unwrap();
|
|
|
|
// Content-validate with the float16 tolerance enabled.
|
|
let summary =
|
|
validate::validate_hdf5(h5_path.to_str().unwrap(), &data, true, true).unwrap();
|
|
assert_eq!(summary.chunks, 1);
|
|
|
|
// Verify float16 values are within tolerance
|
|
for &v in &emb {
|
|
let f16 = half::f16::from_f32(v);
|
|
let roundtrip = f16.to_f32();
|
|
assert!(
|
|
(v - roundtrip).abs() < 0.01,
|
|
"f16 roundtrip too lossy for {v}"
|
|
);
|
|
}
|
|
}
|
|
|
|
// ---------- Test 7: Compression produces valid file ----------
|
|
#[test]
|
|
fn test_compression() {
|
|
let dir = TempDir::new().unwrap();
|
|
let db_path = create_test_db(&dir);
|
|
let h5_compressed = dir.path().join("compressed.h5");
|
|
let h5_uncompressed = dir.path().join("uncompressed.h5");
|
|
|
|
let conn = Connection::open(&db_path).unwrap();
|
|
// Insert enough data so compression can be effective
|
|
for i in 0..100 {
|
|
insert_chunk(&conn, i, &format!("chunk {i}"), &make_embedding(32, 0.0), 0);
|
|
}
|
|
drop(conn);
|
|
|
|
let data =
|
|
sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
|
|
|
let opts_compressed = hdf5_writer::WriteOptions {
|
|
agent_id: "t".into(),
|
|
embedder: "t".into(),
|
|
compression: true,
|
|
compression_level: 6,
|
|
float16: false,
|
|
};
|
|
hdf5_writer::write_hdf5(h5_compressed.to_str().unwrap(), &data, &opts_compressed).unwrap();
|
|
|
|
let opts_plain = hdf5_writer::WriteOptions {
|
|
agent_id: "t".into(),
|
|
embedder: "t".into(),
|
|
compression: false,
|
|
compression_level: 4,
|
|
float16: false,
|
|
};
|
|
hdf5_writer::write_hdf5(h5_uncompressed.to_str().unwrap(), &data, &opts_plain).unwrap();
|
|
|
|
let sz_c = std::fs::metadata(&h5_compressed).unwrap().len();
|
|
let sz_u = std::fs::metadata(&h5_uncompressed).unwrap().len();
|
|
assert!(
|
|
sz_c < sz_u,
|
|
"Compressed ({sz_c}) should be smaller than uncompressed ({sz_u})"
|
|
);
|
|
}
|
|
|
|
// ---------- Test 8: Dry run doesn't create file ----------
|
|
#[test]
|
|
fn test_dry_run() {
|
|
let dir = TempDir::new().unwrap();
|
|
let db_path = create_test_db(&dir);
|
|
let h5_path = dir.path().join("should_not_exist.h5");
|
|
|
|
let conn = Connection::open(&db_path).unwrap();
|
|
insert_chunk(&conn, 1, "test", &make_embedding(4, 1.0), 0);
|
|
drop(conn);
|
|
|
|
// Simulate dry-run: read data but don't write
|
|
let data =
|
|
sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
|
assert_eq!(data.chunks.len(), 1);
|
|
assert!(!h5_path.exists());
|
|
}
|
|
|
|
// ---------- Test 9: Empty database migration ----------
|
|
#[test]
|
|
fn test_empty_db() {
|
|
let dir = TempDir::new().unwrap();
|
|
let db_path = create_test_db(&dir);
|
|
let h5_path = dir.path().join("out.h5");
|
|
|
|
let data =
|
|
sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
|
assert_eq!(data.chunks.len(), 0);
|
|
assert_eq!(data.sessions.len(), 0);
|
|
assert_eq!(data.entities.len(), 0);
|
|
assert_eq!(data.relations.len(), 0);
|
|
|
|
let opts = hdf5_writer::WriteOptions {
|
|
agent_id: "t".into(),
|
|
embedder: "t".into(),
|
|
compression: false,
|
|
compression_level: 4,
|
|
float16: false,
|
|
};
|
|
hdf5_writer::write_hdf5(h5_path.to_str().unwrap(), &data, &opts).unwrap();
|
|
|
|
let summary =
|
|
validate::validate_hdf5(h5_path.to_str().unwrap(), &data, false, false).unwrap();
|
|
assert_eq!(summary.chunks, 0);
|
|
}
|
|
|
|
// ---------- Test 10: Large migration (1000 entries) ----------
|
|
#[test]
|
|
fn test_large_migration() {
|
|
let dir = TempDir::new().unwrap();
|
|
let db_path = create_test_db(&dir);
|
|
let h5_path = dir.path().join("out.h5");
|
|
|
|
let conn = Connection::open(&db_path).unwrap();
|
|
for i in 0..1000 {
|
|
insert_chunk(
|
|
&conn,
|
|
i,
|
|
&format!("chunk number {i} with some text"),
|
|
&make_embedding(64, i as f32 * 0.01),
|
|
0,
|
|
);
|
|
}
|
|
drop(conn);
|
|
|
|
let data =
|
|
sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
|
assert_eq!(data.chunks.len(), 1000);
|
|
|
|
let opts = hdf5_writer::WriteOptions {
|
|
agent_id: "t".into(),
|
|
embedder: "t".into(),
|
|
compression: false,
|
|
compression_level: 4,
|
|
float16: false,
|
|
};
|
|
hdf5_writer::write_hdf5(h5_path.to_str().unwrap(), &data, &opts).unwrap();
|
|
|
|
let summary =
|
|
validate::validate_hdf5(h5_path.to_str().unwrap(), &data, false, false).unwrap();
|
|
assert_eq!(summary.chunks, 1000);
|
|
}
|
|
|
|
// ---------- Test 11: Session migration ----------
|
|
#[test]
|
|
fn test_session_migration() {
|
|
let dir = TempDir::new().unwrap();
|
|
let db_path = create_test_db(&dir);
|
|
let h5_path = dir.path().join("out.h5");
|
|
|
|
let conn = Connection::open(&db_path).unwrap();
|
|
insert_session(&conn, "session-alpha", 0, 10);
|
|
insert_session(&conn, "session-beta", 11, 20);
|
|
insert_session(&conn, "session-gamma", 21, 30);
|
|
drop(conn);
|
|
|
|
let data =
|
|
sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
|
assert_eq!(data.sessions.len(), 3);
|
|
|
|
let opts = hdf5_writer::WriteOptions {
|
|
agent_id: "t".into(),
|
|
embedder: "t".into(),
|
|
compression: false,
|
|
compression_level: 4,
|
|
float16: false,
|
|
};
|
|
hdf5_writer::write_hdf5(h5_path.to_str().unwrap(), &data, &opts).unwrap();
|
|
|
|
let summary =
|
|
validate::validate_hdf5(h5_path.to_str().unwrap(), &data, false, false).unwrap();
|
|
assert_eq!(summary.sessions, 3);
|
|
}
|
|
|
|
// ---------- Test 12: Knowledge graph (entities + relations) ----------
|
|
#[test]
|
|
fn test_knowledge_graph_migration() {
|
|
let dir = TempDir::new().unwrap();
|
|
let db_path = create_test_db(&dir);
|
|
let h5_path = dir.path().join("out.h5");
|
|
|
|
let conn = Connection::open(&db_path).unwrap();
|
|
insert_entity(&conn, 1, "Alice", "person");
|
|
insert_entity(&conn, 2, "Bob", "person");
|
|
insert_entity(&conn, 3, "Rust", "language");
|
|
insert_relation(&conn, 1, 2, "knows");
|
|
insert_relation(&conn, 1, 3, "uses");
|
|
insert_relation(&conn, 2, 3, "uses");
|
|
drop(conn);
|
|
|
|
let data =
|
|
sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
|
assert_eq!(data.entities.len(), 3);
|
|
assert_eq!(data.relations.len(), 3);
|
|
|
|
let opts = hdf5_writer::WriteOptions {
|
|
agent_id: "t".into(),
|
|
embedder: "t".into(),
|
|
compression: false,
|
|
compression_level: 4,
|
|
float16: false,
|
|
};
|
|
hdf5_writer::write_hdf5(h5_path.to_str().unwrap(), &data, &opts).unwrap();
|
|
|
|
let summary =
|
|
validate::validate_hdf5(h5_path.to_str().unwrap(), &data, false, false).unwrap();
|
|
assert_eq!(summary.entities, 3);
|
|
assert_eq!(summary.relations, 3);
|
|
}
|
|
|
|
// ---------- Test 13: Validation catches chunk count mismatch ----------
|
|
#[test]
|
|
fn test_validation_catches_count_mismatch() {
|
|
let dir = TempDir::new().unwrap();
|
|
let db_path = create_test_db(&dir);
|
|
let h5_path = dir.path().join("out.h5");
|
|
|
|
let conn = Connection::open(&db_path).unwrap();
|
|
insert_chunk(&conn, 1, "test", &make_embedding(4, 1.0), 0);
|
|
drop(conn);
|
|
|
|
let data =
|
|
sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
|
let opts = hdf5_writer::WriteOptions {
|
|
agent_id: "t".into(),
|
|
embedder: "t".into(),
|
|
compression: false,
|
|
compression_level: 4,
|
|
float16: false,
|
|
};
|
|
hdf5_writer::write_hdf5(h5_path.to_str().unwrap(), &data, &opts).unwrap();
|
|
|
|
// Validating against a source with an extra (unwritten) chunk must fail.
|
|
let mut bigger =
|
|
sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
|
let mut extra = bigger.chunks[0].clone();
|
|
extra.id = 999;
|
|
bigger.chunks.push(extra);
|
|
let result = validate::validate_hdf5(h5_path.to_str().unwrap(), &bigger, false, false);
|
|
assert!(result.is_err());
|
|
assert!(result.unwrap_err().to_string().contains("count mismatch"));
|
|
}
|
|
|
|
// ---------- Test 14: Metadata attributes are stored ----------
|
|
#[test]
|
|
fn test_metadata_attributes() {
|
|
let dir = TempDir::new().unwrap();
|
|
let db_path = create_test_db(&dir);
|
|
let h5_path = dir.path().join("out.h5");
|
|
|
|
let conn = Connection::open(&db_path).unwrap();
|
|
insert_chunk(&conn, 1, "test", &make_embedding(8, 1.0), 0);
|
|
drop(conn);
|
|
|
|
let data =
|
|
sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
|
let opts = hdf5_writer::WriteOptions {
|
|
agent_id: "my-agent-42".into(),
|
|
embedder: "openai-ada".into(),
|
|
compression: false,
|
|
compression_level: 4,
|
|
float16: false,
|
|
};
|
|
hdf5_writer::write_hdf5(h5_path.to_str().unwrap(), &data, &opts).unwrap();
|
|
|
|
let file = clawhdf5::File::open(h5_path.to_str().unwrap()).unwrap();
|
|
let root = file.root();
|
|
let attrs = root.attrs().unwrap();
|
|
|
|
match attrs.get("agent_id") {
|
|
Some(clawhdf5_format::type_builders::AttrValue::String(s)) => {
|
|
assert_eq!(s, "my-agent-42");
|
|
}
|
|
other => panic!("Expected String agent_id, got {other:?}"),
|
|
}
|
|
match attrs.get("embedder") {
|
|
Some(clawhdf5_format::type_builders::AttrValue::String(s)) => {
|
|
assert_eq!(s, "openai-ada");
|
|
}
|
|
other => panic!("Expected String embedder, got {other:?}"),
|
|
}
|
|
match attrs.get("embedding_dim") {
|
|
Some(clawhdf5_format::type_builders::AttrValue::I64(d)) => {
|
|
assert_eq!(*d, 8);
|
|
}
|
|
other => panic!("Expected I64 embedding_dim, got {other:?}"),
|
|
}
|
|
}
|
|
|
|
// ---------- Test 15: Embedding values roundtrip correctly ----------
|
|
#[test]
|
|
fn test_embedding_roundtrip() {
|
|
let dir = TempDir::new().unwrap();
|
|
let db_path = create_test_db(&dir);
|
|
let h5_path = dir.path().join("out.h5");
|
|
|
|
let conn = Connection::open(&db_path).unwrap();
|
|
let emb = vec![0.1, 0.2, 0.3, 0.4];
|
|
insert_chunk(&conn, 1, "test", &emb, 0);
|
|
drop(conn);
|
|
|
|
let data =
|
|
sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
|
let opts = hdf5_writer::WriteOptions {
|
|
agent_id: "t".into(),
|
|
embedder: "t".into(),
|
|
compression: false,
|
|
compression_level: 4,
|
|
float16: false,
|
|
};
|
|
hdf5_writer::write_hdf5(h5_path.to_str().unwrap(), &data, &opts).unwrap();
|
|
|
|
let file = clawhdf5::File::open(h5_path.to_str().unwrap()).unwrap();
|
|
let chunks_group = file.group("chunks").unwrap();
|
|
let emb_ds = chunks_group.dataset("embeddings").unwrap();
|
|
let read_back = emb_ds.read_f32().unwrap();
|
|
|
|
assert_eq!(read_back.len(), 4);
|
|
for (a, b) in emb.iter().zip(read_back.iter()) {
|
|
assert!((a - b).abs() < 1e-6, "Embedding mismatch: {a} vs {b}");
|
|
}
|
|
}
|
|
|
|
// ---------- Test 16: Full combined migration ----------
|
|
#[test]
|
|
fn test_full_combined_migration() {
|
|
let dir = TempDir::new().unwrap();
|
|
let db_path = create_test_db(&dir);
|
|
let h5_path = dir.path().join("out.h5");
|
|
|
|
let conn = Connection::open(&db_path).unwrap();
|
|
for i in 0..5 {
|
|
insert_chunk(
|
|
&conn,
|
|
i,
|
|
&format!("chunk {i}"),
|
|
&make_embedding(16, i as f32),
|
|
if i == 3 { 1 } else { 0 },
|
|
);
|
|
}
|
|
insert_session(&conn, "s1", 0, 2);
|
|
insert_session(&conn, "s2", 3, 4);
|
|
insert_entity(&conn, 1, "Alice", "person");
|
|
insert_entity(&conn, 2, "Bob", "person");
|
|
insert_relation(&conn, 1, 2, "knows");
|
|
drop(conn);
|
|
|
|
// Skip deleted
|
|
let data =
|
|
sqlite_reader::read_sqlite(&db_path, true, None, &SchemaConfig::default()).unwrap();
|
|
assert_eq!(data.chunks.len(), 4); // chunk 3 is deleted
|
|
|
|
let opts = hdf5_writer::WriteOptions {
|
|
agent_id: "combined-test".into(),
|
|
embedder: "test-embedder".into(),
|
|
compression: true,
|
|
compression_level: 4,
|
|
float16: false,
|
|
};
|
|
hdf5_writer::write_hdf5(h5_path.to_str().unwrap(), &data, &opts).unwrap();
|
|
|
|
let summary =
|
|
validate::validate_hdf5(h5_path.to_str().unwrap(), &data, false, false).unwrap();
|
|
assert_eq!(summary.chunks, 4);
|
|
assert_eq!(summary.sessions, 2);
|
|
assert_eq!(summary.entities, 2);
|
|
assert_eq!(summary.relations, 1);
|
|
assert_eq!(summary.embedding_dim, 16);
|
|
}
|
|
|
|
// ---------- Test 17: Validation catches session mismatch ----------
|
|
#[test]
|
|
fn test_validation_session_mismatch() {
|
|
let dir = TempDir::new().unwrap();
|
|
let db_path = create_test_db(&dir);
|
|
let h5_path = dir.path().join("out.h5");
|
|
|
|
let conn = Connection::open(&db_path).unwrap();
|
|
insert_session(&conn, "s1", 0, 10);
|
|
drop(conn);
|
|
|
|
let data =
|
|
sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
|
let opts = hdf5_writer::WriteOptions {
|
|
agent_id: "t".into(),
|
|
embedder: "t".into(),
|
|
compression: false,
|
|
compression_level: 4,
|
|
float16: false,
|
|
};
|
|
hdf5_writer::write_hdf5(h5_path.to_str().unwrap(), &data, &opts).unwrap();
|
|
|
|
// Validating against a source whose session content differs must fail.
|
|
let mut tampered =
|
|
sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
|
tampered.sessions[0].summary = "DIFFERENT".into();
|
|
let result = validate::validate_hdf5(h5_path.to_str().unwrap(), &tampered, false, false);
|
|
assert!(result.is_err());
|
|
assert!(result.unwrap_err().to_string().contains("session"));
|
|
}
|
|
|
|
// ---------- Real content validation catches corrupt embeddings ----------
|
|
#[test]
|
|
fn test_content_validation_catches_embedding_corruption() {
|
|
let dir = TempDir::new().unwrap();
|
|
let db_path = create_test_db(&dir);
|
|
let h5_path = dir.path().join("out.h5");
|
|
|
|
let conn = Connection::open(&db_path).unwrap();
|
|
insert_chunk(&conn, 1, "hello", &make_embedding(8, 1.0), 0);
|
|
drop(conn);
|
|
|
|
let data =
|
|
sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
|
let opts = hdf5_writer::WriteOptions {
|
|
agent_id: "t".into(),
|
|
embedder: "t".into(),
|
|
compression: false,
|
|
compression_level: 4,
|
|
float16: false,
|
|
};
|
|
hdf5_writer::write_hdf5(h5_path.to_str().unwrap(), &data, &opts).unwrap();
|
|
|
|
// A source whose embedding differs (but counts match) must fail validation.
|
|
let mut tampered =
|
|
sqlite_reader::read_sqlite(&db_path, false, None, &SchemaConfig::default()).unwrap();
|
|
tampered.chunks[0].embedding[3] += 9.0;
|
|
let result = validate::validate_hdf5(h5_path.to_str().unwrap(), &tampered, true, false);
|
|
assert!(result.is_err());
|
|
assert!(result.unwrap_err().to_string().contains("embedding"));
|
|
}
|
|
|
|
// ---------- Configurable schema: custom table names ----------
|
|
#[test]
|
|
fn test_configurable_table_names() {
|
|
let dir = TempDir::new().unwrap();
|
|
let db_path = dir.path().join("custom.db");
|
|
let path_str = db_path.to_str().unwrap().to_string();
|
|
let conn = Connection::open(&path_str).unwrap();
|
|
// Chunks live in a differently-named table; the others use defaults.
|
|
conn.execute_batch(
|
|
"CREATE TABLE my_chunks (
|
|
id INTEGER PRIMARY KEY, chunk TEXT, embedding BLOB,
|
|
source_channel TEXT, timestamp REAL, session_id TEXT, tags TEXT, deleted INTEGER
|
|
);
|
|
CREATE TABLE sessions (id TEXT, start_idx INTEGER, end_idx INTEGER, channel TEXT, timestamp REAL, summary TEXT);
|
|
CREATE TABLE entities (id INTEGER, name TEXT, type TEXT, embedding_idx INTEGER);
|
|
CREATE TABLE relations (src INTEGER, tgt INTEGER, relation TEXT, weight REAL, timestamp REAL);",
|
|
)
|
|
.unwrap();
|
|
let blob: Vec<u8> = make_embedding(4, 1.0)
|
|
.iter()
|
|
.flat_map(|v| v.to_le_bytes())
|
|
.collect();
|
|
conn.execute(
|
|
"INSERT INTO my_chunks VALUES (1, 'hi', ?1, 'api', 1.0, 's', '', 0)",
|
|
rusqlite::params![blob],
|
|
)
|
|
.unwrap();
|
|
drop(conn);
|
|
|
|
let mut schema = SchemaConfig::default();
|
|
schema.chunks.table = "my_chunks".into();
|
|
let data = sqlite_reader::read_sqlite(&path_str, false, None, &schema).unwrap();
|
|
assert_eq!(data.chunks.len(), 1);
|
|
assert_eq!(data.chunks[0].chunk, "hi");
|
|
assert_eq!(data.embedding_dim, 4);
|
|
|
|
// Counts pass should also honor the custom table name.
|
|
let counts = sqlite_reader::read_counts(&path_str, false, &schema).unwrap();
|
|
assert_eq!(counts.chunks, 1);
|
|
}
|
|
|
|
// ---------- Incremental migration appends only new rows ----------
|
|
#[test]
|
|
fn test_incremental_migration() {
|
|
let dir = TempDir::new().unwrap();
|
|
let db_path = create_test_db(&dir);
|
|
let h5_path = dir.path().join("out.h5");
|
|
let cfg = SchemaConfig::default();
|
|
let opts = hdf5_writer::WriteOptions {
|
|
agent_id: "t".into(),
|
|
embedder: "t".into(),
|
|
compression: false,
|
|
compression_level: 4,
|
|
float16: false,
|
|
};
|
|
|
|
// First migration: 2 chunks.
|
|
let conn = Connection::open(&db_path).unwrap();
|
|
insert_chunk(&conn, 1, "one", &make_embedding(4, 1.0), 0);
|
|
insert_chunk(&conn, 2, "two", &make_embedding(4, 2.0), 0);
|
|
drop(conn);
|
|
let data = sqlite_reader::read_sqlite(&db_path, false, None, &cfg).unwrap();
|
|
hdf5_writer::write_hdf5(h5_path.to_str().unwrap(), &data, &opts).unwrap();
|
|
|
|
// Add two more rows, then migrate incrementally.
|
|
let conn = Connection::open(&db_path).unwrap();
|
|
insert_chunk(&conn, 3, "three", &make_embedding(4, 3.0), 0);
|
|
insert_chunk(&conn, 4, "four", &make_embedding(4, 4.0), 0);
|
|
drop(conn);
|
|
|
|
let base = hdf5_reader::read_hdf5(h5_path.to_str().unwrap()).unwrap();
|
|
let max_id = base.chunks.iter().map(|c| c.id).max().unwrap_or(0);
|
|
assert_eq!(max_id, 2);
|
|
let new =
|
|
sqlite_reader::read_sqlite_filtered(&db_path, false, Some(4), &cfg, max_id).unwrap();
|
|
assert_eq!(new.chunks.len(), 2); // only id 3 and 4
|
|
|
|
let mut merged = base;
|
|
merged.chunks.extend(new.chunks);
|
|
hdf5_writer::write_hdf5(h5_path.to_str().unwrap(), &merged, &opts).unwrap();
|
|
|
|
let final_data = hdf5_reader::read_hdf5(h5_path.to_str().unwrap()).unwrap();
|
|
assert_eq!(final_data.chunks.len(), 4);
|
|
let texts: Vec<&str> = final_data.chunks.iter().map(|c| c.chunk.as_str()).collect();
|
|
assert_eq!(texts, vec!["one", "two", "three", "four"]);
|
|
}
|
|
}
|