1use crate::knowledge::{BeliefStatus, EntityId, EntityTypeRef, StatementId};
11use oxibrain_index::{Direction, PredicateFilter};
12use oxibrain_ports::{TIME_MAX, TIME_MIN, Timestamp};
13use serde::{Deserialize, Serialize};
14use std::collections::{BTreeMap, HashMap};
15
16#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
21#[serde(rename_all = "snake_case")]
22pub enum TargetKind {
23 Statement,
24 Entity,
25 Episode,
26 Chunk,
27 Community,
28}
29
30#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
35#[serde(rename_all = "snake_case")]
36pub enum TargetSet {
37 #[default]
38 Statement,
39 Entity,
40 Episode,
41 Chunk,
42 Community,
43}
44
45#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
47#[serde(rename_all = "snake_case")]
48pub enum LexIndex {
49 Word,
51 Ngram,
54}
55
56#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
59#[serde(rename_all = "snake_case")]
60pub enum VecSpace {
61 Entity,
62 Statement,
63 Chunk,
64}
65
66#[derive(Debug, Clone, Serialize, Deserialize)]
68#[serde(tag = "kind", rename_all = "snake_case")]
69pub enum SeedPolicy {
70 Explicit { entities: Vec<EntityId> },
72 FromHits { top_k: usize },
74}
75
76#[derive(Debug, Clone, Serialize, Deserialize)]
79#[serde(tag = "kind", rename_all = "snake_case")]
80pub enum Channel {
81 Lexical { index: LexIndex },
82 Vector { space: VecSpace },
83 GraphExpand { seed: SeedPolicy, depth: u8 },
84 CommunityExpand { seed: SeedPolicy },
85}
86
87#[derive(Debug, Clone, Serialize, Deserialize)]
90#[serde(tag = "kind", rename_all = "snake_case")]
91pub enum Fusion {
92 Rrf { k: u32 },
94 Weighted { weights: Vec<f64> },
96}
97
98impl Default for Fusion {
99 fn default() -> Self {
100 Fusion::Rrf { k: 60 }
101 }
102}
103
104#[derive(Debug, Clone, Default, Serialize, Deserialize)]
106#[serde(tag = "kind", rename_all = "snake_case")]
107pub enum Rerank {
108 #[default]
109 None,
110 GraphDistance { from: Vec<EntityId> },
112 Corroboration,
115 Mmr {
122 lambda: f32,
123 max_similarity: Option<f32>,
124 },
125 Chain(Vec<Rerank>),
127}
128
129#[derive(Debug, Clone, Default, Serialize, Deserialize)]
132#[serde(tag = "kind", rename_all = "snake_case")]
133pub enum TrustPolicy {
134 #[default]
135 All,
136 Exclude(Vec<crate::TrustTier>),
138}
139
140#[derive(Debug, Clone, Serialize, Deserialize)]
144#[serde(rename_all = "snake_case")]
145pub struct Filters {
146 pub space: String,
147 pub as_of: Option<Timestamp>,
149 pub known_at: Option<Timestamp>,
151 pub min_confidence: f32,
152 pub trust: TrustPolicy,
153 pub predicates: PredicateFilter,
154 pub entity_types: Option<Vec<EntityTypeRef>>,
155}
156
157impl Filters {
158 pub fn open(space: impl Into<String>) -> Self {
161 Self {
162 space: space.into(),
163 as_of: None,
164 known_at: None,
165 min_confidence: 0.0,
166 trust: TrustPolicy::All,
167 predicates: PredicateFilter::AllowAll,
168 entity_types: None,
169 }
170 }
171}
172
173#[derive(Debug, Clone, Serialize, Deserialize)]
175#[serde(rename_all = "snake_case")]
176pub struct Retrieval {
177 pub targets: TargetSet,
178 pub channels: Vec<Channel>,
179 pub fusion: Fusion,
180 pub rerank: Rerank,
181 pub filters: Filters,
182 pub limit: usize,
183 pub explain: bool,
184}
185
186impl Retrieval {
187 pub fn hybrid(space: impl Into<String>) -> Self {
190 Self {
191 targets: TargetSet::Statement,
192 channels: vec![
193 Channel::Lexical {
194 index: LexIndex::Word,
195 },
196 Channel::Lexical {
197 index: LexIndex::Ngram,
198 },
199 Channel::Vector {
200 space: VecSpace::Entity,
201 },
202 Channel::GraphExpand {
203 seed: SeedPolicy::FromHits { top_k: 5 },
204 depth: 1,
205 },
206 ],
207 fusion: Fusion::Rrf { k: 60 },
208 rerank: Rerank::Corroboration,
209 filters: Filters::open(space),
210 limit: 20,
211 explain: false,
212 }
213 }
214
215 pub fn lexical(space: impl Into<String>) -> Self {
217 Self {
218 targets: TargetSet::Statement,
219 channels: vec![
220 Channel::Lexical {
221 index: LexIndex::Word,
222 },
223 Channel::Lexical {
224 index: LexIndex::Ngram,
225 },
226 ],
227 fusion: Fusion::Rrf { k: 60 },
228 rerank: Rerank::None,
229 filters: Filters::open(space),
230 limit: 20,
231 explain: false,
232 }
233 }
234
235 pub fn semantic(space: impl Into<String>) -> Self {
237 Self {
238 targets: TargetSet::Statement,
239 channels: vec![Channel::Vector {
240 space: VecSpace::Entity,
241 }],
242 fusion: Fusion::Rrf { k: 60 },
243 rerank: Rerank::Mmr {
244 lambda: 0.5,
245 max_similarity: Some(0.9),
246 },
247 filters: Filters::open(space),
248 limit: 20,
249 explain: false,
250 }
251 }
252
253 pub fn graph(space: impl Into<String>, seeds: Vec<EntityId>) -> Self {
255 Self {
256 targets: TargetSet::Statement,
257 channels: vec![Channel::GraphExpand {
258 seed: SeedPolicy::Explicit { entities: seeds },
259 depth: 2,
260 }],
261 fusion: Fusion::Rrf { k: 60 },
262 rerank: Rerank::GraphDistance { from: Vec::new() },
263 filters: Filters::open(space),
264 limit: 50,
265 explain: false,
266 }
267 }
268
269 pub fn community(space: impl Into<String>, seeds: Vec<EntityId>) -> Self {
271 Self {
272 targets: TargetSet::Statement,
273 channels: vec![Channel::CommunityExpand {
274 seed: SeedPolicy::Explicit { entities: seeds },
275 }],
276 fusion: Fusion::Rrf { k: 60 },
277 rerank: Rerank::None,
278 filters: Filters::open(space),
279 limit: 20,
280 explain: false,
281 }
282 }
283}
284
285#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
290pub struct ChannelRank {
291 pub channel: u8,
292 pub rank: u32,
293}
294
295#[derive(Debug, Clone, Serialize, Deserialize)]
299pub struct TargetFacts {
300 pub target: TargetId,
301 pub confidence: f32,
302 pub valid_from: Timestamp,
303 pub valid_to: Timestamp,
304 pub recorded_at: Timestamp,
305 pub retracted_at: Option<Timestamp>,
306 pub trust: crate::TrustTier,
307 pub status: BeliefStatus,
308 pub predicate: String,
309 pub salience: f64,
310 pub distinct_episodes: u32,
311 pub channels: Vec<ChannelRank>,
312 pub channel_scores: Vec<f64>,
314}
315
316#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
319#[serde(tag = "kind", rename_all = "snake_case")]
320pub enum TargetId {
321 Episode { id: String },
322 Statement { id: StatementId },
323 Entity { id: EntityId },
324 Chunk { id: String },
325 Community { id: String },
326}
327
328impl TargetId {
329 pub fn rrf_key(&self) -> String {
332 match self {
333 TargetId::Episode { id } => format!("episode:{id}"),
334 TargetId::Statement { id } => format!("statement:{id}"),
335 TargetId::Entity { id } => format!("entity:{id}"),
336 TargetId::Chunk { id } => format!("chunk:{id}"),
337 TargetId::Community { id } => format!("community:{id}"),
338 }
339 }
340}
341
342#[derive(Debug, Clone, Serialize, Deserialize)]
345pub struct ChannelResult {
346 pub channel: u8,
347 pub hits: Vec<(TargetId, f64)>,
348}
349
350#[derive(Debug, Clone, Default, Serialize, Deserialize)]
353pub struct RetrievalInput {
354 pub channels: Vec<ChannelResult>,
355 pub facts: HashMap<TargetId, TargetFacts>,
356 #[serde(default)]
360 pub entity_vectors: HashMap<String, Vec<f32>>,
361}
362
363#[derive(Debug, Clone, Serialize, Deserialize)]
366pub struct RankedItem {
367 pub target: TargetId,
368 pub fused_score: f64,
369 pub rank: usize,
370 pub channels: Vec<ChannelRank>,
371 pub salience: f64,
372 pub facts: TargetFacts,
374}
375
376#[derive(Debug, Clone, Serialize, Deserialize)]
379pub struct DroppedItem {
380 pub target: TargetId,
381 pub reason: DropReason,
382 pub score: Option<f64>,
384}
385
386#[derive(Debug, Clone, Serialize, Deserialize)]
389#[serde(tag = "kind", rename_all = "snake_case")]
390pub enum DropReason {
391 BelowConfidenceFloor {
392 actual: f32,
393 floor: f32,
394 },
395 OutsideValidWindow {
396 valid_at: Timestamp,
397 },
398 BeforeKnownAt {
399 known_at: Timestamp,
400 recorded_at: Timestamp,
401 },
402 TrustExcluded {
403 tier: crate::TrustTier,
404 },
405 PredicateDenied {
406 predicate: String,
407 },
408 EntityTypeMismatch {
409 expected: Vec<EntityTypeRef>,
410 },
411 TruncatedByBudget {
412 position: usize,
413 },
414}
415
416#[derive(Debug, Clone, Serialize, Deserialize)]
418pub struct RankingResult {
419 pub items: Vec<RankedItem>,
420 pub dropped: Vec<DroppedItem>,
421 pub total_candidates: usize,
422 pub spec: Retrieval,
424}
425
426pub fn rank(input: &RetrievalInput, spec: &Retrieval) -> RankingResult {
440 let mut candidate_facts: HashMap<String, TargetFacts> = HashMap::new();
443 let mut channel_ranks: HashMap<String, Vec<ChannelRank>> = HashMap::new();
444 let mut channel_scores: HashMap<String, Vec<f64>> = HashMap::new();
445
446 for cr in &input.channels {
447 for (rank_idx, (target, score)) in cr.hits.iter().enumerate() {
448 let key = target.rrf_key();
449 let entry = candidate_facts.entry(key.clone());
455 match entry {
456 std::collections::hash_map::Entry::Vacant(v) => {
457 if let Some(facts) = input.facts.get(target) {
458 v.insert(facts.clone());
459 } else {
460 v.insert(minimal_facts(target));
463 }
464 }
465 std::collections::hash_map::Entry::Occupied(mut o) => {
466 if let Some(newer) = input.facts.get(target) {
467 if newer.confidence > o.get().confidence {
468 o.insert(newer.clone());
469 }
470 }
471 }
472 }
473 channel_ranks
474 .entry(key.clone())
475 .or_default()
476 .push(ChannelRank {
477 channel: cr.channel,
478 rank: rank_idx as u32,
479 });
480 channel_scores.entry(key).or_default().push(*score);
481 }
482 }
483
484 for (target, facts) in &input.facts {
488 let key = target.rrf_key();
489 candidate_facts.entry(key).or_insert_with(|| facts.clone());
490 }
491
492 let total_candidates = candidate_facts.len();
493
494 let mut items: Vec<RankedItem> = Vec::with_capacity(candidate_facts.len());
497 let mut dropped: Vec<DroppedItem> = Vec::new();
498 let mut ordered_keys: Vec<&String> = candidate_facts.keys().collect();
500 ordered_keys.sort();
501
502 for key in &ordered_keys {
503 let facts = &candidate_facts[*key];
504 if let Some(reason) = check_filters(facts, &spec.filters) {
505 dropped.push(DroppedItem {
506 target: facts.target.clone(),
507 reason,
508 score: None,
509 });
510 continue;
511 }
512 let fused = fuse(channel_scores.get(*key), &spec.fusion);
513 let channels = channel_ranks.get(*key).cloned().unwrap_or_default();
514 items.push(RankedItem {
515 target: facts.target.clone(),
516 fused_score: fused,
517 rank: 0, channels,
519 salience: facts.salience,
520 facts: facts.clone(),
521 });
522 }
523
524 apply_rerank(&mut items, &spec.rerank, &input.entity_vectors);
526
527 items.sort_by(|a, b| {
531 b.fused_score
532 .partial_cmp(&a.fused_score)
533 .unwrap_or(std::cmp::Ordering::Equal)
534 .then_with(|| target_type_rank(&a.target).cmp(&target_type_rank(&b.target)))
535 .then_with(|| a.target.rrf_key().cmp(&b.target.rrf_key()))
536 });
537
538 items.truncate(spec.limit);
541 for (i, item) in items.iter_mut().enumerate() {
542 item.rank = i;
543 }
544 let kept_keys: std::collections::HashSet<String> =
545 items.iter().map(|i| i.target.rrf_key()).collect();
546 let already_dropped: std::collections::HashSet<String> =
547 dropped.iter().map(|d| d.target.rrf_key()).collect();
548 let mut post_truncate_keys: Vec<&String> = candidate_facts
549 .keys()
550 .filter(|k| !kept_keys.contains(*k) && !already_dropped.contains(*k))
551 .collect();
552 post_truncate_keys.sort();
553 let mut post_truncate_drops: Vec<DroppedItem> = post_truncate_keys
554 .into_iter()
555 .map(|k| {
556 let facts = &candidate_facts[k];
557 DroppedItem {
558 target: facts.target.clone(),
559 reason: DropReason::TruncatedByBudget {
560 position: ordered_keys.iter().position(|x| *x == k).unwrap_or(0),
561 },
562 score: None,
563 }
564 })
565 .collect();
566
567 dropped.append(&mut post_truncate_drops);
571
572 RankingResult {
573 items,
574 dropped,
575 total_candidates,
576 spec: spec.clone(),
577 }
578}
579
580fn minimal_facts(target: &TargetId) -> TargetFacts {
584 TargetFacts {
585 target: target.clone(),
586 confidence: 0.0,
587 valid_from: TIME_MIN,
588 valid_to: TIME_MAX,
589 recorded_at: TIME_MIN,
590 retracted_at: None,
591 trust: crate::TrustTier::SemiTrusted,
592 status: BeliefStatus::Active,
593 predicate: String::new(),
594 salience: 0.0,
595 distinct_episodes: 0,
596 channels: Vec::new(),
597 channel_scores: Vec::new(),
598 }
599}
600
601fn target_type_rank(t: &TargetId) -> u8 {
604 match t {
605 TargetId::Statement { .. } => 0,
606 TargetId::Episode { .. } => 1,
607 TargetId::Entity { .. } => 2,
608 TargetId::Chunk { .. } => 3,
609 TargetId::Community { .. } => 4,
610 }
611}
612
613fn check_filters(facts: &TargetFacts, filters: &Filters) -> Option<DropReason> {
616 if let Some(t) = filters.as_of {
618 if t < facts.valid_from || t > facts.valid_to {
619 return Some(DropReason::OutsideValidWindow { valid_at: t });
620 }
621 }
622 if let Some(t) = filters.known_at {
625 if facts.recorded_at > t {
626 return Some(DropReason::BeforeKnownAt {
627 known_at: t,
628 recorded_at: facts.recorded_at,
629 });
630 }
631 if let Some(retracted_at) = facts.retracted_at {
632 if retracted_at <= t {
633 return Some(DropReason::BeforeKnownAt {
634 known_at: t,
635 recorded_at: retracted_at,
636 });
637 }
638 }
639 }
640 if facts.confidence < filters.min_confidence {
643 return Some(DropReason::BelowConfidenceFloor {
644 actual: facts.confidence,
645 floor: filters.min_confidence,
646 });
647 }
648 if let TrustPolicy::Exclude(excluded) = &filters.trust {
650 if excluded.contains(&facts.trust) {
651 return Some(DropReason::TrustExcluded { tier: facts.trust });
652 }
653 }
654 if !filters.predicates.allows(&facts.predicate) {
656 return Some(DropReason::PredicateDenied {
657 predicate: facts.predicate.clone(),
658 });
659 }
660 if let Some(expected) = &filters.entity_types {
664 if !expected.is_empty() && !expected.iter().any(|t| t == &facts.predicate) {
665 return Some(DropReason::EntityTypeMismatch {
666 expected: expected.clone(),
667 });
668 }
669 }
670 None
671}
672
673fn fuse(scores: Option<&Vec<f64>>, fusion: &Fusion) -> f64 {
675 let Some(scores) = scores else { return 0.0 };
676 if scores.is_empty() {
677 return 0.0;
678 }
679 match fusion {
680 Fusion::Rrf { k } => {
681 let mut indexed: Vec<(usize, f64)> = scores.iter().copied().enumerate().collect();
690 indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
691 let k = *k as f64;
692 indexed
693 .iter()
694 .enumerate()
695 .map(|(pos, (_, _))| 1.0 / (k + (pos as f64) + 1.0))
696 .sum()
697 }
698 Fusion::Weighted { weights } => {
699 let min = scores.iter().cloned().fold(f64::INFINITY, f64::min);
703 let max = scores.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
704 let span = (max - min).max(f64::EPSILON);
705 let n = scores.len().max(weights.len());
706 let w_per = if weights.is_empty() {
707 1.0 / n as f64
708 } else {
709 let s: f64 = weights.iter().sum();
710 if s > 0.0 { 1.0 / s } else { 1.0 / n as f64 }
711 };
712 scores
713 .iter()
714 .map(|s| (s - min) / span)
715 .zip(weights.iter().chain(std::iter::repeat(&w_per)))
716 .map(|(n, w)| n * w)
717 .sum()
718 }
719 }
720}
721
722fn cosine_sim(a: &[f32], b: &[f32]) -> f64 {
725 if a.len() != b.len() || a.is_empty() {
726 return 0.0;
727 }
728 let dot: f64 = a
729 .iter()
730 .zip(b.iter())
731 .map(|(x, y)| (*x as f64) * (*y as f64))
732 .sum();
733 let na: f64 = a.iter().map(|x| (*x as f64).powi(2)).sum::<f64>().sqrt();
734 let nb: f64 = b.iter().map(|x| (*x as f64).powi(2)).sum::<f64>().sqrt();
735 if na == 0.0 || nb == 0.0 {
736 return 0.0;
737 }
738 dot / (na * nb)
739}
740
741fn item_vector<'a>(
745 item: &RankedItem,
746 vectors: &'a HashMap<String, Vec<f32>>,
747) -> Option<&'a Vec<f32>> {
748 match &item.target {
749 TargetId::Entity { id } => vectors.get(id),
750 _ => None,
751 }
752}
753
754fn apply_rerank(
758 items: &mut [RankedItem],
759 rerank: &Rerank,
760 entity_vectors: &HashMap<String, Vec<f32>>,
761) {
762 match rerank {
763 Rerank::None => {}
764 Rerank::Corroboration => {
765 for item in items.iter_mut() {
769 let boost = 1.0 + (1.0 + item.facts.distinct_episodes as f64).ln();
770 item.fused_score *= boost;
771 }
772 }
773 Rerank::GraphDistance { from } => {
774 let _ = from;
780 items.sort_by(|a, b| {
781 b.salience
782 .partial_cmp(&a.salience)
783 .unwrap_or(std::cmp::Ordering::Equal)
784 });
785 }
786 Rerank::Mmr {
787 lambda,
788 max_similarity,
789 } => {
790 if items.is_empty() {
804 return;
805 }
806 let lambda = *lambda as f64;
807 let ceiling = max_similarity.map(|c| c as f64);
808 let mut reordered: Vec<RankedItem> = Vec::with_capacity(items.len());
809 let mut pool: Vec<RankedItem> = items.to_vec();
810 pool.sort_by(|a, b| {
812 b.fused_score
813 .partial_cmp(&a.fused_score)
814 .unwrap_or(std::cmp::Ordering::Equal)
815 });
816 reordered.push(pool.remove(0));
817 while !pool.is_empty() {
818 let mut best_idx: Option<usize> = None;
819 let mut best_score = f64::NEG_INFINITY;
820 for (i, cand) in pool.iter().enumerate() {
821 let cand_vec = item_vector(cand, entity_vectors);
822 let mut max_sim: f64 = 0.0;
824 let mut used_cosine = false;
825 for sel in &reordered {
826 let sel_vec = item_vector(sel, entity_vectors);
827 if let (Some(cv), Some(sv)) = (cand_vec, sel_vec) {
828 let cos = cosine_sim(cv, sv);
829 if cos > max_sim {
830 max_sim = cos;
831 }
832 used_cosine = true;
833 }
834 }
835 let sim = if used_cosine {
837 max_sim
838 } else {
839 let last = reordered.last().unwrap().fused_score;
840 (last - cand.fused_score).abs()
841 };
842 if let Some(c) = ceiling {
844 if used_cosine && max_sim > c {
845 continue;
846 }
847 }
848 let mmr = lambda * cand.fused_score - (1.0 - lambda) * sim;
849 if mmr > best_score {
850 best_score = mmr;
851 best_idx = Some(i);
852 }
853 }
854 match best_idx {
855 Some(i) => reordered.push(pool.remove(i)),
856 None => {
859 pool.sort_by(|a, b| {
860 b.fused_score
861 .partial_cmp(&a.fused_score)
862 .unwrap_or(std::cmp::Ordering::Equal)
863 });
864 reordered.append(&mut pool);
865 break;
866 }
867 }
868 }
869 items.clone_from_slice(&reordered);
870 }
871 Rerank::Chain(reranks) => {
872 for r in reranks {
873 apply_rerank(items, r, entity_vectors);
874 }
875 }
876 }
877}
878
879pub fn direction_allows(direction: Direction, from_subject: bool) -> bool {
882 match direction {
883 Direction::Both => true,
884 Direction::Out => from_subject,
885 Direction::In => !from_subject,
886 }
887}
888
889pub fn explain_item(item: &RankedItem) -> BTreeMap<String, String> {
891 let mut out = BTreeMap::new();
892 out.insert("target".into(), item.target.rrf_key());
893 out.insert("fused_score".into(), format!("{:.6}", item.fused_score));
894 out.insert("salience".into(), format!("{:.6}", item.salience));
895 out.insert(
896 "channels".into(),
897 item.channels
898 .iter()
899 .map(|c| format!("ch{}:{}", c.channel, c.rank))
900 .collect::<Vec<_>>()
901 .join(","),
902 );
903 out
904}
905
906#[cfg(test)]
907mod tests {
908 use super::*;
909 use crate::TrustTier;
910
911 fn facts(target: TargetId, confidence: f32, predicate: &str) -> TargetFacts {
912 TargetFacts {
913 target: target.clone(),
914 confidence,
915 valid_from: TIME_MIN,
916 valid_to: TIME_MAX,
917 recorded_at: Timestamp(1000),
918 retracted_at: None,
919 trust: TrustTier::Trusted,
920 status: BeliefStatus::Active,
921 predicate: predicate.into(),
922 salience: 0.5,
923 distinct_episodes: 1,
924 channels: vec![],
925 channel_scores: vec![],
926 }
927 }
928
929 #[test]
930 fn rrf_key_is_stable_and_distinguishes_kinds() {
931 let e = TargetId::Episode { id: "x".into() };
932 let s = TargetId::Statement { id: "x".into() };
933 assert_ne!(e.rrf_key(), s.rrf_key());
934 assert_eq!(e.rrf_key(), e.rrf_key());
935 }
936
937 #[test]
938 fn filter_drops_below_floor() {
939 let target = TargetId::Statement { id: "s1".into() };
940 let f = facts(target.clone(), 0.1, "works_on");
941 let mut filters = Filters::open("default");
942 filters.min_confidence = 0.5;
943 assert!(matches!(
944 check_filters(&f, &filters),
945 Some(DropReason::BelowConfidenceFloor { .. })
946 ));
947 }
948
949 #[test]
950 fn filter_keeps_at_or_above_floor() {
951 let target = TargetId::Statement { id: "s1".into() };
952 let f = facts(target, 0.5, "works_on");
953 let filters = Filters::open("default");
954 assert!(check_filters(&f, &filters).is_none());
955 }
956
957 #[test]
958 fn filter_drops_outside_as_of() {
959 let target = TargetId::Statement { id: "s1".into() };
960 let mut f = facts(target, 0.9, "works_on");
961 f.valid_from = Timestamp(100);
962 f.valid_to = Timestamp(200);
963 let mut filters = Filters::open("default");
964 filters.as_of = Some(Timestamp(300));
965 assert!(matches!(
966 check_filters(&f, &filters),
967 Some(DropReason::OutsideValidWindow { .. })
968 ));
969 }
970
971 #[test]
972 fn filter_drops_before_known_at() {
973 let target = TargetId::Statement { id: "s1".into() };
974 let mut f = facts(target, 0.9, "works_on");
975 f.recorded_at = Timestamp(1000);
976 let mut filters = Filters::open("default");
977 filters.known_at = Some(Timestamp(500));
978 assert!(matches!(
979 check_filters(&f, &filters),
980 Some(DropReason::BeforeKnownAt { .. })
981 ));
982 }
983
984 #[test]
985 fn filter_drops_retracted_before_known_at() {
986 let target = TargetId::Statement { id: "s1".into() };
987 let mut f = facts(target, 0.9, "works_on");
988 f.recorded_at = Timestamp(100);
989 f.retracted_at = Some(Timestamp(200));
990 let mut filters = Filters::open("default");
991 filters.known_at = Some(Timestamp(300));
992 assert!(matches!(
993 check_filters(&f, &filters),
994 Some(DropReason::BeforeKnownAt { .. })
995 ));
996 }
997
998 #[test]
999 fn rank_conservation_simple() {
1000 let mut input = RetrievalInput::default();
1002 let a = TargetId::Statement { id: "a".into() };
1003 let b = TargetId::Statement { id: "b".into() };
1004 input.facts.insert(a.clone(), facts(a.clone(), 0.9, "p"));
1005 input.facts.insert(b.clone(), facts(b.clone(), 0.8, "p"));
1006 input.channels.push(ChannelResult {
1007 channel: 0,
1008 hits: vec![(a.clone(), 0.9), (b.clone(), 0.8)],
1009 });
1010 let spec = Retrieval::hybrid("default");
1011 let r = rank(&input, &spec);
1012 assert_eq!(r.items.len() + r.dropped.len(), 2);
1014 }
1015
1016 #[test]
1017 fn rank_filter_totality() {
1018 let mut input = RetrievalInput::default();
1020 let a = TargetId::Statement { id: "a".into() };
1021 input.facts.insert(a.clone(), facts(a.clone(), 0.1, "p"));
1022 input.channels.push(ChannelResult {
1023 channel: 0,
1024 hits: vec![(a.clone(), 0.5)],
1025 });
1026 let mut spec = Retrieval::hybrid("default");
1027 spec.filters.min_confidence = 0.5;
1028 let r = rank(&input, &spec);
1029 assert!(r.items.is_empty());
1030 assert_eq!(r.dropped.len(), 1);
1031 assert!(matches!(
1032 r.dropped[0].reason,
1033 DropReason::BelowConfidenceFloor { .. }
1034 ));
1035 }
1036
1037 #[test]
1038 fn rank_determinism() {
1039 let make = || {
1040 let mut input = RetrievalInput::default();
1041 for (i, score) in [(0u8, 0.9f64), (1, 0.7), (2, 0.5)] {
1042 let id = format!("s{i}");
1043 let t = TargetId::Statement { id: id.clone() };
1044 input.facts.insert(t.clone(), facts(t, score as f32, "p"));
1045 input.channels.push(ChannelResult {
1046 channel: 0,
1047 hits: vec![(TargetId::Statement { id }, score)],
1048 });
1049 }
1050 input
1051 };
1052 let spec = Retrieval::hybrid("default");
1053 let r1 = rank(&make(), &spec);
1054 let r2 = rank(&make(), &spec);
1055 let keys1: Vec<_> = r1
1057 .items
1058 .iter()
1059 .map(|i| (i.target.rrf_key(), i.fused_score, i.rank))
1060 .collect();
1061 let keys2: Vec<_> = r2
1062 .items
1063 .iter()
1064 .map(|i| (i.target.rrf_key(), i.fused_score, i.rank))
1065 .collect();
1066 assert_eq!(keys1, keys2);
1067 }
1068
1069 #[test]
1072 fn corroboration_invariance_equal_episodes_preserves_order() {
1073 let mut items: Vec<RankedItem> = [0.9f64, 0.7, 0.5]
1076 .iter()
1077 .enumerate()
1078 .map(|(i, score)| RankedItem {
1079 target: TargetId::Statement {
1080 id: format!("s{i}"),
1081 },
1082 fused_score: *score,
1083 rank: i,
1084 channels: vec![],
1085 salience: 0.5,
1086 facts: {
1087 let mut f = facts(
1088 TargetId::Statement {
1089 id: format!("s{i}"),
1090 },
1091 0.8,
1092 "works_on",
1093 );
1094 f.distinct_episodes = 3; f
1096 },
1097 })
1098 .collect();
1099 apply_rerank(&mut items, &Rerank::Corroboration, &HashMap::new());
1100 let boosted: Vec<f64> = items.iter().map(|i| i.fused_score).collect();
1103 for w in boosted.windows(2) {
1104 assert!(
1105 w[0] >= w[1],
1106 "ordering broken after equal-boost corroboration"
1107 );
1108 }
1109 }
1110
1111 #[test]
1112 fn corroboration_monotonicity_higher_distinct_ranks_higher() {
1113 let mut items: Vec<RankedItem> = vec![
1116 RankedItem {
1117 target: TargetId::Statement { id: "low".into() },
1118 fused_score: 0.5,
1119 rank: 0,
1120 channels: vec![],
1121 salience: 0.5,
1122 facts: {
1123 let mut f = facts(TargetId::Statement { id: "low".into() }, 0.8, "works_on");
1124 f.distinct_episodes = 1;
1125 f
1126 },
1127 },
1128 RankedItem {
1129 target: TargetId::Statement { id: "high".into() },
1130 fused_score: 0.5,
1131 rank: 1,
1132 channels: vec![],
1133 salience: 0.5,
1134 facts: {
1135 let mut f = facts(TargetId::Statement { id: "high".into() }, 0.8, "works_on");
1136 f.distinct_episodes = 10;
1137 f
1138 },
1139 },
1140 ];
1141 apply_rerank(&mut items, &Rerank::Corroboration, &HashMap::new());
1142 let high_score = items
1143 .iter()
1144 .find(|i| i.target == TargetId::Statement { id: "high".into() })
1145 .unwrap()
1146 .fused_score;
1147 let low_score = items
1148 .iter()
1149 .find(|i| i.target == TargetId::Statement { id: "low".into() })
1150 .unwrap()
1151 .fused_score;
1152 assert!(
1153 high_score > low_score,
1154 "higher corroboration must rank higher: {high_score} vs {low_score}"
1155 );
1156 }
1157
1158 use proptest::prelude::*;
1164
1165 prop_compose! {
1166 fn arb_target_id()(i in any::<u8>()) -> TargetId {
1167 let id = format!("t{i:03}");
1168 match i % 3 {
1169 0 => TargetId::Statement { id },
1170 1 => TargetId::Entity { id },
1171 _ => TargetId::Episode { id },
1172 }
1173 }
1174 }
1175
1176 prop_compose! {
1177 fn arb_facts()(
1178 target in arb_target_id(),
1179 confidence in 0.0f32..1.0f32,
1180 vf in 0i64..10_000,
1181 vt in 0i64..10_000,
1182 recorded_at in 0i64..10_000,
1183 retracted_at in prop::option::of(0i64..10_000),
1184 trust in prop::sample::select(vec![
1185 TrustTier::Trusted, TrustTier::SemiTrusted, TrustTier::Untrusted,
1186 ]),
1187 status in prop::sample::select(vec![
1188 BeliefStatus::Active,
1189 BeliefStatus::Superseded,
1190 BeliefStatus::Contradicted,
1191 BeliefStatus::Retracted,
1192 ]),
1193 predicate in prop::sample::select(vec![
1194 "works_on".to_string(), "knows".to_string(), "likes".to_string(),
1195 ]),
1196 salience in 0.0f64..1.0f64,
1197 distinct in 0u32..10,
1198 ) -> TargetFacts {
1199 TargetFacts {
1200 target,
1201 confidence,
1202 valid_from: Timestamp(vf),
1203 valid_to: Timestamp(vt),
1204 recorded_at: Timestamp(recorded_at),
1205 retracted_at: retracted_at.map(Timestamp),
1206 trust,
1207 status,
1208 predicate,
1209 salience,
1210 distinct_episodes: distinct,
1211 channels: vec![],
1212 channel_scores: vec![],
1213 }
1214 }
1215 }
1216
1217 prop_compose! {
1218 fn arb_filters()(
1219 as_of in prop::option::of(0i64..10_000),
1220 known_at in prop::option::of(0i64..10_000),
1221 min_confidence in 0.0f32..1.0f32,
1222 ) -> Filters {
1223 Filters {
1224 space: "default".into(),
1225 as_of: as_of.map(Timestamp),
1226 known_at: known_at.map(Timestamp),
1227 min_confidence,
1228 trust: TrustPolicy::All,
1229 predicates: PredicateFilter::AllowAll,
1230 entity_types: None,
1231 }
1232 }
1233 }
1234
1235 prop_compose! {
1236 fn arb_input()(entries in prop::collection::vec((arb_facts(), 0.0f64..1.0f64), 1..16)) -> RetrievalInput {
1237 let mut input = RetrievalInput::default();
1238 let mut hits: Vec<(TargetId, f64)> = Vec::new();
1239 for (f, score) in entries {
1240 let target = f.target.clone();
1241 input.facts.insert(target.clone(), f);
1242 hits.push((target, score));
1243 }
1244 input.channels.push(ChannelResult { channel: 0, hits });
1245 input
1246 }
1247 }
1248
1249 proptest! {
1250 #![proptest_config(ProptestConfig::with_cases(64))]
1251
1252 #[test]
1254 fn prop_conservation(input in arb_input(), filters in arb_filters()) {
1255 let mut spec = Retrieval::hybrid("default");
1256 spec.filters = filters;
1257 let total_candidates = input.facts.len();
1258 let r = rank(&input, &spec);
1259 prop_assert_eq!(r.items.len() + r.dropped.len(), total_candidates);
1260 let item_keys: std::collections::HashSet<_> =
1261 r.items.iter().map(|i| i.target.rrf_key()).collect();
1262 let drop_keys: std::collections::HashSet<_> =
1263 r.dropped.iter().map(|d| d.target.rrf_key()).collect();
1264 prop_assert!(item_keys.is_disjoint(&drop_keys));
1265 prop_assert_eq!(r.total_candidates, total_candidates);
1266 }
1267
1268 #[test]
1270 fn prop_filter_totality(input in arb_input(), filters in arb_filters()) {
1271 let mut spec = Retrieval::hybrid("default");
1272 spec.filters = filters.clone();
1273 let r = rank(&input, &spec);
1274 for item in &r.items {
1275 let f = &item.facts;
1276 prop_assert!(f.confidence >= filters.min_confidence,
1277 "item {} violates min_confidence: {} < {}",
1278 item.target.rrf_key(), f.confidence, filters.min_confidence);
1279 if let Some(t) = filters.as_of {
1280 prop_assert!(t >= f.valid_from && t <= f.valid_to,
1281 "item {} violates as_of {}", item.target.rrf_key(), t.0);
1282 }
1283 if let Some(t) = filters.known_at {
1284 prop_assert!(f.recorded_at <= t,
1285 "item {} violates known_at {}", item.target.rrf_key(), t.0);
1286 if let Some(r_at) = f.retracted_at {
1287 prop_assert!(r_at > t,
1288 "item {} violates known_at (retracted)", item.target.rrf_key());
1289 }
1290 }
1291 if let TrustPolicy::Exclude(excluded) = &filters.trust {
1292 prop_assert!(!excluded.contains(&f.trust),
1293 "item {} has excluded trust tier", item.target.rrf_key());
1294 }
1295 }
1296 }
1297
1298 #[test]
1300 fn prop_determinism(input in arb_input(), filters in arb_filters()) {
1301 let mut spec = Retrieval::hybrid("default");
1302 spec.filters = filters;
1303 let r1 = rank(&input, &spec);
1304 let r2 = rank(&input, &spec);
1305 let j = |r: &RankingResult| serde_json::to_string(r).expect("serialize");
1306 prop_assert_eq!(j(&r1), j(&r2));
1307 }
1308 }
1309
1310 #[test]
1322 fn mmr_ceiling_defers_near_duplicates_in_top_10() {
1323 use std::collections::HashMap;
1324 let mut vectors: HashMap<String, Vec<f32>> = HashMap::new();
1325 let mut names: Vec<String> = Vec::new();
1326
1327 for i in 0..10 {
1329 let name = format!("d{i}");
1330 names.push(name.clone());
1331 let angle = (i as f64) * std::f64::consts::PI / 5.0;
1332 vectors.insert(name, vec![angle.cos() as f32, angle.sin() as f32]);
1333 }
1334 for i in 0..5 {
1336 let name = format!("dup{i}");
1337 names.push(name.clone());
1338 let angle = (i as f64) * std::f64::consts::PI / 5.0;
1339 vectors.insert(
1341 name,
1342 vec![(angle.cos() * 1.01) as f32, (angle.sin() * 1.01) as f32],
1343 );
1344 }
1345
1346 let mut items: Vec<RankedItem> = names
1348 .iter()
1349 .enumerate()
1350 .map(|(i, e)| {
1351 let score = if e.starts_with("dup") {
1352 1.5 - 0.01 * i as f64 } else {
1354 1.0 - 0.01 * i as f64 };
1356 RankedItem {
1357 target: TargetId::Entity { id: e.clone() },
1358 facts: facts(TargetId::Entity { id: e.clone() }, 0.8, "works_on"),
1359 fused_score: score,
1360 salience: 0.8,
1361 rank: i,
1362 channels: vec![],
1363 }
1364 })
1365 .collect();
1366
1367 apply_rerank(
1368 &mut items,
1369 &Rerank::Mmr {
1370 lambda: 0.5,
1371 max_similarity: Some(0.9),
1372 },
1373 &vectors,
1374 );
1375
1376 for (i, a) in items.iter().take(10).enumerate() {
1378 let a_vec = match &a.target {
1379 TargetId::Entity { id } => vectors.get(id).unwrap(),
1380 _ => unreachable!(),
1381 };
1382 for (j, b) in items.iter().take(10).enumerate() {
1383 if i >= j {
1384 continue;
1385 }
1386 let b_vec = match &b.target {
1387 TargetId::Entity { id } => vectors.get(id).unwrap(),
1388 _ => unreachable!(),
1389 };
1390 let sim = cosine_sim(a_vec, b_vec);
1391 assert!(
1392 sim <= 0.9,
1393 "MMR kept a >0.9 pair at top-10 positions {i}/{j} \
1394 ({:?} / {:?}): cosine = {sim:.4}",
1395 a.target,
1396 b.target,
1397 );
1398 }
1399 }
1400
1401 assert_eq!(items.len(), 15);
1403 }
1404
1405 #[test]
1410 fn mmr_ceiling_never_drops_items() {
1411 use std::collections::HashMap;
1412 let mut vectors: HashMap<String, Vec<f32>> = HashMap::new();
1413 let names: Vec<String> = (0..8).map(|i| format!("e{i}")).collect();
1414 for name in &names {
1417 vectors.insert(name.clone(), vec![1.0, 0.01, 0.0]);
1418 }
1419
1420 let mut items: Vec<RankedItem> = names
1421 .iter()
1422 .enumerate()
1423 .map(|(i, e)| RankedItem {
1424 target: TargetId::Entity { id: e.clone() },
1425 facts: facts(TargetId::Entity { id: e.clone() }, 0.8, "works_on"),
1426 fused_score: 1.0 - 0.01 * i as f64,
1427 salience: 0.8,
1428 rank: i,
1429 channels: vec![],
1430 })
1431 .collect();
1432
1433 let count_before = items.len();
1434 apply_rerank(
1435 &mut items,
1436 &Rerank::Mmr {
1437 lambda: 0.5,
1438 max_similarity: Some(0.9),
1439 },
1440 &vectors,
1441 );
1442 assert_eq!(
1443 items.len(),
1444 count_before,
1445 "MMR ceiling must not drop items — conservation invariant"
1446 );
1447 }
1448}