haematite-cli 0.21.1

//! V8b: a question's vector, checked against the field's declaration.

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

use haem_frames::embedding::Distance;

use super::{Ask, question_for};
use crate::serve::embed::{Answer, Embedded, Longer, ModelId};
use crate::serve::protocol::Metric;
use crate::serve::storage::Failure;
use crate::serve::storage::embedded::{Declaration, ModelName, VectorKind};

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

/// A worker that answers every question with `answers`, as `model`.
struct StandIn {
    model: &'static str,
    answers: Vec<Answer>,
}

impl Ask for StandIn {
    fn question(&self, _text: &str) -> Result<Embedded, Failure> {
        Ok(Embedded {
            model: ModelId {
                name: self.model.to_owned(),
                revision: "1".to_owned(),
                manifest_sha256: "abc".to_owned(),
                dimensions: 2,
                distance: Distance::Cosine,
                normalised: true,
            },
            answers: self.answers.clone(),
        })
    }
}

fn declaration() -> Result<Declaration, Box<dyn Error>> {
    Ok(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,
    })
}

fn code(outcome: Result<Vec<f32>, Failure>) -> Option<String> {
    outcome.err().map(|failure| failure.code)
}

/// The control: the declared model's one vector of the declared length is
/// the question.
#[test]
fn the_declared_models_one_vector_is_the_question() -> TestResult {
    let asker = StandIn {
        model: "stand-in",
        answers: vec![Answer::Vector(vec![0.6, 0.8])],
    };
    assert_eq!(
        question_for(&asker, "chest pain", &declaration()?)?,
        [0.6_f32, 0.8]
    );
    Ok(())
}

/// Fails while the model is not compared: another model's vector refuses
/// `embed_model_mismatch`.
#[test]
fn another_models_question_is_refused() -> TestResult {
    let asker = StandIn {
        model: "another",
        answers: vec![Answer::Vector(vec![0.6, 0.8])],
    };
    let outcome = question_for(&asker, "chest pain", &declaration()?);
    assert_eq!(code(outcome).as_deref(), Some("embed_model_mismatch"));
    Ok(())
}

/// Fails while the count is not checked: two vectors for one question
/// refuse `embed_vector_count`.
#[test]
fn two_vectors_for_one_question_are_refused() -> TestResult {
    let asker = StandIn {
        model: "stand-in",
        answers: vec![
            Answer::Vector(vec![0.6, 0.8]),
            Answer::Vector(vec![0.8, 0.6]),
        ],
    };
    let outcome = question_for(&asker, "chest pain", &declaration()?);
    assert_eq!(code(outcome).as_deref(), Some("embed_vector_count"));
    Ok(())
}

/// Fails while the length is not checked: a vector of three members for a
/// field of two refuses `embed_vector_length`.
#[test]
fn a_question_of_another_length_is_refused() -> TestResult {
    let asker = StandIn {
        model: "stand-in",
        answers: vec![Answer::Vector(vec![0.6, 0.8, 0.0])],
    };
    let outcome = question_for(&asker, "chest pain", &declaration()?);
    assert_eq!(code(outcome).as_deref(), Some("embed_vector_length"));
    Ok(())
}

/// EMB1: a question of more tokens than the model reads is refused by name,
/// with its tokens and the limit, never searched by a part of it.
#[test]
fn a_question_longer_than_the_model_is_refused_by_name() -> TestResult {
    let asker = StandIn {
        model: "stand-in",
        answers: vec![Answer::Longer(Longer {
            tokens_total: 600,
            tokens_limit: 512,
        })],
    };
    let refused = question_for(&asker, "words", &declaration()?)
        .err()
        .ok_or("a longer question was searched")?;
    assert_eq!(refused.code, "text_longer_than_model");
    assert_eq!(
        refused.message,
        "the question is 600 tokens; the model's limit is 512"
    );
    Ok(())
}