Skip to main content

askit_std_agents/
string.rs

1use agent_stream_kit::{
2    ASKit, Agent, AgentContext, AgentData, AgentError, AgentOutput, AgentSpec, AgentValue, AsAgent,
3    askit_agent, async_trait,
4};
5use handlebars::Handlebars;
6use im::vector;
7use serde_json::json;
8
9const CATEGORY: &str = "Std/String";
10
11const PIN_STRING: &str = "string";
12const PIN_STRINGS: &str = "strings";
13const PIN_VALUE: &str = "value";
14const PIN_T: &str = "t";
15const PIN_F: &str = "f";
16
17const CONFIG_LEN: &str = "len";
18const CONFIG_OVERLAP: &str = "overlap";
19const CONFIG_SEP: &str = "sep";
20const CONFIG_TEMPLATE: &str = "template";
21
22/// Check if the input is a string.
23#[askit_agent(
24    title = "IsString",
25    category = CATEGORY,
26    inputs = [PIN_VALUE],
27    outputs = [PIN_T, PIN_F],
28)]
29struct IsStringAgent {
30    data: AgentData,
31}
32
33#[async_trait]
34impl AsAgent for IsStringAgent {
35    fn new(askit: ASKit, id: String, spec: AgentSpec) -> Result<Self, AgentError> {
36        Ok(Self {
37            data: AgentData::new(askit, id, spec),
38        })
39    }
40
41    async fn process(
42        &mut self,
43        ctx: AgentContext,
44        _pin: String,
45        value: AgentValue,
46    ) -> Result<(), AgentError> {
47        if value.is_string() {
48            self.output(ctx, PIN_T, value).await
49        } else {
50            self.output(ctx, PIN_F, value).await
51        }
52    }
53}
54
55/// Check if the input string is empty.
56#[askit_agent(
57    title = "IsEmptyString",
58    category = CATEGORY,
59    inputs = [PIN_STRING],
60    outputs = [PIN_T, PIN_F],
61)]
62struct IsEmptyStringAgent {
63    data: AgentData,
64}
65
66#[async_trait]
67impl AsAgent for IsEmptyStringAgent {
68    fn new(askit: ASKit, id: String, spec: AgentSpec) -> Result<Self, AgentError> {
69        Ok(Self {
70            data: AgentData::new(askit, id, spec),
71        })
72    }
73
74    async fn process(
75        &mut self,
76        ctx: AgentContext,
77        _pin: String,
78        value: AgentValue,
79    ) -> Result<(), AgentError> {
80        let is_empty = if let Some(s) = value.as_str() {
81            s.is_empty()
82        } else {
83            false
84        };
85        if is_empty {
86            self.output(ctx, PIN_T, value).await
87        } else {
88            self.output(ctx, PIN_F, value).await
89        }
90    }
91}
92
93/// The `StringJoinAgent` is responsible for joining an array of strings into a single string
94/// using a specified separator. It processes input value, applies transformations to handle
95/// escape sequences (e.g., `\n`, `\t`), and outputs the resulting string.
96///
97/// # Configuration
98/// - `CONFIG_SEP`: Specifies the separator to use when joining strings. Defaults to an empty string.
99///
100/// # Input
101/// - Expects an array of strings as input value.
102///
103/// # Output
104/// - Produces a single joined string as output.
105///
106/// # Example
107/// Given the input `["Hello", "World"]` and `CONFIG_SEP` set to `" "`, the output will be `"Hello World"`.
108#[askit_agent(
109    title = "String Join",
110    category = CATEGORY,
111    inputs = [PIN_STRINGS],
112    outputs = [PIN_STRING],
113    string_config(name = CONFIG_SEP, default = "\\n")
114)]
115struct StringJoinAgent {
116    data: AgentData,
117}
118
119#[async_trait]
120impl AsAgent for StringJoinAgent {
121    fn new(askit: ASKit, id: String, spec: AgentSpec) -> Result<Self, AgentError> {
122        Ok(Self {
123            data: AgentData::new(askit, id, spec),
124        })
125    }
126
127    async fn process(
128        &mut self,
129        ctx: AgentContext,
130        _pin: String,
131        value: AgentValue,
132    ) -> Result<(), AgentError> {
133        let config = self.configs()?;
134
135        let sep = config.get_string_or_default(CONFIG_SEP);
136
137        if value.is_array() {
138            let mut out = Vec::new();
139            for v in value
140                .as_array()
141                .ok_or_else(|| AgentError::InvalidArrayValue("Expected array".into()))?
142            {
143                out.push(v.as_str().unwrap_or_default());
144            }
145            let mut out = out.join(&sep);
146            out = out.replace("\\n", "\n");
147            out = out.replace("\\t", "\t");
148            out = out.replace("\\r", "\r");
149            out = out.replace("\\\\", "\\");
150            let out_value = AgentValue::string(out);
151            self.output(ctx, PIN_STRING, out_value).await
152        } else {
153            self.output(ctx, PIN_STRING, value).await
154        }
155    }
156}
157
158#[askit_agent(
159    title = "String Length Split",
160    category = CATEGORY,
161    inputs = [PIN_STRING],
162    outputs = [PIN_STRINGS],
163    integer_config(name = CONFIG_LEN, default = 65536),
164    integer_config(name = CONFIG_OVERLAP, default = 1024),
165)]
166struct StringLengthSplitAgent {
167    data: AgentData,
168}
169
170#[async_trait]
171impl AsAgent for StringLengthSplitAgent {
172    fn new(askit: ASKit, id: String, spec: AgentSpec) -> Result<Self, AgentError> {
173        Ok(Self {
174            data: AgentData::new(askit, id, spec),
175        })
176    }
177
178    async fn process(
179        &mut self,
180        ctx: AgentContext,
181        _pin: String,
182        value: AgentValue,
183    ) -> Result<(), AgentError> {
184        let config = self.configs()?;
185
186        let n = config.get_integer_or_default(CONFIG_LEN) as usize;
187        if n <= 0 {
188            return Err(AgentError::InvalidConfig("n must be greater than 0".into()));
189        }
190
191        let overlap = config.get_integer_or_default(CONFIG_OVERLAP) as usize;
192        if overlap >= n {
193            return Err(AgentError::InvalidConfig(
194                "overlap must be less than n".into(),
195            ));
196        }
197
198        let s = value
199            .as_str()
200            .ok_or_else(|| AgentError::InvalidValue("Input value must be a string".into()))?;
201
202        let mut out = Vec::new();
203        let mut start = 0;
204        let len = s.len();
205        while start < len {
206            let mut end = usize::min(start + n, len);
207            while !s.is_char_boundary(end) {
208                end -= 1;
209            }
210            if end <= start {
211                end = start + s[start..].chars().next().map(|c| c.len_utf8()).unwrap_or(1);
212            }
213
214            out.push(AgentValue::string(s[start..end].to_string()));
215
216            if end == len {
217                break;
218            }
219
220            let mut next_start = end.saturating_sub(overlap);
221            while next_start < len && !s.is_char_boundary(next_start) {
222                next_start += 1;
223            }
224            start = next_start;
225        }
226        self.output(ctx, PIN_STRINGS, AgentValue::array(out.into()))
227            .await
228    }
229}
230
231// Template String Agent
232#[askit_agent(
233    title = "Template String",
234    category = CATEGORY,
235    inputs = [PIN_VALUE],
236    outputs = [PIN_STRING],
237    string_config(name = CONFIG_TEMPLATE, default = "{{value}}")
238)]
239struct TemplateStringAgent {
240    data: AgentData,
241}
242
243#[async_trait]
244impl AsAgent for TemplateStringAgent {
245    fn new(askit: ASKit, id: String, spec: AgentSpec) -> Result<Self, AgentError> {
246        Ok(Self {
247            data: AgentData::new(askit, id, spec),
248        })
249    }
250
251    async fn process(
252        &mut self,
253        ctx: AgentContext,
254        _pin: String,
255        value: AgentValue,
256    ) -> Result<(), AgentError> {
257        let config = self.configs()?;
258
259        let template = config.get_string_or_default(CONFIG_TEMPLATE);
260        if template.is_empty() {
261            return Err(AgentError::InvalidConfig("template is not set".into()));
262        }
263
264        let reg = handlebars_new();
265
266        if value.is_array() {
267            let mut out_arr = Vec::new();
268            for v in value
269                .as_array()
270                .ok_or_else(|| AgentError::InvalidArrayValue("Expected array".into()))?
271            {
272                let data = json!({"value": v});
273                let rendered_string = reg.render_template(&template, &data).map_err(|e| {
274                    AgentError::InvalidValue(format!("Failed to render template: {}", e))
275                })?;
276                out_arr.push(rendered_string.into());
277            }
278            self.output(ctx, PIN_STRING, AgentValue::array(out_arr.into()))
279                .await
280        } else {
281            let data = json!({"value": value});
282            let rendered_string = reg.render_template(&template, &data).map_err(|e| {
283                AgentError::InvalidValue(format!("Failed to render template: {}", e))
284            })?;
285            let out_value = AgentValue::string(rendered_string);
286            self.output(ctx, PIN_STRING, out_value).await
287        }
288    }
289}
290
291// Template Text Agent
292#[askit_agent(
293    title = "Template Text",
294    category = CATEGORY,
295    inputs = [PIN_VALUE],
296    outputs = [PIN_STRING],
297    text_config(name = CONFIG_TEMPLATE, default = "{{value}}")
298)]
299struct TemplateTextAgent {
300    data: AgentData,
301}
302
303#[async_trait]
304impl AsAgent for TemplateTextAgent {
305    fn new(askit: ASKit, id: String, spec: AgentSpec) -> Result<Self, AgentError> {
306        Ok(Self {
307            data: AgentData::new(askit, id, spec),
308        })
309    }
310
311    async fn process(
312        &mut self,
313        ctx: AgentContext,
314        _pin: String,
315        value: AgentValue,
316    ) -> Result<(), AgentError> {
317        let config = self.configs()?;
318
319        let template = config.get_string_or_default(CONFIG_TEMPLATE);
320        if template.is_empty() {
321            return Err(AgentError::InvalidConfig("template is not set".into()));
322        }
323
324        let reg = handlebars_new();
325
326        if value.is_array() {
327            let mut out_arr = Vec::new();
328            for v in value
329                .as_array()
330                .ok_or_else(|| AgentError::InvalidArrayValue("Expected array".into()))?
331            {
332                let data = json!({"value": v});
333                let rendered_string = reg.render_template(&template, &data).map_err(|e| {
334                    AgentError::InvalidValue(format!("Failed to render template: {}", e))
335                })?;
336                out_arr.push(rendered_string.into());
337            }
338            self.output(ctx, PIN_STRING, AgentValue::array(out_arr.into()))
339                .await
340        } else {
341            let data = json!({"value": value});
342            let rendered_string = reg.render_template(&template, &data).map_err(|e| {
343                AgentError::InvalidValue(format!("Failed to render template: {}", e))
344            })?;
345            let out_value = AgentValue::string(rendered_string);
346            self.output(ctx, PIN_STRING, out_value).await
347        }
348    }
349}
350
351// Template Array Agent
352#[askit_agent(
353    title = "Template Array",
354    category = CATEGORY,
355    inputs = [PIN_VALUE],
356    outputs = [PIN_STRING],
357    text_config(name = CONFIG_TEMPLATE, default = "{{value}}")
358)]
359struct TemplateArrayAgent {
360    data: AgentData,
361}
362
363#[async_trait]
364impl AsAgent for TemplateArrayAgent {
365    fn new(askit: ASKit, id: String, spec: AgentSpec) -> Result<Self, AgentError> {
366        Ok(Self {
367            data: AgentData::new(askit, id, spec),
368        })
369    }
370
371    async fn process(
372        &mut self,
373        ctx: AgentContext,
374        _pin: String,
375        value: AgentValue,
376    ) -> Result<(), AgentError> {
377        let config = self.configs()?;
378
379        let template = config.get_string_or_default(CONFIG_TEMPLATE);
380        if template.is_empty() {
381            return Err(AgentError::InvalidConfig("template is not set".into()));
382        }
383
384        let reg = handlebars_new();
385
386        if value.is_array() {
387            let rendered_string = reg.render_template(&template, &value).map_err(|e| {
388                AgentError::InvalidValue(format!("Failed to render template: {}", e))
389            })?;
390            self.output(ctx, PIN_STRING, AgentValue::string(rendered_string))
391                .await
392        } else {
393            let d = AgentValue::array(vector![value.clone()]);
394            let rendered_string = reg.render_template(&template, &d).map_err(|e| {
395                AgentError::InvalidValue(format!("Failed to render template: {}", e))
396            })?;
397            let out_value = AgentValue::string(rendered_string);
398            self.output(ctx, PIN_STRING, out_value).await
399        }
400    }
401}
402
403fn handlebars_new<'a>() -> Handlebars<'a> {
404    let mut reg = Handlebars::new();
405    reg.register_escape_fn(handlebars::no_escape);
406    reg.register_helper("to_json", Box::new(to_json_helper));
407
408    #[cfg(feature = "yaml")]
409    reg.register_helper("to_yaml", Box::new(to_yaml_helper));
410
411    reg
412}
413
414fn to_json_helper(
415    h: &handlebars::Helper<'_>,
416    _: &handlebars::Handlebars<'_>,
417    _: &handlebars::Context,
418    _: &mut handlebars::RenderContext<'_, '_>,
419    out: &mut dyn handlebars::Output,
420) -> handlebars::HelperResult {
421    if let Some(value) = h.param(0) {
422        let json_str = serde_json::to_string_pretty(&value.value()).map_err(|e| {
423            handlebars::RenderErrorReason::Other(format!("Failed to serialize to JSON: {}", e))
424        })?;
425        out.write(&json_str)?;
426    }
427    Ok(())
428}
429
430#[cfg(feature = "yaml")]
431fn to_yaml_helper(
432    h: &handlebars::Helper<'_>,
433    _: &handlebars::Handlebars<'_>,
434    _: &handlebars::Context,
435    _: &mut handlebars::RenderContext<'_, '_>,
436    out: &mut dyn handlebars::Output,
437) -> handlebars::HelperResult {
438    if let Some(value) = h.param(0) {
439        let yaml_str = serde_yaml_ng::to_string(&value.value()).map_err(|e| {
440            handlebars::RenderErrorReason::Other(format!("Failed to serialize to YAML: {}", e))
441        })?;
442        out.write(&yaml_str)?;
443    }
444    Ok(())
445}