clawhdf5/crates/clawhdf5-agent/tests/float16_store.rs

//! `MemoryConfig::float16`: embeddings stored as IEEE half precision.
//!
//! The setting used to be recorded in `/meta` and otherwise ignored — the
//! embeddings dataset was always `f32`. These tests pin what it now does: the
//! dataset is `float16`, the in-memory cache holds exactly the values the file
//! holds (so search results survive a reopen bit for bit), and a value half
//! precision cannot represent is refused rather than stored as infinity.

use std::path::{Path, PathBuf};

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

const DIM: usize = 64;

/// Deterministic, embedding-like unit vectors.
fn embedding(seed: u64) -> Vec<f32> {
    let mut x = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1;
    let v: Vec<f32> = (0..DIM)
        .map(|_| {
            x ^= x << 13;
            x ^= x >> 7;
            x ^= x << 17;
            (x >> 40) as f32 / (1u64 << 24) as f32 - 0.5
        })
        .collect();
    let norm = v.iter().map(|a| a * a).sum::<f32>().sqrt();
    v.iter().map(|a| a / norm).collect()
}

fn entry(i: u64) -> MemoryEntry {
    MemoryEntry {
        chunk: format!("memory number {i} about topic {}", i % 7),
        embedding: embedding(i),
        source_channel: "test".into(),
        timestamp: i as f64,
        session_id: "s".into(),
        tags: format!("t{i}"),
    }
}

fn config(dir: &TempDir, name: &str, float16: bool) -> MemoryConfig {
    let mut c = MemoryConfig::new(dir.path().join(name), "agent", DIM);
    c.float16 = float16;
    c
}

fn embeddings_dtype_and_values(path: &Path) -> (String, Vec<f32>) {
    let file = clawhdf5::File::open(path).unwrap();
    let ds = file.dataset("memory/embeddings").unwrap();
    (format!("{:?}", ds.dtype().unwrap()), ds.read_f32().unwrap())
}

fn search_bits(m: &mut HDF5Memory, q: u64) -> Vec<(usize, u32)> {
    m.hybrid_search(&embedding(q), "memory topic 3", 0.4, 0.6, 10)
        .iter()
        .map(|r| (r.index, r.score.to_bits()))
        .collect()
}

#[test]
fn float16_store_writes_half_precision_and_reopens_identically() {
    let dir = TempDir::new().unwrap();
    // Two identical stores. Search is not read-only (it boosts the Hebbian
    // activation of what it returns, and checkpoints persist that), so each
    // is queried exactly once: one live, one after a checkpoint and reopen.
    let live_cfg = config(&dir, "live.h5", true);
    let cfg = config(&dir, "f16.h5", true);
    let path: PathBuf = cfg.path.clone();

    let mut live = HDF5Memory::create(live_cfg).unwrap();
    live.save_batch((0..200).map(entry).collect()).unwrap();
    let mut m = HDF5Memory::create(cfg).unwrap();
    m.save_batch((0..200).map(entry).collect()).unwrap();
    drop(m);

    // On disk: a genuine float16 dataset holding the rounded inputs.
    let (dtype, values) = embeddings_dtype_and_values(&path);
    assert_eq!(dtype, "Other(\"float16\")");
    let expected: Vec<u32> = (0..200)
        .flat_map(|i| embedding(i).into_iter().map(|v| round_to_f16(v).to_bits()))
        .collect();
    let got: Vec<u32> = values.iter().map(|v| v.to_bits()).collect();
    assert_eq!(got, expected);

    // Reopened, the store answers exactly as the live one does: the cache
    // held the half-rounded values before the checkpoint.
    let mut reopened = HDF5Memory::open(&path).unwrap();
    for q in 0..5 {
        assert_eq!(
            search_bits(&mut live, 1000 + q),
            search_bits(&mut reopened, 1000 + q),
            "query {q}"
        );
    }
}

#[test]
fn float16_halves_the_embeddings_on_disk() {
    let dir = TempDir::new().unwrap();
    let mut sizes = Vec::new();
    for float16 in [false, true] {
        let cfg = config(&dir, &format!("s{float16}.h5"), float16);
        let path = cfg.path.clone();
        let mut m = HDF5Memory::create(cfg).unwrap();
        m.save_batch((0..2000).map(entry).collect()).unwrap();
        drop(m);
        sizes.push(std::fs::metadata(&path).unwrap().len());
    }
    let embedding_bytes_f32 = (2000 * DIM * 4) as u64;
    let saved = sizes[0] - sizes[1];
    // Half of the f32 embeddings, give or take metadata and alignment.
    assert!(
        saved.abs_diff(embedding_bytes_f32 / 2) < 16 * 1024,
        "f32 {} B, f16 {} B, saved {saved} B, expected ~{} B",
        sizes[0],
        sizes[1],
        embedding_bytes_f32 / 2
    );
}

#[test]
fn f32_store_is_unchanged() {
    let dir = TempDir::new().unwrap();
    let cfg = config(&dir, "f32.h5", false);
    let path = cfg.path.clone();
    let mut m = HDF5Memory::create(cfg).unwrap();
    m.save_batch((0..50).map(entry).collect()).unwrap();
    drop(m);
    let (dtype, values) = embeddings_dtype_and_values(&path);
    assert_eq!(dtype, "F32");
    let expected: Vec<f32> = (0..50).flat_map(embedding).collect();
    assert_eq!(values, expected);
}

#[test]
fn out_of_range_values_are_refused_not_stored_as_infinity() {
    let dir = TempDir::new().unwrap();
    let mut cfg = config(&dir, "range.h5", true);
    cfg.wal_enabled = true;
    let path = cfg.path.clone();
    let mut m = HDF5Memory::create(cfg).unwrap();
    m.save(entry(1)).unwrap();

    let mut bad = entry(2);
    bad.embedding[5] = 70_000.0;
    match m.save(bad.clone()) {
        Err(MemoryError::InvalidEntry(msg)) => assert!(msg.contains("embedding[5]"), "{msg}"),
        other => panic!("expected InvalidEntry, got {other:?}"),
    }
    assert!(matches!(
        m.save_or_update(bad.clone()),
        Err(MemoryError::InvalidEntry(_))
    ));
    // A batch is all or nothing.
    assert!(matches!(
        m.save_batch(vec![entry(3), bad.clone(), entry(4)]),
        Err(MemoryError::InvalidEntry(_))
    ));
    assert_eq!(m.count(), 1);

    // The largest finite half, and values that round down to it, are fine.
    let mut edge = entry(5);
    edge.embedding[0] = 65504.0;
    edge.embedding[1] = -65519.0;
    m.save(edge).unwrap();
    assert_eq!(m.count(), 2);
    drop(m);

    // Nothing rejected reached the WAL or the file.
    let m = HDF5Memory::open(&path).unwrap();
    assert_eq!(m.count(), 2);

    // An f32 store takes the same value as it always did.
    let mut m32 = HDF5Memory::create(config(&dir, "range32.h5", false)).unwrap();
    m32.save(bad).unwrap();
}

#[test]
fn wal_replay_rounds_like_a_live_save() {
    let dir = TempDir::new().unwrap();
    let mut cfg = config(&dir, "wal.h5", true);
    cfg.wal_enabled = true;
    cfg.wal_max_entries = 10_000; // keep everything in the WAL
    let path = cfg.path.clone();
    let mut m = HDF5Memory::create(cfg).unwrap();
    for i in 0..30 {
        m.save(entry(i)).unwrap();
    }
    let live = search_bits(&mut m, 77);

    // Crash image: the .h5 is still the empty checkpoint; everything is in
    // the WAL, which holds the caller's f32 values.
    let crash = TempDir::new().unwrap();
    let image = crash.path().join("image.h5");
    std::fs::copy(&path, &image).unwrap();
    std::fs::copy(
        path.with_extension("h5.wal"),
        image.with_extension("h5.wal"),
    )
    .unwrap();
    drop(m);

    let mut recovered = HDF5Memory::open(&image).unwrap();
    assert_eq!(recovered.count(), 30);
    assert_eq!(search_bits(&mut recovered, 77), live);
}