langchain_rust/prompt/
chat.rs

1use crate::schemas::{messages::Message, prompt::PromptValue};
2
3use super::{
4    FormatPrompter, MessageFormatter, PromptArgs, PromptError, PromptFromatter, PromptTemplate,
5};
6
7/// Struct `HumanMessagePromptTemplate` defines a template for creating human (user) messages.
8/// `PromptTemplate` is used to generate the message template.
9///
10/// # Usage
11/// ```rust,ignore
12/// let human_message_prompt = HumanMessagePromptTemplate::new(template_fstring!(
13///    "User says: {content}",
14///    "content",
15/// ));
16/// ```
17#[derive(Clone)]
18pub struct HumanMessagePromptTemplate {
19    prompt: PromptTemplate,
20}
21
22impl HumanMessagePromptTemplate {
23    pub fn new(prompt: PromptTemplate) -> Self {
24        Self { prompt }
25    }
26}
27impl MessageFormatter for HumanMessagePromptTemplate {
28    fn format_messages(&self, input_variables: PromptArgs) -> Result<Vec<Message>, PromptError> {
29        let message = Message::new_human_message(self.prompt.format(input_variables)?);
30        log::debug!("message: {:?}", message);
31        Ok(vec![message])
32    }
33    fn input_variables(&self) -> Vec<String> {
34        self.prompt.variables().clone()
35    }
36}
37
38impl FormatPrompter for HumanMessagePromptTemplate {
39    fn format_prompt(&self, input_variables: PromptArgs) -> Result<PromptValue, PromptError> {
40        let messages = self.format_messages(input_variables)?;
41        Ok(PromptValue::from_messages(messages))
42    }
43    fn get_input_variables(&self) -> Vec<String> {
44        self.input_variables()
45    }
46}
47
48/// Struct `SystemMessagePromptTemplate` defines a template for creating system-level messages.
49/// `PromptTemplate` is used to generate the message template.
50///
51/// # Usage
52/// ```rust,ignore
53/// let system_message_prompt = SystemMessagePromptTemplate::new(template_fstring!(
54///    "System alert: {alert_type} {alert_detail}",
55///    "alert_type",
56///    "alert_detail"
57/// ));
58/// ```
59#[derive(Clone)]
60pub struct SystemMessagePromptTemplate {
61    prompt: PromptTemplate,
62}
63
64impl SystemMessagePromptTemplate {
65    pub fn new(prompt: PromptTemplate) -> Self {
66        Self { prompt }
67    }
68}
69
70impl FormatPrompter for SystemMessagePromptTemplate {
71    fn format_prompt(&self, input_variables: PromptArgs) -> Result<PromptValue, PromptError> {
72        let messages = self.format_messages(input_variables)?;
73        Ok(PromptValue::from_messages(messages))
74    }
75    fn get_input_variables(&self) -> Vec<String> {
76        self.input_variables()
77    }
78}
79
80impl MessageFormatter for SystemMessagePromptTemplate {
81    fn format_messages(&self, input_variables: PromptArgs) -> Result<Vec<Message>, PromptError> {
82        let message = Message::new_system_message(self.prompt.format(input_variables)?);
83        log::debug!("message: {:?}", message);
84        Ok(vec![message])
85    }
86    fn input_variables(&self) -> Vec<String> {
87        self.prompt.variables().clone()
88    }
89}
90
91/// Struct `AIMessagePromptTemplate` defines a template for creating AI (assistant) messages.
92/// `PromptTemplate` is used to generate the message template.
93///
94/// # Usage
95/// ```rust,ignore
96/// let ai_message_prompt = AIMessagePromptTemplate::new(template_fstring!(
97///    "AI response: {content} {additional_info}",
98///    "content",
99///    "additional_info"
100/// ));
101#[derive(Clone)]
102pub struct AIMessagePromptTemplate {
103    prompt: PromptTemplate,
104}
105
106impl FormatPrompter for AIMessagePromptTemplate {
107    fn format_prompt(&self, input_variables: PromptArgs) -> Result<PromptValue, PromptError> {
108        let messages = self.format_messages(input_variables)?;
109        Ok(PromptValue::from_messages(messages))
110    }
111    fn get_input_variables(&self) -> Vec<String> {
112        self.input_variables()
113    }
114}
115
116impl MessageFormatter for AIMessagePromptTemplate {
117    fn format_messages(&self, input_variables: PromptArgs) -> Result<Vec<Message>, PromptError> {
118        let message = Message::new_ai_message(self.prompt.format(input_variables)?);
119        log::debug!("message: {:?}", message);
120        Ok(vec![message])
121    }
122    fn input_variables(&self) -> Vec<String> {
123        self.prompt.variables().clone()
124    }
125}
126
127impl AIMessagePromptTemplate {
128    pub fn new(prompt: PromptTemplate) -> Self {
129        Self { prompt }
130    }
131}
132
133pub enum MessageOrTemplate {
134    Message(Message),
135    Template(Box<dyn MessageFormatter>),
136    MessagesPlaceholder(String),
137}
138
139/// `fmt_message` is a utility macro used to create a `MessageOrTemplate::Message` variant.
140///
141/// # Usage
142/// The macro is called with a `Message` object. For example:
143/// ```rust,ignore
144/// let message = Message::new_human_message("Hello World");
145/// fmt_message!(message) // Returns a `MessageOrTemplate::Message` variant that wraps the `Message` object
146/// ```
147#[macro_export]
148macro_rules! fmt_message {
149    ($msg:expr) => {
150        $crate::prompt::MessageOrTemplate::Message($msg)
151    };
152}
153
154/// `fmt_template` is a utility macro used to create a `MessageOrTemplate::Template` variant.
155///
156/// # Usage
157/// The macro is called with a `MessageFormatter` object, for instance `HumanMessagePromptTemplate`,
158/// `SystemMessagePromptTemplate`, `AIMessagePromptTemplate` or any other implementation of `MessageFormatter`.
159///
160/// ```rust,ignore
161/// let prompt_template = HumanMessagePromptTemplate::new(template);
162/// fmt_template!(prompt_template)
163/// ```
164/// This returns a `MessageOrTemplate::Template` variant that wraps the `MessageFormatter` object within a Box.
165#[macro_export]
166macro_rules! fmt_template {
167    ($template:expr) => {
168        $crate::prompt::MessageOrTemplate::Template(Box::new($template))
169    };
170}
171
172/// `fmt_placeholder` is a utility macro used to create a `MessageOrTemplate::MessagesPlaceholder` variant.
173///
174/// # Usage
175/// The macro is called with a string literal or a String object:
176/// ```rust,ignore
177/// fmt_placeholder!("Placeholder message")
178/// ```
179/// This returns a `MessageOrTemplate::MessagesPlaceholder` variant that wraps the given string.
180#[macro_export]
181macro_rules! fmt_placeholder {
182    ($placeholder:expr) => {
183        $crate::prompt::MessageOrTemplate::MessagesPlaceholder($placeholder.into())
184    };
185}
186
187pub struct MessageFormatterStruct {
188    items: Vec<MessageOrTemplate>,
189}
190
191impl MessageFormatterStruct {
192    pub fn new() -> Self {
193        Self { items: Vec::new() }
194    }
195
196    pub fn add_message(&mut self, message: Message) {
197        self.items.push(MessageOrTemplate::Message(message));
198    }
199
200    pub fn add_template(&mut self, template: Box<dyn MessageFormatter>) {
201        self.items.push(MessageOrTemplate::Template(template));
202    }
203
204    pub fn add_messages_placeholder(&mut self, placeholder: &str) {
205        self.items.push(MessageOrTemplate::MessagesPlaceholder(
206            placeholder.to_string(),
207        ));
208    }
209
210    fn format(&self, input_variables: PromptArgs) -> Result<Vec<Message>, PromptError> {
211        let mut result: Vec<Message> = Vec::new();
212        for item in &self.items {
213            match item {
214                MessageOrTemplate::Message(msg) => result.push(msg.clone()),
215                MessageOrTemplate::Template(tmpl) => {
216                    result.extend(tmpl.format_messages(input_variables.clone())?)
217                }
218                MessageOrTemplate::MessagesPlaceholder(placeholder) => {
219                    let messages = input_variables
220                        .get(placeholder)
221                        .ok_or_else(|| PromptError::MissingVariable(placeholder.clone()))?;
222                    result.extend(Message::messages_from_value(messages)?);
223                }
224            }
225        }
226        Ok(result)
227    }
228}
229
230impl MessageFormatter for MessageFormatterStruct {
231    fn format_messages(&self, input_variables: PromptArgs) -> Result<Vec<Message>, PromptError> {
232        self.format(input_variables)
233    }
234    fn input_variables(&self) -> Vec<String> {
235        let mut variables = Vec::new();
236        for item in &self.items {
237            match item {
238                MessageOrTemplate::Message(_) => {}
239                MessageOrTemplate::Template(tmpl) => {
240                    variables.extend(tmpl.input_variables());
241                }
242                MessageOrTemplate::MessagesPlaceholder(placeholder) => {
243                    variables.extend(vec![placeholder.clone()]);
244                }
245            }
246        }
247        variables
248    }
249}
250
251impl FormatPrompter for MessageFormatterStruct {
252    fn format_prompt(&self, input_variables: PromptArgs) -> Result<PromptValue, PromptError> {
253        let messages = self.format(input_variables)?;
254        Ok(PromptValue::from_messages(messages))
255    }
256    fn get_input_variables(&self) -> Vec<String> {
257        self.input_variables()
258    }
259}
260
261#[macro_export]
262// A macro for creating a new MessageFormatterStruct with various types of messages.
263///
264///# Example
265/// ```rust,ignore
266/// // Create an AI message prompt template
267/// let ai_message_prompt = AIMessagePromptTemplate::new(
268/// template_fstring!(
269///     "AI response: {content} {test}",
270///     "content",
271///     "test"
272/// ));
273///
274///
275/// let human_msg = Message::new_human_message("Hello from user");
276///
277/// // Use the `message_formatter` macro to construct the formatter.
278/// let formatter = message_formatter![
279///     fmt_message!(human_msg),
280///     fmt_template!(ai_message_prompt),
281///     fmt_placeholder!("history")
282/// ];
283/// ```
284macro_rules! message_formatter {
285($($item:expr),* $(,)?) => {{
286    let mut formatter = $crate::prompt::MessageFormatterStruct::new();
287    $(
288        match $item {
289            $crate::prompt::MessageOrTemplate::Message(msg) => formatter.add_message(msg),
290            $crate::prompt::MessageOrTemplate::Template(tmpl) => formatter.add_template(tmpl),
291            $crate::prompt::MessageOrTemplate::MessagesPlaceholder(placeholder) => formatter.add_messages_placeholder(&placeholder.clone()),
292        }
293    )*
294    formatter
295}};
296}
297
298#[cfg(test)]
299mod tests {
300    use crate::{
301        message_formatter,
302        prompt::{chat::AIMessagePromptTemplate, FormatPrompter},
303        prompt_args,
304        schemas::messages::Message,
305        template_fstring,
306    };
307
308    #[test]
309    fn test_message_formatter_macro() {
310        // Create a human message and system message
311        let human_msg = Message::new_human_message("Hello from user");
312
313        // Create an AI message prompt template
314        let ai_message_prompt = AIMessagePromptTemplate::new(template_fstring!(
315            "AI response: {content} {test}",
316            "content",
317            "test"
318        ));
319
320        // Use the `message_formatter` macro to construct the formatter
321        let formatter = message_formatter![
322            fmt_message!(human_msg),
323            fmt_template!(ai_message_prompt),
324            fmt_placeholder!("history")
325        ];
326
327        // Define input variables for the AI message template
328        let input_variables = prompt_args! {
329            "content" => "This is a test",
330            "test" => "test2",
331            "history" => vec![
332                Message::new_human_message("Placeholder message 1"),
333                Message::new_ai_message("Placeholder message 2"),
334            ],
335
336
337        };
338
339        // Format messages
340        let formatted_messages = formatter
341            .format_prompt(input_variables)
342            .unwrap()
343            .to_chat_messages();
344
345        // Verify the number of messages
346        assert_eq!(formatted_messages.len(), 4);
347
348        // Verify the content of each message
349        assert_eq!(formatted_messages[0].content, "Hello from user");
350        assert_eq!(
351            formatted_messages[1].content,
352            "AI response: This is a test test2"
353        );
354        assert_eq!(formatted_messages[2].content, "Placeholder message 1");
355        assert_eq!(formatted_messages[3].content, "Placeholder message 2");
356    }
357
358    #[test]
359    fn message_templates_roles_and_variables() {
360        use crate::{
361            prompt::{
362                HumanMessagePromptTemplate, MessageFormatter, PromptError,
363                SystemMessagePromptTemplate,
364            },
365            schemas::MessageType,
366        };
367        let formatter = message_formatter![
368            fmt_template!(SystemMessagePromptTemplate::new(template_fstring!(
369                "You are {role}",
370                "role"
371            ))),
372            fmt_placeholder!("history"),
373            fmt_template!(HumanMessagePromptTemplate::new(template_fstring!(
374                "{q}", "q"
375            ))),
376        ];
377        assert_eq!(
378            formatter.input_variables(),
379            vec!["role".to_string(), "history".into(), "q".into()]
380        );
381
382        let messages = formatter
383            .format_messages(prompt_args! {
384                "role" => "a bot", "q" => "hi", "history" => Vec::<Message>::new()
385            })
386            .unwrap();
387        assert_eq!(messages.len(), 2);
388        assert_eq!(messages[0].message_type, MessageType::SystemMessage);
389        assert_eq!(messages[0].content, "You are a bot");
390        assert_eq!(messages[1].message_type, MessageType::HumanMessage);
391
392        // missing placeholder is an error, not a panic
393        let err = formatter
394            .format_messages(prompt_args! {"role" => "x", "q" => "y"})
395            .unwrap_err();
396        assert!(matches!(err, PromptError::MissingVariable(v) if v == "history"));
397
398        // missing template variable
399        let err = formatter
400            .format_messages(prompt_args! {"history" => Vec::<Message>::new(), "q" => "y"})
401            .unwrap_err();
402        assert!(matches!(err, PromptError::MissingVariable(v) if v == "role"));
403
404        // placeholder with wrong type
405        let err = formatter
406            .format_messages(prompt_args! {"role" => "x", "q" => "y", "history" => "nope"})
407            .unwrap_err();
408        assert!(matches!(err, PromptError::SerializationError(_)));
409    }
410}
langchain_rust/prompt/chat.rs

langchain_rust/prompt/
chat.rs
