Skip to main content

ai_agents_runtime/spec/
memory.rs

1//! Memory configuration types
2
3use 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    /// Cross-session actor memory configuration.
33    #[serde(default, skip_serializing_if = "Option::is_none")]
34    pub actor_memory: Option<ActorMemoryConfig>,
35
36    /// Key facts extraction configuration.
37    #[serde(default, skip_serializing_if = "Option::is_none")]
38    pub facts: Option<FactsConfig>,
39
40    /// Session-level metadata defaults.
41    #[serde(default, skip_serializing_if = "Option::is_none")]
42    pub session: Option<SessionConfig>,
43
44    /// Actor-scoped relationship memory configuration.
45    #[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    /// Check if actor memory is enabled.
81    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    /// Check if facts extraction is enabled.
89    pub fn has_facts(&self) -> bool {
90        self.facts.as_ref().map(|f| f.enabled).unwrap_or(false)
91    }
92
93    /// Check if relationship memory is enabled.
94    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}