History

assistant edit tool: Track read buffers and notify model of user edits (#26952)

When the model reads file, we'll track the version it read, and let it know if the user makes edits to the buffer. This helps prevent edit failures because it'll know to re-read the file before. Release Notes: - N/A

Agus Zubiaga committed 1 year ago

Commit a05066c1 parent cb439e67 files · +83 −23Browse files

Unified · Split

7 files
  1. Cargo.lock +1 −0
  2. crates/assistant2/src/active_thread.rs +1 −1
  3. crates/assistant2/src/thread.rs +31 −0
  4. crates/assistant_tool/Cargo.toml +1 −0
  5. crates/assistant_tool/src/assistant_tool.rs +37 −15
  6. crates/assistant_tools/src/edit_files_tool.rs +1 −3
  7. crates/assistant_tools/src/read_file_tool.rs +11 −4
Cargo.lock+1−0
@@ -692,6 +692,7 @@ name = "assistant_tool" · 691 unmodified lines
692version = "0.1.0"692version = "0.1.0"
693dependencies = [693dependencies = [
694 "anyhow",694 "anyhow",
695 "clock",
695 "collections",696 "collections",
696 "derive_more",697 "derive_more",
697 "gpui",698 "gpui",
crates/assistant2/src/active_thread.rs+1−1
@@ -361,7 +361,7 @@ impl ActiveThread { · 360 unmodified lines
361 if self.thread.read(cx).all_tools_finished() {361 if self.thread.read(cx).all_tools_finished() {
362 let pending_refresh_buffers = self.thread.update(cx, |thread, cx| {362 let pending_refresh_buffers = self.thread.update(cx, |thread, cx| {
363 thread.action_log().update(cx, |action_log, _cx| {363 thread.action_log().update(cx, |action_log, _cx| {
364 action_log.take_pending_refresh_buffers()364 action_log.take_stale_buffers_in_context()
365 })365 })
366 });366 });
367 367
crates/assistant2/src/thread.rs+31−0
@@ -1,3 +1,4 @@
1use std::fmt::Write as _;
1use std::io::Write;2use std::io::Write;
2use std::sync::Arc;3use std::sync::Arc;
3 4
@@ -560,9 +561,39 @@ impl Thread { · 556 unmodified lines
560 request.messages.push(context_message);561 request.messages.push(context_message);
561 }562 }
562 563
564 self.attach_stale_files(&mut request.messages, cx);
565
563 request566 request
564 }567 }
565 568
569 fn attach_stale_files(&self, messages: &mut Vec<LanguageModelRequestMessage>, cx: &App) {
570 const STALE_FILES_HEADER: &str = "These files changed since last read:";
571
572 let mut stale_message = String::new();
573
574 for stale_file in self.action_log.read(cx).stale_buffers(cx) {
575 let Some(file) = stale_file.read(cx).file() else {
576 continue;
577 };
578
579 if stale_message.is_empty() {
580 write!(&mut stale_message, "{}", STALE_FILES_HEADER).ok();
581 }
582
583 writeln!(&mut stale_message, "- {}", file.path().display()).ok();
584 }
585
586 if !stale_message.is_empty() {
587 let context_message = LanguageModelRequestMessage {
588 role: Role::User,
589 content: vec![stale_message.into()],
590 cache: false,
591 };
592
593 messages.push(context_message);
594 }
595 }
596
566 pub fn stream_completion(597 pub fn stream_completion(
567 &mut self,598 &mut self,
568 request: LanguageModelRequest,599 request: LanguageModelRequest,
crates/assistant_tool/Cargo.toml+1−0
@@ -14,6 +14,7 @@ path = "src/assistant_tool.rs" · 13 unmodified lines
14[dependencies]14[dependencies]
15anyhow.workspace = true15anyhow.workspace = true
16collections.workspace = true16collections.workspace = true
17clock.workspace = true
17derive_more.workspace = true18derive_more.workspace = true
18gpui.workspace = true19gpui.workspace = true
19language.workspace = true20language.workspace = true
crates/assistant_tool/src/assistant_tool.rs+37−15
@@ -4,7 +4,7 @@ mod tool_working_set; · 3 unmodified lines
4use std::sync::Arc;4use std::sync::Arc;
5 5
6use anyhow::Result;6use anyhow::Result;
7use collections::HashSet;7use collections::{HashMap, HashSet};
8use gpui::Context;8use gpui::Context;
9use gpui::{App, Entity, SharedString, Task};9use gpui::{App, Entity, SharedString, Task};
10use language::Buffer;10use language::Buffer;
@@ -58,31 +58,53 @@ pub trait Tool: 'static + Send + Sync { · 47 unmodified lines
58/// Tracks actions performed by tools in a thread58/// Tracks actions performed by tools in a thread
59#[derive(Debug)]59#[derive(Debug)]
60pub struct ActionLog {60pub struct ActionLog {
61 changed_buffers: HashSet<Entity<Buffer>>,61 /// Buffers that user manually added to the context, and whose content has
62 pending_refresh: HashSet<Entity<Buffer>>,62 /// changed since the model last saw them.
63 stale_buffers_in_context: HashSet<Entity<Buffer>>,
64 /// Buffers that we want to notify the model about when they change.
65 tracked_buffers: HashMap<Entity<Buffer>, TrackedBuffer>,
66}
67
68#[derive(Debug, Default)]
69struct TrackedBuffer {
70 version: clock::Global,
63}71}
64 72
65impl ActionLog {73impl ActionLog {
66 /// Creates a new, empty action log.74 /// Creates a new, empty action log.
67 pub fn new() -> Self {75 pub fn new() -> Self {
68 Self {76 Self {
69 changed_buffers: HashSet::default(),77 stale_buffers_in_context: HashSet::default(),
70 pending_refresh: HashSet::default(),78 tracked_buffers: HashMap::default(),
79 }
80 }
81
82 /// Track a buffer as read, so we can notify the model about user edits.
83 pub fn buffer_read(&mut self, buffer: Entity<Buffer>, cx: &mut Context<Self>) {
84 let tracked_buffer = self.tracked_buffers.entry(buffer.clone()).or_default();
85 tracked_buffer.version = buffer.read(cx).version();
86 }
87
88 /// Mark a buffer as edited, so we can refresh it in the context
89 pub fn buffer_edited(&mut self, buffers: HashSet<Entity<Buffer>>, cx: &mut Context<Self>) {
90 for buffer in &buffers {
91 let tracked_buffer = self.tracked_buffers.entry(buffer.clone()).or_default();
92 tracked_buffer.version = buffer.read(cx).version();
71 }93 }
94
95 self.stale_buffers_in_context.extend(buffers);
72 }96 }
73 97
74 /// Registers buffers that have changed and need refreshing.98 /// Iterate over buffers changed since last read or edited by the model
75 pub fn notify_buffers_changed(99 pub fn stale_buffers<'a>(&'a self, cx: &'a App) -> impl Iterator<Item = &'a Entity<Buffer>> {
76 &mut self,100 self.tracked_buffers
77 buffers: HashSet<Entity<Buffer>>,101 .iter()
78 _cx: &mut Context<Self>,102 .filter(|(buffer, tracked)| tracked.version != buffer.read(cx).version)
79 ) {103 .map(|(buffer, _)| buffer)
80 self.changed_buffers.extend(buffers.clone());
81 self.pending_refresh.extend(buffers);
82 }104 }
83 105
84 /// Takes and returns the set of buffers pending refresh, clearing internal state.106 /// Takes and returns the set of buffers pending refresh, clearing internal state.
85 pub fn take_pending_refresh_buffers(&mut self) -> HashSet<Entity<Buffer>> {107 pub fn take_stale_buffers_in_context(&mut self) -> HashSet<Entity<Buffer>> {
86 std::mem::take(&mut self.pending_refresh)108 std::mem::take(&mut self.stale_buffers_in_context)
87 }109 }
88}110}
crates/assistant_tools/src/edit_files_tool.rs+1−3
@@ -309,9 +309,7 @@ impl EditToolRequest { · 308 unmodified lines
309 }309 }
310 310
311 self.action_log311 self.action_log
312 .update(cx, |log, cx| {312 .update(cx, |log, cx| log.buffer_edited(self.changed_buffers, cx))
313 log.notify_buffers_changed(self.changed_buffers, cx)
314 })
315 .log_err();313 .log_err();
316 314
317 let errors = self.parser.errors();315 let errors = self.parser.errors();
crates/assistant_tools/src/read_file_tool.rs+11−4
@@ -49,7 +49,7 @@ impl Tool for ReadFileTool { · 48 unmodified lines
49 input: serde_json::Value,49 input: serde_json::Value,
50 _messages: &[LanguageModelRequestMessage],50 _messages: &[LanguageModelRequestMessage],
51 project: Entity<Project>,51 project: Entity<Project>,
52 _action_log: Entity<ActionLog>,52 action_log: Entity<ActionLog>,
53 cx: &mut App,53 cx: &mut App,
54 ) -> Task<Result<String>> {54 ) -> Task<Result<String>> {
55 let input = match serde_json::from_value::<ReadFileToolInput>(input) {55 let input = match serde_json::from_value::<ReadFileToolInput>(input) {
@@ -60,14 +60,15 @@ impl Tool for ReadFileTool { · 4 unmodified lines
60 let Some(project_path) = project.read(cx).find_project_path(&input.path, cx) else {60 let Some(project_path) = project.read(cx).find_project_path(&input.path, cx) else {
61 return Task::ready(Err(anyhow!("Path not found in project")));61 return Task::ready(Err(anyhow!("Path not found in project")));
62 };62 };
63 cx.spawn(|cx| async move {63
64 cx.spawn(|mut cx| async move {
64 let buffer = cx65 let buffer = cx
65 .update(|cx| {66 .update(|cx| {
66 project.update(cx, |project, cx| project.open_buffer(project_path, cx))67 project.update(cx, |project, cx| project.open_buffer(project_path, cx))
67 })?68 })?
68 .await?;69 .await?;
69 70
70 buffer.read_with(&cx, |buffer, _cx| {71 let result = buffer.read_with(&cx, |buffer, _cx| {
71 if buffer72 if buffer
72 .file()73 .file()
73 .map_or(false, |file| file.disk_state().exists())74 .map_or(false, |file| file.disk_state().exists())
@@ -76,7 +77,13 @@ impl Tool for ReadFileTool { · 2 unmodified lines
76 } else {77 } else {
77 Err(anyhow!("File does not exist"))78 Err(anyhow!("File does not exist"))
78 }79 }
79 })?80 })??;
81
82 action_log.update(&mut cx, |log, cx| {
83 log.buffer_read(buffer, cx);
84 })?;
85
86 anyhow::Ok(result)
80 })87 })
81 }88 }
82}89}

Keyboard shortcuts

General

Show keyboard shortcuts
?
Focus the current search
/

Issues and pull requests

Open filters
f

Issues, pull requests and Inbox

Select next item
j↓
Select previous item
k↑
Open selected item
Enter
Dismiss filters or clear selection
Esc

Repository code

Go to file
t