agent_chain_core/messages/
base.rs

1//! Base message types.
2//!
3//! This module contains the core `BaseMessage` enum and related traits,
4//! mirroring `langchain_core.messages.base`.
5
6use serde::{Deserialize, Serialize};
7use serde_json::Value;
8use std::collections::HashMap;
9
10#[cfg(feature = "specta")]
11use specta::Type;
12
13use super::ai::{AIMessage, AIMessageChunk};
14use super::chat::{ChatMessage, ChatMessageChunk};
15use super::content::ReasoningContentBlock;
16use super::function::{FunctionMessage, FunctionMessageChunk};
17use super::human::{HumanMessage, HumanMessageChunk};
18use super::modifier::RemoveMessage;
19use super::system::{SystemMessage, SystemMessageChunk};
20use super::tool::{ToolCall, ToolMessage, ToolMessageChunk};
21use crate::utils::merge::merge_lists;
22
23/// A unified message type that can represent any message role.
24///
25/// This corresponds to `BaseMessage` in LangChain Python.
26#[cfg_attr(feature = "specta", derive(Type))]
27#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
28#[serde(tag = "type")]
29pub enum BaseMessage {
30    /// A human message
31    #[serde(rename = "human")]
32    Human(HumanMessage),
33    /// A system message
34    #[serde(rename = "system")]
35    System(SystemMessage),
36    /// An AI message
37    #[serde(rename = "ai")]
38    AI(AIMessage),
39    /// A tool result message
40    #[serde(rename = "tool")]
41    Tool(ToolMessage),
42    /// A chat message with arbitrary role
43    #[serde(rename = "chat")]
44    Chat(ChatMessage),
45    /// A function message (deprecated, use Tool)
46    #[serde(rename = "function")]
47    Function(FunctionMessage),
48    /// A remove message (for message deletion)
49    #[serde(rename = "remove")]
50    Remove(RemoveMessage),
51}
52
53impl BaseMessage {
54    /// Get the message content as a string reference.
55    ///
56    /// For messages with multimodal content, this returns the first text content
57    /// or an empty string.
58    pub fn content(&self) -> &str {
59        match self {
60            BaseMessage::Human(m) => m.content(),
61            BaseMessage::System(m) => m.content(),
62            BaseMessage::AI(m) => m.content(),
63            BaseMessage::Tool(m) => m.content(),
64            BaseMessage::Chat(m) => m.content(),
65            BaseMessage::Function(m) => m.content(),
66            BaseMessage::Remove(_) => "",
67        }
68    }
69
70    /// Get the text content of the message as a string.
71    ///
72    /// This extracts text from both simple string content and list content
73    /// (filtering for text blocks). Corresponds to the `text` property in Python.
74    pub fn text(&self) -> String {
75        match self {
76            BaseMessage::Human(m) => m.message_content().as_text(),
77            BaseMessage::System(m) => m.content().to_string(),
78            BaseMessage::AI(m) => m.content().to_string(),
79            BaseMessage::Tool(m) => m.content().to_string(),
80            BaseMessage::Chat(m) => m.content().to_string(),
81            BaseMessage::Function(m) => m.content().to_string(),
82            BaseMessage::Remove(_) => String::new(),
83        }
84    }
85
86    /// Get the message ID.
87    pub fn id(&self) -> Option<&str> {
88        match self {
89            BaseMessage::Human(m) => m.id(),
90            BaseMessage::System(m) => m.id(),
91            BaseMessage::AI(m) => m.id(),
92            BaseMessage::Tool(m) => m.id(),
93            BaseMessage::Chat(m) => m.id(),
94            BaseMessage::Function(m) => m.id(),
95            BaseMessage::Remove(m) => m.id(),
96        }
97    }
98
99    /// Get the message name if present.
100    pub fn name(&self) -> Option<&str> {
101        match self {
102            BaseMessage::Human(m) => m.name(),
103            BaseMessage::System(m) => m.name(),
104            BaseMessage::AI(m) => m.name(),
105            BaseMessage::Tool(m) => m.name(),
106            BaseMessage::Chat(m) => m.name(),
107            BaseMessage::Function(_) => None,
108            BaseMessage::Remove(_) => None,
109        }
110    }
111
112    /// Get tool calls if this is an AI message.
113    pub fn tool_calls(&self) -> &[ToolCall] {
114        match self {
115            BaseMessage::AI(m) => m.tool_calls(),
116            _ => &[],
117        }
118    }
119
120    /// Get the message type as a string.
121    pub fn message_type(&self) -> &'static str {
122        match self {
123            BaseMessage::Human(_) => "human",
124            BaseMessage::System(_) => "system",
125            BaseMessage::AI(_) => "ai",
126            BaseMessage::Tool(_) => "tool",
127            BaseMessage::Chat(_) => "chat",
128            BaseMessage::Function(_) => "function",
129            BaseMessage::Remove(_) => "remove",
130        }
131    }
132
133    /// Get additional kwargs if present.
134    pub fn additional_kwargs(&self) -> Option<&HashMap<String, serde_json::Value>> {
135        match self {
136            BaseMessage::Human(m) => Some(m.additional_kwargs()),
137            BaseMessage::System(m) => Some(m.additional_kwargs()),
138            BaseMessage::AI(m) => Some(m.additional_kwargs()),
139            BaseMessage::Tool(m) => Some(m.additional_kwargs()),
140            BaseMessage::Chat(m) => Some(m.additional_kwargs()),
141            BaseMessage::Function(m) => Some(m.additional_kwargs()),
142            BaseMessage::Remove(_) => None,
143        }
144    }
145
146    /// Get response metadata if present.
147    pub fn response_metadata(&self) -> Option<&HashMap<String, serde_json::Value>> {
148        match self {
149            BaseMessage::AI(m) => Some(m.response_metadata()),
150            BaseMessage::Chat(m) => Some(m.response_metadata()),
151            BaseMessage::Function(m) => Some(m.response_metadata()),
152            BaseMessage::Tool(m) => Some(m.response_metadata()),
153            _ => None,
154        }
155    }
156
157    /// Pretty print the message to stdout.
158    /// This mimics LangChain's pretty_print() method for messages.
159    pub fn pretty_print(&self) {
160        let (role, content) = match self {
161            BaseMessage::Human(m) => ("Human", m.content()),
162            BaseMessage::System(m) => ("System", m.content()),
163            BaseMessage::AI(m) => {
164                let tool_calls = m.tool_calls();
165                if tool_calls.is_empty() {
166                    ("AI", m.content())
167                } else {
168                    println!(
169                        "================================== AI Message =================================="
170                    );
171                    if !m.content().is_empty() {
172                        println!("{}", m.content());
173                    }
174                    for tc in tool_calls {
175                        println!("Tool Call: {} ({})", tc.name(), tc.id());
176                        println!("  Args: {}", tc.args());
177                    }
178                    return;
179                }
180            }
181            BaseMessage::Tool(m) => {
182                println!(
183                    "================================= Tool Message ================================="
184                );
185                println!("[{}] {}", m.tool_call_id(), m.content());
186                return;
187            }
188            BaseMessage::Chat(m) => (m.role(), m.content()),
189            BaseMessage::Function(m) => {
190                println!(
191                    "=============================== Function Message ==============================="
192                );
193                println!("[{}] {}", m.name(), m.content());
194                return;
195            }
196            BaseMessage::Remove(m) => {
197                println!(
198                    "================================ Remove Message ================================"
199                );
200                if let Some(id) = m.id() {
201                    println!("Remove message with id: {}", id);
202                }
203                return;
204            }
205        };
206
207        let header = format!("=== {} Message ===", role);
208        let padding = (80 - header.len()) / 2;
209        println!(
210            "{:=>padding$}{}{:=>padding$}",
211            "",
212            header,
213            "",
214            padding = padding
215        );
216        println!("{}", content);
217    }
218
219    /// Get a pretty representation of the message.
220    ///
221    /// # Arguments
222    ///
223    /// * `html` - Whether to format the message with bold text (using ANSI codes).
224    ///   Named `html` for Python compatibility but actually uses terminal codes.
225    pub fn pretty_repr(&self, html: bool) -> String {
226        let msg_type = self.message_type();
227        let title_cased = title_case(msg_type);
228        let title = format!("{} Message", title_cased);
229        let title = get_msg_title_repr(&title, html);
230
231        let name_line = if let Some(name) = self.name() {
232            format!("\nName: {}", name)
233        } else {
234            String::new()
235        };
236
237        format!("{}{}\n\n{}", title, name_line, self.content())
238    }
239}
240
241impl From<HumanMessage> for BaseMessage {
242    fn from(msg: HumanMessage) -> Self {
243        BaseMessage::Human(msg)
244    }
245}
246
247impl From<SystemMessage> for BaseMessage {
248    fn from(msg: SystemMessage) -> Self {
249        BaseMessage::System(msg)
250    }
251}
252
253impl From<AIMessage> for BaseMessage {
254    fn from(msg: AIMessage) -> Self {
255        BaseMessage::AI(msg)
256    }
257}
258
259impl From<ToolMessage> for BaseMessage {
260    fn from(msg: ToolMessage) -> Self {
261        BaseMessage::Tool(msg)
262    }
263}
264
265impl From<ChatMessage> for BaseMessage {
266    fn from(msg: ChatMessage) -> Self {
267        BaseMessage::Chat(msg)
268    }
269}
270
271impl From<FunctionMessage> for BaseMessage {
272    fn from(msg: FunctionMessage) -> Self {
273        BaseMessage::Function(msg)
274    }
275}
276
277impl From<RemoveMessage> for BaseMessage {
278    fn from(msg: RemoveMessage) -> Self {
279        BaseMessage::Remove(msg)
280    }
281}
282
283/// Trait for types that have an optional ID.
284/// Used for message merging operations.
285pub trait HasId {
286    /// Get the ID if present.
287    fn get_id(&self) -> Option<&str>;
288}
289
290impl HasId for BaseMessage {
291    fn get_id(&self) -> Option<&str> {
292        self.id()
293    }
294}
295
296/// A message chunk enum that represents streaming message chunks.
297///
298/// This corresponds to `BaseMessageChunk` in LangChain Python.
299#[cfg_attr(feature = "specta", derive(Type))]
300#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
301#[serde(tag = "type")]
302pub enum BaseMessageChunk {
303    /// An AI message chunk
304    #[serde(rename = "AIMessageChunk")]
305    AI(AIMessageChunk),
306    /// A human message chunk
307    #[serde(rename = "HumanMessageChunk")]
308    Human(HumanMessageChunk),
309    /// A system message chunk
310    #[serde(rename = "SystemMessageChunk")]
311    System(SystemMessageChunk),
312    /// A tool message chunk
313    #[serde(rename = "ToolMessageChunk")]
314    Tool(ToolMessageChunk),
315    /// A chat message chunk
316    #[serde(rename = "ChatMessageChunk")]
317    Chat(ChatMessageChunk),
318    /// A function message chunk
319    #[serde(rename = "FunctionMessageChunk")]
320    Function(FunctionMessageChunk),
321}
322
323impl BaseMessageChunk {
324    /// Get the message content.
325    pub fn content(&self) -> &str {
326        match self {
327            BaseMessageChunk::AI(m) => m.content(),
328            BaseMessageChunk::Human(m) => m.content(),
329            BaseMessageChunk::System(m) => m.content(),
330            BaseMessageChunk::Tool(m) => m.content(),
331            BaseMessageChunk::Chat(m) => m.content(),
332            BaseMessageChunk::Function(m) => m.content(),
333        }
334    }
335
336    /// Get the message ID.
337    pub fn id(&self) -> Option<&str> {
338        match self {
339            BaseMessageChunk::AI(m) => m.id(),
340            BaseMessageChunk::Human(m) => m.id(),
341            BaseMessageChunk::System(m) => m.id(),
342            BaseMessageChunk::Tool(m) => m.id(),
343            BaseMessageChunk::Chat(m) => m.id(),
344            BaseMessageChunk::Function(m) => m.id(),
345        }
346    }
347
348    /// Get the message type as a string.
349    pub fn message_type(&self) -> &'static str {
350        match self {
351            BaseMessageChunk::AI(_) => "AIMessageChunk",
352            BaseMessageChunk::Human(_) => "HumanMessageChunk",
353            BaseMessageChunk::System(_) => "SystemMessageChunk",
354            BaseMessageChunk::Tool(_) => "ToolMessageChunk",
355            BaseMessageChunk::Chat(_) => "ChatMessageChunk",
356            BaseMessageChunk::Function(_) => "FunctionMessageChunk",
357        }
358    }
359
360    /// Convert this chunk to a complete message.
361    pub fn to_message(&self) -> BaseMessage {
362        match self {
363            BaseMessageChunk::AI(m) => BaseMessage::AI(m.to_message()),
364            BaseMessageChunk::Human(m) => BaseMessage::Human(m.to_message()),
365            BaseMessageChunk::System(m) => BaseMessage::System(m.to_message()),
366            BaseMessageChunk::Tool(m) => BaseMessage::Tool(m.to_message()),
367            BaseMessageChunk::Chat(m) => BaseMessage::Chat(m.to_message()),
368            BaseMessageChunk::Function(m) => BaseMessage::Function(m.to_message()),
369        }
370    }
371}
372
373impl From<AIMessageChunk> for BaseMessageChunk {
374    fn from(chunk: AIMessageChunk) -> Self {
375        BaseMessageChunk::AI(chunk)
376    }
377}
378
379impl From<HumanMessageChunk> for BaseMessageChunk {
380    fn from(chunk: HumanMessageChunk) -> Self {
381        BaseMessageChunk::Human(chunk)
382    }
383}
384
385impl From<SystemMessageChunk> for BaseMessageChunk {
386    fn from(chunk: SystemMessageChunk) -> Self {
387        BaseMessageChunk::System(chunk)
388    }
389}
390
391impl From<ToolMessageChunk> for BaseMessageChunk {
392    fn from(chunk: ToolMessageChunk) -> Self {
393        BaseMessageChunk::Tool(chunk)
394    }
395}
396
397impl From<ChatMessageChunk> for BaseMessageChunk {
398    fn from(chunk: ChatMessageChunk) -> Self {
399        BaseMessageChunk::Chat(chunk)
400    }
401}
402
403impl From<FunctionMessageChunk> for BaseMessageChunk {
404    fn from(chunk: FunctionMessageChunk) -> Self {
405        BaseMessageChunk::Function(chunk)
406    }
407}
408
409/// Content type for merge operations.
410///
411/// Represents message content that can be either a string or a list of values.
412#[derive(Debug, Clone, PartialEq)]
413pub enum MergeableContent {
414    /// String content.
415    Text(String),
416    /// List content (strings or dicts).
417    List(Vec<Value>),
418}
419
420impl From<String> for MergeableContent {
421    fn from(s: String) -> Self {
422        MergeableContent::Text(s)
423    }
424}
425
426impl From<&str> for MergeableContent {
427    fn from(s: &str) -> Self {
428        MergeableContent::Text(s.to_string())
429    }
430}
431
432impl From<Vec<Value>> for MergeableContent {
433    fn from(v: Vec<Value>) -> Self {
434        MergeableContent::List(v)
435    }
436}
437
438/// Merge multiple message contents (simple string version).
439///
440/// Concatenates two strings together. This is the simple version
441/// that corresponds to the basic case of `merge_content` in LangChain Python.
442///
443/// For more complex merging with lists, use `merge_content_complex`.
444pub fn merge_content(first: &str, second: &str) -> String {
445    let mut result = first.to_string();
446    result.push_str(second);
447    result
448}
449
450/// Merge multiple message contents with support for both strings and lists.
451///
452/// This function handles merging string contents and list contents together.
453/// If both contents are strings, they are concatenated.
454/// If one is a string and one is a list, the string is prepended/appended.
455/// If both are lists, the lists are concatenated with smart merging.
456///
457/// This corresponds to the full `merge_content` function in LangChain Python.
458///
459/// # Arguments
460///
461/// * `first_content` - The first content to merge.
462/// * `contents` - Additional contents to merge.
463///
464/// # Returns
465///
466/// The merged content.
467pub fn merge_content_complex(
468    first_content: Option<MergeableContent>,
469    contents: Vec<MergeableContent>,
470) -> MergeableContent {
471    let mut merged = first_content.unwrap_or(MergeableContent::Text(String::new()));
472
473    for content in contents {
474        merged = match (merged, content) {
475            (MergeableContent::Text(mut left), MergeableContent::Text(right)) => {
476                left.push_str(&right);
477                MergeableContent::Text(left)
478            }
479            (MergeableContent::Text(left), MergeableContent::List(right)) => {
480                let mut new_list = vec![Value::String(left)];
481                new_list.extend(right);
482                MergeableContent::List(new_list)
483            }
484            (MergeableContent::List(mut left), MergeableContent::List(right)) => {
485                if let Ok(Some(merged_list)) =
486                    merge_lists(Some(left.clone()), vec![Some(right.clone())])
487                {
488                    MergeableContent::List(merged_list)
489                } else {
490                    left.extend(right);
491                    MergeableContent::List(left)
492                }
493            }
494            (MergeableContent::List(mut left), MergeableContent::Text(right)) => {
495                if !right.is_empty() {
496                    if let Some(Value::String(last)) = left.last_mut() {
497                        last.push_str(&right);
498                    } else if !left.is_empty() {
499                        left.push(Value::String(right));
500                    }
501                }
502                MergeableContent::List(left)
503            }
504        };
505    }
506
507    merged
508}
509
510/// Merge content vectors (for multimodal content).
511pub fn merge_content_vec(first: Vec<Value>, second: Vec<Value>) -> Vec<Value> {
512    let mut result = first;
513    result.extend(second);
514    result
515}
516
517/// Convert a Message to a dictionary.
518///
519/// This corresponds to `message_to_dict` in LangChain Python.
520/// The dict will have a `type` key with the message type and a `data` key
521/// with the message data as a dict (all fields serialized).
522pub fn message_to_dict(message: &BaseMessage) -> Value {
523    let data = serde_json::to_value(message).unwrap_or_default();
524    serde_json::json!({
525        "type": message.message_type(),
526        "data": data
527    })
528}
529
530/// Convert a sequence of Messages to a list of dictionaries.
531///
532/// This corresponds to `messages_to_dict` in LangChain Python.
533pub fn messages_to_dict(messages: &[BaseMessage]) -> Vec<serde_json::Value> {
534    messages.iter().map(message_to_dict).collect()
535}
536
537/// Get a title representation for a message.
538///
539/// # Arguments
540///
541/// * `title` - The title to format.
542/// * `bold` - Whether to bold the title using ANSI escape codes.
543///
544/// # Returns
545///
546/// The formatted title representation.
547pub fn get_msg_title_repr(title: &str, bold: bool) -> String {
548    let padded = format!(" {} ", title);
549    let sep_len = (80 - padded.len()) / 2;
550    let sep: String = "=".repeat(sep_len);
551    let second_sep = if padded.len() % 2 == 0 {
552        sep.clone()
553    } else {
554        format!("{}=", sep)
555    };
556
557    if bold {
558        let bolded = get_bolded_text(&padded);
559        format!("{}{}{}", sep, bolded, second_sep)
560    } else {
561        format!("{}{}{}", sep, padded, second_sep)
562    }
563}
564
565/// Get bolded text using ANSI escape codes.
566///
567/// Corresponds to `get_bolded_text` in Python's `langchain_core.utils.input`.
568pub fn get_bolded_text(text: &str) -> String {
569    format!("\x1b[1m{}\x1b[0m", text)
570}
571
572/// Convert a string to title case (capitalize first letter of each word).
573fn title_case(s: &str) -> String {
574    s.split('_')
575        .map(|word| {
576            let mut chars = word.chars();
577            match chars.next() {
578                Some(first) => {
579                    let upper = first.to_uppercase().to_string();
580                    upper + &chars.as_str().to_lowercase()
581                }
582                None => String::new(),
583            }
584        })
585        .collect::<Vec<_>>()
586        .join(" ")
587}
588
589/// Extract `reasoning_content` from `additional_kwargs`.
590///
591/// Handles reasoning content stored in various formats:
592/// - `additional_kwargs["reasoning_content"]` (string) - Ollama, DeepSeek, XAI, Groq
593///
594/// Corresponds to `_extract_reasoning_from_additional_kwargs` in Python.
595///
596/// # Arguments
597///
598/// * `additional_kwargs` - The additional_kwargs dictionary from a message.
599///
600/// # Returns
601///
602/// A `ReasoningContentBlock` if reasoning content is found, None otherwise.
603pub fn extract_reasoning_from_additional_kwargs(
604    additional_kwargs: &HashMap<String, Value>,
605) -> Option<ReasoningContentBlock> {
606    if let Some(Value::String(reasoning_content)) = additional_kwargs.get("reasoning_content") {
607        Some(ReasoningContentBlock::new(reasoning_content.clone()))
608    } else {
609        None
610    }
611}
612
613/// Check if running in an interactive environment.
614///
615/// In Rust, this always returns false as we don't have the same
616/// IPython/Jupyter detection available. Applications can override
617/// behavior based on their own environment detection.
618pub fn is_interactive_env() -> bool {
619    false
620}