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::{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    /// 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#[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    /// A compacted summary of previous conversation history.
94    /// This replaces multiple messages with a structured summary to reduce context usage.
95    Summary {
96        message_id: MessageId,
97        content: String,
98        timestamp: IsoString,
99        /// Number of messages that were compacted into this summary
100        messages_compacted: usize,
101    },
102}
103
104impl ChatMessage {
105    /// A user message carrying a single block of text, stamped now.
106    pub fn user(text: impl Into<String>) -> Self {
107        Self::user_with_id(MessageId::new(), vec![ContentBlock::text(text)])
108    }
109
110    /// A system prompt, stamped now.
111    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    /// Identity of every message that can be addressed by clients; system prompts have none.
120    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    /// Returns true if this message is a tool call result
133    pub fn is_tool_result(&self) -> bool {
134        matches!(self, ChatMessage::ToolCallResult(_))
135    }
136
137    /// Returns true if this message is a system prompt
138    pub fn is_system(&self) -> bool {
139        matches!(self, ChatMessage::System { .. })
140    }
141
142    /// Returns true if this message is a compacted summary
143    pub fn is_summary(&self) -> bool {
144        matches!(self, ChatMessage::Summary { .. })
145    }
146
147    /// Rough byte-size estimate of the message content for pre-flight context checks.
148    /// Not meant to be exact — just close enough to detect overflow before calling the LLM.
149    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    /// Returns the timestamp of this message, if it has one
164    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}