synaptic_prompts/
few_shot.rs1use 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#[derive(Debug, Clone)]
11pub struct FewShotExample {
12 pub input: String,
13 pub output: String,
14}
15
16pub 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}