ai_agents_context/
source.rs

1use serde::{Deserialize, Serialize};
2use std::collections::HashMap;
3
4#[derive(Debug, Clone, Serialize, Deserialize, Default, PartialEq)]
5#[serde(rename_all = "snake_case")]
6pub enum RefreshPolicy {
7    Once,
8    #[default]
9    PerSession,
10    PerTurn,
11}
12
13#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
14#[serde(rename_all = "lowercase")]
15pub enum BuiltinSource {
16    Datetime,
17    Session,
18    Agent,
19}
20
21#[derive(Debug, Clone, Serialize, Deserialize)]
22#[serde(tag = "type", rename_all = "lowercase", deny_unknown_fields)]
23pub enum ContextSource {
24    Runtime {
25        #[serde(default)]
26        required: bool,
27        #[serde(default)]
28        schema: Option<serde_json::Value>,
29        #[serde(default)]
30        default: Option<serde_json::Value>,
31    },
32    Builtin {
33        source: BuiltinSource,
34        #[serde(default)]
35        refresh: RefreshPolicy,
36    },
37    File {
38        path: String,
39        #[serde(default)]
40        refresh: RefreshPolicy,
41        #[serde(default)]
42        fallback: Option<String>,
43    },
44    Http {
45        url: String,
46        #[serde(default = "default_method")]
47        method: String,
48        #[serde(default)]
49        headers: HashMap<String, String>,
50        #[serde(default)]
51        refresh: RefreshPolicy,
52        #[serde(default)]
53        timeout_ms: Option<u64>,
54        #[serde(default)]
55        fallback: Option<serde_json::Value>,
56    },
57    Env {
58        name: String,
59    },
60    Callback {
61        name: String,
62        #[serde(default)]
63        refresh: RefreshPolicy,
64    },
65}
66
67fn default_method() -> String {
68    "GET".to_string()
69}
70
71impl ContextSource {
72    pub fn refresh_policy(&self) -> RefreshPolicy {
73        match self {
74            ContextSource::Runtime { .. } => RefreshPolicy::Once,
75            ContextSource::Builtin { refresh, .. } => refresh.clone(),
76            ContextSource::File { refresh, .. } => refresh.clone(),
77            ContextSource::Http { refresh, .. } => refresh.clone(),
78            ContextSource::Env { .. } => RefreshPolicy::Once,
79            ContextSource::Callback { refresh, .. } => refresh.clone(),
80        }
81    }
82
83    pub fn is_required(&self) -> bool {
84        matches!(self, ContextSource::Runtime { required: true, .. })
85    }
86}
87
88#[cfg(test)]
89mod tests {
90    use super::*;
91
92    #[test]
93    fn test_runtime_source() {
94        let yaml = r#"
95type: runtime
96required: true
97default:
98  name: "Guest"
99"#;
100        let source: ContextSource = serde_yaml::from_str(yaml).unwrap();
101        assert!(source.is_required());
102        assert_eq!(source.refresh_policy(), RefreshPolicy::Once);
103    }
104
105    #[test]
106    fn test_builtin_source() {
107        let yaml = r#"
108type: builtin
109source: datetime
110refresh: per_turn
111"#;
112        let source: ContextSource = serde_yaml::from_str(yaml).unwrap();
113        assert_eq!(source.refresh_policy(), RefreshPolicy::PerTurn);
114    }
115
116    #[test]
117    fn test_file_source() {
118        let yaml = r#"
119type: file
120path: "./rules/{{ context.user.language }}/support.txt"
121refresh: per_session
122fallback: "./rules/en/support.txt"
123"#;
124        let source: ContextSource = serde_yaml::from_str(yaml).unwrap();
125        if let ContextSource::File { path, fallback, .. } = source {
126            assert!(path.contains("{{ context.user.language }}"));
127            assert_eq!(fallback, Some("./rules/en/support.txt".into()));
128        } else {
129            panic!("Expected File source");
130        }
131    }
132
133    #[test]
134    fn test_http_source() {
135        let yaml = r#"
136type: http
137url: "https://api.example.com/users/{{ context.user.id }}"
138method: GET
139headers:
140  Authorization: "Bearer {{ env.API_TOKEN }}"
141refresh: per_session
142timeout_ms: 5000
143fallback:
144  theme: "default"
145"#;
146        let source: ContextSource = serde_yaml::from_str(yaml).unwrap();
147        if let ContextSource::Http {
148            url,
149            method,
150            headers,
151            timeout_ms,
152            ..
153        } = source
154        {
155            assert!(url.contains("{{ context.user.id }}"));
156            assert_eq!(method, "GET");
157            assert!(headers.contains_key("Authorization"));
158            assert_eq!(timeout_ms, Some(5000));
159        } else {
160            panic!("Expected Http source");
161        }
162    }
163
164    #[test]
165    fn test_http_source_rejects_removed_cache_ttl() {
166        let yaml = r#"
167type: http
168url: "https://api.example.com/context"
169cache_ttl: 300
170"#;
171        let error = serde_yaml::from_str::<ContextSource>(yaml).unwrap_err();
172        assert!(error.to_string().contains("cache_ttl"));
173    }
174
175    #[test]
176    fn test_env_source() {
177        let yaml = r#"
178type: env
179name: API_TOKEN
180"#;
181        let source: ContextSource = serde_yaml::from_str(yaml).unwrap();
182        if let ContextSource::Env { name } = source {
183            assert_eq!(name, "API_TOKEN");
184        } else {
185            panic!("Expected Env source");
186        }
187    }
188
189    #[test]
190    fn test_callback_source() {
191        let yaml = r#"
192type: callback
193name: get_user_analytics
194refresh: per_session
195"#;
196        let source: ContextSource = serde_yaml::from_str(yaml).unwrap();
197        if let ContextSource::Callback { name, refresh } = source {
198            assert_eq!(name, "get_user_analytics");
199            assert_eq!(refresh, RefreshPolicy::PerSession);
200        } else {
201            panic!("Expected Callback source");
202        }
203    }
204
205    #[test]
206    fn test_context_source_rejects_variant_field_typo() {
207        let yaml = r#"
208type: http
209url: "https://api.example.com"
210timeout_mz: 5000
211"#;
212        assert!(serde_yaml::from_str::<ContextSource>(yaml).is_err());
213    }
214}
ai_agents_context/source.rs

ai_agents_context/
source.rs
