use crate::knowledge::{BeliefStatus, EntityId, EntityTypeRef, StatementId};
use oxibrain_index::{Direction, PredicateFilter};
use oxibrain_ports::{TIME_MAX, TIME_MIN, Timestamp};
use serde::{Deserialize, Serialize};
use std::collections::{BTreeMap, HashMap};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum TargetKind {
Statement,
Entity,
Episode,
Chunk,
Community,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum TargetSet {
#[default]
Statement,
Entity,
Episode,
Chunk,
Community,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum LexIndex {
Word,
Ngram,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum VecSpace {
Entity,
Statement,
Chunk,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum SeedPolicy {
Explicit { entities: Vec<EntityId> },
FromHits { top_k: usize },
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum Channel {
Lexical { index: LexIndex },
Vector { space: VecSpace },
GraphExpand { seed: SeedPolicy, depth: u8 },
CommunityExpand { seed: SeedPolicy },
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum Fusion {
Rrf { k: u32 },
Weighted { weights: Vec<f64> },
}
impl Default for Fusion {
fn default() -> Self {
Fusion::Rrf { k: 60 }
}
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum Rerank {
#[default]
None,
GraphDistance { from: Vec<EntityId> },
Corroboration,
Mmr {
lambda: f32,
max_similarity: Option<f32>,
},
Chain(Vec<Rerank>),
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum TrustPolicy {
#[default]
All,
Exclude(Vec<crate::TrustTier>),
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub struct Filters {
pub space: String,
pub as_of: Option<Timestamp>,
pub known_at: Option<Timestamp>,
pub min_confidence: f32,
pub trust: TrustPolicy,
pub predicates: PredicateFilter,
pub entity_types: Option<Vec<EntityTypeRef>>,
}
impl Filters {
pub fn open(space: impl Into<String>) -> Self {
Self {
space: space.into(),
as_of: None,
known_at: None,
min_confidence: 0.0,
trust: TrustPolicy::All,
predicates: PredicateFilter::AllowAll,
entity_types: None,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub struct Retrieval {
pub targets: TargetSet,
pub channels: Vec<Channel>,
pub fusion: Fusion,
pub rerank: Rerank,
pub filters: Filters,
pub limit: usize,
pub explain: bool,
}
impl Retrieval {
pub fn hybrid(space: impl Into<String>) -> Self {
Self {
targets: TargetSet::Statement,
channels: vec![
Channel::Lexical {
index: LexIndex::Word,
},
Channel::Lexical {
index: LexIndex::Ngram,
},
Channel::Vector {
space: VecSpace::Entity,
},
Channel::GraphExpand {
seed: SeedPolicy::FromHits { top_k: 5 },
depth: 1,
},
],
fusion: Fusion::Rrf { k: 60 },
rerank: Rerank::Corroboration,
filters: Filters::open(space),
limit: 20,
explain: false,
}
}
pub fn lexical(space: impl Into<String>) -> Self {
Self {
targets: TargetSet::Statement,
channels: vec![
Channel::Lexical {
index: LexIndex::Word,
},
Channel::Lexical {
index: LexIndex::Ngram,
},
],
fusion: Fusion::Rrf { k: 60 },
rerank: Rerank::None,
filters: Filters::open(space),
limit: 20,
explain: false,
}
}
pub fn semantic(space: impl Into<String>) -> Self {
Self {
targets: TargetSet::Statement,
channels: vec![Channel::Vector {
space: VecSpace::Entity,
}],
fusion: Fusion::Rrf { k: 60 },
rerank: Rerank::Mmr {
lambda: 0.5,
max_similarity: Some(0.9),
},
filters: Filters::open(space),
limit: 20,
explain: false,
}
}
pub fn graph(space: impl Into<String>, seeds: Vec<EntityId>) -> Self {
Self {
targets: TargetSet::Statement,
channels: vec![Channel::GraphExpand {
seed: SeedPolicy::Explicit { entities: seeds },
depth: 2,
}],
fusion: Fusion::Rrf { k: 60 },
rerank: Rerank::GraphDistance { from: Vec::new() },
filters: Filters::open(space),
limit: 50,
explain: false,
}
}
pub fn community(space: impl Into<String>, seeds: Vec<EntityId>) -> Self {
Self {
targets: TargetSet::Statement,
channels: vec![Channel::CommunityExpand {
seed: SeedPolicy::Explicit { entities: seeds },
}],
fusion: Fusion::Rrf { k: 60 },
rerank: Rerank::None,
filters: Filters::open(space),
limit: 20,
explain: false,
}
}
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub struct ChannelRank {
pub channel: u8,
pub rank: u32,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TargetFacts {
pub target: TargetId,
pub confidence: f32,
pub valid_from: Timestamp,
pub valid_to: Timestamp,
pub recorded_at: Timestamp,
pub retracted_at: Option<Timestamp>,
pub trust: crate::TrustTier,
pub status: BeliefStatus,
pub predicate: String,
pub salience: f64,
pub distinct_episodes: u32,
pub channels: Vec<ChannelRank>,
pub channel_scores: Vec<f64>,
}
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum TargetId {
Episode { id: String },
Statement { id: StatementId },
Entity { id: EntityId },
Chunk { id: String },
Community { id: String },
}
impl TargetId {
pub fn rrf_key(&self) -> String {
match self {
TargetId::Episode { id } => format!("episode:{id}"),
TargetId::Statement { id } => format!("statement:{id}"),
TargetId::Entity { id } => format!("entity:{id}"),
TargetId::Chunk { id } => format!("chunk:{id}"),
TargetId::Community { id } => format!("community:{id}"),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChannelResult {
pub channel: u8,
pub hits: Vec<(TargetId, f64)>,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct RetrievalInput {
pub channels: Vec<ChannelResult>,
pub facts: HashMap<TargetId, TargetFacts>,
#[serde(default)]
pub entity_vectors: HashMap<String, Vec<f32>>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RankedItem {
pub target: TargetId,
pub fused_score: f64,
pub rank: usize,
pub channels: Vec<ChannelRank>,
pub salience: f64,
pub facts: TargetFacts,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DroppedItem {
pub target: TargetId,
pub reason: DropReason,
pub score: Option<f64>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum DropReason {
BelowConfidenceFloor {
actual: f32,
floor: f32,
},
OutsideValidWindow {
valid_at: Timestamp,
},
BeforeKnownAt {
known_at: Timestamp,
recorded_at: Timestamp,
},
TrustExcluded {
tier: crate::TrustTier,
},
PredicateDenied {
predicate: String,
},
EntityTypeMismatch {
expected: Vec<EntityTypeRef>,
},
TruncatedByBudget {
position: usize,
},
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RankingResult {
pub items: Vec<RankedItem>,
pub dropped: Vec<DroppedItem>,
pub total_candidates: usize,
pub spec: Retrieval,
}
pub fn rank(input: &RetrievalInput, spec: &Retrieval) -> RankingResult {
let mut candidate_facts: HashMap<String, TargetFacts> = HashMap::new();
let mut channel_ranks: HashMap<String, Vec<ChannelRank>> = HashMap::new();
let mut channel_scores: HashMap<String, Vec<f64>> = HashMap::new();
for cr in &input.channels {
for (rank_idx, (target, score)) in cr.hits.iter().enumerate() {
let key = target.rrf_key();
let entry = candidate_facts.entry(key.clone());
match entry {
std::collections::hash_map::Entry::Vacant(v) => {
if let Some(facts) = input.facts.get(target) {
v.insert(facts.clone());
} else {
v.insert(minimal_facts(target));
}
}
std::collections::hash_map::Entry::Occupied(mut o) => {
if let Some(newer) = input.facts.get(target) {
if newer.confidence > o.get().confidence {
o.insert(newer.clone());
}
}
}
}
channel_ranks
.entry(key.clone())
.or_default()
.push(ChannelRank {
channel: cr.channel,
rank: rank_idx as u32,
});
channel_scores.entry(key).or_default().push(*score);
}
}
for (target, facts) in &input.facts {
let key = target.rrf_key();
candidate_facts.entry(key).or_insert_with(|| facts.clone());
}
let total_candidates = candidate_facts.len();
let mut items: Vec<RankedItem> = Vec::with_capacity(candidate_facts.len());
let mut dropped: Vec<DroppedItem> = Vec::new();
let mut ordered_keys: Vec<&String> = candidate_facts.keys().collect();
ordered_keys.sort();
for key in &ordered_keys {
let facts = &candidate_facts[*key];
if let Some(reason) = check_filters(facts, &spec.filters) {
dropped.push(DroppedItem {
target: facts.target.clone(),
reason,
score: None,
});
continue;
}
let fused = fuse(channel_scores.get(*key), &spec.fusion);
let channels = channel_ranks.get(*key).cloned().unwrap_or_default();
items.push(RankedItem {
target: facts.target.clone(),
fused_score: fused,
rank: 0, channels,
salience: facts.salience,
facts: facts.clone(),
});
}
apply_rerank(&mut items, &spec.rerank, &input.entity_vectors);
items.sort_by(|a, b| {
b.fused_score
.partial_cmp(&a.fused_score)
.unwrap_or(std::cmp::Ordering::Equal)
.then_with(|| target_type_rank(&a.target).cmp(&target_type_rank(&b.target)))
.then_with(|| a.target.rrf_key().cmp(&b.target.rrf_key()))
});
items.truncate(spec.limit);
for (i, item) in items.iter_mut().enumerate() {
item.rank = i;
}
let kept_keys: std::collections::HashSet<String> =
items.iter().map(|i| i.target.rrf_key()).collect();
let already_dropped: std::collections::HashSet<String> =
dropped.iter().map(|d| d.target.rrf_key()).collect();
let mut post_truncate_keys: Vec<&String> = candidate_facts
.keys()
.filter(|k| !kept_keys.contains(*k) && !already_dropped.contains(*k))
.collect();
post_truncate_keys.sort();
let mut post_truncate_drops: Vec<DroppedItem> = post_truncate_keys
.into_iter()
.map(|k| {
let facts = &candidate_facts[k];
DroppedItem {
target: facts.target.clone(),
reason: DropReason::TruncatedByBudget {
position: ordered_keys.iter().position(|x| *x == k).unwrap_or(0),
},
score: None,
}
})
.collect();
dropped.append(&mut post_truncate_drops);
RankingResult {
items,
dropped,
total_candidates,
spec: spec.clone(),
}
}
fn minimal_facts(target: &TargetId) -> TargetFacts {
TargetFacts {
target: target.clone(),
confidence: 0.0,
valid_from: TIME_MIN,
valid_to: TIME_MAX,
recorded_at: TIME_MIN,
retracted_at: None,
trust: crate::TrustTier::SemiTrusted,
status: BeliefStatus::Active,
predicate: String::new(),
salience: 0.0,
distinct_episodes: 0,
channels: Vec::new(),
channel_scores: Vec::new(),
}
}
fn target_type_rank(t: &TargetId) -> u8 {
match t {
TargetId::Statement { .. } => 0,
TargetId::Episode { .. } => 1,
TargetId::Entity { .. } => 2,
TargetId::Chunk { .. } => 3,
TargetId::Community { .. } => 4,
}
}
fn check_filters(facts: &TargetFacts, filters: &Filters) -> Option<DropReason> {
if let Some(t) = filters.as_of {
if t < facts.valid_from || t > facts.valid_to {
return Some(DropReason::OutsideValidWindow { valid_at: t });
}
}
if let Some(t) = filters.known_at {
if facts.recorded_at > t {
return Some(DropReason::BeforeKnownAt {
known_at: t,
recorded_at: facts.recorded_at,
});
}
if let Some(retracted_at) = facts.retracted_at {
if retracted_at <= t {
return Some(DropReason::BeforeKnownAt {
known_at: t,
recorded_at: retracted_at,
});
}
}
}
if facts.confidence < filters.min_confidence {
return Some(DropReason::BelowConfidenceFloor {
actual: facts.confidence,
floor: filters.min_confidence,
});
}
if let TrustPolicy::Exclude(excluded) = &filters.trust {
if excluded.contains(&facts.trust) {
return Some(DropReason::TrustExcluded { tier: facts.trust });
}
}
if !filters.predicates.allows(&facts.predicate) {
return Some(DropReason::PredicateDenied {
predicate: facts.predicate.clone(),
});
}
if let Some(expected) = &filters.entity_types {
if !expected.is_empty() && !expected.iter().any(|t| t == &facts.predicate) {
return Some(DropReason::EntityTypeMismatch {
expected: expected.clone(),
});
}
}
None
}
fn fuse(scores: Option<&Vec<f64>>, fusion: &Fusion) -> f64 {
let Some(scores) = scores else { return 0.0 };
if scores.is_empty() {
return 0.0;
}
match fusion {
Fusion::Rrf { k } => {
let mut indexed: Vec<(usize, f64)> = scores.iter().copied().enumerate().collect();
indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
let k = *k as f64;
indexed
.iter()
.enumerate()
.map(|(pos, (_, _))| 1.0 / (k + (pos as f64) + 1.0))
.sum()
}
Fusion::Weighted { weights } => {
let min = scores.iter().cloned().fold(f64::INFINITY, f64::min);
let max = scores.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
let span = (max - min).max(f64::EPSILON);
let n = scores.len().max(weights.len());
let w_per = if weights.is_empty() {
1.0 / n as f64
} else {
let s: f64 = weights.iter().sum();
if s > 0.0 { 1.0 / s } else { 1.0 / n as f64 }
};
scores
.iter()
.map(|s| (s - min) / span)
.zip(weights.iter().chain(std::iter::repeat(&w_per)))
.map(|(n, w)| n * w)
.sum()
}
}
}
fn cosine_sim(a: &[f32], b: &[f32]) -> f64 {
if a.len() != b.len() || a.is_empty() {
return 0.0;
}
let dot: f64 = a
.iter()
.zip(b.iter())
.map(|(x, y)| (*x as f64) * (*y as f64))
.sum();
let na: f64 = a.iter().map(|x| (*x as f64).powi(2)).sum::<f64>().sqrt();
let nb: f64 = b.iter().map(|x| (*x as f64).powi(2)).sum::<f64>().sqrt();
if na == 0.0 || nb == 0.0 {
return 0.0;
}
dot / (na * nb)
}
fn item_vector<'a>(
item: &RankedItem,
vectors: &'a HashMap<String, Vec<f32>>,
) -> Option<&'a Vec<f32>> {
match &item.target {
TargetId::Entity { id } => vectors.get(id),
_ => None,
}
}
fn apply_rerank(
items: &mut [RankedItem],
rerank: &Rerank,
entity_vectors: &HashMap<String, Vec<f32>>,
) {
match rerank {
Rerank::None => {}
Rerank::Corroboration => {
for item in items.iter_mut() {
let boost = 1.0 + (1.0 + item.facts.distinct_episodes as f64).ln();
item.fused_score *= boost;
}
}
Rerank::GraphDistance { from } => {
let _ = from;
items.sort_by(|a, b| {
b.salience
.partial_cmp(&a.salience)
.unwrap_or(std::cmp::Ordering::Equal)
});
}
Rerank::Mmr {
lambda,
max_similarity,
} => {
if items.is_empty() {
return;
}
let lambda = *lambda as f64;
let ceiling = max_similarity.map(|c| c as f64);
let mut reordered: Vec<RankedItem> = Vec::with_capacity(items.len());
let mut pool: Vec<RankedItem> = items.to_vec();
pool.sort_by(|a, b| {
b.fused_score
.partial_cmp(&a.fused_score)
.unwrap_or(std::cmp::Ordering::Equal)
});
reordered.push(pool.remove(0));
while !pool.is_empty() {
let mut best_idx: Option<usize> = None;
let mut best_score = f64::NEG_INFINITY;
for (i, cand) in pool.iter().enumerate() {
let cand_vec = item_vector(cand, entity_vectors);
let mut max_sim: f64 = 0.0;
let mut used_cosine = false;
for sel in &reordered {
let sel_vec = item_vector(sel, entity_vectors);
if let (Some(cv), Some(sv)) = (cand_vec, sel_vec) {
let cos = cosine_sim(cv, sv);
if cos > max_sim {
max_sim = cos;
}
used_cosine = true;
}
}
let sim = if used_cosine {
max_sim
} else {
let last = reordered.last().unwrap().fused_score;
(last - cand.fused_score).abs()
};
if let Some(c) = ceiling {
if used_cosine && max_sim > c {
continue;
}
}
let mmr = lambda * cand.fused_score - (1.0 - lambda) * sim;
if mmr > best_score {
best_score = mmr;
best_idx = Some(i);
}
}
match best_idx {
Some(i) => reordered.push(pool.remove(i)),
None => {
pool.sort_by(|a, b| {
b.fused_score
.partial_cmp(&a.fused_score)
.unwrap_or(std::cmp::Ordering::Equal)
});
reordered.append(&mut pool);
break;
}
}
}
items.clone_from_slice(&reordered);
}
Rerank::Chain(reranks) => {
for r in reranks {
apply_rerank(items, r, entity_vectors);
}
}
}
}
pub fn direction_allows(direction: Direction, from_subject: bool) -> bool {
match direction {
Direction::Both => true,
Direction::Out => from_subject,
Direction::In => !from_subject,
}
}
pub fn explain_item(item: &RankedItem) -> BTreeMap<String, String> {
let mut out = BTreeMap::new();
out.insert("target".into(), item.target.rrf_key());
out.insert("fused_score".into(), format!("{:.6}", item.fused_score));
out.insert("salience".into(), format!("{:.6}", item.salience));
out.insert(
"channels".into(),
item.channels
.iter()
.map(|c| format!("ch{}:{}", c.channel, c.rank))
.collect::<Vec<_>>()
.join(","),
);
out
}
#[cfg(test)]
mod tests {
use super::*;
use crate::TrustTier;
fn facts(target: TargetId, confidence: f32, predicate: &str) -> TargetFacts {
TargetFacts {
target: target.clone(),
confidence,
valid_from: TIME_MIN,
valid_to: TIME_MAX,
recorded_at: Timestamp(1000),
retracted_at: None,
trust: TrustTier::Trusted,
status: BeliefStatus::Active,
predicate: predicate.into(),
salience: 0.5,
distinct_episodes: 1,
channels: vec![],
channel_scores: vec![],
}
}
#[test]
fn rrf_key_is_stable_and_distinguishes_kinds() {
let e = TargetId::Episode { id: "x".into() };
let s = TargetId::Statement { id: "x".into() };
assert_ne!(e.rrf_key(), s.rrf_key());
assert_eq!(e.rrf_key(), e.rrf_key());
}
#[test]
fn filter_drops_below_floor() {
let target = TargetId::Statement { id: "s1".into() };
let f = facts(target.clone(), 0.1, "works_on");
let mut filters = Filters::open("default");
filters.min_confidence = 0.5;
assert!(matches!(
check_filters(&f, &filters),
Some(DropReason::BelowConfidenceFloor { .. })
));
}
#[test]
fn filter_keeps_at_or_above_floor() {
let target = TargetId::Statement { id: "s1".into() };
let f = facts(target, 0.5, "works_on");
let filters = Filters::open("default");
assert!(check_filters(&f, &filters).is_none());
}
#[test]
fn filter_drops_outside_as_of() {
let target = TargetId::Statement { id: "s1".into() };
let mut f = facts(target, 0.9, "works_on");
f.valid_from = Timestamp(100);
f.valid_to = Timestamp(200);
let mut filters = Filters::open("default");
filters.as_of = Some(Timestamp(300));
assert!(matches!(
check_filters(&f, &filters),
Some(DropReason::OutsideValidWindow { .. })
));
}
#[test]
fn filter_drops_before_known_at() {
let target = TargetId::Statement { id: "s1".into() };
let mut f = facts(target, 0.9, "works_on");
f.recorded_at = Timestamp(1000);
let mut filters = Filters::open("default");
filters.known_at = Some(Timestamp(500));
assert!(matches!(
check_filters(&f, &filters),
Some(DropReason::BeforeKnownAt { .. })
));
}
#[test]
fn filter_drops_retracted_before_known_at() {
let target = TargetId::Statement { id: "s1".into() };
let mut f = facts(target, 0.9, "works_on");
f.recorded_at = Timestamp(100);
f.retracted_at = Some(Timestamp(200));
let mut filters = Filters::open("default");
filters.known_at = Some(Timestamp(300));
assert!(matches!(
check_filters(&f, &filters),
Some(DropReason::BeforeKnownAt { .. })
));
}
#[test]
fn rank_conservation_simple() {
let mut input = RetrievalInput::default();
let a = TargetId::Statement { id: "a".into() };
let b = TargetId::Statement { id: "b".into() };
input.facts.insert(a.clone(), facts(a.clone(), 0.9, "p"));
input.facts.insert(b.clone(), facts(b.clone(), 0.8, "p"));
input.channels.push(ChannelResult {
channel: 0,
hits: vec![(a.clone(), 0.9), (b.clone(), 0.8)],
});
let spec = Retrieval::hybrid("default");
let r = rank(&input, &spec);
assert_eq!(r.items.len() + r.dropped.len(), 2);
}
#[test]
fn rank_filter_totality() {
let mut input = RetrievalInput::default();
let a = TargetId::Statement { id: "a".into() };
input.facts.insert(a.clone(), facts(a.clone(), 0.1, "p"));
input.channels.push(ChannelResult {
channel: 0,
hits: vec![(a.clone(), 0.5)],
});
let mut spec = Retrieval::hybrid("default");
spec.filters.min_confidence = 0.5;
let r = rank(&input, &spec);
assert!(r.items.is_empty());
assert_eq!(r.dropped.len(), 1);
assert!(matches!(
r.dropped[0].reason,
DropReason::BelowConfidenceFloor { .. }
));
}
#[test]
fn rank_determinism() {
let make = || {
let mut input = RetrievalInput::default();
for (i, score) in [(0u8, 0.9f64), (1, 0.7), (2, 0.5)] {
let id = format!("s{i}");
let t = TargetId::Statement { id: id.clone() };
input.facts.insert(t.clone(), facts(t, score as f32, "p"));
input.channels.push(ChannelResult {
channel: 0,
hits: vec![(TargetId::Statement { id }, score)],
});
}
input
};
let spec = Retrieval::hybrid("default");
let r1 = rank(&make(), &spec);
let r2 = rank(&make(), &spec);
let keys1: Vec<_> = r1
.items
.iter()
.map(|i| (i.target.rrf_key(), i.fused_score, i.rank))
.collect();
let keys2: Vec<_> = r2
.items
.iter()
.map(|i| (i.target.rrf_key(), i.fused_score, i.rank))
.collect();
assert_eq!(keys1, keys2);
}
#[test]
fn corroboration_invariance_equal_episodes_preserves_order() {
let mut items: Vec<RankedItem> = [0.9f64, 0.7, 0.5]
.iter()
.enumerate()
.map(|(i, score)| RankedItem {
target: TargetId::Statement {
id: format!("s{i}"),
},
fused_score: *score,
rank: i,
channels: vec![],
salience: 0.5,
facts: {
let mut f = facts(
TargetId::Statement {
id: format!("s{i}"),
},
0.8,
"works_on",
);
f.distinct_episodes = 3; f
},
})
.collect();
apply_rerank(&mut items, &Rerank::Corroboration, &HashMap::new());
let boosted: Vec<f64> = items.iter().map(|i| i.fused_score).collect();
for w in boosted.windows(2) {
assert!(
w[0] >= w[1],
"ordering broken after equal-boost corroboration"
);
}
}
#[test]
fn corroboration_monotonicity_higher_distinct_ranks_higher() {
let mut items: Vec<RankedItem> = vec![
RankedItem {
target: TargetId::Statement { id: "low".into() },
fused_score: 0.5,
rank: 0,
channels: vec![],
salience: 0.5,
facts: {
let mut f = facts(TargetId::Statement { id: "low".into() }, 0.8, "works_on");
f.distinct_episodes = 1;
f
},
},
RankedItem {
target: TargetId::Statement { id: "high".into() },
fused_score: 0.5,
rank: 1,
channels: vec![],
salience: 0.5,
facts: {
let mut f = facts(TargetId::Statement { id: "high".into() }, 0.8, "works_on");
f.distinct_episodes = 10;
f
},
},
];
apply_rerank(&mut items, &Rerank::Corroboration, &HashMap::new());
let high_score = items
.iter()
.find(|i| i.target == TargetId::Statement { id: "high".into() })
.unwrap()
.fused_score;
let low_score = items
.iter()
.find(|i| i.target == TargetId::Statement { id: "low".into() })
.unwrap()
.fused_score;
assert!(
high_score > low_score,
"higher corroboration must rank higher: {high_score} vs {low_score}"
);
}
use proptest::prelude::*;
prop_compose! {
fn arb_target_id()(i in any::<u8>()) -> TargetId {
let id = format!("t{i:03}");
match i % 3 {
0 => TargetId::Statement { id },
1 => TargetId::Entity { id },
_ => TargetId::Episode { id },
}
}
}
prop_compose! {
fn arb_facts()(
target in arb_target_id(),
confidence in 0.0f32..1.0f32,
vf in 0i64..10_000,
vt in 0i64..10_000,
recorded_at in 0i64..10_000,
retracted_at in prop::option::of(0i64..10_000),
trust in prop::sample::select(vec![
TrustTier::Trusted, TrustTier::SemiTrusted, TrustTier::Untrusted,
]),
status in prop::sample::select(vec![
BeliefStatus::Active,
BeliefStatus::Superseded,
BeliefStatus::Contradicted,
BeliefStatus::Retracted,
]),
predicate in prop::sample::select(vec![
"works_on".to_string(), "knows".to_string(), "likes".to_string(),
]),
salience in 0.0f64..1.0f64,
distinct in 0u32..10,
) -> TargetFacts {
TargetFacts {
target,
confidence,
valid_from: Timestamp(vf),
valid_to: Timestamp(vt),
recorded_at: Timestamp(recorded_at),
retracted_at: retracted_at.map(Timestamp),
trust,
status,
predicate,
salience,
distinct_episodes: distinct,
channels: vec![],
channel_scores: vec![],
}
}
}
prop_compose! {
fn arb_filters()(
as_of in prop::option::of(0i64..10_000),
known_at in prop::option::of(0i64..10_000),
min_confidence in 0.0f32..1.0f32,
) -> Filters {
Filters {
space: "default".into(),
as_of: as_of.map(Timestamp),
known_at: known_at.map(Timestamp),
min_confidence,
trust: TrustPolicy::All,
predicates: PredicateFilter::AllowAll,
entity_types: None,
}
}
}
prop_compose! {
fn arb_input()(entries in prop::collection::vec((arb_facts(), 0.0f64..1.0f64), 1..16)) -> RetrievalInput {
let mut input = RetrievalInput::default();
let mut hits: Vec<(TargetId, f64)> = Vec::new();
for (f, score) in entries {
let target = f.target.clone();
input.facts.insert(target.clone(), f);
hits.push((target, score));
}
input.channels.push(ChannelResult { channel: 0, hits });
input
}
}
proptest! {
#![proptest_config(ProptestConfig::with_cases(64))]
#[test]
fn prop_conservation(input in arb_input(), filters in arb_filters()) {
let mut spec = Retrieval::hybrid("default");
spec.filters = filters;
let total_candidates = input.facts.len();
let r = rank(&input, &spec);
prop_assert_eq!(r.items.len() + r.dropped.len(), total_candidates);
let item_keys: std::collections::HashSet<_> =
r.items.iter().map(|i| i.target.rrf_key()).collect();
let drop_keys: std::collections::HashSet<_> =
r.dropped.iter().map(|d| d.target.rrf_key()).collect();
prop_assert!(item_keys.is_disjoint(&drop_keys));
prop_assert_eq!(r.total_candidates, total_candidates);
}
#[test]
fn prop_filter_totality(input in arb_input(), filters in arb_filters()) {
let mut spec = Retrieval::hybrid("default");
spec.filters = filters.clone();
let r = rank(&input, &spec);
for item in &r.items {
let f = &item.facts;
prop_assert!(f.confidence >= filters.min_confidence,
"item {} violates min_confidence: {} < {}",
item.target.rrf_key(), f.confidence, filters.min_confidence);
if let Some(t) = filters.as_of {
prop_assert!(t >= f.valid_from && t <= f.valid_to,
"item {} violates as_of {}", item.target.rrf_key(), t.0);
}
if let Some(t) = filters.known_at {
prop_assert!(f.recorded_at <= t,
"item {} violates known_at {}", item.target.rrf_key(), t.0);
if let Some(r_at) = f.retracted_at {
prop_assert!(r_at > t,
"item {} violates known_at (retracted)", item.target.rrf_key());
}
}
if let TrustPolicy::Exclude(excluded) = &filters.trust {
prop_assert!(!excluded.contains(&f.trust),
"item {} has excluded trust tier", item.target.rrf_key());
}
}
}
#[test]
fn prop_determinism(input in arb_input(), filters in arb_filters()) {
let mut spec = Retrieval::hybrid("default");
spec.filters = filters;
let r1 = rank(&input, &spec);
let r2 = rank(&input, &spec);
let j = |r: &RankingResult| serde_json::to_string(r).expect("serialize");
prop_assert_eq!(j(&r1), j(&r2));
}
}
#[test]
fn mmr_ceiling_defers_near_duplicates_in_top_10() {
use std::collections::HashMap;
let mut vectors: HashMap<String, Vec<f32>> = HashMap::new();
let mut names: Vec<String> = Vec::new();
for i in 0..10 {
let name = format!("d{i}");
names.push(name.clone());
let angle = (i as f64) * std::f64::consts::PI / 5.0;
vectors.insert(name, vec![angle.cos() as f32, angle.sin() as f32]);
}
for i in 0..5 {
let name = format!("dup{i}");
names.push(name.clone());
let angle = (i as f64) * std::f64::consts::PI / 5.0;
vectors.insert(
name,
vec![(angle.cos() * 1.01) as f32, (angle.sin() * 1.01) as f32],
);
}
let mut items: Vec<RankedItem> = names
.iter()
.enumerate()
.map(|(i, e)| {
let score = if e.starts_with("dup") {
1.5 - 0.01 * i as f64 } else {
1.0 - 0.01 * i as f64 };
RankedItem {
target: TargetId::Entity { id: e.clone() },
facts: facts(TargetId::Entity { id: e.clone() }, 0.8, "works_on"),
fused_score: score,
salience: 0.8,
rank: i,
channels: vec![],
}
})
.collect();
apply_rerank(
&mut items,
&Rerank::Mmr {
lambda: 0.5,
max_similarity: Some(0.9),
},
&vectors,
);
for (i, a) in items.iter().take(10).enumerate() {
let a_vec = match &a.target {
TargetId::Entity { id } => vectors.get(id).unwrap(),
_ => unreachable!(),
};
for (j, b) in items.iter().take(10).enumerate() {
if i >= j {
continue;
}
let b_vec = match &b.target {
TargetId::Entity { id } => vectors.get(id).unwrap(),
_ => unreachable!(),
};
let sim = cosine_sim(a_vec, b_vec);
assert!(
sim <= 0.9,
"MMR kept a >0.9 pair at top-10 positions {i}/{j} \
({:?} / {:?}): cosine = {sim:.4}",
a.target,
b.target,
);
}
}
assert_eq!(items.len(), 15);
}
#[test]
fn mmr_ceiling_never_drops_items() {
use std::collections::HashMap;
let mut vectors: HashMap<String, Vec<f32>> = HashMap::new();
let names: Vec<String> = (0..8).map(|i| format!("e{i}")).collect();
for name in &names {
vectors.insert(name.clone(), vec![1.0, 0.01, 0.0]);
}
let mut items: Vec<RankedItem> = names
.iter()
.enumerate()
.map(|(i, e)| RankedItem {
target: TargetId::Entity { id: e.clone() },
facts: facts(TargetId::Entity { id: e.clone() }, 0.8, "works_on"),
fused_score: 1.0 - 0.01 * i as f64,
salience: 0.8,
rank: i,
channels: vec![],
})
.collect();
let count_before = items.len();
apply_rerank(
&mut items,
&Rerank::Mmr {
lambda: 0.5,
max_similarity: Some(0.9),
},
&vectors,
);
assert_eq!(
items.len(),
count_before,
"MMR ceiling must not drop items — conservation invariant"
);
}
}