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