From 26b5f340462174544151a464a288b7da2f97a27a Mon Sep 17 00:00:00 2001 From: Marshall Bowers Date: Tue, 14 May 2024 18:39:52 -0400 Subject: [PATCH] assistant: Add basic current project context (#11828) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This PR adds the beginnings of current project context to the Assistant. Currently it supports reading a `Cargo.toml` file and using that to get some basic information about the project, and its dependencies: Screenshot 2024-05-14 at 6 17 03 PM Release Notes: - N/A --------- Co-authored-by: Nate --- Cargo.lock | 2 + crates/assistant/Cargo.toml | 2 + crates/assistant/src/ambient_context.rs | 3 + .../src/ambient_context/current_project.rs | 150 ++++++++++++++++++ .../src/ambient_context/recent_buffers.rs | 12 ++ crates/assistant/src/assistant_panel.rs | 93 +++++++---- 6 files changed, 228 insertions(+), 34 deletions(-) create mode 100644 crates/assistant/src/ambient_context/current_project.rs diff --git a/Cargo.lock b/Cargo.lock index 3109d990ce12678303a4b1bf76038209a9a5b24d..60db91d4edc45012711207f01947c5e50662dbfb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -336,6 +336,7 @@ version = "0.1.0" dependencies = [ "anthropic", "anyhow", + "cargo_toml", "chrono", "client", "collections", @@ -368,6 +369,7 @@ dependencies = [ "telemetry_events", "theme", "tiktoken-rs", + "toml 0.8.10", "ui", "util", "uuid", diff --git a/crates/assistant/Cargo.toml b/crates/assistant/Cargo.toml index eb9caef25ba5257a7ce1bed8c9a339d163e91aa9..35f06fbc959d2cdc88642085f81011c23752beeb 100644 --- a/crates/assistant/Cargo.toml +++ b/crates/assistant/Cargo.toml @@ -12,6 +12,7 @@ doctest = false [dependencies] anyhow.workspace = true anthropic = { workspace = true, features = ["schemars"] } +cargo_toml.workspace = true chrono.workspace = true client.workspace = true collections.workspace = true @@ -41,6 +42,7 @@ smol.workspace = true telemetry_events.workspace = true theme.workspace = true tiktoken-rs.workspace = true +toml.workspace = true ui.workspace = true util.workspace = true uuid.workspace = true diff --git a/crates/assistant/src/ambient_context.rs b/crates/assistant/src/ambient_context.rs index 55fe3e605115aa97baf967587a031ac842e2d7fb..de4195a1dbe257c52a2b015b736e4334b38630fc 100644 --- a/crates/assistant/src/ambient_context.rs +++ b/crates/assistant/src/ambient_context.rs @@ -1,8 +1,11 @@ +mod current_project; mod recent_buffers; +pub use current_project::*; pub use recent_buffers::*; #[derive(Default)] pub struct AmbientContext { pub recent_buffers: RecentBuffersContext, + pub current_project: CurrentProjectContext, } diff --git a/crates/assistant/src/ambient_context/current_project.rs b/crates/assistant/src/ambient_context/current_project.rs new file mode 100644 index 0000000000000000000000000000000000000000..cc776740236715e34f9f1c3a077191a21e6cd6ab --- /dev/null +++ b/crates/assistant/src/ambient_context/current_project.rs @@ -0,0 +1,150 @@ +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::time::Duration; + +use anyhow::{anyhow, Result}; +use fs::Fs; +use gpui::{AsyncAppContext, ModelContext, Task, WeakModel}; +use project::{Project, ProjectPath}; +use util::ResultExt; + +use crate::assistant_panel::Conversation; +use crate::{LanguageModelRequestMessage, Role}; + +/// Ambient context about the current project. +pub struct CurrentProjectContext { + pub enabled: bool, + pub message: String, + pub pending_message: Option>, +} + +#[allow(clippy::derivable_impls)] +impl Default for CurrentProjectContext { + fn default() -> Self { + Self { + enabled: false, + message: String::new(), + pending_message: None, + } + } +} + +impl CurrentProjectContext { + /// Returns the [`CurrentProjectContext`] as a message to the language model. + pub fn to_message(&self) -> Option { + self.enabled.then(|| LanguageModelRequestMessage { + role: Role::System, + content: self.message.clone(), + }) + } + + /// Updates the [`CurrentProjectContext`] for the given [`Project`]. + pub fn update( + &mut self, + fs: Arc, + project: WeakModel, + cx: &mut ModelContext, + ) { + if !self.enabled { + self.message.clear(); + self.pending_message = None; + cx.notify(); + return; + } + + self.pending_message = Some(cx.spawn(|conversation, mut cx| async move { + const DEBOUNCE_TIMEOUT: Duration = Duration::from_millis(100); + cx.background_executor().timer(DEBOUNCE_TIMEOUT).await; + + let Some(path_to_cargo_toml) = Self::path_to_cargo_toml(project, &mut cx).log_err() + else { + return; + }; + + let Some(path_to_cargo_toml) = path_to_cargo_toml + .ok_or_else(|| anyhow!("no Cargo.toml")) + .log_err() + else { + return; + }; + + let message_task = cx + .background_executor() + .spawn(async move { Self::build_message(fs, &path_to_cargo_toml).await }); + + if let Some(message) = message_task.await.log_err() { + conversation + .update(&mut cx, |conversation, _cx| { + conversation.ambient_context.current_project.message = message; + }) + .log_err(); + } + })); + } + + async fn build_message(fs: Arc, path_to_cargo_toml: &Path) -> Result { + let buffer = fs.load(path_to_cargo_toml).await?; + let cargo_toml: cargo_toml::Manifest = toml::from_str(&buffer)?; + + let mut message = String::new(); + + let name = cargo_toml + .package + .as_ref() + .map(|package| package.name.as_str()); + if let Some(name) = name { + message.push_str(&format!(" named \"{name}\"")); + } + message.push_str(". "); + + let description = cargo_toml + .package + .as_ref() + .and_then(|package| package.description.as_ref()) + .and_then(|description| description.get().ok().cloned()); + if let Some(description) = description.as_ref() { + message.push_str("It describes itself as "); + message.push_str(&format!("\"{description}\"")); + message.push_str(". "); + } + + let dependencies = cargo_toml.dependencies.keys().cloned().collect::>(); + if !dependencies.is_empty() { + message.push_str("The following dependencies are installed: "); + message.push_str(&dependencies.join(", ")); + message.push_str(". "); + } + + Ok(message) + } + + fn path_to_cargo_toml( + project: WeakModel, + cx: &mut AsyncAppContext, + ) -> Result> { + cx.update(|cx| { + let worktree = project.update(cx, |project, _cx| { + project + .worktrees() + .next() + .ok_or_else(|| anyhow!("no worktree")) + })??; + + let path_to_cargo_toml = worktree.update(cx, |worktree, _cx| { + let cargo_toml = worktree.entry_for_path("Cargo.toml")?; + Some(ProjectPath { + worktree_id: worktree.id(), + path: cargo_toml.path.clone(), + }) + }); + let path_to_cargo_toml = path_to_cargo_toml.and_then(|path| { + project + .update(cx, |project, cx| project.absolute_path(&path, cx)) + .ok() + .flatten() + }); + + Ok(path_to_cargo_toml) + })? + } +} diff --git a/crates/assistant/src/ambient_context/recent_buffers.rs b/crates/assistant/src/ambient_context/recent_buffers.rs index 4f6280b43c14c2c41d883d207c2cfd1eb9f89232..eceecd6056d904a3ae915cb4fdd95e85c2f08b02 100644 --- a/crates/assistant/src/ambient_context/recent_buffers.rs +++ b/crates/assistant/src/ambient_context/recent_buffers.rs @@ -1,6 +1,8 @@ use gpui::{Subscription, Task, WeakModel}; use language::Buffer; +use crate::{LanguageModelRequestMessage, Role}; + pub struct RecentBuffersContext { pub enabled: bool, pub buffers: Vec, @@ -23,3 +25,13 @@ impl Default for RecentBuffersContext { } } } + +impl RecentBuffersContext { + /// Returns the [`RecentBuffersContext`] as a message to the language model. + pub fn to_message(&self) -> Option { + self.enabled.then(|| LanguageModelRequestMessage { + role: Role::System, + content: self.message.clone(), + }) + } +} diff --git a/crates/assistant/src/assistant_panel.rs b/crates/assistant/src/assistant_panel.rs index b8b5c8b8adab288a162a1d039353ef495acb5bc5..9f65558808ed3c7b2f97e88318c23300f7acca12 100644 --- a/crates/assistant/src/assistant_panel.rs +++ b/crates/assistant/src/assistant_panel.rs @@ -778,6 +778,7 @@ impl AssistantPanel { cx, ) }); + self.show_conversation(editor.clone(), cx); Some(editor) } @@ -1351,7 +1352,7 @@ struct Summary { pub struct Conversation { id: Option, buffer: Model, - ambient_context: AmbientContext, + pub(crate) ambient_context: AmbientContext, message_anchors: Vec, messages_metadata: HashMap, next_message_id: MessageId, @@ -1521,6 +1522,17 @@ impl Conversation { self.update_recent_buffers_context(cx); } + fn toggle_current_project_context( + &mut self, + fs: Arc, + project: WeakModel, + cx: &mut ModelContext, + ) { + self.ambient_context.current_project.enabled = + !self.ambient_context.current_project.enabled; + self.ambient_context.current_project.update(fs, project, cx); + } + fn set_recent_buffers( &mut self, buffers: impl IntoIterator>, @@ -1887,15 +1899,12 @@ impl Conversation { } fn to_completion_request(&self, cx: &mut ModelContext) -> LanguageModelRequest { - let messages = self - .ambient_context - .recent_buffers - .enabled - .then(|| LanguageModelRequestMessage { - role: Role::System, - content: self.ambient_context.recent_buffers.message.clone(), - }) + let recent_buffers_context = self.ambient_context.recent_buffers.to_message(); + let current_project_context = self.ambient_context.current_project.to_message(); + + let messages = recent_buffers_context .into_iter() + .chain(current_project_context) .chain( self.messages(cx) .filter(|message| matches!(message.status, MessageStatus::Done)) @@ -2533,6 +2542,11 @@ impl ConversationEditor { } fn update_message_headers(&mut self, cx: &mut ViewContext) { + let project = self + .workspace + .update(cx, |workspace, _cx| workspace.project().downgrade()) + .unwrap(); + self.editor.update(cx, |editor, cx| { let buffer = editor.buffer().read(cx).snapshot(cx); let excerpt_id = *buffer.as_singleton().unwrap().0; @@ -2549,6 +2563,8 @@ impl ConversationEditor { height: 2, style: BlockStyle::Sticky, render: Box::new({ + let fs = self.fs.clone(); + let project = project.clone(); let conversation = self.conversation.clone(); move |cx| { let message_id = message.id; @@ -2630,31 +2646,40 @@ impl ConversationEditor { Tooltip::text("Include Open Files", cx) }), ) - // .child( - // IconButton::new("include_terminal", IconName::Terminal) - // .icon_size(IconSize::Small) - // .tooltip(|cx| { - // Tooltip::text("Include Terminal", cx) - // }), - // ) - // .child( - // IconButton::new( - // "include_edit_history", - // IconName::FileGit, - // ) - // .icon_size(IconSize::Small) - // .tooltip( - // |cx| Tooltip::text("Include Edit History", cx), - // ), - // ) - // .child( - // IconButton::new( - // "include_file_trees", - // IconName::FileTree, - // ) - // .icon_size(IconSize::Small) - // .tooltip(|cx| Tooltip::text("Include File Trees", cx)), - // ) + .child( + IconButton::new( + "include_current_project", + IconName::FileTree, + ) + .icon_size(IconSize::Small) + .selected( + conversation + .read(cx) + .ambient_context + .current_project + .enabled, + ) + .on_click({ + let fs = fs.clone(); + let project = project.clone(); + let conversation = conversation.downgrade(); + move |_, cx| { + let fs = fs.clone(); + let project = project.clone(); + conversation + .update(cx, |conversation, cx| { + conversation + .toggle_current_project_context( + fs, project, cx, + ); + }) + .ok(); + } + }) + .tooltip( + |cx| Tooltip::text("Include Current Project", cx), + ), + ) .into_any() })) .into_any_element()