ai_agents_runtime/spec/
memory.rs

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

ai_agents_runtime/spec/
memory.rs
