Skip to main content

mentedb_cognitive/
write_inference.rs

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    /// Set valid_until on the old memory instead of deleting it.
22    InvalidateMemory {
23        memory: MemoryId,
24        superseded_by: MemoryId,
25        valid_until: u64,
26    },
27    /// A byte-identical re-save: invalidate the duplicate and link it to the
28    /// surviving copy with a `Derived` edge (deduplication lineage), NOT a
29    /// `Supersedes` edge. A Supersedes edge here reads as a memory superseding an
30    /// exact copy of itself, which only shows up as noise in the conflict view.
31    DeduplicateExact {
32        duplicate: MemoryId,
33        keeper: MemoryId,
34    },
35    /// Update the content of an existing memory with merged information.
36    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/// What write inference does when it finds that a new memory contradicts an
74/// older one (a value update on the same sentence frame, or an LLM verdict that
75/// one memory replaces the other).
76#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
77#[serde(rename_all = "snake_case")]
78pub enum ContradictionPolicy {
79    /// Close the losing memory: set its valid_until, record a Supersedes edge
80    /// and decay its confidence, so recall serves only the winner.
81    #[default]
82    Supersede,
83    /// Hint mode: keep both memories valid and record a Contradicts edge
84    /// between them, so the caller can surface the conflict and decide. Under
85    /// this policy recall keeps serving both sides of a Contradicts edge.
86    /// Memories stored with the explicit Correction type still supersede,
87    /// because the caller already made that decision.
88    Flag,
89}
90
91impl std::str::FromStr for ContradictionPolicy {
92    type Err = String;
93
94    fn from_str(s: &str) -> Result<Self, Self::Err> {
95        match s.trim().to_ascii_lowercase().as_str() {
96            "supersede" => Ok(Self::Supersede),
97            "flag" => Ok(Self::Flag),
98            other => Err(format!(
99                "unknown contradiction policy '{other}' (expected supersede or flag)"
100            )),
101        }
102    }
103}
104
105/// Configuration for write-time inference thresholds.
106#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
107#[serde(default)]
108pub struct WriteInferenceConfig {
109    /// Similarity above which two memories may contradict (default: 0.95).
110    pub contradiction_threshold: f32,
111    /// Similarity above which an older memory is marked obsolete (default: 0.85).
112    pub obsolete_threshold: f32,
113    /// Minimum similarity for creating a Related edge (default: 0.6).
114    pub related_min: f32,
115    /// Maximum similarity for creating a Related edge (default: 0.85).
116    pub related_max: f32,
117    /// Minimum similarity for a Correction to supersede (default: 0.5).
118    pub correction_threshold: f32,
119    /// Multiplier applied to original confidence on correction (default: 0.5).
120    pub confidence_decay_factor: f32,
121    /// Minimum confidence after decay (default: 0.1).
122    pub confidence_floor: f32,
123    /// Whether the value-update rule runs: a new memory that shares an older
124    /// same-owner memory's sentence frame (long common token prefix) but ends
125    /// in a different value supersedes it, so "favorite coffee is a cortado"
126    /// followed by "favorite coffee is a flat white" keeps only the correction.
127    /// This is the cheap, LLM-free slice of contradiction handling; unlike bare
128    /// cosine (which cannot tell paraphrase from contradiction and was removed
129    /// at ~0% precision), the shared-frame test only fires on the update shape.
130    pub value_update_enabled: bool,
131    /// Minimum embedding similarity for the value-update rule (default: 0.3).
132    ///
133    /// This is a loose topical floor, not the primary gate. A value correction
134    /// inherently drives cosine down (the changed value is most of a short
135    /// fact's meaning), so real corrections land well below a high cosine bar:
136    /// measured on Titan V2 at 256 dims, "favorite coffee is a cortado" vs
137    /// "... a flat white" is 0.61 and "phone is 1234" vs "... 4321" is 0.51;
138    /// on OpenAI text-embedding-3-small they are 0.80 and 0.79. A 0.88 gate
139    /// (the first default) was unreachable on either embedder, and even 0.5
140    /// sat inside the real-pair range and silently gated identical-frame
141    /// pairs in production (observed on Cohere 1024: "uses Next.js for the
142    /// front end" vs "uses React for the front end" never superseded). The
143    /// embedder-independent `is_value_update` shared frame test carries the
144    /// precision here; the floor only rejects genuinely off-topic pairs that
145    /// a shared token frame collided on, which sit near 0.2 and below.
146    pub value_update_min_similarity: f32,
147    /// Minimum fraction of the longer memory's tokens that must be shared
148    /// frame, the common prefix plus the common suffix around the changed
149    /// value (default: 0.6).
150    pub value_update_prefix_share: f32,
151    /// Maximum differing value tokens per side (default: 4). Larger middles
152    /// mean genuinely different statements, not a value change.
153    pub value_update_max_tail: usize,
154    /// What to do when a new memory contradicts an older one (default:
155    /// supersede the older one).
156    pub contradiction_policy: ContradictionPolicy,
157}
158
159impl Default for WriteInferenceConfig {
160    fn default() -> Self {
161        Self {
162            contradiction_threshold: 0.95,
163            obsolete_threshold: 0.85,
164            related_min: 0.6,
165            related_max: 0.85,
166            correction_threshold: 0.5,
167            confidence_decay_factor: 0.5,
168            confidence_floor: 0.1,
169            value_update_enabled: true,
170            value_update_min_similarity: 0.3,
171            value_update_prefix_share: 0.6,
172            value_update_max_tail: 4,
173            contradiction_policy: ContradictionPolicy::Supersede,
174        }
175    }
176}
177
178/// Lowercased alphanumeric tokens of a memory's content, the unit of the
179/// value-update frame comparison.
180fn frame_tokens(s: &str) -> Vec<String> {
181    s.split(|c: char| !c.is_alphanumeric())
182        .filter(|t| !t.is_empty())
183        .map(|t| t.to_lowercase())
184        .collect()
185}
186
187/// True when `new` and `old` share the same sentence frame with a short
188/// differing value somewhere inside it: a common token prefix and a common
189/// token suffix that together cover most of the sentence, with a small
190/// non-identical middle on each side. The value may sit at the end ("the
191/// user's favorite coffee order is a cortado" vs "... is a flat white") or
192/// mid sentence ("i use next.js for the front end" vs "i use react for the
193/// front end", where the value is wrapped by a shared tail). A subject swap
194/// at the head ("my dog is allergic to chicken" vs "my cat is allergic to
195/// chicken") shares at most one leading token, so the two token minimum
196/// prefix keeps it out: changing the subject is a different fact, not a
197/// value update. The original prefix-only test could not see shared tails,
198/// which made every mid sentence value change invisible and left
199/// contradicting facts coexisting.
200///
201/// Mid sentence middles that are purely numeric on both sides ("rule 3" vs
202/// "rule 4", "room 2" vs "room 5") are identifiers naming different things,
203/// not changed values, and such facts legitimately coexist; only the end
204/// position keeps numeric updates ("phone is 1234" vs "phone is 4321").
205/// Precision beats recall here on purpose: a wrong supersession silently
206/// deletes a true fact, while a missed one is caught by the LLM
207/// contradiction layer.
208fn is_value_update(new: &str, old: &str, frame_share: f32, max_value: usize) -> bool {
209    let nt = frame_tokens(new);
210    let ot = frame_tokens(old);
211    if nt.is_empty() || ot.is_empty() || nt == ot {
212        return false;
213    }
214    // Very short sentences cannot carry enough frame to tell an update from
215    // an enumeration: "directive number 1" and "directive number 2" are two
216    // items, not one changed value. Real value sentences name a relation and
217    // a value and run longer ("my phone number is 4321" is five tokens).
218    if nt.len().min(ot.len()) < 4 {
219        return false;
220    }
221    let prefix = nt.iter().zip(ot.iter()).take_while(|(a, b)| a == b).count();
222    if prefix < 2 {
223        return false;
224    }
225    let suffix = nt
226        .iter()
227        .rev()
228        .zip(ot.iter().rev())
229        .take_while(|(a, b)| a == b)
230        .count()
231        // The suffix may not reach back into the prefix on the shorter side.
232        .min(nt.len().min(ot.len()) - prefix);
233    let longest = nt.len().max(ot.len());
234    let value_new = nt.len() - prefix - suffix;
235    let value_old = ot.len() - prefix - suffix;
236    if suffix > 0 {
237        let all_numeric =
238            |tokens: &[String]| tokens.iter().all(|t| t.chars().all(|c| c.is_ascii_digit()));
239        if all_numeric(&nt[prefix..nt.len() - suffix])
240            && all_numeric(&ot[prefix..ot.len() - suffix])
241        {
242            return false;
243        }
244    }
245    (prefix + suffix) as f32 / longest as f32 >= frame_share
246        && value_new.max(value_old) > 0
247        && value_new <= max_value
248        && value_old <= max_value
249}
250
251/// Decide which of two value-frame-matching memories supersedes the other, and
252/// which is left stale. Order-independent: a `Correction` supersedes a plain
253/// fact whichever was stored first (the extractor types a correction explicitly,
254/// so it is the authoritative value); otherwise the later-created memory wins.
255/// Returns `(winner_id, loser_id, loser_confidence)`, or `None` when there is no
256/// correction signal and the timestamps are equal, in which case the two are
257/// kept as coexisting facts rather than guessing a direction.
258fn value_update_resolution(
259    new: &MemoryNode,
260    existing: &MemoryNode,
261) -> Option<(MemoryId, MemoryId, f32)> {
262    let new_correction = new.memory_type == MemoryType::Correction;
263    let old_correction = existing.memory_type == MemoryType::Correction;
264    if new_correction && !old_correction {
265        return Some((new.id, existing.id, existing.confidence));
266    }
267    if old_correction && !new_correction {
268        return Some((existing.id, new.id, new.confidence));
269    }
270    if new.created_at > existing.created_at {
271        return Some((new.id, existing.id, existing.confidence));
272    }
273    if existing.created_at > new.created_at {
274        return Some((existing.id, new.id, new.confidence));
275    }
276    None
277}
278
279pub struct WriteInferenceEngine {
280    config: WriteInferenceConfig,
281}
282
283impl WriteInferenceEngine {
284    pub fn new() -> Self {
285        Self {
286            config: WriteInferenceConfig::default(),
287        }
288    }
289
290    pub fn with_config(config: WriteInferenceConfig) -> Self {
291        Self { config }
292    }
293
294    /// Whether `new` and `old` look like a value update (same sentence frame,
295    /// changed value tail) under this engine's configured thresholds. A
296    /// deduplication pass that merges near-duplicates uses this as a guard so a
297    /// correction ("coffee is a cortado" then "coffee is a flat white") is
298    /// never collapsed into the fact it corrects; the two must stay distinct
299    /// for supersession to resolve them.
300    pub fn looks_like_value_update(&self, new: &str, old: &str) -> bool {
301        is_value_update(
302            new,
303            old,
304            self.config.value_update_prefix_share,
305            self.config.value_update_max_tail,
306        )
307    }
308
309    pub fn infer_on_write(
310        &self,
311        new_memory: &MemoryNode,
312        existing_memories: &[MemoryNode],
313        existing_edges: &[(MemoryId, MemoryId, EdgeType)],
314    ) -> Vec<InferredAction> {
315        let _ = existing_edges; // reserved for future graph-aware inference
316        let mut actions = Vec::new();
317        // Value-frame-matched pairs are owned by the value-update rule below; the
318        // broad correction handler skips them so a pair is never resolved twice.
319        let mut value_update_handled: Vec<MemoryId> = Vec::new();
320
321        for existing in existing_memories {
322            if existing.id == new_memory.id {
323                continue;
324            }
325
326            let sim = cosine_similarity(&new_memory.embedding, &existing.embedding);
327
328            // Cosine similarity alone cannot tell a reworded duplicate from a
329            // genuine contradiction: paraphrases ("uses Postgres" / "uses
330            // PostgreSQL") and opposites ("prefers tabs" / "prefers spaces")
331            // both land in the high-similarity band. So the cheap write-time
332            // heuristic makes only the two judgments similarity CAN support:
333            //   1. byte-identical content -> a true duplicate, supersede it
334            //   2. moderate similarity    -> a Related edge
335            // Every semantic decision (merge, paraphrase dedup, real
336            // contradiction, which-fact-wins supersession) is deferred to the
337            // LLM paths (`consolidate_memories`, `detect_conflicts_with_llm`),
338            // which read the actual text. Flagging contradictions from bare
339            // similarity here produced ~0% precision (every reworded re-save
340            // looked like a contradiction), so it was removed.
341            if existing.content == new_memory.content
342                && sim > self.config.obsolete_threshold
343                && new_memory.created_at > existing.created_at
344            {
345                actions.push(InferredAction::DeduplicateExact {
346                    duplicate: existing.id,
347                    keeper: new_memory.id,
348                });
349            } else if self.config.value_update_enabled
350                && existing.agent_id == new_memory.agent_id
351                && existing.user_id == new_memory.user_id
352                && is_value_update(
353                    &new_memory.content,
354                    &existing.content,
355                    self.config.value_update_prefix_share,
356                    self.config.value_update_max_tail,
357                )
358            {
359                if sim < self.config.value_update_min_similarity {
360                    // A frame matched pair below the floor is exactly the
361                    // signal that the floor sits wrong for the embedding
362                    // space in use; log it so production data keeps the
363                    // default honest.
364                    tracing::info!(
365                        similarity = sim,
366                        "value update frame matched below the cosine floor, skipped"
367                    );
368                    continue;
369                }
370                // Same sentence frame, different value: one supersedes the other.
371                // The value-update rule owns this pair (so the correction handler
372                // skips it), and the winner is decided deterministically rather
373                // than by arrival order: a correction beats a plain fact, else the
374                // later-created wins. Batch enrichment stores extracted facts in
375                // arbitrary order with near-equal timestamps, so an order-dependent
376                // rule fired only about half the time; this resolves the same way
377                // whether the correction arrives before or after the original.
378                value_update_handled.push(existing.id);
379                if let Some((winner, loser, loser_confidence)) =
380                    value_update_resolution(new_memory, existing)
381                {
382                    // Invalidate the loser, not just edge it: recall's
383                    // validity filter reads valid_until, so a Supersedes edge
384                    // alone left the stale value recallable and the agent
385                    // kept answering with it. InvalidateMemory sets
386                    // valid_until and creates the Supersedes edge.
387                    let now = std::time::SystemTime::now()
388                        .duration_since(std::time::UNIX_EPOCH)
389                        .unwrap_or_default()
390                        .as_micros() as u64;
391                    self.resolve_contradiction(
392                        &mut actions,
393                        existing.id,
394                        new_memory.id,
395                        loser,
396                        winner,
397                        loser_confidence,
398                        now,
399                        "value update: same sentence frame with a different value",
400                    );
401                }
402                // else: no correction signal and equal timestamps, so the two are
403                // kept as coexisting facts (a multi-valued attribute) rather than
404                // one silently superseding the other.
405            } else if sim > self.config.related_min && sim <= self.config.related_max {
406                actions.push(InferredAction::CreateEdge {
407                    source: new_memory.id,
408                    target: existing.id,
409                    edge_type: EdgeType::Related,
410                    weight: sim,
411                });
412            }
413        }
414
415        // Correction type: find the most similar existing memory and supersede
416        // it, skipping any pair the value-update rule already resolved above.
417        if new_memory.memory_type == MemoryType::Correction
418            && let Some(original) = existing_memories
419                .iter()
420                .filter(|m| m.id != new_memory.id && !value_update_handled.contains(&m.id))
421                .max_by(|a, b| {
422                    cosine_similarity(&new_memory.embedding, &a.embedding)
423                        .partial_cmp(&cosine_similarity(&new_memory.embedding, &b.embedding))
424                        .unwrap_or(std::cmp::Ordering::Equal)
425                })
426        {
427            let sim = cosine_similarity(&new_memory.embedding, &original.embedding);
428            if sim > self.config.correction_threshold {
429                actions.push(InferredAction::CreateEdge {
430                    source: new_memory.id,
431                    target: original.id,
432                    edge_type: EdgeType::Supersedes,
433                    weight: 1.0,
434                });
435                actions.push(InferredAction::UpdateConfidence {
436                    memory: original.id,
437                    new_confidence: (original.confidence * self.config.confidence_decay_factor)
438                        .max(self.config.confidence_floor),
439                });
440            }
441        }
442
443        actions
444    }
445
446    /// LLM enhanced write inference. Uses the CognitiveLlmService when available
447    /// to make smarter invalidation and contradiction decisions. Falls back to
448    /// cosine similarity heuristics for moderate similarity (Related edges) and
449    /// correction handling.
450    pub async fn infer_on_write_with_llm<J: LlmJudge>(
451        &self,
452        new_memory: &MemoryNode,
453        existing_memories: &[MemoryNode],
454        existing_edges: &[(MemoryId, MemoryId, EdgeType)],
455        llm: &CognitiveLlmService<J>,
456    ) -> Vec<InferredAction> {
457        let _ = existing_edges;
458        let mut actions = Vec::new();
459
460        let new_summary = memory_to_summary(new_memory);
461
462        for existing in existing_memories {
463            if existing.id == new_memory.id {
464                continue;
465            }
466
467            let sim = cosine_similarity(&new_memory.embedding, &existing.embedding);
468
469            // Only consult LLM for memories with moderate+ similarity (> 0.5)
470            // to avoid burning tokens on completely unrelated pairs
471            if sim > 0.5 && existing.agent_id == new_memory.agent_id {
472                let old_summary = memory_to_summary(existing);
473
474                // LLM invalidation check
475                if let Ok(verdict) = llm.judge_invalidation(&old_summary, &new_summary).await {
476                    match verdict {
477                        InvalidationVerdict::Invalidate { reason } => {
478                            // InvalidateMemory sets valid_until and creates the
479                            // Supersedes edge; no separate MarkObsolete/CreateEdge.
480                            self.resolve_contradiction(
481                                &mut actions,
482                                existing.id,
483                                new_memory.id,
484                                existing.id,
485                                new_memory.id,
486                                existing.confidence,
487                                new_memory.created_at,
488                                &reason,
489                            );
490                            // Skip contradiction check since we already know the relationship
491                            continue;
492                        }
493                        InvalidationVerdict::Update {
494                            merged_content,
495                            reason,
496                        } => {
497                            actions.push(InferredAction::UpdateContent {
498                                memory: existing.id,
499                                new_content: merged_content,
500                                reason,
501                            });
502                            continue;
503                        }
504                        InvalidationVerdict::Keep { .. } => {
505                            // Fall through to contradiction check
506                        }
507                    }
508                }
509
510                // LLM contradiction check for high similarity pairs that weren't invalidated
511                if sim > 0.7
512                    && existing.content != new_memory.content
513                    && let Ok(verdict) = llm.detect_contradiction(&old_summary, &new_summary).await
514                {
515                    match verdict {
516                        ContradictionVerdict::Contradicts { reason } => {
517                            actions.push(InferredAction::FlagContradiction {
518                                existing: existing.id,
519                                new: new_memory.id,
520                                reason,
521                            });
522                        }
523                        ContradictionVerdict::Supersedes { winner, reason } => {
524                            let winner_is_new = winner == new_memory.id.to_string();
525                            let (obsolete, superseder) = if winner_is_new {
526                                (existing.id, new_memory.id)
527                            } else {
528                                (new_memory.id, existing.id)
529                            };
530                            match self.config.contradiction_policy {
531                                ContradictionPolicy::Supersede => {
532                                    actions.push(InferredAction::InvalidateMemory {
533                                        memory: obsolete,
534                                        superseded_by: superseder,
535                                        valid_until: new_memory.created_at,
536                                    });
537                                }
538                                ContradictionPolicy::Flag => {
539                                    actions.push(InferredAction::FlagContradiction {
540                                        existing: existing.id,
541                                        new: new_memory.id,
542                                        reason,
543                                    });
544                                }
545                            }
546                        }
547                        ContradictionVerdict::Compatible { .. } => {}
548                    }
549                }
550            }
551
552            // Moderate similarity: create Related edge (heuristic, no LLM needed)
553            if sim > self.config.related_min && sim <= self.config.related_max {
554                actions.push(InferredAction::CreateEdge {
555                    source: new_memory.id,
556                    target: existing.id,
557                    edge_type: EdgeType::Related,
558                    weight: sim,
559                });
560            }
561        }
562
563        // Correction type still uses heuristic (find most similar and supersede)
564        if new_memory.memory_type == MemoryType::Correction
565            && let Some(original) = existing_memories
566                .iter()
567                .filter(|m| m.id != new_memory.id)
568                .max_by(|a, b| {
569                    cosine_similarity(&new_memory.embedding, &a.embedding)
570                        .partial_cmp(&cosine_similarity(&new_memory.embedding, &b.embedding))
571                        .unwrap_or(std::cmp::Ordering::Equal)
572                })
573        {
574            let sim = cosine_similarity(&new_memory.embedding, &original.embedding);
575            if sim > self.config.correction_threshold {
576                actions.push(InferredAction::InvalidateMemory {
577                    memory: original.id,
578                    superseded_by: new_memory.id,
579                    valid_until: new_memory.created_at,
580                });
581                actions.push(InferredAction::UpdateConfidence {
582                    memory: original.id,
583                    new_confidence: (original.confidence * self.config.confidence_decay_factor)
584                        .max(self.config.confidence_floor),
585                });
586            }
587        }
588
589        actions
590    }
591}
592
593impl WriteInferenceEngine {
594    /// Record a detected contradiction between `existing` and `new` according
595    /// to the configured policy: supersede the loser (close it and decay its
596    /// confidence) or flag the pair and keep both valid.
597    #[allow(clippy::too_many_arguments)]
598    fn resolve_contradiction(
599        &self,
600        actions: &mut Vec<InferredAction>,
601        existing: MemoryId,
602        new: MemoryId,
603        loser: MemoryId,
604        winner: MemoryId,
605        loser_confidence: f32,
606        valid_until: u64,
607        reason: &str,
608    ) {
609        match self.config.contradiction_policy {
610            ContradictionPolicy::Supersede => {
611                actions.push(InferredAction::InvalidateMemory {
612                    memory: loser,
613                    superseded_by: winner,
614                    valid_until,
615                });
616                actions.push(InferredAction::UpdateConfidence {
617                    memory: loser,
618                    new_confidence: (loser_confidence * self.config.confidence_decay_factor)
619                        .max(self.config.confidence_floor),
620                });
621            }
622            ContradictionPolicy::Flag => {
623                actions.push(InferredAction::FlagContradiction {
624                    existing,
625                    new,
626                    reason: reason.to_string(),
627                });
628            }
629        }
630    }
631}
632
633impl Default for WriteInferenceEngine {
634    fn default() -> Self {
635        Self::new()
636    }
637}
638
639fn memory_to_summary(m: &MemoryNode) -> MemorySummary {
640    MemorySummary {
641        id: m.id,
642        content: m.content.clone(),
643        memory_type: m.memory_type,
644        confidence: m.confidence,
645        created_at: m.created_at,
646    }
647}
648
649#[cfg(test)]
650mod tests {
651    use super::*;
652    use mentedb_core::memory::MemoryType;
653    use mentedb_core::types::AgentId;
654
655    use crate::llm::MockLlmJudge;
656
657    fn make_memory(content: &str, embedding: Vec<f32>, mem_type: MemoryType) -> MemoryNode {
658        let mut m = MemoryNode::new(AgentId::new(), mem_type, content.to_string(), embedding);
659        m.created_at = 1000;
660        m
661    }
662
663    #[test]
664    fn test_heuristic_does_not_flag_contradiction_from_similarity() {
665        // Two facts with high embedding similarity but different content look
666        // identical to cosine similarity whether they are paraphrases or
667        // opposites. The cheap heuristic must NOT guess "contradiction" here
668        // (that produced ~0% precision in production); real contradiction
669        // detection is the LLM path's job. High-sim different-content yields no
670        // action from the heuristic.
671        let agent = AgentId::new();
672        let mut existing =
673            make_memory("uses PostgreSQL", vec![1.0, 0.0, 0.0], MemoryType::Semantic);
674        existing.agent_id = agent;
675
676        let mut new_mem = make_memory("uses MySQL", vec![0.99, 0.01, 0.0], MemoryType::Semantic);
677        new_mem.agent_id = agent;
678        new_mem.created_at = 2000;
679
680        let engine = WriteInferenceEngine::new();
681        let actions = engine.infer_on_write(&new_mem, &[existing], &[]);
682        assert!(
683            !actions
684                .iter()
685                .any(|a| matches!(a, InferredAction::FlagContradiction { .. })),
686            "Heuristic must not flag contradictions from bare similarity, got: {:?}",
687            actions
688        );
689    }
690
691    #[test]
692    fn value_update_supersedes_the_old_fact() {
693        // The canonical correction shape: same sentence frame, new value at the
694        // tail. The newer memory must supersede the older one so recall serves
695        // only the corrected fact (live bug: both coffee facts coexisted).
696        let agent = AgentId::new();
697        let mut existing = make_memory(
698            "The user's favorite coffee order is a cortado",
699            vec![1.0, 0.0, 0.0],
700            MemoryType::Semantic,
701        );
702        existing.agent_id = agent;
703        let mut new_mem = make_memory(
704            "The user's favorite coffee order is a flat white",
705            vec![0.98, 0.02, 0.0],
706            MemoryType::Semantic,
707        );
708        new_mem.agent_id = agent;
709        new_mem.created_at = 2000;
710
711        let engine = WriteInferenceEngine::new();
712        let actions = engine.infer_on_write(&new_mem, &[existing.clone()], &[]);
713        assert!(
714            actions.iter().any(|a| matches!(
715                a,
716                InferredAction::InvalidateMemory { memory, superseded_by, .. }
717                    if *superseded_by == new_mem.id && *memory == existing.id
718            )),
719            "value update must invalidate the stale fact, got: {actions:?}"
720        );
721        assert!(
722            actions
723                .iter()
724                .any(|a| matches!(a, InferredAction::UpdateConfidence { memory, .. } if *memory == existing.id)),
725            "superseded fact's confidence must decay"
726        );
727    }
728
729    #[test]
730    fn value_update_fires_at_realistic_embedder_cosine() {
731        // Regression for the production-dead rule: real embedders put a value
732        // correction far below a high cosine bar (measured: Titan V2 0.61,
733        // OpenAI 3-small 0.80 for cortado vs flat white; 0.51 / 0.79 for a phone
734        // number change). The old 0.88 default was unreachable on either, so the
735        // rule never fired in prod. A frame-matched correction at cosine 0.6 must
736        // now supersede; the structural shared-frame test, not the cosine, is
737        // what proves it is a correction.
738        let agent = AgentId::new();
739        let mut existing = make_memory(
740            "The user's favorite coffee order is a cortado",
741            vec![1.0, 0.0, 0.0],
742            MemoryType::Semantic,
743        );
744        existing.agent_id = agent;
745        // cosine([1,0,0], [0.6,0.8,0]) = 0.6: below the old 0.88 gate, above 0.5.
746        let mut new_mem = make_memory(
747            "The user's favorite coffee order is a flat white",
748            vec![0.6, 0.8, 0.0],
749            MemoryType::Semantic,
750        );
751        new_mem.agent_id = agent;
752        new_mem.created_at = 2000;
753
754        let actions =
755            WriteInferenceEngine::new().infer_on_write(&new_mem, &[existing.clone()], &[]);
756        assert!(
757            actions.iter().any(|a| matches!(
758                a,
759                InferredAction::InvalidateMemory { memory, superseded_by, .. }
760                    if *superseded_by == new_mem.id && *memory == existing.id
761            )),
762            "a frame-matched correction at realistic cosine 0.6 must supersede, got: {actions:?}"
763        );
764    }
765
766    #[test]
767    fn value_update_fires_below_the_old_floor() {
768        // The production miss on Cohere 1024: an identical frame pair whose
769        // cosine sits in the 0.3 to 0.5 band was silently gated by the 0.5
770        // floor. cosine([1,0,0],[0.4,0.9165,0]) = 0.4 must fire now.
771        let agent = AgentId::new();
772        let mut existing = make_memory(
773            "The user uses Next.js for the front end",
774            vec![1.0, 0.0, 0.0],
775            MemoryType::Semantic,
776        );
777        existing.agent_id = agent;
778        let mut new_mem = make_memory(
779            "The user uses React for the front end",
780            vec![0.4, 0.9165, 0.0],
781            MemoryType::Semantic,
782        );
783        new_mem.agent_id = agent;
784        new_mem.created_at = 2000;
785
786        let actions =
787            WriteInferenceEngine::new().infer_on_write(&new_mem, &[existing.clone()], &[]);
788        assert!(
789            actions.iter().any(|a| matches!(
790                a,
791                InferredAction::InvalidateMemory { memory, superseded_by, .. }
792                    if *superseded_by == new_mem.id && *memory == existing.id
793            )),
794            "an identical frame pair at cosine 0.4 must supersede, got: {actions:?}"
795        );
796    }
797
798    #[test]
799    fn value_update_below_cosine_floor_does_not_supersede() {
800        // The cosine floor still guards: a shared token frame that collided on
801        // genuinely off-topic content (cosine 0.2, below the 0.3 floor) must not
802        // supersede, even though the frame test alone would pass. Structure is
803        // primary, but the floor rejects pathological frame collisions.
804        let agent = AgentId::new();
805        let mut existing = make_memory(
806            "The user's favorite coffee order is a cortado",
807            vec![1.0, 0.0, 0.0],
808            MemoryType::Semantic,
809        );
810        existing.agent_id = agent;
811        // cosine([1,0,0], [0.2,0.9798,0]) = 0.2: below the 0.3 floor.
812        let mut new_mem = make_memory(
813            "The user's favorite coffee order is a flat white",
814            vec![0.2, 0.9798, 0.0],
815            MemoryType::Semantic,
816        );
817        new_mem.agent_id = agent;
818        new_mem.created_at = 2000;
819
820        let actions = WriteInferenceEngine::new().infer_on_write(&new_mem, &[existing], &[]);
821        assert!(
822            !actions
823                .iter()
824                .any(|a| matches!(a, InferredAction::InvalidateMemory { .. })),
825            "below the cosine floor a frame match must not supersede, got: {actions:?}"
826        );
827    }
828
829    #[test]
830    fn correction_supersedes_regardless_of_store_order() {
831        // Batch enrichment can store the correction BEFORE the original it
832        // corrects. The original arriving after the existing correction must
833        // still be the one marked stale: the correction wins either way. This is
834        // the case the old order-dependent rule missed (both facts coexisted).
835        let agent = AgentId::new();
836        let mut correction = make_memory(
837            "The user's favorite coffee is a flat white",
838            vec![1.0, 0.0, 0.0],
839            MemoryType::Correction,
840        );
841        correction.agent_id = agent;
842        correction.created_at = 5000;
843        let mut original = make_memory(
844            "The user's favorite coffee is a cortado",
845            vec![0.98, 0.02, 0.0],
846            MemoryType::Semantic,
847        );
848        original.agent_id = agent;
849        original.created_at = 5000; // same batch timestamp
850
851        // The original arrives as the new memory; the correction already exists.
852        let actions =
853            WriteInferenceEngine::new().infer_on_write(&original, &[correction.clone()], &[]);
854        assert!(
855            actions.iter().any(|a| matches!(
856                a,
857                InferredAction::InvalidateMemory { memory, superseded_by, .. }
858                    if *superseded_by == correction.id && *memory == original.id
859            )),
860            "the correction must invalidate the later-arriving original, got: {actions:?}"
861        );
862    }
863
864    #[test]
865    fn correction_supersedes_original_at_equal_timestamps() {
866        // The other store order: original first, correction second, same batch
867        // timestamp. The correction still wins.
868        let agent = AgentId::new();
869        let mut original = make_memory(
870            "The user's favorite coffee is a cortado",
871            vec![1.0, 0.0, 0.0],
872            MemoryType::Semantic,
873        );
874        original.agent_id = agent;
875        original.created_at = 5000;
876        let mut correction = make_memory(
877            "The user's favorite coffee is a flat white",
878            vec![0.98, 0.02, 0.0],
879            MemoryType::Correction,
880        );
881        correction.agent_id = agent;
882        correction.created_at = 5000;
883
884        let actions =
885            WriteInferenceEngine::new().infer_on_write(&correction, &[original.clone()], &[]);
886        assert!(
887            actions.iter().any(|a| matches!(
888                a,
889                InferredAction::InvalidateMemory { memory, superseded_by, .. }
890                    if *superseded_by == correction.id && *memory == original.id
891            )),
892            "the correction must invalidate the original at equal timestamps, got: {actions:?}"
893        );
894    }
895
896    #[test]
897    fn multivalued_facts_at_equal_timestamps_are_both_kept() {
898        // Two plain facts sharing a frame that are not a correction of each other
899        // (a multi-valued attribute) must not supersede when timestamps are equal,
900        // so making the rule order-independent does not start eating coexisting
901        // facts. Only a correction signal or a clear time order triggers it.
902        let agent = AgentId::new();
903        let mut a = make_memory(
904            "The user knows Rust",
905            vec![1.0, 0.0, 0.0],
906            MemoryType::Semantic,
907        );
908        a.agent_id = agent;
909        a.created_at = 5000;
910        let mut b = make_memory(
911            "The user knows Python",
912            vec![0.97, 0.03, 0.0],
913            MemoryType::Semantic,
914        );
915        b.agent_id = agent;
916        b.created_at = 5000;
917
918        let actions = WriteInferenceEngine::new().infer_on_write(&b, &[a], &[]);
919        assert!(
920            !actions.iter().any(|x| matches!(
921                x,
922                InferredAction::CreateEdge {
923                    edge_type: EdgeType::Supersedes,
924                    ..
925                }
926            )),
927            "coexisting facts at equal timestamps must not supersede, got: {actions:?}"
928        );
929    }
930
931    #[test]
932    fn subject_change_is_not_a_value_update() {
933        // "my dog ..." vs "my cat ..." differ at the head: two facts about two
934        // subjects, both must survive.
935        let agent = AgentId::new();
936        let mut existing = make_memory(
937            "My dog is allergic to chicken",
938            vec![1.0, 0.0, 0.0],
939            MemoryType::Semantic,
940        );
941        existing.agent_id = agent;
942        let mut new_mem = make_memory(
943            "My cat is allergic to chicken",
944            vec![0.99, 0.01, 0.0],
945            MemoryType::Semantic,
946        );
947        new_mem.agent_id = agent;
948        new_mem.created_at = 2000;
949
950        let engine = WriteInferenceEngine::new();
951        let actions = engine.infer_on_write(&new_mem, &[existing], &[]);
952        assert!(
953            !actions.iter().any(|a| matches!(
954                a,
955                InferredAction::CreateEdge {
956                    edge_type: EdgeType::Supersedes,
957                    ..
958                }
959            )),
960            "a subject change must not supersede, got: {actions:?}"
961        );
962    }
963
964    #[test]
965    fn value_update_respects_owner_and_config() {
966        // Cross-owner pairs never supersede, and the rule can be disabled.
967        let mut existing = make_memory(
968            "The user's favorite coffee order is a cortado",
969            vec![1.0, 0.0, 0.0],
970            MemoryType::Semantic,
971        );
972        existing.agent_id = AgentId::new();
973        let mut new_mem = make_memory(
974            "The user's favorite coffee order is a flat white",
975            vec![0.98, 0.02, 0.0],
976            MemoryType::Semantic,
977        );
978        new_mem.agent_id = AgentId::new(); // different owner
979        new_mem.created_at = 2000;
980
981        let engine = WriteInferenceEngine::new();
982        let actions = engine.infer_on_write(&new_mem, &[existing.clone()], &[]);
983        assert!(
984            !actions.iter().any(|a| matches!(
985                a,
986                InferredAction::CreateEdge {
987                    edge_type: EdgeType::Supersedes,
988                    ..
989                }
990            )),
991            "cross-owner memories must never supersede each other"
992        );
993
994        // Same owner but rule disabled: no supersession either.
995        let agent = AgentId::new();
996        existing.agent_id = agent;
997        new_mem.agent_id = agent;
998        let engine_off = WriteInferenceEngine::with_config(WriteInferenceConfig {
999            value_update_enabled: false,
1000            ..WriteInferenceConfig::default()
1001        });
1002        let actions = engine_off.infer_on_write(&new_mem, &[existing], &[]);
1003        assert!(
1004            !actions.iter().any(|a| matches!(
1005                a,
1006                InferredAction::CreateEdge {
1007                    edge_type: EdgeType::Supersedes,
1008                    ..
1009                }
1010            )),
1011            "disabled rule must not supersede"
1012        );
1013    }
1014
1015    #[test]
1016    fn is_value_update_frame_analysis() {
1017        // Positive: long shared prefix, short differing tail.
1018        assert!(is_value_update(
1019            "my phone number is 4321",
1020            "my phone number is 1234",
1021            0.6,
1022            4
1023        ));
1024        // Negative: head differs (different subject).
1025        assert!(!is_value_update(
1026            "my cat is allergic to chicken",
1027            "my dog is allergic to chicken",
1028            0.6,
1029            4
1030        ));
1031        // Negative: identical content is not an update.
1032        assert!(!is_value_update("same fact", "same fact", 0.6, 4));
1033        // Positive: the value sits mid sentence, wrapped by a shared tail.
1034        // This is the reported production miss: "i use next.js" then "now i
1035        // use react" coexisted forever because the old prefix-only test
1036        // could not see the shared "for the front end" tail.
1037        assert!(is_value_update(
1038            "The user uses React for the front end",
1039            "The user uses Next.js for the front end",
1040            0.6,
1041            4
1042        ));
1043        assert!(is_value_update(
1044            "i use react for the front end",
1045            "i use next.js for the front end",
1046            0.6,
1047            4
1048        ));
1049        // Negative: subject swap with a long shared tail must still not fire;
1050        // only a one token head is shared.
1051        assert!(!is_value_update(
1052            "the cat is allergic to chicken and needs the special food",
1053            "the dog is allergic to chicken and needs the special food",
1054            0.6,
1055            4
1056        ));
1057        // Negative: purely numeric mid sentence middles are identifiers
1058        // naming different things (rule 3 and rule 4 coexist), never a
1059        // changed value.
1060        assert!(!is_value_update(
1061            "API sibling rule 1 about payload naming",
1062            "API sibling rule 0 about payload naming",
1063            0.6,
1064            4
1065        ));
1066        // Positive is unaffected at the end position: a numeric tail is a
1067        // value ("phone is 1234" to "phone is 4321"), covered above.
1068        // Negative: three token enumerations are items, never updates.
1069        assert!(!is_value_update(
1070            "Directive number 1",
1071            "Directive number 0",
1072            0.6,
1073            4
1074        ));
1075        // Negative: tail too long (a genuinely different statement).
1076        assert!(!is_value_update(
1077            "the deploy runs every friday and sarah reviews it after standup",
1078            "the deploy runs every friday unless the release train is frozen for the quarter end audit",
1079            0.6,
1080            4
1081        ));
1082    }
1083
1084    #[test]
1085    fn test_byte_identical_duplicate_is_deduplicated() {
1086        // Identical text is deduplicated (invalidate + Derived lineage), not a
1087        // supersession, so it never surfaces as a conflict.
1088        let agent = AgentId::new();
1089        let mut existing = make_memory(
1090            "Ran command: ls -la",
1091            vec![1.0, 0.0, 0.0],
1092            MemoryType::Episodic,
1093        );
1094        existing.agent_id = agent;
1095
1096        let mut new_mem = make_memory(
1097            "Ran command: ls -la",
1098            vec![1.0, 0.0, 0.0],
1099            MemoryType::Episodic,
1100        );
1101        new_mem.agent_id = agent;
1102        new_mem.created_at = 2000;
1103
1104        let engine = WriteInferenceEngine::new();
1105        let actions = engine.infer_on_write(&new_mem, &[existing], &[]);
1106        assert!(
1107            actions
1108                .iter()
1109                .any(|a| matches!(a, InferredAction::DeduplicateExact { .. })),
1110            "Expected identical duplicate to be deduplicated, got: {:?}",
1111            actions
1112        );
1113    }
1114
1115    #[test]
1116    fn test_moderate_similarity_creates_edge() {
1117        let existing = make_memory("topic A", vec![1.0, 0.0, 0.0], MemoryType::Semantic);
1118        // ~0.7 similarity
1119        let new_mem = make_memory("topic B", vec![0.7, 0.714, 0.0], MemoryType::Semantic);
1120
1121        let engine = WriteInferenceEngine::new();
1122        let actions = engine.infer_on_write(&new_mem, &[existing], &[]);
1123        assert!(
1124            actions.iter().any(|a| matches!(
1125                a,
1126                InferredAction::CreateEdge {
1127                    edge_type: EdgeType::Related,
1128                    ..
1129                }
1130            )),
1131            "Expected CreateEdge Related, got: {:?}",
1132            actions
1133        );
1134    }
1135
1136    #[tokio::test]
1137    async fn test_llm_invalidation_emits_temporal_actions() {
1138        let agent = AgentId::new();
1139        let mut existing = make_memory(
1140            "Alice works at Acme",
1141            vec![0.8, 0.6, 0.0],
1142            MemoryType::Semantic,
1143        );
1144        existing.agent_id = agent;
1145
1146        let mut new_mem = make_memory(
1147            "Alice joined Google last week",
1148            vec![0.75, 0.65, 0.1],
1149            MemoryType::Semantic,
1150        );
1151        new_mem.agent_id = agent;
1152        new_mem.created_at = 2000;
1153
1154        let judge =
1155            MockLlmJudge::new(r#"{"verdict": "invalidate", "reason": "Alice changed jobs"}"#);
1156        let llm = CognitiveLlmService::new(judge);
1157        let engine = WriteInferenceEngine::new();
1158        let actions = engine
1159            .infer_on_write_with_llm(&new_mem, &[existing.clone()], &[], &llm)
1160            .await;
1161
1162        assert!(
1163            actions
1164                .iter()
1165                .any(|a| matches!(a, InferredAction::InvalidateMemory { .. })),
1166            "Expected InvalidateMemory from LLM verdict, got: {:?}",
1167            actions
1168        );
1169        // Exactly one invalidation action: InvalidateMemory both sets
1170        // valid_until and creates the Supersedes edge, so a MarkObsolete
1171        // alongside it would duplicate the edge.
1172        let invalidation_count = actions
1173            .iter()
1174            .filter(|a| {
1175                matches!(
1176                    a,
1177                    InferredAction::InvalidateMemory { .. } | InferredAction::MarkObsolete { .. }
1178                )
1179            })
1180            .count();
1181        assert_eq!(
1182            invalidation_count, 1,
1183            "Expected a single invalidation action, got: {:?}",
1184            actions
1185        );
1186    }
1187
1188    #[tokio::test]
1189    async fn test_llm_update_emits_update_content() {
1190        let agent = AgentId::new();
1191        let mut existing = make_memory(
1192            "Project uses React",
1193            vec![0.8, 0.6, 0.0],
1194            MemoryType::Semantic,
1195        );
1196        existing.agent_id = agent;
1197
1198        let mut new_mem = make_memory(
1199            "Project migrated from React to Vue",
1200            vec![0.75, 0.65, 0.1],
1201            MemoryType::Semantic,
1202        );
1203        new_mem.agent_id = agent;
1204        new_mem.created_at = 2000;
1205
1206        let judge = MockLlmJudge::new(
1207            r#"{"verdict": "update", "merged_content": "Project migrated from React to Vue in Q2", "reason": "adds temporal context"}"#,
1208        );
1209        let llm = CognitiveLlmService::new(judge);
1210        let engine = WriteInferenceEngine::new();
1211        let actions = engine
1212            .infer_on_write_with_llm(&new_mem, &[existing.clone()], &[], &llm)
1213            .await;
1214
1215        assert!(
1216            actions
1217                .iter()
1218                .any(|a| matches!(a, InferredAction::UpdateContent { .. })),
1219            "Expected UpdateContent from LLM update verdict, got: {:?}",
1220            actions
1221        );
1222    }
1223
1224    #[tokio::test]
1225    async fn test_llm_keep_falls_through_to_contradiction_check() {
1226        let agent = AgentId::new();
1227        // High similarity pair where LLM says keep but content differs
1228        let mut existing = make_memory("Prefers tabs", vec![0.9, 0.44, 0.0], MemoryType::Semantic);
1229        existing.agent_id = agent;
1230
1231        let mut new_mem = make_memory(
1232            "Prefers spaces",
1233            vec![0.88, 0.47, 0.0],
1234            MemoryType::Semantic,
1235        );
1236        new_mem.agent_id = agent;
1237        new_mem.created_at = 2000;
1238
1239        let judge = MockLlmJudge::new(
1240            r#"{"verdict": "compatible", "reason": "different formatting preferences"}"#,
1241        );
1242        let llm = CognitiveLlmService::new(judge);
1243        let engine = WriteInferenceEngine::new();
1244        let actions = engine
1245            .infer_on_write_with_llm(&new_mem, &[existing.clone()], &[], &llm)
1246            .await;
1247
1248        // MockLlmJudge returns the same response for all calls, so the keep
1249        // from invalidation falls through, then contradiction returns compatible
1250        assert!(
1251            !actions
1252                .iter()
1253                .any(|a| matches!(a, InferredAction::FlagContradiction { .. })),
1254            "Should not flag contradiction when LLM says compatible",
1255        );
1256    }
1257}