aether_core/testing/
agent_message_builder.rs

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

aether_core/testing/
agent_message_builder.rs
