1use std::collections::HashMap;
16
17#[derive(Debug, Clone)]
23pub struct RemoteNode {
24 pub id: String,
26 pub weight: f64,
28 pub last_latency_ms: u64,
30 pub success_rate: f64,
32}
33
34impl RemoteNode {
35 pub fn new(
37 id: impl Into<String>,
38 weight: f64,
39 last_latency_ms: u64,
40 success_rate: f64,
41 ) -> Self {
42 Self {
43 id: id.into(),
44 weight,
45 last_latency_ms,
46 success_rate,
47 }
48 }
49
50 pub fn effective_weight(&self) -> f64 {
56 self.weight * self.success_rate / (1.0 + self.last_latency_ms as f64 / 1000.0)
57 }
58}
59
60#[derive(Debug, Clone)]
66pub struct RemoteResult {
67 pub cid: String,
69 pub score: f64,
71 pub node_id: String,
73 pub metadata: Vec<(String, String)>,
75}
76
77impl RemoteResult {
78 pub fn new(
80 cid: impl Into<String>,
81 score: f64,
82 node_id: impl Into<String>,
83 metadata: Vec<(String, String)>,
84 ) -> Self {
85 Self {
86 cid: cid.into(),
87 score,
88 node_id: node_id.into(),
89 metadata,
90 }
91 }
92}
93
94#[derive(Debug, Clone)]
100pub struct FederatedQuery {
101 pub query_embedding: Vec<f64>,
103 pub top_k: usize,
105 pub min_score: f64,
107 pub required_quorum: usize,
109}
110
111impl Default for FederatedQuery {
112 fn default() -> Self {
113 Self {
114 query_embedding: Vec::new(),
115 top_k: 10,
116 min_score: 0.0,
117 required_quorum: 1,
118 }
119 }
120}
121
122impl FederatedQuery {
123 pub fn new(query_embedding: Vec<f64>) -> Self {
125 Self {
126 query_embedding,
127 ..Self::default()
128 }
129 }
130
131 pub fn with_top_k(mut self, top_k: usize) -> Self {
133 self.top_k = top_k;
134 self
135 }
136
137 pub fn with_min_score(mut self, min_score: f64) -> Self {
139 self.min_score = min_score;
140 self
141 }
142
143 pub fn with_quorum(mut self, required_quorum: usize) -> Self {
145 self.required_quorum = required_quorum;
146 self
147 }
148}
149
150#[derive(Debug, Clone)]
156pub struct NodeResponse {
157 pub node_id: String,
159 pub results: Vec<RemoteResult>,
161 pub latency_ms: u64,
163 pub success: bool,
165}
166
167impl NodeResponse {
168 pub fn success(
170 node_id: impl Into<String>,
171 results: Vec<RemoteResult>,
172 latency_ms: u64,
173 ) -> Self {
174 Self {
175 node_id: node_id.into(),
176 results,
177 latency_ms,
178 success: true,
179 }
180 }
181
182 pub fn failure(node_id: impl Into<String>, latency_ms: u64) -> Self {
184 Self {
185 node_id: node_id.into(),
186 results: Vec::new(),
187 latency_ms,
188 success: false,
189 }
190 }
191}
192
193#[derive(Debug, Clone)]
199pub enum MergeStrategy {
200 SimpleUnion,
202
203 WeightedMerge,
207
208 QuorumIntersect(usize),
211
212 RankFusion,
216}
217
218#[derive(Debug, Clone)]
224pub struct FederatedResult {
225 pub cid: String,
227 pub merged_score: f64,
229 pub contributing_nodes: Vec<String>,
231 pub appearance_count: usize,
233}
234
235impl FederatedResult {
236 pub fn new(
238 cid: impl Into<String>,
239 merged_score: f64,
240 contributing_nodes: Vec<String>,
241 appearance_count: usize,
242 ) -> Self {
243 Self {
244 cid: cid.into(),
245 merged_score,
246 contributing_nodes,
247 appearance_count,
248 }
249 }
250}
251
252#[derive(Debug, Clone, Default)]
258pub struct FederatedStats {
259 pub queries: u64,
261 pub avg_latency_ms: f64,
263 pub avg_results_returned: f64,
265 pub partial_failures: u64,
267 pub zero_result_queries: u64,
269}
270
271impl FederatedStats {
272 fn update_ewma(&mut self, latency_ms: f64, results_returned: usize) {
274 const ALPHA: f64 = 0.1;
275 if self.queries == 1 {
276 self.avg_latency_ms = latency_ms;
278 self.avg_results_returned = results_returned as f64;
279 } else {
280 self.avg_latency_ms = ALPHA * latency_ms + (1.0 - ALPHA) * self.avg_latency_ms;
281 self.avg_results_returned =
282 ALPHA * results_returned as f64 + (1.0 - ALPHA) * self.avg_results_returned;
283 }
284 }
285}
286
287#[derive(Debug)]
293struct CidAccumulator {
294 entries: Vec<(String, f64)>,
296}
297
298impl CidAccumulator {
299 fn new(node_id: String, score: f64) -> Self {
300 Self {
301 entries: vec![(node_id, score)],
302 }
303 }
304
305 fn push(&mut self, node_id: String, score: f64) {
306 self.entries.push((node_id, score));
307 }
308
309 fn appearance_count(&self) -> usize {
310 self.entries.len()
311 }
312
313 fn max_score(&self) -> f64 {
314 self.entries
315 .iter()
316 .map(|(_, s)| *s)
317 .fold(f64::NEG_INFINITY, f64::max)
318 }
319
320 fn weighted_score(&self, nodes: &HashMap<String, RemoteNode>) -> f64 {
321 let mut total_weight = 0.0_f64;
322 let mut weighted_sum = 0.0_f64;
323 for (node_id, score) in &self.entries {
324 let w = nodes
325 .get(node_id)
326 .map(|n| n.effective_weight())
327 .unwrap_or(1.0);
328 weighted_sum += w * score;
329 total_weight += w;
330 }
331 if total_weight == 0.0 {
332 0.0
333 } else {
334 weighted_sum / total_weight
335 }
336 }
337
338 fn contributing_nodes(&self) -> Vec<String> {
339 self.entries.iter().map(|(id, _)| id.clone()).collect()
340 }
341}
342
343pub struct SemanticFederatedSearch {
353 nodes: HashMap<String, RemoteNode>,
355 strategy: MergeStrategy,
357 stats: FederatedStats,
359}
360
361impl SemanticFederatedSearch {
362 pub fn new(strategy: MergeStrategy) -> Self {
364 Self {
365 nodes: HashMap::new(),
366 strategy,
367 stats: FederatedStats::default(),
368 }
369 }
370
371 pub fn register_node(&mut self, node: RemoteNode) -> bool {
378 if self.nodes.contains_key(&node.id) {
379 return false;
380 }
381 self.nodes.insert(node.id.clone(), node);
382 true
383 }
384
385 pub fn remove_node(&mut self, node_id: &str) -> bool {
387 self.nodes.remove(node_id).is_some()
388 }
389
390 pub fn update_node_stats(&mut self, node_id: &str, latency_ms: u64, success: bool) {
397 if let Some(node) = self.nodes.get_mut(node_id) {
398 node.last_latency_ms = latency_ms;
399 const SR_ALPHA: f64 = 0.2;
400 let outcome = if success { 1.0 } else { 0.0 };
401 node.success_rate = SR_ALPHA * outcome + (1.0 - SR_ALPHA) * node.success_rate;
402 }
403 }
404
405 pub fn merge_responses(
414 &mut self,
415 query: &FederatedQuery,
416 responses: Vec<NodeResponse>,
417 ) -> Vec<FederatedResult> {
418 let successful_responses: Vec<&NodeResponse> =
420 responses.iter().filter(|r| r.success).collect();
421 let successful_count = successful_responses.len();
422
423 self.stats.queries += 1;
424
425 let avg_latency = if successful_count > 0 {
427 successful_responses
428 .iter()
429 .map(|r| r.latency_ms as f64)
430 .sum::<f64>()
431 / successful_count as f64
432 } else {
433 0.0
434 };
435
436 if successful_count < query.required_quorum {
437 self.stats.partial_failures += 1;
438 self.stats.update_ewma(avg_latency, 0);
439 return Vec::new();
440 }
441
442 let accumulators: HashMap<String, CidAccumulator> =
444 self.build_accumulators(&successful_responses);
445
446 let mut merged: Vec<FederatedResult> = match &self.strategy.clone() {
448 MergeStrategy::SimpleUnion => self.apply_simple_union(&accumulators),
449 MergeStrategy::WeightedMerge => self.apply_weighted_merge(&accumulators),
450 MergeStrategy::QuorumIntersect(n) => self.apply_quorum_intersect(&accumulators, *n),
451 MergeStrategy::RankFusion => self.apply_rank_fusion(&successful_responses),
452 };
453
454 merged.retain(|r| r.merged_score >= query.min_score);
456 merged.sort_by(|a, b| {
457 b.merged_score
458 .partial_cmp(&a.merged_score)
459 .unwrap_or(std::cmp::Ordering::Equal)
460 });
461 merged.truncate(query.top_k);
462
463 if merged.is_empty() {
465 self.stats.zero_result_queries += 1;
466 }
467 self.stats.update_ewma(avg_latency, merged.len());
468
469 merged
470 }
471
472 pub fn deduplicate(results: &[FederatedResult]) -> Vec<FederatedResult> {
479 let mut best: HashMap<&str, &FederatedResult> = HashMap::new();
480 for r in results {
481 let entry = best.entry(r.cid.as_str()).or_insert(r);
482 if r.merged_score > entry.merged_score {
483 *entry = r;
484 }
485 }
486 let mut out: Vec<FederatedResult> = best.values().map(|r| (*r).clone()).collect();
488 out.sort_by(|a, b| {
489 b.merged_score
490 .partial_cmp(&a.merged_score)
491 .unwrap_or(std::cmp::Ordering::Equal)
492 });
493 out
494 }
495
496 pub fn stats(&self) -> &FederatedStats {
502 &self.stats
503 }
504
505 pub fn node_count(&self) -> usize {
507 self.nodes.len()
508 }
509
510 pub fn active_nodes(&self) -> Vec<&RemoteNode> {
512 self.nodes
513 .values()
514 .filter(|n| n.success_rate > 0.0)
515 .collect()
516 }
517
518 pub fn best_nodes(&self, n: usize) -> Vec<&RemoteNode> {
520 let mut nodes: Vec<&RemoteNode> = self.nodes.values().collect();
521 nodes.sort_by(|a, b| {
522 b.effective_weight()
523 .partial_cmp(&a.effective_weight())
524 .unwrap_or(std::cmp::Ordering::Equal)
525 });
526 nodes.truncate(n);
527 nodes
528 }
529
530 fn build_accumulators(&self, responses: &[&NodeResponse]) -> HashMap<String, CidAccumulator> {
535 let mut map: HashMap<String, CidAccumulator> = HashMap::new();
536 for response in responses {
537 for result in &response.results {
538 map.entry(result.cid.clone())
539 .and_modify(|acc| acc.push(result.node_id.clone(), result.score))
540 .or_insert_with(|| CidAccumulator::new(result.node_id.clone(), result.score));
541 }
542 }
543 map
544 }
545
546 fn apply_simple_union(
547 &self,
548 accumulators: &HashMap<String, CidAccumulator>,
549 ) -> Vec<FederatedResult> {
550 accumulators
551 .iter()
552 .map(|(cid, acc)| FederatedResult {
553 cid: cid.clone(),
554 merged_score: acc.max_score(),
555 appearance_count: acc.appearance_count(),
556 contributing_nodes: acc.contributing_nodes(),
557 })
558 .collect()
559 }
560
561 fn apply_weighted_merge(
562 &self,
563 accumulators: &HashMap<String, CidAccumulator>,
564 ) -> Vec<FederatedResult> {
565 accumulators
566 .iter()
567 .map(|(cid, acc)| FederatedResult {
568 cid: cid.clone(),
569 merged_score: acc.weighted_score(&self.nodes),
570 appearance_count: acc.appearance_count(),
571 contributing_nodes: acc.contributing_nodes(),
572 })
573 .collect()
574 }
575
576 fn apply_quorum_intersect(
577 &self,
578 accumulators: &HashMap<String, CidAccumulator>,
579 min_appearances: usize,
580 ) -> Vec<FederatedResult> {
581 accumulators
582 .iter()
583 .filter(|(_, acc)| acc.appearance_count() >= min_appearances)
584 .map(|(cid, acc)| FederatedResult {
585 cid: cid.clone(),
586 merged_score: acc.max_score(),
587 appearance_count: acc.appearance_count(),
588 contributing_nodes: acc.contributing_nodes(),
589 })
590 .collect()
591 }
592
593 fn apply_rank_fusion(&self, responses: &[&NodeResponse]) -> Vec<FederatedResult> {
597 const K: f64 = 60.0;
598
599 let mut rrf_scores: HashMap<String, f64> = HashMap::new();
601 let mut rrf_nodes: HashMap<String, Vec<String>> = HashMap::new();
602 let mut rrf_counts: HashMap<String, usize> = HashMap::new();
603
604 for response in responses {
605 let mut sorted: Vec<&RemoteResult> = response.results.iter().collect();
607 sorted.sort_by(|a, b| {
608 b.score
609 .partial_cmp(&a.score)
610 .unwrap_or(std::cmp::Ordering::Equal)
611 });
612
613 for (rank_zero, result) in sorted.iter().enumerate() {
614 let rank = (rank_zero + 1) as f64; let contribution = 1.0 / (K + rank);
616 *rrf_scores.entry(result.cid.clone()).or_insert(0.0) += contribution;
617 rrf_nodes
618 .entry(result.cid.clone())
619 .or_default()
620 .push(result.node_id.clone());
621 *rrf_counts.entry(result.cid.clone()).or_insert(0) += 1;
622 }
623 }
624
625 rrf_scores
626 .into_iter()
627 .map(|(cid, score)| {
628 let contributing = rrf_nodes.get(&cid).cloned().unwrap_or_default();
629 let count = rrf_counts.get(&cid).copied().unwrap_or(0);
630 FederatedResult {
631 cid,
632 merged_score: score,
633 contributing_nodes: contributing,
634 appearance_count: count,
635 }
636 })
637 .collect()
638 }
639}
640
641#[cfg(test)]
646mod tests {
647 use super::*;
648
649 fn make_node(id: &str, weight: f64, latency_ms: u64, success_rate: f64) -> RemoteNode {
654 RemoteNode::new(id, weight, latency_ms, success_rate)
655 }
656
657 fn make_result(cid: &str, score: f64, node_id: &str) -> RemoteResult {
658 RemoteResult::new(cid, score, node_id, vec![])
659 }
660
661 fn make_result_meta(
662 cid: &str,
663 score: f64,
664 node_id: &str,
665 meta: Vec<(&str, &str)>,
666 ) -> RemoteResult {
667 RemoteResult::new(
668 cid,
669 score,
670 node_id,
671 meta.into_iter()
672 .map(|(k, v)| (k.to_string(), v.to_string()))
673 .collect(),
674 )
675 }
676
677 fn make_query(top_k: usize, min_score: f64, quorum: usize) -> FederatedQuery {
678 FederatedQuery::new(vec![0.1, 0.2, 0.3])
679 .with_top_k(top_k)
680 .with_min_score(min_score)
681 .with_quorum(quorum)
682 }
683
684 fn coordinator_with_nodes(
685 strategy: MergeStrategy,
686 nodes: Vec<RemoteNode>,
687 ) -> SemanticFederatedSearch {
688 let mut sfs = SemanticFederatedSearch::new(strategy);
689 for n in nodes {
690 sfs.register_node(n);
691 }
692 sfs
693 }
694
695 #[test]
700 fn test_effective_weight_basic() {
701 let node = make_node("n1", 2.0, 1000, 0.5);
702 let ew = node.effective_weight();
705 assert!((ew - 0.5).abs() < 1e-12);
706 }
707
708 #[test]
709 fn test_effective_weight_zero_latency() {
710 let node = make_node("n1", 1.0, 0, 1.0);
711 assert!((node.effective_weight() - 1.0).abs() < 1e-12);
713 }
714
715 #[test]
716 fn test_effective_weight_zero_success_rate() {
717 let node = make_node("n1", 10.0, 100, 0.0);
718 assert_eq!(node.effective_weight(), 0.0);
719 }
720
721 #[test]
722 fn test_effective_weight_high_latency_reduces_weight() {
723 let low = make_node("a", 1.0, 0, 1.0);
724 let high = make_node("b", 1.0, 10_000, 1.0);
725 assert!(low.effective_weight() > high.effective_weight());
726 }
727
728 #[test]
733 fn test_register_node_new() {
734 let mut sfs = SemanticFederatedSearch::new(MergeStrategy::SimpleUnion);
735 assert!(sfs.register_node(make_node("n1", 1.0, 0, 1.0)));
736 assert_eq!(sfs.node_count(), 1);
737 }
738
739 #[test]
740 fn test_register_node_duplicate_returns_false() {
741 let mut sfs = SemanticFederatedSearch::new(MergeStrategy::SimpleUnion);
742 sfs.register_node(make_node("n1", 1.0, 0, 1.0));
743 assert!(!sfs.register_node(make_node("n1", 2.0, 0, 1.0)));
744 assert_eq!(sfs.node_count(), 1);
745 }
746
747 #[test]
748 fn test_remove_node_existing() {
749 let mut sfs = SemanticFederatedSearch::new(MergeStrategy::SimpleUnion);
750 sfs.register_node(make_node("n1", 1.0, 0, 1.0));
751 assert!(sfs.remove_node("n1"));
752 assert_eq!(sfs.node_count(), 0);
753 }
754
755 #[test]
756 fn test_remove_node_nonexistent() {
757 let mut sfs = SemanticFederatedSearch::new(MergeStrategy::SimpleUnion);
758 assert!(!sfs.remove_node("ghost"));
759 }
760
761 #[test]
766 fn test_update_node_stats_latency() {
767 let mut sfs = SemanticFederatedSearch::new(MergeStrategy::SimpleUnion);
768 sfs.register_node(make_node("n1", 1.0, 0, 1.0));
769 sfs.update_node_stats("n1", 250, true);
770 assert_eq!(sfs.nodes["n1"].last_latency_ms, 250);
771 }
772
773 #[test]
774 fn test_update_node_stats_success_rate_increases() {
775 let mut sfs = SemanticFederatedSearch::new(MergeStrategy::SimpleUnion);
776 sfs.register_node(make_node("n1", 1.0, 0, 0.5));
777 let before = sfs.nodes["n1"].success_rate;
778 sfs.update_node_stats("n1", 10, true);
779 assert!(sfs.nodes["n1"].success_rate > before);
780 }
781
782 #[test]
783 fn test_update_node_stats_success_rate_decreases() {
784 let mut sfs = SemanticFederatedSearch::new(MergeStrategy::SimpleUnion);
785 sfs.register_node(make_node("n1", 1.0, 0, 1.0));
786 sfs.update_node_stats("n1", 10, false);
787 assert!(sfs.nodes["n1"].success_rate < 1.0);
788 }
789
790 #[test]
791 fn test_update_node_stats_unknown_node_noop() {
792 let mut sfs = SemanticFederatedSearch::new(MergeStrategy::SimpleUnion);
793 sfs.update_node_stats("ghost", 100, true);
795 }
796
797 #[test]
802 fn test_quorum_failure_increments_partial_failures() {
803 let mut sfs = coordinator_with_nodes(
804 MergeStrategy::SimpleUnion,
805 vec![make_node("n1", 1.0, 0, 1.0)],
806 );
807 let query = make_query(10, 0.0, 2); let responses = vec![NodeResponse::success(
809 "n1",
810 vec![make_result("cid1", 0.9, "n1")],
811 50,
812 )];
813 let results = sfs.merge_responses(&query, responses);
814 assert!(results.is_empty());
815 assert_eq!(sfs.stats().partial_failures, 1);
816 }
817
818 #[test]
819 fn test_all_failures_returns_empty() {
820 let mut sfs = coordinator_with_nodes(
821 MergeStrategy::SimpleUnion,
822 vec![make_node("n1", 1.0, 0, 1.0)],
823 );
824 let query = make_query(10, 0.0, 1);
825 let responses = vec![NodeResponse::failure("n1", 100)];
826 let results = sfs.merge_responses(&query, responses);
827 assert!(results.is_empty());
828 assert_eq!(sfs.stats().partial_failures, 1);
829 }
830
831 #[test]
836 fn test_simple_union_takes_max_score() {
837 let mut sfs = coordinator_with_nodes(
838 MergeStrategy::SimpleUnion,
839 vec![make_node("n1", 1.0, 0, 1.0), make_node("n2", 1.0, 0, 1.0)],
840 );
841 let query = make_query(10, 0.0, 1);
842 let responses = vec![
843 NodeResponse::success("n1", vec![make_result("cid1", 0.7, "n1")], 10),
844 NodeResponse::success("n2", vec![make_result("cid1", 0.9, "n2")], 10),
845 ];
846 let results = sfs.merge_responses(&query, responses);
847 assert_eq!(results.len(), 1);
848 assert!((results[0].merged_score - 0.9).abs() < 1e-12);
849 }
850
851 #[test]
852 fn test_simple_union_includes_all_cids() {
853 let mut sfs = coordinator_with_nodes(
854 MergeStrategy::SimpleUnion,
855 vec![make_node("n1", 1.0, 0, 1.0), make_node("n2", 1.0, 0, 1.0)],
856 );
857 let query = make_query(10, 0.0, 1);
858 let responses = vec![
859 NodeResponse::success("n1", vec![make_result("cidA", 0.8, "n1")], 10),
860 NodeResponse::success("n2", vec![make_result("cidB", 0.6, "n2")], 10),
861 ];
862 let results = sfs.merge_responses(&query, responses);
863 assert_eq!(results.len(), 2);
864 }
865
866 #[test]
867 fn test_simple_union_sorted_descending() {
868 let mut sfs = coordinator_with_nodes(
869 MergeStrategy::SimpleUnion,
870 vec![make_node("n1", 1.0, 0, 1.0)],
871 );
872 let query = make_query(10, 0.0, 1);
873 let responses = vec![NodeResponse::success(
874 "n1",
875 vec![
876 make_result("cidA", 0.3, "n1"),
877 make_result("cidB", 0.9, "n1"),
878 make_result("cidC", 0.6, "n1"),
879 ],
880 10,
881 )];
882 let results = sfs.merge_responses(&query, responses);
883 assert_eq!(results.len(), 3);
884 assert!(results[0].merged_score >= results[1].merged_score);
885 assert!(results[1].merged_score >= results[2].merged_score);
886 }
887
888 #[test]
893 fn test_weighted_merge_single_node() {
894 let mut sfs = coordinator_with_nodes(
895 MergeStrategy::WeightedMerge,
896 vec![make_node("n1", 1.0, 0, 1.0)],
897 );
898 let query = make_query(10, 0.0, 1);
899 let responses = vec![NodeResponse::success(
900 "n1",
901 vec![make_result("cid1", 0.75, "n1")],
902 10,
903 )];
904 let results = sfs.merge_responses(&query, responses);
905 assert_eq!(results.len(), 1);
906 assert!((results[0].merged_score - 0.75).abs() < 1e-9);
908 }
909
910 #[test]
911 fn test_weighted_merge_two_equal_weight_nodes() {
912 let mut sfs = coordinator_with_nodes(
913 MergeStrategy::WeightedMerge,
914 vec![make_node("n1", 1.0, 0, 1.0), make_node("n2", 1.0, 0, 1.0)],
915 );
916 let query = make_query(10, 0.0, 1);
917 let responses = vec![
918 NodeResponse::success("n1", vec![make_result("cid1", 0.8, "n1")], 10),
919 NodeResponse::success("n2", vec![make_result("cid1", 0.6, "n2")], 10),
920 ];
921 let results = sfs.merge_responses(&query, responses);
922 assert_eq!(results.len(), 1);
923 assert!((results[0].merged_score - 0.7).abs() < 1e-9);
925 }
926
927 #[test]
928 fn test_weighted_merge_higher_weight_dominates() {
929 let mut sfs = coordinator_with_nodes(
930 MergeStrategy::WeightedMerge,
931 vec![
932 make_node("n1", 10.0, 0, 1.0), make_node("n2", 1.0, 0, 1.0), ],
935 );
936 let query = make_query(10, 0.0, 1);
937 let responses = vec![
938 NodeResponse::success("n1", vec![make_result("cid1", 1.0, "n1")], 10),
939 NodeResponse::success("n2", vec![make_result("cid1", 0.0, "n2")], 10),
940 ];
941 let results = sfs.merge_responses(&query, responses);
942 assert_eq!(results.len(), 1);
943 assert!(results[0].merged_score > 0.5);
945 }
946
947 #[test]
952 fn test_quorum_intersect_filters_rare_cids() {
953 let mut sfs = coordinator_with_nodes(
954 MergeStrategy::QuorumIntersect(2),
955 vec![make_node("n1", 1.0, 0, 1.0), make_node("n2", 1.0, 0, 1.0)],
956 );
957 let query = make_query(10, 0.0, 1);
958 let responses = vec![
959 NodeResponse::success(
960 "n1",
961 vec![
962 make_result("cidCommon", 0.8, "n1"),
963 make_result("cidRare", 0.7, "n1"),
964 ],
965 10,
966 ),
967 NodeResponse::success("n2", vec![make_result("cidCommon", 0.6, "n2")], 10),
968 ];
969 let results = sfs.merge_responses(&query, responses);
970 assert!(results.iter().all(|r| r.cid != "cidRare"));
972 assert!(results.iter().any(|r| r.cid == "cidCommon"));
973 }
974
975 #[test]
976 fn test_quorum_intersect_n1_equals_simple_union() {
977 let nodes = vec![make_node("n1", 1.0, 0, 1.0), make_node("n2", 1.0, 0, 1.0)];
978 let mut sfs_union = coordinator_with_nodes(MergeStrategy::SimpleUnion, nodes.clone());
979 let mut sfs_qi = coordinator_with_nodes(MergeStrategy::QuorumIntersect(1), nodes);
980
981 let query = make_query(10, 0.0, 1);
982 let responses = || {
983 vec![
984 NodeResponse::success(
985 "n1",
986 vec![
987 make_result("cidA", 0.8, "n1"),
988 make_result("cidB", 0.5, "n1"),
989 ],
990 10,
991 ),
992 NodeResponse::success("n2", vec![make_result("cidC", 0.6, "n2")], 10),
993 ]
994 };
995
996 let r_union = sfs_union.merge_responses(&query, responses());
997 let r_qi = sfs_qi.merge_responses(&query, responses());
998 assert_eq!(r_union.len(), r_qi.len());
999 }
1000
1001 #[test]
1006 fn test_rank_fusion_top_ranked_from_all_nodes_gets_highest_score() {
1007 let mut sfs = coordinator_with_nodes(
1008 MergeStrategy::RankFusion,
1009 vec![make_node("n1", 1.0, 0, 1.0), make_node("n2", 1.0, 0, 1.0)],
1010 );
1011 let query = make_query(10, 0.0, 1);
1012 let responses = vec![
1013 NodeResponse::success(
1014 "n1",
1015 vec![
1016 make_result("cidTop", 0.95, "n1"),
1017 make_result("cidLow", 0.4, "n1"),
1018 ],
1019 10,
1020 ),
1021 NodeResponse::success(
1022 "n2",
1023 vec![
1024 make_result("cidTop", 0.90, "n2"),
1025 make_result("cidLow", 0.3, "n2"),
1026 ],
1027 10,
1028 ),
1029 ];
1030 let results = sfs.merge_responses(&query, responses);
1031 assert_eq!(results[0].cid, "cidTop");
1032 }
1033
1034 #[test]
1035 fn test_rank_fusion_score_formula() {
1036 let mut sfs = coordinator_with_nodes(
1038 MergeStrategy::RankFusion,
1039 vec![make_node("n1", 1.0, 0, 1.0)],
1040 );
1041 let query = make_query(10, 0.0, 1);
1042 let responses = vec![NodeResponse::success(
1043 "n1",
1044 vec![make_result("cid1", 0.9, "n1")],
1045 10,
1046 )];
1047 let results = sfs.merge_responses(&query, responses);
1048 let expected = 1.0 / (60.0 + 1.0);
1049 assert!((results[0].merged_score - expected).abs() < 1e-12);
1050 }
1051
1052 #[test]
1053 fn test_rank_fusion_two_nodes_contribute() {
1054 let mut sfs = coordinator_with_nodes(
1056 MergeStrategy::RankFusion,
1057 vec![make_node("n1", 1.0, 0, 1.0), make_node("n2", 1.0, 0, 1.0)],
1058 );
1059 let query = make_query(10, 0.0, 1);
1060 let responses = vec![
1061 NodeResponse::success("n1", vec![make_result("cid1", 0.9, "n1")], 10),
1062 NodeResponse::success("n2", vec![make_result("cid1", 0.8, "n2")], 10),
1063 ];
1064 let results = sfs.merge_responses(&query, responses);
1065 let expected = 2.0 / 61.0;
1066 assert!((results[0].merged_score - expected).abs() < 1e-12);
1067 }
1068
1069 #[test]
1074 fn test_min_score_filters_low_results() {
1075 let mut sfs = coordinator_with_nodes(
1076 MergeStrategy::SimpleUnion,
1077 vec![make_node("n1", 1.0, 0, 1.0)],
1078 );
1079 let query = make_query(10, 0.5, 1);
1080 let responses = vec![NodeResponse::success(
1081 "n1",
1082 vec![
1083 make_result("cidHigh", 0.8, "n1"),
1084 make_result("cidLow", 0.3, "n1"),
1085 ],
1086 10,
1087 )];
1088 let results = sfs.merge_responses(&query, responses);
1089 assert_eq!(results.len(), 1);
1090 assert_eq!(results[0].cid, "cidHigh");
1091 }
1092
1093 #[test]
1098 fn test_top_k_truncation() {
1099 let mut sfs = coordinator_with_nodes(
1100 MergeStrategy::SimpleUnion,
1101 vec![make_node("n1", 1.0, 0, 1.0)],
1102 );
1103 let query = make_query(2, 0.0, 1);
1104 let responses = vec![NodeResponse::success(
1105 "n1",
1106 vec![
1107 make_result("c1", 0.9, "n1"),
1108 make_result("c2", 0.8, "n1"),
1109 make_result("c3", 0.7, "n1"),
1110 make_result("c4", 0.6, "n1"),
1111 ],
1112 10,
1113 )];
1114 let results = sfs.merge_responses(&query, responses);
1115 assert_eq!(results.len(), 2);
1116 assert_eq!(results[0].cid, "c1");
1118 assert_eq!(results[1].cid, "c2");
1119 }
1120
1121 #[test]
1126 fn test_deduplicate_keeps_highest_score() {
1127 let results = vec![
1128 FederatedResult::new("cid1", 0.9, vec!["n1".into()], 1),
1129 FederatedResult::new("cid1", 0.5, vec!["n2".into()], 1),
1130 FederatedResult::new("cid2", 0.7, vec!["n1".into()], 1),
1131 ];
1132 let deduped = SemanticFederatedSearch::deduplicate(&results);
1133 assert_eq!(deduped.len(), 2);
1134 let cid1 = deduped
1135 .iter()
1136 .find(|r| r.cid == "cid1")
1137 .expect("cid1 missing");
1138 assert!((cid1.merged_score - 0.9).abs() < 1e-12);
1139 }
1140
1141 #[test]
1142 fn test_deduplicate_sorted_descending() {
1143 let results = vec![
1144 FederatedResult::new("cid2", 0.3, vec![], 1),
1145 FederatedResult::new("cid1", 0.9, vec![], 1),
1146 FederatedResult::new("cid3", 0.6, vec![], 1),
1147 ];
1148 let deduped = SemanticFederatedSearch::deduplicate(&results);
1149 assert_eq!(deduped.len(), 3);
1150 assert!(deduped[0].merged_score >= deduped[1].merged_score);
1151 assert!(deduped[1].merged_score >= deduped[2].merged_score);
1152 }
1153
1154 #[test]
1155 fn test_deduplicate_empty() {
1156 let results: Vec<FederatedResult> = vec![];
1157 let deduped = SemanticFederatedSearch::deduplicate(&results);
1158 assert!(deduped.is_empty());
1159 }
1160
1161 #[test]
1162 fn test_deduplicate_no_duplicates_unchanged_count() {
1163 let results = vec![
1164 FederatedResult::new("cid1", 0.9, vec![], 1),
1165 FederatedResult::new("cid2", 0.7, vec![], 1),
1166 ];
1167 let deduped = SemanticFederatedSearch::deduplicate(&results);
1168 assert_eq!(deduped.len(), 2);
1169 }
1170
1171 #[test]
1176 fn test_stats_query_count_increments() {
1177 let mut sfs = coordinator_with_nodes(
1178 MergeStrategy::SimpleUnion,
1179 vec![make_node("n1", 1.0, 0, 1.0)],
1180 );
1181 let query = make_query(10, 0.0, 1);
1182 sfs.merge_responses(
1183 &query,
1184 vec![NodeResponse::success(
1185 "n1",
1186 vec![make_result("cid1", 0.9, "n1")],
1187 50,
1188 )],
1189 );
1190 sfs.merge_responses(
1191 &query,
1192 vec![NodeResponse::success(
1193 "n1",
1194 vec![make_result("cid2", 0.8, "n1")],
1195 60,
1196 )],
1197 );
1198 assert_eq!(sfs.stats().queries, 2);
1199 }
1200
1201 #[test]
1202 fn test_stats_zero_result_query() {
1203 let mut sfs = coordinator_with_nodes(
1204 MergeStrategy::SimpleUnion,
1205 vec![make_node("n1", 1.0, 0, 1.0)],
1206 );
1207 let query = make_query(10, 0.99, 1);
1209 sfs.merge_responses(
1210 &query,
1211 vec![NodeResponse::success(
1212 "n1",
1213 vec![make_result("cid1", 0.5, "n1")],
1214 10,
1215 )],
1216 );
1217 assert_eq!(sfs.stats().zero_result_queries, 1);
1218 }
1219
1220 #[test]
1221 fn test_stats_avg_latency_ewma_seeded_on_first_query() {
1222 let mut sfs = coordinator_with_nodes(
1223 MergeStrategy::SimpleUnion,
1224 vec![make_node("n1", 1.0, 0, 1.0)],
1225 );
1226 let query = make_query(10, 0.0, 1);
1227 sfs.merge_responses(
1228 &query,
1229 vec![NodeResponse::success(
1230 "n1",
1231 vec![make_result("cid1", 0.9, "n1")],
1232 100,
1233 )],
1234 );
1235 assert!((sfs.stats().avg_latency_ms - 100.0).abs() < 1e-9);
1237 }
1238
1239 #[test]
1244 fn test_active_nodes_excludes_zero_success_rate() {
1245 let sfs = coordinator_with_nodes(
1246 MergeStrategy::SimpleUnion,
1247 vec![
1248 make_node("active", 1.0, 0, 0.9),
1249 make_node("inactive", 1.0, 0, 0.0),
1250 ],
1251 );
1252 let active = sfs.active_nodes();
1253 assert_eq!(active.len(), 1);
1254 assert_eq!(active[0].id, "active");
1255 }
1256
1257 #[test]
1258 fn test_best_nodes_ordered_by_effective_weight() {
1259 let sfs = coordinator_with_nodes(
1260 MergeStrategy::SimpleUnion,
1261 vec![
1262 make_node("low", 1.0, 5000, 0.5),
1263 make_node("high", 2.0, 0, 1.0),
1264 make_node("mid", 1.0, 0, 0.8),
1265 ],
1266 );
1267 let best = sfs.best_nodes(3);
1268 for i in 0..best.len() - 1 {
1270 assert!(best[i].effective_weight() >= best[i + 1].effective_weight());
1271 }
1272 }
1273
1274 #[test]
1275 fn test_best_nodes_limited_by_n() {
1276 let nodes: Vec<RemoteNode> = (0..5)
1277 .map(|i| make_node(&format!("n{i}"), 1.0, 0, 1.0))
1278 .collect();
1279 let sfs = coordinator_with_nodes(MergeStrategy::SimpleUnion, nodes);
1280 assert_eq!(sfs.best_nodes(3).len(), 3);
1281 }
1282
1283 #[test]
1284 fn test_best_nodes_n_larger_than_count() {
1285 let sfs = coordinator_with_nodes(
1286 MergeStrategy::SimpleUnion,
1287 vec![make_node("n1", 1.0, 0, 1.0)],
1288 );
1289 assert_eq!(sfs.best_nodes(100).len(), 1);
1290 }
1291
1292 #[test]
1297 fn test_appearance_count_two_nodes() {
1298 let mut sfs = coordinator_with_nodes(
1299 MergeStrategy::SimpleUnion,
1300 vec![make_node("n1", 1.0, 0, 1.0), make_node("n2", 1.0, 0, 1.0)],
1301 );
1302 let query = make_query(10, 0.0, 1);
1303 let responses = vec![
1304 NodeResponse::success("n1", vec![make_result("cid1", 0.8, "n1")], 10),
1305 NodeResponse::success("n2", vec![make_result("cid1", 0.7, "n2")], 10),
1306 ];
1307 let results = sfs.merge_responses(&query, responses);
1308 assert_eq!(results[0].appearance_count, 2);
1309 assert_eq!(results[0].contributing_nodes.len(), 2);
1310 }
1311
1312 #[test]
1317 fn test_metadata_preserved() {
1318 let r = make_result_meta("cid1", 0.9, "n1", vec![("type", "doc"), ("lang", "en")]);
1319 assert_eq!(r.metadata.len(), 2);
1320 assert_eq!(r.metadata[0], ("type".to_string(), "doc".to_string()));
1321 }
1322
1323 #[test]
1328 fn test_node_response_success_constructor() {
1329 let r = NodeResponse::success("n1", vec![make_result("c1", 0.5, "n1")], 99);
1330 assert!(r.success);
1331 assert_eq!(r.latency_ms, 99);
1332 assert_eq!(r.results.len(), 1);
1333 }
1334
1335 #[test]
1336 fn test_node_response_failure_constructor() {
1337 let r = NodeResponse::failure("n1", 500);
1338 assert!(!r.success);
1339 assert!(r.results.is_empty());
1340 assert_eq!(r.latency_ms, 500);
1341 }
1342
1343 #[test]
1348 fn test_federated_query_default_values() {
1349 let q = FederatedQuery::new(vec![0.1]);
1350 assert_eq!(q.top_k, 10);
1351 assert_eq!(q.min_score, 0.0);
1352 assert_eq!(q.required_quorum, 1);
1353 }
1354
1355 #[test]
1356 fn test_federated_query_builders() {
1357 let q = FederatedQuery::new(vec![])
1358 .with_top_k(5)
1359 .with_min_score(0.3)
1360 .with_quorum(3);
1361 assert_eq!(q.top_k, 5);
1362 assert!((q.min_score - 0.3).abs() < 1e-12);
1363 assert_eq!(q.required_quorum, 3);
1364 }
1365
1366 #[test]
1371 fn test_merge_with_no_registered_nodes_but_synthetic_responses() {
1372 let mut sfs = SemanticFederatedSearch::new(MergeStrategy::WeightedMerge);
1374 let query = make_query(10, 0.0, 1);
1375 let responses = vec![NodeResponse::success(
1376 "unknown_node",
1377 vec![make_result("cid1", 0.6, "unknown_node")],
1378 20,
1379 )];
1380 let results = sfs.merge_responses(&query, responses);
1381 assert_eq!(results.len(), 1);
1382 }
1383
1384 #[test]
1389 fn test_zero_quorum_accepts_all_failure_responses() {
1390 let mut sfs = coordinator_with_nodes(
1391 MergeStrategy::SimpleUnion,
1392 vec![make_node("n1", 1.0, 0, 1.0)],
1393 );
1394 let query = make_query(10, 0.0, 0);
1396 let responses = vec![NodeResponse::failure("n1", 10)];
1397 let results = sfs.merge_responses(&query, responses);
1399 assert!(results.is_empty());
1400 assert_eq!(sfs.stats().partial_failures, 0);
1402 }
1403
1404 #[test]
1409 fn test_stats_partial_failures_multiple() {
1410 let mut sfs = coordinator_with_nodes(
1411 MergeStrategy::SimpleUnion,
1412 vec![make_node("n1", 1.0, 0, 1.0)],
1413 );
1414 let query = make_query(10, 0.0, 2); for _ in 0..3 {
1416 sfs.merge_responses(
1417 &query,
1418 vec![NodeResponse::success(
1419 "n1",
1420 vec![make_result("c1", 0.9, "n1")],
1421 10,
1422 )],
1423 );
1424 }
1425 assert_eq!(sfs.stats().partial_failures, 3);
1426 }
1427
1428 #[test]
1433 fn test_rank_fusion_rank_order_correct() {
1434 let mut sfs = coordinator_with_nodes(
1438 MergeStrategy::RankFusion,
1439 vec![make_node("n1", 1.0, 0, 1.0)],
1440 );
1441 let query = make_query(10, 0.0, 1);
1442 let responses = vec![NodeResponse::success(
1443 "n1",
1444 vec![
1445 make_result("cidC", 0.1, "n1"),
1446 make_result("cidA", 0.9, "n1"),
1447 make_result("cidB", 0.5, "n1"),
1448 ],
1449 10,
1450 )];
1451 let results = sfs.merge_responses(&query, responses);
1452 assert_eq!(results[0].cid, "cidA");
1453 assert!(results[0].merged_score > results[1].merged_score);
1454 }
1455
1456 #[test]
1461 fn test_federated_result_new() {
1462 let r = FederatedResult::new("mycid", 0.42, vec!["n1".into(), "n2".into()], 2);
1463 assert_eq!(r.cid, "mycid");
1464 assert!((r.merged_score - 0.42).abs() < 1e-12);
1465 assert_eq!(r.contributing_nodes.len(), 2);
1466 assert_eq!(r.appearance_count, 2);
1467 }
1468
1469 #[test]
1474 fn test_weighted_merge_unregistered_node_defaults_to_weight_one() {
1475 let mut sfs = SemanticFederatedSearch::new(MergeStrategy::WeightedMerge);
1476 let query = make_query(10, 0.0, 1);
1477 let responses = vec![NodeResponse::success(
1478 "unregistered",
1479 vec![make_result("cid1", 0.8, "unregistered")],
1480 10,
1481 )];
1482 let results = sfs.merge_responses(&query, responses);
1483 assert_eq!(results.len(), 1);
1484 assert!((results[0].merged_score - 0.8).abs() < 1e-9);
1486 }
1487
1488 #[test]
1493 fn test_failed_responses_excluded_from_merge() {
1494 let mut sfs = coordinator_with_nodes(
1495 MergeStrategy::SimpleUnion,
1496 vec![make_node("n1", 1.0, 0, 1.0), make_node("n2", 1.0, 0, 1.0)],
1497 );
1498 let query = make_query(10, 0.0, 1);
1499 let responses = vec![
1500 NodeResponse::success("n1", vec![make_result("cid_good", 0.9, "n1")], 10),
1501 NodeResponse::failure("n2", 5000),
1502 ];
1503 let results = sfs.merge_responses(&query, responses);
1504 assert!(results
1506 .iter()
1507 .all(|r| r.contributing_nodes.contains(&"n1".to_string())));
1508 assert!(!results
1509 .iter()
1510 .any(|r| r.contributing_nodes.contains(&"n2".to_string())));
1511 }
1512
1513 #[test]
1518 fn test_node_count_after_add_remove() {
1519 let mut sfs = SemanticFederatedSearch::new(MergeStrategy::SimpleUnion);
1520 sfs.register_node(make_node("n1", 1.0, 0, 1.0));
1521 sfs.register_node(make_node("n2", 1.0, 0, 1.0));
1522 assert_eq!(sfs.node_count(), 2);
1523 sfs.remove_node("n1");
1524 assert_eq!(sfs.node_count(), 1);
1525 }
1526}