mermaid_cli/models/
types.rs

1use crate::agents::ActionDisplay;
2use serde::{Deserialize, Serialize};
3
4/// Represents a chat message
5#[derive(Debug, Clone, Serialize, Deserialize)]
6pub struct ChatMessage {
7    pub role: MessageRole,
8    pub content: String,
9    pub timestamp: chrono::DateTime<chrono::Local>,
10    /// Actions performed during this message (for display purposes)
11    #[serde(default)]
12    pub actions: Vec<ActionDisplay>,
13    /// Thinking/reasoning content (for models that expose their thought process)
14    #[serde(default)]
15    pub thinking: Option<String>,
16    /// Base64-encoded images/PDFs for multimodal models
17    #[serde(default)]
18    pub images: Option<Vec<String>>,
19    /// Tool calls from the model (Ollama native function calling)
20    #[serde(default)]
21    pub tool_calls: Option<Vec<crate::models::tool_call::ToolCall>>,
22    /// Tool call ID for tool result messages (OpenAI-compatible format)
23    /// This links the tool result back to the original tool_call from the assistant
24    #[serde(default)]
25    pub tool_call_id: Option<String>,
26    /// Tool name for tool result messages (required by Ollama API)
27    /// This tells the model which function's result is being returned
28    #[serde(default)]
29    pub tool_name: Option<String>,
30    /// Anthropic thinking-block signature — encrypted server state that
31    /// MUST round-trip back into the next request when extended thinking
32    /// is enabled, or the API returns 400 `invalid_request_error`. Set
33    /// only by the Anthropic adapter; other adapters leave it `None`
34    /// and other providers ignore it on the wire.
35    #[serde(default)]
36    pub thinking_signature: Option<String>,
37}
38
39impl ChatMessage {
40    /// Create a user message
41    pub fn user(content: impl Into<String>) -> Self {
42        Self::new(MessageRole::User, content.into())
43    }
44
45    /// Create an assistant message
46    pub fn assistant(content: impl Into<String>) -> Self {
47        Self::new(MessageRole::Assistant, content.into())
48    }
49
50    /// Create a system message
51    pub fn system(content: impl Into<String>) -> Self {
52        Self::new(MessageRole::System, content.into())
53    }
54
55    /// Create a tool result message
56    pub fn tool(
57        tool_call_id: impl Into<String>,
58        tool_name: impl Into<String>,
59        content: impl Into<String>,
60    ) -> Self {
61        Self {
62            role: MessageRole::Tool,
63            content: content.into(),
64            timestamp: chrono::Local::now(),
65            actions: Vec::new(),
66            thinking: None,
67            images: None,
68            tool_calls: None,
69            tool_call_id: Some(tool_call_id.into()),
70            tool_name: Some(tool_name.into()),
71            thinking_signature: None,
72        }
73    }
74
75    /// Base constructor with role and content
76    fn new(role: MessageRole, content: String) -> Self {
77        Self {
78            role,
79            content,
80            timestamp: chrono::Local::now(),
81            actions: Vec::new(),
82            thinking: None,
83            images: None,
84            tool_calls: None,
85            tool_call_id: None,
86            tool_name: None,
87            thinking_signature: None,
88        }
89    }
90
91    /// Builder: attach images
92    pub fn with_images(mut self, images: Vec<String>) -> Self {
93        self.images = Some(images);
94        self
95    }
96
97    /// Builder: attach tool calls
98    pub fn with_tool_calls(mut self, tool_calls: Vec<crate::models::tool_call::ToolCall>) -> Self {
99        self.tool_calls = if tool_calls.is_empty() {
100            None
101        } else {
102            Some(tool_calls)
103        };
104        self
105    }
106
107    /// Builder: attach an Anthropic thinking signature. Used by the
108    /// Anthropic adapter when committing assistant messages so the
109    /// signature can round-trip on the next request.
110    pub fn with_thinking_signature(mut self, signature: impl Into<String>) -> Self {
111        self.thinking_signature = Some(signature.into());
112        self
113    }
114
115    /// Extract thinking blocks from message content.
116    /// Returns `(thinking_content, answer_content)`.
117    ///
118    /// Performs a single `find` for the start marker; the previous version
119    /// scanned twice (`contains` + `find`) and called `find("Thinking...")`
120    /// again inside the if-let-chain.
121    ///
122    /// Safety: `str::find()` returns byte offsets. The markers `"Thinking..."`
123    /// and `"...done thinking."` are pure ASCII, so adding their `.len()`
124    /// always lands on a valid UTF-8 char boundary.
125    pub fn extract_thinking(text: &str) -> (Option<String>, String) {
126        let Some(thinking_start) = text.find("Thinking...") else {
127            return (None, text.to_string());
128        };
129        let content_start = thinking_start + "Thinking...".len();
130
131        if let Some(thinking_end) = text.find("...done thinking.") {
132            let thinking_text = text[content_start..thinking_end].trim().to_string();
133            let answer_start = thinking_end + "...done thinking.".len();
134            let answer_text = text[answer_start..].trim().to_string();
135            return (Some(thinking_text), answer_text);
136        }
137
138        // Start marker without end marker — thinking is still in progress.
139        let thinking_text = text[content_start..].trim().to_string();
140        (Some(thinking_text), String::new())
141    }
142}
143
144#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
145pub enum MessageRole {
146    User,
147    Assistant,
148    System,
149    /// Tool result message (OpenAI-compatible format for function calling)
150    Tool,
151}
152
153/// Response from a model
154#[derive(Debug, Clone)]
155pub struct ModelResponse {
156    /// The actual response text
157    pub content: String,
158    /// Usage statistics if available
159    pub usage: Option<TokenUsage>,
160    /// Model that generated the response
161    pub model_name: String,
162    /// Thinking/reasoning content (for models that expose their thought process)
163    pub thinking: Option<String>,
164    /// Tool calls from the model (Ollama native function calling)
165    pub tool_calls: Option<Vec<crate::models::tool_call::ToolCall>>,
166    /// Anthropic thinking-block signature (encrypted server state). Set
167    /// only by `AnthropicAdapter`; other adapters leave it `None`. The
168    /// agent loop's commit step copies this onto the resulting assistant
169    /// `ChatMessage::thinking_signature` so it round-trips on the next
170    /// turn — without this, multi-turn Claude conversations with
171    /// extended thinking 400 with `invalid_request_error`.
172    pub thinking_signature: Option<String>,
173}
174
175/// Token usage statistics
176#[derive(Debug, Clone)]
177pub struct TokenUsage {
178    pub prompt_tokens: usize,
179    pub completion_tokens: usize,
180    pub total_tokens: usize,
181}
182
183#[cfg(test)]
184mod tests {
185    use super::*;
186
187    #[test]
188    fn test_message_role_equality() {
189        let user1 = MessageRole::User;
190        let user2 = MessageRole::User;
191        let assistant = MessageRole::Assistant;
192
193        assert_eq!(user1, user2, "User roles should be equal");
194        assert_ne!(user1, assistant, "Different roles should not be equal");
195    }
196
197    #[test]
198    fn test_chat_message_constructors() {
199        let user = ChatMessage::user("Hello!");
200        assert_eq!(user.role, MessageRole::User);
201        assert_eq!(user.content, "Hello!");
202        assert!(user.tool_calls.is_none());
203
204        let assistant = ChatMessage::assistant("Hi there");
205        assert_eq!(assistant.role, MessageRole::Assistant);
206
207        let system = ChatMessage::system("You are helpful");
208        assert_eq!(system.role, MessageRole::System);
209
210        let tool = ChatMessage::tool("call_1", "read_file", "file contents");
211        assert_eq!(tool.role, MessageRole::Tool);
212        assert_eq!(tool.tool_call_id, Some("call_1".to_string()));
213        assert_eq!(tool.tool_name, Some("read_file".to_string()));
214    }
215
216    #[test]
217    fn test_chat_message_builders() {
218        let msg = ChatMessage::user("test").with_images(vec!["base64data".to_string()]);
219        assert_eq!(msg.images, Some(vec!["base64data".to_string()]));
220    }
221
222    #[test]
223    fn test_token_usage_structure() {
224        let usage = TokenUsage {
225            prompt_tokens: 100,
226            completion_tokens: 50,
227            total_tokens: 150,
228        };
229
230        assert_eq!(usage.prompt_tokens, 100);
231        assert_eq!(usage.completion_tokens, 50);
232        assert_eq!(usage.total_tokens, 150);
233    }
234
235    // --- extract_thinking ---
236
237    #[test]
238    fn extract_thinking_no_marker_returns_text_unchanged() {
239        let (thinking, answer) = ChatMessage::extract_thinking("just a plain answer");
240        assert_eq!(thinking, None);
241        assert_eq!(answer, "just a plain answer");
242    }
243
244    #[test]
245    fn extract_thinking_complete_block() {
246        let raw = "Thinking...\n  reasoning here\n...done thinking.\n\nFinal answer";
247        let (thinking, answer) = ChatMessage::extract_thinking(raw);
248        assert_eq!(thinking.as_deref(), Some("reasoning here"));
249        assert_eq!(answer, "Final answer");
250    }
251
252    #[test]
253    fn thinking_signature_round_trips_through_serde() {
254        // Anthropic encrypted server state — must survive
255        // serialize/deserialize so saved conversations resume cleanly.
256        let msg = ChatMessage::assistant("Step 3 lives.")
257            .with_thinking_signature("sig_abc123_encrypted_blob");
258        let json = serde_json::to_string(&msg).expect("serialize");
259        let back: ChatMessage = serde_json::from_str(&json).expect("deserialize");
260        assert_eq!(
261            back.thinking_signature.as_deref(),
262            Some("sig_abc123_encrypted_blob")
263        );
264        assert_eq!(back.content, "Step 3 lives.");
265    }
266
267    #[test]
268    fn thinking_signature_defaults_to_none() {
269        // Backward compat: messages saved before Step 3 won't have the
270        // field. Serde default kicks in — None — and deserialize
271        // succeeds without errors.
272        let pre_step3_json = r#"{
273            "role": "Assistant",
274            "content": "hello",
275            "timestamp": "2026-04-16T12:00:00-04:00"
276        }"#;
277        let msg: ChatMessage = serde_json::from_str(pre_step3_json).expect("backward compat");
278        assert!(msg.thinking_signature.is_none());
279    }
280
281    #[test]
282    fn extract_thinking_in_progress_no_end_marker() {
283        let raw = "Thinking...\n  partial reasoning so far";
284        let (thinking, answer) = ChatMessage::extract_thinking(raw);
285        assert_eq!(thinking.as_deref(), Some("partial reasoning so far"));
286        assert_eq!(answer, "");
287    }
288
289    #[test]
290    fn test_model_response_creation() {
291        let usage = TokenUsage {
292            prompt_tokens: 100,
293            completion_tokens: 50,
294            total_tokens: 150,
295        };
296
297        let response = ModelResponse {
298            content: "Hello, world!".to_string(),
299            usage: Some(usage),
300            model_name: "ollama/tinyllama".to_string(),
301            thinking: None,
302            tool_calls: None,
303            thinking_signature: None,
304        };
305
306        assert_eq!(response.content, "Hello, world!");
307        assert!(response.usage.is_some());
308        assert_eq!(response.model_name, "ollama/tinyllama");
309        assert_eq!(response.usage.unwrap().total_tokens, 150);
310        assert!(response.tool_calls.is_none());
311    }
312}
mermaid_cli/models/types.rs

mermaid_cli/models/
types.rs
