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

//! `HDF5Memory::search` with `SearchOptions`: source filtering, re-ranking and
//! confidence rejection in the store's own search path.

use std::collections::HashSet;

use clawhdf5_agent::confidence::ConfidenceConfig;
use clawhdf5_agent::reranker::ReRankConfig;
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry, SearchOptions, hybrid};
use tempfile::TempDir;

const DIM: usize = 32;
const N: usize = 3000;
const CLUSTERS: usize = 20;

struct Rng(u64);
impl Rng {
    fn next(&mut self) -> u64 {
        self.0 = self.0.wrapping_add(0x9E37_79B9_7F4A_7C15);
        let mut z = self.0;
        z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
        z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
        z ^ (z >> 31)
    }
    fn unit(&mut self) -> f32 {
        (self.next() >> 40) as f32 / (1u64 << 24) as f32 - 0.5
    }
}

fn normalize(v: &mut [f32]) {
    let n = v.iter().map(|x| x * x).sum::<f32>().sqrt();
    v.iter_mut().for_each(|x| *x /= n);
}

struct Data {
    vectors: Vec<Vec<f32>>,
    cluster: Vec<usize>,
    centres: Vec<Vec<f32>>,
}

fn data() -> Data {
    let mut rng = Rng(42);
    let centres: Vec<Vec<f32>> = (0..CLUSTERS)
        .map(|_| {
            let mut c: Vec<f32> = (0..DIM).map(|_| rng.unit()).collect();
            normalize(&mut c);
            c
        })
        .collect();
    let mut vectors = Vec::new();
    let mut cluster = Vec::new();
    for i in 0..N {
        let c = i % CLUSTERS;
        let mut v: Vec<f32> = centres[c].iter().map(|x| x + rng.unit() * 0.3).collect();
        normalize(&mut v);
        vectors.push(v);
        cluster.push(c);
    }
    Data {
        vectors,
        cluster,
        centres,
    }
}

/// Channel of record `i` for a filter keeping `percent`% of the store at
/// random (independent of the vectors).
fn random_channel(i: usize, rng_seed: u64, percent: u64) -> String {
    let mut r = Rng(rng_seed ^ (i as u64 * 7919));
    if r.next() % 100 < percent {
        "keep".into()
    } else {
        "other".into()
    }
}

fn build(data: &Data, channel: impl Fn(usize) -> String) -> (TempDir, HDF5Memory) {
    let dir = TempDir::new().unwrap();
    let mut cfg = MemoryConfig::new(dir.path().join("s.h5"), "agent", DIM);
    cfg.hebbian_boost = 0.0; // every query sees the same store
    let mut m = HDF5Memory::create(cfg).unwrap();
    let entries = data
        .vectors
        .iter()
        .enumerate()
        .map(|(i, v)| MemoryEntry {
            chunk: format!("record {i} cluster {}", data.cluster[i]),
            embedding: v.clone(),
            source_channel: channel(i),
            timestamp: i as f64,
            session_id: "s".into(),
            tags: format!("t{i}"),
        })
        .collect();
    m.save_batch(entries).unwrap();
    (dir, m)
}

/// Exact top-k by cosine among the records `allowed` keeps.
fn exact_top(data: &Data, q: &[f32], k: usize, allowed: impl Fn(usize) -> bool) -> Vec<usize> {
    let mut s: Vec<(usize, f32)> = (0..N)
        .filter(|&i| allowed(i))
        .map(|i| (i, data.vectors[i].iter().zip(q).map(|(a, b)| a * b).sum()))
        .collect();
    s.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0)));
    s.into_iter().take(k).map(|(i, _)| i).collect()
}

fn query(data: &Data, i: usize) -> Vec<f32> {
    let mut rng = Rng(1000 + i as u64);
    let mut q: Vec<f32> = data.centres[i % CLUSTERS]
        .iter()
        .map(|x| x + rng.unit() * 0.3)
        .collect();
    normalize(&mut q);
    q
}

fn vector_only(k: usize) -> SearchOptions {
    SearchOptions::new(k).with_fusion(hybrid::Fusion::Weighted {
        vector: 1.0,
        keyword: 0.0,
    })
}

#[test]
fn source_filter_returns_only_allowed_records_and_a_full_page() {
    let d = data();
    // At N = 3000 and k = 10 the index serves a filter only when that is
    // cheaper than scanning the allowed records: pool = 80 * N / allowed
    // candidates at ~M = 16 distances each, against `allowed` distances. So
    // 90% goes through the index, 50% and 1% to the exact scan.
    for percent in [90, 50, 1] {
        let (_dir, mut m) = build(&d, |i| random_channel(i, 5, percent));
        let allowed = |i: usize| random_channel(i, 5, percent) == "keep";
        let mut hits = 0;
        for qi in 0..40 {
            let q = query(&d, qi);
            let got = m.search(&q, "", &vector_only(10).with_sources(["keep"]));
            assert_eq!(got.len(), 10, "{percent}%: short page");
            assert!(got.iter().all(|r| r.source_channel == "keep"));
            let want: HashSet<usize> = exact_top(&d, &q, 10, allowed).into_iter().collect();
            hits += got.iter().filter(|r| want.contains(&r.index)).count();
        }
        let recall = hits as f64 / 400.0;
        let floor = if percent == 90 { 0.95 } else { 1.0 };
        assert!(recall >= floor, "{percent}%: recall@10 {recall}");
    }
}

#[test]
fn filter_away_from_the_query_falls_back_to_an_exact_scan() {
    // Channel = cluster, and the filter keeps two clusters (10% of the
    // store) that are not the query's: the index's neighbourhood of the
    // query holds none of them. The search must still return the exact
    // top 10 among the allowed records, not a short or empty page.
    let d = data();
    let (_dir, mut m) = build(&d, |i| format!("c{}", d.cluster[i]));
    for qi in 0..20 {
        let q = query(&d, qi);
        let a = format!("c{}", (qi + 7) % CLUSTERS);
        let b = format!("c{}", (qi + 13) % CLUSTERS);
        let got: Vec<usize> = m
            .search(
                &q,
                "",
                &vector_only(10).with_sources([a.clone(), b.clone()]),
            )
            .iter()
            .map(|r| r.index)
            .collect();
        let want = exact_top(&d, &q, 10, |i| {
            let c = format!("c{}", d.cluster[i]);
            c == a || c == b
        });
        assert_eq!(got, want, "query {qi}");
    }
}

#[test]
fn filter_edge_cases() {
    let d = data();
    let (_dir, mut m) = build(&d, |i| random_channel(i, 9, 50));
    let q = query(&d, 0);
    assert!(
        m.search(
            &q,
            "cluster",
            &SearchOptions::new(10).with_sources(Vec::<String>::new())
        )
        .is_empty()
    );
    assert!(
        m.search(
            &q,
            "cluster",
            &SearchOptions::new(10).with_sources(["nope"])
        )
        .is_empty()
    );
    // Keyword matches from other channels are filtered too.
    let got = m.search(
        &q,
        "record cluster",
        &SearchOptions::new(50).with_sources(["keep"]),
    );
    assert_eq!(got.len(), 50);
    assert!(got.iter().all(|r| r.source_channel == "keep"));
    // Deleted records never come back, filtered or not.
    let first = got[0].index;
    m.delete(first).unwrap();
    let again = m.search(
        &q,
        "record cluster",
        &SearchOptions::new(50).with_sources(["keep"]),
    );
    assert!(again.iter().all(|r| r.index != first));
}

#[test]
fn plain_options_equal_hybrid_search_with() {
    // Two identical stores, so neither query sees the other's boosts.
    let d = data();
    let (_a, mut a) = build(&d, |i| random_channel(i, 3, 50));
    let (_b, mut b) = build(&d, |i| random_channel(i, 3, 50));
    for qi in 0..10 {
        let q = query(&d, qi);
        let x: Vec<(usize, u32)> = a
            .search(&q, "record cluster 3", &SearchOptions::new(10))
            .iter()
            .map(|r| (r.index, r.score.to_bits()))
            .collect();
        let y: Vec<(usize, u32)> = b
            .hybrid_search_with(&q, "record cluster 3", hybrid::DEFAULT_FUSION, 10)
            .iter()
            .map(|r| (r.index, r.score.to_bits()))
            .collect();
        assert_eq!(x, y);
    }
}

fn small_store(entries: &[(&str, &str, f64)]) -> (TempDir, HDF5Memory) {
    let dir = TempDir::new().unwrap();
    let mut m = HDF5Memory::create(MemoryConfig::new(dir.path().join("r.h5"), "a", 4)).unwrap();
    m.save_batch(
        entries
            .iter()
            .map(|(chunk, channel, ts)| MemoryEntry {
                chunk: chunk.to_string(),
                embedding: vec![1.0, 0.0, 0.0, 0.0],
                source_channel: channel.to_string(),
                timestamp: *ts,
                session_id: "s".into(),
                tags: String::new(),
            })
            .collect(),
    )
    .unwrap();
    (dir, m)
}

#[test]
fn rerank_breaks_relevance_ties_by_recency() {
    // Identical text and vectors, so retrieval ties; re-ranking must put the
    // newer record first and report the combined score.
    let now = 1_000_000.0;
    let (_d, mut m) = small_store(&[
        ("user prefers dark mode", "chat", now - 30.0 * 86_400.0),
        ("user prefers dark mode", "chat", now - 60.0),
    ]);
    let q = [1.0, 0.0, 0.0, 0.0];
    let plain = m.search(&q, "dark mode", &SearchOptions::new(2));
    assert_eq!(plain[0].index, 0, "ties break by index without re-ranking");
    let reranked = m.search(
        &q,
        "dark mode",
        &SearchOptions::new(2)
            .with_rerank(ReRankConfig::default())
            .at_time(now),
    );
    assert_eq!(reranked[0].index, 1);
    assert!(reranked[0].score > reranked[1].score);
    assert_ne!(reranked[0].score.to_bits(), plain[0].score.to_bits());
}

#[test]
fn confidence_rejects_when_nothing_is_good_enough() {
    let (_d, mut m) = small_store(&[("alpha", "chat", 0.0), ("beta", "chat", 0.0)]);
    let q = [1.0, 0.0, 0.0, 0.0];
    let strict = ConfidenceConfig {
        min_score: 10.0,
        ..ConfidenceConfig::default()
    };
    assert!(
        m.search(&q, "alpha", &SearchOptions::new(2).with_confidence(strict))
            .is_empty()
    );
    let lenient = ConfidenceConfig {
        min_score: 0.0,
        min_gap: f32::INFINITY,
        max_results: 1,
    };
    assert_eq!(
        m.search(&q, "alpha", &SearchOptions::new(2).with_confidence(lenient))
            .len(),
        1
    );
}

#[test]
fn only_returned_results_are_reinforced() {
    // With re-ranking, a pool of max(3k, 10) candidates is retrieved; only
    // the k returned should gain activation.
    let d = data();
    let dir = TempDir::new().unwrap();
    let path = dir.path().join("h.h5");
    let mut m = HDF5Memory::create(MemoryConfig::new(path, "a", DIM)).unwrap();
    m.save_batch(
        (0..200)
            .map(|i| MemoryEntry {
                chunk: format!("record {i}"),
                embedding: d.vectors[i].clone(),
                source_channel: "chat".into(),
                timestamp: i as f64,
                session_id: "s".into(),
                tags: String::new(),
            })
            .collect(),
    )
    .unwrap();
    let q = query(&d, 0);
    let got = m.search(
        &q,
        "record",
        &SearchOptions::new(3).with_rerank(ReRankConfig::default()),
    );
    assert_eq!(got.len(), 3);
    let returned: HashSet<usize> = got.iter().map(|r| r.index).collect();
    // A second plain search reports each record's current activation.
    let all = m.search(&q, "record", &SearchOptions::new(200));
    for r in &all {
        let boosted = r.activation > 1.0;
        assert_eq!(boosted, returned.contains(&r.index), "record {}", r.index);
    }
}