ds_api/conversation/
summarizer.rs

1//! Summarizer submodule for conversation.
2//!
3//! Contains the `Summarizer` trait and a default `TokenBasedSummarizer`.
4//! This module is intended to be used by the conversation implementation.
5
6use crate::raw::request::message::{Message, Role};
7
8/// Summarizer trait:
9/// - `should_summarize` checks whether current history should be summarized
10/// - `summarize` mutates the history to replace older messages with a single short `system` summary message
11pub trait Summarizer: Send + Sync {
12    /// Return true if history should be summarized now.
13    fn should_summarize(&self, history: &[Message]) -> bool;
14
15    /// Perform summarization by mutating `history`.
16    /// If a summary was created and applied, return `Some(Message)` representing the inserted summary message.
17    /// If nothing was done, return `None`.
18    fn summarize(&self, history: &mut Vec<Message>) -> Option<Message>;
19}
20
21/// Default token-based summarizer:
22/// - Estimates tokens roughly as total characters / 4
23/// - Triggers when estimated tokens exceed `threshold` (default 100_000)
24/// - Keeps `retain_last` most recent messages and summarizes the older ones into a single system message.
25/// - The summary is a concatenation/truncation of older messages, not an LLM semantic summary.
26#[derive(Clone, Debug)]
27pub struct TokenBasedSummarizer {
28    pub threshold: usize,
29    pub retain_last: usize,
30    pub max_summary_chars: usize,
31}
32
33impl Default for TokenBasedSummarizer {
34    fn default() -> Self {
35        Self {
36            threshold: 100_000,
37            retain_last: 10,
38            max_summary_chars: 2_000, // cap the summary length
39        }
40    }
41}
42
43impl TokenBasedSummarizer {
44    /// A simple heuristic to estimate tokens from message history.
45    /// Uses chars / 4 as a rough token estimate.
46    fn estimate_tokens(history: &[Message]) -> usize {
47        // Skip counting system messages (system prompts) when estimating tokens.
48        history
49            .iter()
50            .filter(|m| !matches!(m.role, Role::System))
51            .filter_map(|m| m.content.as_ref())
52            .map(|s| s.len())
53            .sum::<usize>()
54            / 4
55    }
56}
57
58impl Summarizer for TokenBasedSummarizer {
59    fn should_summarize(&self, history: &[Message]) -> bool {
60        let est = Self::estimate_tokens(history);
61        est >= self.threshold
62    }
63
64    fn summarize(&self, history: &mut Vec<Message>) -> Option<Message> {
65        if history.len() <= self.retain_last {
66            return None;
67        }
68
69        // Determine slice to summarize (old messages)
70        let split = history.len().saturating_sub(self.retain_last);
71        if split == 0 {
72            return None;
73        }
74
75        // Collect older messages for summarization
76        let older: Vec<String> = history.drain(0..split).filter_map(|m| m.content).collect();
77
78        if older.is_empty() {
79            // nothing meaningful to summarize; nothing to do
80            return None;
81        }
82
83        // Simple concatenation with newlines; truncate if too long.
84        let joined = older.join("\n");
85        let summary_text = if joined.len() > self.max_summary_chars {
86            let mut s = joined;
87            s.truncate(self.max_summary_chars);
88            format!("Short summary of earlier conversation:\n{}\n(Truncated)", s)
89        } else {
90            format!("Short summary of earlier conversation:\n{}", joined)
91        };
92
93        // Create a system message as the summary and insert at the front.
94        let mut summary_msg = Message::new(Role::System, summary_text.as_str());
95        // Optionally tag name to indicate it's an auto-summary
96        summary_msg.name = Some("[auto-summary]".to_string());
97
98        history.insert(0, summary_msg.clone());
99        Some(summary_msg)
100    }
101}
102
103#[cfg(test)]
104mod tests {
105    use super::*;
106    use crate::raw::request::message::{Message, Role};
107
108    #[test]
109    fn summarizer_should_trigger_and_replace() {
110        let mut hist = vec![
111            Message::new(Role::User, "User message 1"),
112            Message::new(Role::Assistant, "Assistant reply 1"),
113            Message::new(Role::User, "User message 2"),
114            Message::new(Role::Assistant, "Assistant reply 2"),
115        ];
116
117        // Use a very low threshold so summarization triggers.
118        let summ = TokenBasedSummarizer {
119            threshold: 0, // everything triggers
120            retain_last: 1,
121            max_summary_chars: 100,
122        };
123
124        assert!(summ.should_summarize(&hist));
125        let maybe_summary = summ.summarize(&mut hist);
126        assert!(maybe_summary.is_some());
127        // After summarization, history length should be <= retain_last + 1 (summary)
128        assert!(hist.len() <= (1 + 1));
129        // First message must be a system message (the summary)
130        assert!(matches!(hist[0].role, Role::System));
131    }
132}
ds_api/conversation/summarizer.rs

ds_api/conversation/
summarizer.rs
