1mod edit_action;
2pub mod log;
3mod resolve_search_block;
4
5use anyhow::{anyhow, Context, Result};
6use assistant_tool::{ActionLog, Tool};
7use collections::HashSet;
8use edit_action::{EditAction, EditActionParser};
9use futures::StreamExt;
10use gpui::{App, AsyncApp, Entity, Task};
11use language::OffsetRangeExt;
12use language_model::{
13 LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage, MessageContent, Role,
14};
15use log::{EditToolLog, EditToolRequestId};
16use project::Project;
17use resolve_search_block::resolve_search_block;
18use schemars::JsonSchema;
19use serde::{Deserialize, Serialize};
20use std::fmt::Write;
21use std::sync::Arc;
22use util::ResultExt;
23
24#[derive(Debug, Serialize, Deserialize, JsonSchema)]
25pub struct EditFilesToolInput {
26 /// High-level edit instructions. These will be interpreted by a smaller
27 /// model, so explain the changes you want that model to make and which
28 /// file paths need changing.
29 ///
30 /// The description should be concise and clear. We will show this
31 /// description to the user as well.
32 ///
33 /// WARNING: When specifying which file paths need changing, you MUST
34 /// start each path with one of the project's root directories.
35 ///
36 /// WARNING: NEVER include code blocks or snippets in edit instructions.
37 /// Only provide natural language descriptions of the changes needed! The tool will
38 /// reject any instructions that contain code blocks or snippets.
39 ///
40 /// The following examples assume we have two root directories in the project:
41 /// - root-1
42 /// - root-2
43 ///
44 /// <example>
45 /// If you want to introduce a new quit function to kill the process, your
46 /// instructions should be: "Add a new `quit` function to
47 /// `root-1/src/main.rs` to kill the process".
48 ///
49 /// Notice how the file path starts with root-1. Without that, the path
50 /// would be ambiguous and the call would fail!
51 /// </example>
52 ///
53 /// <example>
54 /// If you want to change documentation to always start with a capital
55 /// letter, your instructions should be: "In `root-2/db.js`,
56 /// `root-2/inMemory.js` and `root-2/sql.js`, change all the documentation
57 /// to start with a capital letter".
58 ///
59 /// Notice how we never specify code snippets in the instructions!
60 /// </example>
61 pub edit_instructions: String,
62}
63
64pub struct EditFilesTool;
65
66impl Tool for EditFilesTool {
67 fn name(&self) -> String {
68 "edit-files".into()
69 }
70
71 fn description(&self) -> String {
72 include_str!("./edit_files_tool/description.md").into()
73 }
74
75 fn input_schema(&self) -> serde_json::Value {
76 let schema = schemars::schema_for!(EditFilesToolInput);
77 serde_json::to_value(&schema).unwrap()
78 }
79
80 fn run(
81 self: Arc<Self>,
82 input: serde_json::Value,
83 messages: &[LanguageModelRequestMessage],
84 project: Entity<Project>,
85 action_log: Entity<ActionLog>,
86 cx: &mut App,
87 ) -> Task<Result<String>> {
88 let input = match serde_json::from_value::<EditFilesToolInput>(input) {
89 Ok(input) => input,
90 Err(err) => return Task::ready(Err(anyhow!(err))),
91 };
92
93 match EditToolLog::try_global(cx) {
94 Some(log) => {
95 let req_id = log.update(cx, |log, cx| {
96 log.new_request(input.edit_instructions.clone(), cx)
97 });
98
99 let task = EditToolRequest::new(
100 input,
101 messages,
102 project,
103 action_log,
104 Some((log.clone(), req_id)),
105 cx,
106 );
107
108 cx.spawn(|mut cx| async move {
109 let result = task.await;
110
111 let str_result = match &result {
112 Ok(out) => Ok(out.clone()),
113 Err(err) => Err(err.to_string()),
114 };
115
116 log.update(&mut cx, |log, cx| {
117 log.set_tool_output(req_id, str_result, cx)
118 })
119 .log_err();
120
121 result
122 })
123 }
124
125 None => EditToolRequest::new(input, messages, project, action_log, None, cx),
126 }
127 }
128}
129
130struct EditToolRequest {
131 parser: EditActionParser,
132 output: String,
133 changed_buffers: HashSet<Entity<language::Buffer>>,
134 project: Entity<Project>,
135 action_log: Entity<ActionLog>,
136 tool_log: Option<(Entity<EditToolLog>, EditToolRequestId)>,
137}
138
139impl EditToolRequest {
140 fn new(
141 input: EditFilesToolInput,
142 messages: &[LanguageModelRequestMessage],
143 project: Entity<Project>,
144 action_log: Entity<ActionLog>,
145 tool_log: Option<(Entity<EditToolLog>, EditToolRequestId)>,
146 cx: &mut App,
147 ) -> Task<Result<String>> {
148 let model_registry = LanguageModelRegistry::read_global(cx);
149 let Some(model) = model_registry.editor_model() else {
150 return Task::ready(Err(anyhow!("No editor model configured")));
151 };
152
153 let mut messages = messages.to_vec();
154 // Remove the last tool use (this run) to prevent an invalid request
155 'outer: for message in messages.iter_mut().rev() {
156 for (index, content) in message.content.iter().enumerate().rev() {
157 match content {
158 MessageContent::ToolUse(_) => {
159 message.content.remove(index);
160 break 'outer;
161 }
162 MessageContent::ToolResult(_) => {
163 // If we find any tool results before a tool use, the request is already valid
164 break 'outer;
165 }
166 MessageContent::Text(_) | MessageContent::Image(_) => {}
167 }
168 }
169 }
170
171 messages.push(LanguageModelRequestMessage {
172 role: Role::User,
173 content: vec![
174 include_str!("./edit_files_tool/edit_prompt.md").into(),
175 input.edit_instructions.into(),
176 ],
177 cache: false,
178 });
179
180 cx.spawn(|mut cx| async move {
181 let llm_request = LanguageModelRequest {
182 messages,
183 tools: vec![],
184 stop: vec![],
185 temperature: Some(0.0),
186 };
187
188 let stream = model.stream_completion_text(llm_request, &cx);
189 let mut chunks = stream.await?;
190
191 let mut request = Self {
192 parser: EditActionParser::new(),
193 // we start with the success header so we don't need to shift the output in the common case
194 output: Self::SUCCESS_OUTPUT_HEADER.to_string(),
195 changed_buffers: HashSet::default(),
196 action_log,
197 project,
198 tool_log,
199 };
200
201 while let Some(chunk) = chunks.stream.next().await {
202 request.process_response_chunk(&chunk?, &mut cx).await?;
203 }
204
205 request.finalize(&mut cx).await
206 })
207 }
208
209 async fn process_response_chunk(&mut self, chunk: &str, cx: &mut AsyncApp) -> Result<()> {
210 let new_actions = self.parser.parse_chunk(chunk);
211
212 if let Some((ref log, req_id)) = self.tool_log {
213 log.update(cx, |log, cx| {
214 log.push_editor_response_chunk(req_id, chunk, &new_actions, cx)
215 })
216 .log_err();
217 }
218
219 for action in new_actions {
220 self.apply_action(action, cx).await?;
221 }
222
223 Ok(())
224 }
225
226 async fn apply_action(
227 &mut self,
228 (action, source): (EditAction, String),
229 cx: &mut AsyncApp,
230 ) -> Result<()> {
231 let project_path = self.project.read_with(cx, |project, cx| {
232 project
233 .find_project_path(action.file_path(), cx)
234 .context("Path not found in project")
235 })??;
236
237 let buffer = self
238 .project
239 .update(cx, |project, cx| project.open_buffer(project_path, cx))?
240 .await?;
241
242 let diff = match action {
243 EditAction::Replace {
244 old,
245 new,
246 file_path: _,
247 } => {
248 let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot())?;
249
250 let diff = cx
251 .background_executor()
252 .spawn(Self::replace_diff(old, new, snapshot))
253 .await;
254
255 anyhow::Ok(diff)
256 }
257 EditAction::Write { content, .. } => Ok(buffer
258 .read_with(cx, |buffer, cx| buffer.diff(content, cx))?
259 .await),
260 }?;
261
262 let _clock = buffer.update(cx, |buffer, cx| buffer.apply_diff(diff, cx))?;
263
264 write!(&mut self.output, "\n\n{}", source)?;
265 self.changed_buffers.insert(buffer);
266
267 Ok(())
268 }
269
270 async fn replace_diff(
271 old: String,
272 new: String,
273 snapshot: language::BufferSnapshot,
274 ) -> language::Diff {
275 let edit_range = resolve_search_block(&snapshot, &old).to_offset(&snapshot);
276 let diff = language::text_diff(&old, &new);
277
278 let edits = diff
279 .into_iter()
280 .map(|(old_range, text)| {
281 let start = edit_range.start + old_range.start;
282 let end = edit_range.start + old_range.end;
283 (start..end, text)
284 })
285 .collect::<Vec<_>>();
286
287 let diff = language::Diff {
288 base_version: snapshot.version().clone(),
289 line_ending: snapshot.line_ending(),
290 edits,
291 };
292
293 diff
294 }
295
296 const SUCCESS_OUTPUT_HEADER: &str = "Successfully applied. Here's a list of changes:";
297 const ERROR_OUTPUT_HEADER_NO_EDITS: &str = "I couldn't apply any edits!";
298 const ERROR_OUTPUT_HEADER_WITH_EDITS: &str =
299 "Errors occurred. First, here's a list of the edits we managed to apply:";
300
301 async fn finalize(self, cx: &mut AsyncApp) -> Result<String> {
302 let changed_buffer_count = self.changed_buffers.len();
303
304 // Save each buffer once at the end
305 for buffer in &self.changed_buffers {
306 self.project
307 .update(cx, |project, cx| project.save_buffer(buffer.clone(), cx))?
308 .await?;
309 }
310
311 self.action_log
312 .update(cx, |log, cx| log.buffer_edited(self.changed_buffers, cx))
313 .log_err();
314
315 let errors = self.parser.errors();
316
317 if errors.is_empty() {
318 if changed_buffer_count == 0 {
319 return Err(anyhow!(
320 "The instructions didn't lead to any changes. You might need to consult the file contents first."
321 ));
322 }
323
324 Ok(self.output)
325 } else {
326 let mut output = self.output;
327
328 if output.is_empty() {
329 output.replace_range(
330 0..Self::SUCCESS_OUTPUT_HEADER.len(),
331 Self::ERROR_OUTPUT_HEADER_NO_EDITS,
332 );
333 } else {
334 output.replace_range(
335 0..Self::SUCCESS_OUTPUT_HEADER.len(),
336 Self::ERROR_OUTPUT_HEADER_WITH_EDITS,
337 );
338 }
339
340 if !errors.is_empty() {
341 writeln!(
342 &mut output,
343 "\n\nThese SEARCH/REPLACE blocks failed to parse:"
344 )?;
345
346 for error in errors {
347 writeln!(&mut output, "- {}", error)?;
348 }
349 }
350
351 writeln!(
352 &mut output,
353 "\nYou can fix errors by running the tool again. You can include instructions, \
354 but errors are part of the conversation so you don't need to repeat them."
355 )?;
356
357 Err(anyhow!(output))
358 }
359 }
360}