1mod builder;
38mod config;
39
40pub use builder::MemoryBuilder;
41pub use config::MemoryConfig;
42
43use crate::constants::{
44 MEMORY_IMPORTANCE_DEFAULT, MEMORY_IMPORTANCE_MAX, MEMORY_IMPORTANCE_MIN,
45 MEMORY_RECALL_LIMIT_DEFAULT, MEMORY_RECALL_LIMIT_MAX, MEMORY_TEXT_BYTES_MAX,
46};
47use crate::embedding::EmbeddingProvider;
48use crate::evolution::{DetectionOptions, EvolutionTracker};
49use crate::extraction::{EntityExtractor, ExtractionOptions};
50use crate::llm::LLMProvider;
51use crate::retrieval::{DualRetriever, SearchOptions};
52use crate::storage::{Entity, EntityType, EvolutionRelation, StorageBackend};
53use thiserror::Error;
54
55#[derive(Debug, Error)]
61pub enum MemoryError {
62 #[error("text is empty")]
64 EmptyText,
65
66 #[error("text too long: {len} bytes (max {max})")]
68 TextTooLong {
69 len: usize,
71 max: usize,
73 },
74
75 #[error("query is empty")]
77 EmptyQuery,
78
79 #[error("invalid importance: {value} (must be {min}-{max})")]
81 InvalidImportance {
82 value: f32,
84 min: f32,
86 max: f32,
88 },
89
90 #[error("invalid limit: {value} (must be 1-{max})")]
92 InvalidLimit {
93 value: usize,
95 max: usize,
97 },
98
99 #[error("storage error: {message}")]
101 Storage {
102 message: String,
104 },
105
106 #[error("embedding generation failed: {message}")]
108 EmbeddingFailed {
109 message: String,
111 },
112
113 #[error("vector search unavailable: {reason}")]
115 VectorSearchUnavailable {
116 reason: String,
118 },
119
120 #[error("embedding dimensions mismatch: expected {expected}, got {actual}")]
122 DimensionMismatch {
123 expected: usize,
125 actual: usize,
127 },
128}
129
130impl From<crate::storage::StorageError> for MemoryError {
131 fn from(err: crate::storage::StorageError) -> Self {
132 MemoryError::Storage {
133 message: err.to_string(),
134 }
135 }
136}
137
138impl From<crate::embedding::EmbeddingError> for MemoryError {
139 fn from(err: crate::embedding::EmbeddingError) -> Self {
140 MemoryError::EmbeddingFailed {
141 message: err.to_string(),
142 }
143 }
144}
145
146#[derive(Debug, Clone)]
154pub struct RememberOptions {
155 pub extract_entities: bool,
157
158 pub track_evolution: bool,
160
161 pub importance: f32,
163
164 pub generate_embeddings: bool,
166}
167
168impl RememberOptions {
169 #[must_use]
171 pub fn new() -> Self {
172 Self::default()
173 }
174
175 #[must_use]
177 pub fn without_extraction(mut self) -> Self {
178 self.extract_entities = false;
179 self
180 }
181
182 #[must_use]
184 pub fn without_evolution(mut self) -> Self {
185 self.track_evolution = false;
186 self
187 }
188
189 #[must_use]
194 pub fn with_importance(mut self, importance: f32) -> Self {
195 debug_assert!(
196 (MEMORY_IMPORTANCE_MIN..=MEMORY_IMPORTANCE_MAX).contains(&importance),
197 "importance must be {}-{}: got {}",
198 MEMORY_IMPORTANCE_MIN,
199 MEMORY_IMPORTANCE_MAX,
200 importance
201 );
202 self.importance = importance;
203 self
204 }
205
206 #[must_use]
208 pub fn with_embeddings(mut self) -> Self {
209 self.generate_embeddings = true;
210 self
211 }
212
213 #[must_use]
215 pub fn without_embeddings(mut self) -> Self {
216 self.generate_embeddings = false;
217 self
218 }
219}
220
221impl Default for RememberOptions {
222 fn default() -> Self {
223 Self {
224 extract_entities: true,
225 track_evolution: true,
226 importance: MEMORY_IMPORTANCE_DEFAULT,
227 generate_embeddings: true,
228 }
229 }
230}
231
232#[derive(Debug, Clone)]
236pub struct RecallOptions {
237 pub limit: usize,
239
240 pub deep_search: Option<bool>,
242
243 pub time_range: Option<(u64, u64)>,
245}
246
247impl RecallOptions {
248 #[must_use]
250 pub fn new() -> Self {
251 Self::default()
252 }
253
254 #[must_use]
259 pub fn with_limit(mut self, limit: usize) -> Self {
260 debug_assert!(
261 limit > 0 && limit <= MEMORY_RECALL_LIMIT_MAX,
262 "limit must be 1-{}: got {}",
263 MEMORY_RECALL_LIMIT_MAX,
264 limit
265 );
266 self.limit = limit;
267 self
268 }
269
270 #[must_use]
272 pub fn with_deep_search(mut self) -> Self {
273 self.deep_search = Some(true);
274 self
275 }
276
277 #[must_use]
279 pub fn fast_only(mut self) -> Self {
280 self.deep_search = Some(false);
281 self
282 }
283
284 #[must_use]
286 pub fn with_time_range(mut self, start_ms: u64, end_ms: u64) -> Self {
287 debug_assert!(start_ms <= end_ms, "start_ms must be <= end_ms");
288 self.time_range = Some((start_ms, end_ms));
289 self
290 }
291}
292
293impl Default for RecallOptions {
294 fn default() -> Self {
295 Self {
296 limit: MEMORY_RECALL_LIMIT_DEFAULT,
297 deep_search: None,
298 time_range: None,
299 }
300 }
301}
302
303#[derive(Debug, Clone)]
309pub struct RememberResult {
310 pub entities: Vec<Entity>,
312
313 pub evolutions: Vec<EvolutionRelation>,
315}
316
317impl RememberResult {
318 #[must_use]
320 pub fn new(entities: Vec<Entity>, evolutions: Vec<EvolutionRelation>) -> Self {
321 Self {
322 entities,
323 evolutions,
324 }
325 }
326
327 #[must_use]
329 pub fn entity_count(&self) -> usize {
330 self.entities.len()
331 }
332
333 #[must_use]
335 pub fn has_evolutions(&self) -> bool {
336 !self.evolutions.is_empty()
337 }
338
339 pub fn iter_entities(&self) -> impl Iterator<Item = &Entity> {
341 self.entities.iter()
342 }
343}
344
345pub struct Memory<
374 L: LLMProvider,
375 E: EmbeddingProvider,
376 S: StorageBackend,
377 V: crate::storage::VectorBackend,
378> {
379 storage: S,
380 extractor: EntityExtractor<L>,
381 retriever: DualRetriever<L, E, V, S>,
382 evolution: EvolutionTracker<L, S>,
383 embedder: E,
384 vector: V,
385}
386
387impl<
388 L: LLMProvider + Clone,
389 E: EmbeddingProvider + Clone,
390 S: StorageBackend + Clone,
391 V: crate::storage::VectorBackend + Clone,
392 > Memory<L, E, S, V>
393{
394 #[must_use]
402 pub fn new(llm: L, embedder: E, vector: V, storage: S) -> Self {
403 let extractor = EntityExtractor::new(llm.clone());
404 let retriever = DualRetriever::new(
405 llm.clone(),
406 embedder.clone(),
407 vector.clone(),
408 storage.clone(),
409 );
410 let evolution = EvolutionTracker::new(llm);
411
412 Self {
413 storage,
414 extractor,
415 retriever,
416 evolution,
417 embedder,
418 vector,
419 }
420 }
421
422 #[must_use]
434 pub fn builder() -> MemoryBuilder<L, E, V, S> {
435 MemoryBuilder::new()
436 }
437
438 pub async fn remember(
455 &mut self,
456 text: &str,
457 options: RememberOptions,
458 ) -> Result<RememberResult, MemoryError> {
459 if text.is_empty() {
461 return Err(MemoryError::EmptyText);
462 }
463 if text.len() > MEMORY_TEXT_BYTES_MAX {
464 return Err(MemoryError::TextTooLong {
465 len: text.len(),
466 max: MEMORY_TEXT_BYTES_MAX,
467 });
468 }
469 if !(MEMORY_IMPORTANCE_MIN..=MEMORY_IMPORTANCE_MAX).contains(&options.importance) {
470 return Err(MemoryError::InvalidImportance {
471 value: options.importance,
472 min: MEMORY_IMPORTANCE_MIN,
473 max: MEMORY_IMPORTANCE_MAX,
474 });
475 }
476
477 let mut entities = Vec::new();
478 let mut evolutions = Vec::new();
479
480 let extracted = if options.extract_entities {
482 match self
483 .extractor
484 .extract(text, ExtractionOptions::default())
485 .await
486 {
487 Ok(result) => result.entities,
488 Err(_) => vec![], }
490 } else {
491 vec![]
492 };
493
494 let mut to_store: Vec<Entity> = if extracted.is_empty() {
496 let name = if text.len() > 50 {
498 format!("Note: {}...", &text[..47])
499 } else {
500 format!("Note: {}", text)
501 };
502 vec![Entity::new(EntityType::Note, name, text.to_string())]
503 } else {
504 extracted
505 .into_iter()
506 .map(|e| {
507 let entity_type = convert_entity_type(&e.entity_type);
508 Entity::new(entity_type, e.name, e.content)
509 })
510 .collect()
511 };
512
513 if options.generate_embeddings && !to_store.is_empty() {
515 let contents: Vec<&str> = to_store.iter().map(|e| e.content.as_str()).collect();
517
518 match self.embedder.embed_batch(&contents).await {
519 Ok(embeddings) => {
520 for (entity, embedding) in to_store.iter_mut().zip(embeddings) {
522 entity.set_embedding(embedding);
523 }
524 }
525 Err(e) => {
526 tracing::warn!(
528 "Failed to generate embeddings: {}. Continuing without embeddings.",
529 e
530 );
531 }
532 }
533 }
534
535 for entity in to_store {
537 let _stored_id = self.storage.store_entity(&entity).await?;
539
540 if let Some(ref embedding) = entity.embedding {
542 if let Err(e) = self.vector.store(&entity.id, embedding).await {
543 tracing::warn!(
544 "Failed to store embedding in vector backend for entity {}: {}. Entity searchable by text only.",
545 entity.id, e
546 );
547 }
548 }
549
550 if options.track_evolution {
552 if let Ok(existing) = self.storage.search(&entity.name, 5).await {
554 let existing: Vec<Entity> =
556 existing.into_iter().filter(|e| e.id != entity.id).collect();
557
558 if !existing.is_empty() {
559 if let Ok(Some(detection)) = self
560 .evolution
561 .detect(&entity, &existing, DetectionOptions::default())
562 .await
563 {
564 evolutions.push(detection.relation);
565 }
566 }
567 }
568 }
569
570 entities.push(entity);
571 }
572
573 debug_assert!(!entities.is_empty(), "must store at least one entity");
575
576 Ok(RememberResult::new(entities, evolutions))
577 }
578
579 pub async fn recall(
593 &self,
594 query: &str,
595 options: RecallOptions,
596 ) -> Result<Vec<Entity>, MemoryError> {
597 if query.is_empty() {
599 return Err(MemoryError::EmptyQuery);
600 }
601 if options.limit == 0 || options.limit > MEMORY_RECALL_LIMIT_MAX {
602 return Err(MemoryError::InvalidLimit {
603 value: options.limit,
604 max: MEMORY_RECALL_LIMIT_MAX,
605 });
606 }
607
608 let mut search_options = SearchOptions::new().with_limit(options.limit);
610
611 if let Some(deep) = options.deep_search {
613 search_options = search_options.with_deep_search(deep);
614 }
615
616 if let Some((start, end)) = options.time_range {
618 search_options = search_options.with_time_range(start, end);
619 }
620
621 let result = self
623 .retriever
624 .search(query, search_options)
625 .await
626 .map_err(|e| MemoryError::Storage {
627 message: e.to_string(),
628 })?;
629
630 debug_assert!(
632 result.len() <= options.limit,
633 "results exceed limit: {} > {}",
634 result.len(),
635 options.limit
636 );
637
638 Ok(result.entities)
639 }
640
641 pub async fn forget(&mut self, entity_id: &str) -> Result<bool, MemoryError> {
649 debug_assert!(!entity_id.is_empty(), "entity_id must not be empty");
650
651 self.storage.delete_entity(entity_id).await?;
652 Ok(true)
653 }
654
655 pub async fn get(&self, entity_id: &str) -> Result<Option<Entity>, MemoryError> {
663 debug_assert!(!entity_id.is_empty(), "entity_id must not be empty");
664
665 Ok(self.storage.get_entity(entity_id).await?)
666 }
667
668 pub async fn count(&self) -> Result<usize, MemoryError> {
670 Ok(self.storage.count_entities(None).await?)
671 }
672
673 #[must_use]
675 pub fn storage(&self) -> &S {
676 &self.storage
677 }
678}
679
680fn convert_entity_type(ext_type: &crate::extraction::EntityType) -> EntityType {
686 use crate::extraction::EntityType as ExtType;
687
688 match ext_type {
689 ExtType::Person => EntityType::Person,
690 ExtType::Organization => EntityType::Note, ExtType::Project => EntityType::Project,
692 ExtType::Topic => EntityType::Topic,
693 ExtType::Preference => EntityType::Note, ExtType::Task => EntityType::Task,
695 ExtType::Event => EntityType::Note, ExtType::Note => EntityType::Note,
697 }
698}
699
700#[cfg(test)]
705mod tests {
706 use super::*;
707 use crate::dst::SimConfig;
708 use crate::embedding::SimEmbeddingProvider;
709 use crate::llm::SimLLMProvider;
710 use crate::storage::{SimStorageBackend, SimVectorBackend};
711
712 fn create_memory(
714 seed: u64,
715 ) -> Memory<SimLLMProvider, SimEmbeddingProvider, SimStorageBackend, SimVectorBackend> {
716 let llm = SimLLMProvider::with_seed(seed);
717 let embedder = SimEmbeddingProvider::with_seed(seed);
718 let vector = SimVectorBackend::new(seed);
719 let storage = SimStorageBackend::new(SimConfig::with_seed(seed));
720 Memory::new(llm, embedder, vector, storage)
721 }
722
723 #[test]
728 fn test_remember_options_default() {
729 let options = RememberOptions::default();
730
731 assert!(options.extract_entities);
732 assert!(options.track_evolution);
733 assert!(options.generate_embeddings);
734 assert!((options.importance - MEMORY_IMPORTANCE_DEFAULT).abs() < f32::EPSILON);
735 }
736
737 #[test]
738 fn test_remember_options_builder() {
739 let options = RememberOptions::new()
740 .without_extraction()
741 .without_evolution()
742 .with_importance(0.8);
743
744 assert!(!options.extract_entities);
745 assert!(!options.track_evolution);
746 assert!((options.importance - 0.8).abs() < f32::EPSILON);
747 }
748
749 #[test]
750 #[should_panic(expected = "importance must be")]
751 fn test_remember_options_invalid_importance() {
752 let _ = RememberOptions::new().with_importance(1.5);
753 }
754
755 #[test]
760 fn test_recall_options_default() {
761 let options = RecallOptions::default();
762
763 assert_eq!(options.limit, MEMORY_RECALL_LIMIT_DEFAULT);
764 assert!(options.deep_search.is_none());
765 assert!(options.time_range.is_none());
766 }
767
768 #[test]
769 fn test_recall_options_builder() {
770 let options = RecallOptions::new()
771 .with_limit(20)
772 .with_deep_search()
773 .with_time_range(1000, 2000);
774
775 assert_eq!(options.limit, 20);
776 assert_eq!(options.deep_search, Some(true));
777 assert_eq!(options.time_range, Some((1000, 2000)));
778 }
779
780 #[test]
781 fn test_recall_options_fast_only() {
782 let options = RecallOptions::new().fast_only();
783
784 assert_eq!(options.deep_search, Some(false));
785 }
786
787 #[test]
788 #[should_panic(expected = "limit must be")]
789 fn test_recall_options_invalid_limit_zero() {
790 let _ = RecallOptions::new().with_limit(0);
791 }
792
793 #[test]
794 #[should_panic(expected = "limit must be")]
795 fn test_recall_options_invalid_limit_too_large() {
796 let _ = RecallOptions::new().with_limit(MEMORY_RECALL_LIMIT_MAX + 1);
797 }
798
799 #[test]
804 fn test_remember_result() {
805 let entities = vec![Entity::new(
806 EntityType::Person,
807 "Alice".to_string(),
808 "Works at Acme".to_string(),
809 )];
810 let result = RememberResult::new(entities, vec![]);
811
812 assert_eq!(result.entity_count(), 1);
813 assert!(!result.has_evolutions());
814 }
815
816 #[test]
821 fn test_memory_creation() {
822 let memory = create_memory(42);
823 let _ = memory;
825 }
826
827 #[tokio::test]
832 async fn test_remember_basic() {
833 let mut memory = create_memory(42);
834
835 let result = memory
836 .remember("Alice works at Acme Corp", RememberOptions::default())
837 .await;
838
839 assert!(result.is_ok());
840 let result = result.unwrap();
841 assert!(!result.entities.is_empty());
842 }
843
844 #[tokio::test]
845 async fn test_remember_without_extraction() {
846 let mut memory = create_memory(42);
847
848 let result = memory
849 .remember(
850 "Some text to store",
851 RememberOptions::new().without_extraction(),
852 )
853 .await;
854
855 assert!(result.is_ok());
856 let result = result.unwrap();
857 assert_eq!(result.entity_count(), 1);
858 assert_eq!(result.entities[0].entity_type, EntityType::Note);
860 }
861
862 #[tokio::test]
863 async fn test_remember_empty_text_error() {
864 let mut memory = create_memory(42);
865
866 let result = memory.remember("", RememberOptions::default()).await;
867
868 assert!(result.is_err());
869 assert!(matches!(result.unwrap_err(), MemoryError::EmptyText));
870 }
871
872 #[tokio::test]
873 async fn test_remember_text_too_long_error() {
874 let mut memory = create_memory(42);
875 let long_text = "a".repeat(MEMORY_TEXT_BYTES_MAX + 1);
876
877 let result = memory
878 .remember(&long_text, RememberOptions::default())
879 .await;
880
881 assert!(result.is_err());
882 assert!(matches!(
883 result.unwrap_err(),
884 MemoryError::TextTooLong { .. }
885 ));
886 }
887
888 #[tokio::test]
893 async fn test_recall_basic() {
894 let mut memory = create_memory(42);
895
896 memory
898 .remember("Alice works at Acme Corp", RememberOptions::default())
899 .await
900 .unwrap();
901
902 let results = memory.recall("Alice", RecallOptions::default()).await;
904
905 assert!(results.is_ok());
906 }
907
908 #[tokio::test]
909 async fn test_recall_empty_query_error() {
910 let memory = create_memory(42);
911
912 let result = memory.recall("", RecallOptions::default()).await;
913
914 assert!(result.is_err());
915 assert!(matches!(result.unwrap_err(), MemoryError::EmptyQuery));
916 }
917
918 #[tokio::test]
919 async fn test_recall_with_limit() {
920 let mut memory = create_memory(42);
921
922 for i in 0..5 {
924 memory
925 .remember(
926 &format!("Item {} is interesting", i),
927 RememberOptions::new().without_extraction(),
928 )
929 .await
930 .unwrap();
931 }
932
933 let results = memory
935 .recall("Item", RecallOptions::new().with_limit(2))
936 .await
937 .unwrap();
938
939 assert!(results.len() <= 2);
940 }
941
942 #[tokio::test]
947 async fn test_get_entity() {
948 let mut memory = create_memory(42);
949
950 let result = memory
951 .remember("Test entity", RememberOptions::new().without_extraction())
952 .await
953 .unwrap();
954
955 let entity_id = &result.entities[0].id;
956 let found = memory.get(entity_id).await.unwrap();
957
958 assert!(found.is_some());
959 assert_eq!(found.unwrap().id, *entity_id);
960 }
961
962 #[tokio::test]
963 async fn test_get_nonexistent() {
964 let memory = create_memory(42);
965
966 let found = memory.get("nonexistent-id").await.unwrap();
967
968 assert!(found.is_none());
969 }
970
971 #[tokio::test]
972 async fn test_forget_entity() {
973 let mut memory = create_memory(42);
974
975 let result = memory
976 .remember("Test entity", RememberOptions::new().without_extraction())
977 .await
978 .unwrap();
979
980 let entity_id = &result.entities[0].id;
981
982 assert!(memory.get(entity_id).await.unwrap().is_some());
984
985 let deleted = memory.forget(entity_id).await.unwrap();
987 assert!(deleted);
988
989 assert!(memory.get(entity_id).await.unwrap().is_none());
991 }
992
993 #[tokio::test]
998 async fn test_count() {
999 let mut memory = create_memory(42);
1000
1001 assert_eq!(memory.count().await.unwrap(), 0);
1002
1003 memory
1004 .remember("First item", RememberOptions::new().without_extraction())
1005 .await
1006 .unwrap();
1007
1008 assert_eq!(memory.count().await.unwrap(), 1);
1009
1010 memory
1011 .remember("Second item", RememberOptions::new().without_extraction())
1012 .await
1013 .unwrap();
1014
1015 assert_eq!(memory.count().await.unwrap(), 2);
1016 }
1017
1018 #[tokio::test]
1023 async fn test_deterministic_same_seed() {
1024 let mut memory1 = create_memory(42);
1025 let mut memory2 = create_memory(42);
1026
1027 let text = "Alice works at Acme Corp as an engineer";
1028
1029 let result1 = memory1
1030 .remember(text, RememberOptions::default())
1031 .await
1032 .unwrap();
1033 let result2 = memory2
1034 .remember(text, RememberOptions::default())
1035 .await
1036 .unwrap();
1037
1038 assert_eq!(result1.entity_count(), result2.entity_count());
1040 }
1041
1042 #[test]
1047 fn test_convert_entity_type() {
1048 use crate::extraction::EntityType as ExtType;
1049
1050 assert_eq!(convert_entity_type(&ExtType::Person), EntityType::Person);
1051 assert_eq!(convert_entity_type(&ExtType::Project), EntityType::Project);
1052 assert_eq!(convert_entity_type(&ExtType::Topic), EntityType::Topic);
1053 assert_eq!(convert_entity_type(&ExtType::Task), EntityType::Task);
1054 assert_eq!(convert_entity_type(&ExtType::Note), EntityType::Note);
1055
1056 assert_eq!(
1058 convert_entity_type(&ExtType::Organization),
1059 EntityType::Note
1060 );
1061 assert_eq!(convert_entity_type(&ExtType::Preference), EntityType::Note);
1062 assert_eq!(convert_entity_type(&ExtType::Event), EntityType::Note);
1063 }
1064}
1065
1066impl Memory<
1071 crate::llm::SimLLMProvider,
1072 crate::embedding::SimEmbeddingProvider,
1073 crate::storage::SimStorageBackend,
1074 crate::storage::SimVectorBackend,
1075> {
1076 #[must_use]
1094 pub fn sim(seed: u64) -> Self {
1095 use crate::dst::SimConfig;
1096 use crate::embedding::SimEmbeddingProvider;
1097 use crate::llm::SimLLMProvider;
1098 use crate::storage::{SimStorageBackend, SimVectorBackend};
1099
1100 let llm = SimLLMProvider::with_seed(seed);
1101 let embedder = SimEmbeddingProvider::with_seed(seed);
1102 let vector = SimVectorBackend::new(seed);
1103 let storage = SimStorageBackend::new(SimConfig::with_seed(seed));
1104
1105 Self::new(llm, embedder, vector, storage)
1106 }
1107
1108 #[must_use]
1126 pub fn sim_with_config(_seed: u64, _config: MemoryConfig) -> Self {
1127 use crate::dst::SimConfig;
1128 use crate::embedding::SimEmbeddingProvider;
1129 use crate::llm::SimLLMProvider;
1130 use crate::storage::{SimStorageBackend, SimVectorBackend};
1131
1132 let llm = SimLLMProvider::with_seed(_seed);
1135 let embedder = SimEmbeddingProvider::with_seed(_seed);
1136 let vector = SimVectorBackend::new(_seed);
1137 let storage = SimStorageBackend::new(SimConfig::with_seed(_seed));
1138
1139 Self::new(llm, embedder, vector, storage)
1140 }
1141}
1142
1143#[cfg(test)]
1149mod dst_tests {
1150 use super::*;
1151 use crate::constants::EMBEDDING_DIMENSIONS_COUNT;
1152 use crate::dst::{FaultConfig, FaultType, SimConfig, Simulation};
1153 use crate::embedding::SimEmbeddingProvider;
1154 use crate::llm::SimLLMProvider;
1155 use crate::storage::{SimStorageBackend, SimVectorBackend};
1156
1157 #[tokio::test]
1158 async fn test_remember_with_embedding_timeout() {
1159 let sim = Simulation::new(SimConfig::with_seed(42))
1161 .with_fault(FaultConfig::new(FaultType::EmbeddingTimeout, 1.0)); sim.run(|env| async move {
1164 let embedder = SimEmbeddingProvider::with_faults(42, env.faults.clone());
1166 let llm = SimLLMProvider::with_seed(42);
1167 let vector = SimVectorBackend::new(42);
1168 let storage = SimStorageBackend::new(SimConfig::with_seed(42));
1169 let mut memory = Memory::new(llm, embedder, vector, storage);
1170
1171 let result = memory
1173 .remember("Alice works at Acme", RememberOptions::default())
1174 .await?;
1175
1176 assert!(!result.entities.is_empty());
1178 assert!(result.entities[0].embedding.is_none());
1180
1181 Ok::<(), MemoryError>(())
1182 })
1183 .await
1184 .unwrap();
1185 }
1186
1187 #[tokio::test]
1188 async fn test_remember_with_embedding_rate_limit() {
1189 let sim = Simulation::new(SimConfig::with_seed(42))
1191 .with_fault(FaultConfig::new(FaultType::EmbeddingRateLimit, 0.5)); sim.run(|env| async move {
1194 let embedder = SimEmbeddingProvider::with_faults(42, env.faults.clone());
1195 let llm = SimLLMProvider::with_seed(42);
1196 let vector = SimVectorBackend::new(42);
1197 let storage = SimStorageBackend::new(SimConfig::with_seed(42));
1198 let mut memory = Memory::new(llm, embedder, vector, storage);
1199
1200 let mut successes = 0;
1202 let mut failures = 0;
1203
1204 for i in 0..10 {
1205 let result = memory
1206 .remember(&format!("Text {}", i), RememberOptions::default())
1207 .await;
1208
1209 assert!(result.is_ok()); let res = result.unwrap();
1211
1212 if res.entities[0].embedding.is_some() {
1213 successes += 1;
1214 } else {
1215 failures += 1;
1216 }
1217 }
1218
1219 assert!(successes > 0, "Should have some successful embeddings");
1221 assert!(failures > 0, "Should have some failed embeddings");
1222
1223 Ok::<(), MemoryError>(())
1224 })
1225 .await
1226 .unwrap();
1227 }
1228
1229 #[tokio::test]
1230 async fn test_remember_without_embeddings_option() {
1231 let sim = Simulation::new(SimConfig::with_seed(42));
1233
1234 sim.run(|_env| async move {
1235 let embedder = SimEmbeddingProvider::with_seed(42);
1236 let llm = SimLLMProvider::with_seed(42);
1237 let vector = SimVectorBackend::new(42);
1238 let storage = SimStorageBackend::new(SimConfig::with_seed(42));
1239 let mut memory = Memory::new(llm, embedder, vector, storage);
1240
1241 let result = memory
1243 .remember(
1244 "Alice works at Acme",
1245 RememberOptions::default().without_embeddings(),
1246 )
1247 .await?;
1248
1249 assert!(!result.entities.is_empty());
1251 assert!(result.entities[0].embedding.is_none());
1252
1253 Ok::<(), MemoryError>(())
1254 })
1255 .await
1256 .unwrap();
1257 }
1258
1259 #[tokio::test]
1260 async fn test_remember_embeddings_deterministic() {
1261 async fn run_with_seed(seed: u64) -> Vec<Vec<f32>> {
1263 let embedder = SimEmbeddingProvider::with_seed(seed);
1264 let llm = SimLLMProvider::with_seed(seed);
1265 let vector = SimVectorBackend::new(seed);
1266 let storage = SimStorageBackend::new(SimConfig::with_seed(seed));
1267 let mut memory = Memory::new(llm, embedder, vector, storage);
1268
1269 let result = memory
1270 .remember("Alice works at Acme", RememberOptions::default())
1271 .await
1272 .unwrap();
1273
1274 result
1275 .entities
1276 .into_iter()
1277 .filter_map(|e| e.embedding)
1278 .collect()
1279 }
1280
1281 let embeddings1 = run_with_seed(12345).await;
1282 let embeddings2 = run_with_seed(12345).await;
1283
1284 assert!(!embeddings1.is_empty());
1285 assert_eq!(embeddings1.len(), embeddings2.len());
1286
1287 for (e1, e2) in embeddings1.iter().zip(embeddings2.iter()) {
1289 assert_eq!(e1, e2, "Same seed must produce same embeddings");
1290 }
1291 }
1292
1293 #[tokio::test]
1294 async fn test_remember_embeddings_stored() {
1295 let sim = Simulation::new(SimConfig::with_seed(42));
1297
1298 sim.run(|_env| async move {
1299 let embedder = SimEmbeddingProvider::with_seed(42);
1300 let llm = SimLLMProvider::with_seed(42);
1301 let vector = SimVectorBackend::new(42);
1302 let storage = SimStorageBackend::new(SimConfig::with_seed(42));
1303 let mut memory = Memory::new(llm, embedder, vector, storage);
1304
1305 let result = memory
1307 .remember("Alice works at Acme", RememberOptions::default())
1308 .await?;
1309
1310 assert!(!result.entities.is_empty());
1311 let entity_id = result.entities[0].id.clone();
1312
1313 let retrieved = memory.storage.get_entity(&entity_id).await?;
1315 assert!(retrieved.is_some());
1316
1317 let entity = retrieved.unwrap();
1318 assert!(entity.embedding.is_some(), "Embedding should be stored");
1319 assert_eq!(
1320 entity.embedding.as_ref().unwrap().len(),
1321 EMBEDDING_DIMENSIONS_COUNT
1322 );
1323
1324 let embedding = entity.embedding.unwrap();
1326 let norm: f32 = embedding.iter().map(|x| x * x).sum::<f32>().sqrt();
1327 assert!((norm - 1.0).abs() < 0.01, "Embedding should be normalized");
1328
1329 Ok::<(), MemoryError>(())
1330 })
1331 .await
1332 .unwrap();
1333 }
1334
1335 #[tokio::test]
1336 async fn test_remember_batch_embeddings() {
1337 let sim = Simulation::new(SimConfig::with_seed(42));
1339
1340 sim.run(|_env| async move {
1341 let embedder = SimEmbeddingProvider::with_seed(42);
1342 let llm = SimLLMProvider::with_seed(42);
1343 let vector = SimVectorBackend::new(42);
1344 let storage = SimStorageBackend::new(SimConfig::with_seed(42));
1345 let mut memory = Memory::new(llm, embedder, vector, storage);
1346
1347 let result = memory
1349 .remember(
1350 "Alice works at Acme. Bob works at TechCo.",
1351 RememberOptions::default(),
1352 )
1353 .await?;
1354
1355 if result.entities.len() > 1 {
1357 for entity in &result.entities {
1358 if entity.embedding.is_some() {
1359 assert_eq!(
1360 entity.embedding.as_ref().unwrap().len(),
1361 EMBEDDING_DIMENSIONS_COUNT
1362 );
1363 }
1364 }
1365 }
1366
1367 Ok::<(), MemoryError>(())
1368 })
1369 .await
1370 .unwrap();
1371 }
1372
1373 #[tokio::test]
1374 async fn test_remember_with_service_unavailable() {
1375 let sim = Simulation::new(SimConfig::with_seed(42)).with_fault(FaultConfig::new(
1377 FaultType::EmbeddingServiceUnavailable,
1378 1.0,
1379 ));
1380
1381 sim.run(|env| async move {
1382 let embedder = SimEmbeddingProvider::with_faults(42, env.faults.clone());
1383 let llm = SimLLMProvider::with_seed(42);
1384 let vector = SimVectorBackend::new(42);
1385 let storage = SimStorageBackend::new(SimConfig::with_seed(42));
1386 let mut memory = Memory::new(llm, embedder, vector, storage);
1387
1388 let result = memory
1390 .remember("Alice works at Acme", RememberOptions::default())
1391 .await?;
1392
1393 assert!(!result.entities.is_empty());
1394 assert!(result.entities[0].embedding.is_none());
1396
1397 Ok::<(), MemoryError>(())
1398 })
1399 .await
1400 .unwrap();
1401 }
1402
1403 #[tokio::test]
1408 async fn test_recall_with_vector_search() {
1409 let sim = Simulation::new(SimConfig::with_seed(42));
1411
1412 sim.run(|_env| async move {
1413 let embedder = SimEmbeddingProvider::with_seed(42);
1414 let llm = SimLLMProvider::with_seed(42);
1415 let vector = SimVectorBackend::new(42);
1416 let storage = SimStorageBackend::new(SimConfig::with_seed(42));
1417 let mut memory = Memory::new(llm, embedder, vector, storage);
1418
1419 memory
1421 .remember("Alice works at Acme Corp", RememberOptions::default())
1422 .await?;
1423 memory
1424 .remember("Bob works at TechCo", RememberOptions::default())
1425 .await?;
1426
1427 let result = memory
1429 .recall("Who works at Acme?", RecallOptions::default())
1430 .await?;
1431
1432 assert!(!result.is_empty());
1434 assert!(result.iter().any(|e| e.name.contains("Alice")));
1436
1437 Ok::<(), MemoryError>(())
1438 })
1439 .await
1440 .unwrap();
1441 }
1442
1443 #[tokio::test]
1444 async fn test_recall_vector_search_timeout() {
1445 let sim = Simulation::new(SimConfig::with_seed(42))
1447 .with_fault(FaultConfig::new(FaultType::VectorSearchTimeout, 1.0));
1448
1449 sim.run(|env| async move {
1450 let embedder = SimEmbeddingProvider::with_seed(42);
1451 let llm = SimLLMProvider::with_seed(42);
1452 let vector = SimVectorBackend::with_faults(42, env.faults.clone());
1453 let storage = SimStorageBackend::new(SimConfig::with_seed(42));
1454 let mut memory = Memory::new(llm, embedder, vector, storage);
1455
1456 memory
1458 .remember("Alice works at Acme Corp", RememberOptions::default())
1459 .await?;
1460
1461 let result = memory.recall("Alice", RecallOptions::default()).await;
1463
1464 assert!(result.is_ok());
1466 let entities = result.unwrap();
1467 assert!(!entities.is_empty());
1468
1469 Ok::<(), MemoryError>(())
1470 })
1471 .await
1472 .unwrap();
1473 }
1474
1475 #[tokio::test]
1476 async fn test_recall_vector_deterministic() {
1477 let seed = 42;
1480
1481 let embedder = SimEmbeddingProvider::with_seed(seed);
1483 let llm = SimLLMProvider::with_seed(seed);
1484 let vector = SimVectorBackend::new(seed);
1485 let storage = SimStorageBackend::new(SimConfig::with_seed(seed));
1486 let mut memory1 = Memory::new(llm, embedder, vector, storage);
1487
1488 memory1
1489 .remember(
1490 "Alice works at Acme Corp",
1491 RememberOptions::default().without_extraction(),
1492 )
1493 .await
1494 .unwrap();
1495 memory1
1496 .remember(
1497 "Bob works at TechCo",
1498 RememberOptions::default().without_extraction(),
1499 )
1500 .await
1501 .unwrap();
1502 memory1
1503 .remember(
1504 "Charlie works at DataInc",
1505 RememberOptions::default().without_extraction(),
1506 )
1507 .await
1508 .unwrap();
1509
1510 let result1 = memory1
1511 .recall("works", RecallOptions::default().fast_only())
1512 .await
1513 .unwrap();
1514 let names1: Vec<String> = result1.iter().map(|e| e.name.clone()).collect();
1515
1516 let embedder2 = SimEmbeddingProvider::with_seed(seed);
1518 let llm2 = SimLLMProvider::with_seed(seed);
1519 let vector2 = SimVectorBackend::new(seed);
1520 let storage2 = SimStorageBackend::new(SimConfig::with_seed(seed));
1521 let mut memory2 = Memory::new(llm2, embedder2, vector2, storage2);
1522
1523 memory2
1524 .remember(
1525 "Alice works at Acme Corp",
1526 RememberOptions::default().without_extraction(),
1527 )
1528 .await
1529 .unwrap();
1530 memory2
1531 .remember(
1532 "Bob works at TechCo",
1533 RememberOptions::default().without_extraction(),
1534 )
1535 .await
1536 .unwrap();
1537 memory2
1538 .remember(
1539 "Charlie works at DataInc",
1540 RememberOptions::default().without_extraction(),
1541 )
1542 .await
1543 .unwrap();
1544
1545 let result2 = memory2
1546 .recall("works", RecallOptions::default().fast_only())
1547 .await
1548 .unwrap();
1549 let names2: Vec<String> = result2.iter().map(|e| e.name.clone()).collect();
1550
1551 assert!(!names1.is_empty(), "Should find results");
1553 assert_eq!(names1, names2, "Same seed must produce same ranking");
1554 }
1555
1556 #[tokio::test]
1557 async fn test_recall_vector_storage_partial_failure() {
1558 let sim = Simulation::new(SimConfig::with_seed(42))
1560 .with_fault(FaultConfig::new(FaultType::VectorStoreFail, 0.5)); sim.run(|env| async move {
1563 let embedder = SimEmbeddingProvider::with_seed(42);
1564 let llm = SimLLMProvider::with_seed(42);
1565 let vector = SimVectorBackend::with_faults(42, env.faults.clone());
1566 let storage = SimStorageBackend::new(SimConfig::with_seed(42));
1567 let mut memory = Memory::new(llm, embedder, vector, storage);
1568
1569 for i in 0..10 {
1572 let opts = RememberOptions::default().without_extraction();
1573 memory
1574 .remember(&format!("Entity number {}", i), opts)
1575 .await?;
1576 }
1577
1578 let result = memory.recall("Entity", RecallOptions::default()).await;
1580
1581 assert!(result.is_ok());
1582 let entities = result.unwrap();
1583 assert!(
1585 !entities.is_empty(),
1586 "Should find entities even with vector storage failures"
1587 );
1588
1589 Ok::<(), MemoryError>(())
1590 })
1591 .await
1592 .unwrap();
1593 }
1594}