poimen-memory/crates/mem-ingest/src/optimizer_sink.rs

//! M3.8 Ingest-Time Context Optimization
//!
//! Provides helpers for optimizing records at ingest time.
//! Designed to be called from rebuild.rs during the embed + index pipeline.

use mem_core::{Record, ContextOptimizer};
use std::sync::{Arc, Mutex};

/// Tracks optimization effectiveness across records.
#[derive(Debug, Clone, Default)]
pub struct OptimizationMetrics {
    /// Total records processed
    pub total_records: usize,
    /// Sum of input bytes (before optimization)
    pub input_bytes_total: usize,
    /// Sum of output bytes (after optimization)
    pub output_bytes_total: usize,
    /// Per-compressor breakdown
    pub per_compressor: std::collections::HashMap<String, CompressorStats>,
}

#[derive(Debug, Clone)]
pub struct CompressorStats {
    pub count: usize,
    pub input_bytes: usize,
    pub output_bytes: usize,
}

impl OptimizationMetrics {
    /// Calculate overall compression ratio as percentage
    pub fn compression_ratio(&self) -> f32 {
        if self.input_bytes_total == 0 {
            0.0
        } else {
            (self.output_bytes_total as f32 / self.input_bytes_total as f32) * 100.0
        }
    }

    /// Get stats for a specific compressor
    pub fn get_compressor_stats(&self, name: &str) -> Option<&CompressorStats> {
        self.per_compressor.get(name)
    }

    /// Log metrics summary
    pub fn log_summary(&self, project: &str) {
        tracing::info!(
            project = project,
            total_records = self.total_records,
            input_bytes = self.input_bytes_total,
            output_bytes = self.output_bytes_total,
            compression_ratio = format!("{:.1}%", self.compression_ratio()),
            "M3.8 ingest optimization complete"
        );

        for (name, stats) in &self.per_compressor {
            let ratio = if stats.input_bytes == 0 {
                0.0
            } else {
                (stats.output_bytes as f32 / stats.input_bytes as f32) * 100.0
            };
            tracing::debug!(
                compressor = name,
                count = stats.count,
                input_bytes = stats.input_bytes,
                output_bytes = stats.output_bytes,
                ratio = format!("{:.1}%", ratio),
                "compressor stats"
            );
        }
    }
}

/// Helper to optimize a single record and track metrics.
///
/// Call this in the rebuild loop when processing each record:
/// ```ignore
/// for record in source.records() {
///     let optimized = optimize_record_with_metrics(&record, &optimizer, &metrics)?;
///     embed_and_index(&optimized)?;
/// }
/// ```
pub fn optimize_record_with_metrics(
    record: Record,
    optimizer: &ContextOptimizer,
    metrics: &Arc<Mutex<OptimizationMetrics>>,
) -> Result<Record, String> {
    let input_bytes = record.text.len();

    let optimized = match optimizer.optimize(&record.text) {
        Ok(opt_chunk) => opt_chunk,
        Err(e) => {
            tracing::warn!(error = ?e, "optimization failed, using original");
            return Ok(record);
        }
    };

    let output_bytes = optimized.compressed.len();
    let compressor_name = format!("{:?}", optimized.content_type);

    {
        let mut m = metrics.lock().unwrap();
        m.total_records += 1;
        m.input_bytes_total += input_bytes;
        m.output_bytes_total += output_bytes;

        m.per_compressor
            .entry(compressor_name.clone())
            .or_insert_with(|| CompressorStats {
                count: 0,
                input_bytes: 0,
                output_bytes: 0,
            })
            .count += 1;

        let stats = m.per_compressor.get_mut(&compressor_name).unwrap();
        stats.input_bytes += input_bytes;
        stats.output_bytes += output_bytes;
    }

    Ok(Record {
        text: optimized.compressed,
        ..record
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use mem_core::{Provenance, Role};
    use time::macros::datetime;

    fn make_test_record(text: &str) -> Record {
        Record {
            role: Role::User,
            text: text.to_string(),
            timestamp: datetime!(2024-08-20 12:00:00 UTC),
            provenance: Provenance {
                source_id: "test".to_string(),
                offset: 0,
            },
        }
    }

    #[test]
    fn test_optimize_record_preserves_structure() {
        let record = make_test_record("ERROR: simple");
        let optimizer = ContextOptimizer::new().unwrap();
        let metrics = Arc::new(Mutex::new(OptimizationMetrics::default()));

        let result = optimize_record_with_metrics(record.clone(), &optimizer, &metrics);
        assert!(result.is_ok());

        let optimized = result.unwrap();
        assert_eq!(optimized.role, Role::User);
        assert_eq!(optimized.provenance.source_id, "test");
    }

    #[test]
    fn test_optimize_record_tracks_bytes() {
        let record = make_test_record("ERROR: failed\nINFO: debug\nERROR: permission");
        let optimizer = ContextOptimizer::new().unwrap();
        let metrics = Arc::new(Mutex::new(OptimizationMetrics::default()));

        let _ = optimize_record_with_metrics(record, &optimizer, &metrics);

        let m = metrics.lock().unwrap();
        assert_eq!(m.total_records, 1);
        assert!(m.input_bytes_total > 0);
        assert!(m.output_bytes_total > 0);
    }

    #[test]
    fn test_optimize_record_disabled() {
        std::env::set_var("MEM_CONTEXT_OPTIMIZER", "off");
        let record = make_test_record("content");
        let optimizer = ContextOptimizer::from_env().unwrap();
        let metrics = Arc::new(Mutex::new(OptimizationMetrics::default()));

        let result = optimize_record_with_metrics(record, &optimizer, &metrics);
        assert!(result.is_ok());
    }

    #[test]
    fn test_compression_ratio_calculation() {
        let metrics = OptimizationMetrics {
            total_records: 10,
            input_bytes_total: 1000,
            output_bytes_total: 500,
            per_compressor: Default::default(),
        };

        let ratio = metrics.compression_ratio();
        assert!((ratio - 50.0).abs() < 0.1, "should be 50%");
    }

    #[test]
    fn test_metrics_aggregation() {
        let metrics = Arc::new(Mutex::new(OptimizationMetrics::default()));
        let optimizer = ContextOptimizer::new().unwrap();

        for i in 0..3 {
            let record = make_test_record(&format!("ERROR: {}", i));
            let _ = optimize_record_with_metrics(record, &optimizer, &metrics);
        }

        let m = metrics.lock().unwrap();
        assert_eq!(m.total_records, 3);
    }
}