assistant: Add basic current project context (#11828)

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

@@ -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",

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

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,
 }

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<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)
+        })?
+    }
+}

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<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(),
+        })
+    }
+}

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<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
+        let recent_buffers_context = self.ambient_context.recent_buffers.to_message();
-            .ambient_context
+        let current_project_context = self.ambient_context.current_project.to_message();
-            .recent_buffers
+
-            .enabled
+        let messages = recent_buffers_context
-            .then(|| LanguageModelRequestMessage {
-                role: Role::System,
-                content: self.ambient_context.recent_buffers.message.clone(),
-            })
             .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(
+                                        .child(
-                                        //     IconButton::new("include_terminal", IconName::Terminal)
+                                            IconButton::new(
-                                        //         .icon_size(IconSize::Small)
+                                                "include_current_project",
-                                        //         .tooltip(|cx| {
+                                                IconName::FileTree,
-                                        //             Tooltip::text("Include Terminal", cx)
+                                            )
-                                        //         }),
+                                            .icon_size(IconSize::Small)
-                                        // )
+                                            .selected(
-                                        // .child(
+                                                conversation
-                                        //     IconButton::new(
+                                                    .read(cx)
-                                        //         "include_edit_history",
+                                                    .ambient_context
-                                        //         IconName::FileGit,
+                                                    .current_project
-                                        //     )
+                                                    .enabled,
-                                        //     .icon_size(IconSize::Small)
+                                            )
-                                        //     .tooltip(
+                                            .on_click({
-                                        //         |cx| Tooltip::text("Include Edit History", cx),
+                                                let fs = fs.clone();
-                                        //     ),
+                                                let project = project.clone();
-                                        // )
+                                                let conversation = conversation.downgrade();
-                                        // .child(
+                                                move |_, cx| {
-                                        //     IconButton::new(
+                                                    let fs = fs.clone();
-                                        //         "include_file_trees",
+                                                    let project = project.clone();
-                                        //         IconName::FileTree,
+                                                    conversation
-                                        //     )
+                                                        .update(cx, |conversation, cx| {
-                                        //     .icon_size(IconSize::Small)
+                                                            conversation
-                                        //     .tooltip(|cx| Tooltip::text("Include File Trees", cx)),
+                                                                .toggle_current_project_context(
-                                        // )
+                                                                    fs, project, cx,
+                                                                );
+                                                        })
+                                                        .ok();
+                                                }
+                                            })
+                                            .tooltip(
+                                                |cx| Tooltip::text("Include Current Project", cx),
+                                            ),
+                                        )
                                         .into_any()
                                 }))
                                 .into_any_element()