Skip to main content

synaptic_prompts/
few_shot_template.rs

1use std::collections::HashMap;
2
3use async_trait::async_trait;
4use synaptic_core::{RunnableConfig, SynapticError};
5use synaptic_runnables::Runnable;
6
7use crate::{FewShotExample, PromptTemplate};
8
9/// A string-based few-shot prompt template (as opposed to `FewShotChatMessagePromptTemplate`
10/// which produces `Vec<Message>`).
11///
12/// Produces a single formatted string with examples embedded.
13pub struct FewShotPromptTemplate {
14    examples: Vec<FewShotExample>,
15    example_prompt: PromptTemplate,
16    prefix: Option<String>,
17    suffix: PromptTemplate,
18    example_separator: String,
19}
20
21impl FewShotPromptTemplate {
22    /// Create a new string-based few-shot template.
23    ///
24    /// - `examples`: the input/output example pairs
25    /// - `example_prompt`: template for each example, e.g. `"Input: {{ input }}\nOutput: {{ output }}"`
26    /// - `suffix`: final template rendered with user-provided variables
27    pub fn new(
28        examples: Vec<FewShotExample>,
29        example_prompt: PromptTemplate,
30        suffix: PromptTemplate,
31    ) -> Self {
32        Self {
33            examples,
34            example_prompt,
35            prefix: None,
36            suffix,
37            example_separator: "\n\n".to_string(),
38        }
39    }
40
41    /// Set an optional prefix string prepended before the examples.
42    pub fn with_prefix(mut self, prefix: impl Into<String>) -> Self {
43        self.prefix = Some(prefix.into());
44        self
45    }
46
47    /// Set the separator used between examples (default: `"\n\n"`).
48    pub fn with_separator(mut self, sep: impl Into<String>) -> Self {
49        self.example_separator = sep.into();
50        self
51    }
52
53    /// Render the template with the given variable values.
54    ///
55    /// 1. Start with prefix (if any)
56    /// 2. For each example, render `example_prompt` with `{"input": ..., "output": ...}`
57    /// 3. Join examples with separator
58    /// 4. Append suffix rendered with provided values
59    /// 5. Join all parts with separator
60    pub fn render(&self, values: &HashMap<String, String>) -> Result<String, SynapticError> {
61        let mut parts: Vec<String> = Vec::new();
62
63        // 1. Prefix
64        if let Some(prefix) = &self.prefix {
65            parts.push(prefix.clone());
66        }
67
68        // 2-3. Render and join examples
69        let mut example_strings = Vec::with_capacity(self.examples.len());
70        for example in &self.examples {
71            let example_values = HashMap::from([
72                ("input".to_string(), example.input.clone()),
73                ("output".to_string(), example.output.clone()),
74            ]);
75            let rendered = self
76                .example_prompt
77                .render(&example_values)
78                .map_err(|e| SynapticError::Prompt(e.to_string()))?;
79            example_strings.push(rendered);
80        }
81
82        if !example_strings.is_empty() {
83            parts.push(example_strings.join(&self.example_separator));
84        }
85
86        // 4. Suffix
87        let suffix_rendered = self
88            .suffix
89            .render(values)
90            .map_err(|e| SynapticError::Prompt(e.to_string()))?;
91        parts.push(suffix_rendered);
92
93        // 5. Join all parts
94        Ok(parts.join(&self.example_separator))
95    }
96}
97
98#[async_trait]
99impl Runnable<HashMap<String, String>, String> for FewShotPromptTemplate {
100    async fn invoke(
101        &self,
102        input: HashMap<String, String>,
103        _config: &RunnableConfig,
104    ) -> Result<String, SynapticError> {
105        self.render(&input)
106    }
107}