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