1use std::collections::HashMap;
8
9use uuid::Uuid;
10
11use khive_fold::objective::{Objective, ObjectiveContext};
12use khive_fold::ordering::HasId;
13
14#[derive(Debug, Clone)]
20pub struct RetrievalCandidate {
21 pub id: Uuid,
23 pub vector_score: Option<f64>,
25 pub text_score: Option<f64>,
27 pub graph_distance: Option<u32>,
29 pub rrf_score: Option<f64>,
31}
32
33impl HasId for RetrievalCandidate {
34 #[inline]
35 fn id(&self) -> Uuid {
36 self.id
37 }
38}
39
40pub struct VectorSimilarityObjective;
46
47impl Objective<RetrievalCandidate> for VectorSimilarityObjective {
48 #[inline]
49 fn score(&self, candidate: &RetrievalCandidate, _context: &ObjectiveContext) -> f64 {
50 candidate.vector_score.unwrap_or(0.0)
51 }
52
53 fn name(&self) -> &str {
54 "VectorSimilarityObjective"
55 }
56}
57
58pub struct TextRelevanceObjective;
64
65impl Objective<RetrievalCandidate> for TextRelevanceObjective {
66 #[inline]
67 fn score(&self, candidate: &RetrievalCandidate, _context: &ObjectiveContext) -> f64 {
68 candidate.text_score.unwrap_or(0.0)
69 }
70
71 fn name(&self) -> &str {
72 "TextRelevanceObjective"
73 }
74}
75
76pub struct GraphProximityObjective {
91 pub max_distance: u32,
93}
94
95impl Objective<RetrievalCandidate> for GraphProximityObjective {
96 fn score(&self, candidate: &RetrievalCandidate, _context: &ObjectiveContext) -> f64 {
97 let d = match candidate.graph_distance {
98 Some(d) => d,
99 None => return 0.0,
100 };
101 if self.max_distance == 0 || d >= self.max_distance {
102 return 0.0;
103 }
104 1.0 - (d as f64 / self.max_distance as f64)
105 }
106
107 fn name(&self) -> &str {
108 "GraphProximityObjective"
109 }
110}
111
112pub struct RrfFusionObjective;
121
122impl Objective<RetrievalCandidate> for RrfFusionObjective {
123 #[inline]
124 fn score(&self, candidate: &RetrievalCandidate, _context: &ObjectiveContext) -> f64 {
125 candidate.rrf_score.unwrap_or(0.0)
126 }
127
128 fn name(&self) -> &str {
129 "RrfFusionObjective"
130 }
131}
132
133impl Objective<NoteCandidate> for RrfFusionObjective {
134 #[inline]
135 fn score(&self, candidate: &NoteCandidate, _context: &ObjectiveContext) -> f64 {
136 candidate.rrf_score.unwrap_or(0.0)
137 }
138
139 fn name(&self) -> &str {
140 "RrfFusionObjective"
141 }
142}
143
144#[derive(Debug, Clone)]
153pub struct NoteCandidate {
154 pub id: Uuid,
156 pub rrf_score: Option<f64>,
158 pub salience: f64,
160 pub decay_factor: f64,
162 pub age_days: f64,
164 pub effective_salience: f64,
170 pub rerank_scores: HashMap<String, f64>,
173}
174
175impl HasId for NoteCandidate {
176 #[inline]
177 fn id(&self) -> Uuid {
178 self.id
179 }
180}
181
182pub struct DecayAwareSalienceObjective {
194 pub decay_rate: f64,
197}
198
199impl DecayAwareSalienceObjective {
200 pub fn new(decay_rate: f64) -> Self {
204 Self { decay_rate }
205 }
206
207 pub fn default_memory() -> Self {
209 Self::new(0.01)
210 }
211}
212
213impl Objective<NoteCandidate> for DecayAwareSalienceObjective {
214 #[inline]
215 fn score(&self, candidate: &NoteCandidate, _context: &ObjectiveContext) -> f64 {
216 candidate.salience * (-candidate.decay_factor * candidate.age_days).exp()
217 }
218
219 fn name(&self) -> &str {
220 "DecayAwareSalienceObjective"
221 }
222}
223
224pub struct AmplifiedDecayAwareSalienceObjective {
239 pub alpha: f64,
241}
242
243impl AmplifiedDecayAwareSalienceObjective {
244 pub fn new(alpha: f64) -> Self {
246 Self { alpha }
247 }
248
249 pub fn default_memory() -> Self {
251 Self::new(1.5)
252 }
253}
254
255impl Objective<NoteCandidate> for AmplifiedDecayAwareSalienceObjective {
256 #[inline]
257 fn score(&self, candidate: &NoteCandidate, _context: &ObjectiveContext) -> f64 {
258 candidate.effective_salience.powf(self.alpha)
261 }
262
263 fn name(&self) -> &str {
264 "AmplifiedDecayAwareSalienceObjective"
265 }
266}
267
268pub struct TemporalRecencyObjective {
280 pub half_life_days: f64,
282}
283
284impl TemporalRecencyObjective {
285 pub fn default_memory() -> Self {
287 Self {
288 half_life_days: 30.0,
289 }
290 }
291}
292
293impl Objective<NoteCandidate> for TemporalRecencyObjective {
294 #[inline]
295 fn score(&self, candidate: &NoteCandidate, _context: &ObjectiveContext) -> f64 {
296 let k = std::f64::consts::LN_2 / self.half_life_days.max(f64::EPSILON);
297 (-k * candidate.age_days).exp()
298 }
299
300 fn name(&self) -> &str {
301 "TemporalRecencyObjective"
302 }
303}
304
305pub struct RerankerObjective {
314 pub reranker_name: String,
316}
317
318impl RerankerObjective {
319 pub fn new(name: impl Into<String>) -> Self {
321 Self {
322 reranker_name: name.into(),
323 }
324 }
325}
326
327impl Objective<NoteCandidate> for RerankerObjective {
328 #[inline]
329 fn score(&self, candidate: &NoteCandidate, _context: &ObjectiveContext) -> f64 {
330 candidate
331 .rerank_scores
332 .get(&self.reranker_name)
333 .copied()
334 .unwrap_or(0.0)
335 }
336
337 fn name(&self) -> &str {
338 "RerankerObjective"
339 }
340}
341
342pub struct MemoryRecallPipeline {
351 pipeline: khive_fold::WeightedObjective<NoteCandidate>,
352}
353
354impl MemoryRecallPipeline {
355 pub fn new(
362 relevance_weight: f64,
363 salience_weight: f64,
364 temporal_weight: f64,
365 half_life_days: f64,
366 salience_alpha: f64,
367 ) -> Self {
368 use khive_fold::WeightedObjective;
369 let pipeline = WeightedObjective::<NoteCandidate>::new()
370 .add(Box::new(RrfFusionObjective), relevance_weight)
371 .add(
372 Box::new(AmplifiedDecayAwareSalienceObjective::new(salience_alpha)),
373 salience_weight,
374 )
375 .add(
376 Box::new(TemporalRecencyObjective { half_life_days }),
377 temporal_weight,
378 );
379 Self { pipeline }
380 }
381
382 pub fn default_memory() -> Self {
386 Self::new(0.70, 0.20, 0.10, 30.0, 1.5)
387 }
388
389 pub fn score(&self, candidate: &NoteCandidate) -> f64 {
394 let ctx = ObjectiveContext::new();
395 use khive_fold::objective::Objective;
396 self.pipeline.score(candidate, &ctx).clamp(0.0, 1.0)
397 }
398}
399
400#[cfg(test)]
405mod tests {
406 use super::*;
407 use khive_fold::objective::{Objective, ObjectiveContext};
408 use khive_fold::WeightedObjective;
409 use uuid::Uuid;
410
411 fn ctx() -> ObjectiveContext {
412 ObjectiveContext::new()
413 }
414
415 fn candidate(
416 vector: Option<f64>,
417 text: Option<f64>,
418 dist: Option<u32>,
419 rrf: Option<f64>,
420 ) -> RetrievalCandidate {
421 RetrievalCandidate {
422 id: Uuid::new_v4(),
423 vector_score: vector,
424 text_score: text,
425 graph_distance: dist,
426 rrf_score: rrf,
427 }
428 }
429
430 fn note_candidate(
431 rrf: Option<f64>,
432 salience: f64,
433 decay_factor: f64,
434 age_days: f64,
435 ) -> NoteCandidate {
436 let effective_salience = salience * (-decay_factor * age_days).exp();
438 NoteCandidate {
439 id: Uuid::new_v4(),
440 rrf_score: rrf,
441 salience,
442 decay_factor,
443 age_days,
444 effective_salience,
445 rerank_scores: HashMap::new(),
446 }
447 }
448
449 #[test]
452 fn vector_present_returns_signal() {
453 let c = candidate(Some(0.85), None, None, None);
454 let score = VectorSimilarityObjective.score(&c, &ctx());
455 assert!((score - 0.85).abs() < 1e-12);
456 }
457
458 #[test]
459 fn vector_absent_returns_zero() {
460 let c = candidate(None, None, None, None);
461 assert_eq!(VectorSimilarityObjective.score(&c, &ctx()), 0.0);
462 }
463
464 #[test]
465 fn vector_zero_score_returns_zero() {
466 let c = candidate(Some(0.0), None, None, None);
467 assert_eq!(VectorSimilarityObjective.score(&c, &ctx()), 0.0);
468 }
469
470 #[test]
473 fn text_present_returns_signal() {
474 let c = candidate(None, Some(0.6), None, None);
475 let score = TextRelevanceObjective.score(&c, &ctx());
476 assert!((score - 0.6).abs() < 1e-12);
477 }
478
479 #[test]
480 fn text_absent_returns_zero() {
481 let c = candidate(None, None, None, None);
482 assert_eq!(TextRelevanceObjective.score(&c, &ctx()), 0.0);
483 }
484
485 #[test]
488 fn graph_anchor_hit_scores_one() {
489 let c = candidate(None, None, Some(0), None);
490 let obj = GraphProximityObjective { max_distance: 3 };
491 assert!((obj.score(&c, &ctx()) - 1.0).abs() < 1e-12);
492 }
493
494 #[test]
495 fn graph_midpoint_scores_half() {
496 let c = candidate(None, None, Some(1), None);
497 let obj = GraphProximityObjective { max_distance: 2 };
498 assert!((obj.score(&c, &ctx()) - 0.5).abs() < 1e-12);
499 }
500
501 #[test]
502 fn graph_at_boundary_scores_zero() {
503 let c = candidate(None, None, Some(3), None);
504 let obj = GraphProximityObjective { max_distance: 3 };
505 assert_eq!(obj.score(&c, &ctx()), 0.0);
506 }
507
508 #[test]
509 fn graph_beyond_boundary_scores_zero() {
510 let c = candidate(None, None, Some(10), None);
511 let obj = GraphProximityObjective { max_distance: 3 };
512 assert_eq!(obj.score(&c, &ctx()), 0.0);
513 }
514
515 #[test]
516 fn graph_absent_scores_zero() {
517 let c = candidate(None, None, None, None);
518 let obj = GraphProximityObjective { max_distance: 3 };
519 assert_eq!(obj.score(&c, &ctx()), 0.0);
520 }
521
522 #[test]
523 fn graph_max_distance_zero_always_scores_zero() {
524 let c = candidate(None, None, Some(0), None);
526 let obj = GraphProximityObjective { max_distance: 0 };
527 assert_eq!(obj.score(&c, &ctx()), 0.0);
528 }
529
530 #[test]
533 fn rrf_present_returns_signal() {
534 let c = candidate(None, None, None, Some(0.0327));
535 let score = RrfFusionObjective.score(&c, &ctx());
536 assert!((score - 0.0327).abs() < 1e-12);
537 }
538
539 #[test]
540 fn rrf_absent_returns_zero() {
541 let c = candidate(None, None, None, None);
542 assert_eq!(RrfFusionObjective.score(&c, &ctx()), 0.0);
543 }
544
545 #[test]
548 fn weighted_composition_vector_and_text() {
549 let c = candidate(Some(0.8), Some(0.6), None, None);
550
551 let obj = WeightedObjective::<RetrievalCandidate>::new()
552 .add(Box::new(VectorSimilarityObjective), 0.5)
553 .add(Box::new(TextRelevanceObjective), 0.5);
554
555 let score = obj.score(&c, &ctx());
556 assert!((score - 0.7).abs() < 1e-12);
558 }
559
560 #[test]
561 fn weighted_composition_with_graph() {
562 let c = candidate(Some(1.0), Some(0.0), Some(1), None);
563
564 let obj = WeightedObjective::<RetrievalCandidate>::new()
565 .add(Box::new(VectorSimilarityObjective), 0.4)
566 .add(Box::new(TextRelevanceObjective), 0.3)
567 .add(Box::new(GraphProximityObjective { max_distance: 4 }), 0.3);
568
569 let score = obj.score(&c, &ctx());
570 assert!((score - 0.625).abs() < 1e-12);
571 }
572
573 #[test]
574 fn weighted_all_absent_returns_zero() {
575 let c = candidate(None, None, None, None);
576
577 let obj = WeightedObjective::<RetrievalCandidate>::new()
578 .add(Box::new(VectorSimilarityObjective), 0.5)
579 .add(Box::new(TextRelevanceObjective), 0.5);
580
581 assert_eq!(obj.score(&c, &ctx()), 0.0);
583 }
584
585 #[test]
588 fn has_id_returns_candidate_uuid() {
589 let id = Uuid::new_v4();
590 let c = RetrievalCandidate {
591 id,
592 vector_score: None,
593 text_score: None,
594 graph_distance: None,
595 rrf_score: None,
596 };
597 assert_eq!(c.id(), id);
598 }
599
600 #[test]
603 fn select_top_orders_by_vector_score() {
604 use khive_fold::DeterministicObjective;
605
606 let candidates = vec![
607 candidate(Some(0.3), None, None, None),
608 candidate(Some(0.9), None, None, None),
609 candidate(Some(0.6), None, None, None),
610 ];
611
612 let top = VectorSimilarityObjective.select_top_deterministic(&candidates, 2, &ctx());
613
614 assert_eq!(top.len(), 2);
615 assert!((top[0].score - 0.9).abs() < 1e-12);
616 assert!((top[1].score - 0.6).abs() < 1e-12);
617 }
618
619 #[test]
622 fn note_candidate_has_id_returns_uuid() {
623 let id = Uuid::new_v4();
624 let c = NoteCandidate {
625 id,
626 rrf_score: None,
627 salience: 0.5,
628 decay_factor: 0.01,
629 age_days: 0.0,
630 effective_salience: 0.5,
631 rerank_scores: HashMap::new(),
632 };
633 assert_eq!(c.id(), id);
634 }
635
636 #[test]
639 fn decay_aware_zero_age_returns_full_salience() {
640 let obj = DecayAwareSalienceObjective::new(0.01);
641 let c = note_candidate(None, 0.8, 0.01, 0.0);
642 let score = obj.score(&c, &ctx());
643 assert!((score - 0.8).abs() < 1e-12, "got {score}");
644 }
645
646 #[test]
647 fn decay_aware_uses_note_decay_factor_not_field() {
648 let obj = DecayAwareSalienceObjective::new(0.99); let c = note_candidate(None, 1.0, 0.01, 100.0);
651 let score = obj.score(&c, &ctx());
652 let expected = (-0.01_f64 * 100.0).exp();
653 assert!(
654 (score - expected).abs() < 1e-12,
655 "got {score}, expected {expected}"
656 );
657 }
658
659 #[test]
660 fn decay_aware_high_decay_reduces_score_faster() {
661 let obj = DecayAwareSalienceObjective::new(0.0);
662 let slow = note_candidate(None, 1.0, 0.001, 100.0);
663 let fast = note_candidate(None, 1.0, 0.1, 100.0);
664 let score_slow = obj.score(&slow, &ctx());
665 let score_fast = obj.score(&fast, &ctx());
666 assert!(
667 score_slow > score_fast,
668 "slow decay should score higher: {score_slow} vs {score_fast}"
669 );
670 }
671
672 #[test]
675 fn temporal_score_one_at_zero_age() {
676 let obj = TemporalRecencyObjective {
677 half_life_days: 30.0,
678 };
679 let c = note_candidate(None, 0.5, 0.01, 0.0);
680 let score = obj.score(&c, &ctx());
681 assert!((score - 1.0).abs() < 1e-12, "got {score}");
682 }
683
684 #[test]
685 fn temporal_score_half_at_half_life() {
686 let half_life = 30.0;
687 let obj = TemporalRecencyObjective {
688 half_life_days: half_life,
689 };
690 let c = note_candidate(None, 0.5, 0.01, half_life);
691 let score = obj.score(&c, &ctx());
692 assert!(
693 (score - 0.5).abs() < 1e-10,
694 "expected 0.5 at half_life, got {score}"
695 );
696 }
697
698 #[test]
699 fn temporal_score_decreases_with_age() {
700 let obj = TemporalRecencyObjective {
701 half_life_days: 30.0,
702 };
703 let young = note_candidate(None, 1.0, 0.01, 10.0);
704 let old = note_candidate(None, 1.0, 0.01, 100.0);
705 let score_young = obj.score(&young, &ctx());
706 let score_old = obj.score(&old, &ctx());
707 assert!(
708 score_young > score_old,
709 "younger note should score higher: {score_young} vs {score_old}"
710 );
711 }
712
713 #[test]
716 fn reranker_returns_named_score() {
717 let mut c = note_candidate(None, 0.5, 0.01, 0.0);
718 c.rerank_scores.insert("cross_encoder".to_string(), 0.9);
719 let obj = RerankerObjective::new("cross_encoder");
720 let score = obj.score(&c, &ctx());
721 assert!((score - 0.9).abs() < 1e-12, "got {score}");
722 }
723
724 #[test]
725 fn reranker_absent_key_returns_zero() {
726 let c = note_candidate(None, 0.5, 0.01, 0.0);
727 let obj = RerankerObjective::new("cross_encoder");
728 let score = obj.score(&c, &ctx());
729 assert_eq!(score, 0.0);
730 }
731
732 #[test]
733 fn reranker_different_keys_independent() {
734 let mut c = note_candidate(None, 0.5, 0.01, 0.0);
735 c.rerank_scores.insert("salience".to_string(), 0.7);
736 let obj_ce = RerankerObjective::new("cross_encoder");
737 let obj_sal = RerankerObjective::new("salience");
738 assert_eq!(obj_ce.score(&c, &ctx()), 0.0);
739 assert!((obj_sal.score(&c, &ctx()) - 0.7).abs() < 1e-12);
740 }
741
742 #[test]
745 fn memory_pipeline_weighted_composition() {
746 let c = NoteCandidate {
748 id: Uuid::new_v4(),
749 rrf_score: Some(0.5),
750 salience: 0.8,
751 decay_factor: 0.01,
752 age_days: 0.0,
753 effective_salience: 0.8,
754 rerank_scores: HashMap::new(),
755 };
756 let pipeline = WeightedObjective::<NoteCandidate>::new()
757 .add(Box::new(RrfFusionObjective), 0.70)
758 .add(Box::new(DecayAwareSalienceObjective::new(0.0)), 0.20)
759 .add(
760 Box::new(TemporalRecencyObjective {
761 half_life_days: 30.0,
762 }),
763 0.10,
764 );
765 let score = pipeline.score(&c, &ctx());
766 assert!((score - 0.61).abs() < 1e-10, "got {score}");
767 }
768}