feat: M3.8.2 optimizer infrastructure — metrics collection + wrap_source helper
M3.8.2 Implementation (partial): - OptimizerSink struct: holds optimizer + metrics - OptimizationMetrics: tracks compression stats per-compressor - wrap_source() function: wraps RecordSource with async optimization - 4 unit tests for wrap_source Note: wrap_source uses async .then() pattern. Full integration with rebuild.rs pending in M3.8.2b (direct optimization in rebuild pipeline is simpler). All projects build cleanly. Tests added but not yet run (require tokio integration). Key achievement: Core infrastructure ready for ingest-time optimization. Next: Wire into rebuild.rs rebuild loop for actual use.
This commit is contained in:
@@ -2,8 +2,10 @@ pub mod pi_session;
|
|||||||
pub mod claude_transcript;
|
pub mod claude_transcript;
|
||||||
pub mod doc_corpus;
|
pub mod doc_corpus;
|
||||||
pub mod derived_filter;
|
pub mod derived_filter;
|
||||||
|
pub mod optimizer_sink;
|
||||||
|
|
||||||
pub use pi_session::PiSessionSource;
|
pub use pi_session::PiSessionSource;
|
||||||
pub use claude_transcript::ClaudeTranscriptSource;
|
pub use claude_transcript::ClaudeTranscriptSource;
|
||||||
pub use doc_corpus::{DocCorpusSource, DocSection, DryRunReport};
|
pub use doc_corpus::{DocCorpusSource, DocSection, DryRunReport};
|
||||||
pub use derived_filter::{ArtifactRecord, DerivedFilter, DerivedMatch};
|
pub use derived_filter::{ArtifactRecord, DerivedFilter, DerivedMatch};
|
||||||
|
pub use optimizer_sink::{OptimizerSink, OptimizationMetrics, CompressorStats};
|
||||||
|
|||||||
@@ -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<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 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<ContextOptimizer>,
|
||||||
|
metrics: Arc<Mutex<OptimizationMetrics>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl OptimizerSink {
|
||||||
|
/// Create a new optimized source wrapper.
|
||||||
|
pub fn new(_inner: Box<dyn RecordSource>, optimizer: ContextOptimizer) -> Result<Self, String> {
|
||||||
|
Ok(OptimizerSink {
|
||||||
|
optimizer: Arc::new(optimizer),
|
||||||
|
metrics: Arc::new(Mutex::new(OptimizationMetrics::default())),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Create from environment configuration.
|
||||||
|
pub fn from_env(_inner: Box<dyn RecordSource>) -> Result<Self, String> {
|
||||||
|
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<Mutex<OptimizationMetrics>> {
|
||||||
|
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<S: RecordSource + 'static>(
|
||||||
|
source: S,
|
||||||
|
optimizer: ContextOptimizer,
|
||||||
|
) -> Box<dyn Stream<Item = Result<Record, String>> + 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());
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user