Add more context to the terminal assistant (#15492)

Kirill Bulatov created 1 year ago

Release Notes:

- N/A

Change summary

crates/assistant/src/assistant.rs                  |  2 +
crates/assistant/src/prompts.rs                    | 17 ++++++++++++++++
crates/assistant/src/slash_command/term_command.rs |  6 ++++
crates/assistant/src/terminal_inline_assistant.rs  | 16 ++++++++------
4 files changed, 33 insertions(+), 8 deletions(-)

Detailed changes

crates/assistant/src/assistant.rs

@@ -55,6 +55,8 @@ actions!(
     ]
 );
 
+const DEFAULT_CONTEXT_LINES: usize = 20;
+
 #[derive(Clone, Default, Deserialize, PartialEq)]
 pub struct InlineAssist {
     prompt: Option<String>,

crates/assistant/src/prompts.rs

@@ -115,11 +115,19 @@ pub fn generate_terminal_assistant_prompt(
     user_prompt: &str,
     shell: Option<&str>,
     working_directory: Option<&str>,
+    latest_output: &[String],
 ) -> String {
     let mut prompt = String::new();
     writeln!(&mut prompt, "You are an expert terminal user.").unwrap();
     writeln!(&mut prompt, "You will be given a description of a command and you need to respond with a command that matches the description.").unwrap();
     writeln!(&mut prompt, "Do not include markdown blocks or any other text formatting in your response, always respond with a single command that can be executed in the given shell.").unwrap();
+    writeln!(
+        &mut prompt,
+        "Current OS name is '{}', architecture is '{}'.",
+        std::env::consts::OS,
+        std::env::consts::ARCH,
+    )
+    .unwrap();
     if let Some(shell) = shell {
         writeln!(&mut prompt, "Current shell is '{shell}'.").unwrap();
     }
@@ -130,6 +138,15 @@ pub fn generate_terminal_assistant_prompt(
         )
         .unwrap();
     }
+    if !latest_output.is_empty() {
+        writeln!(
+            &mut prompt,
+            "Latest non-empty {} lines of the terminal output: {:?}",
+            latest_output.len(),
+            latest_output
+        )
+        .unwrap();
+    }
     writeln!(&mut prompt, "Here is the description of the command:").unwrap();
     prompt.push_str(user_prompt);
     prompt

crates/assistant/src/slash_command/term_command.rs

@@ -11,6 +11,8 @@ use terminal_view::{terminal_panel::TerminalPanel, TerminalView};
 use ui::prelude::*;
 use workspace::{dock::Panel, Workspace};
 
+use crate::DEFAULT_CONTEXT_LINES;
+
 use super::create_label_for_command;
 
 pub(crate) struct TermSlashCommand;
@@ -73,7 +75,9 @@ impl SlashCommand for TermSlashCommand {
             return Task::ready(Err(anyhow::anyhow!("no active terminal")));
         };
 
-        let line_count = argument.and_then(|a| parse_argument(a)).unwrap_or(20);
+        let line_count = argument
+            .and_then(|a| parse_argument(a))
+            .unwrap_or(DEFAULT_CONTEXT_LINES);
 
         let lines = active_terminal
             .read(cx)

crates/assistant/src/terminal_inline_assistant.rs

@@ -1,6 +1,6 @@
 use crate::{
     humanize_token_count, prompts::generate_terminal_assistant_prompt, AssistantPanel,
-    AssistantPanelEvent, ModelSelector,
+    AssistantPanelEvent, ModelSelector, DEFAULT_CONTEXT_LINES,
 };
 use anyhow::{Context as _, Result};
 use client::telemetry::Telemetry;
@@ -217,17 +217,18 @@ impl TerminalInlineAssistant {
         let assist = self.assists.get(&assist_id).context("invalid assist")?;
 
         let shell = std::env::var("SHELL").ok();
-        let working_directory = assist
+        let (latest_output, working_directory) = assist
             .terminal
             .update(cx, |terminal, cx| {
-                terminal
-                    .model()
-                    .read(cx)
+                let terminal = terminal.model().read(cx);
+                let latest_output = terminal.last_n_non_empty_lines(DEFAULT_CONTEXT_LINES);
+                let working_directory = terminal
                     .working_directory()
-                    .map(|path| path.to_string_lossy().to_string())
+                    .map(|path| path.to_string_lossy().to_string());
+                (latest_output, working_directory)
             })
             .ok()
-            .flatten();
+            .unwrap_or_default();
 
         let context_request = if assist.include_context {
             assist.workspace.as_ref().and_then(|workspace| {
@@ -254,6 +255,7 @@ impl TerminalInlineAssistant {
                 .prompt(cx),
             shell.as_deref(),
             working_directory.as_deref(),
+            &latest_output,
         );
 
         let mut messages = Vec::new();