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