Skip to main content

llm/
chat_message.rs

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    /// Joins all text blocks with newlines, ignoring non-text content.
44    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    /// Returns a `data:{mime};base64,{data}` URI for image/audio blocks, `None` for text.
56    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/// Opaque encrypted reasoning content from an LLM response.
67///
68/// This is model-specific: encrypted content from one model cannot be replayed
69/// to a different model. Use [`Context::filter_encrypted_reasoning`](crate::Context::filter_encrypted_reasoning)
70/// to strip content that doesn't match the target model.
71#[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/// Reasoning metadata from an assistant response.
80///
81/// Contains an optional human-readable summary and optional encrypted content
82/// that can be replayed to the same model in future turns.
83#[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    /// A compacted summary of previous conversation history.
127    /// This replaces multiple messages with a structured summary to reduce context usage.
128    Summary {
129        content: String,
130        timestamp: IsoString,
131        /// Number of messages that were compacted into this summary
132        messages_compacted: usize,
133    },
134}
135
136impl ChatMessage {
137    /// Returns true if this message is a tool call result
138    pub fn is_tool_result(&self) -> bool {
139        matches!(self, ChatMessage::ToolCallResult(_))
140    }
141
142    /// Returns true if this message is a system prompt
143    pub fn is_system(&self) -> bool {
144        matches!(self, ChatMessage::System { .. })
145    }
146
147    /// Returns true if this message is a compacted summary
148    pub fn is_summary(&self) -> bool {
149        matches!(self, ChatMessage::Summary { .. })
150    }
151
152    /// Rough byte-size estimate of the message content for pre-flight context checks.
153    /// Not meant to be exact — just close enough to detect overflow before calling the LLM.
154    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    /// Returns the timestamp of this message, if it has one
174    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}