aether_core/testing/
agent_event_builder.rs

1use crate::events::{AgentEvent, StreamState, ToolEvent};
2use llm::{ToolCallError, ToolCallRequest, ToolCallResult};
3use serde::Serialize;
4
5pub fn agent_event(message_id: &str) -> AgentEventBuilder {
6    AgentEventBuilder::new(message_id)
7}
8
9pub struct AgentEventBuilder {
10    message_id: String,
11    chunks: Vec<AgentEvent>,
12    full_text: String,
13}
14
15impl AgentEventBuilder {
16    pub fn new(message_id: &str) -> Self {
17        Self { message_id: message_id.to_string(), chunks: Vec::new(), full_text: String::new() }
18    }
19
20    pub fn text(mut self, chunks: &[&str]) -> Self {
21        for chunk in chunks {
22            self.chunks.push(AgentEvent::text(&self.message_id, chunk, StreamState::Partial));
23            self.full_text.push_str(chunk);
24        }
25        self
26    }
27
28    pub fn tool_call<T: Serialize, U: Serialize>(
29        mut self,
30        tool_call_id: &str,
31        name: &str,
32        request: &T,
33        result: &U,
34    ) -> Self {
35        let request_json = serde_json::to_string(request).expect("Failed to serialize request");
36        let result_value = serde_json::to_value(result).expect("Failed to serialize result");
37        let result_yaml = serde_yml::to_string(&result_value).unwrap_or_else(|_| result_value.to_string());
38
39        self.push_tool_call_start(tool_call_id, name);
40        self.push_tool_call_chunk(tool_call_id, &request_json);
41
42        self.chunks.push(AgentEvent::Tool(ToolEvent::Result {
43            result: ToolCallResult {
44                id: tool_call_id.to_string(),
45                name: name.to_string(),
46                arguments: request_json,
47                result: result_yaml,
48            },
49            result_meta: None,
50        }));
51
52        self
53    }
54
55    pub fn tool_call_with_error<T: Serialize>(
56        mut self,
57        tool_call_id: &str,
58        name: &str,
59        request: &T,
60        error_message: &str,
61    ) -> Self {
62        let request_json = serde_json::to_string(request).expect("Failed to serialize request");
63
64        let error_result = format!("Tool execution error: {error_message}");
65
66        self.push_tool_call_start(tool_call_id, name);
67        self.push_tool_call_chunk(tool_call_id, &request_json);
68
69        self.chunks.push(AgentEvent::Tool(ToolEvent::Error {
70            error: ToolCallError {
71                id: tool_call_id.to_string(),
72                name: name.to_string(),
73                arguments: Some(request_json),
74                error: error_result,
75            },
76        }));
77
78        self
79    }
80
81    pub fn build(mut self) -> Vec<AgentEvent> {
82        self.chunks.push(AgentEvent::text(&self.message_id, &self.full_text, StreamState::Complete));
83
84        self.chunks
85    }
86
87    fn push_tool_call_start(&mut self, tool_call_id: &str, name: &str) {
88        self.chunks.push(AgentEvent::Tool(ToolEvent::Call {
89            request: ToolCallRequest { id: tool_call_id.to_string(), name: name.to_string(), arguments: String::new() },
90        }));
91    }
92
93    fn push_tool_call_chunk(&mut self, tool_call_id: &str, chunk: &str) {
94        self.chunks.push(AgentEvent::Tool(ToolEvent::CallUpdate {
95            tool_call_id: tool_call_id.to_string(),
96            chunk: chunk.to_string(),
97        }));
98    }
99}
aether_core/testing/agent_event_builder.rs

aether_core/testing/
agent_event_builder.rs
