1use std::sync::Arc;
10use std::time::Duration;
11
12use zeph_llm::any::AnyProvider;
13use zeph_llm::provider::LlmProvider as _;
14
15use crate::embedding_store::EmbeddingStore;
16use zeph_common::math::cosine_similarity;
17
18#[derive(Debug, Clone, serde::Serialize)]
20pub struct AdmissionFactors {
21 pub future_utility: f32,
23 pub factual_confidence: f32,
25 pub semantic_novelty: f32,
27 pub temporal_recency: f32,
29 pub content_type_prior: f32,
31 pub goal_utility: f32,
34}
35
36#[derive(Debug, Clone)]
38pub struct AdmissionDecision {
39 pub admitted: bool,
40 pub composite_score: f32,
41 pub factors: AdmissionFactors,
42}
43
44#[derive(Debug, Clone, Copy)]
46pub struct AdmissionWeights {
47 pub future_utility: f32,
48 pub factual_confidence: f32,
49 pub semantic_novelty: f32,
50 pub temporal_recency: f32,
51 pub content_type_prior: f32,
52 pub goal_utility: f32,
54}
55
56impl AdmissionWeights {
57 #[must_use]
61 pub fn normalized(&self) -> Self {
62 let fu = self.future_utility.max(0.0);
63 let fc = self.factual_confidence.max(0.0);
64 let sn = self.semantic_novelty.max(0.0);
65 let tr = self.temporal_recency.max(0.0);
66 let cp = self.content_type_prior.max(0.0);
67 let gu = self.goal_utility.max(0.0);
68 let sum = fu + fc + sn + tr + cp + gu;
69 if sum <= f32::EPSILON {
70 return Self {
72 future_utility: 0.2,
73 factual_confidence: 0.2,
74 semantic_novelty: 0.2,
75 temporal_recency: 0.2,
76 content_type_prior: 0.2,
77 goal_utility: 0.0,
78 };
79 }
80 Self {
81 future_utility: fu / sum,
82 factual_confidence: fc / sum,
83 semantic_novelty: sn / sum,
84 temporal_recency: tr / sum,
85 content_type_prior: cp / sum,
86 goal_utility: gu / sum,
87 }
88 }
89}
90
91#[derive(Debug, Clone)]
93pub struct GoalGateConfig {
94 pub threshold: f32,
96 pub provider: Option<AnyProvider>,
98 pub weight: f32,
100}
101
102pub struct AdmissionControl {
128 threshold: f32,
129 fast_path_margin: f32,
130 weights: AdmissionWeights,
131 provider: Option<AnyProvider>,
134 goal_gate: Option<GoalGateConfig>,
136 embed_timeout: Duration,
138}
139
140impl AdmissionControl {
141 #[must_use]
148 pub fn new(threshold: f32, fast_path_margin: f32, weights: AdmissionWeights) -> Self {
149 Self {
150 threshold,
151 fast_path_margin,
152 weights: weights.normalized(),
153 provider: None,
154 goal_gate: None,
155 embed_timeout: Duration::from_secs(5),
156 }
157 }
158
159 #[must_use]
163 pub fn with_embed_timeout(mut self, timeout_secs: u64) -> Self {
164 self.embed_timeout = Duration::from_secs(timeout_secs.max(1));
165 self
166 }
167
168 #[must_use]
172 pub fn with_provider(mut self, provider: AnyProvider) -> Self {
173 self.provider = Some(provider);
174 self
175 }
176
177 #[must_use]
179 pub fn with_goal_gate(mut self, config: GoalGateConfig) -> Self {
180 let gu = config.weight.clamp(0.0, 1.0);
182 let mut weights = self.weights;
183 weights.goal_utility = gu;
184 weights.future_utility = (weights.future_utility - gu).max(0.0);
186 self.weights = weights.normalized();
187 self.goal_gate = Some(config);
188 self
189 }
190
191 #[must_use]
193 pub fn threshold(&self) -> f32 {
194 self.threshold
195 }
196
197 #[cfg_attr(
207 feature = "profiling",
208 tracing::instrument(name = "memory.admission", skip_all)
209 )]
210 pub async fn evaluate(
211 &self,
212 content: &str,
213 role: &str,
214 fallback_provider: &AnyProvider,
215 qdrant: Option<&Arc<EmbeddingStore>>,
216 goal_text: Option<&str>,
217 ) -> AdmissionDecision {
218 let effective_provider = self.provider.as_ref().unwrap_or(fallback_provider);
219 let factual_confidence = compute_factual_confidence(content);
220 let temporal_recency = 1.0f32;
221 let content_type_prior = compute_content_type_prior(role);
222
223 let semantic_novelty =
225 compute_semantic_novelty(content, effective_provider, qdrant, self.embed_timeout).await;
226
227 let goal_utility = match &self.goal_gate {
229 Some(gate) => {
230 let effective_goal = goal_text.filter(|t| t.trim().len() >= 10);
231 if let Some(goal) = effective_goal {
232 compute_goal_utility(
233 content,
234 goal,
235 gate,
236 effective_provider,
237 qdrant,
238 self.embed_timeout,
239 )
240 .await
241 } else {
242 0.0
243 }
244 }
245 None => 0.0,
246 };
247
248 let heuristic_score = self.weighted_score(
250 0.5,
251 factual_confidence,
252 semantic_novelty,
253 temporal_recency,
254 content_type_prior,
255 goal_utility,
256 );
257
258 let future_utility = if heuristic_score >= self.threshold + self.fast_path_margin {
260 0.5 } else {
262 compute_future_utility(content, role, effective_provider).await
263 };
264
265 let composite_score = self.weighted_score(
266 future_utility,
267 factual_confidence,
268 semantic_novelty,
269 temporal_recency,
270 content_type_prior,
271 goal_utility,
272 );
273
274 let admitted = composite_score >= self.threshold
275 || heuristic_score >= self.threshold + self.fast_path_margin;
276
277 AdmissionDecision {
278 admitted,
279 composite_score,
280 factors: AdmissionFactors {
281 future_utility,
282 factual_confidence,
283 semantic_novelty,
284 temporal_recency,
285 content_type_prior,
286 goal_utility,
287 },
288 }
289 }
290
291 fn weighted_score(
292 &self,
293 future_utility: f32,
294 factual_confidence: f32,
295 semantic_novelty: f32,
296 temporal_recency: f32,
297 content_type_prior: f32,
298 goal_utility: f32,
299 ) -> f32 {
300 future_utility * self.weights.future_utility
301 + factual_confidence * self.weights.factual_confidence
302 + semantic_novelty * self.weights.semantic_novelty
303 + temporal_recency * self.weights.temporal_recency
304 + content_type_prior * self.weights.content_type_prior
305 + goal_utility * self.weights.goal_utility
306 }
307}
308
309#[must_use]
313pub fn compute_factual_confidence(content: &str) -> f32 {
314 const HEDGING_MARKERS: &[&str] = &[
316 "maybe",
317 "might",
318 "perhaps",
319 "i think",
320 "i believe",
321 "not sure",
322 "could be",
323 "possibly",
324 "probably",
325 "uncertain",
326 "not certain",
327 "i'm not sure",
328 "im not sure",
329 "not confident",
330 ];
331 let lower = content.to_lowercase();
332 let matches = HEDGING_MARKERS
333 .iter()
334 .filter(|&&m| lower.contains(m))
335 .count();
336 #[allow(clippy::cast_precision_loss)]
338 let penalty = (matches as f32) * 0.1;
339 (1.0 - penalty).max(0.2)
340}
341
342#[must_use]
347pub fn compute_content_type_prior(role: &str) -> f32 {
348 match role {
349 "user" => 0.7,
350 "assistant" => 0.6,
351 "tool" | "tool_result" => 0.8,
352 "system" => 0.3,
353 _ => 0.5,
354 }
355}
356
357#[tracing::instrument(name = "memory.admission.semantic_novelty", skip_all)]
361async fn compute_semantic_novelty(
362 content: &str,
363 provider: &AnyProvider,
364 qdrant: Option<&Arc<EmbeddingStore>>,
365 embed_timeout: Duration,
366) -> f32 {
367 let Some(store) = qdrant else {
368 return 1.0;
369 };
370 if !provider.supports_embeddings() {
371 return 1.0;
372 }
373 let vector = match crate::llm_judge::embed_with_timeout_fail_open(
374 provider,
375 content,
376 embed_timeout,
377 1.0,
378 "A-MAC: semantic_novelty",
379 )
380 .await
381 {
382 Ok(v) => v,
383 Err(fail_open) => return fail_open,
384 };
385 if let Err(e) = store.ensure_collection_for_vector(&vector).await {
386 tracing::debug!(error = %e, "A-MAC: collection not ready for novelty check");
387 return 1.0;
388 }
389 let results = match store.search(&vector, 3, None).await {
390 Ok(r) => r,
391 Err(e) => {
392 tracing::debug!(error = %e, "A-MAC: novelty search failed, using 1.0");
393 return 1.0;
394 }
395 };
396 let max_sim = results.iter().map(|r| r.score).fold(0.0f32, f32::max);
397 (1.0 - max_sim).max(0.0)
398}
399
400#[tracing::instrument(name = "memory.admission.future_utility_llm", skip_all)]
404async fn compute_future_utility(content: &str, role: &str, provider: &AnyProvider) -> f32 {
405 let system = "You are a memory relevance judge. Rate how likely this message will be \
406 referenced in future conversations on a scale of 0.0 to 1.0. \
407 Respond with ONLY a decimal number between 0.0 and 1.0, nothing else.";
408
409 let user = format!(
410 "Role: {role}\nContent: {}\n\nFuture utility score (0.0-1.0):",
411 content.chars().take(500).collect::<String>()
412 );
413
414 crate::llm_judge::llm_judge_score(
415 provider,
416 system,
417 user,
418 Duration::from_secs(8),
419 0.5,
420 "A-MAC: future_utility",
421 |s| s.trim().parse::<f32>().ok(),
422 )
423 .await
424}
425
426async fn compute_goal_utility(
437 content: &str,
438 goal_text: &str,
439 gate: &GoalGateConfig,
440 provider: &AnyProvider,
441 qdrant: Option<&Arc<EmbeddingStore>>,
442 embed_timeout: Duration,
443) -> f32 {
444 use zeph_llm::provider::LlmProvider as _;
445
446 if !provider.supports_embeddings() {
447 return 0.0;
448 }
449
450 let goal_emb = match crate::llm_judge::embed_with_timeout_fail_open(
451 provider,
452 goal_text,
453 embed_timeout,
454 0.0,
455 "goal_utility: goal text",
456 )
457 .await
458 {
459 Ok(v) => v,
460 Err(fail_open) => return fail_open,
461 };
462 let content_emb = match crate::llm_judge::embed_with_timeout_fail_open(
463 provider,
464 content,
465 embed_timeout,
466 0.0,
467 "goal_utility: content",
468 )
469 .await
470 {
471 Ok(v) => v,
472 Err(fail_open) => return fail_open,
473 };
474
475 let _ = qdrant; let similarity = cosine_similarity(&goal_emb, &content_emb);
479
480 let borderline_lo = gate.threshold - 0.1;
482 let borderline_hi = gate.threshold + 0.1;
483 let in_borderline = similarity >= borderline_lo && similarity <= borderline_hi;
484
485 let final_similarity = if in_borderline {
486 if let Some(ref goal_provider) = gate.provider {
487 refine_goal_utility_llm(content, goal_text, similarity, goal_provider).await
488 } else {
489 similarity
490 }
491 } else {
492 similarity
493 };
494
495 if final_similarity < gate.threshold {
497 0.0
498 } else {
499 final_similarity.max(0.1)
500 }
501}
502
503#[tracing::instrument(name = "memory.admission.goal_utility_refine_llm", skip_all)]
507async fn refine_goal_utility_llm(
508 content: &str,
509 goal_text: &str,
510 embedding_sim: f32,
511 provider: &AnyProvider,
512) -> f32 {
513 let system = "You are a memory relevance judge. Given a task goal and a candidate memory, \
514 rate how relevant the memory is to the goal on a scale of 0.0 to 1.0. \
515 Respond with ONLY a decimal number between 0.0 and 1.0, nothing else.";
516
517 let user = format!(
518 "Goal: {}\nMemory: {}\n\nRelevance score (0.0-1.0):",
519 goal_text.chars().take(200).collect::<String>(),
520 content.chars().take(300).collect::<String>(),
521 );
522
523 crate::llm_judge::llm_judge_score(
524 provider,
525 system,
526 user,
527 Duration::from_secs(6),
528 embedding_sim,
529 "goal_utility LLM refinement",
530 |s| s.trim().parse::<f32>().ok(),
531 )
532 .await
533}
534
535pub fn log_admission_decision(
539 decision: &AdmissionDecision,
540 content_preview: &str,
541 role: &str,
542 threshold: f32,
543) {
544 if decision.admitted {
545 tracing::trace!(
546 role,
547 composite_score = decision.composite_score,
548 threshold,
549 content_preview,
550 "A-MAC: admitted"
551 );
552 } else {
553 tracing::debug!(
554 role,
555 composite_score = decision.composite_score,
556 threshold,
557 future_utility = decision.factors.future_utility,
558 factual_confidence = decision.factors.factual_confidence,
559 semantic_novelty = decision.factors.semantic_novelty,
560 content_type_prior = decision.factors.content_type_prior,
561 content_preview,
562 "A-MAC: rejected"
563 );
564 }
565}
566
567#[derive(Debug)]
569pub struct AdmissionRejected {
570 pub composite_score: f32,
571 pub threshold: f32,
572}
573
574impl std::fmt::Display for AdmissionRejected {
575 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
576 write!(
577 f,
578 "A-MAC admission rejected (score={:.3} < threshold={:.3})",
579 self.composite_score, self.threshold
580 )
581 }
582}
583
584#[cfg(test)]
585mod tests {
586 use super::*;
587
588 #[test]
589 fn factual_confidence_no_hedging() {
590 assert!((compute_factual_confidence("The server uses TLS 1.3.") - 1.0).abs() < 0.01);
591 }
592
593 #[test]
594 fn factual_confidence_with_one_marker() {
595 let score = compute_factual_confidence("Maybe we should use TLS 1.3.");
596 assert!((score - 0.9).abs() < 0.01);
597 }
598
599 #[test]
600 fn factual_confidence_many_markers_floors_at_0_2() {
601 let content = "maybe i think perhaps possibly might not sure i believe";
602 let score = compute_factual_confidence(content);
603 assert!(score >= 0.2);
604 assert!(score < 0.5);
605 }
606
607 #[test]
608 fn content_type_prior_values() {
609 assert!((compute_content_type_prior("user") - 0.7).abs() < 0.01);
610 assert!((compute_content_type_prior("assistant") - 0.6).abs() < 0.01);
611 assert!((compute_content_type_prior("tool") - 0.8).abs() < 0.01);
612 assert!((compute_content_type_prior("system") - 0.3).abs() < 0.01);
613 assert!((compute_content_type_prior("unknown") - 0.5).abs() < 0.01);
614 }
615
616 #[test]
617 fn admission_control_admits_high_score() {
618 let weights = AdmissionWeights {
619 future_utility: 0.30,
620 factual_confidence: 0.15,
621 semantic_novelty: 0.30,
622 temporal_recency: 0.10,
623 content_type_prior: 0.15,
624 goal_utility: 0.0,
625 };
626 let ctrl = AdmissionControl::new(0.40, 0.15, weights);
627 let score = ctrl.weighted_score(1.0, 1.0, 1.0, 1.0, 1.0, 0.0);
629 assert!(score >= 0.99);
630 let admitted = score >= ctrl.threshold;
632 assert!(admitted);
633 }
634
635 #[test]
636 fn admission_control_rejects_low_score() {
637 let weights = AdmissionWeights {
638 future_utility: 0.30,
639 factual_confidence: 0.15,
640 semantic_novelty: 0.30,
641 temporal_recency: 0.10,
642 content_type_prior: 0.15,
643 goal_utility: 0.0,
644 };
645 let ctrl = AdmissionControl::new(0.40, 0.15, weights);
646 let score = ctrl.weighted_score(0.0, 0.0, 0.0, 0.0, 0.0, 0.0);
648 assert!(score < ctrl.threshold);
649 }
650
651 #[test]
654 fn fast_path_admits_when_heuristic_above_threshold_plus_margin() {
655 let weights = AdmissionWeights {
656 future_utility: 0.20,
657 factual_confidence: 0.20,
658 semantic_novelty: 0.20,
659 temporal_recency: 0.20,
660 content_type_prior: 0.20,
661 goal_utility: 0.0,
662 };
663 let threshold = 0.40f32;
664 let margin = 0.15f32;
665 let ctrl = AdmissionControl::new(threshold, margin, weights);
666
667 let heuristic = ctrl.weighted_score(0.5, 1.0, 1.0, 1.0, 1.0, 0.0);
669 assert!(
671 heuristic >= threshold + margin,
672 "heuristic {heuristic} must exceed threshold+margin {}",
673 threshold + margin
674 );
675 let admitted = heuristic >= threshold + margin;
677 assert!(admitted, "fast path must admit without LLM call");
678 }
679
680 #[test]
682 fn slow_path_required_when_heuristic_below_threshold_plus_margin() {
683 let weights = AdmissionWeights {
684 future_utility: 0.40,
685 factual_confidence: 0.15,
686 semantic_novelty: 0.15,
687 temporal_recency: 0.15,
688 content_type_prior: 0.15,
689 goal_utility: 0.0,
690 };
691 let threshold = 0.50f32;
692 let margin = 0.20f32;
693 let ctrl = AdmissionControl::new(threshold, margin, weights);
694
695 let heuristic = ctrl.weighted_score(0.5, 0.3, 0.3, 0.3, 0.3, 0.0);
697 assert!(
698 heuristic < threshold + margin,
699 "heuristic {heuristic} must be below threshold+margin {}",
700 threshold + margin
701 );
702 }
703
704 #[test]
706 fn log_admission_decision_does_not_panic() {
707 let admitted_decision = AdmissionDecision {
708 admitted: true,
709 composite_score: 0.75,
710 factors: AdmissionFactors {
711 future_utility: 0.8,
712 factual_confidence: 0.9,
713 semantic_novelty: 0.7,
714 temporal_recency: 1.0,
715 content_type_prior: 0.7,
716 goal_utility: 0.0,
717 },
718 };
719 log_admission_decision(&admitted_decision, "preview text", "user", 0.40);
720
721 let rejected_decision = AdmissionDecision {
722 admitted: false,
723 composite_score: 0.20,
724 factors: AdmissionFactors {
725 future_utility: 0.1,
726 factual_confidence: 0.2,
727 semantic_novelty: 0.3,
728 temporal_recency: 1.0,
729 content_type_prior: 0.3,
730 goal_utility: 0.0,
731 },
732 };
733 log_admission_decision(&rejected_decision, "maybe short content", "assistant", 0.40);
734 }
735
736 #[test]
738 fn admission_rejected_display() {
739 let err = AdmissionRejected {
740 composite_score: 0.25,
741 threshold: 0.45,
742 };
743 let msg = format!("{err}");
744 assert!(msg.contains("0.250"));
745 assert!(msg.contains("0.450"));
746 }
747
748 #[test]
750 fn threshold_accessor() {
751 let weights = AdmissionWeights {
752 future_utility: 0.20,
753 factual_confidence: 0.20,
754 semantic_novelty: 0.20,
755 temporal_recency: 0.20,
756 content_type_prior: 0.20,
757 goal_utility: 0.0,
758 };
759 let ctrl = AdmissionControl::new(0.55, 0.10, weights);
760 assert!((ctrl.threshold() - 0.55).abs() < 0.001);
761 }
762
763 #[test]
765 fn content_type_prior_tool_result_alias() {
766 assert!((compute_content_type_prior("tool_result") - 0.8).abs() < 0.01);
767 }
768
769 #[test]
772 fn cosine_similarity_identical_vectors() {
773 let v = vec![1.0f32, 0.0, 0.0];
774 assert!((cosine_similarity(&v, &v) - 1.0).abs() < 1e-6);
775 }
776
777 #[test]
778 fn cosine_similarity_orthogonal_vectors() {
779 let a = vec![1.0f32, 0.0];
780 let b = vec![0.0f32, 1.0];
781 assert!(cosine_similarity(&a, &b).abs() < 1e-6);
782 }
783
784 #[test]
785 fn cosine_similarity_zero_vector_returns_zero() {
786 let z = vec![0.0f32, 0.0, 0.0];
787 let v = vec![1.0f32, 2.0, 3.0];
788 assert!(cosine_similarity(&z, &v).abs() < f32::EPSILON);
789 }
790
791 #[test]
792 fn cosine_similarity_length_mismatch_returns_zero() {
793 let a = vec![1.0f32, 0.0];
794 let b = vec![1.0f32, 0.0, 0.0];
795 assert!(cosine_similarity(&a, &b).abs() < f32::EPSILON);
796 }
797
798 #[test]
801 fn with_goal_gate_sets_goal_utility_weight() {
802 let weights = AdmissionWeights {
803 future_utility: 0.30,
804 factual_confidence: 0.15,
805 semantic_novelty: 0.30,
806 temporal_recency: 0.10,
807 content_type_prior: 0.15,
808 goal_utility: 0.0,
809 };
810 let ctrl = AdmissionControl::new(0.40, 0.15, weights);
811 let config = GoalGateConfig {
812 weight: 0.20,
813 threshold: 0.5,
814 provider: None,
815 };
816 let ctrl = ctrl.with_goal_gate(config);
817 assert!(
818 ctrl.weights.goal_utility > 0.0,
819 "goal_utility must be nonzero after with_goal_gate"
820 );
821 let w = &ctrl.weights;
823 let total = w.future_utility
824 + w.factual_confidence
825 + w.semantic_novelty
826 + w.temporal_recency
827 + w.content_type_prior
828 + w.goal_utility;
829 assert!(
830 (total - 1.0).abs() < 0.01,
831 "normalized weights must sum to 1.0, got {total}"
832 );
833 }
834
835 #[test]
836 fn with_goal_gate_zero_weight_leaves_goal_utility_at_zero() {
837 let weights = AdmissionWeights {
838 future_utility: 0.30,
839 factual_confidence: 0.15,
840 semantic_novelty: 0.30,
841 temporal_recency: 0.10,
842 content_type_prior: 0.15,
843 goal_utility: 0.0,
844 };
845 let ctrl = AdmissionControl::new(0.40, 0.15, weights);
846 let config = GoalGateConfig {
847 weight: 0.0,
848 threshold: 0.5,
849 provider: None,
850 };
851 let ctrl = ctrl.with_goal_gate(config);
852 assert!(ctrl.weights.goal_utility.abs() < f32::EPSILON);
853 }
854
855 #[tokio::test]
858 async fn compute_semantic_novelty_returns_one_on_embed_timeout() {
859 tokio::time::pause();
860 let mock = zeph_llm::mock::MockProvider::default()
861 .with_embed_delay(10_000)
862 .with_embedding(vec![0.0; 4]);
863 let provider = zeph_llm::any::AnyProvider::Mock(mock);
864 let novelty_fut = async move {
865 compute_semantic_novelty("hello", &provider, None, Duration::from_secs(5)).await
866 };
867 let handle = tokio::spawn(novelty_fut); tokio::time::advance(std::time::Duration::from_secs(6)).await;
869 let result = handle.await.expect("task panicked");
870 assert!(
871 (result - 1.0).abs() < f32::EPSILON,
872 "expected 1.0 on embed timeout, got {result}"
873 );
874 }
875
876 #[tokio::test]
877 async fn compute_goal_utility_returns_zero_on_embed_timeout() {
878 tokio::time::pause();
879 let mock = zeph_llm::mock::MockProvider::default()
880 .with_embed_delay(10_000)
881 .with_embedding(vec![0.0; 4]);
882 let provider = zeph_llm::any::AnyProvider::Mock(mock);
883 let gate = GoalGateConfig {
884 weight: 0.5,
885 threshold: 0.5,
886 provider: None,
887 };
888 let goal_fut = async move {
889 compute_goal_utility(
890 "content",
891 "goal",
892 &gate,
893 &provider,
894 None,
895 Duration::from_secs(5),
896 )
897 .await
898 };
899 let handle = tokio::spawn(goal_fut); tokio::time::advance(std::time::Duration::from_secs(6)).await;
901 let result = handle.await.expect("task panicked");
902 assert!(
903 result.abs() < f32::EPSILON,
904 "expected 0.0 on embed timeout, got {result}"
905 );
906 }
907}