Skip to main content

llm/
chat_message.rs

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    /// Joins all text blocks with newlines, ignoring non-text content.
45    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    /// Returns a `data:{mime};base64,{data}` URI for image/audio blocks, `None` for text.
57    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/// Opaque encrypted reasoning content from an LLM response.
68///
69/// This is model-specific: encrypted content from one model cannot be replayed
70/// to a different model. Use [`Context::filter_encrypted_reasoning`](crate::Context::filter_encrypted_reasoning)
71/// to strip content that doesn't match the target model.
72#[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/// Reasoning metadata from an assistant response.
81///
82/// Contains an optional human-readable summary and optional encrypted content
83/// that can be replayed to the same model in future turns.
84#[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    /// A compacted summary of previous conversation history.
131    /// This replaces multiple messages with a structured summary to reduce context usage.
132    Summary {
133        message_id: MessageId,
134        content: String,
135        timestamp: IsoString,
136        /// Number of messages that were compacted into this summary
137        messages_compacted: usize,
138    },
139}
140
141impl ChatMessage {
142    /// A user message carrying a single block of text, stamped now.
143    pub fn user(text: impl Into<String>) -> Self {
144        Self::user_with_id(MessageId::new(), vec![ContentBlock::text(text)])
145    }
146
147    /// A system prompt, stamped now.
148    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    /// Identity of every message that can be addressed by clients; system prompts have none.
157    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    /// Returns true if this message is a tool call result
170    pub fn is_tool_result(&self) -> bool {
171        matches!(self, ChatMessage::ToolCallResult(_))
172    }
173
174    /// Returns true if this message is a system prompt
175    pub fn is_system(&self) -> bool {
176        matches!(self, ChatMessage::System { .. })
177    }
178
179    /// Returns true if this message is a compacted summary
180    pub fn is_summary(&self) -> bool {
181        matches!(self, ChatMessage::Summary { .. })
182    }
183
184    /// Rough byte-size estimate of the message content for pre-flight context checks.
185    /// Not meant to be exact — just close enough to detect overflow before calling the LLM.
186    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    /// Returns the timestamp of this message, if it has one
206    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    pub(crate) fn as_user_note(&self) -> Option<String> {
218        match self {
219            ChatMessage::Summary { content, .. } => Some(format!("[Previous conversation handoff]\n\n{content}")),
220            ChatMessage::Error { message, .. } => Some(format!("Error: {message}")),
221            _ => None,
222        }
223    }
224}
225
226fn serialize_llm_model<S: serde::Serializer>(model: &LlmModel, s: S) -> Result<S::Ok, S::Error> {
227    s.serialize_str(&model.to_string())
228}
229
230fn deserialize_llm_model<'de, D: serde::Deserializer<'de>>(d: D) -> Result<LlmModel, D::Error> {
231    let s = String::deserialize(d)?;
232    s.parse::<LlmModel>().map_err(serde::de::Error::custom)
233}
234
235#[cfg(test)]
236mod tests {
237    use super::*;
238
239    fn make_model() -> LlmModel {
240        "anthropic:claude-opus-4-6".parse().unwrap()
241    }
242
243    #[test]
244    fn assistant_reasoning_is_empty_when_default() {
245        let r = AssistantReasoning::default();
246        assert!(r.is_empty());
247    }
248
249    #[test]
250    fn assistant_reasoning_not_empty_with_summary() {
251        let r = AssistantReasoning::from_parts("thinking".to_string(), None);
252        assert!(!r.is_empty());
253    }
254
255    #[test]
256    fn assistant_reasoning_not_empty_with_encrypted() {
257        let r = AssistantReasoning {
258            summary_text: None,
259            encrypted_content: Some(EncryptedReasoningContent {
260                id: "r_test".to_string(),
261                model: make_model(),
262                content: "blob".to_string(),
263            }),
264        };
265        assert!(!r.is_empty());
266    }
267
268    #[test]
269    fn from_parts_empty_summary_is_none() {
270        let r = AssistantReasoning::from_parts(String::new(), None);
271        assert!(r.summary_text.is_none());
272        assert!(r.is_empty());
273    }
274
275    #[test]
276    fn first_text_returns_first_non_empty_text_block() {
277        let parts = vec![
278            ContentBlock::Image { data: "a".to_string(), mime_type: "image/png".to_string() },
279            ContentBlock::text(" "),
280            ContentBlock::text("hello"),
281        ];
282
283        assert_eq!(ContentBlock::first_text(&parts), Some("hello"));
284    }
285
286    #[test]
287    fn encrypted_reasoning_content_serde_roundtrip() {
288        let model = make_model();
289        let ec = EncryptedReasoningContent {
290            id: "r_test".to_string(),
291            model: model.clone(),
292            content: "encrypted-data".to_string(),
293        };
294        let json = serde_json::to_string(&ec).unwrap();
295        let parsed: EncryptedReasoningContent = serde_json::from_str(&json).unwrap();
296        assert_eq!(parsed.model, model);
297        assert_eq!(parsed.content, "encrypted-data");
298    }
299
300    #[test]
301    fn assistant_reasoning_serde_roundtrip() {
302        let model = make_model();
303        let r = AssistantReasoning {
304            summary_text: Some("thought".to_string()),
305            encrypted_content: Some(EncryptedReasoningContent {
306                id: "r_test".to_string(),
307                model,
308                content: "blob".to_string(),
309            }),
310        };
311        let json = serde_json::to_string(&r).unwrap();
312        let parsed: AssistantReasoning = serde_json::from_str(&json).unwrap();
313        assert_eq!(parsed, r);
314    }
315
316    #[test]
317    fn assistant_reasoning_serde_empty_roundtrip() {
318        let r = AssistantReasoning::default();
319        let json = serde_json::to_string(&r).unwrap();
320        assert_eq!(json, "{}");
321        let parsed: AssistantReasoning = serde_json::from_str(&json).unwrap();
322        assert_eq!(parsed, r);
323    }
324
325    #[test]
326    fn chat_message_assistant_serde_roundtrip_with_reasoning() {
327        let model = make_model();
328        let msg = ChatMessage::Assistant {
329            message_id: MessageId::new(),
330            content: "response".to_string(),
331            reasoning: AssistantReasoning {
332                summary_text: Some("plan".to_string()),
333                encrypted_content: Some(EncryptedReasoningContent {
334                    id: "r_test".to_string(),
335                    model,
336                    content: "enc".to_string(),
337                }),
338            },
339            timestamp: IsoString::now(),
340            tool_calls: vec![],
341        };
342        let json = serde_json::to_string(&msg).unwrap();
343        let parsed: ChatMessage = serde_json::from_str(&json).unwrap();
344        assert_eq!(parsed, msg);
345    }
346
347    #[test]
348    fn estimated_bytes_includes_encrypted_content() {
349        let model = make_model();
350        let msg_with = ChatMessage::Assistant {
351            message_id: MessageId::new(),
352            content: "hi".to_string(),
353            reasoning: AssistantReasoning {
354                summary_text: Some("think".to_string()),
355                encrypted_content: Some(EncryptedReasoningContent {
356                    id: "r_test".to_string(),
357                    model,
358                    content: "x".repeat(100),
359                }),
360            },
361            timestamp: IsoString::now(),
362            tool_calls: vec![],
363        };
364        let msg_without = ChatMessage::Assistant {
365            message_id: MessageId::new(),
366            content: "hi".to_string(),
367            reasoning: AssistantReasoning { summary_text: Some("think".to_string()), encrypted_content: None },
368            timestamp: IsoString::now(),
369            tool_calls: vec![],
370        };
371        assert!(msg_with.estimated_bytes() > msg_without.estimated_bytes());
372    }
373}