258 lines
8.8 KiB
Rust
258 lines
8.8 KiB
Rust
|
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());
|
||
|
|
}
|
||
|
|
}
|