Skip to main content

relay_knowledge/retrieval/
mod.rs

1//! Retrieval request planning and optional derived backend adapters.
2//!
3//! Retrieval owns query-shape validation and budgets before the application
4//! service asks storage and derived indexes for data.
5
6use std::time::Duration;
7use std::{error::Error, fmt};
8
9pub mod provider;
10mod rerank;
11pub(crate) mod terms;
12
13use crate::domain::{
14    FreshnessPolicy, GraphVersion, IndexKind, IndexStatus, RerankDiagnostics, RerankMode,
15    RetrievalBackendState, RetrievalBackendStatus, RetrievalHit, RetrieverSource,
16};
17
18pub const LOCAL_SEMANTIC_MODEL: &str = "relay-local-token-semantic-v1";
19pub const LOCAL_VECTOR_MODEL: &str = "relay-local-hash-ann-v1";
20pub const LOCAL_RERANK_MODEL: &str = "relay-local-deterministic-rerank-v1";
21pub const LOCAL_VECTOR_DIMENSION: u32 = 16;
22pub const DEFAULT_EMBEDDING_BATCH_SIZE: usize = 32;
23pub const DEFAULT_EMBEDDING_TIMEOUT: Duration = Duration::from_secs(30);
24pub const DEFAULT_EMBEDDING_MAX_CONCURRENCY: usize = 4;
25pub const DEFAULT_RERANK_TIMEOUT: Duration = Duration::from_millis(100);
26pub const DEFAULT_RERANK_CANDIDATE_MULTIPLIER: usize = 4;
27pub const DEFAULT_RERANK_MAX_CANDIDATES: usize = 64;
28
29/// Validated retrieval request with bounded result count.
30#[derive(Debug, Clone, PartialEq, Eq)]
31pub struct RetrievalPlan {
32    pub query: String,
33    pub source_scope: Option<String>,
34    pub limit: usize,
35    pub freshness: FreshnessPolicy,
36}
37
38/// Configured owner of a semantic or vector read model.
39#[derive(Debug, Clone, Copy, PartialEq, Eq)]
40pub enum ReadModelBackendMode {
41    Local,
42    External,
43    Disabled,
44}
45
46/// Remote LLM provider family used for embedding calls.
47#[derive(Debug, Clone, Copy, PartialEq, Eq)]
48pub enum EmbeddingProviderKind {
49    OpenAiCompatible,
50    Echo,
51}
52
53impl EmbeddingProviderKind {
54    /// Parses a stable environment/config value.
55    pub fn parse(value: &str) -> Result<Self, EmbeddingProviderKindError> {
56        match value.trim().to_ascii_lowercase().as_str() {
57            "openai_compatible" => Ok(Self::OpenAiCompatible),
58            "echo" => Ok(Self::Echo),
59            other => Err(EmbeddingProviderKindError {
60                value: other.to_owned(),
61            }),
62        }
63    }
64
65    /// Stable configuration label.
66    pub const fn as_str(self) -> &'static str {
67        match self {
68            Self::OpenAiCompatible => "openai_compatible",
69            Self::Echo => "echo",
70        }
71    }
72}
73
74/// Invalid embedding provider kind supplied by runtime configuration.
75#[derive(Debug, Clone, PartialEq, Eq)]
76pub struct EmbeddingProviderKindError {
77    pub value: String,
78}
79
80impl fmt::Display for EmbeddingProviderKindError {
81    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
82        write!(
83            formatter,
84            "embedding provider '{}' must be openai_compatible or echo",
85            self.value
86        )
87    }
88}
89
90impl Error for EmbeddingProviderKindError {}
91
92impl ReadModelBackendMode {
93    /// Parses a stable environment/config value.
94    pub fn parse(value: &str) -> Result<Self, ReadModelBackendModeError> {
95        match value.trim().to_ascii_lowercase().as_str() {
96            "local" => Ok(Self::Local),
97            "external" => Ok(Self::External),
98            "disabled" => Ok(Self::Disabled),
99            other => Err(ReadModelBackendModeError {
100                value: other.to_owned(),
101            }),
102        }
103    }
104
105    /// Stable configuration label.
106    pub const fn as_str(self) -> &'static str {
107        match self {
108            Self::Local => "local",
109            Self::External => "external",
110            Self::Disabled => "disabled",
111        }
112    }
113}
114
115/// Invalid read model backend mode supplied by runtime configuration.
116#[derive(Debug, Clone, PartialEq, Eq)]
117pub struct ReadModelBackendModeError {
118    pub value: String,
119}
120
121impl fmt::Display for ReadModelBackendModeError {
122    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
123        write!(
124            formatter,
125            "retrieval backend '{}' must be local, external, or disabled",
126            self.value
127        )
128    }
129}
130
131impl Error for ReadModelBackendModeError {}
132
133/// Model metadata used by semantic/vector refresh workers and diagnostics.
134#[derive(Debug, Clone, PartialEq, Eq)]
135pub struct ReadModelMetadata {
136    pub name: String,
137    pub dimension: u32,
138}
139
140/// Runtime configuration for post-fusion result reranking.
141#[derive(Debug, Clone, PartialEq, Eq)]
142pub struct RerankConfig {
143    pub mode: RerankMode,
144    pub model: Option<String>,
145    pub timeout: Duration,
146    pub candidate_multiplier: usize,
147    pub max_candidates: usize,
148}
149
150impl RerankConfig {
151    /// Uses the built-in deterministic reranker on an expanded local candidate pool.
152    pub fn local() -> Self {
153        Self {
154            mode: RerankMode::Local,
155            model: Some(LOCAL_RERANK_MODEL.to_owned()),
156            timeout: DEFAULT_RERANK_TIMEOUT,
157            candidate_multiplier: DEFAULT_RERANK_CANDIDATE_MULTIPLIER,
158            max_candidates: DEFAULT_RERANK_MAX_CANDIDATES,
159        }
160    }
161
162    /// Returns the storage candidate budget required before final truncation.
163    pub fn candidate_limit(&self, requested_limit: usize) -> usize {
164        let truncation_probe_limit = requested_limit.saturating_add(1);
165        if self.mode == RerankMode::Disabled {
166            return truncation_probe_limit;
167        }
168
169        let expanded = requested_limit
170            .saturating_mul(self.candidate_multiplier)
171            .max(truncation_probe_limit);
172        expanded.min(self.max_candidates.max(truncation_probe_limit))
173    }
174
175    /// Applies the configured reranking policy to post-fusion candidates.
176    pub fn rerank(
177        &self,
178        query: &str,
179        hits: Vec<RetrievalHit>,
180    ) -> (Vec<RetrievalHit>, RerankDiagnostics) {
181        rerank::rerank_hits(query, hits, self)
182    }
183}
184
185/// Runtime configuration for a remote embedding provider.
186#[derive(Debug, Clone, PartialEq, Eq)]
187pub struct RemoteEmbeddingConfig {
188    pub provider: EmbeddingProviderKind,
189    pub base_url: String,
190    pub api_key: String,
191    pub batch_size: usize,
192    pub timeout: Duration,
193    pub max_concurrency: usize,
194}
195
196impl RemoteEmbeddingConfig {
197    /// Returns a URL label that is safe to expose in diagnostics.
198    pub fn redacted_base_url(&self) -> String {
199        redacted_url(&self.base_url)
200    }
201}
202
203/// Runtime read model configuration shared by refresh and retrieval status.
204#[derive(Debug, Clone, PartialEq, Eq)]
205pub struct ReadModelBackendConfig {
206    pub semantic_mode: ReadModelBackendMode,
207    pub vector_mode: ReadModelBackendMode,
208    pub semantic_model: ReadModelMetadata,
209    pub vector_model: ReadModelMetadata,
210    pub image_model: ReadModelMetadata,
211    pub remote_embedding: Option<RemoteEmbeddingConfig>,
212    pub rerank: RerankConfig,
213}
214
215impl ReadModelBackendConfig {
216    /// Uses the built-in deterministic read models.
217    pub fn local() -> Self {
218        Self {
219            semantic_mode: ReadModelBackendMode::Local,
220            vector_mode: ReadModelBackendMode::Local,
221            semantic_model: ReadModelMetadata {
222                name: LOCAL_SEMANTIC_MODEL.to_owned(),
223                dimension: LOCAL_VECTOR_DIMENSION,
224            },
225            vector_model: ReadModelMetadata {
226                name: LOCAL_VECTOR_MODEL.to_owned(),
227                dimension: LOCAL_VECTOR_DIMENSION,
228            },
229            image_model: ReadModelMetadata {
230                name: "relay-local-image-hash-v1".to_owned(),
231                dimension: LOCAL_VECTOR_DIMENSION,
232            },
233            remote_embedding: None,
234            rerank: RerankConfig::local(),
235        }
236    }
237
238    /// Returns whether local index refresh should maintain an index family.
239    pub fn refreshes_index(&self, kind: IndexKind) -> bool {
240        match kind {
241            IndexKind::Bm25 => true,
242            IndexKind::Semantic => self.semantic_mode != ReadModelBackendMode::Disabled,
243            IndexKind::Vector => self.vector_mode != ReadModelBackendMode::Disabled,
244        }
245    }
246
247    /// Returns read-model retrievers that must not execute for a request.
248    pub fn disabled_retriever_sources(&self) -> Vec<RetrieverSource> {
249        let mut disabled = Vec::new();
250        if self.semantic_mode == ReadModelBackendMode::Disabled {
251            disabled.push(RetrieverSource::Semantic);
252        }
253        if self.vector_mode == ReadModelBackendMode::Disabled {
254            disabled.push(RetrieverSource::Vector);
255        }
256
257        disabled
258    }
259}
260
261fn redacted_url(value: &str) -> String {
262    let trimmed = value.trim();
263    let Some((scheme, rest)) = trimmed.split_once("://") else {
264        return trimmed.to_owned();
265    };
266    let authority = rest.split('/').next().unwrap_or(rest);
267    let host = authority
268        .rsplit_once('@')
269        .map_or(authority, |(_, host)| host);
270    if host.is_empty() {
271        return scheme.to_owned();
272    }
273
274    format!("{scheme}://{host}")
275}
276
277/// Builds semantic/vector backend status from configured read models and index cursors.
278pub fn read_model_backend_statuses(
279    plan: &RetrievalPlan,
280    graph_version: GraphVersion,
281    indexes: &[IndexStatus],
282    config: &ReadModelBackendConfig,
283) -> Vec<RetrievalBackendStatus> {
284    [
285        (
286            RetrieverSource::Semantic,
287            IndexKind::Semantic,
288            config.semantic_mode,
289            &config.semantic_model,
290        ),
291        (
292            RetrieverSource::Vector,
293            IndexKind::Vector,
294            config.vector_mode,
295            &config.vector_model,
296        ),
297    ]
298    .into_iter()
299    .map(|(source, kind, mode, metadata)| {
300        read_model_backend_status(source, kind, mode, metadata, plan, graph_version, indexes)
301    })
302    .collect()
303}
304
305fn read_model_backend_status(
306    source: RetrieverSource,
307    kind: IndexKind,
308    mode: ReadModelBackendMode,
309    metadata: &ReadModelMetadata,
310    plan: &RetrievalPlan,
311    graph_version: GraphVersion,
312    indexes: &[IndexStatus],
313) -> RetrievalBackendStatus {
314    if mode == ReadModelBackendMode::Disabled {
315        return RetrievalBackendStatus {
316            source,
317            state: RetrievalBackendState::Unavailable,
318            scope_post_filter: plan.source_scope.is_some(),
319            indexed_graph_version: None,
320            reason: Some(format!(
321                "{} read model disabled by configuration",
322                source.as_str()
323            )),
324        };
325    }
326
327    let Some(index) = indexes.iter().find(|status| status.kind == kind) else {
328        return RetrievalBackendStatus {
329            source,
330            state: RetrievalBackendState::Unavailable,
331            scope_post_filter: plan.source_scope.is_some(),
332            indexed_graph_version: None,
333            reason: Some(format!("{} index metadata is unavailable", source.as_str())),
334        };
335    };
336    let stale = index.is_stale_for(graph_version);
337    let reason = if stale {
338        format!(
339            "{} read model index is stale at graph version {} while graph is {}; configured {} backend model={} dimension={}",
340            source.as_str(),
341            index.indexed_graph_version.get(),
342            graph_version.get(),
343            mode.as_str(),
344            metadata.name,
345            metadata.dimension
346        )
347    } else {
348        format!(
349            "{} read model available through {} backend model={} dimension={}",
350            source.as_str(),
351            mode.as_str(),
352            metadata.name,
353            metadata.dimension
354        )
355    };
356
357    RetrievalBackendStatus {
358        source,
359        state: if stale {
360            RetrievalBackendState::Degraded
361        } else {
362            RetrievalBackendState::Available
363        },
364        scope_post_filter: plan.source_scope.is_some(),
365        indexed_graph_version: Some(index.indexed_graph_version),
366        reason: Some(reason),
367    }
368}
369
370impl RetrievalPlan {
371    /// Validates query text and result limits.
372    pub fn new(
373        query: impl Into<String>,
374        source_scope: Option<String>,
375        limit: usize,
376        freshness: FreshnessPolicy,
377    ) -> Result<Self, RetrievalPlanError> {
378        let query = query.into();
379        let trimmed = query.trim();
380        if trimmed.is_empty() {
381            return Err(RetrievalPlanError::EmptyQuery);
382        }
383
384        let limit = match limit {
385            1..=50 => limit,
386            0 => return Err(RetrievalPlanError::ZeroLimit),
387            _ => return Err(RetrievalPlanError::LimitTooLarge { max: 50 }),
388        };
389
390        Ok(Self {
391            query: trimmed.to_owned(),
392            source_scope,
393            limit,
394            freshness,
395        })
396    }
397}
398
399/// Retrieval planning error mapped to stable API errors by application.
400#[derive(Debug, Clone, PartialEq, Eq)]
401pub enum RetrievalPlanError {
402    EmptyQuery,
403    ZeroLimit,
404    LimitTooLarge { max: usize },
405}
406
407impl fmt::Display for RetrievalPlanError {
408    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
409        match self {
410            Self::EmptyQuery => write!(formatter, "query must not be empty"),
411            Self::ZeroLimit => write!(formatter, "limit must be greater than zero"),
412            Self::LimitTooLarge { max } => write!(formatter, "limit must be {max} or less"),
413        }
414    }
415}
416
417impl Error for RetrievalPlanError {}
418
419#[cfg(test)]
420mod tests {
421    use super::*;
422
423    #[test]
424    fn trims_query_and_preserves_retrieval_policy() {
425        let plan = RetrievalPlan::new(
426            " SQLite ",
427            Some("docs".to_owned()),
428            5,
429            FreshnessPolicy::GraphOnly,
430        )
431        .expect("plan should validate");
432
433        assert_eq!(plan.query, "SQLite");
434        assert_eq!(plan.source_scope, Some("docs".to_owned()));
435        assert_eq!(plan.limit, 5);
436        assert_eq!(plan.freshness, FreshnessPolicy::GraphOnly);
437    }
438
439    #[test]
440    fn rejects_empty_and_unbounded_queries() {
441        let empty = RetrievalPlan::new(" ", None, 1, FreshnessPolicy::AllowStale)
442            .expect_err("empty query should fail");
443        let zero = RetrievalPlan::new("x", None, 0, FreshnessPolicy::AllowStale)
444            .expect_err("zero limit should fail");
445        let too_large = RetrievalPlan::new("x", None, 51, FreshnessPolicy::AllowStale)
446            .expect_err("large limit should fail");
447
448        assert_eq!(empty.to_string(), "query must not be empty");
449        assert_eq!(zero.to_string(), "limit must be greater than zero");
450        assert_eq!(too_large.to_string(), "limit must be 50 or less");
451    }
452
453    #[test]
454    fn read_model_statuses_report_available_local_backends() {
455        let plan = RetrievalPlan::new(
456            "SQLite",
457            Some("docs".to_owned()),
458            5,
459            FreshnessPolicy::AllowStale,
460        )
461        .expect("plan should validate");
462        let statuses = read_model_backend_statuses(
463            &plan,
464            GraphVersion::new(7),
465            &[
466                IndexStatus {
467                    kind: IndexKind::Semantic,
468                    index_version: 1,
469                    indexed_graph_version: GraphVersion::new(7),
470                    state: crate::domain::IndexState::Fresh,
471                    last_error: None,
472                },
473                IndexStatus {
474                    kind: IndexKind::Vector,
475                    index_version: 1,
476                    indexed_graph_version: GraphVersion::new(7),
477                    state: crate::domain::IndexState::Fresh,
478                    last_error: None,
479                },
480            ],
481            &ReadModelBackendConfig::local(),
482        );
483
484        assert_eq!(statuses[0].state, RetrievalBackendState::Available);
485        assert_eq!(statuses[1].state, RetrievalBackendState::Available);
486        assert!(statuses.iter().all(|status| status.scope_post_filter));
487    }
488
489    #[test]
490    fn read_model_statuses_report_stale_or_disabled_backends() {
491        let plan = RetrievalPlan::new("SQLite", None, 5, FreshnessPolicy::AllowStale)
492            .expect("plan should validate");
493        let mut config = ReadModelBackendConfig::local();
494        config.vector_mode = ReadModelBackendMode::Disabled;
495
496        let statuses = read_model_backend_statuses(
497            &plan,
498            GraphVersion::new(9),
499            &[IndexStatus {
500                kind: IndexKind::Semantic,
501                index_version: 1,
502                indexed_graph_version: GraphVersion::new(8),
503                state: crate::domain::IndexState::Fresh,
504                last_error: None,
505            }],
506            &config,
507        );
508
509        assert_eq!(statuses[0].state, RetrievalBackendState::Degraded);
510        assert_eq!(statuses[1].state, RetrievalBackendState::Unavailable);
511    }
512
513    #[test]
514    fn redacted_remote_url_strips_userinfo_and_path() {
515        let config = RemoteEmbeddingConfig {
516            provider: EmbeddingProviderKind::OpenAiCompatible,
517            base_url: "https://user:pass@embeddings.example/v1".to_owned(),
518            api_key: "secret".to_owned(),
519            batch_size: DEFAULT_EMBEDDING_BATCH_SIZE,
520            timeout: DEFAULT_EMBEDDING_TIMEOUT,
521            max_concurrency: DEFAULT_EMBEDDING_MAX_CONCURRENCY,
522        };
523
524        assert_eq!(config.redacted_base_url(), "https://embeddings.example");
525    }
526}