From 1b614ef63be00be15aaef17d2437cbf955fca1cc Mon Sep 17 00:00:00 2001 From: Marshall Bowers Date: Thu, 25 Apr 2024 22:21:18 -0400 Subject: [PATCH] Add an Assistant example that can interact with the filesystem (#11027) This PR adds a new Assistant example that is able to interact with the filesystem using a tool. Release Notes: - N/A --- Cargo.lock | 1 + crates/assistant2/Cargo.toml | 5 +- .../assistant2/examples/file_interactions.rs | 221 ++++++++++++++++++ 3 files changed, 225 insertions(+), 2 deletions(-) create mode 100644 crates/assistant2/examples/file_interactions.rs diff --git a/Cargo.lock b/Cargo.lock index 4c206e28c8..d36db0c3a3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -382,6 +382,7 @@ dependencies = [ "editor", "env_logger", "feature_flags", + "fs", "futures 0.3.28", "gpui", "language", diff --git a/crates/assistant2/Cargo.toml b/crates/assistant2/Cargo.toml index 886a84c863..4a0703f27b 100644 --- a/crates/assistant2/Cargo.toml +++ b/crates/assistant2/Cargo.toml @@ -19,15 +19,17 @@ assistant_tooling.workspace = true client.workspace = true editor.workspace = true feature_flags.workspace = true +fs.workspace = true futures.workspace = true gpui.workspace = true language.workspace = true log.workspace = true +nanoid = "0.4" open_ai.workspace = true project.workspace = true rich_text.workspace = true -semantic_index.workspace = true schemars.workspace = true +semantic_index.workspace = true serde.workspace = true serde_json.workspace = true settings.workspace = true @@ -35,7 +37,6 @@ theme.workspace = true ui.workspace = true util.workspace = true workspace.workspace = true -nanoid = "0.4" [dev-dependencies] assets.workspace = true diff --git a/crates/assistant2/examples/file_interactions.rs b/crates/assistant2/examples/file_interactions.rs new file mode 100644 index 0000000000..c810085b86 --- /dev/null +++ b/crates/assistant2/examples/file_interactions.rs @@ -0,0 +1,221 @@ +//! This example creates a basic Chat UI for interacting with the filesystem. + +use anyhow::{Context as _, Result}; +use assets::Assets; +use assistant2::AssistantPanel; +use assistant_tooling::{LanguageModelTool, ToolRegistry}; +use client::Client; +use fs::Fs; +use futures::StreamExt; +use gpui::{actions, App, AppContext, KeyBinding, Task, View, WindowOptions}; +use language::LanguageRegistry; +use project::Project; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; +use settings::{KeymapFile, DEFAULT_KEYMAP_PATH}; +use std::path::PathBuf; +use std::sync::Arc; +use theme::LoadThemes; +use ui::{div, prelude::*, Render}; +use util::ResultExt as _; + +actions!(example, [Quit]); + +struct FileBrowserTool { + fs: Arc, + root_dir: PathBuf, +} + +impl FileBrowserTool { + fn new(fs: Arc, root_dir: PathBuf) -> Self { + Self { fs, root_dir } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)] +struct FileBrowserParams { + command: FileBrowserCommand, +} + +#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)] +enum FileBrowserCommand { + Ls { path: PathBuf }, + Cat { path: PathBuf }, +} + +#[derive(Serialize, Deserialize)] +enum FileBrowserOutput { + Ls { entries: Vec }, + Cat { content: String }, +} + +pub struct FileBrowserView { + result: Result, +} + +impl Render for FileBrowserView { + fn render(&mut self, cx: &mut ViewContext) -> impl IntoElement { + let Ok(output) = self.result.as_ref() else { + return h_flex().child("Failed to perform operation"); + }; + + match output { + FileBrowserOutput::Ls { entries } => v_flex().children( + entries + .into_iter() + .map(|entry| h_flex().text_ui(cx).child(entry.clone())), + ), + FileBrowserOutput::Cat { content } => h_flex().child(content.clone()), + } + } +} + +impl LanguageModelTool for FileBrowserTool { + type Input = FileBrowserParams; + type Output = FileBrowserOutput; + type View = FileBrowserView; + + fn name(&self) -> String { + "file_browser".to_string() + } + + fn description(&self) -> String { + "A tool for browsing the filesystem.".to_string() + } + + fn execute(&self, input: &Self::Input, cx: &AppContext) -> Task> { + cx.spawn({ + let fs = self.fs.clone(); + let root_dir = self.root_dir.clone(); + let input = input.clone(); + |_cx| async move { + match input.command { + FileBrowserCommand::Ls { path } => { + let path = root_dir.join(path); + + let mut output = fs.read_dir(&path).await?; + + let mut entries = Vec::new(); + while let Some(entry) = output.next().await { + let entry = entry?; + entries.push(entry.display().to_string()); + } + + Ok(FileBrowserOutput::Ls { entries }) + } + FileBrowserCommand::Cat { path } => { + let path = root_dir.join(path); + + let output = fs.load(&path).await?; + + Ok(FileBrowserOutput::Cat { content: output }) + } + } + } + }) + } + + fn new_view( + _tool_call_id: String, + _input: Self::Input, + result: Result, + cx: &mut WindowContext, + ) -> gpui::View { + cx.new_view(|_cx| FileBrowserView { result }) + } + + fn format(_input: &Self::Input, output: &Result) -> String { + let Ok(output) = output else { + return "Failed to perform command: {input:?}".to_string(); + }; + + match output { + FileBrowserOutput::Ls { entries } => entries.join("\n"), + FileBrowserOutput::Cat { content } => content.to_owned(), + } + } +} + +fn main() { + env_logger::init(); + App::new().with_assets(Assets).run(|cx| { + cx.bind_keys(Some(KeyBinding::new("cmd-q", Quit, None))); + cx.on_action(|_: &Quit, cx: &mut AppContext| { + cx.quit(); + }); + + settings::init(cx); + language::init(cx); + Project::init_settings(cx); + editor::init(cx); + theme::init(LoadThemes::JustBase, cx); + Assets.load_fonts(cx).unwrap(); + KeymapFile::load_asset(DEFAULT_KEYMAP_PATH, cx).unwrap(); + client::init_settings(cx); + release_channel::init("0.130.0", cx); + + let client = Client::production(cx); + { + let client = client.clone(); + cx.spawn(|cx| async move { client.authenticate_and_connect(false, &cx).await }) + .detach_and_log_err(cx); + } + assistant2::init(client.clone(), cx); + + let language_registry = Arc::new(LanguageRegistry::new( + Task::ready(()), + cx.background_executor().clone(), + )); + let node_runtime = node_runtime::RealNodeRuntime::new(client.http_client()); + languages::init(language_registry.clone(), node_runtime, cx); + + cx.spawn(|cx| async move { + cx.update(|cx| { + let fs = Arc::new(fs::RealFs::new(None)); + let cwd = std::env::current_dir().expect("Failed to get current working directory"); + + let mut tool_registry = ToolRegistry::new(); + tool_registry + .register(FileBrowserTool::new(fs, cwd)) + .context("failed to register FileBrowserTool") + .log_err(); + + let tool_registry = Arc::new(tool_registry); + + println!("Tools registered"); + for definition in tool_registry.definitions() { + println!("{}", definition); + } + + cx.open_window(WindowOptions::default(), |cx| { + cx.new_view(|cx| Example::new(language_registry, tool_registry, cx)) + }); + cx.activate(true); + }) + }) + .detach_and_log_err(cx); + }) +} + +struct Example { + assistant_panel: View, +} + +impl Example { + fn new( + language_registry: Arc, + tool_registry: Arc, + cx: &mut ViewContext, + ) -> Self { + Self { + assistant_panel: cx + .new_view(|cx| AssistantPanel::new(language_registry, tool_registry, cx)), + } + } +} + +impl Render for Example { + fn render(&mut self, _cx: &mut ViewContext) -> impl ui::prelude::IntoElement { + div().size_full().child(self.assistant_panel.clone()) + } +}