langchain_rust/prompt/
chat.rs

1use serde_json::Value;
2
3use crate::schemas::{
4    messages::{ImageContent, Message},
5    prompt::PromptValue,
6};
7
8use super::{
9    FormatPrompter, MessageFormatter, PromptArgs, PromptError, PromptFromatter, PromptTemplate,
10};
11
12/// Struct `HumanMessagePromptTemplate` defines a template for creating human (user) messages.
13/// `PromptTemplate` is used to generate the message template.
14///
15/// # Usage
16/// ```rust,ignore
17/// let human_message_prompt = HumanMessagePromptTemplate::new(template_fstring!(
18///    "User says: {content}",
19///    "content",
20/// ));
21/// ```
22///
23/// # Images
24///
25/// [`with_images`](Self::with_images) names an input variable that holds images to
26/// attach to the message, so a chain or agent can take them like any other input:
27///
28/// ```rust,ignore
29/// let prompt = HumanMessagePromptTemplate::new(template_fstring!("{input}", "input"))
30///     .with_images("images");
31/// chain.invoke(prompt_args! {
32///     "input" => "What is in this picture?",
33///     "images" => ["https://example.com/cat.png"],
34/// }).await?;
35/// ```
36#[derive(Clone)]
37pub struct HumanMessagePromptTemplate {
38    prompt: PromptTemplate,
39    images_key: Option<String>,
40}
41
42impl HumanMessagePromptTemplate {
43    pub fn new(prompt: PromptTemplate) -> Self {
44        Self {
45            prompt,
46            images_key: None,
47        }
48    }
49
50    /// Attaches the images found in the `key` input variable to the message.
51    ///
52    /// The variable can be one image or a list of them. An image is either a string
53    /// (an `http(s)` URL, a `data:` URL or bare base64 data) or an
54    /// [`ImageContent`] object (`{"image_url": "...", "detail": "low"}`).
55    /// A missing variable, `null` or an empty list means no images.
56    pub fn with_images<S: Into<String>>(mut self, key: S) -> Self {
57        self.images_key = Some(key.into());
58        self
59    }
60}
61
62/// Reads the images held in an input variable (see
63/// [`HumanMessagePromptTemplate::with_images`]).
64pub(crate) fn images_from_value(value: &Value) -> Result<Vec<ImageContent>, PromptError> {
65    fn one(value: &Value) -> Result<ImageContent, PromptError> {
66        match value {
67            Value::String(url) if !url.trim().is_empty() => Ok(ImageContent::from(url)),
68            Value::Object(_) => Ok(serde_json::from_value(value.clone())?),
69            other => Err(PromptError::OtherError(format!(
70                "expected an image URL, base64 string or {{\"image_url\": ...}} object, got {other}"
71            ))),
72        }
73    }
74    match value {
75        Value::Null => Ok(Vec::new()),
76        Value::Array(items) => items.iter().map(one).collect(),
77        other => Ok(vec![one(other)?]),
78    }
79}
80
81/// Input variable the built-in chain and agent prompts read images from.
82pub const DEFAULT_IMAGES_KEY: &str = "images";
83
84/// The human turn of a chain or agent call, as it is saved to memory: the text of
85/// `input_key` plus any images in `images_key`.
86pub(crate) fn human_message_from_inputs(
87    inputs: &PromptArgs,
88    input_key: &str,
89    images_key: &str,
90) -> Result<Message, PromptError> {
91    let input = inputs
92        .get(input_key)
93        .ok_or_else(|| PromptError::MissingVariable(input_key.to_string()))?;
94    let message = Message::new_human_message(crate::schemas::messages::value_to_text(input));
95    let images = match inputs.get(images_key) {
96        Some(value) => images_from_value(value)?,
97        None => Vec::new(),
98    };
99    Ok(if images.is_empty() {
100        message
101    } else {
102        message.with_images(images)
103    })
104}
105
106impl MessageFormatter for HumanMessagePromptTemplate {
107    fn format_messages(&self, input_variables: PromptArgs) -> Result<Vec<Message>, PromptError> {
108        let images = match &self.images_key {
109            Some(key) => match input_variables.get(key) {
110                Some(value) => images_from_value(value)?,
111                None => Vec::new(),
112            },
113            None => Vec::new(),
114        };
115        let mut message = Message::new_human_message(self.prompt.format(input_variables)?);
116        if !images.is_empty() {
117            message = message.with_images(images);
118        }
119        tracing::debug!("message: {:?}", message);
120        Ok(vec![message])
121    }
122    fn input_variables(&self) -> Vec<String> {
123        self.prompt.variables().clone()
124    }
125}
126
127impl FormatPrompter for HumanMessagePromptTemplate {
128    fn format_prompt(&self, input_variables: PromptArgs) -> Result<PromptValue, PromptError> {
129        let messages = self.format_messages(input_variables)?;
130        Ok(PromptValue::from_messages(messages))
131    }
132    fn get_input_variables(&self) -> Vec<String> {
133        self.input_variables()
134    }
135}
136
137/// Struct `SystemMessagePromptTemplate` defines a template for creating system-level messages.
138/// `PromptTemplate` is used to generate the message template.
139///
140/// # Usage
141/// ```rust,ignore
142/// let system_message_prompt = SystemMessagePromptTemplate::new(template_fstring!(
143///    "System alert: {alert_type} {alert_detail}",
144///    "alert_type",
145///    "alert_detail"
146/// ));
147/// ```
148#[derive(Clone)]
149pub struct SystemMessagePromptTemplate {
150    prompt: PromptTemplate,
151}
152
153impl SystemMessagePromptTemplate {
154    pub fn new(prompt: PromptTemplate) -> Self {
155        Self { prompt }
156    }
157}
158
159impl FormatPrompter for SystemMessagePromptTemplate {
160    fn format_prompt(&self, input_variables: PromptArgs) -> Result<PromptValue, PromptError> {
161        let messages = self.format_messages(input_variables)?;
162        Ok(PromptValue::from_messages(messages))
163    }
164    fn get_input_variables(&self) -> Vec<String> {
165        self.input_variables()
166    }
167}
168
169impl MessageFormatter for SystemMessagePromptTemplate {
170    fn format_messages(&self, input_variables: PromptArgs) -> Result<Vec<Message>, PromptError> {
171        let message = Message::new_system_message(self.prompt.format(input_variables)?);
172        tracing::debug!("message: {:?}", message);
173        Ok(vec![message])
174    }
175    fn input_variables(&self) -> Vec<String> {
176        self.prompt.variables().clone()
177    }
178}
179
180/// Struct `AIMessagePromptTemplate` defines a template for creating AI (assistant) messages.
181/// `PromptTemplate` is used to generate the message template.
182///
183/// # Usage
184/// ```rust,ignore
185/// let ai_message_prompt = AIMessagePromptTemplate::new(template_fstring!(
186///    "AI response: {content} {additional_info}",
187///    "content",
188///    "additional_info"
189/// ));
190#[derive(Clone)]
191pub struct AIMessagePromptTemplate {
192    prompt: PromptTemplate,
193}
194
195impl FormatPrompter for AIMessagePromptTemplate {
196    fn format_prompt(&self, input_variables: PromptArgs) -> Result<PromptValue, PromptError> {
197        let messages = self.format_messages(input_variables)?;
198        Ok(PromptValue::from_messages(messages))
199    }
200    fn get_input_variables(&self) -> Vec<String> {
201        self.input_variables()
202    }
203}
204
205impl MessageFormatter for AIMessagePromptTemplate {
206    fn format_messages(&self, input_variables: PromptArgs) -> Result<Vec<Message>, PromptError> {
207        let message = Message::new_ai_message(self.prompt.format(input_variables)?);
208        tracing::debug!("message: {:?}", message);
209        Ok(vec![message])
210    }
211    fn input_variables(&self) -> Vec<String> {
212        self.prompt.variables().clone()
213    }
214}
215
216impl AIMessagePromptTemplate {
217    pub fn new(prompt: PromptTemplate) -> Self {
218        Self { prompt }
219    }
220}
221
222pub enum MessageOrTemplate {
223    Message(Message),
224    Template(Box<dyn MessageFormatter>),
225    MessagesPlaceholder(String),
226}
227
228/// `fmt_message` is a utility macro used to create a `MessageOrTemplate::Message` variant.
229///
230/// # Usage
231/// The macro is called with a `Message` object. For example:
232/// ```rust,ignore
233/// let message = Message::new_human_message("Hello World");
234/// fmt_message!(message) // Returns a `MessageOrTemplate::Message` variant that wraps the `Message` object
235/// ```
236#[macro_export]
237macro_rules! fmt_message {
238    ($msg:expr) => {
239        $crate::prompt::MessageOrTemplate::Message($msg)
240    };
241}
242
243/// `fmt_template` is a utility macro used to create a `MessageOrTemplate::Template` variant.
244///
245/// # Usage
246/// The macro is called with a `MessageFormatter` object, for instance `HumanMessagePromptTemplate`,
247/// `SystemMessagePromptTemplate`, `AIMessagePromptTemplate` or any other implementation of `MessageFormatter`.
248///
249/// ```rust,ignore
250/// let prompt_template = HumanMessagePromptTemplate::new(template);
251/// fmt_template!(prompt_template)
252/// ```
253/// This returns a `MessageOrTemplate::Template` variant that wraps the `MessageFormatter` object within a Box.
254#[macro_export]
255macro_rules! fmt_template {
256    ($template:expr) => {
257        $crate::prompt::MessageOrTemplate::Template(Box::new($template))
258    };
259}
260
261/// `fmt_placeholder` is a utility macro used to create a `MessageOrTemplate::MessagesPlaceholder` variant.
262///
263/// # Usage
264/// The macro is called with a string literal or a String object:
265/// ```rust,ignore
266/// fmt_placeholder!("Placeholder message")
267/// ```
268/// This returns a `MessageOrTemplate::MessagesPlaceholder` variant that wraps the given string.
269#[macro_export]
270macro_rules! fmt_placeholder {
271    ($placeholder:expr) => {
272        $crate::prompt::MessageOrTemplate::MessagesPlaceholder($placeholder.into())
273    };
274}
275
276pub struct MessageFormatterStruct {
277    items: Vec<MessageOrTemplate>,
278}
279
280impl MessageFormatterStruct {
281    pub fn new() -> Self {
282        Self { items: Vec::new() }
283    }
284
285    pub fn add_message(&mut self, message: Message) {
286        self.items.push(MessageOrTemplate::Message(message));
287    }
288
289    pub fn add_template(&mut self, template: Box<dyn MessageFormatter>) {
290        self.items.push(MessageOrTemplate::Template(template));
291    }
292
293    pub fn add_messages_placeholder(&mut self, placeholder: &str) {
294        self.items.push(MessageOrTemplate::MessagesPlaceholder(
295            placeholder.to_string(),
296        ));
297    }
298
299    fn format(&self, input_variables: PromptArgs) -> Result<Vec<Message>, PromptError> {
300        let mut result: Vec<Message> = Vec::new();
301        for item in &self.items {
302            match item {
303                MessageOrTemplate::Message(msg) => result.push(msg.clone()),
304                MessageOrTemplate::Template(tmpl) => {
305                    result.extend(tmpl.format_messages(input_variables.clone())?)
306                }
307                MessageOrTemplate::MessagesPlaceholder(placeholder) => {
308                    let messages = input_variables
309                        .get(placeholder)
310                        .ok_or_else(|| PromptError::MissingVariable(placeholder.clone()))?;
311                    result.extend(Message::messages_from_value(messages)?);
312                }
313            }
314        }
315        Ok(result)
316    }
317}
318
319impl MessageFormatter for MessageFormatterStruct {
320    fn format_messages(&self, input_variables: PromptArgs) -> Result<Vec<Message>, PromptError> {
321        self.format(input_variables)
322    }
323    fn input_variables(&self) -> Vec<String> {
324        let mut variables = Vec::new();
325        for item in &self.items {
326            match item {
327                MessageOrTemplate::Message(_) => {}
328                MessageOrTemplate::Template(tmpl) => {
329                    variables.extend(tmpl.input_variables());
330                }
331                MessageOrTemplate::MessagesPlaceholder(placeholder) => {
332                    variables.extend(vec![placeholder.clone()]);
333                }
334            }
335        }
336        variables
337    }
338}
339
340impl FormatPrompter for MessageFormatterStruct {
341    fn format_prompt(&self, input_variables: PromptArgs) -> Result<PromptValue, PromptError> {
342        let messages = self.format(input_variables)?;
343        Ok(PromptValue::from_messages(messages))
344    }
345    fn get_input_variables(&self) -> Vec<String> {
346        self.input_variables()
347    }
348}
349
350#[macro_export]
351// A macro for creating a new MessageFormatterStruct with various types of messages.
352///
353///# Example
354/// ```rust,ignore
355/// // Create an AI message prompt template
356/// let ai_message_prompt = AIMessagePromptTemplate::new(
357/// template_fstring!(
358///     "AI response: {content} {test}",
359///     "content",
360///     "test"
361/// ));
362///
363///
364/// let human_msg = Message::new_human_message("Hello from user");
365///
366/// // Use the `message_formatter` macro to construct the formatter.
367/// let formatter = message_formatter![
368///     fmt_message!(human_msg),
369///     fmt_template!(ai_message_prompt),
370///     fmt_placeholder!("history")
371/// ];
372/// ```
373macro_rules! message_formatter {
374($($item:expr),* $(,)?) => {{
375    let mut formatter = $crate::prompt::MessageFormatterStruct::new();
376    $(
377        match $item {
378            $crate::prompt::MessageOrTemplate::Message(msg) => formatter.add_message(msg),
379            $crate::prompt::MessageOrTemplate::Template(tmpl) => formatter.add_template(tmpl),
380            $crate::prompt::MessageOrTemplate::MessagesPlaceholder(placeholder) => formatter.add_messages_placeholder(&placeholder.clone()),
381        }
382    )*
383    formatter
384}};
385}
386
387#[cfg(test)]
388mod tests {
389    use crate::{
390        message_formatter,
391        prompt::{chat::AIMessagePromptTemplate, FormatPrompter, PromptError},
392        prompt_args,
393        schemas::messages::{Message, MessageType},
394        template_fstring,
395    };
396
397    #[test]
398    fn test_message_formatter_macro() {
399        // Create a human message and system message
400        let human_msg = Message::new_human_message("Hello from user");
401
402        // Create an AI message prompt template
403        let ai_message_prompt = AIMessagePromptTemplate::new(template_fstring!(
404            "AI response: {content} {test}",
405            "content",
406            "test"
407        ));
408
409        // Use the `message_formatter` macro to construct the formatter
410        let formatter = message_formatter![
411            fmt_message!(human_msg),
412            fmt_template!(ai_message_prompt),
413            fmt_placeholder!("history")
414        ];
415
416        // Define input variables for the AI message template
417        let input_variables = prompt_args! {
418            "content" => "This is a test",
419            "test" => "test2",
420            "history" => vec![
421                Message::new_human_message("Placeholder message 1"),
422                Message::new_ai_message("Placeholder message 2"),
423            ],
424
425
426        };
427
428        // Format messages
429        let formatted_messages = formatter
430            .format_prompt(input_variables)
431            .unwrap()
432            .to_chat_messages();
433
434        // Verify the number of messages
435        assert_eq!(formatted_messages.len(), 4);
436
437        // Verify the content of each message
438        assert_eq!(formatted_messages[0].content, "Hello from user");
439        assert_eq!(
440            formatted_messages[1].content,
441            "AI response: This is a test test2"
442        );
443        assert_eq!(formatted_messages[2].content, "Placeholder message 1");
444        assert_eq!(formatted_messages[3].content, "Placeholder message 2");
445    }
446
447    #[test]
448    fn message_templates_roles_and_variables() {
449        use crate::{
450            prompt::{
451                HumanMessagePromptTemplate, MessageFormatter, PromptError,
452                SystemMessagePromptTemplate,
453            },
454            schemas::MessageType,
455        };
456        let formatter = message_formatter![
457            fmt_template!(SystemMessagePromptTemplate::new(template_fstring!(
458                "You are {role}",
459                "role"
460            ))),
461            fmt_placeholder!("history"),
462            fmt_template!(HumanMessagePromptTemplate::new(template_fstring!(
463                "{q}", "q"
464            ))),
465        ];
466        assert_eq!(
467            formatter.input_variables(),
468            vec!["role".to_string(), "history".into(), "q".into()]
469        );
470
471        let messages = formatter
472            .format_messages(prompt_args! {
473                "role" => "a bot", "q" => "hi", "history" => Vec::<Message>::new()
474            })
475            .unwrap();
476        assert_eq!(messages.len(), 2);
477        assert_eq!(messages[0].message_type, MessageType::SystemMessage);
478        assert_eq!(messages[0].content, "You are a bot");
479        assert_eq!(messages[1].message_type, MessageType::HumanMessage);
480
481        // missing placeholder is an error, not a panic
482        let err = formatter
483            .format_messages(prompt_args! {"role" => "x", "q" => "y"})
484            .unwrap_err();
485        assert!(matches!(err, PromptError::MissingVariable(v) if v == "history"));
486
487        // missing template variable
488        let err = formatter
489            .format_messages(prompt_args! {"history" => Vec::<Message>::new(), "q" => "y"})
490            .unwrap_err();
491        assert!(matches!(err, PromptError::MissingVariable(v) if v == "role"));
492
493        // placeholder with wrong type
494        let err = formatter
495            .format_messages(prompt_args! {"role" => "x", "q" => "y", "history" => "nope"})
496            .unwrap_err();
497        assert!(matches!(err, PromptError::SerializationError(_)));
498    }
499
500    #[test]
501    fn human_template_attaches_images_from_a_variable() {
502        use crate::prompt::{HumanMessagePromptTemplate, MessageFormatter};
503        use crate::schemas::messages::ImageContent;
504        use serde_json::json;
505
506        let template = HumanMessagePromptTemplate::new(template_fstring!("Q: {input}", "input"))
507            .with_images("images");
508        // Images are not prompt variables: the template only needs `input`.
509        assert_eq!(template.input_variables(), vec!["input".to_string()]);
510
511        let urls = |messages: &[Message]| -> Vec<String> {
512            messages[0]
513                .images
514                .iter()
515                .flatten()
516                .map(|i| i.image_url.clone())
517                .collect()
518        };
519
520        // A list mixing strings and ImageContent objects.
521        let messages = template
522            .format_messages(prompt_args! {
523                "input" => "what is it?",
524                "images" => json!(["https://a/cat.png", {"image_url": "iVBORw0KGgo", "detail": "low"}]),
525            })
526            .unwrap();
527        assert_eq!(messages.len(), 1);
528        assert_eq!(messages[0].content, "Q: what is it?");
529        assert_eq!(urls(&messages), vec!["https://a/cat.png", "iVBORw0KGgo"]);
530        assert_eq!(
531            messages[0].images.as_ref().unwrap()[1],
532            ImageContent::from("iVBORw0KGgo").with_detail("low")
533        );
534
535        // A single string works too.
536        let messages = template
537            .format_messages(prompt_args! {"input" => "x", "images" => "https://a/b.png"})
538            .unwrap();
539        assert_eq!(urls(&messages), vec!["https://a/b.png"]);
540
541        // Missing, null or empty: a plain text message (images stays None).
542        for args in [
543            prompt_args! {"input" => "x"},
544            prompt_args! {"input" => "x", "images" => json!(null)},
545            prompt_args! {"input" => "x", "images" => json!([])},
546        ] {
547            assert_eq!(template.format_messages(args).unwrap()[0].images, None);
548        }
549
550        // Without with_images the variable is ignored.
551        let plain = HumanMessagePromptTemplate::new(template_fstring!("{input}", "input"));
552        let messages = plain
553            .format_messages(prompt_args! {"input" => "x", "images" => "https://a/b.png"})
554            .unwrap();
555        assert_eq!(messages[0].images, None);
556
557        // Bad values are errors, not silently dropped.
558        for bad in [
559            json!(42),
560            json!([""]),
561            json!([{"detail": "low"}]),
562            json!([[]]),
563        ] {
564            let err = template
565                .format_messages(prompt_args! {"input" => "x", "images" => bad.clone()})
566                .unwrap_err();
567            assert!(
568                matches!(
569                    err,
570                    PromptError::OtherError(_) | PromptError::SerializationError(_)
571                ),
572                "{bad}: {err:?}"
573            );
574        }
575    }
576
577    #[test]
578    fn human_message_from_inputs_reads_text_and_images() {
579        use crate::prompt::human_message_from_inputs;
580        use serde_json::json;
581
582        let m = human_message_from_inputs(
583            &prompt_args! {"input" => "hi", "images" => json!(["https://a/b.png"])},
584            "input",
585            "images",
586        )
587        .unwrap();
588        assert_eq!(m.message_type, MessageType::HumanMessage);
589        assert_eq!(m.content, "hi");
590        assert_eq!(m.images.unwrap()[0].image_url, "https://a/b.png");
591
592        // Non-string input is rendered without JSON quotes for strings only.
593        let m = human_message_from_inputs(
594            &prompt_args! {"input" => json!({"a": 1})},
595            "input",
596            "images",
597        )
598        .unwrap();
599        assert_eq!(m.content, r#"{"a":1}"#);
600        assert_eq!(m.images, None);
601
602        let err = human_message_from_inputs(&prompt_args! {}, "input", "images").unwrap_err();
603        assert!(matches!(err, PromptError::MissingVariable(v) if v == "input"));
604    }
605}
langchain_rust/prompt/chat.rs

langchain_rust/prompt/
chat.rs
