1use 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#[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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
40pub enum ReadModelBackendMode {
41 Local,
42 External,
43 Disabled,
44}
45
46#[derive(Debug, Clone, Copy, PartialEq, Eq)]
48pub enum EmbeddingProviderKind {
49 OpenAiCompatible,
50 Echo,
51}
52
53impl EmbeddingProviderKind {
54 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 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#[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 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 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#[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#[derive(Debug, Clone, PartialEq, Eq)]
135pub struct ReadModelMetadata {
136 pub name: String,
137 pub dimension: u32,
138}
139
140#[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 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 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 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#[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 pub fn redacted_base_url(&self) -> String {
199 redacted_url(&self.base_url)
200 }
201}
202
203#[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 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 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 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
277pub 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 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#[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)]
420#[path = "mod_tests.rs"]
421mod tests;