notescraft/crates/knowledge/tests/rank.rs

//! Behaviour of `retrieve`, `related` and `budget` on a tiny vault.

use std::collections::HashMap;

use notescraft_knowledge::{Index, MapSignals, NoSignals, Options, SignalKind, Vocab, budget, related, retrieve};

struct Fixture {
    index: Index,
    vocab: Vocab,
    files: HashMap<String, String>,
}

fn fixture(notes: &[(&str, &str)]) -> Fixture {
    let (mut index, mut vocab, mut files) = (Index::default(), Vocab::default(), HashMap::new());
    for (p, t) in notes {
        index.update(p, t).unwrap();
        vocab.update(p, t).unwrap();
        files.insert((*p).to_string(), (*t).to_string());
    }
    Fixture { index, vocab, files }
}

fn ask(f: &Fixture, s: &MapSignals, q: &str, o: &Options) -> notescraft_knowledge::Retrieval {
    retrieve(&f.index, &f.vocab, s, |p| f.files.get(p).cloned(), q, o).unwrap()
}

fn paths(r: &notescraft_knowledge::Retrieval) -> Vec<&str> {
    r.passages.iter().map(|p| p.path.as_str()).collect()
}

/// Several notes where "notify" travels with "reminder", plus filler.
fn reminder_vault() -> Fixture {
    let mut notes: Vec<(String, String)> = vec![
        ("a.md".into(), "# Alerts\n\nA reminder will notify you.\n\n# Sound\n\nPick the notify sound for each reminder.".into()),
        ("b.md".into(), "# Settings\n\nThe reminder can notify by email.".into()),
        ("c.md".into(), "# Tasks\n\nEvery task has a reminder at nine.".into()),
    ];
    for i in 0..30 {
        notes.push((format!("filler{i}.md"), format!("# Filler {i}\n\nGarden topic number{i} about soil and water.")));
    }
    let refs: Vec<(&str, &str)> = notes.iter().map(|(p, t)| (p.as_str(), t.as_str())).collect();
    fixture(&refs)
}

#[test]
fn vocabulary_finds_note_with_other_words_and_explains_it() {
    let f = reminder_vault();
    let s = MapSignals::default();
    let plain = ask(&f, &s, "notify", &Options { learn_vocab: false, ..Options::default() });
    assert!(!paths(&plain).contains(&"c.md"));
    let hybrid = ask(&f, &s, "notify", &Options::default());
    let c = hybrid.passages.iter().find(|p| p.path == "c.md").expect("c.md found through the reminder link");
    let signal = c.signals.iter().find(|s| s.kind == SignalKind::Vocab).expect("vocab signal");
    assert_eq!(signal.detail, "added words: reminder");
    assert!(c.text.contains("reminder at nine"));
    assert!(hybrid.passages.iter().any(|p| p.signals.iter().any(|s| s.kind == SignalKind::Text && s.detail.starts_with("matched words: notify"))));
}

#[test]
fn links_and_tags_boost_and_say_so() {
    let f = fixture(&[
        ("top.md", "# Lisbon trip\n\nlisbon room lisbon room lisbon room"),
        ("aa.md", "# Aa\n\nlisbon room lisbon room lisbon room"),
        ("ab.md", "# Ab\n\nlisbon room lisbon room lisbon room"),
        ("hotel.md", "# Hotel\n\nlisbon room near the river"),
        ("other.md", "# Other\n\nlisbon room near the river"),
        ("tagged.md", "# Tagged\n\nlisbon room near the river"),
    ]);
    let mut s = MapSignals::default();
    s.links.insert("top.md".into(), vec!["hotel.md".into()]);
    s.tags.insert("tagged.md".into(), vec!["#lisbon".into()]);
    let r = ask(&f, &s, "lisbon room", &Options::default());
    let pos = |name: &str| paths(&r).iter().position(|p| *p == name);
    assert!(pos("hotel.md") < pos("other.md"));
    assert!(pos("tagged.md") < pos("other.md"));
    let hotel = r.passages.iter().find(|p| p.path == "hotel.md").unwrap();
    assert!(hotel.signals.iter().any(|s| s.kind == SignalKind::Link && s.detail == "linked from top.md"));
    let tagged = r.passages.iter().find(|p| p.path == "tagged.md").unwrap();
    assert!(tagged.signals.iter().any(|s| s.kind == SignalKind::Tag && s.detail.contains("#lisbon")));
    let off = ask(&f, &s, "lisbon room", &Options { follow_links: false, use_tags: false, ..Options::default() });
    assert!(off.passages.iter().all(|p| p.signals.iter().all(|s| ![SignalKind::Link, SignalKind::Tag].contains(&s.kind))));
}

#[test]
fn recency_only_when_asked() {
    let f = fixture(&[("old.md", "# Old\n\nmeeting notes agenda"), ("new.md", "# New\n\nmeeting notes agenda")]);
    let mut s = MapSignals::default();
    s.modified.insert("old.md".into(), 100);
    s.modified.insert("new.md".into(), 900);
    assert_eq!(paths(&ask(&f, &s, "meeting agenda", &Options::default())), ["new.md", "old.md"]);
    let r = ask(&f, &s, "meeting agenda", &Options { prefer_recent: true, ..Options::default() });
    assert_eq!(paths(&r), ["new.md", "old.md"]);
    assert!(r.passages.first().is_some_and(|p| p.signals.iter().any(|s| s.detail == "recently modified")));
}

#[test]
fn at_most_two_passages_per_note_unless_short() {
    let big = "# One\n\nzebra alpha\n\n# Two\n\nzebra beta\n\n# Three\n\nzebra gamma";
    let f = fixture(&[("big.md", big), ("small.md", "# S\n\nzebra")]);
    let s = MapSignals::default();
    let r = ask(&f, &s, "zebra", &Options { limit: 4, ..Options::default() });
    assert_eq!(r.passages.len(), 4);
    assert!(!r.truncated);
    let r = ask(&f, &s, "zebra", &Options { limit: 2, ..Options::default() });
    assert_eq!(r.passages.len(), 2);
    assert!(r.truncated);
    let many = fixture(&[("big.md", big), ("a.md", "# A\n\nzebra"), ("b.md", "# B\n\nzebra")]);
    let r = ask(&many, &s, "zebra", &Options { limit: 3, ..Options::default() });
    assert!(r.passages.iter().filter(|p| p.path == "big.md").count() <= 2);
}

#[test]
fn deterministic_and_hostile_inputs() {
    let f = reminder_vault();
    let s = MapSignals::default();
    let a = ask(&f, &s, "notify reminder", &Options::default());
    assert_eq!(a, ask(&f, &s, "notify reminder", &Options::default()));
    assert_eq!(ask(&f, &s, "", &Options::default()).passages.len(), 0);
    assert_eq!(ask(&f, &s, "the and of", &Options::default()).passages.len(), 0);
    assert!(ask(&f, &s, &"word ".repeat(100_000), &Options { limit: usize::MAX, ..Options::default() }).passages.len() <= 50);
    let gone = retrieve(&f.index, &f.vocab, &NoSignals, |_| None, "reminder", &Options::default()).unwrap();
    assert_eq!(gone.passages.len(), 0);
    let bad = retrieve(&f.index, &f.vocab, &NoSignals, |_| Some("short".into()), "reminder", &Options::default()).unwrap();
    assert_eq!(bad.passages.len(), 0);
}

#[test]
fn related_uses_own_words_links_and_excludes_self() {
    let f = fixture(&[
        ("me.md", "# Bread\n\nsourdough starter flour hydration crumb sourdough starter"),
        ("same.md", "# Starter\n\nfeed the sourdough starter with flour daily"),
        ("linked.md", "# Oven\n\npreheat the stone for an hour"),
        ("noise.md", "# Noise\n\ngarden roses and soil"),
    ]);
    let mut s = MapSignals::default();
    s.links.insert("me.md".into(), vec!["linked.md".into()]);
    let r = related(&f.index, &f.vocab, &s, |p| f.files.get(p).cloned(), "me.md", &Options::default()).unwrap();
    let names: Vec<&str> = r.iter().map(|n| n.path.as_str()).collect();
    assert_eq!(names.first(), Some(&"same.md"));
    assert!(!names.contains(&"me.md") && !names.contains(&"noise.md"));
    assert!(r.first().is_some_and(|n| n.snippet.contains("sourdough") && n.signals.iter().any(|s| s.detail.starts_with("matched words"))));
    assert_eq!(related(&f.index, &f.vocab, &s, |_| None, "me.md", &Options::default()).unwrap(), []);
    assert_eq!(related(&f.index, &f.vocab, &s, |p| f.files.get(p).cloned(), "missing.md", &Options::default()).unwrap().len(), 0);
}

#[test]
fn budget_keeps_best_passages_within_tokens() {
    let f = reminder_vault();
    let r = ask(&f, &MapSignals::default(), "reminder notify", &Options::default());
    let all: usize = r.passages.iter().map(|p| p.text.chars().count().div_ceil(4)).sum();
    let b = budget(r.passages.clone(), all);
    assert_eq!((b.passages.len(), b.truncated), (r.passages.len(), false));
    let b = budget(r.passages, 8);
    assert!(b.tokens <= 8 && b.truncated);
}