1use schemars::JsonSchema;
2use serde::{Deserialize, Serialize};
3
4use crate::catalog::LlmModel;
5use crate::types::IsoString;
6
7use super::{AssistantReply, MessageId, ToolCallError, ToolCallResult};
8
9#[doc = include_str!("docs/content_block.md")]
10#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
11#[serde(tag = "type", rename_all = "camelCase")]
12pub enum ContentBlock {
13 Text { text: String },
14 Image { data: String, mime_type: String },
15 Audio { data: String, mime_type: String },
16}
17
18impl ContentBlock {
19 pub fn text(s: impl Into<String>) -> Self {
20 ContentBlock::Text { text: s.into() }
21 }
22
23 pub fn estimated_bytes(&self) -> usize {
24 match self {
25 ContentBlock::Text { text } => text.len(),
26 ContentBlock::Image { data, .. } | ContentBlock::Audio { data, .. } => data.len(),
27 }
28 }
29
30 pub fn is_image(&self) -> bool {
31 matches!(self, ContentBlock::Image { .. })
32 }
33
34 pub fn first_text(parts: &[ContentBlock]) -> Option<&str> {
35 parts.iter().find_map(|part| match part {
36 ContentBlock::Text { text } => {
37 let trimmed = text.trim();
38 (!trimmed.is_empty()).then_some(trimmed)
39 }
40 _ => None,
41 })
42 }
43
44 pub fn join_text(parts: &[ContentBlock]) -> String {
46 parts
47 .iter()
48 .filter_map(|p| match p {
49 ContentBlock::Text { text } => Some(text.as_str()),
50 _ => None,
51 })
52 .collect::<Vec<_>>()
53 .join("\n")
54 }
55
56 pub fn as_data_uri(&self) -> Option<String> {
58 match self {
59 ContentBlock::Image { data, mime_type } | ContentBlock::Audio { data, mime_type } => {
60 Some(format!("data:{mime_type};base64,{data}"))
61 }
62 ContentBlock::Text { .. } => None,
63 }
64 }
65}
66
67#[doc = include_str!("docs/chat_message.md")]
68#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
69#[serde(tag = "type", rename_all = "camelCase")]
70pub enum ChatMessage {
71 System {
72 content: String,
73 timestamp: IsoString,
74 },
75 User {
76 message_id: MessageId,
77 content: Vec<ContentBlock>,
78 timestamp: IsoString,
79 },
80 Assistant {
81 message_id: MessageId,
82 #[serde(default, skip_serializing_if = "Option::is_none")]
83 model: Option<LlmModel>,
84 content: AssistantReply,
85 timestamp: IsoString,
86 },
87 ToolCallResult(Result<ToolCallResult, ToolCallError>),
88 Error {
89 message_id: MessageId,
90 message: String,
91 timestamp: IsoString,
92 },
93 Summary {
96 message_id: MessageId,
97 content: String,
98 timestamp: IsoString,
99 messages_compacted: usize,
101 },
102}
103
104impl ChatMessage {
105 pub fn user(text: impl Into<String>) -> Self {
107 Self::user_with_id(MessageId::new(), vec![ContentBlock::text(text)])
108 }
109
110 pub fn system(content: impl Into<String>) -> Self {
112 ChatMessage::System { content: content.into(), timestamp: IsoString::now() }
113 }
114
115 pub fn user_with_id(message_id: MessageId, content: Vec<ContentBlock>) -> Self {
116 Self::User { message_id, content, timestamp: IsoString::now() }
117 }
118
119 pub fn message_id(&self) -> Option<MessageId> {
121 match self {
122 Self::User { message_id, .. }
123 | Self::Assistant { message_id, .. }
124 | Self::Error { message_id, .. }
125 | Self::Summary { message_id, .. } => Some(message_id.clone()),
126 Self::ToolCallResult(Ok(result)) => Some(MessageId::tool_result(&result.id)),
127 Self::ToolCallResult(Err(error)) => Some(MessageId::tool_result(&error.id)),
128 Self::System { .. } => None,
129 }
130 }
131
132 pub fn is_tool_result(&self) -> bool {
134 matches!(self, ChatMessage::ToolCallResult(_))
135 }
136
137 pub fn is_system(&self) -> bool {
139 matches!(self, ChatMessage::System { .. })
140 }
141
142 pub fn is_summary(&self) -> bool {
144 matches!(self, ChatMessage::Summary { .. })
145 }
146
147 pub fn estimated_bytes(&self) -> usize {
150 match self {
151 ChatMessage::System { content, .. }
152 | ChatMessage::Error { message: content, .. }
153 | ChatMessage::Summary { content, .. } => content.len(),
154 ChatMessage::User { content, .. } => content.iter().map(ContentBlock::estimated_bytes).sum(),
155 ChatMessage::Assistant { content, .. } => content.estimated_bytes(),
156 ChatMessage::ToolCallResult(Ok(result)) => result.name.len() + result.arguments.len() + result.result.len(),
157 ChatMessage::ToolCallResult(Err(error)) => {
158 error.name.len() + error.arguments.as_ref().map_or(0, String::len) + error.error.len()
159 }
160 }
161 }
162
163 pub fn timestamp(&self) -> Option<&IsoString> {
165 match self {
166 ChatMessage::System { timestamp, .. }
167 | ChatMessage::User { timestamp, .. }
168 | ChatMessage::Assistant { timestamp, .. }
169 | ChatMessage::Error { timestamp, .. }
170 | ChatMessage::Summary { timestamp, .. } => Some(timestamp),
171 ChatMessage::ToolCallResult(_) => None,
172 }
173 }
174
175 pub(crate) fn as_user_note(&self) -> Option<String> {
176 match self {
177 ChatMessage::Summary { content, .. } => Some(format!("[Previous conversation handoff]\n\n{content}")),
178 ChatMessage::Error { message, .. } => Some(format!("Error: {message}")),
179 _ => None,
180 }
181 }
182}
183
184#[cfg(test)]
185mod tests {
186 use super::*;
187
188 #[test]
189 fn first_text_returns_first_non_empty_text_block() {
190 let parts = vec![
191 ContentBlock::Image { data: "a".to_string(), mime_type: "image/png".to_string() },
192 ContentBlock::text(" "),
193 ContentBlock::text("hello"),
194 ];
195
196 assert_eq!(ContentBlock::first_text(&parts), Some("hello"));
197 }
198
199 #[test]
200 fn estimated_bytes_counts_reasoning_in_the_form_it_is_sent_back() {
201 let summary = "s".repeat(10);
202 let encrypted = "e".repeat(100);
203 let assistant = crate::testing::assistant_message().text("hi");
204
205 assert_eq!(assistant.clone().reasoning(&summary).build().estimated_bytes(), 12);
206 assert_eq!(assistant.reasoning(&summary).encrypted_reasoning("r_1", &encrypted).build().estimated_bytes(), 102);
207 }
208
209 #[test]
210 fn assistant_messages_roundtrip_with_their_author() {
211 let message = crate::testing::assistant_message()
212 .model("anthropic:claude-opus-4-6")
213 .reasoning("plan")
214 .encrypted_reasoning("r_1", "blob")
215 .text("answer")
216 .build();
217
218 let json = serde_json::to_value(&message).unwrap();
219
220 assert_eq!(json["model"], "anthropic:claude-opus-4-6");
221 assert_eq!(serde_json::from_value::<ChatMessage>(json).unwrap(), message);
222 }
223}