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