1use crate::{diff::Diff, stream_completion, OpenAIRequest, RequestMessage, Role};
2use collections::HashMap;
3use editor::{Editor, ToOffset};
4use futures::{channel::mpsc, SinkExt, StreamExt};
5use gpui::{
6 actions, elements::*, platform::MouseButton, AnyViewHandle, AppContext, Entity, Task, View,
7 ViewContext, ViewHandle, WeakViewHandle,
8};
9use menu::{Cancel, Confirm};
10use std::{env, sync::Arc};
11use util::TryFutureExt;
12use workspace::{Modal, Workspace};
13
14actions!(assistant, [Refactor]);
15
16pub fn init(cx: &mut AppContext) {
17 cx.set_global(RefactoringAssistant::new());
18 cx.add_action(RefactoringModal::deploy);
19 cx.add_action(RefactoringModal::confirm);
20 cx.add_action(RefactoringModal::cancel);
21}
22
23pub struct RefactoringAssistant {
24 pending_edits_by_editor: HashMap<usize, Task<Option<()>>>,
25}
26
27impl RefactoringAssistant {
28 fn new() -> Self {
29 Self {
30 pending_edits_by_editor: Default::default(),
31 }
32 }
33
34 fn refactor(&mut self, editor: &ViewHandle<Editor>, prompt: &str, cx: &mut AppContext) {
35 let snapshot = editor.read(cx).buffer().read(cx).snapshot(cx);
36 let selection = editor.read(cx).selections.newest_anchor().clone();
37 let selected_text = snapshot
38 .text_for_range(selection.start..selection.end)
39 .collect::<String>();
40 let language_name = snapshot
41 .language_at(selection.start)
42 .map(|language| language.name());
43 let language_name = language_name.as_deref().unwrap_or("");
44 let request = OpenAIRequest {
45 model: "gpt-4".into(),
46 messages: vec![
47 RequestMessage {
48 role: Role::User,
49 content: format!(
50 "Given the following {language_name} snippet:\n{selected_text}\n{prompt}. Never make remarks and reply only with the new code. Never change the leading whitespace on each line."
51 ),
52 }],
53 stream: true,
54 };
55 let api_key = env::var("OPENAI_API_KEY").unwrap();
56 let response = stream_completion(api_key, cx.background().clone(), request);
57 let editor = editor.downgrade();
58 self.pending_edits_by_editor.insert(
59 editor.id(),
60 cx.spawn(|mut cx| {
61 async move {
62 let mut edit_start = selection.start.to_offset(&snapshot);
63
64 let (mut hunks_tx, mut hunks_rx) = mpsc::channel(1);
65 let diff = cx.background().spawn(async move {
66 let mut messages = response.await?.ready_chunks(4);
67 let mut diff = Diff::new(selected_text);
68
69 while let Some(messages) = messages.next().await {
70 let mut new_text = String::new();
71 for message in messages {
72 let mut message = message?;
73 if let Some(choice) = message.choices.pop() {
74 if let Some(text) = choice.delta.content {
75 new_text.push_str(&text);
76 }
77 }
78 }
79
80 let hunks = diff.push_new(&new_text);
81 hunks_tx.send(hunks).await?;
82 }
83 hunks_tx.send(diff.finish()).await?;
84
85 anyhow::Ok(())
86 });
87
88 let mut first_transaction = None;
89 while let Some(hunks) = hunks_rx.next().await {
90 editor.update(&mut cx, |editor, cx| {
91 let mut highlights = Vec::new();
92
93 editor.buffer().update(cx, |buffer, cx| {
94 buffer.finalize_last_transaction(cx);
95
96 buffer.start_transaction(cx);
97 buffer.edit(
98 hunks.into_iter().filter_map(|hunk| match hunk {
99 crate::diff::Hunk::Insert { text } => {
100 let edit_start = snapshot.anchor_after(edit_start);
101 Some((edit_start..edit_start, text))
102 }
103 crate::diff::Hunk::Remove { len } => {
104 let edit_end = edit_start + len;
105 let edit_range = snapshot.anchor_after(edit_start)
106 ..snapshot.anchor_before(edit_end);
107 edit_start = edit_end;
108 Some((edit_range, String::new()))
109 }
110 crate::diff::Hunk::Keep { len } => {
111 let edit_end = edit_start + len;
112 let edit_range = snapshot.anchor_after(edit_start)
113 ..snapshot.anchor_before(edit_end);
114 edit_start += len;
115 highlights.push(edit_range);
116 None
117 }
118 }),
119 None,
120 cx,
121 );
122 if let Some(transaction) = buffer.end_transaction(cx) {
123 if let Some(first_transaction) = first_transaction {
124 buffer.merge_transaction_into(
125 transaction,
126 first_transaction,
127 cx,
128 );
129 } else {
130 first_transaction = Some(transaction);
131 }
132 buffer.finalize_last_transaction(cx);
133 }
134 });
135
136 editor.highlight_text::<Self>(
137 highlights,
138 gpui::fonts::HighlightStyle {
139 fade_out: Some(0.6),
140 ..Default::default()
141 },
142 cx,
143 );
144 })?;
145 }
146
147 diff.await?;
148 editor.update(&mut cx, |editor, cx| {
149 editor.clear_text_highlights::<Self>(cx);
150 })?;
151
152 anyhow::Ok(())
153 }
154 .log_err()
155 }),
156 );
157 }
158}
159
160enum Event {
161 Dismissed,
162}
163
164struct RefactoringModal {
165 active_editor: WeakViewHandle<Editor>,
166 prompt_editor: ViewHandle<Editor>,
167 has_focus: bool,
168}
169
170impl Entity for RefactoringModal {
171 type Event = Event;
172}
173
174impl View for RefactoringModal {
175 fn ui_name() -> &'static str {
176 "RefactoringModal"
177 }
178
179 fn render(&mut self, cx: &mut ViewContext<Self>) -> AnyElement<Self> {
180 let theme = theme::current(cx);
181
182 ChildView::new(&self.prompt_editor, cx)
183 .constrained()
184 .with_width(theme.assistant.modal.width)
185 .contained()
186 .with_style(theme.assistant.modal.container)
187 .mouse::<Self>(0)
188 .on_click_out(MouseButton::Left, |_, _, cx| cx.emit(Event::Dismissed))
189 .on_click_out(MouseButton::Right, |_, _, cx| cx.emit(Event::Dismissed))
190 .aligned()
191 .right()
192 .into_any()
193 }
194
195 fn focus_in(&mut self, _: AnyViewHandle, cx: &mut ViewContext<Self>) {
196 self.has_focus = true;
197 cx.focus(&self.prompt_editor);
198 }
199
200 fn focus_out(&mut self, _: AnyViewHandle, _: &mut ViewContext<Self>) {
201 self.has_focus = false;
202 }
203}
204
205impl Modal for RefactoringModal {
206 fn has_focus(&self) -> bool {
207 self.has_focus
208 }
209
210 fn dismiss_on_event(event: &Self::Event) -> bool {
211 matches!(event, Self::Event::Dismissed)
212 }
213}
214
215impl RefactoringModal {
216 fn deploy(workspace: &mut Workspace, _: &Refactor, cx: &mut ViewContext<Workspace>) {
217 if let Some(active_editor) = workspace
218 .active_item(cx)
219 .and_then(|item| Some(item.act_as::<Editor>(cx)?.downgrade()))
220 {
221 workspace.toggle_modal(cx, |_, cx| {
222 let prompt_editor = cx.add_view(|cx| {
223 let mut editor = Editor::auto_height(
224 theme::current(cx).assistant.modal.editor_max_lines,
225 Some(Arc::new(|theme| theme.assistant.modal.editor.clone())),
226 cx,
227 );
228 editor
229 .set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
230 editor
231 });
232 cx.add_view(|_| RefactoringModal {
233 active_editor,
234 prompt_editor,
235 has_focus: false,
236 })
237 });
238 }
239 }
240
241 fn cancel(&mut self, _: &Cancel, cx: &mut ViewContext<Self>) {
242 cx.emit(Event::Dismissed);
243 }
244
245 fn confirm(&mut self, _: &Confirm, cx: &mut ViewContext<Self>) {
246 if let Some(editor) = self.active_editor.upgrade(cx) {
247 let prompt = self.prompt_editor.read(cx).text(cx);
248 cx.update_global(|assistant: &mut RefactoringAssistant, cx| {
249 assistant.refactor(&editor, &prompt, cx);
250 });
251 cx.emit(Event::Dismissed);
252 }
253 }
254}