use std::error::Error as StdError;
use haem_frames::embedding::{
Distance, Embed, FailedCode, Input, ItemCode, Kind, Limits, Model, Outcome, PROTOCOL,
PROTOCOL_WINDOWS, Ready, Start, Tag, Vectors, Window, Windows, check_embed, check_model_limits,
check_ready, check_start, check_vectors, check_vectors_windowed, decode_vector, encode_vector,
input_refusal, vectors_message_size,
};
use serde_json::json;
type TestResult = Result<(), Box<dyn StdError>>;
const fn limits() -> Limits {
Limits {
memory_bytes: 1024,
threads: 2,
batch_items: 2,
input_bytes: 32,
tokens: 16,
frame_bytes: 8192,
}
}
fn start() -> Start {
Start {
protocol: PROTOCOL,
model_dir: "/models/test".to_owned(),
limits: limits(),
windows: None,
}
}
fn ready() -> Ready {
Ready {
protocol: PROTOCOL,
worker: "worker/1".to_owned(),
core_limit: 0,
descriptors_closed: 0,
environment: 0,
model: Model {
name: "dense".to_owned(),
revision: "r1".to_owned(),
manifest_sha256: "a".repeat(64),
dimensions: 2,
max_tokens: 32,
special_tokens: 2,
prefix_tokens: 0,
normalised: true,
distance: Distance::Cosine,
},
}
}
fn embed(text: &str) -> Embed {
Embed {
kind: Kind::Document,
items: vec![Input {
id: "one".to_owned(),
text: text.to_owned(),
}],
}
}
fn vectors() -> Result<Vectors, Box<dyn StdError>> {
Ok(Vectors {
items: vec![Outcome::Vector {
id: "one".to_owned(),
vector: encode_vector(&[1.0, 0.0])?,
tokens_read: 3,
tokens_total: 3,
}],
})
}
fn refused(code: ItemCode) -> Vectors {
Vectors {
items: vec![Outcome::Refused {
id: "one".to_owned(),
code,
tokens_total: None,
tokens_limit: None,
}],
}
}
fn longer(tokens_total: Option<u64>, tokens_limit: Option<u64>) -> Vectors {
Vectors {
items: vec![Outcome::Refused {
id: "one".to_owned(),
code: ItemCode::TextLongerThanModel,
tokens_total,
tokens_limit,
}],
}
}
fn code<T>(result: Result<T, haem_frames::embedding::Error>, expected: FailedCode) -> TestResult {
assert_eq!(result.err().ok_or("unexpected acceptance")?.code, expected);
Ok(())
}
#[test]
fn tags_are_exact_and_do_not_overlap_query_tags() -> TestResult {
for (tag, number) in [
(Tag::Start, 16),
(Tag::Ready, 17),
(Tag::Embed, 18),
(Tag::Vectors, 19),
(Tag::Failed, 20),
] {
assert_eq!(tag as u8, number);
assert_eq!(Tag::try_from(number)?, tag);
assert!(!(1..=6).contains(&number));
}
for tag in 1..=6 {
code(Tag::try_from(tag), FailedCode::EmbedProtocol)?;
}
Ok(())
}
#[test]
fn hexadecimal_round_trips_original_float_bits_and_little_endian_bytes() -> TestResult {
let values = [0.0_f32, -0.0, f32::MIN_POSITIVE, f32::MAX];
let text = encode_vector(&values)?;
assert_eq!(text.len(), 32);
assert!(text.starts_with("0000000000000080"));
let decoded = decode_vector(&text)?;
assert_eq!(
decoded
.iter()
.copied()
.map(f32::to_bits)
.collect::<Vec<_>>(),
values.into_iter().map(f32::to_bits).collect::<Vec<_>>()
);
assert_eq!(encode_vector(&[1.0])?, "0000803f");
Ok(())
}
#[test]
fn hexadecimal_refuses_uppercase_wrong_length_and_other_characters() -> TestResult {
for text in [
"0000803F",
"0",
"0000000g",
"é000000",
"0000 000",
"00000000\n",
] {
code(decode_vector(text), FailedCode::EmbedProtocol)?;
}
Ok(())
}
#[test]
fn hexadecimal_refuses_non_finite_numbers_on_both_sides() -> TestResult {
for member in [f32::INFINITY, f32::NEG_INFINITY, f32::NAN] {
code(encode_vector(&[member]), FailedCode::EmbedOutputInvalid)?;
}
for text in ["0000807f", "000080ff", "0000c07f"] {
code(decode_vector(text), FailedCode::EmbedOutputInvalid)?;
}
Ok(())
}
#[test]
fn valid_start_and_model_limits_pass() -> TestResult {
check_start(&start())?;
check_ready(&ready())?;
check_model_limits(&limits(), &ready().model)?;
Ok(())
}
#[test]
fn every_zero_limit_and_small_frame_is_embed_limits() -> TestResult {
for member in [
"memory_bytes",
"threads",
"batch_items",
"input_bytes",
"tokens",
"frame_bytes",
] {
let mut value = serde_json::to_value(start())?;
value["limits"][member] = json!(0);
let value: Start = serde_json::from_value(value)?;
code(check_start(&value), FailedCode::EmbedLimits)?;
}
let mut value = start();
value.limits.frame_bytes = 8191;
code(check_start(&value), FailedCode::EmbedLimits)
}
#[test]
fn relative_model_directory_and_wrong_start_protocol_are_named() -> TestResult {
let mut value = start();
value.model_dir = "relative".to_owned();
code(check_start(&value), FailedCode::EmbedLimits)?;
value = start();
value.protocol = 2;
code(check_start(&value), FailedCode::EmbedProtocol)
}
#[test]
fn model_token_and_frame_limits_use_the_same_shared_check() -> TestResult {
let model = ready().model;
let mut value = limits();
value.tokens = model.max_tokens + 1;
code(check_model_limits(&value, &model), FailedCode::EmbedLimits)?;
for tokens in [0, model.special_tokens] {
value = limits();
value.tokens = tokens;
code(check_model_limits(&value, &model), FailedCode::EmbedLimits)?;
}
value = limits();
value.frame_bytes = 1;
code(check_model_limits(&value, &model), FailedCode::EmbedLimits)
}
#[test]
fn model_token_budget_reserves_prefix_and_one_text_token() -> TestResult {
let mut model = ready().model;
model.prefix_tokens = 3;
let mut value = limits();
value.tokens = 5;
let error = check_model_limits(&value, &model)
.err()
.ok_or("token budget without text was accepted")?;
assert_eq!(error.code, FailedCode::EmbedLimits);
assert_eq!(
error.message,
"tokens 5 must be at least special_tokens 2 + prefix_tokens 3 + 1 and not exceed max_tokens 32"
);
value.tokens = 6;
check_model_limits(&value, &model)?;
Ok(())
}
#[test]
fn model_token_reservation_overflow_is_embed_limits() -> TestResult {
for (special, prefix) in [(u64::MAX - 1, 2), (u64::MAX - 1, 1)] {
let mut model = ready().model;
model.max_tokens = u64::MAX;
model.special_tokens = special;
model.prefix_tokens = prefix;
let mut value = limits();
value.tokens = u64::MAX;
let error = check_model_limits(&value, &model)
.err()
.ok_or("overflowing token reservation was accepted")?;
assert_eq!(error.code, FailedCode::EmbedLimits);
assert_eq!(
error.message,
format!(
"tokens {} must be at least special_tokens {special} + prefix_tokens {prefix} + 1 and not exceed max_tokens {}",
u64::MAX,
u64::MAX
)
);
}
Ok(())
}
#[test]
fn ready_rejects_each_malformed_model_fact_and_protocol() -> TestResult {
for (field, invalid) in [
("name", json!("")),
("manifest_sha256", json!("A".repeat(64))),
("manifest_sha256", json!("a".repeat(63))),
("dimensions", json!(0)),
("max_tokens", json!(0)),
("special_tokens", json!(32)),
] {
let mut value = serde_json::to_value(ready())?;
value["model"][field] = invalid;
code(
check_ready(&serde_json::from_value(value)?),
FailedCode::EmbedProtocol,
)?;
}
let mut value = ready();
value.protocol = PROTOCOL_WINDOWS + 1;
code(check_ready(&value), FailedCode::EmbedProtocol)
}
#[test]
fn frame_bound_matches_actual_maximum_escaping_and_integer_widths() -> TestResult {
let message = Vectors {
items: [char::from(1_u8), char::from(2_u8)]
.into_iter()
.map(|last| {
let mut id = "\u{1}".repeat(511);
id.push(last);
Outcome::Vector {
id,
vector: "00000000".repeat(3),
tokens_read: u64::MAX,
tokens_total: u64::MAX,
}
})
.collect(),
};
assert_eq!(
vectors_message_size(2, 3)?,
u64::try_from(serde_json::to_vec(&message)?.len())? + 1
);
code(vectors_message_size(0, 3), FailedCode::EmbedLimits)?;
code(vectors_message_size(2, 0), FailedCode::EmbedLimits)?;
code(vectors_message_size(u64::MAX, 3), FailedCode::EmbedLimits)?;
code(vectors_message_size(2, u64::MAX), FailedCode::EmbedLimits)
}
#[test]
fn model_frame_bound_accepts_exact_size_and_refuses_one_byte_less() -> TestResult {
let model = ready().model;
let mut value = limits();
value.frame_bytes = u32::try_from(vectors_message_size(2, model.dimensions)?)?;
check_model_limits(&value, &model)?;
value.frame_bytes -= 1;
code(check_model_limits(&value, &model), FailedCode::EmbedLimits)
}
#[test]
fn embed_accepts_legal_unique_byte_ids_and_both_kinds() -> TestResult {
let mut value = embed("text");
value.items[0].id = "é".repeat(256);
check_embed(&value, &limits())?;
value.kind = Kind::Query;
value.items.push(Input {
id: "two".to_owned(),
text: "more".to_owned(),
});
check_embed(&value, &limits())?;
Ok(())
}
#[test]
fn embed_refuses_empty_batch_and_bad_or_duplicate_ids() -> TestResult {
code(
check_embed(
&Embed {
kind: Kind::Query,
items: vec![],
},
&limits(),
),
FailedCode::EmbedProtocol,
)?;
for id in [String::new(), "\0".to_owned(), "é".repeat(257)] {
let mut value = embed("text");
value.items[0].id = id;
code(check_embed(&value, &limits()), FailedCode::EmbedProtocol)?;
}
let mut value = embed("text");
value.items.push(Input {
id: "one".to_owned(),
text: "more".to_owned(),
});
code(check_embed(&value, &limits()), FailedCode::EmbedProtocol)
}
#[test]
fn embed_batch_limit_is_embed_batch_too_large() -> TestResult {
let mut value = limits();
value.batch_items = 0;
code(
check_embed(&embed("text"), &value),
FailedCode::EmbedBatchTooLarge,
)
}
#[test]
fn classifier_uses_unicode_whitespace_utf8_bytes_and_empty_precedence() {
assert_eq!(input_refusal("", 32), Some(ItemCode::EmbedInputEmpty));
assert_eq!(
input_refusal("\u{2003}".repeat(20).as_str(), 1),
Some(ItemCode::EmbedInputEmpty)
);
assert_eq!(input_refusal("é", 1), Some(ItemCode::EmbedInputTooLarge));
assert_eq!(input_refusal("é", 2), None);
assert_eq!(input_refusal("\u{200b}", 32), None);
}
#[test]
fn service_accepts_a_valid_vector_and_matching_item_refusals() -> TestResult {
check_vectors(&vectors()?, &embed("text"), &ready(), &limits())?;
check_vectors(
&refused(ItemCode::EmbedInputEmpty),
&embed(" "),
&ready(),
&limits(),
)?;
check_vectors(
&refused(ItemCode::EmbedInputTooLarge),
&embed(&"x".repeat(33)),
&ready(),
&limits(),
)?;
Ok(())
}
#[test]
fn service_refuses_wrong_count_wrong_id_and_wrong_order_as_protocol() -> TestResult {
code(
check_vectors(
&Vectors { items: vec![] },
&embed("text"),
&ready(),
&limits(),
),
FailedCode::EmbedProtocol,
)?;
let mut value = vectors()?;
let Outcome::Vector { id, .. } = &mut value.items[0] else {
return Err("not a vector".into());
};
*id = "other".to_owned();
code(
check_vectors(&value, &embed("text"), &ready(), &limits()),
FailedCode::EmbedProtocol,
)?;
let inputs = Embed {
kind: Kind::Document,
items: vec![
Input {
id: "one".to_owned(),
text: " ".to_owned(),
},
Input {
id: "two".to_owned(),
text: " ".to_owned(),
},
],
};
let values = Vectors {
items: vec![
Outcome::Refused {
id: "two".to_owned(),
code: ItemCode::EmbedInputEmpty,
tokens_total: None,
tokens_limit: None,
},
Outcome::Refused {
id: "one".to_owned(),
code: ItemCode::EmbedInputEmpty,
tokens_total: None,
tokens_limit: None,
},
],
};
code(
check_vectors(&values, &inputs, &ready(), &limits()),
FailedCode::EmbedProtocol,
)
}
#[test]
fn service_distinguishes_hex_form_from_vector_content() -> TestResult {
for (text, expected) in [
("0000803F", FailedCode::EmbedProtocol),
("0", FailedCode::EmbedProtocol),
("0000000g", FailedCode::EmbedProtocol),
("0000807f00000000", FailedCode::EmbedOutputInvalid),
("0000803f", FailedCode::EmbedOutputInvalid),
] {
let value = Vectors {
items: vec![Outcome::Vector {
id: "one".to_owned(),
vector: text.to_owned(),
tokens_read: 3,
tokens_total: 3,
}],
};
code(
check_vectors(&value, &embed("text"), &ready(), &limits()),
expected,
)?;
}
Ok(())
}
#[test]
fn service_checks_each_token_count_boundary() -> TestResult {
for (read, total) in [(0, 3), (4, 3), (17, 20)] {
let value = Vectors {
items: vec![Outcome::Vector {
id: "one".to_owned(),
vector: encode_vector(&[1.0, 0.0])?,
tokens_read: read,
tokens_total: total,
}],
};
code(
check_vectors(&value, &embed("text"), &ready(), &limits()),
FailedCode::EmbedOutputInvalid,
)?;
}
let value = Vectors {
items: vec![Outcome::Vector {
id: "one".to_owned(),
vector: encode_vector(&[1.0, 0.0])?,
tokens_read: 16,
tokens_total: 100,
}],
};
check_vectors(&value, &embed("text"), &ready(), &limits())?;
Ok(())
}
#[test]
fn service_checks_norm_only_when_ready_says_normalised() -> TestResult {
let value = Vectors {
items: vec![Outcome::Vector {
id: "one".to_owned(),
vector: encode_vector(&[2.0, 0.0])?,
tokens_read: 3,
tokens_total: 3,
}],
};
code(
check_vectors(&value, &embed("text"), &ready(), &limits()),
FailedCode::EmbedOutputInvalid,
)?;
let mut model = ready();
model.model.normalised = false;
check_vectors(&value, &embed("text"), &model, &limits())?;
let near = Vectors {
items: vec![Outcome::Vector {
id: "one".to_owned(),
vector: encode_vector(&[1.00005, 0.0])?,
tokens_read: 3,
tokens_total: 3,
}],
};
check_vectors(&near, &embed("text"), &ready(), &limits())?;
Ok(())
}
#[test]
fn service_refuses_every_wrong_classifier_outcome() -> TestResult {
for text in ["", &"x".repeat(33)] {
code(
check_vectors(&vectors()?, &embed(text), &ready(), &limits()),
FailedCode::EmbedOutputInvalid,
)?;
}
code(
check_vectors(
&refused(ItemCode::EmbedInputEmpty),
&embed("text"),
&ready(),
&limits(),
),
FailedCode::EmbedOutputInvalid,
)?;
code(
check_vectors(
&refused(ItemCode::EmbedInputTooLarge),
&embed(" "),
&ready(),
&limits(),
),
FailedCode::EmbedOutputInvalid,
)
}
#[test]
fn service_accepts_a_text_longer_than_the_model_only_beyond_what_the_model_reads() -> TestResult {
check_vectors(
&longer(Some(31), Some(32)),
&embed("text"),
&ready(),
&limits(),
)?;
for (total, limit) in [(30, 32), (31, 16), (31, 30)] {
code(
check_vectors(
&longer(Some(total), Some(limit)),
&embed("text"),
&ready(),
&limits(),
),
FailedCode::EmbedOutputInvalid,
)?;
}
Ok(())
}
#[test]
fn service_refuses_text_longer_than_model_without_both_counts() -> TestResult {
for (total, limit) in [(None, None), (Some(31), None), (None, Some(32))] {
code(
check_vectors(&longer(total, limit), &embed("text"), &ready(), &limits()),
FailedCode::EmbedOutputInvalid,
)?;
}
Ok(())
}
#[test]
fn service_refuses_text_longer_than_model_for_a_text_the_classifier_refused() -> TestResult {
for text in [" ", &"x".repeat(33)] {
code(
check_vectors(
&longer(Some(31), Some(32)),
&embed(text),
&ready(),
&limits(),
),
FailedCode::EmbedOutputInvalid,
)?;
}
Ok(())
}
#[test]
fn service_refuses_token_counts_on_any_other_refusal_code() -> TestResult {
for (total, limit) in [(Some(31), Some(32)), (Some(31), None), (None, Some(32))] {
code(
check_vectors(
&Vectors {
items: vec![Outcome::Refused {
id: "one".to_owned(),
code: ItemCode::EmbedInputEmpty,
tokens_total: total,
tokens_limit: limit,
}],
},
&embed(" "),
&ready(),
&limits(),
),
FailedCode::EmbedOutputInvalid,
)?;
}
Ok(())
}
const ASKED: Windows = Windows {
overlap_tokens: 4,
max_windows: 3,
};
fn window(start: u64, end: u64) -> Result<Window, Box<dyn StdError>> {
Ok(Window {
start,
end,
vector: encode_vector(&[1.0, 0.0])?,
tokens_read: 8,
})
}
fn windowed(windows: Vec<Window>) -> Vectors {
Vectors {
items: vec![Outcome::Windows {
id: "one".to_owned(),
windows,
tokens_total: 40,
}],
}
}
fn check_windowed(
vectors: &Vectors,
text: &str,
asked: Option<Windows>,
) -> Result<(), haem_frames::embedding::Error> {
check_vectors_windowed(vectors, &embed(text), &ready(), &limits(), asked)
}
#[test]
fn windows_are_asked_for_by_protocol_two_and_only_by_it() -> TestResult {
let windowed_start = Start {
protocol: PROTOCOL_WINDOWS,
windows: Some(ASKED),
..start()
};
check_start(&windowed_start)?;
let no_windows = Start {
windows: None,
..start_with(PROTOCOL_WINDOWS)
};
code(check_start(&no_windows), FailedCode::EmbedProtocol)?;
let old_with_windows = Start {
windows: Some(ASKED),
..start()
};
code(check_start(&old_with_windows), FailedCode::EmbedProtocol)?;
for asked in [
Windows {
max_windows: 1,
..ASKED
},
Windows {
overlap_tokens: 16,
..ASKED
},
] {
let bad = Start {
windows: Some(asked),
..start_with(PROTOCOL_WINDOWS)
};
code(check_start(&bad), FailedCode::EmbedLimits)?;
}
Ok(())
}
fn start_with(protocol: u32) -> Start {
Start {
protocol,
..start()
}
}
#[test]
fn a_text_three_windows_long_is_covered_whole() -> TestResult {
let text = "abcdefghij";
let whole = vec![window(0, 4)?, window(3, 7)?, window(6, 10)?];
check_windowed(&windowed(whole), text, Some(ASKED))?;
let touching = vec![window(0, 4)?, window(4, 10)?];
check_windowed(&windowed(touching), text, Some(ASKED))?;
let covers = [
vec![window(0, 4)?, window(3, 7)?],
vec![window(0, 4)?, window(5, 10)?],
vec![window(1, 6)?, window(5, 10)?],
vec![window(0, 6)?, window(5, 11)?],
vec![window(0, 3)?, window(2, 5)?, window(4, 8)?, window(7, 10)?],
vec![window(0, 10)?],
vec![window(0, 4)?, window(0, 10)?],
vec![window(0, 10)?, window(5, 10)?],
];
for cover in covers {
code(
check_windowed(&windowed(cover), text, Some(ASKED)),
FailedCode::EmbedOutputInvalid,
)?;
}
Ok(())
}
#[test]
fn a_window_bound_inside_a_character_is_refused() -> TestResult {
let text = "abcdéfghi";
let split = vec![window(0, 5)?, window(3, 10)?];
code(
check_windowed(&windowed(split), text, Some(ASKED)),
FailedCode::EmbedOutputInvalid,
)?;
let whole = vec![window(0, 6)?, window(3, 10)?];
Ok(check_windowed(&windowed(whole), text, Some(ASKED))?)
}
#[test]
fn windows_and_over_cap_answer_only_documents_when_asked() -> TestResult {
let text = "abcdefghij";
let whole = || -> Result<Vectors, Box<dyn StdError>> {
Ok(windowed(vec![window(0, 6)?, window(4, 10)?]))
};
code(
check_windowed(&whole()?, text, None),
FailedCode::EmbedOutputInvalid,
)?;
let question = Embed {
kind: Kind::Query,
..embed(text)
};
code(
check_vectors_windowed(&whole()?, &question, &ready(), &limits(), Some(ASKED)),
FailedCode::EmbedOutputInvalid,
)?;
let over = Vectors {
items: vec![Outcome::Refused {
id: "one".to_owned(),
code: ItemCode::WindowsOverCap,
tokens_total: Some(400),
tokens_limit: Some(32),
}],
};
check_windowed(&over, text, Some(ASKED))?;
code(
check_windowed(&over, text, None),
FailedCode::EmbedOutputInvalid,
)
}
#[test]
fn a_window_read_past_the_token_limit_is_refused() -> TestResult {
let text = "abcdefghij";
let cut = Window {
tokens_read: 17,
..window(4, 10)?
};
code(
check_windowed(&windowed(vec![window(0, 6)?, cut]), text, Some(ASKED)),
FailedCode::EmbedOutputInvalid,
)?;
let read = Window {
tokens_read: 16,
..window(4, 10)?
};
Ok(check_windowed(
&windowed(vec![window(0, 6)?, read]),
text,
Some(ASKED),
)?)
}