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