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