haematite-cli 0.21.7

//! V7a: one page of the embedding job, with an in-process stand-in for the
//! worker. No process is started.

use std::cell::RefCell;
use std::error::Error;
use std::num::NonZeroUsize;

use haem_frames::embedding::{Distance, input_refusal};
use haematite::db::DatabaseConfig;
use haematite::store::NodeCacheBudget;
use serde_json::value::RawValue;

use super::{Bounds, Counts, Embed, LongerThanModel, Step, make, step};
use crate::serve::embed::{Answer, Embedded, Longer, ModelId, WindowVector};
use crate::serve::protocol::{
    Change, CollectionWrite, DocumentWrite, Metric, Mutation, Selection, Version,
};
use crate::serve::storage::embed_pending::Waiting;
use crate::serve::storage::embedded::{self, Declaration, ModelName, Position, VectorKind};
use crate::serve::storage::{Failure, Storage, reads};

type TestResult = Result<(), Box<dyn Error>>;

const FIELD: &str = "body_vec";

/// A worker in this process: answers each text with `[1, 0]`, records every
/// batch it was sent, fails its call numbered `fail_on`, and answers one
/// vector short when `short`. A text in `longer` is one the model reads only
/// part of: 600 tokens past the limit of 512. As the embedder's `asked`
/// does, one text it refuses refuses the whole call.
#[derive(Default)]
struct StandIn {
    sent: RefCell<Vec<Vec<String>>>,
    fail_on: Option<usize>,
    short: bool,
    /// Answer every vector one member too long.
    long: bool,
    longer: Vec<&'static str>,
    /// EMB-E2: texts answered as two windows over their bytes.
    windowed: Vec<&'static str>,
    /// EMB-E2: texts whose cover needs more windows than the cap.
    over_cap: Vec<&'static str>,
}

impl Embed for StandIn {
    fn documents(&self, texts: &[String]) -> Result<Embedded, Failure> {
        let call = {
            let mut sent = self.sent.borrow_mut();
            sent.push(texts.to_vec());
            sent.len()
        };
        if self.fail_on == Some(call) {
            return Err(Failure::new("embed_worker_exited", "scripted to fail"));
        }
        if let Some(index) = texts
            .iter()
            .position(|text| input_refusal(text, 64).is_some())
        {
            return Err(Failure::new(
                "embed_input_empty",
                format!("text {index} cannot be embedded"),
            ));
        }
        let one = if self.long {
            vec![1.0, 0.0, 0.0]
        } else {
            vec![1.0, 0.0]
        };
        let mut answers: Vec<Answer> = texts
            .iter()
            .map(|text| {
                if self.longer.contains(&text.as_str()) {
                    Answer::Longer(LONGER)
                } else if self.over_cap.contains(&text.as_str()) {
                    Answer::OverCap(LONGER)
                } else if self.windowed.contains(&text.as_str()) {
                    Answer::Windows(halves(text, &one))
                } else {
                    Answer::Vector(one.clone())
                }
            })
            .collect();
        if self.short {
            answers.pop();
        }
        Ok(Embedded {
            model: ModelId {
                name: "stand-in".to_owned(),
                revision: "1".to_owned(),
                manifest_sha256: "abc".to_owned(),
                dimensions: 2,
                distance: Distance::Cosine,
                normalised: false,
            },
            answers,
        })
    }
}

/// Two windows over `text`'s bytes, overlapping by one.
fn halves(text: &str, vector: &[f32]) -> Vec<WindowVector> {
    let length = u64::try_from(text.len()).unwrap_or(u64::MAX);
    let middle = length / 2;
    [(0, middle + 1), (middle, length)]
        .into_iter()
        .map(|(start, end)| WindowVector {
            start,
            end,
            vector: vector.to_vec(),
        })
        .collect()
}

const LONGER: Longer = Longer {
    tokens_total: 600,
    tokens_limit: 512,
};

impl StandIn {
    fn calls(&self) -> usize {
        self.sent.borrow().len()
    }
}

fn create() -> Result<(tempfile::TempDir, Storage), Box<dyn Error>> {
    let directory = tempfile::tempdir()?;
    let mut storage = Storage::create(
        DatabaseConfig {
            data_dir: directory.path().join("database"),
            shard_count: 2,
            distributed: None,
            executor_threads: Some(1),
            node_cache_budget: Some(NodeCacheBudget::Unlimited),
            retention: haematite::Retention::reclaim(0),
        },
        NodeCacheBudget::Unlimited,
        fixture_long_wait_interval(),
    )?;
    mutate(
        &mut storage,
        Change::Collection(CollectionWrite {
            collection: "notes".to_owned(),
            schema: None,
        }),
    )?;
    let declaration = Declaration {
        kind: VectorKind::Dense,
        source: "body".to_owned(),
        model: ModelName {
            name: "stand-in".to_owned(),
            revision: "1".to_owned(),
            manifest_sha256: "abc".to_owned(),
        },
        dimensions: NonZeroUsize::new(2).ok_or("two is not zero")?,
        distance: Metric::Cosine,
    };
    storage.declare_for_tests("main", "notes", FIELD, &declaration)?;
    Ok((directory, storage))
}

fn head(storage: &Storage) -> Result<Version, Failure> {
    let at = storage.capture(&Selection {
        branch: "main".to_owned(),
        version: None,
    })?;
    Ok(at.version)
}

fn mutate(storage: &mut Storage, change: Change) -> Result<Version, Failure> {
    let at = head(storage)?;
    let receipt = storage.mutate(&Mutation {
        request_id: format!("job-{}-{}", at.generation, at.sequence),
        branch: "main".to_owned(),
        expected: at,
        change,
    })?;
    Ok(receipt.version)
}

fn write(storage: &mut Storage, id: &str, document: &str) -> Result<Version, Box<dyn Error>> {
    let change = Change::Put(DocumentWrite {
        collection: "notes".to_owned(),
        document_id: id.to_owned(),
        document: RawValue::from_string(document.to_owned())?,
    });
    Ok(mutate(storage, change)?)
}

/// Records `n-1` to `n-{count}`, each with its own text.
fn records(storage: &mut Storage, count: u8) -> TestResult {
    for index in 1..=count {
        let document = format!(r#"{{"body":"text {index}"}}"#);
        write(storage, &format!("n-{index}"), &document)?;
    }
    Ok(())
}

fn bounds(page_records: usize) -> Result<Bounds, Box<dyn Error>> {
    Ok(Bounds {
        page_records: NonZeroUsize::new(page_records).ok_or("no page")?,
        batch: NonZeroUsize::new(64).ok_or("no batch")?,
        input_bytes: NonZeroUsize::new(64).ok_or("no input")?,
        call: std::time::Duration::from_secs(1),
    })
}

fn position(storage: &Storage) -> Result<Option<Position>, Failure> {
    let companion = embedded::companion_name("main");
    storage.position(&companion, &embedded::position_key("notes", FIELD))
}

fn has_vector(storage: &Storage, id: &str) -> Result<bool, Box<dyn Error>> {
    let at = storage.capture(&Selection {
        branch: "main".to_owned(),
        version: None,
    })?;
    let (_declaration, companion) = storage.embedded_field(&at, "notes", FIELD)?;
    let Some(companion) = companion else {
        return Ok(false);
    };
    let key = embedded::vector_key("notes", FIELD, id);
    Ok(reads::get_bytes(&storage.reader()?, &companion, &key)?.is_some())
}

fn run(
    storage: &mut Storage,
    worker: &StandIn,
    page_records: usize,
) -> Result<usize, Box<dyn Error>> {
    let mut pages_with_text = 0;
    for _step in 0..100 {
        match step(
            storage,
            "main",
            "notes",
            FIELD,
            worker,
            bounds(page_records)?,
        )? {
            Step::Idle => return Ok(pages_with_text),
            Step::Committed(counts) if counts.made > 0 => pages_with_text += 1,
            Step::Committed(_) => {}
            Step::Failed(failure) => return Err(failure.into()),
        }
    }
    Err("the job did not finish in 100 steps".into())
}

/// Fails while the steps do not carry a walk to its end, or call the worker
/// more than once a page: seven records, two a page, each has its vector and
/// the worker was called once for each page that had text.
#[test]
fn seven_records_two_a_page_are_embedded_whole() -> TestResult {
    let (_directory, mut storage) = create()?;
    records(&mut storage, 7)?;
    let worker = StandIn::default();
    let pages = run(&mut storage, &worker, 2)?;
    assert_eq!(worker.calls(), pages);
    for index in 1..=7 {
        assert!(has_vector(&storage, &format!("n-{index}"))?, "n-{index}");
    }
    Ok(())
}

/// Fails while a failed worker call commits anything, or moves the job on:
/// the first page commits, the second fails by name with the position where
/// the first left it, and the next step sends the same texts again.
#[test]
fn a_failed_call_commits_nothing_and_the_next_step_takes_the_same_page() -> TestResult {
    let (_directory, mut storage) = create()?;
    records(&mut storage, 4)?;
    let worker = StandIn {
        fail_on: Some(2),
        ..StandIn::default()
    };
    let first = step(&mut storage, "main", "notes", FIELD, &worker, bounds(2)?)?;
    assert!(matches!(first, Step::Committed(_)), "{first:?}");
    let after_first = position(&storage)?;
    let failed = step(&mut storage, "main", "notes", FIELD, &worker, bounds(2)?)?;
    let Step::Failed(failure) = failed else {
        return Err(format!("the failed call gave {failed:?}").into());
    };
    assert_eq!(failure.code, "embed_worker_exited");
    assert_eq!(position(&storage)?, after_first);
    let again = step(&mut storage, "main", "notes", FIELD, &worker, bounds(2)?)?;
    assert!(matches!(again, Step::Committed(_)), "{again:?}");
    let sent = worker.sent.borrow();
    assert_eq!(sent.get(1), sent.get(2), "the next step sent another page");
    Ok(())
}

/// Fails while a page with nothing to send calls the worker or stays put:
/// records with no text make no call, and the walk still ends at the head.
#[test]
fn nothing_to_send_makes_no_call_and_moves_the_position() -> TestResult {
    let (_directory, mut storage) = create()?;
    write(&mut storage, "n-1", r#"{"body":7}"#)?;
    write(&mut storage, "n-2", r#"{"title":"no body"}"#)?;
    let worker = StandIn::default();
    let done = step(&mut storage, "main", "notes", FIELD, &worker, bounds(64)?)?;
    let Step::Committed(counts) = done else {
        return Err(format!("the page gave {done:?}").into());
    };
    assert_eq!((counts.made, counts.no_text, counts.ended), (0, 2, true));
    assert_eq!(worker.calls(), 0);
    let through = position(&storage)?.and_then(|at| at.through);
    assert_eq!(through, Some(head(&storage)?));
    Ok(())
}

/// Fails while an answer with too few vectors is committed: it is refused by
/// name, and nothing is written.
#[test]
fn one_vector_too_few_is_refused_and_nothing_committed() -> TestResult {
    let (_directory, mut storage) = create()?;
    records(&mut storage, 2)?;
    let worker = StandIn {
        short: true,
        ..StandIn::default()
    };
    let refused = step(&mut storage, "main", "notes", FIELD, &worker, bounds(64)?)?;
    let Step::Failed(failure) = refused else {
        return Err(format!("the short answer gave {refused:?}").into());
    };
    assert_eq!(failure.code, "embed_vector_count");
    assert_eq!(position(&storage)?, None);
    assert!(!has_vector(&storage, "n-1")?);
    Ok(())
}

/// Fails while a worker answering vectors of the wrong length is reported as
/// the store failing: it is the worker's failure by name, and nothing is
/// committed.
#[test]
fn vectors_of_the_wrong_length_are_the_workers_failure() -> TestResult {
    let (_directory, mut storage) = create()?;
    records(&mut storage, 2)?;
    let worker = StandIn {
        long: true,
        ..StandIn::default()
    };
    let refused = step(&mut storage, "main", "notes", FIELD, &worker, bounds(64)?)?;
    let Step::Failed(failure) = refused else {
        return Err(format!("the long answer gave {refused:?}").into());
    };
    assert_eq!(failure.code, "embed_vector_length");
    assert_eq!(position(&storage)?, None);
    assert!(!has_vector(&storage, "n-1")?);
    Ok(())
}

/// The one page a step commits, or what it did instead.
fn committed(step: Step) -> Result<Counts, Box<dyn Error>> {
    match step {
        Step::Committed(counts) => Ok(counts),
        other => Err(format!("the page gave {other:?}").into()),
    }
}

/// EMB1: one text of three the model reads only part of: the other two are
/// stored, it is named and nothing is stored for it, the walk goes past it,
/// and a later walk does not send it again.
#[test]
fn a_text_longer_than_the_model_is_named_and_its_batch_stored() -> TestResult {
    let (_directory, mut storage) = create()?;
    records(&mut storage, 3)?;
    let worker = StandIn {
        longer: vec!["text 2"],
        ..StandIn::default()
    };
    let done = step(&mut storage, "main", "notes", FIELD, &worker, bounds(3)?)?;
    let counts = committed(done)?;
    assert_eq!(
        counts,
        Counts {
            made: 2,
            removed: 0,
            no_text: 0,
            too_long: 0,
            longer_than_model: 1,
            empty_input: 0,
            windowed: 0,
            windows_over_cap: 0,
            ended: true,
        }
    );
    assert_eq!(worker.calls(), 1);
    assert!(has_vector(&storage, "n-1")?);
    assert!(!has_vector(&storage, "n-2")?);
    assert!(has_vector(&storage, "n-3")?);
    let through = position(&storage)?.and_then(|at| at.through);
    assert_eq!(through, Some(head(&storage)?));
    write(&mut storage, "n-4", r#"{"body":"text 4"}"#)?;
    run(&mut storage, &worker, 3)?;
    let sent = worker.sent.borrow();
    assert_eq!(sent.len(), 2);
    assert_eq!(*sent.last().ok_or("no call")?, ["text 4"]);
    Ok(())
}

/// EMB1: the page names a longer text with its tokens and the model's
/// limit, and holds no vector for it.
#[test]
fn a_made_page_names_its_longer_text_with_its_tokens() -> TestResult {
    let (_directory, mut storage) = create()?;
    records(&mut storage, 3)?;
    let worker = StandIn {
        longer: vec!["text 2"],
        ..StandIn::default()
    };
    let page = NonZeroUsize::new(3).ok_or("no page")?;
    let input = NonZeroUsize::new(64).ok_or("no input")?;
    let pending = storage
        .embed_pending("main", "notes", FIELD, page, input)?
        .ok_or("nothing pending")?;
    let at = storage.capture(&Selection {
        branch: "main".to_owned(),
        version: None,
    })?;
    let (declaration, _companion) = storage.embedded_field(&at, "notes", FIELD)?;
    let made = make(&pending, "notes", FIELD, &declaration, &worker, bounds(3)?)?;
    assert_eq!(
        made.longer_than_model,
        [LongerThanModel {
            id: "n-2".to_owned(),
            tokens_total: 600,
            tokens_limit: 512,
        }]
    );
    let ids: Vec<&str> = made
        .page
        .made
        .iter()
        .map(|record| record.id.as_str())
        .collect();
    assert_eq!(ids, ["n-1", "n-3"]);
    Ok(())
}

/// EMB-E1: a page listing texts the worker would refuse, first, between and
/// last (as a walk read under another check could), sends only the rest, in
/// one call. The refused are named, passed and keep no vector. Today the
/// first refuses the whole call and nothing of the page is made.
#[test]
fn a_refused_text_never_refuses_its_batch_mates() -> TestResult {
    let (_directory, mut storage) = create()?;
    records(&mut storage, 3)?;
    let worker = StandIn::default();
    let page = NonZeroUsize::new(3).ok_or("no page")?;
    let input = NonZeroUsize::new(64).ok_or("no input")?;
    let mut pending = storage
        .embed_pending("main", "notes", FIELD, page, input)?
        .ok_or("nothing pending")?;
    for (at, id, text) in [(0, "n-0", ""), (2, "n-1b", " \n"), (5, "n-9", "\t")] {
        let waiting = Waiting::Text {
            id: id.to_owned(),
            made_from: pending.toward,
            text: text.to_owned(),
        };
        pending.waiting.insert(at, waiting);
    }
    let at = storage.capture(&Selection {
        branch: "main".to_owned(),
        version: None,
    })?;
    let (declaration, _companion) = storage.embedded_field(&at, "notes", FIELD)?;
    let made = make(&pending, "notes", FIELD, &declaration, &worker, bounds(3)?)?;
    assert_eq!(*worker.sent.borrow(), [["text 1", "text 2", "text 3"]]);
    assert_eq!(made.empty_input, ["n-0", "n-1b", "n-9"]);
    let ids: Vec<&str> = made
        .page
        .made
        .iter()
        .map(|record| record.id.as_str())
        .collect();
    assert_eq!(ids, ["n-1", "n-2", "n-3"]);
    assert_eq!(made.page.cleared, ["n-0", "n-1b", "n-9"]);
    Ok(())
}

/// EMB1: a page whose every text the model reads only part of stores
/// nothing, names each, and still ends its walk.
#[test]
fn a_page_of_only_longer_texts_stores_nothing_and_ends_its_walk() -> TestResult {
    let (_directory, mut storage) = create()?;
    records(&mut storage, 2)?;
    let worker = StandIn {
        longer: vec!["text 1", "text 2"],
        ..StandIn::default()
    };
    let done = step(&mut storage, "main", "notes", FIELD, &worker, bounds(64)?)?;
    let counts = committed(done)?;
    assert_eq!(
        (counts.made, counts.longer_than_model, counts.ended),
        (0, 2, true)
    );
    assert!(!has_vector(&storage, "n-1")?);
    assert!(!has_vector(&storage, "n-2")?);
    let idle = step(&mut storage, "main", "notes", FIELD, &worker, bounds(64)?)?;
    assert!(matches!(idle, Step::Idle), "{idle:?}");
    assert_eq!(worker.calls(), 1);
    Ok(())
}

/// RED4: a record whose text grows past what the model reads keeps no
/// vector made from its old text.
#[test]
fn a_text_grown_past_the_model_keeps_no_vector_of_its_old_text() -> TestResult {
    let (_directory, mut storage) = create()?;
    records(&mut storage, 1)?;
    let worker = StandIn {
        longer: vec!["grown"],
        ..StandIn::default()
    };
    run(&mut storage, &worker, 64)?;
    assert!(has_vector(&storage, "n-1")?);
    write(&mut storage, "n-1", r#"{"body":"grown"}"#)?;
    run(&mut storage, &worker, 64)?;
    assert!(!has_vector(&storage, "n-1")?);
    Ok(())
}

/// RED4: a record whose text grows past `input_bytes` (64 here) keeps no
/// vector made from its old text, and the worker is not sent it.
#[test]
fn a_text_grown_past_its_bytes_keeps_no_vector_of_its_old_text() -> TestResult {
    let (_directory, mut storage) = create()?;
    records(&mut storage, 1)?;
    let worker = StandIn::default();
    run(&mut storage, &worker, 64)?;
    assert!(has_vector(&storage, "n-1")?);
    let long = format!(r#"{{"body":"{}"}}"#, "x".repeat(65));
    write(&mut storage, "n-1", &long)?;
    run(&mut storage, &worker, 64)?;
    assert!(!has_vector(&storage, "n-1")?);
    assert_eq!(worker.calls(), 1);
    Ok(())
}

/// EMB-E1: a record whose text is emptied keeps no vector made from its
/// old text, and the worker is not sent it.
#[test]
fn a_text_emptied_keeps_no_vector_of_its_old_text() -> TestResult {
    let (_directory, mut storage) = create()?;
    records(&mut storage, 1)?;
    let worker = StandIn::default();
    run(&mut storage, &worker, 64)?;
    assert!(has_vector(&storage, "n-1")?);
    write(&mut storage, "n-1", r#"{"body":" "}"#)?;
    run(&mut storage, &worker, 64)?;
    assert!(!has_vector(&storage, "n-1")?);
    assert_eq!(worker.calls(), 1);
    Ok(())
}

/// P2a: a record whose source stops being a string keeps no vector made
/// from its old text, and the worker is not sent it.
#[test]
fn a_text_no_longer_a_string_keeps_no_vector_of_its_old_text() -> TestResult {
    let (_directory, mut storage) = create()?;
    records(&mut storage, 1)?;
    let worker = StandIn::default();
    run(&mut storage, &worker, 64)?;
    assert!(has_vector(&storage, "n-1")?);
    write(&mut storage, "n-1", r#"{"body":7}"#)?;
    run(&mut storage, &worker, 64)?;
    assert!(!has_vector(&storage, "n-1")?);
    assert_eq!(worker.calls(), 1);
    Ok(())
}

fn has_window(storage: &Storage, id: &str) -> Result<bool, Box<dyn Error>> {
    let at = storage.capture(&Selection {
        branch: "main".to_owned(),
        version: None,
    })?;
    let (_declaration, companion) = storage.embedded_field(&at, "notes", FIELD)?;
    let Some(companion) = companion else {
        return Ok(false);
    };
    let key = embedded::window_key("notes", FIELD, id, 0);
    Ok(reads::get_bytes(&storage.reader()?, &companion, &key)?.is_some())
}

/// EMB-E2 (d), (f): one text of three answered as windows is stored as its
/// windows, never as a `v\0` vector, and counted `windowed`.
#[test]
fn a_long_text_is_stored_as_its_windows_and_counted() -> TestResult {
    let (_directory, mut storage) = create()?;
    records(&mut storage, 3)?;
    let worker = StandIn {
        windowed: vec!["text 2"],
        ..StandIn::default()
    };
    let done = step(&mut storage, "main", "notes", FIELD, &worker, bounds(3)?)?;
    let counts = committed(done)?;
    assert_eq!((counts.made, counts.windowed), (2, 1));
    assert!(has_window(&storage, "n-2")?);
    assert!(!has_vector(&storage, "n-2")?);
    assert!(has_vector(&storage, "n-1")?);
    Ok(())
}

/// EMB-E2 (e): fails if the first N windows are kept: a text past the cap
/// gets no vector and no window, is counted `windows_over_cap`, and the walk
/// goes past it.
#[test]
fn a_text_over_the_cap_gets_no_vector_and_is_counted() -> TestResult {
    let (_directory, mut storage) = create()?;
    records(&mut storage, 3)?;
    let worker = StandIn {
        over_cap: vec!["text 2"],
        ..StandIn::default()
    };
    let done = step(&mut storage, "main", "notes", FIELD, &worker, bounds(3)?)?;
    let counts = committed(done)?;
    assert_eq!((counts.made, counts.windows_over_cap), (2, 1));
    assert_eq!(
        counts.longer_than_model, 1,
        "(i): it is one the model reads part of"
    );
    assert!(counts.ended);
    assert!(!has_vector(&storage, "n-2")?);
    assert!(!has_window(&storage, "n-2")?);
    Ok(())
}

/// EMB-E2 (d): fails on a partial set: a call answering windows that fails
/// commits nothing, the position stays where it was and no `vw\0` key
/// exists.
#[test]
fn a_window_call_failing_mid_record_commits_nothing() -> TestResult {
    let (_directory, mut storage) = create()?;
    records(&mut storage, 2)?;
    let worker = StandIn {
        windowed: vec!["text 1", "text 2"],
        fail_on: Some(1),
        ..StandIn::default()
    };
    let failed = step(&mut storage, "main", "notes", FIELD, &worker, bounds(2)?)?;
    assert!(matches!(failed, Step::Failed(_)), "{failed:?}");
    assert_eq!(position(&storage)?, None);
    assert!(!has_window(&storage, "n-1")?);
    assert!(!has_window(&storage, "n-2")?);
    Ok(())
}

const fn fixture_long_wait_interval() -> std::time::Duration {
    std::time::Duration::from_secs(1)
}