swiftide_core/chat_completion/
chat_message.rs

1use super::tools::{ToolCall, ToolOutput};
2
3#[derive(Clone, strum_macros::EnumIs, PartialEq, Debug)]
4pub enum ChatMessage {
5    System(String),
6    User(String),
7    Assistant(Option<String>, Option<Vec<ToolCall>>),
8    ToolOutput(ToolCall, ToolOutput),
9
10    // A summary of the chat. If encountered all previous messages are ignored, except the system
11    // prompt
12    Summary(String),
13}
14
15impl std::fmt::Display for ChatMessage {
16    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
17        match self {
18            ChatMessage::System(s) => write!(f, "System: \"{s}\""),
19            ChatMessage::User(s) => write!(f, "User: \"{s}\""),
20            ChatMessage::Assistant(message, tool_calls) => write!(
21                f,
22                "Assistant: \"{}\", tools: {}",
23                message.as_deref().unwrap_or("None"),
24                tool_calls.as_deref().map_or("None".to_string(), |tc| {
25                    tc.iter()
26                        .map(ToString::to_string)
27                        .collect::<Vec<_>>()
28                        .join(", ")
29                })
30            ),
31            ChatMessage::ToolOutput(tc, to) => write!(f, "ToolOutput: \"{tc}\": \"{to}\""),
32            ChatMessage::Summary(s) => write!(f, "Summary: \"{s}\""),
33        }
34    }
35}
36
37impl ChatMessage {
38    pub fn new_system(message: impl Into<String>) -> Self {
39        ChatMessage::System(message.into())
40    }
41
42    pub fn new_user(message: impl Into<String>) -> Self {
43        ChatMessage::User(message.into())
44    }
45
46    pub fn new_assistant(
47        message: Option<impl Into<String>>,
48        tool_calls: Option<Vec<ToolCall>>,
49    ) -> Self {
50        ChatMessage::Assistant(message.map(Into::into), tool_calls)
51    }
52
53    pub fn new_tool_output(tool_call: impl Into<ToolCall>, output: impl Into<ToolOutput>) -> Self {
54        ChatMessage::ToolOutput(tool_call.into(), output.into())
55    }
56
57    pub fn new_summary(message: impl Into<String>) -> Self {
58        ChatMessage::Summary(message.into())
59    }
60}