askit_std_agents/
data.rs

1use std::vec;
2
3use agent_stream_kit::{
4    ASKit, Agent, AgentConfigs, AgentContext, AgentData, AgentError, AgentOutput, AgentValue,
5    AsAgent, async_trait,
6};
7use askit_macros::askit_agent;
8
9static CATEGORY: &str = "Std/Data";
10
11static PIN_JSON: &str = "json";
12static PIN_VALUE: &str = "value";
13
14static CONFIG_KEY: &str = "key";
15static CONFIG_VALUE: &str = "value";
16
17// Get Value
18#[askit_agent(
19    title = "Get Value",
20    category = CATEGORY,
21    inputs = [PIN_VALUE],
22    outputs = [PIN_VALUE],
23    string_config(name = CONFIG_KEY)
24)]
25struct GetValueAgent {
26    data: AgentData,
27}
28
29#[async_trait]
30impl AsAgent for GetValueAgent {
31    fn new(
32        askit: ASKit,
33        id: String,
34        def_name: String,
35        config: Option<AgentConfigs>,
36    ) -> Result<Self, AgentError> {
37        Ok(Self {
38            data: AgentData::new(askit, id, def_name, config),
39        })
40    }
41
42    async fn process(
43        &mut self,
44        ctx: AgentContext,
45        _pin: String,
46        value: AgentValue,
47    ) -> Result<(), AgentError> {
48        let key = self.configs()?.get_string(CONFIG_KEY)?;
49        if key.is_empty() {
50            return Ok(());
51        }
52        let keys = key.split('.').collect::<Vec<_>>();
53
54        if value.is_object() {
55            if let Some(value) = get_nested_value(&value, &keys) {
56                self.try_output(ctx, PIN_VALUE, value.to_owned())?;
57            } else {
58                self.try_output(ctx, PIN_VALUE, AgentValue::unit())?;
59            }
60        } else if value.is_array() {
61            let mut out_arr = Vec::new();
62            for v in value
63                .as_array()
64                .ok_or_else(|| AgentError::InvalidValue("failed as_array".to_string()))?
65            {
66                let value = get_nested_value(v, &keys);
67                if let Some(v) = value {
68                    out_arr.push(v.to_owned());
69                } else {
70                    out_arr.push(AgentValue::unit());
71                }
72            }
73            self.try_output(ctx, PIN_VALUE, AgentValue::array(out_arr))?;
74        }
75
76        Ok(())
77    }
78}
79
80// Set Value
81#[askit_agent(
82    title = "Set Value",
83    category = CATEGORY,
84    inputs = [PIN_VALUE],
85    outputs = [PIN_VALUE],
86    string_config(name = CONFIG_KEY),
87    object_config(name = CONFIG_VALUE),
88)]
89struct SetValueAgent {
90    data: AgentData,
91}
92
93#[async_trait]
94impl AsAgent for SetValueAgent {
95    fn new(
96        askit: ASKit,
97        id: String,
98        def_name: String,
99        config: Option<AgentConfigs>,
100    ) -> Result<Self, AgentError> {
101        Ok(Self {
102            data: AgentData::new(askit, id, def_name, config),
103        })
104    }
105
106    async fn process(
107        &mut self,
108        ctx: AgentContext,
109        _pin: String,
110        value: AgentValue,
111    ) -> Result<(), AgentError> {
112        // parse key
113        let key = self.configs()?.get_string(CONFIG_KEY)?;
114        if key.is_empty() {
115            return Ok(());
116        }
117        let keys = key.split('.').collect::<Vec<_>>();
118
119        let v = self.configs()?.get(CONFIG_VALUE)?;
120        let mut value = value;
121        set_nested_value(&mut value, keys, v.clone());
122
123        self.try_output(ctx, PIN_VALUE, value)?;
124
125        Ok(())
126    }
127}
128
129// To Object
130#[askit_agent(
131    title = "To Object",
132    category = CATEGORY,
133    inputs = [PIN_VALUE],
134    outputs = [PIN_VALUE],
135    string_config(name = CONFIG_KEY)
136)]
137struct ToObjectAgent {
138    data: AgentData,
139}
140
141#[async_trait]
142impl AsAgent for ToObjectAgent {
143    fn new(
144        askit: ASKit,
145        id: String,
146        def_name: String,
147        config: Option<AgentConfigs>,
148    ) -> Result<Self, AgentError> {
149        Ok(Self {
150            data: AgentData::new(askit, id, def_name, config),
151        })
152    }
153
154    async fn process(
155        &mut self,
156        ctx: AgentContext,
157        _pin: String,
158        value: AgentValue,
159    ) -> Result<(), AgentError> {
160        let key = self.configs()?.get_string(CONFIG_KEY)?;
161        if key.is_empty() {
162            return Ok(());
163        }
164
165        let keys = key.split('.').collect::<Vec<_>>();
166        let mut new_value = AgentValue::object_default();
167        set_nested_value(&mut new_value, keys, value);
168
169        self.try_output(ctx, PIN_VALUE, new_value)?;
170        Ok(())
171    }
172}
173
174// To JSON
175#[askit_agent(
176    title = "To JSON",
177    category = CATEGORY,
178    inputs = [PIN_VALUE],
179    outputs = [PIN_JSON]
180)]
181struct ToJsonAgent {
182    data: AgentData,
183}
184
185#[async_trait]
186impl AsAgent for ToJsonAgent {
187    fn new(
188        askit: ASKit,
189        id: String,
190        def_name: String,
191        config: Option<AgentConfigs>,
192    ) -> Result<Self, AgentError> {
193        Ok(Self {
194            data: AgentData::new(askit, id, def_name, config),
195        })
196    }
197
198    async fn process(
199        &mut self,
200        ctx: AgentContext,
201        _pin: String,
202        value: AgentValue,
203    ) -> Result<(), AgentError> {
204        let json = serde_json::to_string_pretty(&value)
205            .map_err(|e| AgentError::InvalidValue(e.to_string()))?;
206        self.try_output(ctx, PIN_JSON, AgentValue::string(json))?;
207        Ok(())
208    }
209}
210
211// From JSON
212#[askit_agent(
213    title = "From JSON",
214    category = CATEGORY,
215    inputs = [PIN_JSON],
216    outputs = [PIN_VALUE]
217)]
218struct FromJsonAgent {
219    data: AgentData,
220}
221
222#[async_trait]
223impl AsAgent for FromJsonAgent {
224    fn new(
225        askit: ASKit,
226        id: String,
227        def_name: String,
228        config: Option<AgentConfigs>,
229    ) -> Result<Self, AgentError> {
230        Ok(Self {
231            data: AgentData::new(askit, id, def_name, config),
232        })
233    }
234
235    async fn process(
236        &mut self,
237        ctx: AgentContext,
238        _pin: String,
239        value: AgentValue,
240    ) -> Result<(), AgentError> {
241        let s = value
242            .as_str()
243            .ok_or_else(|| AgentError::InvalidValue("not a string".to_string()))?;
244        let json_value: serde_json::Value =
245            serde_json::from_str(s).map_err(|e| AgentError::InvalidValue(e.to_string()))?;
246        let value = AgentValue::from_json(json_value)?;
247        self.try_output(ctx, PIN_VALUE, value)?;
248        Ok(())
249    }
250}
251
252fn get_nested_value<'a>(value: &'a AgentValue, keys: &[&str]) -> Option<&'a AgentValue> {
253    let mut current_value = value;
254    for key in keys {
255        let obj = current_value.as_object()?;
256        current_value = obj.get(*key)?;
257    }
258    Some(current_value)
259}
260
261fn set_nested_value<'a>(value: &'a mut AgentValue, keys: Vec<&str>, new_value: AgentValue) {
262    let mut current_value = value;
263
264    if keys.is_empty() {
265        return;
266    }
267
268    for key in keys[..keys.len() - 1].iter() {
269        if !current_value.is_object() {
270            return;
271        }
272
273        if current_value.get(*key).is_none() {
274            let _ = current_value.set((*key).to_string(), AgentValue::object_default());
275        }
276
277        if let Some(v) = current_value.get_mut(*key) {
278            current_value = v;
279        } else {
280            // just in case
281            return;
282        }
283    }
284
285    let last_key = keys.last().unwrap();
286    if let Some(obj) = current_value.as_object_mut() {
287        obj.insert((*last_key).to_string(), new_value);
288    }
289}
290
291#[cfg(test)]
292mod tests {
293    use super::*;
294
295    #[test]
296    fn test_get_nested_value() {
297        // Setup data: { "users": { "admin": { "name": "Alice" } } }
298        let mut root = AgentValue::object_default();
299        let mut users = AgentValue::object_default();
300        let mut admin = AgentValue::object_default();
301
302        admin
303            .set("name".to_string(), AgentValue::string("Alice"))
304            .unwrap();
305        users.set("admin".to_string(), admin).unwrap();
306        root.set("users".to_string(), users).unwrap();
307
308        // Case 1: Successfully retrieve an existing value
309        let keys = vec!["users", "admin", "name"];
310        let result = get_nested_value(&root, &keys);
311        assert_eq!(result, Some(&AgentValue::string("Alice")));
312
313        // Case 2: Intermediate key does not exist (users -> guest)
314        let keys_missing = vec!["users", "guest", "name"];
315        let result_missing = get_nested_value(&root, &keys_missing);
316        assert_eq!(result_missing, None);
317
318        // Case 3: Intermediate path is not an object (users -> admin -> name -> something)
319        // "name" is a string, so we cannot traverse deeper -> Should return None
320        let keys_not_obj = vec!["users", "admin", "name", "length"];
321        let result_not_obj = get_nested_value(&root, &keys_not_obj);
322        assert_eq!(result_not_obj, None); // Filtered out by as_object()?
323
324        // Case 4: Empty keys (Should return the root object)
325        let keys_empty: Vec<&str> = vec![];
326        let result_root = get_nested_value(&root, &keys_empty);
327        assert_eq!(result_root, Some(&root));
328    }
329
330    /// Test 1: Verify if a deeply nested structure (a.b.c) can be auto-generated from an empty state.
331    /// This confirms the fix for the previous bug (failure to traverse down levels).
332    #[test]
333    fn test_create_deeply_nested_structure() {
334        let mut root = AgentValue::object_default();
335        let keys = vec!["users", "admin", "name"];
336        let value = AgentValue::string("Alice");
337
338        set_nested_value(&mut root, keys, value);
339
340        // Verify: root["users"]["admin"]["name"] == "Alice"
341        if let Some(users) = root.get_mut("users") {
342            if let Some(admin) = users.get_mut("admin") {
343                if let Some(name) = admin.get_mut("name") {
344                    assert_eq!(*name, AgentValue::string("Alice"));
345                    return;
346                }
347            }
348        }
349        panic!("Nested structure was not created correctly: {:?}", root);
350    }
351
352    /// Test 2: Verify if a new key can be added without breaking existing structures.
353    #[test]
354    fn test_add_to_existing_structure() {
355        let mut root = AgentValue::object_default();
356        // Pre-create { "config": {} }
357        root.set("config".to_string(), AgentValue::object_default())
358            .unwrap();
359
360        let keys = vec!["config", "timeout"];
361        let value = AgentValue::string("30s");
362
363        set_nested_value(&mut root, keys, value);
364
365        // Verify
366        let config = root.get_mut("config").unwrap();
367        let timeout = config.get_mut("timeout").unwrap();
368        assert_eq!(*timeout, AgentValue::string("30s"));
369    }
370
371    /// Test 3: Verify if an existing value can be overwritten.
372    #[test]
373    fn test_overwrite_existing_value() {
374        let mut root = AgentValue::object_default();
375        // Pre-create { "app": { "version": "v1" } }
376        let mut app = AgentValue::object_default();
377        app.set("version".to_string(), AgentValue::string("v1"))
378            .unwrap();
379        root.set("app".to_string(), app).unwrap();
380
381        // Execute overwrite
382        let keys = vec!["app", "version"];
383        let new_val = AgentValue::string("v2");
384        set_nested_value(&mut root, keys, new_val);
385
386        // Verify
387        let app = root.get_mut("app").unwrap();
388        let version = app.get_mut("version").unwrap();
389        assert_eq!(*version, AgentValue::string("v2"));
390    }
391
392    /// Test 4: Verify if the operation stops safely when an intermediate path is not an object.
393    /// Example: Try setting ["tags", "new_key"] against { "tags": "immutable_string" }
394    #[test]
395    fn test_stop_if_path_is_not_object() {
396        let mut root = AgentValue::object_default();
397        // "tags" is a string, not an object
398        root.set("tags".to_string(), AgentValue::string("some_string"))
399            .unwrap();
400
401        let keys = vec!["tags", "new_key"];
402        let value = AgentValue::string("value");
403
404        // Ensure it returns without crashing
405        set_nested_value(&mut root, keys, value);
406
407        // Verify that "tags" remains a string
408        let tags = root.get_mut("tags").unwrap();
409        assert_eq!(*tags, AgentValue::string("some_string"));
410    }
411}
askit_std_agents/data.rs

askit_std_agents/
data.rs
