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