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 let (mut results, mut rerank) = self
132 .runtime
133 .retrieval
134 .rerank
135 .rerank(&plan.query, search_outcome.hits);
136 let result_truncated = results.len() > plan.limit;
137 results.truncate(plan.limit);
138 rerank.returned_count = results.len();
139 if rerank.degraded {
140 if let Some(reason) = &rerank.reason {
141 degraded_reasons.push(reason.clone());
142 }
143 }
144 let degraded_reason = (!degraded_reasons.is_empty()).then(|| degraded_reasons.join("; "));
145 let mut provenance_trace = search_outcome.trace;
146 provenance_trace.mark_citations_for_hits(results.iter());
147 provenance_trace.stale = degraded_reasons
148 .iter()
149 .any(|reason| reason.contains("behind the graph version"));
150 provenance_trace.degraded_reason = degraded_reason.clone();
151 provenance_trace.truncated |= result_truncated;
152 provenance_trace.apply_budget(plan.limit.saturating_mul(4).max(plan.limit + 8).min(64));
153 let truncated = result_truncated || provenance_trace.truncated;
154
155 let context_pack = RetrievedContextPack {
156 graph_version,
157 source_scope: plan.source_scope.clone(),
158 freshness: plan.freshness,
159 truncated,
160 backend_statuses: backend_statuses.clone(),
161 provenance_trace: Some(provenance_trace),
162 items: results
163 .iter()
164 .map(|hit| ContextPackItem {
165 result_id: hit.evidence_id.clone(),
166 source_scope: hit.source_scope.clone(),
167 source_path: hit.source_path.clone(),
168 source_span: hit.source_span,
169 entities: hit.entities.clone(),
170 graph_facts: hit.graph_facts.clone(),
171 graph_paths: hit
172 .graph_facts
173 .iter()
174 .map(ContextGraphPath::from_fact)
175 .collect(),
176 code_artifact: hit.code_artifact.clone(),
177 retriever_sources: hit.retriever_sources.clone(),
178 ranking: hit.ranking.clone(),
179 rerank: hit.rerank.clone(),
180 })
181 .collect(),
182 };
183 let budget_used = RetrievalBudgetUsed {
184 limit: plan.limit,
185 candidate_count: rerank.candidate_count,
186 returned_count: results.len(),
187 context_bytes: retrieval_context_bytes(&results, &context_pack, &backend_statuses),
188 };
189 let fusion = FusionDiagnostics {
190 algorithm: "reciprocal_rank_fusion".to_owned(),
191 k: RECIPROCAL_RANK_FUSION_K,
192 candidate_count: budget_used.candidate_count,
193 };
194
195 Ok(HybridRetrievalResponse {
196 metadata,
197 context_pack,
198 retrieval_mode,
199 source_scope: plan.source_scope,
200 freshness: plan.freshness,
201 results,
202 fusion,
203 rerank,
204 backend_statuses,
205 truncated,
206 budget_used,
207 degraded_reason,
208 indexes,
209 index_cursors,
210 index_refresh,
211 })
212 }
213}
214
215async fn retrieval_index_freshness_snapshot(
216 store: &Arc<dyn KnowledgeStore>,
217) -> Result<IndexRefreshOutcome, ApiError> {
218 let indexes = store.index_statuses().await.map_err(storage_api_error)?;
219 let cursors = match store.index_cursors().await {
220 Ok(cursors) => cursors,
221 Err(StorageError::InvalidInput(message))
222 if message == "index cursor storage is unavailable" =>
223 {
224 Vec::new()
225 }
226 Err(error) => return Err(storage_api_error(error)),
227 };
228 let diagnostics = match store.index_refresh_diagnostics(current_time_millis()).await {
229 Ok(diagnostics) => diagnostics,
230 Err(StorageError::InvalidInput(message))
231 if message == "index refresh diagnostics are unavailable" =>
232 {
233 IndexRefreshDiagnostics::default()
234 }
235 Err(error) => return Err(storage_api_error(error)),
236 };
237
238 Ok(IndexRefreshOutcome {
239 indexes,
240 cursors,
241 diagnostics,
242 })
243}
244
245fn normalize_optional_source_scope(value: Option<String>) -> Result<Option<String>, String> {
246 value
247 .map(|scope| {
248 SourceScope::parse(scope)
249 .map(String::from)
250 .map_err(|error| error.to_string())
251 })
252 .transpose()
253}
254
255fn retrieval_context_bytes(
256 results: &[RetrievalHit],
257 context_pack: &RetrievedContextPack,
258 backend_statuses: &[RetrievalBackendStatus],
259) -> usize {
260 serialized_context_bytes(&context_pack.backend_statuses)
261 .saturating_add(serialized_context_bytes(backend_statuses))
262 .saturating_add(
263 context_pack
264 .provenance_trace
265 .as_ref()
266 .map(serialized_context_bytes)
267 .unwrap_or_default(),
268 )
269 .saturating_add(results.iter().map(serialized_context_bytes).sum::<usize>())
270 .saturating_add(
271 context_pack
272 .items
273 .iter()
274 .map(serialized_context_bytes)
275 .sum::<usize>(),
276 )
277}
278
279fn serialized_context_bytes<T: Serialize + ?Sized>(value: &T) -> usize {
280 serde_json::to_vec(value)
281 .map(|bytes| bytes.len())
282 .unwrap_or(usize::MAX / 4)
283}
284
285#[cfg(test)]
286#[path = "mod_tests.rs"]
287mod tests;