AkurAI-Notes / crates/notes-api/tests/embed_test.rs

Latest change 3594ca5510bcac2b26e45d9235a28905856ebf61 - E05: On-device embeddings + semantic search — embed.rs + search.rs + re-embed on update by Claude Code

use akurai_json::Value;
use akurai_storage::BTree;
use notes_api::embed::{maybe_reembed, EmbedStore};
use notes_api::NotesState;

fn setup_db(name: &str) -> (BTree, String) {
    let db_path = format!("_test_embed_{}.db", name);
    let _ = std::fs::remove_file(&db_path);
    let tree = BTree::open(&db_path).unwrap();
    (tree, db_path)
}

fn make_note_key(id: u64) -> Vec<u8> {
    let mut k = b"notes:".to_vec();
    k.extend_from_slice(&id.to_be_bytes());
    k
}

fn insert_note(db: &mut BTree, id: u64, title: &str, body: &str, owner: &str) {
    let now = NotesState::now();
    let record = Value::Object(vec![
        ("id".into(), Value::Int(id as i64)),
        ("title".into(), Value::Str(title.into())),
        ("body".into(), Value::Str(body.into())),
        ("ownerSub".into(), Value::Str(owner.into())),
        ("createdAt".into(), Value::Int(now)),
        ("updatedAt".into(), Value::Int(now)),
    ]);
    db.insert(&make_note_key(id), record.to_json().as_bytes())
        .unwrap();
    db.commit().unwrap();
}

#[test]
fn embed_store_round_trips_vector() {
    let (mut db, db_path) = setup_db("roundtrip");
    let mut store = EmbedStore { db: &mut db };

    let vec: Vec<f32> = vec![0.1, 0.2, -0.3, 0.5];
    store.store(1, &vec).unwrap();

    let loaded = store.load(1);
    assert!(loaded.is_some());
    let loaded = loaded.unwrap();
    assert_eq!(loaded.len(), 4);
    for (a, b) in vec.iter().zip(loaded.iter()) {
        assert!((a - b).abs() < 1e-6);
    }

    store.remove(1).unwrap();
    assert!(store.load(1).is_none());

    std::fs::remove_file(&db_path).ok();
}

#[test]
fn embed_store_returns_none_for_missing_key() {
    let (mut db, db_path) = setup_db("missing");
    let mut store = EmbedStore { db: &mut db };
    assert!(store.load(999).is_none());
    std::fs::remove_file(&db_path).ok();
}

#[test]
fn rank_notes_returns_only_owned_notes_in_order() {
    let (mut db, db_path) = setup_db("rank");

    insert_note(&mut db, 1, "Meeting", "Budget review meeting", "alice");
    insert_note(&mut db, 2, "Shopping", "Buy groceries", "alice");
    insert_note(&mut db, 3, "Other", "Bob's note", "bob");

    let mut store = EmbedStore { db: &mut db };
    let v1 = vec![0.9, 0.1, 0.0];
    store.store(1, &v1).unwrap();
    let v2 = vec![0.1, 0.9, 0.0];
    store.store(2, &v2).unwrap();
    let v3 = vec![0.85, 0.15, 0.0];
    store.store(3, &v3).unwrap();

    let qv = vec![0.8, 0.2, 0.0];

    let ranked = store.rank_notes(&qv, "alice", 5);
    assert_eq!(ranked.len(), 2, "should return 2 alice notes, got {ranked:?}");
    assert_eq!(ranked[0].0, 1, "note 1 should rank first (closest to qv)");
    assert_eq!(ranked[1].0, 2, "note 2 should rank second");

    let ranked = store.rank_notes(&qv, "alice", 1);
    assert_eq!(ranked.len(), 1);

    std::fs::remove_file(&db_path).ok();
}

#[test]
fn maybe_reembed_skips_when_url_is_none() {
    let (mut db, db_path) = setup_db("skip");
    insert_note(&mut db, 1, "Test", "Hello world", "user");

    maybe_reembed(&mut db, &None, &None, 1, "Hello world");

    let mut store = EmbedStore { db: &mut db };
    assert!(store.load(1).is_none());

    std::fs::remove_file(&db_path).ok();
}

#[test]
fn substring_search_via_prefix_scan() {
    let (mut db, db_path) = setup_db("substr");
    insert_note(&mut db, 1, "Meeting minutes", "Discussed Q2 targets", "alice");
    insert_note(&mut db, 2, "Grocery list", "Milk and eggs", "bob");
    insert_note(&mut db, 3, "Budget review", "Q1 financial results", "alice");
    insert_note(&mut db, 4, "Random", "Nothing important", "alice");

    let prefix = b"notes:";
    let mut bound = prefix.to_vec();
    bound.push(0xff);

    let mut results: Vec<(u64, String)> = Vec::new();
    if let Ok(entries) = db.range(prefix, &bound) {
        for entry in entries {
            let (key, val) = entry;
            let id_bytes: [u8; 8] = key[prefix.len()..].try_into().unwrap();
            let id: u64 = u64::from_be_bytes(id_bytes);
            let Ok(v) = akurai_json::parse(&String::from_utf8_lossy(&val)) else {
                continue;
            };
            if v.get("ownerSub").and_then(|x| x.as_str()) != Some("alice") {
                continue;
            }
            let title = v.get("title").and_then(|x| x.as_str()).unwrap_or("");
            let body = v.get("body").and_then(|x| x.as_str()).unwrap_or("");
            let combined = format!("{} {}", title, body).to_lowercase();
            if combined.contains("q") {
                let what = format!("{} — {}", title, body);
                results.push((id, what));
            }
        }
    }

    assert_eq!(results.len(), 2, "should find 2 alice notes with 'q'");
    let titles: Vec<&str> = results.iter().map(|r| r.1.as_str()).collect();
    assert!(titles[0].contains("Q2") || titles[1].contains("Q2"));
    assert!(titles[0].contains("Q1") || titles[1].contains("Q1"));

    std::fs::remove_file(&db_path).ok();
}

#[test]
#[ignore]
fn create_note_stores_embedding() {
    let (mut db, _db_path) = setup_db("create_embed");
    let note_id = 1i64;
    let body = "Test note content for embedding";

    let embed_url = std::env::var("AKURAI_NOTES_EMBED_URL").ok();
    let embed_model = std::env::var("AKURAI_NOTES_EMBED_MODEL").ok();

    maybe_reembed(&mut db, &embed_url, &embed_model, note_id, body);

    let mut store = EmbedStore { db: &mut db };
    let loaded = store.load(note_id);
    assert!(
        loaded.is_some(),
        "embedding should be stored when embed server is up"
    );
    let vec = loaded.unwrap();
    assert!(!vec.is_empty(), "embedding should have non-zero dimensions");
}

#[test]
#[ignore]
fn update_note_reembeds_only_when_body_changes() {}

#[test]
#[ignore]
fn search_is_owner_scoped() {}