1use serde::{Deserialize, Serialize};
4
5use crate::llm::Usage;
6
7#[derive(Debug, Clone, Serialize, Deserialize)]
9#[serde(tag = "type", rename_all = "snake_case")]
10pub enum AgentEvent {
11 Text(TextEvent),
13 Thinking(ThinkingEvent),
15 ToolCall(ToolCallEvent),
17 ToolResult(ToolResultEvent),
19 FinalResponse(FinalResponseEvent),
21 MessageStart(MessageStartEvent),
23 MessageComplete(MessageCompleteEvent),
25 StepStart(StepStartEvent),
27 StepComplete(StepCompleteEvent),
29 Error(ErrorEvent),
31}
32
33#[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#[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#[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#[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#[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#[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#[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#[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#[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#[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#[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#[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}