Marshall Bowers and Nate created
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:
<img width="1264" alt="Screenshot 2024-05-14 at 6 17 03 PM"
src="https://github.com/zed-industries/zed/assets/1486634/cc8ed5ad-0ccb-45da-9c07-c96af84a14e3">
Release Notes:
- N/A
---------
Co-authored-by: Nate <nate@zed.dev>Cargo.lock | 2
crates/assistant/Cargo.toml | 2
crates/assistant/src/ambient_context.rs | 3
crates/assistant/src/ambient_context/current_project.rs | 150 +++++++++++
crates/assistant/src/ambient_context/recent_buffers.rs | 12
crates/assistant/src/assistant_panel.rs | 93 ++++--
6 files changed, 228 insertions(+), 34 deletions(-)
@@ -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",
@@ -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
@@ -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,
}
@@ -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<Task<()>>,
+}
+
+#[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<LanguageModelRequestMessage> {
+ 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<dyn Fs>,
+ project: WeakModel<Project>,
+ cx: &mut ModelContext<Conversation>,
+ ) {
+ 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<dyn Fs>, path_to_cargo_toml: &Path) -> Result<String> {
+ 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::<Vec<_>>();
+ 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<Project>,
+ cx: &mut AsyncAppContext,
+ ) -> Result<Option<PathBuf>> {
+ 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)
+ })?
+ }
+}
@@ -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<RecentBuffer>,
@@ -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<LanguageModelRequestMessage> {
+ self.enabled.then(|| LanguageModelRequestMessage {
+ role: Role::System,
+ content: self.message.clone(),
+ })
+ }
+}
@@ -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<String>,
buffer: Model<Buffer>,
- ambient_context: AmbientContext,
+ pub(crate) ambient_context: AmbientContext,
message_anchors: Vec<MessageAnchor>,
messages_metadata: HashMap<MessageId, MessageMetadata>,
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<dyn Fs>,
+ project: WeakModel<Project>,
+ cx: &mut ModelContext<Self>,
+ ) {
+ 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<Item = Model<Buffer>>,
@@ -1887,15 +1899,12 @@ impl Conversation {
}
fn to_completion_request(&self, cx: &mut ModelContext<Conversation>) -> 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<Self>) {
+ 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()