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