1use 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#[cfg_attr(feature = "specta", derive(Type))]
27#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
28#[serde(tag = "type")]
29pub enum BaseMessage {
30 #[serde(rename = "human")]
32 Human(HumanMessage),
33 #[serde(rename = "system")]
35 System(SystemMessage),
36 #[serde(rename = "ai")]
38 AI(AIMessage),
39 #[serde(rename = "tool")]
41 Tool(ToolMessage),
42 #[serde(rename = "chat")]
44 Chat(ChatMessage),
45 #[serde(rename = "function")]
47 Function(FunctionMessage),
48 #[serde(rename = "remove")]
50 Remove(RemoveMessage),
51}
52
53impl BaseMessage {
54 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 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 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 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 pub fn tool_calls(&self) -> &[ToolCall] {
114 match self {
115 BaseMessage::AI(m) => m.tool_calls(),
116 _ => &[],
117 }
118 }
119
120 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 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 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 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 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
283pub trait HasId {
286 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#[cfg_attr(feature = "specta", derive(Type))]
300#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
301#[serde(tag = "type")]
302pub enum BaseMessageChunk {
303 #[serde(rename = "AIMessageChunk")]
305 AI(AIMessageChunk),
306 #[serde(rename = "HumanMessageChunk")]
308 Human(HumanMessageChunk),
309 #[serde(rename = "SystemMessageChunk")]
311 System(SystemMessageChunk),
312 #[serde(rename = "ToolMessageChunk")]
314 Tool(ToolMessageChunk),
315 #[serde(rename = "ChatMessageChunk")]
317 Chat(ChatMessageChunk),
318 #[serde(rename = "FunctionMessageChunk")]
320 Function(FunctionMessageChunk),
321}
322
323impl BaseMessageChunk {
324 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 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 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 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#[derive(Debug, Clone, PartialEq)]
413pub enum MergeableContent {
414 Text(String),
416 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
438pub fn merge_content(first: &str, second: &str) -> String {
445 let mut result = first.to_string();
446 result.push_str(second);
447 result
448}
449
450pub 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
510pub 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
517pub 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
530pub fn messages_to_dict(messages: &[BaseMessage]) -> Vec<serde_json::Value> {
534 messages.iter().map(message_to_dict).collect()
535}
536
537pub 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
565pub fn get_bolded_text(text: &str) -> String {
569 format!("\x1b[1m{}\x1b[0m", text)
570}
571
572fn 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
589pub 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
613pub fn is_interactive_env() -> bool {
619 false
620}