mnemo_amp/
router.rs

1//! Fan-out router + cross-adapter fusion primitives.
2//!
3//! [`AmpRouter`] is the single entry point a transport hands an
4//! [`AmpEnvelope`] to: it dispatches to the backing [`MemoryStore`].
5//! When more than one store is registered it **fans out** writes to
6//! all of them and fuses reads, which is the shape the AMP
7//! cross-adapter conformance suite exercises (one logical memory
8//! surface backed by several adapters).
9//!
10//! The fusion lane ships two combiners — [`rrf_fuse`] (Reciprocal Rank
11//! Fusion) and [`max_fuse`] (max-score) — used by the conformance test
12//! to demonstrate RRF's robustness to a rank-0 adversarial injection
13//! that max-fusion is fooled by.
14
15use std::sync::Arc;
16
17use crate::error::AmpError;
18use crate::store::MemoryStore;
19use crate::wire::{AmpEnvelope, AmpHit, AmpOp, AmpResult};
20
21/// Routes AMP envelopes to one or more [`MemoryStore`] backends.
22#[derive(Clone)]
23pub struct AmpRouter {
24    stores: Vec<Arc<dyn MemoryStore>>,
25    /// RRF rank constant `k`. 60 is the canonical TREC default.
26    rrf_k: f32,
27}
28
29impl AmpRouter {
30    /// Single-backend router (the common case).
31    pub fn new(store: Arc<dyn MemoryStore>) -> Self {
32        Self {
33            stores: vec![store],
34            rrf_k: 60.0,
35        }
36    }
37
38    /// Fan-out router over several backends.
39    pub fn fan_out(stores: Vec<Arc<dyn MemoryStore>>) -> Self {
40        Self {
41            stores,
42            rrf_k: 60.0,
43        }
44    }
45
46    pub fn with_rrf_k(mut self, k: f32) -> Self {
47        self.rrf_k = k;
48        self
49    }
50
51    /// Route one envelope. Writes (`remember` / `forget` / `merge` /
52    /// `expire`) fan out to every backend; the first backend's result
53    /// is returned, with any backend error surfaced. `recall` fans out
54    /// and fuses the per-backend hit lists with RRF.
55    pub async fn route(&self, env: &AmpEnvelope) -> Result<AmpResult, AmpError> {
56        match env.op {
57            AmpOp::Recall => self.route_recall(env).await,
58            _ => {
59                let mut first: Option<AmpResult> = None;
60                for store in &self.stores {
61                    let r = store.dispatch(env).await?;
62                    if first.is_none() {
63                        first = Some(r);
64                    }
65                }
66                first.ok_or_else(|| AmpError::Validation("router has no backends".into()))
67            }
68        }
69    }
70
71    async fn route_recall(&self, env: &AmpEnvelope) -> Result<AmpResult, AmpError> {
72        if self.stores.len() == 1 {
73            return self.stores[0].recall(env).await;
74        }
75        let mut lists: Vec<Vec<AmpHit>> = Vec::with_capacity(self.stores.len());
76        for store in &self.stores {
77            lists.push(store.recall(env).await?.hits);
78        }
79        let fused = rrf_fuse(&lists, self.rrf_k);
80        let mut out = AmpResult::ok(AmpOp::Recall);
81        out.hits = fused;
82        Ok(out)
83    }
84}
85
86/// Reciprocal Rank Fusion over several ranked hit lists.
87///
88/// Each hit contributes `1 / (k + rank)` (rank is 0-based within its
89/// list) to its id's fused score; identical ids across lists sum. The
90/// `k` damping is what makes RRF robust to a single adversarial top
91/// rank: a rank-0 injection in one list adds only `1/(k+0)`, which a
92/// genuinely-relevant item ranked highly across *multiple* lists still
93/// beats. Returns hits sorted by fused score, highest first.
94pub fn rrf_fuse(lists: &[Vec<AmpHit>], k: f32) -> Vec<AmpHit> {
95    use std::collections::HashMap;
96    let mut score: HashMap<String, f32> = HashMap::new();
97    let mut repr: HashMap<String, AmpHit> = HashMap::new();
98    for list in lists {
99        for (rank, hit) in list.iter().enumerate() {
100            *score.entry(hit.id.clone()).or_insert(0.0) += 1.0 / (k + rank as f32);
101            repr.entry(hit.id.clone()).or_insert_with(|| hit.clone());
102        }
103    }
104    sort_by_fused(score, repr)
105}
106
107/// Max-score fusion: an id's fused score is the single best score it
108/// earned in any list. Simple, but a rank-0 adversarial injection with
109/// an inflated score wins outright — the failure mode the conformance
110/// test contrasts against RRF.
111pub fn max_fuse(lists: &[Vec<AmpHit>]) -> Vec<AmpHit> {
112    use std::collections::HashMap;
113    let mut score: HashMap<String, f32> = HashMap::new();
114    let mut repr: HashMap<String, AmpHit> = HashMap::new();
115    for list in lists {
116        for hit in list {
117            let e = score.entry(hit.id.clone()).or_insert(f32::MIN);
118            if hit.score > *e {
119                *e = hit.score;
120            }
121            repr.entry(hit.id.clone()).or_insert_with(|| hit.clone());
122        }
123    }
124    sort_by_fused(score, repr)
125}
126
127fn sort_by_fused(
128    score: std::collections::HashMap<String, f32>,
129    repr: std::collections::HashMap<String, AmpHit>,
130) -> Vec<AmpHit> {
131    let mut fused: Vec<AmpHit> = repr
132        .into_iter()
133        .map(|(id, mut hit)| {
134            hit.score = score.get(&id).copied().unwrap_or(0.0);
135            hit
136        })
137        .collect();
138    fused.sort_by(|a, b| {
139        // Deterministic: score desc, then id asc to break ties.
140        b.score
141            .partial_cmp(&a.score)
142            .unwrap_or(std::cmp::Ordering::Equal)
143            .then_with(|| a.id.cmp(&b.id))
144    });
145    fused
146}
147
148#[cfg(test)]
149mod tests {
150    use super::*;
151    use crate::wire::AmpMemoryType;
152
153    fn hit(id: &str, score: f32) -> AmpHit {
154        AmpHit {
155            id: id.to_string(),
156            content: format!("content-{id}"),
157            memory_type: AmpMemoryType::Semantic,
158            score,
159            tags: vec![],
160        }
161    }
162
163    #[test]
164    fn rrf_holds_under_rank0_injection_but_max_is_fooled() {
165        // `TRUE` is the genuinely-relevant item: ranked high in BOTH
166        // lists. `ADV` is an adversarial injection sitting at rank 0 of
167        // list A only, with an inflated raw score.
168        let list_a = vec![
169            hit("ADV", 999.0), // rank 0, adversarial, huge score
170            hit("TRUE", 0.9),  // rank 1
171            hit("x1", 0.5),
172        ];
173        let list_b = vec![
174            hit("TRUE", 0.95), // rank 0 in the honest list
175            hit("y1", 0.6),
176            hit("y2", 0.4),
177        ];
178
179        // RRF: TRUE scores 1/(60+1) + 1/(60+0); ADV scores 1/(60+0)
180        // only. TRUE wins.
181        let rrf = rrf_fuse(&[list_a.clone(), list_b.clone()], 60.0);
182        assert_eq!(rrf[0].id, "TRUE", "RRF must rank the true item first");
183
184        // Max-fusion: ADV's 999.0 is the single best score anywhere, so
185        // the injection wins — the failure mode.
186        let max = max_fuse(&[list_a, list_b]);
187        assert_eq!(
188            max[0].id, "ADV",
189            "max-fusion is fooled by the rank-0 injection"
190        );
191    }
192
193    #[test]
194    fn rrf_is_deterministic() {
195        let a = vec![hit("a", 0.9), hit("b", 0.8)];
196        let b = vec![hit("b", 0.7), hit("a", 0.6)];
197        let r1 = rrf_fuse(&[a.clone(), b.clone()], 60.0);
198        let r2 = rrf_fuse(&[a, b], 60.0);
199        assert_eq!(r1, r2);
200    }
201}
mnemo_amp/router.rs

mnemo_amp/
router.rs
