Skip to main content

agent_io/agent/
events.rs

1//! Agent events for streaming responses
2
3use serde::{Deserialize, Serialize};
4
5use crate::llm::Usage;
6
7/// Event emitted during agent execution
8#[derive(Debug, Clone, Serialize, Deserialize)]
9#[serde(tag = "type", rename_all = "snake_case")]
10pub enum AgentEvent {
11    /// Text content from the LLM
12    Text(TextEvent),
13    /// Thinking/reasoning content
14    Thinking(ThinkingEvent),
15    /// Tool is being called
16    ToolCall(ToolCallEvent),
17    /// Tool execution result
18    ToolResult(ToolResultEvent),
19    /// Final response ready
20    FinalResponse(FinalResponseEvent),
21    /// Message started
22    MessageStart(MessageStartEvent),
23    /// Message completed
24    MessageComplete(MessageCompleteEvent),
25    /// Step started
26    StepStart(StepStartEvent),
27    /// Step completed
28    StepComplete(StepCompleteEvent),
29    /// Error occurred
30    Error(ErrorEvent),
31}
32
33/// Text content event
34#[derive(Debug, Clone, Serialize, Deserialize)]
35pub struct TextEvent {
36    pub content: String,
37    pub delta: bool,
38}
39
40impl TextEvent {
41    pub fn new(content: impl Into<String>) -> Self {
42        Self {
43            content: content.into(),
44            delta: false,
45        }
46    }
47
48    pub fn delta(content: impl Into<String>) -> Self {
49        Self {
50            content: content.into(),
51            delta: true,
52        }
53    }
54}
55
56/// Thinking/reasoning content event
57#[derive(Debug, Clone, Serialize, Deserialize)]
58pub struct ThinkingEvent {
59    pub content: String,
60    pub delta: bool,
61}
62
63impl ThinkingEvent {
64    pub fn new(content: impl Into<String>) -> Self {
65        Self {
66            content: content.into(),
67            delta: false,
68        }
69    }
70
71    pub fn delta(content: impl Into<String>) -> Self {
72        Self {
73            content: content.into(),
74            delta: true,
75        }
76    }
77}
78
79/// Tool call event
80#[derive(Debug, Clone, Serialize, Deserialize)]
81pub struct ToolCallEvent {
82    pub tool_call_id: String,
83    pub name: String,
84    pub arguments: String,
85    pub step: usize,
86}
87
88impl ToolCallEvent {
89    pub fn new(tool_call: &crate::llm::ToolCall, step: usize) -> Self {
90        Self {
91            tool_call_id: tool_call.id.clone(),
92            name: tool_call.function.name.clone(),
93            arguments: tool_call.function.arguments.clone(),
94            step,
95        }
96    }
97}
98
99/// Tool result event
100#[derive(Debug, Clone, Serialize, Deserialize)]
101pub struct ToolResultEvent {
102    pub tool_call_id: String,
103    pub name: String,
104    pub result: String,
105    pub step: usize,
106    pub ephemeral: bool,
107}
108
109impl ToolResultEvent {
110    pub fn new(
111        tool_call_id: impl Into<String>,
112        name: impl Into<String>,
113        result: impl Into<String>,
114        step: usize,
115    ) -> Self {
116        Self {
117            tool_call_id: tool_call_id.into(),
118            name: name.into(),
119            result: result.into(),
120            step,
121            ephemeral: false,
122        }
123    }
124
125    pub fn with_ephemeral(mut self, ephemeral: bool) -> Self {
126        self.ephemeral = ephemeral;
127        self
128    }
129}
130
131/// Final response event
132#[derive(Debug, Clone, Serialize, Deserialize)]
133pub struct FinalResponseEvent {
134    pub content: String,
135    pub usage: Option<UsageSummary>,
136    pub steps: usize,
137}
138
139impl FinalResponseEvent {
140    pub fn new(content: impl Into<String>) -> Self {
141        Self {
142            content: content.into(),
143            usage: None,
144            steps: 0,
145        }
146    }
147
148    pub fn with_usage(mut self, usage: UsageSummary) -> Self {
149        self.usage = Some(usage);
150        self
151    }
152
153    pub fn with_steps(mut self, steps: usize) -> Self {
154        self.steps = steps;
155        self
156    }
157}
158
159/// Message start event
160#[derive(Debug, Clone, Serialize, Deserialize)]
161pub struct MessageStartEvent {
162    pub role: String,
163}
164
165impl MessageStartEvent {
166    pub fn user() -> Self {
167        Self {
168            role: "user".to_string(),
169        }
170    }
171
172    pub fn assistant() -> Self {
173        Self {
174            role: "assistant".to_string(),
175        }
176    }
177}
178
179/// Message complete event
180#[derive(Debug, Clone, Serialize, Deserialize)]
181pub struct MessageCompleteEvent {
182    pub role: String,
183}
184
185impl MessageCompleteEvent {
186    pub fn user() -> Self {
187        Self {
188            role: "user".to_string(),
189        }
190    }
191
192    pub fn assistant() -> Self {
193        Self {
194            role: "assistant".to_string(),
195        }
196    }
197}
198
199/// Step start event
200#[derive(Debug, Clone, Serialize, Deserialize)]
201pub struct StepStartEvent {
202    pub step: usize,
203}
204
205impl StepStartEvent {
206    pub fn new(step: usize) -> Self {
207        Self { step }
208    }
209}
210
211/// Step complete event
212#[derive(Debug, Clone, Serialize, Deserialize)]
213pub struct StepCompleteEvent {
214    pub step: usize,
215}
216
217impl StepCompleteEvent {
218    pub fn new(step: usize) -> Self {
219        Self { step }
220    }
221}
222
223/// Error event
224#[derive(Debug, Clone, Serialize, Deserialize)]
225pub struct ErrorEvent {
226    pub message: String,
227    pub code: Option<String>,
228}
229
230impl ErrorEvent {
231    pub fn new(message: impl Into<String>) -> Self {
232        Self {
233            message: message.into(),
234            code: None,
235        }
236    }
237
238    pub fn with_code(mut self, code: impl Into<String>) -> Self {
239        self.code = Some(code.into());
240        self
241    }
242}
243
244/// Usage summary for the session
245#[derive(Debug, Clone, Default, Serialize, Deserialize)]
246pub struct UsageSummary {
247    pub total_prompt_tokens: u64,
248    pub total_completion_tokens: u64,
249    pub total_tokens: u64,
250    pub total_cost: Option<f64>,
251    pub by_model: std::collections::HashMap<String, ModelUsage>,
252}
253
254impl UsageSummary {
255    pub fn new() -> Self {
256        Self::default()
257    }
258
259    pub fn add_usage(&mut self, model: &str, usage: &Usage) {
260        self.total_prompt_tokens += usage.prompt_tokens;
261        self.total_completion_tokens += usage.completion_tokens;
262        self.total_tokens += usage.total_tokens;
263
264        let model_usage = self.by_model.entry(model.to_string()).or_default();
265        model_usage.prompt_tokens += usage.prompt_tokens;
266        model_usage.completion_tokens += usage.completion_tokens;
267        model_usage.total_tokens += usage.total_tokens;
268        model_usage.calls += 1;
269    }
270}
271
272/// Per-model usage statistics
273#[derive(Debug, Clone, Default, Serialize, Deserialize)]
274pub struct ModelUsage {
275    pub prompt_tokens: u64,
276    pub completion_tokens: u64,
277    pub total_tokens: u64,
278    pub calls: u64,
279}
280
281#[cfg(test)]
282mod tests {
283    use super::*;
284
285    #[test]
286    fn test_text_event() {
287        let event = TextEvent::new("Hello");
288        assert_eq!(event.content, "Hello");
289        assert!(!event.delta);
290
291        let delta = TextEvent::delta("Hello");
292        assert!(delta.delta);
293    }
294
295    #[test]
296    fn test_usage_summary() {
297        let mut summary = UsageSummary::new();
298        let usage = Usage::new(100, 50);
299
300        summary.add_usage("gpt-4o", &usage);
301
302        assert_eq!(summary.total_prompt_tokens, 100);
303        assert_eq!(summary.total_completion_tokens, 50);
304        assert_eq!(summary.total_tokens, 150);
305        assert!(summary.by_model.contains_key("gpt-4o"));
306    }
307}