1use serde::{Deserialize, Serialize};
4
5#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
7pub struct TextBlock {
8 pub text: String,
9}
10
11#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
13pub struct ThinkingBlock {
14 pub thinking: String,
15 pub redacted: Option<String>,
17}
18
19#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
21pub struct ImageSource {
22 pub data: String,
24 pub media_type: String,
26}
27
28#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
30pub struct ToolCall {
31 pub id: String,
32 pub name: String,
33 pub arguments: serde_json::Value,
34}
35
36#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
39#[serde(tag = "type", rename_all = "snake_case")]
40pub enum ContentBlock {
41 Text(TextBlock),
42 Thinking(ThinkingBlock),
43 Image { source: ImageSource },
44 ToolCall(ToolCall),
45}
46
47impl ContentBlock {
48 pub fn text(s: String) -> Self {
49 ContentBlock::Text(TextBlock { text: s })
50 }
51
52 pub fn as_text(&self) -> Option<&str> {
53 match self {
54 ContentBlock::Text(block) => Some(&block.text),
55 _ => None,
56 }
57 }
58}
59
60#[derive(Debug, Clone, Serialize, Deserialize)]
62#[serde(tag = "type", rename_all = "snake_case")]
63pub enum Message {
64 System {
65 content: Vec<ContentBlock>,
66 },
67 User {
68 content: Vec<ContentBlock>,
69 },
70 Assistant {
71 content: Vec<ContentBlock>,
72 },
73 ToolResult {
74 tool_call_id: String,
75 content: Vec<ContentBlock>,
76 },
77}
78
79impl Message {
80 pub fn role(&self) -> &str {
83 match self {
84 Message::System { .. } => "system",
85 Message::User { .. } => "user",
86 Message::Assistant { .. } => "assistant",
87 Message::ToolResult { .. } => "tool_result",
88 }
89 }
90
91 pub fn content(&self) -> &Vec<ContentBlock> {
93 match self {
94 Message::System { content }
95 | Message::User { content }
96 | Message::Assistant { content }
97 | Message::ToolResult { content, .. } => content,
98 }
99 }
100
101 pub fn extract_text(&self) -> String {
103 match self {
104 Message::System { content } => Self::join_text(content),
105 Message::User { content } => Self::join_text(content),
106 Message::Assistant { content } => Self::join_text(content),
107 Message::ToolResult { content, .. } => Self::join_text(content),
108 }
109 }
110
111 fn join_text(blocks: &[ContentBlock]) -> String {
112 blocks
113 .iter()
114 .filter_map(|b| b.as_text().map(|s| s.to_string()))
115 .collect::<Vec<_>>()
116 .join("")
117 }
118
119 pub fn extract_tool_calls(&self) -> Vec<ToolCall> {
121 match self {
122 Message::Assistant { content } => content
123 .iter()
124 .filter_map(|b| {
125 if let ContentBlock::ToolCall(tc) = b {
126 Some(tc.clone())
127 } else {
128 None
129 }
130 })
131 .collect(),
132 _ => Vec::new(),
133 }
134 }
135}
136
137pub fn text_block(s: String) -> Vec<ContentBlock> {
139 vec![ContentBlock::text(s)]
140}
141
142#[cfg(test)]
143mod tests {
144 use super::*;
145
146 #[test]
147 fn test_content_block_text() {
148 let block = ContentBlock::text("hello".to_string());
149 assert_eq!(block.as_text(), Some("hello"));
150 }
151
152 #[test]
153 fn test_content_block_tool_call_no_as_text() {
154 let block = ContentBlock::ToolCall(ToolCall {
155 id: "1".into(),
156 name: "test".into(),
157 arguments: serde_json::json!({}),
158 });
159 assert_eq!(block.as_text(), None);
160 }
161
162 #[test]
163 fn test_message_extract_text() {
164 let msg = Message::User {
165 content: text_block("hello world".to_string()),
166 };
167 assert_eq!(msg.extract_text(), "hello world");
168 }
169
170 #[test]
171 fn test_message_extract_tool_calls() {
172 let tc = ToolCall {
173 id: "1".into(),
174 name: "test".into(),
175 arguments: serde_json::json!({}),
176 };
177 let msg = Message::Assistant {
178 content: vec![ContentBlock::ToolCall(tc.clone())],
179 };
180 let calls = msg.extract_tool_calls();
181 assert_eq!(calls.len(), 1);
182 assert_eq!(calls[0].name, "test");
183 }
184}