Skip to main content

agentik_sdk/types/
streaming.rs

1//! Streaming response types for real-time message generation.
2//!
3//! This module provides types for handling Server-Sent Events (SSE) from the Anthropic API
4//! when streaming is enabled. It includes all event types and delta structures needed to
5//! process incremental responses from Claude.
6
7use serde::{Deserialize, Serialize};
8use crate::types::{Message, ContentBlock, StopReason, ServerToolUsage};
9
10/// Main stream event type that encompasses all possible streaming events.
11///
12/// This is the primary type you'll work with when processing streaming responses.
13/// Each event represents a different stage of the message generation process.
14#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
15#[serde(tag = "type")]
16pub enum MessageStreamEvent {
17    /// Initial event when a message starts being generated.
18    ///
19    /// Contains the initial message structure with metadata but no content yet.
20    #[serde(rename = "message_start")]
21    MessageStart {
22        /// The initial message structure
23        message: Message,
24    },
25
26    /// Event when the message metadata is updated during generation.
27    ///
28    /// Contains updates to stop reason, stop sequence, and usage information.
29    #[serde(rename = "message_delta")]
30    MessageDelta {
31        /// The delta containing updated fields
32        delta: MessageDelta,
33        /// Current usage statistics
34        usage: MessageDeltaUsage,
35    },
36
37    /// Final event when message generation is complete.
38    #[serde(rename = "message_stop")]
39    MessageStop,
40
41    /// Event when a new content block starts being generated.
42    ///
43    /// This occurs when Claude begins generating a new piece of content
44    /// (text, tool use, etc.).
45    #[serde(rename = "content_block_start")]
46    ContentBlockStart {
47        /// The content block being started
48        content_block: ContentBlock,
49        /// Index of this content block in the message
50        index: usize,
51    },
52
53    /// Event containing incremental updates to a content block.
54    ///
55    /// This is where you'll receive the actual text being generated,
56    /// tool input being parsed, etc.
57    #[serde(rename = "content_block_delta")]
58    ContentBlockDelta {
59        /// The incremental update
60        delta: ContentBlockDelta,
61        /// Index of the content block being updated
62        index: usize,
63    },
64
65    /// Event when a content block finishes generation.
66    #[serde(rename = "content_block_stop")]
67    ContentBlockStop {
68        /// Index of the content block that finished
69        index: usize,
70    },
71}
72
73/// Delta updates for message-level information during streaming.
74#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
75pub struct MessageDelta {
76    /// Updated stop reason, if the message has stopped
77    pub stop_reason: Option<StopReason>,
78    /// The stop sequence that triggered stopping, if any
79    pub stop_sequence: Option<String>,
80}
81
82/// Usage statistics for streaming responses.
83///
84/// These are cumulative totals that get updated throughout the stream.
85#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
86pub struct MessageDeltaUsage {
87    /// Cumulative output tokens generated so far
88    pub output_tokens: u32,
89    /// Cumulative input tokens used (may be null during streaming)
90    pub input_tokens: Option<u32>,
91    /// Cumulative cache creation tokens (may be null)
92    pub cache_creation_input_tokens: Option<u32>,
93    /// Cumulative cache read tokens (may be null)
94    pub cache_read_input_tokens: Option<u32>,
95    /// Server tool usage statistics (may be null)
96    pub server_tool_use: Option<ServerToolUsage>,
97}
98
99
100
101/// Delta updates for content blocks during streaming.
102///
103/// This enum contains all possible types of incremental updates
104/// that can occur within a content block.
105#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
106#[serde(tag = "type")]
107pub enum ContentBlockDelta {
108    /// Incremental text content being generated.
109    #[serde(rename = "text_delta")]
110    TextDelta {
111        /// The new text to append
112        text: String,
113    },
114
115    /// Incremental JSON input for tool use blocks.
116    ///
117    /// This provides partial JSON as it's being parsed,
118    /// useful for streaming tool input generation.
119    #[serde(rename = "input_json_delta")]
120    InputJsonDelta {
121        /// Partial JSON string
122        partial_json: String,
123    },
124
125    /// Citations being added to text blocks.
126    #[serde(rename = "citations_delta")]
127    CitationsDelta {
128        /// The citation being added
129        citation: TextCitation,
130    },
131
132    /// Incremental thinking content (extended reasoning).
133    #[serde(rename = "thinking_delta")]
134    ThinkingDelta {
135        /// The thinking text being generated
136        thinking: String,
137    },
138
139    /// Signature updates for thinking blocks.
140    #[serde(rename = "signature_delta")]
141    SignatureDelta {
142        /// The signature string
143        signature: String,
144    },
145}
146
147/// Citation information for text blocks.
148///
149/// Citations provide source attribution for generated text,
150/// with different types depending on the source document.
151#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
152#[serde(tag = "type")]
153pub enum TextCitation {
154    /// Citation pointing to character locations in plain text.
155    #[serde(rename = "char_location")]
156    CharLocation {
157        /// The text being cited
158        cited_text: String,
159        /// Index of the source document
160        document_index: usize,
161        /// Title of the source document (may be null)
162        document_title: Option<String>,
163        /// Starting character index
164        start_char_index: usize,
165        /// Ending character index
166        end_char_index: usize,
167    },
168
169    /// Citation pointing to page locations in PDFs.
170    #[serde(rename = "page_location")]
171    PageLocation {
172        /// The text being cited
173        cited_text: String,
174        /// Index of the source document
175        document_index: usize,
176        /// Title of the source document (may be null)
177        document_title: Option<String>,
178        /// Starting page number
179        start_page_number: usize,
180        /// Ending page number
181        end_page_number: usize,
182    },
183
184    /// Citation pointing to content block locations.
185    #[serde(rename = "content_block_location")]
186    ContentBlockLocation {
187        /// The text being cited
188        cited_text: String,
189        /// Index of the source document
190        document_index: usize,
191        /// Title of the source document (may be null)
192        document_title: Option<String>,
193        /// Starting content block index
194        start_block_index: usize,
195        /// Ending content block index
196        end_block_index: usize,
197    },
198
199    /// Citation pointing to web search results.
200    #[serde(rename = "web_search_result_location")]
201    WebSearchResultLocation {
202        /// The text being cited
203        cited_text: String,
204        /// Encrypted index for the search result
205        encrypted_index: String,
206        /// Title of the web page (may be null)
207        title: Option<String>,
208        /// URL of the web page
209        url: String,
210    },
211}
212
213/// Type aliases for clarity and compatibility with the main API types.
214pub 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}