Skip to main content

aether_evals/agents/
fake_agent.rs

1use super::agent::{Agent, AgentRunResult, RunError};
2use crate::Task;
3use aether_core::events::AgentMessage;
4use async_stream::try_stream;
5use futures::Stream;
6use llm::ToolCallResult;
7use std::path::{Path, PathBuf};
8use tokio::fs::{create_dir_all, write};
9
10#[derive(Clone)]
11pub struct FakeAgent {
12    messages: Vec<AgentMessage>,
13    file_writes: Vec<(PathBuf, String)>,
14    workspace: Option<PathBuf>,
15}
16
17impl FakeAgent {
18    pub fn new(messages: Vec<AgentMessage>) -> Self {
19        Self { messages, file_writes: Vec::new(), workspace: None }
20    }
21
22    pub fn success() -> Self {
23        Self::new(vec![AgentMessage::text("fake_1", "Task completed successfully", true, "fake"), AgentMessage::Done])
24    }
25
26    pub fn with_tool_call(tool_name: impl Into<String>, result: impl Into<String>) -> Self {
27        let tool_name = tool_name.into();
28        Self::new(vec![
29            AgentMessage::ToolResult {
30                result: ToolCallResult {
31                    id: "fake_call_1".to_string(),
32                    name: tool_name,
33                    arguments: "{}".to_string(),
34                    result: result.into(),
35                },
36                result_meta: None,
37                model_name: "fake".to_string(),
38            },
39            AgentMessage::text("fake_2", "Task completed using tools", true, "fake"),
40            AgentMessage::Done,
41        ])
42    }
43
44    pub fn writes_file(path: impl Into<PathBuf>, contents: impl Into<String>) -> Self {
45        Self::success().with_file_write(path, contents)
46    }
47
48    pub fn with_file_write(mut self, path: impl Into<PathBuf>, contents: impl Into<String>) -> Self {
49        self.file_writes.push((path.into(), contents.into()));
50        self
51    }
52
53    pub fn with_workspace(mut self, workspace: impl AsRef<Path>) -> Self {
54        self.workspace = Some(workspace.as_ref().to_path_buf());
55        self
56    }
57}
58
59impl Agent for FakeAgent {
60    fn run(&self, _task: Task) -> impl Stream<Item = AgentRunResult> + Send + '_ {
61        try_stream! {
62            if !self.file_writes.is_empty() && self.workspace.is_none() {
63                Err(RunError::ConfigurationError("FakeAgent file writes require a bound workspace".to_string()))?;
64            }
65
66            for (path, contents) in &self.file_writes {
67                let file_path = self.workspace.as_ref().expect("workspace checked above").join(path);
68                if let Some(parent) = file_path.parent() {
69                    create_dir_all(parent).await.map_err(|error| {
70                        RunError::ExecutionFailed(format!("failed to create fake file parent: {error}"))
71                    })?;
72                }
73                write(&file_path, contents)
74                    .await
75                    .map_err(|error| RunError::ExecutionFailed(format!("failed to write fake file: {error}")))?;
76            }
77
78            for message in &self.messages {
79                yield message.clone();
80            }
81        }
82    }
83}
84
85#[cfg(test)]
86mod tests {
87    use super::*;
88
89    #[tokio::test]
90    async fn fake_agent_success_sends_done() {
91        let agent = FakeAgent::success();
92        let transcript = crate::Transcript::from_stream(agent.run(Task::new("test task"))).await.unwrap();
93        let messages = transcript.messages();
94
95        assert!(matches!(messages.first(), Some(AgentMessage::Text { .. })));
96        assert!(matches!(messages.get(1), Some(AgentMessage::Done)));
97    }
98
99    #[tokio::test]
100    async fn fake_agent_with_tool_call_sends_tool_messages() {
101        let agent = FakeAgent::with_tool_call("bash", "success");
102        let transcript = crate::Transcript::from_stream(agent.run(Task::new("test task"))).await.unwrap();
103        let messages = transcript.messages();
104
105        assert_eq!(messages.len(), 3);
106        assert!(matches!(messages.last(), Some(AgentMessage::Done)));
107    }
108
109    #[tokio::test]
110    async fn fake_agent_writes_file_in_workspace() {
111        let temp_dir = tempfile::tempdir().unwrap();
112        let agent = FakeAgent::writes_file("nested/hello.txt", "hello").with_workspace(temp_dir.path());
113        crate::Transcript::from_stream(agent.run(Task::new("test task"))).await.unwrap();
114        assert_eq!(std::fs::read_to_string(temp_dir.path().join("nested/hello.txt")).unwrap(), "hello");
115    }
116}