1use chrono::{DateTime, Utc};
8use enact_core::kernel::StepId;
9use serde::{Deserialize, Serialize};
10
11#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
15#[serde(rename_all = "snake_case")]
16pub enum ContextSegmentType {
17 System,
19 History,
21 WorkingMemory,
23 ToolResults,
25 RagContext,
27 UserInput,
29 AgentScratchpad,
31 ChildSummary,
33 Guidance,
35}
36
37impl ContextSegmentType {
38 pub fn is_compressible(&self) -> bool {
40 match self {
41 Self::System => false,
42 Self::History => true,
43 Self::WorkingMemory => true,
44 Self::ToolResults => true,
45 Self::RagContext => true,
46 Self::UserInput => false,
47 Self::AgentScratchpad => true,
48 Self::ChildSummary => false,
49 Self::Guidance => false,
50 }
51 }
52
53 pub fn default_priority(&self) -> ContextPriority {
55 match self {
56 Self::System => ContextPriority::Critical,
57 Self::UserInput => ContextPriority::Critical,
58 Self::Guidance => ContextPriority::High,
59 Self::ChildSummary => ContextPriority::High,
60 Self::History => ContextPriority::Medium,
61 Self::WorkingMemory => ContextPriority::Medium,
62 Self::ToolResults => ContextPriority::Medium,
63 Self::RagContext => ContextPriority::Low,
64 Self::AgentScratchpad => ContextPriority::Low,
65 }
66 }
67}
68
69#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
74#[serde(rename_all = "snake_case")]
75pub enum ContextPriority {
76 Low = 0,
78 Medium = 1,
80 High = 2,
82 Critical = 3,
84}
85
86#[derive(Debug, Clone, Serialize, Deserialize)]
90#[serde(rename_all = "camelCase")]
91pub struct ContextSegment {
92 pub id: String,
94
95 #[serde(rename = "type")]
97 pub segment_type: ContextSegmentType,
98
99 pub content: String,
101
102 pub token_count: usize,
104
105 pub priority: ContextPriority,
107
108 pub compressible: bool,
110
111 #[serde(skip_serializing_if = "Option::is_none")]
113 pub source_step_id: Option<StepId>,
114
115 pub added_at: DateTime<Utc>,
117
118 pub sequence: u64,
120}
121
122impl ContextSegment {
123 pub fn new(
125 segment_type: ContextSegmentType,
126 content: String,
127 token_count: usize,
128 sequence: u64,
129 ) -> Self {
130 Self {
131 id: format!("seg_{}", uuid::Uuid::new_v4()),
132 segment_type,
133 content,
134 token_count,
135 priority: segment_type.default_priority(),
136 compressible: segment_type.is_compressible(),
137 source_step_id: None,
138 added_at: Utc::now(),
139 sequence,
140 }
141 }
142
143 pub fn system(content: impl Into<String>, token_count: usize) -> Self {
145 Self::new(ContextSegmentType::System, content.into(), token_count, 0)
146 }
147
148 pub fn user_input(content: impl Into<String>, token_count: usize, sequence: u64) -> Self {
150 Self::new(
151 ContextSegmentType::UserInput,
152 content.into(),
153 token_count,
154 sequence,
155 )
156 }
157
158 pub fn history(content: impl Into<String>, token_count: usize, sequence: u64) -> Self {
160 Self::new(
161 ContextSegmentType::History,
162 content.into(),
163 token_count,
164 sequence,
165 )
166 }
167
168 pub fn tool_results(
170 content: impl Into<String>,
171 token_count: usize,
172 sequence: u64,
173 step_id: StepId,
174 ) -> Self {
175 let mut segment = Self::new(
176 ContextSegmentType::ToolResults,
177 content.into(),
178 token_count,
179 sequence,
180 );
181 segment.source_step_id = Some(step_id);
182 segment
183 }
184
185 pub fn rag_context(content: impl Into<String>, token_count: usize, sequence: u64) -> Self {
187 Self::new(
188 ContextSegmentType::RagContext,
189 content.into(),
190 token_count,
191 sequence,
192 )
193 }
194
195 pub fn child_summary(
197 content: impl Into<String>,
198 token_count: usize,
199 sequence: u64,
200 step_id: StepId,
201 ) -> Self {
202 let mut segment = Self::new(
203 ContextSegmentType::ChildSummary,
204 content.into(),
205 token_count,
206 sequence,
207 );
208 segment.source_step_id = Some(step_id);
209 segment
210 }
211
212 pub fn guidance(content: impl Into<String>, token_count: usize, sequence: u64) -> Self {
214 Self::new(
215 ContextSegmentType::Guidance,
216 content.into(),
217 token_count,
218 sequence,
219 )
220 }
221
222 pub fn with_priority(mut self, priority: ContextPriority) -> Self {
224 self.priority = priority;
225 self
226 }
227
228 pub fn non_compressible(mut self) -> Self {
230 self.compressible = false;
231 self
232 }
233}
234
235#[cfg(test)]
236mod tests {
237 use super::*;
238
239 #[test]
240 fn test_segment_type_compressibility() {
241 assert!(!ContextSegmentType::System.is_compressible());
242 assert!(ContextSegmentType::History.is_compressible());
243 assert!(!ContextSegmentType::UserInput.is_compressible());
244 assert!(ContextSegmentType::ToolResults.is_compressible());
245 }
246
247 #[test]
248 fn test_segment_type_priority() {
249 assert_eq!(
250 ContextSegmentType::System.default_priority(),
251 ContextPriority::Critical
252 );
253 assert_eq!(
254 ContextSegmentType::History.default_priority(),
255 ContextPriority::Medium
256 );
257 assert_eq!(
258 ContextSegmentType::RagContext.default_priority(),
259 ContextPriority::Low
260 );
261 }
262
263 #[test]
264 fn test_priority_ordering() {
265 assert!(ContextPriority::Critical > ContextPriority::High);
266 assert!(ContextPriority::High > ContextPriority::Medium);
267 assert!(ContextPriority::Medium > ContextPriority::Low);
268 }
269
270 #[test]
271 fn test_create_segment() {
272 let segment = ContextSegment::system("You are helpful", 10);
273 assert_eq!(segment.segment_type, ContextSegmentType::System);
274 assert_eq!(segment.priority, ContextPriority::Critical);
275 assert!(!segment.compressible);
276 }
277}