ai_agents_runtime/spec/
memory.rs

1//! Memory configuration types
2
3use serde::{Deserialize, Serialize};
4use std::collections::HashMap;
5
6use ai_agents_memory::{CompactingMemoryConfig, MemoryTokenBudget};
7
8#[derive(Debug, Clone, Serialize, Deserialize)]
9pub struct MemoryConfig {
10    #[serde(rename = "type", default = "default_memory_type")]
11    pub memory_type: String,
12
13    #[serde(default = "default_max_messages")]
14    pub max_messages: usize,
15
16    #[serde(default)]
17    pub max_recent_messages: Option<usize>,
18
19    #[serde(default)]
20    pub compress_threshold: Option<usize>,
21
22    #[serde(default)]
23    pub summarize_batch_size: Option<usize>,
24
25    #[serde(default)]
26    pub token_budget: Option<MemoryTokenBudget>,
27
28    #[serde(default)]
29    pub summarizer_llm: Option<String>,
30
31    #[serde(flatten)]
32    pub extra: HashMap<String, serde_json::Value>,
33}
34
35fn default_memory_type() -> String {
36    "in-memory".to_string()
37}
38
39fn default_max_messages() -> usize {
40    100
41}
42
43impl Default for MemoryConfig {
44    fn default() -> Self {
45        Self {
46            memory_type: default_memory_type(),
47            max_messages: default_max_messages(),
48            max_recent_messages: None,
49            compress_threshold: None,
50            summarize_batch_size: None,
51            token_budget: None,
52            summarizer_llm: None,
53            extra: HashMap::new(),
54        }
55    }
56}
57
58impl MemoryConfig {
59    pub fn is_compacting(&self) -> bool {
60        self.memory_type == "compacting"
61    }
62
63    pub fn to_compacting_config(&self) -> CompactingMemoryConfig {
64        CompactingMemoryConfig {
65            max_recent_messages: self.max_recent_messages.unwrap_or(50),
66            compress_threshold: self.compress_threshold.unwrap_or(30),
67            summarize_batch_size: self.summarize_batch_size.unwrap_or(10),
68            max_summary_length: 2000,
69        }
70    }
71}
72
73#[cfg(test)]
74mod tests {
75    use super::*;
76
77    #[test]
78    fn test_memory_config_default() {
79        let config = MemoryConfig::default();
80        assert_eq!(config.memory_type, "in-memory");
81        assert_eq!(config.max_messages, 100);
82        assert!(!config.is_compacting());
83    }
84
85    #[test]
86    fn test_memory_config_deserialize() {
87        let yaml = r#"
88type: in-memory
89max_messages: 50
90"#;
91        let config: MemoryConfig = serde_yaml::from_str(yaml).unwrap();
92        assert_eq!(config.memory_type, "in-memory");
93        assert_eq!(config.max_messages, 50);
94    }
95
96    #[test]
97    fn test_memory_config_with_defaults() {
98        let yaml = r#"
99type: sqlite
100"#;
101        let config: MemoryConfig = serde_yaml::from_str(yaml).unwrap();
102        assert_eq!(config.memory_type, "sqlite");
103        assert_eq!(config.max_messages, 100);
104    }
105
106    #[test]
107    fn test_memory_config_extra_fields() {
108        let yaml = r#"
109type: sqlite
110max_messages: 200
111db_path: "/path/to/db.sqlite"
112"#;
113        let config: MemoryConfig = serde_yaml::from_str(yaml).unwrap();
114        assert!(config.extra.contains_key("db_path"));
115    }
116
117    #[test]
118    fn test_compacting_memory_config() {
119        let yaml = r#"
120type: compacting
121max_messages: 100
122max_recent_messages: 20
123compress_threshold: 30
124summarize_batch_size: 10
125summarizer_llm: router
126"#;
127        let config: MemoryConfig = serde_yaml::from_str(yaml).unwrap();
128        assert!(config.is_compacting());
129        assert_eq!(config.max_recent_messages, Some(20));
130        assert_eq!(config.compress_threshold, Some(30));
131        assert_eq!(config.summarize_batch_size, Some(10));
132        assert_eq!(config.summarizer_llm, Some("router".to_string()));
133
134        let compacting_config = config.to_compacting_config();
135        assert_eq!(compacting_config.max_recent_messages, 20);
136        assert_eq!(compacting_config.compress_threshold, 30);
137    }
138
139    #[test]
140    fn test_memory_config_with_token_budget() {
141        let yaml = r#"
142type: compacting
143max_messages: 100
144token_budget:
145  total: 8192
146  allocation:
147    summary: 2048
148    recent_messages: 4096
149    facts: 1024
150  overflow_strategy: summarize_more
151  warn_at_percent: 75
152"#;
153        let config: MemoryConfig = serde_yaml::from_str(yaml).unwrap();
154        assert!(config.token_budget.is_some());
155        let budget = config.token_budget.unwrap();
156        assert_eq!(budget.total, 8192);
157        assert_eq!(budget.allocation.summary, 2048);
158        assert_eq!(budget.warn_at_percent, 75);
159    }
160
161    #[test]
162    fn test_to_compacting_config_defaults() {
163        let config = MemoryConfig {
164            memory_type: "compacting".to_string(),
165            ..Default::default()
166        };
167        let compacting = config.to_compacting_config();
168        assert_eq!(compacting.max_recent_messages, 50);
169        assert_eq!(compacting.compress_threshold, 30);
170        assert_eq!(compacting.summarize_batch_size, 10);
171    }
172}
ai_agents_runtime/spec/memory.rs

ai_agents_runtime/spec/
memory.rs
