Skip to main content

synaptic_prompts/
few_shot.rs

1use std::collections::HashMap;
2
3use async_trait::async_trait;
4use synaptic_core::{Message, RunnableConfig, SynapticError};
5use synaptic_runnables::Runnable;
6
7use crate::PromptTemplate;
8
9/// An example for few-shot prompting, consisting of input/output pairs.
10#[derive(Debug, Clone)]
11pub struct FewShotExample {
12    pub input: String,
13    pub output: String,
14}
15
16/// A few-shot chat message prompt template that injects examples before the user query.
17///
18/// Each example is formatted as a Human message (input) followed by an AI message (output).
19/// An optional prefix system message can be provided.
20pub struct FewShotChatMessagePromptTemplate {
21    examples: Vec<FewShotExample>,
22    prefix: Option<PromptTemplate>,
23    suffix: PromptTemplate,
24}
25
26impl FewShotChatMessagePromptTemplate {
27    pub fn new(examples: Vec<FewShotExample>, suffix: PromptTemplate) -> Self {
28        Self {
29            examples,
30            prefix: None,
31            suffix,
32        }
33    }
34
35    pub fn with_prefix(mut self, prefix: PromptTemplate) -> Self {
36        self.prefix = Some(prefix);
37        self
38    }
39
40    pub fn format(&self, values: &HashMap<String, String>) -> Result<Vec<Message>, SynapticError> {
41        let mut messages = Vec::new();
42
43        if let Some(prefix) = &self.prefix {
44            let content = prefix
45                .render(values)
46                .map_err(|e| SynapticError::Prompt(e.to_string()))?;
47            messages.push(Message::system(content));
48        }
49
50        for example in &self.examples {
51            messages.push(Message::human(&example.input));
52            messages.push(Message::ai(&example.output));
53        }
54
55        let content = self
56            .suffix
57            .render(values)
58            .map_err(|e| SynapticError::Prompt(e.to_string()))?;
59        messages.push(Message::human(content));
60
61        Ok(messages)
62    }
63}
64
65#[async_trait]
66impl Runnable<HashMap<String, String>, Vec<Message>> for FewShotChatMessagePromptTemplate {
67    async fn invoke(
68        &self,
69        input: HashMap<String, String>,
70        _config: &RunnableConfig,
71    ) -> Result<Vec<Message>, SynapticError> {
72        self.format(&input)
73    }
74}