aether_core/testing/
agent_event_builder.rs

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

aether_core/testing/
agent_event_builder.rs
