1#![allow(clippy::missing_const_for_fn)]
33
34use std::collections::HashMap;
35use std::sync::{Arc, Mutex};
36
37use uuid::Uuid;
38use wm_core::{CoreError, Galaxy, Result};
39
40use crate::associations::AssociationStore;
41use crate::embedder::Embedder;
42use crate::memory::content_hash;
43use crate::search::{SearchEngine, SearchResult};
44use crate::store::MemoryStore;
45use crate::vector::{VectorSearchResult, VectorStore};
46
47#[derive(Debug, Clone)]
51pub struct RecallResult {
52 pub memory_id: Uuid,
54 pub galaxy: Galaxy,
56 pub score: f32,
58 pub bm25_score: f32,
60 pub raw_bm25_score: f32,
64 pub vector_score: f32,
66 pub importance: f32,
68 pub graph_score: f32,
72 pub trust_factor: f32,
75 pub corroboration: u32,
79 pub in_conformal_set: bool,
82 pub content: String,
84}
85
86#[derive(Debug, Clone)]
90pub struct RecallConfig {
91 pub bm25_weight: f32,
93 pub vector_weight: f32,
95 pub importance_weight: f32,
97 pub graph_weight: f32,
104 pub trust_weight: f32,
112 pub corroboration_weight: f32,
118 pub conformal_alpha: Option<f32>,
124 pub cache_embeddings: bool,
126 pub max_cache_entries: usize,
128 #[allow(dead_code)]
131 pub writer_heap_size: usize,
132 pub promotion_on_read: bool,
137 pub association_rerank: bool,
141}
142
143impl Default for RecallConfig {
144 fn default() -> Self {
145 Self {
146 bm25_weight: 0.5,
147 vector_weight: 0.3,
148 importance_weight: 0.2,
149 graph_weight: 0.0,
150 trust_weight: 0.0,
151 corroboration_weight: 0.0,
152 conformal_alpha: None,
153 cache_embeddings: true,
154 max_cache_entries: 1000,
155 writer_heap_size: 50_000_000,
156 promotion_on_read: false,
157 association_rerank: false,
158 }
159 }
160}
161
162impl RecallConfig {
163 #[must_use]
168 pub fn from_env() -> Self {
169 let mut config = Self::default();
170
171 if let Ok(v) = std::env::var("WM_RECALL_BM25_WEIGHT") {
172 if let Ok(w) = v.parse::<f32>() {
173 if w.is_finite() && w >= 0.0 {
174 config.bm25_weight = w.min(1.0);
175 }
176 }
177 }
178 if let Ok(v) = std::env::var("WM_RECALL_VECTOR_WEIGHT") {
179 if let Ok(w) = v.parse::<f32>() {
180 if w.is_finite() && w >= 0.0 {
181 config.vector_weight = w.min(1.0);
182 }
183 }
184 }
185 if let Ok(v) = std::env::var("WM_RECALL_IMPORTANCE_WEIGHT") {
186 if let Ok(w) = v.parse::<f32>() {
187 if w.is_finite() && w >= 0.0 {
188 config.importance_weight = w.min(1.0);
189 }
190 }
191 }
192 if let Ok(v) = std::env::var("WM_RECALL_GRAPH_WEIGHT") {
195 if let Ok(w) = v.parse::<f32>() {
196 if w.is_finite() && w >= 0.0 {
197 config.graph_weight = w.min(1.0);
198 }
199 }
200 }
201 if let Ok(v) = std::env::var("WM_TRUST_WEIGHT") {
205 if let Ok(w) = v.parse::<f32>() {
206 if w.is_finite() && w >= 0.0 {
207 config.trust_weight = w.min(1.0);
208 }
209 }
210 }
211 if let Ok(v) = std::env::var("WM_CORROBORATION_WEIGHT") {
214 if let Ok(w) = v.parse::<f32>() {
215 if w.is_finite() && w >= 0.0 {
216 config.corroboration_weight = w.min(1.0);
217 }
218 }
219 }
220 if let Ok(v) = std::env::var("WM_RECALL_CONFORMAL_ALPHA") {
223 if let Ok(a) = v.parse::<f32>() {
224 if a.is_finite() && a > 0.0 && a < 1.0 {
225 config.conformal_alpha = Some(a);
226 }
227 }
228 }
229 if let Ok(v) = std::env::var("WM_PROMOTION_ON_READ") {
232 if v == "1" || v.eq_ignore_ascii_case("true") {
233 config.promotion_on_read = true;
234 }
235 }
236 if let Ok(v) = std::env::var("WM_ASSOCIATION_RERANK") {
239 if v == "1" || v.eq_ignore_ascii_case("true") {
240 config.association_rerank = true;
241 }
242 }
243
244 let sum = config.bm25_weight + config.vector_weight + config.importance_weight;
246 if sum > 0.0 && (sum - 1.0).abs() > 0.01 {
247 config.bm25_weight /= sum;
248 config.vector_weight /= sum;
249 config.importance_weight /= sum;
250 }
251
252 config
253 }
254
255 #[must_use]
257 pub fn weights_normalized(&self) -> bool {
258 let sum = self.bm25_weight + self.vector_weight + self.importance_weight;
259 (sum - 1.0).abs() < 0.01
260 }
261}
262
263pub struct RecallEngine {
270 store: Arc<MemoryStore>,
271 search_engine: Arc<SearchEngine>,
272 vector_store: Mutex<VectorStore>,
273 embedder: Arc<dyn Embedder>,
274 config: RecallConfig,
275 embedding_cache: Mutex<HashMap<String, Vec<f32>>>,
276 conformal: Mutex<Option<crate::recall_conformal::RecallConformal>>,
279}
280
281#[derive(Debug, Clone, Default, serde::Serialize)]
283pub struct BackfillReport {
284 pub scanned: usize,
286 pub candidates: usize,
288 pub embedded: usize,
290 pub already_embedded: usize,
292 pub skipped_empty: usize,
295 pub errors: usize,
297 pub dry_run: bool,
299}
300
301impl RecallEngine {
302 pub fn new(
306 store: Arc<MemoryStore>,
307 search_engine: Arc<SearchEngine>,
308 vector_store: VectorStore,
309 embedder: Arc<dyn Embedder>,
310 config: RecallConfig,
311 ) -> Result<Self> {
312 let conformal =
313 Mutex::new(config.conformal_alpha.and_then(|alpha| {
314 crate::recall_conformal::RecallConformal::new(alpha, store.clone())
315 }));
316 Ok(Self {
317 store,
318 search_engine,
319 vector_store: Mutex::new(vector_store),
320 embedder,
321 config,
322 embedding_cache: Mutex::new(HashMap::new()),
323 conformal,
324 })
325 }
326
327 pub fn record_relevance_feedback(&self, score: f32, relevant: bool) -> Result<usize> {
331 let mut guard = self
332 .conformal
333 .lock()
334 .map_err(|e| CoreError::Memory(format!("recall conformal lock: {e}")))?;
335 match guard.as_mut() {
336 Some(rc) => Ok(rc.record_feedback(score, relevant)),
337 None => Err(CoreError::InvalidArgs(
338 "conformal calibration is not enabled — set WM_RECALL_CONFORMAL_ALPHA in (0,1)"
339 .into(),
340 )),
341 }
342 }
343
344 #[allow(clippy::significant_drop_tightening)]
352 pub fn conformal_disclosure(
353 &self,
354 results: &mut [RecallResult],
355 ) -> Result<Option<crate::recall_conformal::ConformalSetInfo>> {
356 use crate::recall_conformal::ConformalSetInfo;
357 let guard = self
358 .conformal
359 .lock()
360 .map_err(|e| CoreError::Memory(format!("recall conformal lock: {e}")))?;
361 let Some(rc) = guard.as_ref() else {
362 return Ok(None);
363 };
364 let coverage_target = Some(f64::from(1.0 - rc.alpha()));
365 let info = if rc.is_fitted() {
366 let threshold = rc.threshold();
367 let mut set_size = 0usize;
368 for r in results.iter_mut() {
369 r.in_conformal_set = rc.membership(r.score) == Some(true);
370 if r.in_conformal_set {
371 set_size += 1;
372 }
373 }
374 ConformalSetInfo {
375 status: "active".into(),
376 alpha: Some(f64::from(rc.alpha())),
377 coverage_target,
378 calibration_samples: Some(rc.sample_count()),
379 threshold,
380 set_size: Some(set_size),
381 hint: None,
382 }
383 } else {
384 ConformalSetInfo {
385 status: "uncalibrated".into(),
386 alpha: Some(f64::from(rc.alpha())),
387 coverage_target,
388 calibration_samples: Some(rc.sample_count()),
389 threshold: None,
390 set_size: None,
391 hint: Some(format!(
392 "record ≥ {} relevance-feedback samples to calibrate",
393 crate::recall_conformal::MIN_SAMPLES
394 )),
395 }
396 };
397 Ok(Some(info))
398 }
399
400 #[must_use]
402 pub fn config(&self) -> &RecallConfig {
403 &self.config
404 }
405
406 #[must_use]
411 pub fn embedder_is_real(&self) -> bool {
412 self.embedder.backend_name() != "stub"
413 }
414
415 #[must_use]
425 fn vector_half_inert(&self) -> bool {
426 self.config.vector_weight <= 0.0 && self.config.importance_weight <= 0.0
427 }
428
429 pub fn embedder_probe(&self) -> Result<usize> {
438 self.embedder
439 .embed("wm embedder probe")
440 .map(|vector| vector.len())
441 }
442
443 fn embedding_cache_key(&self, embedded_text: &str) -> String {
448 format!(
449 "{}:{}",
450 self.embedder.cache_namespace(),
451 content_hash(embedded_text)
452 )
453 }
454
455 fn embed_content(&self, content: &str) -> Result<Vec<f32>> {
461 let hash = content_hash(content);
462
463 if self.config.cache_embeddings {
465 let cache = self
466 .embedding_cache
467 .lock()
468 .map_err(|e| CoreError::Memory(format!("embedding cache lock: {e}")))?;
469 if let Some(vec) = cache.get(&hash) {
470 return Ok(vec.clone());
471 }
472 }
473
474 let cache_key = self.embedding_cache_key(content);
476 match self.store.get_embedding_cache(&cache_key) {
477 Ok(Some(vector)) => {
478 if self.config.cache_embeddings {
479 if let Ok(mut cache) = self.embedding_cache.lock() {
480 cache.insert(hash, vector.clone());
481 }
482 }
483 return Ok(vector);
484 }
485 Ok(None) => {}
486 Err(error) => tracing::warn!("embedding cache read failed: {error}"),
487 }
488
489 let embedding = self.embedder.embed(content)?;
491
492 if let Err(error) = self.store.put_embedding_cache(&cache_key, &embedding) {
494 tracing::warn!("embedding cache write failed: {error}");
495 }
496
497 if self.config.cache_embeddings {
499 let mut cache = self
500 .embedding_cache
501 .lock()
502 .map_err(|e| CoreError::Memory(format!("embedding cache lock: {e}")))?;
503 if cache.len() >= self.config.max_cache_entries {
504 let to_remove: Vec<String> = cache.keys().take(cache.len() / 10).cloned().collect();
506 for key in to_remove {
507 cache.remove(&key);
508 }
509 }
510 cache.insert(hash, embedding.clone());
511 }
512
513 Ok(embedding)
514 }
515
516 pub fn store_with_embedding(&self, galaxy: Galaxy, memory: &crate::Memory) -> Result<()> {
524 let embedding = self.embed_content(&memory.content)?;
526
527 self.store.put(galaxy, memory)?;
529
530 self.store.put_embedding(memory.metadata.id, &embedding)?;
532
533 {
535 let mut vs = self
536 .vector_store
537 .lock()
538 .map_err(|e| CoreError::Memory(format!("vector store lock: {e}")))?;
539 vs.add(memory.metadata.id, galaxy, embedding);
540 }
541
542 let timestamp = memory.metadata.created_at.timestamp();
544 let tags: Vec<String> = memory.metadata.tags.clone();
545 {
546 let mut writer = self.search_engine.writer()?;
547 self.search_engine.add_document(
548 &mut writer,
549 &memory.metadata.id.to_string(),
550 galaxy.db_name(),
551 &memory.content,
552 &tags,
553 timestamp,
554 )?;
555 self.search_engine.commit(&mut writer)?;
556 }
557
558 Ok(())
559 }
560
561 pub fn store_batch_with_embedding(
572 &self,
573 entries: &[(Galaxy, &crate::Memory)],
574 ) -> Result<usize> {
575 if entries.is_empty() {
576 return Ok(0);
577 }
578
579 const MAX_CHARS_PER_CHUNK: usize = 1500; const MAX_CHARS_PER_ITEM: usize = 1500; let contents: Vec<String> = entries
586 .iter()
587 .map(|(_, m)| {
588 if m.content.len() > MAX_CHARS_PER_ITEM {
589 m.content.chars().take(MAX_CHARS_PER_ITEM).collect()
590 } else {
591 m.content.clone()
592 }
593 })
594 .collect();
595 let content_refs: Vec<&str> = contents.iter().map(String::as_str).collect();
596 let cache_keys: Vec<String> = content_refs
597 .iter()
598 .map(|c| self.embedding_cache_key(c))
599 .collect();
600 let mut embeddings: Vec<Option<Vec<f32>>> =
601 match self.store.get_embedding_cache_batch(&cache_keys) {
602 Ok(cached) => cached,
603 Err(error) => {
604 tracing::warn!("embedding cache batch read failed: {error}");
605 vec![None; cache_keys.len()]
606 }
607 };
608 let misses: Vec<(usize, &str)> = content_refs
609 .iter()
610 .enumerate()
611 .filter(|(i, _)| embeddings[*i].is_none())
612 .map(|(i, content)| (i, *content))
613 .collect();
614
615 let max_batch_texts = self.embedder.preferred_max_batch_texts();
620 let mut chunk: Vec<&str> = Vec::new();
621 let mut chunk_positions: Vec<usize> = Vec::new();
622 let mut chunk_chars: usize = 0;
623 let mut embed_chunk = |chunk: &[&str], positions: &[usize], store: &Self| -> Result<()> {
624 let chunk_vecs = store.embedder.embed_batch(chunk)?;
625 if chunk_vecs.len() != chunk.len() {
626 return Err(CoreError::Memory(format!(
627 "embed_batch returned {} vectors for {} inputs (chunk)",
628 chunk_vecs.len(),
629 chunk.len()
630 )));
631 }
632 for (pos, vector) in positions.iter().zip(chunk_vecs) {
633 embeddings[*pos] = Some(vector);
634 }
635 Ok(())
636 };
637 for &(position, content) in &misses {
638 let content_chars = content.len();
639 let flush = !chunk.is_empty()
640 && (chunk_chars + content_chars > MAX_CHARS_PER_CHUNK
641 || chunk.len() >= max_batch_texts);
642 if flush {
643 embed_chunk(&chunk, &chunk_positions, self)?;
644 chunk.clear();
645 chunk_positions.clear();
646 chunk_chars = 0;
647 }
648 chunk.push(content);
649 chunk_positions.push(position);
650 chunk_chars += content_chars;
651 }
652 if !chunk.is_empty() {
653 embed_chunk(&chunk, &chunk_positions, self)?;
654 }
655
656 let fresh: Vec<(String, Vec<f32>)> = misses
659 .iter()
660 .filter_map(|&(position, _)| {
661 embeddings[position]
662 .clone()
663 .map(|vector| (cache_keys[position].clone(), vector))
664 })
665 .collect();
666 if let Err(error) = self.store.put_embedding_cache_batch(&fresh) {
667 tracing::warn!("embedding cache batch write failed: {error}");
668 }
669
670 for (i, (galaxy, memory)) in entries.iter().enumerate() {
672 let Some(ref embedding) = embeddings[i] else {
673 return Err(CoreError::Memory(format!(
674 "embedding missing for entry {i} after cache resolution"
675 )));
676 };
677 self.store.put(*galaxy, memory)?;
678 self.store.put_embedding(memory.metadata.id, embedding)?;
679 }
680
681 {
683 let mut vs = self
684 .vector_store
685 .lock()
686 .map_err(|e| CoreError::Memory(format!("vector store lock: {e}")))?;
687 for (i, (galaxy, memory)) in entries.iter().enumerate() {
688 if let Some(ref embedding) = embeddings[i] {
689 vs.add(memory.metadata.id, *galaxy, embedding.clone());
690 }
691 }
692 }
693
694 {
696 let mut writer = self.search_engine.writer()?;
697 for (galaxy, memory) in entries {
698 let timestamp = memory.metadata.created_at.timestamp();
699 let tags: Vec<String> = memory.metadata.tags.clone();
700 self.search_engine.add_document(
701 &mut writer,
702 &memory.metadata.id.to_string(),
703 galaxy.db_name(),
704 &memory.content,
705 &tags,
706 timestamp,
707 )?;
708 }
709 self.search_engine.commit(&mut writer)?;
710 }
711
712 if self.config.cache_embeddings {
716 let mut cache = self
717 .embedding_cache
718 .lock()
719 .map_err(|e| CoreError::Memory(format!("embedding cache lock: {e}")))?;
720 for &(position, _) in &misses {
721 let Some(ref vector) = embeddings[position] else {
722 continue;
723 };
724 let hash = content_hash(&entries[position].1.content);
725 if cache.len() >= self.config.max_cache_entries {
726 let to_remove: Vec<String> =
727 cache.keys().take(cache.len() / 10).cloned().collect();
728 for key in to_remove {
729 cache.remove(&key);
730 }
731 }
732 cache.insert(hash, vector.clone());
733 }
734 }
735
736 Ok(entries.len())
737 }
738
739 pub fn backfill_embeddings(
755 &self,
756 galaxy: Option<Galaxy>,
757 limit: usize,
758 dry_run: bool,
759 ) -> Result<BackfillReport> {
760 if !self.embedder_is_real() {
761 return Err(CoreError::InvalidArgs(
762 "no real embedder configured — memory.reembed requires WM_EMBEDDER_ENDPOINT or the onnx backend".into(),
763 ));
764 }
765 use lmdb::{Cursor as _, Transaction as _};
766 let limit = if limit == 0 { usize::MAX } else { limit };
767 let galaxies: Vec<Galaxy> = match galaxy {
768 Some(g) => vec![g],
769 None => Galaxy::memory_galaxies().to_vec(),
770 };
771 let mut report = BackfillReport {
772 dry_run,
773 ..Default::default()
774 };
775 let mut candidates: Vec<crate::Memory> = Vec::new();
776
777 'galaxy: for g in galaxies {
778 let db = self.store.galaxy_db(g)?;
779 let embeddings_db = self.store.galaxy_db(Galaxy::Embeddings)?;
780 let tx = self
781 .store
782 .env()
783 .begin_ro_txn()
784 .map_err(|e| CoreError::Memory(format!("LMDB ro_txn failed: {e}")))?;
785 {
786 let mut cursor = tx
787 .open_ro_cursor(db)
788 .map_err(|e| CoreError::Memory(format!("LMDB cursor failed: {e}")))?;
789 for (key, value) in cursor.iter() {
790 report.scanned += 1;
791 let memory = match self.store.decode_record_value(g, key, value) {
792 Ok(memory) => memory,
793 Err(error) => {
794 report.errors += 1;
795 tracing::warn!("reembed scan skipped undecodable entry: {error}");
796 continue;
797 }
798 };
799 let has_vector = match tx.get(embeddings_db, memory.metadata.id.as_bytes()) {
800 Ok(_) => true,
801 Err(lmdb::Error::NotFound) => false,
802 Err(e) => {
803 return Err(CoreError::Memory(format!("LMDB get failed: {e}")));
804 }
805 };
806 if has_vector {
807 report.already_embedded += 1;
808 continue;
809 }
810 if memory.content.trim().is_empty() {
811 report.skipped_empty += 1;
815 continue;
816 }
817 candidates.push(memory);
818 if candidates.len() >= limit {
819 break;
820 }
821 }
822 }
823 tx.commit()
824 .map_err(|e| CoreError::Memory(format!("LMDB commit failed: {e}")))?;
825 if candidates.len() >= limit {
826 break 'galaxy;
827 }
828 }
829
830 report.candidates = candidates.len();
831 if dry_run {
832 return Ok(report);
833 }
834
835 let chunk_size = std::env::var("WM_REEMBED_BATCH")
843 .ok()
844 .and_then(|value| value.parse::<usize>().ok())
845 .map_or(32, |size| size.clamp(1, 64));
846 for chunk in candidates.chunks(chunk_size) {
847 let keys: Vec<String> = chunk
848 .iter()
849 .map(|memory| self.embedding_cache_key(&memory.content))
850 .collect();
851 let mut vectors: Vec<Option<Vec<f32>>> =
852 match self.store.get_embedding_cache_batch(&keys) {
853 Ok(cached) if cached.len() == chunk.len() => cached,
854 Ok(_) | Err(_) => vec![None; chunk.len()],
855 };
856
857 let miss_idx: Vec<usize> = vectors
858 .iter()
859 .enumerate()
860 .filter(|(_, vector)| vector.is_none())
861 .map(|(i, _)| i)
862 .collect();
863 if !miss_idx.is_empty() {
864 let texts: Vec<&str> = miss_idx
865 .iter()
866 .map(|&i| chunk[i].content.as_str())
867 .collect();
868 match self.embedder.embed_batch(&texts) {
869 Ok(embedded) if embedded.len() == texts.len() => {
870 let cache_entries: Vec<(String, Vec<f32>)> = miss_idx
871 .iter()
872 .zip(embedded.iter())
873 .map(|(&i, vector)| (keys[i].clone(), vector.clone()))
874 .collect();
875 if let Err(error) = self.store.put_embedding_cache_batch(&cache_entries) {
876 tracing::warn!("reembed cache persist failed: {error}");
877 }
878 for (&i, vector) in miss_idx.iter().zip(embedded) {
879 vectors[i] = Some(vector);
880 }
881 }
882 Ok(embedded) => {
883 tracing::warn!(
884 expected = texts.len(),
885 got = embedded.len(),
886 "reembed batch length mismatch — falling back per item"
887 );
888 self.embed_misses_per_item(chunk, &miss_idx, &mut vectors, &mut report);
889 }
890 Err(error) => {
891 tracing::warn!("reembed batch failed ({error}) — falling back per item");
892 self.embed_misses_per_item(chunk, &miss_idx, &mut vectors, &mut report);
893 }
894 }
895 }
896
897 for (i, memory) in chunk.iter().enumerate() {
898 let Some(vector) = vectors[i].take() else {
899 continue;
902 };
903 if let Err(error) = self.store.put_embedding(memory.metadata.id, &vector) {
904 report.errors += 1;
905 tracing::warn!(
906 memory = %memory.metadata.id,
907 "reembed persist failed: {error}"
908 );
909 continue;
910 }
911 if let Ok(mut vs) = self.vector_store.lock() {
912 vs.add(memory.metadata.id, memory.metadata.galaxy, vector);
913 }
914 report.embedded += 1;
915 }
916 }
917 Ok(report)
918 }
919
920 fn embed_misses_per_item(
923 &self,
924 chunk: &[crate::Memory],
925 miss_idx: &[usize],
926 vectors: &mut [Option<Vec<f32>>],
927 report: &mut BackfillReport,
928 ) {
929 for &i in miss_idx {
930 match self.embed_content(&chunk[i].content) {
931 Ok(vector) => vectors[i] = Some(vector),
932 Err(error) => {
933 report.errors += 1;
934 tracing::warn!(
935 memory = %chunk[i].metadata.id,
936 "reembed embed failed: {error}"
937 );
938 }
939 }
940 }
941 }
942
943 fn ensure_vectors_loaded(&self) -> Result<()> {
951 let mut vs = self
952 .vector_store
953 .lock()
954 .map_err(|e| CoreError::Memory(format!("vector store lock: {e}")))?;
955 if vs.is_loaded() {
956 return Ok(());
957 }
958 vs.load(&self.store)
959 }
960
961 #[must_use]
965 pub fn hybrid_search(
966 &self,
967 query: &str,
968 limit: usize,
969 galaxy_filter: Option<Galaxy>,
970 ) -> Vec<RecallResult> {
971 self.hybrid_search_with_disclosure(query, limit, galaxy_filter)
972 .0
973 }
974
975 pub fn hybrid_search_with_disclosure(
979 &self,
980 query: &str,
981 limit: usize,
982 galaxy_filter: Option<Galaxy>,
983 ) -> (
984 Vec<RecallResult>,
985 Option<crate::recall_conformal::ConformalSetInfo>,
986 ) {
987 let bm25_limit = limit * 3;
989 let bm25_results = self
990 .search_engine
991 .search_in_galaxy(query, galaxy_filter, bm25_limit)
992 .unwrap_or_default();
993
994 let vector_results = if self.vector_half_inert() {
1005 Vec::new()
1006 } else {
1007 let query_vec = match self.embedder.embed_query(query) {
1009 Ok(v) => v,
1010 Err(_) => return (Vec::new(), None),
1011 };
1012
1013 if let Err(error) = self.ensure_vectors_loaded() {
1017 tracing::warn!(
1018 "vector store rehydration failed ({error}) — hybrid vector half degraded"
1019 );
1020 }
1021
1022 let Ok(vs) = self.vector_store.lock() else {
1023 return (Vec::new(), None);
1024 };
1025 vs.search(&query_vec, bm25_limit, galaxy_filter)
1026 };
1027
1028 let fused = self.fuse_results(&bm25_results, &vector_results, limit);
1030
1031 let fused = if crate::memory::validity_enforced() {
1035 fused
1036 .into_iter()
1037 .filter(|r| {
1038 self.find_memory_anywhere(r.memory_id)
1039 .is_none_or(|mem| mem.metadata.validity.is_current())
1040 })
1041 .collect()
1042 } else {
1043 fused
1044 };
1045
1046 let mut expanded = self.expand_with_graph(fused, limit);
1049
1050 if self.config.corroboration_weight > 0.0 {
1055 for r in &mut expanded {
1056 if let Some(mem) = self.find_memory_anywhere(r.memory_id) {
1057 let n = mem.metadata.corroborated_by.len();
1058 r.corroboration = n.min(u32::MAX as usize) as u32;
1059 r.score = crate::memory::corroboration_boost(
1060 r.score,
1061 n,
1062 self.config.corroboration_weight,
1063 );
1064 }
1065 }
1066 expanded.sort_by(|a, b| {
1067 b.score
1068 .partial_cmp(&a.score)
1069 .unwrap_or(std::cmp::Ordering::Equal)
1070 });
1071 }
1072
1073 if self.config.association_rerank {
1078 if let Ok(assoc_store) = AssociationStore::open(self.store.env()) {
1079 let env = self.store.env();
1080 for r in &mut expanded {
1081 let outgoing = assoc_store.find_from(env, r.memory_id).unwrap_or_default();
1082 let incoming = assoc_store.find_to(env, r.memory_id).unwrap_or_default();
1083 let active_edges = outgoing
1084 .iter()
1085 .chain(incoming.iter())
1086 .filter(|e| e.weight >= 0.2)
1087 .count();
1088 if active_edges > 0 {
1089 let boost = (active_edges as f32 * 0.05).min(0.25);
1090 r.score *= 1.0 + boost;
1091 }
1092 }
1093 expanded.sort_by(|a, b| {
1094 b.score
1095 .partial_cmp(&a.score)
1096 .unwrap_or(std::cmp::Ordering::Equal)
1097 });
1098 }
1099 }
1100
1101 if self.config.promotion_on_read {
1105 for r in expanded.iter().take(limit) {
1106 if let Err(e) = self.promote_memory(r.galaxy, r.memory_id) {
1107 tracing::warn!(
1108 error = %e,
1109 memory_id = %r.memory_id,
1110 galaxy = %r.galaxy.db_name(),
1111 "promotion on read failed"
1112 );
1113 }
1114 }
1115 }
1116
1117 match self.conformal_disclosure(&mut expanded) {
1120 Ok(info) => (expanded, info),
1121 Err(e) => {
1122 tracing::warn!(error = %e, "recall conformal disclosure failed");
1123 (expanded, None)
1124 }
1125 }
1126 }
1127
1128 #[must_use]
1130 pub fn vector_search(
1131 &self,
1132 query: &str,
1133 limit: usize,
1134 galaxy_filter: Option<Galaxy>,
1135 ) -> Vec<RecallResult> {
1136 let query_vec = match self.embedder.embed_query(query) {
1137 Ok(v) => v,
1138 Err(_) => return Vec::new(),
1139 };
1140
1141 let vector_results = {
1142 let Ok(vs) = self.vector_store.lock() else {
1143 return Vec::new();
1144 };
1145 vs.search(&query_vec, limit, galaxy_filter)
1146 };
1147
1148 vector_results
1149 .into_iter()
1150 .map(|vr| {
1151 let content = self.get_memory_content(vr.memory_id, vr.galaxy);
1152 RecallResult {
1153 memory_id: vr.memory_id,
1154 galaxy: vr.galaxy,
1155 score: vr.score,
1156 bm25_score: 0.0,
1157 raw_bm25_score: 0.0,
1158 vector_score: vr.score,
1159 importance: 0.0,
1160 graph_score: 0.0,
1161 trust_factor: 1.0,
1162 in_conformal_set: false,
1163 corroboration: 0,
1164 content,
1165 }
1166 })
1167 .collect()
1168 }
1169
1170 #[must_use]
1172 pub fn text_search(&self, query: &str, limit: usize) -> Vec<RecallResult> {
1173 let bm25_results = self.search_engine.search(query, limit).unwrap_or_default();
1174
1175 bm25_results
1176 .into_iter()
1177 .filter_map(|sr| {
1178 let memory_id = Uuid::parse_str(&sr.memory_id).ok()?;
1179 let galaxy = Galaxy::from_db_name(&sr.galaxy)?;
1180 Some(RecallResult {
1181 memory_id,
1182 galaxy,
1183 score: sr.score,
1184 bm25_score: sr.score,
1185 raw_bm25_score: sr.score,
1186 vector_score: 0.0,
1187 importance: 0.0,
1188 graph_score: 0.0,
1189 trust_factor: 1.0,
1190 in_conformal_set: false,
1191 corroboration: 0,
1192 content: sr.content,
1193 })
1194 })
1195 .collect()
1196 }
1197
1198 fn fuse_results(
1202 &self,
1203 bm25_results: &[SearchResult],
1204 vector_results: &[VectorSearchResult],
1205 limit: usize,
1206 ) -> Vec<RecallResult> {
1207 fuse_results_inner(
1208 bm25_results,
1209 vector_results,
1210 limit,
1211 self.config.bm25_weight,
1212 self.config.vector_weight,
1213 self.config.importance_weight,
1214 self.config.trust_weight,
1215 |id, galaxy| self.get_memory_content(id, galaxy),
1216 |id, galaxy| self.get_memory_importance(id, galaxy),
1217 |id, galaxy| self.get_memory_source_trust(id, galaxy),
1218 )
1219 }
1220
1221 fn expand_with_graph(&self, mut results: Vec<RecallResult>, limit: usize) -> Vec<RecallResult> {
1231 if self.config.graph_weight <= 0.0 || results.is_empty() {
1232 return results;
1233 }
1234 let Ok(assoc_store) = AssociationStore::open(self.store.env()) else {
1235 return results;
1236 };
1237 let env = self.store.env();
1238 let seeds: Vec<RecallResult> = results.iter().take(3).cloned().collect();
1239 for seed in seeds {
1240 let outgoing = assoc_store
1241 .find_from(env, seed.memory_id)
1242 .unwrap_or_default();
1243 let incoming = assoc_store.find_to(env, seed.memory_id).unwrap_or_default();
1244 for edge in outgoing.into_iter().chain(incoming) {
1245 if edge.weight < 0.2 {
1246 continue;
1247 }
1248 let neighbor_id = if edge.source == seed.memory_id {
1249 edge.target
1250 } else {
1251 edge.source
1252 };
1253 if neighbor_id == seed.memory_id {
1254 continue;
1255 }
1256 if crate::memory::validity_enforced()
1260 && self
1261 .find_memory_anywhere(neighbor_id)
1262 .is_some_and(|mem| !mem.metadata.validity.is_current())
1263 {
1264 continue;
1265 }
1266 let contribution = seed.score * edge.weight * self.config.graph_weight;
1267 if self.config.promotion_on_read {
1268 let mut activated_edge = edge.clone();
1269 activated_edge.activate();
1270 let _ = assoc_store.put(env, &activated_edge);
1271 }
1272 if let Some(existing) = results.iter_mut().find(|r| r.memory_id == neighbor_id) {
1273 existing.score += contribution;
1274 existing.graph_score += contribution;
1275 } else if let Some(mem) = self.find_memory_anywhere(neighbor_id) {
1276 if mem.metadata.is_private {
1280 continue;
1281 }
1282 if crate::memory::validity_enforced() && !mem.metadata.validity.is_current() {
1283 continue;
1284 }
1285 results.push(RecallResult {
1286 memory_id: neighbor_id,
1287 galaxy: mem.metadata.galaxy,
1288 score: contribution,
1289 bm25_score: 0.0,
1290 raw_bm25_score: 0.0,
1291 vector_score: 0.0,
1292 importance: mem.metadata.importance,
1293 graph_score: contribution,
1294 trust_factor: 1.0,
1295 in_conformal_set: false,
1296 corroboration: 0,
1297 content: mem.content.chars().take(400).collect(),
1298 });
1299 }
1300 }
1301 }
1302 results.sort_by(|a, b| {
1303 b.score
1304 .partial_cmp(&a.score)
1305 .unwrap_or(std::cmp::Ordering::Equal)
1306 });
1307 results.truncate(limit.max(3));
1308 results
1309 }
1310
1311 fn find_memory_anywhere(&self, id: Uuid) -> Option<crate::memory::Memory> {
1313 self.store
1314 .find_across_galaxies(id)
1315 .ok()
1316 .flatten()
1317 .map(|(_, m)| m)
1318 }
1319
1320 pub fn promote_memory(&self, galaxy: Galaxy, id: Uuid) -> Result<bool> {
1323 if let Some(mut mem) = self.store.get(galaxy, id)? {
1324 mem.recall();
1325 self.store.put(galaxy, &mem)?;
1326 Ok(true)
1327 } else {
1328 Ok(false)
1329 }
1330 }
1331
1332 fn get_memory_content(&self, id: Uuid, galaxy: Galaxy) -> String {
1336 self.store
1337 .get(galaxy, id)
1338 .ok()
1339 .flatten()
1340 .map(|m| m.content)
1341 .unwrap_or_default()
1342 }
1343
1344 #[must_use]
1347 pub fn is_private(&self, id: Uuid, galaxy: Galaxy) -> bool {
1348 self.store
1349 .get(galaxy, id)
1350 .ok()
1351 .flatten()
1352 .is_none_or(|m| m.metadata.is_private)
1353 }
1354
1355 fn get_memory_importance(&self, id: Uuid, galaxy: Galaxy) -> f32 {
1357 self.store
1358 .get(galaxy, id)
1359 .ok()
1360 .flatten()
1361 .map_or(0.0, |m| m.metadata.importance)
1362 }
1363
1364 fn get_memory_source_trust(&self, id: Uuid, galaxy: Galaxy) -> f32 {
1368 self.store
1369 .get(galaxy, id)
1370 .ok()
1371 .flatten()
1372 .map_or(0.7, |m| m.metadata.source_trust)
1373 }
1374
1375 #[must_use]
1377 pub fn cache_size(&self) -> usize {
1378 self.embedding_cache.lock().map_or(0, |c| c.len())
1379 }
1380
1381 pub fn clear_cache(&self) {
1383 if let Ok(mut c) = self.embedding_cache.lock() {
1384 c.clear();
1385 }
1386 }
1387
1388 #[must_use]
1390 pub fn vector_count(&self) -> usize {
1391 self.vector_store.lock().map_or(0, |c| c.len())
1392 }
1393}
1394
1395#[allow(clippy::too_many_arguments)]
1399fn fuse_results_inner(
1400 bm25_results: &[SearchResult],
1401 vector_results: &[VectorSearchResult],
1402 limit: usize,
1403 bm25_weight: f32,
1404 vector_weight: f32,
1405 importance_weight: f32,
1406 trust_weight: f32,
1407 mut get_content: impl FnMut(Uuid, Galaxy) -> String,
1408 mut get_importance: impl FnMut(Uuid, Galaxy) -> f32,
1409 mut get_source_trust: impl FnMut(Uuid, Galaxy) -> f32,
1410) -> Vec<RecallResult> {
1411 let max_bm25 = bm25_results
1413 .iter()
1414 .map(|r| r.score)
1415 .fold(0.0_f32, f32::max)
1416 .max(0.001);
1417
1418 let mut bm25_map: HashMap<Uuid, (f32, f32, String, Galaxy)> = HashMap::new();
1420 for sr in bm25_results {
1421 if let Ok(id) = Uuid::parse_str(&sr.memory_id) {
1422 match Galaxy::from_db_name(&sr.galaxy) {
1423 Some(galaxy) => {
1424 let normalized = sr.score / max_bm25;
1425 bm25_map.insert(id, (normalized, sr.score, sr.content.clone(), galaxy));
1426 }
1427 None => {
1428 tracing::warn!(
1429 "Skipping BM25 result with unknown galaxy '{}' (memory_id={})",
1430 sr.galaxy,
1431 sr.memory_id
1432 );
1433 }
1434 }
1435 }
1436 }
1437
1438 let mut vector_map: HashMap<Uuid, (f32, Galaxy)> = HashMap::new();
1439 for vr in vector_results {
1440 vector_map.insert(vr.memory_id, (vr.score, vr.galaxy));
1441 }
1442
1443 let mut all_ids: std::collections::HashSet<Uuid> = std::collections::HashSet::new();
1445 all_ids.extend(bm25_map.keys());
1446 all_ids.extend(vector_map.keys());
1447
1448 let mut results: Vec<RecallResult> = all_ids
1450 .into_iter()
1451 .map(|id| {
1452 let (bm25_score, raw_bm25_score, content, galaxy_bm25) = bm25_map.get(&id).map_or(
1453 (0.0, 0.0, String::new(), Galaxy::Codex),
1454 |(s, raw, c, g)| (*s, *raw, c.clone(), *g),
1455 );
1456
1457 let (vector_score, galaxy_vec) = vector_map
1458 .get(&id)
1459 .map_or((0.0, Galaxy::Codex), |(s, g)| (*s, *g));
1460
1461 let galaxy = if bm25_score > 0.0 {
1462 galaxy_bm25
1463 } else {
1464 galaxy_vec
1465 };
1466
1467 let content = if content.is_empty() {
1468 get_content(id, galaxy)
1469 } else {
1470 content
1471 };
1472
1473 let importance = get_importance(id, galaxy);
1474
1475 let fused = bm25_weight.mul_add(
1476 bm25_score,
1477 vector_weight.mul_add(vector_score, importance_weight * importance),
1478 );
1479
1480 #[allow(clippy::suboptimal_flops)]
1487 let (score, trust_factor) = if trust_weight > 0.0 {
1488 let source_trust = get_source_trust(id, galaxy);
1489 let factor = (1.0 + trust_weight * (source_trust.clamp(0.0, 1.0) - 0.7)).max(0.0);
1490 (fused * factor, factor)
1491 } else {
1492 (fused, 1.0)
1493 };
1494
1495 RecallResult {
1496 memory_id: id,
1497 galaxy,
1498 score,
1499 bm25_score,
1500 raw_bm25_score,
1501 vector_score,
1502 importance,
1503 graph_score: 0.0,
1504 trust_factor,
1505 in_conformal_set: false,
1506 corroboration: 0,
1507 content,
1508 }
1509 })
1510 .collect();
1511
1512 results.sort_by(|a, b| {
1514 b.score
1515 .partial_cmp(&a.score)
1516 .unwrap_or(std::cmp::Ordering::Equal)
1517 });
1518 results.truncate(limit);
1519 results
1520}
1521
1522#[cfg(test)]
1525mod tests {
1526 use super::*;
1527 use crate::associations::{Association, LinkType};
1528 use crate::embedder::StubEmbedder;
1529
1530 struct GraphHarness {
1534 _dir: tempfile::TempDir,
1535 store: Arc<MemoryStore>,
1536 engine_with_graph: RecallEngine,
1537 engine_plain: RecallEngine,
1538 }
1539
1540 fn graph_harness() -> GraphHarness {
1541 let dir = tempfile::tempdir().unwrap();
1542 let lmdb = dir.path().join("lmdb");
1543 std::fs::create_dir_all(&lmdb).unwrap();
1544 let store = Arc::new(MemoryStore::open_default(&lmdb).unwrap());
1545 let tantivy = dir.path().join("tantivy");
1546 std::fs::create_dir_all(&tantivy).unwrap();
1547 let search = Arc::new(SearchEngine::open(&tantivy).unwrap());
1548
1549 let a = Memory::new(Galaxy::Codex, "kumquat governance ratchet".into());
1552 let mut b = Memory::new(Galaxy::Codex, "the follow-up decision".into());
1553 let c = Memory::new(Galaxy::Codex, "kumquat harvest notes".into());
1554 b.metadata.is_private = false;
1555 let (id_a, id_b, id_c) = (a.metadata.id, b.metadata.id, c.metadata.id);
1556 store.put(Galaxy::Codex, &a).unwrap();
1557 store.put(Galaxy::Codex, &b).unwrap();
1558 store.put(Galaxy::Codex, &c).unwrap();
1559
1560 let mut writer = search.writer().unwrap();
1561 for (id, content) in [
1562 (id_a, "kumquat governance ratchet"),
1563 (id_c, "kumquat harvest notes"),
1564 ] {
1565 search
1566 .add_document(
1567 &mut writer,
1568 &id.to_string(),
1569 "codex",
1570 content,
1571 &[],
1572 1_700_000_000,
1573 )
1574 .unwrap();
1575 }
1576 search.commit(&mut writer).unwrap();
1577
1578 let env = store.env();
1579 let assocs = AssociationStore::open(env).unwrap();
1580 assocs
1581 .put(env, &Association::new(id_a, id_b, LinkType::Related, 0.8))
1582 .unwrap();
1583
1584 let store_for_engine = store.clone();
1585 let search_for_engine = search.clone();
1586 let mk_engine = move |graph_weight: f32| {
1587 let config = RecallConfig {
1588 bm25_weight: 1.0,
1589 vector_weight: 0.0,
1590 importance_weight: 0.0,
1591 graph_weight,
1592 ..RecallConfig::default()
1593 };
1594 RecallEngine::new(
1595 store_for_engine.clone(),
1596 search_for_engine.clone(),
1597 VectorStore::new(),
1598 Arc::new(StubEmbedder::default()),
1599 config,
1600 )
1601 .unwrap()
1602 };
1603 GraphHarness {
1604 _dir: dir,
1605 store,
1606 engine_with_graph: mk_engine(0.5),
1607 engine_plain: mk_engine(0.0),
1608 }
1609 }
1610
1611 struct CountingEmbedder {
1614 inner: StubEmbedder,
1615 calls: std::sync::atomic::AtomicUsize,
1616 }
1617
1618 impl CountingEmbedder {
1619 fn new() -> Self {
1620 Self {
1621 inner: StubEmbedder::default(),
1622 calls: std::sync::atomic::AtomicUsize::new(0),
1623 }
1624 }
1625
1626 fn call_count(&self) -> usize {
1627 self.calls.load(std::sync::atomic::Ordering::SeqCst)
1628 }
1629 }
1630
1631 impl Embedder for CountingEmbedder {
1632 fn embed_batch(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
1633 self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
1634 self.inner.embed_batch(texts)
1635 }
1636 fn dimension(&self) -> usize {
1637 self.inner.dimension()
1638 }
1639 fn is_available(&self) -> bool {
1640 true
1641 }
1642 fn backend_name(&self) -> &'static str {
1643 "stub-counting"
1644 }
1645 }
1646
1647 fn engine_fixture() -> (tempfile::TempDir, Arc<MemoryStore>, Arc<SearchEngine>) {
1648 let dir = tempfile::tempdir().unwrap();
1649 let lmdb = dir.path().join("lmdb");
1650 std::fs::create_dir_all(&lmdb).unwrap();
1651 let store = Arc::new(MemoryStore::open_default(&lmdb).unwrap());
1652 let tantivy = dir.path().join("tantivy");
1653 std::fs::create_dir_all(&tantivy).unwrap();
1654 let search = Arc::new(SearchEngine::open(&tantivy).unwrap());
1655 (dir, store, search)
1656 }
1657
1658 fn mk_engine(
1659 store: &Arc<MemoryStore>,
1660 search: &Arc<SearchEngine>,
1661 embedder: Arc<dyn Embedder>,
1662 ) -> RecallEngine {
1663 RecallEngine::new(
1664 store.clone(),
1665 search.clone(),
1666 VectorStore::new(),
1667 embedder,
1668 RecallConfig::default(),
1669 )
1670 .unwrap()
1671 }
1672
1673 #[test]
1674 fn embedding_cache_warm_starts_reingest_across_engine_restart() {
1675 let (_dir, store, search) = engine_fixture();
1679
1680 let contents: Vec<String> = (0..12)
1681 .map(|i| format!("cache warm-start probe number {i} with distinct wording {i}"))
1682 .collect();
1683 let entries: Vec<(Galaxy, crate::Memory)> = contents
1684 .iter()
1685 .map(|c| (Galaxy::Codex, crate::Memory::new(Galaxy::Codex, c.clone())))
1686 .collect();
1687 let refs: Vec<(Galaxy, &crate::Memory)> = entries.iter().map(|(g, m)| (*g, m)).collect();
1688
1689 let first = Arc::new(CountingEmbedder::new());
1690 let engine = mk_engine(&store, &search, first.clone());
1691 assert_eq!(engine.store_batch_with_embedding(&refs).unwrap(), 12);
1692 let first_calls = first.call_count();
1693 assert!(first_calls > 0, "first ingest must embed");
1694 assert_eq!(store.embedding_cache_count().unwrap(), 12);
1695
1696 let second = Arc::new(CountingEmbedder::new());
1699 let engine2 = mk_engine(&store, &search, second.clone());
1700 let entries2: Vec<(Galaxy, crate::Memory)> = contents
1701 .iter()
1702 .map(|c| (Galaxy::Codex, crate::Memory::new(Galaxy::Codex, c.clone())))
1703 .collect();
1704 let refs2: Vec<(Galaxy, &crate::Memory)> = entries2.iter().map(|(g, m)| (*g, m)).collect();
1705 assert_eq!(engine2.store_batch_with_embedding(&refs2).unwrap(), 12);
1706 assert_eq!(
1707 second.call_count(),
1708 0,
1709 "re-ingest of identical content must serve from the persistent cache"
1710 );
1711 assert_eq!(store.embedding_cache_count().unwrap(), 12);
1712 }
1713
1714 #[test]
1715 fn embedding_cache_scopes_vectors_by_embedder_namespace() {
1716 let (_dir, store, search) = engine_fixture();
1719
1720 let content = "namespace isolation probe";
1721 let first = Arc::new(CountingEmbedder::new());
1722 let engine = mk_engine(&store, &search, first.clone());
1723 let mem = crate::Memory::new(Galaxy::Codex, content.into());
1724 engine.store_with_embedding(Galaxy::Codex, &mem).unwrap();
1725 assert_eq!(first.call_count(), 1);
1726
1727 struct OtherNamespaceEmbedder(StubEmbedder);
1730 impl Embedder for OtherNamespaceEmbedder {
1731 fn embed_batch(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
1732 self.0.embed_batch(texts)
1733 }
1734 fn dimension(&self) -> usize {
1735 self.0.dimension()
1736 }
1737 fn is_available(&self) -> bool {
1738 true
1739 }
1740 fn backend_name(&self) -> &'static str {
1741 "stub-other"
1742 }
1743 }
1744 let second = Arc::new(OtherNamespaceEmbedder(StubEmbedder::default()));
1745 let engine2 = mk_engine(&store, &search, second);
1746 let mem2 = crate::Memory::new(Galaxy::Codex, content.into());
1747 engine2.store_with_embedding(Galaxy::Codex, &mem2).unwrap();
1748
1749 assert_eq!(store.embedding_cache_count().unwrap(), 2);
1751 }
1752
1753 #[test]
1754 fn graph_expansion_injects_unindexed_neighbors_and_boosts_connected() {
1755 let h = graph_harness();
1756
1757 let plain = h.engine_plain.hybrid_search("kumquat", 10, None);
1759 assert!(plain.iter().all(|r| r.memory_id != {
1760 h.store
1761 .find_by_content_hash(Galaxy::Codex, &content_hash("the follow-up decision"))
1762 .unwrap()
1763 .unwrap()
1764 }));
1765 assert!(plain.iter().all(|r| r.graph_score == 0.0));
1766
1767 let expanded = h.engine_with_graph.hybrid_search("kumquat", 10, None);
1770 let id_b = h
1771 .store
1772 .find_by_content_hash(Galaxy::Codex, &content_hash("the follow-up decision"))
1773 .unwrap()
1774 .unwrap();
1775 let b = expanded
1776 .iter()
1777 .find(|r| r.memory_id == id_b)
1778 .expect("graph expansion must surface the unindexed neighbor");
1779 assert!(b.graph_score > 0.0, "injected neighbor: {b:?}");
1780 assert_eq!(b.bm25_score, 0.0, "B had no BM25 hit — pure graph entry");
1781 let a_score = expanded
1782 .iter()
1783 .find(|r| r.content.contains("ratchet"))
1784 .unwrap()
1785 .score;
1786 assert!(a_score >= b.score, "seed outranks its 1-hop neighbor");
1787 }
1788
1789 #[test]
1790 fn graph_expansion_honors_the_privacy_flag() {
1791 let h = graph_harness();
1792 let id_b = h
1793 .store
1794 .find_by_content_hash(Galaxy::Codex, &content_hash("the follow-up decision"))
1795 .unwrap()
1796 .unwrap();
1797 let mut b = h.store.get(Galaxy::Codex, id_b).unwrap().unwrap();
1799 b.metadata.is_private = true;
1800 h.store.put(Galaxy::Codex, &b).unwrap();
1801 let expanded = h.engine_with_graph.hybrid_search("kumquat", 10, None);
1802 assert!(
1803 expanded.iter().all(|r| r.memory_id != id_b),
1804 "private memory must not be graph-injected"
1805 );
1806 }
1807
1808 #[test]
1809 fn config_default_graph_weight_is_off() {
1810 let config = RecallConfig::default();
1811 assert_eq!(config.graph_weight, 0.0, "evidence-gated: default off");
1812 assert!(config.weights_normalized());
1813 }
1814
1815 #[test]
1818 fn config_default_weights() {
1819 let config = RecallConfig::default();
1820 assert!(config.weights_normalized());
1821 assert_eq!(config.bm25_weight, 0.5);
1822 assert_eq!(config.vector_weight, 0.3);
1823 assert_eq!(config.importance_weight, 0.2);
1824 }
1825
1826 #[test]
1827 fn config_custom_weights() {
1828 let config = RecallConfig {
1829 bm25_weight: 0.6,
1830 vector_weight: 0.3,
1831 importance_weight: 0.1,
1832 ..Default::default()
1833 };
1834 assert!(config.weights_normalized());
1835 }
1836
1837 #[test]
1838 fn config_unnormalized_weights() {
1839 let config = RecallConfig {
1840 bm25_weight: 0.7,
1841 vector_weight: 0.5,
1842 importance_weight: 0.2,
1843 ..Default::default()
1844 };
1845 assert!(!config.weights_normalized());
1846 }
1847
1848 #[test]
1849 fn config_from_env_uses_defaults() {
1850 let config = RecallConfig::from_env();
1852 assert_eq!(config.bm25_weight, 0.5);
1853 assert_eq!(config.vector_weight, 0.3);
1854 assert_eq!(config.importance_weight, 0.2);
1855 }
1856
1857 #[test]
1861 fn corroboration_boost_is_knob_gated_and_disclosed() {
1862 let dir = tempfile::tempdir().unwrap();
1863 let lmdb = dir.path().join("lmdb");
1864 std::fs::create_dir_all(&lmdb).unwrap();
1865 let store = Arc::new(MemoryStore::open_default(&lmdb).unwrap());
1866 let tantivy = dir.path().join("tantivy");
1867 std::fs::create_dir_all(&tantivy).unwrap();
1868 let search = Arc::new(SearchEngine::open(&tantivy).unwrap());
1869
1870 let mut backed = Memory::new(Galaxy::Codex, "zanzibar treaty terms".into());
1871 backed.metadata.corroborated_by = vec![Uuid::new_v4(), Uuid::new_v4(), Uuid::new_v4()];
1872 let plain = Memory::new(Galaxy::Codex, "zanzibar treaty terms".into());
1873 let (id_backed, id_plain) = (backed.metadata.id, plain.metadata.id);
1874 store.put(Galaxy::Codex, &backed).unwrap();
1875 store.put(Galaxy::Codex, &plain).unwrap();
1876 let mut writer = search.writer().unwrap();
1877 for (id, content) in [
1878 (id_backed, "zanzibar treaty terms"),
1879 (id_plain, "zanzibar treaty terms"),
1880 ] {
1881 search
1882 .add_document(
1883 &mut writer,
1884 &id.to_string(),
1885 "codex",
1886 content,
1887 &[],
1888 1_700_000_000,
1889 )
1890 .unwrap();
1891 }
1892 search.commit(&mut writer).unwrap();
1893
1894 let mk = |weight: f32| {
1895 RecallEngine::new(
1896 store.clone(),
1897 search.clone(),
1898 VectorStore::new(),
1899 Arc::new(StubEmbedder::default()),
1900 RecallConfig {
1901 bm25_weight: 1.0,
1902 vector_weight: 0.0,
1903 importance_weight: 0.0,
1904 corroboration_weight: weight,
1905 ..RecallConfig::default()
1906 },
1907 )
1908 .unwrap()
1909 };
1910 let off = mk(0.0).hybrid_search("zanzibar treaty", 10, None);
1912 let (b_off, p_off) = (
1913 off.iter().find(|r| r.memory_id == id_backed).unwrap(),
1914 off.iter().find(|r| r.memory_id == id_plain).unwrap(),
1915 );
1916 assert!((b_off.score - p_off.score).abs() < 1e-6);
1917 assert_eq!(b_off.corroboration, 0);
1918 let on = mk(1.0).hybrid_search("zanzibar treaty", 10, None);
1920 let (b_on, p_on) = (
1921 on.iter().find(|r| r.memory_id == id_backed).unwrap(),
1922 on.iter().find(|r| r.memory_id == id_plain).unwrap(),
1923 );
1924 assert_eq!(b_on.corroboration, 3);
1925 assert_eq!(p_on.corroboration, 0);
1926 let expected = b_off.score * 1.6;
1927 assert!((b_on.score - expected).abs() < 1e-4, "{b_on:?}");
1928 assert!((p_on.score - p_off.score).abs() < 1e-6);
1929 assert!(on[0].memory_id == id_backed, "boosted memory ranks first");
1930 }
1931
1932 #[test]
1934 fn recall_result_fields() {
1935 let result = RecallResult {
1936 memory_id: Uuid::new_v4(),
1937 galaxy: Galaxy::Codex,
1938 score: 0.85,
1939 bm25_score: 0.7,
1940 raw_bm25_score: 3.5,
1941 vector_score: 0.9,
1942 importance: 0.5,
1943 graph_score: 0.0,
1944 trust_factor: 1.0,
1945 in_conformal_set: false,
1946 corroboration: 0,
1947 content: "test content".into(),
1948 };
1949 assert_eq!(result.score, 0.85);
1950 assert_eq!(result.bm25_score, 0.7);
1951 assert_eq!(result.raw_bm25_score, 3.5);
1952 assert_eq!(result.vector_score, 0.9);
1953 }
1954
1955 fn fuse(
1958 bm25: &[SearchResult],
1959 vector: &[VectorSearchResult],
1960 limit: usize,
1961 ) -> Vec<RecallResult> {
1962 fuse_results_inner(
1963 bm25,
1964 vector,
1965 limit,
1966 0.5,
1967 0.3,
1968 0.2,
1969 0.0,
1970 |_, _| String::new(),
1971 |_, _| 0.0,
1972 |_, _| 0.7,
1973 )
1974 }
1975
1976 #[test]
1977 fn engine_config_default() {
1978 let config = RecallConfig::default();
1979 assert_eq!(config.bm25_weight, 0.5);
1980 }
1981
1982 #[test]
1983 fn engine_cache_concept() {
1984 let config = RecallConfig::default();
1986 assert!(config.cache_embeddings);
1987 }
1988
1989 #[test]
1992 fn fuse_results_empty() {
1993 let results = fuse(&[], &[], 10);
1994 assert!(results.is_empty());
1995 }
1996
1997 #[test]
1998 fn fuse_results_bm25_only() {
1999 let id = Uuid::new_v4();
2000 let bm25 = vec![SearchResult {
2001 memory_id: id.to_string(),
2002 galaxy: Galaxy::Codex.db_name().to_string(),
2003 score: 5.0,
2004 normalized_score: 0.0,
2005 content: "test".into(),
2006 }];
2007 let results = fuse(&bm25, &[], 10);
2008 assert_eq!(results.len(), 1);
2009 assert!(results[0].bm25_score > 0.0);
2010 assert_eq!(results[0].vector_score, 0.0);
2011 }
2012
2013 #[test]
2014 fn fuse_results_vector_only() {
2015 let id = Uuid::new_v4();
2016 let vector = vec![VectorSearchResult {
2017 memory_id: id,
2018 galaxy: Galaxy::Codex,
2019 score: 0.85,
2020 }];
2021 let results = fuse(&[], &vector, 10);
2022 assert_eq!(results.len(), 1);
2023 assert_eq!(results[0].bm25_score, 0.0);
2024 assert!(results[0].vector_score > 0.0);
2025 }
2026
2027 #[test]
2028 fn fuse_results_both_sources() {
2029 let id = Uuid::new_v4();
2030 let bm25 = vec![SearchResult {
2031 memory_id: id.to_string(),
2032 galaxy: Galaxy::Codex.db_name().to_string(),
2033 score: 5.0,
2034 normalized_score: 0.0,
2035 content: "test content".into(),
2036 }];
2037 let vector = vec![VectorSearchResult {
2038 memory_id: id,
2039 galaxy: Galaxy::Codex,
2040 score: 0.85,
2041 }];
2042 let results = fuse(&bm25, &vector, 10);
2043 assert_eq!(results.len(), 1);
2044 assert!(results[0].bm25_score > 0.0);
2045 assert!(results[0].vector_score > 0.0);
2046 assert!(results[0].score > results[0].bm25_score * 0.5);
2047 }
2048
2049 #[test]
2050 fn fuse_results_sorted_by_score() {
2051 let id1 = Uuid::new_v4();
2052 let id2 = Uuid::new_v4();
2053 let bm25 = vec![
2054 SearchResult {
2055 memory_id: id1.to_string(),
2056 galaxy: Galaxy::Codex.db_name().to_string(),
2057 score: 3.0,
2058 normalized_score: 0.0,
2059 content: "lower".into(),
2060 },
2061 SearchResult {
2062 memory_id: id2.to_string(),
2063 galaxy: Galaxy::Codex.db_name().to_string(),
2064 score: 8.0,
2065 normalized_score: 0.0,
2066 content: "higher".into(),
2067 },
2068 ];
2069 let results = fuse(&bm25, &[], 10);
2070 assert_eq!(results.len(), 2);
2071 assert!(results[0].score >= results[1].score);
2072 }
2073
2074 #[test]
2075 fn fuse_results_truncated_to_limit() {
2076 let bm25: Vec<SearchResult> = (0..20)
2077 .map(|i| SearchResult {
2078 memory_id: Uuid::new_v4().to_string(),
2079 galaxy: Galaxy::Codex.db_name().to_string(),
2080 score: 1.0 + i as f32,
2081 normalized_score: 0.0,
2082 content: format!("content {i}"),
2083 })
2084 .collect();
2085 let results = fuse(&bm25, &[], 5);
2086 assert_eq!(results.len(), 5);
2087 }
2088
2089 #[test]
2090 fn fuse_results_normalizes_bm25() {
2091 let id = Uuid::new_v4();
2092 let bm25 = vec![SearchResult {
2093 memory_id: id.to_string(),
2094 galaxy: Galaxy::Codex.db_name().to_string(),
2095 score: 100.0,
2096 normalized_score: 0.0,
2097 content: "test".into(),
2098 }];
2099 let results = fuse(&bm25, &[], 10);
2100 assert!((results[0].bm25_score - 1.0).abs() < 0.01);
2101 }
2102
2103 #[test]
2104 fn fuse_results_retains_raw_bm25_for_absolute_floors() {
2105 let id = Uuid::new_v4();
2106 let bm25 = vec![SearchResult {
2107 memory_id: id.to_string(),
2108 galaxy: Galaxy::Codex.db_name().to_string(),
2109 score: 3.5,
2110 normalized_score: 0.0,
2111 content: "test".into(),
2112 }];
2113 let results = fuse(&bm25, &[], 10);
2114 assert!((results[0].bm25_score - 1.0).abs() < 0.01);
2116 assert!((results[0].raw_bm25_score - 3.5).abs() < f32::EPSILON);
2117 }
2118
2119 #[test]
2122 fn embed_content_caches_result() {
2123 let embedder = StubEmbedder::new(384);
2124 let content = "test content for caching";
2125 let vec1 = embedder.embed(content).unwrap();
2126 let vec2 = embedder.embed(content).unwrap();
2127 assert_eq!(vec1, vec2);
2128 }
2129
2130 #[test]
2131 fn embed_content_different_content_different_result() {
2132 let embedder = StubEmbedder::new(384);
2133 let vec1 = embedder.embed("content one").unwrap();
2134 let vec2 = embedder.embed("content two").unwrap();
2135 assert_ne!(vec1, vec2);
2136 }
2137
2138 #[test]
2141 fn fuse_with_zero_bm25_weight() {
2142 let id = Uuid::new_v4();
2143 let bm25 = vec![SearchResult {
2144 memory_id: id.to_string(),
2145 galaxy: Galaxy::Codex.db_name().to_string(),
2146 score: 5.0,
2147 normalized_score: 0.0,
2148 content: "test".into(),
2149 }];
2150 let results = fuse_results_inner(
2151 &bm25,
2152 &[],
2153 10,
2154 0.5,
2155 0.3,
2156 0.2,
2157 0.0,
2158 |_, _| String::new(),
2159 |_, _| 0.0,
2160 |_, _| 0.7,
2161 );
2162 assert!((results[0].score - 0.5).abs() < 0.01);
2163 }
2164
2165 #[test]
2166 fn fuse_with_zero_vector_weight() {
2167 let id = Uuid::new_v4();
2168 let vector = vec![VectorSearchResult {
2169 memory_id: id,
2170 galaxy: Galaxy::Codex,
2171 score: 0.9,
2172 }];
2173 let results = fuse_results_inner(
2174 &[],
2175 &vector,
2176 10,
2177 0.5,
2178 0.3,
2179 0.2,
2180 0.0,
2181 |_, _| String::new(),
2182 |_, _| 0.0,
2183 |_, _| 0.7,
2184 );
2185 assert!((results[0].score - 0.27).abs() < 0.01);
2186 }
2187
2188 #[test]
2189 fn trust_weight_zero_is_byte_identical_to_no_weight() {
2190 let id = Uuid::new_v4();
2191 let bm25 = vec![SearchResult {
2192 memory_id: id.to_string(),
2193 galaxy: Galaxy::Codex.db_name().to_string(),
2194 score: 5.0,
2195 normalized_score: 0.0,
2196 content: "test".into(),
2197 }];
2198 let results = fuse_results_inner(
2201 &bm25,
2202 &[],
2203 10,
2204 0.5,
2205 0.3,
2206 0.2,
2207 0.0,
2208 |_, _| String::new(),
2209 |_, _| 0.0,
2210 |_, _| 0.4,
2211 );
2212 assert!((results[0].score - 0.5).abs() < 0.01);
2213 assert!((results[0].trust_factor - 1.0).abs() < f32::EPSILON);
2214 assert!(!results[0].in_conformal_set);
2215 }
2216
2217 #[test]
2218 fn trust_weight_orders_high_trust_above_low() {
2219 let high = Uuid::new_v4();
2222 let low = Uuid::new_v4();
2223 let mk = |id: &Uuid| SearchResult {
2224 memory_id: id.to_string(),
2225 galaxy: Galaxy::Codex.db_name().to_string(),
2226 score: 5.0,
2227 normalized_score: 0.0,
2228 content: "test".into(),
2229 };
2230 let bm25 = vec![mk(&high), mk(&low)];
2231 let mut trust_calls = 0;
2232 let results = fuse_results_inner(
2233 &bm25,
2234 &[],
2235 10,
2236 0.5,
2237 0.3,
2238 0.2,
2239 0.5,
2240 |_, _| String::new(),
2241 |_, _| 0.0,
2242 |id, _| {
2243 trust_calls += 1;
2244 if id == high { 1.0 } else { 0.4 }
2245 },
2246 );
2247 let hi = results.iter().find(|r| r.memory_id == high).unwrap();
2248 let lo = results.iter().find(|r| r.memory_id == low).unwrap();
2249 assert!(
2250 hi.score > lo.score,
2251 "user-confirmed (1.0) must outrank low-trust (0.4) at equal fused score"
2252 );
2253 assert!((hi.trust_factor - 1.15).abs() < 0.001);
2255 assert!((lo.trust_factor - 0.85).abs() < 0.001);
2256 assert!(trust_calls >= 2, "getter consulted per candidate");
2257 let neutral = Uuid::new_v4();
2259 let bm25_neutral = vec![SearchResult {
2260 memory_id: neutral.to_string(),
2261 galaxy: Galaxy::Codex.db_name().to_string(),
2262 score: 5.0,
2263 normalized_score: 0.0,
2264 content: "test".into(),
2265 }];
2266 let res_n = fuse_results_inner(
2267 &bm25_neutral,
2268 &[],
2269 10,
2270 0.5,
2271 0.3,
2272 0.2,
2273 0.5,
2274 |_, _| String::new(),
2275 |_, _| 0.0,
2276 |_, _| 0.7,
2277 );
2278 assert!((res_n[0].trust_factor - 1.0).abs() < 0.001);
2279 }
2280
2281 #[test]
2282 fn config_defaults_keep_both_s8_knobs_off() {
2283 let cfg = RecallConfig::default();
2288 assert_eq!(cfg.trust_weight, 0.0);
2289 assert_eq!(cfg.conformal_alpha, None);
2290 let env_cfg = RecallConfig::from_env();
2291 assert_eq!(env_cfg.trust_weight, 0.0, "unset env stays off");
2292 assert_eq!(env_cfg.conformal_alpha, None, "unset env stays off");
2293 }
2294
2295 #[test]
2298 fn galaxy_from_db_name_valid() {
2299 assert_eq!(Galaxy::from_db_name("codex"), Some(Galaxy::Codex));
2300 }
2301
2302 #[test]
2303 fn galaxy_from_db_name_invalid() {
2304 assert_eq!(Galaxy::from_db_name("nonexistent"), None);
2305 }
2306
2307 use crate::Memory;
2310 use tempfile::tempdir;
2311
2312 fn setup_engine() -> (tempfile::TempDir, RecallEngine) {
2313 let tmp = tempdir().unwrap();
2314 let store = Arc::new(MemoryStore::open_default(tmp.path()).unwrap());
2315 let tantivy_path = tmp.path().join("tantivy");
2316 std::fs::create_dir_all(&tantivy_path).unwrap();
2317 let search = Arc::new(SearchEngine::open(&tantivy_path).unwrap());
2318 let vector_store = VectorStore::new();
2319 let embedder: Arc<dyn Embedder> = Arc::new(StubEmbedder::new(384));
2320 let engine = RecallEngine::new(
2321 store,
2322 search,
2323 vector_store,
2324 embedder,
2325 RecallConfig::default(),
2326 )
2327 .unwrap();
2328 (tmp, engine)
2329 }
2330
2331 #[test]
2332 fn integration_store_and_hybrid_search_roundtrip() {
2333 let (_tmp, engine) = setup_engine();
2334
2335 let mem1 = Memory::new(
2336 Galaxy::Codex,
2337 "Rust programming language is fast and safe".into(),
2338 )
2339 .with_importance(0.8)
2340 .with_tags(vec!["rust".into(), "programming".into()]);
2341 let mem2 = Memory::new(Galaxy::Codex, "Python is great for data science".into())
2342 .with_importance(0.5)
2343 .with_tags(vec!["python".into(), "data".into()]);
2344 let mem3 = Memory::new(
2345 Galaxy::Codex,
2346 "The Rust ownership model prevents memory leaks".into(),
2347 )
2348 .with_importance(0.9)
2349 .with_tags(vec!["rust".into(), "memory".into()]);
2350
2351 engine.store_with_embedding(Galaxy::Codex, &mem1).unwrap();
2352 engine.store_with_embedding(Galaxy::Codex, &mem2).unwrap();
2353 engine.store_with_embedding(Galaxy::Codex, &mem3).unwrap();
2354
2355 let results = engine.hybrid_search("rust", 10, None);
2357 assert!(!results.is_empty(), "hybrid search should return results");
2358
2359 let top_contents: Vec<&str> = results.iter().map(|r| r.content.as_str()).collect();
2361 assert!(
2362 top_contents.iter().any(|c| c.contains("Rust")),
2363 "top results should include Rust content, got: {top_contents:?}"
2364 );
2365 }
2366
2367 #[test]
2368 fn integration_bm25_and_vector_both_contribute() {
2369 let (_tmp, engine) = setup_engine();
2370
2371 for i in 0..5 {
2373 let mem = Memory::new(
2374 Galaxy::Codex,
2375 format!("memory about topic {i} with unique content"),
2376 )
2377 .with_importance(0.5);
2378 engine.store_with_embedding(Galaxy::Codex, &mem).unwrap();
2379 }
2380
2381 let results = engine.hybrid_search("memory", 10, None);
2383 assert!(!results.is_empty(), "should find memories");
2384
2385 let has_bm25 = results.iter().any(|r| r.bm25_score > 0.0);
2387 assert!(has_bm25, "BM25 should contribute to fused results");
2388 }
2389
2390 #[test]
2391 fn integration_vector_search_only() {
2392 let (_tmp, engine) = setup_engine();
2393
2394 let content = "unique searchable content for vector test";
2395 let mem = Memory::new(Galaxy::Codex, content.into()).with_importance(0.7);
2396 engine.store_with_embedding(Galaxy::Codex, &mem).unwrap();
2397
2398 let results = engine.vector_search(content, 10, None);
2400 assert_eq!(results.len(), 1);
2401 assert_eq!(results[0].memory_id, mem.metadata.id);
2402 assert!(results[0].vector_score > 0.0);
2403 }
2404
2405 #[test]
2406 fn integration_text_search_only() {
2407 let (_tmp, engine) = setup_engine();
2408
2409 let mem = Memory::new(Galaxy::Codex, "specific text about rust ownership".into())
2410 .with_tags(vec!["rust".into()]);
2411 engine.store_with_embedding(Galaxy::Codex, &mem).unwrap();
2412
2413 let results = engine.text_search("rust", 10);
2414 assert!(!results.is_empty(), "text search should find results");
2415 assert!(results.iter().any(|r| r.bm25_score > 0.0));
2416 }
2417
2418 #[test]
2419 fn integration_batch_store_with_embedding() {
2420 let (_tmp, engine) = setup_engine();
2421
2422 let mem1 = Memory::new(Galaxy::Codex, "alpha beta gamma".into());
2423 let mem2 = Memory::new(Galaxy::Codex, "delta epsilon zeta".into());
2424 let mem3 = Memory::new(Galaxy::Codex, "eta theta iota".into());
2425
2426 let entries = vec![
2427 (Galaxy::Codex, &mem1),
2428 (Galaxy::Codex, &mem2),
2429 (Galaxy::Codex, &mem3),
2430 ];
2431
2432 let count = engine.store_batch_with_embedding(&entries).unwrap();
2433 assert_eq!(count, 3);
2434
2435 let results = engine.text_search("alpha", 10);
2437 assert!(
2438 !results.is_empty(),
2439 "batch-stored memory should be searchable"
2440 );
2441
2442 let vresults = engine.vector_search("alpha beta gamma", 10, None);
2444 assert_eq!(
2445 vresults.len(),
2446 1,
2447 "vector search should find the exact match"
2448 );
2449 assert_eq!(vresults[0].memory_id, mem1.metadata.id);
2450 }
2451
2452 #[test]
2453 fn integration_batch_store_empty() {
2454 let (_tmp, engine) = setup_engine();
2455 let entries: Vec<(Galaxy, &Memory)> = vec![];
2456 let count = engine.store_batch_with_embedding(&entries).unwrap();
2457 assert_eq!(count, 0);
2458 }
2459
2460 #[test]
2461 fn integration_galaxy_filter() {
2462 let (_tmp, engine) = setup_engine();
2463
2464 let mem_codex = Memory::new(Galaxy::Codex, "codex memory about rust".into());
2465 let mem_research = Memory::new(Galaxy::Research, "research memory about rust".into());
2466
2467 engine
2468 .store_with_embedding(Galaxy::Codex, &mem_codex)
2469 .unwrap();
2470 engine
2471 .store_with_embedding(Galaxy::Research, &mem_research)
2472 .unwrap();
2473
2474 let results = engine.hybrid_search("rust", 10, Some(Galaxy::Codex));
2475 assert!(!results.is_empty());
2476 assert!(
2477 results.iter().all(|r| r.galaxy == Galaxy::Codex),
2478 "all results should be from Codex galaxy"
2479 );
2480 }
2481
2482 #[test]
2483 fn integration_empty_search() {
2484 let (_tmp, engine) = setup_engine();
2485 let results = engine.hybrid_search("nonexistent", 10, None);
2486 assert!(results.is_empty());
2487 }
2488
2489 #[test]
2490 fn integration_cache_populated_after_store() {
2491 let (_tmp, engine) = setup_engine();
2492
2493 let mem = Memory::new(Galaxy::Codex, "content to be cached".into());
2494 engine.store_with_embedding(Galaxy::Codex, &mem).unwrap();
2495
2496 assert_eq!(engine.cache_size(), 1);
2498 }
2499
2500 #[test]
2501 fn integration_vector_count_tracks_stores() {
2502 let (_tmp, engine) = setup_engine();
2503
2504 assert_eq!(engine.vector_count(), 0);
2505
2506 for i in 0..3 {
2507 let mem = Memory::new(Galaxy::Codex, format!("memory {i}"));
2508 engine.store_with_embedding(Galaxy::Codex, &mem).unwrap();
2509 }
2510
2511 assert_eq!(engine.vector_count(), 3);
2512 }
2513
2514 #[test]
2515 fn integration_importance_affects_ranking() {
2516 let (_tmp, engine) = setup_engine();
2517
2518 let mem_low =
2520 Memory::new(Galaxy::Codex, "rust programming basics".into()).with_importance(0.1);
2521 let mem_high =
2522 Memory::new(Galaxy::Codex, "rust programming advanced".into()).with_importance(0.9);
2523
2524 engine
2525 .store_with_embedding(Galaxy::Codex, &mem_low)
2526 .unwrap();
2527 engine
2528 .store_with_embedding(Galaxy::Codex, &mem_high)
2529 .unwrap();
2530
2531 let results = engine.hybrid_search("rust", 10, None);
2532 assert_eq!(results.len(), 2);
2533
2534 let high_idx = results
2537 .iter()
2538 .position(|r| r.memory_id == mem_high.metadata.id)
2539 .unwrap();
2540 let low_idx = results
2541 .iter()
2542 .position(|r| r.memory_id == mem_low.metadata.id)
2543 .unwrap();
2544 assert!(
2545 high_idx < low_idx,
2546 "higher importance memory should rank higher"
2547 );
2548 }
2549
2550 #[test]
2551 fn config_from_env_rejects_nan_weights() {
2552 let mut config = RecallConfig::default();
2555 let w: f32 = "NaN".parse().unwrap();
2556 if w.is_finite() && w >= 0.0 {
2557 config.bm25_weight = w.min(1.0);
2558 }
2559 assert_eq!(
2560 config.bm25_weight, 0.5,
2561 "NaN should be rejected, default kept"
2562 );
2563 }
2564
2565 #[test]
2566 fn config_from_env_rejects_negative_weights() {
2567 let mut config = RecallConfig::default();
2568 let w: f32 = "-0.5".parse().unwrap();
2569 if w.is_finite() && w >= 0.0 {
2570 config.vector_weight = w.min(1.0);
2571 }
2572 assert_eq!(
2573 config.vector_weight, 0.3,
2574 "Negative should be rejected, default kept"
2575 );
2576 }
2577
2578 #[test]
2579 fn config_from_env_clamps_weights_to_1() {
2580 let mut config = RecallConfig::default();
2581 let w: f32 = "5.0".parse().unwrap();
2582 if w.is_finite() && w >= 0.0 {
2583 config.importance_weight = w.min(1.0);
2584 }
2585 assert_eq!(
2586 config.importance_weight, 1.0,
2587 "Weight should be clamped to 1.0"
2588 );
2589 }
2590
2591 #[test]
2592 fn config_from_env_normalizes_weights() {
2593 let mut config = RecallConfig {
2594 bm25_weight: 0.8,
2595 vector_weight: 0.8,
2596 importance_weight: 0.8,
2597 ..Default::default()
2598 };
2599 let sum = config.bm25_weight + config.vector_weight + config.importance_weight;
2600 if sum > 0.0 && (sum - 1.0).abs() > 0.01 {
2601 config.bm25_weight /= sum;
2602 config.vector_weight /= sum;
2603 config.importance_weight /= sum;
2604 }
2605 assert!(
2606 config.weights_normalized(),
2607 "Weights should be normalized to sum to 1.0"
2608 );
2609 }
2610
2611 #[test]
2612 fn config_from_env_rejects_infinity() {
2613 let mut config = RecallConfig::default();
2614 let w: f32 = "inf".parse().unwrap();
2615 if w.is_finite() && w >= 0.0 {
2616 config.bm25_weight = w.min(1.0);
2617 }
2618 assert_eq!(
2619 config.bm25_weight, 0.5,
2620 "Infinity should be rejected, default kept"
2621 );
2622 }
2623
2624 #[test]
2625 fn test_promotion_on_read_config_default() {
2626 let default_config = RecallConfig::default();
2627 assert!(!default_config.promotion_on_read);
2628
2629 let custom_config = RecallConfig {
2630 promotion_on_read: true,
2631 ..Default::default()
2632 };
2633 assert!(custom_config.promotion_on_read);
2634 }
2635
2636 #[test]
2637 fn test_promote_memory_updates_hebbian_score_and_counts() {
2638 let tmp = tempfile::tempdir().unwrap();
2639 let store_dir = tmp.path().join("store");
2640 std::fs::create_dir_all(&store_dir).unwrap();
2641 let store = Arc::new(MemoryStore::open_default(&store_dir).unwrap());
2642 let index_dir = tmp.path().join("index");
2643 std::fs::create_dir_all(&index_dir).unwrap();
2644 let search_engine = Arc::new(SearchEngine::open(&index_dir).unwrap());
2645 let vector_store = VectorStore::new();
2646 let embedder = Arc::new(crate::embedder::StubEmbedder::default());
2647 let config = RecallConfig {
2648 promotion_on_read: true,
2649 ..Default::default()
2650 };
2651 let engine =
2652 RecallEngine::new(store.clone(), search_engine, vector_store, embedder, config)
2653 .unwrap();
2654
2655 let mut mem = crate::Memory::new(Galaxy::Codex, "promotion on read test".to_string());
2656 mem.metadata.neuro_score = 0.5;
2657 mem.metadata.novelty_score = 1.0;
2658 let mem_id = mem.metadata.id;
2659 store.put(Galaxy::Codex, &mem).unwrap();
2660
2661 let promoted = engine.promote_memory(Galaxy::Codex, mem_id).unwrap();
2663 assert!(promoted);
2664
2665 let reloaded = store.get(Galaxy::Codex, mem_id).unwrap().unwrap();
2666 assert_eq!(reloaded.metadata.recall_count, 1);
2667 assert_eq!(reloaded.metadata.access_count, 1);
2668 assert!(
2669 reloaded.metadata.neuro_score > 0.5,
2670 "neuro_score should increase via Hebbian boost"
2671 );
2672 assert!(
2673 reloaded.metadata.novelty_score < 1.0,
2674 "novelty_score should decay on recall"
2675 );
2676 }
2677
2678 #[test]
2679 fn test_hybrid_search_triggers_promotion_on_read() {
2680 let tmp = tempfile::tempdir().unwrap();
2681 let store_dir = tmp.path().join("store");
2682 std::fs::create_dir_all(&store_dir).unwrap();
2683 let store = Arc::new(MemoryStore::open_default(&store_dir).unwrap());
2684 let index_dir = tmp.path().join("index");
2685 std::fs::create_dir_all(&index_dir).unwrap();
2686 let search_engine = Arc::new(SearchEngine::open(&index_dir).unwrap());
2687 let vector_store = VectorStore::new();
2688 let embedder = Arc::new(crate::embedder::StubEmbedder::default());
2689 let config = RecallConfig {
2690 promotion_on_read: true,
2691 ..Default::default()
2692 };
2693 let engine = RecallEngine::new(
2694 store.clone(),
2695 search_engine.clone(),
2696 vector_store,
2697 embedder,
2698 config,
2699 )
2700 .unwrap();
2701
2702 let mut mem = crate::Memory::new(Galaxy::Codex, "tokio army swarm tactics".to_string());
2703 mem.metadata.neuro_score = 0.5;
2704 mem.metadata.novelty_score = 1.0;
2705 let mem_id = mem.metadata.id;
2706 store.put(Galaxy::Codex, &mem).unwrap();
2707
2708 let mut writer = search_engine.writer().unwrap();
2709 search_engine
2710 .add_document(
2711 &mut writer,
2712 &mem_id.to_string(),
2713 "codex",
2714 "tokio army swarm tactics",
2715 &[],
2716 1_700_000_000,
2717 )
2718 .unwrap();
2719 search_engine.commit(&mut writer).unwrap();
2720
2721 let (results, _) =
2723 engine.hybrid_search_with_disclosure("tokio army", 5, Some(Galaxy::Codex));
2724 assert!(!results.is_empty());
2725 assert_eq!(results[0].memory_id, mem_id);
2726
2727 let reloaded = store.get(Galaxy::Codex, mem_id).unwrap().unwrap();
2728 assert_eq!(reloaded.metadata.recall_count, 1);
2729 assert!(reloaded.metadata.neuro_score > 0.5);
2730 }
2731
2732 #[test]
2733 fn hybrid_search_rehydrates_vectors_across_restart() {
2734 let tmp = tempfile::tempdir().unwrap();
2735 let store_dir = tmp.path().join("store");
2736 std::fs::create_dir_all(&store_dir).unwrap();
2737 let store = Arc::new(MemoryStore::open_default(&store_dir).unwrap());
2738 let index_dir = tmp.path().join("index");
2739 std::fs::create_dir_all(&index_dir).unwrap();
2740 let search_engine = Arc::new(SearchEngine::open(&index_dir).unwrap());
2741 let embedder = Arc::new(crate::embedder::StubEmbedder::default());
2742 let dim = embedder.dimension();
2743
2744 let mem = crate::Memory::new(Galaxy::Codex, "persisted vector canary".to_string());
2746 let mem_id = mem.metadata.id;
2747 store.put(Galaxy::Codex, &mem).unwrap();
2748 store.put_embedding(mem_id, &vec![0.5_f32; dim]).unwrap();
2749
2750 let engine = RecallEngine::new(
2752 store,
2753 search_engine,
2754 VectorStore::new(),
2755 embedder,
2756 RecallConfig::default(),
2757 )
2758 .unwrap();
2759 assert!(!engine.vector_store.lock().unwrap().is_loaded());
2760
2761 let _ = engine.hybrid_search_with_disclosure("rehydration probe", 5, None);
2764
2765 let vs = engine.vector_store.lock().unwrap();
2766 assert!(
2767 vs.is_loaded(),
2768 "vector store should be loaded after the first hybrid search"
2769 );
2770 assert_eq!(vs.len(), 1, "persisted embedding should be indexed");
2771 }
2772
2773 #[test]
2774 fn bm25_weighted_search_skips_the_query_embed() {
2775 struct FailingEmbedder;
2782 impl crate::embedder::Embedder for FailingEmbedder {
2783 fn embed_batch(&self, _texts: &[&str]) -> Result<Vec<Vec<f32>>> {
2784 Err(CoreError::Memory("embedder offline".into()))
2785 }
2786 fn dimension(&self) -> usize {
2787 16
2788 }
2789 fn is_available(&self) -> bool {
2790 false
2791 }
2792 fn backend_name(&self) -> &'static str {
2793 "failing-test"
2794 }
2795 }
2796
2797 let tmp = tempfile::tempdir().unwrap();
2798 let store_dir = tmp.path().join("store");
2799 std::fs::create_dir_all(&store_dir).unwrap();
2800 let store = Arc::new(MemoryStore::open_default(&store_dir).unwrap());
2801 let index_dir = tmp.path().join("index");
2802 std::fs::create_dir_all(&index_dir).unwrap();
2803 let search_engine = Arc::new(SearchEngine::open(&index_dir).unwrap());
2804
2805 let config = RecallConfig {
2806 bm25_weight: 1.0,
2807 vector_weight: 0.0,
2808 importance_weight: 0.0,
2809 ..RecallConfig::default()
2810 };
2811 let engine = RecallEngine::new(
2812 store.clone(),
2813 search_engine.clone(),
2814 VectorStore::new(),
2815 Arc::new(FailingEmbedder),
2816 config,
2817 )
2818 .unwrap();
2819
2820 let mem = crate::Memory::new(
2821 Galaxy::Codex,
2822 "kotlin coroutine budget meeting notes".to_string(),
2823 );
2824 let mem_id = mem.metadata.id;
2825 store.put(Galaxy::Codex, &mem).unwrap();
2826 let mut writer = search_engine.writer().unwrap();
2827 search_engine
2828 .add_document(
2829 &mut writer,
2830 &mem_id.to_string(),
2831 "codex",
2832 "kotlin coroutine budget meeting notes",
2833 &[],
2834 1_700_000_000,
2835 )
2836 .unwrap();
2837 search_engine.commit(&mut writer).unwrap();
2838
2839 let (results, _) =
2840 engine.hybrid_search_with_disclosure("kotlin coroutine budget", 5, Some(Galaxy::Codex));
2841 assert_eq!(
2842 results.len(),
2843 1,
2844 "the BM25 half must answer without the embedder"
2845 );
2846 assert_eq!(results[0].memory_id, mem_id);
2847 assert!(
2848 !engine.vector_store.lock().unwrap().is_loaded(),
2849 "an inert vector half must not rehydrate the vector index"
2850 );
2851 }
2852
2853 #[test]
2854 fn backfill_embeddings_dry_run_then_apply() {
2855 struct TestEmbedder;
2856 impl crate::embedder::Embedder for TestEmbedder {
2857 fn embed_batch(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
2858 Ok(texts.iter().map(|_| vec![0.25_f32; 8]).collect())
2859 }
2860 fn dimension(&self) -> usize {
2861 8
2862 }
2863 fn is_available(&self) -> bool {
2864 true
2865 }
2866 fn backend_name(&self) -> &'static str {
2867 "test"
2868 }
2869 }
2870
2871 let tmp = tempfile::tempdir().unwrap();
2872 let store_dir = tmp.path().join("store");
2873 std::fs::create_dir_all(&store_dir).unwrap();
2874 let store = Arc::new(MemoryStore::open_default(&store_dir).unwrap());
2875 let index_dir = tmp.path().join("index");
2876 std::fs::create_dir_all(&index_dir).unwrap();
2877 let search_engine = Arc::new(SearchEngine::open(&index_dir).unwrap());
2878
2879 let mem_a = crate::Memory::new(Galaxy::Codex, "alpha unique content".to_string());
2880 let mem_b = crate::Memory::new(Galaxy::Codex, "beta unique content".to_string());
2881 let (id_a, id_b) = (mem_a.metadata.id, mem_b.metadata.id);
2882 store.put(Galaxy::Codex, &mem_a).unwrap();
2883 store.put(Galaxy::Codex, &mem_b).unwrap();
2884
2885 let engine = RecallEngine::new(
2886 store.clone(),
2887 search_engine,
2888 VectorStore::new(),
2889 Arc::new(TestEmbedder),
2890 RecallConfig::default(),
2891 )
2892 .unwrap();
2893
2894 let plan = engine
2896 .backfill_embeddings(Some(Galaxy::Codex), 0, true)
2897 .unwrap();
2898 assert!(plan.dry_run);
2899 assert_eq!(plan.scanned, 2);
2900 assert_eq!(plan.candidates, 2);
2901 assert_eq!(plan.embedded, 0);
2902 assert!(store.get_embedding(id_a).unwrap().is_none());
2903
2904 let applied = engine
2906 .backfill_embeddings(Some(Galaxy::Codex), 0, false)
2907 .unwrap();
2908 assert_eq!(applied.embedded, 2);
2909 assert!(store.get_embedding(id_a).unwrap().is_some());
2910 assert!(store.get_embedding(id_b).unwrap().is_some());
2911 assert_eq!(engine.vector_store.lock().unwrap().len(), 2);
2912
2913 let again = engine
2915 .backfill_embeddings(Some(Galaxy::Codex), 0, false)
2916 .unwrap();
2917 assert_eq!(again.candidates, 0);
2918 assert_eq!(again.already_embedded, 2);
2919 }
2920
2921 #[test]
2922 fn backfill_chunks_large_batches() {
2923 struct TestEmbedder;
2924 impl crate::embedder::Embedder for TestEmbedder {
2925 fn embed_batch(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
2926 Ok(texts.iter().map(|_| vec![0.1_f32; 4]).collect())
2927 }
2928 fn dimension(&self) -> usize {
2929 4
2930 }
2931 fn is_available(&self) -> bool {
2932 true
2933 }
2934 fn backend_name(&self) -> &'static str {
2935 "test"
2936 }
2937 }
2938
2939 let tmp = tempfile::tempdir().unwrap();
2940 let store_dir = tmp.path().join("store");
2941 std::fs::create_dir_all(&store_dir).unwrap();
2942 let store = Arc::new(MemoryStore::open_default(&store_dir).unwrap());
2943 let index_dir = tmp.path().join("index");
2944 std::fs::create_dir_all(&index_dir).unwrap();
2945 let search_engine = Arc::new(SearchEngine::open(&index_dir).unwrap());
2946 for i in 0..40 {
2947 let mem = crate::Memory::new(Galaxy::Codex, format!("chunked memory {i}"));
2948 store.put(Galaxy::Codex, &mem).unwrap();
2949 }
2950 let engine = RecallEngine::new(
2951 store,
2952 search_engine,
2953 VectorStore::new(),
2954 Arc::new(TestEmbedder),
2955 RecallConfig::default(),
2956 )
2957 .unwrap();
2958
2959 let report = engine
2961 .backfill_embeddings(Some(Galaxy::Codex), 0, false)
2962 .unwrap();
2963 assert_eq!(report.embedded, 40);
2964 assert_eq!(report.errors, 0);
2965 assert_eq!(engine.vector_store.lock().unwrap().len(), 40);
2966 }
2967
2968 #[test]
2969 fn backfill_skips_empty_content_and_counts_failures_once() {
2970 struct FailEmbedder;
2971 impl crate::embedder::Embedder for FailEmbedder {
2972 fn embed_batch(&self, _texts: &[&str]) -> Result<Vec<Vec<f32>>> {
2973 Err(CoreError::Memory("simulated embedder failure".into()))
2974 }
2975 fn dimension(&self) -> usize {
2976 4
2977 }
2978 fn is_available(&self) -> bool {
2979 true
2980 }
2981 fn backend_name(&self) -> &'static str {
2982 "test-fail"
2983 }
2984 }
2985
2986 let tmp = tempfile::tempdir().unwrap();
2987 let store_dir = tmp.path().join("store");
2988 std::fs::create_dir_all(&store_dir).unwrap();
2989 let store = Arc::new(MemoryStore::open_default(&store_dir).unwrap());
2990 let index_dir = tmp.path().join("index");
2991 std::fs::create_dir_all(&index_dir).unwrap();
2992 let search_engine = Arc::new(SearchEngine::open(&index_dir).unwrap());
2993
2994 let empty = crate::Memory::new(Galaxy::Codex, " ".to_string());
2995 let real = crate::Memory::new(Galaxy::Codex, "real content".to_string());
2996 store.put(Galaxy::Codex, &empty).unwrap();
2997 store.put(Galaxy::Codex, &real).unwrap();
2998
2999 let engine = RecallEngine::new(
3000 store,
3001 search_engine,
3002 VectorStore::new(),
3003 Arc::new(FailEmbedder),
3004 RecallConfig::default(),
3005 )
3006 .unwrap();
3007 let report = engine
3008 .backfill_embeddings(Some(Galaxy::Codex), 0, false)
3009 .unwrap();
3010 assert_eq!(report.skipped_empty, 1, "whitespace-only memory is skipped");
3011 assert_eq!(report.candidates, 1);
3012 assert_eq!(
3013 report.errors, 1,
3014 "a failed memory must be counted once (batch fallback), not twice"
3015 );
3016 }
3017
3018 #[test]
3019 fn backfill_refuses_stub_embedder() {
3020 let tmp = tempfile::tempdir().unwrap();
3021 let store_dir = tmp.path().join("store");
3022 std::fs::create_dir_all(&store_dir).unwrap();
3023 let store = Arc::new(MemoryStore::open_default(&store_dir).unwrap());
3024 let index_dir = tmp.path().join("index");
3025 std::fs::create_dir_all(&index_dir).unwrap();
3026 let search_engine = Arc::new(SearchEngine::open(&index_dir).unwrap());
3027 let engine = RecallEngine::new(
3028 store,
3029 search_engine,
3030 VectorStore::new(),
3031 Arc::new(crate::embedder::StubEmbedder::default()),
3032 RecallConfig::default(),
3033 )
3034 .unwrap();
3035 let err = engine
3036 .backfill_embeddings(Some(Galaxy::Codex), 10, true)
3037 .unwrap_err();
3038 assert!(err.to_string().contains("no real embedder"));
3039 }
3040
3041 #[test]
3042 fn embedder_probe_returns_vector_len() {
3043 let (_tmp, engine) = setup_engine();
3044 let dim = engine.embedder_probe().unwrap();
3045 assert!(dim > 0, "probe must return the embedder dimension");
3046 }
3047
3048 #[test]
3049 fn test_s10_association_rerank() {
3050 let (_tmp, mut engine) = setup_engine();
3051 let env = engine.store.env();
3052 let assoc_store = AssociationStore::open(env).unwrap();
3053
3054 let mem_a = Memory::new(Galaxy::Codex, "alpha query topic node".into());
3056 engine.store_with_embedding(Galaxy::Codex, &mem_a).unwrap();
3057
3058 let mem_b = Memory::new(Galaxy::Codex, "beta query topic node".into());
3060 let id_b = mem_b.metadata.id;
3061 engine.store_with_embedding(Galaxy::Codex, &mem_b).unwrap();
3062
3063 let mem_c = Memory::new(Galaxy::Research, "gamma target node".into());
3065 let id_c = mem_c.metadata.id;
3066 engine.store.put(Galaxy::Research, &mem_c).unwrap();
3067
3068 let edge = crate::associations::Association::new(
3069 id_b,
3070 id_c,
3071 crate::associations::LinkType::Related,
3072 0.8,
3073 );
3074 assoc_store.put(env, &edge).unwrap();
3075
3076 let results_default = engine.hybrid_search("query topic", 10, None);
3078 assert!(!results_default.is_empty());
3079
3080 engine.config.association_rerank = true;
3082 let results_rerank = engine.hybrid_search("query topic", 10, None);
3083 assert!(!results_rerank.is_empty());
3084
3085 let score_b_default = results_default
3087 .iter()
3088 .find(|r| r.memory_id == id_b)
3089 .unwrap()
3090 .score;
3091 let score_b_rerank = results_rerank
3092 .iter()
3093 .find(|r| r.memory_id == id_b)
3094 .unwrap()
3095 .score;
3096 assert!(score_b_rerank > score_b_default);
3097 }
3098}