use crate::error::RunLogError;
use crate::types::{RunLog, Step, StepKind};
use async_trait::async_trait;
use std::time::Duration;
#[async_trait]
pub trait CompactionPolicy: Send + Sync {
async fn compact(&self, run_log: &mut RunLog) -> Result<(), RunLogError>;
}
pub struct DeterministicCompaction {
pub max_steps: u32,
pub drop_kinds: Vec<StepKind>,
}
impl Default for DeterministicCompaction {
fn default() -> Self {
Self {
max_steps: 100,
drop_kinds: vec![],
}
}
}
#[async_trait]
impl CompactionPolicy for DeterministicCompaction {
async fn compact(&self, run_log: &mut RunLog) -> Result<(), RunLogError> {
if !self.drop_kinds.is_empty() {
run_log.steps.retain(|s| !self.drop_kinds.contains(&s.kind));
}
let max = self.max_steps as usize;
if run_log.steps.len() > max {
let total = run_log.steps.len();
let dropped = total - max;
let kept_tail: Vec<Step> = run_log.steps.split_off(dropped);
run_log.steps = vec![Step {
idx: 0,
kind: StepKind::ToolCall,
name: Some("compacted".into()),
prompt_tokens: None,
completion_tokens: None,
cost_usd: None,
input: serde_json::json!({ "dropped_count": dropped }),
output: serde_json::Value::Null,
error: None,
latency: Duration::ZERO,
span_id: None,
}];
run_log.steps.extend(kept_tail);
}
for (i, s) in run_log.steps.iter_mut().enumerate() {
s.idx = i as u32;
}
Ok(())
}
}
#[cfg(feature = "compaction-llm")]
mod llm {
use super::{CompactionPolicy, RunLogError, Step, StepKind};
use crate::types::RunLog;
use async_trait::async_trait;
use klieo_core::llm::{ChatRequest, LlmClient, Message, Role};
use std::sync::Arc;
use std::time::Duration;
pub struct LlmCompaction {
pub llm: Arc<dyn LlmClient>,
pub max_steps: u32,
}
#[async_trait]
impl CompactionPolicy for LlmCompaction {
async fn compact(&self, run_log: &mut RunLog) -> Result<(), RunLogError> {
let max = self.max_steps as usize;
if run_log.steps.len() <= max {
return Ok(());
}
let total = run_log.steps.len();
let dropped = total - max;
let older_json = serde_json::to_string(&run_log.steps[..dropped])
.map_err(|e| RunLogError::Compaction(e.to_string()))?;
let prompt = format!(
"Summarise the following {dropped} agent steps in one sentence:\n{older_json}"
);
let req = ChatRequest::new(vec![Message {
role: Role::User,
content: prompt,
tool_calls: vec![],
tool_call_id: None,
}]);
let resp = self
.llm
.complete(req)
.await
.map_err(|e| RunLogError::Compaction(e.to_string()))?;
let summary = resp.message.content;
let kept_tail: Vec<Step> = run_log.steps.split_off(dropped);
run_log.steps = vec![Step {
idx: 0,
kind: StepKind::ToolCall,
name: Some("compacted".into()),
prompt_tokens: None,
completion_tokens: None,
cost_usd: None,
input: serde_json::json!({ "dropped_count": dropped }),
output: serde_json::Value::String(summary),
error: None,
latency: Duration::ZERO,
span_id: None,
}];
run_log.steps.extend(kept_tail);
for (i, s) in run_log.steps.iter_mut().enumerate() {
s.idx = i as u32;
}
Ok(())
}
}
}
#[cfg(feature = "compaction-llm")]
pub use llm::LlmCompaction;