1use serde::{Deserialize, Serialize};
2
3use crate::catalog::LlmModel;
4use crate::types::IsoString;
5
6use super::{ToolCallError, ToolCallRequest, ToolCallResult};
7
8#[doc = include_str!("docs/content_block.md")]
9#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
10#[serde(tag = "type", rename_all = "camelCase")]
11pub enum ContentBlock {
12 Text { text: String },
13 Image { data: String, mime_type: String },
14 Audio { data: String, mime_type: String },
15}
16
17impl ContentBlock {
18 pub fn text(s: impl Into<String>) -> Self {
19 ContentBlock::Text { text: s.into() }
20 }
21
22 pub fn estimated_bytes(&self) -> usize {
23 match self {
24 ContentBlock::Text { text } => text.len(),
25 ContentBlock::Image { data, .. } | ContentBlock::Audio { data, .. } => data.len(),
26 }
27 }
28
29 pub fn is_image(&self) -> bool {
30 matches!(self, ContentBlock::Image { .. })
31 }
32
33 pub fn first_text(parts: &[ContentBlock]) -> Option<&str> {
34 parts.iter().find_map(|part| match part {
35 ContentBlock::Text { text } => {
36 let trimmed = text.trim();
37 (!trimmed.is_empty()).then_some(trimmed)
38 }
39 _ => None,
40 })
41 }
42
43 pub fn join_text(parts: &[ContentBlock]) -> String {
45 parts
46 .iter()
47 .filter_map(|p| match p {
48 ContentBlock::Text { text } => Some(text.as_str()),
49 _ => None,
50 })
51 .collect::<Vec<_>>()
52 .join("\n")
53 }
54
55 pub fn as_data_uri(&self) -> Option<String> {
57 match self {
58 ContentBlock::Image { data, mime_type } | ContentBlock::Audio { data, mime_type } => {
59 Some(format!("data:{mime_type};base64,{data}"))
60 }
61 ContentBlock::Text { .. } => None,
62 }
63 }
64}
65
66#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
72pub struct EncryptedReasoningContent {
73 pub id: String,
74 #[serde(serialize_with = "serialize_llm_model", deserialize_with = "deserialize_llm_model")]
75 pub model: LlmModel,
76 pub content: String,
77}
78
79#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
84pub struct AssistantReasoning {
85 #[serde(default, skip_serializing_if = "Option::is_none")]
86 pub summary_text: Option<String>,
87
88 #[serde(default, skip_serializing_if = "Option::is_none")]
89 pub encrypted_content: Option<EncryptedReasoningContent>,
90}
91
92impl AssistantReasoning {
93 pub fn from_parts(summary_text: String, encrypted: Option<EncryptedReasoningContent>) -> Self {
94 Self { summary_text: (!summary_text.is_empty()).then_some(summary_text), encrypted_content: encrypted }
95 }
96
97 pub fn is_empty(&self) -> bool {
98 self.summary_text.is_none() && self.encrypted_content.is_none()
99 }
100}
101
102#[doc = include_str!("docs/chat_message.md")]
103#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
104#[serde(tag = "type", rename_all = "camelCase")]
105pub enum ChatMessage {
106 System {
107 content: String,
108 timestamp: IsoString,
109 },
110 User {
111 content: Vec<ContentBlock>,
112 timestamp: IsoString,
113 },
114 Assistant {
115 content: String,
116 #[serde(default)]
117 reasoning: AssistantReasoning,
118 timestamp: IsoString,
119 tool_calls: Vec<ToolCallRequest>,
120 },
121 ToolCallResult(Result<ToolCallResult, ToolCallError>),
122 Error {
123 message: String,
124 timestamp: IsoString,
125 },
126 Summary {
129 content: String,
130 timestamp: IsoString,
131 messages_compacted: usize,
133 },
134}
135
136impl ChatMessage {
137 pub fn is_tool_result(&self) -> bool {
139 matches!(self, ChatMessage::ToolCallResult(_))
140 }
141
142 pub fn is_system(&self) -> bool {
144 matches!(self, ChatMessage::System { .. })
145 }
146
147 pub fn is_summary(&self) -> bool {
149 matches!(self, ChatMessage::Summary { .. })
150 }
151
152 pub fn estimated_bytes(&self) -> usize {
155 match self {
156 ChatMessage::System { content, .. }
157 | ChatMessage::Error { message: content, .. }
158 | ChatMessage::Summary { content, .. } => content.len(),
159 ChatMessage::User { content, .. } => content.iter().map(ContentBlock::estimated_bytes).sum(),
160 ChatMessage::Assistant { content, reasoning, tool_calls, .. } => {
161 content.len()
162 + reasoning.summary_text.as_ref().map_or(0, String::len)
163 + reasoning.encrypted_content.as_ref().map_or(0, |ec| ec.content.len())
164 + tool_calls.iter().map(|tc| tc.name.len() + tc.arguments.len()).sum::<usize>()
165 }
166 ChatMessage::ToolCallResult(Ok(result)) => result.name.len() + result.arguments.len() + result.result.len(),
167 ChatMessage::ToolCallResult(Err(error)) => {
168 error.name.len() + error.arguments.as_ref().map_or(0, String::len) + error.error.len()
169 }
170 }
171 }
172
173 pub fn timestamp(&self) -> Option<&IsoString> {
175 match self {
176 ChatMessage::System { timestamp, .. }
177 | ChatMessage::User { timestamp, .. }
178 | ChatMessage::Assistant { timestamp, .. }
179 | ChatMessage::Error { timestamp, .. }
180 | ChatMessage::Summary { timestamp, .. } => Some(timestamp),
181 ChatMessage::ToolCallResult(_) => None,
182 }
183 }
184}
185
186fn serialize_llm_model<S: serde::Serializer>(model: &LlmModel, s: S) -> Result<S::Ok, S::Error> {
187 s.serialize_str(&model.to_string())
188}
189
190fn deserialize_llm_model<'de, D: serde::Deserializer<'de>>(d: D) -> Result<LlmModel, D::Error> {
191 let s = String::deserialize(d)?;
192 s.parse::<LlmModel>().map_err(serde::de::Error::custom)
193}
194
195#[cfg(test)]
196mod tests {
197 use super::*;
198
199 fn make_model() -> LlmModel {
200 "anthropic:claude-opus-4-6".parse().unwrap()
201 }
202
203 #[test]
204 fn assistant_reasoning_is_empty_when_default() {
205 let r = AssistantReasoning::default();
206 assert!(r.is_empty());
207 }
208
209 #[test]
210 fn assistant_reasoning_not_empty_with_summary() {
211 let r = AssistantReasoning::from_parts("thinking".to_string(), None);
212 assert!(!r.is_empty());
213 }
214
215 #[test]
216 fn assistant_reasoning_not_empty_with_encrypted() {
217 let r = AssistantReasoning {
218 summary_text: None,
219 encrypted_content: Some(EncryptedReasoningContent {
220 id: "r_test".to_string(),
221 model: make_model(),
222 content: "blob".to_string(),
223 }),
224 };
225 assert!(!r.is_empty());
226 }
227
228 #[test]
229 fn from_parts_empty_summary_is_none() {
230 let r = AssistantReasoning::from_parts(String::new(), None);
231 assert!(r.summary_text.is_none());
232 assert!(r.is_empty());
233 }
234
235 #[test]
236 fn first_text_returns_first_non_empty_text_block() {
237 let parts = vec![
238 ContentBlock::Image { data: "a".to_string(), mime_type: "image/png".to_string() },
239 ContentBlock::text(" "),
240 ContentBlock::text("hello"),
241 ];
242
243 assert_eq!(ContentBlock::first_text(&parts), Some("hello"));
244 }
245
246 #[test]
247 fn encrypted_reasoning_content_serde_roundtrip() {
248 let model = make_model();
249 let ec = EncryptedReasoningContent {
250 id: "r_test".to_string(),
251 model: model.clone(),
252 content: "encrypted-data".to_string(),
253 };
254 let json = serde_json::to_string(&ec).unwrap();
255 let parsed: EncryptedReasoningContent = serde_json::from_str(&json).unwrap();
256 assert_eq!(parsed.model, model);
257 assert_eq!(parsed.content, "encrypted-data");
258 }
259
260 #[test]
261 fn assistant_reasoning_serde_roundtrip() {
262 let model = make_model();
263 let r = AssistantReasoning {
264 summary_text: Some("thought".to_string()),
265 encrypted_content: Some(EncryptedReasoningContent {
266 id: "r_test".to_string(),
267 model,
268 content: "blob".to_string(),
269 }),
270 };
271 let json = serde_json::to_string(&r).unwrap();
272 let parsed: AssistantReasoning = serde_json::from_str(&json).unwrap();
273 assert_eq!(parsed, r);
274 }
275
276 #[test]
277 fn assistant_reasoning_serde_empty_roundtrip() {
278 let r = AssistantReasoning::default();
279 let json = serde_json::to_string(&r).unwrap();
280 assert_eq!(json, "{}");
281 let parsed: AssistantReasoning = serde_json::from_str(&json).unwrap();
282 assert_eq!(parsed, r);
283 }
284
285 #[test]
286 fn chat_message_assistant_serde_roundtrip_with_reasoning() {
287 let model = make_model();
288 let msg = ChatMessage::Assistant {
289 content: "response".to_string(),
290 reasoning: AssistantReasoning {
291 summary_text: Some("plan".to_string()),
292 encrypted_content: Some(EncryptedReasoningContent {
293 id: "r_test".to_string(),
294 model,
295 content: "enc".to_string(),
296 }),
297 },
298 timestamp: IsoString::now(),
299 tool_calls: vec![],
300 };
301 let json = serde_json::to_string(&msg).unwrap();
302 let parsed: ChatMessage = serde_json::from_str(&json).unwrap();
303 assert_eq!(parsed, msg);
304 }
305
306 #[test]
307 fn estimated_bytes_includes_encrypted_content() {
308 let model = make_model();
309 let msg_with = ChatMessage::Assistant {
310 content: "hi".to_string(),
311 reasoning: AssistantReasoning {
312 summary_text: Some("think".to_string()),
313 encrypted_content: Some(EncryptedReasoningContent {
314 id: "r_test".to_string(),
315 model,
316 content: "x".repeat(100),
317 }),
318 },
319 timestamp: IsoString::now(),
320 tool_calls: vec![],
321 };
322 let msg_without = ChatMessage::Assistant {
323 content: "hi".to_string(),
324 reasoning: AssistantReasoning { summary_text: Some("think".to_string()), encrypted_content: None },
325 timestamp: IsoString::now(),
326 tool_calls: vec![],
327 };
328 assert!(msg_with.estimated_bytes() > msg_without.estimated_bytes());
329 }
330}