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