ai_agents_runtime/spec/
memory.rs1use serde::{Deserialize, Serialize};
4
5use ai_agents_facts::{ActorMemoryConfig, FactsConfig, SessionConfig};
6use ai_agents_memory::{CompactingMemoryConfig, MemoryTokenBudget};
7use ai_agents_relationships::RelationshipConfig;
8
9#[derive(Debug, Clone, Serialize, Deserialize)]
10pub struct MemoryConfig {
11 #[serde(rename = "type", default = "default_memory_type")]
12 pub memory_type: String,
13
14 #[serde(default = "default_max_messages")]
15 pub max_messages: usize,
16
17 #[serde(default)]
18 pub max_recent_messages: Option<usize>,
19
20 #[serde(default)]
21 pub compress_threshold: Option<usize>,
22
23 #[serde(default)]
24 pub summarize_batch_size: Option<usize>,
25
26 #[serde(default)]
27 pub token_budget: Option<MemoryTokenBudget>,
28
29 #[serde(default)]
30 pub summarizer_llm: Option<String>,
31
32 #[serde(default, skip_serializing_if = "Option::is_none")]
34 pub actor_memory: Option<ActorMemoryConfig>,
35
36 #[serde(default, skip_serializing_if = "Option::is_none")]
38 pub facts: Option<FactsConfig>,
39
40 #[serde(default, skip_serializing_if = "Option::is_none")]
42 pub session: Option<SessionConfig>,
43
44 #[serde(default, skip_serializing_if = "Option::is_none")]
46 pub relationships: Option<RelationshipConfig>,
47}
48
49fn default_memory_type() -> String {
50 "in-memory".to_string()
51}
52
53fn default_max_messages() -> usize {
54 100
55}
56
57impl Default for MemoryConfig {
58 fn default() -> Self {
59 Self {
60 memory_type: default_memory_type(),
61 max_messages: default_max_messages(),
62 max_recent_messages: None,
63 compress_threshold: None,
64 summarize_batch_size: None,
65 token_budget: None,
66 summarizer_llm: None,
67 actor_memory: None,
68 facts: None,
69 session: None,
70 relationships: None,
71 }
72 }
73}
74
75impl MemoryConfig {
76 pub fn is_compacting(&self) -> bool {
77 self.memory_type == "compacting"
78 }
79
80 pub fn has_actor_memory(&self) -> bool {
82 self.actor_memory
83 .as_ref()
84 .map(|am| am.enabled)
85 .unwrap_or(false)
86 }
87
88 pub fn has_facts(&self) -> bool {
90 self.facts.as_ref().map(|f| f.enabled).unwrap_or(false)
91 }
92
93 pub fn has_relationships(&self) -> bool {
95 self.relationships
96 .as_ref()
97 .map(|r| r.enabled)
98 .unwrap_or(false)
99 }
100
101 pub fn to_compacting_config(&self) -> CompactingMemoryConfig {
102 CompactingMemoryConfig {
103 max_recent_messages: self.max_recent_messages.unwrap_or(50),
104 compress_threshold: self.compress_threshold.unwrap_or(30),
105 summarize_batch_size: self.summarize_batch_size.unwrap_or(10),
106 max_summary_length: 2000,
107 }
108 }
109}
110
111#[cfg(test)]
112mod tests {
113 use super::*;
114
115 #[test]
116 fn test_memory_config_default() {
117 let config = MemoryConfig::default();
118 assert_eq!(config.memory_type, "in-memory");
119 assert_eq!(config.max_messages, 100);
120 assert!(!config.is_compacting());
121 }
122
123 #[test]
124 fn test_memory_config_deserialize() {
125 let yaml = r#"
126type: in-memory
127max_messages: 50
128"#;
129 let config: MemoryConfig = serde_yaml::from_str(yaml).unwrap();
130 assert_eq!(config.memory_type, "in-memory");
131 assert_eq!(config.max_messages, 50);
132 }
133
134 #[test]
135 fn test_memory_config_with_defaults() {
136 let yaml = r#"
137type: sqlite
138"#;
139 let config: MemoryConfig = serde_yaml::from_str(yaml).unwrap();
140 assert_eq!(config.memory_type, "sqlite");
141 assert_eq!(config.max_messages, 100);
142 }
143
144 #[test]
145 fn test_memory_config_with_actor_memory() {
146 let yaml = r#"
147type: compacting
148max_messages: 100
149actor_memory:
150 enabled: true
151 identification:
152 method: from_context
153 context_path: user.id
154 injection:
155 mode: all
156 max_tokens: 800
157 privacy:
158 retention_days: 365
159 allow_deletion: true
160facts:
161 enabled: true
162 extractor_llm: router
163 auto_extract: true
164 categories:
165 - user_preference
166 - user_context
167 max_facts: 30
168session:
169 tags: [support]
170 ttl_seconds: 86400
171relationships:
172 enabled: true
173 dimensions:
174 - trust
175 - sentiment
176 auto_update:
177 enabled: true
178 llm: router
179"#;
180 let config: MemoryConfig = serde_yaml::from_str(yaml).unwrap();
181 assert!(config.has_actor_memory());
182 assert!(config.has_facts());
183 assert!(config.has_relationships());
184 let am = config.actor_memory.unwrap();
185 assert!(am.enabled);
186 assert_eq!(
187 am.identification.method,
188 ai_agents_facts::IdentificationMethod::FromContext
189 );
190 assert_eq!(am.identification.context_path.as_deref(), Some("user.id"));
191 let facts = config.facts.unwrap();
192 assert!(facts.enabled);
193 assert_eq!(facts.extractor_llm.as_deref(), Some("router"));
194 assert_eq!(facts.max_facts, 30);
195 let session = config.session.unwrap();
196 assert_eq!(session.tags, vec!["support"]);
197 assert_eq!(session.ttl_seconds, Some(86400));
198 }
199
200 #[test]
201 fn test_compacting_memory_config() {
202 let yaml = r#"
203type: compacting
204max_messages: 100
205max_recent_messages: 20
206compress_threshold: 30
207summarize_batch_size: 10
208summarizer_llm: router
209"#;
210 let config: MemoryConfig = serde_yaml::from_str(yaml).unwrap();
211 assert!(config.is_compacting());
212 assert_eq!(config.max_recent_messages, Some(20));
213 assert_eq!(config.compress_threshold, Some(30));
214 assert_eq!(config.summarize_batch_size, Some(10));
215 assert_eq!(config.summarizer_llm, Some("router".to_string()));
216
217 let compacting_config = config.to_compacting_config();
218 assert_eq!(compacting_config.max_recent_messages, 20);
219 assert_eq!(compacting_config.compress_threshold, 30);
220 }
221
222 #[test]
223 fn test_memory_config_with_token_budget() {
224 let yaml = r#"
225type: compacting
226max_messages: 100
227token_budget:
228 total: 8192
229 allocation:
230 summary: 2048
231 recent_messages: 4096
232 facts: 1024
233 relationships: 512
234 overflow_strategy: summarize_more
235 warn_at_percent: 75
236"#;
237 let config: MemoryConfig = serde_yaml::from_str(yaml).unwrap();
238 assert!(config.token_budget.is_some());
239 let budget = config.token_budget.unwrap();
240 assert_eq!(budget.total, 8192);
241 assert_eq!(budget.allocation.summary, 2048);
242 assert_eq!(budget.allocation.relationships, 512);
243 assert_eq!(budget.warn_at_percent, 75);
244 }
245
246 #[test]
247 fn test_to_compacting_config_defaults() {
248 let config = MemoryConfig {
249 memory_type: "compacting".to_string(),
250 ..Default::default()
251 };
252 let compacting = config.to_compacting_config();
253 assert_eq!(compacting.max_recent_messages, 50);
254 assert_eq!(compacting.compress_threshold, 30);
255 assert_eq!(compacting.summarize_batch_size, 10);
256 }
257}