1use crate::{
2 assistant_settings::{AssistantDockPosition, AssistantSettings},
3 OpenAIRequest, OpenAIResponseStreamEvent, RequestMessage, Role,
4};
5use anyhow::{anyhow, Result};
6use chrono::{DateTime, Local};
7use collections::{HashMap, HashSet};
8use editor::{
9 display_map::ToDisplayPoint,
10 scroll::{
11 autoscroll::{Autoscroll, AutoscrollStrategy},
12 ScrollAnchor,
13 },
14 Anchor, DisplayPoint, Editor, ExcerptId, ExcerptRange, MultiBuffer,
15};
16use fs::Fs;
17use futures::{io::BufReader, AsyncBufReadExt, AsyncReadExt, Stream, StreamExt};
18use gpui::{
19 actions,
20 elements::*,
21 executor::Background,
22 geometry::vector::vec2f,
23 platform::{CursorStyle, MouseButton},
24 Action, AppContext, AsyncAppContext, ClipboardItem, Entity, ModelContext, ModelHandle,
25 Subscription, Task, View, ViewContext, ViewHandle, WeakViewHandle, WindowContext,
26};
27use isahc::{http::StatusCode, Request, RequestExt};
28use language::{language_settings::SoftWrap, Buffer, LanguageRegistry};
29use serde::Deserialize;
30use settings::SettingsStore;
31use std::{borrow::Cow, cell::RefCell, cmp, fmt::Write, io, rc::Rc, sync::Arc, time::Duration};
32use util::{post_inc, truncate_and_trailoff, ResultExt, TryFutureExt};
33use workspace::{
34 dock::{DockPosition, Panel},
35 item::Item,
36 pane, Pane, Workspace,
37};
38
39const OPENAI_API_URL: &'static str = "https://api.openai.com/v1";
40
41actions!(
42 assistant,
43 [NewContext, Assist, QuoteSelection, ToggleFocus, ResetKey]
44);
45
46pub fn init(cx: &mut AppContext) {
47 settings::register::<AssistantSettings>(cx);
48 cx.add_action(
49 |workspace: &mut Workspace, _: &NewContext, cx: &mut ViewContext<Workspace>| {
50 if let Some(this) = workspace.panel::<AssistantPanel>(cx) {
51 this.update(cx, |this, cx| this.add_context(cx))
52 }
53
54 workspace.focus_panel::<AssistantPanel>(cx);
55 },
56 );
57 cx.add_action(AssistantEditor::assist);
58 cx.capture_action(AssistantEditor::cancel_last_assist);
59 cx.add_action(AssistantEditor::quote_selection);
60 cx.capture_action(AssistantEditor::copy);
61 cx.add_action(AssistantPanel::save_api_key);
62 cx.add_action(AssistantPanel::reset_api_key);
63}
64
65pub enum AssistantPanelEvent {
66 ZoomIn,
67 ZoomOut,
68 Focus,
69 Close,
70 DockPositionChanged,
71}
72
73pub struct AssistantPanel {
74 width: Option<f32>,
75 height: Option<f32>,
76 pane: ViewHandle<Pane>,
77 api_key: Rc<RefCell<Option<String>>>,
78 api_key_editor: Option<ViewHandle<Editor>>,
79 has_read_credentials: bool,
80 languages: Arc<LanguageRegistry>,
81 fs: Arc<dyn Fs>,
82 subscriptions: Vec<Subscription>,
83}
84
85impl AssistantPanel {
86 pub fn load(
87 workspace: WeakViewHandle<Workspace>,
88 cx: AsyncAppContext,
89 ) -> Task<Result<ViewHandle<Self>>> {
90 cx.spawn(|mut cx| async move {
91 // TODO: deserialize state.
92 workspace.update(&mut cx, |workspace, cx| {
93 cx.add_view::<Self, _>(|cx| {
94 let weak_self = cx.weak_handle();
95 let pane = cx.add_view(|cx| {
96 let mut pane = Pane::new(
97 workspace.weak_handle(),
98 workspace.project().clone(),
99 workspace.app_state().background_actions,
100 Default::default(),
101 cx,
102 );
103 pane.set_can_split(false, cx);
104 pane.set_can_navigate(false, cx);
105 pane.on_can_drop(move |_, _| false);
106 pane.set_render_tab_bar_buttons(cx, move |pane, cx| {
107 let weak_self = weak_self.clone();
108 Flex::row()
109 .with_child(Pane::render_tab_bar_button(
110 0,
111 "icons/plus_12.svg",
112 false,
113 Some(("New Context".into(), Some(Box::new(NewContext)))),
114 cx,
115 move |_, cx| {
116 let weak_self = weak_self.clone();
117 cx.window_context().defer(move |cx| {
118 if let Some(this) = weak_self.upgrade(cx) {
119 this.update(cx, |this, cx| this.add_context(cx));
120 }
121 })
122 },
123 None,
124 ))
125 .with_child(Pane::render_tab_bar_button(
126 1,
127 if pane.is_zoomed() {
128 "icons/minimize_8.svg"
129 } else {
130 "icons/maximize_8.svg"
131 },
132 pane.is_zoomed(),
133 Some((
134 "Toggle Zoom".into(),
135 Some(Box::new(workspace::ToggleZoom)),
136 )),
137 cx,
138 move |pane, cx| pane.toggle_zoom(&Default::default(), cx),
139 None,
140 ))
141 .into_any()
142 });
143 let buffer_search_bar = cx.add_view(search::BufferSearchBar::new);
144 pane.toolbar()
145 .update(cx, |toolbar, cx| toolbar.add_item(buffer_search_bar, cx));
146 pane
147 });
148
149 let mut this = Self {
150 pane,
151 api_key: Rc::new(RefCell::new(None)),
152 api_key_editor: None,
153 has_read_credentials: false,
154 languages: workspace.app_state().languages.clone(),
155 fs: workspace.app_state().fs.clone(),
156 width: None,
157 height: None,
158 subscriptions: Default::default(),
159 };
160
161 let mut old_dock_position = this.position(cx);
162 this.subscriptions = vec![
163 cx.observe(&this.pane, |_, _, cx| cx.notify()),
164 cx.subscribe(&this.pane, Self::handle_pane_event),
165 cx.observe_global::<SettingsStore, _>(move |this, cx| {
166 let new_dock_position = this.position(cx);
167 if new_dock_position != old_dock_position {
168 old_dock_position = new_dock_position;
169 cx.emit(AssistantPanelEvent::DockPositionChanged);
170 }
171 }),
172 ];
173
174 this
175 })
176 })
177 })
178 }
179
180 fn handle_pane_event(
181 &mut self,
182 _pane: ViewHandle<Pane>,
183 event: &pane::Event,
184 cx: &mut ViewContext<Self>,
185 ) {
186 match event {
187 pane::Event::ZoomIn => cx.emit(AssistantPanelEvent::ZoomIn),
188 pane::Event::ZoomOut => cx.emit(AssistantPanelEvent::ZoomOut),
189 pane::Event::Focus => cx.emit(AssistantPanelEvent::Focus),
190 pane::Event::Remove => cx.emit(AssistantPanelEvent::Close),
191 _ => {}
192 }
193 }
194
195 fn add_context(&mut self, cx: &mut ViewContext<Self>) {
196 let focus = self.has_focus(cx);
197 let editor = cx
198 .add_view(|cx| AssistantEditor::new(self.api_key.clone(), self.languages.clone(), cx));
199 self.subscriptions
200 .push(cx.subscribe(&editor, Self::handle_assistant_editor_event));
201 self.pane.update(cx, |pane, cx| {
202 pane.add_item(Box::new(editor), true, focus, None, cx)
203 });
204 }
205
206 fn handle_assistant_editor_event(
207 &mut self,
208 _: ViewHandle<AssistantEditor>,
209 event: &AssistantEditorEvent,
210 cx: &mut ViewContext<Self>,
211 ) {
212 match event {
213 AssistantEditorEvent::TabContentChanged => self.pane.update(cx, |_, cx| cx.notify()),
214 }
215 }
216
217 fn save_api_key(&mut self, _: &menu::Confirm, cx: &mut ViewContext<Self>) {
218 if let Some(api_key) = self
219 .api_key_editor
220 .as_ref()
221 .map(|editor| editor.read(cx).text(cx))
222 {
223 if !api_key.is_empty() {
224 cx.platform()
225 .write_credentials(OPENAI_API_URL, "Bearer", api_key.as_bytes())
226 .log_err();
227 *self.api_key.borrow_mut() = Some(api_key);
228 self.api_key_editor.take();
229 cx.focus_self();
230 cx.notify();
231 }
232 } else {
233 cx.propagate_action();
234 }
235 }
236
237 fn reset_api_key(&mut self, _: &ResetKey, cx: &mut ViewContext<Self>) {
238 cx.platform().delete_credentials(OPENAI_API_URL).log_err();
239 self.api_key.take();
240 self.api_key_editor = Some(build_api_key_editor(cx));
241 cx.focus_self();
242 cx.notify();
243 }
244}
245
246fn build_api_key_editor(cx: &mut ViewContext<AssistantPanel>) -> ViewHandle<Editor> {
247 cx.add_view(|cx| {
248 let mut editor = Editor::single_line(
249 Some(Arc::new(|theme| theme.assistant.api_key_editor.clone())),
250 cx,
251 );
252 editor.set_placeholder_text("sk-000000000000000000000000000000000000000000000000", cx);
253 editor
254 })
255}
256
257impl Entity for AssistantPanel {
258 type Event = AssistantPanelEvent;
259}
260
261impl View for AssistantPanel {
262 fn ui_name() -> &'static str {
263 "AssistantPanel"
264 }
265
266 fn render(&mut self, cx: &mut ViewContext<Self>) -> AnyElement<Self> {
267 let style = &theme::current(cx).assistant;
268 if let Some(api_key_editor) = self.api_key_editor.as_ref() {
269 Flex::column()
270 .with_child(
271 Text::new(
272 "Paste your OpenAI API key and press Enter to use the assistant",
273 style.api_key_prompt.text.clone(),
274 )
275 .aligned(),
276 )
277 .with_child(
278 ChildView::new(api_key_editor, cx)
279 .contained()
280 .with_style(style.api_key_editor.container)
281 .aligned(),
282 )
283 .contained()
284 .with_style(style.api_key_prompt.container)
285 .aligned()
286 .into_any()
287 } else {
288 ChildView::new(&self.pane, cx).into_any()
289 }
290 }
291
292 fn focus_in(&mut self, _: gpui::AnyViewHandle, cx: &mut ViewContext<Self>) {
293 if cx.is_self_focused() {
294 if let Some(api_key_editor) = self.api_key_editor.as_ref() {
295 cx.focus(api_key_editor);
296 } else {
297 cx.focus(&self.pane);
298 }
299 }
300 }
301}
302
303impl Panel for AssistantPanel {
304 fn position(&self, cx: &WindowContext) -> DockPosition {
305 match settings::get::<AssistantSettings>(cx).dock {
306 AssistantDockPosition::Left => DockPosition::Left,
307 AssistantDockPosition::Bottom => DockPosition::Bottom,
308 AssistantDockPosition::Right => DockPosition::Right,
309 }
310 }
311
312 fn position_is_valid(&self, _: DockPosition) -> bool {
313 true
314 }
315
316 fn set_position(&mut self, position: DockPosition, cx: &mut ViewContext<Self>) {
317 settings::update_settings_file::<AssistantSettings>(self.fs.clone(), cx, move |settings| {
318 let dock = match position {
319 DockPosition::Left => AssistantDockPosition::Left,
320 DockPosition::Bottom => AssistantDockPosition::Bottom,
321 DockPosition::Right => AssistantDockPosition::Right,
322 };
323 settings.dock = Some(dock);
324 });
325 }
326
327 fn size(&self, cx: &WindowContext) -> f32 {
328 let settings = settings::get::<AssistantSettings>(cx);
329 match self.position(cx) {
330 DockPosition::Left | DockPosition::Right => {
331 self.width.unwrap_or_else(|| settings.default_width)
332 }
333 DockPosition::Bottom => self.height.unwrap_or_else(|| settings.default_height),
334 }
335 }
336
337 fn set_size(&mut self, size: f32, cx: &mut ViewContext<Self>) {
338 match self.position(cx) {
339 DockPosition::Left | DockPosition::Right => self.width = Some(size),
340 DockPosition::Bottom => self.height = Some(size),
341 }
342 cx.notify();
343 }
344
345 fn should_zoom_in_on_event(event: &AssistantPanelEvent) -> bool {
346 matches!(event, AssistantPanelEvent::ZoomIn)
347 }
348
349 fn should_zoom_out_on_event(event: &AssistantPanelEvent) -> bool {
350 matches!(event, AssistantPanelEvent::ZoomOut)
351 }
352
353 fn is_zoomed(&self, cx: &WindowContext) -> bool {
354 self.pane.read(cx).is_zoomed()
355 }
356
357 fn set_zoomed(&mut self, zoomed: bool, cx: &mut ViewContext<Self>) {
358 self.pane.update(cx, |pane, cx| pane.set_zoomed(zoomed, cx));
359 }
360
361 fn set_active(&mut self, active: bool, cx: &mut ViewContext<Self>) {
362 if active {
363 if self.api_key.borrow().is_none() && !self.has_read_credentials {
364 self.has_read_credentials = true;
365 let api_key = if let Some((_, api_key)) = cx
366 .platform()
367 .read_credentials(OPENAI_API_URL)
368 .log_err()
369 .flatten()
370 {
371 String::from_utf8(api_key).log_err()
372 } else {
373 None
374 };
375 if let Some(api_key) = api_key {
376 *self.api_key.borrow_mut() = Some(api_key);
377 } else if self.api_key_editor.is_none() {
378 self.api_key_editor = Some(build_api_key_editor(cx));
379 cx.notify();
380 }
381 }
382
383 if self.pane.read(cx).items_len() == 0 {
384 self.add_context(cx);
385 }
386 }
387 }
388
389 fn icon_path(&self) -> &'static str {
390 "icons/speech_bubble_12.svg"
391 }
392
393 fn icon_tooltip(&self) -> (String, Option<Box<dyn Action>>) {
394 ("Assistant Panel".into(), Some(Box::new(ToggleFocus)))
395 }
396
397 fn should_change_position_on_event(event: &Self::Event) -> bool {
398 matches!(event, AssistantPanelEvent::DockPositionChanged)
399 }
400
401 fn should_activate_on_event(_: &Self::Event) -> bool {
402 false
403 }
404
405 fn should_close_on_event(event: &AssistantPanelEvent) -> bool {
406 matches!(event, AssistantPanelEvent::Close)
407 }
408
409 fn has_focus(&self, cx: &WindowContext) -> bool {
410 self.pane.read(cx).has_focus()
411 || self
412 .api_key_editor
413 .as_ref()
414 .map_or(false, |editor| editor.is_focused(cx))
415 }
416
417 fn is_focus_event(event: &Self::Event) -> bool {
418 matches!(event, AssistantPanelEvent::Focus)
419 }
420}
421
422enum AssistantEvent {
423 MessagesEdited { ids: Vec<ExcerptId> },
424 SummaryChanged,
425 StreamedCompletion,
426}
427
428struct Assistant {
429 buffer: ModelHandle<MultiBuffer>,
430 messages: Vec<Message>,
431 messages_metadata: HashMap<ExcerptId, MessageMetadata>,
432 summary: Option<String>,
433 pending_summary: Task<Option<()>>,
434 completion_count: usize,
435 pending_completions: Vec<PendingCompletion>,
436 languages: Arc<LanguageRegistry>,
437 model: String,
438 token_count: Option<usize>,
439 max_token_count: usize,
440 pending_token_count: Task<Option<()>>,
441 api_key: Rc<RefCell<Option<String>>>,
442 _subscriptions: Vec<Subscription>,
443}
444
445impl Entity for Assistant {
446 type Event = AssistantEvent;
447}
448
449impl Assistant {
450 fn new(
451 api_key: Rc<RefCell<Option<String>>>,
452 language_registry: Arc<LanguageRegistry>,
453 cx: &mut ModelContext<Self>,
454 ) -> Self {
455 let model = "gpt-3.5-turbo";
456 let buffer = cx.add_model(|_| MultiBuffer::new(0));
457 let mut this = Self {
458 messages: Default::default(),
459 messages_metadata: Default::default(),
460 summary: None,
461 pending_summary: Task::ready(None),
462 completion_count: Default::default(),
463 pending_completions: Default::default(),
464 languages: language_registry,
465 token_count: None,
466 max_token_count: tiktoken_rs::model::get_context_size(model),
467 pending_token_count: Task::ready(None),
468 model: model.into(),
469 _subscriptions: vec![cx.subscribe(&buffer, Self::handle_buffer_event)],
470 api_key,
471 buffer,
472 };
473 this.insert_message_after(ExcerptId::max(), Role::User, cx);
474 this.count_remaining_tokens(cx);
475 this
476 }
477
478 fn handle_buffer_event(
479 &mut self,
480 _: ModelHandle<MultiBuffer>,
481 event: &editor::multi_buffer::Event,
482 cx: &mut ModelContext<Self>,
483 ) {
484 match event {
485 editor::multi_buffer::Event::ExcerptsAdded { .. }
486 | editor::multi_buffer::Event::ExcerptsRemoved { .. }
487 | editor::multi_buffer::Event::Edited => self.count_remaining_tokens(cx),
488 editor::multi_buffer::Event::ExcerptsEdited { ids } => {
489 cx.emit(AssistantEvent::MessagesEdited { ids: ids.clone() });
490 }
491 _ => {}
492 }
493 }
494
495 fn count_remaining_tokens(&mut self, cx: &mut ModelContext<Self>) {
496 let messages = self
497 .messages
498 .iter()
499 .filter_map(|message| {
500 Some(tiktoken_rs::ChatCompletionRequestMessage {
501 role: match self.messages_metadata.get(&message.excerpt_id)?.role {
502 Role::User => "user".into(),
503 Role::Assistant => "assistant".into(),
504 Role::System => "system".into(),
505 },
506 content: message.content.read(cx).text(),
507 name: None,
508 })
509 })
510 .collect::<Vec<_>>();
511 let model = self.model.clone();
512 self.pending_token_count = cx.spawn_weak(|this, mut cx| {
513 async move {
514 cx.background().timer(Duration::from_millis(200)).await;
515 let token_count = cx
516 .background()
517 .spawn(async move { tiktoken_rs::num_tokens_from_messages(&model, &messages) })
518 .await?;
519
520 this.upgrade(&cx)
521 .ok_or_else(|| anyhow!("assistant was dropped"))?
522 .update(&mut cx, |this, cx| {
523 this.max_token_count = tiktoken_rs::model::get_context_size(&this.model);
524 this.token_count = Some(token_count);
525 cx.notify()
526 });
527 anyhow::Ok(())
528 }
529 .log_err()
530 });
531 }
532
533 fn remaining_tokens(&self) -> Option<isize> {
534 Some(self.max_token_count as isize - self.token_count? as isize)
535 }
536
537 fn set_model(&mut self, model: String, cx: &mut ModelContext<Self>) {
538 self.model = model;
539 self.count_remaining_tokens(cx);
540 cx.notify();
541 }
542
543 fn assist(&mut self, cx: &mut ModelContext<Self>) -> Option<(Message, Message)> {
544 let messages = self
545 .messages
546 .iter()
547 .filter_map(|message| {
548 Some(RequestMessage {
549 role: self.messages_metadata.get(&message.excerpt_id)?.role,
550 content: message.content.read(cx).text(),
551 })
552 })
553 .collect();
554 let request = OpenAIRequest {
555 model: self.model.clone(),
556 messages,
557 stream: true,
558 };
559
560 let api_key = self.api_key.borrow().clone()?;
561 let stream = stream_completion(api_key, cx.background().clone(), request);
562 let assistant_message = self.insert_message_after(ExcerptId::max(), Role::Assistant, cx);
563 let user_message = self.insert_message_after(ExcerptId::max(), Role::User, cx);
564 let task = cx.spawn_weak({
565 let assistant_message = assistant_message.clone();
566 |this, mut cx| async move {
567 let assistant_message = assistant_message;
568 let stream_completion = async {
569 let mut messages = stream.await?;
570
571 while let Some(message) = messages.next().await {
572 let mut message = message?;
573 if let Some(choice) = message.choices.pop() {
574 assistant_message.content.update(&mut cx, |content, cx| {
575 let text: Arc<str> = choice.delta.content?.into();
576 content.edit([(content.len()..content.len(), text)], None, cx);
577 Some(())
578 });
579 this.upgrade(&cx)
580 .ok_or_else(|| anyhow!("assistant was dropped"))?
581 .update(&mut cx, |_, cx| {
582 cx.emit(AssistantEvent::StreamedCompletion);
583 });
584 }
585 }
586
587 this.upgrade(&cx)
588 .ok_or_else(|| anyhow!("assistant was dropped"))?
589 .update(&mut cx, |this, cx| {
590 this.pending_completions
591 .retain(|completion| completion.id != this.completion_count);
592 this.summarize(cx);
593 });
594
595 anyhow::Ok(())
596 };
597
598 let result = stream_completion.await;
599 if let Some(this) = this.upgrade(&cx) {
600 this.update(&mut cx, |this, cx| {
601 if let Err(error) = result {
602 if let Some(metadata) = this
603 .messages_metadata
604 .get_mut(&assistant_message.excerpt_id)
605 {
606 metadata.error = Some(error.to_string().trim().into());
607 cx.notify();
608 }
609 }
610 });
611 }
612 }
613 });
614
615 self.pending_completions.push(PendingCompletion {
616 id: post_inc(&mut self.completion_count),
617 _task: task,
618 });
619 Some((assistant_message, user_message))
620 }
621
622 fn cancel_last_assist(&mut self) -> bool {
623 self.pending_completions.pop().is_some()
624 }
625
626 fn remove_empty_messages<'a>(
627 &mut self,
628 excerpts: HashSet<ExcerptId>,
629 protected_offsets: HashSet<usize>,
630 cx: &mut ModelContext<Self>,
631 ) {
632 let mut offset = 0;
633 let mut excerpts_to_remove = Vec::new();
634 self.messages.retain(|message| {
635 let range = offset..offset + message.content.read(cx).len();
636 offset = range.end + 1;
637 if range.is_empty()
638 && !protected_offsets.contains(&range.start)
639 && excerpts.contains(&message.excerpt_id)
640 {
641 excerpts_to_remove.push(message.excerpt_id);
642 self.messages_metadata.remove(&message.excerpt_id);
643 false
644 } else {
645 true
646 }
647 });
648
649 if !excerpts_to_remove.is_empty() {
650 self.buffer.update(cx, |buffer, cx| {
651 buffer.remove_excerpts(excerpts_to_remove, cx)
652 });
653 cx.notify();
654 }
655 }
656
657 fn cycle_message_role(&mut self, excerpt_id: ExcerptId, cx: &mut ModelContext<Self>) {
658 if let Some(metadata) = self.messages_metadata.get_mut(&excerpt_id) {
659 metadata.role.cycle();
660 cx.notify();
661 }
662 }
663
664 fn insert_message_after(
665 &mut self,
666 excerpt_id: ExcerptId,
667 role: Role,
668 cx: &mut ModelContext<Self>,
669 ) -> Message {
670 let content = cx.add_model(|cx| {
671 let mut buffer = Buffer::new(0, "", cx);
672 let markdown = self.languages.language_for_name("Markdown");
673 cx.spawn_weak(|buffer, mut cx| async move {
674 let markdown = markdown.await?;
675 let buffer = buffer
676 .upgrade(&cx)
677 .ok_or_else(|| anyhow!("buffer was dropped"))?;
678 buffer.update(&mut cx, |buffer, cx| {
679 buffer.set_language(Some(markdown), cx)
680 });
681 anyhow::Ok(())
682 })
683 .detach_and_log_err(cx);
684 buffer.set_language_registry(self.languages.clone());
685 buffer
686 });
687 let new_excerpt_id = self.buffer.update(cx, |buffer, cx| {
688 buffer
689 .insert_excerpts_after(
690 excerpt_id,
691 content.clone(),
692 vec![ExcerptRange {
693 context: 0..0,
694 primary: None,
695 }],
696 cx,
697 )
698 .pop()
699 .unwrap()
700 });
701
702 let ix = self
703 .messages
704 .iter()
705 .position(|message| message.excerpt_id == excerpt_id)
706 .map_or(self.messages.len(), |ix| ix + 1);
707 let message = Message {
708 excerpt_id: new_excerpt_id,
709 content: content.clone(),
710 };
711 self.messages.insert(ix, message.clone());
712 self.messages_metadata.insert(
713 new_excerpt_id,
714 MessageMetadata {
715 role,
716 sent_at: Local::now(),
717 error: None,
718 },
719 );
720 message
721 }
722
723 fn summarize(&mut self, cx: &mut ModelContext<Self>) {
724 if self.messages.len() >= 2 && self.summary.is_none() {
725 let api_key = self.api_key.borrow().clone();
726 if let Some(api_key) = api_key {
727 let messages = self
728 .messages
729 .iter()
730 .take(2)
731 .filter_map(|message| {
732 Some(RequestMessage {
733 role: self.messages_metadata.get(&message.excerpt_id)?.role,
734 content: message.content.read(cx).text(),
735 })
736 })
737 .chain(Some(RequestMessage {
738 role: Role::User,
739 content:
740 "Summarize the conversation into a short title without punctuation"
741 .into(),
742 }))
743 .collect();
744 let request = OpenAIRequest {
745 model: self.model.clone(),
746 messages,
747 stream: true,
748 };
749
750 let stream = stream_completion(api_key, cx.background().clone(), request);
751 self.pending_summary = cx.spawn(|this, mut cx| {
752 async move {
753 let mut messages = stream.await?;
754
755 while let Some(message) = messages.next().await {
756 let mut message = message?;
757 if let Some(choice) = message.choices.pop() {
758 let text = choice.delta.content.unwrap_or_default();
759 this.update(&mut cx, |this, cx| {
760 this.summary.get_or_insert(String::new()).push_str(&text);
761 cx.emit(AssistantEvent::SummaryChanged);
762 });
763 }
764 }
765
766 anyhow::Ok(())
767 }
768 .log_err()
769 });
770 }
771 }
772 }
773}
774
775struct PendingCompletion {
776 id: usize,
777 _task: Task<()>,
778}
779
780enum AssistantEditorEvent {
781 TabContentChanged,
782}
783
784struct AssistantEditor {
785 assistant: ModelHandle<Assistant>,
786 editor: ViewHandle<Editor>,
787 scroll_bottom: ScrollAnchor,
788 _subscriptions: Vec<Subscription>,
789}
790
791impl AssistantEditor {
792 fn new(
793 api_key: Rc<RefCell<Option<String>>>,
794 language_registry: Arc<LanguageRegistry>,
795 cx: &mut ViewContext<Self>,
796 ) -> Self {
797 let assistant = cx.add_model(|cx| Assistant::new(api_key, language_registry, cx));
798 let editor = cx.add_view(|cx| {
799 let mut editor = Editor::for_multibuffer(assistant.read(cx).buffer.clone(), None, cx);
800 editor.set_soft_wrap_mode(SoftWrap::EditorWidth, cx);
801 editor.set_show_gutter(false, cx);
802 editor.set_render_excerpt_header(
803 {
804 let assistant = assistant.clone();
805 move |_editor, params: editor::RenderExcerptHeaderParams, cx| {
806 enum Sender {}
807 enum ErrorTooltip {}
808
809 let theme = theme::current(cx);
810 let style = &theme.assistant;
811 let excerpt_id = params.id;
812 if let Some(metadata) = assistant
813 .read(cx)
814 .messages_metadata
815 .get(&excerpt_id)
816 .cloned()
817 {
818 let sender = MouseEventHandler::<Sender, _>::new(
819 params.id.into(),
820 cx,
821 |state, _| match metadata.role {
822 Role::User => {
823 let style = style.user_sender.style_for(state, false);
824 Label::new("You", style.text.clone())
825 .contained()
826 .with_style(style.container)
827 }
828 Role::Assistant => {
829 let style = style.assistant_sender.style_for(state, false);
830 Label::new("Assistant", style.text.clone())
831 .contained()
832 .with_style(style.container)
833 }
834 Role::System => {
835 let style = style.system_sender.style_for(state, false);
836 Label::new("System", style.text.clone())
837 .contained()
838 .with_style(style.container)
839 }
840 },
841 )
842 .with_cursor_style(CursorStyle::PointingHand)
843 .on_down(MouseButton::Left, {
844 let assistant = assistant.clone();
845 move |_, _, cx| {
846 assistant.update(cx, |assistant, cx| {
847 assistant.cycle_message_role(excerpt_id, cx)
848 })
849 }
850 });
851
852 Flex::row()
853 .with_child(sender.aligned())
854 .with_child(
855 Label::new(
856 metadata.sent_at.format("%I:%M%P").to_string(),
857 style.sent_at.text.clone(),
858 )
859 .contained()
860 .with_style(style.sent_at.container)
861 .aligned(),
862 )
863 .with_children(metadata.error.map(|error| {
864 Svg::new("icons/circle_x_mark_12.svg")
865 .with_color(style.error_icon.color)
866 .constrained()
867 .with_width(style.error_icon.width)
868 .contained()
869 .with_style(style.error_icon.container)
870 .with_tooltip::<ErrorTooltip>(
871 params.id.into(),
872 error,
873 None,
874 theme.tooltip.clone(),
875 cx,
876 )
877 .aligned()
878 }))
879 .aligned()
880 .left()
881 .contained()
882 .with_style(style.header)
883 .into_any()
884 } else {
885 Empty::new().into_any()
886 }
887 }
888 },
889 cx,
890 );
891 editor
892 });
893
894 let _subscriptions = vec![
895 cx.observe(&assistant, |_, _, cx| cx.notify()),
896 cx.subscribe(&assistant, Self::handle_assistant_event),
897 cx.subscribe(&editor, Self::handle_editor_event),
898 ];
899
900 Self {
901 assistant,
902 editor,
903 scroll_bottom: ScrollAnchor {
904 offset: Default::default(),
905 anchor: Anchor::max(),
906 },
907 _subscriptions,
908 }
909 }
910
911 fn assist(&mut self, _: &Assist, cx: &mut ViewContext<Self>) {
912 let user_message = self.assistant.update(cx, |assistant, cx| {
913 let editor = self.editor.read(cx);
914 let newest_selection = editor.selections.newest_anchor();
915 let excerpt_id = if newest_selection.head() == Anchor::min() {
916 assistant
917 .messages
918 .first()
919 .map(|message| message.excerpt_id)?
920 } else if newest_selection.head() == Anchor::max() {
921 assistant
922 .messages
923 .last()
924 .map(|message| message.excerpt_id)?
925 } else {
926 newest_selection.head().excerpt_id()
927 };
928
929 let metadata = assistant.messages_metadata.get(&excerpt_id)?;
930 let user_message = if metadata.role == Role::User {
931 let (_, user_message) = assistant.assist(cx)?;
932 user_message
933 } else {
934 let user_message = assistant.insert_message_after(excerpt_id, Role::User, cx);
935 user_message
936 };
937 Some(user_message)
938 });
939
940 if let Some(user_message) = user_message {
941 self.editor.update(cx, |editor, cx| {
942 let cursor = editor
943 .buffer()
944 .read(cx)
945 .snapshot(cx)
946 .anchor_in_excerpt(user_message.excerpt_id, language::Anchor::MIN);
947 editor.change_selections(
948 Some(Autoscroll::Strategy(AutoscrollStrategy::Fit)),
949 cx,
950 |selections| selections.select_anchor_ranges([cursor..cursor]),
951 );
952 });
953 self.update_scroll_bottom(cx);
954 }
955 }
956
957 fn cancel_last_assist(&mut self, _: &editor::Cancel, cx: &mut ViewContext<Self>) {
958 if !self
959 .assistant
960 .update(cx, |assistant, _| assistant.cancel_last_assist())
961 {
962 cx.propagate_action();
963 }
964 }
965
966 fn handle_assistant_event(
967 &mut self,
968 _: ModelHandle<Assistant>,
969 event: &AssistantEvent,
970 cx: &mut ViewContext<Self>,
971 ) {
972 match event {
973 AssistantEvent::MessagesEdited { ids } => {
974 let selections = self.editor.read(cx).selections.all::<usize>(cx);
975 let selection_heads = selections
976 .iter()
977 .map(|selection| selection.head())
978 .collect::<HashSet<usize>>();
979 let ids = ids.iter().copied().collect::<HashSet<_>>();
980 self.assistant.update(cx, |assistant, cx| {
981 assistant.remove_empty_messages(ids, selection_heads, cx)
982 });
983 }
984 AssistantEvent::SummaryChanged => {
985 cx.emit(AssistantEditorEvent::TabContentChanged);
986 }
987 AssistantEvent::StreamedCompletion => {
988 self.editor.update(cx, |editor, cx| {
989 let snapshot = editor.snapshot(cx);
990 let scroll_bottom_row = self
991 .scroll_bottom
992 .anchor
993 .to_display_point(&snapshot.display_snapshot)
994 .row();
995
996 let scroll_bottom = scroll_bottom_row as f32 + self.scroll_bottom.offset.y();
997 let visible_line_count = editor.visible_line_count().unwrap_or(0.);
998 let scroll_top = scroll_bottom - visible_line_count;
999 editor
1000 .set_scroll_position(vec2f(self.scroll_bottom.offset.x(), scroll_top), cx);
1001 });
1002 }
1003 }
1004 }
1005
1006 fn handle_editor_event(
1007 &mut self,
1008 _: ViewHandle<Editor>,
1009 event: &editor::Event,
1010 cx: &mut ViewContext<Self>,
1011 ) {
1012 match event {
1013 editor::Event::ScrollPositionChanged { .. } => self.update_scroll_bottom(cx),
1014 _ => {}
1015 }
1016 }
1017
1018 fn update_scroll_bottom(&mut self, cx: &mut ViewContext<Self>) {
1019 self.editor.update(cx, |editor, cx| {
1020 let snapshot = editor.snapshot(cx);
1021 let scroll_position = editor
1022 .scroll_manager
1023 .anchor()
1024 .scroll_position(&snapshot.display_snapshot);
1025 let scroll_bottom = scroll_position.y() + editor.visible_line_count().unwrap_or(0.);
1026 let scroll_bottom_point = cmp::min(
1027 DisplayPoint::new(scroll_bottom.floor() as u32, 0),
1028 snapshot.display_snapshot.max_point(),
1029 );
1030 let scroll_bottom_anchor = snapshot
1031 .buffer_snapshot
1032 .anchor_after(scroll_bottom_point.to_point(&snapshot.display_snapshot));
1033 let scroll_bottom_offset = vec2f(
1034 scroll_position.x(),
1035 scroll_bottom - scroll_bottom_point.row() as f32,
1036 );
1037 self.scroll_bottom = ScrollAnchor {
1038 anchor: scroll_bottom_anchor,
1039 offset: scroll_bottom_offset,
1040 };
1041 });
1042 }
1043
1044 fn quote_selection(
1045 workspace: &mut Workspace,
1046 _: &QuoteSelection,
1047 cx: &mut ViewContext<Workspace>,
1048 ) {
1049 let Some(panel) = workspace.panel::<AssistantPanel>(cx) else {
1050 return;
1051 };
1052 let Some(editor) = workspace.active_item(cx).and_then(|item| item.downcast::<Editor>()) else {
1053 return;
1054 };
1055
1056 let text = editor.read_with(cx, |editor, cx| {
1057 let range = editor.selections.newest::<usize>(cx).range();
1058 let buffer = editor.buffer().read(cx).snapshot(cx);
1059 let start_language = buffer.language_at(range.start);
1060 let end_language = buffer.language_at(range.end);
1061 let language_name = if start_language == end_language {
1062 start_language.map(|language| language.name())
1063 } else {
1064 None
1065 };
1066 let language_name = language_name.as_deref().unwrap_or("").to_lowercase();
1067
1068 let selected_text = buffer.text_for_range(range).collect::<String>();
1069 if selected_text.is_empty() {
1070 None
1071 } else {
1072 Some(if language_name == "markdown" {
1073 selected_text
1074 .lines()
1075 .map(|line| format!("> {}", line))
1076 .collect::<Vec<_>>()
1077 .join("\n")
1078 } else {
1079 format!("```{language_name}\n{selected_text}\n```")
1080 })
1081 }
1082 });
1083
1084 // Activate the panel
1085 if !panel.read(cx).has_focus(cx) {
1086 workspace.toggle_panel_focus::<AssistantPanel>(cx);
1087 }
1088
1089 if let Some(text) = text {
1090 panel.update(cx, |panel, cx| {
1091 if let Some(assistant) = panel
1092 .pane
1093 .read(cx)
1094 .active_item()
1095 .and_then(|item| item.downcast::<AssistantEditor>())
1096 .ok_or_else(|| anyhow!("no active context"))
1097 .log_err()
1098 {
1099 assistant.update(cx, |assistant, cx| {
1100 assistant
1101 .editor
1102 .update(cx, |editor, cx| editor.insert(&text, cx))
1103 });
1104 }
1105 });
1106 }
1107 }
1108
1109 fn copy(&mut self, _: &editor::Copy, cx: &mut ViewContext<Self>) {
1110 let editor = self.editor.read(cx);
1111 let assistant = self.assistant.read(cx);
1112 if editor.selections.count() == 1 {
1113 let selection = editor.selections.newest::<usize>(cx);
1114 let mut offset = 0;
1115 let mut copied_text = String::new();
1116 let mut spanned_messages = 0;
1117 for message in &assistant.messages {
1118 let message_range = offset..offset + message.content.read(cx).len() + 1;
1119
1120 if message_range.start >= selection.range().end {
1121 break;
1122 } else if message_range.end >= selection.range().start {
1123 let range = cmp::max(message_range.start, selection.range().start)
1124 ..cmp::min(message_range.end, selection.range().end);
1125 if !range.is_empty() {
1126 if let Some(metadata) = assistant.messages_metadata.get(&message.excerpt_id)
1127 {
1128 spanned_messages += 1;
1129 write!(&mut copied_text, "## {}\n\n", metadata.role).unwrap();
1130 for chunk in
1131 assistant.buffer.read(cx).snapshot(cx).text_for_range(range)
1132 {
1133 copied_text.push_str(&chunk);
1134 }
1135 copied_text.push('\n');
1136 }
1137 }
1138 }
1139
1140 offset = message_range.end;
1141 }
1142
1143 if spanned_messages > 1 {
1144 cx.platform()
1145 .write_to_clipboard(ClipboardItem::new(copied_text));
1146 return;
1147 }
1148 }
1149
1150 cx.propagate_action();
1151 }
1152
1153 fn cycle_model(&mut self, cx: &mut ViewContext<Self>) {
1154 self.assistant.update(cx, |assistant, cx| {
1155 let new_model = match assistant.model.as_str() {
1156 "gpt-4" => "gpt-3.5-turbo",
1157 _ => "gpt-4",
1158 };
1159 assistant.set_model(new_model.into(), cx);
1160 });
1161 }
1162
1163 fn title(&self, cx: &AppContext) -> String {
1164 self.assistant
1165 .read(cx)
1166 .summary
1167 .clone()
1168 .unwrap_or_else(|| "New Context".into())
1169 }
1170}
1171
1172impl Entity for AssistantEditor {
1173 type Event = AssistantEditorEvent;
1174}
1175
1176impl View for AssistantEditor {
1177 fn ui_name() -> &'static str {
1178 "AssistantEditor"
1179 }
1180
1181 fn render(&mut self, cx: &mut ViewContext<Self>) -> AnyElement<Self> {
1182 enum Model {}
1183 let theme = &theme::current(cx).assistant;
1184 let assistant = &self.assistant.read(cx);
1185 let model = assistant.model.clone();
1186 let remaining_tokens = assistant.remaining_tokens().map(|remaining_tokens| {
1187 let remaining_tokens_style = if remaining_tokens <= 0 {
1188 &theme.no_remaining_tokens
1189 } else {
1190 &theme.remaining_tokens
1191 };
1192 Label::new(
1193 remaining_tokens.to_string(),
1194 remaining_tokens_style.text.clone(),
1195 )
1196 .contained()
1197 .with_style(remaining_tokens_style.container)
1198 });
1199
1200 Stack::new()
1201 .with_child(
1202 ChildView::new(&self.editor, cx)
1203 .contained()
1204 .with_style(theme.container),
1205 )
1206 .with_child(
1207 Flex::row()
1208 .with_child(
1209 MouseEventHandler::<Model, _>::new(0, cx, |state, _| {
1210 let style = theme.model.style_for(state, false);
1211 Label::new(model, style.text.clone())
1212 .contained()
1213 .with_style(style.container)
1214 })
1215 .with_cursor_style(CursorStyle::PointingHand)
1216 .on_click(MouseButton::Left, |_, this, cx| this.cycle_model(cx)),
1217 )
1218 .with_children(remaining_tokens)
1219 .contained()
1220 .with_style(theme.model_info_container)
1221 .aligned()
1222 .top()
1223 .right(),
1224 )
1225 .into_any()
1226 }
1227
1228 fn focus_in(&mut self, _: gpui::AnyViewHandle, cx: &mut ViewContext<Self>) {
1229 if cx.is_self_focused() {
1230 cx.focus(&self.editor);
1231 }
1232 }
1233}
1234
1235impl Item for AssistantEditor {
1236 fn tab_content<V: View>(
1237 &self,
1238 _: Option<usize>,
1239 style: &theme::Tab,
1240 cx: &gpui::AppContext,
1241 ) -> AnyElement<V> {
1242 let title = truncate_and_trailoff(&self.title(cx), editor::MAX_TAB_TITLE_LEN);
1243 Label::new(title, style.label.clone()).into_any()
1244 }
1245
1246 fn tab_tooltip_text(&self, cx: &AppContext) -> Option<Cow<str>> {
1247 Some(self.title(cx).into())
1248 }
1249
1250 fn as_searchable(
1251 &self,
1252 _: &ViewHandle<Self>,
1253 ) -> Option<Box<dyn workspace::searchable::SearchableItemHandle>> {
1254 Some(Box::new(self.editor.clone()))
1255 }
1256}
1257
1258#[derive(Clone, Debug)]
1259struct Message {
1260 excerpt_id: ExcerptId,
1261 content: ModelHandle<Buffer>,
1262}
1263
1264#[derive(Clone, Debug)]
1265struct MessageMetadata {
1266 role: Role,
1267 sent_at: DateTime<Local>,
1268 error: Option<String>,
1269}
1270
1271async fn stream_completion(
1272 api_key: String,
1273 executor: Arc<Background>,
1274 mut request: OpenAIRequest,
1275) -> Result<impl Stream<Item = Result<OpenAIResponseStreamEvent>>> {
1276 request.stream = true;
1277
1278 let (tx, rx) = futures::channel::mpsc::unbounded::<Result<OpenAIResponseStreamEvent>>();
1279
1280 let json_data = serde_json::to_string(&request)?;
1281 let mut response = Request::post(format!("{OPENAI_API_URL}/chat/completions"))
1282 .header("Content-Type", "application/json")
1283 .header("Authorization", format!("Bearer {}", api_key))
1284 .body(json_data)?
1285 .send_async()
1286 .await?;
1287
1288 let status = response.status();
1289 if status == StatusCode::OK {
1290 executor
1291 .spawn(async move {
1292 let mut lines = BufReader::new(response.body_mut()).lines();
1293
1294 fn parse_line(
1295 line: Result<String, io::Error>,
1296 ) -> Result<Option<OpenAIResponseStreamEvent>> {
1297 if let Some(data) = line?.strip_prefix("data: ") {
1298 let event = serde_json::from_str(&data)?;
1299 Ok(Some(event))
1300 } else {
1301 Ok(None)
1302 }
1303 }
1304
1305 while let Some(line) = lines.next().await {
1306 if let Some(event) = parse_line(line).transpose() {
1307 let done = event.as_ref().map_or(false, |event| {
1308 event
1309 .choices
1310 .last()
1311 .map_or(false, |choice| choice.finish_reason.is_some())
1312 });
1313 if tx.unbounded_send(event).is_err() {
1314 break;
1315 }
1316
1317 if done {
1318 break;
1319 }
1320 }
1321 }
1322
1323 anyhow::Ok(())
1324 })
1325 .detach();
1326
1327 Ok(rx)
1328 } else {
1329 let mut body = String::new();
1330 response.body_mut().read_to_string(&mut body).await?;
1331
1332 #[derive(Deserialize)]
1333 struct OpenAIResponse {
1334 error: OpenAIError,
1335 }
1336
1337 #[derive(Deserialize)]
1338 struct OpenAIError {
1339 message: String,
1340 }
1341
1342 match serde_json::from_str::<OpenAIResponse>(&body) {
1343 Ok(response) if !response.error.message.is_empty() => Err(anyhow!(
1344 "Failed to connect to OpenAI API: {}",
1345 response.error.message,
1346 )),
1347
1348 _ => Err(anyhow!(
1349 "Failed to connect to OpenAI API: {} {}",
1350 response.status(),
1351 body,
1352 )),
1353 }
1354 }
1355}
1356
1357#[cfg(test)]
1358mod tests {
1359 use super::*;
1360 use gpui::AppContext;
1361
1362 #[gpui::test]
1363 fn test_inserting_and_removing_messages(cx: &mut AppContext) {
1364 let registry = Arc::new(LanguageRegistry::test());
1365
1366 cx.add_model(|cx| {
1367 let mut assistant = Assistant::new(Default::default(), registry, cx);
1368 let message_1 = assistant.messages[0].clone();
1369 let message_2 = assistant.insert_message_after(ExcerptId::max(), Role::Assistant, cx);
1370 let message_3 = assistant.insert_message_after(message_2.excerpt_id, Role::User, cx);
1371 let message_4 = assistant.insert_message_after(message_2.excerpt_id, Role::User, cx);
1372 assistant.remove_empty_messages(
1373 HashSet::from_iter([message_3.excerpt_id, message_4.excerpt_id]),
1374 Default::default(),
1375 cx,
1376 );
1377 assert_eq!(assistant.messages.len(), 2);
1378 assert_eq!(assistant.messages[0].excerpt_id, message_1.excerpt_id);
1379 assert_eq!(assistant.messages[1].excerpt_id, message_2.excerpt_id);
1380 assistant
1381 });
1382 }
1383}