agent_chain_core/messages/
chat.rs

1//! Chat message type.
2//!
3//! This module contains the `ChatMessage` and `ChatMessageChunk` types which represent
4//! messages with an arbitrary speaker role. Mirrors `langchain_core.messages.chat`.
5
6use crate::utils::uuid7;
7use serde::{Deserialize, Serialize};
8use std::collections::HashMap;
9
10#[cfg(feature = "specta")]
11use specta::Type;
12
13use super::base::merge_content;
14
15/// A chat message that can be assigned an arbitrary speaker (role).
16///
17/// Use this when you need to specify a custom role that isn't covered
18/// by the standard message types (Human, AI, System, Tool).
19///
20/// This corresponds to `ChatMessage` in LangChain Python.
21#[cfg_attr(feature = "specta", derive(Type))]
22#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
23pub struct ChatMessage {
24    /// The message content
25    content: String,
26    /// The speaker / role of the message
27    role: String,
28    /// Optional unique identifier
29    id: Option<String>,
30    /// Optional name for the message
31    #[serde(skip_serializing_if = "Option::is_none")]
32    name: Option<String>,
33    /// Additional metadata
34    #[serde(default)]
35    additional_kwargs: HashMap<String, serde_json::Value>,
36    /// Response metadata
37    #[serde(default)]
38    response_metadata: HashMap<String, serde_json::Value>,
39}
40
41impl ChatMessage {
42    /// Create a new chat message with the given role.
43    pub fn new(role: impl Into<String>, content: impl Into<String>) -> Self {
44        Self {
45            content: content.into(),
46            role: role.into(),
47            id: Some(uuid7(None).to_string()),
48            name: None,
49            additional_kwargs: HashMap::new(),
50            response_metadata: HashMap::new(),
51        }
52    }
53
54    /// Create a new chat message with an explicit ID.
55    ///
56    /// Use this when deserializing or reconstructing messages where the ID must be preserved.
57    pub fn with_id(
58        id: impl Into<String>,
59        role: impl Into<String>,
60        content: impl Into<String>,
61    ) -> Self {
62        Self {
63            content: content.into(),
64            role: role.into(),
65            id: Some(id.into()),
66            name: None,
67            additional_kwargs: HashMap::new(),
68            response_metadata: HashMap::new(),
69        }
70    }
71
72    /// Set the name for this message.
73    pub fn with_name(mut self, name: impl Into<String>) -> Self {
74        self.name = Some(name.into());
75        self
76    }
77
78    /// Get the message content.
79    pub fn content(&self) -> &str {
80        &self.content
81    }
82
83    /// Get the message role.
84    pub fn role(&self) -> &str {
85        &self.role
86    }
87
88    /// Get the message ID.
89    pub fn id(&self) -> Option<&str> {
90        self.id.as_deref()
91    }
92
93    /// Get the message name.
94    pub fn name(&self) -> Option<&str> {
95        self.name.as_deref()
96    }
97
98    /// Get additional kwargs.
99    pub fn additional_kwargs(&self) -> &HashMap<String, serde_json::Value> {
100        &self.additional_kwargs
101    }
102
103    /// Get response metadata.
104    pub fn response_metadata(&self) -> &HashMap<String, serde_json::Value> {
105        &self.response_metadata
106    }
107
108    /// Set response metadata.
109    pub fn with_response_metadata(
110        mut self,
111        response_metadata: HashMap<String, serde_json::Value>,
112    ) -> Self {
113        self.response_metadata = response_metadata;
114        self
115    }
116}
117
118/// Chat message chunk (yielded when streaming).
119///
120/// This corresponds to `ChatMessageChunk` in LangChain Python.
121#[cfg_attr(feature = "specta", derive(Type))]
122#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
123pub struct ChatMessageChunk {
124    /// The message content (may be partial during streaming)
125    content: String,
126    /// The speaker / role of the message
127    role: String,
128    /// Optional unique identifier
129    id: Option<String>,
130    /// Optional name for the message
131    #[serde(skip_serializing_if = "Option::is_none")]
132    name: Option<String>,
133    /// Additional metadata
134    #[serde(default)]
135    additional_kwargs: HashMap<String, serde_json::Value>,
136    /// Response metadata
137    #[serde(default)]
138    response_metadata: HashMap<String, serde_json::Value>,
139}
140
141impl ChatMessageChunk {
142    /// Create a new chat message chunk with the given role.
143    pub fn new(role: impl Into<String>, content: impl Into<String>) -> Self {
144        Self {
145            content: content.into(),
146            role: role.into(),
147            id: None,
148            name: None,
149            additional_kwargs: HashMap::new(),
150            response_metadata: HashMap::new(),
151        }
152    }
153
154    /// Create a new chat message chunk with an ID.
155    pub fn with_id(
156        id: impl Into<String>,
157        role: impl Into<String>,
158        content: impl Into<String>,
159    ) -> Self {
160        Self {
161            content: content.into(),
162            role: role.into(),
163            id: Some(id.into()),
164            name: None,
165            additional_kwargs: HashMap::new(),
166            response_metadata: HashMap::new(),
167        }
168    }
169
170    /// Get the message content.
171    pub fn content(&self) -> &str {
172        &self.content
173    }
174
175    /// Get the message role.
176    pub fn role(&self) -> &str {
177        &self.role
178    }
179
180    /// Get the message ID.
181    pub fn id(&self) -> Option<&str> {
182        self.id.as_deref()
183    }
184
185    /// Get the message name.
186    pub fn name(&self) -> Option<&str> {
187        self.name.as_deref()
188    }
189
190    /// Get additional kwargs.
191    pub fn additional_kwargs(&self) -> &HashMap<String, serde_json::Value> {
192        &self.additional_kwargs
193    }
194
195    /// Get response metadata.
196    pub fn response_metadata(&self) -> &HashMap<String, serde_json::Value> {
197        &self.response_metadata
198    }
199
200    /// Concatenate this chunk with another chunk.
201    ///
202    /// # Panics
203    ///
204    /// Panics if the roles are different.
205    pub fn concat(&self, other: &ChatMessageChunk) -> ChatMessageChunk {
206        if self.role != other.role {
207            panic!("Cannot concatenate ChatMessageChunks with different roles");
208        }
209
210        let content = merge_content(&self.content, &other.content);
211
212        // Merge additional_kwargs
213        let mut additional_kwargs = self.additional_kwargs.clone();
214        for (k, v) in &other.additional_kwargs {
215            additional_kwargs.insert(k.clone(), v.clone());
216        }
217
218        // Merge response_metadata
219        let mut response_metadata = self.response_metadata.clone();
220        for (k, v) in &other.response_metadata {
221            response_metadata.insert(k.clone(), v.clone());
222        }
223
224        ChatMessageChunk {
225            content,
226            role: self.role.clone(),
227            id: self.id.clone().or_else(|| other.id.clone()),
228            name: self.name.clone().or_else(|| other.name.clone()),
229            additional_kwargs,
230            response_metadata,
231        }
232    }
233
234    /// Convert this chunk to a complete ChatMessage.
235    pub fn to_message(&self) -> ChatMessage {
236        ChatMessage {
237            content: self.content.clone(),
238            role: self.role.clone(),
239            id: self.id.clone(),
240            name: self.name.clone(),
241            additional_kwargs: self.additional_kwargs.clone(),
242            response_metadata: self.response_metadata.clone(),
243        }
244    }
245}
246
247impl std::ops::Add for ChatMessageChunk {
248    type Output = ChatMessageChunk;
249
250    fn add(self, other: ChatMessageChunk) -> ChatMessageChunk {
251        self.concat(&other)
252    }
253}