1#[derive(Clone, Debug, PartialEq)]
8pub enum DriftSignal {
9 NoDrift,
11 MinorDrift { magnitude: f64 },
13 MajorDrift { magnitude: f64 },
15 InsufficientData,
17}
18
19#[derive(Clone, Debug)]
21pub struct BaselineStats {
22 pub mean: Vec<f32>,
24 pub std: Vec<f32>,
26 pub sample_count: usize,
28}
29
30impl BaselineStats {
31 pub fn compute(samples: &[Vec<f32>]) -> Option<Self> {
35 if samples.is_empty() {
36 return None;
37 }
38 let dims = samples[0].len();
39 if dims == 0 {
40 return None;
41 }
42 let n = samples.len() as f64;
43
44 let mut mean = vec![0.0_f64; dims];
46 for vec in samples {
47 for (d, &v) in vec.iter().enumerate() {
48 mean[d] += v as f64;
49 }
50 }
51 for m in &mut mean {
52 *m /= n;
53 }
54
55 let mut variance = vec![0.0_f64; dims];
57 for vec in samples {
58 for (d, &v) in vec.iter().enumerate() {
59 let diff = v as f64 - mean[d];
60 variance[d] += diff * diff;
61 }
62 }
63 let mean_f32: Vec<f32> = mean.iter().map(|&m| m as f32).collect();
64 let std_f32: Vec<f32> = variance
65 .iter()
66 .map(|&v| {
67 let s = (v / n).sqrt();
68 s.max(1e-6_f64) as f32
69 })
70 .collect();
71
72 Some(BaselineStats {
73 mean: mean_f32,
74 std: std_f32,
75 sample_count: samples.len(),
76 })
77 }
78}
79
80#[derive(Clone, Debug)]
82pub struct DriftMonitorConfig {
83 pub min_baseline_samples: usize,
85 pub minor_drift_threshold: f64,
87 pub major_drift_threshold: f64,
89 pub window_size: usize,
91 pub baseline_update_interval: usize,
93}
94
95impl Default for DriftMonitorConfig {
96 fn default() -> Self {
97 Self {
98 min_baseline_samples: 50,
99 minor_drift_threshold: 0.1,
100 major_drift_threshold: 0.3,
101 window_size: 100,
102 baseline_update_interval: 25,
103 }
104 }
105}
106
107#[derive(Clone, Debug, Default)]
109pub struct DriftMonitorStats {
110 pub total_checked: u64,
112 pub minor_drift_count: u64,
114 pub major_drift_count: u64,
116 pub baseline_updates: u64,
118}
119
120impl DriftMonitorStats {
121 pub fn drift_rate(&self) -> f64 {
125 let total = self.total_checked.max(1);
126 (self.minor_drift_count + self.major_drift_count) as f64 / total as f64
127 }
128}
129
130pub struct EmbeddingDriftMonitor {
136 pub baseline: Option<BaselineStats>,
138 pub recent: Vec<Vec<f32>>,
140 pub config: DriftMonitorConfig,
142 pub stats: DriftMonitorStats,
144 pub update_counter: usize,
146}
147
148impl EmbeddingDriftMonitor {
149 pub fn new(config: DriftMonitorConfig) -> Self {
151 Self {
152 baseline: None,
153 recent: Vec::new(),
154 config,
155 stats: DriftMonitorStats::default(),
156 update_counter: 0,
157 }
158 }
159
160 pub fn add_vector(&mut self, vec: Vec<f32>) -> DriftSignal {
170 self.recent.push(vec.clone());
172 if self.recent.len() > self.config.window_size {
173 self.recent.remove(0);
174 }
175
176 self.update_counter += 1;
178 if self.update_counter >= self.config.baseline_update_interval
179 && self.recent.len() >= self.config.min_baseline_samples
180 {
181 if let Some(new_baseline) = BaselineStats::compute(&self.recent) {
182 self.baseline = Some(new_baseline);
183 self.update_counter = 0;
184 self.stats.baseline_updates += 1;
185 }
186 }
187
188 let baseline = match &self.baseline {
190 Some(b) => b,
191 None => return DriftSignal::InsufficientData,
192 };
193
194 let dims = baseline.mean.len();
196 if dims == 0 || vec.len() != dims {
197 return DriftSignal::InsufficientData;
198 }
199
200 let magnitude: f64 = vec
201 .iter()
202 .enumerate()
203 .map(|(i, &v)| {
204 let deviation = (v as f64 - baseline.mean[i] as f64) / baseline.std[i] as f64;
205 deviation.abs()
206 })
207 .sum::<f64>()
208 / dims as f64;
209
210 self.stats.total_checked += 1;
211
212 if magnitude >= self.config.major_drift_threshold {
213 self.stats.major_drift_count += 1;
214 DriftSignal::MajorDrift { magnitude }
215 } else if magnitude >= self.config.minor_drift_threshold {
216 self.stats.minor_drift_count += 1;
217 DriftSignal::MinorDrift { magnitude }
218 } else {
219 DriftSignal::NoDrift
220 }
221 }
222
223 pub fn baseline(&self) -> Option<&BaselineStats> {
225 self.baseline.as_ref()
226 }
227
228 pub fn stats(&self) -> &DriftMonitorStats {
230 &self.stats
231 }
232
233 pub fn reset_baseline(&mut self) {
235 self.baseline = None;
236 self.recent.clear();
237 self.update_counter = 0;
238 }
239}
240
241#[cfg(test)]
242mod tests {
243 use super::*;
244
245 fn small_config(min_baseline: usize, window: usize, interval: usize) -> DriftMonitorConfig {
247 DriftMonitorConfig {
248 min_baseline_samples: min_baseline,
249 window_size: window,
250 baseline_update_interval: interval,
251 minor_drift_threshold: 0.1,
252 major_drift_threshold: 0.3,
253 }
254 }
255
256 fn const_vec(dim: usize, value: f32) -> Vec<f32> {
258 vec![value; dim]
259 }
260
261 fn seed_monitor(monitor: &mut EmbeddingDriftMonitor, n: usize, dim: usize, value: f32) {
263 for _ in 0..n {
264 monitor.add_vector(const_vec(dim, value));
265 }
266 }
267
268 #[test]
272 fn test_insufficient_data_initially() {
273 let cfg = small_config(5, 20, 5);
274 let mut monitor = EmbeddingDriftMonitor::new(cfg);
275 let signal = monitor.add_vector(const_vec(4, 0.5));
276 assert_eq!(signal, DriftSignal::InsufficientData);
277 }
278
279 #[test]
280 fn test_insufficient_data_before_threshold() {
281 let cfg = small_config(5, 20, 5);
282 let mut monitor = EmbeddingDriftMonitor::new(cfg);
283 for _ in 0..4 {
285 let s = monitor.add_vector(const_vec(4, 0.5));
286 assert_eq!(s, DriftSignal::InsufficientData);
287 }
288 }
289
290 #[test]
294 fn test_baseline_computed_after_enough_samples() {
295 let cfg = small_config(5, 20, 5);
296 let mut monitor = EmbeddingDriftMonitor::new(cfg);
297 seed_monitor(&mut monitor, 5, 4, 1.0);
298 assert!(monitor.baseline().is_some());
299 }
300
301 #[test]
302 fn test_baseline_mean_values() {
303 let cfg = small_config(4, 20, 4);
304 let mut monitor = EmbeddingDriftMonitor::new(cfg);
305 seed_monitor(&mut monitor, 4, 3, 2.0);
307 let baseline = monitor.baseline().expect("baseline should exist");
308 for &m in &baseline.mean {
309 assert!((m - 2.0_f32).abs() < 1e-4, "mean={m}");
310 }
311 }
312
313 #[test]
317 fn test_no_drift_for_identical_vector() {
318 let cfg = small_config(5, 20, 5);
319 let mut monitor = EmbeddingDriftMonitor::new(cfg);
320 for _ in 0..4 {
322 monitor.add_vector(const_vec(4, 0.0));
323 }
324 let signal = monitor.add_vector(const_vec(4, 0.0));
326 assert_eq!(signal, DriftSignal::NoDrift);
328 }
329
330 #[test]
331 fn test_no_drift_small_deviation() {
332 let cfg = small_config(5, 20, 5);
333 let mut monitor = EmbeddingDriftMonitor::new(cfg);
334 seed_monitor(&mut monitor, 5, 4, 1.0);
335 let signal = monitor.add_vector(const_vec(4, 1.0));
339 assert_eq!(signal, DriftSignal::NoDrift);
340 }
341
342 #[test]
346 fn test_minor_drift_detected() {
347 let cfg = small_config(4, 20, 4);
349 let mut monitor = EmbeddingDriftMonitor::new(cfg);
350 monitor.add_vector(vec![0.0, 0.0, 0.0, 0.0]);
352 monitor.add_vector(vec![2.0, 2.0, 2.0, 2.0]);
353 monitor.add_vector(vec![0.0, 0.0, 0.0, 0.0]);
354 monitor.add_vector(vec![2.0, 2.0, 2.0, 2.0]);
356 let signal = monitor.add_vector(vec![1.15, 1.15, 1.15, 1.15]);
358 match signal {
359 DriftSignal::MinorDrift { magnitude } => {
360 assert!((0.1..0.3).contains(&magnitude), "magnitude={magnitude}");
361 }
362 other => panic!("expected MinorDrift, got {other:?}"),
363 }
364 }
365
366 #[test]
367 fn test_minor_drift_stats_increment() {
368 let cfg = small_config(4, 20, 4);
369 let mut monitor = EmbeddingDriftMonitor::new(cfg);
370 monitor.add_vector(vec![0.0; 4]);
371 monitor.add_vector(vec![2.0; 4]);
372 monitor.add_vector(vec![0.0; 4]);
373 monitor.add_vector(vec![2.0; 4]);
374 monitor.add_vector(vec![1.15; 4]);
375 assert_eq!(monitor.stats().minor_drift_count, 1);
376 }
377
378 #[test]
382 fn test_major_drift_detected() {
383 let cfg = small_config(4, 20, 4);
384 let mut monitor = EmbeddingDriftMonitor::new(cfg);
385 monitor.add_vector(vec![0.0; 4]);
386 monitor.add_vector(vec![2.0; 4]);
387 monitor.add_vector(vec![0.0; 4]);
388 monitor.add_vector(vec![2.0; 4]);
389 let signal = monitor.add_vector(vec![5.0; 4]);
391 match signal {
392 DriftSignal::MajorDrift { magnitude } => {
393 assert!(magnitude >= 0.3, "magnitude={magnitude}");
394 }
395 other => panic!("expected MajorDrift, got {other:?}"),
396 }
397 }
398
399 #[test]
400 fn test_major_drift_stats_increment() {
401 let cfg = small_config(4, 20, 4);
402 let mut monitor = EmbeddingDriftMonitor::new(cfg);
403 monitor.add_vector(vec![0.0; 4]);
406 monitor.add_vector(vec![2.0; 4]);
407 monitor.add_vector(vec![0.0; 4]);
408 monitor.add_vector(vec![2.0; 4]);
409 monitor.add_vector(vec![5.0; 4]);
411 assert_eq!(monitor.stats().major_drift_count, 2);
413 }
414
415 #[test]
419 fn test_sliding_window_size_capped() {
420 let cfg = small_config(3, 5, 3);
421 let mut monitor = EmbeddingDriftMonitor::new(cfg);
422 for i in 0..10_u8 {
423 monitor.add_vector(vec![i as f32; 2]);
424 }
425 assert!(
426 monitor.recent.len() <= 5,
427 "window exceeded: {}",
428 monitor.recent.len()
429 );
430 }
431
432 #[test]
433 fn test_sliding_window_oldest_evicted() {
434 let cfg = small_config(2, 3, 2);
435 let mut monitor = EmbeddingDriftMonitor::new(cfg);
436 monitor.add_vector(vec![1.0, 1.0]);
437 monitor.add_vector(vec![2.0, 2.0]);
438 monitor.add_vector(vec![3.0, 3.0]);
439 monitor.add_vector(vec![4.0, 4.0]);
441 assert!(!monitor.recent.contains(&vec![1.0, 1.0]));
443 assert_eq!(monitor.recent.len(), 3);
444 }
445
446 #[test]
450 fn test_baseline_update_increments_counter() {
451 let cfg = small_config(3, 20, 3);
453 let mut monitor = EmbeddingDriftMonitor::new(cfg);
454 seed_monitor(&mut monitor, 3, 2, 1.0);
456 assert_eq!(monitor.stats().baseline_updates, 1);
457 seed_monitor(&mut monitor, 3, 2, 1.0);
459 assert_eq!(monitor.stats().baseline_updates, 2);
460 }
461
462 #[test]
463 fn test_no_baseline_update_before_interval() {
464 let cfg = small_config(3, 20, 5);
465 let mut monitor = EmbeddingDriftMonitor::new(cfg);
466 seed_monitor(&mut monitor, 3, 2, 1.0);
468 assert_eq!(monitor.stats().baseline_updates, 0);
469 assert!(monitor.baseline().is_none());
471 }
472
473 #[test]
477 fn test_drift_rate_zero_initially() {
478 let stats = DriftMonitorStats::default();
479 assert_eq!(stats.drift_rate(), 0.0);
480 }
481
482 #[test]
483 fn test_drift_rate_calculation() {
484 let stats = DriftMonitorStats {
485 total_checked: 10,
486 minor_drift_count: 2,
487 major_drift_count: 3,
488 baseline_updates: 1,
489 };
490 let rate = stats.drift_rate();
491 assert!((rate - 0.5).abs() < 1e-9, "rate={rate}");
492 }
493
494 #[test]
495 fn test_drift_rate_all_drift() {
496 let stats = DriftMonitorStats {
497 total_checked: 4,
498 minor_drift_count: 4,
499 major_drift_count: 0,
500 baseline_updates: 0,
501 };
502 assert!((stats.drift_rate() - 1.0).abs() < 1e-9);
503 }
504
505 #[test]
509 fn test_reset_baseline_clears_baseline() {
510 let cfg = small_config(3, 20, 3);
511 let mut monitor = EmbeddingDriftMonitor::new(cfg);
512 seed_monitor(&mut monitor, 3, 2, 1.0);
513 assert!(monitor.baseline().is_some());
514 monitor.reset_baseline();
515 assert!(monitor.baseline().is_none());
516 }
517
518 #[test]
519 fn test_reset_baseline_clears_recent() {
520 let cfg = small_config(3, 20, 3);
521 let mut monitor = EmbeddingDriftMonitor::new(cfg);
522 seed_monitor(&mut monitor, 3, 2, 1.0);
523 monitor.reset_baseline();
524 assert!(monitor.recent.is_empty());
525 }
526
527 #[test]
528 fn test_reset_baseline_resets_update_counter() {
529 let cfg = small_config(3, 20, 3);
530 let mut monitor = EmbeddingDriftMonitor::new(cfg);
531 seed_monitor(&mut monitor, 3, 2, 1.0);
532 monitor.reset_baseline();
533 assert_eq!(monitor.update_counter, 0);
534 }
535
536 #[test]
540 fn test_baseline_updates_exact_count() {
541 let cfg = small_config(2, 20, 2);
543 let mut monitor = EmbeddingDriftMonitor::new(cfg);
544 for _ in 0..8 {
545 monitor.add_vector(const_vec(2, 1.0));
546 }
547 assert_eq!(monitor.stats().baseline_updates, 4);
549 }
550
551 #[test]
555 fn test_total_checked_only_after_baseline() {
556 let cfg = small_config(3, 20, 3);
557 let mut monitor = EmbeddingDriftMonitor::new(cfg);
558 seed_monitor(&mut monitor, 3, 2, 1.0);
560 assert_eq!(monitor.stats().total_checked, 1);
562 }
563}