semtree-rag 0.5.0

//! Deterministic stand-ins for the real backends.
//!
//! The point is to exercise the index lifecycle - reopening, incremental
//! updates, rebuilds - without downloading an ONNX model, so these tests run
//! in CI in milliseconds. Both fakes persist to disk, because "does an index
//! survive being closed and reopened" is exactly what is under test.

use std::collections::HashMap;
use std::path::Path;
use std::sync::RwLock;

use async_trait::async_trait;
use semtree_embed::{EmbedError, Embedder, Embedding};
use semtree_store::{Hit, Metric, StoreError, VectorStore};

/// Hashes words into a small vector space, so texts sharing vocabulary end up
/// near each other under cosine similarity. Crude, but enough for a test to
/// assert that a query about "authentication" ranks the auth code first.
pub struct BagOfWords {
    dimension: usize,
    model_id: String,
}

impl BagOfWords {
    pub fn new(model_id: &str) -> Self {
        Self {
            dimension: 32,
            model_id: model_id.to_string(),
        }
    }

    fn vector(&self, text: &str) -> Embedding {
        let mut v = vec![0.0f32; self.dimension];
        for word in text
            .split(|c: char| !c.is_alphanumeric())
            .filter(|w| !w.is_empty())
        {
            let bucket = word
                .to_lowercase()
                .bytes()
                .fold(0u64, |acc, b| acc.wrapping_mul(31).wrapping_add(b as u64));
            v[bucket as usize % self.dimension] += 1.0;
        }

        let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt();
        if norm > 0.0 {
            for x in &mut v {
                *x /= norm;
            }
        }
        v
    }
}

#[async_trait]
impl Embedder for BagOfWords {
    async fn embed(&self, texts: &[&str]) -> Result<Vec<Embedding>, EmbedError> {
        Ok(texts.iter().map(|t| self.vector(t)).collect())
    }

    fn dimension(&self) -> usize {
        self.dimension
    }

    fn model_id(&self) -> &str {
        &self.model_id
    }
}

/// An exhaustive-search vector store that persists as JSON.
#[derive(Default)]
pub struct MemStore {
    vectors: RwLock<HashMap<String, Embedding>>,
}

impl MemStore {
    fn path(dir: &Path) -> std::path::PathBuf {
        dir.join("memstore.json")
    }
}

#[async_trait]
impl VectorStore for MemStore {
    async fn insert(&self, id: &str, embedding: &Embedding) -> Result<(), StoreError> {
        self.vectors
            .write()
            .unwrap()
            .insert(id.to_string(), embedding.clone());
        Ok(())
    }

    async fn search(&self, query: &Embedding, top_k: usize) -> Result<Vec<Hit>, StoreError> {
        let vectors = self.vectors.read().unwrap();
        let mut hits: Vec<Hit> = vectors
            .iter()
            .map(|(id, v)| Hit {
                id: id.clone(),
                score: v.iter().zip(query).map(|(a, b)| a * b).sum(),
            })
            .collect();
        hits.sort_by(|a, b| b.score.total_cmp(&a.score));
        hits.truncate(top_k);
        Ok(hits)
    }

    async fn delete(&self, id: &str) -> Result<(), StoreError> {
        self.vectors.write().unwrap().remove(id);
        Ok(())
    }

    async fn clear(&self) -> Result<(), StoreError> {
        self.vectors.write().unwrap().clear();
        Ok(())
    }

    fn save(&self, path: &Path) -> Result<(), StoreError> {
        let json = serde_json::to_string(&*self.vectors.read().unwrap())
            .map_err(|e| StoreError::Init(e.to_string()))?;
        std::fs::write(Self::path(path), json).map_err(|e| StoreError::Init(e.to_string()))
    }

    fn load(&self, path: &Path) -> Result<(), StoreError> {
        let raw = std::fs::read_to_string(Self::path(path))
            .map_err(|e| StoreError::Init(e.to_string()))?;
        *self.vectors.write().unwrap() =
            serde_json::from_str(&raw).map_err(|e| StoreError::Init(e.to_string()))?;
        Ok(())
    }

    fn len(&self) -> usize {
        self.vectors.read().unwrap().len()
    }

    fn metric(&self) -> Metric {
        Metric::Cosine
    }
}