1mod edit_action;
2pub mod log;
3
4use anyhow::{anyhow, Context, Result};
5use assistant_tool::{ActionLog, Tool};
6use collections::HashSet;
7use edit_action::{EditAction, EditActionParser};
8use futures::StreamExt;
9use gpui::{App, AsyncApp, Entity, Task};
10use language_model::{
11 LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage, MessageContent, Role,
12};
13use log::{EditToolLog, EditToolRequestId};
14use project::{search::SearchQuery, Project};
15use schemars::JsonSchema;
16use serde::{Deserialize, Serialize};
17use std::fmt::Write;
18use std::sync::Arc;
19use util::paths::PathMatcher;
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.
27 ///
28 /// The description should be concise and clear. We will show this
29 /// description to the user as well.
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
62pub struct EditFilesTool;
63
64impl Tool for EditFilesTool {
65 fn name(&self) -> String {
66 "edit-files".into()
67 }
68
69 fn description(&self) -> String {
70 include_str!("./edit_files_tool/description.md").into()
71 }
72
73 fn input_schema(&self) -> serde_json::Value {
74 let schema = schemars::schema_for!(EditFilesToolInput);
75 serde_json::to_value(&schema).unwrap()
76 }
77
78 fn run(
79 self: Arc<Self>,
80 input: serde_json::Value,
81 messages: &[LanguageModelRequestMessage],
82 project: Entity<Project>,
83 action_log: Entity<ActionLog>,
84 cx: &mut App,
85 ) -> Task<Result<String>> {
86 let input = match serde_json::from_value::<EditFilesToolInput>(input) {
87 Ok(input) => input,
88 Err(err) => return Task::ready(Err(anyhow!(err))),
89 };
90
91 match EditToolLog::try_global(cx) {
92 Some(log) => {
93 let req_id = log.update(cx, |log, cx| {
94 log.new_request(input.edit_instructions.clone(), cx)
95 });
96
97 let task = EditToolRequest::new(
98 input,
99 messages,
100 project,
101 action_log,
102 Some((log.clone(), req_id)),
103 cx,
104 );
105
106 cx.spawn(|mut cx| async move {
107 let result = task.await;
108
109 let str_result = match &result {
110 Ok(out) => Ok(out.clone()),
111 Err(err) => Err(err.to_string()),
112 };
113
114 log.update(&mut cx, |log, cx| {
115 log.set_tool_output(req_id, str_result, cx)
116 })
117 .log_err();
118
119 result
120 })
121 }
122
123 None => EditToolRequest::new(input, messages, project, action_log, None, cx),
124 }
125 }
126}
127
128struct EditToolRequest {
129 parser: EditActionParser,
130 output: String,
131 changed_buffers: HashSet<Entity<language::Buffer>>,
132 bad_searches: Vec<BadSearch>,
133 project: Entity<Project>,
134 action_log: Entity<ActionLog>,
135 tool_log: Option<(Entity<EditToolLog>, EditToolRequestId)>,
136}
137
138#[derive(Debug)]
139enum DiffResult {
140 BadSearch(BadSearch),
141 Diff(language::Diff),
142}
143
144#[derive(Debug)]
145struct BadSearch {
146 file_path: String,
147 search: String,
148}
149
150impl EditToolRequest {
151 fn new(
152 input: EditFilesToolInput,
153 messages: &[LanguageModelRequestMessage],
154 project: Entity<Project>,
155 action_log: Entity<ActionLog>,
156 tool_log: Option<(Entity<EditToolLog>, EditToolRequestId)>,
157 cx: &mut App,
158 ) -> Task<Result<String>> {
159 let model_registry = LanguageModelRegistry::read_global(cx);
160 let Some(model) = model_registry.editor_model() else {
161 return Task::ready(Err(anyhow!("No editor model configured")));
162 };
163
164 let mut messages = messages.to_vec();
165 // Remove the last tool use (this run) to prevent an invalid request
166 'outer: for message in messages.iter_mut().rev() {
167 for (index, content) in message.content.iter().enumerate().rev() {
168 match content {
169 MessageContent::ToolUse(_) => {
170 message.content.remove(index);
171 break 'outer;
172 }
173 MessageContent::ToolResult(_) => {
174 // If we find any tool results before a tool use, the request is already valid
175 break 'outer;
176 }
177 MessageContent::Text(_) | MessageContent::Image(_) => {}
178 }
179 }
180 }
181
182 messages.push(LanguageModelRequestMessage {
183 role: Role::User,
184 content: vec![
185 include_str!("./edit_files_tool/edit_prompt.md").into(),
186 input.edit_instructions.into(),
187 ],
188 cache: false,
189 });
190
191 cx.spawn(|mut cx| async move {
192 let llm_request = LanguageModelRequest {
193 messages,
194 tools: vec![],
195 stop: vec![],
196 temperature: Some(0.0),
197 };
198
199 let stream = model.stream_completion_text(llm_request, &cx);
200 let mut chunks = stream.await?;
201
202 let mut request = Self {
203 parser: EditActionParser::new(),
204 // we start with the success header so we don't need to shift the output in the common case
205 output: Self::SUCCESS_OUTPUT_HEADER.to_string(),
206 changed_buffers: HashSet::default(),
207 bad_searches: Vec::new(),
208 action_log,
209 project,
210 tool_log,
211 };
212
213 while let Some(chunk) = chunks.stream.next().await {
214 request.process_response_chunk(&chunk?, &mut cx).await?;
215 }
216
217 request.finalize(&mut cx).await
218 })
219 }
220
221 async fn process_response_chunk(&mut self, chunk: &str, cx: &mut AsyncApp) -> Result<()> {
222 let new_actions = self.parser.parse_chunk(chunk);
223
224 if let Some((ref log, req_id)) = self.tool_log {
225 log.update(cx, |log, cx| {
226 log.push_editor_response_chunk(req_id, chunk, &new_actions, cx)
227 })
228 .log_err();
229 }
230
231 for action in new_actions {
232 self.apply_action(action, cx).await?;
233 }
234
235 Ok(())
236 }
237
238 async fn apply_action(
239 &mut self,
240 (action, source): (EditAction, String),
241 cx: &mut AsyncApp,
242 ) -> Result<()> {
243 let project_path = self.project.read_with(cx, |project, cx| {
244 project
245 .find_project_path(action.file_path(), cx)
246 .context("Path not found in project")
247 })??;
248
249 let buffer = self
250 .project
251 .update(cx, |project, cx| project.open_buffer(project_path, cx))?
252 .await?;
253
254 let result = match action {
255 EditAction::Replace {
256 old,
257 new,
258 file_path,
259 } => {
260 let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot())?;
261
262 cx.background_executor()
263 .spawn(Self::replace_diff(old, new, file_path, snapshot))
264 .await
265 }
266 EditAction::Write { content, .. } => Ok(DiffResult::Diff(
267 buffer
268 .read_with(cx, |buffer, cx| buffer.diff(content, cx))?
269 .await,
270 )),
271 }?;
272
273 match result {
274 DiffResult::BadSearch(invalid_replace) => {
275 self.bad_searches.push(invalid_replace);
276 }
277 DiffResult::Diff(diff) => {
278 let _clock = buffer.update(cx, |buffer, cx| buffer.apply_diff(diff, cx))?;
279
280 write!(&mut self.output, "\n\n{}", source)?;
281 self.changed_buffers.insert(buffer);
282 }
283 }
284
285 Ok(())
286 }
287
288 async fn replace_diff(
289 old: String,
290 new: String,
291 file_path: std::path::PathBuf,
292 snapshot: language::BufferSnapshot,
293 ) -> Result<DiffResult> {
294 let query = SearchQuery::text(
295 old.clone(),
296 false,
297 true,
298 true,
299 PathMatcher::new(&[])?,
300 PathMatcher::new(&[])?,
301 None,
302 )?;
303
304 let matches = query.search(&snapshot, None).await;
305
306 if matches.is_empty() {
307 return Ok(DiffResult::BadSearch(BadSearch {
308 search: new.clone(),
309 file_path: file_path.display().to_string(),
310 }));
311 }
312
313 let edit_range = matches[0].clone();
314 let diff = language::text_diff(&old, &new);
315
316 let edits = diff
317 .into_iter()
318 .map(|(old_range, text)| {
319 let start = edit_range.start + old_range.start;
320 let end = edit_range.start + old_range.end;
321 (start..end, text)
322 })
323 .collect::<Vec<_>>();
324
325 let diff = language::Diff {
326 base_version: snapshot.version().clone(),
327 line_ending: snapshot.line_ending(),
328 edits,
329 };
330
331 anyhow::Ok(DiffResult::Diff(diff))
332 }
333
334 const SUCCESS_OUTPUT_HEADER: &str = "Successfully applied. Here's a list of changes:";
335 const ERROR_OUTPUT_HEADER_NO_EDITS: &str = "I couldn't apply any edits!";
336 const ERROR_OUTPUT_HEADER_WITH_EDITS: &str =
337 "Errors occurred. First, here's a list of the edits we managed to apply:";
338
339 async fn finalize(self, cx: &mut AsyncApp) -> Result<String> {
340 let changed_buffer_count = self.changed_buffers.len();
341
342 // Save each buffer once at the end
343 for buffer in &self.changed_buffers {
344 self.project
345 .update(cx, |project, cx| project.save_buffer(buffer.clone(), cx))?
346 .await?;
347 }
348
349 self.action_log
350 .update(cx, |log, cx| log.buffer_edited(self.changed_buffers, cx))
351 .log_err();
352
353 let errors = self.parser.errors();
354
355 if errors.is_empty() && self.bad_searches.is_empty() {
356 if changed_buffer_count == 0 {
357 return Err(anyhow!(
358 "The instructions didn't lead to any changes. You might need to consult the file contents first."
359 ));
360 }
361
362 Ok(self.output)
363 } else {
364 let mut output = self.output;
365
366 if output.is_empty() {
367 output.replace_range(
368 0..Self::SUCCESS_OUTPUT_HEADER.len(),
369 Self::ERROR_OUTPUT_HEADER_NO_EDITS,
370 );
371 } else {
372 output.replace_range(
373 0..Self::SUCCESS_OUTPUT_HEADER.len(),
374 Self::ERROR_OUTPUT_HEADER_WITH_EDITS,
375 );
376 }
377
378 if !self.bad_searches.is_empty() {
379 writeln!(
380 &mut output,
381 "\n\nThese searches failed because they didn't match any strings:"
382 )?;
383
384 for replace in self.bad_searches {
385 writeln!(
386 &mut output,
387 "- '{}' does not appear in `{}`",
388 replace.search.replace("\r", "\\r").replace("\n", "\\n"),
389 replace.file_path
390 )?;
391 }
392
393 write!(&mut output, "Make sure to use exact searches.")?;
394 }
395
396 if !errors.is_empty() {
397 writeln!(
398 &mut output,
399 "\n\nThese SEARCH/REPLACE blocks failed to parse:"
400 )?;
401
402 for error in errors {
403 writeln!(&mut output, "- {}", error)?;
404 }
405 }
406
407 writeln!(
408 &mut output,
409 "\nYou can fix errors by running the tool again. You can include instructions, \
410 but errors are part of the conversation so you don't need to repeat them."
411 )?;
412
413 Err(anyhow!(output))
414 }
415 }
416}