Skip to main content

aether_evals/agents/
fake_agent.rs

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