1use serde::{Deserialize, Serialize};
8use crate::types::{Message, ContentBlock, StopReason, ServerToolUsage};
9
10#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
15#[serde(tag = "type")]
16pub enum MessageStreamEvent {
17 #[serde(rename = "message_start")]
21 MessageStart {
22 message: Message,
24 },
25
26 #[serde(rename = "message_delta")]
30 MessageDelta {
31 delta: MessageDelta,
33 usage: MessageDeltaUsage,
35 },
36
37 #[serde(rename = "message_stop")]
39 MessageStop,
40
41 #[serde(rename = "content_block_start")]
46 ContentBlockStart {
47 content_block: ContentBlock,
49 index: usize,
51 },
52
53 #[serde(rename = "content_block_delta")]
58 ContentBlockDelta {
59 delta: ContentBlockDelta,
61 index: usize,
63 },
64
65 #[serde(rename = "content_block_stop")]
67 ContentBlockStop {
68 index: usize,
70 },
71}
72
73#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
75pub struct MessageDelta {
76 pub stop_reason: Option<StopReason>,
78 pub stop_sequence: Option<String>,
80}
81
82#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
86pub struct MessageDeltaUsage {
87 pub output_tokens: u32,
89 pub input_tokens: Option<u32>,
91 pub cache_creation_input_tokens: Option<u32>,
93 pub cache_read_input_tokens: Option<u32>,
95 pub server_tool_use: Option<ServerToolUsage>,
97}
98
99
100
101#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
106#[serde(tag = "type")]
107pub enum ContentBlockDelta {
108 #[serde(rename = "text_delta")]
110 TextDelta {
111 text: String,
113 },
114
115 #[serde(rename = "input_json_delta")]
120 InputJsonDelta {
121 partial_json: String,
123 },
124
125 #[serde(rename = "citations_delta")]
127 CitationsDelta {
128 citation: TextCitation,
130 },
131
132 #[serde(rename = "thinking_delta")]
134 ThinkingDelta {
135 thinking: String,
137 },
138
139 #[serde(rename = "signature_delta")]
141 SignatureDelta {
142 signature: String,
144 },
145}
146
147#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
152#[serde(tag = "type")]
153pub enum TextCitation {
154 #[serde(rename = "char_location")]
156 CharLocation {
157 cited_text: String,
159 document_index: usize,
161 document_title: Option<String>,
163 start_char_index: usize,
165 end_char_index: usize,
167 },
168
169 #[serde(rename = "page_location")]
171 PageLocation {
172 cited_text: String,
174 document_index: usize,
176 document_title: Option<String>,
178 start_page_number: usize,
180 end_page_number: usize,
182 },
183
184 #[serde(rename = "content_block_location")]
186 ContentBlockLocation {
187 cited_text: String,
189 document_index: usize,
191 document_title: Option<String>,
193 start_block_index: usize,
195 end_block_index: usize,
197 },
198
199 #[serde(rename = "web_search_result_location")]
201 WebSearchResultLocation {
202 cited_text: String,
204 encrypted_index: String,
206 title: Option<String>,
208 url: String,
210 },
211}
212
213pub type MessageStartEvent = MessageStreamEvent;
215pub type MessageDeltaEvent = MessageStreamEvent;
216pub type MessageStopEvent = MessageStreamEvent;
217pub type ContentBlockStartEvent = MessageStreamEvent;
218pub type ContentBlockDeltaEvent = MessageStreamEvent;
219pub type ContentBlockStopEvent = MessageStreamEvent;
220
221#[cfg(test)]
222mod tests {
223 use super::*;
224 use serde_json;
225 use crate::types::Usage;
226
227 #[test]
228 fn test_message_start_event_serialization() {
229 let event = MessageStreamEvent::MessageStart {
230 message: Message {
231 id: "msg_123".to_string(),
232 type_: "message".to_string(),
233 role: crate::types::Role::Assistant,
234 content: vec![],
235 model: "claude-3-5-sonnet-latest".to_string(),
236 stop_reason: None,
237 stop_sequence: None,
238 usage: Usage {
239 input_tokens: 10,
240 output_tokens: 0,
241 cache_creation_input_tokens: None,
242 cache_read_input_tokens: None,
243 server_tool_use: None,
244 service_tier: None,
245 },
246 request_id: None,
247 },
248 };
249
250 let json = serde_json::to_string(&event).unwrap();
251 let parsed: MessageStreamEvent = serde_json::from_str(&json).unwrap();
252 assert_eq!(event, parsed);
253 }
254
255 #[test]
256 fn test_content_block_delta_serialization() {
257 let event = MessageStreamEvent::ContentBlockDelta {
258 delta: ContentBlockDelta::TextDelta {
259 text: "Hello".to_string(),
260 },
261 index: 0,
262 };
263
264 let json = serde_json::to_string(&event).unwrap();
265 let parsed: MessageStreamEvent = serde_json::from_str(&json).unwrap();
266 assert_eq!(event, parsed);
267 }
268
269 #[test]
270 fn test_message_delta_event_serialization() {
271 let event = MessageStreamEvent::MessageDelta {
272 delta: MessageDelta {
273 stop_reason: Some(StopReason::EndTurn),
274 stop_sequence: None,
275 },
276 usage: MessageDeltaUsage {
277 output_tokens: 25,
278 input_tokens: Some(10),
279 cache_creation_input_tokens: None,
280 cache_read_input_tokens: None,
281 server_tool_use: None,
282 },
283 };
284
285 let json = serde_json::to_string(&event).unwrap();
286 let parsed: MessageStreamEvent = serde_json::from_str(&json).unwrap();
287 assert_eq!(event, parsed);
288 }
289
290 #[test]
291 fn test_citation_serialization() {
292 let citation = TextCitation::CharLocation {
293 cited_text: "Example text".to_string(),
294 document_index: 0,
295 document_title: Some("Document Title".to_string()),
296 start_char_index: 10,
297 end_char_index: 23,
298 };
299
300 let json = serde_json::to_string(&citation).unwrap();
301 let parsed: TextCitation = serde_json::from_str(&json).unwrap();
302 assert_eq!(citation, parsed);
303 }
304
305 #[test]
306 fn test_all_delta_types() {
307 let deltas = vec![
308 ContentBlockDelta::TextDelta {
309 text: "Hello world".to_string(),
310 },
311 ContentBlockDelta::InputJsonDelta {
312 partial_json: r#"{"key": "val"#.to_string(),
313 },
314 ContentBlockDelta::ThinkingDelta {
315 thinking: "Let me think...".to_string(),
316 },
317 ContentBlockDelta::SignatureDelta {
318 signature: "signature_123".to_string(),
319 },
320 ];
321
322 for delta in deltas {
323 let json = serde_json::to_string(&delta).unwrap();
324 let parsed: ContentBlockDelta = serde_json::from_str(&json).unwrap();
325 assert_eq!(delta, parsed);
326 }
327 }
328}