agent_chain_core/utils/
mustache.rs

1//! Mustache template rendering.
2//!
3//! Adapted from langchain_core/utils/mustache.py (which is based on chevron)
4//! MIT License.
5
6use std::collections::HashMap;
7
8/// Error type for mustache template operations.
9#[derive(Debug, Clone, PartialEq)]
10pub enum MustacheError {
11    /// Syntax error in template.
12    SyntaxError(String),
13    /// Unclosed tag.
14    UnclosedTag { line: usize },
15    /// Unclosed section.
16    UnclosedSection { tag: String, line: usize },
17    /// Mismatched section closing tag.
18    MismatchedSection {
19        expected: String,
20        got: String,
21        line: usize,
22    },
23    /// Empty tag.
24    EmptyTag { line: usize },
25}
26
27impl std::fmt::Display for MustacheError {
28    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
29        match self {
30            MustacheError::SyntaxError(msg) => write!(f, "Syntax error: {}", msg),
31            MustacheError::UnclosedTag { line } => write!(f, "Unclosed tag at line {}", line),
32            MustacheError::UnclosedSection { tag, line } => {
33                write!(f, "Unclosed section '{}' opened at line {}", tag, line)
34            }
35            MustacheError::MismatchedSection {
36                expected,
37                got,
38                line,
39            } => {
40                write!(
41                    f,
42                    "Mismatched section at line {}: expected '{}', got '{}'",
43                    line, expected, got
44                )
45            }
46            MustacheError::EmptyTag { line } => write!(f, "Empty tag at line {}", line),
47        }
48    }
49}
50
51impl std::error::Error for MustacheError {}
52
53/// A mustache template value.
54#[derive(Debug, Clone)]
55pub enum MustacheValue {
56    /// A string value.
57    String(String),
58    /// A boolean value.
59    Bool(bool),
60    /// An integer value.
61    Int(i64),
62    /// A floating point value.
63    Float(f64),
64    /// A list of values.
65    List(Vec<MustacheValue>),
66    /// A map of values.
67    Map(HashMap<String, MustacheValue>),
68    /// A null value.
69    Null,
70}
71
72impl MustacheValue {
73    /// Check if the value is truthy.
74    pub fn is_truthy(&self) -> bool {
75        match self {
76            MustacheValue::String(s) => !s.is_empty(),
77            MustacheValue::Bool(b) => *b,
78            MustacheValue::Int(i) => *i != 0,
79            MustacheValue::Float(f) => *f != 0.0,
80            MustacheValue::List(l) => !l.is_empty(),
81            MustacheValue::Map(_) => true,
82            MustacheValue::Null => false,
83        }
84    }
85
86    /// Convert to string for output.
87    pub fn to_output_string(&self) -> String {
88        match self {
89            MustacheValue::String(s) => s.clone(),
90            MustacheValue::Bool(b) => b.to_string(),
91            MustacheValue::Int(i) => i.to_string(),
92            MustacheValue::Float(f) => f.to_string(),
93            MustacheValue::List(_) => String::new(),
94            MustacheValue::Map(_) => String::new(),
95            MustacheValue::Null => String::new(),
96        }
97    }
98}
99
100impl From<String> for MustacheValue {
101    fn from(s: String) -> Self {
102        MustacheValue::String(s)
103    }
104}
105
106impl From<&str> for MustacheValue {
107    fn from(s: &str) -> Self {
108        MustacheValue::String(s.to_string())
109    }
110}
111
112impl From<bool> for MustacheValue {
113    fn from(b: bool) -> Self {
114        MustacheValue::Bool(b)
115    }
116}
117
118impl From<i64> for MustacheValue {
119    fn from(i: i64) -> Self {
120        MustacheValue::Int(i)
121    }
122}
123
124impl From<i32> for MustacheValue {
125    fn from(i: i32) -> Self {
126        MustacheValue::Int(i as i64)
127    }
128}
129
130impl From<f64> for MustacheValue {
131    fn from(f: f64) -> Self {
132        MustacheValue::Float(f)
133    }
134}
135
136impl From<Vec<MustacheValue>> for MustacheValue {
137    fn from(v: Vec<MustacheValue>) -> Self {
138        MustacheValue::List(v)
139    }
140}
141
142impl From<HashMap<String, MustacheValue>> for MustacheValue {
143    fn from(m: HashMap<String, MustacheValue>) -> Self {
144        MustacheValue::Map(m)
145    }
146}
147
148/// Token types for mustache parsing.
149#[derive(Debug, Clone, PartialEq)]
150enum TokenType {
151    Literal,
152    Variable,
153    NoEscape,
154    Section,
155    InvertedSection,
156    End,
157    Partial,
158    Comment,
159}
160
161/// A token from mustache parsing.
162#[derive(Debug, Clone)]
163struct Token {
164    token_type: TokenType,
165    key: String,
166}
167
168/// Tokenize a mustache template.
169fn tokenize(template: &str, l_del: &str, r_del: &str) -> Result<Vec<Token>, MustacheError> {
170    let mut tokens = Vec::new();
171    let mut current_line = 1;
172    let mut open_sections: Vec<(String, usize)> = Vec::new();
173    let mut remaining = template;
174
175    while !remaining.is_empty() {
176        if let Some(pos) = remaining.find(l_del) {
177            if pos > 0 {
178                let literal = &remaining[..pos];
179                current_line += literal.matches('\n').count();
180                tokens.push(Token {
181                    token_type: TokenType::Literal,
182                    key: literal.to_string(),
183                });
184            }
185            remaining = &remaining[pos + l_del.len()..];
186
187            if let Some(end_pos) = remaining.find(r_del) {
188                let tag = &remaining[..end_pos];
189                remaining = &remaining[end_pos + r_del.len()..];
190
191                if tag.is_empty() {
192                    return Err(MustacheError::EmptyTag { line: current_line });
193                }
194
195                let first_char = tag.chars().next().unwrap();
196                let (token_type, key) = match first_char {
197                    '!' => (TokenType::Comment, tag[1..].trim().to_string()),
198                    '#' => {
199                        let key = tag[1..].trim().to_string();
200                        open_sections.push((key.clone(), current_line));
201                        (TokenType::Section, key)
202                    }
203                    '^' => {
204                        let key = tag[1..].trim().to_string();
205                        open_sections.push((key.clone(), current_line));
206                        (TokenType::InvertedSection, key)
207                    }
208                    '/' => {
209                        let key = tag[1..].trim().to_string();
210                        if let Some((expected, _)) = open_sections.pop()
211                            && expected != key
212                        {
213                            return Err(MustacheError::MismatchedSection {
214                                expected,
215                                got: key,
216                                line: current_line,
217                            });
218                        }
219                        (TokenType::End, key)
220                    }
221                    '>' => (TokenType::Partial, tag[1..].trim().to_string()),
222                    '&' => (TokenType::NoEscape, tag[1..].trim().to_string()),
223                    '{' => {
224                        let tag = tag[1..].trim();
225                        let tag = if let Some(stripped) = tag.strip_suffix('}') {
226                            stripped.trim()
227                        } else {
228                            if remaining.starts_with('}') {
229                                remaining = &remaining[1..];
230                            }
231                            tag
232                        };
233                        (TokenType::NoEscape, tag.to_string())
234                    }
235                    _ => (TokenType::Variable, tag.trim().to_string()),
236                };
237
238                if token_type != TokenType::Comment {
239                    tokens.push(Token { token_type, key });
240                }
241            } else {
242                return Err(MustacheError::UnclosedTag { line: current_line });
243            }
244        } else {
245            if !remaining.is_empty() {
246                tokens.push(Token {
247                    token_type: TokenType::Literal,
248                    key: remaining.to_string(),
249                });
250            }
251            break;
252        }
253    }
254
255    if let Some((tag, line)) = open_sections.pop() {
256        return Err(MustacheError::UnclosedSection { tag, line });
257    }
258
259    Ok(tokens)
260}
261
262/// HTML escape a string.
263fn html_escape(s: &str) -> String {
264    let mut result = String::with_capacity(s.len());
265    for c in s.chars() {
266        match c {
267            '&' => result.push_str("&amp;"),
268            '<' => result.push_str("&lt;"),
269            '>' => result.push_str("&gt;"),
270            '"' => result.push_str("&quot;"),
271            _ => result.push(c),
272        }
273    }
274    result
275}
276
277/// Get a key from the scope stack.
278fn get_key(key: &str, scopes: &[&MustacheValue]) -> MustacheValue {
279    if key == "." {
280        return scopes
281            .first()
282            .map(|v| (*v).clone())
283            .unwrap_or(MustacheValue::Null);
284    }
285
286    for scope in scopes {
287        if let MustacheValue::Map(map) = scope {
288            let parts: Vec<&str> = key.split('.').collect();
289            let mut current: Option<&MustacheValue> = map.get(parts[0]);
290
291            for part in parts.iter().skip(1) {
292                if let Some(MustacheValue::Map(m)) = current {
293                    current = m.get(*part);
294                } else {
295                    current = None;
296                    break;
297                }
298            }
299
300            if let Some(value) = current {
301                return value.clone();
302            }
303        }
304    }
305
306    MustacheValue::Null
307}
308
309/// Render a mustache template.
310///
311/// # Arguments
312///
313/// * `template` - The template string.
314/// * `data` - The data to render with.
315/// * `partials` - Optional partial templates.
316///
317/// # Returns
318///
319/// The rendered string, or an error if rendering fails.
320///
321/// # Example
322///
323/// ```
324/// use std::collections::HashMap;
325/// use agent_chain_core::utils::mustache::{render, MustacheValue};
326///
327/// let mut data = HashMap::new();
328/// data.insert("name".to_string(), MustacheValue::String("World".to_string()));
329///
330/// let result = render("Hello, {{name}}!", &MustacheValue::Map(data), None);
331/// assert_eq!(result.unwrap(), "Hello, World!");
332/// ```
333pub fn render(
334    template: &str,
335    data: &MustacheValue,
336    partials: Option<&HashMap<String, String>>,
337) -> Result<String, MustacheError> {
338    render_with_delimiters(template, data, partials, "{{", "}}")
339}
340
341/// Render a mustache template with custom delimiters.
342pub fn render_with_delimiters(
343    template: &str,
344    data: &MustacheValue,
345    partials: Option<&HashMap<String, String>>,
346    l_del: &str,
347    r_del: &str,
348) -> Result<String, MustacheError> {
349    let tokens = tokenize(template, l_del, r_del)?;
350    let scopes = vec![data];
351    render_tokens(&tokens, &scopes, partials, l_del, r_del)
352}
353
354fn render_tokens(
355    tokens: &[Token],
356    scopes: &[&MustacheValue],
357    partials: Option<&HashMap<String, String>>,
358    l_del: &str,
359    r_del: &str,
360) -> Result<String, MustacheError> {
361    let mut output = String::new();
362    let mut i = 0;
363
364    while i < tokens.len() {
365        let token = &tokens[i];
366
367        match token.token_type {
368            TokenType::Literal => {
369                output.push_str(&token.key);
370            }
371            TokenType::Variable => {
372                let value = get_key(&token.key, scopes);
373                output.push_str(&html_escape(&value.to_output_string()));
374            }
375            TokenType::NoEscape => {
376                let value = get_key(&token.key, scopes);
377                output.push_str(&value.to_output_string());
378            }
379            TokenType::Section => {
380                let value = get_key(&token.key, scopes);
381                let end_index = find_section_end(tokens, i, &token.key);
382
383                if value.is_truthy() {
384                    let section_tokens = &tokens[i + 1..end_index];
385                    match &value {
386                        MustacheValue::List(items) => {
387                            for item in items {
388                                let mut new_scopes = vec![item];
389                                new_scopes.extend(scopes.iter());
390                                output.push_str(&render_tokens(
391                                    section_tokens,
392                                    &new_scopes,
393                                    partials,
394                                    l_del,
395                                    r_del,
396                                )?);
397                            }
398                        }
399                        _ => {
400                            let mut new_scopes = vec![&value];
401                            new_scopes.extend(scopes.iter());
402                            output.push_str(&render_tokens(
403                                section_tokens,
404                                &new_scopes,
405                                partials,
406                                l_del,
407                                r_del,
408                            )?);
409                        }
410                    }
411                }
412
413                i = end_index;
414            }
415            TokenType::InvertedSection => {
416                let value = get_key(&token.key, scopes);
417                let end_index = find_section_end(tokens, i, &token.key);
418
419                if !value.is_truthy() {
420                    let section_tokens = &tokens[i + 1..end_index];
421                    output.push_str(&render_tokens(
422                        section_tokens,
423                        scopes,
424                        partials,
425                        l_del,
426                        r_del,
427                    )?);
428                }
429
430                i = end_index;
431            }
432            TokenType::Partial => {
433                if let Some(partials_map) = partials
434                    && let Some(partial_template) = partials_map.get(&token.key)
435                {
436                    output.push_str(&render_with_delimiters(
437                        partial_template,
438                        scopes[0],
439                        partials,
440                        l_del,
441                        r_del,
442                    )?);
443                }
444            }
445            TokenType::End | TokenType::Comment => {}
446        }
447
448        i += 1;
449    }
450
451    Ok(output)
452}
453
454fn find_section_end(tokens: &[Token], start: usize, key: &str) -> usize {
455    let mut depth = 1;
456    for (i, token) in tokens.iter().enumerate().skip(start + 1) {
457        match &token.token_type {
458            TokenType::Section | TokenType::InvertedSection if token.key == key => {
459                depth += 1;
460            }
461            TokenType::End if token.key == key => {
462                depth -= 1;
463                if depth == 0 {
464                    return i;
465                }
466            }
467            _ => {}
468        }
469    }
470    tokens.len()
471}
472
473#[cfg(test)]
474mod tests {
475    use super::*;
476
477    fn make_data(pairs: &[(&str, MustacheValue)]) -> MustacheValue {
478        let mut map = HashMap::new();
479        for (k, v) in pairs {
480            map.insert(k.to_string(), v.clone());
481        }
482        MustacheValue::Map(map)
483    }
484
485    #[test]
486    fn test_simple_variable() {
487        let data = make_data(&[("name", "World".into())]);
488        let result = render("Hello, {{name}}!", &data, None).unwrap();
489        assert_eq!(result, "Hello, World!");
490    }
491
492    #[test]
493    fn test_html_escape() {
494        let data = make_data(&[("html", "<b>bold</b>".into())]);
495        let result = render("{{html}}", &data, None).unwrap();
496        assert_eq!(result, "&lt;b&gt;bold&lt;/b&gt;");
497    }
498
499    #[test]
500    fn test_no_escape() {
501        let data = make_data(&[("html", "<b>bold</b>".into())]);
502        let result = render("{{{html}}}", &data, None).unwrap();
503        assert_eq!(result, "<b>bold</b>");
504    }
505
506    #[test]
507    fn test_section() {
508        let data = make_data(&[("show", true.into())]);
509        let result = render("{{#show}}Shown{{/show}}", &data, None).unwrap();
510        assert_eq!(result, "Shown");
511    }
512
513    #[test]
514    fn test_section_false() {
515        let data = make_data(&[("show", false.into())]);
516        let result = render("{{#show}}Hidden{{/show}}", &data, None).unwrap();
517        assert_eq!(result, "");
518    }
519
520    #[test]
521    fn test_inverted_section() {
522        let data = make_data(&[("show", false.into())]);
523        let result = render("{{^show}}Shown{{/show}}", &data, None).unwrap();
524        assert_eq!(result, "Shown");
525    }
526
527    #[test]
528    fn test_list() {
529        let items = [
530            make_data(&[("name", "Alice".into())]),
531            make_data(&[("name", "Bob".into())]),
532        ];
533        let data = make_data(&[(
534            "items",
535            MustacheValue::List(vec![items[0].clone(), items[1].clone()]),
536        )]);
537        let result = render("{{#items}}{{name}} {{/items}}", &data, None).unwrap();
538        assert_eq!(result, "Alice Bob ");
539    }
540
541    #[test]
542    fn test_dot_notation() {
543        let person = make_data(&[("name", "John".into())]);
544        let data = make_data(&[("person", person)]);
545        let result = render("{{person.name}}", &data, None).unwrap();
546        assert_eq!(result, "John");
547    }
548
549    #[test]
550    fn test_partial() {
551        let data = make_data(&[("name", "World".into())]);
552        let mut partials = HashMap::new();
553        partials.insert("greeting".to_string(), "Hello, {{name}}!".to_string());
554        let result = render("{{>greeting}}", &data, Some(&partials)).unwrap();
555        assert_eq!(result, "Hello, World!");
556    }
557}
agent_chain_core/utils/mustache.rs

agent_chain_core/utils/
mustache.rs
