use crate::types::Message;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
pub const SESSION_COMPACTION_CADENCE_KEY: &str = "session_compaction_cadence";
#[derive(Debug, Clone, Serialize, Deserialize, Default, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub struct SessionCompactionCadence {
pub session_boundary_index: u64,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub last_compaction_boundary_index: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub last_compaction_attempt_boundary_index: Option<u64>,
}
#[derive(Debug, Clone)]
pub struct CompactionContext {
pub last_input_tokens: u64,
pub message_count: usize,
pub estimated_history_tokens: u64,
pub last_compaction_boundary_index: Option<u64>,
pub session_boundary_index: u64,
}
#[derive(Debug, Clone)]
pub struct CompactionResult {
pub messages: Vec<Message>,
pub summary: CompactionSummary,
pub retained: Vec<CompactionRetained>,
pub discarded: Vec<CompactionDiscard>,
}
#[derive(Debug, Clone, PartialEq)]
pub struct CompactionSummary {
pub rebuilt_offset: u64,
pub message: Message,
}
impl CompactionSummary {
pub fn new(rebuilt_offset: u64, message: Message) -> Self {
Self {
rebuilt_offset,
message,
}
}
}
pub const COMPACTION_SUMMARY_PREFIX: &str = "\
[Context compacted] A previous context produced the following summary of work so far. \
The current tool and session state is preserved. Use this summary to continue without \
duplicating work:\n\n";
#[derive(Debug, Clone, PartialEq)]
pub struct CompactionDiscard {
pub source_offset: u64,
pub message: Message,
}
impl CompactionDiscard {
pub fn new(source_offset: u64, message: Message) -> Self {
Self {
source_offset,
message,
}
}
}
#[derive(Debug, Clone, PartialEq)]
pub struct CompactionRetained {
pub source_offset: u64,
pub rebuilt_offset: u64,
pub message: Message,
}
impl CompactionRetained {
pub fn new(source_offset: u64, rebuilt_offset: u64, message: Message) -> Self {
Self {
source_offset,
rebuilt_offset,
message,
}
}
}
#[derive(Debug, Clone)]
pub struct CompactionConfig {
pub auto_compact_threshold: u64,
pub recent_turn_budget: usize,
pub max_summary_tokens: u32,
pub min_turns_between_compactions: u32,
}
impl Default for CompactionConfig {
fn default() -> Self {
Self {
auto_compact_threshold: 100_000,
recent_turn_budget: 4,
max_summary_tokens: 4096,
min_turns_between_compactions: 3,
}
}
}
pub trait Compactor: Send + Sync {
fn should_compact(&self, ctx: &CompactionContext) -> bool;
fn compaction_prompt(&self) -> &str;
fn max_summary_tokens(&self) -> u32;
fn prepare_for_summarization(&self, messages: &[Message]) -> Vec<Message> {
messages.to_vec()
}
fn rebuild_history(&self, messages: &[Message], summary: &str) -> CompactionResult;
}
#[derive(Debug, Clone, Copy)]
pub struct CompactionWindow<'a> {
pub messages: &'a [Message],
pub last_input_tokens: u64,
pub session_boundary_index: u64,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CuratedCompactionSummary(String);
impl CuratedCompactionSummary {
pub fn new(text: impl Into<String>) -> Result<Self, CompactionCuratorError> {
let text = text.into();
if text.trim().is_empty() {
return Err(CompactionCuratorError::EmptySummary);
}
Ok(Self(text))
}
pub fn as_str(&self) -> &str {
&self.0
}
pub fn into_string(self) -> String {
self.0
}
}
#[derive(Debug, thiserror::Error)]
pub enum CompactionCuratorError {
#[error("curator produced an empty compaction summary")]
EmptySummary,
#[error("curator failed to produce a compaction summary: {0}")]
Failed(String),
}
#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
pub trait CompactionCurator: Send + Sync {
async fn curate_summary(
&self,
window: CompactionWindow<'_>,
) -> Result<CuratedCompactionSummary, CompactionCuratorError>;
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
use super::*;
#[test]
fn curated_compaction_summary_rejects_empty_text() {
assert!(matches!(
CuratedCompactionSummary::new(""),
Err(CompactionCuratorError::EmptySummary)
));
assert!(matches!(
CuratedCompactionSummary::new(" \n\t"),
Err(CompactionCuratorError::EmptySummary)
));
}
#[test]
fn curated_compaction_summary_round_trips_text() {
let summary = CuratedCompactionSummary::new("curated summary").unwrap();
assert_eq!(summary.as_str(), "curated summary");
assert_eq!(summary.into_string(), "curated summary");
}
}