assistant edit tool: Reliability improvements (#27431)

Agus Zubiaga created 9 months ago

- Add instructions in description to read before editing
- Add instructions in edit prefix to explicitly ask for reads
- Fix `connection error: delay between messages too long` by processing
chunks off a channel

Release Notes:

- N/A

Change summary

crates/assistant_tools/src/edit_files_tool.rs             | 26 ++++++--
crates/assistant_tools/src/edit_files_tool/description.md |  2 
crates/assistant_tools/src/edit_files_tool/edit_prompt.md |  2 
3 files changed, 23 insertions(+), 7 deletions(-)

Detailed changes

crates/assistant_tools/src/edit_files_tool.rs

@@ -6,8 +6,8 @@ use anyhow::{anyhow, Context, Result};
 use assistant_tool::{ActionLog, Tool};
 use collections::HashSet;
 use edit_action::{EditAction, EditActionParser};
-use futures::StreamExt;
-use gpui::{App, AsyncApp, Entity, Task};
+use futures::{channel::mpsc, SinkExt, StreamExt};
+use gpui::{App, AppContext, AsyncApp, Entity, Task};
 use language_model::{
     LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage, MessageContent, Role,
 };
@@ -225,12 +225,24 @@ impl EditToolRequest {
                 temperature: Some(0.0),
             };
 
+            let (mut tx, mut rx) = mpsc::channel::<String>(32);
             let stream = model.stream_completion_text(llm_request, &cx);
-            let mut chunks = stream.await?;
+            let reader_task = cx.background_spawn(async move {
+                let mut chunks = stream.await?;
+
+                while let Some(chunk) = chunks.stream.next().await {
+                    if let Some(chunk) = chunk.log_err() {
+                        // we don't process here because the API fails
+                        // if we take too long between reads
+                        tx.send(chunk).await?
+                    }
+                }
+                tx.close().await?;
+                anyhow::Ok(())
+            });
 
             let mut request = Self {
                 parser: EditActionParser::new(),
-                // we start with the success header so we don't need to shift the output in the common case
                 output: Self::SUCCESS_OUTPUT_HEADER.to_string(),
                 changed_buffers: HashSet::default(),
                 bad_searches: Vec::new(),
@@ -239,10 +251,12 @@ impl EditToolRequest {
                 tool_log,
             };
 
-            while let Some(chunk) = chunks.stream.next().await {
-                request.process_response_chunk(&chunk?, cx).await?;
+            while let Some(chunk) = rx.next().await {
+                request.process_response_chunk(&chunk, cx).await?;
             }
 
+            reader_task.await?;
+
             request.finalize(cx).await
         })
     }

crates/assistant_tools/src/edit_files_tool/description.md

@@ -5,3 +5,5 @@ When using this tool, you should suggest one coherent edit that can be made to t
 When the set of edits you want to make is large or complex, feel free to invoke this tool multiple times, each time focusing on a specific change you wanna make.
 
 You should use this tool when you want to edit a subset of a file's contents, but not the entire file. You should not use this tool when you want to replace the entire contents of a file with completely different contents, and you absolutely must never use this tool to create new files from scratch. If you ever consider using this tool to create a new file from scratch, for any reason, instead you must reconsider and choose a different approach.
+
+DO NOT call this tool until the code to be edited appears in the conversation! You must use the `read-files` tool or ask the user to add it to context first.

crates/assistant_tools/src/edit_files_tool/edit_prompt.md

@@ -120,7 +120,7 @@ Break large *SEARCH/REPLACE* blocks into a series of smaller blocks that each ch
 Include just the changing lines, and a few surrounding lines if needed for uniqueness.
 Do not include long runs of unchanging lines in *SEARCH/REPLACE* blocks.
 
-Only create *SEARCH/REPLACE* blocks for files that the user has added to the chat!
+Only create *SEARCH/REPLACE* blocks for files that have been read! Even though the conversation includes `read-file` tool results, you *CANNOT* issue your own reads. If the conversation doesn't include the code you need to edit, ask for it to be read explicitly.
 
 To move code within a file, use 2 *SEARCH/REPLACE* blocks: 1 to delete it from its current location, 1 to insert it in the new location.