1use serde::{Deserialize, Serialize};
2use serde_json::Value;
3
4use crate::types::{ModelId, ToolId, ToolName};
5
6#[derive(Serialize, Deserialize, Debug, Clone)]
7pub struct Message {
8 pub role: String,
9 pub content: Vec<ContentBlock>,
10}
11
12impl Message {
13 pub fn user(text: impl Into<String>) -> Self {
14 Self {
15 role: "user".to_string(),
16 content: vec![ContentBlock::Text { text: text.into() }],
17 }
18 }
19
20 pub fn assistant(content: Vec<ContentBlock>) -> Self {
21 Self {
22 role: "assistant".to_string(),
23 content,
24 }
25 }
26
27 pub fn tool_results(results: Vec<ContentBlock>) -> Self {
28 Self {
29 role: "user".to_string(),
30 content: results,
31 }
32 }
33}
34
35#[derive(Serialize, Deserialize, Debug, Clone)]
36#[serde(tag = "type", rename_all = "snake_case")]
37pub enum ContentBlock {
38 Text {
39 text: String,
40 },
41 ToolUse {
42 id: ToolId,
43 name: ToolName,
44 input: Value,
45 },
46 ToolResult {
47 tool_use_id: ToolId,
48 content: String,
49 },
50}
51
52#[allow(dead_code)]
53#[derive(Serialize)]
54pub struct ApiRequest {
55 pub model: ModelId,
56 pub max_tokens: u32,
57 pub system: String,
58 pub messages: Vec<Message>,
59 pub tools: Vec<ToolDefinition>,
60}
61
62#[derive(Serialize, Clone, Debug)]
63pub struct ToolDefinition {
64 pub name: String,
65 pub description: String,
66 pub input_schema: Value,
67}
68
69#[cfg(test)]
70mod tests {
71 use super::*;
72 use serde_json::json;
73
74 #[test]
75 fn message_user() {
76 let msg = Message::user("Hello world");
77 assert_eq!(msg.role, "user");
78 assert_eq!(msg.content.len(), 1);
79 match &msg.content[0] {
80 ContentBlock::Text { text } => assert_eq!(text, "Hello world"),
81 other => panic!("Expected Text content block, got: {other:?}"),
82 }
83 }
84
85 #[test]
86 fn message_assistant() {
87 let content = vec![
88 ContentBlock::Text {
89 text: "Response".to_string(),
90 },
91 ContentBlock::ToolUse {
92 id: ToolId::new("tool_1"),
93 name: ToolName::new("read"),
94 input: json!({"path": "/tmp/file"}),
95 },
96 ];
97 let msg = Message::assistant(content);
98 assert_eq!(msg.role, "assistant");
99 assert_eq!(msg.content.len(), 2);
100 }
101
102 #[test]
103 fn message_tool_results() {
104 let results = vec![ContentBlock::ToolResult {
105 tool_use_id: ToolId::new("tool_1"),
106 content: "File contents".to_string(),
107 }];
108 let msg = Message::tool_results(results);
109 assert_eq!(msg.role, "user");
110 assert_eq!(msg.content.len(), 1);
111 }
112
113 #[test]
114 fn content_block_serde_roundtrip() {
115 let block = ContentBlock::Text {
116 text: "Hello".to_string(),
117 };
118 let json = serde_json::to_value(&block).expect("serialize");
119 assert_eq!(json["type"], "text");
120 assert_eq!(json["text"], "Hello");
121
122 let json_str = json!({"type": "text", "text": "Hello"});
123 let block: ContentBlock = serde_json::from_value(json_str).expect("deserialize");
124 match block {
125 ContentBlock::Text { text } => assert_eq!(text, "Hello"),
126 other => panic!("Expected Text block, got: {other:?}"),
127 }
128 }
129
130 #[test]
131 fn message_user_string() {
132 let msg = Message::user(String::from("Test"));
133 assert_eq!(msg.role, "user");
134 match &msg.content[0] {
135 ContentBlock::Text { text } => assert_eq!(text, "Test"),
136 other => panic!("Expected Text content block, got: {other:?}"),
137 }
138 }
139
140 #[test]
141 fn message_tool_results_detailed() {
142 let results = vec![ContentBlock::ToolResult {
143 tool_use_id: ToolId::new("tool_1"),
144 content: "File contents".to_string(),
145 }];
146 let msg = Message::tool_results(results);
147 assert_eq!(msg.role, "user");
148 assert_eq!(msg.content.len(), 1);
149 match &msg.content[0] {
150 ContentBlock::ToolResult {
151 tool_use_id,
152 content,
153 } => {
154 assert_eq!(tool_use_id.as_str(), "tool_1");
155 assert_eq!(content, "File contents");
156 }
157 other => panic!("Expected ToolResult content block, got: {other:?}"),
158 }
159 }
160
161 #[test]
162 fn content_block_tool_use_serialization() {
163 let block = ContentBlock::ToolUse {
164 id: ToolId::new("123"),
165 name: ToolName::new("bash"),
166 input: json!({"command": "ls"}),
167 };
168 let json = serde_json::to_value(&block).expect("serialize ContentBlock");
169 assert_eq!(json["type"], "tool_use");
170 assert_eq!(json["id"], "123");
171 assert_eq!(json["name"], "bash");
172 assert_eq!(json["input"]["command"], "ls");
173 }
174
175 #[test]
176 fn content_block_tool_result_serialization() {
177 let block = ContentBlock::ToolResult {
178 tool_use_id: ToolId::new("123"),
179 content: "output".to_string(),
180 };
181 let json = serde_json::to_value(&block).expect("serialize ContentBlock");
182 assert_eq!(json["type"], "tool_result");
183 assert_eq!(json["tool_use_id"], "123");
184 assert_eq!(json["content"], "output");
185 }
186
187 #[test]
188 fn message_serialization_roundtrip() {
189 let original = Message::user("Test message");
190 let json = serde_json::to_string(&original).expect("serialize Message");
191 let deserialized: Message = serde_json::from_str(&json).expect("deserialize Message");
192 assert_eq!(deserialized.role, "user");
193 assert_eq!(deserialized.content.len(), 1);
194 }
195
196 #[test]
197 fn message_with_multiple_content_blocks() {
198 let content = vec![
199 ContentBlock::Text {
200 text: "I'll read the file".to_string(),
201 },
202 ContentBlock::ToolUse {
203 id: ToolId::new("1"),
204 name: ToolName::new("read"),
205 input: json!({"path": "test.txt"}),
206 },
207 ];
208 let msg = Message::assistant(content);
209 assert_eq!(msg.content.len(), 2);
210 match &msg.content[0] {
211 ContentBlock::Text { text } => assert!(text.contains("read")),
212 other => panic!("Expected Text block first, got: {other:?}"),
213 }
214 match &msg.content[1] {
215 ContentBlock::ToolUse { name, .. } => assert_eq!(name.as_str(), "read"),
216 other => panic!("Expected ToolUse block second, got: {other:?}"),
217 }
218 }
219
220 #[test]
221 fn api_request_serialization() {
222 let req = ApiRequest {
223 model: ModelId::new("test-model"),
224 max_tokens: 123,
225 system: "system".to_string(),
226 messages: vec![Message::user("hello")],
227 tools: vec![ToolDefinition {
228 name: "read".to_string(),
229 description: "Read a file".to_string(),
230 input_schema: json!({
231 "type": "object",
232 "properties": {"path": {"type": "string"}},
233 }),
234 }],
235 };
236 let json = serde_json::to_value(&req).expect("serialize ApiRequest");
237 assert_eq!(json["model"], "test-model");
238 assert_eq!(json["max_tokens"], 123);
239 assert_eq!(json["system"], "system");
240 assert!(json["messages"].is_array());
241 assert!(json["tools"].is_array());
242 }
243
244 #[test]
245 fn tool_definition_serialization() {
246 let tool = ToolDefinition {
247 name: "read".to_string(),
248 description: "Read a file".to_string(),
249 input_schema: json!({"type": "object", "properties": {"path": {"type": "string"}}}),
250 };
251 let json = serde_json::to_value(&tool).expect("serialize");
252 assert_eq!(json["name"], "read");
253 }
254}