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: 80_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| {
53                s.chars()
54                    .map(|c| if c.is_ascii() { 1 } else { 4 })
55                    .sum::<usize>()
56            })
57            .sum::<usize>()
58            / 4
59    }
60}
61
62impl Summarizer for TokenBasedSummarizer {
63    fn should_summarize(&self, history: &[Message]) -> bool {
64        let est = Self::estimate_tokens(history);
65        est >= self.threshold
66    }
67
68    fn summarize(&self, history: &mut Vec<Message>) -> Option<Message> {
69        if history.len() <= self.retain_last {
70            return None;
71        }
72
73        // Determine slice to summarize (old messages)
74        let split = history.len().saturating_sub(self.retain_last);
75        if split == 0 {
76            return None;
77        }
78
79        // Collect older messages for summarization
80        let older: Vec<String> = history.drain(0..split).filter_map(|m| m.content).collect();
81
82        if older.is_empty() {
83            // nothing meaningful to summarize; nothing to do
84            return None;
85        }
86
87        // Simple concatenation with newlines; truncate if too long.
88        let joined = older.join("\n");
89        let summary_text = if joined.len() > self.max_summary_chars {
90            let mut s = joined;
91            s.truncate(self.max_summary_chars);
92            format!("Short summary of earlier conversation:\n{}\n(Truncated)", s)
93        } else {
94            format!("Short summary of earlier conversation:\n{}", joined)
95        };
96
97        // Create a system message as the summary and insert at the front.
98        let mut summary_msg = Message::new(Role::System, summary_text.as_str());
99        // Optionally tag name to indicate it's an auto-summary
100        summary_msg.name = Some("[auto-summary]".to_string());
101
102        history.insert(0, summary_msg.clone());
103        Some(summary_msg)
104    }
105}
106
107#[cfg(test)]
108mod tests {
109    use super::*;
110    use crate::raw::request::message::{Message, Role};
111
112    #[test]
113    fn summarizer_should_trigger_and_replace() {
114        let mut hist = vec![
115            Message::new(Role::User, "User message 1"),
116            Message::new(Role::Assistant, "Assistant reply 1"),
117            Message::new(Role::User, "User message 2"),
118            Message::new(Role::Assistant, "Assistant reply 2"),
119        ];
120
121        // Use a very low threshold so summarization triggers.
122        let summ = TokenBasedSummarizer {
123            threshold: 0, // everything triggers
124            retain_last: 1,
125            max_summary_chars: 100,
126        };
127
128        assert!(summ.should_summarize(&hist));
129        let maybe_summary = summ.summarize(&mut hist);
130        assert!(maybe_summary.is_some());
131        // After summarization, history length should be <= retain_last + 1 (summary)
132        assert!(hist.len() <= (1 + 1));
133        // First message must be a system message (the summary)
134        assert!(matches!(hist[0].role, Role::System));
135    }
136}
ds_api/conversation/summarizer.rs

ds_api/conversation/
summarizer.rs
