Skip to main content

ai_agents_tools/builtin/
template.rs

1use async_trait::async_trait;
2use minijinja::Environment;
3use schemars::JsonSchema;
4use serde::{Deserialize, Serialize};
5use serde_json::Value;
6
7use crate::generate_schema;
8use ai_agents_core::{Tool, ToolResult};
9
10pub struct TemplateTool {
11    env: Environment<'static>,
12}
13
14impl TemplateTool {
15    pub fn new() -> Self {
16        let mut env = Environment::new();
17        env.set_trim_blocks(true);
18        env.set_lstrip_blocks(true);
19        Self { env }
20    }
21
22    fn render_template(&self, template: &str, data: &Value) -> Result<String, String> {
23        let tmpl = self
24            .env
25            .template_from_str(template)
26            .map_err(|e| format!("Template parse error: {}", e))?;
27
28        tmpl.render(data)
29            .map_err(|e| format!("Render error: {}", e))
30    }
31}
32
33impl Default for TemplateTool {
34    fn default() -> Self {
35        Self::new()
36    }
37}
38
39#[derive(Debug, Deserialize, JsonSchema)]
40struct TemplateInput {
41    /// Operation: render (inline template), render_file (from file)
42    operation: String,
43    /// Template string (for render operation)
44    #[serde(default)]
45    template: Option<String>,
46    /// Template file path (for render_file operation)
47    #[serde(default)]
48    path: Option<String>,
49    /// Data to render into template
50    data: Value,
51}
52
53#[derive(Debug, Serialize, Deserialize)]
54struct RenderOutput {
55    rendered: String,
56}
57
58#[derive(Debug, Serialize, Deserialize)]
59struct RenderFileOutput {
60    rendered: String,
61    template_path: String,
62}
63
64#[async_trait]
65impl Tool for TemplateTool {
66    fn id(&self) -> &str {
67        "template"
68    }
69
70    fn name(&self) -> &str {
71        "Template Renderer"
72    }
73
74    fn description(&self) -> &str {
75        "Render Jinja2-style templates with data. Operations: render (inline template), render_file (from file). Supports variables ({{ var }}), filters ({{ name|upper }}), conditionals ({% if %}), and loops ({% for %})."
76    }
77
78    fn input_schema(&self) -> Value {
79        generate_schema::<TemplateInput>()
80    }
81
82    async fn execute(&self, args: Value) -> ToolResult {
83        let input: TemplateInput = match serde_json::from_value(args) {
84            Ok(input) => input,
85            Err(e) => return ToolResult::error(format!("Invalid input: {}", e)),
86        };
87
88        match input.operation.to_lowercase().as_str() {
89            "render" => self.handle_render(&input),
90            "render_file" => self.handle_render_file(&input),
91            _ => ToolResult::error(format!(
92                "Unknown operation: {}. Valid: render, render_file",
93                input.operation
94            )),
95        }
96    }
97}
98
99impl TemplateTool {
100    fn handle_render(&self, input: &TemplateInput) -> ToolResult {
101        let template = match &input.template {
102            Some(t) => t,
103            None => return ToolResult::error("'template' is required for render operation"),
104        };
105
106        match self.render_template(template, &input.data) {
107            Ok(rendered) => {
108                let output = RenderOutput { rendered };
109                self.to_result(&output)
110            }
111            Err(e) => ToolResult::error(e),
112        }
113    }
114
115    fn handle_render_file(&self, input: &TemplateInput) -> ToolResult {
116        let path = match &input.path {
117            Some(p) => p,
118            None => return ToolResult::error("'path' is required for render_file operation"),
119        };
120
121        let template = match std::fs::read_to_string(path) {
122            Ok(content) => content,
123            Err(e) => return ToolResult::error(format!("File read error: {}", e)),
124        };
125
126        match self.render_template(&template, &input.data) {
127            Ok(rendered) => {
128                let output = RenderFileOutput {
129                    rendered,
130                    template_path: path.clone(),
131                };
132                self.to_result(&output)
133            }
134            Err(e) => ToolResult::error(e),
135        }
136    }
137
138    fn to_result<T: Serialize>(&self, output: &T) -> ToolResult {
139        match serde_json::to_string(output) {
140            Ok(json) => ToolResult::ok(json),
141            Err(e) => ToolResult::error(format!("Serialization error: {}", e)),
142        }
143    }
144}
145
146#[cfg(test)]
147mod tests {
148    use super::*;
149    use tempfile::tempdir;
150
151    #[tokio::test]
152    async fn test_render_simple() {
153        let tool = TemplateTool::new();
154        let result = tool
155            .execute(serde_json::json!({
156                "operation": "render",
157                "template": "Hello {{ name }}!",
158                "data": {"name": "World"}
159            }))
160            .await;
161        assert!(result.success);
162        let output: RenderOutput = serde_json::from_str(&result.output).unwrap();
163        assert_eq!(output.rendered, "Hello World!");
164    }
165
166    #[tokio::test]
167    async fn test_render_with_filter() {
168        let tool = TemplateTool::new();
169        let result = tool
170            .execute(serde_json::json!({
171                "operation": "render",
172                "template": "Hello {{ name|upper }}!",
173                "data": {"name": "world"}
174            }))
175            .await;
176        assert!(result.success);
177        let output: RenderOutput = serde_json::from_str(&result.output).unwrap();
178        assert_eq!(output.rendered, "Hello WORLD!");
179    }
180
181    #[tokio::test]
182    async fn test_render_with_loop() {
183        let tool = TemplateTool::new();
184        let result = tool
185            .execute(serde_json::json!({
186                "operation": "render",
187                "template": "{% for item in items %}{{ item }}{% if not loop.last %}, {% endif %}{% endfor %}",
188                "data": {"items": ["a", "b", "c"]}
189            }))
190            .await;
191        assert!(result.success);
192        let output: RenderOutput = serde_json::from_str(&result.output).unwrap();
193        assert_eq!(output.rendered, "a, b, c");
194    }
195
196    #[tokio::test]
197    async fn test_render_with_conditional() {
198        let tool = TemplateTool::new();
199        let result = tool
200            .execute(serde_json::json!({
201                "operation": "render",
202                "template": "{% if show %}visible{% else %}hidden{% endif %}",
203                "data": {"show": true}
204            }))
205            .await;
206        assert!(result.success);
207        let output: RenderOutput = serde_json::from_str(&result.output).unwrap();
208        assert_eq!(output.rendered, "visible");
209    }
210
211    #[tokio::test]
212    async fn test_render_nested_data() {
213        let tool = TemplateTool::new();
214        let result = tool
215            .execute(serde_json::json!({
216                "operation": "render",
217                "template": "Order #{{ order.id }}: {{ order.items|length }} items",
218                "data": {
219                    "order": {
220                        "id": "12345",
221                        "items": ["item1", "item2", "item3"]
222                    }
223                }
224            }))
225            .await;
226        assert!(result.success);
227        let output: RenderOutput = serde_json::from_str(&result.output).unwrap();
228        assert_eq!(output.rendered, "Order #12345: 3 items");
229    }
230
231    #[tokio::test]
232    async fn test_render_file() {
233        let dir = tempdir().unwrap();
234        let file_path = dir.path().join("template.txt");
235        std::fs::write(&file_path, "Hello {{ name }}!").unwrap();
236
237        let tool = TemplateTool::new();
238        let result = tool
239            .execute(serde_json::json!({
240                "operation": "render_file",
241                "path": file_path.to_str().unwrap(),
242                "data": {"name": "File"}
243            }))
244            .await;
245        assert!(result.success);
246        let output: RenderFileOutput = serde_json::from_str(&result.output).unwrap();
247        assert_eq!(output.rendered, "Hello File!");
248    }
249
250    #[tokio::test]
251    async fn test_render_missing_template() {
252        let tool = TemplateTool::new();
253        let result = tool
254            .execute(serde_json::json!({
255                "operation": "render",
256                "data": {}
257            }))
258            .await;
259        assert!(!result.success);
260    }
261
262    #[tokio::test]
263    async fn test_render_invalid_template() {
264        let tool = TemplateTool::new();
265        let result = tool
266            .execute(serde_json::json!({
267                "operation": "render",
268                "template": "{{ unclosed",
269                "data": {}
270            }))
271            .await;
272        assert!(!result.success);
273    }
274
275    #[tokio::test]
276    async fn test_invalid_operation() {
277        let tool = TemplateTool::new();
278        let result = tool
279            .execute(serde_json::json!({
280                "operation": "invalid",
281                "data": {}
282            }))
283            .await;
284        assert!(!result.success);
285    }
286}