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