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}