1use uuid::Uuid;
10
11use khive_fold::objective::{Objective, ObjectiveContext};
12use khive_fold::ordering::HasId;
13
14#[derive(Debug, Clone)]
20pub struct RetrievalCandidate {
21 pub id: Uuid,
23 pub vector_score: Option<f64>,
25 pub text_score: Option<f64>,
27 pub graph_distance: Option<u32>,
29 pub rrf_score: Option<f64>,
31}
32
33impl HasId for RetrievalCandidate {
34 #[inline]
35 fn id(&self) -> Uuid {
36 self.id
37 }
38}
39
40pub struct VectorSimilarityObjective;
46
47impl Objective<RetrievalCandidate> for VectorSimilarityObjective {
48 #[inline]
49 fn score(&self, candidate: &RetrievalCandidate, _context: &ObjectiveContext) -> f64 {
50 candidate.vector_score.unwrap_or(0.0)
51 }
52
53 fn name(&self) -> &str {
54 "VectorSimilarityObjective"
55 }
56}
57
58pub struct TextRelevanceObjective;
64
65impl Objective<RetrievalCandidate> for TextRelevanceObjective {
66 #[inline]
67 fn score(&self, candidate: &RetrievalCandidate, _context: &ObjectiveContext) -> f64 {
68 candidate.text_score.unwrap_or(0.0)
69 }
70
71 fn name(&self) -> &str {
72 "TextRelevanceObjective"
73 }
74}
75
76pub struct GraphProximityObjective {
91 pub max_distance: u32,
93}
94
95impl Objective<RetrievalCandidate> for GraphProximityObjective {
96 fn score(&self, candidate: &RetrievalCandidate, _context: &ObjectiveContext) -> f64 {
97 let d = match candidate.graph_distance {
98 Some(d) => d,
99 None => return 0.0,
100 };
101 if self.max_distance == 0 || d >= self.max_distance {
102 return 0.0;
103 }
104 1.0 - (d as f64 / self.max_distance as f64)
105 }
106
107 fn name(&self) -> &str {
108 "GraphProximityObjective"
109 }
110}
111
112pub struct RrfFusionObjective;
118
119impl Objective<RetrievalCandidate> for RrfFusionObjective {
120 #[inline]
121 fn score(&self, candidate: &RetrievalCandidate, _context: &ObjectiveContext) -> f64 {
122 candidate.rrf_score.unwrap_or(0.0)
123 }
124
125 fn name(&self) -> &str {
126 "RrfFusionObjective"
127 }
128}
129
130#[cfg(test)]
133mod tests {
134 use super::*;
135 use khive_fold::objective::{Objective, ObjectiveContext};
136 use khive_fold::WeightedObjective;
137 use uuid::Uuid;
138
139 fn ctx() -> ObjectiveContext {
140 ObjectiveContext::new()
141 }
142
143 fn candidate(
144 vector: Option<f64>,
145 text: Option<f64>,
146 dist: Option<u32>,
147 rrf: Option<f64>,
148 ) -> RetrievalCandidate {
149 RetrievalCandidate {
150 id: Uuid::new_v4(),
151 vector_score: vector,
152 text_score: text,
153 graph_distance: dist,
154 rrf_score: rrf,
155 }
156 }
157
158 #[test]
161 fn vector_present_returns_signal() {
162 let c = candidate(Some(0.85), None, None, None);
163 let score = VectorSimilarityObjective.score(&c, &ctx());
164 assert!((score - 0.85).abs() < 1e-12);
165 }
166
167 #[test]
168 fn vector_absent_returns_zero() {
169 let c = candidate(None, None, None, None);
170 assert_eq!(VectorSimilarityObjective.score(&c, &ctx()), 0.0);
171 }
172
173 #[test]
174 fn vector_zero_score_returns_zero() {
175 let c = candidate(Some(0.0), None, None, None);
176 assert_eq!(VectorSimilarityObjective.score(&c, &ctx()), 0.0);
177 }
178
179 #[test]
182 fn text_present_returns_signal() {
183 let c = candidate(None, Some(0.6), None, None);
184 let score = TextRelevanceObjective.score(&c, &ctx());
185 assert!((score - 0.6).abs() < 1e-12);
186 }
187
188 #[test]
189 fn text_absent_returns_zero() {
190 let c = candidate(None, None, None, None);
191 assert_eq!(TextRelevanceObjective.score(&c, &ctx()), 0.0);
192 }
193
194 #[test]
197 fn graph_anchor_hit_scores_one() {
198 let c = candidate(None, None, Some(0), None);
200 let obj = GraphProximityObjective { max_distance: 3 };
201 assert!((obj.score(&c, &ctx()) - 1.0).abs() < 1e-12);
202 }
203
204 #[test]
205 fn graph_midpoint_scores_half() {
206 let c = candidate(None, None, Some(1), None);
208 let obj = GraphProximityObjective { max_distance: 2 };
209 assert!((obj.score(&c, &ctx()) - 0.5).abs() < 1e-12);
210 }
211
212 #[test]
213 fn graph_at_boundary_scores_zero() {
214 let c = candidate(None, None, Some(3), None);
216 let obj = GraphProximityObjective { max_distance: 3 };
217 assert_eq!(obj.score(&c, &ctx()), 0.0);
218 }
219
220 #[test]
221 fn graph_beyond_boundary_scores_zero() {
222 let c = candidate(None, None, Some(10), None);
223 let obj = GraphProximityObjective { max_distance: 3 };
224 assert_eq!(obj.score(&c, &ctx()), 0.0);
225 }
226
227 #[test]
228 fn graph_absent_scores_zero() {
229 let c = candidate(None, None, None, None);
230 let obj = GraphProximityObjective { max_distance: 3 };
231 assert_eq!(obj.score(&c, &ctx()), 0.0);
232 }
233
234 #[test]
235 fn graph_max_distance_zero_always_scores_zero() {
236 let c = candidate(None, None, Some(0), None);
238 let obj = GraphProximityObjective { max_distance: 0 };
239 assert_eq!(obj.score(&c, &ctx()), 0.0);
240 }
241
242 #[test]
245 fn rrf_present_returns_signal() {
246 let c = candidate(None, None, None, Some(0.0327));
247 let score = RrfFusionObjective.score(&c, &ctx());
248 assert!((score - 0.0327).abs() < 1e-12);
249 }
250
251 #[test]
252 fn rrf_absent_returns_zero() {
253 let c = candidate(None, None, None, None);
254 assert_eq!(RrfFusionObjective.score(&c, &ctx()), 0.0);
255 }
256
257 #[test]
260 fn weighted_composition_vector_and_text() {
261 let c = candidate(Some(0.8), Some(0.6), None, None);
264
265 let obj = WeightedObjective::<RetrievalCandidate>::new()
266 .add(Box::new(VectorSimilarityObjective), 0.5)
267 .add(Box::new(TextRelevanceObjective), 0.5);
268
269 let score = obj.score(&c, &ctx());
270 assert!((score - 0.7).abs() < 1e-12);
272 }
273
274 #[test]
275 fn weighted_composition_with_graph() {
276 let c = candidate(Some(1.0), Some(0.0), Some(1), None);
280
281 let obj = WeightedObjective::<RetrievalCandidate>::new()
282 .add(Box::new(VectorSimilarityObjective), 0.4)
283 .add(Box::new(TextRelevanceObjective), 0.3)
284 .add(Box::new(GraphProximityObjective { max_distance: 4 }), 0.3);
285
286 let score = obj.score(&c, &ctx());
287 assert!((score - 0.625).abs() < 1e-12);
288 }
289
290 #[test]
291 fn weighted_all_absent_returns_zero() {
292 let c = candidate(None, None, None, None);
293
294 let obj = WeightedObjective::<RetrievalCandidate>::new()
295 .add(Box::new(VectorSimilarityObjective), 0.5)
296 .add(Box::new(TextRelevanceObjective), 0.5);
297
298 assert_eq!(obj.score(&c, &ctx()), 0.0);
300 }
301
302 #[test]
305 fn has_id_returns_candidate_uuid() {
306 let id = Uuid::new_v4();
307 let c = RetrievalCandidate {
308 id,
309 vector_score: None,
310 text_score: None,
311 graph_distance: None,
312 rrf_score: None,
313 };
314 assert_eq!(c.id(), id);
315 }
316
317 #[test]
320 fn select_top_orders_by_vector_score() {
321 use khive_fold::DeterministicObjective;
322
323 let candidates = vec![
324 candidate(Some(0.3), None, None, None),
325 candidate(Some(0.9), None, None, None),
326 candidate(Some(0.6), None, None, None),
327 ];
328
329 let top = VectorSimilarityObjective.select_top_deterministic(&candidates, 2, &ctx());
330
331 assert_eq!(top.len(), 2);
332 assert!((top[0].score - 0.9).abs() < 1e-12);
333 assert!((top[1].score - 0.6).abs() < 1e-12);
334 }
335}