1use schemars::JsonSchema;
2use serde::{Deserialize, Serialize};
3
4use crate::catalog::LlmModel;
5use crate::types::IsoString;
6
7use super::{MessageId, ToolCallError, ToolCallRequest, 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#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
73pub struct EncryptedReasoningContent {
74 pub id: String,
75 #[serde(serialize_with = "serialize_llm_model", deserialize_with = "deserialize_llm_model")]
76 pub model: LlmModel,
77 pub content: String,
78}
79
80#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
85pub struct AssistantReasoning {
86 #[serde(default, skip_serializing_if = "Option::is_none")]
87 pub summary_text: Option<String>,
88
89 #[serde(default, skip_serializing_if = "Option::is_none")]
90 pub encrypted_content: Option<EncryptedReasoningContent>,
91}
92
93impl AssistantReasoning {
94 pub fn from_parts(summary_text: String, encrypted: Option<EncryptedReasoningContent>) -> Self {
95 Self { summary_text: (!summary_text.is_empty()).then_some(summary_text), encrypted_content: encrypted }
96 }
97
98 pub fn is_empty(&self) -> bool {
99 self.summary_text.is_none() && self.encrypted_content.is_none()
100 }
101}
102
103#[doc = include_str!("docs/chat_message.md")]
104#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
105#[serde(tag = "type", rename_all = "camelCase")]
106pub enum ChatMessage {
107 System {
108 content: String,
109 timestamp: IsoString,
110 },
111 User {
112 message_id: MessageId,
113 content: Vec<ContentBlock>,
114 timestamp: IsoString,
115 },
116 Assistant {
117 message_id: MessageId,
118 content: String,
119 #[serde(default)]
120 reasoning: AssistantReasoning,
121 timestamp: IsoString,
122 tool_calls: Vec<ToolCallRequest>,
123 },
124 ToolCallResult(Result<ToolCallResult, ToolCallError>),
125 Error {
126 message_id: MessageId,
127 message: String,
128 timestamp: IsoString,
129 },
130 Summary {
133 message_id: MessageId,
134 content: String,
135 timestamp: IsoString,
136 messages_compacted: usize,
138 },
139}
140
141impl ChatMessage {
142 pub fn user(text: impl Into<String>) -> Self {
144 Self::user_with_id(MessageId::new(), vec![ContentBlock::text(text)])
145 }
146
147 pub fn system(content: impl Into<String>) -> Self {
149 ChatMessage::System { content: content.into(), timestamp: IsoString::now() }
150 }
151
152 pub fn user_with_id(message_id: MessageId, content: Vec<ContentBlock>) -> Self {
153 Self::User { message_id, content, timestamp: IsoString::now() }
154 }
155
156 pub fn message_id(&self) -> Option<MessageId> {
158 match self {
159 Self::User { message_id, .. }
160 | Self::Assistant { message_id, .. }
161 | Self::Error { message_id, .. }
162 | Self::Summary { message_id, .. } => Some(message_id.clone()),
163 Self::ToolCallResult(Ok(result)) => Some(MessageId::tool_result(&result.id)),
164 Self::ToolCallResult(Err(error)) => Some(MessageId::tool_result(&error.id)),
165 Self::System { .. } => None,
166 }
167 }
168
169 pub fn is_tool_result(&self) -> bool {
171 matches!(self, ChatMessage::ToolCallResult(_))
172 }
173
174 pub fn is_system(&self) -> bool {
176 matches!(self, ChatMessage::System { .. })
177 }
178
179 pub fn is_summary(&self) -> bool {
181 matches!(self, ChatMessage::Summary { .. })
182 }
183
184 pub fn estimated_bytes(&self) -> usize {
187 match self {
188 ChatMessage::System { content, .. }
189 | ChatMessage::Error { message: content, .. }
190 | ChatMessage::Summary { content, .. } => content.len(),
191 ChatMessage::User { content, .. } => content.iter().map(ContentBlock::estimated_bytes).sum(),
192 ChatMessage::Assistant { content, reasoning, tool_calls, .. } => {
193 content.len()
194 + reasoning.summary_text.as_ref().map_or(0, String::len)
195 + reasoning.encrypted_content.as_ref().map_or(0, |ec| ec.content.len())
196 + tool_calls.iter().map(|tc| tc.name.len() + tc.arguments.len()).sum::<usize>()
197 }
198 ChatMessage::ToolCallResult(Ok(result)) => result.name.len() + result.arguments.len() + result.result.len(),
199 ChatMessage::ToolCallResult(Err(error)) => {
200 error.name.len() + error.arguments.as_ref().map_or(0, String::len) + error.error.len()
201 }
202 }
203 }
204
205 pub fn timestamp(&self) -> Option<&IsoString> {
207 match self {
208 ChatMessage::System { timestamp, .. }
209 | ChatMessage::User { timestamp, .. }
210 | ChatMessage::Assistant { timestamp, .. }
211 | ChatMessage::Error { timestamp, .. }
212 | ChatMessage::Summary { timestamp, .. } => Some(timestamp),
213 ChatMessage::ToolCallResult(_) => None,
214 }
215 }
216}
217
218fn serialize_llm_model<S: serde::Serializer>(model: &LlmModel, s: S) -> Result<S::Ok, S::Error> {
219 s.serialize_str(&model.to_string())
220}
221
222fn deserialize_llm_model<'de, D: serde::Deserializer<'de>>(d: D) -> Result<LlmModel, D::Error> {
223 let s = String::deserialize(d)?;
224 s.parse::<LlmModel>().map_err(serde::de::Error::custom)
225}
226
227#[cfg(test)]
228mod tests {
229 use super::*;
230
231 fn make_model() -> LlmModel {
232 "anthropic:claude-opus-4-6".parse().unwrap()
233 }
234
235 #[test]
236 fn assistant_reasoning_is_empty_when_default() {
237 let r = AssistantReasoning::default();
238 assert!(r.is_empty());
239 }
240
241 #[test]
242 fn assistant_reasoning_not_empty_with_summary() {
243 let r = AssistantReasoning::from_parts("thinking".to_string(), None);
244 assert!(!r.is_empty());
245 }
246
247 #[test]
248 fn assistant_reasoning_not_empty_with_encrypted() {
249 let r = AssistantReasoning {
250 summary_text: None,
251 encrypted_content: Some(EncryptedReasoningContent {
252 id: "r_test".to_string(),
253 model: make_model(),
254 content: "blob".to_string(),
255 }),
256 };
257 assert!(!r.is_empty());
258 }
259
260 #[test]
261 fn from_parts_empty_summary_is_none() {
262 let r = AssistantReasoning::from_parts(String::new(), None);
263 assert!(r.summary_text.is_none());
264 assert!(r.is_empty());
265 }
266
267 #[test]
268 fn first_text_returns_first_non_empty_text_block() {
269 let parts = vec![
270 ContentBlock::Image { data: "a".to_string(), mime_type: "image/png".to_string() },
271 ContentBlock::text(" "),
272 ContentBlock::text("hello"),
273 ];
274
275 assert_eq!(ContentBlock::first_text(&parts), Some("hello"));
276 }
277
278 #[test]
279 fn encrypted_reasoning_content_serde_roundtrip() {
280 let model = make_model();
281 let ec = EncryptedReasoningContent {
282 id: "r_test".to_string(),
283 model: model.clone(),
284 content: "encrypted-data".to_string(),
285 };
286 let json = serde_json::to_string(&ec).unwrap();
287 let parsed: EncryptedReasoningContent = serde_json::from_str(&json).unwrap();
288 assert_eq!(parsed.model, model);
289 assert_eq!(parsed.content, "encrypted-data");
290 }
291
292 #[test]
293 fn assistant_reasoning_serde_roundtrip() {
294 let model = make_model();
295 let r = AssistantReasoning {
296 summary_text: Some("thought".to_string()),
297 encrypted_content: Some(EncryptedReasoningContent {
298 id: "r_test".to_string(),
299 model,
300 content: "blob".to_string(),
301 }),
302 };
303 let json = serde_json::to_string(&r).unwrap();
304 let parsed: AssistantReasoning = serde_json::from_str(&json).unwrap();
305 assert_eq!(parsed, r);
306 }
307
308 #[test]
309 fn assistant_reasoning_serde_empty_roundtrip() {
310 let r = AssistantReasoning::default();
311 let json = serde_json::to_string(&r).unwrap();
312 assert_eq!(json, "{}");
313 let parsed: AssistantReasoning = serde_json::from_str(&json).unwrap();
314 assert_eq!(parsed, r);
315 }
316
317 #[test]
318 fn chat_message_assistant_serde_roundtrip_with_reasoning() {
319 let model = make_model();
320 let msg = ChatMessage::Assistant {
321 message_id: MessageId::new(),
322 content: "response".to_string(),
323 reasoning: AssistantReasoning {
324 summary_text: Some("plan".to_string()),
325 encrypted_content: Some(EncryptedReasoningContent {
326 id: "r_test".to_string(),
327 model,
328 content: "enc".to_string(),
329 }),
330 },
331 timestamp: IsoString::now(),
332 tool_calls: vec![],
333 };
334 let json = serde_json::to_string(&msg).unwrap();
335 let parsed: ChatMessage = serde_json::from_str(&json).unwrap();
336 assert_eq!(parsed, msg);
337 }
338
339 #[test]
340 fn estimated_bytes_includes_encrypted_content() {
341 let model = make_model();
342 let msg_with = ChatMessage::Assistant {
343 message_id: MessageId::new(),
344 content: "hi".to_string(),
345 reasoning: AssistantReasoning {
346 summary_text: Some("think".to_string()),
347 encrypted_content: Some(EncryptedReasoningContent {
348 id: "r_test".to_string(),
349 model,
350 content: "x".repeat(100),
351 }),
352 },
353 timestamp: IsoString::now(),
354 tool_calls: vec![],
355 };
356 let msg_without = ChatMessage::Assistant {
357 message_id: MessageId::new(),
358 content: "hi".to_string(),
359 reasoning: AssistantReasoning { summary_text: Some("think".to_string()), encrypted_content: None },
360 timestamp: IsoString::now(),
361 tool_calls: vec![],
362 };
363 assert!(msg_with.estimated_bytes() > msg_without.estimated_bytes());
364 }
365}