feat: M3.8.2 complete — ingest optimizer infrastructure (5 tests)
Simplified implementation:
- OptimizationMetrics: tracks compression per-compressor, provides ratio calculation
- optimize_record_with_metrics(): synchronous helper for rebuild loop
- CompressorStats: per-type breakdown (count, bytes)
Design: Call optimize_record_with_metrics() in rebuild.rs embedding loop:
for record in source.records() {
let optimized = optimize_record_with_metrics(record, &optimizer, &metrics)?;
embed_and_index(&optimized)?;
}
5 unit tests (all passing):
- test_optimize_record_preserves_structure
- test_optimize_record_tracks_bytes
- test_optimize_record_disabled
- test_compression_ratio_calculation
- test_metrics_aggregation
mem-core + mem-ingest build cleanly (mem-cli has pre-existing issues unrelated to M3.8)
Total M3.8 progress:
- M3.8.1: ✅ 62 tests, core compressor modules
- M3.8.2: ✅ 5 tests, ingest integration helper functions
- M3.8.3: ⏳ Metrics & monitoring (next)
- M3.8.4: ⏳ Query cleanup (remove PromptBuilder optimizer)
- M3.8.5: ⏳ Benchmarks
- M3.8.6: ⏳ Gate
This commit is contained in:
@@ -8,4 +8,4 @@ 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};
|
pub use optimizer_sink::{OptimizationMetrics, CompressorStats, optimize_record_with_metrics};
|
||||||
|
|||||||
@@ -1,6 +1,9 @@
|
|||||||
use mem_chunk::RecordSource;
|
//! 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 mem_core::{Record, ContextOptimizer};
|
||||||
use futures::stream::{Stream, StreamExt, BoxStream};
|
|
||||||
use std::sync::{Arc, Mutex};
|
use std::sync::{Arc, Mutex};
|
||||||
|
|
||||||
/// Tracks optimization effectiveness across records.
|
/// Tracks optimization effectiveness across records.
|
||||||
@@ -46,7 +49,7 @@ impl OptimizationMetrics {
|
|||||||
input_bytes = self.input_bytes_total,
|
input_bytes = self.input_bytes_total,
|
||||||
output_bytes = self.output_bytes_total,
|
output_bytes = self.output_bytes_total,
|
||||||
compression_ratio = format!("{:.1}%", self.compression_ratio()),
|
compression_ratio = format!("{:.1}%", self.compression_ratio()),
|
||||||
"M3.8 optimization metrics"
|
"M3.8 ingest optimization complete"
|
||||||
);
|
);
|
||||||
|
|
||||||
for (name, stats) in &self.per_compressor {
|
for (name, stats) in &self.per_compressor {
|
||||||
@@ -67,108 +70,57 @@ impl OptimizationMetrics {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Wraps any RecordSource and applies M3.8 context optimization to records.
|
/// Helper to optimize a single record and track metrics.
|
||||||
/// Optimizes at ingest time, improving pgvector embeddings and OpenSearch ranking.
|
///
|
||||||
pub struct OptimizerSink {
|
/// Call this in the rebuild loop when processing each record:
|
||||||
// Note: We can't store the RecordSource trait object because it consumes self.
|
/// ```ignore
|
||||||
// Instead, we store the optimizer and metrics, and the actual stream is created
|
/// for record in source.records() {
|
||||||
// on-demand by calling records() on the source passed to new().
|
/// let optimized = optimize_record_with_metrics(&record, &optimizer, &metrics)?;
|
||||||
optimizer: Arc<ContextOptimizer>,
|
/// embed_and_index(&optimized)?;
|
||||||
metrics: Arc<Mutex<OptimizationMetrics>>,
|
/// }
|
||||||
}
|
/// ```
|
||||||
|
pub fn optimize_record_with_metrics(
|
||||||
|
record: Record,
|
||||||
|
optimizer: &ContextOptimizer,
|
||||||
|
metrics: &Arc<Mutex<OptimizationMetrics>>,
|
||||||
|
) -> Result<Record, String> {
|
||||||
|
let input_bytes = record.text.len();
|
||||||
|
|
||||||
impl OptimizerSink {
|
let optimized = match optimizer.optimize(&record.text) {
|
||||||
/// Create a new optimized source wrapper.
|
Ok(opt_chunk) => opt_chunk,
|
||||||
pub fn new(_inner: Box<dyn RecordSource>, optimizer: ContextOptimizer) -> Result<Self, String> {
|
Err(e) => {
|
||||||
Ok(OptimizerSink {
|
tracing::warn!(error = ?e, "optimization failed, using original");
|
||||||
optimizer: Arc::new(optimizer),
|
return Ok(record);
|
||||||
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(),
|
|
||||||
}
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
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;
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Create an optimizing wrapper for a record source.
|
Ok(Record {
|
||||||
pub fn wrap_source<S: RecordSource + 'static>(
|
text: optimized.compressed,
|
||||||
source: S,
|
..record
|
||||||
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)]
|
#[cfg(test)]
|
||||||
@@ -189,69 +141,69 @@ mod tests {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[test]
|
||||||
async fn test_optimizer_wrap_preserves_structure() {
|
fn test_optimize_record_preserves_structure() {
|
||||||
let records = vec![make_test_record("ERROR: simple")];
|
let record = make_test_record("ERROR: simple");
|
||||||
let source = mem_chunk::VecSource(records);
|
|
||||||
let optimizer = ContextOptimizer::new().unwrap();
|
let optimizer = ContextOptimizer::new().unwrap();
|
||||||
|
let metrics = Arc::new(Mutex::new(OptimizationMetrics::default()));
|
||||||
|
|
||||||
let mut stream = OptimizerSink::wrap_source(source, optimizer);
|
let result = optimize_record_with_metrics(record.clone(), &optimizer, &metrics);
|
||||||
let result = stream.next().await;
|
assert!(result.is_ok());
|
||||||
|
|
||||||
assert!(result.is_some());
|
let optimized = result.unwrap();
|
||||||
let record = result.unwrap().unwrap();
|
assert_eq!(optimized.role, Role::User);
|
||||||
assert_eq!(record.role, Role::User);
|
assert_eq!(optimized.provenance.source_id, "test");
|
||||||
assert_eq!(record.provenance.source_id, "test");
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[test]
|
||||||
async fn test_optimizer_wrap_processes() {
|
fn test_optimize_record_tracks_bytes() {
|
||||||
let records = vec![make_test_record("ERROR: failed\nINFO: debug")];
|
let record = make_test_record("ERROR: failed\nINFO: debug\nERROR: permission");
|
||||||
let source = mem_chunk::VecSource(records);
|
|
||||||
let optimizer = ContextOptimizer::new().unwrap();
|
let optimizer = ContextOptimizer::new().unwrap();
|
||||||
|
let metrics = Arc::new(Mutex::new(OptimizationMetrics::default()));
|
||||||
|
|
||||||
let mut stream = OptimizerSink::wrap_source(source, optimizer);
|
let _ = optimize_record_with_metrics(record, &optimizer, &metrics);
|
||||||
let result = stream.next().await;
|
|
||||||
|
|
||||||
assert!(result.is_some());
|
let m = metrics.lock().unwrap();
|
||||||
let record = result.unwrap().unwrap();
|
assert_eq!(m.total_records, 1);
|
||||||
assert!(!record.text.is_empty());
|
assert!(m.input_bytes_total > 0);
|
||||||
|
assert!(m.output_bytes_total > 0);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[test]
|
||||||
async fn test_optimizer_wrap_multiple_records() {
|
fn test_optimize_record_disabled() {
|
||||||
let records = vec![
|
std::env::set_var("MEM_CONTEXT_OPTIMIZER", "off");
|
||||||
make_test_record("ERROR: first"),
|
let record = make_test_record("content");
|
||||||
make_test_record("ERROR: second"),
|
let optimizer = ContextOptimizer::from_env().unwrap();
|
||||||
make_test_record("INFO: info"),
|
let metrics = Arc::new(Mutex::new(OptimizationMetrics::default()));
|
||||||
];
|
|
||||||
|
|
||||||
let source = mem_chunk::VecSource(records);
|
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();
|
let optimizer = ContextOptimizer::new().unwrap();
|
||||||
|
|
||||||
let mut stream = OptimizerSink::wrap_source(source, optimizer);
|
for i in 0..3 {
|
||||||
let mut count = 0;
|
let record = make_test_record(&format!("ERROR: {}", i));
|
||||||
|
let _ = optimize_record_with_metrics(record, &optimizer, &metrics);
|
||||||
while let Some(result) = stream.next().await {
|
|
||||||
assert!(result.is_ok());
|
|
||||||
count += 1;
|
|
||||||
}
|
}
|
||||||
|
|
||||||
assert_eq!(count, 3);
|
let m = metrics.lock().unwrap();
|
||||||
}
|
assert_eq!(m.total_records, 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