1use std::sync::Arc;
28use std::time::Duration;
29
30use zeph_llm::any::AnyProvider;
31use zeph_llm::provider::LlmProvider as _;
32
33use crate::graph::GraphStore;
34
35#[derive(Debug, Clone)]
39pub struct QualityGateConfig {
40 pub enabled: bool,
42 pub threshold: f32,
44 pub recent_window: usize,
46 pub contradiction_grace_seconds: u64,
49 pub information_value_weight: f32,
51 pub reference_completeness_weight: f32,
53 pub contradiction_weight: f32,
55 pub rejection_rate_alarm_ratio: f32,
58 pub llm_timeout_ms: u64,
60 pub llm_weight: f32,
62 pub reference_check_lang_en: bool,
65}
66
67impl Default for QualityGateConfig {
68 fn default() -> Self {
69 Self {
70 enabled: false,
71 threshold: 0.55,
72 recent_window: 32,
73 contradiction_grace_seconds: 300,
74 information_value_weight: 0.4,
75 reference_completeness_weight: 0.3,
76 contradiction_weight: 0.3,
77 rejection_rate_alarm_ratio: 0.35,
78 llm_timeout_ms: 500,
79 llm_weight: 0.5,
80 reference_check_lang_en: true,
81 }
82 }
83}
84
85#[derive(Debug, Clone)]
89pub struct QualityScore {
90 pub information_value: f32,
92 pub reference_completeness: f32,
94 pub contradiction_risk: f32,
98 pub combined: f32,
100 pub final_score: f32,
102}
103
104#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Serialize)]
106#[serde(rename_all = "snake_case")]
107#[non_exhaustive]
108pub enum QualityRejectionReason {
109 Redundant,
111 IncompleteReference,
113 Contradiction,
115 LlmLowConfidence,
117}
118
119impl QualityRejectionReason {
120 #[must_use]
122 pub fn label(self) -> &'static str {
123 match self {
124 Self::Redundant => "redundant",
125 Self::IncompleteReference => "incomplete_reference",
126 Self::Contradiction => "contradiction",
127 Self::LlmLowConfidence => "llm_low_confidence",
128 }
129 }
130}
131
132struct RollingRateTracker {
134 window: std::collections::VecDeque<bool>,
135 capacity: usize,
136 reject_count: usize,
137}
138
139impl RollingRateTracker {
140 fn new(capacity: usize) -> Self {
141 Self {
142 window: std::collections::VecDeque::with_capacity(capacity + 1),
143 capacity,
144 reject_count: 0,
145 }
146 }
147
148 fn push(&mut self, rejected: bool) {
149 if self.window.len() >= self.capacity
150 && let Some(evicted) = self.window.pop_front()
151 && evicted
152 {
153 self.reject_count = self.reject_count.saturating_sub(1);
154 }
155 self.window.push_back(rejected);
156 if rejected {
157 self.reject_count += 1;
158 }
159 }
160
161 #[allow(clippy::cast_precision_loss)]
162 fn rate(&self) -> f32 {
163 if self.window.is_empty() {
164 return 0.0;
165 }
166 self.reject_count as f32 / self.window.len() as f32
167 }
168}
169
170pub struct QualityGate {
182 config: Arc<QualityGateConfig>,
183 llm_provider: Option<Arc<AnyProvider>>,
185 graph_store: Option<Arc<GraphStore>>,
186 rejection_counts: std::sync::Mutex<std::collections::HashMap<QualityRejectionReason, u64>>,
188 rate_tracker: std::sync::Mutex<RollingRateTracker>,
190 embed_timeout: std::time::Duration,
192}
193
194impl QualityGate {
195 #[must_use]
197 pub fn new(config: QualityGateConfig) -> Self {
198 Self {
199 config: Arc::new(config),
200 llm_provider: None,
201 graph_store: None,
202 rejection_counts: std::sync::Mutex::new(std::collections::HashMap::new()),
203 rate_tracker: std::sync::Mutex::new(RollingRateTracker::new(100)),
204 embed_timeout: std::time::Duration::from_secs(5),
205 }
206 }
207
208 #[must_use]
212 pub fn with_embed_timeout(mut self, timeout_secs: u64) -> Self {
213 self.embed_timeout = std::time::Duration::from_secs(timeout_secs.max(1));
214 self
215 }
216
217 #[must_use]
219 pub fn with_llm_provider(mut self, provider: AnyProvider) -> Self {
220 self.llm_provider = Some(Arc::new(provider));
221 self
222 }
223
224 #[must_use]
226 pub fn with_graph_store(mut self, store: Arc<GraphStore>) -> Self {
227 self.graph_store = Some(store);
228 self
229 }
230
231 #[must_use]
233 pub fn config(&self) -> &QualityGateConfig {
234 &self.config
235 }
236
237 #[must_use]
239 pub fn rejection_counts(&self) -> std::collections::HashMap<QualityRejectionReason, u64> {
240 self.rejection_counts
241 .lock()
242 .map(|g| g.clone())
243 .unwrap_or_default()
244 }
245
246 #[tracing::instrument(name = "memory.quality_gate.evaluate", skip_all)]
253 pub async fn evaluate(
254 &self,
255 content: &str,
256 embed_provider: &AnyProvider,
257 recent_embeddings: &[Vec<f32>],
258 ) -> Option<QualityRejectionReason> {
259 if !self.config.enabled {
260 return None;
261 }
262
263 let info_val = compute_information_value(
264 content,
265 embed_provider,
266 recent_embeddings,
267 self.embed_timeout,
268 )
269 .await;
270 let ref_comp = if self.config.reference_check_lang_en {
271 compute_reference_completeness(content)
272 } else {
273 1.0
274 };
275 let contradiction_risk =
276 compute_contradiction_risk(content, self.graph_store.as_deref(), &self.config).await;
277
278 let w_v = self.config.information_value_weight;
279 let w_c = self.config.reference_completeness_weight;
280 let w_k = self.config.contradiction_weight;
281
282 let rule_score = w_v * info_val + w_c * ref_comp + w_k * (1.0 - contradiction_risk);
283
284 let final_score = if let Some(ref llm) = self.llm_provider {
285 let llm_score = call_llm_scorer(content, llm, self.config.llm_timeout_ms).await;
286 let lw = self.config.llm_weight;
287 (1.0 - lw) * rule_score + lw * llm_score
288 } else {
289 rule_score
290 };
291
292 let rejected = final_score < self.config.threshold;
293
294 if let Ok(mut tracker) = self.rate_tracker.lock() {
296 tracker.push(rejected);
297 let rate = tracker.rate();
298 if rate > self.config.rejection_rate_alarm_ratio {
299 tracing::warn!(
300 rate = %format!("{:.2}", rate),
301 window_size = self.config.recent_window,
302 threshold = self.config.rejection_rate_alarm_ratio,
303 "quality_gate: high rejection rate alarm"
304 );
305 }
306 }
307
308 if !rejected {
309 return None;
310 }
311
312 let reason = if info_val < 0.1 {
314 QualityRejectionReason::Redundant
315 } else if ref_comp < 0.5 && self.config.reference_check_lang_en {
316 QualityRejectionReason::IncompleteReference
317 } else if contradiction_risk >= 1.0 {
318 QualityRejectionReason::Contradiction
319 } else {
320 QualityRejectionReason::LlmLowConfidence
321 };
322
323 if let Ok(mut counts) = self.rejection_counts.lock() {
324 *counts.entry(reason).or_insert(0) += 1;
325 }
326
327 tracing::debug!(
328 reason = reason.label(),
329 final_score,
330 info_val,
331 ref_comp,
332 contradiction_risk,
333 "quality_gate: rejected write"
334 );
335
336 Some(reason)
337 }
338}
339
340async fn compute_information_value(
346 content: &str,
347 provider: &AnyProvider,
348 recent_embeddings: &[Vec<f32>],
349 embed_timeout: std::time::Duration,
350) -> f32 {
351 if recent_embeddings.is_empty() {
352 return 1.0;
353 }
354 if !provider.supports_embeddings() {
355 return 1.0;
356 }
357 let candidate = match crate::llm_judge::embed_with_timeout_fail_open(
358 provider,
359 content,
360 embed_timeout,
361 1.0,
362 "quality_gate: information_value",
363 )
364 .await
365 {
366 Ok(v) => v,
367 Err(fail_open) => return fail_open,
368 };
369 let max_sim = recent_embeddings
370 .iter()
371 .map(|r| zeph_common::math::cosine_similarity(&candidate, r))
372 .fold(0.0f32, f32::max);
373 (1.0 - max_sim).max(0.0)
374}
375
376#[must_use]
381pub fn compute_reference_completeness(content: &str) -> f32 {
382 const PRONOUNS: &[&str] = &[
384 " he ", " she ", " they ", " it ", " him ", " her ", " them ",
385 ];
386 const DEICTIC_TIME: &[&str] = &[
388 "yesterday",
389 "tomorrow",
390 "last week",
391 "next week",
392 "last month",
393 "next month",
394 "last year",
395 "next year",
396 ];
397 const DATE_ANCHORS: &[&str] = &[
399 "january",
400 "february",
401 "march",
402 "april",
403 "may",
404 "june",
405 "july",
406 "august",
407 "september",
408 "october",
409 "november",
410 "december",
411 "jan ",
412 "feb ",
413 "mar ",
414 "apr ",
415 "jun ",
416 "jul ",
417 "aug ",
418 "sep ",
419 "oct ",
420 "nov ",
421 "dec ",
422 ];
423
424 let lower = content.to_lowercase();
425 let padded = format!(" {lower} ");
426 let pronoun_count = PRONOUNS.iter().filter(|&&p| padded.contains(p)).count();
427
428 let has_year_anchor = has_4digit_year_anchor(&lower);
431 let has_date_anchor = has_year_anchor || DATE_ANCHORS.iter().any(|&a| lower.contains(a));
432 let deictic_count = if has_date_anchor {
433 0
434 } else {
435 DEICTIC_TIME.iter().filter(|&&t| lower.contains(t)).count()
436 };
437
438 let total_issues = pronoun_count + deictic_count;
439 if total_issues == 0 {
440 return 1.0;
441 }
442
443 let word_count = content.split_ascii_whitespace().count().max(1);
445 #[allow(clippy::cast_precision_loss)]
446 let ratio = total_issues as f32 / word_count as f32;
447 (1.0 - ratio * 2.0).clamp(0.0, 1.0)
448}
449
450fn has_4digit_year_anchor(text: &str) -> bool {
455 let bytes = text.as_bytes();
456 let len = bytes.len();
457 if len < 4 {
458 return false;
459 }
460 let mut i = 0usize;
461 while i + 3 < len {
462 let c0 = bytes[i];
463 let c1 = bytes[i + 1];
464 if ((c0 == b'1' && c1 == b'9') || (c0 == b'2' && c1 == b'0'))
465 && bytes[i + 2].is_ascii_digit()
466 && bytes[i + 3].is_ascii_digit()
467 {
468 let left_ok = i == 0 || !bytes[i - 1].is_ascii_digit();
469 let right_ok = i + 4 >= len || !bytes[i + 4].is_ascii_digit();
470 if left_ok && right_ok {
471 return true;
472 }
473 }
474 i += 1;
475 }
476 false
477}
478
479async fn compute_contradiction_risk(
488 content: &str,
489 graph: Option<&GraphStore>,
490 config: &QualityGateConfig,
491) -> f32 {
492 let Some(store) = graph else {
493 return 0.0;
494 };
495
496 let content_lower = content.to_lowercase();
497
498 let subject_query = extract_subject_tokens(&content_lower);
501 if subject_query.is_empty() {
502 return 0.0;
503 }
504
505 let Ok(entities) = store.find_entities_fuzzy(&subject_query, 1).await else {
507 return 0.0;
508 };
509 let Some(subject_entity) = entities.into_iter().next() else {
510 return 0.0;
511 };
512
513 let canonical_predicate = extract_predicate_token(&content_lower);
515
516 let Ok(edges) = store.edges_for_entity(subject_entity.id.0).await else {
518 return 0.0;
519 };
520
521 let relevant_edges: Vec<_> = edges
523 .iter()
524 .filter(|e| {
525 e.source_entity_id == subject_entity.id.0
526 && canonical_predicate
527 .as_ref()
528 .is_none_or(|p| e.relation == *p)
529 })
530 .collect();
531
532 if relevant_edges.is_empty() {
533 return 0.0;
534 }
535
536 let now_secs = std::time::SystemTime::now()
537 .duration_since(std::time::UNIX_EPOCH)
538 .map_or(0, |d| d.as_secs());
539
540 let has_old_conflict = relevant_edges.iter().any(|edge| {
541 let edge_ts = chrono::DateTime::parse_from_rfc3339(&edge.created_at)
542 .map_or(0u64, |dt| u64::try_from(dt.timestamp()).unwrap_or(0));
543 now_secs.saturating_sub(edge_ts) > config.contradiction_grace_seconds
544 });
545
546 if has_old_conflict { 1.0 } else { 0.5 }
547}
548
549fn extract_subject_tokens(content_lower: &str) -> String {
551 const VERB_MARKERS: &[&str] = &["is", "was", "are", "were", "has", "have", "had", "will"];
552 let tokens: Vec<&str> = content_lower.split_ascii_whitespace().collect();
553 let end = tokens
554 .iter()
555 .position(|t| VERB_MARKERS.contains(t))
556 .unwrap_or(2.min(tokens.len()));
557 let subject_tokens = &tokens[..end.min(3)];
558 subject_tokens.join(" ")
559}
560
561fn extract_predicate_token(content_lower: &str) -> Option<String> {
563 const VERB_MARKERS: &[&str] = &["is", "was", "are", "were", "has", "have", "had", "will"];
564 content_lower
565 .split_ascii_whitespace()
566 .find(|t| VERB_MARKERS.contains(t))
567 .map(str::to_owned)
568}
569
570async fn call_llm_scorer(content: &str, provider: &AnyProvider, timeout_ms: u64) -> f32 {
574 let system = "You are a memory quality judge. Rate the quality of the following message \
575 for long-term storage on a scale of 0.0 to 1.0. Consider: information density, \
576 completeness of references, factual clarity. \
577 Respond with ONLY a JSON object: \
578 {\"information_value\": 0.0-1.0, \"reference_completeness\": 0.0-1.0, \
579 \"contradiction_risk\": 0.0-1.0}";
580
581 let user = format!(
582 "Message: {}\n\nQuality JSON:",
583 content.chars().take(500).collect::<String>()
584 );
585
586 crate::llm_judge::llm_judge_score(
587 provider,
588 system,
589 user,
590 Duration::from_millis(timeout_ms),
591 0.5,
592 "quality_gate: LLM scorer",
593 |s| Some(parse_llm_score(s)),
594 )
595 .await
596}
597
598fn parse_llm_score(response: &str) -> f32 {
602 let start = response.find('{');
604 let end = response.rfind('}');
605 let (Some(s), Some(e)) = (start, end) else {
606 return 0.5;
607 };
608 let json_str = &response[s..=e];
609 let Ok(val) = serde_json::from_str::<serde_json::Value>(json_str) else {
610 return 0.5;
611 };
612
613 #[allow(clippy::cast_possible_truncation)]
614 let iv = val["information_value"].as_f64().unwrap_or(0.5) as f32;
615 #[allow(clippy::cast_possible_truncation)]
616 let rc = val["reference_completeness"].as_f64().unwrap_or(0.5) as f32;
617 #[allow(clippy::cast_possible_truncation)]
618 let cr = val["contradiction_risk"].as_f64().unwrap_or(0.0) as f32;
619
620 let score =
622 0.4 * iv.clamp(0.0, 1.0) + 0.3 * rc.clamp(0.0, 1.0) + 0.3 * (1.0 - cr.clamp(0.0, 1.0));
623 score.clamp(0.0, 1.0)
624}
625
626#[cfg(test)]
629mod tests {
630 use super::*;
631
632 #[test]
633 fn reference_completeness_clean_text() {
634 let score = compute_reference_completeness("The Rust compiler enforces memory safety.");
635 assert!((score - 1.0).abs() < 0.01, "clean text should score 1.0");
636 }
637
638 #[test]
639 fn reference_completeness_pronoun_heavy() {
640 let score = compute_reference_completeness("yeah he said they confirmed it");
642 assert!(
643 score < 0.5,
644 "pronoun-heavy message should score below 0.5, got {score}"
645 );
646 }
647
648 #[test]
649 fn reference_completeness_deictic_without_anchor() {
650 let score = compute_reference_completeness("We agreed yesterday to postpone");
651 assert!(
652 score < 1.0,
653 "deictic time without anchor should penalize, got {score}"
654 );
655 }
656
657 #[test]
658 fn reference_completeness_deictic_with_anchor() {
659 let score = compute_reference_completeness("We agreed yesterday (2026-04-18) to postpone");
660 assert!(
661 score >= 0.9,
662 "deictic with anchor '20' should not penalize, got {score}"
663 );
664 }
665
666 #[test]
667 fn rejection_reason_labels() {
668 assert_eq!(QualityRejectionReason::Redundant.label(), "redundant");
669 assert_eq!(
670 QualityRejectionReason::IncompleteReference.label(),
671 "incomplete_reference"
672 );
673 assert_eq!(
674 QualityRejectionReason::Contradiction.label(),
675 "contradiction"
676 );
677 assert_eq!(
678 QualityRejectionReason::LlmLowConfidence.label(),
679 "llm_low_confidence"
680 );
681 }
682
683 #[test]
684 fn rolling_rate_tracker_basic() {
685 let mut tracker = RollingRateTracker::new(4);
686 tracker.push(true);
687 tracker.push(true);
688 tracker.push(false);
689 tracker.push(false);
690 let rate = tracker.rate();
691 assert!((rate - 0.5).abs() < 0.01, "rate should be 0.5, got {rate}");
692 }
693
694 #[test]
695 fn rolling_rate_tracker_evicts_oldest() {
696 let mut tracker = RollingRateTracker::new(3);
697 tracker.push(true); tracker.push(false);
699 tracker.push(false);
700 tracker.push(false); let rate = tracker.rate();
702 assert!(
703 rate < 0.01,
704 "evicted rejection should not count, rate={rate}"
705 );
706 }
707
708 #[test]
709 fn parse_llm_score_valid_json() {
710 let json = r#"{"information_value": 0.8, "reference_completeness": 0.9, "contradiction_risk": 0.1}"#;
711 let score = parse_llm_score(json);
712 assert!(
713 score > 0.7,
714 "high-quality JSON should yield high score, got {score}"
715 );
716 }
717
718 #[test]
719 fn parse_llm_score_malformed_returns_neutral() {
720 let score = parse_llm_score("not json");
721 assert!(
722 (score - 0.5).abs() < 0.01,
723 "malformed JSON should return 0.5"
724 );
725 }
726
727 fn mock_provider() -> zeph_llm::any::AnyProvider {
728 zeph_llm::any::AnyProvider::Mock(zeph_llm::mock::MockProvider::default())
729 }
730
731 #[tokio::test]
732 async fn gate_disabled_always_passes() {
733 let config = QualityGateConfig {
734 enabled: false,
735 ..QualityGateConfig::default()
736 };
737 let gate = QualityGate::new(config);
738 let provider = mock_provider();
739
740 let result = gate.evaluate("yeah he confirmed it", &provider, &[]).await;
741 assert!(result.is_none(), "disabled gate must always pass");
742 }
743
744 #[tokio::test]
745 async fn gate_admits_novel_clean_content() {
746 let config = QualityGateConfig {
747 enabled: true,
748 threshold: 0.3, ..QualityGateConfig::default()
750 };
751 let gate = QualityGate::new(config);
752 let provider = mock_provider();
753
754 let result = gate
756 .evaluate(
757 "The Rust compiler enforces memory safety through the borrow checker.",
758 &provider,
759 &[],
760 )
761 .await;
762 assert!(result.is_none(), "clean novel content should be admitted");
763 }
764
765 #[tokio::test]
766 async fn gate_rejects_pronoun_only_at_low_threshold() {
767 let config = QualityGateConfig {
768 enabled: true,
769 threshold: 0.75, reference_completeness_weight: 0.9,
771 information_value_weight: 0.05,
772 contradiction_weight: 0.05,
773 ..QualityGateConfig::default()
774 };
775 let gate = QualityGate::new(config);
776 let provider = mock_provider();
777
778 let result = gate
779 .evaluate("yeah he confirmed it they said so", &provider, &[])
780 .await;
781 assert!(
782 result == Some(QualityRejectionReason::IncompleteReference),
783 "pronoun-heavy message should be rejected as IncompleteReference, got {result:?}"
784 );
785 }
786
787 #[test]
788 fn quality_gate_counts_rejections() {
789 let config = QualityGateConfig {
790 enabled: true,
791 threshold: 0.99, ..QualityGateConfig::default()
793 };
794 let gate = QualityGate::new(config);
795
796 if let Ok(mut counts) = gate.rejection_counts.lock() {
798 *counts.entry(QualityRejectionReason::Redundant).or_insert(0) += 1;
799 }
800
801 let counts = gate.rejection_counts();
802 assert_eq!(counts.get(&QualityRejectionReason::Redundant), Some(&1));
803 }
804
805 #[tokio::test]
807 async fn gate_fail_open_on_embed_error() {
808 let config = QualityGateConfig {
809 enabled: true,
810 threshold: 0.5,
811 ..QualityGateConfig::default()
812 };
813 let gate = QualityGate::new(config);
814
815 let provider = zeph_llm::any::AnyProvider::Mock(
817 zeph_llm::mock::MockProvider::default().with_embed_invalid_input(),
818 );
819
820 let result = gate
821 .evaluate(
822 "Alice confirmed the meeting at 3pm.",
823 &provider,
824 &[], )
826 .await;
827 assert!(
828 result.is_none(),
829 "embed error must be treated as fail-open (admitted), got {result:?}"
830 );
831 }
832
833 #[tokio::test]
835 async fn gate_rejects_redundant_with_populated_embeddings() {
836 let config = QualityGateConfig {
837 enabled: true,
838 threshold: 0.5,
839 information_value_weight: 0.9,
841 reference_completeness_weight: 0.05,
842 contradiction_weight: 0.05,
843 ..QualityGateConfig::default()
844 };
845 let gate = QualityGate::new(config);
846
847 let fixed_embedding = vec![0.1_f32; 384];
849 let provider = zeph_llm::any::AnyProvider::Mock(
850 zeph_llm::mock::MockProvider::default().with_embedding(fixed_embedding.clone()),
851 );
852
853 let result = gate
855 .evaluate(
856 "The Rust compiler enforces memory safety through the borrow checker.",
857 &provider,
858 &[fixed_embedding],
859 )
860 .await;
861 assert_eq!(
862 result,
863 Some(QualityRejectionReason::Redundant),
864 "identical recent embedding must trigger Redundant rejection"
865 );
866 }
867
868 #[tokio::test]
871 async fn gate_fail_open_on_embed_timeout() {
872 tokio::time::pause();
873
874 let config = QualityGateConfig {
875 enabled: true,
876 threshold: 0.5,
877 information_value_weight: 0.9,
878 reference_completeness_weight: 0.05,
879 contradiction_weight: 0.05,
880 ..QualityGateConfig::default()
881 };
882 let gate = QualityGate::new(config);
883
884 let provider = zeph_llm::any::AnyProvider::Mock(
886 zeph_llm::mock::MockProvider::default().with_embed_delay(10_000),
887 );
888
889 let recent = vec![vec![0.1_f32; 384]];
892
893 let fut = gate.evaluate("Alice confirmed the meeting at 3pm.", &provider, &recent);
894 let (result, ()) = tokio::join!(fut, async {
896 tokio::time::advance(std::time::Duration::from_secs(6)).await;
897 });
898
899 assert!(
900 result.is_none(),
901 "embed timeout must be treated as fail-open (info_val=1.0, admitted), got {result:?}"
902 );
903 }
904
905 #[tokio::test]
908 async fn gate_llm_timeout_falls_back_to_rule_score() {
909 let config = QualityGateConfig {
910 enabled: true,
911 threshold: 0.3, llm_timeout_ms: 50, llm_weight: 0.5,
914 ..QualityGateConfig::default()
915 };
916 let gate = QualityGate::new(config);
917
918 let slow_provider = zeph_llm::any::AnyProvider::Mock(
920 zeph_llm::mock::MockProvider::default().with_delay(600),
921 );
922 let gate = gate.with_llm_provider(slow_provider);
923
924 let embed_provider = mock_provider(); let result = gate
927 .evaluate(
928 "The release is scheduled for next Friday.",
929 &embed_provider,
930 &[],
931 )
932 .await;
933 assert!(
936 result.is_none(),
937 "LLM timeout must fall back to rule score and admit clean content, got {result:?}"
938 );
939 }
940}