weavatrix_memory/context/
retrieval.rs

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

weavatrix_memory/context/
retrieval.rs
