use crate::compression::{CompressionInput, CompressionOutput, ContextCompressor};
use echo_core::error::Result;
use futures::future::BoxFuture;
pub struct SlidingWindowCompressor {
window_size: usize,
}
impl SlidingWindowCompressor {
pub fn new(window_size: usize) -> Self {
Self { window_size }
}
}
impl ContextCompressor for SlidingWindowCompressor {
fn compress(&self, input: CompressionInput) -> BoxFuture<'_, Result<CompressionOutput>> {
Box::pin(async move {
let (system_msgs, conv_msgs): (Vec<_>, Vec<_>) =
input.messages.into_iter().partition(|m| m.role == "system");
if conv_msgs.len() <= self.window_size {
let mut messages = system_msgs;
messages.extend(conv_msgs);
return Ok(CompressionOutput {
messages,
evicted: vec![],
});
}
let split_at = conv_msgs.len() - self.window_size;
let evicted = conv_msgs[..split_at].to_vec();
let kept = conv_msgs[split_at..].to_vec();
let mut messages = system_msgs;
messages.extend(kept);
Ok(CompressionOutput { messages, evicted })
})
}
}