agent_chain_core/prompts/
message.rs

1//! Message prompt templates.
2//!
3//! This module provides the base trait for message prompt templates,
4//! mirroring `langchain_core.prompts.message` in Python.
5
6use std::collections::HashMap;
7
8use crate::error::Result;
9use crate::messages::BaseMessage;
10use crate::utils::interactive_env::is_interactive_env;
11
12/// Base trait for message prompt templates.
13///
14/// Message prompt templates format into a list of messages rather than a single string.
15/// They are used in chat-based models where conversations consist of multiple messages.
16pub trait BaseMessagePromptTemplate: Send + Sync {
17    /// Get the input variables for this template.
18    ///
19    /// Returns a list of variable names that are required to format this template.
20    fn input_variables(&self) -> Vec<String>;
21
22    /// Format messages from kwargs.
23    ///
24    /// # Arguments
25    ///
26    /// * `kwargs` - Keyword arguments to use for formatting.
27    ///
28    /// # Returns
29    ///
30    /// A list of formatted `BaseMessage` objects.
31    fn format_messages(&self, kwargs: &HashMap<String, String>) -> Result<Vec<BaseMessage>>;
32
33    /// Async format messages from kwargs.
34    ///
35    /// Default implementation calls the sync version.
36    ///
37    /// # Arguments
38    ///
39    /// * `kwargs` - Keyword arguments to use for formatting.
40    ///
41    /// # Returns
42    ///
43    /// A list of formatted `BaseMessage` objects.
44    fn aformat_messages(
45        &self,
46        kwargs: &HashMap<String, String>,
47    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<Vec<BaseMessage>>> + Send + '_>>
48    {
49        let result = self.format_messages(kwargs);
50        Box::pin(async move { result })
51    }
52
53    /// Get a pretty representation of the template.
54    ///
55    /// # Arguments
56    ///
57    /// * `html` - Whether to format as HTML.
58    ///
59    /// # Returns
60    ///
61    /// A human-readable representation of the template.
62    fn pretty_repr(&self, html: bool) -> String;
63
64    /// Print a human-readable representation.
65    fn pretty_print(&self) {
66        println!("{}", self.pretty_repr(is_interactive_env()));
67    }
68}
69
70/// Helper function to get a title representation for a message.
71pub fn get_msg_title_repr(title: &str, bold: bool) -> String {
72    let padded = format!(" {} ", title);
73    let sep_len = (80_usize).saturating_sub(padded.len()) / 2;
74    let sep: String = "=".repeat(sep_len);
75    let second_sep = if padded.len() % 2 == 0 {
76        sep.clone()
77    } else {
78        format!("{}=", sep)
79    };
80
81    if bold {
82        format!("{}\x1b[1m{}\x1b[0m{}", sep, padded, second_sep)
83    } else {
84        format!("{}{}{}", sep, padded, second_sep)
85    }
86}
87
88#[cfg(test)]
89mod tests {
90    use super::*;
91
92    #[test]
93    fn test_get_msg_title_repr() {
94        let title = get_msg_title_repr("Test", false);
95        assert!(title.contains("Test"));
96        assert!(title.contains("="));
97    }
98
99    #[test]
100    fn test_get_msg_title_repr_bold() {
101        let title = get_msg_title_repr("Test", true);
102        assert!(title.contains("Test"));
103        assert!(title.contains("\x1b[1m"));
104        assert!(title.contains("\x1b[0m"));
105    }
106}
agent_chain_core/prompts/message.rs

agent_chain_core/prompts/
message.rs
