weavatrix_memory/context/
retrieval.rs

1use crate::EntityId;
2use core::fmt;
3use serde::{Deserialize, Serialize};
4use std::collections::{BTreeMap, BTreeSet};
5
6const RRF_K: u64 = 60;
7const RRF_SCALE: u64 = 1_000_000_000;
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
10#[serde(rename_all = "snake_case")]
11pub enum RetrievalChannel {
12    Literal,
13    Lexical,
14    Semantic,
15    Hybrid,
16}
17
18#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
19pub struct RetrievalQuery {
20    pub text: String,
21    pub limit: usize,
22    pub channels: BTreeSet<RetrievalChannel>,
23}
24
25impl RetrievalQuery {
26    /// Creates a provider-neutral retrieval request.
27    ///
28    /// # Errors
29    ///
30    /// Rejects empty text and a zero result limit.
31    pub fn new(text: impl Into<String>, limit: usize) -> RetrievalResult<Self> {
32        let text = text.into();
33        if text.is_empty() || text.trim() != text {
34            return Err(RetrievalError::new(
35                "query",
36                "text must be non-empty and trimmed",
37            ));
38        }
39        if limit == 0 {
40            return Err(RetrievalError::new(
41                "query",
42                "limit must be greater than zero",
43            ));
44        }
45        Ok(Self {
46            text,
47            limit,
48            channels: BTreeSet::new(),
49        })
50    }
51
52    #[must_use]
53    pub fn include(mut self, channel: RetrievalChannel) -> Self {
54        self.channels.insert(channel);
55        self
56    }
57}
58
59#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
60pub struct RetrievalHit {
61    pub entity: EntityId,
62    pub channel: RetrievalChannel,
63    pub score: u32,
64}
65
66impl RetrievalHit {
67    #[must_use]
68    pub const fn new(entity: EntityId, channel: RetrievalChannel, score: u32) -> Self {
69        Self {
70            entity,
71            channel,
72            score,
73        }
74    }
75}
76
77#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
78pub struct RetrievalSource {
79    pub provider: String,
80    pub channel: RetrievalChannel,
81    pub rank: usize,
82    pub raw_score: u32,
83}
84
85#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
86pub struct FusedRetrievalHit {
87    pub entity: EntityId,
88    pub fused_score: u64,
89    pub sources: Vec<RetrievalSource>,
90}
91
92pub trait RetrievalProvider: Sync {
93    fn name(&self) -> &str;
94
95    /// Finds exact memory entity identifiers for a free-text query.
96    ///
97    /// # Errors
98    ///
99    /// Returns a provider-specific retrieval failure.
100    fn retrieve(&self, query: &RetrievalQuery) -> RetrievalResult<Vec<RetrievalHit>>;
101}
102
103#[derive(Debug, Clone, PartialEq, Eq)]
104pub struct RetrievalError {
105    pub provider: String,
106    pub message: String,
107}
108
109impl RetrievalError {
110    #[must_use]
111    pub fn new(provider: impl Into<String>, message: impl Into<String>) -> Self {
112        Self {
113            provider: provider.into(),
114            message: message.into(),
115        }
116    }
117}
118
119impl fmt::Display for RetrievalError {
120    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
121        write!(
122            formatter,
123            "retrieval provider {} failed: {}",
124            self.provider, self.message
125        )
126    }
127}
128
129impl std::error::Error for RetrievalError {}
130
131pub type RetrievalResult<T> = core::result::Result<T, RetrievalError>;
132
133/// Fuses provider ranks with deterministic integer reciprocal-rank fusion.
134///
135/// Scores from BM25 and vector stores are deliberately not compared directly.
136/// Providers, ties, duplicates, and final output all have stable ordering.
137///
138/// # Errors
139///
140/// Propagates provider failures and rejects duplicate or empty provider names.
141pub fn fuse_retrieval(
142    providers: &[&dyn RetrievalProvider],
143    query: &RetrievalQuery,
144) -> RetrievalResult<Vec<FusedRetrievalHit>> {
145    let mut providers = providers.to_vec();
146    providers.sort_by(|left, right| left.name().cmp(right.name()));
147    validate_providers(&providers)?;
148    let mut fused = BTreeMap::<EntityId, FusedRetrievalHit>::new();
149    for provider in providers {
150        let mut hits = provider.retrieve(query)?;
151        hits.retain(|hit| query.channels.is_empty() || query.channels.contains(&hit.channel));
152        hits.sort_by(|left, right| {
153            right
154                .score
155                .cmp(&left.score)
156                .then_with(|| left.entity.cmp(&right.entity))
157                .then_with(|| left.channel.cmp(&right.channel))
158        });
159        let mut seen = BTreeSet::new();
160        hits.retain(|hit| seen.insert(hit.entity.clone()));
161        for (offset, hit) in hits.into_iter().take(query.limit).enumerate() {
162            let rank = offset + 1;
163            let contribution = RRF_SCALE / (RRF_K + rank as u64);
164            let entry = fused
165                .entry(hit.entity.clone())
166                .or_insert_with(|| FusedRetrievalHit {
167                    entity: hit.entity,
168                    fused_score: 0,
169                    sources: Vec::new(),
170                });
171            entry.fused_score = entry.fused_score.saturating_add(contribution);
172            entry.sources.push(RetrievalSource {
173                provider: provider.name().to_owned(),
174                channel: hit.channel,
175                rank,
176                raw_score: hit.score,
177            });
178        }
179    }
180    let mut output = fused.into_values().collect::<Vec<_>>();
181    output.sort_by(|left, right| {
182        right
183            .fused_score
184            .cmp(&left.fused_score)
185            .then_with(|| left.entity.cmp(&right.entity))
186    });
187    output.truncate(query.limit);
188    Ok(output)
189}
190
191fn validate_providers(providers: &[&dyn RetrievalProvider]) -> RetrievalResult<()> {
192    let mut names = BTreeSet::new();
193    for provider in providers {
194        let name = provider.name();
195        if name.is_empty() || name.trim() != name {
196            return Err(RetrievalError::new(
197                "provider",
198                "provider names must be non-empty and trimmed",
199            ));
200        }
201        if !names.insert(name) {
202            return Err(RetrievalError::new(name, "provider names must be unique"));
203        }
204    }
205    Ok(())
206}
weavatrix_memory/context/retrieval.rs

weavatrix_memory/context/
retrieval.rs
