Skip to main content

ai_agents_facts/
config.rs

1//! Configuration types for session management and key facts extraction.
2
3use serde::{Deserialize, Serialize};
4
5/// Config for `memory.actor_memory:` YAML block.
6#[derive(Debug, Clone, Default, Serialize, Deserialize)]
7pub struct ActorMemoryConfig {
8    #[serde(default)]
9    pub enabled: bool,
10    #[serde(default)]
11    pub identification: IdentificationConfig,
12    #[serde(default)]
13    pub injection: InjectionConfig,
14    #[serde(default)]
15    pub privacy: PrivacyConfig,
16}
17
18/// How to resolve the current actor ID.
19#[derive(Debug, Clone, Serialize, Deserialize)]
20pub struct IdentificationConfig {
21    /// explicit = set via set_actor_id API. from_context = read from context path.
22    #[serde(default = "default_explicit")]
23    pub method: IdentificationMethod,
24    /// Context path to read actor_id from (when method = from_context).
25    #[serde(default, skip_serializing_if = "Option::is_none")]
26    pub context_path: Option<String>,
27}
28
29#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
30pub enum IdentificationMethod {
31    #[serde(rename = "explicit")]
32    Explicit,
33    #[serde(rename = "from_context")]
34    FromContext,
35}
36
37fn default_explicit() -> IdentificationMethod {
38    IdentificationMethod::Explicit
39}
40
41impl Default for IdentificationConfig {
42    fn default() -> Self {
43        Self {
44            method: IdentificationMethod::Explicit,
45            context_path: None,
46        }
47    }
48}
49
50/// How facts are injected into the prompt context.
51#[derive(Debug, Clone, Serialize, Deserialize)]
52pub struct InjectionConfig {
53    /// all = inject all facts. category = inject only specified categories. on_demand = no auto-inject.
54    #[serde(default = "default_all")]
55    pub mode: InjectionMode,
56    /// Max tokens for injected facts.
57    #[serde(default = "default_injection_tokens")]
58    pub max_tokens: usize,
59    /// Categories to inject (when mode = category).
60    #[serde(default)]
61    pub categories: Vec<String>,
62}
63
64#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
65pub enum InjectionMode {
66    #[serde(rename = "all")]
67    All,
68    #[serde(rename = "category")]
69    Category,
70    #[serde(rename = "on_demand")]
71    OnDemand,
72}
73
74fn default_all() -> InjectionMode {
75    InjectionMode::All
76}
77
78fn default_injection_tokens() -> usize {
79    800
80}
81
82impl Default for InjectionConfig {
83    fn default() -> Self {
84        Self {
85            mode: InjectionMode::All,
86            max_tokens: 800,
87            categories: vec![],
88        }
89    }
90}
91
92/// Privacy and data retention settings.
93#[derive(Debug, Clone, Serialize, Deserialize)]
94pub struct PrivacyConfig {
95    /// Number of days to retain actor facts. None = no expiry.
96    #[serde(default, skip_serializing_if = "Option::is_none")]
97    pub retention_days: Option<u32>,
98    /// Whether actors can request deletion of all their data.
99    #[serde(default = "default_true")]
100    pub allow_deletion: bool,
101}
102
103fn default_true() -> bool {
104    true
105}
106
107impl Default for PrivacyConfig {
108    fn default() -> Self {
109        Self {
110            retention_days: None,
111            allow_deletion: true,
112        }
113    }
114}
115
116/// Config for `memory.facts:` YAML block.
117#[derive(Debug, Clone, Serialize, Deserialize)]
118pub struct FactsConfig {
119    #[serde(default)]
120    pub enabled: bool,
121    #[serde(default, skip_serializing_if = "Option::is_none")]
122    pub extractor_llm: Option<String>,
123    #[serde(default = "default_true")]
124    pub auto_extract: bool,
125    #[serde(default)]
126    pub categories: Vec<String>,
127    #[serde(default)]
128    pub custom_categories: Vec<CategoryDefinition>,
129    #[serde(default = "default_true")]
130    pub inject_in_context: bool,
131    #[serde(default = "default_max_facts")]
132    pub max_facts: usize,
133    #[serde(default)]
134    pub dedup: DedupConfig,
135    /// Custom extraction prompt. None = use default.
136    #[serde(default, skip_serializing_if = "Option::is_none")]
137    pub extraction_prompt: Option<String>,
138}
139
140fn default_max_facts() -> usize {
141    50
142}
143
144impl Default for FactsConfig {
145    fn default() -> Self {
146        Self {
147            enabled: false,
148            extractor_llm: None,
149            auto_extract: true,
150            categories: vec![],
151            custom_categories: vec![],
152            inject_in_context: true,
153            max_facts: 50,
154            dedup: DedupConfig::default(),
155            extraction_prompt: None,
156        }
157    }
158}
159
160/// User-defined fact category with a description for the LLM.
161#[derive(Debug, Clone, Serialize, Deserialize)]
162pub struct CategoryDefinition {
163    pub name: String,
164    pub description: String,
165}
166
167/// Deduplication strategy for extracted facts.
168#[derive(Debug, Clone, Serialize, Deserialize)]
169pub struct DedupConfig {
170    #[serde(default = "default_true")]
171    pub enabled: bool,
172    /// exact = normalized string match. llm = LLM-based semantic dedup.
173    #[serde(default = "default_dedup_method")]
174    pub method: DedupMethod,
175}
176
177#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
178pub enum DedupMethod {
179    #[serde(rename = "llm")]
180    Llm,
181    #[serde(rename = "exact")]
182    Exact,
183}
184
185fn default_dedup_method() -> DedupMethod {
186    DedupMethod::Exact
187}
188
189impl Default for DedupConfig {
190    fn default() -> Self {
191        Self {
192            enabled: true,
193            method: DedupMethod::Exact,
194        }
195    }
196}
197
198/// Config for `memory.session:` YAML block.
199#[derive(Debug, Clone, Default, Serialize, Deserialize)]
200pub struct SessionConfig {
201    #[serde(default)]
202    pub tags: Vec<String>,
203    #[serde(default, skip_serializing_if = "Option::is_none")]
204    pub ttl_seconds: Option<u64>,
205}
206
207#[cfg(test)]
208mod tests {
209    use super::*;
210
211    #[test]
212    fn test_actor_memory_config_default() {
213        let config = ActorMemoryConfig::default();
214        assert!(!config.enabled);
215        assert_eq!(config.identification.method, IdentificationMethod::Explicit);
216        assert_eq!(config.injection.mode, InjectionMode::All);
217        assert_eq!(config.injection.max_tokens, 800);
218        assert!(config.privacy.allow_deletion);
219    }
220
221    #[test]
222    fn test_facts_config_default() {
223        let config = FactsConfig::default();
224        assert!(!config.enabled);
225        assert!(config.auto_extract);
226        assert!(config.inject_in_context);
227        assert_eq!(config.max_facts, 50);
228        assert!(config.dedup.enabled);
229        assert_eq!(config.dedup.method, DedupMethod::Exact);
230    }
231
232    #[test]
233    fn test_facts_config_deserialize() {
234        let yaml = r#"
235enabled: true
236extractor_llm: router
237auto_extract: true
238categories:
239  - user_preference
240  - user_context
241  - decision
242custom_categories:
243  - name: suspicion
244    description: "Suspicious behavior observed"
245inject_in_context: true
246max_facts: 30
247dedup:
248  enabled: true
249  method: llm
250"#;
251        let config: FactsConfig = serde_yaml::from_str(yaml).unwrap();
252        assert!(config.enabled);
253        assert_eq!(config.extractor_llm.as_deref(), Some("router"));
254        assert_eq!(config.categories.len(), 3);
255        assert_eq!(config.custom_categories.len(), 1);
256        assert_eq!(config.custom_categories[0].name, "suspicion");
257        assert_eq!(config.max_facts, 30);
258        assert_eq!(config.dedup.method, DedupMethod::Llm);
259    }
260
261    #[test]
262    fn test_actor_memory_config_deserialize() {
263        let yaml = r#"
264enabled: true
265identification:
266  method: from_context
267  context_path: user.id
268injection:
269  mode: category
270  max_tokens: 500
271  categories:
272    - user_preference
273privacy:
274  retention_days: 365
275  allow_deletion: true
276"#;
277        let config: ActorMemoryConfig = serde_yaml::from_str(yaml).unwrap();
278        assert!(config.enabled);
279        assert_eq!(
280            config.identification.method,
281            IdentificationMethod::FromContext
282        );
283        assert_eq!(
284            config.identification.context_path.as_deref(),
285            Some("user.id")
286        );
287        assert_eq!(config.injection.mode, InjectionMode::Category);
288        assert_eq!(config.injection.max_tokens, 500);
289        assert_eq!(config.injection.categories, vec!["user_preference"]);
290        assert_eq!(config.privacy.retention_days, Some(365));
291    }
292
293    #[test]
294    fn test_session_config_deserialize() {
295        let yaml = r#"
296tags: [support, tier-1]
297ttl_seconds: 86400
298"#;
299        let config: SessionConfig = serde_yaml::from_str(yaml).unwrap();
300        assert_eq!(config.tags, vec!["support", "tier-1"]);
301        assert_eq!(config.ttl_seconds, Some(86400));
302    }
303
304    #[test]
305    fn test_session_config_default() {
306        let config = SessionConfig::default();
307        assert!(config.tags.is_empty());
308        assert!(config.ttl_seconds.is_none());
309    }
310}