1use mentedb_core::MemoryNode;
2use mentedb_core::edge::EdgeType;
3use mentedb_core::memory::MemoryType;
4use mentedb_core::types::MemoryId;
5
6use crate::llm::{
7 CognitiveLlmService, ContradictionVerdict, InvalidationVerdict, LlmJudge, MemorySummary,
8};
9
10#[derive(Debug, Clone)]
11pub enum InferredAction {
12 FlagContradiction {
13 existing: MemoryId,
14 new: MemoryId,
15 reason: String,
16 },
17 MarkObsolete {
18 memory: MemoryId,
19 superseded_by: MemoryId,
20 },
21 InvalidateMemory {
23 memory: MemoryId,
24 superseded_by: MemoryId,
25 valid_until: u64,
26 },
27 DeduplicateExact {
32 duplicate: MemoryId,
33 keeper: MemoryId,
34 },
35 UpdateContent {
37 memory: MemoryId,
38 new_content: String,
39 reason: String,
40 },
41 CreateEdge {
42 source: MemoryId,
43 target: MemoryId,
44 edge_type: EdgeType,
45 weight: f32,
46 },
47 UpdateConfidence {
48 memory: MemoryId,
49 new_confidence: f32,
50 },
51 PropagateBeliefChange {
52 root: MemoryId,
53 delta: f32,
54 },
55}
56
57fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
58 if a.len() != b.len() || a.is_empty() {
59 return 0.0;
60 }
61 let mut dot = 0.0f32;
62 let mut norm_a = 0.0f32;
63 let mut norm_b = 0.0f32;
64 for i in 0..a.len() {
65 dot += a[i] * b[i];
66 norm_a += a[i] * a[i];
67 norm_b += b[i] * b[i];
68 }
69 let denom = norm_a.sqrt() * norm_b.sqrt();
70 if denom == 0.0 { 0.0 } else { dot / denom }
71}
72
73#[derive(Debug, Clone)]
75pub struct WriteInferenceConfig {
76 pub contradiction_threshold: f32,
78 pub obsolete_threshold: f32,
80 pub related_min: f32,
82 pub related_max: f32,
84 pub correction_threshold: f32,
86 pub confidence_decay_factor: f32,
88 pub confidence_floor: f32,
90 pub value_update_enabled: bool,
98 pub value_update_min_similarity: f32,
100 pub value_update_prefix_share: f32,
103 pub value_update_max_tail: usize,
106}
107
108impl Default for WriteInferenceConfig {
109 fn default() -> Self {
110 Self {
111 contradiction_threshold: 0.95,
112 obsolete_threshold: 0.85,
113 related_min: 0.6,
114 related_max: 0.85,
115 correction_threshold: 0.5,
116 confidence_decay_factor: 0.5,
117 confidence_floor: 0.1,
118 value_update_enabled: true,
119 value_update_min_similarity: 0.88,
120 value_update_prefix_share: 0.6,
121 value_update_max_tail: 4,
122 }
123 }
124}
125
126fn frame_tokens(s: &str) -> Vec<String> {
129 s.split(|c: char| !c.is_alphanumeric())
130 .filter(|t| !t.is_empty())
131 .map(|t| t.to_lowercase())
132 .collect()
133}
134
135fn is_value_update(new: &str, old: &str, prefix_share: f32, max_tail: usize) -> bool {
141 let nt = frame_tokens(new);
142 let ot = frame_tokens(old);
143 if nt.is_empty() || ot.is_empty() || nt == ot {
144 return false;
145 }
146 let prefix = nt.iter().zip(ot.iter()).take_while(|(a, b)| a == b).count();
147 let longest = nt.len().max(ot.len());
148 let tail = longest - prefix;
149 prefix as f32 / longest as f32 >= prefix_share && tail > 0 && tail <= max_tail
150}
151
152pub struct WriteInferenceEngine {
153 config: WriteInferenceConfig,
154}
155
156impl WriteInferenceEngine {
157 pub fn new() -> Self {
158 Self {
159 config: WriteInferenceConfig::default(),
160 }
161 }
162
163 pub fn with_config(config: WriteInferenceConfig) -> Self {
164 Self { config }
165 }
166
167 pub fn infer_on_write(
168 &self,
169 new_memory: &MemoryNode,
170 existing_memories: &[MemoryNode],
171 existing_edges: &[(MemoryId, MemoryId, EdgeType)],
172 ) -> Vec<InferredAction> {
173 let _ = existing_edges; let mut actions = Vec::new();
175
176 for existing in existing_memories {
177 if existing.id == new_memory.id {
178 continue;
179 }
180
181 let sim = cosine_similarity(&new_memory.embedding, &existing.embedding);
182
183 if existing.content == new_memory.content
197 && sim > self.config.obsolete_threshold
198 && new_memory.created_at > existing.created_at
199 {
200 actions.push(InferredAction::DeduplicateExact {
201 duplicate: existing.id,
202 keeper: new_memory.id,
203 });
204 } else if self.config.value_update_enabled
205 && sim >= self.config.value_update_min_similarity
206 && new_memory.created_at > existing.created_at
207 && existing.agent_id == new_memory.agent_id
208 && existing.user_id == new_memory.user_id
209 && is_value_update(
210 &new_memory.content,
211 &existing.content,
212 self.config.value_update_prefix_share,
213 self.config.value_update_max_tail,
214 )
215 {
216 actions.push(InferredAction::CreateEdge {
219 source: new_memory.id,
220 target: existing.id,
221 edge_type: EdgeType::Supersedes,
222 weight: sim,
223 });
224 actions.push(InferredAction::UpdateConfidence {
225 memory: existing.id,
226 new_confidence: (existing.confidence * self.config.confidence_decay_factor)
227 .max(self.config.confidence_floor),
228 });
229 } else if sim > self.config.related_min && sim <= self.config.related_max {
230 actions.push(InferredAction::CreateEdge {
231 source: new_memory.id,
232 target: existing.id,
233 edge_type: EdgeType::Related,
234 weight: sim,
235 });
236 }
237 }
238
239 if new_memory.memory_type == MemoryType::Correction
241 && let Some(original) = existing_memories
242 .iter()
243 .filter(|m| m.id != new_memory.id)
244 .max_by(|a, b| {
245 cosine_similarity(&new_memory.embedding, &a.embedding)
246 .partial_cmp(&cosine_similarity(&new_memory.embedding, &b.embedding))
247 .unwrap_or(std::cmp::Ordering::Equal)
248 })
249 {
250 let sim = cosine_similarity(&new_memory.embedding, &original.embedding);
251 if sim > self.config.correction_threshold {
252 actions.push(InferredAction::CreateEdge {
253 source: new_memory.id,
254 target: original.id,
255 edge_type: EdgeType::Supersedes,
256 weight: 1.0,
257 });
258 actions.push(InferredAction::UpdateConfidence {
259 memory: original.id,
260 new_confidence: (original.confidence * self.config.confidence_decay_factor)
261 .max(self.config.confidence_floor),
262 });
263 }
264 }
265
266 actions
267 }
268
269 pub async fn infer_on_write_with_llm<J: LlmJudge>(
274 &self,
275 new_memory: &MemoryNode,
276 existing_memories: &[MemoryNode],
277 existing_edges: &[(MemoryId, MemoryId, EdgeType)],
278 llm: &CognitiveLlmService<J>,
279 ) -> Vec<InferredAction> {
280 let _ = existing_edges;
281 let mut actions = Vec::new();
282
283 let new_summary = memory_to_summary(new_memory);
284
285 for existing in existing_memories {
286 if existing.id == new_memory.id {
287 continue;
288 }
289
290 let sim = cosine_similarity(&new_memory.embedding, &existing.embedding);
291
292 if sim > 0.5 && existing.agent_id == new_memory.agent_id {
295 let old_summary = memory_to_summary(existing);
296
297 if let Ok(verdict) = llm.judge_invalidation(&old_summary, &new_summary).await {
299 match verdict {
300 InvalidationVerdict::Invalidate { reason: _ } => {
301 actions.push(InferredAction::InvalidateMemory {
304 memory: existing.id,
305 superseded_by: new_memory.id,
306 valid_until: new_memory.created_at,
307 });
308 actions.push(InferredAction::UpdateConfidence {
309 memory: existing.id,
310 new_confidence: (existing.confidence
311 * self.config.confidence_decay_factor)
312 .max(self.config.confidence_floor),
313 });
314 continue;
316 }
317 InvalidationVerdict::Update {
318 merged_content,
319 reason,
320 } => {
321 actions.push(InferredAction::UpdateContent {
322 memory: existing.id,
323 new_content: merged_content,
324 reason,
325 });
326 continue;
327 }
328 InvalidationVerdict::Keep { .. } => {
329 }
331 }
332 }
333
334 if sim > 0.7
336 && existing.content != new_memory.content
337 && let Ok(verdict) = llm.detect_contradiction(&old_summary, &new_summary).await
338 {
339 match verdict {
340 ContradictionVerdict::Contradicts { reason } => {
341 actions.push(InferredAction::FlagContradiction {
342 existing: existing.id,
343 new: new_memory.id,
344 reason,
345 });
346 }
347 ContradictionVerdict::Supersedes { winner, reason: _ } => {
348 let winner_is_new = winner == new_memory.id.to_string();
349 let (obsolete, superseder) = if winner_is_new {
350 (existing.id, new_memory.id)
351 } else {
352 (new_memory.id, existing.id)
353 };
354 actions.push(InferredAction::InvalidateMemory {
355 memory: obsolete,
356 superseded_by: superseder,
357 valid_until: new_memory.created_at,
358 });
359 }
360 ContradictionVerdict::Compatible { .. } => {}
361 }
362 }
363 }
364
365 if sim > self.config.related_min && sim <= self.config.related_max {
367 actions.push(InferredAction::CreateEdge {
368 source: new_memory.id,
369 target: existing.id,
370 edge_type: EdgeType::Related,
371 weight: sim,
372 });
373 }
374 }
375
376 if new_memory.memory_type == MemoryType::Correction
378 && let Some(original) = existing_memories
379 .iter()
380 .filter(|m| m.id != new_memory.id)
381 .max_by(|a, b| {
382 cosine_similarity(&new_memory.embedding, &a.embedding)
383 .partial_cmp(&cosine_similarity(&new_memory.embedding, &b.embedding))
384 .unwrap_or(std::cmp::Ordering::Equal)
385 })
386 {
387 let sim = cosine_similarity(&new_memory.embedding, &original.embedding);
388 if sim > self.config.correction_threshold {
389 actions.push(InferredAction::InvalidateMemory {
390 memory: original.id,
391 superseded_by: new_memory.id,
392 valid_until: new_memory.created_at,
393 });
394 actions.push(InferredAction::UpdateConfidence {
395 memory: original.id,
396 new_confidence: (original.confidence * self.config.confidence_decay_factor)
397 .max(self.config.confidence_floor),
398 });
399 }
400 }
401
402 actions
403 }
404}
405
406impl Default for WriteInferenceEngine {
407 fn default() -> Self {
408 Self::new()
409 }
410}
411
412fn memory_to_summary(m: &MemoryNode) -> MemorySummary {
413 MemorySummary {
414 id: m.id,
415 content: m.content.clone(),
416 memory_type: m.memory_type,
417 confidence: m.confidence,
418 created_at: m.created_at,
419 }
420}
421
422#[cfg(test)]
423mod tests {
424 use super::*;
425 use mentedb_core::memory::MemoryType;
426 use mentedb_core::types::AgentId;
427
428 use crate::llm::MockLlmJudge;
429
430 fn make_memory(content: &str, embedding: Vec<f32>, mem_type: MemoryType) -> MemoryNode {
431 let mut m = MemoryNode::new(AgentId::new(), mem_type, content.to_string(), embedding);
432 m.created_at = 1000;
433 m
434 }
435
436 #[test]
437 fn test_heuristic_does_not_flag_contradiction_from_similarity() {
438 let agent = AgentId::new();
445 let mut existing =
446 make_memory("uses PostgreSQL", vec![1.0, 0.0, 0.0], MemoryType::Semantic);
447 existing.agent_id = agent;
448
449 let mut new_mem = make_memory("uses MySQL", vec![0.99, 0.01, 0.0], MemoryType::Semantic);
450 new_mem.agent_id = agent;
451 new_mem.created_at = 2000;
452
453 let engine = WriteInferenceEngine::new();
454 let actions = engine.infer_on_write(&new_mem, &[existing], &[]);
455 assert!(
456 !actions
457 .iter()
458 .any(|a| matches!(a, InferredAction::FlagContradiction { .. })),
459 "Heuristic must not flag contradictions from bare similarity, got: {:?}",
460 actions
461 );
462 }
463
464 #[test]
465 fn value_update_supersedes_the_old_fact() {
466 let agent = AgentId::new();
470 let mut existing = make_memory(
471 "The user's favorite coffee order is a cortado",
472 vec![1.0, 0.0, 0.0],
473 MemoryType::Semantic,
474 );
475 existing.agent_id = agent;
476 let mut new_mem = make_memory(
477 "The user's favorite coffee order is a flat white",
478 vec![0.98, 0.02, 0.0],
479 MemoryType::Semantic,
480 );
481 new_mem.agent_id = agent;
482 new_mem.created_at = 2000;
483
484 let engine = WriteInferenceEngine::new();
485 let actions = engine.infer_on_write(&new_mem, &[existing.clone()], &[]);
486 assert!(
487 actions.iter().any(|a| matches!(
488 a,
489 InferredAction::CreateEdge { source, target, edge_type: EdgeType::Supersedes, .. }
490 if *source == new_mem.id && *target == existing.id
491 )),
492 "value update must create a Supersedes edge, got: {actions:?}"
493 );
494 assert!(
495 actions
496 .iter()
497 .any(|a| matches!(a, InferredAction::UpdateConfidence { memory, .. } if *memory == existing.id)),
498 "superseded fact's confidence must decay"
499 );
500 }
501
502 #[test]
503 fn subject_change_is_not_a_value_update() {
504 let agent = AgentId::new();
507 let mut existing = make_memory(
508 "My dog is allergic to chicken",
509 vec![1.0, 0.0, 0.0],
510 MemoryType::Semantic,
511 );
512 existing.agent_id = agent;
513 let mut new_mem = make_memory(
514 "My cat is allergic to chicken",
515 vec![0.99, 0.01, 0.0],
516 MemoryType::Semantic,
517 );
518 new_mem.agent_id = agent;
519 new_mem.created_at = 2000;
520
521 let engine = WriteInferenceEngine::new();
522 let actions = engine.infer_on_write(&new_mem, &[existing], &[]);
523 assert!(
524 !actions.iter().any(|a| matches!(
525 a,
526 InferredAction::CreateEdge {
527 edge_type: EdgeType::Supersedes,
528 ..
529 }
530 )),
531 "a subject change must not supersede, got: {actions:?}"
532 );
533 }
534
535 #[test]
536 fn value_update_respects_owner_and_config() {
537 let mut existing = make_memory(
539 "The user's favorite coffee order is a cortado",
540 vec![1.0, 0.0, 0.0],
541 MemoryType::Semantic,
542 );
543 existing.agent_id = AgentId::new();
544 let mut new_mem = make_memory(
545 "The user's favorite coffee order is a flat white",
546 vec![0.98, 0.02, 0.0],
547 MemoryType::Semantic,
548 );
549 new_mem.agent_id = AgentId::new(); new_mem.created_at = 2000;
551
552 let engine = WriteInferenceEngine::new();
553 let actions = engine.infer_on_write(&new_mem, &[existing.clone()], &[]);
554 assert!(
555 !actions.iter().any(|a| matches!(
556 a,
557 InferredAction::CreateEdge {
558 edge_type: EdgeType::Supersedes,
559 ..
560 }
561 )),
562 "cross-owner memories must never supersede each other"
563 );
564
565 let agent = AgentId::new();
567 existing.agent_id = agent;
568 new_mem.agent_id = agent;
569 let engine_off = WriteInferenceEngine::with_config(WriteInferenceConfig {
570 value_update_enabled: false,
571 ..WriteInferenceConfig::default()
572 });
573 let actions = engine_off.infer_on_write(&new_mem, &[existing], &[]);
574 assert!(
575 !actions.iter().any(|a| matches!(
576 a,
577 InferredAction::CreateEdge {
578 edge_type: EdgeType::Supersedes,
579 ..
580 }
581 )),
582 "disabled rule must not supersede"
583 );
584 }
585
586 #[test]
587 fn is_value_update_frame_analysis() {
588 assert!(is_value_update(
590 "my phone number is 4321",
591 "my phone number is 1234",
592 0.6,
593 4
594 ));
595 assert!(!is_value_update(
597 "my cat is allergic to chicken",
598 "my dog is allergic to chicken",
599 0.6,
600 4
601 ));
602 assert!(!is_value_update("same fact", "same fact", 0.6, 4));
604 assert!(!is_value_update(
606 "the deploy runs every friday and sarah reviews it after standup",
607 "the deploy runs every friday unless the release train is frozen for the quarter end audit",
608 0.6,
609 4
610 ));
611 }
612
613 #[test]
614 fn test_byte_identical_duplicate_is_deduplicated() {
615 let agent = AgentId::new();
618 let mut existing = make_memory(
619 "Ran command: ls -la",
620 vec![1.0, 0.0, 0.0],
621 MemoryType::Episodic,
622 );
623 existing.agent_id = agent;
624
625 let mut new_mem = make_memory(
626 "Ran command: ls -la",
627 vec![1.0, 0.0, 0.0],
628 MemoryType::Episodic,
629 );
630 new_mem.agent_id = agent;
631 new_mem.created_at = 2000;
632
633 let engine = WriteInferenceEngine::new();
634 let actions = engine.infer_on_write(&new_mem, &[existing], &[]);
635 assert!(
636 actions
637 .iter()
638 .any(|a| matches!(a, InferredAction::DeduplicateExact { .. })),
639 "Expected identical duplicate to be deduplicated, got: {:?}",
640 actions
641 );
642 }
643
644 #[test]
645 fn test_moderate_similarity_creates_edge() {
646 let existing = make_memory("topic A", vec![1.0, 0.0, 0.0], MemoryType::Semantic);
647 let new_mem = make_memory("topic B", vec![0.7, 0.714, 0.0], MemoryType::Semantic);
649
650 let engine = WriteInferenceEngine::new();
651 let actions = engine.infer_on_write(&new_mem, &[existing], &[]);
652 assert!(
653 actions.iter().any(|a| matches!(
654 a,
655 InferredAction::CreateEdge {
656 edge_type: EdgeType::Related,
657 ..
658 }
659 )),
660 "Expected CreateEdge Related, got: {:?}",
661 actions
662 );
663 }
664
665 #[tokio::test]
666 async fn test_llm_invalidation_emits_temporal_actions() {
667 let agent = AgentId::new();
668 let mut existing = make_memory(
669 "Alice works at Acme",
670 vec![0.8, 0.6, 0.0],
671 MemoryType::Semantic,
672 );
673 existing.agent_id = agent;
674
675 let mut new_mem = make_memory(
676 "Alice joined Google last week",
677 vec![0.75, 0.65, 0.1],
678 MemoryType::Semantic,
679 );
680 new_mem.agent_id = agent;
681 new_mem.created_at = 2000;
682
683 let judge =
684 MockLlmJudge::new(r#"{"verdict": "invalidate", "reason": "Alice changed jobs"}"#);
685 let llm = CognitiveLlmService::new(judge);
686 let engine = WriteInferenceEngine::new();
687 let actions = engine
688 .infer_on_write_with_llm(&new_mem, &[existing.clone()], &[], &llm)
689 .await;
690
691 assert!(
692 actions
693 .iter()
694 .any(|a| matches!(a, InferredAction::InvalidateMemory { .. })),
695 "Expected InvalidateMemory from LLM verdict, got: {:?}",
696 actions
697 );
698 let invalidation_count = actions
702 .iter()
703 .filter(|a| {
704 matches!(
705 a,
706 InferredAction::InvalidateMemory { .. } | InferredAction::MarkObsolete { .. }
707 )
708 })
709 .count();
710 assert_eq!(
711 invalidation_count, 1,
712 "Expected a single invalidation action, got: {:?}",
713 actions
714 );
715 }
716
717 #[tokio::test]
718 async fn test_llm_update_emits_update_content() {
719 let agent = AgentId::new();
720 let mut existing = make_memory(
721 "Project uses React",
722 vec![0.8, 0.6, 0.0],
723 MemoryType::Semantic,
724 );
725 existing.agent_id = agent;
726
727 let mut new_mem = make_memory(
728 "Project migrated from React to Vue",
729 vec![0.75, 0.65, 0.1],
730 MemoryType::Semantic,
731 );
732 new_mem.agent_id = agent;
733 new_mem.created_at = 2000;
734
735 let judge = MockLlmJudge::new(
736 r#"{"verdict": "update", "merged_content": "Project migrated from React to Vue in Q2", "reason": "adds temporal context"}"#,
737 );
738 let llm = CognitiveLlmService::new(judge);
739 let engine = WriteInferenceEngine::new();
740 let actions = engine
741 .infer_on_write_with_llm(&new_mem, &[existing.clone()], &[], &llm)
742 .await;
743
744 assert!(
745 actions
746 .iter()
747 .any(|a| matches!(a, InferredAction::UpdateContent { .. })),
748 "Expected UpdateContent from LLM update verdict, got: {:?}",
749 actions
750 );
751 }
752
753 #[tokio::test]
754 async fn test_llm_keep_falls_through_to_contradiction_check() {
755 let agent = AgentId::new();
756 let mut existing = make_memory("Prefers tabs", vec![0.9, 0.44, 0.0], MemoryType::Semantic);
758 existing.agent_id = agent;
759
760 let mut new_mem = make_memory(
761 "Prefers spaces",
762 vec![0.88, 0.47, 0.0],
763 MemoryType::Semantic,
764 );
765 new_mem.agent_id = agent;
766 new_mem.created_at = 2000;
767
768 let judge = MockLlmJudge::new(
769 r#"{"verdict": "compatible", "reason": "different formatting preferences"}"#,
770 );
771 let llm = CognitiveLlmService::new(judge);
772 let engine = WriteInferenceEngine::new();
773 let actions = engine
774 .infer_on_write_with_llm(&new_mem, &[existing.clone()], &[], &llm)
775 .await;
776
777 assert!(
780 !actions
781 .iter()
782 .any(|a| matches!(a, InferredAction::FlagContradiction { .. })),
783 "Should not flag contradiction when LLM says compatible",
784 );
785 }
786}