1#[derive(Clone, Debug, PartialEq)]
9pub enum ContextEntry {
10 Query {
12 text: String,
14 embedding: Vec<f32>,
16 tick: u64,
18 },
19 Result {
21 doc_id: u64,
23 embedding: Vec<f32>,
25 relevance: f32,
27 tick: u64,
29 },
30 Feedback {
32 doc_id: u64,
34 positive: bool,
36 tick: u64,
38 },
39}
40
41#[derive(Clone, Debug, PartialEq)]
43pub struct WindowConfig {
44 pub max_entries: usize,
46 pub decay_factor: f32,
48 pub embedding_dim: usize,
50}
51
52impl Default for WindowConfig {
53 fn default() -> Self {
54 Self {
55 max_entries: 20,
56 decay_factor: 0.9,
57 embedding_dim: 128,
58 }
59 }
60}
61
62#[derive(Clone, Debug, PartialEq)]
64pub struct ContextStats {
65 pub total_entries_added: u64,
67 pub queries_in_window: usize,
69 pub results_in_window: usize,
71 pub feedbacks_in_window: usize,
73 pub positive_feedback_ratio: f32,
76}
77
78pub struct SemanticContextWindow {
83 pub entries: Vec<ContextEntry>,
85 pub config: WindowConfig,
87 pub total_added: u64,
89}
90
91impl SemanticContextWindow {
92 pub fn new(config: WindowConfig) -> Self {
94 Self {
95 entries: Vec::new(),
96 config,
97 total_added: 0,
98 }
99 }
100
101 pub fn add(&mut self, entry: ContextEntry) {
105 self.entries.push(entry);
106 if self.entries.len() > self.config.max_entries {
107 self.entries.remove(0);
108 }
109 self.total_added += 1;
110 }
111
112 pub fn context_embedding(&self) -> Vec<f32> {
123 let dim = self.config.embedding_dim;
124 let decay = self.config.decay_factor;
125 let window_len = self.entries.len();
126
127 let embeddable: Vec<(usize, &Vec<f32>)> = self
129 .entries
130 .iter()
131 .enumerate()
132 .filter_map(|(i, e)| match e {
133 ContextEntry::Query { embedding, .. } => Some((i, embedding)),
134 ContextEntry::Result { embedding, .. } => Some((i, embedding)),
135 ContextEntry::Feedback { .. } => None,
136 })
137 .filter(|(_, emb)| emb.len() == dim)
138 .collect();
139
140 if embeddable.is_empty() {
141 return vec![0.0_f32; dim];
142 }
143
144 let weight_of = |idx: usize| -> f32 {
146 let age = (window_len - 1).saturating_sub(idx); decay.powi(age as i32)
148 };
149
150 let weight_sum: f32 = embeddable.iter().map(|(i, _)| weight_of(*i)).sum();
151
152 let mut result = vec![0.0_f32; dim];
153 for (i, emb) in &embeddable {
154 let w = weight_of(*i) / weight_sum;
155 for (r, v) in result.iter_mut().zip(emb.iter()) {
156 *r += w * v;
157 }
158 }
159 result
160 }
161
162 pub fn recent_queries(&self, n: usize) -> Vec<&ContextEntry> {
164 self.entries
165 .iter()
166 .rev()
167 .filter(|e| matches!(e, ContextEntry::Query { .. }))
168 .take(n)
169 .collect()
170 }
171
172 pub fn recent_results(&self, n: usize) -> Vec<&ContextEntry> {
174 self.entries
175 .iter()
176 .rev()
177 .filter(|e| matches!(e, ContextEntry::Result { .. }))
178 .take(n)
179 .collect()
180 }
181
182 pub fn positive_docs(&self) -> Vec<u64> {
184 let mut seen = std::collections::HashSet::new();
185 let mut result = Vec::new();
186 for e in &self.entries {
187 if let ContextEntry::Feedback {
188 doc_id,
189 positive: true,
190 ..
191 } = e
192 {
193 if seen.insert(*doc_id) {
194 result.push(*doc_id);
195 }
196 }
197 }
198 result
199 }
200
201 pub fn negative_docs(&self) -> Vec<u64> {
203 let mut seen = std::collections::HashSet::new();
204 let mut result = Vec::new();
205 for e in &self.entries {
206 if let ContextEntry::Feedback {
207 doc_id,
208 positive: false,
209 ..
210 } = e
211 {
212 if seen.insert(*doc_id) {
213 result.push(*doc_id);
214 }
215 }
216 }
217 result
218 }
219
220 pub fn clear(&mut self) {
222 self.entries.clear();
223 }
224
225 pub fn stats(&self) -> ContextStats {
227 let mut queries_in_window = 0usize;
228 let mut results_in_window = 0usize;
229 let mut feedbacks_in_window = 0usize;
230 let mut positive_count = 0usize;
231
232 for e in &self.entries {
233 match e {
234 ContextEntry::Query { .. } => queries_in_window += 1,
235 ContextEntry::Result { .. } => results_in_window += 1,
236 ContextEntry::Feedback { positive, .. } => {
237 feedbacks_in_window += 1;
238 if *positive {
239 positive_count += 1;
240 }
241 }
242 }
243 }
244
245 let positive_feedback_ratio = if feedbacks_in_window == 0 {
246 0.0
247 } else {
248 positive_count as f32 / feedbacks_in_window as f32
249 };
250
251 ContextStats {
252 total_entries_added: self.total_added,
253 queries_in_window,
254 results_in_window,
255 feedbacks_in_window,
256 positive_feedback_ratio,
257 }
258 }
259}
260
261#[cfg(test)]
266mod tests {
267 use super::*;
268
269 fn make_embedding(val: f32, dim: usize) -> Vec<f32> {
270 vec![val; dim]
271 }
272
273 fn default_window() -> SemanticContextWindow {
274 SemanticContextWindow::new(WindowConfig::default())
275 }
276
277 #[test]
282 fn test_add_single_entry() {
283 let mut w = default_window();
284 w.add(ContextEntry::Feedback {
285 doc_id: 1,
286 positive: true,
287 tick: 0,
288 });
289 assert_eq!(w.entries.len(), 1);
290 assert_eq!(w.total_added, 1);
291 }
292
293 #[test]
294 fn test_sliding_window_evicts_oldest() {
295 let config = WindowConfig {
296 max_entries: 3,
297 ..Default::default()
298 };
299 let mut w = SemanticContextWindow::new(config);
300 for tick in 0..4u64 {
301 w.add(ContextEntry::Feedback {
302 doc_id: tick,
303 positive: true,
304 tick,
305 });
306 }
307 assert_eq!(w.entries.len(), 3);
309 assert!(matches!(
310 &w.entries[0],
311 ContextEntry::Feedback { doc_id: 1, .. }
312 ));
313 assert!(matches!(
314 &w.entries[2],
315 ContextEntry::Feedback { doc_id: 3, .. }
316 ));
317 }
318
319 #[test]
320 fn test_total_added_counts_all_including_evicted() {
321 let config = WindowConfig {
322 max_entries: 2,
323 ..Default::default()
324 };
325 let mut w = SemanticContextWindow::new(config);
326 for tick in 0..5u64 {
327 w.add(ContextEntry::Feedback {
328 doc_id: tick,
329 positive: false,
330 tick,
331 });
332 }
333 assert_eq!(w.total_added, 5);
334 assert_eq!(w.entries.len(), 2);
335 }
336
337 #[test]
338 fn test_add_up_to_capacity_no_eviction() {
339 let config = WindowConfig {
340 max_entries: 5,
341 ..Default::default()
342 };
343 let mut w = SemanticContextWindow::new(config);
344 for tick in 0..5u64 {
345 w.add(ContextEntry::Feedback {
346 doc_id: tick,
347 positive: true,
348 tick,
349 });
350 }
351 assert_eq!(w.entries.len(), 5);
352 assert_eq!(w.total_added, 5);
353 }
354
355 #[test]
360 fn test_context_embedding_empty_returns_zeros() {
361 let w = default_window();
362 let emb = w.context_embedding();
363 assert_eq!(emb.len(), 128);
364 assert!(emb.iter().all(|&v| v == 0.0));
365 }
366
367 #[test]
368 fn test_context_embedding_only_feedback_returns_zeros() {
369 let mut w = default_window();
370 w.add(ContextEntry::Feedback {
371 doc_id: 1,
372 positive: true,
373 tick: 0,
374 });
375 let emb = w.context_embedding();
376 assert!(emb.iter().all(|&v| v == 0.0));
377 }
378
379 #[test]
380 fn test_context_embedding_single_query_returns_that_embedding() {
381 let mut w = default_window();
382 let emb_val = 0.5_f32;
383 w.add(ContextEntry::Query {
384 text: "hello".into(),
385 embedding: make_embedding(emb_val, 128),
386 tick: 0,
387 });
388 let result = w.context_embedding();
389 assert_eq!(result.len(), 128);
390 for v in &result {
392 assert!((v - emb_val).abs() < 1e-6, "expected {emb_val} got {v}");
393 }
394 }
395
396 #[test]
397 fn test_context_embedding_single_result_returns_that_embedding() {
398 let mut w = default_window();
399 let emb_val = 0.7_f32;
400 w.add(ContextEntry::Result {
401 doc_id: 42,
402 embedding: make_embedding(emb_val, 128),
403 relevance: 1.0,
404 tick: 0,
405 });
406 let result = w.context_embedding();
407 for v in &result {
408 assert!((v - emb_val).abs() < 1e-6);
409 }
410 }
411
412 #[test]
413 fn test_context_embedding_newer_has_higher_weight() {
414 let mut w = default_window();
420 w.add(ContextEntry::Query {
421 text: "old".into(),
422 embedding: make_embedding(0.0, 128),
423 tick: 0,
424 });
425 w.add(ContextEntry::Query {
426 text: "new".into(),
427 embedding: make_embedding(1.0, 128),
428 tick: 1,
429 });
430 let result = w.context_embedding();
431 let expected = 1.0_f32 / 1.9_f32;
432 for v in &result {
433 assert!((v - expected).abs() < 1e-5, "expected {expected} got {v}");
434 }
435 }
436
437 #[test]
438 fn test_context_embedding_skips_wrong_dimension() {
439 let config = WindowConfig {
440 embedding_dim: 4,
441 ..Default::default()
442 };
443 let mut w = SemanticContextWindow::new(config);
444 w.add(ContextEntry::Query {
446 text: "bad".into(),
447 embedding: vec![99.0; 8],
448 tick: 0,
449 });
450 w.add(ContextEntry::Query {
452 text: "good".into(),
453 embedding: vec![0.5; 4],
454 tick: 1,
455 });
456 let result = w.context_embedding();
457 assert_eq!(result.len(), 4);
459 for v in &result {
460 assert!((v - 0.5).abs() < 1e-6, "got {v}");
461 }
462 }
463
464 #[test]
465 fn test_context_embedding_all_wrong_dimension_returns_zeros() {
466 let config = WindowConfig {
467 embedding_dim: 4,
468 ..Default::default()
469 };
470 let mut w = SemanticContextWindow::new(config);
471 w.add(ContextEntry::Query {
472 text: "bad".into(),
473 embedding: vec![1.0; 8],
474 tick: 0,
475 });
476 let result = w.context_embedding();
477 assert_eq!(result.len(), 4);
478 assert!(result.iter().all(|&v| v == 0.0));
479 }
480
481 #[test]
482 fn test_context_embedding_mixed_entries_feedback_ignored() {
483 let config = WindowConfig {
484 embedding_dim: 2,
485 decay_factor: 1.0, ..Default::default()
487 };
488 let mut w = SemanticContextWindow::new(config);
489 w.add(ContextEntry::Feedback {
490 doc_id: 1,
491 positive: true,
492 tick: 0,
493 });
494 w.add(ContextEntry::Query {
495 text: "q".into(),
496 embedding: vec![2.0, 4.0],
497 tick: 1,
498 });
499 w.add(ContextEntry::Feedback {
500 doc_id: 2,
501 positive: false,
502 tick: 2,
503 });
504 w.add(ContextEntry::Result {
505 doc_id: 10,
506 embedding: vec![6.0, 8.0],
507 relevance: 0.9,
508 tick: 3,
509 });
510 let result = w.context_embedding();
512 assert!((result[0] - 4.0).abs() < 1e-5, "got {}", result[0]);
513 assert!((result[1] - 6.0).abs() < 1e-5, "got {}", result[1]);
514 }
515
516 #[test]
521 fn test_recent_queries_newest_first() {
522 let mut w = default_window();
523 for tick in 0..4u64 {
524 w.add(ContextEntry::Query {
525 text: format!("q{tick}"),
526 embedding: make_embedding(0.0, 128),
527 tick,
528 });
529 }
530 let recent = w.recent_queries(3);
531 assert_eq!(recent.len(), 3);
532 assert!(matches!(recent[0], ContextEntry::Query { tick: 3, .. }));
534 assert!(matches!(recent[1], ContextEntry::Query { tick: 2, .. }));
535 assert!(matches!(recent[2], ContextEntry::Query { tick: 1, .. }));
536 }
537
538 #[test]
539 fn test_recent_queries_fewer_than_n() {
540 let mut w = default_window();
541 w.add(ContextEntry::Query {
542 text: "only".into(),
543 embedding: make_embedding(0.1, 128),
544 tick: 0,
545 });
546 let recent = w.recent_queries(5);
547 assert_eq!(recent.len(), 1);
548 }
549
550 #[test]
551 fn test_recent_queries_ignores_non_query_entries() {
552 let mut w = default_window();
553 w.add(ContextEntry::Feedback {
554 doc_id: 1,
555 positive: true,
556 tick: 0,
557 });
558 w.add(ContextEntry::Query {
559 text: "q".into(),
560 embedding: make_embedding(0.0, 128),
561 tick: 1,
562 });
563 let recent = w.recent_queries(10);
564 assert_eq!(recent.len(), 1);
565 }
566
567 #[test]
572 fn test_recent_results_newest_first() {
573 let mut w = default_window();
574 for tick in 0..3u64 {
575 w.add(ContextEntry::Result {
576 doc_id: tick,
577 embedding: make_embedding(0.0, 128),
578 relevance: 0.5,
579 tick,
580 });
581 }
582 let recent = w.recent_results(2);
583 assert_eq!(recent.len(), 2);
584 assert!(matches!(recent[0], ContextEntry::Result { tick: 2, .. }));
585 assert!(matches!(recent[1], ContextEntry::Result { tick: 1, .. }));
586 }
587
588 #[test]
589 fn test_recent_results_ignores_non_result_entries() {
590 let mut w = default_window();
591 w.add(ContextEntry::Query {
592 text: "q".into(),
593 embedding: make_embedding(0.0, 128),
594 tick: 0,
595 });
596 let recent = w.recent_results(10);
597 assert_eq!(recent.len(), 0);
598 }
599
600 #[test]
605 fn test_positive_docs_deduplicated_oldest_first() {
606 let mut w = default_window();
607 w.add(ContextEntry::Feedback {
608 doc_id: 5,
609 positive: true,
610 tick: 0,
611 });
612 w.add(ContextEntry::Feedback {
613 doc_id: 3,
614 positive: true,
615 tick: 1,
616 });
617 w.add(ContextEntry::Feedback {
619 doc_id: 5,
620 positive: true,
621 tick: 2,
622 });
623 let docs = w.positive_docs();
624 assert_eq!(docs, vec![5, 3]);
626 }
627
628 #[test]
629 fn test_positive_docs_excludes_negatives() {
630 let mut w = default_window();
631 w.add(ContextEntry::Feedback {
632 doc_id: 1,
633 positive: false,
634 tick: 0,
635 });
636 w.add(ContextEntry::Feedback {
637 doc_id: 2,
638 positive: true,
639 tick: 1,
640 });
641 let docs = w.positive_docs();
642 assert_eq!(docs, vec![2]);
643 }
644
645 #[test]
646 fn test_negative_docs_deduplicated_oldest_first() {
647 let mut w = default_window();
648 w.add(ContextEntry::Feedback {
649 doc_id: 7,
650 positive: false,
651 tick: 0,
652 });
653 w.add(ContextEntry::Feedback {
654 doc_id: 8,
655 positive: false,
656 tick: 1,
657 });
658 w.add(ContextEntry::Feedback {
659 doc_id: 7,
660 positive: false,
661 tick: 2,
662 });
663 let docs = w.negative_docs();
664 assert_eq!(docs, vec![7, 8]);
665 }
666
667 #[test]
668 fn test_negative_docs_excludes_positives() {
669 let mut w = default_window();
670 w.add(ContextEntry::Feedback {
671 doc_id: 1,
672 positive: true,
673 tick: 0,
674 });
675 let docs = w.negative_docs();
676 assert!(docs.is_empty());
677 }
678
679 #[test]
684 fn test_stats_counts() {
685 let mut w = default_window();
686 w.add(ContextEntry::Query {
687 text: "a".into(),
688 embedding: make_embedding(0.0, 128),
689 tick: 0,
690 });
691 w.add(ContextEntry::Query {
692 text: "b".into(),
693 embedding: make_embedding(0.0, 128),
694 tick: 1,
695 });
696 w.add(ContextEntry::Result {
697 doc_id: 1,
698 embedding: make_embedding(0.0, 128),
699 relevance: 0.8,
700 tick: 2,
701 });
702 w.add(ContextEntry::Feedback {
703 doc_id: 1,
704 positive: true,
705 tick: 3,
706 });
707 w.add(ContextEntry::Feedback {
708 doc_id: 2,
709 positive: false,
710 tick: 4,
711 });
712 let s = w.stats();
713 assert_eq!(s.total_entries_added, 5);
714 assert_eq!(s.queries_in_window, 2);
715 assert_eq!(s.results_in_window, 1);
716 assert_eq!(s.feedbacks_in_window, 2);
717 }
718
719 #[test]
720 fn test_stats_positive_feedback_ratio() {
721 let mut w = default_window();
722 for i in 0..3u64 {
724 w.add(ContextEntry::Feedback {
725 doc_id: i,
726 positive: true,
727 tick: i,
728 });
729 }
730 w.add(ContextEntry::Feedback {
731 doc_id: 99,
732 positive: false,
733 tick: 3,
734 });
735 let s = w.stats();
736 assert!((s.positive_feedback_ratio - 0.75).abs() < 1e-6);
737 }
738
739 #[test]
740 fn test_stats_no_feedback_ratio_is_zero() {
741 let mut w = default_window();
742 w.add(ContextEntry::Query {
743 text: "q".into(),
744 embedding: make_embedding(0.0, 128),
745 tick: 0,
746 });
747 let s = w.stats();
748 assert_eq!(s.positive_feedback_ratio, 0.0);
749 }
750
751 #[test]
752 fn test_stats_total_added_includes_evicted() {
753 let config = WindowConfig {
754 max_entries: 2,
755 ..Default::default()
756 };
757 let mut w = SemanticContextWindow::new(config);
758 for tick in 0..10u64 {
759 w.add(ContextEntry::Feedback {
760 doc_id: tick,
761 positive: true,
762 tick,
763 });
764 }
765 let s = w.stats();
766 assert_eq!(s.total_entries_added, 10);
767 assert_eq!(s.feedbacks_in_window, 2);
768 }
769
770 #[test]
775 fn test_clear_resets_entries_preserves_config() {
776 let config = WindowConfig {
777 max_entries: 5,
778 decay_factor: 0.8,
779 embedding_dim: 64,
780 };
781 let mut w = SemanticContextWindow::new(config.clone());
782 w.add(ContextEntry::Feedback {
783 doc_id: 1,
784 positive: true,
785 tick: 0,
786 });
787 w.add(ContextEntry::Query {
788 text: "q".into(),
789 embedding: make_embedding(0.0, 64),
790 tick: 1,
791 });
792 w.clear();
793 assert!(w.entries.is_empty());
794 assert_eq!(w.config, config);
795 assert_eq!(w.total_added, 2);
797 }
798
799 #[test]
800 fn test_clear_then_add_works() {
801 let mut w = default_window();
802 w.add(ContextEntry::Feedback {
803 doc_id: 1,
804 positive: true,
805 tick: 0,
806 });
807 w.clear();
808 w.add(ContextEntry::Feedback {
809 doc_id: 2,
810 positive: false,
811 tick: 1,
812 });
813 assert_eq!(w.entries.len(), 1);
814 assert_eq!(w.total_added, 2);
815 }
816
817 #[test]
822 fn test_recent_queries_zero_n_returns_empty() {
823 let mut w = default_window();
824 w.add(ContextEntry::Query {
825 text: "q".into(),
826 embedding: make_embedding(0.0, 128),
827 tick: 0,
828 });
829 assert!(w.recent_queries(0).is_empty());
830 }
831
832 #[test]
833 fn test_recent_results_zero_n_returns_empty() {
834 let mut w = default_window();
835 w.add(ContextEntry::Result {
836 doc_id: 1,
837 embedding: make_embedding(0.0, 128),
838 relevance: 1.0,
839 tick: 0,
840 });
841 assert!(w.recent_results(0).is_empty());
842 }
843
844 #[test]
845 fn test_context_embedding_three_queries_recency_weighted() {
846 let config = WindowConfig {
855 embedding_dim: 2,
856 decay_factor: 0.5,
857 ..Default::default()
858 };
859 let mut w = SemanticContextWindow::new(config);
860 for (i, val) in [(0.0_f32), (1.0), (2.0)].iter().enumerate() {
861 w.add(ContextEntry::Query {
862 text: format!("q{i}"),
863 embedding: vec![*val; 2],
864 tick: i as u64,
865 });
866 }
867 let result = w.context_embedding();
868 let expected = 2.5_f32 / 1.75_f32;
869 assert!((result[0] - expected).abs() < 1e-5, "got {}", result[0]);
870 assert!((result[1] - expected).abs() < 1e-5, "got {}", result[1]);
871 }
872
873 #[test]
874 fn test_positive_and_negative_docs_same_doc_id() {
875 let mut w = default_window();
877 w.add(ContextEntry::Feedback {
878 doc_id: 5,
879 positive: true,
880 tick: 0,
881 });
882 w.add(ContextEntry::Feedback {
883 doc_id: 5,
884 positive: false,
885 tick: 1,
886 });
887 let pos = w.positive_docs();
888 let neg = w.negative_docs();
889 assert_eq!(pos, vec![5]);
890 assert_eq!(neg, vec![5]);
891 }
892
893 #[test]
894 fn test_window_default_config() {
895 let c = WindowConfig::default();
896 assert_eq!(c.max_entries, 20);
897 assert!((c.decay_factor - 0.9).abs() < 1e-6);
898 assert_eq!(c.embedding_dim, 128);
899 }
900}