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]
2120 fn fuse_results_near_ties_keep_absolute_raw_scores() {
2121 let id_hi = Uuid::new_v4();
2125 let id_lo = Uuid::new_v4();
2126 let bm25 = vec![
2127 SearchResult {
2128 memory_id: id_hi.to_string(),
2129 galaxy: Galaxy::Codex.db_name().to_string(),
2130 score: 1.45,
2131 normalized_score: 0.0,
2132 content: "top".into(),
2133 },
2134 SearchResult {
2135 memory_id: id_lo.to_string(),
2136 galaxy: Galaxy::Codex.db_name().to_string(),
2137 score: 1.18,
2138 normalized_score: 0.0,
2139 content: "near tie".into(),
2140 },
2141 ];
2142 let results = fuse(&bm25, &[], 10);
2143 assert_eq!(results.len(), 2);
2144 assert!((results[0].bm25_score - 1.0).abs() < 0.01);
2145 assert!((results[1].bm25_score - (1.18 / 1.45)).abs() < 0.01);
2146 assert!((results[0].raw_bm25_score - 1.45).abs() < f32::EPSILON);
2147 assert!((results[1].raw_bm25_score - 1.18).abs() < f32::EPSILON);
2148 }
2149
2150 #[test]
2153 fn embed_content_caches_result() {
2154 let embedder = StubEmbedder::new(384);
2155 let content = "test content for caching";
2156 let vec1 = embedder.embed(content).unwrap();
2157 let vec2 = embedder.embed(content).unwrap();
2158 assert_eq!(vec1, vec2);
2159 }
2160
2161 #[test]
2162 fn embed_content_different_content_different_result() {
2163 let embedder = StubEmbedder::new(384);
2164 let vec1 = embedder.embed("content one").unwrap();
2165 let vec2 = embedder.embed("content two").unwrap();
2166 assert_ne!(vec1, vec2);
2167 }
2168
2169 #[test]
2172 fn fuse_with_zero_bm25_weight() {
2173 let id = Uuid::new_v4();
2174 let bm25 = vec![SearchResult {
2175 memory_id: id.to_string(),
2176 galaxy: Galaxy::Codex.db_name().to_string(),
2177 score: 5.0,
2178 normalized_score: 0.0,
2179 content: "test".into(),
2180 }];
2181 let results = fuse_results_inner(
2182 &bm25,
2183 &[],
2184 10,
2185 0.5,
2186 0.3,
2187 0.2,
2188 0.0,
2189 |_, _| String::new(),
2190 |_, _| 0.0,
2191 |_, _| 0.7,
2192 );
2193 assert!((results[0].score - 0.5).abs() < 0.01);
2194 }
2195
2196 #[test]
2197 fn fuse_with_zero_vector_weight() {
2198 let id = Uuid::new_v4();
2199 let vector = vec![VectorSearchResult {
2200 memory_id: id,
2201 galaxy: Galaxy::Codex,
2202 score: 0.9,
2203 }];
2204 let results = fuse_results_inner(
2205 &[],
2206 &vector,
2207 10,
2208 0.5,
2209 0.3,
2210 0.2,
2211 0.0,
2212 |_, _| String::new(),
2213 |_, _| 0.0,
2214 |_, _| 0.7,
2215 );
2216 assert!((results[0].score - 0.27).abs() < 0.01);
2217 }
2218
2219 #[test]
2220 fn trust_weight_zero_is_byte_identical_to_no_weight() {
2221 let id = Uuid::new_v4();
2222 let bm25 = vec![SearchResult {
2223 memory_id: id.to_string(),
2224 galaxy: Galaxy::Codex.db_name().to_string(),
2225 score: 5.0,
2226 normalized_score: 0.0,
2227 content: "test".into(),
2228 }];
2229 let results = fuse_results_inner(
2232 &bm25,
2233 &[],
2234 10,
2235 0.5,
2236 0.3,
2237 0.2,
2238 0.0,
2239 |_, _| String::new(),
2240 |_, _| 0.0,
2241 |_, _| 0.4,
2242 );
2243 assert!((results[0].score - 0.5).abs() < 0.01);
2244 assert!((results[0].trust_factor - 1.0).abs() < f32::EPSILON);
2245 assert!(!results[0].in_conformal_set);
2246 }
2247
2248 #[test]
2249 fn trust_weight_orders_high_trust_above_low() {
2250 let high = Uuid::new_v4();
2253 let low = Uuid::new_v4();
2254 let mk = |id: &Uuid| SearchResult {
2255 memory_id: id.to_string(),
2256 galaxy: Galaxy::Codex.db_name().to_string(),
2257 score: 5.0,
2258 normalized_score: 0.0,
2259 content: "test".into(),
2260 };
2261 let bm25 = vec![mk(&high), mk(&low)];
2262 let mut trust_calls = 0;
2263 let results = fuse_results_inner(
2264 &bm25,
2265 &[],
2266 10,
2267 0.5,
2268 0.3,
2269 0.2,
2270 0.5,
2271 |_, _| String::new(),
2272 |_, _| 0.0,
2273 |id, _| {
2274 trust_calls += 1;
2275 if id == high { 1.0 } else { 0.4 }
2276 },
2277 );
2278 let hi = results.iter().find(|r| r.memory_id == high).unwrap();
2279 let lo = results.iter().find(|r| r.memory_id == low).unwrap();
2280 assert!(
2281 hi.score > lo.score,
2282 "user-confirmed (1.0) must outrank low-trust (0.4) at equal fused score"
2283 );
2284 assert!((hi.trust_factor - 1.15).abs() < 0.001);
2286 assert!((lo.trust_factor - 0.85).abs() < 0.001);
2287 assert!(trust_calls >= 2, "getter consulted per candidate");
2288 let neutral = Uuid::new_v4();
2290 let bm25_neutral = vec![SearchResult {
2291 memory_id: neutral.to_string(),
2292 galaxy: Galaxy::Codex.db_name().to_string(),
2293 score: 5.0,
2294 normalized_score: 0.0,
2295 content: "test".into(),
2296 }];
2297 let res_n = fuse_results_inner(
2298 &bm25_neutral,
2299 &[],
2300 10,
2301 0.5,
2302 0.3,
2303 0.2,
2304 0.5,
2305 |_, _| String::new(),
2306 |_, _| 0.0,
2307 |_, _| 0.7,
2308 );
2309 assert!((res_n[0].trust_factor - 1.0).abs() < 0.001);
2310 }
2311
2312 #[test]
2313 fn config_defaults_keep_both_s8_knobs_off() {
2314 let cfg = RecallConfig::default();
2319 assert_eq!(cfg.trust_weight, 0.0);
2320 assert_eq!(cfg.conformal_alpha, None);
2321 let env_cfg = RecallConfig::from_env();
2322 assert_eq!(env_cfg.trust_weight, 0.0, "unset env stays off");
2323 assert_eq!(env_cfg.conformal_alpha, None, "unset env stays off");
2324 }
2325
2326 #[test]
2329 fn galaxy_from_db_name_valid() {
2330 assert_eq!(Galaxy::from_db_name("codex"), Some(Galaxy::Codex));
2331 }
2332
2333 #[test]
2334 fn galaxy_from_db_name_invalid() {
2335 assert_eq!(Galaxy::from_db_name("nonexistent"), None);
2336 }
2337
2338 use crate::Memory;
2341 use tempfile::tempdir;
2342
2343 fn setup_engine() -> (tempfile::TempDir, RecallEngine) {
2344 let tmp = tempdir().unwrap();
2345 let store = Arc::new(MemoryStore::open_default(tmp.path()).unwrap());
2346 let tantivy_path = tmp.path().join("tantivy");
2347 std::fs::create_dir_all(&tantivy_path).unwrap();
2348 let search = Arc::new(SearchEngine::open(&tantivy_path).unwrap());
2349 let vector_store = VectorStore::new();
2350 let embedder: Arc<dyn Embedder> = Arc::new(StubEmbedder::new(384));
2351 let engine = RecallEngine::new(
2352 store,
2353 search,
2354 vector_store,
2355 embedder,
2356 RecallConfig::default(),
2357 )
2358 .unwrap();
2359 (tmp, engine)
2360 }
2361
2362 #[test]
2363 fn integration_store_and_hybrid_search_roundtrip() {
2364 let (_tmp, engine) = setup_engine();
2365
2366 let mem1 = Memory::new(
2367 Galaxy::Codex,
2368 "Rust programming language is fast and safe".into(),
2369 )
2370 .with_importance(0.8)
2371 .with_tags(vec!["rust".into(), "programming".into()]);
2372 let mem2 = Memory::new(Galaxy::Codex, "Python is great for data science".into())
2373 .with_importance(0.5)
2374 .with_tags(vec!["python".into(), "data".into()]);
2375 let mem3 = Memory::new(
2376 Galaxy::Codex,
2377 "The Rust ownership model prevents memory leaks".into(),
2378 )
2379 .with_importance(0.9)
2380 .with_tags(vec!["rust".into(), "memory".into()]);
2381
2382 engine.store_with_embedding(Galaxy::Codex, &mem1).unwrap();
2383 engine.store_with_embedding(Galaxy::Codex, &mem2).unwrap();
2384 engine.store_with_embedding(Galaxy::Codex, &mem3).unwrap();
2385
2386 let results = engine.hybrid_search("rust", 10, None);
2388 assert!(!results.is_empty(), "hybrid search should return results");
2389
2390 let top_contents: Vec<&str> = results.iter().map(|r| r.content.as_str()).collect();
2392 assert!(
2393 top_contents.iter().any(|c| c.contains("Rust")),
2394 "top results should include Rust content, got: {top_contents:?}"
2395 );
2396 }
2397
2398 #[test]
2399 fn integration_bm25_and_vector_both_contribute() {
2400 let (_tmp, engine) = setup_engine();
2401
2402 for i in 0..5 {
2404 let mem = Memory::new(
2405 Galaxy::Codex,
2406 format!("memory about topic {i} with unique content"),
2407 )
2408 .with_importance(0.5);
2409 engine.store_with_embedding(Galaxy::Codex, &mem).unwrap();
2410 }
2411
2412 let results = engine.hybrid_search("memory", 10, None);
2414 assert!(!results.is_empty(), "should find memories");
2415
2416 let has_bm25 = results.iter().any(|r| r.bm25_score > 0.0);
2418 assert!(has_bm25, "BM25 should contribute to fused results");
2419 }
2420
2421 #[test]
2422 fn integration_vector_search_only() {
2423 let (_tmp, engine) = setup_engine();
2424
2425 let content = "unique searchable content for vector test";
2426 let mem = Memory::new(Galaxy::Codex, content.into()).with_importance(0.7);
2427 engine.store_with_embedding(Galaxy::Codex, &mem).unwrap();
2428
2429 let results = engine.vector_search(content, 10, None);
2431 assert_eq!(results.len(), 1);
2432 assert_eq!(results[0].memory_id, mem.metadata.id);
2433 assert!(results[0].vector_score > 0.0);
2434 }
2435
2436 #[test]
2437 fn integration_text_search_only() {
2438 let (_tmp, engine) = setup_engine();
2439
2440 let mem = Memory::new(Galaxy::Codex, "specific text about rust ownership".into())
2441 .with_tags(vec!["rust".into()]);
2442 engine.store_with_embedding(Galaxy::Codex, &mem).unwrap();
2443
2444 let results = engine.text_search("rust", 10);
2445 assert!(!results.is_empty(), "text search should find results");
2446 assert!(results.iter().any(|r| r.bm25_score > 0.0));
2447 }
2448
2449 #[test]
2450 fn integration_batch_store_with_embedding() {
2451 let (_tmp, engine) = setup_engine();
2452
2453 let mem1 = Memory::new(Galaxy::Codex, "alpha beta gamma".into());
2454 let mem2 = Memory::new(Galaxy::Codex, "delta epsilon zeta".into());
2455 let mem3 = Memory::new(Galaxy::Codex, "eta theta iota".into());
2456
2457 let entries = vec![
2458 (Galaxy::Codex, &mem1),
2459 (Galaxy::Codex, &mem2),
2460 (Galaxy::Codex, &mem3),
2461 ];
2462
2463 let count = engine.store_batch_with_embedding(&entries).unwrap();
2464 assert_eq!(count, 3);
2465
2466 let results = engine.text_search("alpha", 10);
2468 assert!(
2469 !results.is_empty(),
2470 "batch-stored memory should be searchable"
2471 );
2472
2473 let vresults = engine.vector_search("alpha beta gamma", 10, None);
2475 assert_eq!(
2476 vresults.len(),
2477 1,
2478 "vector search should find the exact match"
2479 );
2480 assert_eq!(vresults[0].memory_id, mem1.metadata.id);
2481 }
2482
2483 #[test]
2484 fn integration_batch_store_empty() {
2485 let (_tmp, engine) = setup_engine();
2486 let entries: Vec<(Galaxy, &Memory)> = vec![];
2487 let count = engine.store_batch_with_embedding(&entries).unwrap();
2488 assert_eq!(count, 0);
2489 }
2490
2491 #[test]
2492 fn integration_galaxy_filter() {
2493 let (_tmp, engine) = setup_engine();
2494
2495 let mem_codex = Memory::new(Galaxy::Codex, "codex memory about rust".into());
2496 let mem_research = Memory::new(Galaxy::Research, "research memory about rust".into());
2497
2498 engine
2499 .store_with_embedding(Galaxy::Codex, &mem_codex)
2500 .unwrap();
2501 engine
2502 .store_with_embedding(Galaxy::Research, &mem_research)
2503 .unwrap();
2504
2505 let results = engine.hybrid_search("rust", 10, Some(Galaxy::Codex));
2506 assert!(!results.is_empty());
2507 assert!(
2508 results.iter().all(|r| r.galaxy == Galaxy::Codex),
2509 "all results should be from Codex galaxy"
2510 );
2511 }
2512
2513 #[test]
2514 fn integration_empty_search() {
2515 let (_tmp, engine) = setup_engine();
2516 let results = engine.hybrid_search("nonexistent", 10, None);
2517 assert!(results.is_empty());
2518 }
2519
2520 #[test]
2521 fn integration_cache_populated_after_store() {
2522 let (_tmp, engine) = setup_engine();
2523
2524 let mem = Memory::new(Galaxy::Codex, "content to be cached".into());
2525 engine.store_with_embedding(Galaxy::Codex, &mem).unwrap();
2526
2527 assert_eq!(engine.cache_size(), 1);
2529 }
2530
2531 #[test]
2532 fn integration_vector_count_tracks_stores() {
2533 let (_tmp, engine) = setup_engine();
2534
2535 assert_eq!(engine.vector_count(), 0);
2536
2537 for i in 0..3 {
2538 let mem = Memory::new(Galaxy::Codex, format!("memory {i}"));
2539 engine.store_with_embedding(Galaxy::Codex, &mem).unwrap();
2540 }
2541
2542 assert_eq!(engine.vector_count(), 3);
2543 }
2544
2545 #[test]
2546 fn integration_importance_affects_ranking() {
2547 let (_tmp, engine) = setup_engine();
2548
2549 let mem_low =
2551 Memory::new(Galaxy::Codex, "rust programming basics".into()).with_importance(0.1);
2552 let mem_high =
2553 Memory::new(Galaxy::Codex, "rust programming advanced".into()).with_importance(0.9);
2554
2555 engine
2556 .store_with_embedding(Galaxy::Codex, &mem_low)
2557 .unwrap();
2558 engine
2559 .store_with_embedding(Galaxy::Codex, &mem_high)
2560 .unwrap();
2561
2562 let results = engine.hybrid_search("rust", 10, None);
2563 assert_eq!(results.len(), 2);
2564
2565 let high_idx = results
2568 .iter()
2569 .position(|r| r.memory_id == mem_high.metadata.id)
2570 .unwrap();
2571 let low_idx = results
2572 .iter()
2573 .position(|r| r.memory_id == mem_low.metadata.id)
2574 .unwrap();
2575 assert!(
2576 high_idx < low_idx,
2577 "higher importance memory should rank higher"
2578 );
2579 }
2580
2581 #[test]
2582 fn config_from_env_rejects_nan_weights() {
2583 let mut config = RecallConfig::default();
2586 let w: f32 = "NaN".parse().unwrap();
2587 if w.is_finite() && w >= 0.0 {
2588 config.bm25_weight = w.min(1.0);
2589 }
2590 assert_eq!(
2591 config.bm25_weight, 0.5,
2592 "NaN should be rejected, default kept"
2593 );
2594 }
2595
2596 #[test]
2597 fn config_from_env_rejects_negative_weights() {
2598 let mut config = RecallConfig::default();
2599 let w: f32 = "-0.5".parse().unwrap();
2600 if w.is_finite() && w >= 0.0 {
2601 config.vector_weight = w.min(1.0);
2602 }
2603 assert_eq!(
2604 config.vector_weight, 0.3,
2605 "Negative should be rejected, default kept"
2606 );
2607 }
2608
2609 #[test]
2610 fn config_from_env_clamps_weights_to_1() {
2611 let mut config = RecallConfig::default();
2612 let w: f32 = "5.0".parse().unwrap();
2613 if w.is_finite() && w >= 0.0 {
2614 config.importance_weight = w.min(1.0);
2615 }
2616 assert_eq!(
2617 config.importance_weight, 1.0,
2618 "Weight should be clamped to 1.0"
2619 );
2620 }
2621
2622 #[test]
2623 fn config_from_env_normalizes_weights() {
2624 let mut config = RecallConfig {
2625 bm25_weight: 0.8,
2626 vector_weight: 0.8,
2627 importance_weight: 0.8,
2628 ..Default::default()
2629 };
2630 let sum = config.bm25_weight + config.vector_weight + config.importance_weight;
2631 if sum > 0.0 && (sum - 1.0).abs() > 0.01 {
2632 config.bm25_weight /= sum;
2633 config.vector_weight /= sum;
2634 config.importance_weight /= sum;
2635 }
2636 assert!(
2637 config.weights_normalized(),
2638 "Weights should be normalized to sum to 1.0"
2639 );
2640 }
2641
2642 #[test]
2643 fn config_from_env_rejects_infinity() {
2644 let mut config = RecallConfig::default();
2645 let w: f32 = "inf".parse().unwrap();
2646 if w.is_finite() && w >= 0.0 {
2647 config.bm25_weight = w.min(1.0);
2648 }
2649 assert_eq!(
2650 config.bm25_weight, 0.5,
2651 "Infinity should be rejected, default kept"
2652 );
2653 }
2654
2655 #[test]
2656 fn test_promotion_on_read_config_default() {
2657 let default_config = RecallConfig::default();
2658 assert!(!default_config.promotion_on_read);
2659
2660 let custom_config = RecallConfig {
2661 promotion_on_read: true,
2662 ..Default::default()
2663 };
2664 assert!(custom_config.promotion_on_read);
2665 }
2666
2667 #[test]
2668 fn test_promote_memory_updates_hebbian_score_and_counts() {
2669 let tmp = tempfile::tempdir().unwrap();
2670 let store_dir = tmp.path().join("store");
2671 std::fs::create_dir_all(&store_dir).unwrap();
2672 let store = Arc::new(MemoryStore::open_default(&store_dir).unwrap());
2673 let index_dir = tmp.path().join("index");
2674 std::fs::create_dir_all(&index_dir).unwrap();
2675 let search_engine = Arc::new(SearchEngine::open(&index_dir).unwrap());
2676 let vector_store = VectorStore::new();
2677 let embedder = Arc::new(crate::embedder::StubEmbedder::default());
2678 let config = RecallConfig {
2679 promotion_on_read: true,
2680 ..Default::default()
2681 };
2682 let engine =
2683 RecallEngine::new(store.clone(), search_engine, vector_store, embedder, config)
2684 .unwrap();
2685
2686 let mut mem = crate::Memory::new(Galaxy::Codex, "promotion on read test".to_string());
2687 mem.metadata.neuro_score = 0.5;
2688 mem.metadata.novelty_score = 1.0;
2689 let mem_id = mem.metadata.id;
2690 store.put(Galaxy::Codex, &mem).unwrap();
2691
2692 let promoted = engine.promote_memory(Galaxy::Codex, mem_id).unwrap();
2694 assert!(promoted);
2695
2696 let reloaded = store.get(Galaxy::Codex, mem_id).unwrap().unwrap();
2697 assert_eq!(reloaded.metadata.recall_count, 1);
2698 assert_eq!(reloaded.metadata.access_count, 1);
2699 assert!(
2700 reloaded.metadata.neuro_score > 0.5,
2701 "neuro_score should increase via Hebbian boost"
2702 );
2703 assert!(
2704 reloaded.metadata.novelty_score < 1.0,
2705 "novelty_score should decay on recall"
2706 );
2707 }
2708
2709 #[test]
2710 fn test_hybrid_search_triggers_promotion_on_read() {
2711 let tmp = tempfile::tempdir().unwrap();
2712 let store_dir = tmp.path().join("store");
2713 std::fs::create_dir_all(&store_dir).unwrap();
2714 let store = Arc::new(MemoryStore::open_default(&store_dir).unwrap());
2715 let index_dir = tmp.path().join("index");
2716 std::fs::create_dir_all(&index_dir).unwrap();
2717 let search_engine = Arc::new(SearchEngine::open(&index_dir).unwrap());
2718 let vector_store = VectorStore::new();
2719 let embedder = Arc::new(crate::embedder::StubEmbedder::default());
2720 let config = RecallConfig {
2721 promotion_on_read: true,
2722 ..Default::default()
2723 };
2724 let engine = RecallEngine::new(
2725 store.clone(),
2726 search_engine.clone(),
2727 vector_store,
2728 embedder,
2729 config,
2730 )
2731 .unwrap();
2732
2733 let mut mem = crate::Memory::new(Galaxy::Codex, "tokio army swarm tactics".to_string());
2734 mem.metadata.neuro_score = 0.5;
2735 mem.metadata.novelty_score = 1.0;
2736 let mem_id = mem.metadata.id;
2737 store.put(Galaxy::Codex, &mem).unwrap();
2738
2739 let mut writer = search_engine.writer().unwrap();
2740 search_engine
2741 .add_document(
2742 &mut writer,
2743 &mem_id.to_string(),
2744 "codex",
2745 "tokio army swarm tactics",
2746 &[],
2747 1_700_000_000,
2748 )
2749 .unwrap();
2750 search_engine.commit(&mut writer).unwrap();
2751
2752 let (results, _) =
2754 engine.hybrid_search_with_disclosure("tokio army", 5, Some(Galaxy::Codex));
2755 assert!(!results.is_empty());
2756 assert_eq!(results[0].memory_id, mem_id);
2757
2758 let reloaded = store.get(Galaxy::Codex, mem_id).unwrap().unwrap();
2759 assert_eq!(reloaded.metadata.recall_count, 1);
2760 assert!(reloaded.metadata.neuro_score > 0.5);
2761 }
2762
2763 #[test]
2764 fn hybrid_search_rehydrates_vectors_across_restart() {
2765 let tmp = tempfile::tempdir().unwrap();
2766 let store_dir = tmp.path().join("store");
2767 std::fs::create_dir_all(&store_dir).unwrap();
2768 let store = Arc::new(MemoryStore::open_default(&store_dir).unwrap());
2769 let index_dir = tmp.path().join("index");
2770 std::fs::create_dir_all(&index_dir).unwrap();
2771 let search_engine = Arc::new(SearchEngine::open(&index_dir).unwrap());
2772 let embedder = Arc::new(crate::embedder::StubEmbedder::default());
2773 let dim = embedder.dimension();
2774
2775 let mem = crate::Memory::new(Galaxy::Codex, "persisted vector canary".to_string());
2777 let mem_id = mem.metadata.id;
2778 store.put(Galaxy::Codex, &mem).unwrap();
2779 store.put_embedding(mem_id, &vec![0.5_f32; dim]).unwrap();
2780
2781 let engine = RecallEngine::new(
2783 store,
2784 search_engine,
2785 VectorStore::new(),
2786 embedder,
2787 RecallConfig::default(),
2788 )
2789 .unwrap();
2790 assert!(!engine.vector_store.lock().unwrap().is_loaded());
2791
2792 let _ = engine.hybrid_search_with_disclosure("rehydration probe", 5, None);
2795
2796 let vs = engine.vector_store.lock().unwrap();
2797 assert!(
2798 vs.is_loaded(),
2799 "vector store should be loaded after the first hybrid search"
2800 );
2801 assert_eq!(vs.len(), 1, "persisted embedding should be indexed");
2802 }
2803
2804 #[test]
2805 fn bm25_weighted_search_skips_the_query_embed() {
2806 struct FailingEmbedder;
2813 impl crate::embedder::Embedder for FailingEmbedder {
2814 fn embed_batch(&self, _texts: &[&str]) -> Result<Vec<Vec<f32>>> {
2815 Err(CoreError::Memory("embedder offline".into()))
2816 }
2817 fn dimension(&self) -> usize {
2818 16
2819 }
2820 fn is_available(&self) -> bool {
2821 false
2822 }
2823 fn backend_name(&self) -> &'static str {
2824 "failing-test"
2825 }
2826 }
2827
2828 let tmp = tempfile::tempdir().unwrap();
2829 let store_dir = tmp.path().join("store");
2830 std::fs::create_dir_all(&store_dir).unwrap();
2831 let store = Arc::new(MemoryStore::open_default(&store_dir).unwrap());
2832 let index_dir = tmp.path().join("index");
2833 std::fs::create_dir_all(&index_dir).unwrap();
2834 let search_engine = Arc::new(SearchEngine::open(&index_dir).unwrap());
2835
2836 let config = RecallConfig {
2837 bm25_weight: 1.0,
2838 vector_weight: 0.0,
2839 importance_weight: 0.0,
2840 ..RecallConfig::default()
2841 };
2842 let engine = RecallEngine::new(
2843 store.clone(),
2844 search_engine.clone(),
2845 VectorStore::new(),
2846 Arc::new(FailingEmbedder),
2847 config,
2848 )
2849 .unwrap();
2850
2851 let mem = crate::Memory::new(
2852 Galaxy::Codex,
2853 "kotlin coroutine budget meeting notes".to_string(),
2854 );
2855 let mem_id = mem.metadata.id;
2856 store.put(Galaxy::Codex, &mem).unwrap();
2857 let mut writer = search_engine.writer().unwrap();
2858 search_engine
2859 .add_document(
2860 &mut writer,
2861 &mem_id.to_string(),
2862 "codex",
2863 "kotlin coroutine budget meeting notes",
2864 &[],
2865 1_700_000_000,
2866 )
2867 .unwrap();
2868 search_engine.commit(&mut writer).unwrap();
2869
2870 let (results, _) =
2871 engine.hybrid_search_with_disclosure("kotlin coroutine budget", 5, Some(Galaxy::Codex));
2872 assert_eq!(
2873 results.len(),
2874 1,
2875 "the BM25 half must answer without the embedder"
2876 );
2877 assert_eq!(results[0].memory_id, mem_id);
2878 assert!(
2879 !engine.vector_store.lock().unwrap().is_loaded(),
2880 "an inert vector half must not rehydrate the vector index"
2881 );
2882 }
2883
2884 #[test]
2885 fn backfill_embeddings_dry_run_then_apply() {
2886 struct TestEmbedder;
2887 impl crate::embedder::Embedder for TestEmbedder {
2888 fn embed_batch(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
2889 Ok(texts.iter().map(|_| vec![0.25_f32; 8]).collect())
2890 }
2891 fn dimension(&self) -> usize {
2892 8
2893 }
2894 fn is_available(&self) -> bool {
2895 true
2896 }
2897 fn backend_name(&self) -> &'static str {
2898 "test"
2899 }
2900 }
2901
2902 let tmp = tempfile::tempdir().unwrap();
2903 let store_dir = tmp.path().join("store");
2904 std::fs::create_dir_all(&store_dir).unwrap();
2905 let store = Arc::new(MemoryStore::open_default(&store_dir).unwrap());
2906 let index_dir = tmp.path().join("index");
2907 std::fs::create_dir_all(&index_dir).unwrap();
2908 let search_engine = Arc::new(SearchEngine::open(&index_dir).unwrap());
2909
2910 let mem_a = crate::Memory::new(Galaxy::Codex, "alpha unique content".to_string());
2911 let mem_b = crate::Memory::new(Galaxy::Codex, "beta unique content".to_string());
2912 let (id_a, id_b) = (mem_a.metadata.id, mem_b.metadata.id);
2913 store.put(Galaxy::Codex, &mem_a).unwrap();
2914 store.put(Galaxy::Codex, &mem_b).unwrap();
2915
2916 let engine = RecallEngine::new(
2917 store.clone(),
2918 search_engine,
2919 VectorStore::new(),
2920 Arc::new(TestEmbedder),
2921 RecallConfig::default(),
2922 )
2923 .unwrap();
2924
2925 let plan = engine
2927 .backfill_embeddings(Some(Galaxy::Codex), 0, true)
2928 .unwrap();
2929 assert!(plan.dry_run);
2930 assert_eq!(plan.scanned, 2);
2931 assert_eq!(plan.candidates, 2);
2932 assert_eq!(plan.embedded, 0);
2933 assert!(store.get_embedding(id_a).unwrap().is_none());
2934
2935 let applied = engine
2937 .backfill_embeddings(Some(Galaxy::Codex), 0, false)
2938 .unwrap();
2939 assert_eq!(applied.embedded, 2);
2940 assert!(store.get_embedding(id_a).unwrap().is_some());
2941 assert!(store.get_embedding(id_b).unwrap().is_some());
2942 assert_eq!(engine.vector_store.lock().unwrap().len(), 2);
2943
2944 let again = engine
2946 .backfill_embeddings(Some(Galaxy::Codex), 0, false)
2947 .unwrap();
2948 assert_eq!(again.candidates, 0);
2949 assert_eq!(again.already_embedded, 2);
2950 }
2951
2952 #[test]
2953 fn backfill_chunks_large_batches() {
2954 struct TestEmbedder;
2955 impl crate::embedder::Embedder for TestEmbedder {
2956 fn embed_batch(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
2957 Ok(texts.iter().map(|_| vec![0.1_f32; 4]).collect())
2958 }
2959 fn dimension(&self) -> usize {
2960 4
2961 }
2962 fn is_available(&self) -> bool {
2963 true
2964 }
2965 fn backend_name(&self) -> &'static str {
2966 "test"
2967 }
2968 }
2969
2970 let tmp = tempfile::tempdir().unwrap();
2971 let store_dir = tmp.path().join("store");
2972 std::fs::create_dir_all(&store_dir).unwrap();
2973 let store = Arc::new(MemoryStore::open_default(&store_dir).unwrap());
2974 let index_dir = tmp.path().join("index");
2975 std::fs::create_dir_all(&index_dir).unwrap();
2976 let search_engine = Arc::new(SearchEngine::open(&index_dir).unwrap());
2977 for i in 0..40 {
2978 let mem = crate::Memory::new(Galaxy::Codex, format!("chunked memory {i}"));
2979 store.put(Galaxy::Codex, &mem).unwrap();
2980 }
2981 let engine = RecallEngine::new(
2982 store,
2983 search_engine,
2984 VectorStore::new(),
2985 Arc::new(TestEmbedder),
2986 RecallConfig::default(),
2987 )
2988 .unwrap();
2989
2990 let report = engine
2992 .backfill_embeddings(Some(Galaxy::Codex), 0, false)
2993 .unwrap();
2994 assert_eq!(report.embedded, 40);
2995 assert_eq!(report.errors, 0);
2996 assert_eq!(engine.vector_store.lock().unwrap().len(), 40);
2997 }
2998
2999 #[test]
3000 fn backfill_skips_empty_content_and_counts_failures_once() {
3001 struct FailEmbedder;
3002 impl crate::embedder::Embedder for FailEmbedder {
3003 fn embed_batch(&self, _texts: &[&str]) -> Result<Vec<Vec<f32>>> {
3004 Err(CoreError::Memory("simulated embedder failure".into()))
3005 }
3006 fn dimension(&self) -> usize {
3007 4
3008 }
3009 fn is_available(&self) -> bool {
3010 true
3011 }
3012 fn backend_name(&self) -> &'static str {
3013 "test-fail"
3014 }
3015 }
3016
3017 let tmp = tempfile::tempdir().unwrap();
3018 let store_dir = tmp.path().join("store");
3019 std::fs::create_dir_all(&store_dir).unwrap();
3020 let store = Arc::new(MemoryStore::open_default(&store_dir).unwrap());
3021 let index_dir = tmp.path().join("index");
3022 std::fs::create_dir_all(&index_dir).unwrap();
3023 let search_engine = Arc::new(SearchEngine::open(&index_dir).unwrap());
3024
3025 let empty = crate::Memory::new(Galaxy::Codex, " ".to_string());
3026 let real = crate::Memory::new(Galaxy::Codex, "real content".to_string());
3027 store.put(Galaxy::Codex, &empty).unwrap();
3028 store.put(Galaxy::Codex, &real).unwrap();
3029
3030 let engine = RecallEngine::new(
3031 store,
3032 search_engine,
3033 VectorStore::new(),
3034 Arc::new(FailEmbedder),
3035 RecallConfig::default(),
3036 )
3037 .unwrap();
3038 let report = engine
3039 .backfill_embeddings(Some(Galaxy::Codex), 0, false)
3040 .unwrap();
3041 assert_eq!(report.skipped_empty, 1, "whitespace-only memory is skipped");
3042 assert_eq!(report.candidates, 1);
3043 assert_eq!(
3044 report.errors, 1,
3045 "a failed memory must be counted once (batch fallback), not twice"
3046 );
3047 }
3048
3049 #[test]
3050 fn backfill_refuses_stub_embedder() {
3051 let tmp = tempfile::tempdir().unwrap();
3052 let store_dir = tmp.path().join("store");
3053 std::fs::create_dir_all(&store_dir).unwrap();
3054 let store = Arc::new(MemoryStore::open_default(&store_dir).unwrap());
3055 let index_dir = tmp.path().join("index");
3056 std::fs::create_dir_all(&index_dir).unwrap();
3057 let search_engine = Arc::new(SearchEngine::open(&index_dir).unwrap());
3058 let engine = RecallEngine::new(
3059 store,
3060 search_engine,
3061 VectorStore::new(),
3062 Arc::new(crate::embedder::StubEmbedder::default()),
3063 RecallConfig::default(),
3064 )
3065 .unwrap();
3066 let err = engine
3067 .backfill_embeddings(Some(Galaxy::Codex), 10, true)
3068 .unwrap_err();
3069 assert!(err.to_string().contains("no real embedder"));
3070 }
3071
3072 #[test]
3073 fn embedder_probe_returns_vector_len() {
3074 let (_tmp, engine) = setup_engine();
3075 let dim = engine.embedder_probe().unwrap();
3076 assert!(dim > 0, "probe must return the embedder dimension");
3077 }
3078
3079 #[test]
3080 fn test_s10_association_rerank() {
3081 let (_tmp, mut engine) = setup_engine();
3082 let env = engine.store.env();
3083 let assoc_store = AssociationStore::open(env).unwrap();
3084
3085 let mem_a = Memory::new(Galaxy::Codex, "alpha query topic node".into());
3087 engine.store_with_embedding(Galaxy::Codex, &mem_a).unwrap();
3088
3089 let mem_b = Memory::new(Galaxy::Codex, "beta query topic node".into());
3091 let id_b = mem_b.metadata.id;
3092 engine.store_with_embedding(Galaxy::Codex, &mem_b).unwrap();
3093
3094 let mem_c = Memory::new(Galaxy::Research, "gamma target node".into());
3096 let id_c = mem_c.metadata.id;
3097 engine.store.put(Galaxy::Research, &mem_c).unwrap();
3098
3099 let edge = crate::associations::Association::new(
3100 id_b,
3101 id_c,
3102 crate::associations::LinkType::Related,
3103 0.8,
3104 );
3105 assoc_store.put(env, &edge).unwrap();
3106
3107 let results_default = engine.hybrid_search("query topic", 10, None);
3109 assert!(!results_default.is_empty());
3110
3111 engine.config.association_rerank = true;
3113 let results_rerank = engine.hybrid_search("query topic", 10, None);
3114 assert!(!results_rerank.is_empty());
3115
3116 let score_b_default = results_default
3118 .iter()
3119 .find(|r| r.memory_id == id_b)
3120 .unwrap()
3121 .score;
3122 let score_b_rerank = results_rerank
3123 .iter()
3124 .find(|r| r.memory_id == id_b)
3125 .unwrap()
3126 .score;
3127 assert!(score_b_rerank > score_b_default);
3128 }
3129}