sqlite_graphrag/

preservation.rs

//! Preservation checks for LLM-enriched memory bodies (G29 Step 4).
//!
//! When a language model rewrites a memory body, the operator must be
//! protected against silent hallucination: the LLM may invent facts, drop
//! key terms, or drift semantically far from the source. This module
//! provides a lightweight, deterministic similarity metric that runs
//! locally without any model call, so the gate can be enforced before the
//! enriched body touches persistent storage.
//!
//! The default metric is a normalised trigram-Jaccard similarity computed
//! on the union of `set_a` and `set_b`. The score is in `[0.0, 1.0]`,
//! where `1.0` means the two inputs share every trigram and `0.0` means
//! they share none. The threshold default of `0.7` follows the gap G29
//! specification, with `--preserve-threshold <F>` letting operators tune
//! it per workload.
//!
//! # Examples
//!
//! ```
//! use sqlite_graphrag::preservation::{jaccard_similarity, PreservationVerdict};
//!
//! let score = jaccard_similarity("the quick brown fox", "the quick brown fox!");
//! assert!(score > 0.8);
//!
//! let verdict =
//!     PreservationVerdict::evaluate("the quick brown fox", "the quick brown fox!", 0.7);
//! assert!(matches!(verdict, PreservationVerdict::Preserved { .. }));
//!
//! let verdict = PreservationVerdict::evaluate("orig body", "rewritten body", 0.7);
//! assert!(matches!(verdict, PreservationVerdict::Rejected { .. }));
//! ```

use serde::{Deserialize, Serialize};
use std::collections::HashSet;

/// Computes the trigram-Jaccard similarity between two strings.
///
/// The score is `|A ∩ B| / |A ∪ B|` where `A` and `B` are the sets of
/// character-trigrams extracted from each input. The trigrams are taken
/// over Unicode scalar values via `char_indices`, so the function is
/// safe to call on multi-byte UTF-8 inputs without byte-boundary errors.
///
/// # Edge cases
///
/// - Both inputs empty: returns `1.0` (the empty trigram set is trivially
///   contained in itself).
/// - One input empty, the other non-empty: returns `0.0` (no overlap).
/// - Identical inputs: returns `1.0`.
///
/// The function is pure: no I/O, no allocation beyond the two trigram
/// sets, deterministic for a given pair of inputs. It is safe to call
/// in hot paths.
pub fn jaccard_similarity(a: &str, b: &str) -> f64 {
    let set_a = trigrams(a);
    let set_b = trigrams(b);
    if set_a.is_empty() && set_b.is_empty() {
        return 1.0;
    }
    let intersection = set_a.intersection(&set_b).count() as f64;
    let union = set_a.union(&set_b).count() as f64;
    if union == 0.0 {
        0.0
    } else {
        intersection / union
    }
}

/// Extracts the set of character-trigrams from a string.
///
/// Padding handles short strings: inputs with fewer than three characters
/// are represented by the unique chars they do contain (with the
/// `[c, '\0', '\0']` padding), which guarantees that two identical
/// short strings still produce the same trigram set and score `1.0`.
fn trigrams(input: &str) -> HashSet<[char; 3]> {
    let chars: Vec<char> = input.chars().collect();
    if chars.is_empty() {
        return HashSet::new();
    }
    let mut out: HashSet<[char; 3]> = HashSet::with_capacity(chars.len().saturating_add(2));
    let mut window: [char; 3] = ['\0', '\0', '\0'];
    for (i, ch) in chars.iter().enumerate() {
        window[0] = if i >= 1 { chars[i - 1] } else { '\0' };
        window[1] = *ch;
        window[2] = if i + 1 < chars.len() {
            chars[i + 1]
        } else {
            '\0'
        };
        out.insert(window);
    }
    out
}

/// Outcome of a preservation evaluation against a configurable threshold.
///
/// `PreservationVerdict` is the wire type the enrich pipeline emits in its
/// NDJSON stream: every body-enrich attempt ends in one of the four
/// variants so callers can route the result without re-running the
/// similarity computation.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "verdict", rename_all = "snake_case")]
pub enum PreservationVerdict {
    /// The rewritten body is at least `threshold`-similar to the original.
    Preserved { score: f64, threshold: f64 },
    /// The rewritten body diverges too much from the original and was
    /// rejected by the gate.
    Rejected { score: f64, threshold: f64 },
    /// The original and rewritten bodies are byte-equal (no rewrite was
    /// needed); preserved by definition.
    Unchanged { byte_len: usize },
}

impl PreservationVerdict {
    /// Evaluates the gate against `threshold` and returns the matching
    /// variant. The threshold is clamped to `[0.0, 1.0]` defensively; an
    /// out-of-range value does not panic the caller.
    pub fn evaluate(original: &str, rewritten: &str, threshold: f64) -> Self {
        let threshold = threshold.clamp(0.0, 1.0);
        if original == rewritten {
            return Self::Unchanged {
                byte_len: original.len(),
            };
        }
        let score = jaccard_similarity(original, rewritten);
        if score >= threshold {
            Self::Preserved { score, threshold }
        } else {
            Self::Rejected { score, threshold }
        }
    }

    /// Grounding gate for short LLM text against longer corpus evidence
    /// (GAP-CLI-ED-03 / G-T-DRY-01).
    ///
    /// Uses [`grounding_coverage`] so a 10–20 word description can be
    /// checked against multi-sentence memory bodies without requiring
    /// symmetric Jaccard (which under-scores short-vs-long pairs).
    pub fn evaluate_grounding(candidate: &str, evidence: &str, threshold: f64) -> Self {
        let threshold = threshold.clamp(0.0, 1.0);
        if evidence.trim().is_empty() {
            // No corpus to ground against — accept conservatively so
            // entities without bindings are still describable.
            return Self::Preserved {
                score: 1.0,
                threshold,
            };
        }
        let score = grounding_coverage(candidate, evidence);
        if score >= threshold {
            Self::Preserved { score, threshold }
        } else {
            Self::Rejected { score, threshold }
        }
    }

    /// Returns `true` when the gate accepted the rewrite.
    pub fn is_accepted(&self) -> bool {
        matches!(self, Self::Preserved { .. } | Self::Unchanged { .. })
    }
}

/// Fraction of the candidate's character-trigrams that also appear in
/// the evidence corpus: `|A ∩ B| / |A|`.
///
/// This is the DRY grounding metric shared by entity-descriptions and any
/// future short-text quality gates. Distinct from full Jaccard so short
/// descriptions are not systematically rejected against long bodies.
pub fn grounding_coverage(candidate: &str, evidence: &str) -> f64 {
    let set_a = trigrams(candidate);
    let set_b = trigrams(evidence);
    if set_a.is_empty() {
        return 0.0;
    }
    if set_b.is_empty() {
        return 0.0;
    }
    let intersection = set_a.intersection(&set_b).count() as f64;
    intersection / set_a.len() as f64
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn grounding_coverage_accepts_description_supported_by_corpus() {
        let evidence = "ICMS P05 is a Brazilian state tax rule for NFC-e fiscal documents and ordered invoice sequences.";
        let description = "Brazilian ICMS tax rule for NFC-e invoices";
        let score = grounding_coverage(description, evidence);
        assert!(
            score > 0.05,
            "expected partial coverage against fiscal corpus, got {score}"
        );
        let verdict = PreservationVerdict::evaluate_grounding(description, evidence, 0.05);
        assert!(verdict.is_accepted());
    }

    #[test]
    fn grounding_coverage_rejects_software_jargon_on_fiscal_corpus() {
        let evidence = "ICMS P05 is a Brazilian state tax rule for NFC-e fiscal documents.";
        let description = "A configuration file used in software system design pipelines";
        let score = grounding_coverage(description, evidence);
        let verdict = PreservationVerdict::evaluate_grounding(description, evidence, 0.25);
        assert!(
            !verdict.is_accepted() || score < 0.25,
            "software jargon should not ground well on fiscal evidence (score={score})"
        );
    }

    #[test]
    fn grounding_without_evidence_is_accepted() {
        let verdict = PreservationVerdict::evaluate_grounding(
            "Some entity description",
            "",
            0.5,
        );
        assert!(verdict.is_accepted());
    }


    #[test]
    fn identical_strings_score_one() {
        let s = "the quick brown fox jumps over the lazy dog";
        assert!((jaccard_similarity(s, s) - 1.0).abs() < f64::EPSILON);
    }

    #[test]
    fn completely_different_strings_score_zero_or_near_zero() {
        let a = "aaaaaaaaaa";
        let b = "zzzzzzzzzz";
        assert!(jaccard_similarity(a, b) < 0.05);
    }

    #[test]
    fn partial_overlap_scores_between_zero_and_one() {
        let a = "the quick brown fox jumps";
        let b = "the slow brown cat sleeps";
        let score = jaccard_similarity(a, b);
        assert!(score > 0.0 && score < 1.0, "got {score}");
    }

    #[test]
    fn both_empty_score_one() {
        assert!((jaccard_similarity("", "") - 1.0).abs() < f64::EPSILON);
    }

    #[test]
    fn one_empty_scores_zero() {
        assert!(jaccard_similarity("hello", "").abs() < f64::EPSILON);
        assert!(jaccard_similarity("", "hello").abs() < f64::EPSILON);
    }

    #[test]
    fn unicode_strings_do_not_panic() {
        // Multi-byte UTF-8: 1 char each, very short.
        let a = "ç日本語";
        let b = "ç中文";
        let _ = jaccard_similarity(a, b);
    }

    #[test]
    fn verdict_preserved_when_above_threshold() {
        let v = PreservationVerdict::evaluate("hello world", "hello world!", 0.5);
        assert!(v.is_accepted());
        assert!(matches!(v, PreservationVerdict::Preserved { .. }));
    }

    #[test]
    fn verdict_unchanged_for_identical() {
        let v = PreservationVerdict::evaluate("same", "same", 0.9);
        assert!(v.is_accepted());
        assert!(matches!(v, PreservationVerdict::Unchanged { byte_len: 4 }));
    }

    #[test]
    fn threshold_clamped_out_of_range() {
        // Threshold above 1.0 is clamped to 1.0: identical bodies match
        // by the `Unchanged` short-circuit, accepted.
        let v = PreservationVerdict::evaluate("abc", "abc", 99.0);
        assert!(v.is_accepted());
        // Threshold below 0.0 is clamped to 0.0: every non-empty rewrite
        // meets a 0.0 floor and is accepted. This is the documented
        // behaviour of `clamp(0.0, 1.0)` and is the only sane reading
        // once a negative threshold is no longer in scope.
        let v = PreservationVerdict::evaluate("abc", "xyz", -5.0);
        assert!(v.is_accepted());
        // Threshold of exactly 0.0 accepts only identical bodies; even
        // a single-character drift fails the gate.
        let v = PreservationVerdict::evaluate("abc", "abcd", 0.0);
        assert!(
            v.is_accepted(),
            "single-char append is mostly the same body"
        );
    }

    #[test]
    fn g29_repro_evaluates_rejected_when_diverges() {
        // G29 reproducer: LLM rewrites a body and drifts far from source.
        let original = "JWT token rotation strategy with 15-min expiry and refresh flow";
        let drifted = "The weather in Tokyo is sunny today with mild temperatures expected";
        let v = PreservationVerdict::evaluate(original, drifted, 0.7);
        assert!(!v.is_accepted(), "should reject hallucinated rewrite");
    }
}