1use std::{collections::BTreeMap, error::Error, fmt};
2
3use sim_lib_numbers_stats::{bayesian_update_binary, entropy};
4use sim_lib_rank::{EmbeddingStore, retrieve_ids};
5
6const DEFAULT_RADAR_LIMIT: usize = 8;
7const DEFAULT_QUERY_TEXT: &str = "atelier radar";
8
9#[derive(Clone, Debug, PartialEq, Eq)]
11pub struct SourceSpan {
12 pub repo: String,
14 pub path: String,
16 pub line: usize,
18}
19
20#[derive(Clone, Debug, PartialEq, Eq)]
22pub struct RadarQuery {
23 pub text: String,
25 pub repo: Option<String>,
27 pub crate_name: Option<String>,
29 pub kind: Option<String>,
31 pub capability: Option<String>,
33 pub codec: Option<String>,
35 pub agent_role: Option<String>,
37 pub limit: usize,
39}
40
41impl RadarQuery {
42 pub fn new(text: impl Into<String>) -> Self {
44 Self {
45 text: text.into(),
46 repo: None,
47 crate_name: None,
48 kind: None,
49 capability: None,
50 codec: None,
51 agent_role: None,
52 limit: DEFAULT_RADAR_LIMIT,
53 }
54 }
55
56 fn search_text(&self) -> String {
57 [
58 self.text.as_str(),
59 self.repo.as_deref().unwrap_or_default(),
60 self.crate_name.as_deref().unwrap_or_default(),
61 self.kind.as_deref().unwrap_or_default(),
62 self.capability.as_deref().unwrap_or_default(),
63 self.codec.as_deref().unwrap_or_default(),
64 self.agent_role.as_deref().unwrap_or_default(),
65 ]
66 .into_iter()
67 .filter(|part| !part.is_empty())
68 .collect::<Vec<_>>()
69 .join(" ")
70 }
71}
72
73#[derive(Clone, Debug, PartialEq, Eq)]
75pub struct RadarChunk {
76 pub chunk_id: String,
78 pub title: String,
80 pub span: SourceSpan,
82 pub crate_name: Option<String>,
84 pub kind: String,
86 pub text: String,
88 pub capabilities: Vec<String>,
90 pub codecs: Vec<String>,
92 pub agent_roles: Vec<String>,
94 pub live: bool,
96}
97
98impl RadarChunk {
99 pub fn new(
101 chunk_id: impl Into<String>,
102 title: impl Into<String>,
103 span: SourceSpan,
104 kind: impl Into<String>,
105 text: impl Into<String>,
106 ) -> Self {
107 Self {
108 chunk_id: chunk_id.into(),
109 title: title.into(),
110 span,
111 crate_name: None,
112 kind: kind.into(),
113 text: text.into(),
114 capabilities: Vec::new(),
115 codecs: Vec::new(),
116 agent_roles: Vec::new(),
117 live: true,
118 }
119 }
120
121 fn matches(&self, query: &RadarQuery) -> bool {
122 matches_field(&self.span.repo, &query.repo)
123 && matches_optional(self.crate_name.as_deref(), &query.crate_name)
124 && matches_field(&self.kind, &query.kind)
125 && matches_list(&self.capabilities, &query.capability)
126 && matches_list(&self.codecs, &query.codec)
127 && matches_list_or_text(&self.agent_roles, &self.text, &query.agent_role)
128 }
129
130 fn search_text(&self) -> String {
131 [
132 self.title.as_str(),
133 self.kind.as_str(),
134 self.crate_name.as_deref().unwrap_or_default(),
135 self.text.as_str(),
136 &self.capabilities.join(" "),
137 &self.codecs.join(" "),
138 &self.agent_roles.join(" "),
139 ]
140 .into_iter()
141 .filter(|part| !part.is_empty())
142 .collect::<Vec<_>>()
143 .join(" ")
144 }
145}
146
147#[derive(Clone, Debug, Default, PartialEq, Eq)]
149pub struct RadarIndex {
150 pub chunks: Vec<RadarChunk>,
152}
153
154#[derive(Clone, Debug, PartialEq)]
156pub struct RadarHint {
157 pub chunk_id: String,
159 pub title: String,
161 pub span: SourceSpan,
163 pub capabilities: Vec<String>,
165 pub preferred_codec: Option<String>,
167 pub confidence: f64,
169}
170
171#[derive(Clone, Debug, PartialEq)]
173pub struct RadarReport {
174 pub hints: Vec<RadarHint>,
176 pub stale_index: bool,
178 pub stale_chunk_ids: Vec<String>,
180}
181
182pub type RadarResult<T> = Result<T, RadarError>;
184
185#[derive(Clone, Debug, PartialEq, Eq)]
187pub struct RadarError {
188 message: String,
189}
190
191impl RadarError {
192 fn new(message: impl Into<String>) -> Self {
193 Self {
194 message: message.into(),
195 }
196 }
197}
198
199impl fmt::Display for RadarError {
200 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
201 f.write_str(&self.message)
202 }
203}
204
205impl Error for RadarError {}
206
207pub fn retrieve_radar_hints(index: &RadarIndex, query: &RadarQuery) -> RadarResult<RadarReport> {
209 let matching = index
210 .chunks
211 .iter()
212 .filter(|chunk| chunk.matches(query))
213 .collect::<Vec<_>>();
214 let stale_chunk_ids = matching
215 .iter()
216 .filter(|chunk| !chunk.live)
217 .map(|chunk| chunk.chunk_id.clone())
218 .collect::<Vec<_>>();
219 let live_chunks = matching
220 .into_iter()
221 .filter(|chunk| chunk.live)
222 .collect::<Vec<_>>();
223 if live_chunks.is_empty() {
224 return Ok(RadarReport {
225 hints: Vec::new(),
226 stale_index: !stale_chunk_ids.is_empty(),
227 stale_chunk_ids,
228 });
229 }
230
231 let mut store = EmbeddingStore::new();
232 let mut by_id = BTreeMap::new();
233 for chunk in &live_chunks {
234 store
235 .insert(chunk.chunk_id.clone(), embedding_vec(&chunk.search_text()))
236 .map_err(|err| RadarError::new(format!("rank index insert failed: {err}")))?;
237 by_id.insert(chunk.chunk_id.as_str(), *chunk);
238 }
239
240 let query_embedding = embedding_vec(&query.search_text());
241 let neighbors = retrieve_ids(
242 &store,
243 &query_embedding,
244 live_chunks.iter().map(|chunk| chunk.chunk_id.as_str()),
245 query.limit,
246 )
247 .map_err(|err| RadarError::new(format!("rank retrieve failed: {err}")))?;
248
249 let mut hints = Vec::new();
250 for neighbor in neighbors {
251 let Some(chunk) = by_id.get(neighbor.id.as_str()) else {
252 continue;
253 };
254 hints.push(RadarHint {
255 chunk_id: chunk.chunk_id.clone(),
256 title: chunk.title.clone(),
257 span: chunk.span.clone(),
258 capabilities: chunk.capabilities.clone(),
259 preferred_codec: preferred_codec(chunk, query),
260 confidence: confidence_from_score(neighbor.score)?,
261 });
262 }
263
264 Ok(RadarReport {
265 hints,
266 stale_index: !stale_chunk_ids.is_empty(),
267 stale_chunk_ids,
268 })
269}
270
271fn embedding_vec(text: &str) -> Vec<f32> {
272 let text = if text
273 .split(|ch: char| !ch.is_alphanumeric())
274 .any(|token| !token.is_empty())
275 {
276 text
277 } else {
278 DEFAULT_QUERY_TEXT
279 };
280 crate::embed(text).to_vec()
281}
282
283fn confidence_from_score(score: f32) -> RadarResult<f64> {
284 let likelihood = ((score as f64 + 1.0) / 2.0).clamp(0.0, 1.0);
285 let posterior = bayesian_update_binary(0.5, likelihood, 0.25)
286 .map_err(|err| RadarError::new(format!("confidence update failed: {err}")))?;
287 let uncertainty = entropy(&[posterior, 1.0 - posterior])
288 .map_err(|err| RadarError::new(format!("confidence entropy failed: {err}")))?;
289 Ok((posterior * 0.85 + (1.0 - uncertainty) * 0.15).clamp(0.0, 1.0))
290}
291
292fn preferred_codec(chunk: &RadarChunk, query: &RadarQuery) -> Option<String> {
293 query
294 .codec
295 .clone()
296 .or_else(|| chunk.codecs.first().cloned())
297}
298
299fn matches_field(value: &str, filter: &Option<String>) -> bool {
300 filter
301 .as_deref()
302 .is_none_or(|filter| value.eq_ignore_ascii_case(filter))
303}
304
305fn matches_optional(value: Option<&str>, filter: &Option<String>) -> bool {
306 filter
307 .as_deref()
308 .is_none_or(|filter| value.is_some_and(|value| value.eq_ignore_ascii_case(filter)))
309}
310
311fn matches_list(values: &[String], filter: &Option<String>) -> bool {
312 filter.as_deref().is_none_or(|filter| {
313 values
314 .iter()
315 .any(|value| value.eq_ignore_ascii_case(filter))
316 })
317}
318
319fn matches_list_or_text(values: &[String], text: &str, filter: &Option<String>) -> bool {
320 filter.as_deref().is_none_or(|filter| {
321 values
322 .iter()
323 .any(|value| value.eq_ignore_ascii_case(filter))
324 || text
325 .to_ascii_lowercase()
326 .contains(&filter.to_ascii_lowercase())
327 })
328}