1use std::collections::HashMap;
2
3use oxinbox_core::{Task, Uuid};
4use serde::{Deserialize, Serialize};
5use tracing::instrument;
6
7#[derive(Debug, Clone, Serialize, Deserialize)]
8pub struct SearchResult {
9 pub task: Task,
10 pub score: f64,
11}
12
13fn bm25_text(task: &Task) -> String {
14 let mut parts = vec![task.description.clone()];
15 for p in &task.projects {
16 parts.push(format!("+{p}"));
17 }
18 for c in &task.contexts {
19 parts.push(format!("@{c}"));
20 }
21 if let Some(p) = task.priority {
22 parts.push(format!("({p})"));
23 }
24 parts.join(" ")
25}
26
27#[derive(Default)]
28pub struct SearchIndex {
29 pub bm25: Bm25Index,
30 vectors: HashMap<Uuid, Vec<f32>>,
31}
32
33impl SearchIndex {
34 pub fn index_task(&mut self, task: &Task) {
35 self.bm25.index(task.id, &bm25_text(task));
36 }
37
38 pub fn remove_task(&mut self, id: Uuid) {
39 self.bm25.remove(id);
40 self.vectors.remove(&id);
41 }
42
43 pub fn store_embedding(&mut self, task_id: Uuid, embedding: Vec<f32>) {
44 self.vectors.insert(task_id, embedding);
45 }
46}
47
48#[derive(Default)]
49pub struct Bm25Index {
50 doc_count: usize,
51 doc_freq: HashMap<String, usize>,
52 term_counts: HashMap<Uuid, HashMap<String, usize>>,
53 doc_lengths: HashMap<Uuid, usize>,
54}
55
56impl Bm25Index {
57 const K1: f64 = 1.5;
58 const B: f64 = 0.75;
59
60 fn tokenize(text: &str) -> Vec<String> {
61 text.to_lowercase()
62 .split_whitespace()
63 .filter(|t| t.len() >= 2)
64 .map(|t| {
65 t.trim_matches(|c: char| c.is_ascii_punctuation())
66 .to_string()
67 })
68 .filter(|t| !t.is_empty())
69 .collect()
70 }
71
72 fn index(&mut self, id: Uuid, text: &str) {
73 self.remove(id);
74
75 let tokens = Self::tokenize(text);
76 let len = tokens.len();
77
78 let mut counts: HashMap<String, usize> = HashMap::new();
79 for token in &tokens {
80 *counts.entry(token.clone()).or_default() += 1;
81 }
82
83 for token in counts.keys() {
84 *self.doc_freq.entry(token.clone()).or_default() += 1;
85 }
86
87 self.term_counts.insert(id, counts);
88 self.doc_lengths.insert(id, len);
89 self.doc_count += 1;
90 }
91
92 fn remove(&mut self, id: Uuid) {
93 if let Some(counts) = self.term_counts.remove(&id) {
94 for token in counts.keys() {
95 if let Some(freq) = self.doc_freq.get_mut(token) {
96 *freq = freq.saturating_sub(1);
97 if *freq == 0 {
98 self.doc_freq.remove(token);
99 }
100 }
101 }
102 self.doc_lengths.remove(&id);
103 self.doc_count = self.doc_count.saturating_sub(1);
104 }
105 }
106
107 #[allow(
108 clippy::cast_precision_loss,
109 clippy::imprecise_flops,
110 clippy::suboptimal_flops
111 )]
112 fn score(&self, id: Uuid, query_tokens: &[String]) -> f64 {
113 let Some(counts) = self.term_counts.get(&id) else {
114 return 0.0;
115 };
116 let Some(&doc_len) = self.doc_lengths.get(&id) else {
117 return 0.0;
118 };
119 let avg_dl =
120 self.doc_lengths.values().copied().sum::<usize>() as f64 / self.doc_count.max(1) as f64;
121
122 let mut score = 0.0;
123 for token in query_tokens {
124 let tf = *counts.get(token).unwrap_or(&0) as f64;
125 let df = *self.doc_freq.get(token).unwrap_or(&1) as f64;
126 let idf = ((self.doc_count as f64 - df + 0.5) / (df + 0.5)).ln_1p();
127 let numerator = tf * (Self::K1 + 1.0);
128 let denominator =
129 Self::K1.mul_add(1.0 - Self::B + Self::B * doc_len as f64 / avg_dl, tf);
130 score += idf * numerator / denominator;
131 }
132 score
133 }
134}
135
136pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f64 {
137 let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
138 let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
139 let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
140 if norm_a == 0.0 || norm_b == 0.0 {
141 return 0.0;
142 }
143 f64::from(dot / (norm_a * norm_b))
144}
145
146#[instrument(skip(index))]
147pub fn hybrid_search(
148 tasks: &[Task],
149 index: &SearchIndex,
150 query_text: &str,
151 query_embedding: Option<&[f32]>,
152 limit: usize,
153 alpha: f64,
154) -> Vec<SearchResult> {
155 let query_tokens = Bm25Index::tokenize(query_text);
156
157 let bm25_scores: HashMap<Uuid, f64> = tasks
158 .iter()
159 .map(|t| (t.id, index.bm25.score(t.id, &query_tokens)))
160 .collect();
161
162 let bm25_max = bm25_scores.values().copied().fold(0.0_f64, f64::max);
163 let bm25_min = bm25_scores.values().copied().fold(f64::MAX, f64::min);
164
165 let vector_scores: HashMap<Uuid, f64> = query_embedding.map_or_else(HashMap::new, |qv| {
166 tasks
167 .iter()
168 .map(|t| {
169 let sim = index
170 .vectors
171 .get(&t.id)
172 .map_or(0.0, |ev| cosine_similarity(qv, ev));
173 (t.id, sim)
174 })
175 .collect()
176 });
177
178 let vec_max = vector_scores.values().copied().fold(0.0_f64, f64::max);
179 let vec_min = vector_scores.values().copied().fold(f64::MAX, f64::min);
180
181 let mut results: Vec<SearchResult> = tasks
182 .iter()
183 .map(|task| {
184 let bm25_norm = if bm25_max > bm25_min {
185 (bm25_scores[&task.id] - bm25_min) / (bm25_max - bm25_min)
186 } else {
187 bm25_scores[&task.id]
188 };
189
190 let vec_norm = if vec_max > vec_min {
191 (vector_scores.get(&task.id).copied().unwrap_or(0.0) - vec_min)
192 / (vec_max - vec_min)
193 } else {
194 vector_scores.get(&task.id).copied().unwrap_or(0.0)
195 };
196
197 let score = alpha.mul_add(bm25_norm, (1.0 - alpha) * vec_norm);
198
199 SearchResult {
200 task: task.clone(),
201 score,
202 }
203 })
204 .filter(|r| r.score > 0.0)
205 .collect();
206
207 results.sort_by(|a, b| {
208 b.score
209 .partial_cmp(&a.score)
210 .unwrap_or(std::cmp::Ordering::Equal)
211 });
212 results.truncate(limit);
213 results
214}
215
216#[cfg(test)]
217mod tests {
218 use super::*;
219 use chrono::Utc;
220
221 fn make_task(id: Uuid, desc: &str, projects: &[&str], contexts: &[&str]) -> Task {
222 Task {
223 id,
224 completed: false,
225 priority: None,
226 description: desc.into(),
227 projects: projects.iter().map(ToString::to_string).collect(),
228 contexts: contexts.iter().map(ToString::to_string).collect(),
229 status: oxinbox_core::TaskStatus::Inbox,
230 created_at: Utc::now(),
231 updated_at: Utc::now(),
232 completed_at: None,
233 due_date: None,
234 }
235 }
236
237 #[test]
238 fn bm25_ranks_relevant_higher() {
239 let t1 = make_task(Uuid::now_v7(), "buy milk from the store", &[], &[]);
240 let t2 = make_task(Uuid::now_v7(), "fix database indexing bug", &[], &[]);
241
242 let mut idx = Bm25Index::default();
243 idx.index(t1.id, &bm25_text(&t1));
244 idx.index(t2.id, &bm25_text(&t2));
245
246 let tokens = Bm25Index::tokenize("milk store");
247 let s1 = idx.score(t1.id, &tokens);
248 let s2 = idx.score(t2.id, &tokens);
249
250 assert!(
251 s1 > s2,
252 "BM25 should rank 'milk store' higher for the milk task"
253 );
254 }
255
256 #[test]
257 fn hybrid_search_returns_results() {
258 let t1 = make_task(Uuid::now_v7(), "buy milk", &["proyecto"], &["casa"]);
259 let t2 = make_task(Uuid::now_v7(), "fix database", &[], &[]);
260 let tasks = vec![t1.clone(), t2.clone()];
261
262 let mut bm25 = Bm25Index::default();
263 bm25.index(t1.id, &bm25_text(&t1));
264 bm25.index(t2.id, &bm25_text(&t2));
265
266 let mut vs = SearchIndex::default();
267 vs.index_task(&t1);
268 vs.index_task(&t2);
269
270 let results = hybrid_search(&tasks, &vs, "milk", None, 10, 1.0);
271 assert!(!results.is_empty());
272 assert_eq!(results[0].task.description, "buy milk");
273 }
274
275 #[test]
276 fn cosine_similarity_identical() {
277 let a = vec![1.0, 2.0, 3.0];
278 let b = vec![1.0, 2.0, 3.0];
279 let sim = cosine_similarity(&a, &b);
280 assert!((sim - 1.0).abs() < 1e-6);
281 }
282
283 #[test]
284 fn cosine_similarity_orthogonal() {
285 let a = vec![1.0, 0.0];
286 let b = vec![0.0, 1.0];
287 let sim = cosine_similarity(&a, &b);
288 assert!((sim - 0.0).abs() < 1e-6);
289 }
290}