1use crate::types::{ContentBlock, LlmMessage, StopReason};
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum DisplayRole {
11 User,
12 Assistant,
13 ToolResult,
14 Error,
15 System,
16}
17
18#[derive(Debug, Clone)]
24pub struct CoreDisplayMessage {
25 pub role: DisplayRole,
26 pub content: String,
27 pub thinking: Option<String>,
28}
29
30pub trait IntoDisplayMessages {
32 fn to_display_messages(&self) -> Vec<CoreDisplayMessage>;
33}
34
35impl IntoDisplayMessages for LlmMessage {
36 fn to_display_messages(&self) -> Vec<CoreDisplayMessage> {
37 match self {
38 Self::User(user) => {
39 vec![CoreDisplayMessage {
40 role: DisplayRole::User,
41 content: ContentBlock::extract_text(&user.content),
42 thinking: None,
43 }]
44 }
45 Self::Assistant(assistant) => {
46 let mut text_parts = Vec::new();
47 let mut thinking_parts = Vec::new();
48 for block in &assistant.content {
49 match block {
50 ContentBlock::Text { text } => text_parts.push(text.as_str()),
51 ContentBlock::Thinking { thinking, .. } => {
52 thinking_parts.push(thinking.as_str());
53 }
54 _ => {}
55 }
56 }
57
58 let content = if !text_parts.is_empty() {
59 text_parts.join("")
60 } else if assistant.stop_reason == StopReason::Error {
61 assistant.error_message.clone().unwrap_or_default()
62 } else {
63 String::new()
64 };
65
66 let thinking = if thinking_parts.is_empty() {
67 None
68 } else {
69 Some(thinking_parts.join(""))
70 };
71
72 let role = if assistant.stop_reason == StopReason::Error {
73 DisplayRole::Error
74 } else {
75 DisplayRole::Assistant
76 };
77
78 vec![CoreDisplayMessage {
79 role,
80 content,
81 thinking,
82 }]
83 }
84 Self::ToolResult(tool_result) => {
85 let content = ContentBlock::extract_text(&tool_result.content);
86 if content.is_empty() {
87 return vec![];
88 }
89 let role = if tool_result.is_error {
90 DisplayRole::Error
91 } else {
92 DisplayRole::ToolResult
93 };
94 vec![CoreDisplayMessage {
95 role,
96 content,
97 thinking: None,
98 }]
99 }
100 }
101 }
102}
103
104impl IntoDisplayMessages for [LlmMessage] {
105 fn to_display_messages(&self) -> Vec<CoreDisplayMessage> {
106 self.iter()
107 .flat_map(IntoDisplayMessages::to_display_messages)
108 .collect()
109 }
110}
111
112impl IntoDisplayMessages for Vec<LlmMessage> {
113 fn to_display_messages(&self) -> Vec<CoreDisplayMessage> {
114 self.as_slice().to_display_messages()
115 }
116}
117
118#[cfg(test)]
119mod tests {
120 use super::*;
121 use crate::types::{AssistantMessage, Cost, Usage, UserMessage};
122
123 #[test]
124 fn user_message_to_display() {
125 let msg = LlmMessage::User(UserMessage {
126 content: vec![ContentBlock::Text {
127 text: "hello".to_string(),
128 }],
129 timestamp: 0,
130 cache_hint: None,
131 });
132 let display = msg.to_display_messages();
133 assert_eq!(display.len(), 1);
134 assert_eq!(display[0].role, DisplayRole::User);
135 assert_eq!(display[0].content, "hello");
136 assert!(display[0].thinking.is_none());
137 }
138
139 #[test]
140 fn assistant_message_with_thinking() {
141 let msg = LlmMessage::Assistant(AssistantMessage {
142 content: vec![
143 ContentBlock::Thinking {
144 thinking: "reasoning".to_string(),
145 signature: None,
146 },
147 ContentBlock::Text {
148 text: "answer".to_string(),
149 },
150 ],
151 provider: String::new(),
152 model_id: String::new(),
153 usage: Usage::default(),
154 cost: Cost::default(),
155 stop_reason: StopReason::Stop,
156 error_message: None,
157 error_kind: None,
158 timestamp: 0,
159 cache_hint: None,
160 });
161 let display = msg.to_display_messages();
162 assert_eq!(display.len(), 1);
163 assert_eq!(display[0].role, DisplayRole::Assistant);
164 assert_eq!(display[0].content, "answer");
165 assert_eq!(display[0].thinking.as_deref(), Some("reasoning"));
166 }
167
168 #[test]
169 fn assistant_error_message() {
170 let msg = LlmMessage::Assistant(AssistantMessage {
171 content: vec![],
172 provider: String::new(),
173 model_id: String::new(),
174 usage: Usage::default(),
175 cost: Cost::default(),
176 stop_reason: StopReason::Error,
177 error_message: Some("something broke".to_string()),
178 error_kind: None,
179 timestamp: 0,
180 cache_hint: None,
181 });
182 let display = msg.to_display_messages();
183 assert_eq!(display.len(), 1);
184 assert_eq!(display[0].role, DisplayRole::Error);
185 assert_eq!(display[0].content, "something broke");
186 }
187
188 #[test]
189 fn empty_tool_result_produces_no_messages() {
190 let msg = LlmMessage::ToolResult(crate::types::ToolResultMessage {
191 tool_call_id: "tc1".to_string(),
192 content: vec![],
193 is_error: false,
194 timestamp: 0,
195 details: serde_json::Value::Null,
196 cache_hint: None,
197 });
198 let display = msg.to_display_messages();
199 assert!(display.is_empty());
200 }
201}