1use crate::{
2 assistant_settings::{AssistantDockPosition, AssistantSettings},
3 MessageId, MessageMetadata, MessageStatus, OpenAIRequest, OpenAIResponseStreamEvent,
4 RequestMessage, Role, SavedConversation, SavedConversationMetadata, SavedMessage,
5};
6use anyhow::{anyhow, Result};
7use chrono::{DateTime, Local};
8use collections::{HashMap, HashSet};
9use editor::{
10 display_map::{BlockDisposition, BlockId, BlockProperties, BlockStyle, ToDisplayPoint},
11 scroll::autoscroll::{Autoscroll, AutoscrollStrategy},
12 Anchor, Editor, ToOffset,
13};
14use fs::Fs;
15use futures::{io::BufReader, AsyncBufReadExt, AsyncReadExt, Stream, StreamExt};
16use gpui::{
17 actions,
18 elements::*,
19 executor::Background,
20 geometry::vector::{vec2f, Vector2F},
21 platform::{CursorStyle, MouseButton},
22 Action, AppContext, AsyncAppContext, ClipboardItem, Entity, ModelContext, ModelHandle,
23 Subscription, Task, View, ViewContext, ViewHandle, WeakViewHandle, WindowContext,
24};
25use isahc::{http::StatusCode, Request, RequestExt};
26use language::{language_settings::SoftWrap, Buffer, LanguageRegistry, ToOffset as _};
27use serde::Deserialize;
28use settings::SettingsStore;
29use std::{
30 borrow::Cow,
31 cell::RefCell,
32 cmp, env,
33 fmt::Write,
34 io, iter,
35 ops::Range,
36 path::{Path, PathBuf},
37 rc::Rc,
38 sync::Arc,
39 time::Duration,
40};
41use theme::{ui::IconStyle, AssistantStyle};
42use util::{
43 channel::ReleaseChannel, paths::CONVERSATIONS_DIR, post_inc, truncate_and_trailoff, ResultExt,
44 TryFutureExt,
45};
46use workspace::{
47 dock::{DockPosition, Panel},
48 item::Item,
49 Save, ToggleZoom, Workspace,
50};
51
52const OPENAI_API_URL: &'static str = "https://api.openai.com/v1";
53
54actions!(
55 assistant,
56 [
57 NewContext,
58 Assist,
59 Split,
60 CycleMessageRole,
61 QuoteSelection,
62 ToggleFocus,
63 ResetKey,
64 ]
65);
66
67pub fn init(cx: &mut AppContext) {
68 if *util::channel::RELEASE_CHANNEL == ReleaseChannel::Stable {
69 cx.update_default_global::<collections::CommandPaletteFilter, _, _>(move |filter, _cx| {
70 filter.filtered_namespaces.insert("assistant");
71 });
72 }
73
74 settings::register::<AssistantSettings>(cx);
75 cx.add_action(
76 |workspace: &mut Workspace, _: &NewContext, cx: &mut ViewContext<Workspace>| {
77 if let Some(this) = workspace.panel::<AssistantPanel>(cx) {
78 this.update(cx, |this, cx| {
79 this.new_conversation(cx);
80 })
81 }
82
83 workspace.focus_panel::<AssistantPanel>(cx);
84 },
85 );
86 cx.add_action(ConversationEditor::assist);
87 cx.capture_action(ConversationEditor::cancel_last_assist);
88 cx.capture_action(ConversationEditor::save);
89 cx.add_action(ConversationEditor::quote_selection);
90 cx.capture_action(ConversationEditor::copy);
91 cx.capture_action(ConversationEditor::split);
92 cx.capture_action(ConversationEditor::cycle_message_role);
93 cx.add_action(AssistantPanel::save_api_key);
94 cx.add_action(AssistantPanel::reset_api_key);
95 cx.add_action(AssistantPanel::toggle_zoom);
96 cx.add_action(
97 |workspace: &mut Workspace, _: &ToggleFocus, cx: &mut ViewContext<Workspace>| {
98 workspace.toggle_panel_focus::<AssistantPanel>(cx);
99 },
100 );
101}
102
103#[derive(Debug)]
104pub enum AssistantPanelEvent {
105 ZoomIn,
106 ZoomOut,
107 Focus,
108 Close,
109 DockPositionChanged,
110}
111
112pub struct AssistantPanel {
113 width: Option<f32>,
114 height: Option<f32>,
115 active_editor_index: Option<usize>,
116 editors: Vec<ViewHandle<ConversationEditor>>,
117 saved_conversations: Vec<SavedConversationMetadata>,
118 saved_conversations_list_state: UniformListState,
119 zoomed: bool,
120 has_focus: bool,
121 api_key: Rc<RefCell<Option<String>>>,
122 api_key_editor: Option<ViewHandle<Editor>>,
123 has_read_credentials: bool,
124 languages: Arc<LanguageRegistry>,
125 fs: Arc<dyn Fs>,
126 subscriptions: Vec<Subscription>,
127 _watch_saved_conversations: Task<Result<()>>,
128}
129
130impl AssistantPanel {
131 pub fn load(
132 workspace: WeakViewHandle<Workspace>,
133 cx: AsyncAppContext,
134 ) -> Task<Result<ViewHandle<Self>>> {
135 cx.spawn(|mut cx| async move {
136 let fs = workspace.read_with(&cx, |workspace, _| workspace.app_state().fs.clone())?;
137 let saved_conversations = SavedConversationMetadata::list(fs.clone())
138 .await
139 .log_err()
140 .unwrap_or_default();
141
142 // TODO: deserialize state.
143 workspace.update(&mut cx, |workspace, cx| {
144 cx.add_view::<Self, _>(|cx| {
145 const CONVERSATION_WATCH_DURATION: Duration = Duration::from_millis(100);
146 let _watch_saved_conversations = cx.spawn(move |this, mut cx| async move {
147 let mut events = fs
148 .watch(&CONVERSATIONS_DIR, CONVERSATION_WATCH_DURATION)
149 .await;
150 while events.next().await.is_some() {
151 let saved_conversations = SavedConversationMetadata::list(fs.clone())
152 .await
153 .log_err()
154 .unwrap_or_default();
155 this.update(&mut cx, |this, _| {
156 this.saved_conversations = saved_conversations
157 })
158 .ok();
159 }
160
161 anyhow::Ok(())
162 });
163
164 let mut this = Self {
165 active_editor_index: Default::default(),
166 editors: Default::default(),
167 saved_conversations,
168 saved_conversations_list_state: Default::default(),
169 zoomed: false,
170 has_focus: false,
171 api_key: Rc::new(RefCell::new(None)),
172 api_key_editor: None,
173 has_read_credentials: false,
174 languages: workspace.app_state().languages.clone(),
175 fs: workspace.app_state().fs.clone(),
176 width: None,
177 height: None,
178 subscriptions: Default::default(),
179 _watch_saved_conversations,
180 };
181
182 let mut old_dock_position = this.position(cx);
183 this.subscriptions =
184 vec![cx.observe_global::<SettingsStore, _>(move |this, cx| {
185 let new_dock_position = this.position(cx);
186 if new_dock_position != old_dock_position {
187 old_dock_position = new_dock_position;
188 cx.emit(AssistantPanelEvent::DockPositionChanged);
189 }
190 })];
191
192 this
193 })
194 })
195 })
196 }
197
198 fn new_conversation(&mut self, cx: &mut ViewContext<Self>) -> ViewHandle<ConversationEditor> {
199 let editor = cx.add_view(|cx| {
200 ConversationEditor::new(
201 self.api_key.clone(),
202 self.languages.clone(),
203 self.fs.clone(),
204 cx,
205 )
206 });
207 self.add_conversation(editor.clone(), cx);
208 editor
209 }
210
211 fn add_conversation(
212 &mut self,
213 editor: ViewHandle<ConversationEditor>,
214 cx: &mut ViewContext<Self>,
215 ) {
216 self.subscriptions
217 .push(cx.subscribe(&editor, Self::handle_conversation_editor_event));
218
219 let conversation = editor.read(cx).conversation.clone();
220 self.subscriptions
221 .push(cx.observe(&conversation, |_, _, cx| cx.notify()));
222
223 self.active_editor_index = Some(self.editors.len());
224 self.editors.push(editor.clone());
225 if self.has_focus(cx) {
226 cx.focus(&editor);
227 }
228 cx.notify();
229 }
230
231 fn handle_conversation_editor_event(
232 &mut self,
233 _: ViewHandle<ConversationEditor>,
234 event: &ConversationEditorEvent,
235 cx: &mut ViewContext<Self>,
236 ) {
237 match event {
238 ConversationEditorEvent::TabContentChanged => cx.notify(),
239 }
240 }
241
242 fn save_api_key(&mut self, _: &menu::Confirm, cx: &mut ViewContext<Self>) {
243 if let Some(api_key) = self
244 .api_key_editor
245 .as_ref()
246 .map(|editor| editor.read(cx).text(cx))
247 {
248 if !api_key.is_empty() {
249 cx.platform()
250 .write_credentials(OPENAI_API_URL, "Bearer", api_key.as_bytes())
251 .log_err();
252 *self.api_key.borrow_mut() = Some(api_key);
253 self.api_key_editor.take();
254 cx.focus_self();
255 cx.notify();
256 }
257 } else {
258 cx.propagate_action();
259 }
260 }
261
262 fn reset_api_key(&mut self, _: &ResetKey, cx: &mut ViewContext<Self>) {
263 cx.platform().delete_credentials(OPENAI_API_URL).log_err();
264 self.api_key.take();
265 self.api_key_editor = Some(build_api_key_editor(cx));
266 cx.focus_self();
267 cx.notify();
268 }
269
270 fn toggle_zoom(&mut self, _: &workspace::ToggleZoom, cx: &mut ViewContext<Self>) {
271 if self.zoomed {
272 cx.emit(AssistantPanelEvent::ZoomOut)
273 } else {
274 cx.emit(AssistantPanelEvent::ZoomIn)
275 }
276 }
277
278 fn active_editor(&self) -> Option<&ViewHandle<ConversationEditor>> {
279 self.editors.get(self.active_editor_index?)
280 }
281
282 fn render_hamburger_button(style: &IconStyle) -> impl Element<Self> {
283 enum ListConversations {}
284 Svg::for_style(style.icon.clone())
285 .contained()
286 .with_style(style.container)
287 .mouse::<ListConversations>(0)
288 .with_cursor_style(CursorStyle::PointingHand)
289 .on_click(MouseButton::Left, |_, this: &mut Self, cx| {
290 this.active_editor_index = None;
291 cx.notify();
292 })
293 }
294
295 fn render_current_model(
296 &self,
297 style: &AssistantStyle,
298 cx: &mut ViewContext<Self>,
299 ) -> Option<impl Element<Self>> {
300 enum Model {}
301
302 let model = self
303 .active_editor()?
304 .read(cx)
305 .conversation
306 .read(cx)
307 .model
308 .clone();
309
310 Some(
311 MouseEventHandler::<Model, _>::new(0, cx, |state, _| {
312 let style = style.model.style_for(state);
313 Label::new(model, style.text.clone())
314 .contained()
315 .with_style(style.container)
316 })
317 .with_cursor_style(CursorStyle::PointingHand)
318 .on_click(MouseButton::Left, |_, this, cx| {
319 if let Some(editor) = this.active_editor() {
320 editor.update(cx, |editor, cx| {
321 editor.cycle_model(cx);
322 });
323 }
324 }),
325 )
326 }
327
328 fn render_remaining_tokens(
329 &self,
330 style: &AssistantStyle,
331 cx: &mut ViewContext<Self>,
332 ) -> Option<impl Element<Self>> {
333 self.active_editor().and_then(|editor| {
334 editor
335 .read(cx)
336 .conversation
337 .read(cx)
338 .remaining_tokens()
339 .map(|remaining_tokens| {
340 let remaining_tokens_style = if remaining_tokens <= 0 {
341 &style.no_remaining_tokens
342 } else {
343 &style.remaining_tokens
344 };
345 Label::new(
346 remaining_tokens.to_string(),
347 remaining_tokens_style.text.clone(),
348 )
349 .contained()
350 .with_style(remaining_tokens_style.container)
351 })
352 })
353 }
354
355 fn render_plus_button(style: &IconStyle) -> impl Element<Self> {
356 enum AddConversation {}
357 Svg::for_style(style.icon.clone())
358 .contained()
359 .with_style(style.container)
360 .mouse::<AddConversation>(0)
361 .with_cursor_style(CursorStyle::PointingHand)
362 .on_click(MouseButton::Left, |_, this: &mut Self, cx| {
363 this.new_conversation(cx);
364 })
365 }
366
367 fn render_zoom_button(
368 &self,
369 style: &AssistantStyle,
370 cx: &mut ViewContext<Self>,
371 ) -> impl Element<Self> {
372 enum ToggleZoomButton {}
373
374 let style = if self.zoomed {
375 &style.zoom_out_button
376 } else {
377 &style.zoom_in_button
378 };
379
380 MouseEventHandler::<ToggleZoomButton, _>::new(0, cx, |_, _| {
381 Svg::for_style(style.icon.clone())
382 .contained()
383 .with_style(style.container)
384 })
385 .with_cursor_style(CursorStyle::PointingHand)
386 .on_click(MouseButton::Left, |_, this, cx| {
387 this.toggle_zoom(&ToggleZoom, cx);
388 })
389 }
390
391 fn render_saved_conversation(
392 &mut self,
393 index: usize,
394 cx: &mut ViewContext<Self>,
395 ) -> impl Element<Self> {
396 let conversation = &self.saved_conversations[index];
397 let path = conversation.path.clone();
398 MouseEventHandler::<SavedConversationMetadata, _>::new(index, cx, move |state, cx| {
399 let style = &theme::current(cx).assistant.saved_conversation;
400 Flex::row()
401 .with_child(
402 Label::new(
403 conversation.mtime.format("%F %I:%M%p").to_string(),
404 style.saved_at.text.clone(),
405 )
406 .aligned()
407 .contained()
408 .with_style(style.saved_at.container),
409 )
410 .with_child(
411 Label::new(conversation.title.clone(), style.title.text.clone())
412 .aligned()
413 .contained()
414 .with_style(style.title.container),
415 )
416 .contained()
417 .with_style(*style.container.style_for(state))
418 })
419 .with_cursor_style(CursorStyle::PointingHand)
420 .on_click(MouseButton::Left, move |_, this, cx| {
421 this.open_conversation(path.clone(), cx)
422 .detach_and_log_err(cx)
423 })
424 }
425
426 fn open_conversation(&mut self, path: PathBuf, cx: &mut ViewContext<Self>) -> Task<Result<()>> {
427 if let Some(ix) = self.editor_index_for_path(&path, cx) {
428 self.active_editor_index = Some(ix);
429 cx.notify();
430 return Task::ready(Ok(()));
431 }
432
433 let fs = self.fs.clone();
434 let conversation = Conversation::load(
435 path.clone(),
436 self.api_key.clone(),
437 self.languages.clone(),
438 self.fs.clone(),
439 cx,
440 );
441 cx.spawn(|this, mut cx| async move {
442 let conversation = conversation.await?;
443 this.update(&mut cx, |this, cx| {
444 // If, by the time we've loaded the conversation, the user has already opened
445 // the same conversation, we don't want to open it again.
446 if let Some(ix) = this.editor_index_for_path(&path, cx) {
447 this.active_editor_index = Some(ix);
448 } else {
449 let editor = cx
450 .add_view(|cx| ConversationEditor::from_conversation(conversation, fs, cx));
451 this.add_conversation(editor, cx);
452 }
453 })?;
454 Ok(())
455 })
456 }
457
458 fn editor_index_for_path(&self, path: &Path, cx: &AppContext) -> Option<usize> {
459 self.editors
460 .iter()
461 .position(|editor| editor.read(cx).conversation.read(cx).path.as_deref() == Some(path))
462 }
463}
464
465fn build_api_key_editor(cx: &mut ViewContext<AssistantPanel>) -> ViewHandle<Editor> {
466 cx.add_view(|cx| {
467 let mut editor = Editor::single_line(
468 Some(Arc::new(|theme| theme.assistant.api_key_editor.clone())),
469 cx,
470 );
471 editor.set_placeholder_text("sk-000000000000000000000000000000000000000000000000", cx);
472 editor
473 })
474}
475
476impl Entity for AssistantPanel {
477 type Event = AssistantPanelEvent;
478}
479
480impl View for AssistantPanel {
481 fn ui_name() -> &'static str {
482 "AssistantPanel"
483 }
484
485 fn render(&mut self, cx: &mut ViewContext<Self>) -> AnyElement<Self> {
486 let theme = &theme::current(cx);
487 let style = &theme.assistant;
488 if let Some(api_key_editor) = self.api_key_editor.as_ref() {
489 Flex::column()
490 .with_child(
491 Text::new(
492 "Paste your OpenAI API key and press Enter to use the assistant",
493 style.api_key_prompt.text.clone(),
494 )
495 .aligned(),
496 )
497 .with_child(
498 ChildView::new(api_key_editor, cx)
499 .contained()
500 .with_style(style.api_key_editor.container)
501 .aligned(),
502 )
503 .contained()
504 .with_style(style.api_key_prompt.container)
505 .aligned()
506 .into_any()
507 } else {
508 let title = self.active_editor().map(|editor| {
509 Label::new(editor.read(cx).title(cx), style.title.text.clone())
510 .contained()
511 .with_style(style.title.container)
512 .aligned()
513 .left()
514 .flex(1., false)
515 });
516
517 Flex::column()
518 .with_child(
519 Flex::row()
520 .with_child(
521 Self::render_hamburger_button(&style.hamburger_button).aligned(),
522 )
523 .with_children(title)
524 .with_children(
525 self.render_current_model(&style, cx)
526 .map(|current_model| current_model.aligned().flex_float()),
527 )
528 .with_children(
529 self.render_remaining_tokens(&style, cx)
530 .map(|remaining_tokens| remaining_tokens.aligned().flex_float()),
531 )
532 .with_child(
533 Self::render_plus_button(&style.plus_button)
534 .aligned()
535 .flex_float(),
536 )
537 .with_child(self.render_zoom_button(&style, cx).aligned().flex_float())
538 .contained()
539 .with_style(theme.workspace.tab_bar.container)
540 .expanded()
541 .constrained()
542 .with_height(theme.workspace.tab_bar.height),
543 )
544 .with_child(if let Some(editor) = self.active_editor() {
545 ChildView::new(editor, cx).flex(1., true).into_any()
546 } else {
547 UniformList::new(
548 self.saved_conversations_list_state.clone(),
549 self.saved_conversations.len(),
550 cx,
551 |this, range, items, cx| {
552 for ix in range {
553 items.push(this.render_saved_conversation(ix, cx).into_any());
554 }
555 },
556 )
557 .flex(1., true)
558 .into_any()
559 })
560 .into_any()
561 }
562 }
563
564 fn focus_in(&mut self, _: gpui::AnyViewHandle, cx: &mut ViewContext<Self>) {
565 self.has_focus = true;
566 if cx.is_self_focused() {
567 if let Some(editor) = self.active_editor() {
568 cx.focus(editor);
569 } else if let Some(api_key_editor) = self.api_key_editor.as_ref() {
570 cx.focus(api_key_editor);
571 }
572 }
573 }
574
575 fn focus_out(&mut self, _: gpui::AnyViewHandle, _: &mut ViewContext<Self>) {
576 self.has_focus = false;
577 }
578}
579
580impl Panel for AssistantPanel {
581 fn position(&self, cx: &WindowContext) -> DockPosition {
582 match settings::get::<AssistantSettings>(cx).dock {
583 AssistantDockPosition::Left => DockPosition::Left,
584 AssistantDockPosition::Bottom => DockPosition::Bottom,
585 AssistantDockPosition::Right => DockPosition::Right,
586 }
587 }
588
589 fn position_is_valid(&self, _: DockPosition) -> bool {
590 true
591 }
592
593 fn set_position(&mut self, position: DockPosition, cx: &mut ViewContext<Self>) {
594 settings::update_settings_file::<AssistantSettings>(self.fs.clone(), cx, move |settings| {
595 let dock = match position {
596 DockPosition::Left => AssistantDockPosition::Left,
597 DockPosition::Bottom => AssistantDockPosition::Bottom,
598 DockPosition::Right => AssistantDockPosition::Right,
599 };
600 settings.dock = Some(dock);
601 });
602 }
603
604 fn size(&self, cx: &WindowContext) -> f32 {
605 let settings = settings::get::<AssistantSettings>(cx);
606 match self.position(cx) {
607 DockPosition::Left | DockPosition::Right => {
608 self.width.unwrap_or_else(|| settings.default_width)
609 }
610 DockPosition::Bottom => self.height.unwrap_or_else(|| settings.default_height),
611 }
612 }
613
614 fn set_size(&mut self, size: f32, cx: &mut ViewContext<Self>) {
615 match self.position(cx) {
616 DockPosition::Left | DockPosition::Right => self.width = Some(size),
617 DockPosition::Bottom => self.height = Some(size),
618 }
619 cx.notify();
620 }
621
622 fn should_zoom_in_on_event(event: &AssistantPanelEvent) -> bool {
623 matches!(event, AssistantPanelEvent::ZoomIn)
624 }
625
626 fn should_zoom_out_on_event(event: &AssistantPanelEvent) -> bool {
627 matches!(event, AssistantPanelEvent::ZoomOut)
628 }
629
630 fn is_zoomed(&self, _: &WindowContext) -> bool {
631 self.zoomed
632 }
633
634 fn set_zoomed(&mut self, zoomed: bool, cx: &mut ViewContext<Self>) {
635 self.zoomed = zoomed;
636 cx.notify();
637 }
638
639 fn set_active(&mut self, active: bool, cx: &mut ViewContext<Self>) {
640 if active {
641 if self.api_key.borrow().is_none() && !self.has_read_credentials {
642 self.has_read_credentials = true;
643 let api_key = if let Ok(api_key) = env::var("OPENAI_API_KEY") {
644 Some(api_key)
645 } else if let Some((_, api_key)) = cx
646 .platform()
647 .read_credentials(OPENAI_API_URL)
648 .log_err()
649 .flatten()
650 {
651 String::from_utf8(api_key).log_err()
652 } else {
653 None
654 };
655 if let Some(api_key) = api_key {
656 *self.api_key.borrow_mut() = Some(api_key);
657 } else if self.api_key_editor.is_none() {
658 self.api_key_editor = Some(build_api_key_editor(cx));
659 cx.notify();
660 }
661 }
662
663 if self.editors.is_empty() {
664 self.new_conversation(cx);
665 }
666 }
667 }
668
669 fn icon_path(&self) -> &'static str {
670 "icons/robot_14.svg"
671 }
672
673 fn icon_tooltip(&self) -> (String, Option<Box<dyn Action>>) {
674 ("Assistant Panel".into(), Some(Box::new(ToggleFocus)))
675 }
676
677 fn should_change_position_on_event(event: &Self::Event) -> bool {
678 matches!(event, AssistantPanelEvent::DockPositionChanged)
679 }
680
681 fn should_activate_on_event(_: &Self::Event) -> bool {
682 false
683 }
684
685 fn should_close_on_event(event: &AssistantPanelEvent) -> bool {
686 matches!(event, AssistantPanelEvent::Close)
687 }
688
689 fn has_focus(&self, _: &WindowContext) -> bool {
690 self.has_focus
691 }
692
693 fn is_focus_event(event: &Self::Event) -> bool {
694 matches!(event, AssistantPanelEvent::Focus)
695 }
696}
697
698enum ConversationEvent {
699 MessagesEdited,
700 SummaryChanged,
701 StreamedCompletion,
702}
703
704#[derive(Default)]
705struct Summary {
706 text: String,
707 done: bool,
708}
709
710struct Conversation {
711 buffer: ModelHandle<Buffer>,
712 message_anchors: Vec<MessageAnchor>,
713 messages_metadata: HashMap<MessageId, MessageMetadata>,
714 next_message_id: MessageId,
715 summary: Option<Summary>,
716 pending_summary: Task<Option<()>>,
717 completion_count: usize,
718 pending_completions: Vec<PendingCompletion>,
719 model: String,
720 token_count: Option<usize>,
721 max_token_count: usize,
722 pending_token_count: Task<Option<()>>,
723 api_key: Rc<RefCell<Option<String>>>,
724 pending_save: Task<Result<()>>,
725 path: Option<PathBuf>,
726 _subscriptions: Vec<Subscription>,
727}
728
729impl Entity for Conversation {
730 type Event = ConversationEvent;
731}
732
733impl Conversation {
734 fn new(
735 api_key: Rc<RefCell<Option<String>>>,
736 language_registry: Arc<LanguageRegistry>,
737 cx: &mut ModelContext<Self>,
738 ) -> Self {
739 let model = "gpt-3.5-turbo-0613";
740 let markdown = language_registry.language_for_name("Markdown");
741 let buffer = cx.add_model(|cx| {
742 let mut buffer = Buffer::new(0, "", cx);
743 buffer.set_language_registry(language_registry);
744 cx.spawn_weak(|buffer, mut cx| async move {
745 let markdown = markdown.await?;
746 let buffer = buffer
747 .upgrade(&cx)
748 .ok_or_else(|| anyhow!("buffer was dropped"))?;
749 buffer.update(&mut cx, |buffer, cx| {
750 buffer.set_language(Some(markdown), cx)
751 });
752 anyhow::Ok(())
753 })
754 .detach_and_log_err(cx);
755 buffer
756 });
757
758 let mut this = Self {
759 message_anchors: Default::default(),
760 messages_metadata: Default::default(),
761 next_message_id: Default::default(),
762 summary: None,
763 pending_summary: Task::ready(None),
764 completion_count: Default::default(),
765 pending_completions: Default::default(),
766 token_count: None,
767 max_token_count: tiktoken_rs::model::get_context_size(model),
768 pending_token_count: Task::ready(None),
769 model: model.into(),
770 _subscriptions: vec![cx.subscribe(&buffer, Self::handle_buffer_event)],
771 pending_save: Task::ready(Ok(())),
772 path: None,
773 api_key,
774 buffer,
775 };
776 let message = MessageAnchor {
777 id: MessageId(post_inc(&mut this.next_message_id.0)),
778 start: language::Anchor::MIN,
779 };
780 this.message_anchors.push(message.clone());
781 this.messages_metadata.insert(
782 message.id,
783 MessageMetadata {
784 role: Role::User,
785 sent_at: Local::now(),
786 status: MessageStatus::Done,
787 },
788 );
789
790 this.count_remaining_tokens(cx);
791 this
792 }
793
794 fn load(
795 path: PathBuf,
796 api_key: Rc<RefCell<Option<String>>>,
797 language_registry: Arc<LanguageRegistry>,
798 fs: Arc<dyn Fs>,
799 cx: &mut AppContext,
800 ) -> Task<Result<ModelHandle<Self>>> {
801 cx.spawn(|mut cx| async move {
802 let saved_conversation = fs.load(&path).await?;
803 let saved_conversation: SavedConversation = serde_json::from_str(&saved_conversation)?;
804
805 let model = saved_conversation.model;
806 let markdown = language_registry.language_for_name("Markdown");
807 let mut message_anchors = Vec::new();
808 let mut next_message_id = MessageId(0);
809 let buffer = cx.add_model(|cx| {
810 let mut buffer = Buffer::new(0, saved_conversation.text, cx);
811 for message in saved_conversation.messages {
812 message_anchors.push(MessageAnchor {
813 id: message.id,
814 start: buffer.anchor_before(message.start),
815 });
816 next_message_id = cmp::max(next_message_id, MessageId(message.id.0 + 1));
817 }
818 buffer.set_language_registry(language_registry);
819 cx.spawn_weak(|buffer, mut cx| async move {
820 let markdown = markdown.await?;
821 let buffer = buffer
822 .upgrade(&cx)
823 .ok_or_else(|| anyhow!("buffer was dropped"))?;
824 buffer.update(&mut cx, |buffer, cx| {
825 buffer.set_language(Some(markdown), cx)
826 });
827 anyhow::Ok(())
828 })
829 .detach_and_log_err(cx);
830 buffer
831 });
832 let conversation = cx.add_model(|cx| {
833 let mut this = Self {
834 message_anchors,
835 messages_metadata: saved_conversation.message_metadata,
836 next_message_id,
837 summary: Some(Summary {
838 text: saved_conversation.summary,
839 done: true,
840 }),
841 pending_summary: Task::ready(None),
842 completion_count: Default::default(),
843 pending_completions: Default::default(),
844 token_count: None,
845 max_token_count: tiktoken_rs::model::get_context_size(&model),
846 pending_token_count: Task::ready(None),
847 model,
848 _subscriptions: vec![cx.subscribe(&buffer, Self::handle_buffer_event)],
849 pending_save: Task::ready(Ok(())),
850 path: Some(path),
851 api_key,
852 buffer,
853 };
854
855 this.count_remaining_tokens(cx);
856 this
857 });
858 Ok(conversation)
859 })
860 }
861
862 fn handle_buffer_event(
863 &mut self,
864 _: ModelHandle<Buffer>,
865 event: &language::Event,
866 cx: &mut ModelContext<Self>,
867 ) {
868 match event {
869 language::Event::Edited => {
870 self.count_remaining_tokens(cx);
871 cx.emit(ConversationEvent::MessagesEdited);
872 }
873 _ => {}
874 }
875 }
876
877 fn count_remaining_tokens(&mut self, cx: &mut ModelContext<Self>) {
878 let messages = self
879 .messages(cx)
880 .into_iter()
881 .filter_map(|message| {
882 Some(tiktoken_rs::ChatCompletionRequestMessage {
883 role: match message.role {
884 Role::User => "user".into(),
885 Role::Assistant => "assistant".into(),
886 Role::System => "system".into(),
887 },
888 content: self.buffer.read(cx).text_for_range(message.range).collect(),
889 name: None,
890 })
891 })
892 .collect::<Vec<_>>();
893 let model = self.model.clone();
894 self.pending_token_count = cx.spawn_weak(|this, mut cx| {
895 async move {
896 cx.background().timer(Duration::from_millis(200)).await;
897 let token_count = cx
898 .background()
899 .spawn(async move { tiktoken_rs::num_tokens_from_messages(&model, &messages) })
900 .await?;
901
902 this.upgrade(&cx)
903 .ok_or_else(|| anyhow!("conversation was dropped"))?
904 .update(&mut cx, |this, cx| {
905 this.max_token_count = tiktoken_rs::model::get_context_size(&this.model);
906 this.token_count = Some(token_count);
907 cx.notify()
908 });
909 anyhow::Ok(())
910 }
911 .log_err()
912 });
913 }
914
915 fn remaining_tokens(&self) -> Option<isize> {
916 Some(self.max_token_count as isize - self.token_count? as isize)
917 }
918
919 fn set_model(&mut self, model: String, cx: &mut ModelContext<Self>) {
920 self.model = model;
921 self.count_remaining_tokens(cx);
922 cx.notify();
923 }
924
925 fn assist(
926 &mut self,
927 selected_messages: HashSet<MessageId>,
928 cx: &mut ModelContext<Self>,
929 ) -> Vec<MessageAnchor> {
930 let mut user_messages = Vec::new();
931 let mut tasks = Vec::new();
932 for selected_message_id in selected_messages {
933 let selected_message_role =
934 if let Some(metadata) = self.messages_metadata.get(&selected_message_id) {
935 metadata.role
936 } else {
937 continue;
938 };
939
940 if selected_message_role == Role::Assistant {
941 if let Some(user_message) = self.insert_message_after(
942 selected_message_id,
943 Role::User,
944 MessageStatus::Done,
945 cx,
946 ) {
947 user_messages.push(user_message);
948 } else {
949 continue;
950 }
951 } else {
952 let request = OpenAIRequest {
953 model: self.model.clone(),
954 messages: self
955 .messages(cx)
956 .filter(|message| matches!(message.status, MessageStatus::Done))
957 .flat_map(|message| {
958 let mut system_message = None;
959 if message.id == selected_message_id {
960 system_message = Some(RequestMessage {
961 role: Role::System,
962 content: concat!(
963 "Treat the following messages as additional knowledge you have learned about, ",
964 "but act as if they were not part of this conversation. That is, treat them ",
965 "as if the user didn't see them and couldn't possibly inquire about them."
966 ).into()
967 });
968 }
969
970 Some(message.to_open_ai_message(self.buffer.read(cx))).into_iter().chain(system_message)
971 })
972 .chain(Some(RequestMessage {
973 role: Role::System,
974 content: format!(
975 "Direct your reply to message with id {}. Do not include a [Message X] header.",
976 selected_message_id.0
977 ),
978 }))
979 .collect(),
980 stream: true,
981 };
982
983 let Some(api_key) = self.api_key.borrow().clone() else { continue };
984 let stream = stream_completion(api_key, cx.background().clone(), request);
985 let assistant_message = self
986 .insert_message_after(
987 selected_message_id,
988 Role::Assistant,
989 MessageStatus::Pending,
990 cx,
991 )
992 .unwrap();
993
994 tasks.push(cx.spawn_weak({
995 |this, mut cx| async move {
996 let assistant_message_id = assistant_message.id;
997 let stream_completion = async {
998 let mut messages = stream.await?;
999
1000 while let Some(message) = messages.next().await {
1001 let mut message = message?;
1002 if let Some(choice) = message.choices.pop() {
1003 this.upgrade(&cx)
1004 .ok_or_else(|| anyhow!("conversation was dropped"))?
1005 .update(&mut cx, |this, cx| {
1006 let text: Arc<str> = choice.delta.content?.into();
1007 let message_ix = this.message_anchors.iter().position(
1008 |message| message.id == assistant_message_id,
1009 )?;
1010 this.buffer.update(cx, |buffer, cx| {
1011 let offset = this.message_anchors[message_ix + 1..]
1012 .iter()
1013 .find(|message| message.start.is_valid(buffer))
1014 .map_or(buffer.len(), |message| {
1015 message
1016 .start
1017 .to_offset(buffer)
1018 .saturating_sub(1)
1019 });
1020 buffer.edit([(offset..offset, text)], None, cx);
1021 });
1022 cx.emit(ConversationEvent::StreamedCompletion);
1023
1024 Some(())
1025 });
1026 }
1027 smol::future::yield_now().await;
1028 }
1029
1030 this.upgrade(&cx)
1031 .ok_or_else(|| anyhow!("conversation was dropped"))?
1032 .update(&mut cx, |this, cx| {
1033 this.pending_completions.retain(|completion| {
1034 completion.id != this.completion_count
1035 });
1036 this.summarize(cx);
1037 });
1038
1039 anyhow::Ok(())
1040 };
1041
1042 let result = stream_completion.await;
1043 if let Some(this) = this.upgrade(&cx) {
1044 this.update(&mut cx, |this, cx| {
1045 if let Some(metadata) =
1046 this.messages_metadata.get_mut(&assistant_message.id)
1047 {
1048 match result {
1049 Ok(_) => {
1050 metadata.status = MessageStatus::Done;
1051 }
1052 Err(error) => {
1053 metadata.status = MessageStatus::Error(
1054 error.to_string().trim().into(),
1055 );
1056 }
1057 }
1058 cx.notify();
1059 }
1060 });
1061 }
1062 }
1063 }));
1064 }
1065 }
1066
1067 if !tasks.is_empty() {
1068 self.pending_completions.push(PendingCompletion {
1069 id: post_inc(&mut self.completion_count),
1070 _tasks: tasks,
1071 });
1072 }
1073
1074 user_messages
1075 }
1076
1077 fn cancel_last_assist(&mut self) -> bool {
1078 self.pending_completions.pop().is_some()
1079 }
1080
1081 fn cycle_message_roles(&mut self, ids: HashSet<MessageId>, cx: &mut ModelContext<Self>) {
1082 for id in ids {
1083 if let Some(metadata) = self.messages_metadata.get_mut(&id) {
1084 metadata.role.cycle();
1085 cx.emit(ConversationEvent::MessagesEdited);
1086 cx.notify();
1087 }
1088 }
1089 }
1090
1091 fn insert_message_after(
1092 &mut self,
1093 message_id: MessageId,
1094 role: Role,
1095 status: MessageStatus,
1096 cx: &mut ModelContext<Self>,
1097 ) -> Option<MessageAnchor> {
1098 if let Some(prev_message_ix) = self
1099 .message_anchors
1100 .iter()
1101 .position(|message| message.id == message_id)
1102 {
1103 let start = self.buffer.update(cx, |buffer, cx| {
1104 let offset = self.message_anchors[prev_message_ix + 1..]
1105 .iter()
1106 .find(|message| message.start.is_valid(buffer))
1107 .map_or(buffer.len(), |message| message.start.to_offset(buffer) - 1);
1108 buffer.edit([(offset..offset, "\n")], None, cx);
1109 buffer.anchor_before(offset + 1)
1110 });
1111 let message = MessageAnchor {
1112 id: MessageId(post_inc(&mut self.next_message_id.0)),
1113 start,
1114 };
1115 self.message_anchors
1116 .insert(prev_message_ix + 1, message.clone());
1117 self.messages_metadata.insert(
1118 message.id,
1119 MessageMetadata {
1120 role,
1121 sent_at: Local::now(),
1122 status,
1123 },
1124 );
1125 cx.emit(ConversationEvent::MessagesEdited);
1126 Some(message)
1127 } else {
1128 None
1129 }
1130 }
1131
1132 fn split_message(
1133 &mut self,
1134 range: Range<usize>,
1135 cx: &mut ModelContext<Self>,
1136 ) -> (Option<MessageAnchor>, Option<MessageAnchor>) {
1137 let start_message = self.message_for_offset(range.start, cx);
1138 let end_message = self.message_for_offset(range.end, cx);
1139 if let Some((start_message, end_message)) = start_message.zip(end_message) {
1140 // Prevent splitting when range spans multiple messages.
1141 if start_message.index != end_message.index {
1142 return (None, None);
1143 }
1144
1145 let message = start_message;
1146 let role = message.role;
1147 let mut edited_buffer = false;
1148
1149 let mut suffix_start = None;
1150 if range.start > message.range.start && range.end < message.range.end - 1 {
1151 if self.buffer.read(cx).chars_at(range.end).next() == Some('\n') {
1152 suffix_start = Some(range.end + 1);
1153 } else if self.buffer.read(cx).reversed_chars_at(range.end).next() == Some('\n') {
1154 suffix_start = Some(range.end);
1155 }
1156 }
1157
1158 let suffix = if let Some(suffix_start) = suffix_start {
1159 MessageAnchor {
1160 id: MessageId(post_inc(&mut self.next_message_id.0)),
1161 start: self.buffer.read(cx).anchor_before(suffix_start),
1162 }
1163 } else {
1164 self.buffer.update(cx, |buffer, cx| {
1165 buffer.edit([(range.end..range.end, "\n")], None, cx);
1166 });
1167 edited_buffer = true;
1168 MessageAnchor {
1169 id: MessageId(post_inc(&mut self.next_message_id.0)),
1170 start: self.buffer.read(cx).anchor_before(range.end + 1),
1171 }
1172 };
1173
1174 self.message_anchors
1175 .insert(message.index + 1, suffix.clone());
1176 self.messages_metadata.insert(
1177 suffix.id,
1178 MessageMetadata {
1179 role,
1180 sent_at: Local::now(),
1181 status: MessageStatus::Done,
1182 },
1183 );
1184
1185 let new_messages = if range.start == range.end || range.start == message.range.start {
1186 (None, Some(suffix))
1187 } else {
1188 let mut prefix_end = None;
1189 if range.start > message.range.start && range.end < message.range.end - 1 {
1190 if self.buffer.read(cx).chars_at(range.start).next() == Some('\n') {
1191 prefix_end = Some(range.start + 1);
1192 } else if self.buffer.read(cx).reversed_chars_at(range.start).next()
1193 == Some('\n')
1194 {
1195 prefix_end = Some(range.start);
1196 }
1197 }
1198
1199 let selection = if let Some(prefix_end) = prefix_end {
1200 cx.emit(ConversationEvent::MessagesEdited);
1201 MessageAnchor {
1202 id: MessageId(post_inc(&mut self.next_message_id.0)),
1203 start: self.buffer.read(cx).anchor_before(prefix_end),
1204 }
1205 } else {
1206 self.buffer.update(cx, |buffer, cx| {
1207 buffer.edit([(range.start..range.start, "\n")], None, cx)
1208 });
1209 edited_buffer = true;
1210 MessageAnchor {
1211 id: MessageId(post_inc(&mut self.next_message_id.0)),
1212 start: self.buffer.read(cx).anchor_before(range.end + 1),
1213 }
1214 };
1215
1216 self.message_anchors
1217 .insert(message.index + 1, selection.clone());
1218 self.messages_metadata.insert(
1219 selection.id,
1220 MessageMetadata {
1221 role,
1222 sent_at: Local::now(),
1223 status: MessageStatus::Done,
1224 },
1225 );
1226 (Some(selection), Some(suffix))
1227 };
1228
1229 if !edited_buffer {
1230 cx.emit(ConversationEvent::MessagesEdited);
1231 }
1232 new_messages
1233 } else {
1234 (None, None)
1235 }
1236 }
1237
1238 fn summarize(&mut self, cx: &mut ModelContext<Self>) {
1239 if self.message_anchors.len() >= 2 && self.summary.is_none() {
1240 let api_key = self.api_key.borrow().clone();
1241 if let Some(api_key) = api_key {
1242 let messages = self
1243 .messages(cx)
1244 .take(2)
1245 .map(|message| message.to_open_ai_message(self.buffer.read(cx)))
1246 .chain(Some(RequestMessage {
1247 role: Role::User,
1248 content:
1249 "Summarize the conversation into a short title without punctuation"
1250 .into(),
1251 }));
1252 let request = OpenAIRequest {
1253 model: self.model.clone(),
1254 messages: messages.collect(),
1255 stream: true,
1256 };
1257
1258 let stream = stream_completion(api_key, cx.background().clone(), request);
1259 self.pending_summary = cx.spawn(|this, mut cx| {
1260 async move {
1261 let mut messages = stream.await?;
1262
1263 while let Some(message) = messages.next().await {
1264 let mut message = message?;
1265 if let Some(choice) = message.choices.pop() {
1266 let text = choice.delta.content.unwrap_or_default();
1267 this.update(&mut cx, |this, cx| {
1268 this.summary
1269 .get_or_insert(Default::default())
1270 .text
1271 .push_str(&text);
1272 cx.emit(ConversationEvent::SummaryChanged);
1273 });
1274 }
1275 }
1276
1277 this.update(&mut cx, |this, cx| {
1278 if let Some(summary) = this.summary.as_mut() {
1279 summary.done = true;
1280 cx.emit(ConversationEvent::SummaryChanged);
1281 }
1282 });
1283
1284 anyhow::Ok(())
1285 }
1286 .log_err()
1287 });
1288 }
1289 }
1290 }
1291
1292 fn message_for_offset(&self, offset: usize, cx: &AppContext) -> Option<Message> {
1293 self.messages_for_offsets([offset], cx).pop()
1294 }
1295
1296 fn messages_for_offsets(
1297 &self,
1298 offsets: impl IntoIterator<Item = usize>,
1299 cx: &AppContext,
1300 ) -> Vec<Message> {
1301 let mut result = Vec::new();
1302
1303 let mut messages = self.messages(cx).peekable();
1304 let mut offsets = offsets.into_iter().peekable();
1305 let mut current_message = messages.next();
1306 while let Some(offset) = offsets.next() {
1307 // Locate the message that contains the offset.
1308 while current_message.as_ref().map_or(false, |message| {
1309 !message.range.contains(&offset) && messages.peek().is_some()
1310 }) {
1311 current_message = messages.next();
1312 }
1313 let Some(message) = current_message.as_ref() else { break };
1314
1315 // Skip offsets that are in the same message.
1316 while offsets.peek().map_or(false, |offset| {
1317 message.range.contains(offset) || messages.peek().is_none()
1318 }) {
1319 offsets.next();
1320 }
1321
1322 result.push(message.clone());
1323 }
1324 result
1325 }
1326
1327 fn messages<'a>(&'a self, cx: &'a AppContext) -> impl 'a + Iterator<Item = Message> {
1328 let buffer = self.buffer.read(cx);
1329 let mut message_anchors = self.message_anchors.iter().enumerate().peekable();
1330 iter::from_fn(move || {
1331 while let Some((ix, message_anchor)) = message_anchors.next() {
1332 let metadata = self.messages_metadata.get(&message_anchor.id)?;
1333 let message_start = message_anchor.start.to_offset(buffer);
1334 let mut message_end = None;
1335 while let Some((_, next_message)) = message_anchors.peek() {
1336 if next_message.start.is_valid(buffer) {
1337 message_end = Some(next_message.start);
1338 break;
1339 } else {
1340 message_anchors.next();
1341 }
1342 }
1343 let message_end = message_end
1344 .unwrap_or(language::Anchor::MAX)
1345 .to_offset(buffer);
1346 return Some(Message {
1347 index: ix,
1348 range: message_start..message_end,
1349 id: message_anchor.id,
1350 anchor: message_anchor.start,
1351 role: metadata.role,
1352 sent_at: metadata.sent_at,
1353 status: metadata.status.clone(),
1354 });
1355 }
1356 None
1357 })
1358 }
1359
1360 fn save(
1361 &mut self,
1362 debounce: Option<Duration>,
1363 fs: Arc<dyn Fs>,
1364 cx: &mut ModelContext<Conversation>,
1365 ) {
1366 self.pending_save = cx.spawn(|this, mut cx| async move {
1367 if let Some(debounce) = debounce {
1368 cx.background().timer(debounce).await;
1369 }
1370
1371 let (old_path, summary) = this.read_with(&cx, |this, _| {
1372 let path = this.path.clone();
1373 let summary = if let Some(summary) = this.summary.as_ref() {
1374 if summary.done {
1375 Some(summary.text.clone())
1376 } else {
1377 None
1378 }
1379 } else {
1380 None
1381 };
1382 (path, summary)
1383 });
1384
1385 if let Some(summary) = summary {
1386 let conversation = this.read_with(&cx, |this, cx| SavedConversation {
1387 zed: "conversation".into(),
1388 version: SavedConversation::VERSION.into(),
1389 text: this.buffer.read(cx).text(),
1390 message_metadata: this.messages_metadata.clone(),
1391 messages: this
1392 .message_anchors
1393 .iter()
1394 .map(|message| SavedMessage {
1395 id: message.id,
1396 start: message.start.to_offset(this.buffer.read(cx)),
1397 })
1398 .collect(),
1399 summary: summary.clone(),
1400 model: this.model.clone(),
1401 });
1402
1403 let path = if let Some(old_path) = old_path {
1404 old_path
1405 } else {
1406 let mut discriminant = 1;
1407 let mut new_path;
1408 loop {
1409 new_path = CONVERSATIONS_DIR.join(&format!(
1410 "{} - {}.zed.json",
1411 summary.trim(),
1412 discriminant
1413 ));
1414 if fs.is_file(&new_path).await {
1415 discriminant += 1;
1416 } else {
1417 break;
1418 }
1419 }
1420 new_path
1421 };
1422
1423 fs.create_dir(CONVERSATIONS_DIR.as_ref()).await?;
1424 fs.atomic_write(path.clone(), serde_json::to_string(&conversation).unwrap())
1425 .await?;
1426 this.update(&mut cx, |this, _| this.path = Some(path));
1427 }
1428
1429 Ok(())
1430 });
1431 }
1432}
1433
1434struct PendingCompletion {
1435 id: usize,
1436 _tasks: Vec<Task<()>>,
1437}
1438
1439enum ConversationEditorEvent {
1440 TabContentChanged,
1441}
1442
1443#[derive(Copy, Clone, Debug, PartialEq)]
1444struct ScrollPosition {
1445 offset_before_cursor: Vector2F,
1446 cursor: Anchor,
1447}
1448
1449struct ConversationEditor {
1450 conversation: ModelHandle<Conversation>,
1451 fs: Arc<dyn Fs>,
1452 editor: ViewHandle<Editor>,
1453 blocks: HashSet<BlockId>,
1454 scroll_position: Option<ScrollPosition>,
1455 _subscriptions: Vec<Subscription>,
1456}
1457
1458impl ConversationEditor {
1459 fn new(
1460 api_key: Rc<RefCell<Option<String>>>,
1461 language_registry: Arc<LanguageRegistry>,
1462 fs: Arc<dyn Fs>,
1463 cx: &mut ViewContext<Self>,
1464 ) -> Self {
1465 let conversation = cx.add_model(|cx| Conversation::new(api_key, language_registry, cx));
1466 Self::from_conversation(conversation, fs, cx)
1467 }
1468
1469 fn from_conversation(
1470 conversation: ModelHandle<Conversation>,
1471 fs: Arc<dyn Fs>,
1472 cx: &mut ViewContext<Self>,
1473 ) -> Self {
1474 let editor = cx.add_view(|cx| {
1475 let mut editor = Editor::for_buffer(conversation.read(cx).buffer.clone(), None, cx);
1476 editor.set_soft_wrap_mode(SoftWrap::EditorWidth, cx);
1477 editor.set_show_gutter(false, cx);
1478 editor
1479 });
1480
1481 let _subscriptions = vec![
1482 cx.observe(&conversation, |_, _, cx| cx.notify()),
1483 cx.subscribe(&conversation, Self::handle_conversation_event),
1484 cx.subscribe(&editor, Self::handle_editor_event),
1485 ];
1486
1487 let mut this = Self {
1488 conversation,
1489 editor,
1490 blocks: Default::default(),
1491 scroll_position: None,
1492 fs,
1493 _subscriptions,
1494 };
1495 this.update_message_headers(cx);
1496 this
1497 }
1498
1499 fn assist(&mut self, _: &Assist, cx: &mut ViewContext<Self>) {
1500 let cursors = self.cursors(cx);
1501
1502 let user_messages = self.conversation.update(cx, |conversation, cx| {
1503 let selected_messages = conversation
1504 .messages_for_offsets(cursors, cx)
1505 .into_iter()
1506 .map(|message| message.id)
1507 .collect();
1508 conversation.assist(selected_messages, cx)
1509 });
1510 let new_selections = user_messages
1511 .iter()
1512 .map(|message| {
1513 let cursor = message
1514 .start
1515 .to_offset(self.conversation.read(cx).buffer.read(cx));
1516 cursor..cursor
1517 })
1518 .collect::<Vec<_>>();
1519 if !new_selections.is_empty() {
1520 self.editor.update(cx, |editor, cx| {
1521 editor.change_selections(
1522 Some(Autoscroll::Strategy(AutoscrollStrategy::Fit)),
1523 cx,
1524 |selections| selections.select_ranges(new_selections),
1525 );
1526 });
1527 }
1528 }
1529
1530 fn cancel_last_assist(&mut self, _: &editor::Cancel, cx: &mut ViewContext<Self>) {
1531 if !self
1532 .conversation
1533 .update(cx, |conversation, _| conversation.cancel_last_assist())
1534 {
1535 cx.propagate_action();
1536 }
1537 }
1538
1539 fn cycle_message_role(&mut self, _: &CycleMessageRole, cx: &mut ViewContext<Self>) {
1540 let cursors = self.cursors(cx);
1541 self.conversation.update(cx, |conversation, cx| {
1542 let messages = conversation
1543 .messages_for_offsets(cursors, cx)
1544 .into_iter()
1545 .map(|message| message.id)
1546 .collect();
1547 conversation.cycle_message_roles(messages, cx)
1548 });
1549 }
1550
1551 fn cursors(&self, cx: &AppContext) -> Vec<usize> {
1552 let selections = self.editor.read(cx).selections.all::<usize>(cx);
1553 selections
1554 .into_iter()
1555 .map(|selection| selection.head())
1556 .collect()
1557 }
1558
1559 fn handle_conversation_event(
1560 &mut self,
1561 _: ModelHandle<Conversation>,
1562 event: &ConversationEvent,
1563 cx: &mut ViewContext<Self>,
1564 ) {
1565 match event {
1566 ConversationEvent::MessagesEdited => {
1567 self.update_message_headers(cx);
1568 self.conversation.update(cx, |conversation, cx| {
1569 conversation.save(Some(Duration::from_millis(500)), self.fs.clone(), cx);
1570 });
1571 }
1572 ConversationEvent::SummaryChanged => {
1573 cx.emit(ConversationEditorEvent::TabContentChanged);
1574 self.conversation.update(cx, |conversation, cx| {
1575 conversation.save(None, self.fs.clone(), cx);
1576 });
1577 }
1578 ConversationEvent::StreamedCompletion => {
1579 self.editor.update(cx, |editor, cx| {
1580 if let Some(scroll_position) = self.scroll_position {
1581 let snapshot = editor.snapshot(cx);
1582 let cursor_point = scroll_position.cursor.to_display_point(&snapshot);
1583 let scroll_top =
1584 cursor_point.row() as f32 - scroll_position.offset_before_cursor.y();
1585 editor.set_scroll_position(
1586 vec2f(scroll_position.offset_before_cursor.x(), scroll_top),
1587 cx,
1588 );
1589 }
1590 });
1591 }
1592 }
1593 }
1594
1595 fn handle_editor_event(
1596 &mut self,
1597 _: ViewHandle<Editor>,
1598 event: &editor::Event,
1599 cx: &mut ViewContext<Self>,
1600 ) {
1601 match event {
1602 editor::Event::ScrollPositionChanged { autoscroll, .. } => {
1603 let cursor_scroll_position = self.cursor_scroll_position(cx);
1604 if *autoscroll {
1605 self.scroll_position = cursor_scroll_position;
1606 } else if self.scroll_position != cursor_scroll_position {
1607 self.scroll_position = None;
1608 }
1609 }
1610 editor::Event::SelectionsChanged { .. } => {
1611 self.scroll_position = self.cursor_scroll_position(cx);
1612 }
1613 _ => {}
1614 }
1615 }
1616
1617 fn cursor_scroll_position(&self, cx: &mut ViewContext<Self>) -> Option<ScrollPosition> {
1618 self.editor.update(cx, |editor, cx| {
1619 let snapshot = editor.snapshot(cx);
1620 let cursor = editor.selections.newest_anchor().head();
1621 let cursor_row = cursor.to_display_point(&snapshot.display_snapshot).row() as f32;
1622 let scroll_position = editor
1623 .scroll_manager
1624 .anchor()
1625 .scroll_position(&snapshot.display_snapshot);
1626
1627 let scroll_bottom = scroll_position.y() + editor.visible_line_count().unwrap_or(0.);
1628 if (scroll_position.y()..scroll_bottom).contains(&cursor_row) {
1629 Some(ScrollPosition {
1630 cursor,
1631 offset_before_cursor: vec2f(
1632 scroll_position.x(),
1633 cursor_row - scroll_position.y(),
1634 ),
1635 })
1636 } else {
1637 None
1638 }
1639 })
1640 }
1641
1642 fn update_message_headers(&mut self, cx: &mut ViewContext<Self>) {
1643 self.editor.update(cx, |editor, cx| {
1644 let buffer = editor.buffer().read(cx).snapshot(cx);
1645 let excerpt_id = *buffer.as_singleton().unwrap().0;
1646 let old_blocks = std::mem::take(&mut self.blocks);
1647 let new_blocks = self
1648 .conversation
1649 .read(cx)
1650 .messages(cx)
1651 .map(|message| BlockProperties {
1652 position: buffer.anchor_in_excerpt(excerpt_id, message.anchor),
1653 height: 2,
1654 style: BlockStyle::Sticky,
1655 render: Arc::new({
1656 let conversation = self.conversation.clone();
1657 // let metadata = message.metadata.clone();
1658 // let message = message.clone();
1659 move |cx| {
1660 enum Sender {}
1661 enum ErrorTooltip {}
1662
1663 let theme = theme::current(cx);
1664 let style = &theme.assistant;
1665 let message_id = message.id;
1666 let sender = MouseEventHandler::<Sender, _>::new(
1667 message_id.0,
1668 cx,
1669 |state, _| match message.role {
1670 Role::User => {
1671 let style = style.user_sender.style_for(state);
1672 Label::new("You", style.text.clone())
1673 .contained()
1674 .with_style(style.container)
1675 }
1676 Role::Assistant => {
1677 let style = style.assistant_sender.style_for(state);
1678 Label::new("Assistant", style.text.clone())
1679 .contained()
1680 .with_style(style.container)
1681 }
1682 Role::System => {
1683 let style = style.system_sender.style_for(state);
1684 Label::new("System", style.text.clone())
1685 .contained()
1686 .with_style(style.container)
1687 }
1688 },
1689 )
1690 .with_cursor_style(CursorStyle::PointingHand)
1691 .on_down(MouseButton::Left, {
1692 let conversation = conversation.clone();
1693 move |_, _, cx| {
1694 conversation.update(cx, |conversation, cx| {
1695 conversation.cycle_message_roles(
1696 HashSet::from_iter(Some(message_id)),
1697 cx,
1698 )
1699 })
1700 }
1701 });
1702
1703 Flex::row()
1704 .with_child(sender.aligned())
1705 .with_child(
1706 Label::new(
1707 message.sent_at.format("%I:%M%P").to_string(),
1708 style.sent_at.text.clone(),
1709 )
1710 .contained()
1711 .with_style(style.sent_at.container)
1712 .aligned(),
1713 )
1714 .with_children(
1715 if let MessageStatus::Error(error) = &message.status {
1716 Some(
1717 Svg::new("icons/circle_x_mark_12.svg")
1718 .with_color(style.error_icon.color)
1719 .constrained()
1720 .with_width(style.error_icon.width)
1721 .contained()
1722 .with_style(style.error_icon.container)
1723 .with_tooltip::<ErrorTooltip>(
1724 message_id.0,
1725 error.to_string(),
1726 None,
1727 theme.tooltip.clone(),
1728 cx,
1729 )
1730 .aligned(),
1731 )
1732 } else {
1733 None
1734 },
1735 )
1736 .aligned()
1737 .left()
1738 .contained()
1739 .with_style(style.message_header)
1740 .into_any()
1741 }
1742 }),
1743 disposition: BlockDisposition::Above,
1744 })
1745 .collect::<Vec<_>>();
1746
1747 editor.remove_blocks(old_blocks, None, cx);
1748 let ids = editor.insert_blocks(new_blocks, None, cx);
1749 self.blocks = HashSet::from_iter(ids);
1750 });
1751 }
1752
1753 fn quote_selection(
1754 workspace: &mut Workspace,
1755 _: &QuoteSelection,
1756 cx: &mut ViewContext<Workspace>,
1757 ) {
1758 let Some(panel) = workspace.panel::<AssistantPanel>(cx) else {
1759 return;
1760 };
1761 let Some(editor) = workspace.active_item(cx).and_then(|item| item.downcast::<Editor>()) else {
1762 return;
1763 };
1764
1765 let text = editor.read_with(cx, |editor, cx| {
1766 let range = editor.selections.newest::<usize>(cx).range();
1767 let buffer = editor.buffer().read(cx).snapshot(cx);
1768 let start_language = buffer.language_at(range.start);
1769 let end_language = buffer.language_at(range.end);
1770 let language_name = if start_language == end_language {
1771 start_language.map(|language| language.name())
1772 } else {
1773 None
1774 };
1775 let language_name = language_name.as_deref().unwrap_or("").to_lowercase();
1776
1777 let selected_text = buffer.text_for_range(range).collect::<String>();
1778 if selected_text.is_empty() {
1779 None
1780 } else {
1781 Some(if language_name == "markdown" {
1782 selected_text
1783 .lines()
1784 .map(|line| format!("> {}", line))
1785 .collect::<Vec<_>>()
1786 .join("\n")
1787 } else {
1788 format!("```{language_name}\n{selected_text}\n```")
1789 })
1790 }
1791 });
1792
1793 // Activate the panel
1794 if !panel.read(cx).has_focus(cx) {
1795 workspace.toggle_panel_focus::<AssistantPanel>(cx);
1796 }
1797
1798 if let Some(text) = text {
1799 panel.update(cx, |panel, cx| {
1800 let conversation = panel
1801 .active_editor()
1802 .cloned()
1803 .unwrap_or_else(|| panel.new_conversation(cx));
1804 conversation.update(cx, |conversation, cx| {
1805 conversation
1806 .editor
1807 .update(cx, |editor, cx| editor.insert(&text, cx))
1808 });
1809 });
1810 }
1811 }
1812
1813 fn copy(&mut self, _: &editor::Copy, cx: &mut ViewContext<Self>) {
1814 let editor = self.editor.read(cx);
1815 let conversation = self.conversation.read(cx);
1816 if editor.selections.count() == 1 {
1817 let selection = editor.selections.newest::<usize>(cx);
1818 let mut copied_text = String::new();
1819 let mut spanned_messages = 0;
1820 for message in conversation.messages(cx) {
1821 if message.range.start >= selection.range().end {
1822 break;
1823 } else if message.range.end >= selection.range().start {
1824 let range = cmp::max(message.range.start, selection.range().start)
1825 ..cmp::min(message.range.end, selection.range().end);
1826 if !range.is_empty() {
1827 spanned_messages += 1;
1828 write!(&mut copied_text, "## {}\n\n", message.role).unwrap();
1829 for chunk in conversation.buffer.read(cx).text_for_range(range) {
1830 copied_text.push_str(&chunk);
1831 }
1832 copied_text.push('\n');
1833 }
1834 }
1835 }
1836
1837 if spanned_messages > 1 {
1838 cx.platform()
1839 .write_to_clipboard(ClipboardItem::new(copied_text));
1840 return;
1841 }
1842 }
1843
1844 cx.propagate_action();
1845 }
1846
1847 fn split(&mut self, _: &Split, cx: &mut ViewContext<Self>) {
1848 self.conversation.update(cx, |conversation, cx| {
1849 let selections = self.editor.read(cx).selections.disjoint_anchors();
1850 for selection in selections.into_iter() {
1851 let buffer = self.editor.read(cx).buffer().read(cx).snapshot(cx);
1852 let range = selection
1853 .map(|endpoint| endpoint.to_offset(&buffer))
1854 .range();
1855 conversation.split_message(range, cx);
1856 }
1857 });
1858 }
1859
1860 fn save(&mut self, _: &Save, cx: &mut ViewContext<Self>) {
1861 self.conversation.update(cx, |conversation, cx| {
1862 conversation.save(None, self.fs.clone(), cx)
1863 });
1864 }
1865
1866 fn cycle_model(&mut self, cx: &mut ViewContext<Self>) {
1867 self.conversation.update(cx, |conversation, cx| {
1868 let new_model = match conversation.model.as_str() {
1869 "gpt-4-0613" => "gpt-3.5-turbo-0613",
1870 _ => "gpt-4-0613",
1871 };
1872 conversation.set_model(new_model.into(), cx);
1873 });
1874 }
1875
1876 fn title(&self, cx: &AppContext) -> String {
1877 self.conversation
1878 .read(cx)
1879 .summary
1880 .as_ref()
1881 .map(|summary| summary.text.clone())
1882 .unwrap_or_else(|| "New Conversation".into())
1883 }
1884}
1885
1886impl Entity for ConversationEditor {
1887 type Event = ConversationEditorEvent;
1888}
1889
1890impl View for ConversationEditor {
1891 fn ui_name() -> &'static str {
1892 "ConversationEditor"
1893 }
1894
1895 fn render(&mut self, cx: &mut ViewContext<Self>) -> AnyElement<Self> {
1896 let theme = &theme::current(cx).assistant;
1897 ChildView::new(&self.editor, cx)
1898 .contained()
1899 .with_style(theme.container)
1900 .into_any()
1901 }
1902
1903 fn focus_in(&mut self, _: gpui::AnyViewHandle, cx: &mut ViewContext<Self>) {
1904 if cx.is_self_focused() {
1905 cx.focus(&self.editor);
1906 }
1907 }
1908}
1909
1910impl Item for ConversationEditor {
1911 fn tab_content<V: View>(
1912 &self,
1913 _: Option<usize>,
1914 style: &theme::Tab,
1915 cx: &gpui::AppContext,
1916 ) -> AnyElement<V> {
1917 let title = truncate_and_trailoff(&self.title(cx), editor::MAX_TAB_TITLE_LEN);
1918 Label::new(title, style.label.clone()).into_any()
1919 }
1920
1921 fn tab_tooltip_text(&self, cx: &AppContext) -> Option<Cow<str>> {
1922 Some(self.title(cx).into())
1923 }
1924
1925 fn as_searchable(
1926 &self,
1927 _: &ViewHandle<Self>,
1928 ) -> Option<Box<dyn workspace::searchable::SearchableItemHandle>> {
1929 Some(Box::new(self.editor.clone()))
1930 }
1931}
1932
1933#[derive(Clone, Debug)]
1934struct MessageAnchor {
1935 id: MessageId,
1936 start: language::Anchor,
1937}
1938
1939#[derive(Clone, Debug)]
1940pub struct Message {
1941 range: Range<usize>,
1942 index: usize,
1943 id: MessageId,
1944 anchor: language::Anchor,
1945 role: Role,
1946 sent_at: DateTime<Local>,
1947 status: MessageStatus,
1948}
1949
1950impl Message {
1951 fn to_open_ai_message(&self, buffer: &Buffer) -> RequestMessage {
1952 let mut content = format!("[Message {}]\n", self.id.0).to_string();
1953 content.extend(buffer.text_for_range(self.range.clone()));
1954 RequestMessage {
1955 role: self.role,
1956 content: content.trim_end().into(),
1957 }
1958 }
1959}
1960
1961async fn stream_completion(
1962 api_key: String,
1963 executor: Arc<Background>,
1964 mut request: OpenAIRequest,
1965) -> Result<impl Stream<Item = Result<OpenAIResponseStreamEvent>>> {
1966 request.stream = true;
1967
1968 let (tx, rx) = futures::channel::mpsc::unbounded::<Result<OpenAIResponseStreamEvent>>();
1969
1970 let json_data = serde_json::to_string(&request)?;
1971 let mut response = Request::post(format!("{OPENAI_API_URL}/chat/completions"))
1972 .header("Content-Type", "application/json")
1973 .header("Authorization", format!("Bearer {}", api_key))
1974 .body(json_data)?
1975 .send_async()
1976 .await?;
1977
1978 let status = response.status();
1979 if status == StatusCode::OK {
1980 executor
1981 .spawn(async move {
1982 let mut lines = BufReader::new(response.body_mut()).lines();
1983
1984 fn parse_line(
1985 line: Result<String, io::Error>,
1986 ) -> Result<Option<OpenAIResponseStreamEvent>> {
1987 if let Some(data) = line?.strip_prefix("data: ") {
1988 let event = serde_json::from_str(&data)?;
1989 Ok(Some(event))
1990 } else {
1991 Ok(None)
1992 }
1993 }
1994
1995 while let Some(line) = lines.next().await {
1996 if let Some(event) = parse_line(line).transpose() {
1997 let done = event.as_ref().map_or(false, |event| {
1998 event
1999 .choices
2000 .last()
2001 .map_or(false, |choice| choice.finish_reason.is_some())
2002 });
2003 if tx.unbounded_send(event).is_err() {
2004 break;
2005 }
2006
2007 if done {
2008 break;
2009 }
2010 }
2011 }
2012
2013 anyhow::Ok(())
2014 })
2015 .detach();
2016
2017 Ok(rx)
2018 } else {
2019 let mut body = String::new();
2020 response.body_mut().read_to_string(&mut body).await?;
2021
2022 #[derive(Deserialize)]
2023 struct OpenAIResponse {
2024 error: OpenAIError,
2025 }
2026
2027 #[derive(Deserialize)]
2028 struct OpenAIError {
2029 message: String,
2030 }
2031
2032 match serde_json::from_str::<OpenAIResponse>(&body) {
2033 Ok(response) if !response.error.message.is_empty() => Err(anyhow!(
2034 "Failed to connect to OpenAI API: {}",
2035 response.error.message,
2036 )),
2037
2038 _ => Err(anyhow!(
2039 "Failed to connect to OpenAI API: {} {}",
2040 response.status(),
2041 body,
2042 )),
2043 }
2044 }
2045}
2046
2047#[cfg(test)]
2048mod tests {
2049 use crate::MessageId;
2050
2051 use super::*;
2052 use fs::FakeFs;
2053 use gpui::{AppContext, TestAppContext};
2054 use project::Project;
2055
2056 fn init_test(cx: &mut TestAppContext) {
2057 cx.foreground().forbid_parking();
2058 cx.update(|cx| {
2059 cx.set_global(SettingsStore::test(cx));
2060 theme::init((), cx);
2061 language::init(cx);
2062 editor::init_settings(cx);
2063 crate::init(cx);
2064 workspace::init_settings(cx);
2065 Project::init_settings(cx);
2066 });
2067 }
2068
2069 #[gpui::test]
2070 async fn test_panel(cx: &mut TestAppContext) {
2071 init_test(cx);
2072
2073 let fs = FakeFs::new(cx.background());
2074 let project = Project::test(fs, [], cx).await;
2075 let (window_id, workspace) = cx.add_window(|cx| Workspace::test_new(project.clone(), cx));
2076 let weak_workspace = workspace.downgrade();
2077
2078 let panel = cx
2079 .spawn(|cx| async move { AssistantPanel::load(weak_workspace, cx).await })
2080 .await
2081 .unwrap();
2082
2083 workspace.update(cx, |workspace, cx| {
2084 workspace.add_panel(panel.clone(), cx);
2085 workspace.toggle_dock(DockPosition::Right, cx);
2086 assert!(workspace.right_dock().read(cx).is_open());
2087 cx.focus(&panel);
2088 });
2089
2090 cx.dispatch_action(window_id, workspace::ToggleZoom);
2091
2092 workspace.read_with(cx, |workspace, cx| {
2093 assert_eq!(workspace.zoomed_view(cx).unwrap(), panel);
2094 })
2095 }
2096
2097 #[gpui::test]
2098 fn test_inserting_and_removing_messages(cx: &mut AppContext) {
2099 let registry = Arc::new(LanguageRegistry::test());
2100 let conversation = cx.add_model(|cx| Conversation::new(Default::default(), registry, cx));
2101 let buffer = conversation.read(cx).buffer.clone();
2102
2103 let message_1 = conversation.read(cx).message_anchors[0].clone();
2104 assert_eq!(
2105 messages(&conversation, cx),
2106 vec![(message_1.id, Role::User, 0..0)]
2107 );
2108
2109 let message_2 = conversation.update(cx, |conversation, cx| {
2110 conversation
2111 .insert_message_after(message_1.id, Role::Assistant, MessageStatus::Done, cx)
2112 .unwrap()
2113 });
2114 assert_eq!(
2115 messages(&conversation, cx),
2116 vec![
2117 (message_1.id, Role::User, 0..1),
2118 (message_2.id, Role::Assistant, 1..1)
2119 ]
2120 );
2121
2122 buffer.update(cx, |buffer, cx| {
2123 buffer.edit([(0..0, "1"), (1..1, "2")], None, cx)
2124 });
2125 assert_eq!(
2126 messages(&conversation, cx),
2127 vec![
2128 (message_1.id, Role::User, 0..2),
2129 (message_2.id, Role::Assistant, 2..3)
2130 ]
2131 );
2132
2133 let message_3 = conversation.update(cx, |conversation, cx| {
2134 conversation
2135 .insert_message_after(message_2.id, Role::User, MessageStatus::Done, cx)
2136 .unwrap()
2137 });
2138 assert_eq!(
2139 messages(&conversation, cx),
2140 vec![
2141 (message_1.id, Role::User, 0..2),
2142 (message_2.id, Role::Assistant, 2..4),
2143 (message_3.id, Role::User, 4..4)
2144 ]
2145 );
2146
2147 let message_4 = conversation.update(cx, |conversation, cx| {
2148 conversation
2149 .insert_message_after(message_2.id, Role::User, MessageStatus::Done, cx)
2150 .unwrap()
2151 });
2152 assert_eq!(
2153 messages(&conversation, cx),
2154 vec![
2155 (message_1.id, Role::User, 0..2),
2156 (message_2.id, Role::Assistant, 2..4),
2157 (message_4.id, Role::User, 4..5),
2158 (message_3.id, Role::User, 5..5),
2159 ]
2160 );
2161
2162 buffer.update(cx, |buffer, cx| {
2163 buffer.edit([(4..4, "C"), (5..5, "D")], None, cx)
2164 });
2165 assert_eq!(
2166 messages(&conversation, cx),
2167 vec![
2168 (message_1.id, Role::User, 0..2),
2169 (message_2.id, Role::Assistant, 2..4),
2170 (message_4.id, Role::User, 4..6),
2171 (message_3.id, Role::User, 6..7),
2172 ]
2173 );
2174
2175 // Deleting across message boundaries merges the messages.
2176 buffer.update(cx, |buffer, cx| buffer.edit([(1..4, "")], None, cx));
2177 assert_eq!(
2178 messages(&conversation, cx),
2179 vec![
2180 (message_1.id, Role::User, 0..3),
2181 (message_3.id, Role::User, 3..4),
2182 ]
2183 );
2184
2185 // Undoing the deletion should also undo the merge.
2186 buffer.update(cx, |buffer, cx| buffer.undo(cx));
2187 assert_eq!(
2188 messages(&conversation, cx),
2189 vec![
2190 (message_1.id, Role::User, 0..2),
2191 (message_2.id, Role::Assistant, 2..4),
2192 (message_4.id, Role::User, 4..6),
2193 (message_3.id, Role::User, 6..7),
2194 ]
2195 );
2196
2197 // Redoing the deletion should also redo the merge.
2198 buffer.update(cx, |buffer, cx| buffer.redo(cx));
2199 assert_eq!(
2200 messages(&conversation, cx),
2201 vec![
2202 (message_1.id, Role::User, 0..3),
2203 (message_3.id, Role::User, 3..4),
2204 ]
2205 );
2206
2207 // Ensure we can still insert after a merged message.
2208 let message_5 = conversation.update(cx, |conversation, cx| {
2209 conversation
2210 .insert_message_after(message_1.id, Role::System, MessageStatus::Done, cx)
2211 .unwrap()
2212 });
2213 assert_eq!(
2214 messages(&conversation, cx),
2215 vec![
2216 (message_1.id, Role::User, 0..3),
2217 (message_5.id, Role::System, 3..4),
2218 (message_3.id, Role::User, 4..5)
2219 ]
2220 );
2221 }
2222
2223 #[gpui::test]
2224 fn test_message_splitting(cx: &mut AppContext) {
2225 let registry = Arc::new(LanguageRegistry::test());
2226 let conversation = cx.add_model(|cx| Conversation::new(Default::default(), registry, cx));
2227 let buffer = conversation.read(cx).buffer.clone();
2228
2229 let message_1 = conversation.read(cx).message_anchors[0].clone();
2230 assert_eq!(
2231 messages(&conversation, cx),
2232 vec![(message_1.id, Role::User, 0..0)]
2233 );
2234
2235 buffer.update(cx, |buffer, cx| {
2236 buffer.edit([(0..0, "aaa\nbbb\nccc\nddd\n")], None, cx)
2237 });
2238
2239 let (_, message_2) =
2240 conversation.update(cx, |conversation, cx| conversation.split_message(3..3, cx));
2241 let message_2 = message_2.unwrap();
2242
2243 // We recycle newlines in the middle of a split message
2244 assert_eq!(buffer.read(cx).text(), "aaa\nbbb\nccc\nddd\n");
2245 assert_eq!(
2246 messages(&conversation, cx),
2247 vec![
2248 (message_1.id, Role::User, 0..4),
2249 (message_2.id, Role::User, 4..16),
2250 ]
2251 );
2252
2253 let (_, message_3) =
2254 conversation.update(cx, |conversation, cx| conversation.split_message(3..3, cx));
2255 let message_3 = message_3.unwrap();
2256
2257 // We don't recycle newlines at the end of a split message
2258 assert_eq!(buffer.read(cx).text(), "aaa\n\nbbb\nccc\nddd\n");
2259 assert_eq!(
2260 messages(&conversation, cx),
2261 vec![
2262 (message_1.id, Role::User, 0..4),
2263 (message_3.id, Role::User, 4..5),
2264 (message_2.id, Role::User, 5..17),
2265 ]
2266 );
2267
2268 let (_, message_4) =
2269 conversation.update(cx, |conversation, cx| conversation.split_message(9..9, cx));
2270 let message_4 = message_4.unwrap();
2271 assert_eq!(buffer.read(cx).text(), "aaa\n\nbbb\nccc\nddd\n");
2272 assert_eq!(
2273 messages(&conversation, cx),
2274 vec![
2275 (message_1.id, Role::User, 0..4),
2276 (message_3.id, Role::User, 4..5),
2277 (message_2.id, Role::User, 5..9),
2278 (message_4.id, Role::User, 9..17),
2279 ]
2280 );
2281
2282 let (_, message_5) =
2283 conversation.update(cx, |conversation, cx| conversation.split_message(9..9, cx));
2284 let message_5 = message_5.unwrap();
2285 assert_eq!(buffer.read(cx).text(), "aaa\n\nbbb\n\nccc\nddd\n");
2286 assert_eq!(
2287 messages(&conversation, cx),
2288 vec![
2289 (message_1.id, Role::User, 0..4),
2290 (message_3.id, Role::User, 4..5),
2291 (message_2.id, Role::User, 5..9),
2292 (message_4.id, Role::User, 9..10),
2293 (message_5.id, Role::User, 10..18),
2294 ]
2295 );
2296
2297 let (message_6, message_7) = conversation.update(cx, |conversation, cx| {
2298 conversation.split_message(14..16, cx)
2299 });
2300 let message_6 = message_6.unwrap();
2301 let message_7 = message_7.unwrap();
2302 assert_eq!(buffer.read(cx).text(), "aaa\n\nbbb\n\nccc\ndd\nd\n");
2303 assert_eq!(
2304 messages(&conversation, cx),
2305 vec![
2306 (message_1.id, Role::User, 0..4),
2307 (message_3.id, Role::User, 4..5),
2308 (message_2.id, Role::User, 5..9),
2309 (message_4.id, Role::User, 9..10),
2310 (message_5.id, Role::User, 10..14),
2311 (message_6.id, Role::User, 14..17),
2312 (message_7.id, Role::User, 17..19),
2313 ]
2314 );
2315 }
2316
2317 #[gpui::test]
2318 fn test_messages_for_offsets(cx: &mut AppContext) {
2319 let registry = Arc::new(LanguageRegistry::test());
2320 let conversation = cx.add_model(|cx| Conversation::new(Default::default(), registry, cx));
2321 let buffer = conversation.read(cx).buffer.clone();
2322
2323 let message_1 = conversation.read(cx).message_anchors[0].clone();
2324 assert_eq!(
2325 messages(&conversation, cx),
2326 vec![(message_1.id, Role::User, 0..0)]
2327 );
2328
2329 buffer.update(cx, |buffer, cx| buffer.edit([(0..0, "aaa")], None, cx));
2330 let message_2 = conversation
2331 .update(cx, |conversation, cx| {
2332 conversation.insert_message_after(message_1.id, Role::User, MessageStatus::Done, cx)
2333 })
2334 .unwrap();
2335 buffer.update(cx, |buffer, cx| buffer.edit([(4..4, "bbb")], None, cx));
2336
2337 let message_3 = conversation
2338 .update(cx, |conversation, cx| {
2339 conversation.insert_message_after(message_2.id, Role::User, MessageStatus::Done, cx)
2340 })
2341 .unwrap();
2342 buffer.update(cx, |buffer, cx| buffer.edit([(8..8, "ccc")], None, cx));
2343
2344 assert_eq!(buffer.read(cx).text(), "aaa\nbbb\nccc");
2345 assert_eq!(
2346 messages(&conversation, cx),
2347 vec![
2348 (message_1.id, Role::User, 0..4),
2349 (message_2.id, Role::User, 4..8),
2350 (message_3.id, Role::User, 8..11)
2351 ]
2352 );
2353
2354 assert_eq!(
2355 message_ids_for_offsets(&conversation, &[0, 4, 9], cx),
2356 [message_1.id, message_2.id, message_3.id]
2357 );
2358 assert_eq!(
2359 message_ids_for_offsets(&conversation, &[0, 1, 11], cx),
2360 [message_1.id, message_3.id]
2361 );
2362
2363 let message_4 = conversation
2364 .update(cx, |conversation, cx| {
2365 conversation.insert_message_after(message_3.id, Role::User, MessageStatus::Done, cx)
2366 })
2367 .unwrap();
2368 assert_eq!(buffer.read(cx).text(), "aaa\nbbb\nccc\n");
2369 assert_eq!(
2370 messages(&conversation, cx),
2371 vec![
2372 (message_1.id, Role::User, 0..4),
2373 (message_2.id, Role::User, 4..8),
2374 (message_3.id, Role::User, 8..12),
2375 (message_4.id, Role::User, 12..12)
2376 ]
2377 );
2378 assert_eq!(
2379 message_ids_for_offsets(&conversation, &[0, 4, 8, 12], cx),
2380 [message_1.id, message_2.id, message_3.id, message_4.id]
2381 );
2382
2383 fn message_ids_for_offsets(
2384 conversation: &ModelHandle<Conversation>,
2385 offsets: &[usize],
2386 cx: &AppContext,
2387 ) -> Vec<MessageId> {
2388 conversation
2389 .read(cx)
2390 .messages_for_offsets(offsets.iter().copied(), cx)
2391 .into_iter()
2392 .map(|message| message.id)
2393 .collect()
2394 }
2395 }
2396
2397 fn messages(
2398 conversation: &ModelHandle<Conversation>,
2399 cx: &AppContext,
2400 ) -> Vec<(MessageId, Role, Range<usize>)> {
2401 conversation
2402 .read(cx)
2403 .messages(cx)
2404 .map(|message| (message.id, message.role, message.range))
2405 .collect()
2406 }
2407}