agent_chain_core/messages/
chat.rs1use 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#[cfg_attr(feature = "specta", derive(Type))]
22#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
23pub struct ChatMessage {
24 content: String,
26 role: String,
28 id: Option<String>,
30 #[serde(skip_serializing_if = "Option::is_none")]
32 name: Option<String>,
33 #[serde(default)]
35 additional_kwargs: HashMap<String, serde_json::Value>,
36 #[serde(default)]
38 response_metadata: HashMap<String, serde_json::Value>,
39}
40
41impl ChatMessage {
42 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 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 pub fn with_name(mut self, name: impl Into<String>) -> Self {
74 self.name = Some(name.into());
75 self
76 }
77
78 pub fn content(&self) -> &str {
80 &self.content
81 }
82
83 pub fn role(&self) -> &str {
85 &self.role
86 }
87
88 pub fn id(&self) -> Option<&str> {
90 self.id.as_deref()
91 }
92
93 pub fn name(&self) -> Option<&str> {
95 self.name.as_deref()
96 }
97
98 pub fn additional_kwargs(&self) -> &HashMap<String, serde_json::Value> {
100 &self.additional_kwargs
101 }
102
103 pub fn response_metadata(&self) -> &HashMap<String, serde_json::Value> {
105 &self.response_metadata
106 }
107
108 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#[cfg_attr(feature = "specta", derive(Type))]
122#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
123pub struct ChatMessageChunk {
124 content: String,
126 role: String,
128 id: Option<String>,
130 #[serde(skip_serializing_if = "Option::is_none")]
132 name: Option<String>,
133 #[serde(default)]
135 additional_kwargs: HashMap<String, serde_json::Value>,
136 #[serde(default)]
138 response_metadata: HashMap<String, serde_json::Value>,
139}
140
141impl ChatMessageChunk {
142 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 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 pub fn content(&self) -> &str {
172 &self.content
173 }
174
175 pub fn role(&self) -> &str {
177 &self.role
178 }
179
180 pub fn id(&self) -> Option<&str> {
182 self.id.as_deref()
183 }
184
185 pub fn name(&self) -> Option<&str> {
187 self.name.as_deref()
188 }
189
190 pub fn additional_kwargs(&self) -> &HashMap<String, serde_json::Value> {
192 &self.additional_kwargs
193 }
194
195 pub fn response_metadata(&self) -> &HashMap<String, serde_json::Value> {
197 &self.response_metadata
198 }
199
200 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 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 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 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}