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::{ContextLoadResult, 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(
119 text.to_string(),
120 ContextLoadResult::default(),
121 None,
122 cx,
123 );
124 })
125 .unwrap();
126 }
127
128 pub fn assert(&mut self, expected: bool, message: impl ToString) -> Result<()> {
129 let message = message.to_string();
130 self.log_assertion(
131 if expected {
132 Ok(())
133 } else {
134 Err(anyhow::Error::from(FailedAssertion(message.clone())))
135 },
136 message,
137 )
138 }
139
140 pub fn assert_some<T>(&mut self, option: Option<T>, message: impl ToString) -> Result<T> {
141 let message = message.to_string();
142 self.log_assertion(
143 match option {
144 Some(value) => Ok(value),
145 None => Err(anyhow::Error::from(FailedAssertion(message.clone()))),
146 },
147 message,
148 )
149 }
150
151 #[allow(dead_code)]
152 pub fn assert_eq<T: PartialEq + Debug>(
153 &mut self,
154 left: T,
155 right: T,
156 message: impl ToString,
157 ) -> Result<()> {
158 let message = message.to_string();
159 self.log_assertion(
160 if left == right {
161 Ok(())
162 } else {
163 println!("{}{:#?} != {:#?}", self.log_prefix, left, right);
164 Err(anyhow::Error::from(FailedAssertion(message.clone())))
165 },
166 message,
167 )
168 }
169
170 fn log_assertion<T>(&mut self, result: Result<T>, message: String) -> Result<T> {
171 if let Some(max) = self.meta.max_assertions {
172 if self.assertions.run_count() > max {
173 return Err(anyhow!(
174 "More assertions were run than the stated max_assertions of {}",
175 max
176 ));
177 }
178 }
179
180 self.assertions.ran.push(RanAssertion {
181 id: message.clone(),
182 result: Ok(RanAssertionResult {
183 analysis: None,
184 passed: result.is_ok(),
185 }),
186 });
187
188 if result.is_ok() {
189 println!("{}✅ {}", self.log_prefix, message);
190 } else {
191 println!("{}❌ {}", self.log_prefix, message);
192 }
193
194 result
195 }
196
197 pub async fn run_to_end(&mut self) -> Result<Response> {
198 self.run_turns(u32::MAX).await
199 }
200
201 pub async fn run_turn(&mut self) -> Result<Response> {
202 self.run_turns(1).await
203 }
204
205 pub async fn run_turns(&mut self, iterations: u32) -> Result<Response> {
206 let (mut tx, mut rx) = mpsc::channel(1);
207
208 let tool_metrics = self.tool_metrics.clone();
209 let log_prefix = self.log_prefix.clone();
210 let _subscription = self.app.subscribe(
211 &self.agent_thread,
212 move |thread, event: &ThreadEvent, cx| match event {
213 ThreadEvent::ShowError(thread_error) => {
214 tx.try_send(Err(anyhow!(thread_error.clone()))).ok();
215 }
216 ThreadEvent::Stopped(reason) => match reason {
217 Ok(StopReason::EndTurn) => {
218 tx.close_channel();
219 }
220 Ok(StopReason::ToolUse) => {
221 if thread.read(cx).remaining_turns() == 0 {
222 tx.close_channel();
223 }
224 }
225 Ok(StopReason::MaxTokens) => {
226 tx.try_send(Err(anyhow!("Exceeded maximum tokens"))).ok();
227 }
228 Err(err) => {
229 tx.try_send(Err(anyhow!(err.clone()))).ok();
230 }
231 },
232 ThreadEvent::StreamedAssistantText(_, _)
233 | ThreadEvent::StreamedAssistantThinking(_, _)
234 | ThreadEvent::UsePendingTools { .. } => {}
235 ThreadEvent::ToolFinished {
236 tool_use_id,
237 pending_tool_use,
238 ..
239 } => {
240 thread.update(cx, |thread, _cx| {
241 if let Some(tool_use) = pending_tool_use {
242 let mut tool_metrics = tool_metrics.lock().unwrap();
243 if let Some(tool_result) = thread.tool_result(&tool_use_id) {
244 let message = if tool_result.is_error {
245 format!("✖︎ {}", tool_use.name)
246 } else {
247 format!("✔︎ {}", tool_use.name)
248 };
249 println!("{log_prefix}{message}");
250 tool_metrics
251 .insert(tool_result.tool_name.clone(), !tool_result.is_error);
252 } else {
253 let message =
254 format!("TOOL FINISHED WITHOUT RESULT: {}", tool_use.name);
255 println!("{log_prefix}{message}");
256 tool_metrics.insert(tool_use.name.clone(), true);
257 }
258 }
259 });
260 }
261 ThreadEvent::InvalidToolInput { .. } => {
262 println!("{log_prefix} invalid tool input");
263 }
264 ThreadEvent::ToolConfirmationNeeded => {
265 panic!(
266 "{}Bug: Tool confirmation should not be required in eval",
267 log_prefix
268 );
269 }
270 ThreadEvent::StreamedCompletion
271 | ThreadEvent::MessageAdded(_)
272 | ThreadEvent::MessageEdited(_)
273 | ThreadEvent::MessageDeleted(_)
274 | ThreadEvent::SummaryChanged
275 | ThreadEvent::SummaryGenerated
276 | ThreadEvent::ReceivedTextChunk
277 | ThreadEvent::StreamedToolUse { .. }
278 | ThreadEvent::CheckpointChanged
279 | ThreadEvent::UsageUpdated(_)
280 | ThreadEvent::CancelEditing => {
281 tx.try_send(Ok(())).ok();
282 if std::env::var("ZED_EVAL_DEBUG").is_ok() {
283 println!("{}Event: {:#?}", log_prefix, event);
284 }
285 }
286 },
287 );
288
289 let model = self.model.clone();
290
291 let message_count_before = self.app.update_entity(&self.agent_thread, |thread, cx| {
292 thread.set_remaining_turns(iterations);
293 thread.send_to_model(model, None, cx);
294 thread.messages().len()
295 })?;
296
297 loop {
298 select_biased! {
299 result = rx.next() => {
300 if let Some(result) = result {
301 result?;
302 } else {
303 break;
304 }
305 }
306 _ = self.app.background_executor().timer(THREAD_EVENT_TIMEOUT).fuse() => {
307 return Err(anyhow!("Agentic loop stalled - waited {:?} without any events", THREAD_EVENT_TIMEOUT));
308 }
309 }
310 }
311
312 let messages = self.app.read_entity(&self.agent_thread, |thread, cx| {
313 let mut messages = Vec::new();
314 for message in thread.messages().skip(message_count_before) {
315 messages.push(Message {
316 _role: message.role,
317 _text: message.to_string(),
318 tool_use: thread
319 .tool_uses_for_message(message.id, cx)
320 .into_iter()
321 .map(|tool_use| ToolUse {
322 name: tool_use.name.to_string(),
323 value: tool_use.input,
324 })
325 .collect(),
326 });
327 }
328 messages
329 })?;
330
331 let response = Response::new(messages);
332
333 Ok(response)
334 }
335
336 pub fn edits(&self) -> HashMap<Arc<Path>, FileEdits> {
337 self.app
338 .read_entity(&self.agent_thread, |thread, cx| {
339 let action_log = thread.action_log().read(cx);
340 HashMap::from_iter(action_log.changed_buffers(cx).into_iter().map(
341 |(buffer, diff)| {
342 let snapshot = buffer.read(cx).snapshot();
343
344 let file = snapshot.file().unwrap();
345 let diff = diff.read(cx);
346 let base_text = diff.base_text().text();
347
348 let hunks = diff
349 .hunks(&snapshot, cx)
350 .map(|hunk| FileEditHunk {
351 base_text: base_text[hunk.diff_base_byte_range.clone()].to_string(),
352 text: snapshot
353 .text_for_range(hunk.range.clone())
354 .collect::<String>(),
355 status: hunk.status(),
356 })
357 .collect();
358
359 (file.path().clone(), FileEdits { hunks })
360 },
361 ))
362 })
363 .unwrap()
364 }
365}
366
367#[derive(Debug)]
368pub struct Response {
369 messages: Vec<Message>,
370}
371
372impl Response {
373 pub fn new(messages: Vec<Message>) -> Self {
374 Self { messages }
375 }
376
377 pub fn expect_tool(
378 &self,
379 tool_name: &'static str,
380 cx: &mut ExampleContext,
381 ) -> Result<&ToolUse> {
382 let result = self.messages.iter().find_map(|msg| {
383 msg.tool_use
384 .iter()
385 .find(|tool_use| tool_use.name == tool_name)
386 });
387 cx.assert_some(result, format!("called `{}`", tool_name))
388 }
389
390 #[allow(dead_code)]
391 pub fn tool_uses(&self) -> impl Iterator<Item = &ToolUse> {
392 self.messages.iter().flat_map(|msg| &msg.tool_use)
393 }
394}
395
396#[derive(Debug)]
397pub struct Message {
398 _role: Role,
399 _text: String,
400 tool_use: Vec<ToolUse>,
401}
402
403#[derive(Debug)]
404pub struct ToolUse {
405 pub name: String,
406 value: serde_json::Value,
407}
408
409impl ToolUse {
410 pub fn parse_input<Input>(&self) -> Result<Input>
411 where
412 Input: for<'de> serde::Deserialize<'de>,
413 {
414 serde_json::from_value::<Input>(self.value.clone()).map_err(|err| anyhow!(err))
415 }
416}
417
418#[derive(Debug)]
419pub struct FileEdits {
420 hunks: Vec<FileEditHunk>,
421}
422
423#[derive(Debug)]
424struct FileEditHunk {
425 base_text: String,
426 text: String,
427 status: DiffHunkStatus,
428}
429
430impl FileEdits {
431 pub fn has_added_line(&self, line: &str) -> bool {
432 self.hunks.iter().any(|hunk| {
433 hunk.status == DiffHunkStatus::added_none()
434 && hunk.base_text.is_empty()
435 && hunk.text.contains(line)
436 })
437 }
438}