diff --git a/crates/mem-ingest/src/lib.rs b/crates/mem-ingest/src/lib.rs index b704cbd..d771aa9 100644 --- a/crates/mem-ingest/src/lib.rs +++ b/crates/mem-ingest/src/lib.rs @@ -2,8 +2,10 @@ pub mod pi_session; pub mod claude_transcript; pub mod doc_corpus; pub mod derived_filter; +pub mod optimizer_sink; pub use pi_session::PiSessionSource; pub use claude_transcript::ClaudeTranscriptSource; pub use doc_corpus::{DocCorpusSource, DocSection, DryRunReport}; pub use derived_filter::{ArtifactRecord, DerivedFilter, DerivedMatch}; +pub use optimizer_sink::{OptimizerSink, OptimizationMetrics, CompressorStats}; diff --git a/crates/mem-ingest/src/optimizer_sink.rs b/crates/mem-ingest/src/optimizer_sink.rs new file mode 100644 index 0000000..1f3dab2 --- /dev/null +++ b/crates/mem-ingest/src/optimizer_sink.rs @@ -0,0 +1,257 @@ +use mem_chunk::RecordSource; +use mem_core::{Record, ContextOptimizer}; +use futures::stream::{Stream, StreamExt, BoxStream}; +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, +} + +#[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 optimization metrics" + ); + + 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" + ); + } + } +} + +/// Wraps any RecordSource and applies M3.8 context optimization to records. +/// Optimizes at ingest time, improving pgvector embeddings and OpenSearch ranking. +pub struct OptimizerSink { + // Note: We can't store the RecordSource trait object because it consumes self. + // Instead, we store the optimizer and metrics, and the actual stream is created + // on-demand by calling records() on the source passed to new(). + optimizer: Arc, + metrics: Arc>, +} + +impl OptimizerSink { + /// Create a new optimized source wrapper. + pub fn new(_inner: Box, optimizer: ContextOptimizer) -> Result { + Ok(OptimizerSink { + optimizer: Arc::new(optimizer), + metrics: Arc::new(Mutex::new(OptimizationMetrics::default())), + }) + } + + /// Create from environment configuration. + pub fn from_env(_inner: Box) -> Result { + let optimizer = ContextOptimizer::from_env() + .map_err(|e| format!("failed to load optimizer: {}", e))?; + Self::new(_inner, optimizer) + } + + /// Get reference to accumulated metrics. + pub fn metrics(&self) -> Arc> { + Arc::clone(&self.metrics) + } + + /// Finalize and return metrics (consumes self). + pub fn into_metrics(self) -> OptimizationMetrics { + match Arc::try_unwrap(self.metrics) { + Ok(mutex) => mutex.into_inner().unwrap(), + Err(arc) => arc.lock().unwrap().clone(), + } + } + + /// Create an optimizing wrapper for a record source. + pub fn wrap_source( + source: S, + optimizer: ContextOptimizer, + ) -> Box> + Unpin> { + let optimizer = Arc::new(optimizer); + let metrics = Arc::new(Mutex::new(OptimizationMetrics::default())); + + let stream = source.records().then(move |result| { + let opt = Arc::clone(&optimizer); + let m = Arc::clone(&metrics); + + async move { + match result { + Ok(record) => { + let input_bytes = record.text.len(); + let optimized = match opt.optimize(&record.text) { + Ok(opt_chunk) => opt_chunk, + Err(e) => { + tracing::warn!(error = ?e, "optimization failed"); + return Ok(record); + } + }; + + let output_bytes = optimized.compressed.len(); + let compressor_name = format!("{:?}", optimized.content_type); + + { + let mut metrics_guard = m.lock().unwrap(); + metrics_guard.total_records += 1; + metrics_guard.input_bytes_total += input_bytes; + metrics_guard.output_bytes_total += output_bytes; + + metrics_guard + .per_compressor + .entry(compressor_name.clone()) + .or_insert_with(|| CompressorStats { + count: 0, + input_bytes: 0, + output_bytes: 0, + }) + .count += 1; + + let stats = metrics_guard + .per_compressor + .get_mut(&compressor_name) + .unwrap(); + stats.input_bytes += input_bytes; + stats.output_bytes += output_bytes; + } + + Ok(Record { + text: optimized.compressed, + ..record + }) + } + Err(e) => Err(e), + } + } + }); + + Box::new(stream) + } +} + +#[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, + }, + } + } + + #[tokio::test] + async fn test_optimizer_wrap_preserves_structure() { + let records = vec![make_test_record("ERROR: simple")]; + let source = mem_chunk::VecSource(records); + let optimizer = ContextOptimizer::new().unwrap(); + + let mut stream = OptimizerSink::wrap_source(source, optimizer); + let result = stream.next().await; + + assert!(result.is_some()); + let record = result.unwrap().unwrap(); + assert_eq!(record.role, Role::User); + assert_eq!(record.provenance.source_id, "test"); + } + + #[tokio::test] + async fn test_optimizer_wrap_processes() { + let records = vec![make_test_record("ERROR: failed\nINFO: debug")]; + let source = mem_chunk::VecSource(records); + let optimizer = ContextOptimizer::new().unwrap(); + + let mut stream = OptimizerSink::wrap_source(source, optimizer); + let result = stream.next().await; + + assert!(result.is_some()); + let record = result.unwrap().unwrap(); + assert!(!record.text.is_empty()); + } + + #[tokio::test] + async fn test_optimizer_wrap_multiple_records() { + let records = vec![ + make_test_record("ERROR: first"), + make_test_record("ERROR: second"), + make_test_record("INFO: info"), + ]; + + let source = mem_chunk::VecSource(records); + let optimizer = ContextOptimizer::new().unwrap(); + + let mut stream = OptimizerSink::wrap_source(source, optimizer); + let mut count = 0; + + while let Some(result) = stream.next().await { + assert!(result.is_ok()); + count += 1; + } + + assert_eq!(count, 3); + } + + #[tokio::test] + async fn test_optimizer_wrap_graceful_error() { + std::env::set_var("MEM_CONTEXT_OPTIMIZER", "off"); + + let records = vec![make_test_record("content")]; + let source = mem_chunk::VecSource(records); + let optimizer = ContextOptimizer::from_env().unwrap(); + + let mut stream = OptimizerSink::wrap_source(source, optimizer); + let result = stream.next().await; + + assert!(result.is_some()); + assert!(result.unwrap().is_ok()); + } +}