1use std::{
2 error::Error,
3 fmt::{self, Debug},
4 path::Path,
5 sync::{Arc, Mutex},
6 time::Duration,
7};
8
9use crate::{
10 ToolMetrics,
11 assertions::{AssertionsReport, RanAssertion, RanAssertionResult},
12};
13use agent::ThreadEvent;
14use anyhow::{Result, anyhow};
15use async_trait::async_trait;
16use buffer_diff::DiffHunkStatus;
17use collections::HashMap;
18use futures::{FutureExt as _, StreamExt, channel::mpsc, select_biased};
19use gpui::{AppContext, AsyncApp, Entity};
20use language_model::{LanguageModel, Role, StopReason};
21
22pub const THREAD_EVENT_TIMEOUT: Duration = Duration::from_secs(60 * 2);
23
24#[async_trait(?Send)]
25pub trait Example {
26 fn meta(&self) -> ExampleMetadata;
27 async fn conversation(&self, cx: &mut ExampleContext) -> Result<()>;
28 fn diff_assertions(&self) -> Vec<JudgeAssertion> {
29 Vec::new()
30 }
31 fn thread_assertions(&self) -> Vec<JudgeAssertion> {
32 Vec::new()
33 }
34}
35
36#[derive(Clone, Debug)]
37pub struct JudgeAssertion {
38 pub id: String,
39 pub description: String,
40}
41
42#[derive(Clone, Debug)]
43pub struct ExampleMetadata {
44 pub name: String,
45 pub url: String,
46 pub revision: String,
47 pub language_server: Option<LanguageServer>,
48 pub max_assertions: Option<usize>,
49}
50
51#[derive(Clone, Debug)]
52pub struct LanguageServer {
53 pub file_extension: String,
54 pub allow_preexisting_diagnostics: bool,
55}
56
57impl ExampleMetadata {
58 pub fn repo_name(&self) -> String {
59 self.url
60 .split('/')
61 .next_back()
62 .unwrap_or(&"")
63 .trim_end_matches(".git")
64 .into()
65 }
66}
67
68pub struct FailedAssertion(pub String);
69
70impl fmt::Debug for FailedAssertion {
71 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
72 write!(f, "Assertion failure: {}", self.0)
73 }
74}
75
76impl fmt::Display for FailedAssertion {
77 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
78 write!(f, "{}", self.0)
79 }
80}
81
82impl Error for FailedAssertion {}
83
84pub struct ExampleContext {
85 meta: ExampleMetadata,
86 log_prefix: String,
87 agent_thread: Entity<agent::Thread>,
88 app: AsyncApp,
89 model: Arc<dyn LanguageModel>,
90 pub assertions: AssertionsReport,
91 pub tool_metrics: Arc<Mutex<ToolMetrics>>,
92}
93
94impl ExampleContext {
95 pub fn new(
96 meta: ExampleMetadata,
97 log_prefix: String,
98 agent_thread: Entity<agent::Thread>,
99 model: Arc<dyn LanguageModel>,
100 app: AsyncApp,
101 ) -> Self {
102 let assertions = AssertionsReport::new(meta.max_assertions);
103
104 Self {
105 meta,
106 log_prefix,
107 agent_thread,
108 assertions,
109 model,
110 app,
111 tool_metrics: Arc::new(Mutex::new(ToolMetrics::default())),
112 }
113 }
114
115 pub fn push_user_message(&mut self, text: impl ToString) {
116 self.app
117 .update_entity(&self.agent_thread, |thread, cx| {
118 thread.insert_user_message(text.to_string(), vec![], None, cx);
119 })
120 .unwrap();
121 }
122
123 pub fn assert(&mut self, expected: bool, message: impl ToString) -> Result<()> {
124 let message = message.to_string();
125 self.log_assertion(
126 if expected {
127 Ok(())
128 } else {
129 Err(anyhow::Error::from(FailedAssertion(message.clone())))
130 },
131 message,
132 )
133 }
134
135 pub fn assert_some<T>(&mut self, option: Option<T>, message: impl ToString) -> Result<T> {
136 let message = message.to_string();
137 self.log_assertion(
138 match option {
139 Some(value) => Ok(value),
140 None => Err(anyhow::Error::from(FailedAssertion(message.clone()))),
141 },
142 message,
143 )
144 }
145
146 #[allow(dead_code)]
147 pub fn assert_eq<T: PartialEq + Debug>(
148 &mut self,
149 left: T,
150 right: T,
151 message: impl ToString,
152 ) -> Result<()> {
153 let message = message.to_string();
154 self.log_assertion(
155 if left == right {
156 Ok(())
157 } else {
158 println!("{}{:#?} != {:#?}", self.log_prefix, left, right);
159 Err(anyhow::Error::from(FailedAssertion(message.clone())))
160 },
161 message,
162 )
163 }
164
165 fn log_assertion<T>(&mut self, result: Result<T>, message: String) -> Result<T> {
166 if let Some(max) = self.meta.max_assertions {
167 if self.assertions.run_count() > max {
168 return Err(anyhow!(
169 "More assertions were run than the stated max_assertions of {}",
170 max
171 ));
172 }
173 }
174
175 self.assertions.ran.push(RanAssertion {
176 id: message.clone(),
177 result: Ok(RanAssertionResult {
178 analysis: None,
179 passed: result.is_ok(),
180 }),
181 });
182
183 if result.is_ok() {
184 println!("{}✅ {}", self.log_prefix, message);
185 } else {
186 println!("{}❌ {}", self.log_prefix, message);
187 }
188
189 result
190 }
191
192 pub async fn run_to_end(&mut self) -> Result<Response> {
193 self.run_turns(u32::MAX).await
194 }
195
196 pub async fn run_turn(&mut self) -> Result<Response> {
197 self.run_turns(1).await
198 }
199
200 pub async fn run_turns(&mut self, iterations: u32) -> Result<Response> {
201 let (mut tx, mut rx) = mpsc::channel(1);
202
203 let tool_metrics = self.tool_metrics.clone();
204 let log_prefix = self.log_prefix.clone();
205 let _subscription = self.app.subscribe(
206 &self.agent_thread,
207 move |thread, event: &ThreadEvent, cx| match event {
208 ThreadEvent::ShowError(thread_error) => {
209 tx.try_send(Err(anyhow!(thread_error.clone()))).ok();
210 }
211 ThreadEvent::Stopped(reason) => match reason {
212 Ok(StopReason::EndTurn) => {
213 tx.close_channel();
214 }
215 Ok(StopReason::ToolUse) => {
216 if thread.read(cx).remaining_turns() == 0 {
217 tx.close_channel();
218 }
219 }
220 Ok(StopReason::MaxTokens) => {
221 tx.try_send(Err(anyhow!("Exceeded maximum tokens"))).ok();
222 }
223 Err(err) => {
224 tx.try_send(Err(anyhow!(err.clone()))).ok();
225 }
226 },
227 ThreadEvent::StreamedAssistantText(_, _)
228 | ThreadEvent::StreamedAssistantThinking(_, _)
229 | ThreadEvent::UsePendingTools { .. } => {}
230 ThreadEvent::ToolFinished {
231 tool_use_id,
232 pending_tool_use,
233 ..
234 } => {
235 thread.update(cx, |thread, _cx| {
236 if let Some(tool_use) = pending_tool_use {
237 let mut tool_metrics = tool_metrics.lock().unwrap();
238 if let Some(tool_result) = thread.tool_result(&tool_use_id) {
239 let message = if tool_result.is_error {
240 format!("✖︎ {}", tool_use.name)
241 } else {
242 format!("✔︎ {}", tool_use.name)
243 };
244 println!("{log_prefix}{message}");
245 tool_metrics
246 .insert(tool_result.tool_name.clone(), !tool_result.is_error);
247 } else {
248 let message =
249 format!("TOOL FINISHED WITHOUT RESULT: {}", tool_use.name);
250 println!("{log_prefix}{message}");
251 tool_metrics.insert(tool_use.name.clone(), true);
252 }
253 }
254 });
255 }
256 ThreadEvent::ToolConfirmationNeeded => {
257 panic!(
258 "{}Bug: Tool confirmation should not be required in eval",
259 log_prefix
260 );
261 }
262 ThreadEvent::StreamedCompletion
263 | ThreadEvent::MessageAdded(_)
264 | ThreadEvent::MessageEdited(_)
265 | ThreadEvent::MessageDeleted(_)
266 | ThreadEvent::SummaryChanged
267 | ThreadEvent::SummaryGenerated
268 | ThreadEvent::ReceivedTextChunk
269 | ThreadEvent::StreamedToolUse { .. }
270 | ThreadEvent::CheckpointChanged
271 | ThreadEvent::UsageUpdated(_) => {
272 tx.try_send(Ok(())).ok();
273 if std::env::var("ZED_EVAL_DEBUG").is_ok() {
274 println!("{}Event: {:#?}", log_prefix, event);
275 }
276 }
277 },
278 );
279
280 let model = self.model.clone();
281
282 let message_count_before = self.app.update_entity(&self.agent_thread, |thread, cx| {
283 thread.set_remaining_turns(iterations);
284 thread.send_to_model(model, None, cx);
285 thread.messages().len()
286 })?;
287
288 loop {
289 select_biased! {
290 result = rx.next() => {
291 if let Some(result) = result {
292 result?;
293 } else {
294 break;
295 }
296 }
297 _ = self.app.background_executor().timer(THREAD_EVENT_TIMEOUT).fuse() => {
298 return Err(anyhow!("Agentic loop stalled - waited {:?} without any events", THREAD_EVENT_TIMEOUT));
299 }
300 }
301 }
302
303 let messages = self.app.read_entity(&self.agent_thread, |thread, cx| {
304 let mut messages = Vec::new();
305 for message in thread.messages().skip(message_count_before) {
306 messages.push(Message {
307 _role: message.role,
308 _text: message.to_string(),
309 tool_use: thread
310 .tool_uses_for_message(message.id, cx)
311 .into_iter()
312 .map(|tool_use| ToolUse {
313 name: tool_use.name.to_string(),
314 value: tool_use.input,
315 })
316 .collect(),
317 });
318 }
319 messages
320 })?;
321
322 let response = Response::new(messages);
323
324 Ok(response)
325 }
326
327 pub fn edits(&self) -> HashMap<Arc<Path>, FileEdits> {
328 self.app
329 .read_entity(&self.agent_thread, |thread, cx| {
330 let action_log = thread.action_log().read(cx);
331 HashMap::from_iter(action_log.changed_buffers(cx).into_iter().map(
332 |(buffer, diff)| {
333 let snapshot = buffer.read(cx).snapshot();
334
335 let file = snapshot.file().unwrap();
336 let diff = diff.read(cx);
337 let base_text = diff.base_text().text();
338
339 let hunks = diff
340 .hunks(&snapshot, cx)
341 .map(|hunk| FileEditHunk {
342 base_text: base_text[hunk.diff_base_byte_range.clone()].to_string(),
343 text: snapshot
344 .text_for_range(hunk.range.clone())
345 .collect::<String>(),
346 status: hunk.status(),
347 })
348 .collect();
349
350 (file.path().clone(), FileEdits { hunks })
351 },
352 ))
353 })
354 .unwrap()
355 }
356}
357
358#[derive(Debug)]
359pub struct Response {
360 messages: Vec<Message>,
361}
362
363impl Response {
364 pub fn new(messages: Vec<Message>) -> Self {
365 Self { messages }
366 }
367
368 pub fn expect_tool(
369 &self,
370 tool_name: &'static str,
371 cx: &mut ExampleContext,
372 ) -> Result<&ToolUse> {
373 let result = self.messages.iter().find_map(|msg| {
374 msg.tool_use
375 .iter()
376 .find(|tool_use| tool_use.name == tool_name)
377 });
378 cx.assert_some(result, format!("called `{}`", tool_name))
379 }
380
381 pub fn tool_uses(&self) -> impl Iterator<Item = &ToolUse> {
382 self.messages.iter().flat_map(|msg| &msg.tool_use)
383 }
384}
385
386#[derive(Debug)]
387pub struct Message {
388 _role: Role,
389 _text: String,
390 tool_use: Vec<ToolUse>,
391}
392
393#[derive(Debug)]
394pub struct ToolUse {
395 pub name: String,
396 value: serde_json::Value,
397}
398
399impl ToolUse {
400 pub fn parse_input<Input>(&self) -> Result<Input>
401 where
402 Input: for<'de> serde::Deserialize<'de>,
403 {
404 serde_json::from_value::<Input>(self.value.clone()).map_err(|err| anyhow!(err))
405 }
406}
407
408#[derive(Debug)]
409pub struct FileEdits {
410 hunks: Vec<FileEditHunk>,
411}
412
413#[derive(Debug)]
414struct FileEditHunk {
415 base_text: String,
416 text: String,
417 status: DiffHunkStatus,
418}
419
420impl FileEdits {
421 pub fn has_added_line(&self, line: &str) -> bool {
422 self.hunks.iter().any(|hunk| {
423 hunk.status == DiffHunkStatus::added_none()
424 && hunk.base_text.is_empty()
425 && hunk.text.contains(line)
426 })
427 }
428}