rig-memvid 0.1.0

//! [`MemvidStore`]: a [`VectorStoreIndex`] backed by a single `.mv2` file.

use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};

#[cfg(feature = "vec")]
use memvid_core::{LocalTextEmbedder, TextEmbedConfig};
use memvid_core::{Memvid, PutOptions, SearchHit, SearchRequest};
use rig::{
    Embed, OneOrMany,
    embeddings::Embedding,
    vector_store::{
        InsertDocuments, VectorSearchRequest, VectorStoreError, VectorStoreIndex,
        request::SearchFilter,
    },
    wasm_compat::WasmCompatSend,
};
use serde::{Deserialize, Serialize};

use crate::error::MemvidError;

/// A persistent, file-backed vector / lexical index over a memvid `.mv2`
/// archive.
///
/// `MemvidStore` is cheap to clone (it shares an `Arc<Mutex<Memvid>>` with
/// every clone) and can be both read from and written to concurrently from
/// multiple async tasks. Writes are serialised through the inner mutex.
///
/// Unlike most rig vector stores, `MemvidStore` is **not** parameterised over
/// an [`EmbeddingModel`]: memvid embeds queries internally using whichever
/// engine its file is configured with (BM25/Tantivy when the `lex` feature is
/// enabled, HNSW + BGE-small when `vec` is enabled). Pass plain text in
/// [`VectorSearchRequest::query`] and let memvid do the rest.
///
/// [`EmbeddingModel`]: rig::embeddings::EmbeddingModel
#[derive(Clone)]
pub struct MemvidStore {
    inner: Arc<Mutex<Memvid>>,
    #[cfg(feature = "vec")]
    embedder: Option<Arc<LocalTextEmbedder>>,
}

impl std::fmt::Debug for MemvidStore {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("MemvidStore").finish_non_exhaustive()
    }
}

impl MemvidStore {
    /// Wraps an already-open [`Memvid`] handle.
    pub fn from_memvid(memvid: Memvid) -> Self {
        Self {
            inner: Arc::new(Mutex::new(memvid)),
            #[cfg(feature = "vec")]
            embedder: None,
        }
    }

    /// Begin building a new store. See [`MemvidStoreBuilder`].
    pub fn builder() -> MemvidStoreBuilder {
        MemvidStoreBuilder::default()
    }

    /// Acquire the inner mutex. Returns [`MemvidError::Poisoned`] if a prior
    /// holder of the lock panicked.
    fn lock(&self) -> Result<std::sync::MutexGuard<'_, Memvid>, MemvidError> {
        self.inner.lock().map_err(|_| MemvidError::Poisoned)
    }

    /// Whether this store will route writes/queries through a local
    /// embedding model.
    #[cfg(feature = "vec")]
    #[must_use]
    pub fn has_embedder(&self) -> bool {
        self.embedder.is_some()
    }

    /// Encode `text` with the configured embedder, if any.
    #[cfg(feature = "vec")]
    fn encode(&self, text: &str) -> Result<Option<Vec<f32>>, MemvidError> {
        match &self.embedder {
            Some(embedder) => Ok(Some(embedder.encode_text(text)?)),
            None => Ok(None),
        }
    }

    /// Append a UTF-8 text payload to the archive and immediately commit.
    ///
    /// Returns the assigned `frame_id`. When the store has been built with
    /// an embedder (`vec` feature), the text is embedded and stored
    /// alongside its frame so that subsequent
    /// [`VectorStoreIndex::top_n`] calls perform semantic search.
    pub fn put_text(&self, text: &str, options: PutOptions) -> Result<u64, MemvidError> {
        #[cfg(feature = "vec")]
        let embedding = self.encode(text)?;
        let mut guard = self.lock()?;
        #[cfg(feature = "vec")]
        let id = if let Some(emb) = embedding {
            guard.put_with_embedding_and_options(text.as_bytes(), emb, options)?
        } else {
            guard.put_bytes_with_options(text.as_bytes(), options)?
        };
        #[cfg(not(feature = "vec"))]
        let id = guard.put_bytes_with_options(text.as_bytes(), options)?;
        guard.commit()?;
        Ok(id)
    }

    /// Append a payload without committing. The caller is responsible for
    /// invoking [`MemvidStore::commit`] before a subsequent search will see
    /// the new frame.
    pub fn put_text_uncommitted(
        &self,
        text: &str,
        options: PutOptions,
    ) -> Result<u64, MemvidError> {
        #[cfg(feature = "vec")]
        let embedding = self.encode(text)?;
        let mut guard = self.lock()?;
        #[cfg(feature = "vec")]
        let id = if let Some(emb) = embedding {
            guard.put_with_embedding_and_options(text.as_bytes(), emb, options)?
        } else {
            guard.put_bytes_with_options(text.as_bytes(), options)?
        };
        #[cfg(not(feature = "vec"))]
        let id = guard.put_bytes_with_options(text.as_bytes(), options)?;
        Ok(id)
    }

    /// Flush any pending writes to disk.
    pub fn commit(&self) -> Result<(), MemvidError> {
        let mut guard = self.lock()?;
        guard.commit()?;
        Ok(())
    }

    /// Run a [`SearchRequest`] directly. Useful for callers that need
    /// memvid-native features (cursors, ACL contexts, etc.) that do not map
    /// onto [`VectorSearchRequest`].
    pub fn search(
        &self,
        request: SearchRequest,
    ) -> Result<memvid_core::SearchResponse, MemvidError> {
        let mut guard = self.lock()?;
        let resp = guard.search(request)?;
        Ok(resp)
    }
}

/// Builder for [`MemvidStore`].
#[derive(Default)]
pub struct MemvidStoreBuilder {
    path: Option<PathBuf>,
    enable_lex: bool,
    #[cfg(feature = "vec")]
    enable_vec: bool,
    #[cfg(feature = "vec")]
    vec_model: Option<String>,
    #[cfg(feature = "vec")]
    embedder: Option<Arc<LocalTextEmbedder>>,
}

impl std::fmt::Debug for MemvidStoreBuilder {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mut d = f.debug_struct("MemvidStoreBuilder");
        d.field("path", &self.path)
            .field("enable_lex", &self.enable_lex);
        #[cfg(feature = "vec")]
        {
            d.field("enable_vec", &self.enable_vec)
                .field("vec_model", &self.vec_model)
                .field("embedder", &self.embedder.as_ref().map(|_| "<embedder>"));
        }
        d.finish()
    }
}

impl MemvidStoreBuilder {
    /// Path to the `.mv2` file.
    pub fn path<P: Into<PathBuf>>(mut self, path: P) -> Self {
        self.path = Some(path.into());
        self
    }

    /// Enable BM25 / Tantivy lexical search on the underlying archive.
    pub fn enable_lex(mut self) -> Self {
        self.enable_lex = true;
        self
    }

    /// Enable HNSW vector search on the underlying archive.
    ///
    /// Available only when this crate is built with the `vec` feature, which
    /// pulls in `memvid-core/vec` (ONNX Runtime + bundled BGE-small).
    /// Mutually compatible with [`Self::enable_lex`]; both can be on at once
    /// for hybrid retrieval.
    #[cfg(feature = "vec")]
    pub fn enable_vec(mut self) -> Self {
        self.enable_vec = true;
        self
    }

    /// Bind (or validate) the embedding model identifier on the vector
    /// index. See [`memvid_core::Memvid::set_vec_model`].
    #[cfg(feature = "vec")]
    pub fn vec_model(mut self, model: impl Into<String>) -> Self {
        self.vec_model = Some(model.into());
        self
    }

    /// Attach a local text embedder. Writes performed via
    /// [`MemvidStore::put_text`] and queries performed via
    /// [`VectorStoreIndex::top_n`] will be embedded with this model and
    /// routed through memvid's HNSW vector index.
    ///
    /// Implies [`Self::enable_vec`]. If [`Self::vec_model`] has not been
    /// set, the model identifier reported by the embedder is bound
    /// automatically.
    #[cfg(feature = "vec")]
    pub fn embedder(mut self, embedder: LocalTextEmbedder) -> Self {
        if self.vec_model.is_none() {
            self.vec_model = Some(embedder.model_info().name.to_string());
        }
        self.embedder = Some(Arc::new(embedder));
        self.enable_vec = true;
        self
    }

    /// Convenience: attach the default local embedder (BGE-small,
    /// 384-dimensional). The model is loaded from
    /// [`TextEmbedConfig::default`]'s on-disk cache; if absent and
    /// `offline` is `false` it will be downloaded.
    #[cfg(feature = "vec")]
    pub fn with_default_embedder(self) -> Result<Self, MemvidError> {
        let embedder = LocalTextEmbedder::new(TextEmbedConfig::bge_small())?;
        Ok(self.embedder(embedder))
    }

    /// Convenience: attach a local embedder built from an explicit
    /// [`TextEmbedConfig`].
    #[cfg(feature = "vec")]
    pub fn with_embedder_config(self, config: TextEmbedConfig) -> Result<Self, MemvidError> {
        let embedder = LocalTextEmbedder::new(config)?;
        Ok(self.embedder(embedder))
    }

    fn require_path(&self) -> Result<&Path, MemvidError> {
        self.path.as_deref().ok_or_else(|| {
            MemvidError::Io(std::io::Error::new(
                std::io::ErrorKind::InvalidInput,
                "MemvidStoreBuilder requires a path",
            ))
        })
    }

    fn finish(self, memvid: Memvid) -> Result<MemvidStore, MemvidError> {
        let mut memvid = memvid;
        if self.enable_lex {
            memvid.enable_lex()?;
        }
        #[cfg(feature = "vec")]
        {
            if self.enable_vec {
                memvid.enable_vec()?;
            }
            if let Some(model) = self.vec_model.as_deref() {
                memvid.set_vec_model(model)?;
            }
        }
        #[cfg_attr(not(feature = "vec"), allow(unused_mut))]
        let mut store = MemvidStore::from_memvid(memvid);
        #[cfg(feature = "vec")]
        {
            store.embedder = self.embedder;
        }
        Ok(store)
    }

    /// Open an existing `.mv2` file. Errors if the file does not exist.
    pub fn open(self) -> Result<MemvidStore, MemvidError> {
        let path = self.require_path()?.to_path_buf();
        let memvid = Memvid::open(&path)?;
        self.finish(memvid)
    }

    /// Create a new `.mv2` file. Errors if the file already exists.
    pub fn create(self) -> Result<MemvidStore, MemvidError> {
        let path = self.require_path()?.to_path_buf();
        let memvid = Memvid::create(&path)?;
        self.finish(memvid)
    }

    /// Open the file if it exists, otherwise create it.
    pub fn open_or_create(self) -> Result<MemvidStore, MemvidError> {
        let path = self.require_path()?.to_path_buf();
        let memvid = if path.exists() {
            Memvid::open(&path)?
        } else {
            Memvid::create(&path)?
        };
        self.finish(memvid)
    }

    /// Open the file read-only.
    pub fn open_read_only(self) -> Result<MemvidStore, MemvidError> {
        let path = self.require_path()?.to_path_buf();
        let memvid = Memvid::open_read_only(&path)?;
        self.finish(memvid)
    }
}

/// A filter clause supported by [`MemvidStore`].
///
/// Memvid's query model does not support arbitrary boolean predicates;
/// this filter only carries the four restriction parameters that map onto
/// fields of [`SearchRequest`]:
///
/// | Predicate                       | Effect on the search request  |
/// | ------------------------------- | ----------------------------- |
/// | `eq("uri", "...")`              | `request.uri = Some(value)`   |
/// | `eq("scope", "...")`            | `request.scope = Some(value)` |
/// | `eq("as_of_frame", n)`          | `request.as_of_frame`         |
/// | `eq("as_of_ts", n)`             | `request.as_of_ts`            |
///
/// `gt`, `lt`, and `or` are not representable; constructing such a filter
/// produces an error at query time
/// ([`MemvidError::UnsupportedFilter`]).
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct MemvidFilter {
    /// Optional URI prefix restriction.
    pub uri: Option<String>,
    /// Optional logical scope.
    pub scope: Option<String>,
    /// Optional point-in-time frame id.
    pub as_of_frame: Option<u64>,
    /// Optional point-in-time unix-millis timestamp.
    pub as_of_ts: Option<i64>,
    /// Reasons this filter cannot be applied. Populated when the user calls
    /// `gt`, `lt`, `or`, or `eq` with an unknown key.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    invalid: Vec<String>,
}

impl MemvidFilter {
    fn unsupported(reason: impl Into<String>) -> Self {
        Self {
            invalid: vec![reason.into()],
            ..Self::default()
        }
    }

    fn merge(mut self, rhs: Self) -> Self {
        if rhs.uri.is_some() {
            self.uri = rhs.uri;
        }
        if rhs.scope.is_some() {
            self.scope = rhs.scope;
        }
        if rhs.as_of_frame.is_some() {
            self.as_of_frame = rhs.as_of_frame;
        }
        if rhs.as_of_ts.is_some() {
            self.as_of_ts = rhs.as_of_ts;
        }
        self.invalid.extend(rhs.invalid);
        self
    }

    fn into_validated(self) -> Result<Self, MemvidError> {
        if self.invalid.is_empty() {
            Ok(self)
        } else {
            Err(MemvidError::UnsupportedFilter(self.invalid.join("; ")))
        }
    }

    fn apply_to(self, request: &mut SearchRequest) {
        request.uri = self.uri;
        request.scope = self.scope;
        request.as_of_frame = self.as_of_frame;
        request.as_of_ts = self.as_of_ts;
    }
}

fn json_as_string(value: &serde_json::Value) -> Option<String> {
    match value {
        serde_json::Value::String(s) => Some(s.clone()),
        other => Some(other.to_string()),
    }
}

impl SearchFilter for MemvidFilter {
    type Value = serde_json::Value;

    fn eq(key: impl AsRef<str>, value: Self::Value) -> Self {
        let key = key.as_ref();
        match key {
            "uri" => Self {
                uri: json_as_string(&value),
                ..Self::default()
            },
            "scope" => Self {
                scope: json_as_string(&value),
                ..Self::default()
            },
            "as_of_frame" => match value.as_u64() {
                Some(n) => Self {
                    as_of_frame: Some(n),
                    ..Self::default()
                },
                None => Self::unsupported(format!("as_of_frame must be a u64, got {value}")),
            },
            "as_of_ts" => match value.as_i64() {
                Some(n) => Self {
                    as_of_ts: Some(n),
                    ..Self::default()
                },
                None => Self::unsupported(format!("as_of_ts must be an i64, got {value}")),
            },
            other => Self::unsupported(format!(
                "unsupported filter key '{other}' (allowed: uri, scope, as_of_frame, as_of_ts)"
            )),
        }
    }

    fn gt(key: impl AsRef<str>, _value: Self::Value) -> Self {
        Self::unsupported(format!(
            "memvid does not support gt() on '{}'",
            key.as_ref()
        ))
    }

    fn lt(key: impl AsRef<str>, _value: Self::Value) -> Self {
        Self::unsupported(format!(
            "memvid does not support lt() on '{}'",
            key.as_ref()
        ))
    }

    fn and(self, rhs: Self) -> Self {
        self.merge(rhs)
    }

    fn or(self, rhs: Self) -> Self {
        self.merge(rhs)
            .merge(Self::unsupported("memvid does not support or() in filters"))
    }
}

/// Default snippet size when memvid is asked for context around a hit.
///
/// Tuned to be roughly one paragraph; callers who want different behaviour
/// should call [`MemvidStore::search`] directly with their own
/// [`SearchRequest`].
const DEFAULT_SNIPPET_CHARS: usize = 400;

fn build_search_request(
    query: String,
    samples: u64,
    filter: Option<MemvidFilter>,
) -> Result<SearchRequest, MemvidError> {
    let filter = match filter {
        Some(f) => f.into_validated()?,
        None => MemvidFilter::default(),
    };
    let mut req = SearchRequest {
        query,
        top_k: usize::try_from(samples).unwrap_or(usize::MAX),
        snippet_chars: DEFAULT_SNIPPET_CHARS,
        uri: None,
        scope: None,
        cursor: None,
        #[cfg(feature = "temporal")]
        temporal: None,
        as_of_frame: None,
        as_of_ts: None,
        no_sketch: false,
        acl_context: None,
        acl_enforcement_mode: memvid_core::AclEnforcementMode::default(),
    };
    filter.apply_to(&mut req);
    Ok(req)
}

fn hit_score(hit: &SearchHit) -> f64 {
    match hit.score {
        Some(s) => f64::from(s),
        // Lexical hits often arrive without a numeric score; fall back to
        // rank-derived order-preserving values so callers can still sort.
        None => 1.0 / (hit.rank as f64 + 1.0),
    }
}

#[cfg(feature = "vec")]
fn ensure_vec_filter_supported(filter: &MemvidFilter) -> Result<(), MemvidError> {
    if filter.uri.is_some() {
        return Err(MemvidError::UnsupportedFilter(
            "`uri` filter is not supported when querying through the embedder; use lex search"
                .into(),
        ));
    }
    if filter.as_of_frame.is_some() || filter.as_of_ts.is_some() {
        return Err(MemvidError::UnsupportedFilter(
            "point-in-time filters (`as_of_frame`, `as_of_ts`) are not supported under vector \
             search; use lex or `MemvidStore::search` directly"
                .into(),
        ));
    }
    Ok(())
}

impl MemvidStore {
    /// Run an embedding-driven search through memvid's HNSW index.
    /// Pre-validated by the caller; returns the raw memvid response.
    #[cfg(feature = "vec")]
    fn vec_search(
        &self,
        query: &str,
        samples: u64,
        filter: &MemvidFilter,
    ) -> Result<memvid_core::SearchResponse, MemvidError> {
        let embedder = self
            .embedder
            .as_ref()
            .ok_or_else(|| MemvidError::UnsupportedFilter("no embedder configured".into()))?;
        let embedding = embedder.encode_text(query)?;
        let top_k = usize::try_from(samples).unwrap_or(usize::MAX);
        let mut guard = self.lock()?;
        let resp = guard.vec_search_with_embedding(
            query,
            &embedding,
            top_k,
            DEFAULT_SNIPPET_CHARS,
            filter.scope.as_deref(),
        )?;
        Ok(resp)
    }
}

impl VectorStoreIndex for MemvidStore {
    type Filter = MemvidFilter;

    async fn top_n<T>(
        &self,
        req: VectorSearchRequest<Self::Filter>,
    ) -> Result<Vec<(f64, String, T)>, VectorStoreError>
    where
        T: for<'a> Deserialize<'a> + WasmCompatSend,
    {
        let query = req.query().to_owned();
        let samples = req.samples();
        let filter = req.filter().clone();

        #[cfg(feature = "vec")]
        let response = if self.embedder.is_some() {
            let validated = match filter.clone() {
                Some(f) => f.into_validated().map_err(VectorStoreError::from)?,
                None => MemvidFilter::default(),
            };
            ensure_vec_filter_supported(&validated).map_err(VectorStoreError::from)?;
            self.vec_search(&query, samples, &validated)
                .map_err(VectorStoreError::from)?
        } else {
            let request =
                build_search_request(query, samples, filter).map_err(VectorStoreError::from)?;
            let mut guard = self
                .inner
                .lock()
                .map_err(|_| VectorStoreError::from(MemvidError::Poisoned))?;
            guard.search(request).map_err(MemvidError::from)?
        };
        #[cfg(not(feature = "vec"))]
        let response = {
            let request =
                build_search_request(query, samples, filter).map_err(VectorStoreError::from)?;
            let mut guard = self
                .inner
                .lock()
                .map_err(|_| VectorStoreError::from(MemvidError::Poisoned))?;
            guard.search(request).map_err(MemvidError::from)?
        };

        let mut out = Vec::with_capacity(response.hits.len());
        for hit in response.hits {
            let score = hit_score(&hit);
            let id = hit.frame_id.to_string();
            let value = serde_json::to_value(&hit).map_err(MemvidError::from)?;
            let doc: T = serde_json::from_value(value).map_err(MemvidError::from)?;
            out.push((score, id, doc));
        }
        Ok(out)
    }

    async fn top_n_ids(
        &self,
        req: VectorSearchRequest<Self::Filter>,
    ) -> Result<Vec<(f64, String)>, VectorStoreError> {
        let query = req.query().to_owned();
        let samples = req.samples();
        let filter = req.filter().clone();

        #[cfg(feature = "vec")]
        let response = if self.embedder.is_some() {
            let validated = match filter.clone() {
                Some(f) => f.into_validated().map_err(VectorStoreError::from)?,
                None => MemvidFilter::default(),
            };
            ensure_vec_filter_supported(&validated).map_err(VectorStoreError::from)?;
            self.vec_search(&query, samples, &validated)
                .map_err(VectorStoreError::from)?
        } else {
            let request =
                build_search_request(query, samples, filter).map_err(VectorStoreError::from)?;
            let mut guard = self
                .inner
                .lock()
                .map_err(|_| VectorStoreError::from(MemvidError::Poisoned))?;
            guard.search(request).map_err(MemvidError::from)?
        };
        #[cfg(not(feature = "vec"))]
        let response = {
            let request =
                build_search_request(query, samples, filter).map_err(VectorStoreError::from)?;
            let mut guard = self
                .inner
                .lock()
                .map_err(|_| VectorStoreError::from(MemvidError::Poisoned))?;
            guard.search(request).map_err(MemvidError::from)?
        };

        Ok(response
            .hits
            .into_iter()
            .map(|hit| (hit_score(&hit), hit.frame_id.to_string()))
            .collect())
    }
}

impl InsertDocuments for MemvidStore {
    async fn insert_documents<Doc>(
        &self,
        documents: Vec<(Doc, OneOrMany<Embedding>)>,
    ) -> Result<(), VectorStoreError>
    where
        Doc: Serialize + Embed + WasmCompatSend,
    {
        // We deliberately ignore the externally-supplied embeddings (rig
        // computes them with its own model, but memvid validates the
        // dimension against its bound model and would reject mismatches).
        // When this store has its own embedder, embed each document with
        // the local model. Round-tripping the document through JSON gives
        // us a stable byte payload that `serde_json::from_value::<T>` can
        // recover during search.
        #[cfg(feature = "vec")]
        let local_embedder = self.embedder.clone();
        let mut prepared: Vec<(Vec<u8>, Option<Vec<f32>>)> = Vec::with_capacity(documents.len());
        for (doc, _embeddings) in documents {
            let bytes = serde_json::to_vec(&doc).map_err(MemvidError::from)?;
            #[cfg(feature = "vec")]
            let emb = match &local_embedder {
                Some(embedder) => {
                    let text = std::str::from_utf8(&bytes).unwrap_or("");
                    Some(embedder.encode_text(text).map_err(MemvidError::from)?)
                }
                None => None,
            };
            #[cfg(not(feature = "vec"))]
            let emb: Option<Vec<f32>> = None;
            prepared.push((bytes, emb));
        }

        let mut guard = self
            .inner
            .lock()
            .map_err(|_| VectorStoreError::from(MemvidError::Poisoned))?;
        for (bytes, emb) in prepared {
            match emb {
                Some(embedding) => {
                    guard
                        .put_with_embedding_and_options(&bytes, embedding, PutOptions::default())
                        .map_err(MemvidError::from)?;
                }
                None => {
                    guard
                        .put_bytes_with_options(&bytes, PutOptions::default())
                        .map_err(MemvidError::from)?;
                }
            }
        }
        guard.commit().map_err(MemvidError::from)?;
        Ok(())
    }
}