1use serde::{Deserialize, Serialize};
6use serde_json::Value as JsonValue;
7use serdes_ai_core::{ModelResponse, RequestUsage};
8use std::fmt;
9
10#[derive(Debug, Clone, Serialize, Deserialize)]
12#[serde(tag = "type", rename_all = "snake_case")]
13pub enum AgentStreamEvent<Output = JsonValue> {
14 RunStart {
16 run_id: String,
18 step: u32,
20 },
21
22 RequestStart {
24 step: u32,
26 },
27
28 TextDelta {
30 content: String,
32 part_index: usize,
34 },
35
36 ToolCallStart {
38 name: String,
40 tool_call_id: Option<String>,
42 index: usize,
44 },
45
46 ToolCallDelta {
48 args_delta: String,
50 index: usize,
52 },
53
54 ToolCallComplete {
56 name: String,
58 args: JsonValue,
60 index: usize,
62 },
63
64 ToolResult {
66 name: String,
68 result: JsonValue,
70 success: bool,
72 index: usize,
74 },
75
76 ThinkingDelta {
78 content: String,
80 index: usize,
82 },
83
84 PartialOutput {
86 output: Output,
88 },
89
90 ResponseComplete {
92 response: ModelResponse,
94 },
95
96 UsageUpdate {
98 usage: RequestUsage,
100 },
101
102 FinalOutput {
104 output: Output,
106 },
107
108 RunComplete {
110 run_id: String,
112 total_steps: u32,
114 },
115
116 Error {
118 message: String,
120 recoverable: bool,
122 },
123}
124
125impl<Output> AgentStreamEvent<Output> {
126 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 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 pub fn error(message: impl Into<String>, recoverable: bool) -> Self {
144 Self::Error {
145 message: message.into(),
146 recoverable,
147 }
148 }
149
150 #[must_use]
152 pub fn is_terminal(&self) -> bool {
153 matches!(self, Self::RunComplete { .. } | Self::Error { .. })
154 }
155
156 #[must_use]
158 pub fn is_error(&self) -> bool {
159 matches!(self, Self::Error { .. })
160 }
161
162 pub fn as_text(&self) -> Option<&str> {
164 match self {
165 Self::TextDelta { content, .. } => Some(content),
166 _ => None,
167 }
168 }
169
170 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 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}