daimon 0.16.0

A Rust-native AI agent framework
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132

//! Few-shot example templates for prompt construction.

/// A collection of input/output example pairs injected into a prompt to guide
/// the model's behaviour through demonstration.
///
/// ```ignore
/// use daimon::prompt::FewShotTemplate;
///
/// let tpl = FewShotTemplate::new()
///     .example("What is 2+2?", "4")
///     .example("What is the capital of France?", "Paris")
///     .with_prefix("Here are some examples of how to respond:");
///
/// let rendered = tpl.render();
/// ```
#[derive(Debug, Clone, Default)]
pub struct FewShotTemplate {
    examples: Vec<(String, String)>,
    prefix: Option<String>,
    input_label: String,
    output_label: String,
}

impl FewShotTemplate {
    pub fn new() -> Self {
        Self {
            examples: Vec::new(),
            prefix: None,
            input_label: "Input".to_string(),
            output_label: "Output".to_string(),
        }
    }

    /// Adds an example input/output pair.
    pub fn example(mut self, input: impl Into<String>, output: impl Into<String>) -> Self {
        self.examples.push((input.into(), output.into()));
        self
    }

    /// Sets text rendered before the examples.
    pub fn with_prefix(mut self, prefix: impl Into<String>) -> Self {
        self.prefix = Some(prefix.into());
        self
    }

    /// Customises the label used for inputs (default: `"Input"`).
    pub fn with_input_label(mut self, label: impl Into<String>) -> Self {
        self.input_label = label.into();
        self
    }

    /// Customises the label used for outputs (default: `"Output"`).
    pub fn with_output_label(mut self, label: impl Into<String>) -> Self {
        self.output_label = label.into();
        self
    }

    /// Renders the examples into a formatted string.
    pub fn render(&self) -> String {
        let mut parts = Vec::new();

        if let Some(prefix) = &self.prefix {
            parts.push(prefix.clone());
            parts.push(String::new());
        }

        for (i, (input, output)) in self.examples.iter().enumerate() {
            if i > 0 {
                parts.push(String::new());
            }
            parts.push(format!("{}: {input}", self.input_label));
            parts.push(format!("{}: {output}", self.output_label));
        }

        parts.join("\n")
    }

    /// Returns the number of examples.
    pub fn len(&self) -> usize {
        self.examples.len()
    }

    /// Returns true if no examples have been added.
    pub fn is_empty(&self) -> bool {
        self.examples.is_empty()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_basic_few_shot() {
        let tpl = FewShotTemplate::new()
            .example("What is 2+2?", "4")
            .example("Capital of France?", "Paris");
        let rendered = tpl.render();
        assert!(rendered.contains("Input: What is 2+2?"));
        assert!(rendered.contains("Output: 4"));
        assert!(rendered.contains("Input: Capital of France?"));
        assert!(rendered.contains("Output: Paris"));
    }

    #[test]
    fn test_with_prefix() {
        let tpl = FewShotTemplate::new()
            .with_prefix("Examples:")
            .example("hi", "hello");
        let rendered = tpl.render();
        assert!(rendered.starts_with("Examples:"));
    }

    #[test]
    fn test_custom_labels() {
        let tpl = FewShotTemplate::new()
            .with_input_label("Q")
            .with_output_label("A")
            .example("question", "answer");
        let rendered = tpl.render();
        assert!(rendered.contains("Q: question"));
        assert!(rendered.contains("A: answer"));
    }

    #[test]
    fn test_empty() {
        let tpl = FewShotTemplate::new();
        assert!(tpl.is_empty());
        assert_eq!(tpl.len(), 0);
        assert_eq!(tpl.render(), "");
    }
}