1mod edit_action;
2pub mod log;
3
4use crate::replace::{replace_exact, replace_with_flexible_indent};
5use anyhow::{anyhow, Context, Result};
6use assistant_tool::{ActionLog, Tool};
7use collections::HashSet;
8use edit_action::{EditAction, EditActionParser};
9use futures::{channel::mpsc, SinkExt, StreamExt};
10use gpui::{App, AppContext, AsyncApp, Entity, Task};
11use language_model::{
12 LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage, MessageContent, Role,
13};
14use log::{EditToolLog, EditToolRequestId};
15use project::Project;
16use schemars::JsonSchema;
17use serde::{Deserialize, Serialize};
18use std::fmt::Write;
19use std::sync::Arc;
20use util::ResultExt;
21
22#[derive(Debug, Serialize, Deserialize, JsonSchema)]
23pub struct EditFilesToolInput {
24 /// High-level edit instructions. These will be interpreted by a smaller
25 /// model, so explain the changes you want that model to make and which
26 /// file paths need changing. The description should be concise and clear.
27 ///
28 /// WARNING: When specifying which file paths need changing, you MUST
29 /// start each path with one of the project's root directories.
30 ///
31 /// WARNING: NEVER include code blocks or snippets in edit instructions.
32 /// Only provide natural language descriptions of the changes needed! The tool will
33 /// reject any instructions that contain code blocks or snippets.
34 ///
35 /// The following examples assume we have two root directories in the project:
36 /// - root-1
37 /// - root-2
38 ///
39 /// <example>
40 /// If you want to introduce a new quit function to kill the process, your
41 /// instructions should be: "Add a new `quit` function to
42 /// `root-1/src/main.rs` to kill the process".
43 ///
44 /// Notice how the file path starts with root-1. Without that, the path
45 /// would be ambiguous and the call would fail!
46 /// </example>
47 ///
48 /// <example>
49 /// If you want to change documentation to always start with a capital
50 /// letter, your instructions should be: "In `root-2/db.js`,
51 /// `root-2/inMemory.js` and `root-2/sql.js`, change all the documentation
52 /// to start with a capital letter".
53 ///
54 /// Notice how we never specify code snippets in the instructions!
55 /// </example>
56 pub edit_instructions: String,
57
58 /// A user-friendly description of what changes are being made.
59 /// 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
60 /// constrained, so make the description extremely terse.
61 ///
62 /// <example>
63 /// For fixing a broken authentication system:
64 /// "Fix auth bug in login flow"
65 /// </example>
66 ///
67 /// <example>
68 /// For adding unit tests to a module:
69 /// "Add tests for user profile logic"
70 /// </example>
71 pub display_description: String,
72}
73
74pub struct EditFilesTool;
75
76impl Tool for EditFilesTool {
77 fn name(&self) -> String {
78 "edit-files".into()
79 }
80
81 fn needs_confirmation(&self) -> bool {
82 true
83 }
84
85 fn description(&self) -> String {
86 include_str!("./edit_files_tool/description.md").into()
87 }
88
89 fn input_schema(&self) -> serde_json::Value {
90 let schema = schemars::schema_for!(EditFilesToolInput);
91 serde_json::to_value(&schema).unwrap()
92 }
93
94 fn ui_text(&self, input: &serde_json::Value) -> String {
95 match serde_json::from_value::<EditFilesToolInput>(input.clone()) {
96 Ok(input) => input.display_description,
97 Err(_) => "Edit files".to_string(),
98 }
99 }
100
101 fn run(
102 self: Arc<Self>,
103 input: serde_json::Value,
104 messages: &[LanguageModelRequestMessage],
105 project: Entity<Project>,
106 action_log: Entity<ActionLog>,
107 cx: &mut App,
108 ) -> Task<Result<String>> {
109 let input = match serde_json::from_value::<EditFilesToolInput>(input) {
110 Ok(input) => input,
111 Err(err) => return Task::ready(Err(anyhow!(err))),
112 };
113
114 match EditToolLog::try_global(cx) {
115 Some(log) => {
116 let req_id = log.update(cx, |log, cx| {
117 log.new_request(input.edit_instructions.clone(), cx)
118 });
119
120 let task = EditToolRequest::new(
121 input,
122 messages,
123 project,
124 action_log,
125 Some((log.clone(), req_id)),
126 cx,
127 );
128
129 cx.spawn(async move |cx| {
130 let result = task.await;
131
132 let str_result = match &result {
133 Ok(out) => Ok(out.clone()),
134 Err(err) => Err(err.to_string()),
135 };
136
137 log.update(cx, |log, cx| log.set_tool_output(req_id, str_result, cx))
138 .log_err();
139
140 result
141 })
142 }
143
144 None => EditToolRequest::new(input, messages, project, action_log, None, cx),
145 }
146 }
147}
148
149struct EditToolRequest {
150 parser: EditActionParser,
151 editor_response: EditorResponse,
152 project: Entity<Project>,
153 action_log: Entity<ActionLog>,
154 tool_log: Option<(Entity<EditToolLog>, EditToolRequestId)>,
155}
156
157enum EditorResponse {
158 /// The editor model hasn't produced any actions yet.
159 /// If we don't have any by the end, we'll return its message to the architect model.
160 Message(String),
161 /// The editor model produced at least one action.
162 Actions {
163 applied: Vec<AppliedAction>,
164 search_errors: Vec<SearchError>,
165 },
166}
167
168struct AppliedAction {
169 source: String,
170 buffer: Entity<language::Buffer>,
171}
172
173#[derive(Debug)]
174enum DiffResult {
175 Diff(language::Diff),
176 SearchError(SearchError),
177}
178
179#[derive(Debug)]
180enum SearchError {
181 NoMatch {
182 file_path: String,
183 search: String,
184 },
185 EmptyBuffer {
186 file_path: String,
187 search: String,
188 exists: bool,
189 },
190}
191
192impl EditToolRequest {
193 fn new(
194 input: EditFilesToolInput,
195 messages: &[LanguageModelRequestMessage],
196 project: Entity<Project>,
197 action_log: Entity<ActionLog>,
198 tool_log: Option<(Entity<EditToolLog>, EditToolRequestId)>,
199 cx: &mut App,
200 ) -> Task<Result<String>> {
201 let model_registry = LanguageModelRegistry::read_global(cx);
202 let Some(model) = model_registry.editor_model() else {
203 return Task::ready(Err(anyhow!("No editor model configured")));
204 };
205
206 let mut messages = messages.to_vec();
207 // Remove the last tool use (this run) to prevent an invalid request
208 'outer: for message in messages.iter_mut().rev() {
209 for (index, content) in message.content.iter().enumerate().rev() {
210 match content {
211 MessageContent::ToolUse(_) => {
212 message.content.remove(index);
213 break 'outer;
214 }
215 MessageContent::ToolResult(_) => {
216 // If we find any tool results before a tool use, the request is already valid
217 break 'outer;
218 }
219 MessageContent::Text(_) | MessageContent::Image(_) => {}
220 }
221 }
222 }
223
224 messages.push(LanguageModelRequestMessage {
225 role: Role::User,
226 content: vec![
227 include_str!("./edit_files_tool/edit_prompt.md").into(),
228 input.edit_instructions.into(),
229 ],
230 cache: false,
231 });
232
233 cx.spawn(async move |cx| {
234 let llm_request = LanguageModelRequest {
235 messages,
236 tools: vec![],
237 stop: vec![],
238 temperature: Some(0.0),
239 };
240
241 let (mut tx, mut rx) = mpsc::channel::<String>(32);
242 let stream = model.stream_completion_text(llm_request, &cx);
243 let reader_task = cx.background_spawn(async move {
244 let mut chunks = stream.await?;
245
246 while let Some(chunk) = chunks.stream.next().await {
247 if let Some(chunk) = chunk.log_err() {
248 // we don't process here because the API fails
249 // if we take too long between reads
250 tx.send(chunk).await?
251 }
252 }
253 tx.close().await?;
254 anyhow::Ok(())
255 });
256
257 let mut request = Self {
258 parser: EditActionParser::new(),
259 editor_response: EditorResponse::Message(String::with_capacity(256)),
260 action_log,
261 project,
262 tool_log,
263 };
264
265 while let Some(chunk) = rx.next().await {
266 request.process_response_chunk(&chunk, cx).await?;
267 }
268
269 reader_task.await?;
270
271 request.finalize(cx).await
272 })
273 }
274
275 async fn process_response_chunk(&mut self, chunk: &str, cx: &mut AsyncApp) -> Result<()> {
276 let new_actions = self.parser.parse_chunk(chunk);
277
278 if let EditorResponse::Message(ref mut message) = self.editor_response {
279 if new_actions.is_empty() {
280 message.push_str(chunk);
281 }
282 }
283
284 if let Some((ref log, req_id)) = self.tool_log {
285 log.update(cx, |log, cx| {
286 log.push_editor_response_chunk(req_id, chunk, &new_actions, cx)
287 })
288 .log_err();
289 }
290
291 for action in new_actions {
292 self.apply_action(action, cx).await?;
293 }
294
295 Ok(())
296 }
297
298 async fn apply_action(
299 &mut self,
300 (action, source): (EditAction, String),
301 cx: &mut AsyncApp,
302 ) -> Result<()> {
303 let project_path = self.project.read_with(cx, |project, cx| {
304 project
305 .find_project_path(action.file_path(), cx)
306 .context("Path not found in project")
307 })??;
308
309 let buffer = self
310 .project
311 .update(cx, |project, cx| project.open_buffer(project_path, cx))?
312 .await?;
313
314 let result = match action {
315 EditAction::Replace {
316 old,
317 new,
318 file_path,
319 } => {
320 let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot())?;
321
322 cx.background_executor()
323 .spawn(Self::replace_diff(old, new, file_path, snapshot))
324 .await
325 }
326 EditAction::Write { content, .. } => Ok(DiffResult::Diff(
327 buffer
328 .read_with(cx, |buffer, cx| buffer.diff(content, cx))?
329 .await,
330 )),
331 }?;
332
333 match result {
334 DiffResult::SearchError(error) => {
335 self.push_search_error(error);
336 }
337 DiffResult::Diff(diff) => {
338 let _clock = buffer.update(cx, |buffer, cx| buffer.apply_diff(diff, cx))?;
339
340 self.push_applied_action(AppliedAction { source, buffer });
341 }
342 }
343
344 anyhow::Ok(())
345 }
346
347 fn push_search_error(&mut self, error: SearchError) {
348 match &mut self.editor_response {
349 EditorResponse::Message(_) => {
350 self.editor_response = EditorResponse::Actions {
351 applied: Vec::new(),
352 search_errors: vec![error],
353 };
354 }
355 EditorResponse::Actions { search_errors, .. } => {
356 search_errors.push(error);
357 }
358 }
359 }
360
361 fn push_applied_action(&mut self, action: AppliedAction) {
362 match &mut self.editor_response {
363 EditorResponse::Message(_) => {
364 self.editor_response = EditorResponse::Actions {
365 applied: vec![action],
366 search_errors: Vec::new(),
367 };
368 }
369 EditorResponse::Actions { applied, .. } => {
370 applied.push(action);
371 }
372 }
373 }
374
375 async fn replace_diff(
376 old: String,
377 new: String,
378 file_path: std::path::PathBuf,
379 snapshot: language::BufferSnapshot,
380 ) -> Result<DiffResult> {
381 if snapshot.is_empty() {
382 let exists = snapshot
383 .file()
384 .map_or(false, |file| file.disk_state().exists());
385
386 let error = SearchError::EmptyBuffer {
387 file_path: file_path.display().to_string(),
388 exists,
389 search: old,
390 };
391
392 return Ok(DiffResult::SearchError(error));
393 }
394
395 let replace_result =
396 // Try to match exactly
397 replace_exact(&old, &new, &snapshot)
398 .await
399 // If that fails, try being flexible about indentation
400 .or_else(|| replace_with_flexible_indent(&old, &new, &snapshot));
401
402 let Some(diff) = replace_result else {
403 let error = SearchError::NoMatch {
404 search: old,
405 file_path: file_path.display().to_string(),
406 };
407
408 return Ok(DiffResult::SearchError(error));
409 };
410
411 Ok(DiffResult::Diff(diff))
412 }
413
414 async fn finalize(self, cx: &mut AsyncApp) -> Result<String> {
415 match self.editor_response {
416 EditorResponse::Message(message) => Err(anyhow!(
417 "No edits were applied! You might need to provide more context.\n\n{}",
418 message
419 )),
420 EditorResponse::Actions {
421 applied,
422 search_errors,
423 } => {
424 let mut output = String::with_capacity(1024);
425
426 let parse_errors = self.parser.errors();
427 let has_errors = !search_errors.is_empty() || !parse_errors.is_empty();
428
429 if has_errors {
430 let error_count = search_errors.len() + parse_errors.len();
431
432 if applied.is_empty() {
433 writeln!(
434 &mut output,
435 "{} errors occurred! No edits were applied.",
436 error_count,
437 )?;
438 } else {
439 writeln!(
440 &mut output,
441 "{} errors occurred, but {} edits were correctly applied.",
442 error_count,
443 applied.len(),
444 )?;
445
446 writeln!(
447 &mut output,
448 "# {} SEARCH/REPLACE block(s) applied:\n\nDo not re-send these since they are already applied!\n",
449 applied.len()
450 )?;
451 }
452 } else {
453 write!(
454 &mut output,
455 "Successfully applied! Here's a list of applied edits:"
456 )?;
457 }
458
459 let mut changed_buffers = HashSet::default();
460
461 for action in applied {
462 changed_buffers.insert(action.buffer);
463 write!(&mut output, "\n\n{}", action.source)?;
464 }
465
466 for buffer in &changed_buffers {
467 self.project
468 .update(cx, |project, cx| project.save_buffer(buffer.clone(), cx))?
469 .await?;
470 }
471
472 self.action_log
473 .update(cx, |log, cx| log.buffer_edited(changed_buffers.clone(), cx))
474 .log_err();
475
476 if !search_errors.is_empty() {
477 writeln!(
478 &mut output,
479 "\n\n## {} SEARCH/REPLACE block(s) failed to match:\n",
480 search_errors.len()
481 )?;
482
483 for error in search_errors {
484 match error {
485 SearchError::NoMatch { file_path, search } => {
486 writeln!(
487 &mut output,
488 "### No exact match in: `{}`\n```\n{}\n```\n",
489 file_path, search,
490 )?;
491 }
492 SearchError::EmptyBuffer {
493 file_path,
494 exists: true,
495 search,
496 } => {
497 writeln!(
498 &mut output,
499 "### No match because `{}` is empty:\n```\n{}\n```\n",
500 file_path, search,
501 )?;
502 }
503 SearchError::EmptyBuffer {
504 file_path,
505 exists: false,
506 search,
507 } => {
508 writeln!(
509 &mut output,
510 "### No match because `{}` does not exist:\n```\n{}\n```\n",
511 file_path, search,
512 )?;
513 }
514 }
515 }
516
517 write!(&mut output,
518 "The SEARCH section must exactly match an existing block of lines including all white \
519 space, comments, indentation, docstrings, etc."
520 )?;
521 }
522
523 if !parse_errors.is_empty() {
524 writeln!(
525 &mut output,
526 "\n\n## {} SEARCH/REPLACE blocks failed to parse:",
527 parse_errors.len()
528 )?;
529
530 for error in parse_errors {
531 writeln!(&mut output, "- {}", error)?;
532 }
533 }
534
535 if has_errors {
536 writeln!(&mut output,
537 "\n\nYou can fix errors by running the tool again. You can include instructions, \
538 but errors are part of the conversation so you don't need to repeat them.",
539 )?;
540
541 Err(anyhow!(output))
542 } else {
543 Ok(output)
544 }
545 }
546 }
547 }
548}