klieo_runlog/

compaction.rs

//! `CompactionPolicy` trait + deterministic and (optional) LLM-summarise impls.
//!
//! Roadmap decision-point #5: deterministic drop-by-rule is the default.
//! LLM-summarise is opt-in behind the `compaction-llm` feature.

use crate::error::RunLogError;
use crate::types::{RunLog, Step, StepKind};
use async_trait::async_trait;
use std::time::Duration;

/// A pluggable run-log compaction strategy.
#[async_trait]
pub trait CompactionPolicy: Send + Sync {
    /// Mutate `run_log` in place to reduce its on-disk footprint.
    async fn compact(&self, run_log: &mut RunLog) -> Result<(), RunLogError>;
}

/// Deterministic drop-by-rule compaction. Default policy.
///
/// - Keeps the most recent `max_steps` verbatim.
/// - Drops any step whose `kind` matches `drop_kinds` first.
/// - When still over `max_steps`, the older overflow steps are replaced by a
///   single summary step (kind = `ToolCall`, name = `"compacted"`,
///   input = `{ "dropped_count": N }`).
/// - Re-indexes `Step::idx` after compaction.
pub struct DeterministicCompaction {
    /// Maximum number of `Step`s kept verbatim. Older overflow becomes a single summary step.
    pub max_steps: u32,
    /// Step kinds dropped unconditionally before the `max_steps` truncation runs.
    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> {
        // 1. Drop kinds unconditionally.
        if !self.drop_kinds.is_empty() {
            run_log.steps.retain(|s| !self.drop_kinds.contains(&s.kind));
        }

        // 2. Truncate older overflow into a summary step.
        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 now holds the older `dropped` steps; replace with summary.
            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);
        }

        // 3. Re-index.
        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;

    /// LLM-summarise compaction. Available behind the `compaction-llm` feature.
    ///
    /// When `run_log.steps.len() > max_steps`, the older overflow is summarised
    /// by calling `llm.complete(...)` with a single `User` message that asks for
    /// a one-sentence summary. The summary replaces the older steps as a single
    /// `ToolCall` step (name `"compacted"`, output = the summary text).
    pub struct LlmCompaction {
        /// LLM client used to produce the summary string.
        pub llm: Arc<dyn LlmClient>,
        /// Maximum number of `Step`s kept verbatim.
        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;
            // Build a compact prompt of older steps.
            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}"
            );
            // klieo-core's LlmClient::complete takes a ChatRequest and returns
            // a ChatResponse — wrap the string prompt as a single User message
            // and pull `message.content` out of the response.
            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;