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