rskit_agent/types/
strategy.rs

1use rskit_errors::AppError;
2use rskit_llm::types::Message;
3
4/// Strategy for compacting the message history when context is too large.
5pub trait ContextStrategy: Send + Sync {
6    /// Compact the message list toward `max_tokens`.
7    ///
8    /// Strategies decide how aggressively to reduce history. Implementations
9    /// that cannot satisfy their own policy should return an error.
10    fn compact(&self, messages: Vec<Message>, max_tokens: usize) -> Result<Vec<Message>, AppError>;
11}
12
13/// Fail immediately if context exceeds the limit (no compaction).
14pub struct FailStrategy;
15
16impl ContextStrategy for FailStrategy {
17    fn compact(
18        &self,
19        _messages: Vec<Message>,
20        _max_tokens: usize,
21    ) -> Result<Vec<Message>, AppError> {
22        Err(AppError::new(
23            rskit_errors::ErrorCode::InvalidInput,
24            "context exceeds token limit and no compaction strategy is configured",
25        ))
26    }
27}
28
29/// Keep only the system prompt (first message) and the last `keep_last` messages.
30///
31/// This is a message-count strategy; `max_tokens` is accepted through the
32/// common strategy trait but is not used for token-aware trimming.
33pub struct TruncateStrategy {
34    /// Number of recent messages to preserve.
35    pub keep_last: usize,
36}
37
38impl ContextStrategy for TruncateStrategy {
39    fn compact(
40        &self,
41        messages: Vec<Message>,
42        _max_tokens: usize,
43    ) -> Result<Vec<Message>, AppError> {
44        let has_system = matches!(messages.first(), Some(Message::System(_)));
45        let max_messages = self.keep_last + usize::from(has_system);
46        if messages.len() <= max_messages {
47            return Ok(messages);
48        }
49
50        let mut result = Vec::with_capacity(self.keep_last + 1);
51        if has_system {
52            result.push(messages[0].clone());
53        }
54
55        let start = messages.len().saturating_sub(self.keep_last);
56        for msg in &messages[start..] {
57            if matches!(msg, Message::System(_)) && !result.is_empty() {
58                continue;
59            }
60            result.push(msg.clone());
61        }
62        Ok(result)
63    }
64}
65
66#[cfg(test)]
67mod tests {
68    use super::*;
69    use rskit_llm::types;
70
71    #[test]
72    fn test_fail_strategy() {
73        let strategy = FailStrategy;
74        let msgs = vec![types::user("hello")];
75        let result = strategy.compact(msgs, 100);
76        assert!(result.is_err());
77    }
78
79    #[test]
80    fn test_truncate_strategy_no_truncation_needed() {
81        let strategy = TruncateStrategy { keep_last: 10 };
82        let msgs = vec![types::user("a"), types::assistant("b")];
83        let result = strategy.compact(msgs, 100).unwrap();
84        assert_eq!(result.len(), 2);
85    }
86
87    #[test]
88    fn test_truncate_strategy_keeps_system_and_last() {
89        let strategy = TruncateStrategy { keep_last: 2 };
90        let msgs = vec![
91            types::system("sys"),
92            types::user("a"),
93            types::assistant("b"),
94            types::user("c"),
95            types::assistant("d"),
96        ];
97        let result = strategy.compact(msgs, 100).unwrap();
98        assert_eq!(result.len(), 3);
99        assert!(matches!(&result[0], Message::System(_)));
100    }
101
102    #[test]
103    fn test_truncate_strategy_no_system() {
104        let strategy = TruncateStrategy { keep_last: 2 };
105        let msgs = vec![
106            types::user("a"),
107            types::assistant("b"),
108            types::user("c"),
109            types::assistant("d"),
110        ];
111        let result = strategy.compact(msgs, 100).unwrap();
112        assert_eq!(result.len(), 2);
113    }
114
115    #[test]
116    fn test_truncate_strategy_no_system_drops_keep_last_plus_one() {
117        let strategy = TruncateStrategy { keep_last: 2 };
118        let msgs = vec![types::user("a"), types::assistant("b"), types::user("c")];
119
120        let result = strategy.compact(msgs, 100).unwrap();
121
122        assert_eq!(result.len(), 2);
123        assert!(matches!(&result[0], Message::Assistant(_)));
124        assert!(matches!(&result[1], Message::User(_)));
125    }
126}
rskit_agent/types/strategy.rs

rskit_agent/types/
strategy.rs
