811 lines
27 KiB
Rust
811 lines
27 KiB
Rust
//! Multi-Modal Memory — Track 6
|
|
//!
|
|
//! Supports storing and searching memories across multiple modalities:
|
|
//! Text, Image, Audio, Video, and Structured data.
|
|
//!
|
|
//! Each record can carry multiple embeddings (e.g. a CLIP image embedding
|
|
//! alongside a text embedding for the same document), enabling both
|
|
//! within-modality and cross-modal retrieval.
|
|
|
|
use std::collections::HashMap;
|
|
use std::fmt;
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Modality
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// The sensory / semantic modality of a memory record.
|
|
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
|
|
pub enum Modality {
|
|
Text,
|
|
Image,
|
|
Audio,
|
|
Video,
|
|
Structured,
|
|
}
|
|
|
|
impl fmt::Display for Modality {
|
|
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
|
match self {
|
|
Modality::Text => write!(f, "text"),
|
|
Modality::Image => write!(f, "image"),
|
|
Modality::Audio => write!(f, "audio"),
|
|
Modality::Video => write!(f, "video"),
|
|
Modality::Structured => write!(f, "structured"),
|
|
}
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// ModalEmbedding
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// A dense float embedding produced by a specific model for a specific modality.
|
|
#[derive(Debug, Clone)]
|
|
pub struct ModalEmbedding {
|
|
pub modality: Modality,
|
|
/// Raw embedding values.
|
|
pub embedding: Vec<f32>,
|
|
/// Expected dimensionality (must equal `embedding.len()`).
|
|
pub dimension: usize,
|
|
/// Identifier of the model that produced this embedding,
|
|
/// e.g. `"clip-vit-large"`, `"whisper-base"`, `"text-embedding-3-small"`.
|
|
pub model_id: String,
|
|
}
|
|
|
|
impl ModalEmbedding {
|
|
/// Create a new embedding, setting `dimension` from the vector length.
|
|
pub fn new(modality: Modality, embedding: Vec<f32>, model_id: impl Into<String>) -> Self {
|
|
let dimension = embedding.len();
|
|
Self {
|
|
modality,
|
|
embedding,
|
|
dimension,
|
|
model_id: model_id.into(),
|
|
}
|
|
}
|
|
|
|
/// L2 norm of the embedding.
|
|
#[inline]
|
|
pub fn norm(&self) -> f32 {
|
|
self.embedding.iter().map(|x| x * x).sum::<f32>().sqrt()
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// MediaRefType / MediaRef
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// Where the raw media bytes live.
|
|
#[derive(Debug, Clone)]
|
|
pub enum MediaRefType {
|
|
/// A path on the local filesystem.
|
|
Path(String),
|
|
/// A remote URL.
|
|
Url(String),
|
|
/// Bytes stored inline.
|
|
Inline(Vec<u8>),
|
|
}
|
|
|
|
/// A reference to the raw media associated with a record.
|
|
#[derive(Debug, Clone)]
|
|
pub struct MediaRef {
|
|
pub ref_type: MediaRefType,
|
|
pub mime_type: String,
|
|
pub size_bytes: Option<u64>,
|
|
/// FNV-1a 64-bit hash of the content (for Inline) or of the path/URL string.
|
|
pub checksum: Option<u64>,
|
|
}
|
|
|
|
impl MediaRef {
|
|
/// Construct a `Path` reference, computing a checksum of the path string.
|
|
pub fn path(path: impl Into<String>, mime_type: impl Into<String>) -> Self {
|
|
let p = path.into();
|
|
let cs = fnv1a_64(p.as_bytes());
|
|
Self {
|
|
ref_type: MediaRefType::Path(p),
|
|
mime_type: mime_type.into(),
|
|
size_bytes: None,
|
|
checksum: Some(cs),
|
|
}
|
|
}
|
|
|
|
/// Construct a `Url` reference, computing a checksum of the URL string.
|
|
pub fn url(url: impl Into<String>, mime_type: impl Into<String>) -> Self {
|
|
let u = url.into();
|
|
let cs = fnv1a_64(u.as_bytes());
|
|
Self {
|
|
ref_type: MediaRefType::Url(u),
|
|
mime_type: mime_type.into(),
|
|
size_bytes: None,
|
|
checksum: Some(cs),
|
|
}
|
|
}
|
|
|
|
/// Construct an `Inline` reference, computing a checksum of the bytes.
|
|
pub fn inline(data: Vec<u8>, mime_type: impl Into<String>) -> Self {
|
|
let cs = fnv1a_64(&data);
|
|
let sz = data.len() as u64;
|
|
Self {
|
|
ref_type: MediaRefType::Inline(data),
|
|
mime_type: mime_type.into(),
|
|
size_bytes: Some(sz),
|
|
checksum: Some(cs),
|
|
}
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// FNV-1a helper (no external deps)
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// 64-bit FNV-1a hash.
|
|
fn fnv1a_64(data: &[u8]) -> u64 {
|
|
const OFFSET: u64 = 14_695_981_039_346_656_037;
|
|
const PRIME: u64 = 1_099_511_628_211;
|
|
let mut h = OFFSET;
|
|
for &b in data {
|
|
h ^= b as u64;
|
|
h = h.wrapping_mul(PRIME);
|
|
}
|
|
h
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Observation
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// What an agent perceived versus what it concluded from a modality.
|
|
#[derive(Debug, Clone)]
|
|
pub struct Observation {
|
|
/// Literal description of what was perceived (e.g. "I see a red stop-sign").
|
|
pub raw_perception: String,
|
|
/// Higher-level interpretation (e.g. "the vehicle must stop").
|
|
pub interpretation: String,
|
|
/// Confidence in the interpretation, clamped to [0.0, 1.0].
|
|
pub confidence: f32,
|
|
pub modality: Modality,
|
|
}
|
|
|
|
impl Observation {
|
|
pub fn new(
|
|
raw_perception: impl Into<String>,
|
|
interpretation: impl Into<String>,
|
|
confidence: f32,
|
|
modality: Modality,
|
|
) -> Self {
|
|
Self {
|
|
raw_perception: raw_perception.into(),
|
|
interpretation: interpretation.into(),
|
|
confidence: confidence.clamp(0.0, 1.0),
|
|
modality,
|
|
}
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// MultiModalRecord
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// A single memory record that may span multiple modalities.
|
|
///
|
|
/// A record can hold embeddings from several models/modalities simultaneously,
|
|
/// enabling cross-modal nearest-neighbour queries.
|
|
#[derive(Debug, Clone)]
|
|
pub struct MultiModalRecord {
|
|
pub id: u64,
|
|
pub primary_modality: Modality,
|
|
/// Textual content (caption, transcript, document text, …).
|
|
pub text_content: Option<String>,
|
|
/// Reference to the raw media artifact.
|
|
pub media_ref: Option<MediaRef>,
|
|
/// One or more embeddings, potentially from different models/modalities.
|
|
pub embeddings: Vec<ModalEmbedding>,
|
|
/// Optional agent observation attached to this record.
|
|
pub observation: Option<Observation>,
|
|
/// Unix timestamp (seconds, float for sub-second precision).
|
|
pub timestamp: f64,
|
|
pub metadata: HashMap<String, String>,
|
|
}
|
|
|
|
impl MultiModalRecord {
|
|
/// Iterate over embeddings that belong to a particular modality.
|
|
pub fn embeddings_for(&self, modality: &Modality) -> impl Iterator<Item = &ModalEmbedding> {
|
|
self.embeddings
|
|
.iter()
|
|
.filter(move |e| &e.modality == modality)
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// MultiModalStore
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// In-memory store for multi-modal memory records with cosine search.
|
|
pub struct MultiModalStore {
|
|
records: Vec<MultiModalRecord>,
|
|
next_id: u64,
|
|
}
|
|
|
|
impl MultiModalStore {
|
|
pub fn new() -> Self {
|
|
Self {
|
|
records: Vec::new(),
|
|
next_id: 1,
|
|
}
|
|
}
|
|
|
|
// ------------------------------------------------------------------
|
|
// Mutations
|
|
// ------------------------------------------------------------------
|
|
|
|
/// Add a record to the store. The `id` field on the record is ignored
|
|
/// and replaced with the store's auto-incremented counter.
|
|
pub fn add_record(&mut self, mut record: MultiModalRecord) -> u64 {
|
|
let id = self.next_id;
|
|
self.next_id += 1;
|
|
record.id = id;
|
|
self.records.push(record);
|
|
id
|
|
}
|
|
|
|
// ------------------------------------------------------------------
|
|
// Queries
|
|
// ------------------------------------------------------------------
|
|
|
|
/// Retrieve a record by id.
|
|
pub fn get_record(&self, id: u64) -> Option<&MultiModalRecord> {
|
|
self.records.iter().find(|r| r.id == id)
|
|
}
|
|
|
|
/// All records whose primary modality matches.
|
|
pub fn get_by_modality(&self, modality: &Modality) -> Vec<&MultiModalRecord> {
|
|
self.records
|
|
.iter()
|
|
.filter(|r| &r.primary_modality == modality)
|
|
.collect()
|
|
}
|
|
|
|
/// All `Observation`s attached to records of a specific modality.
|
|
pub fn get_observations(&self, modality: &Modality) -> Vec<&Observation> {
|
|
self.records
|
|
.iter()
|
|
.filter_map(|r| r.observation.as_ref().filter(|o| &o.modality == modality))
|
|
.collect()
|
|
}
|
|
|
|
/// Total number of records.
|
|
pub fn count(&self) -> usize {
|
|
self.records.len()
|
|
}
|
|
|
|
/// Number of records whose primary modality matches.
|
|
pub fn count_by_modality(&self, modality: &Modality) -> usize {
|
|
self.records
|
|
.iter()
|
|
.filter(|r| &r.primary_modality == modality)
|
|
.count()
|
|
}
|
|
|
|
// ------------------------------------------------------------------
|
|
// Vector search
|
|
// ------------------------------------------------------------------
|
|
|
|
/// Cosine nearest-neighbour search restricted to embeddings of `modality`.
|
|
///
|
|
/// For each record, the highest cosine similarity across all embeddings
|
|
/// that match `modality` is used as the record's score.
|
|
///
|
|
/// Returns up to `k` `(record_id, similarity)` pairs sorted descending.
|
|
pub fn search_by_modality(
|
|
&self,
|
|
modality: &Modality,
|
|
query: &[f32],
|
|
k: usize,
|
|
) -> Vec<(u64, f32)> {
|
|
let q_norm = l2_norm(query);
|
|
let mut scored: Vec<(u64, f32)> = self
|
|
.records
|
|
.iter()
|
|
.filter_map(|r| {
|
|
let best = r
|
|
.embeddings_for(modality)
|
|
.map(|e| cosine_sim_prenorm(query, q_norm, &e.embedding))
|
|
.fold(f32::NEG_INFINITY, f32::max);
|
|
if best.is_finite() {
|
|
Some((r.id, best))
|
|
} else {
|
|
None
|
|
}
|
|
})
|
|
.collect();
|
|
|
|
scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
|
scored.truncate(k);
|
|
scored
|
|
}
|
|
|
|
/// Cross-modal cosine search — considers all embeddings regardless of modality.
|
|
///
|
|
/// Returns up to `k` `(record_id, similarity)` pairs sorted descending.
|
|
pub fn search_cross_modal(&self, query: &[f32], k: usize) -> Vec<(u64, f32)> {
|
|
let q_norm = l2_norm(query);
|
|
let mut scored: Vec<(u64, f32)> = self
|
|
.records
|
|
.iter()
|
|
.filter_map(|r| {
|
|
let best = r
|
|
.embeddings
|
|
.iter()
|
|
.map(|e| cosine_sim_prenorm(query, q_norm, &e.embedding))
|
|
.fold(f32::NEG_INFINITY, f32::max);
|
|
if best.is_finite() {
|
|
Some((r.id, best))
|
|
} else {
|
|
None
|
|
}
|
|
})
|
|
.collect();
|
|
|
|
scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
|
scored.truncate(k);
|
|
scored
|
|
}
|
|
}
|
|
|
|
impl Default for MultiModalStore {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Internal helpers
|
|
// ---------------------------------------------------------------------------
|
|
|
|
#[inline]
|
|
fn l2_norm(v: &[f32]) -> f32 {
|
|
v.iter().map(|x| x * x).sum::<f32>().sqrt()
|
|
}
|
|
|
|
#[inline]
|
|
fn cosine_sim_prenorm(query: &[f32], q_norm: f32, candidate: &[f32]) -> f32 {
|
|
let c_norm = l2_norm(candidate);
|
|
let denom = q_norm * c_norm;
|
|
if denom == 0.0 {
|
|
return 0.0;
|
|
}
|
|
let dot: f32 = query.iter().zip(candidate.iter()).map(|(a, b)| a * b).sum();
|
|
dot / denom
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Tests
|
|
// ---------------------------------------------------------------------------
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
// -----------------------------------------------------------------------
|
|
// Modality
|
|
// -----------------------------------------------------------------------
|
|
|
|
#[test]
|
|
fn modality_display() {
|
|
assert_eq!(Modality::Text.to_string(), "text");
|
|
assert_eq!(Modality::Image.to_string(), "image");
|
|
assert_eq!(Modality::Audio.to_string(), "audio");
|
|
assert_eq!(Modality::Video.to_string(), "video");
|
|
assert_eq!(Modality::Structured.to_string(), "structured");
|
|
}
|
|
|
|
#[test]
|
|
fn modality_equality() {
|
|
assert_eq!(Modality::Text, Modality::Text);
|
|
assert_ne!(Modality::Text, Modality::Image);
|
|
}
|
|
|
|
#[test]
|
|
fn modality_clone_debug() {
|
|
let m = Modality::Audio;
|
|
let c = m.clone();
|
|
assert_eq!(m, c);
|
|
let _ = format!("{:?}", Modality::Video);
|
|
}
|
|
|
|
// -----------------------------------------------------------------------
|
|
// ModalEmbedding
|
|
// -----------------------------------------------------------------------
|
|
|
|
#[test]
|
|
fn modal_embedding_dimension() {
|
|
let e = ModalEmbedding::new(Modality::Text, vec![1.0, 0.0, 0.0], "text-embed-small");
|
|
assert_eq!(e.dimension, 3);
|
|
assert_eq!(e.model_id, "text-embed-small");
|
|
}
|
|
|
|
#[test]
|
|
fn modal_embedding_norm() {
|
|
let e = ModalEmbedding::new(Modality::Image, vec![3.0, 4.0], "clip-vit-large");
|
|
assert!((e.norm() - 5.0).abs() < 1e-6);
|
|
}
|
|
|
|
#[test]
|
|
fn modal_embedding_zero_norm() {
|
|
let e = ModalEmbedding::new(Modality::Text, vec![0.0, 0.0], "m");
|
|
assert_eq!(e.norm(), 0.0);
|
|
}
|
|
|
|
// -----------------------------------------------------------------------
|
|
// MediaRef
|
|
// -----------------------------------------------------------------------
|
|
|
|
#[test]
|
|
fn media_ref_path_checksum_is_set() {
|
|
let r = MediaRef::path("/tmp/photo.jpg", "image/jpeg");
|
|
assert!(r.checksum.is_some());
|
|
assert_eq!(r.mime_type, "image/jpeg");
|
|
assert!(r.size_bytes.is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn media_ref_url_checksum_is_set() {
|
|
let r = MediaRef::url("https://example.com/audio.mp3", "audio/mpeg");
|
|
assert!(r.checksum.is_some());
|
|
}
|
|
|
|
#[test]
|
|
fn media_ref_inline_size_and_checksum() {
|
|
let bytes = vec![1u8, 2, 3, 4, 5];
|
|
let r = MediaRef::inline(bytes, "application/octet-stream");
|
|
assert_eq!(r.size_bytes, Some(5));
|
|
assert!(r.checksum.is_some());
|
|
}
|
|
|
|
#[test]
|
|
fn media_ref_inline_checksum_deterministic() {
|
|
let b1 = vec![42u8; 16];
|
|
let b2 = vec![42u8; 16];
|
|
let r1 = MediaRef::inline(b1, "application/octet-stream");
|
|
let r2 = MediaRef::inline(b2, "application/octet-stream");
|
|
assert_eq!(r1.checksum, r2.checksum);
|
|
}
|
|
|
|
#[test]
|
|
fn media_ref_inline_checksum_differs() {
|
|
let r1 = MediaRef::inline(vec![1u8], "application/octet-stream");
|
|
let r2 = MediaRef::inline(vec![2u8], "application/octet-stream");
|
|
assert_ne!(r1.checksum, r2.checksum);
|
|
}
|
|
|
|
// -----------------------------------------------------------------------
|
|
// FNV-1a
|
|
// -----------------------------------------------------------------------
|
|
|
|
#[test]
|
|
fn fnv1a_known_value() {
|
|
// FNV-1a 64-bit hash of empty string is the offset basis
|
|
assert_eq!(fnv1a_64(b""), 14_695_981_039_346_656_037);
|
|
}
|
|
|
|
#[test]
|
|
fn fnv1a_deterministic() {
|
|
assert_eq!(fnv1a_64(b"hello"), fnv1a_64(b"hello"));
|
|
}
|
|
|
|
#[test]
|
|
fn fnv1a_different_inputs() {
|
|
assert_ne!(fnv1a_64(b"foo"), fnv1a_64(b"bar"));
|
|
}
|
|
|
|
// -----------------------------------------------------------------------
|
|
// Observation
|
|
// -----------------------------------------------------------------------
|
|
|
|
#[test]
|
|
fn observation_confidence_clamped_high() {
|
|
let o = Observation::new("saw fire", "fire detected", 2.5, Modality::Image);
|
|
assert!((o.confidence - 1.0).abs() < 1e-6);
|
|
}
|
|
|
|
#[test]
|
|
fn observation_confidence_clamped_low() {
|
|
let o = Observation::new("heard something", "noise detected", -0.5, Modality::Audio);
|
|
assert!((o.confidence - 0.0).abs() < 1e-6);
|
|
}
|
|
|
|
#[test]
|
|
fn observation_normal_confidence() {
|
|
let o = Observation::new("text block", "english paragraph", 0.85, Modality::Text);
|
|
assert!((o.confidence - 0.85).abs() < 1e-6);
|
|
}
|
|
|
|
// -----------------------------------------------------------------------
|
|
// MultiModalStore — basic ops
|
|
// -----------------------------------------------------------------------
|
|
|
|
fn make_text_record(text: &str, emb: Vec<f32>) -> MultiModalRecord {
|
|
MultiModalRecord {
|
|
id: 0,
|
|
primary_modality: Modality::Text,
|
|
text_content: Some(text.to_string()),
|
|
media_ref: None,
|
|
embeddings: vec![ModalEmbedding::new(Modality::Text, emb, "text-embed-small")],
|
|
observation: None,
|
|
timestamp: 0.0,
|
|
metadata: HashMap::new(),
|
|
}
|
|
}
|
|
|
|
fn make_image_record(emb: Vec<f32>) -> MultiModalRecord {
|
|
MultiModalRecord {
|
|
id: 0,
|
|
primary_modality: Modality::Image,
|
|
text_content: None,
|
|
media_ref: Some(MediaRef::path("/img/cat.jpg", "image/jpeg")),
|
|
embeddings: vec![ModalEmbedding::new(Modality::Image, emb, "clip-vit-large")],
|
|
observation: Some(Observation::new(
|
|
"cat on mat",
|
|
"domestic cat",
|
|
0.9,
|
|
Modality::Image,
|
|
)),
|
|
timestamp: 1.0,
|
|
metadata: HashMap::new(),
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn add_record_assigns_ids() {
|
|
let mut store = MultiModalStore::new();
|
|
let id1 = store.add_record(make_text_record("hello", vec![1.0, 0.0]));
|
|
let id2 = store.add_record(make_text_record("world", vec![0.0, 1.0]));
|
|
assert_eq!(id1, 1);
|
|
assert_eq!(id2, 2);
|
|
assert_eq!(store.count(), 2);
|
|
}
|
|
|
|
#[test]
|
|
fn get_record_found_and_not_found() {
|
|
let mut store = MultiModalStore::new();
|
|
let id = store.add_record(make_text_record("hello", vec![1.0, 0.0]));
|
|
assert!(store.get_record(id).is_some());
|
|
assert!(store.get_record(id + 99).is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn get_record_id_is_correct() {
|
|
let mut store = MultiModalStore::new();
|
|
let id = store.add_record(make_text_record("hi", vec![1.0]));
|
|
assert_eq!(store.get_record(id).unwrap().id, id);
|
|
}
|
|
|
|
#[test]
|
|
fn get_by_modality_filters_correctly() {
|
|
let mut store = MultiModalStore::new();
|
|
store.add_record(make_text_record("a", vec![1.0]));
|
|
store.add_record(make_text_record("b", vec![0.5]));
|
|
store.add_record(make_image_record(vec![1.0, 0.0]));
|
|
|
|
let texts = store.get_by_modality(&Modality::Text);
|
|
let images = store.get_by_modality(&Modality::Image);
|
|
let audios = store.get_by_modality(&Modality::Audio);
|
|
|
|
assert_eq!(texts.len(), 2);
|
|
assert_eq!(images.len(), 1);
|
|
assert_eq!(audios.len(), 0);
|
|
}
|
|
|
|
#[test]
|
|
fn count_by_modality() {
|
|
let mut store = MultiModalStore::new();
|
|
store.add_record(make_text_record("a", vec![1.0]));
|
|
store.add_record(make_image_record(vec![0.0, 1.0]));
|
|
assert_eq!(store.count_by_modality(&Modality::Text), 1);
|
|
assert_eq!(store.count_by_modality(&Modality::Image), 1);
|
|
assert_eq!(store.count_by_modality(&Modality::Audio), 0);
|
|
}
|
|
|
|
#[test]
|
|
fn get_observations_filters_by_modality() {
|
|
let mut store = MultiModalStore::new();
|
|
store.add_record(make_image_record(vec![1.0, 0.0])); // has observation
|
|
store.add_record(make_text_record("no obs", vec![0.0, 1.0])); // no observation
|
|
let obs = store.get_observations(&Modality::Image);
|
|
assert_eq!(obs.len(), 1);
|
|
assert_eq!(obs[0].interpretation, "domestic cat");
|
|
|
|
let text_obs = store.get_observations(&Modality::Text);
|
|
assert_eq!(text_obs.len(), 0);
|
|
}
|
|
|
|
// -----------------------------------------------------------------------
|
|
// Vector search
|
|
// -----------------------------------------------------------------------
|
|
|
|
#[test]
|
|
fn search_by_modality_returns_top_k() {
|
|
let mut store = MultiModalStore::new();
|
|
// query = [1, 0]; record A is [1,0] (perfect match), B is [0,1] (orthogonal)
|
|
store.add_record(make_text_record("A", vec![1.0, 0.0]));
|
|
store.add_record(make_text_record("B", vec![0.0, 1.0]));
|
|
|
|
let results = store.search_by_modality(&Modality::Text, &[1.0, 0.0], 2);
|
|
assert_eq!(results.len(), 2);
|
|
// Best match should be record A (sim ≈ 1.0)
|
|
assert_eq!(results[0].0, 1);
|
|
assert!((results[0].1 - 1.0).abs() < 1e-5);
|
|
// Second should be B (sim ≈ 0.0)
|
|
assert_eq!(results[1].0, 2);
|
|
}
|
|
|
|
#[test]
|
|
fn search_by_modality_k_limits_results() {
|
|
let mut store = MultiModalStore::new();
|
|
for i in 0..5 {
|
|
store.add_record(make_text_record("x", vec![i as f32, 1.0]));
|
|
}
|
|
let results = store.search_by_modality(&Modality::Text, &[1.0, 1.0], 3);
|
|
assert_eq!(results.len(), 3);
|
|
}
|
|
|
|
#[test]
|
|
fn search_by_modality_ignores_other_modalities() {
|
|
let mut store = MultiModalStore::new();
|
|
// Image record with identical embedding to query — should NOT appear
|
|
store.add_record(make_image_record(vec![1.0, 0.0]));
|
|
// Text record
|
|
store.add_record(make_text_record("txt", vec![0.5, 0.5]));
|
|
|
|
let results = store.search_by_modality(&Modality::Text, &[1.0, 0.0], 5);
|
|
// Only the text record should be returned
|
|
assert_eq!(results.len(), 1);
|
|
}
|
|
|
|
#[test]
|
|
fn search_cross_modal_sees_all_modalities() {
|
|
let mut store = MultiModalStore::new();
|
|
store.add_record(make_image_record(vec![1.0, 0.0]));
|
|
store.add_record(make_text_record("txt", vec![0.0, 1.0]));
|
|
|
|
let results = store.search_cross_modal(&[1.0, 0.0], 5);
|
|
assert_eq!(results.len(), 2);
|
|
// Image should score higher (sim ≈ 1.0)
|
|
assert_eq!(results[0].0, 1);
|
|
}
|
|
|
|
#[test]
|
|
fn search_cross_modal_k_limits() {
|
|
let mut store = MultiModalStore::new();
|
|
for i in 0..10 {
|
|
store.add_record(make_text_record("x", vec![i as f32, 0.0]));
|
|
}
|
|
let results = store.search_cross_modal(&[1.0, 0.0], 4);
|
|
assert_eq!(results.len(), 4);
|
|
}
|
|
|
|
#[test]
|
|
fn search_empty_store_returns_empty() {
|
|
let store = MultiModalStore::new();
|
|
assert!(
|
|
store
|
|
.search_by_modality(&Modality::Text, &[1.0, 0.0], 5)
|
|
.is_empty()
|
|
);
|
|
assert!(store.search_cross_modal(&[1.0, 0.0], 5).is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn search_zero_query_returns_zero_similarity() {
|
|
let mut store = MultiModalStore::new();
|
|
store.add_record(make_text_record("a", vec![1.0, 0.0]));
|
|
let results = store.search_by_modality(&Modality::Text, &[0.0, 0.0], 5);
|
|
assert_eq!(results.len(), 1);
|
|
assert_eq!(results[0].1, 0.0);
|
|
}
|
|
|
|
#[test]
|
|
fn search_sorted_descending() {
|
|
let mut store = MultiModalStore::new();
|
|
store.add_record(make_text_record("low", vec![0.0, 1.0])); // sim ≈ 0
|
|
store.add_record(make_text_record("high", vec![1.0, 0.0])); // sim ≈ 1
|
|
store.add_record(make_text_record("mid", vec![1.0, 1.0])); // sim ≈ 0.707
|
|
|
|
let results = store.search_by_modality(&Modality::Text, &[1.0, 0.0], 3);
|
|
assert_eq!(results.len(), 3);
|
|
assert!(results[0].1 >= results[1].1);
|
|
assert!(results[1].1 >= results[2].1);
|
|
}
|
|
|
|
// -----------------------------------------------------------------------
|
|
// MultiModalRecord with multiple embeddings
|
|
// -----------------------------------------------------------------------
|
|
|
|
#[test]
|
|
fn record_with_multiple_embeddings() {
|
|
let mut store = MultiModalStore::new();
|
|
let record = MultiModalRecord {
|
|
id: 0,
|
|
primary_modality: Modality::Image,
|
|
text_content: Some("a cat sitting on a mat".to_string()),
|
|
media_ref: None,
|
|
embeddings: vec![
|
|
ModalEmbedding::new(Modality::Image, vec![1.0, 0.0, 0.0], "clip-vit-large"),
|
|
ModalEmbedding::new(
|
|
Modality::Text,
|
|
vec![0.0, 1.0, 0.0],
|
|
"text-embedding-3-small",
|
|
),
|
|
],
|
|
observation: None,
|
|
timestamp: 42.0,
|
|
metadata: HashMap::new(),
|
|
};
|
|
let id = store.add_record(record);
|
|
|
|
// Cross-modal query aligned with image embedding
|
|
let img_results = store.search_cross_modal(&[1.0, 0.0, 0.0], 5);
|
|
assert_eq!(img_results.len(), 1);
|
|
assert_eq!(img_results[0].0, id);
|
|
assert!((img_results[0].1 - 1.0).abs() < 1e-5);
|
|
|
|
// Per-modality: text embedding search
|
|
let txt_results = store.search_by_modality(&Modality::Text, &[0.0, 1.0, 0.0], 5);
|
|
assert_eq!(txt_results.len(), 1);
|
|
assert!((txt_results[0].1 - 1.0).abs() < 1e-5);
|
|
}
|
|
|
|
#[test]
|
|
fn embeddings_for_iterator() {
|
|
let record = MultiModalRecord {
|
|
id: 1,
|
|
primary_modality: Modality::Image,
|
|
text_content: None,
|
|
media_ref: None,
|
|
embeddings: vec![
|
|
ModalEmbedding::new(Modality::Image, vec![1.0], "clip"),
|
|
ModalEmbedding::new(Modality::Text, vec![0.5], "text"),
|
|
ModalEmbedding::new(Modality::Image, vec![0.8], "clip-v2"),
|
|
],
|
|
observation: None,
|
|
timestamp: 0.0,
|
|
metadata: HashMap::new(),
|
|
};
|
|
let image_embs: Vec<_> = record.embeddings_for(&Modality::Image).collect();
|
|
assert_eq!(image_embs.len(), 2);
|
|
let text_embs: Vec<_> = record.embeddings_for(&Modality::Text).collect();
|
|
assert_eq!(text_embs.len(), 1);
|
|
}
|
|
|
|
// -----------------------------------------------------------------------
|
|
// Default / metadata
|
|
// -----------------------------------------------------------------------
|
|
|
|
#[test]
|
|
fn store_default() {
|
|
let store = MultiModalStore::default();
|
|
assert_eq!(store.count(), 0);
|
|
}
|
|
|
|
#[test]
|
|
fn record_metadata() {
|
|
let mut store = MultiModalStore::new();
|
|
let mut meta = HashMap::new();
|
|
meta.insert("source".to_string(), "camera-1".to_string());
|
|
let record = MultiModalRecord {
|
|
id: 0,
|
|
primary_modality: Modality::Image,
|
|
text_content: None,
|
|
media_ref: None,
|
|
embeddings: vec![],
|
|
observation: None,
|
|
timestamp: 0.0,
|
|
metadata: meta,
|
|
};
|
|
let id = store.add_record(record);
|
|
let r = store.get_record(id).unwrap();
|
|
assert_eq!(r.metadata.get("source").unwrap(), "camera-1");
|
|
}
|
|
}
|