Skip to main content

serdes_ai_streaming/
events.rs

1//! Streaming event types.
2//!
3//! This module defines the events emitted during agent streaming.
4
5use serde::{Deserialize, Serialize};
6use serde_json::Value as JsonValue;
7use serdes_ai_core::{ModelResponse, RequestUsage};
8use std::fmt;
9
10/// Events emitted during streaming.
11#[derive(Debug, Clone, Serialize, Deserialize)]
12#[serde(tag = "type", rename_all = "snake_case")]
13pub enum AgentStreamEvent<Output = JsonValue> {
14    /// Run started.
15    RunStart {
16        /// The run ID.
17        run_id: String,
18        /// Step number (0 for start).
19        step: u32,
20    },
21
22    /// Model request started.
23    RequestStart {
24        /// Current step number.
25        step: u32,
26    },
27
28    /// Text delta received.
29    TextDelta {
30        /// The text content.
31        content: String,
32        /// Index of the part this delta belongs to.
33        part_index: usize,
34    },
35
36    /// Tool call started.
37    ToolCallStart {
38        /// Tool name.
39        name: String,
40        /// Tool call ID (if available).
41        tool_call_id: Option<String>,
42        /// Index of this tool call.
43        index: usize,
44    },
45
46    /// Tool call arguments delta.
47    ToolCallDelta {
48        /// Arguments delta (JSON string).
49        args_delta: String,
50        /// Index of this tool call.
51        index: usize,
52    },
53
54    /// Tool call completed.
55    ToolCallComplete {
56        /// Tool name.
57        name: String,
58        /// Complete arguments.
59        args: JsonValue,
60        /// Index of this tool call.
61        index: usize,
62    },
63
64    /// Tool result received.
65    ToolResult {
66        /// Tool name.
67        name: String,
68        /// Result content.
69        result: JsonValue,
70        /// Whether the tool succeeded.
71        success: bool,
72        /// Index of this tool call.
73        index: usize,
74    },
75
76    /// Thinking delta (for Claude extended thinking).
77    ThinkingDelta {
78        /// The thinking content.
79        content: String,
80        /// Index of this thinking part.
81        index: usize,
82    },
83
84    /// Partial output available (validated incrementally).
85    PartialOutput {
86        /// The partial output.
87        output: Output,
88    },
89
90    /// Model response complete.
91    ResponseComplete {
92        /// The complete response.
93        response: ModelResponse,
94    },
95
96    /// Usage update.
97    UsageUpdate {
98        /// Current usage.
99        usage: RequestUsage,
100    },
101
102    /// Final output ready.
103    FinalOutput {
104        /// The final output.
105        output: Output,
106    },
107
108    /// Run complete.
109    RunComplete {
110        /// The run ID.
111        run_id: String,
112        /// Total steps.
113        total_steps: u32,
114    },
115
116    /// Error occurred.
117    Error {
118        /// Error message.
119        message: String,
120        /// Whether the error is recoverable.
121        recoverable: bool,
122    },
123}
124
125impl<Output> AgentStreamEvent<Output> {
126    /// Create a run start event.
127    pub fn run_start(run_id: impl Into<String>, step: u32) -> Self {
128        Self::RunStart {
129            run_id: run_id.into(),
130            step,
131        }
132    }
133
134    /// Create a text delta event.
135    pub fn text_delta(content: impl Into<String>, part_index: usize) -> Self {
136        Self::TextDelta {
137            content: content.into(),
138            part_index,
139        }
140    }
141
142    /// Create an error event.
143    pub fn error(message: impl Into<String>, recoverable: bool) -> Self {
144        Self::Error {
145            message: message.into(),
146            recoverable,
147        }
148    }
149
150    /// Check if this is the final event.
151    #[must_use]
152    pub fn is_terminal(&self) -> bool {
153        matches!(self, Self::RunComplete { .. } | Self::Error { .. })
154    }
155
156    /// Check if this is an error event.
157    #[must_use]
158    pub fn is_error(&self) -> bool {
159        matches!(self, Self::Error { .. })
160    }
161
162    /// Get the text content if this is a text delta.
163    pub fn as_text(&self) -> Option<&str> {
164        match self {
165            Self::TextDelta { content, .. } => Some(content),
166            _ => None,
167        }
168    }
169
170    /// Get the output if this is a final output event.
171    pub fn as_output(&self) -> Option<&Output> {
172        match self {
173            Self::FinalOutput { output } => Some(output),
174            Self::PartialOutput { output } => Some(output),
175            _ => None,
176        }
177    }
178
179    /// Map the output type.
180    pub fn map_output<U, F>(self, f: F) -> AgentStreamEvent<U>
181    where
182        F: FnOnce(Output) -> U,
183    {
184        match self {
185            Self::RunStart { run_id, step } => AgentStreamEvent::RunStart { run_id, step },
186            Self::RequestStart { step } => AgentStreamEvent::RequestStart { step },
187            Self::TextDelta {
188                content,
189                part_index,
190            } => AgentStreamEvent::TextDelta {
191                content,
192                part_index,
193            },
194            Self::ToolCallStart {
195                name,
196                tool_call_id,
197                index,
198            } => AgentStreamEvent::ToolCallStart {
199                name,
200                tool_call_id,
201                index,
202            },
203            Self::ToolCallDelta { args_delta, index } => {
204                AgentStreamEvent::ToolCallDelta { args_delta, index }
205            }
206            Self::ToolCallComplete { name, args, index } => {
207                AgentStreamEvent::ToolCallComplete { name, args, index }
208            }
209            Self::ToolResult {
210                name,
211                result,
212                success,
213                index,
214            } => AgentStreamEvent::ToolResult {
215                name,
216                result,
217                success,
218                index,
219            },
220            Self::ThinkingDelta { content, index } => {
221                AgentStreamEvent::ThinkingDelta { content, index }
222            }
223            Self::PartialOutput { output } => AgentStreamEvent::PartialOutput { output: f(output) },
224            Self::ResponseComplete { response } => AgentStreamEvent::ResponseComplete { response },
225            Self::UsageUpdate { usage } => AgentStreamEvent::UsageUpdate { usage },
226            Self::FinalOutput { output } => AgentStreamEvent::FinalOutput { output: f(output) },
227            Self::RunComplete {
228                run_id,
229                total_steps,
230            } => AgentStreamEvent::RunComplete {
231                run_id,
232                total_steps,
233            },
234            Self::Error {
235                message,
236                recoverable,
237            } => AgentStreamEvent::Error {
238                message,
239                recoverable,
240            },
241        }
242    }
243}
244
245impl<Output: fmt::Display> fmt::Display for AgentStreamEvent<Output> {
246    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
247        match self {
248            Self::RunStart { run_id, .. } => write!(f, "[run_start] {}", run_id),
249            Self::RequestStart { step } => write!(f, "[request_start] step {}", step),
250            Self::TextDelta { content, .. } => write!(f, "{}", content),
251            Self::ToolCallStart { name, .. } => write!(f, "[tool_start] {}", name),
252            Self::ToolCallDelta { args_delta, .. } => write!(f, "{}", args_delta),
253            Self::ToolCallComplete { name, .. } => write!(f, "[tool_complete] {}", name),
254            Self::ToolResult { name, success, .. } => {
255                write!(
256                    f,
257                    "[tool_result] {} ({})",
258                    name,
259                    if *success { "ok" } else { "error" }
260                )
261            }
262            Self::ThinkingDelta { content, .. } => write!(f, "[thinking] {}", content),
263            Self::PartialOutput { output } => write!(f, "[partial] {}", output),
264            Self::ResponseComplete { .. } => write!(f, "[response_complete]"),
265            Self::UsageUpdate { .. } => write!(f, "[usage_update]"),
266            Self::FinalOutput { output } => write!(f, "[output] {}", output),
267            Self::RunComplete { run_id, .. } => write!(f, "[run_complete] {}", run_id),
268            Self::Error { message, .. } => write!(f, "[error] {}", message),
269        }
270    }
271}
272
273#[cfg(test)]
274mod tests {
275    use super::*;
276
277    #[test]
278    fn test_event_creation() {
279        let event: AgentStreamEvent<String> = AgentStreamEvent::run_start("run-123", 0);
280        assert!(!event.is_terminal());
281        assert!(!event.is_error());
282    }
283
284    #[test]
285    fn test_text_delta() {
286        let event: AgentStreamEvent<String> = AgentStreamEvent::text_delta("Hello", 0);
287        assert_eq!(event.as_text(), Some("Hello"));
288    }
289
290    #[test]
291    fn test_terminal_events() {
292        let complete: AgentStreamEvent<String> = AgentStreamEvent::RunComplete {
293            run_id: "run-123".to_string(),
294            total_steps: 1,
295        };
296        assert!(complete.is_terminal());
297
298        let error: AgentStreamEvent<String> = AgentStreamEvent::error("oops", false);
299        assert!(error.is_terminal());
300        assert!(error.is_error());
301    }
302
303    #[test]
304    fn test_map_output() {
305        let event: AgentStreamEvent<i32> = AgentStreamEvent::FinalOutput { output: 42 };
306        let mapped = event.map_output(|n| n.to_string());
307
308        if let AgentStreamEvent::FinalOutput { output } = mapped {
309            assert_eq!(output, "42");
310        } else {
311            panic!("Expected FinalOutput");
312        }
313    }
314
315    #[test]
316    fn test_display() {
317        let event: AgentStreamEvent<String> = AgentStreamEvent::text_delta("test", 0);
318        assert_eq!(format!("{}", event), "test");
319    }
320}