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