1201 lines · 43.4 KB · Rust
RawDownload
use agent::{ThreadEvent, ThreadStore};use anyhow::{Context as _, Result, anyhow};use assistant_tool::ToolWorkingSet;use client::proto::LspWorkProgress;use collections::HashMap;use dap::DapRegistry;use futures::channel::mpsc;use futures::{FutureExt, StreamExt as _, select_biased};use gpui::{App, AppContext as _, AsyncApp, Entity, Task};use handlebars::Handlebars;use language::{Buffer, DiagnosticSeverity, OffsetRangeExt};use language_model::{    LanguageModel, LanguageModelCompletionEvent, LanguageModelRequest, LanguageModelRequestMessage,    MessageContent, Role, StopReason, TokenUsage,};use project::{Project, ProjectPath};use serde::{Deserialize, Serialize};use std::cell::RefCell;use std::fmt::Write as _;use std::fs::File;use std::io::Write as _;use std::rc::Rc;use std::sync::{Arc, Mutex};use std::time::Duration;use std::{    fs,    path::{Path, PathBuf},};use unindent::Unindent as _;use util::ResultExt as _;use util::command::new_smol_command;use util::markdown::MarkdownString;use util::serde::default_true; use crate::AgentAppState; pub const EXAMPLES_DIR: &str = "./crates/eval/examples";pub const REPOS_DIR: &str = "./crates/eval/repos";pub const WORKTREES_DIR: &str = "./crates/eval/worktrees"; const THREAD_EVENT_TIMEOUT: Duration = Duration::from_secs(60 * 2); #[derive(Clone, Debug, Deserialize)]pub struct ExampleBase {    pub url: String,    pub revision: String,    pub language_extension: Option<String>,    pub insert_id: Option<String>,    #[serde(default = "default_true")]    pub require_lsp: bool,    #[serde(default)]    pub allow_preexisting_diagnostics: bool,} impl ExampleBase {    pub fn repo_name(&self) -> String {        self.url            .split('/')            .next_back()            .unwrap_or(&"")            .trim_end_matches(".git")            .into()    }} #[derive(Clone, Debug)]pub struct Example {    pub name: String,    /// Content of `base.toml`    pub base: ExampleBase,    /// Content of `prompt.md`    pub prompt: String,    /// Content of `diff_criteria.md`    pub diff_criteria: String,    /// Content of `thread_criteria.md`, if that file exists (it's optional)    pub thread_criteria: Option<String>,    /// Path to the directory containing the requests and responses for the agentic loop    pub run_directory_path: PathBuf,    /// Prefix used for logging that identifies this example    pub log_prefix: String,} #[derive(Debug, Serialize, Deserialize, Clone)]pub struct RunOutput {    pub repository_diff: String,    pub ran_diagnostics_check: bool,    pub diagnostics_before: Option<String>,    pub diagnostics_after: Option<String>,    pub response_count: usize,    pub token_usage: TokenUsage,    pub tool_use_counts: HashMap<Arc<str>, u32>,    pub last_request: LanguageModelRequest,} #[derive(Debug, Clone, Serialize, Deserialize)]pub struct JudgeDiffInput {    pub repository_diff: String,    pub ran_diagnostics_check: bool,    #[serde(skip_serializing_if = "Option::is_none")]    pub diagnostics_before: Option<String>,    #[serde(skip_serializing_if = "Option::is_none")]    pub diagnostics_after: Option<String>,    pub criteria: String,} #[derive(Debug, Clone, Serialize, Deserialize)]pub struct JudgeThreadInput {    pub messages: String,    pub criteria: String,} #[derive(Debug, Clone, Serialize, Deserialize)]pub struct JudgeResponse {    pub analysis: String,    pub score: u32,} #[derive(Debug, Clone, Serialize, Deserialize)]pub struct JudgeOutput {    pub thread: Option<JudgeResponse>,    pub diff: JudgeResponse,} impl Example {    /// Load an example from a directory containing base.toml, prompt.md, and criteria.md    pub fn load_from_directory(dir_path: &Path, run_dir: &Path) -> Result<Self> {        let name = Self::name_from_path(dir_path);        let base_path = dir_path.join("base.toml");        let prompt_path = dir_path.join("prompt.md");        let diff_criteria_path = dir_path.join("diff_criteria.md");        let thread_criteria_path = dir_path.join("thread_criteria.md");        let thread_criteria = if thread_criteria_path.exists() {            Some(fs::read_to_string(thread_criteria_path.clone())?)        } else {            None        };         Ok(Example {            name: name.clone(),            base: toml::from_str(&fs::read_to_string(&base_path)?)?,            prompt: fs::read_to_string(prompt_path.clone())?,            thread_criteria,            diff_criteria: fs::read_to_string(diff_criteria_path.clone())?,            run_directory_path: run_dir.to_path_buf(),            log_prefix: name,        })    }     pub fn set_repetition_number(&mut self, repetition_number: u32) {        if repetition_number > 0 {            self.name = format!("{}-{}", self.name, repetition_number);        }    }     pub fn example_output_directory(&self) -> PathBuf {        self.run_directory_path.join(&self.name)    }     pub fn set_log_prefix_style(&mut self, color: &str, name_width: usize) {        self.log_prefix = format!(            "{}{:<width$}\x1b[0m | ",            color,            self.name,            width = name_width        );    }     pub fn name_from_path(path: &Path) -> String {        path.file_name().unwrap().to_string_lossy().to_string()    }     pub fn worktree_path(&self) -> PathBuf {        Path::new(WORKTREES_DIR)            .canonicalize()            .context(format!("No such directory {WORKTREES_DIR}"))            .unwrap()            .join(&self.name)            .join(self.base.repo_name())    }     /// Set up the example by checking out the specified Git revision    pub async fn setup(&mut self) -> Result<()> {        let repo_path = repo_path_for_url(&self.base.url);         let revision_exists = run_git(            &repo_path,            &["rev-parse", &format!("{}^{{commit}}", self.base.revision)],        )        .await        .is_ok();         if !revision_exists {            println!(                "{}Fetching revision {}",                self.log_prefix, &self.base.revision            );            run_git(                &repo_path,                &["fetch", "--depth", "1", "origin", &self.base.revision],            )            .await?;        }         let worktree_path = self.worktree_path();         if worktree_path.is_dir() {            println!("{}Resetting existing worktree", self.log_prefix);             // TODO: consider including "-x" to remove ignored files. The downside of this is that            // it will also remove build artifacts, and so prevent incremental reuse there.            run_git(&worktree_path, &["clean", "--force", "-d"]).await?;            run_git(&worktree_path, &["reset", "--hard", "HEAD"]).await?;            run_git(&worktree_path, &["checkout", &self.base.revision]).await?;        } else {            println!("{}Creating worktree", self.log_prefix);             let worktree_path_string = worktree_path.to_string_lossy().to_string();             run_git(                &repo_path,                &[                    "worktree",                    "add",                    "-f",                    &worktree_path_string,                    &self.base.revision,                ],            )            .await?;        }         std::fs::create_dir_all(self.example_output_directory())?;         Ok(())    }     pub fn run(        &self,        model: Arc<dyn LanguageModel>,        app_state: Arc<AgentAppState>,        cx: &mut App,    ) -> Task<Result<RunOutput>> {        let project = Project::local(            app_state.client.clone(),            app_state.node_runtime.clone(),            app_state.user_store.clone(),            app_state.languages.clone(),            Arc::new(DapRegistry::default()),            app_state.fs.clone(),            None,            cx,        );         let worktree_path = self.worktree_path();        let worktree = project.update(cx, |project, cx| {            project.create_worktree(&worktree_path, true, cx)        });         let tools = cx.new(|_| ToolWorkingSet::default());        let thread_store =            ThreadStore::load(project.clone(), tools, app_state.prompt_builder.clone(), cx);        let this = self.clone();         cx.spawn(async move |cx| {            let worktree = worktree.await?;             // Wait for worktree scan to finish before choosing a file to open.            worktree                .update(cx, |worktree, _cx| {                    worktree.as_local().unwrap().scan_complete()                })?                .await;             let lsp = if this.base.require_lsp {                let language_extension = this.base.language_extension.as_deref().context(                    "language_extension field is required in base.toml when `require_lsp == true`",                )?;                 // Open a file that matches the language to cause LSP to start.                let language_file = worktree.read_with(cx, |worktree, _cx| {                    worktree                        .files(false, 0)                        .find_map(|e| {                            if e.path.clone().extension().and_then(|ext| ext.to_str())                                == Some(language_extension)                            {                                Some(ProjectPath {                                    worktree_id: worktree.id(),                                    path: e.path.clone(),                                })                            } else {                                None                            }                        })                        .context("Failed to find a file for example language")                })??;                 let open_language_file_buffer_task = project.update(cx, |project, cx| {                    project.open_buffer(language_file.clone(), cx)                })?;                 let language_file_buffer = open_language_file_buffer_task.await?;                 let lsp_open_handle = project.update(cx, |project, cx| {                    project.register_buffer_with_language_servers(&language_file_buffer, cx)                })?;                 wait_for_lang_server(&project, &language_file_buffer, this.log_prefix.clone(), cx).await?;                 Some((lsp_open_handle, language_file_buffer))            } else {                None            };             let diagnostics_before = query_lsp_diagnostics(project.clone(), cx).await?;            if diagnostics_before.is_some() && !this.base.allow_preexisting_diagnostics {                return Err(anyhow!("Example has pre-existing diagnostics. If you want to run this example regardless, set `allow_preexisting_diagnostics` to `true` in `base.toml`"));            }             if std::env::var("ZED_EVAL_SETUP_ONLY").is_ok() {                return Err(anyhow!("Setup only mode"));            }             let thread_store = thread_store.await?;            let thread =                thread_store.update(cx, |thread_store, cx| thread_store.create_thread(cx))?;            let last_request = Rc::new(RefCell::new(None));             thread.update(cx, |thread, _cx| {                let mut request_count = 0;                let example_dir_path = this.example_output_directory();                 let last_request = Rc::clone(&last_request);                thread.set_request_callback(move |request, response_events| {                    *last_request.borrow_mut() = Some(request.clone());                     request_count += 1;                    let messages_file_path = example_dir_path.join(format!("{request_count}.messages.md"));                    let last_messages_file_path = example_dir_path.join("last.messages.md");                    let request_markdown = RequestMarkdown::new(request);                    let response_events_markdown = response_events_to_markdown(response_events);                     let messages = format!("{}\n\n{}", request_markdown.messages, response_events_markdown);                    fs::write(messages_file_path, messages.clone()).expect("failed to write messages file");                    fs::write(last_messages_file_path, messages).expect("failed to write last messages file");                     if request_count == 1 {                        let tools_file_path = example_dir_path.join("tools.md");                        fs::write(tools_file_path, request_markdown.tools).expect("failed to write tools file");                    }                });            })?;             let tool_use_counts: Arc<Mutex<HashMap<Arc<str>, u32>>> =                Mutex::new(HashMap::default()).into();             let (thread_event_tx, mut thread_event_rx) = mpsc::unbounded();             let subscription = cx.subscribe(&thread, move |_thread, event: &ThreadEvent, _cx| {                thread_event_tx.unbounded_send(event.clone()).log_err();            });             let event_handler_task = cx.spawn({                let log_prefix = this.log_prefix.clone();                let tool_use_counts = tool_use_counts.clone();                let thread = thread.downgrade();                async move |cx| {                    loop {                        let event = select_biased! {                            event = thread_event_rx.next() => event,                            _ = cx.background_executor().timer(THREAD_EVENT_TIMEOUT).fuse() => {                                return Err(anyhow!("Agentic loop stalled - waited {:?} without any events", THREAD_EVENT_TIMEOUT));                            }                        };                        let Some(event) = event else {                            return Err(anyhow!("ThreadEvent channel ended early"));                        };                         match event {                            ThreadEvent::Stopped(reason) => match reason {                                Ok(StopReason::EndTurn) => {                                    return Ok(());                                }                                Ok(StopReason::MaxTokens) => {                                    return Err(anyhow!("Exceeded maximum tokens"));                                }                                Ok(StopReason::ToolUse) => {                                    if std::env::var("ZED_EVAL_DEBUG").is_ok() {                                        println!("{}StopReason: Tool use", log_prefix);                                    }                                }                                Err(error) => {                                    return Err(anyhow!(error.clone()));                                }                            },                            ThreadEvent::ShowError(thread_error) => {                                break Err(anyhow!(thread_error.clone()));                            }                            ThreadEvent::StreamedAssistantText(_, _)| ThreadEvent::StreamedAssistantThinking(_, _) | ThreadEvent::UsePendingTools { .. } => {                            }                            ThreadEvent::ToolFinished {                                tool_use_id,                                pending_tool_use,                                ..                            } => {                                thread.update(cx, |thread, _cx| {                                    if let Some(tool_use) = pending_tool_use {                                        if let Some(tool_result) = thread.tool_result(&tool_use_id) {                                            let message = if tool_result.is_error {                                                format!("TOOL FAILED: {}", tool_use.name)                                            } else {                                                format!("TOOL FINISHED: {}", tool_use.name)                                            };                                            println!("{log_prefix}{message}");                                            let mut tool_use_counts = tool_use_counts.lock().unwrap();                                            *tool_use_counts                                                .entry(tool_result.tool_name.clone())                                                .or_insert(0) += 1;                                        } else {                                            let message = format!("TOOL FINISHED WITHOUT RESULT: {}", tool_use.name);                                            println!("{log_prefix}{message}");                                        }                                    }                                })?;                            }                            ThreadEvent::ToolConfirmationNeeded => {                                panic!("{}Bug: Tool confirmation should not be required in eval", log_prefix);                            },                            ThreadEvent::StreamedCompletion |                            ThreadEvent::MessageAdded(_) |                            ThreadEvent::MessageEdited(_) |                            ThreadEvent::MessageDeleted(_) |                            ThreadEvent::SummaryChanged |                            ThreadEvent::SummaryGenerated |                            ThreadEvent::CheckpointChanged |                            ThreadEvent::UsageUpdated(_) => {                                if std::env::var("ZED_EVAL_DEBUG").is_ok() {                                    println!("{}Event: {:#?}", log_prefix, event);                                }                            }                        }                    }                }            });             thread.update(cx, |thread, cx| {                let context = vec![];                thread.insert_user_message(this.prompt.clone(), context, None, cx);                thread.send_to_model(model, cx);            })?;             event_handler_task.await?;             println!("{}Stopped", this.log_prefix);             if let Some((_, language_file_buffer)) = lsp.as_ref() {                wait_for_lang_server(&project, &language_file_buffer, this.log_prefix.clone(), cx).await?;            }             println!("{}Getting repository diff", this.log_prefix);            let repository_diff = this.repository_diff().await?;             let example_output_dir = this.example_output_directory();            let repository_diff_path = example_output_dir.join("patch.diff");            let mut repository_diff_output_file = File::create(&repository_diff_path)?;            writeln!(&mut repository_diff_output_file, "{}", &repository_diff).log_err();             println!("{}Getting diagnostics", this.log_prefix);            let diagnostics_after = cx                .update(move |cx| {                    cx.spawn(async move |cx| query_lsp_diagnostics(project, cx).await)                })?                .await?;            println!("{}Got diagnostics", this.log_prefix);             let Some(last_request) = last_request.borrow_mut().take() else {                return Err(anyhow!("No requests ran."));            };             drop(subscription);            drop(lsp);             if let Some(diagnostics_before) = &diagnostics_before {                fs::write(example_output_dir.join("diagnostics_before.txt"), diagnostics_before)?;            }             if let Some(diagnostics_after) = &diagnostics_after {                fs::write(example_output_dir.join("diagnostics_after.txt"), diagnostics_after)?;            }              thread.update(cx, |thread, _cx| {                let response_count = thread                    .messages()                    .filter(|message| message.role == language_model::Role::Assistant)                    .count();                RunOutput {                    repository_diff,                    ran_diagnostics_check: this.base.require_lsp,                    diagnostics_before,

Lines 1–500 of 1201.

Keyboard shortcuts

General

Show keyboard shortcuts
?
Focus the current search
/

Issues and pull requests

Open filters
f

Issues, pull requests and Inbox

Select next item
j↓
Select previous item
k↑
Open selected item
Enter
Dismiss filters or clear selection
Esc

Repository code

Go to file
t