Skip to main content

rig_memvid/
store.rs

1//! [`MemvidStore`]: a [`VectorStoreIndex`] backed by a single `.mv2` file.
2
3use std::path::{Path, PathBuf};
4use std::sync::{Arc, Mutex};
5
6use memvid_core::{AclContext, AclEnforcementMode, Memvid, PutOptions, SearchHit, SearchRequest};
7#[cfg(feature = "vec")]
8use memvid_core::{LocalTextEmbedder, TextEmbedConfig};
9use rig::{
10    Embed, OneOrMany,
11    embeddings::Embedding,
12    vector_store::{
13        InsertDocuments, VectorSearchRequest, VectorStoreError, VectorStoreIndex,
14        request::SearchFilter,
15    },
16    wasm_compat::WasmCompatSend,
17};
18use serde::{Deserialize, Serialize};
19
20use crate::error::MemvidError;
21
22/// A persistent, file-backed vector / lexical index over a memvid `.mv2`
23/// archive.
24///
25/// `MemvidStore` is cheap to clone (it shares an `Arc<Mutex<Memvid>>` with
26/// every clone) and can be both read from and written to concurrently from
27/// multiple async tasks. Writes are serialised through the inner mutex.
28///
29/// ## Concurrency
30///
31/// Every public method on the underlying [`Memvid`] handle — including
32/// `search`, `vec_search_with_embedding`, `frame_count`, and the various
33/// `put_*` writers — takes `&mut self`. Reads cannot run in parallel with
34/// other reads, so the inner lock is a [`Mutex`] rather than an
35/// `RwLock`. Workloads that require concurrent reads should open separate
36/// read-only handles via [`MemvidStoreBuilder::open_read_only`].
37///
38/// The lock is [`std::sync::Mutex`] (not `tokio::sync::Mutex`): the crate
39/// is intentionally runtime-agnostic and the clippy `await_holding_lock`
40/// lint enforces that no `.await` ever happens while a guard is live. Every
41/// guard in this module is scope-dropped before any async boundary.
42///
43/// Unlike most rig vector stores, `MemvidStore` is **not** parameterised over
44/// an [`EmbeddingModel`]: memvid embeds queries internally using whichever
45/// engine its file is configured with (BM25/Tantivy when the `lex` feature is
46/// enabled, HNSW + BGE-small when `vec` is enabled). Pass plain text in
47/// [`VectorSearchRequest::query`] and let memvid do the rest.
48///
49/// [`EmbeddingModel`]: rig::embeddings::EmbeddingModel
50#[derive(Clone)]
51pub struct MemvidStore {
52    inner: Arc<Mutex<Memvid>>,
53    #[cfg(feature = "vec")]
54    embedder: Option<Arc<LocalTextEmbedder>>,
55    /// Default `snippet_chars` applied to [`VectorStoreIndex`] queries.
56    /// Configurable via [`MemvidStoreBuilder::snippet_chars`].
57    snippet_chars: usize,
58    /// Default ACL context applied to every search. `None` means no ACL
59    /// filtering. Configurable via [`MemvidStoreBuilder::acl_context`].
60    acl_context: Option<AclContext>,
61    /// ACL enforcement mode (`Audit` or `Enforce`).
62    acl_enforcement_mode: AclEnforcementMode,
63}
64
65impl std::fmt::Debug for MemvidStore {
66    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
67        f.debug_struct("MemvidStore").finish_non_exhaustive()
68    }
69}
70
71impl MemvidStore {
72    /// Wraps an already-open [`Memvid`] handle.
73    pub fn from_memvid(memvid: Memvid) -> Self {
74        Self {
75            inner: Arc::new(Mutex::new(memvid)),
76            #[cfg(feature = "vec")]
77            embedder: None,
78            snippet_chars: DEFAULT_SNIPPET_CHARS,
79            acl_context: None,
80            acl_enforcement_mode: AclEnforcementMode::default(),
81        }
82    }
83
84    /// Number of frames currently stored in the underlying `.mv2` file.
85    pub fn frame_count(&self) -> Result<usize, MemvidError> {
86        Ok(self.lock()?.frame_count())
87    }
88
89    /// Aggregate statistics for the underlying memory.
90    pub fn stats(&self) -> Result<memvid_core::types::frame::Stats, MemvidError> {
91        Ok(self.lock()?.stats()?)
92    }
93
94    /// Begin building a new store. See [`MemvidStoreBuilder`].
95    pub fn builder() -> MemvidStoreBuilder {
96        MemvidStoreBuilder::default()
97    }
98
99    /// Acquire the inner mutex. Returns [`MemvidError::Poisoned`] if a prior
100    /// holder of the lock panicked.
101    fn lock(&self) -> Result<std::sync::MutexGuard<'_, Memvid>, MemvidError> {
102        self.inner.lock().map_err(|_| MemvidError::Poisoned)
103    }
104
105    /// Whether this store will route writes/queries through a local
106    /// embedding model.
107    #[cfg(feature = "vec")]
108    #[must_use]
109    pub fn has_embedder(&self) -> bool {
110        self.embedder.is_some()
111    }
112
113    /// Encode `text` with the configured embedder, if any.
114    #[cfg(feature = "vec")]
115    fn encode(&self, text: &str) -> Result<Option<Vec<f32>>, MemvidError> {
116        match &self.embedder {
117            Some(embedder) => Ok(Some(embedder.encode_text(text)?)),
118            None => Ok(None),
119        }
120    }
121
122    /// Append a UTF-8 text payload to the archive and immediately commit.
123    ///
124    /// Returns the assigned `frame_id`. When the store has been built with
125    /// an embedder (`vec` feature), the text is embedded and stored
126    /// alongside its frame so that subsequent
127    /// [`VectorStoreIndex::top_n`] calls perform semantic search.
128    pub fn put_text(&self, text: &str, options: PutOptions) -> Result<u64, MemvidError> {
129        #[cfg(feature = "vec")]
130        let embedding = self.encode(text)?;
131        let mut guard = self.lock()?;
132        #[cfg(feature = "vec")]
133        let id = if let Some(emb) = embedding {
134            guard.put_with_embedding_and_options(text.as_bytes(), emb, options)?
135        } else {
136            guard.put_bytes_with_options(text.as_bytes(), options)?
137        };
138        #[cfg(not(feature = "vec"))]
139        let id = guard.put_bytes_with_options(text.as_bytes(), options)?;
140        guard.commit()?;
141        Ok(id)
142    }
143
144    /// Append a payload without committing. The caller is responsible for
145    /// invoking [`MemvidStore::commit`] before a subsequent search will see
146    /// the new frame.
147    pub fn put_text_uncommitted(
148        &self,
149        text: &str,
150        options: PutOptions,
151    ) -> Result<u64, MemvidError> {
152        #[cfg(feature = "vec")]
153        let embedding = self.encode(text)?;
154        let mut guard = self.lock()?;
155        #[cfg(feature = "vec")]
156        let id = if let Some(emb) = embedding {
157            guard.put_with_embedding_and_options(text.as_bytes(), emb, options)?
158        } else {
159            guard.put_bytes_with_options(text.as_bytes(), options)?
160        };
161        #[cfg(not(feature = "vec"))]
162        let id = guard.put_bytes_with_options(text.as_bytes(), options)?;
163        Ok(id)
164    }
165
166    /// Flush any pending writes to disk.
167    pub fn commit(&self) -> Result<(), MemvidError> {
168        let mut guard = self.lock()?;
169        guard.commit()?;
170        Ok(())
171    }
172
173    /// Run a [`SearchRequest`] directly. Useful for callers that need
174    /// memvid-native features (cursors, ACL contexts, etc.) that do not map
175    /// onto [`VectorSearchRequest`].
176    ///
177    /// # Concurrency
178    ///
179    /// Acquires the store's inner [`Mutex`] for the duration of the call.
180    /// Do **not** invoke this (or any other `MemvidStore` method) from
181    /// within a [`crate::WriteTransform`] closure: hook writes already hold
182    /// a path through `put_text` and a re-entrant call would deadlock.
183    pub fn search(
184        &self,
185        request: SearchRequest,
186    ) -> Result<memvid_core::SearchResponse, MemvidError> {
187        let mut guard = self.lock()?;
188        let resp = guard.search(request)?;
189        Ok(resp)
190    }
191
192    /// Total number of [`memvid_core::MemoryCard`]s currently stored on
193    /// the memories track.
194    ///
195    /// Cards are produced automatically when frames are written with
196    /// [`memvid_core::PutOptions::extract_triplets`] enabled (the default,
197    /// also exposed through [`crate::MemoryConfig::extract_triplets`]).
198    /// They form a structured Subject-Predicate-Object index over the
199    /// underlying free-text frames.
200    pub fn memory_card_count(&self) -> Result<usize, MemvidError> {
201        Ok(self.lock()?.memory_card_count())
202    }
203
204    /// Snapshot of every [`memvid_core::MemoryCard`] currently on the
205    /// memories track, cloned to owned values so the inner lock is
206    /// released before returning.
207    ///
208    /// Useful for callers that need to filter / sort across the entire
209    /// card set (for example
210    /// [`crate::MemoryCardContext`]'s `EntityMentions` selection
211    /// strategy). Avoid in hot paths against very large archives:
212    /// returns one allocation per card.
213    pub fn all_memory_cards(&self) -> Result<Vec<memvid_core::MemoryCard>, MemvidError> {
214        let guard = self.lock()?;
215        Ok(guard.memories().cards().to_vec())
216    }
217
218    /// Cards whose `entity` mentions appear (case-insensitive,
219    /// whole-word) in `query`. Filters behind the inner mutex so only
220    /// matching cards are cloned out, avoiding the full-archive
221    /// snapshot that [`MemvidStore::all_memory_cards`] performs.
222    pub fn cards_for_query(
223        &self,
224        query: &str,
225    ) -> Result<Vec<memvid_core::MemoryCard>, MemvidError> {
226        let needle = query.to_lowercase();
227        let guard = self.lock()?;
228        Ok(guard
229            .memories()
230            .cards()
231            .iter()
232            .filter(|card| {
233                let entity = card.entity.to_lowercase();
234                !entity.is_empty() && crate::cards_context::contains_word(&needle, &entity)
235            })
236            .cloned()
237            .collect())
238    }
239
240    /// Insert a fully-built [`memvid_core::MemoryCard`] onto the memories
241    /// track. The card's `id` field is overwritten with a freshly assigned
242    /// [`memvid_core::MemoryCardId`], which is returned.
243    ///
244    /// Useful for tests, deterministic seeding, or callers that have their
245    /// own structured-extraction pipeline upstream of memvid's.
246    pub fn put_memory_card(
247        &self,
248        card: memvid_core::MemoryCard,
249    ) -> Result<memvid_core::MemoryCardId, MemvidError> {
250        let mut guard = self.lock()?;
251        Ok(guard.put_memory_card(card)?)
252    }
253
254    /// All memory cards associated with `entity`, returned as owned
255    /// values (the underlying lock is released before returning).
256    ///
257    /// Returns an empty `Vec` if the entity is unknown. Pair with
258    /// [`MemvidStore::current_memory`] when only the latest non-retracted
259    /// value of a single slot is needed.
260    pub fn entity_memories(
261        &self,
262        entity: &str,
263    ) -> Result<Vec<memvid_core::MemoryCard>, MemvidError> {
264        let guard = self.lock()?;
265        Ok(guard
266            .get_entity_memories(entity)
267            .into_iter()
268            .cloned()
269            .collect())
270    }
271
272    /// The most recent non-retracted card for the given `entity` and
273    /// `slot`, if any. Mirrors
274    /// [`memvid_core::Memvid::get_current_memory`].
275    pub fn current_memory(
276        &self,
277        entity: &str,
278        slot: &str,
279    ) -> Result<Option<memvid_core::MemoryCard>, MemvidError> {
280        let guard = self.lock()?;
281        Ok(guard.get_current_memory(entity, slot).cloned())
282    }
283
284    /// All preference-kind cards for `entity`, in insertion order.
285    pub fn entity_preferences(
286        &self,
287        entity: &str,
288    ) -> Result<Vec<memvid_core::MemoryCard>, MemvidError> {
289        let guard = self.lock()?;
290        Ok(guard.get_preferences(entity).into_iter().cloned().collect())
291    }
292
293    /// Aggregate every distinct value recorded for `entity`/`slot` across
294    /// all sessions. Useful for slots that legitimately accumulate (lists
295    /// of hobbies, places lived in, etc.).
296    pub fn aggregate_memory_slot(
297        &self,
298        entity: &str,
299        slot: &str,
300    ) -> Result<Vec<String>, MemvidError> {
301        Ok(self.lock()?.aggregate_memory_slot(entity, slot))
302    }
303
304    /// Event-kind cards for `entity` in chronological order.
305    pub fn memory_timeline(
306        &self,
307        entity: &str,
308    ) -> Result<Vec<memvid_core::MemoryCard>, MemvidError> {
309        let guard = self.lock()?;
310        Ok(guard
311            .get_memory_timeline(entity)
312            .into_iter()
313            .cloned()
314            .collect())
315    }
316
317    // ---- Logic-Mesh (graph) pass-through ---------------------------------
318
319    /// Number of entity nodes in the underlying memvid Logic-Mesh.
320    ///
321    /// The Logic-Mesh is memvid's graph track: typed entity nodes
322    /// ([`memvid_core::MeshNode`]) connected by relationship edges
323    /// ([`memvid_core::MeshEdge`]). Populated automatically when frames
324    /// are written with NER-style enrichment (controlled by
325    /// [`memvid_core::PutOptions`]).
326    pub fn mesh_node_count(&self) -> Result<usize, MemvidError> {
327        Ok(self.lock()?.mesh_node_count())
328    }
329
330    /// Number of relationship edges in the Logic-Mesh.
331    pub fn mesh_edge_count(&self) -> Result<usize, MemvidError> {
332        Ok(self.lock()?.mesh_edge_count())
333    }
334
335    /// Find an entity node by canonical or display name (case-insensitive).
336    pub fn find_entity(&self, name: &str) -> Result<Option<memvid_core::MeshNode>, MemvidError> {
337        let guard = self.lock()?;
338        Ok(guard.find_entity(name).cloned())
339    }
340
341    /// All entity nodes mentioned in `frame_id`. Returns owned values
342    /// so the inner lock is released before returning.
343    pub fn frame_entities(&self, frame_id: u64) -> Result<Vec<memvid_core::MeshNode>, MemvidError> {
344        let guard = self.lock()?;
345        Ok(guard
346            .frame_entities(frame_id)
347            .into_iter()
348            .cloned()
349            .collect())
350    }
351
352    /// All entity nodes of the given [`memvid_core::EntityKind`].
353    pub fn entities_by_kind(
354        &self,
355        kind: memvid_core::EntityKind,
356    ) -> Result<Vec<memvid_core::MeshNode>, MemvidError> {
357        let guard = self.lock()?;
358        Ok(guard.entities_by_kind(kind).into_iter().cloned().collect())
359    }
360
361    /// Traverse the Logic-Mesh starting from `start`, following edges
362    /// of `link_type` up to `hops` deep. Wraps
363    /// [`memvid_core::Memvid::follow`].
364    ///
365    /// Useful for "who reports to alice's manager?"-style relationship
366    /// queries. Returns the result list directly; callers that want
367    /// streaming traversal should call memvid's `logic_mesh()` API by
368    /// holding their own clone of the inner [`memvid_core::Memvid`]
369    /// handle.
370    pub fn follow_relationship(
371        &self,
372        start: &str,
373        link_type: &str,
374        hops: usize,
375    ) -> Result<Vec<memvid_core::FollowResult>, MemvidError> {
376        let guard = self.lock()?;
377        Ok(guard.follow(start, link_type, hops))
378    }
379}
380
381/// Builder for [`MemvidStore`].
382#[derive(Default)]
383pub struct MemvidStoreBuilder {
384    path: Option<PathBuf>,
385    enable_lex: bool,
386    snippet_chars: Option<usize>,
387    acl_context: Option<AclContext>,
388    acl_enforcement_mode: Option<AclEnforcementMode>,
389    #[cfg(feature = "vec")]
390    enable_vec: bool,
391    #[cfg(feature = "vec")]
392    vec_model: Option<String>,
393    #[cfg(feature = "vec")]
394    embedder: Option<Arc<LocalTextEmbedder>>,
395}
396
397impl std::fmt::Debug for MemvidStoreBuilder {
398    // L6: hand-rolled to avoid leaking the boxed `embedder` closure
399    // through `#[derive(Debug)]` (which would require Debug on every
400    // captured value behind the `dyn Embedder` trait object). The
401    // placeholder `<embedder>` keeps the output stable and redacts a
402    // surface that may hold API keys or model handles by value.
403    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
404        let mut d = f.debug_struct("MemvidStoreBuilder");
405        d.field("path", &self.path)
406            .field("enable_lex", &self.enable_lex);
407        #[cfg(feature = "vec")]
408        {
409            d.field("enable_vec", &self.enable_vec)
410                .field("vec_model", &self.vec_model)
411                .field("embedder", &self.embedder.as_ref().map(|_| "<embedder>"));
412        }
413        d.finish()
414    }
415}
416
417impl MemvidStoreBuilder {
418    /// Path to the `.mv2` file.
419    pub fn path<P: Into<PathBuf>>(mut self, path: P) -> Self {
420        self.path = Some(path.into());
421        self
422    }
423
424    /// Enable BM25 / Tantivy lexical search on the underlying archive.
425    pub fn enable_lex(mut self) -> Self {
426        self.enable_lex = true;
427        self
428    }
429
430    /// Number of context characters to capture around each search hit.
431    /// Defaults to 400 characters. Applies to queries issued
432    /// via [`VectorStoreIndex::top_n`] and the `vec` search path; callers
433    /// who need per-query control should use [`MemvidStore::search`]
434    /// directly with a hand-built [`SearchRequest`].
435    pub fn snippet_chars(mut self, n: usize) -> Self {
436        self.snippet_chars = Some(n);
437        self
438    }
439
440    /// Default [`AclContext`] attached to every search performed through
441    /// the [`VectorStoreIndex`] / vector-search interfaces. When unset,
442    /// ACL filtering is disabled.
443    pub fn acl_context(mut self, ctx: AclContext) -> Self {
444        self.acl_context = Some(ctx);
445        self
446    }
447
448    /// ACL enforcement mode for default-attached contexts. Defaults to
449    /// [`AclEnforcementMode::Audit`].
450    pub fn acl_enforcement_mode(mut self, mode: AclEnforcementMode) -> Self {
451        self.acl_enforcement_mode = Some(mode);
452        self
453    }
454
455    /// Enable HNSW vector search on the underlying archive.
456    ///
457    /// Available only when this crate is built with the `vec` feature, which
458    /// pulls in `memvid-core/vec` (ONNX Runtime + bundled BGE-small).
459    /// Mutually compatible with [`Self::enable_lex`]; both can be on at once
460    /// for hybrid retrieval.
461    #[cfg(feature = "vec")]
462    pub fn enable_vec(mut self) -> Self {
463        self.enable_vec = true;
464        self
465    }
466
467    /// Bind (or validate) the embedding model identifier on the vector
468    /// index. See [`memvid_core::Memvid::set_vec_model`].
469    #[cfg(feature = "vec")]
470    pub fn vec_model(mut self, model: impl Into<String>) -> Self {
471        self.vec_model = Some(model.into());
472        self
473    }
474
475    /// Attach a local text embedder. Writes performed via
476    /// [`MemvidStore::put_text`] and queries performed via
477    /// [`VectorStoreIndex::top_n`] will be embedded with this model and
478    /// routed through memvid's HNSW vector index.
479    ///
480    /// Implies [`Self::enable_vec`]. If [`Self::vec_model`] has not been
481    /// set, the model identifier reported by the embedder is bound
482    /// automatically.
483    #[cfg(feature = "vec")]
484    pub fn embedder(mut self, embedder: LocalTextEmbedder) -> Self {
485        if self.vec_model.is_none() {
486            self.vec_model = Some(embedder.model_info().name.to_string());
487        }
488        self.embedder = Some(Arc::new(embedder));
489        self.enable_vec = true;
490        self
491    }
492
493    /// Convenience: attach the default local embedder (BGE-small,
494    /// 384-dimensional). The model is loaded from
495    /// [`TextEmbedConfig::default`]'s on-disk cache; if absent and
496    /// `offline` is `false` it will be downloaded.
497    #[cfg(feature = "vec")]
498    pub fn with_default_embedder(self) -> Result<Self, MemvidError> {
499        let embedder = LocalTextEmbedder::new(TextEmbedConfig::bge_small())?;
500        Ok(self.embedder(embedder))
501    }
502
503    /// Convenience: attach a local embedder built from an explicit
504    /// [`TextEmbedConfig`].
505    #[cfg(feature = "vec")]
506    pub fn with_embedder_config(self, config: TextEmbedConfig) -> Result<Self, MemvidError> {
507        let embedder = LocalTextEmbedder::new(config)?;
508        Ok(self.embedder(embedder))
509    }
510
511    fn require_path(&self) -> Result<&Path, MemvidError> {
512        self.path.as_deref().ok_or_else(|| {
513            MemvidError::Io(std::io::Error::new(
514                std::io::ErrorKind::InvalidInput,
515                "MemvidStoreBuilder requires a path",
516            ))
517        })
518    }
519
520    fn finish(self, memvid: Memvid) -> Result<MemvidStore, MemvidError> {
521        let mut memvid = memvid;
522        if self.enable_lex {
523            memvid.enable_lex()?;
524        }
525        #[cfg(feature = "vec")]
526        {
527            if self.enable_vec {
528                memvid.enable_vec()?;
529            }
530            if let Some(model) = self.vec_model.as_deref() {
531                memvid.set_vec_model(model)?;
532            }
533        }
534        #[cfg_attr(not(feature = "vec"), allow(unused_mut))]
535        let mut store = MemvidStore::from_memvid(memvid);
536        if let Some(s) = self.snippet_chars {
537            store.snippet_chars = s;
538        }
539        if let Some(ctx) = self.acl_context {
540            store.acl_context = Some(ctx);
541        }
542        if let Some(mode) = self.acl_enforcement_mode {
543            store.acl_enforcement_mode = mode;
544        }
545        #[cfg(feature = "vec")]
546        {
547            store.embedder = self.embedder;
548        }
549        Ok(store)
550    }
551
552    /// Open an existing `.mv2` file. Errors if the file does not exist.
553    pub fn open(self) -> Result<MemvidStore, MemvidError> {
554        let path = self.require_path()?.to_path_buf();
555        let memvid = Memvid::open(&path)?;
556        self.finish(memvid)
557    }
558
559    /// Create a new `.mv2` file. Errors if the file already exists.
560    pub fn create(self) -> Result<MemvidStore, MemvidError> {
561        let path = self.require_path()?.to_path_buf();
562        let memvid = Memvid::create(&path)?;
563        self.finish(memvid)
564    }
565
566    /// Open the file if it exists, otherwise create it.
567    pub fn open_or_create(self) -> Result<MemvidStore, MemvidError> {
568        let path = self.require_path()?.to_path_buf();
569        let memvid = if path.exists() {
570            Memvid::open(&path)?
571        } else {
572            Memvid::create(&path)?
573        };
574        self.finish(memvid)
575    }
576
577    /// Open the file read-only.
578    pub fn open_read_only(self) -> Result<MemvidStore, MemvidError> {
579        let path = self.require_path()?.to_path_buf();
580        let memvid = Memvid::open_read_only(&path)?;
581        self.finish(memvid)
582    }
583}
584
585/// A filter clause supported by [`MemvidStore`].
586///
587/// Memvid's query model does not support arbitrary boolean predicates;
588/// this filter only carries the restriction parameters that map onto
589/// fields of [`SearchRequest`]:
590///
591/// | Predicate                       | Effect on the search request    |
592/// | ------------------------------- | ------------------------------- |
593/// | `eq("uri", "...")`              | `request.uri = Some(value)`     |
594/// | `eq("scope", "...")`            | `request.scope = Some(value)`   |
595/// | `eq("as_of_frame", n)`          | `request.as_of_frame`           |
596/// | `eq("as_of_ts", n)`             | `request.as_of_ts`              |
597/// | `eq("cursor", "...")`           | `request.cursor` (pagination)   |
598/// | `eq("no_sketch", true/false)`   | disable sketch pre-filtering    |
599///
600/// `gt`, `lt`, and `or` are not representable; constructing such a filter
601/// produces an error at query time
602/// ([`MemvidError::UnsupportedFilter`]).
603#[derive(Debug, Clone, Default, Serialize, Deserialize)]
604pub struct MemvidFilter {
605    /// Optional URI prefix restriction.
606    pub uri: Option<String>,
607    /// Optional logical scope.
608    pub scope: Option<String>,
609    /// Optional point-in-time frame id.
610    pub as_of_frame: Option<u64>,
611    /// Optional point-in-time unix-millis timestamp.
612    pub as_of_ts: Option<i64>,
613    /// Optional pagination cursor (opaque token returned by a prior search).
614    #[serde(default, skip_serializing_if = "Option::is_none")]
615    pub cursor: Option<String>,
616    /// If `Some(true)`, disable the sketch pre-filter for this query.
617    #[serde(default, skip_serializing_if = "Option::is_none")]
618    pub no_sketch: Option<bool>,
619    /// Reasons this filter cannot be applied. Populated when the user calls
620    /// `gt`, `lt`, `or`, or `eq` with an unknown key.
621    #[serde(default, skip_serializing_if = "Vec::is_empty")]
622    invalid: Vec<String>,
623}
624
625impl MemvidFilter {
626    fn unsupported(reason: impl Into<String>) -> Self {
627        Self {
628            invalid: vec![reason.into()],
629            ..Self::default()
630        }
631    }
632
633    fn merge(mut self, rhs: Self) -> Self {
634        if rhs.uri.is_some() {
635            self.uri = rhs.uri;
636        }
637        if rhs.scope.is_some() {
638            self.scope = rhs.scope;
639        }
640        if rhs.as_of_frame.is_some() {
641            self.as_of_frame = rhs.as_of_frame;
642        }
643        if rhs.as_of_ts.is_some() {
644            self.as_of_ts = rhs.as_of_ts;
645        }
646        if rhs.cursor.is_some() {
647            self.cursor = rhs.cursor;
648        }
649        if rhs.no_sketch.is_some() {
650            self.no_sketch = rhs.no_sketch;
651        }
652        self.invalid.extend(rhs.invalid);
653        self
654    }
655
656    fn into_validated(self) -> Result<Self, MemvidError> {
657        if self.invalid.is_empty() {
658            Ok(self)
659        } else {
660            Err(MemvidError::UnsupportedFilter(self.invalid.join("; ")))
661        }
662    }
663
664    fn apply_to(self, request: &mut SearchRequest) {
665        request.uri = self.uri;
666        request.scope = self.scope;
667        request.as_of_frame = self.as_of_frame;
668        request.as_of_ts = self.as_of_ts;
669        if let Some(c) = self.cursor {
670            request.cursor = Some(c);
671        }
672        if let Some(b) = self.no_sketch {
673            request.no_sketch = b;
674        }
675    }
676
677    /// Returns `true` when this filter has no recorded validity
678    /// problems. Filters with `is_valid() == false` are rejected by
679    /// the search path with [`MemvidError::UnsupportedFilter`].
680    ///
681    /// Callers that build a [`MemvidFilter`] programmatically (for
682    /// example through Rig's `SearchFilter` combinators) can use this
683    /// pair with [`MemvidFilter::errors`] to surface the failure
684    /// before issuing the query.
685    pub fn is_valid(&self) -> bool {
686        self.invalid.is_empty()
687    }
688
689    /// Human-readable reasons why this filter cannot be applied, or
690    /// an empty slice when [`MemvidFilter::is_valid`] returns `true`.
691    pub fn errors(&self) -> &[String] {
692        &self.invalid
693    }
694}
695
696fn json_as_string(value: &serde_json::Value) -> Option<String> {
697    match value {
698        serde_json::Value::String(s) => Some(s.clone()),
699        other => Some(other.to_string()),
700    }
701}
702
703/// Coerce a JSON value into an `i64` for `as_of_ts`.
704///
705/// Accepts integer JSON numbers and integer-valued floats (which is the
706/// default representation for many JSON producers).
707fn as_of_ts_from_value(value: &serde_json::Value) -> Option<i64> {
708    if let Some(n) = value.as_i64() {
709        return Some(n);
710    }
711    let f = value.as_f64()?;
712    if f.is_finite() && f.fract() == 0.0 && f >= i64::MIN as f64 && f <= i64::MAX as f64 {
713        Some(f as i64)
714    } else {
715        None
716    }
717}
718
719impl SearchFilter for MemvidFilter {
720    type Value = serde_json::Value;
721
722    fn eq(key: impl AsRef<str>, value: Self::Value) -> Self {
723        let key = key.as_ref();
724        match key {
725            "uri" => Self {
726                uri: json_as_string(&value),
727                ..Self::default()
728            },
729            "scope" => Self {
730                scope: json_as_string(&value),
731                ..Self::default()
732            },
733            "as_of_frame" => match value.as_u64() {
734                Some(n) => Self {
735                    as_of_frame: Some(n),
736                    ..Self::default()
737                },
738                None => Self::unsupported(format!("as_of_frame must be a u64, got {value}")),
739            },
740            "as_of_ts" => match as_of_ts_from_value(&value) {
741                Some(n) => Self {
742                    as_of_ts: Some(n),
743                    ..Self::default()
744                },
745                None => Self::unsupported(format!("as_of_ts must be an i64, got {value}")),
746            },
747            "cursor" => Self {
748                cursor: json_as_string(&value),
749                ..Self::default()
750            },
751            "no_sketch" => match value.as_bool() {
752                Some(b) => Self {
753                    no_sketch: Some(b),
754                    ..Self::default()
755                },
756                None => Self::unsupported(format!("no_sketch must be a bool, got {value}")),
757            },
758            other => Self::unsupported(format!(
759                "unsupported filter key '{other}' (allowed: uri, scope, as_of_frame, as_of_ts, \
760                 cursor, no_sketch)"
761            )),
762        }
763    }
764
765    fn gt(key: impl AsRef<str>, _value: Self::Value) -> Self {
766        Self::unsupported(format!(
767            "memvid does not support gt() on '{}'",
768            key.as_ref()
769        ))
770    }
771
772    fn lt(key: impl AsRef<str>, _value: Self::Value) -> Self {
773        Self::unsupported(format!(
774            "memvid does not support lt() on '{}'",
775            key.as_ref()
776        ))
777    }
778
779    fn and(self, rhs: Self) -> Self {
780        self.merge(rhs)
781    }
782
783    fn or(self, _rhs: Self) -> Self {
784        // Memvid's filter model is a flat conjunction; representing a true
785        // disjunction would require widening the search request. Discard
786        // both operands and return a bare unsupported marker — the
787        // resulting filter is rejected by `into_validated()` regardless.
788        // Warn so callers using `SearchFilter::or` through Rig's generic
789        // combinator surface notice the silent rejection at runtime
790        // rather than only seeing the eventual `UnsupportedFilter` error.
791        tracing::warn!(
792            target: "rig_memvid::filter",
793            "SearchFilter::or is not supported by MemvidFilter; the resulting filter will be \
794             rejected by the search path with MemvidError::UnsupportedFilter"
795        );
796        let _ = self;
797        Self::unsupported("memvid does not support or() in filters")
798    }
799}
800
801/// Default snippet size when memvid is asked for context around a hit.
802///
803/// Tuned to be roughly one paragraph; callers who want different behaviour
804/// should call [`MemvidStore::search`] directly with their own
805/// [`SearchRequest`].
806const DEFAULT_SNIPPET_CHARS: usize = 400;
807
808/// Hard cap applied to `samples` (a.k.a. `top_k`) so callers cannot request
809/// `usize::MAX` worth of hits — both as a defensive measure on 32-bit
810/// targets where `u64 -> usize` may saturate, and to keep memvid from
811/// allocating absurdly large result vectors.
812const MAX_SAMPLES: usize = 1024;
813
814fn samples_to_top_k(samples: u64) -> usize {
815    let n = usize::try_from(samples).unwrap_or(MAX_SAMPLES);
816    n.min(MAX_SAMPLES)
817}
818
819fn build_search_request(
820    query: String,
821    samples: u64,
822    snippet_chars: usize,
823    filter: Option<MemvidFilter>,
824    acl_context: Option<AclContext>,
825    acl_enforcement_mode: AclEnforcementMode,
826) -> Result<SearchRequest, MemvidError> {
827    let filter = match filter {
828        Some(f) => f.into_validated()?,
829        None => MemvidFilter::default(),
830    };
831    let mut req = SearchRequest {
832        query,
833        top_k: samples_to_top_k(samples),
834        snippet_chars,
835        uri: None,
836        scope: None,
837        cursor: None,
838        #[cfg(feature = "temporal")]
839        temporal: None,
840        as_of_frame: None,
841        as_of_ts: None,
842        no_sketch: false,
843        acl_context,
844        acl_enforcement_mode,
845    };
846    filter.apply_to(&mut req);
847    Ok(req)
848}
849
850fn hit_score(hit: &SearchHit) -> f64 {
851    match hit.score {
852        Some(s) => f64::from(s),
853        // Lexical hits often arrive without a numeric score; fall back to
854        // rank-derived order-preserving values so callers can still sort.
855        // `hit.rank` is `usize`; cap at `u32::MAX` before promoting to f64
856        // to avoid lossy `as` casts that clippy would otherwise reject.
857        None => {
858            let rank = u32::try_from(hit.rank).unwrap_or(u32::MAX);
859            1.0 / (f64::from(rank) + 1.0)
860        }
861    }
862}
863
864#[cfg(feature = "vec")]
865fn ensure_vec_filter_supported(filter: &MemvidFilter) -> Result<(), MemvidError> {
866    if filter.uri.is_some() {
867        return Err(MemvidError::UnsupportedFilter(
868            "`uri` filter is not supported when querying through the embedder; use lex search"
869                .into(),
870        ));
871    }
872    if filter.as_of_frame.is_some() || filter.as_of_ts.is_some() {
873        return Err(MemvidError::UnsupportedFilter(
874            "point-in-time filters (`as_of_frame`, `as_of_ts`) are not supported under vector \
875             search; use lex or `MemvidStore::search` directly"
876                .into(),
877        ));
878    }
879    Ok(())
880}
881
882impl MemvidStore {
883    /// Run an embedding-driven search through memvid's HNSW index.
884    /// Pre-validated by the caller; returns the raw memvid response.
885    #[cfg(feature = "vec")]
886    fn vec_search(
887        &self,
888        query: &str,
889        samples: u64,
890        filter: &MemvidFilter,
891    ) -> Result<memvid_core::SearchResponse, MemvidError> {
892        let embedder = self
893            .embedder
894            .as_ref()
895            .ok_or_else(|| MemvidError::UnsupportedFilter("no embedder configured".into()))?;
896        let embedding = embedder.encode_text(query)?;
897        let top_k = samples_to_top_k(samples);
898        let mut guard = self.lock()?;
899        let resp = if self.acl_context.is_some() {
900            guard.vec_search_with_embedding_acl(
901                query,
902                &embedding,
903                top_k,
904                self.snippet_chars,
905                filter.scope.as_deref(),
906                self.acl_context.as_ref(),
907                self.acl_enforcement_mode,
908            )?
909        } else {
910            guard.vec_search_with_embedding(
911                query,
912                &embedding,
913                top_k,
914                self.snippet_chars,
915                filter.scope.as_deref(),
916            )?
917        };
918        Ok(resp)
919    }
920}
921
922impl MemvidStore {
923    /// Internal: dispatch a `VectorSearchRequest` to either the embedder-driven
924    /// vector path (if a local embedder is configured) or the lex/raw search
925    /// path. Centralises the `cfg(feature = "vec")` plumbing so the public
926    /// `VectorStoreIndex` methods stay small and free of duplication.
927    fn run_search(
928        &self,
929        query: String,
930        samples: u64,
931        filter: Option<MemvidFilter>,
932    ) -> Result<memvid_core::SearchResponse, MemvidError> {
933        #[cfg(feature = "vec")]
934        {
935            if self.embedder.is_some() {
936                let validated = match filter {
937                    Some(f) => f.into_validated()?,
938                    None => MemvidFilter::default(),
939                };
940                ensure_vec_filter_supported(&validated)?;
941                return self.vec_search(&query, samples, &validated);
942            }
943        }
944        let request = build_search_request(
945            query,
946            samples,
947            self.snippet_chars,
948            filter,
949            self.acl_context.clone(),
950            self.acl_enforcement_mode,
951        )?;
952        let mut guard = self.lock()?;
953        Ok(guard.search(request)?)
954    }
955}
956
957impl VectorStoreIndex for MemvidStore {
958    type Filter = MemvidFilter;
959
960    /// Run a search and deserialise each hit's JSON representation into `T`.
961    ///
962    /// # Contract
963    ///
964    /// The type `T` must be deserialisable from a [`SearchHit`] JSON object —
965    /// i.e. either `T = SearchHit` itself, or a struct whose fields are a
966    /// subset of `SearchHit`'s public fields (`frame_id`, `text`, `score`,
967    /// `metadata`, …). Use `serde_json::Value` for an opaque view.
968    ///
969    /// **This method does not round-trip user-defined document types.**
970    /// If you persisted JSON documents through [`InsertDocuments`] and want
971    /// them back, use [`VectorStoreIndex::top_n_ids`] for the frame ids and
972    /// then [`MemvidStore::search`] for full-fidelity access via the
973    /// memvid-native [`SearchRequest`] API.
974    ///
975    /// # Example
976    ///
977    /// ```rust,no_run
978    /// use memvid_core::SearchHit;
979    /// use rig::vector_store::{
980    ///     VectorSearchRequest, VectorStoreIndex,
981    ///     request::VectorSearchRequestBuilder,
982    /// };
983    /// use rig_memvid::{MemvidFilter, MemvidStore};
984    ///
985    /// # async fn run(store: MemvidStore) -> anyhow::Result<()> {
986    /// let req: VectorSearchRequest<MemvidFilter> =
987    ///     VectorSearchRequestBuilder::<MemvidFilter>::default()
988    ///         .query("hello")
989    ///         .samples(5)
990    ///         .build();
991    /// let hits: Vec<(f64, String, SearchHit)> = store.top_n(req).await?;
992    /// # Ok(())
993    /// # }
994    /// ```
995    async fn top_n<T>(
996        &self,
997        req: VectorSearchRequest<Self::Filter>,
998    ) -> Result<Vec<(f64, String, T)>, VectorStoreError>
999    where
1000        T: for<'a> Deserialize<'a> + WasmCompatSend,
1001    {
1002        let query = req.query().to_owned();
1003        let samples = req.samples();
1004        let filter = req.filter().clone();
1005
1006        let response = self.run_search(query, samples, filter)?;
1007
1008        let mut out = Vec::with_capacity(response.hits.len());
1009        for hit in response.hits {
1010            let score = hit_score(&hit);
1011            let id = hit.frame_id.to_string();
1012            let value = serde_json::to_value(&hit).map_err(MemvidError::from)?;
1013            let doc: T = serde_json::from_value(value).map_err(MemvidError::from)?;
1014            out.push((score, id, doc));
1015        }
1016        Ok(out)
1017    }
1018
1019    async fn top_n_ids(
1020        &self,
1021        req: VectorSearchRequest<Self::Filter>,
1022    ) -> Result<Vec<(f64, String)>, VectorStoreError> {
1023        let query = req.query().to_owned();
1024        let samples = req.samples();
1025        let filter = req.filter().clone();
1026
1027        let response = self.run_search(query, samples, filter)?;
1028
1029        Ok(response
1030            .hits
1031            .into_iter()
1032            .map(|hit| (hit_score(&hit), hit.frame_id.to_string()))
1033            .collect())
1034    }
1035}
1036
1037impl InsertDocuments for MemvidStore {
1038    /// Persist `documents` into the underlying `.mv2` file.
1039    ///
1040    /// **Note:** caller-supplied embeddings are intentionally ignored.
1041    /// On the lex-only path the document JSON is written as bytes and
1042    /// embeddings are dropped. When this store is configured with a
1043    /// local embedder (`vec` feature) every document is **re-embedded**
1044    /// with that model so memvid's vector index stays consistent with its
1045    /// bound model identifier.
1046    async fn insert_documents<Doc>(
1047        &self,
1048        documents: Vec<(Doc, OneOrMany<Embedding>)>,
1049    ) -> Result<(), VectorStoreError>
1050    where
1051        Doc: Serialize + Embed + WasmCompatSend,
1052    {
1053        // We deliberately ignore the externally-supplied embeddings (rig
1054        // computes them with its own model, but memvid validates the
1055        // dimension against its bound model and would reject mismatches).
1056        // When this store has its own embedder, embed each document with
1057        // the local model. Round-tripping the document through JSON gives
1058        // us a stable byte payload that `serde_json::from_value::<T>` can
1059        // recover during search.
1060        #[cfg(feature = "vec")]
1061        let local_embedder = self.embedder.clone();
1062        let mut prepared: Vec<(Vec<u8>, Option<Vec<f32>>)> = Vec::with_capacity(documents.len());
1063        for (doc, _embeddings) in documents {
1064            let bytes = serde_json::to_vec(&doc).map_err(MemvidError::from)?;
1065            #[cfg(feature = "vec")]
1066            let emb = match &local_embedder {
1067                Some(embedder) => {
1068                    // `serde_json::to_vec` always returns valid UTF-8, so this
1069                    // path is fully infallible today. Use `from_utf8` (not
1070                    // `from_utf8_unchecked`) to keep the invariant explicit:
1071                    // if a future refactor swaps the encoder, we surface the
1072                    // problem as a typed error instead of silently embedding
1073                    // an empty string.
1074                    let text = std::str::from_utf8(&bytes).map_err(|e| {
1075                        MemvidError::Io(std::io::Error::new(std::io::ErrorKind::InvalidData, e))
1076                    })?;
1077                    Some(embedder.encode_text(text).map_err(MemvidError::from)?)
1078                }
1079                None => None,
1080            };
1081            #[cfg(not(feature = "vec"))]
1082            let emb: Option<Vec<f32>> = None;
1083            prepared.push((bytes, emb));
1084        }
1085
1086        let mut guard = self
1087            .inner
1088            .lock()
1089            .map_err(|_| VectorStoreError::from(MemvidError::Poisoned))?;
1090        for (bytes, emb) in prepared {
1091            match emb {
1092                Some(embedding) => {
1093                    guard
1094                        .put_with_embedding_and_options(&bytes, embedding, PutOptions::default())
1095                        .map_err(MemvidError::from)?;
1096                }
1097                None => {
1098                    guard
1099                        .put_bytes_with_options(&bytes, PutOptions::default())
1100                        .map_err(MemvidError::from)?;
1101                }
1102            }
1103        }
1104        guard.commit().map_err(MemvidError::from)?;
1105        Ok(())
1106    }
1107}