nitrite_vector 0.7.0

//! A thin retrieval-augmented-generation (RAG) store over a Nitrite collection.
//!
//! Each record is a document holding a `text` field, an `embedding` field (the
//! vector), and any caller-supplied metadata. The collection carries a vector
//! index on `embedding`, so [`RagStore::search`] performs kNN and can combine
//! the result with arbitrary Nitrite metadata filters.
//!
//! Embeddings are supplied by the caller; this store does not generate them.
//!
//! The database must be built with a
//! [`VectorModule`](crate::VectorModule) whose
//! [`VectorIndexConfig`](crate::VectorIndexConfig) sets the dimension and the
//! same metric passed here.

use nitrite::collection::{Document, NitriteCollection, NitriteId};
use nitrite::errors::{ErrorKind, NitriteError, NitriteResult};
use nitrite::filter::{by_id, Filter};
use nitrite::nitrite::Nitrite;

use crate::distance::Metric;
use crate::filter::{value_to_vector, vector_to_value};
use crate::fluent::vector_field;
use crate::vector_index_options;

/// Document field holding the raw text of a record.
pub const TEXT_FIELD: &str = "text";
/// Document field holding the embedding vector.
pub const EMBEDDING_FIELD: &str = "embedding";

/// How many extra candidates to over-fetch per requested result so that a
/// metadata filter and score cutoff still leave enough hits.
const DEFAULT_OVERSAMPLE: usize = 4;

/// A single search result.
#[derive(Debug, Clone)]
pub struct SearchHit {
    /// The document id.
    pub id: NitriteId,
    /// The stored text.
    pub text: String,
    /// Similarity score (higher is more similar; see [`Metric::score`]).
    pub score: f32,
    /// The full stored document (text, embedding, and metadata).
    pub document: Document,
}

/// A vector store tuned for RAG usage.
#[derive(Clone)]
pub struct RagStore {
    collection: NitriteCollection,
    metric: Metric,
}

impl RagStore {
    /// Opens (or creates) a RAG store backed by the named collection, ensuring
    /// a vector index exists on the `embedding` field.
    ///
    /// `metric` **must** match the one configured on the loaded
    /// [`VectorModule`](crate::VectorModule) for this collection's `embedding`
    /// index; it is used to convert distances into similarity scores, and a
    /// mismatch silently produces wrong scores/`min_score` cutoffs (the index
    /// itself still ranks by its own configured metric).
    pub fn create(db: &Nitrite, name: &str, metric: Metric) -> NitriteResult<Self> {
        let collection = db.collection(name)?;
        if !collection.has_index(vec![EMBEDDING_FIELD])? {
            collection.create_index(vec![EMBEDDING_FIELD], &vector_index_options())?;
        }
        Ok(RagStore { collection, metric })
    }

    /// Adds a record and returns its id.
    pub fn add(
        &self,
        text: impl Into<String>,
        embedding: Vec<f32>,
        metadata: Document,
    ) -> NitriteResult<NitriteId> {
        let mut doc = metadata;
        doc.put(TEXT_FIELD, text.into())?;
        doc.put(EMBEDDING_FIELD, vector_to_value(&embedding))?;
        let result = self.collection.insert(doc)?;
        result
            .affected_nitrite_ids()
            .first()
            .cloned()
            .ok_or_else(|| NitriteError::new("Insert returned no id", ErrorKind::InvalidOperation))
    }

    /// Adds many records, returning their ids in order.
    pub fn add_many(
        &self,
        records: Vec<(String, Vec<f32>, Document)>,
    ) -> NitriteResult<Vec<NitriteId>> {
        records
            .into_iter()
            .map(|(text, emb, meta)| self.add(text, emb, meta))
            .collect()
    }

    /// Fetches a record by id.
    pub fn get(&self, id: &NitriteId) -> NitriteResult<Option<Document>> {
        self.collection.get_by_id(id)
    }

    /// Deletes a record by id, returning whether anything was removed.
    pub fn delete(&self, id: &NitriteId) -> NitriteResult<bool> {
        let result = self.collection.remove(by_id(*id), true)?;
        Ok(!result.affected_nitrite_ids().is_empty())
    }

    /// Number of records in the store.
    pub fn len(&self) -> NitriteResult<u64> {
        self.collection.size()
    }

    /// Whether the store has no records.
    pub fn is_empty(&self) -> NitriteResult<bool> {
        Ok(self.len()? == 0)
    }

    /// The underlying collection, for advanced use.
    pub fn collection(&self) -> &NitriteCollection {
        &self.collection
    }

    /// Begins a nearest-neighbor search for the `k` best matches to `query`.
    pub fn search(&self, query: Vec<f32>, k: usize) -> SearchQuery<'_> {
        SearchQuery {
            store: self,
            query,
            k,
            ef: None,
            min_score: None,
            meta_filter: None,
            oversample: DEFAULT_OVERSAMPLE,
        }
    }
}

/// Builder for a RAG search. Finalize with [`SearchQuery::run`].
pub struct SearchQuery<'a> {
    store: &'a RagStore,
    query: Vec<f32>,
    k: usize,
    ef: Option<usize>,
    min_score: Option<f32>,
    meta_filter: Option<Filter>,
    oversample: usize,
}

impl<'a> SearchQuery<'a> {
    /// Restricts results to those also matching a metadata filter.
    ///
    /// The filter is evaluated after the ANN traversal, so a *selective* filter
    /// (the usual RAG shape — "search only this document's chunks") can reject
    /// every one of the nearest neighbours. To avoid silently returning nothing,
    /// [`run`](Self::run) widens the traversal and re-queries until `k` matching
    /// hits are found or the index is exhausted; a filter matching very few
    /// documents therefore costs a full index scan.
    pub fn filter(mut self, filter: Filter) -> Self {
        self.meta_filter = Some(filter);
        self
    }

    /// Sets the query-time `ef` (search width; larger = higher recall).
    pub fn ef(mut self, ef: usize) -> Self {
        self.ef = Some(ef);
        self
    }

    /// Drops hits whose similarity score is below `min_score`.
    pub fn min_score(mut self, min_score: f32) -> Self {
        self.min_score = Some(min_score);
        self
    }

    /// Sets how many extra candidates to over-fetch per requested result when a
    /// metadata filter or score cutoff is in play (default 4).
    ///
    /// With a metadata filter this is only the *starting* window; see
    /// [`filter`](Self::filter).
    pub fn oversample(mut self, factor: usize) -> Self {
        self.oversample = factor.max(1);
        self
    }

    /// Executes the search, returning up to `k` hits ordered by descending
    /// score.
    ///
    /// With a metadata filter the traversal widens and repeats until `k`
    /// matching hits are found or the index is exhausted, so a selective filter
    /// never silently returns fewer hits than exist.
    pub fn run(self) -> NitriteResult<Vec<SearchHit>> {
        if self.k == 0 {
            return Ok(Vec::new());
        }
        // Over-fetch so metadata filtering / score cutoff still leaves k hits.
        let mut fetch = if self.meta_filter.is_some() || self.min_score.is_some() {
            self.k.saturating_mul(self.oversample)
        } else {
            self.k
        };

        loop {
            let (mut hits, candidates) = self.probe(fetch)?;
            // Widening only helps a metadata filter: a `min_score` cutoff is
            // monotone in the ANN ranking, so hits it drops are the tail and
            // fetching further out cannot recover any.
            let exhausted = candidates < fetch;
            if self.meta_filter.is_none() || hits.len() >= self.k || exhausted {
                // Rank by score (descending) and trim to k, independent of the
                // order the cursor produced.
                hits.sort_by(|a, b| b.score.total_cmp(&a.score));
                hits.truncate(self.k);
                return Ok(hits);
            }
            // Terminates: `fetch` saturates, at which point the index can never
            // return `fetch` candidates and `exhausted` ends the loop.
            fetch = fetch.saturating_mul(2);
        }
    }

    /// One ANN probe: the `fetch` nearest neighbours, filtered and scored.
    ///
    /// Returns the surviving hits (unordered) and how many candidates the index
    /// produced — fewer than `fetch` means the index has nothing further out.
    fn probe(&self, fetch: usize) -> NitriteResult<(Vec<SearchHit>, usize)> {
        let metric = self.store.metric;
        let mut builder = vector_field(EMBEDDING_FIELD).nearest(self.query.clone(), fetch);
        if let Some(ef) = self.ef {
            builder = builder.ef(ef);
        }
        let mut cursor = self.store.collection.find(builder.build())?;

        // Prepare the query once for scoring (normalizes for cosine).
        let prepared_query = metric.prepare(self.query.clone());

        let mut hits: Vec<SearchHit> = Vec::new();
        let mut candidates = 0usize;
        for entry in cursor.iter_with_id() {
            let (id, document) = entry?;
            candidates += 1;

            if let Some(mf) = &self.meta_filter {
                if !mf.apply(&document)? {
                    continue;
                }
            }

            let embedding = match document.get(EMBEDDING_FIELD).ok().and_then(|v| value_to_vector(&v)) {
                Some(e) if e.len() == prepared_query.len() => e,
                _ => continue,
            };
            let prepared = metric.prepare(embedding);
            let score = metric.score(metric.distance(&prepared_query, &prepared));

            if let Some(min) = self.min_score {
                if score < min {
                    continue;
                }
            }

            let text = match document.get(TEXT_FIELD) {
                Ok(nitrite::common::Value::String(s)) => s,
                _ => String::new(),
            };
            hits.push(SearchHit { id, text, score, document });
        }

        Ok((hits, candidates))
    }
}