relay_knowledge/application/service/retrieval/
mod.rs1use std::sync::Arc;
4
5use serde::Serialize;
6
7use crate::{
8 api::{ApiError, ApiMetadata, HybridRetrievalRequest, HybridRetrievalResponse, RequestContext},
9 domain::{
10 ContextGraphPath, ContextPackItem, FreshnessPolicy, FusionDiagnostics,
11 RECIPROCAL_RANK_FUSION_K, RetrievalBackendStatus, RetrievalBudgetUsed, RetrievalHit,
12 RetrievalMode, RetrievedContextPack, RetrieverSource, SourceScope,
13 },
14 retrieval::{RetrievalPlan, read_model_backend_statuses},
15 storage::{GraphSearchRequest, IndexRefreshDiagnostics, KnowledgeStore, StorageError},
16};
17
18use super::{
19 super::knowledge::index_refresh::{
20 IndexRefreshOutcome, metadata_for_indexes, refresh_index_kinds,
21 },
22 RelayKnowledgeService, current_time_millis, storage_api_error,
23};
24
25impl RelayKnowledgeService {
26 pub async fn retrieve_context(
28 &self,
29 request: HybridRetrievalRequest,
30 context: RequestContext,
31 ) -> Result<HybridRetrievalResponse, ApiError> {
32 let source_scope = normalize_optional_source_scope(request.source_scope)
33 .map_err(ApiError::invalid_argument)?;
34 let plan = RetrievalPlan::new(
35 request.query,
36 source_scope,
37 request.limit,
38 request.freshness,
39 )
40 .map_err(|error| ApiError::invalid_argument(error.to_string()))?;
41 let store = self.storage.get().await.map_err(storage_api_error)?;
42 let graph_version = store
43 .current_graph_version()
44 .await
45 .map_err(storage_api_error)?;
46
47 let mut retrieval_mode = RetrievalMode::Hybrid;
48 let mut indexes = Vec::new();
49 let mut index_cursors = Vec::new();
50 let mut index_refresh = IndexRefreshDiagnostics::default();
51 let mut metadata = ApiMetadata::graph_only(&context, graph_version);
52 let mut degraded_reasons = Vec::new();
53 let backend_statuses = if plan.freshness == FreshnessPolicy::GraphOnly {
54 retrieval_mode = RetrievalMode::GraphOnly;
55 degraded_reasons.push("graph_only freshness policy selected".to_owned());
56 Vec::new()
57 } else {
58 let mut index_outcome = retrieval_index_freshness_snapshot(&store).await?;
59 indexes = index_outcome.indexes;
60 index_cursors = index_outcome.cursors;
61 index_refresh = index_outcome.diagnostics;
62 let mut active_indexes = indexes
63 .iter()
64 .filter(|status| self.runtime.retrieval.refreshes_index(status.kind))
65 .cloned()
66 .collect::<Vec<_>>();
67 if plan.freshness == FreshnessPolicy::WaitUntilFresh {
68 let stale_kinds = active_indexes
69 .iter()
70 .filter(|status| status.is_stale_for(graph_version))
71 .map(|status| status.kind)
72 .collect::<Vec<_>>();
73 if !stale_kinds.is_empty() {
74 refresh_index_kinds(
75 &store,
76 stale_kinds,
77 graph_version,
78 &self.runtime.retrieval,
79 )
80 .await?;
81 index_outcome = retrieval_index_freshness_snapshot(&store).await?;
82 indexes = index_outcome.indexes;
83 index_cursors = index_outcome.cursors;
84 index_refresh = index_outcome.diagnostics;
85 active_indexes = indexes
86 .iter()
87 .filter(|status| self.runtime.retrieval.refreshes_index(status.kind))
88 .cloned()
89 .collect();
90 }
91 }
92
93 let stale = active_indexes
94 .iter()
95 .any(|status| status.is_stale_for(graph_version));
96 metadata = metadata_for_indexes(&context, graph_version, &active_indexes);
97 if plan.freshness == FreshnessPolicy::AllowStale && stale {
98 degraded_reasons
99 .push("one or more indexes are behind the graph version".to_owned());
100 }
101 read_model_backend_statuses(&plan, graph_version, &indexes, &self.runtime.retrieval)
102 };
103 if backend_statuses
104 .iter()
105 .any(|status| status.state == crate::domain::RetrievalBackendState::Unavailable)
106 {
107 degraded_reasons.push(
108 "semantic/vector retrieval backends unavailable; using bm25, graph evidence, and code graph fallback"
109 .to_owned(),
110 );
111 }
112 let mut disabled_retriever_sources = self.runtime.retrieval.disabled_retriever_sources();
113 if plan.freshness == FreshnessPolicy::GraphOnly {
114 for source in [RetrieverSource::Semantic, RetrieverSource::Vector] {
115 if !disabled_retriever_sources.contains(&source) {
116 disabled_retriever_sources.push(source);
117 }
118 }
119 }
120 let candidate_limit = self.runtime.retrieval.rerank.candidate_limit(plan.limit);
121 let search_outcome = store
122 .search(GraphSearchRequest {
123 query: plan.query.clone(),
124 source_scope: plan.source_scope.clone(),
125 graph_version,
126 limit: candidate_limit,
127 disabled_retriever_sources,
128 })
129 .await
130 .map_err(storage_api_error)?;
131 if let Some(reason) = search_outcome.trace.degraded_reason.clone() {
132 degraded_reasons.push(reason);
133 }
134 let (mut results, mut rerank) = self
135 .runtime
136 .retrieval
137 .rerank
138 .rerank(&plan.query, search_outcome.hits);
139 let result_truncated = results.len() > plan.limit;
140 results.truncate(plan.limit);
141 rerank.returned_count = results.len();
142 if rerank.degraded {
143 if let Some(reason) = &rerank.reason {
144 degraded_reasons.push(reason.clone());
145 }
146 }
147 let degraded_reason = (!degraded_reasons.is_empty()).then(|| degraded_reasons.join("; "));
148 let mut provenance_trace = search_outcome.trace;
149 provenance_trace.mark_citations_for_hits(results.iter());
150 provenance_trace.stale = degraded_reasons
151 .iter()
152 .any(|reason| reason.contains("behind the graph version"));
153 provenance_trace.degraded_reason = degraded_reason.clone();
154 provenance_trace.truncated |= result_truncated;
155 provenance_trace.apply_budget(plan.limit.saturating_mul(4).max(plan.limit + 8).min(64));
156 let truncated = result_truncated || provenance_trace.truncated;
157
158 let context_pack = RetrievedContextPack {
159 graph_version,
160 source_scope: plan.source_scope.clone(),
161 freshness: plan.freshness,
162 truncated,
163 backend_statuses: backend_statuses.clone(),
164 provenance_trace: Some(provenance_trace),
165 items: results
166 .iter()
167 .map(|hit| ContextPackItem {
168 result_id: hit.evidence_id.clone(),
169 source_scope: hit.source_scope.clone(),
170 source_path: hit.source_path.clone(),
171 source_span: hit.source_span,
172 entities: hit.entities.clone(),
173 graph_facts: hit.graph_facts.clone(),
174 graph_paths: hit
175 .graph_facts
176 .iter()
177 .map(ContextGraphPath::from_fact)
178 .collect(),
179 code_artifact: hit.code_artifact.clone(),
180 retriever_sources: hit.retriever_sources.clone(),
181 ranking: hit.ranking.clone(),
182 rerank: hit.rerank.clone(),
183 })
184 .collect(),
185 };
186 let budget_used = RetrievalBudgetUsed {
187 limit: plan.limit,
188 candidate_count: rerank.candidate_count,
189 returned_count: results.len(),
190 context_bytes: retrieval_context_bytes(&results, &context_pack, &backend_statuses),
191 };
192 let fusion = FusionDiagnostics {
193 algorithm: "reciprocal_rank_fusion".to_owned(),
194 k: RECIPROCAL_RANK_FUSION_K,
195 candidate_count: budget_used.candidate_count,
196 };
197
198 Ok(HybridRetrievalResponse {
199 metadata,
200 context_pack,
201 retrieval_mode,
202 source_scope: plan.source_scope,
203 freshness: plan.freshness,
204 results,
205 fusion,
206 rerank,
207 backend_statuses,
208 truncated,
209 budget_used,
210 degraded_reason,
211 indexes,
212 index_cursors,
213 index_refresh,
214 })
215 }
216}
217
218async fn retrieval_index_freshness_snapshot(
219 store: &Arc<dyn KnowledgeStore>,
220) -> Result<IndexRefreshOutcome, ApiError> {
221 let indexes = store.index_statuses().await.map_err(storage_api_error)?;
222 let cursors = match store.index_cursors().await {
223 Ok(cursors) => cursors,
224 Err(StorageError::InvalidInput(message))
225 if message == "index cursor storage is unavailable" =>
226 {
227 Vec::new()
228 }
229 Err(error) => return Err(storage_api_error(error)),
230 };
231 let diagnostics = match store.index_refresh_diagnostics(current_time_millis()).await {
232 Ok(diagnostics) => diagnostics,
233 Err(StorageError::InvalidInput(message))
234 if message == "index refresh diagnostics are unavailable" =>
235 {
236 IndexRefreshDiagnostics::default()
237 }
238 Err(error) => return Err(storage_api_error(error)),
239 };
240
241 Ok(IndexRefreshOutcome {
242 indexes,
243 cursors,
244 diagnostics,
245 })
246}
247
248fn normalize_optional_source_scope(value: Option<String>) -> Result<Option<String>, String> {
249 value
250 .map(|scope| {
251 SourceScope::parse(scope)
252 .map(String::from)
253 .map_err(|error| error.to_string())
254 })
255 .transpose()
256}
257
258fn retrieval_context_bytes(
259 results: &[RetrievalHit],
260 context_pack: &RetrievedContextPack,
261 backend_statuses: &[RetrievalBackendStatus],
262) -> usize {
263 serialized_context_bytes(&context_pack.backend_statuses)
264 .saturating_add(serialized_context_bytes(backend_statuses))
265 .saturating_add(
266 context_pack
267 .provenance_trace
268 .as_ref()
269 .map(serialized_context_bytes)
270 .unwrap_or_default(),
271 )
272 .saturating_add(results.iter().map(serialized_context_bytes).sum::<usize>())
273 .saturating_add(
274 context_pack
275 .items
276 .iter()
277 .map(serialized_context_bytes)
278 .sum::<usize>(),
279 )
280}
281
282fn serialized_context_bytes<T: Serialize + ?Sized>(value: &T) -> usize {
283 serde_json::to_vec(value)
284 .map(|bytes| bytes.len())
285 .unwrap_or(usize::MAX / 4)
286}
287
288#[cfg(test)]
289#[path = "mod_tests.rs"]
290mod tests;