1use crate::{
2 assistant_settings::{AssistantDockPosition, AssistantSettings, OpenAIModel},
3 stream_completion,
4 streaming_diff::{Hunk, StreamingDiff},
5 MessageId, MessageMetadata, MessageStatus, OpenAIRequest, RequestMessage, Role,
6 SavedConversation, SavedConversationMetadata, SavedMessage, OPENAI_API_URL,
7};
8use anyhow::{anyhow, Result};
9use chrono::{DateTime, Local};
10use collections::{hash_map, HashMap, HashSet};
11use editor::{
12 display_map::{
13 BlockContext, BlockDisposition, BlockId, BlockProperties, BlockStyle, ToDisplayPoint,
14 },
15 scroll::autoscroll::{Autoscroll, AutoscrollStrategy},
16 Anchor, Editor, ToOffset, ToPoint,
17};
18use fs::Fs;
19use futures::{channel::mpsc, SinkExt, StreamExt};
20use gpui::{
21 actions,
22 elements::*,
23 geometry::vector::{vec2f, Vector2F},
24 platform::{CursorStyle, MouseButton},
25 Action, AppContext, AsyncAppContext, ClipboardItem, Entity, ModelContext, ModelHandle,
26 Subscription, Task, View, ViewContext, ViewHandle, WeakViewHandle, WindowContext,
27};
28use language::{
29 language_settings::SoftWrap, Buffer, LanguageRegistry, Point, Rope, ToOffset as _,
30 TransactionId,
31};
32use search::BufferSearchBar;
33use settings::SettingsStore;
34use std::{
35 cell::RefCell,
36 cmp, env,
37 fmt::Write,
38 iter,
39 ops::Range,
40 path::{Path, PathBuf},
41 rc::Rc,
42 sync::Arc,
43 time::Duration,
44};
45use theme::AssistantStyle;
46use util::{paths::CONVERSATIONS_DIR, post_inc, ResultExt, TryFutureExt};
47use workspace::{
48 dock::{DockPosition, Panel},
49 searchable::Direction,
50 Save, ToggleZoom, Toolbar, Workspace,
51};
52
53actions!(
54 assistant,
55 [
56 NewConversation,
57 Assist,
58 Split,
59 CycleMessageRole,
60 QuoteSelection,
61 ToggleFocus,
62 ResetKey,
63 InlineAssist
64 ]
65);
66
67pub fn init(cx: &mut AppContext) {
68 settings::register::<AssistantSettings>(cx);
69 cx.add_action(
70 |this: &mut AssistantPanel,
71 _: &workspace::NewFile,
72 cx: &mut ViewContext<AssistantPanel>| {
73 this.new_conversation(cx);
74 },
75 );
76 cx.add_action(ConversationEditor::assist);
77 cx.capture_action(ConversationEditor::cancel_last_assist);
78 cx.capture_action(ConversationEditor::save);
79 cx.add_action(ConversationEditor::quote_selection);
80 cx.capture_action(ConversationEditor::copy);
81 cx.add_action(ConversationEditor::split);
82 cx.capture_action(ConversationEditor::cycle_message_role);
83 cx.add_action(AssistantPanel::save_api_key);
84 cx.add_action(AssistantPanel::reset_api_key);
85 cx.add_action(AssistantPanel::toggle_zoom);
86 cx.add_action(AssistantPanel::deploy);
87 cx.add_action(AssistantPanel::select_next_match);
88 cx.add_action(AssistantPanel::select_prev_match);
89 cx.add_action(AssistantPanel::handle_editor_cancel);
90 cx.add_action(
91 |workspace: &mut Workspace, _: &ToggleFocus, cx: &mut ViewContext<Workspace>| {
92 workspace.toggle_panel_focus::<AssistantPanel>(cx);
93 },
94 );
95 cx.add_action(AssistantPanel::inline_assist);
96 cx.add_action(AssistantPanel::cancel_last_inline_assist);
97 cx.add_action(InlineAssistant::confirm);
98 cx.add_action(InlineAssistant::cancel);
99}
100
101#[derive(Debug)]
102pub enum AssistantPanelEvent {
103 ZoomIn,
104 ZoomOut,
105 Focus,
106 Close,
107 DockPositionChanged,
108}
109
110pub struct AssistantPanel {
111 workspace: WeakViewHandle<Workspace>,
112 width: Option<f32>,
113 height: Option<f32>,
114 active_editor_index: Option<usize>,
115 prev_active_editor_index: Option<usize>,
116 editors: Vec<ViewHandle<ConversationEditor>>,
117 saved_conversations: Vec<SavedConversationMetadata>,
118 saved_conversations_list_state: UniformListState,
119 zoomed: bool,
120 has_focus: bool,
121 toolbar: ViewHandle<Toolbar>,
122 api_key: Rc<RefCell<Option<String>>>,
123 api_key_editor: Option<ViewHandle<Editor>>,
124 has_read_credentials: bool,
125 languages: Arc<LanguageRegistry>,
126 fs: Arc<dyn Fs>,
127 subscriptions: Vec<Subscription>,
128 next_inline_assist_id: usize,
129 pending_inline_assists: HashMap<usize, PendingInlineAssist>,
130 pending_inline_assist_ids_by_editor: HashMap<WeakViewHandle<Editor>, Vec<usize>>,
131 _watch_saved_conversations: Task<Result<()>>,
132}
133
134impl AssistantPanel {
135 pub fn load(
136 workspace: WeakViewHandle<Workspace>,
137 cx: AsyncAppContext,
138 ) -> Task<Result<ViewHandle<Self>>> {
139 cx.spawn(|mut cx| async move {
140 let fs = workspace.read_with(&cx, |workspace, _| workspace.app_state().fs.clone())?;
141 let saved_conversations = SavedConversationMetadata::list(fs.clone())
142 .await
143 .log_err()
144 .unwrap_or_default();
145
146 // TODO: deserialize state.
147 let workspace_handle = workspace.clone();
148 workspace.update(&mut cx, |workspace, cx| {
149 cx.add_view::<Self, _>(|cx| {
150 const CONVERSATION_WATCH_DURATION: Duration = Duration::from_millis(100);
151 let _watch_saved_conversations = cx.spawn(move |this, mut cx| async move {
152 let mut events = fs
153 .watch(&CONVERSATIONS_DIR, CONVERSATION_WATCH_DURATION)
154 .await;
155 while events.next().await.is_some() {
156 let saved_conversations = SavedConversationMetadata::list(fs.clone())
157 .await
158 .log_err()
159 .unwrap_or_default();
160 this.update(&mut cx, |this, cx| {
161 this.saved_conversations = saved_conversations;
162 cx.notify();
163 })
164 .ok();
165 }
166
167 anyhow::Ok(())
168 });
169
170 let toolbar = cx.add_view(|cx| {
171 let mut toolbar = Toolbar::new();
172 toolbar.set_can_navigate(false, cx);
173 toolbar.add_item(cx.add_view(|cx| BufferSearchBar::new(cx)), cx);
174 toolbar
175 });
176 let mut this = Self {
177 workspace: workspace_handle,
178 active_editor_index: Default::default(),
179 prev_active_editor_index: Default::default(),
180 editors: Default::default(),
181 saved_conversations,
182 saved_conversations_list_state: Default::default(),
183 zoomed: false,
184 has_focus: false,
185 toolbar,
186 api_key: Rc::new(RefCell::new(None)),
187 api_key_editor: None,
188 has_read_credentials: false,
189 languages: workspace.app_state().languages.clone(),
190 fs: workspace.app_state().fs.clone(),
191 width: None,
192 height: None,
193 subscriptions: Default::default(),
194 next_inline_assist_id: 0,
195 pending_inline_assists: Default::default(),
196 pending_inline_assist_ids_by_editor: Default::default(),
197 _watch_saved_conversations,
198 };
199
200 let mut old_dock_position = this.position(cx);
201 this.subscriptions =
202 vec![cx.observe_global::<SettingsStore, _>(move |this, cx| {
203 let new_dock_position = this.position(cx);
204 if new_dock_position != old_dock_position {
205 old_dock_position = new_dock_position;
206 cx.emit(AssistantPanelEvent::DockPositionChanged);
207 }
208 cx.notify();
209 })];
210
211 this
212 })
213 })
214 })
215 }
216
217 fn inline_assist(workspace: &mut Workspace, _: &InlineAssist, cx: &mut ViewContext<Workspace>) {
218 let this = if let Some(this) = workspace.panel::<AssistantPanel>(cx) {
219 if this
220 .update(cx, |assistant, cx| assistant.load_api_key(cx))
221 .is_some()
222 {
223 this
224 } else {
225 workspace.focus_panel::<AssistantPanel>(cx);
226 return;
227 }
228 } else {
229 return;
230 };
231
232 let active_editor = if let Some(active_editor) = workspace
233 .active_item(cx)
234 .and_then(|item| item.act_as::<Editor>(cx))
235 {
236 active_editor
237 } else {
238 return;
239 };
240
241 this.update(cx, |assistant, cx| {
242 assistant.new_inline_assist(&active_editor, cx)
243 });
244 }
245
246 fn new_inline_assist(&mut self, editor: &ViewHandle<Editor>, cx: &mut ViewContext<Self>) {
247 let id = post_inc(&mut self.next_inline_assist_id);
248 let snapshot = editor.read(cx).buffer().read(cx).snapshot(cx);
249 let selection = editor.read(cx).selections.newest_anchor().clone();
250 let range = selection.start.bias_left(&snapshot)..selection.end.bias_right(&snapshot);
251 let assist_kind = if editor.read(cx).selections.newest::<usize>(cx).is_empty() {
252 InlineAssistKind::Generate
253 } else {
254 InlineAssistKind::Transform
255 };
256 let prompt_editor = cx.add_view(|cx| {
257 let mut editor = Editor::single_line(
258 Some(Arc::new(|theme| theme.assistant.inline.editor.clone())),
259 cx,
260 );
261 let placeholder = match assist_kind {
262 InlineAssistKind::Transform => "Enter transformation prompt…",
263 InlineAssistKind::Generate => "Enter generation prompt…",
264 };
265 editor.set_placeholder_text(placeholder, cx);
266 editor
267 });
268 let inline_assistant = cx.add_view(|cx| {
269 let assistant = InlineAssistant {
270 id,
271 prompt_editor,
272 confirmed: false,
273 has_focus: false,
274 };
275 cx.focus_self();
276 assistant
277 });
278 let block_id = editor.update(cx, |editor, cx| {
279 editor.highlight_background::<Self>(
280 vec![range.clone()],
281 |theme| theme.assistant.inline.pending_edit_background,
282 cx,
283 );
284 editor.insert_blocks(
285 [BlockProperties {
286 style: BlockStyle::Flex,
287 position: selection.head().bias_left(&snapshot),
288 height: 2,
289 render: Arc::new({
290 let inline_assistant = inline_assistant.clone();
291 move |cx: &mut BlockContext| {
292 let theme = theme::current(cx);
293 ChildView::new(&inline_assistant, cx)
294 .contained()
295 .with_padding_left(cx.anchor_x)
296 .contained()
297 .with_style(theme.assistant.inline.container)
298 .into_any()
299 }
300 }),
301 disposition: if selection.reversed {
302 BlockDisposition::Above
303 } else {
304 BlockDisposition::Below
305 },
306 }],
307 Some(Autoscroll::Strategy(AutoscrollStrategy::Newest)),
308 cx,
309 )[0]
310 });
311
312 self.pending_inline_assists.insert(
313 id,
314 PendingInlineAssist {
315 kind: assist_kind,
316 editor: editor.downgrade(),
317 range,
318 inline_assistant_block_id: Some(block_id),
319 code_generation: Task::ready(None),
320 transaction_id: None,
321 _subscriptions: vec![
322 cx.subscribe(&inline_assistant, Self::handle_inline_assistant_event),
323 cx.subscribe(editor, {
324 let inline_assistant = inline_assistant.downgrade();
325 move |_, editor, event, cx| {
326 if let Some(inline_assistant) = inline_assistant.upgrade(cx) {
327 if let editor::Event::SelectionsChanged { local } = event {
328 if *local && inline_assistant.read(cx).has_focus {
329 cx.focus(&editor);
330 }
331 }
332 }
333 }
334 }),
335 ],
336 },
337 );
338 self.pending_inline_assist_ids_by_editor
339 .entry(editor.downgrade())
340 .or_default()
341 .push(id);
342 }
343
344 fn handle_inline_assistant_event(
345 &mut self,
346 inline_assistant: ViewHandle<InlineAssistant>,
347 event: &InlineAssistantEvent,
348 cx: &mut ViewContext<Self>,
349 ) {
350 let assist_id = inline_assistant.read(cx).id;
351 match event {
352 InlineAssistantEvent::Confirmed { prompt } => {
353 self.confirm_inline_assist(assist_id, prompt, cx);
354 }
355 InlineAssistantEvent::Canceled => {
356 self.close_inline_assist(assist_id, true, cx);
357 }
358 InlineAssistantEvent::Dismissed => {
359 self.hide_inline_assist(assist_id, cx);
360 }
361 }
362 }
363
364 fn cancel_last_inline_assist(
365 workspace: &mut Workspace,
366 _: &editor::Cancel,
367 cx: &mut ViewContext<Workspace>,
368 ) {
369 let panel = if let Some(panel) = workspace.panel::<AssistantPanel>(cx) {
370 panel
371 } else {
372 return;
373 };
374 let editor = if let Some(editor) = workspace
375 .active_item(cx)
376 .and_then(|item| item.downcast::<Editor>())
377 {
378 editor
379 } else {
380 return;
381 };
382
383 let handled = panel.update(cx, |panel, cx| {
384 if let Some(assist_id) = panel
385 .pending_inline_assist_ids_by_editor
386 .get(&editor.downgrade())
387 .and_then(|assist_ids| assist_ids.last().copied())
388 {
389 panel.close_inline_assist(assist_id, true, cx);
390 true
391 } else {
392 false
393 }
394 });
395
396 if !handled {
397 cx.propagate_action();
398 }
399 }
400
401 fn close_inline_assist(&mut self, assist_id: usize, cancel: bool, cx: &mut ViewContext<Self>) {
402 self.hide_inline_assist(assist_id, cx);
403
404 if let Some(pending_assist) = self.pending_inline_assists.remove(&assist_id) {
405 if let hash_map::Entry::Occupied(mut entry) = self
406 .pending_inline_assist_ids_by_editor
407 .entry(pending_assist.editor)
408 {
409 entry.get_mut().retain(|id| *id != assist_id);
410 if entry.get().is_empty() {
411 entry.remove();
412 }
413 }
414
415 if let Some(editor) = pending_assist.editor.upgrade(cx) {
416 editor.update(cx, |editor, cx| {
417 editor.clear_background_highlights::<Self>(cx);
418 editor.clear_text_highlights::<Self>(cx);
419 });
420
421 if cancel {
422 if let Some(transaction_id) = pending_assist.transaction_id {
423 editor.update(cx, |editor, cx| {
424 editor.buffer().update(cx, |buffer, cx| {
425 buffer.undo_and_forget(transaction_id, cx)
426 });
427 });
428 }
429 }
430 }
431 }
432 }
433
434 fn hide_inline_assist(&mut self, assist_id: usize, cx: &mut ViewContext<Self>) {
435 if let Some(pending_assist) = self.pending_inline_assists.get_mut(&assist_id) {
436 if let Some(editor) = pending_assist.editor.upgrade(cx) {
437 if let Some(block_id) = pending_assist.inline_assistant_block_id.take() {
438 editor.update(cx, |editor, cx| {
439 editor.remove_blocks(HashSet::from_iter([block_id]), None, cx);
440 });
441 }
442 }
443 }
444 }
445
446 fn confirm_inline_assist(
447 &mut self,
448 inline_assist_id: usize,
449 user_prompt: &str,
450 cx: &mut ViewContext<Self>,
451 ) {
452 let api_key = if let Some(api_key) = self.api_key.borrow().clone() {
453 api_key
454 } else {
455 return;
456 };
457
458 let pending_assist =
459 if let Some(pending_assist) = self.pending_inline_assists.get_mut(&inline_assist_id) {
460 pending_assist
461 } else {
462 return;
463 };
464
465 let editor = if let Some(editor) = pending_assist.editor.upgrade(cx) {
466 editor
467 } else {
468 return;
469 };
470
471 let range = pending_assist.range.clone();
472 let snapshot = editor.read(cx).buffer().read(cx).snapshot(cx);
473 let selected_text = snapshot
474 .text_for_range(range.start..range.end)
475 .collect::<Rope>();
476
477 let mut base_indentation: Option<language::IndentSize> = None;
478 let selection_start = range.start.to_point(&snapshot);
479 let selection_end = range.end.to_point(&snapshot);
480 let mut start_row = selection_start.row;
481 if snapshot.is_line_blank(start_row) {
482 if let Some(prev_non_blank_row) = snapshot.prev_non_blank_row(start_row) {
483 start_row = prev_non_blank_row;
484 }
485 }
486
487 for row in start_row..=selection_end.row {
488 if snapshot.is_line_blank(row) {
489 continue;
490 }
491
492 let line_indentation = snapshot.indent_size_for_line(row);
493 if let Some(base_indentation) = base_indentation.as_mut() {
494 if line_indentation.len < base_indentation.len {
495 *base_indentation = line_indentation;
496 }
497 } else {
498 base_indentation = Some(line_indentation);
499 }
500 }
501
502 let mut normalized_selected_text = selected_text.clone();
503 if let Some(base_indentation) = base_indentation {
504 for row in selection_start.row..=selection_end.row {
505 let selection_row = row - selection_start.row;
506 let line_start =
507 normalized_selected_text.point_to_offset(Point::new(selection_row, 0));
508 let indentation_len = if row == selection_start.row {
509 base_indentation.len.saturating_sub(selection_start.column)
510 } else {
511 let line_len = normalized_selected_text.line_len(selection_row);
512 cmp::min(line_len, base_indentation.len)
513 };
514 let indentation_end = cmp::min(
515 line_start + indentation_len as usize,
516 normalized_selected_text.len(),
517 );
518 normalized_selected_text.replace(line_start..indentation_end, "");
519 }
520 }
521
522 let language_name = snapshot
523 .language_at(range.start)
524 .map(|language| language.name());
525 let language_name = language_name.as_deref().unwrap_or("");
526 let model = settings::get::<AssistantSettings>(cx)
527 .default_open_ai_model
528 .clone();
529
530 let mut prompt = String::new();
531 writeln!(prompt, "You're an expert {language_name} engineer.").unwrap();
532 match pending_assist.kind {
533 InlineAssistKind::Transform => {
534 writeln!(
535 prompt,
536 "You're currently working inside an editor on this code:"
537 )
538 .unwrap();
539 writeln!(prompt, "```{language_name}").unwrap();
540 for chunk in snapshot.text_for_range(Anchor::min()..Anchor::max()) {
541 write!(prompt, "{chunk}").unwrap();
542 }
543 writeln!(prompt, "```").unwrap();
544
545 writeln!(
546 prompt,
547 "In particular, the user has selected the following code:"
548 )
549 .unwrap();
550 writeln!(prompt, "```{language_name}").unwrap();
551 writeln!(prompt, "{normalized_selected_text}").unwrap();
552 writeln!(prompt, "```").unwrap();
553 writeln!(prompt).unwrap();
554 writeln!(
555 prompt,
556 "Modify the selected code given the user prompt: {user_prompt}"
557 )
558 .unwrap();
559 writeln!(
560 prompt,
561 "You MUST reply only with the edited selected code, not the entire file."
562 )
563 .unwrap();
564 }
565 InlineAssistKind::Generate => {
566 writeln!(
567 prompt,
568 "You're currently working inside an editor on this code:"
569 )
570 .unwrap();
571 writeln!(prompt, "```{language_name}").unwrap();
572 for chunk in snapshot.text_for_range(Anchor::min()..range.start) {
573 write!(prompt, "{chunk}").unwrap();
574 }
575 write!(prompt, "<|>").unwrap();
576 for chunk in snapshot.text_for_range(range.start..Anchor::max()) {
577 write!(prompt, "{chunk}").unwrap();
578 }
579 writeln!(prompt).unwrap();
580 writeln!(prompt, "```").unwrap();
581 writeln!(
582 prompt,
583 "Assume the cursor is located where the `<|>` marker is."
584 )
585 .unwrap();
586 writeln!(prompt, "Assume your answer will be inserted at the cursor.").unwrap();
587 writeln!(
588 prompt,
589 "Complete the code given the user prompt: {user_prompt}"
590 )
591 .unwrap();
592 }
593 }
594 writeln!(prompt, "Your answer MUST always be valid {language_name}.").unwrap();
595 writeln!(prompt, "Always wrap your response in a Markdown codeblock.").unwrap();
596 writeln!(prompt, "Never make remarks, always output code.").unwrap();
597
598 let request = OpenAIRequest {
599 model: model.full_name().into(),
600 messages: vec![RequestMessage {
601 role: Role::User,
602 content: prompt,
603 }],
604 stream: true,
605 };
606 let response = stream_completion(api_key, cx.background().clone(), request);
607 let editor = editor.downgrade();
608
609 pending_assist.code_generation = cx.spawn(|this, mut cx| {
610 async move {
611 let _cleanup = util::defer({
612 let mut cx = cx.clone();
613 let this = this.clone();
614 move || {
615 let _ = this.update(&mut cx, |this, cx| {
616 this.close_inline_assist(inline_assist_id, false, cx)
617 });
618 }
619 });
620
621 let mut edit_start = range.start.to_offset(&snapshot);
622
623 let (mut hunks_tx, mut hunks_rx) = mpsc::channel(1);
624 let diff = cx.background().spawn(async move {
625 let mut messages = response.await?;
626 let mut diff = StreamingDiff::new(selected_text.to_string());
627
628 let indentation_len;
629 let indentation_text;
630 if let Some(base_indentation) = base_indentation {
631 indentation_len = base_indentation.len;
632 indentation_text = match base_indentation.kind {
633 language::IndentKind::Space => " ",
634 language::IndentKind::Tab => "\t",
635 };
636 } else {
637 indentation_len = 0;
638 indentation_text = "";
639 };
640
641 let mut inside_first_line = true;
642 let mut starts_with_fenced_code_block = None;
643 let mut has_pending_newline = false;
644 let mut new_text = String::new();
645
646 while let Some(message) = messages.next().await {
647 let mut message = message?;
648 if let Some(mut choice) = message.choices.pop() {
649 if has_pending_newline {
650 has_pending_newline = false;
651 choice
652 .delta
653 .content
654 .get_or_insert(String::new())
655 .insert(0, '\n');
656 }
657
658 // Buffer a trailing codeblock fence. Note that we don't stop
659 // right away because this may be an inner fence that we need
660 // to insert into the editor.
661 if starts_with_fenced_code_block.is_some()
662 && choice.delta.content.as_deref() == Some("\n```")
663 {
664 new_text.push_str("\n```");
665 continue;
666 }
667
668 // If this was the last completion and we started with a codeblock
669 // fence and we ended with another codeblock fence, then we can
670 // stop right away. Otherwise, whatever text we buffered will be
671 // processed normally.
672 if choice.finish_reason.is_some()
673 && starts_with_fenced_code_block.unwrap_or(false)
674 && new_text == "\n```"
675 {
676 break;
677 }
678
679 if let Some(text) = choice.delta.content {
680 // Never push a newline if there's nothing after it. This is
681 // useful to detect if the newline was pushed because of a
682 // trailing codeblock fence.
683 let text = if let Some(prefix) = text.strip_suffix('\n') {
684 has_pending_newline = true;
685 prefix
686 } else {
687 text.as_str()
688 };
689
690 if text.is_empty() {
691 continue;
692 }
693
694 let mut lines = text.split('\n');
695 if let Some(line) = lines.next() {
696 if starts_with_fenced_code_block.is_none() {
697 starts_with_fenced_code_block =
698 Some(line.starts_with("```"));
699 }
700
701 // Avoid pushing the first line if it's the start of a fenced code block.
702 if !inside_first_line || !starts_with_fenced_code_block.unwrap()
703 {
704 new_text.push_str(&line);
705 }
706 }
707
708 for line in lines {
709 if inside_first_line && starts_with_fenced_code_block.unwrap() {
710 // If we were inside the first line and that line was the
711 // start of a fenced code block, we just need to push the
712 // leading indentation of the original selection.
713 new_text.push_str(&indentation_text.repeat(
714 indentation_len.saturating_sub(selection_start.column)
715 as usize,
716 ));
717 } else {
718 // Otherwise, we need to push a newline and the base indentation.
719 new_text.push('\n');
720 new_text.push_str(
721 &indentation_text.repeat(indentation_len as usize),
722 );
723 }
724
725 new_text.push_str(line);
726 inside_first_line = false;
727 }
728 }
729 }
730
731 let hunks = diff.push_new(&new_text);
732 hunks_tx.send(hunks).await?;
733 new_text.clear();
734 }
735 hunks_tx.send(diff.finish()).await?;
736
737 anyhow::Ok(())
738 });
739
740 while let Some(hunks) = hunks_rx.next().await {
741 let this = this
742 .upgrade(&cx)
743 .ok_or_else(|| anyhow!("assistant was dropped"))?;
744 editor.update(&mut cx, |editor, cx| {
745 let mut highlights = Vec::new();
746
747 let transaction = editor.buffer().update(cx, |buffer, cx| {
748 // Avoid grouping assistant edits with user edits.
749 buffer.finalize_last_transaction(cx);
750
751 buffer.start_transaction(cx);
752 buffer.edit(
753 hunks.into_iter().filter_map(|hunk| match hunk {
754 Hunk::Insert { text } => {
755 let edit_start = snapshot.anchor_after(edit_start);
756 Some((edit_start..edit_start, text))
757 }
758 Hunk::Remove { len } => {
759 let edit_end = edit_start + len;
760 let edit_range = snapshot.anchor_after(edit_start)
761 ..snapshot.anchor_before(edit_end);
762 edit_start = edit_end;
763 Some((edit_range, String::new()))
764 }
765 Hunk::Keep { len } => {
766 let edit_end = edit_start + len;
767 let edit_range = snapshot.anchor_after(edit_start)
768 ..snapshot.anchor_before(edit_end);
769 edit_start += len;
770 highlights.push(edit_range);
771 None
772 }
773 }),
774 None,
775 cx,
776 );
777
778 buffer.end_transaction(cx)
779 });
780
781 if let Some(transaction) = transaction {
782 this.update(cx, |this, cx| {
783 if let Some(pending_assist) =
784 this.pending_inline_assists.get_mut(&inline_assist_id)
785 {
786 if let Some(first_transaction) = pending_assist.transaction_id {
787 // Group all assistant edits into the first transaction.
788 editor.buffer().update(cx, |buffer, cx| {
789 buffer.merge_transactions(
790 transaction,
791 first_transaction,
792 cx,
793 )
794 });
795 } else {
796 pending_assist.transaction_id = Some(transaction);
797 editor.buffer().update(cx, |buffer, cx| {
798 buffer.finalize_last_transaction(cx)
799 });
800 }
801 }
802 });
803 }
804
805 editor.highlight_text::<Self>(
806 highlights,
807 gpui::fonts::HighlightStyle {
808 fade_out: Some(0.6),
809 ..Default::default()
810 },
811 cx,
812 );
813 })?;
814 }
815 diff.await?;
816
817 anyhow::Ok(())
818 }
819 .log_err()
820 });
821 }
822
823 fn new_conversation(&mut self, cx: &mut ViewContext<Self>) -> ViewHandle<ConversationEditor> {
824 let editor = cx.add_view(|cx| {
825 ConversationEditor::new(
826 self.api_key.clone(),
827 self.languages.clone(),
828 self.fs.clone(),
829 cx,
830 )
831 });
832 self.add_conversation(editor.clone(), cx);
833 editor
834 }
835
836 fn add_conversation(
837 &mut self,
838 editor: ViewHandle<ConversationEditor>,
839 cx: &mut ViewContext<Self>,
840 ) {
841 self.subscriptions
842 .push(cx.subscribe(&editor, Self::handle_conversation_editor_event));
843
844 let conversation = editor.read(cx).conversation.clone();
845 self.subscriptions
846 .push(cx.observe(&conversation, |_, _, cx| cx.notify()));
847
848 let index = self.editors.len();
849 self.editors.push(editor);
850 self.set_active_editor_index(Some(index), cx);
851 }
852
853 fn set_active_editor_index(&mut self, index: Option<usize>, cx: &mut ViewContext<Self>) {
854 self.prev_active_editor_index = self.active_editor_index;
855 self.active_editor_index = index;
856 if let Some(editor) = self.active_editor() {
857 let editor = editor.read(cx).editor.clone();
858 self.toolbar.update(cx, |toolbar, cx| {
859 toolbar.set_active_item(Some(&editor), cx);
860 });
861 if self.has_focus(cx) {
862 cx.focus(&editor);
863 }
864 } else {
865 self.toolbar.update(cx, |toolbar, cx| {
866 toolbar.set_active_item(None, cx);
867 });
868 }
869
870 cx.notify();
871 }
872
873 fn handle_conversation_editor_event(
874 &mut self,
875 _: ViewHandle<ConversationEditor>,
876 event: &ConversationEditorEvent,
877 cx: &mut ViewContext<Self>,
878 ) {
879 match event {
880 ConversationEditorEvent::TabContentChanged => cx.notify(),
881 }
882 }
883
884 fn save_api_key(&mut self, _: &menu::Confirm, cx: &mut ViewContext<Self>) {
885 if let Some(api_key) = self
886 .api_key_editor
887 .as_ref()
888 .map(|editor| editor.read(cx).text(cx))
889 {
890 if !api_key.is_empty() {
891 cx.platform()
892 .write_credentials(OPENAI_API_URL, "Bearer", api_key.as_bytes())
893 .log_err();
894 *self.api_key.borrow_mut() = Some(api_key);
895 self.api_key_editor.take();
896 cx.focus_self();
897 cx.notify();
898 }
899 } else {
900 cx.propagate_action();
901 }
902 }
903
904 fn reset_api_key(&mut self, _: &ResetKey, cx: &mut ViewContext<Self>) {
905 cx.platform().delete_credentials(OPENAI_API_URL).log_err();
906 self.api_key.take();
907 self.api_key_editor = Some(build_api_key_editor(cx));
908 cx.focus_self();
909 cx.notify();
910 }
911
912 fn toggle_zoom(&mut self, _: &workspace::ToggleZoom, cx: &mut ViewContext<Self>) {
913 if self.zoomed {
914 cx.emit(AssistantPanelEvent::ZoomOut)
915 } else {
916 cx.emit(AssistantPanelEvent::ZoomIn)
917 }
918 }
919
920 fn deploy(&mut self, action: &search::buffer_search::Deploy, cx: &mut ViewContext<Self>) {
921 let mut propagate_action = true;
922 if let Some(search_bar) = self.toolbar.read(cx).item_of_type::<BufferSearchBar>() {
923 search_bar.update(cx, |search_bar, cx| {
924 if search_bar.show(cx) {
925 search_bar.search_suggested(cx);
926 if action.focus {
927 search_bar.select_query(cx);
928 cx.focus_self();
929 }
930 propagate_action = false
931 }
932 });
933 }
934 if propagate_action {
935 cx.propagate_action();
936 }
937 }
938
939 fn handle_editor_cancel(&mut self, _: &editor::Cancel, cx: &mut ViewContext<Self>) {
940 if let Some(search_bar) = self.toolbar.read(cx).item_of_type::<BufferSearchBar>() {
941 if !search_bar.read(cx).is_dismissed() {
942 search_bar.update(cx, |search_bar, cx| {
943 search_bar.dismiss(&Default::default(), cx)
944 });
945 return;
946 }
947 }
948 cx.propagate_action();
949 }
950
951 fn select_next_match(&mut self, _: &search::SelectNextMatch, cx: &mut ViewContext<Self>) {
952 if let Some(search_bar) = self.toolbar.read(cx).item_of_type::<BufferSearchBar>() {
953 search_bar.update(cx, |bar, cx| bar.select_match(Direction::Next, 1, cx));
954 }
955 }
956
957 fn select_prev_match(&mut self, _: &search::SelectPrevMatch, cx: &mut ViewContext<Self>) {
958 if let Some(search_bar) = self.toolbar.read(cx).item_of_type::<BufferSearchBar>() {
959 search_bar.update(cx, |bar, cx| bar.select_match(Direction::Prev, 1, cx));
960 }
961 }
962
963 fn active_editor(&self) -> Option<&ViewHandle<ConversationEditor>> {
964 self.editors.get(self.active_editor_index?)
965 }
966
967 fn render_hamburger_button(cx: &mut ViewContext<Self>) -> impl Element<Self> {
968 enum History {}
969 let theme = theme::current(cx);
970 let tooltip_style = theme::current(cx).tooltip.clone();
971 MouseEventHandler::new::<History, _>(0, cx, |state, _| {
972 let style = theme.assistant.hamburger_button.style_for(state);
973 Svg::for_style(style.icon.clone())
974 .contained()
975 .with_style(style.container)
976 })
977 .with_cursor_style(CursorStyle::PointingHand)
978 .on_click(MouseButton::Left, |_, this: &mut Self, cx| {
979 if this.active_editor().is_some() {
980 this.set_active_editor_index(None, cx);
981 } else {
982 this.set_active_editor_index(this.prev_active_editor_index, cx);
983 }
984 })
985 .with_tooltip::<History>(1, "History", None, tooltip_style, cx)
986 }
987
988 fn render_editor_tools(&self, cx: &mut ViewContext<Self>) -> Vec<AnyElement<Self>> {
989 if self.active_editor().is_some() {
990 vec![
991 Self::render_split_button(cx).into_any(),
992 Self::render_quote_button(cx).into_any(),
993 Self::render_assist_button(cx).into_any(),
994 ]
995 } else {
996 Default::default()
997 }
998 }
999
1000 fn render_split_button(cx: &mut ViewContext<Self>) -> impl Element<Self> {
1001 let theme = theme::current(cx);
1002 let tooltip_style = theme::current(cx).tooltip.clone();
1003 MouseEventHandler::new::<Split, _>(0, cx, |state, _| {
1004 let style = theme.assistant.split_button.style_for(state);
1005 Svg::for_style(style.icon.clone())
1006 .contained()
1007 .with_style(style.container)
1008 })
1009 .with_cursor_style(CursorStyle::PointingHand)
1010 .on_click(MouseButton::Left, |_, this: &mut Self, cx| {
1011 if let Some(active_editor) = this.active_editor() {
1012 active_editor.update(cx, |editor, cx| editor.split(&Default::default(), cx));
1013 }
1014 })
1015 .with_tooltip::<Split>(
1016 1,
1017 "Split Message",
1018 Some(Box::new(Split)),
1019 tooltip_style,
1020 cx,
1021 )
1022 }
1023
1024 fn render_assist_button(cx: &mut ViewContext<Self>) -> impl Element<Self> {
1025 let theme = theme::current(cx);
1026 let tooltip_style = theme::current(cx).tooltip.clone();
1027 MouseEventHandler::new::<Assist, _>(0, cx, |state, _| {
1028 let style = theme.assistant.assist_button.style_for(state);
1029 Svg::for_style(style.icon.clone())
1030 .contained()
1031 .with_style(style.container)
1032 })
1033 .with_cursor_style(CursorStyle::PointingHand)
1034 .on_click(MouseButton::Left, |_, this: &mut Self, cx| {
1035 if let Some(active_editor) = this.active_editor() {
1036 active_editor.update(cx, |editor, cx| editor.assist(&Default::default(), cx));
1037 }
1038 })
1039 .with_tooltip::<Assist>(1, "Assist", Some(Box::new(Assist)), tooltip_style, cx)
1040 }
1041
1042 fn render_quote_button(cx: &mut ViewContext<Self>) -> impl Element<Self> {
1043 let theme = theme::current(cx);
1044 let tooltip_style = theme::current(cx).tooltip.clone();
1045 MouseEventHandler::new::<QuoteSelection, _>(0, cx, |state, _| {
1046 let style = theme.assistant.quote_button.style_for(state);
1047 Svg::for_style(style.icon.clone())
1048 .contained()
1049 .with_style(style.container)
1050 })
1051 .with_cursor_style(CursorStyle::PointingHand)
1052 .on_click(MouseButton::Left, |_, this: &mut Self, cx| {
1053 if let Some(workspace) = this.workspace.upgrade(cx) {
1054 cx.window_context().defer(move |cx| {
1055 workspace.update(cx, |workspace, cx| {
1056 ConversationEditor::quote_selection(workspace, &Default::default(), cx)
1057 });
1058 });
1059 }
1060 })
1061 .with_tooltip::<QuoteSelection>(
1062 1,
1063 "Quote Selection",
1064 Some(Box::new(QuoteSelection)),
1065 tooltip_style,
1066 cx,
1067 )
1068 }
1069
1070 fn render_plus_button(cx: &mut ViewContext<Self>) -> impl Element<Self> {
1071 let theme = theme::current(cx);
1072 let tooltip_style = theme::current(cx).tooltip.clone();
1073 MouseEventHandler::new::<NewConversation, _>(0, cx, |state, _| {
1074 let style = theme.assistant.plus_button.style_for(state);
1075 Svg::for_style(style.icon.clone())
1076 .contained()
1077 .with_style(style.container)
1078 })
1079 .with_cursor_style(CursorStyle::PointingHand)
1080 .on_click(MouseButton::Left, |_, this: &mut Self, cx| {
1081 this.new_conversation(cx);
1082 })
1083 .with_tooltip::<NewConversation>(
1084 1,
1085 "New Conversation",
1086 Some(Box::new(NewConversation)),
1087 tooltip_style,
1088 cx,
1089 )
1090 }
1091
1092 fn render_zoom_button(&self, cx: &mut ViewContext<Self>) -> impl Element<Self> {
1093 enum ToggleZoomButton {}
1094
1095 let theme = theme::current(cx);
1096 let tooltip_style = theme::current(cx).tooltip.clone();
1097 let style = if self.zoomed {
1098 &theme.assistant.zoom_out_button
1099 } else {
1100 &theme.assistant.zoom_in_button
1101 };
1102
1103 MouseEventHandler::new::<ToggleZoomButton, _>(0, cx, |state, _| {
1104 let style = style.style_for(state);
1105 Svg::for_style(style.icon.clone())
1106 .contained()
1107 .with_style(style.container)
1108 })
1109 .with_cursor_style(CursorStyle::PointingHand)
1110 .on_click(MouseButton::Left, |_, this, cx| {
1111 this.toggle_zoom(&ToggleZoom, cx);
1112 })
1113 .with_tooltip::<ToggleZoom>(
1114 0,
1115 if self.zoomed { "Zoom Out" } else { "Zoom In" },
1116 Some(Box::new(ToggleZoom)),
1117 tooltip_style,
1118 cx,
1119 )
1120 }
1121
1122 fn render_saved_conversation(
1123 &mut self,
1124 index: usize,
1125 cx: &mut ViewContext<Self>,
1126 ) -> impl Element<Self> {
1127 let conversation = &self.saved_conversations[index];
1128 let path = conversation.path.clone();
1129 MouseEventHandler::new::<SavedConversationMetadata, _>(index, cx, move |state, cx| {
1130 let style = &theme::current(cx).assistant.saved_conversation;
1131 Flex::row()
1132 .with_child(
1133 Label::new(
1134 conversation.mtime.format("%F %I:%M%p").to_string(),
1135 style.saved_at.text.clone(),
1136 )
1137 .aligned()
1138 .contained()
1139 .with_style(style.saved_at.container),
1140 )
1141 .with_child(
1142 Label::new(conversation.title.clone(), style.title.text.clone())
1143 .aligned()
1144 .contained()
1145 .with_style(style.title.container),
1146 )
1147 .contained()
1148 .with_style(*style.container.style_for(state))
1149 })
1150 .with_cursor_style(CursorStyle::PointingHand)
1151 .on_click(MouseButton::Left, move |_, this, cx| {
1152 this.open_conversation(path.clone(), cx)
1153 .detach_and_log_err(cx)
1154 })
1155 }
1156
1157 fn open_conversation(&mut self, path: PathBuf, cx: &mut ViewContext<Self>) -> Task<Result<()>> {
1158 if let Some(ix) = self.editor_index_for_path(&path, cx) {
1159 self.set_active_editor_index(Some(ix), cx);
1160 return Task::ready(Ok(()));
1161 }
1162
1163 let fs = self.fs.clone();
1164 let api_key = self.api_key.clone();
1165 let languages = self.languages.clone();
1166 cx.spawn(|this, mut cx| async move {
1167 let saved_conversation = fs.load(&path).await?;
1168 let saved_conversation = serde_json::from_str(&saved_conversation)?;
1169 let conversation = cx.add_model(|cx| {
1170 Conversation::deserialize(saved_conversation, path.clone(), api_key, languages, cx)
1171 });
1172 this.update(&mut cx, |this, cx| {
1173 // If, by the time we've loaded the conversation, the user has already opened
1174 // the same conversation, we don't want to open it again.
1175 if let Some(ix) = this.editor_index_for_path(&path, cx) {
1176 this.set_active_editor_index(Some(ix), cx);
1177 } else {
1178 let editor = cx
1179 .add_view(|cx| ConversationEditor::for_conversation(conversation, fs, cx));
1180 this.add_conversation(editor, cx);
1181 }
1182 })?;
1183 Ok(())
1184 })
1185 }
1186
1187 fn editor_index_for_path(&self, path: &Path, cx: &AppContext) -> Option<usize> {
1188 self.editors
1189 .iter()
1190 .position(|editor| editor.read(cx).conversation.read(cx).path.as_deref() == Some(path))
1191 }
1192
1193 fn load_api_key(&mut self, cx: &mut ViewContext<Self>) -> Option<String> {
1194 if self.api_key.borrow().is_none() && !self.has_read_credentials {
1195 self.has_read_credentials = true;
1196 let api_key = if let Ok(api_key) = env::var("OPENAI_API_KEY") {
1197 Some(api_key)
1198 } else if let Some((_, api_key)) = cx
1199 .platform()
1200 .read_credentials(OPENAI_API_URL)
1201 .log_err()
1202 .flatten()
1203 {
1204 String::from_utf8(api_key).log_err()
1205 } else {
1206 None
1207 };
1208 if let Some(api_key) = api_key {
1209 *self.api_key.borrow_mut() = Some(api_key);
1210 } else if self.api_key_editor.is_none() {
1211 self.api_key_editor = Some(build_api_key_editor(cx));
1212 cx.notify();
1213 }
1214 }
1215
1216 self.api_key.borrow().clone()
1217 }
1218}
1219
1220fn build_api_key_editor(cx: &mut ViewContext<AssistantPanel>) -> ViewHandle<Editor> {
1221 cx.add_view(|cx| {
1222 let mut editor = Editor::single_line(
1223 Some(Arc::new(|theme| theme.assistant.api_key_editor.clone())),
1224 cx,
1225 );
1226 editor.set_placeholder_text("sk-000000000000000000000000000000000000000000000000", cx);
1227 editor
1228 })
1229}
1230
1231impl Entity for AssistantPanel {
1232 type Event = AssistantPanelEvent;
1233}
1234
1235impl View for AssistantPanel {
1236 fn ui_name() -> &'static str {
1237 "AssistantPanel"
1238 }
1239
1240 fn render(&mut self, cx: &mut ViewContext<Self>) -> AnyElement<Self> {
1241 let theme = &theme::current(cx);
1242 let style = &theme.assistant;
1243 if let Some(api_key_editor) = self.api_key_editor.as_ref() {
1244 Flex::column()
1245 .with_child(
1246 Text::new(
1247 "Paste your OpenAI API key and press Enter to use the assistant",
1248 style.api_key_prompt.text.clone(),
1249 )
1250 .aligned(),
1251 )
1252 .with_child(
1253 ChildView::new(api_key_editor, cx)
1254 .contained()
1255 .with_style(style.api_key_editor.container)
1256 .aligned(),
1257 )
1258 .contained()
1259 .with_style(style.api_key_prompt.container)
1260 .aligned()
1261 .into_any()
1262 } else {
1263 let title = self.active_editor().map(|editor| {
1264 Label::new(editor.read(cx).title(cx), style.title.text.clone())
1265 .contained()
1266 .with_style(style.title.container)
1267 .aligned()
1268 .left()
1269 .flex(1., false)
1270 });
1271 let mut header = Flex::row()
1272 .with_child(Self::render_hamburger_button(cx).aligned())
1273 .with_children(title);
1274 if self.has_focus {
1275 header.add_children(
1276 self.render_editor_tools(cx)
1277 .into_iter()
1278 .map(|tool| tool.aligned().flex_float()),
1279 );
1280 header.add_child(Self::render_plus_button(cx).aligned().flex_float());
1281 header.add_child(self.render_zoom_button(cx).aligned());
1282 }
1283
1284 Flex::column()
1285 .with_child(
1286 header
1287 .contained()
1288 .with_style(theme.workspace.tab_bar.container)
1289 .expanded()
1290 .constrained()
1291 .with_height(theme.workspace.tab_bar.height),
1292 )
1293 .with_children(if self.toolbar.read(cx).hidden() {
1294 None
1295 } else {
1296 Some(ChildView::new(&self.toolbar, cx).expanded())
1297 })
1298 .with_child(if let Some(editor) = self.active_editor() {
1299 ChildView::new(editor, cx).flex(1., true).into_any()
1300 } else {
1301 UniformList::new(
1302 self.saved_conversations_list_state.clone(),
1303 self.saved_conversations.len(),
1304 cx,
1305 |this, range, items, cx| {
1306 for ix in range {
1307 items.push(this.render_saved_conversation(ix, cx).into_any());
1308 }
1309 },
1310 )
1311 .flex(1., true)
1312 .into_any()
1313 })
1314 .into_any()
1315 }
1316 }
1317
1318 fn focus_in(&mut self, _: gpui::AnyViewHandle, cx: &mut ViewContext<Self>) {
1319 self.has_focus = true;
1320 self.toolbar
1321 .update(cx, |toolbar, cx| toolbar.focus_changed(true, cx));
1322 cx.notify();
1323 if cx.is_self_focused() {
1324 if let Some(editor) = self.active_editor() {
1325 cx.focus(editor);
1326 } else if let Some(api_key_editor) = self.api_key_editor.as_ref() {
1327 cx.focus(api_key_editor);
1328 }
1329 }
1330 }
1331
1332 fn focus_out(&mut self, _: gpui::AnyViewHandle, cx: &mut ViewContext<Self>) {
1333 self.has_focus = false;
1334 self.toolbar
1335 .update(cx, |toolbar, cx| toolbar.focus_changed(false, cx));
1336 cx.notify();
1337 }
1338}
1339
1340impl Panel for AssistantPanel {
1341 fn position(&self, cx: &WindowContext) -> DockPosition {
1342 match settings::get::<AssistantSettings>(cx).dock {
1343 AssistantDockPosition::Left => DockPosition::Left,
1344 AssistantDockPosition::Bottom => DockPosition::Bottom,
1345 AssistantDockPosition::Right => DockPosition::Right,
1346 }
1347 }
1348
1349 fn position_is_valid(&self, _: DockPosition) -> bool {
1350 true
1351 }
1352
1353 fn set_position(&mut self, position: DockPosition, cx: &mut ViewContext<Self>) {
1354 settings::update_settings_file::<AssistantSettings>(self.fs.clone(), cx, move |settings| {
1355 let dock = match position {
1356 DockPosition::Left => AssistantDockPosition::Left,
1357 DockPosition::Bottom => AssistantDockPosition::Bottom,
1358 DockPosition::Right => AssistantDockPosition::Right,
1359 };
1360 settings.dock = Some(dock);
1361 });
1362 }
1363
1364 fn size(&self, cx: &WindowContext) -> f32 {
1365 let settings = settings::get::<AssistantSettings>(cx);
1366 match self.position(cx) {
1367 DockPosition::Left | DockPosition::Right => {
1368 self.width.unwrap_or_else(|| settings.default_width)
1369 }
1370 DockPosition::Bottom => self.height.unwrap_or_else(|| settings.default_height),
1371 }
1372 }
1373
1374 fn set_size(&mut self, size: Option<f32>, cx: &mut ViewContext<Self>) {
1375 match self.position(cx) {
1376 DockPosition::Left | DockPosition::Right => self.width = size,
1377 DockPosition::Bottom => self.height = size,
1378 }
1379 cx.notify();
1380 }
1381
1382 fn should_zoom_in_on_event(event: &AssistantPanelEvent) -> bool {
1383 matches!(event, AssistantPanelEvent::ZoomIn)
1384 }
1385
1386 fn should_zoom_out_on_event(event: &AssistantPanelEvent) -> bool {
1387 matches!(event, AssistantPanelEvent::ZoomOut)
1388 }
1389
1390 fn is_zoomed(&self, _: &WindowContext) -> bool {
1391 self.zoomed
1392 }
1393
1394 fn set_zoomed(&mut self, zoomed: bool, cx: &mut ViewContext<Self>) {
1395 self.zoomed = zoomed;
1396 cx.notify();
1397 }
1398
1399 fn set_active(&mut self, active: bool, cx: &mut ViewContext<Self>) {
1400 if active {
1401 self.load_api_key(cx);
1402
1403 if self.editors.is_empty() {
1404 self.new_conversation(cx);
1405 }
1406 }
1407 }
1408
1409 fn icon_path(&self, cx: &WindowContext) -> Option<&'static str> {
1410 settings::get::<AssistantSettings>(cx)
1411 .button
1412 .then(|| "icons/ai.svg")
1413 }
1414
1415 fn icon_tooltip(&self) -> (String, Option<Box<dyn Action>>) {
1416 ("Assistant Panel".into(), Some(Box::new(ToggleFocus)))
1417 }
1418
1419 fn should_change_position_on_event(event: &Self::Event) -> bool {
1420 matches!(event, AssistantPanelEvent::DockPositionChanged)
1421 }
1422
1423 fn should_activate_on_event(_: &Self::Event) -> bool {
1424 false
1425 }
1426
1427 fn should_close_on_event(event: &AssistantPanelEvent) -> bool {
1428 matches!(event, AssistantPanelEvent::Close)
1429 }
1430
1431 fn has_focus(&self, _: &WindowContext) -> bool {
1432 self.has_focus
1433 }
1434
1435 fn is_focus_event(event: &Self::Event) -> bool {
1436 matches!(event, AssistantPanelEvent::Focus)
1437 }
1438}
1439
1440enum ConversationEvent {
1441 MessagesEdited,
1442 SummaryChanged,
1443 StreamedCompletion,
1444}
1445
1446#[derive(Default)]
1447struct Summary {
1448 text: String,
1449 done: bool,
1450}
1451
1452struct Conversation {
1453 buffer: ModelHandle<Buffer>,
1454 message_anchors: Vec<MessageAnchor>,
1455 messages_metadata: HashMap<MessageId, MessageMetadata>,
1456 next_message_id: MessageId,
1457 summary: Option<Summary>,
1458 pending_summary: Task<Option<()>>,
1459 completion_count: usize,
1460 pending_completions: Vec<PendingCompletion>,
1461 model: OpenAIModel,
1462 token_count: Option<usize>,
1463 max_token_count: usize,
1464 pending_token_count: Task<Option<()>>,
1465 api_key: Rc<RefCell<Option<String>>>,
1466 pending_save: Task<Result<()>>,
1467 path: Option<PathBuf>,
1468 _subscriptions: Vec<Subscription>,
1469}
1470
1471impl Entity for Conversation {
1472 type Event = ConversationEvent;
1473}
1474
1475impl Conversation {
1476 fn new(
1477 api_key: Rc<RefCell<Option<String>>>,
1478 language_registry: Arc<LanguageRegistry>,
1479 cx: &mut ModelContext<Self>,
1480 ) -> Self {
1481 let markdown = language_registry.language_for_name("Markdown");
1482 let buffer = cx.add_model(|cx| {
1483 let mut buffer = Buffer::new(0, "", cx);
1484 buffer.set_language_registry(language_registry);
1485 cx.spawn_weak(|buffer, mut cx| async move {
1486 let markdown = markdown.await?;
1487 let buffer = buffer
1488 .upgrade(&cx)
1489 .ok_or_else(|| anyhow!("buffer was dropped"))?;
1490 buffer.update(&mut cx, |buffer, cx| {
1491 buffer.set_language(Some(markdown), cx)
1492 });
1493 anyhow::Ok(())
1494 })
1495 .detach_and_log_err(cx);
1496 buffer
1497 });
1498
1499 let settings = settings::get::<AssistantSettings>(cx);
1500 let model = settings.default_open_ai_model.clone();
1501
1502 let mut this = Self {
1503 message_anchors: Default::default(),
1504 messages_metadata: Default::default(),
1505 next_message_id: Default::default(),
1506 summary: None,
1507 pending_summary: Task::ready(None),
1508 completion_count: Default::default(),
1509 pending_completions: Default::default(),
1510 token_count: None,
1511 max_token_count: tiktoken_rs::model::get_context_size(&model.full_name()),
1512 pending_token_count: Task::ready(None),
1513 model: model.clone(),
1514 _subscriptions: vec![cx.subscribe(&buffer, Self::handle_buffer_event)],
1515 pending_save: Task::ready(Ok(())),
1516 path: None,
1517 api_key,
1518 buffer,
1519 };
1520 let message = MessageAnchor {
1521 id: MessageId(post_inc(&mut this.next_message_id.0)),
1522 start: language::Anchor::MIN,
1523 };
1524 this.message_anchors.push(message.clone());
1525 this.messages_metadata.insert(
1526 message.id,
1527 MessageMetadata {
1528 role: Role::User,
1529 sent_at: Local::now(),
1530 status: MessageStatus::Done,
1531 },
1532 );
1533
1534 this.count_remaining_tokens(cx);
1535 this
1536 }
1537
1538 fn serialize(&self, cx: &AppContext) -> SavedConversation {
1539 SavedConversation {
1540 zed: "conversation".into(),
1541 version: SavedConversation::VERSION.into(),
1542 text: self.buffer.read(cx).text(),
1543 message_metadata: self.messages_metadata.clone(),
1544 messages: self
1545 .messages(cx)
1546 .map(|message| SavedMessage {
1547 id: message.id,
1548 start: message.offset_range.start,
1549 })
1550 .collect(),
1551 summary: self
1552 .summary
1553 .as_ref()
1554 .map(|summary| summary.text.clone())
1555 .unwrap_or_default(),
1556 model: self.model.clone(),
1557 }
1558 }
1559
1560 fn deserialize(
1561 saved_conversation: SavedConversation,
1562 path: PathBuf,
1563 api_key: Rc<RefCell<Option<String>>>,
1564 language_registry: Arc<LanguageRegistry>,
1565 cx: &mut ModelContext<Self>,
1566 ) -> Self {
1567 let model = saved_conversation.model;
1568 let markdown = language_registry.language_for_name("Markdown");
1569 let mut message_anchors = Vec::new();
1570 let mut next_message_id = MessageId(0);
1571 let buffer = cx.add_model(|cx| {
1572 let mut buffer = Buffer::new(0, saved_conversation.text, cx);
1573 for message in saved_conversation.messages {
1574 message_anchors.push(MessageAnchor {
1575 id: message.id,
1576 start: buffer.anchor_before(message.start),
1577 });
1578 next_message_id = cmp::max(next_message_id, MessageId(message.id.0 + 1));
1579 }
1580 buffer.set_language_registry(language_registry);
1581 cx.spawn_weak(|buffer, mut cx| async move {
1582 let markdown = markdown.await?;
1583 let buffer = buffer
1584 .upgrade(&cx)
1585 .ok_or_else(|| anyhow!("buffer was dropped"))?;
1586 buffer.update(&mut cx, |buffer, cx| {
1587 buffer.set_language(Some(markdown), cx)
1588 });
1589 anyhow::Ok(())
1590 })
1591 .detach_and_log_err(cx);
1592 buffer
1593 });
1594
1595 let mut this = Self {
1596 message_anchors,
1597 messages_metadata: saved_conversation.message_metadata,
1598 next_message_id,
1599 summary: Some(Summary {
1600 text: saved_conversation.summary,
1601 done: true,
1602 }),
1603 pending_summary: Task::ready(None),
1604 completion_count: Default::default(),
1605 pending_completions: Default::default(),
1606 token_count: None,
1607 max_token_count: tiktoken_rs::model::get_context_size(&model.full_name()),
1608 pending_token_count: Task::ready(None),
1609 model,
1610 _subscriptions: vec![cx.subscribe(&buffer, Self::handle_buffer_event)],
1611 pending_save: Task::ready(Ok(())),
1612 path: Some(path),
1613 api_key,
1614 buffer,
1615 };
1616 this.count_remaining_tokens(cx);
1617 this
1618 }
1619
1620 fn handle_buffer_event(
1621 &mut self,
1622 _: ModelHandle<Buffer>,
1623 event: &language::Event,
1624 cx: &mut ModelContext<Self>,
1625 ) {
1626 match event {
1627 language::Event::Edited => {
1628 self.count_remaining_tokens(cx);
1629 cx.emit(ConversationEvent::MessagesEdited);
1630 }
1631 _ => {}
1632 }
1633 }
1634
1635 fn count_remaining_tokens(&mut self, cx: &mut ModelContext<Self>) {
1636 let messages = self
1637 .messages(cx)
1638 .into_iter()
1639 .filter_map(|message| {
1640 Some(tiktoken_rs::ChatCompletionRequestMessage {
1641 role: match message.role {
1642 Role::User => "user".into(),
1643 Role::Assistant => "assistant".into(),
1644 Role::System => "system".into(),
1645 },
1646 content: self
1647 .buffer
1648 .read(cx)
1649 .text_for_range(message.offset_range)
1650 .collect(),
1651 name: None,
1652 })
1653 })
1654 .collect::<Vec<_>>();
1655 let model = self.model.clone();
1656 self.pending_token_count = cx.spawn_weak(|this, mut cx| {
1657 async move {
1658 cx.background().timer(Duration::from_millis(200)).await;
1659 let token_count = cx
1660 .background()
1661 .spawn(async move {
1662 tiktoken_rs::num_tokens_from_messages(&model.full_name(), &messages)
1663 })
1664 .await?;
1665
1666 this.upgrade(&cx)
1667 .ok_or_else(|| anyhow!("conversation was dropped"))?
1668 .update(&mut cx, |this, cx| {
1669 this.max_token_count =
1670 tiktoken_rs::model::get_context_size(&this.model.full_name());
1671 this.token_count = Some(token_count);
1672 cx.notify()
1673 });
1674 anyhow::Ok(())
1675 }
1676 .log_err()
1677 });
1678 }
1679
1680 fn remaining_tokens(&self) -> Option<isize> {
1681 Some(self.max_token_count as isize - self.token_count? as isize)
1682 }
1683
1684 fn set_model(&mut self, model: OpenAIModel, cx: &mut ModelContext<Self>) {
1685 self.model = model;
1686 self.count_remaining_tokens(cx);
1687 cx.notify();
1688 }
1689
1690 fn assist(
1691 &mut self,
1692 selected_messages: HashSet<MessageId>,
1693 cx: &mut ModelContext<Self>,
1694 ) -> Vec<MessageAnchor> {
1695 let mut user_messages = Vec::new();
1696 let mut tasks = Vec::new();
1697
1698 let last_message_id = self.message_anchors.iter().rev().find_map(|message| {
1699 message
1700 .start
1701 .is_valid(self.buffer.read(cx))
1702 .then_some(message.id)
1703 });
1704
1705 for selected_message_id in selected_messages {
1706 let selected_message_role =
1707 if let Some(metadata) = self.messages_metadata.get(&selected_message_id) {
1708 metadata.role
1709 } else {
1710 continue;
1711 };
1712
1713 if selected_message_role == Role::Assistant {
1714 if let Some(user_message) = self.insert_message_after(
1715 selected_message_id,
1716 Role::User,
1717 MessageStatus::Done,
1718 cx,
1719 ) {
1720 user_messages.push(user_message);
1721 } else {
1722 continue;
1723 }
1724 } else {
1725 let request = OpenAIRequest {
1726 model: self.model.full_name().to_string(),
1727 messages: self
1728 .messages(cx)
1729 .filter(|message| matches!(message.status, MessageStatus::Done))
1730 .flat_map(|message| {
1731 let mut system_message = None;
1732 if message.id == selected_message_id {
1733 system_message = Some(RequestMessage {
1734 role: Role::System,
1735 content: concat!(
1736 "Treat the following messages as additional knowledge you have learned about, ",
1737 "but act as if they were not part of this conversation. That is, treat them ",
1738 "as if the user didn't see them and couldn't possibly inquire about them."
1739 ).into()
1740 });
1741 }
1742
1743 Some(message.to_open_ai_message(self.buffer.read(cx))).into_iter().chain(system_message)
1744 })
1745 .chain(Some(RequestMessage {
1746 role: Role::System,
1747 content: format!(
1748 "Direct your reply to message with id {}. Do not include a [Message X] header.",
1749 selected_message_id.0
1750 ),
1751 }))
1752 .collect(),
1753 stream: true,
1754 };
1755
1756 let Some(api_key) = self.api_key.borrow().clone() else { continue };
1757 let stream = stream_completion(api_key, cx.background().clone(), request);
1758 let assistant_message = self
1759 .insert_message_after(
1760 selected_message_id,
1761 Role::Assistant,
1762 MessageStatus::Pending,
1763 cx,
1764 )
1765 .unwrap();
1766
1767 // Queue up the user's next reply
1768 if Some(selected_message_id) == last_message_id {
1769 let user_message = self
1770 .insert_message_after(
1771 assistant_message.id,
1772 Role::User,
1773 MessageStatus::Done,
1774 cx,
1775 )
1776 .unwrap();
1777 user_messages.push(user_message);
1778 }
1779
1780 tasks.push(cx.spawn_weak({
1781 |this, mut cx| async move {
1782 let assistant_message_id = assistant_message.id;
1783 let stream_completion = async {
1784 let mut messages = stream.await?;
1785
1786 while let Some(message) = messages.next().await {
1787 let mut message = message?;
1788 if let Some(choice) = message.choices.pop() {
1789 this.upgrade(&cx)
1790 .ok_or_else(|| anyhow!("conversation was dropped"))?
1791 .update(&mut cx, |this, cx| {
1792 let text: Arc<str> = choice.delta.content?.into();
1793 let message_ix = this.message_anchors.iter().position(
1794 |message| message.id == assistant_message_id,
1795 )?;
1796 this.buffer.update(cx, |buffer, cx| {
1797 let offset = this.message_anchors[message_ix + 1..]
1798 .iter()
1799 .find(|message| message.start.is_valid(buffer))
1800 .map_or(buffer.len(), |message| {
1801 message
1802 .start
1803 .to_offset(buffer)
1804 .saturating_sub(1)
1805 });
1806 buffer.edit([(offset..offset, text)], None, cx);
1807 });
1808 cx.emit(ConversationEvent::StreamedCompletion);
1809
1810 Some(())
1811 });
1812 }
1813 smol::future::yield_now().await;
1814 }
1815
1816 this.upgrade(&cx)
1817 .ok_or_else(|| anyhow!("conversation was dropped"))?
1818 .update(&mut cx, |this, cx| {
1819 this.pending_completions.retain(|completion| {
1820 completion.id != this.completion_count
1821 });
1822 this.summarize(cx);
1823 });
1824
1825 anyhow::Ok(())
1826 };
1827
1828 let result = stream_completion.await;
1829 if let Some(this) = this.upgrade(&cx) {
1830 this.update(&mut cx, |this, cx| {
1831 if let Some(metadata) =
1832 this.messages_metadata.get_mut(&assistant_message.id)
1833 {
1834 match result {
1835 Ok(_) => {
1836 metadata.status = MessageStatus::Done;
1837 }
1838 Err(error) => {
1839 metadata.status = MessageStatus::Error(
1840 error.to_string().trim().into(),
1841 );
1842 }
1843 }
1844 cx.notify();
1845 }
1846 });
1847 }
1848 }
1849 }));
1850 }
1851 }
1852
1853 if !tasks.is_empty() {
1854 self.pending_completions.push(PendingCompletion {
1855 id: post_inc(&mut self.completion_count),
1856 _tasks: tasks,
1857 });
1858 }
1859
1860 user_messages
1861 }
1862
1863 fn cancel_last_assist(&mut self) -> bool {
1864 self.pending_completions.pop().is_some()
1865 }
1866
1867 fn cycle_message_roles(&mut self, ids: HashSet<MessageId>, cx: &mut ModelContext<Self>) {
1868 for id in ids {
1869 if let Some(metadata) = self.messages_metadata.get_mut(&id) {
1870 metadata.role.cycle();
1871 cx.emit(ConversationEvent::MessagesEdited);
1872 cx.notify();
1873 }
1874 }
1875 }
1876
1877 fn insert_message_after(
1878 &mut self,
1879 message_id: MessageId,
1880 role: Role,
1881 status: MessageStatus,
1882 cx: &mut ModelContext<Self>,
1883 ) -> Option<MessageAnchor> {
1884 if let Some(prev_message_ix) = self
1885 .message_anchors
1886 .iter()
1887 .position(|message| message.id == message_id)
1888 {
1889 // Find the next valid message after the one we were given.
1890 let mut next_message_ix = prev_message_ix + 1;
1891 while let Some(next_message) = self.message_anchors.get(next_message_ix) {
1892 if next_message.start.is_valid(self.buffer.read(cx)) {
1893 break;
1894 }
1895 next_message_ix += 1;
1896 }
1897
1898 let start = self.buffer.update(cx, |buffer, cx| {
1899 let offset = self
1900 .message_anchors
1901 .get(next_message_ix)
1902 .map_or(buffer.len(), |message| message.start.to_offset(buffer) - 1);
1903 buffer.edit([(offset..offset, "\n")], None, cx);
1904 buffer.anchor_before(offset + 1)
1905 });
1906 let message = MessageAnchor {
1907 id: MessageId(post_inc(&mut self.next_message_id.0)),
1908 start,
1909 };
1910 self.message_anchors
1911 .insert(next_message_ix, message.clone());
1912 self.messages_metadata.insert(
1913 message.id,
1914 MessageMetadata {
1915 role,
1916 sent_at: Local::now(),
1917 status,
1918 },
1919 );
1920 cx.emit(ConversationEvent::MessagesEdited);
1921 Some(message)
1922 } else {
1923 None
1924 }
1925 }
1926
1927 fn split_message(
1928 &mut self,
1929 range: Range<usize>,
1930 cx: &mut ModelContext<Self>,
1931 ) -> (Option<MessageAnchor>, Option<MessageAnchor>) {
1932 let start_message = self.message_for_offset(range.start, cx);
1933 let end_message = self.message_for_offset(range.end, cx);
1934 if let Some((start_message, end_message)) = start_message.zip(end_message) {
1935 // Prevent splitting when range spans multiple messages.
1936 if start_message.id != end_message.id {
1937 return (None, None);
1938 }
1939
1940 let message = start_message;
1941 let role = message.role;
1942 let mut edited_buffer = false;
1943
1944 let mut suffix_start = None;
1945 if range.start > message.offset_range.start && range.end < message.offset_range.end - 1
1946 {
1947 if self.buffer.read(cx).chars_at(range.end).next() == Some('\n') {
1948 suffix_start = Some(range.end + 1);
1949 } else if self.buffer.read(cx).reversed_chars_at(range.end).next() == Some('\n') {
1950 suffix_start = Some(range.end);
1951 }
1952 }
1953
1954 let suffix = if let Some(suffix_start) = suffix_start {
1955 MessageAnchor {
1956 id: MessageId(post_inc(&mut self.next_message_id.0)),
1957 start: self.buffer.read(cx).anchor_before(suffix_start),
1958 }
1959 } else {
1960 self.buffer.update(cx, |buffer, cx| {
1961 buffer.edit([(range.end..range.end, "\n")], None, cx);
1962 });
1963 edited_buffer = true;
1964 MessageAnchor {
1965 id: MessageId(post_inc(&mut self.next_message_id.0)),
1966 start: self.buffer.read(cx).anchor_before(range.end + 1),
1967 }
1968 };
1969
1970 self.message_anchors
1971 .insert(message.index_range.end + 1, suffix.clone());
1972 self.messages_metadata.insert(
1973 suffix.id,
1974 MessageMetadata {
1975 role,
1976 sent_at: Local::now(),
1977 status: MessageStatus::Done,
1978 },
1979 );
1980
1981 let new_messages =
1982 if range.start == range.end || range.start == message.offset_range.start {
1983 (None, Some(suffix))
1984 } else {
1985 let mut prefix_end = None;
1986 if range.start > message.offset_range.start
1987 && range.end < message.offset_range.end - 1
1988 {
1989 if self.buffer.read(cx).chars_at(range.start).next() == Some('\n') {
1990 prefix_end = Some(range.start + 1);
1991 } else if self.buffer.read(cx).reversed_chars_at(range.start).next()
1992 == Some('\n')
1993 {
1994 prefix_end = Some(range.start);
1995 }
1996 }
1997
1998 let selection = if let Some(prefix_end) = prefix_end {
1999 cx.emit(ConversationEvent::MessagesEdited);
2000 MessageAnchor {
2001 id: MessageId(post_inc(&mut self.next_message_id.0)),
2002 start: self.buffer.read(cx).anchor_before(prefix_end),
2003 }
2004 } else {
2005 self.buffer.update(cx, |buffer, cx| {
2006 buffer.edit([(range.start..range.start, "\n")], None, cx)
2007 });
2008 edited_buffer = true;
2009 MessageAnchor {
2010 id: MessageId(post_inc(&mut self.next_message_id.0)),
2011 start: self.buffer.read(cx).anchor_before(range.end + 1),
2012 }
2013 };
2014
2015 self.message_anchors
2016 .insert(message.index_range.end + 1, selection.clone());
2017 self.messages_metadata.insert(
2018 selection.id,
2019 MessageMetadata {
2020 role,
2021 sent_at: Local::now(),
2022 status: MessageStatus::Done,
2023 },
2024 );
2025 (Some(selection), Some(suffix))
2026 };
2027
2028 if !edited_buffer {
2029 cx.emit(ConversationEvent::MessagesEdited);
2030 }
2031 new_messages
2032 } else {
2033 (None, None)
2034 }
2035 }
2036
2037 fn summarize(&mut self, cx: &mut ModelContext<Self>) {
2038 if self.message_anchors.len() >= 2 && self.summary.is_none() {
2039 let api_key = self.api_key.borrow().clone();
2040 if let Some(api_key) = api_key {
2041 let messages = self
2042 .messages(cx)
2043 .take(2)
2044 .map(|message| message.to_open_ai_message(self.buffer.read(cx)))
2045 .chain(Some(RequestMessage {
2046 role: Role::User,
2047 content:
2048 "Summarize the conversation into a short title without punctuation"
2049 .into(),
2050 }));
2051 let request = OpenAIRequest {
2052 model: self.model.full_name().to_string(),
2053 messages: messages.collect(),
2054 stream: true,
2055 };
2056
2057 let stream = stream_completion(api_key, cx.background().clone(), request);
2058 self.pending_summary = cx.spawn(|this, mut cx| {
2059 async move {
2060 let mut messages = stream.await?;
2061
2062 while let Some(message) = messages.next().await {
2063 let mut message = message?;
2064 if let Some(choice) = message.choices.pop() {
2065 let text = choice.delta.content.unwrap_or_default();
2066 this.update(&mut cx, |this, cx| {
2067 this.summary
2068 .get_or_insert(Default::default())
2069 .text
2070 .push_str(&text);
2071 cx.emit(ConversationEvent::SummaryChanged);
2072 });
2073 }
2074 }
2075
2076 this.update(&mut cx, |this, cx| {
2077 if let Some(summary) = this.summary.as_mut() {
2078 summary.done = true;
2079 cx.emit(ConversationEvent::SummaryChanged);
2080 }
2081 });
2082
2083 anyhow::Ok(())
2084 }
2085 .log_err()
2086 });
2087 }
2088 }
2089 }
2090
2091 fn message_for_offset(&self, offset: usize, cx: &AppContext) -> Option<Message> {
2092 self.messages_for_offsets([offset], cx).pop()
2093 }
2094
2095 fn messages_for_offsets(
2096 &self,
2097 offsets: impl IntoIterator<Item = usize>,
2098 cx: &AppContext,
2099 ) -> Vec<Message> {
2100 let mut result = Vec::new();
2101
2102 let mut messages = self.messages(cx).peekable();
2103 let mut offsets = offsets.into_iter().peekable();
2104 let mut current_message = messages.next();
2105 while let Some(offset) = offsets.next() {
2106 // Locate the message that contains the offset.
2107 while current_message.as_ref().map_or(false, |message| {
2108 !message.offset_range.contains(&offset) && messages.peek().is_some()
2109 }) {
2110 current_message = messages.next();
2111 }
2112 let Some(message) = current_message.as_ref() else { break };
2113
2114 // Skip offsets that are in the same message.
2115 while offsets.peek().map_or(false, |offset| {
2116 message.offset_range.contains(offset) || messages.peek().is_none()
2117 }) {
2118 offsets.next();
2119 }
2120
2121 result.push(message.clone());
2122 }
2123 result
2124 }
2125
2126 fn messages<'a>(&'a self, cx: &'a AppContext) -> impl 'a + Iterator<Item = Message> {
2127 let buffer = self.buffer.read(cx);
2128 let mut message_anchors = self.message_anchors.iter().enumerate().peekable();
2129 iter::from_fn(move || {
2130 while let Some((start_ix, message_anchor)) = message_anchors.next() {
2131 let metadata = self.messages_metadata.get(&message_anchor.id)?;
2132 let message_start = message_anchor.start.to_offset(buffer);
2133 let mut message_end = None;
2134 let mut end_ix = start_ix;
2135 while let Some((_, next_message)) = message_anchors.peek() {
2136 if next_message.start.is_valid(buffer) {
2137 message_end = Some(next_message.start);
2138 break;
2139 } else {
2140 end_ix += 1;
2141 message_anchors.next();
2142 }
2143 }
2144 let message_end = message_end
2145 .unwrap_or(language::Anchor::MAX)
2146 .to_offset(buffer);
2147 return Some(Message {
2148 index_range: start_ix..end_ix,
2149 offset_range: message_start..message_end,
2150 id: message_anchor.id,
2151 anchor: message_anchor.start,
2152 role: metadata.role,
2153 sent_at: metadata.sent_at,
2154 status: metadata.status.clone(),
2155 });
2156 }
2157 None
2158 })
2159 }
2160
2161 fn save(
2162 &mut self,
2163 debounce: Option<Duration>,
2164 fs: Arc<dyn Fs>,
2165 cx: &mut ModelContext<Conversation>,
2166 ) {
2167 self.pending_save = cx.spawn(|this, mut cx| async move {
2168 if let Some(debounce) = debounce {
2169 cx.background().timer(debounce).await;
2170 }
2171
2172 let (old_path, summary) = this.read_with(&cx, |this, _| {
2173 let path = this.path.clone();
2174 let summary = if let Some(summary) = this.summary.as_ref() {
2175 if summary.done {
2176 Some(summary.text.clone())
2177 } else {
2178 None
2179 }
2180 } else {
2181 None
2182 };
2183 (path, summary)
2184 });
2185
2186 if let Some(summary) = summary {
2187 let conversation = this.read_with(&cx, |this, cx| this.serialize(cx));
2188 let path = if let Some(old_path) = old_path {
2189 old_path
2190 } else {
2191 let mut discriminant = 1;
2192 let mut new_path;
2193 loop {
2194 new_path = CONVERSATIONS_DIR.join(&format!(
2195 "{} - {}.zed.json",
2196 summary.trim(),
2197 discriminant
2198 ));
2199 if fs.is_file(&new_path).await {
2200 discriminant += 1;
2201 } else {
2202 break;
2203 }
2204 }
2205 new_path
2206 };
2207
2208 fs.create_dir(CONVERSATIONS_DIR.as_ref()).await?;
2209 fs.atomic_write(path.clone(), serde_json::to_string(&conversation).unwrap())
2210 .await?;
2211 this.update(&mut cx, |this, _| this.path = Some(path));
2212 }
2213
2214 Ok(())
2215 });
2216 }
2217}
2218
2219struct PendingCompletion {
2220 id: usize,
2221 _tasks: Vec<Task<()>>,
2222}
2223
2224enum ConversationEditorEvent {
2225 TabContentChanged,
2226}
2227
2228#[derive(Copy, Clone, Debug, PartialEq)]
2229struct ScrollPosition {
2230 offset_before_cursor: Vector2F,
2231 cursor: Anchor,
2232}
2233
2234struct ConversationEditor {
2235 conversation: ModelHandle<Conversation>,
2236 fs: Arc<dyn Fs>,
2237 editor: ViewHandle<Editor>,
2238 blocks: HashSet<BlockId>,
2239 scroll_position: Option<ScrollPosition>,
2240 _subscriptions: Vec<Subscription>,
2241}
2242
2243impl ConversationEditor {
2244 fn new(
2245 api_key: Rc<RefCell<Option<String>>>,
2246 language_registry: Arc<LanguageRegistry>,
2247 fs: Arc<dyn Fs>,
2248 cx: &mut ViewContext<Self>,
2249 ) -> Self {
2250 let conversation = cx.add_model(|cx| Conversation::new(api_key, language_registry, cx));
2251 Self::for_conversation(conversation, fs, cx)
2252 }
2253
2254 fn for_conversation(
2255 conversation: ModelHandle<Conversation>,
2256 fs: Arc<dyn Fs>,
2257 cx: &mut ViewContext<Self>,
2258 ) -> Self {
2259 let editor = cx.add_view(|cx| {
2260 let mut editor = Editor::for_buffer(conversation.read(cx).buffer.clone(), None, cx);
2261 editor.set_soft_wrap_mode(SoftWrap::EditorWidth, cx);
2262 editor.set_show_gutter(false, cx);
2263 editor.set_show_wrap_guides(false, cx);
2264 editor
2265 });
2266
2267 let _subscriptions = vec![
2268 cx.observe(&conversation, |_, _, cx| cx.notify()),
2269 cx.subscribe(&conversation, Self::handle_conversation_event),
2270 cx.subscribe(&editor, Self::handle_editor_event),
2271 ];
2272
2273 let mut this = Self {
2274 conversation,
2275 editor,
2276 blocks: Default::default(),
2277 scroll_position: None,
2278 fs,
2279 _subscriptions,
2280 };
2281 this.update_message_headers(cx);
2282 this
2283 }
2284
2285 fn assist(&mut self, _: &Assist, cx: &mut ViewContext<Self>) {
2286 let cursors = self.cursors(cx);
2287
2288 let user_messages = self.conversation.update(cx, |conversation, cx| {
2289 let selected_messages = conversation
2290 .messages_for_offsets(cursors, cx)
2291 .into_iter()
2292 .map(|message| message.id)
2293 .collect();
2294 conversation.assist(selected_messages, cx)
2295 });
2296 let new_selections = user_messages
2297 .iter()
2298 .map(|message| {
2299 let cursor = message
2300 .start
2301 .to_offset(self.conversation.read(cx).buffer.read(cx));
2302 cursor..cursor
2303 })
2304 .collect::<Vec<_>>();
2305 if !new_selections.is_empty() {
2306 self.editor.update(cx, |editor, cx| {
2307 editor.change_selections(
2308 Some(Autoscroll::Strategy(AutoscrollStrategy::Fit)),
2309 cx,
2310 |selections| selections.select_ranges(new_selections),
2311 );
2312 });
2313 // Avoid scrolling to the new cursor position so the assistant's output is stable.
2314 cx.defer(|this, _| this.scroll_position = None);
2315 }
2316 }
2317
2318 fn cancel_last_assist(&mut self, _: &editor::Cancel, cx: &mut ViewContext<Self>) {
2319 if !self
2320 .conversation
2321 .update(cx, |conversation, _| conversation.cancel_last_assist())
2322 {
2323 cx.propagate_action();
2324 }
2325 }
2326
2327 fn cycle_message_role(&mut self, _: &CycleMessageRole, cx: &mut ViewContext<Self>) {
2328 let cursors = self.cursors(cx);
2329 self.conversation.update(cx, |conversation, cx| {
2330 let messages = conversation
2331 .messages_for_offsets(cursors, cx)
2332 .into_iter()
2333 .map(|message| message.id)
2334 .collect();
2335 conversation.cycle_message_roles(messages, cx)
2336 });
2337 }
2338
2339 fn cursors(&self, cx: &AppContext) -> Vec<usize> {
2340 let selections = self.editor.read(cx).selections.all::<usize>(cx);
2341 selections
2342 .into_iter()
2343 .map(|selection| selection.head())
2344 .collect()
2345 }
2346
2347 fn handle_conversation_event(
2348 &mut self,
2349 _: ModelHandle<Conversation>,
2350 event: &ConversationEvent,
2351 cx: &mut ViewContext<Self>,
2352 ) {
2353 match event {
2354 ConversationEvent::MessagesEdited => {
2355 self.update_message_headers(cx);
2356 self.conversation.update(cx, |conversation, cx| {
2357 conversation.save(Some(Duration::from_millis(500)), self.fs.clone(), cx);
2358 });
2359 }
2360 ConversationEvent::SummaryChanged => {
2361 cx.emit(ConversationEditorEvent::TabContentChanged);
2362 self.conversation.update(cx, |conversation, cx| {
2363 conversation.save(None, self.fs.clone(), cx);
2364 });
2365 }
2366 ConversationEvent::StreamedCompletion => {
2367 self.editor.update(cx, |editor, cx| {
2368 if let Some(scroll_position) = self.scroll_position {
2369 let snapshot = editor.snapshot(cx);
2370 let cursor_point = scroll_position.cursor.to_display_point(&snapshot);
2371 let scroll_top =
2372 cursor_point.row() as f32 - scroll_position.offset_before_cursor.y();
2373 editor.set_scroll_position(
2374 vec2f(scroll_position.offset_before_cursor.x(), scroll_top),
2375 cx,
2376 );
2377 }
2378 });
2379 }
2380 }
2381 }
2382
2383 fn handle_editor_event(
2384 &mut self,
2385 _: ViewHandle<Editor>,
2386 event: &editor::Event,
2387 cx: &mut ViewContext<Self>,
2388 ) {
2389 match event {
2390 editor::Event::ScrollPositionChanged { autoscroll, .. } => {
2391 let cursor_scroll_position = self.cursor_scroll_position(cx);
2392 if *autoscroll {
2393 self.scroll_position = cursor_scroll_position;
2394 } else if self.scroll_position != cursor_scroll_position {
2395 self.scroll_position = None;
2396 }
2397 }
2398 editor::Event::SelectionsChanged { .. } => {
2399 self.scroll_position = self.cursor_scroll_position(cx);
2400 }
2401 _ => {}
2402 }
2403 }
2404
2405 fn cursor_scroll_position(&self, cx: &mut ViewContext<Self>) -> Option<ScrollPosition> {
2406 self.editor.update(cx, |editor, cx| {
2407 let snapshot = editor.snapshot(cx);
2408 let cursor = editor.selections.newest_anchor().head();
2409 let cursor_row = cursor.to_display_point(&snapshot.display_snapshot).row() as f32;
2410 let scroll_position = editor
2411 .scroll_manager
2412 .anchor()
2413 .scroll_position(&snapshot.display_snapshot);
2414
2415 let scroll_bottom = scroll_position.y() + editor.visible_line_count().unwrap_or(0.);
2416 if (scroll_position.y()..scroll_bottom).contains(&cursor_row) {
2417 Some(ScrollPosition {
2418 cursor,
2419 offset_before_cursor: vec2f(
2420 scroll_position.x(),
2421 cursor_row - scroll_position.y(),
2422 ),
2423 })
2424 } else {
2425 None
2426 }
2427 })
2428 }
2429
2430 fn update_message_headers(&mut self, cx: &mut ViewContext<Self>) {
2431 self.editor.update(cx, |editor, cx| {
2432 let buffer = editor.buffer().read(cx).snapshot(cx);
2433 let excerpt_id = *buffer.as_singleton().unwrap().0;
2434 let old_blocks = std::mem::take(&mut self.blocks);
2435 let new_blocks = self
2436 .conversation
2437 .read(cx)
2438 .messages(cx)
2439 .map(|message| BlockProperties {
2440 position: buffer.anchor_in_excerpt(excerpt_id, message.anchor),
2441 height: 2,
2442 style: BlockStyle::Sticky,
2443 render: Arc::new({
2444 let conversation = self.conversation.clone();
2445 // let metadata = message.metadata.clone();
2446 // let message = message.clone();
2447 move |cx| {
2448 enum Sender {}
2449 enum ErrorTooltip {}
2450
2451 let theme = theme::current(cx);
2452 let style = &theme.assistant;
2453 let message_id = message.id;
2454 let sender = MouseEventHandler::new::<Sender, _>(
2455 message_id.0,
2456 cx,
2457 |state, _| match message.role {
2458 Role::User => {
2459 let style = style.user_sender.style_for(state);
2460 Label::new("You", style.text.clone())
2461 .contained()
2462 .with_style(style.container)
2463 }
2464 Role::Assistant => {
2465 let style = style.assistant_sender.style_for(state);
2466 Label::new("Assistant", style.text.clone())
2467 .contained()
2468 .with_style(style.container)
2469 }
2470 Role::System => {
2471 let style = style.system_sender.style_for(state);
2472 Label::new("System", style.text.clone())
2473 .contained()
2474 .with_style(style.container)
2475 }
2476 },
2477 )
2478 .with_cursor_style(CursorStyle::PointingHand)
2479 .on_down(MouseButton::Left, {
2480 let conversation = conversation.clone();
2481 move |_, _, cx| {
2482 conversation.update(cx, |conversation, cx| {
2483 conversation.cycle_message_roles(
2484 HashSet::from_iter(Some(message_id)),
2485 cx,
2486 )
2487 })
2488 }
2489 });
2490
2491 Flex::row()
2492 .with_child(sender.aligned())
2493 .with_child(
2494 Label::new(
2495 message.sent_at.format("%I:%M%P").to_string(),
2496 style.sent_at.text.clone(),
2497 )
2498 .contained()
2499 .with_style(style.sent_at.container)
2500 .aligned(),
2501 )
2502 .with_children(
2503 if let MessageStatus::Error(error) = &message.status {
2504 Some(
2505 Svg::new("icons/circle_x_mark_12.svg")
2506 .with_color(style.error_icon.color)
2507 .constrained()
2508 .with_width(style.error_icon.width)
2509 .contained()
2510 .with_style(style.error_icon.container)
2511 .with_tooltip::<ErrorTooltip>(
2512 message_id.0,
2513 error.to_string(),
2514 None,
2515 theme.tooltip.clone(),
2516 cx,
2517 )
2518 .aligned(),
2519 )
2520 } else {
2521 None
2522 },
2523 )
2524 .aligned()
2525 .left()
2526 .contained()
2527 .with_style(style.message_header)
2528 .into_any()
2529 }
2530 }),
2531 disposition: BlockDisposition::Above,
2532 })
2533 .collect::<Vec<_>>();
2534
2535 editor.remove_blocks(old_blocks, None, cx);
2536 let ids = editor.insert_blocks(new_blocks, None, cx);
2537 self.blocks = HashSet::from_iter(ids);
2538 });
2539 }
2540
2541 fn quote_selection(
2542 workspace: &mut Workspace,
2543 _: &QuoteSelection,
2544 cx: &mut ViewContext<Workspace>,
2545 ) {
2546 let Some(panel) = workspace.panel::<AssistantPanel>(cx) else {
2547 return;
2548 };
2549 let Some(editor) = workspace.active_item(cx).and_then(|item| item.act_as::<Editor>(cx)) else {
2550 return;
2551 };
2552
2553 let text = editor.read_with(cx, |editor, cx| {
2554 let range = editor.selections.newest::<usize>(cx).range();
2555 let buffer = editor.buffer().read(cx).snapshot(cx);
2556 let start_language = buffer.language_at(range.start);
2557 let end_language = buffer.language_at(range.end);
2558 let language_name = if start_language == end_language {
2559 start_language.map(|language| language.name())
2560 } else {
2561 None
2562 };
2563 let language_name = language_name.as_deref().unwrap_or("").to_lowercase();
2564
2565 let selected_text = buffer.text_for_range(range).collect::<String>();
2566 if selected_text.is_empty() {
2567 None
2568 } else {
2569 Some(if language_name == "markdown" {
2570 selected_text
2571 .lines()
2572 .map(|line| format!("> {}", line))
2573 .collect::<Vec<_>>()
2574 .join("\n")
2575 } else {
2576 format!("```{language_name}\n{selected_text}\n```")
2577 })
2578 }
2579 });
2580
2581 // Activate the panel
2582 if !panel.read(cx).has_focus(cx) {
2583 workspace.toggle_panel_focus::<AssistantPanel>(cx);
2584 }
2585
2586 if let Some(text) = text {
2587 panel.update(cx, |panel, cx| {
2588 let conversation = panel
2589 .active_editor()
2590 .cloned()
2591 .unwrap_or_else(|| panel.new_conversation(cx));
2592 conversation.update(cx, |conversation, cx| {
2593 conversation
2594 .editor
2595 .update(cx, |editor, cx| editor.insert(&text, cx))
2596 });
2597 });
2598 }
2599 }
2600
2601 fn copy(&mut self, _: &editor::Copy, cx: &mut ViewContext<Self>) {
2602 let editor = self.editor.read(cx);
2603 let conversation = self.conversation.read(cx);
2604 if editor.selections.count() == 1 {
2605 let selection = editor.selections.newest::<usize>(cx);
2606 let mut copied_text = String::new();
2607 let mut spanned_messages = 0;
2608 for message in conversation.messages(cx) {
2609 if message.offset_range.start >= selection.range().end {
2610 break;
2611 } else if message.offset_range.end >= selection.range().start {
2612 let range = cmp::max(message.offset_range.start, selection.range().start)
2613 ..cmp::min(message.offset_range.end, selection.range().end);
2614 if !range.is_empty() {
2615 spanned_messages += 1;
2616 write!(&mut copied_text, "## {}\n\n", message.role).unwrap();
2617 for chunk in conversation.buffer.read(cx).text_for_range(range) {
2618 copied_text.push_str(&chunk);
2619 }
2620 copied_text.push('\n');
2621 }
2622 }
2623 }
2624
2625 if spanned_messages > 1 {
2626 cx.platform()
2627 .write_to_clipboard(ClipboardItem::new(copied_text));
2628 return;
2629 }
2630 }
2631
2632 cx.propagate_action();
2633 }
2634
2635 fn split(&mut self, _: &Split, cx: &mut ViewContext<Self>) {
2636 self.conversation.update(cx, |conversation, cx| {
2637 let selections = self.editor.read(cx).selections.disjoint_anchors();
2638 for selection in selections.into_iter() {
2639 let buffer = self.editor.read(cx).buffer().read(cx).snapshot(cx);
2640 let range = selection
2641 .map(|endpoint| endpoint.to_offset(&buffer))
2642 .range();
2643 conversation.split_message(range, cx);
2644 }
2645 });
2646 }
2647
2648 fn save(&mut self, _: &Save, cx: &mut ViewContext<Self>) {
2649 self.conversation.update(cx, |conversation, cx| {
2650 conversation.save(None, self.fs.clone(), cx)
2651 });
2652 }
2653
2654 fn cycle_model(&mut self, cx: &mut ViewContext<Self>) {
2655 self.conversation.update(cx, |conversation, cx| {
2656 let new_model = conversation.model.cycle();
2657 conversation.set_model(new_model, cx);
2658 });
2659 }
2660
2661 fn title(&self, cx: &AppContext) -> String {
2662 self.conversation
2663 .read(cx)
2664 .summary
2665 .as_ref()
2666 .map(|summary| summary.text.clone())
2667 .unwrap_or_else(|| "New Conversation".into())
2668 }
2669
2670 fn render_current_model(
2671 &self,
2672 style: &AssistantStyle,
2673 cx: &mut ViewContext<Self>,
2674 ) -> impl Element<Self> {
2675 enum Model {}
2676
2677 MouseEventHandler::new::<Model, _>(0, cx, |state, cx| {
2678 let style = style.model.style_for(state);
2679 let model_display_name = self.conversation.read(cx).model.short_name();
2680 Label::new(model_display_name, style.text.clone())
2681 .contained()
2682 .with_style(style.container)
2683 })
2684 .with_cursor_style(CursorStyle::PointingHand)
2685 .on_click(MouseButton::Left, |_, this, cx| this.cycle_model(cx))
2686 }
2687
2688 fn render_remaining_tokens(
2689 &self,
2690 style: &AssistantStyle,
2691 cx: &mut ViewContext<Self>,
2692 ) -> Option<impl Element<Self>> {
2693 let remaining_tokens = self.conversation.read(cx).remaining_tokens()?;
2694 let remaining_tokens_style = if remaining_tokens <= 0 {
2695 &style.no_remaining_tokens
2696 } else if remaining_tokens <= 500 {
2697 &style.low_remaining_tokens
2698 } else {
2699 &style.remaining_tokens
2700 };
2701 Some(
2702 Label::new(
2703 remaining_tokens.to_string(),
2704 remaining_tokens_style.text.clone(),
2705 )
2706 .contained()
2707 .with_style(remaining_tokens_style.container),
2708 )
2709 }
2710}
2711
2712impl Entity for ConversationEditor {
2713 type Event = ConversationEditorEvent;
2714}
2715
2716impl View for ConversationEditor {
2717 fn ui_name() -> &'static str {
2718 "ConversationEditor"
2719 }
2720
2721 fn render(&mut self, cx: &mut ViewContext<Self>) -> AnyElement<Self> {
2722 let theme = &theme::current(cx).assistant;
2723 Stack::new()
2724 .with_child(
2725 ChildView::new(&self.editor, cx)
2726 .contained()
2727 .with_style(theme.container),
2728 )
2729 .with_child(
2730 Flex::row()
2731 .with_child(self.render_current_model(theme, cx))
2732 .with_children(self.render_remaining_tokens(theme, cx))
2733 .aligned()
2734 .top()
2735 .right(),
2736 )
2737 .into_any()
2738 }
2739
2740 fn focus_in(&mut self, _: gpui::AnyViewHandle, cx: &mut ViewContext<Self>) {
2741 if cx.is_self_focused() {
2742 cx.focus(&self.editor);
2743 }
2744 }
2745}
2746
2747#[derive(Clone, Debug)]
2748struct MessageAnchor {
2749 id: MessageId,
2750 start: language::Anchor,
2751}
2752
2753#[derive(Clone, Debug)]
2754pub struct Message {
2755 offset_range: Range<usize>,
2756 index_range: Range<usize>,
2757 id: MessageId,
2758 anchor: language::Anchor,
2759 role: Role,
2760 sent_at: DateTime<Local>,
2761 status: MessageStatus,
2762}
2763
2764impl Message {
2765 fn to_open_ai_message(&self, buffer: &Buffer) -> RequestMessage {
2766 let mut content = format!("[Message {}]\n", self.id.0).to_string();
2767 content.extend(buffer.text_for_range(self.offset_range.clone()));
2768 RequestMessage {
2769 role: self.role,
2770 content: content.trim_end().into(),
2771 }
2772 }
2773}
2774
2775enum InlineAssistantEvent {
2776 Confirmed { prompt: String },
2777 Canceled,
2778 Dismissed,
2779}
2780
2781#[derive(Copy, Clone)]
2782enum InlineAssistKind {
2783 Transform,
2784 Generate,
2785}
2786
2787struct InlineAssistant {
2788 id: usize,
2789 prompt_editor: ViewHandle<Editor>,
2790 confirmed: bool,
2791 has_focus: bool,
2792}
2793
2794impl Entity for InlineAssistant {
2795 type Event = InlineAssistantEvent;
2796}
2797
2798impl View for InlineAssistant {
2799 fn ui_name() -> &'static str {
2800 "InlineAssistant"
2801 }
2802
2803 fn render(&mut self, cx: &mut ViewContext<Self>) -> AnyElement<Self> {
2804 ChildView::new(&self.prompt_editor, cx)
2805 .aligned()
2806 .left()
2807 .into_any()
2808 }
2809
2810 fn focus_in(&mut self, _: gpui::AnyViewHandle, cx: &mut ViewContext<Self>) {
2811 cx.focus(&self.prompt_editor);
2812 self.has_focus = true;
2813 }
2814
2815 fn focus_out(&mut self, _: gpui::AnyViewHandle, _: &mut ViewContext<Self>) {
2816 self.has_focus = false;
2817 }
2818}
2819
2820impl InlineAssistant {
2821 fn cancel(&mut self, _: &editor::Cancel, cx: &mut ViewContext<Self>) {
2822 cx.emit(InlineAssistantEvent::Canceled);
2823 }
2824
2825 fn confirm(&mut self, _: &menu::Confirm, cx: &mut ViewContext<Self>) {
2826 if self.confirmed {
2827 cx.emit(InlineAssistantEvent::Dismissed);
2828 } else {
2829 let prompt = self.prompt_editor.read(cx).text(cx);
2830 self.prompt_editor
2831 .update(cx, |editor, _| editor.set_read_only(true));
2832 cx.emit(InlineAssistantEvent::Confirmed { prompt });
2833 self.confirmed = true;
2834 }
2835 }
2836}
2837
2838struct PendingInlineAssist {
2839 kind: InlineAssistKind,
2840 editor: WeakViewHandle<Editor>,
2841 range: Range<Anchor>,
2842 inline_assistant_block_id: Option<BlockId>,
2843 code_generation: Task<Option<()>>,
2844 transaction_id: Option<TransactionId>,
2845 _subscriptions: Vec<Subscription>,
2846}
2847
2848#[cfg(test)]
2849mod tests {
2850 use super::*;
2851 use crate::MessageId;
2852 use gpui::AppContext;
2853
2854 #[gpui::test]
2855 fn test_inserting_and_removing_messages(cx: &mut AppContext) {
2856 cx.set_global(SettingsStore::test(cx));
2857 init(cx);
2858 let registry = Arc::new(LanguageRegistry::test());
2859 let conversation = cx.add_model(|cx| Conversation::new(Default::default(), registry, cx));
2860 let buffer = conversation.read(cx).buffer.clone();
2861
2862 let message_1 = conversation.read(cx).message_anchors[0].clone();
2863 assert_eq!(
2864 messages(&conversation, cx),
2865 vec![(message_1.id, Role::User, 0..0)]
2866 );
2867
2868 let message_2 = conversation.update(cx, |conversation, cx| {
2869 conversation
2870 .insert_message_after(message_1.id, Role::Assistant, MessageStatus::Done, cx)
2871 .unwrap()
2872 });
2873 assert_eq!(
2874 messages(&conversation, cx),
2875 vec![
2876 (message_1.id, Role::User, 0..1),
2877 (message_2.id, Role::Assistant, 1..1)
2878 ]
2879 );
2880
2881 buffer.update(cx, |buffer, cx| {
2882 buffer.edit([(0..0, "1"), (1..1, "2")], None, cx)
2883 });
2884 assert_eq!(
2885 messages(&conversation, cx),
2886 vec![
2887 (message_1.id, Role::User, 0..2),
2888 (message_2.id, Role::Assistant, 2..3)
2889 ]
2890 );
2891
2892 let message_3 = conversation.update(cx, |conversation, cx| {
2893 conversation
2894 .insert_message_after(message_2.id, Role::User, MessageStatus::Done, cx)
2895 .unwrap()
2896 });
2897 assert_eq!(
2898 messages(&conversation, cx),
2899 vec![
2900 (message_1.id, Role::User, 0..2),
2901 (message_2.id, Role::Assistant, 2..4),
2902 (message_3.id, Role::User, 4..4)
2903 ]
2904 );
2905
2906 let message_4 = conversation.update(cx, |conversation, cx| {
2907 conversation
2908 .insert_message_after(message_2.id, Role::User, MessageStatus::Done, cx)
2909 .unwrap()
2910 });
2911 assert_eq!(
2912 messages(&conversation, cx),
2913 vec![
2914 (message_1.id, Role::User, 0..2),
2915 (message_2.id, Role::Assistant, 2..4),
2916 (message_4.id, Role::User, 4..5),
2917 (message_3.id, Role::User, 5..5),
2918 ]
2919 );
2920
2921 buffer.update(cx, |buffer, cx| {
2922 buffer.edit([(4..4, "C"), (5..5, "D")], None, cx)
2923 });
2924 assert_eq!(
2925 messages(&conversation, cx),
2926 vec![
2927 (message_1.id, Role::User, 0..2),
2928 (message_2.id, Role::Assistant, 2..4),
2929 (message_4.id, Role::User, 4..6),
2930 (message_3.id, Role::User, 6..7),
2931 ]
2932 );
2933
2934 // Deleting across message boundaries merges the messages.
2935 buffer.update(cx, |buffer, cx| buffer.edit([(1..4, "")], None, cx));
2936 assert_eq!(
2937 messages(&conversation, cx),
2938 vec![
2939 (message_1.id, Role::User, 0..3),
2940 (message_3.id, Role::User, 3..4),
2941 ]
2942 );
2943
2944 // Undoing the deletion should also undo the merge.
2945 buffer.update(cx, |buffer, cx| buffer.undo(cx));
2946 assert_eq!(
2947 messages(&conversation, cx),
2948 vec![
2949 (message_1.id, Role::User, 0..2),
2950 (message_2.id, Role::Assistant, 2..4),
2951 (message_4.id, Role::User, 4..6),
2952 (message_3.id, Role::User, 6..7),
2953 ]
2954 );
2955
2956 // Redoing the deletion should also redo the merge.
2957 buffer.update(cx, |buffer, cx| buffer.redo(cx));
2958 assert_eq!(
2959 messages(&conversation, cx),
2960 vec![
2961 (message_1.id, Role::User, 0..3),
2962 (message_3.id, Role::User, 3..4),
2963 ]
2964 );
2965
2966 // Ensure we can still insert after a merged message.
2967 let message_5 = conversation.update(cx, |conversation, cx| {
2968 conversation
2969 .insert_message_after(message_1.id, Role::System, MessageStatus::Done, cx)
2970 .unwrap()
2971 });
2972 assert_eq!(
2973 messages(&conversation, cx),
2974 vec![
2975 (message_1.id, Role::User, 0..3),
2976 (message_5.id, Role::System, 3..4),
2977 (message_3.id, Role::User, 4..5)
2978 ]
2979 );
2980 }
2981
2982 #[gpui::test]
2983 fn test_message_splitting(cx: &mut AppContext) {
2984 cx.set_global(SettingsStore::test(cx));
2985 init(cx);
2986 let registry = Arc::new(LanguageRegistry::test());
2987 let conversation = cx.add_model(|cx| Conversation::new(Default::default(), registry, cx));
2988 let buffer = conversation.read(cx).buffer.clone();
2989
2990 let message_1 = conversation.read(cx).message_anchors[0].clone();
2991 assert_eq!(
2992 messages(&conversation, cx),
2993 vec![(message_1.id, Role::User, 0..0)]
2994 );
2995
2996 buffer.update(cx, |buffer, cx| {
2997 buffer.edit([(0..0, "aaa\nbbb\nccc\nddd\n")], None, cx)
2998 });
2999
3000 let (_, message_2) =
3001 conversation.update(cx, |conversation, cx| conversation.split_message(3..3, cx));
3002 let message_2 = message_2.unwrap();
3003
3004 // We recycle newlines in the middle of a split message
3005 assert_eq!(buffer.read(cx).text(), "aaa\nbbb\nccc\nddd\n");
3006 assert_eq!(
3007 messages(&conversation, cx),
3008 vec![
3009 (message_1.id, Role::User, 0..4),
3010 (message_2.id, Role::User, 4..16),
3011 ]
3012 );
3013
3014 let (_, message_3) =
3015 conversation.update(cx, |conversation, cx| conversation.split_message(3..3, cx));
3016 let message_3 = message_3.unwrap();
3017
3018 // We don't recycle newlines at the end of a split message
3019 assert_eq!(buffer.read(cx).text(), "aaa\n\nbbb\nccc\nddd\n");
3020 assert_eq!(
3021 messages(&conversation, cx),
3022 vec![
3023 (message_1.id, Role::User, 0..4),
3024 (message_3.id, Role::User, 4..5),
3025 (message_2.id, Role::User, 5..17),
3026 ]
3027 );
3028
3029 let (_, message_4) =
3030 conversation.update(cx, |conversation, cx| conversation.split_message(9..9, cx));
3031 let message_4 = message_4.unwrap();
3032 assert_eq!(buffer.read(cx).text(), "aaa\n\nbbb\nccc\nddd\n");
3033 assert_eq!(
3034 messages(&conversation, cx),
3035 vec![
3036 (message_1.id, Role::User, 0..4),
3037 (message_3.id, Role::User, 4..5),
3038 (message_2.id, Role::User, 5..9),
3039 (message_4.id, Role::User, 9..17),
3040 ]
3041 );
3042
3043 let (_, message_5) =
3044 conversation.update(cx, |conversation, cx| conversation.split_message(9..9, cx));
3045 let message_5 = message_5.unwrap();
3046 assert_eq!(buffer.read(cx).text(), "aaa\n\nbbb\n\nccc\nddd\n");
3047 assert_eq!(
3048 messages(&conversation, cx),
3049 vec![
3050 (message_1.id, Role::User, 0..4),
3051 (message_3.id, Role::User, 4..5),
3052 (message_2.id, Role::User, 5..9),
3053 (message_4.id, Role::User, 9..10),
3054 (message_5.id, Role::User, 10..18),
3055 ]
3056 );
3057
3058 let (message_6, message_7) = conversation.update(cx, |conversation, cx| {
3059 conversation.split_message(14..16, cx)
3060 });
3061 let message_6 = message_6.unwrap();
3062 let message_7 = message_7.unwrap();
3063 assert_eq!(buffer.read(cx).text(), "aaa\n\nbbb\n\nccc\ndd\nd\n");
3064 assert_eq!(
3065 messages(&conversation, cx),
3066 vec![
3067 (message_1.id, Role::User, 0..4),
3068 (message_3.id, Role::User, 4..5),
3069 (message_2.id, Role::User, 5..9),
3070 (message_4.id, Role::User, 9..10),
3071 (message_5.id, Role::User, 10..14),
3072 (message_6.id, Role::User, 14..17),
3073 (message_7.id, Role::User, 17..19),
3074 ]
3075 );
3076 }
3077
3078 #[gpui::test]
3079 fn test_messages_for_offsets(cx: &mut AppContext) {
3080 cx.set_global(SettingsStore::test(cx));
3081 init(cx);
3082 let registry = Arc::new(LanguageRegistry::test());
3083 let conversation = cx.add_model(|cx| Conversation::new(Default::default(), registry, cx));
3084 let buffer = conversation.read(cx).buffer.clone();
3085
3086 let message_1 = conversation.read(cx).message_anchors[0].clone();
3087 assert_eq!(
3088 messages(&conversation, cx),
3089 vec![(message_1.id, Role::User, 0..0)]
3090 );
3091
3092 buffer.update(cx, |buffer, cx| buffer.edit([(0..0, "aaa")], None, cx));
3093 let message_2 = conversation
3094 .update(cx, |conversation, cx| {
3095 conversation.insert_message_after(message_1.id, Role::User, MessageStatus::Done, cx)
3096 })
3097 .unwrap();
3098 buffer.update(cx, |buffer, cx| buffer.edit([(4..4, "bbb")], None, cx));
3099
3100 let message_3 = conversation
3101 .update(cx, |conversation, cx| {
3102 conversation.insert_message_after(message_2.id, Role::User, MessageStatus::Done, cx)
3103 })
3104 .unwrap();
3105 buffer.update(cx, |buffer, cx| buffer.edit([(8..8, "ccc")], None, cx));
3106
3107 assert_eq!(buffer.read(cx).text(), "aaa\nbbb\nccc");
3108 assert_eq!(
3109 messages(&conversation, cx),
3110 vec![
3111 (message_1.id, Role::User, 0..4),
3112 (message_2.id, Role::User, 4..8),
3113 (message_3.id, Role::User, 8..11)
3114 ]
3115 );
3116
3117 assert_eq!(
3118 message_ids_for_offsets(&conversation, &[0, 4, 9], cx),
3119 [message_1.id, message_2.id, message_3.id]
3120 );
3121 assert_eq!(
3122 message_ids_for_offsets(&conversation, &[0, 1, 11], cx),
3123 [message_1.id, message_3.id]
3124 );
3125
3126 let message_4 = conversation
3127 .update(cx, |conversation, cx| {
3128 conversation.insert_message_after(message_3.id, Role::User, MessageStatus::Done, cx)
3129 })
3130 .unwrap();
3131 assert_eq!(buffer.read(cx).text(), "aaa\nbbb\nccc\n");
3132 assert_eq!(
3133 messages(&conversation, cx),
3134 vec![
3135 (message_1.id, Role::User, 0..4),
3136 (message_2.id, Role::User, 4..8),
3137 (message_3.id, Role::User, 8..12),
3138 (message_4.id, Role::User, 12..12)
3139 ]
3140 );
3141 assert_eq!(
3142 message_ids_for_offsets(&conversation, &[0, 4, 8, 12], cx),
3143 [message_1.id, message_2.id, message_3.id, message_4.id]
3144 );
3145
3146 fn message_ids_for_offsets(
3147 conversation: &ModelHandle<Conversation>,
3148 offsets: &[usize],
3149 cx: &AppContext,
3150 ) -> Vec<MessageId> {
3151 conversation
3152 .read(cx)
3153 .messages_for_offsets(offsets.iter().copied(), cx)
3154 .into_iter()
3155 .map(|message| message.id)
3156 .collect()
3157 }
3158 }
3159
3160 #[gpui::test]
3161 fn test_serialization(cx: &mut AppContext) {
3162 cx.set_global(SettingsStore::test(cx));
3163 init(cx);
3164 let registry = Arc::new(LanguageRegistry::test());
3165 let conversation =
3166 cx.add_model(|cx| Conversation::new(Default::default(), registry.clone(), cx));
3167 let buffer = conversation.read(cx).buffer.clone();
3168 let message_0 = conversation.read(cx).message_anchors[0].id;
3169 let message_1 = conversation.update(cx, |conversation, cx| {
3170 conversation
3171 .insert_message_after(message_0, Role::Assistant, MessageStatus::Done, cx)
3172 .unwrap()
3173 });
3174 let message_2 = conversation.update(cx, |conversation, cx| {
3175 conversation
3176 .insert_message_after(message_1.id, Role::System, MessageStatus::Done, cx)
3177 .unwrap()
3178 });
3179 buffer.update(cx, |buffer, cx| {
3180 buffer.edit([(0..0, "a"), (1..1, "b\nc")], None, cx);
3181 buffer.finalize_last_transaction();
3182 });
3183 let _message_3 = conversation.update(cx, |conversation, cx| {
3184 conversation
3185 .insert_message_after(message_2.id, Role::System, MessageStatus::Done, cx)
3186 .unwrap()
3187 });
3188 buffer.update(cx, |buffer, cx| buffer.undo(cx));
3189 assert_eq!(buffer.read(cx).text(), "a\nb\nc\n");
3190 assert_eq!(
3191 messages(&conversation, cx),
3192 [
3193 (message_0, Role::User, 0..2),
3194 (message_1.id, Role::Assistant, 2..6),
3195 (message_2.id, Role::System, 6..6),
3196 ]
3197 );
3198
3199 let deserialized_conversation = cx.add_model(|cx| {
3200 Conversation::deserialize(
3201 conversation.read(cx).serialize(cx),
3202 Default::default(),
3203 Default::default(),
3204 registry.clone(),
3205 cx,
3206 )
3207 });
3208 let deserialized_buffer = deserialized_conversation.read(cx).buffer.clone();
3209 assert_eq!(deserialized_buffer.read(cx).text(), "a\nb\nc\n");
3210 assert_eq!(
3211 messages(&deserialized_conversation, cx),
3212 [
3213 (message_0, Role::User, 0..2),
3214 (message_1.id, Role::Assistant, 2..6),
3215 (message_2.id, Role::System, 6..6),
3216 ]
3217 );
3218 }
3219
3220 fn messages(
3221 conversation: &ModelHandle<Conversation>,
3222 cx: &AppContext,
3223 ) -> Vec<(MessageId, Role, Range<usize>)> {
3224 conversation
3225 .read(cx)
3226 .messages(cx)
3227 .map(|message| (message.id, message.role, message.offset_range))
3228 .collect()
3229 }
3230}