463 lines · 15.6 KB · Rust
RawDownload
use crate::{schema::json_schema_for, ui::ToolCallCardHeader};use anyhow::{Result, anyhow};use assistant_tool::{    ActionLog, Tool, ToolCard, ToolResult, ToolResultContent, ToolResultOutput, ToolUseStatus,};use editor::Editor;use futures::channel::oneshot::{self, Receiver};use gpui::{    AnyWindowHandle, App, AppContext, Context, Entity, IntoElement, Task, WeakEntity, Window,};use language;use language_model::{LanguageModel, LanguageModelRequest, LanguageModelToolSchemaFormat};use project::Project;use schemars::JsonSchema;use serde::{Deserialize, Serialize};use std::fmt::Write;use std::{cmp, path::PathBuf, sync::Arc};use ui::{Disclosure, Tooltip, prelude::*};use util::{ResultExt, paths::PathMatcher};use workspace::Workspace; #[derive(Debug, Serialize, Deserialize, JsonSchema)]pub struct FindPathToolInput {    /// The glob to match against every path in the project.    ///    /// <example>    /// If the project has the following root directories:    ///    /// - directory1/a/something.txt    /// - directory2/a/things.txt    /// - directory3/a/other.txt    ///    /// You can get back the first two paths by providing a glob of "*thing*.txt"    /// </example>    pub glob: String,     /// Optional starting position for paginated results (0-based).    /// When not provided, starts from the beginning.    #[serde(default)]    pub offset: usize,} #[derive(Debug, Serialize, Deserialize)]struct FindPathToolOutput {    glob: String,    paths: Vec<PathBuf>,} const RESULTS_PER_PAGE: usize = 50; pub struct FindPathTool; impl Tool for FindPathTool {    fn name(&self) -> String {        "find_path".into()    }     fn needs_confirmation(&self, _: &serde_json::Value, _: &Entity<Project>, _: &App) -> bool {        false    }     fn may_perform_edits(&self) -> bool {        false    }     fn description(&self) -> String {        include_str!("./find_path_tool/description.md").into()    }     fn icon(&self) -> IconName {        IconName::ToolSearch    }     fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> Result<serde_json::Value> {        json_schema_for::<FindPathToolInput>(format)    }     fn ui_text(&self, input: &serde_json::Value) -> String {        match serde_json::from_value::<FindPathToolInput>(input.clone()) {            Ok(input) => format!("Find paths matching “`{}`”", input.glob),            Err(_) => "Search paths".to_string(),        }    }     fn run(        self: Arc<Self>,        input: serde_json::Value,        _request: Arc<LanguageModelRequest>,        project: Entity<Project>,        _action_log: Entity<ActionLog>,        _model: Arc<dyn LanguageModel>,        _window: Option<AnyWindowHandle>,        cx: &mut App,    ) -> ToolResult {        let (offset, glob) = match serde_json::from_value::<FindPathToolInput>(input) {            Ok(input) => (input.offset, input.glob),            Err(err) => return Task::ready(Err(anyhow!(err))).into(),        };         let (sender, receiver) = oneshot::channel();         let card = cx.new(|cx| FindPathToolCard::new(glob.clone(), receiver, cx));         let search_paths_task = search_paths(&glob, project, cx);         let task = cx.background_spawn(async move {            let matches = search_paths_task.await?;            let paginated_matches: &[PathBuf] = &matches[cmp::min(offset, matches.len())                ..cmp::min(offset + RESULTS_PER_PAGE, matches.len())];             sender.send(paginated_matches.to_vec()).log_err();             if matches.is_empty() {                Ok("No matches found".to_string().into())            } else {                let mut message = format!("Found {} total matches.", matches.len());                if matches.len() > RESULTS_PER_PAGE {                    write!(                        &mut message,                        "\nShowing results {}-{} (provide 'offset' parameter for more results):",                        offset + 1,                        offset + paginated_matches.len()                    )                    .unwrap();                }                 for mat in matches.iter().skip(offset).take(RESULTS_PER_PAGE) {                    write!(&mut message, "\n{}", mat.display()).unwrap();                }                 let output = FindPathToolOutput {                    glob,                    paths: matches,                };                 Ok(ToolResultOutput {                    content: ToolResultContent::Text(message),                    output: Some(serde_json::to_value(output)?),                })            }        });         ToolResult {            output: task,            card: Some(card.into()),        }    }     fn deserialize_card(        self: Arc<Self>,        output: serde_json::Value,        _project: Entity<Project>,        _window: &mut Window,        cx: &mut App,    ) -> Option<assistant_tool::AnyToolCard> {        let output = serde_json::from_value::<FindPathToolOutput>(output).ok()?;        let card = cx.new(|_| FindPathToolCard::from_output(output));        Some(card.into())    }} fn search_paths(glob: &str, project: Entity<Project>, cx: &mut App) -> Task<Result<Vec<PathBuf>>> {    let path_matcher = match PathMatcher::new([        // Sometimes models try to search for "". In this case, return all paths in the project.        if glob.is_empty() { "*" } else { glob },    ]) {        Ok(matcher) => matcher,        Err(err) => return Task::ready(Err(anyhow!("Invalid glob: {err}"))),    };    let snapshots: Vec<_> = project        .read(cx)        .worktrees(cx)        .map(|worktree| worktree.read(cx).snapshot())        .collect();     cx.background_spawn(async move {        Ok(snapshots            .iter()            .flat_map(|snapshot| {                let root_name = PathBuf::from(snapshot.root_name());                snapshot                    .entries(false, 0)                    .map(move |entry| root_name.join(&entry.path))                    .filter(|path| path_matcher.is_match(&path))            })            .collect())    })} struct FindPathToolCard {    paths: Vec<PathBuf>,    expanded: bool,    glob: String,    _receiver_task: Option<Task<Result<()>>>,} impl FindPathToolCard {    fn new(glob: String, receiver: Receiver<Vec<PathBuf>>, cx: &mut Context<Self>) -> Self {        let _receiver_task = cx.spawn(async move |this, cx| {            let paths = receiver.await?;             this.update(cx, |this, _cx| {                this.paths = paths;            })            .log_err();             Ok(())        });         Self {            paths: Vec::new(),            expanded: false,            glob,            _receiver_task: Some(_receiver_task),        }    }     fn from_output(output: FindPathToolOutput) -> Self {        Self {            glob: output.glob,            paths: output.paths,            expanded: false,            _receiver_task: None,        }    }} impl ToolCard for FindPathToolCard {    fn render(        &mut self,        _status: &ToolUseStatus,        _window: &mut Window,        workspace: WeakEntity<Workspace>,        cx: &mut Context<Self>,    ) -> impl IntoElement {        let matches_label: SharedString = if self.paths.len() == 0 {            "No matches".into()        } else if self.paths.len() == 1 {            "1 match".into()        } else {            format!("{} matches", self.paths.len()).into()        };         let content = if !self.paths.is_empty() && self.expanded {            Some(                v_flex()                    .relative()                    .ml_1p5()                    .px_1p5()                    .gap_0p5()                    .border_l_1()                    .border_color(cx.theme().colors().border_variant)                    .children(self.paths.iter().enumerate().map(|(index, path)| {                        let path_clone = path.clone();                        let workspace_clone = workspace.clone();                        let button_label = path.to_string_lossy().to_string();                         Button::new(("path", index), button_label)                            .icon(IconName::ArrowUpRight)                            .icon_size(IconSize::XSmall)                            .icon_position(IconPosition::End)                            .label_size(LabelSize::Small)                            .color(Color::Muted)                            .tooltip(Tooltip::text("Jump to File"))                            .on_click(move |_, window, cx| {                                workspace_clone                                    .update(cx, |workspace, cx| {                                        let path = PathBuf::from(&path_clone);                                        let Some(project_path) = workspace                                            .project()                                            .read(cx)                                            .find_project_path(&path, cx)                                        else {                                            return;                                        };                                        let open_task = workspace.open_path(                                            project_path,                                            None,                                            true,                                            window,                                            cx,                                        );                                        window                                            .spawn(cx, async move |cx| {                                                let item = open_task.await?;                                                if let Some(active_editor) =                                                    item.downcast::<Editor>()                                                {                                                    active_editor                                                        .update_in(cx, |editor, window, cx| {                                                            editor.go_to_singleton_buffer_point(                                                                language::Point::new(0, 0),                                                                window,                                                                cx,                                                            );                                                        })                                                        .log_err();                                                }                                                anyhow::Ok(())                                            })                                            .detach_and_log_err(cx);                                    })                                    .ok();                            })                    }))                    .into_any(),            )        } else {            None        };         v_flex()            .mb_2()            .gap_1()            .child(                ToolCallCardHeader::new(IconName::ToolSearch, matches_label)                    .with_code_path(&self.glob)                    .disclosure_slot(                        Disclosure::new("path-search-disclosure", self.expanded)                            .opened_icon(IconName::ChevronUp)                            .closed_icon(IconName::ChevronDown)                            .disabled(self.paths.is_empty())                            .on_click(cx.listener(move |this, _, _, _cx| {                                this.expanded = !this.expanded;                            })),                    ),            )            .children(content)    }} impl Component for FindPathTool {    fn scope() -> ComponentScope {        ComponentScope::Agent    }     fn sort_name() -> &'static str {        "FindPathTool"    }     fn preview(window: &mut Window, cx: &mut App) -> Option<AnyElement> {        let successful_card = cx.new(|_| FindPathToolCard {            paths: vec![                PathBuf::from("src/main.rs"),                PathBuf::from("src/lib.rs"),                PathBuf::from("tests/test.rs"),            ],            expanded: true,            glob: "*.rs".to_string(),            _receiver_task: None,        });         let empty_card = cx.new(|_| FindPathToolCard {            paths: Vec::new(),            expanded: false,            glob: "*.nonexistent".to_string(),            _receiver_task: None,        });         Some(            v_flex()                .gap_6()                .children(vec![example_group(vec![                    single_example(                        "With Paths",                        div()                            .size_full()                            .child(successful_card.update(cx, |tool, cx| {                                tool.render(                                    &ToolUseStatus::Finished("".into()),                                    window,                                    WeakEntity::new_invalid(),                                    cx,                                )                                .into_any_element()                            }))                            .into_any_element(),                    ),                    single_example(                        "No Paths",                        div()                            .size_full()                            .child(empty_card.update(cx, |tool, cx| {                                tool.render(                                    &ToolUseStatus::Finished("".into()),                                    window,                                    WeakEntity::new_invalid(),                                    cx,                                )                                .into_any_element()                            }))                            .into_any_element(),                    ),                ])])                .into_any_element(),        )    }} #[cfg(test)]mod test {    use super::*;    use gpui::TestAppContext;    use project::{FakeFs, Project};    use settings::SettingsStore;    use util::path;     #[gpui::test]    async fn test_find_path_tool(cx: &mut TestAppContext) {        init_test(cx);         let fs = FakeFs::new(cx.executor());        fs.insert_tree(            "/root",            serde_json::json!({                "apple": {                    "banana": {                        "carrot": "1",                    },                    "bandana": {                        "carbonara": "2",                    },                    "endive": "3"                }            }),        )        .await;        let project = Project::test(fs.clone(), [path!("/root").as_ref()], cx).await;         let matches = cx            .update(|cx| search_paths("root/**/car*", project.clone(), cx))            .await            .unwrap();        assert_eq!(            matches,            &[                PathBuf::from("root/apple/banana/carrot"),                PathBuf::from("root/apple/bandana/carbonara")            ]        );         let matches = cx            .update(|cx| search_paths("**/car*", project.clone(), cx))            .await            .unwrap();        assert_eq!(            matches,            &[                PathBuf::from("root/apple/banana/carrot"),                PathBuf::from("root/apple/bandana/carbonara")            ]        );    }     fn init_test(cx: &mut TestAppContext) {        cx.update(|cx| {            let settings_store = SettingsStore::test(cx);            cx.set_global(settings_store);            language::init(cx);            Project::init_settings(cx);        });    }}