swiftide_core/chat_completion/
chat_message.rs

1use serde::{Deserialize, Serialize};
2
3use super::tools::{ToolCall, ToolOutput};
4
5#[derive(Clone, strum_macros::EnumIs, PartialEq, Debug, Serialize, Deserialize)]
6pub enum ChatMessage {
7    System(String),
8    User(String),
9    Assistant(Option<String>, Option<Vec<ToolCall>>),
10    ToolOutput(ToolCall, ToolOutput),
11
12    // A summary of the chat. If encountered all previous messages are ignored, except the system
13    // prompt
14    Summary(String),
15}
16
17impl std::fmt::Display for ChatMessage {
18    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
19        match self {
20            ChatMessage::System(s) => write!(f, "System: \"{s}\""),
21            ChatMessage::User(s) => write!(f, "User: \"{s}\""),
22            ChatMessage::Assistant(message, tool_calls) => write!(
23                f,
24                "Assistant: \"{}\", tools: {}",
25                message.as_deref().unwrap_or("None"),
26                tool_calls.as_deref().map_or("None".to_string(), |tc| {
27                    tc.iter()
28                        .map(ToString::to_string)
29                        .collect::<Vec<_>>()
30                        .join(", ")
31                })
32            ),
33            ChatMessage::ToolOutput(tc, to) => write!(f, "ToolOutput: \"{tc}\": \"{to}\""),
34            ChatMessage::Summary(s) => write!(f, "Summary: \"{s}\""),
35        }
36    }
37}
38
39impl ChatMessage {
40    pub fn new_system(message: impl Into<String>) -> Self {
41        ChatMessage::System(message.into())
42    }
43
44    pub fn new_user(message: impl Into<String>) -> Self {
45        ChatMessage::User(message.into())
46    }
47
48    pub fn new_assistant(
49        message: Option<impl Into<String>>,
50        tool_calls: Option<Vec<ToolCall>>,
51    ) -> Self {
52        ChatMessage::Assistant(message.map(Into::into), tool_calls)
53    }
54
55    pub fn new_tool_output(tool_call: impl Into<ToolCall>, output: impl Into<ToolOutput>) -> Self {
56        ChatMessage::ToolOutput(tool_call.into(), output.into())
57    }
58
59    pub fn new_summary(message: impl Into<String>) -> Self {
60        ChatMessage::Summary(message.into())
61    }
62}
63
64/// Returns the content of the message as a string slice.
65///
66/// Note that this omits the tool calls from the assistant message.
67///
68/// If used for estimating tokens, consider this a very rought estimate
69impl AsRef<str> for ChatMessage {
70    fn as_ref(&self) -> &str {
71        match self {
72            ChatMessage::System(s) | ChatMessage::User(s) | ChatMessage::Summary(s) => s,
73            ChatMessage::Assistant(message, _) => message.as_deref().unwrap_or(""),
74            ChatMessage::ToolOutput(_, output) => output.content().unwrap_or(""),
75        }
76    }
77}
swiftide_core/chat_completion/chat_message.rs

swiftide_core/chat_completion/
chat_message.rs
