use serde::{Deserialize, Serialize};
pub mod defaults {
pub const RETENTION_WINDOW: usize = 10;
pub const EVICTION_WINDOW: f64 = 0.6;
pub const TOKEN_THRESHOLD: usize = 80_000;
pub const TURN_THRESHOLD: usize = 20;
pub const MESSAGE_THRESHOLD: usize = 50;
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CompactThresholds {
pub token_threshold: Option<usize>,
pub turn_threshold: Option<usize>,
pub message_threshold: Option<usize>,
pub on_turn_end: Option<bool>,
}
impl Default for CompactThresholds {
fn default() -> Self {
Self {
token_threshold: Some(defaults::TOKEN_THRESHOLD),
turn_threshold: Some(defaults::TURN_THRESHOLD),
message_threshold: Some(defaults::MESSAGE_THRESHOLD),
on_turn_end: None,
}
}
}
impl CompactThresholds {
pub fn aggressive() -> Self {
Self {
token_threshold: Some(40_000),
turn_threshold: Some(10),
message_threshold: Some(25),
on_turn_end: Some(true),
}
}
pub fn relaxed() -> Self {
Self {
token_threshold: Some(150_000),
turn_threshold: Some(50),
message_threshold: Some(100),
on_turn_end: None,
}
}
pub fn disabled() -> Self {
Self {
token_threshold: None,
turn_threshold: None,
message_threshold: None,
on_turn_end: None,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CompactConfig {
pub retention_window: usize,
pub eviction_window: f64,
pub thresholds: CompactThresholds,
}
impl Default for CompactConfig {
fn default() -> Self {
Self {
retention_window: defaults::RETENTION_WINDOW,
eviction_window: defaults::EVICTION_WINDOW,
thresholds: CompactThresholds::default(),
}
}
}
impl CompactConfig {
pub fn with_retention(retention: usize) -> Self {
Self {
retention_window: retention,
..Default::default()
}
}
pub fn with_thresholds(thresholds: CompactThresholds) -> Self {
Self {
thresholds,
..Default::default()
}
}
pub fn should_compact(
&self,
token_count: usize,
turn_count: usize,
message_count: usize,
last_is_user: bool,
) -> bool {
if let Some(threshold) = self.thresholds.token_threshold
&& token_count >= threshold
{
return true;
}
if let Some(threshold) = self.thresholds.turn_threshold
&& turn_count >= threshold
{
return true;
}
if let Some(threshold) = self.thresholds.message_threshold
&& message_count >= threshold
{
return true;
}
if let Some(true) = self.thresholds.on_turn_end
&& last_is_user
{
let near_token = self
.thresholds
.token_threshold
.map(|t| token_count >= t / 2)
.unwrap_or(false);
let near_turn = self
.thresholds
.turn_threshold
.map(|t| turn_count >= t / 2)
.unwrap_or(false);
if near_token || near_turn {
return true;
}
}
false
}
pub fn compaction_reason(
&self,
token_count: usize,
turn_count: usize,
message_count: usize,
) -> Option<String> {
if let Some(threshold) = self.thresholds.token_threshold
&& token_count >= threshold
{
return Some(format!(
"token count ({}) >= threshold ({})",
token_count, threshold
));
}
if let Some(threshold) = self.thresholds.turn_threshold
&& turn_count >= threshold
{
return Some(format!(
"turn count ({}) >= threshold ({})",
turn_count, threshold
));
}
if let Some(threshold) = self.thresholds.message_threshold
&& message_count >= threshold
{
return Some(format!(
"message count ({}) >= threshold ({})",
message_count, threshold
));
}
None
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_default_config() {
let config = CompactConfig::default();
assert_eq!(config.retention_window, defaults::RETENTION_WINDOW);
assert!((config.eviction_window - defaults::EVICTION_WINDOW).abs() < f64::EPSILON);
}
#[test]
fn test_should_compact_tokens() {
let config = CompactConfig::default();
assert!(!config.should_compact(50_000, 5, 10, false));
assert!(config.should_compact(100_000, 5, 10, false));
}
#[test]
fn test_should_compact_turns() {
let config = CompactConfig::default();
assert!(!config.should_compact(10_000, 10, 20, false));
assert!(config.should_compact(10_000, 25, 50, false));
}
#[test]
fn test_should_compact_messages() {
let config = CompactConfig::default();
assert!(!config.should_compact(10_000, 10, 30, false));
assert!(config.should_compact(10_000, 10, 60, false));
}
#[test]
fn test_aggressive_thresholds() {
let thresholds = CompactThresholds::aggressive();
assert_eq!(thresholds.token_threshold, Some(40_000));
assert_eq!(thresholds.turn_threshold, Some(10));
}
#[test]
fn test_disabled_thresholds() {
let config = CompactConfig::with_thresholds(CompactThresholds::disabled());
assert!(!config.should_compact(1_000_000, 1000, 10000, true));
}
}