1mod assistant_settings;
   2mod attachments;
   3mod completion_provider;
   4mod saved_conversation;
   5mod saved_conversations;
   6mod tools;
   7pub mod ui;
   8
   9use crate::saved_conversation::SavedConversationMetadata;
  10use crate::ui::UserOrAssistant;
  11use ::ui::{div, prelude::*, Color, Tooltip, ViewContext};
  12use anyhow::{Context, Result};
  13use assistant_tooling::{
  14    AttachmentRegistry, ProjectContext, ToolFunctionCall, ToolRegistry, UserAttachment,
  15};
  16use attachments::ActiveEditorAttachmentTool;
  17use client::{proto, Client, UserStore};
  18use collections::HashMap;
  19use completion_provider::*;
  20use editor::Editor;
  21use feature_flags::FeatureFlagAppExt as _;
  22use fs::Fs;
  23use futures::{future::join_all, StreamExt};
  24use gpui::{
  25    list, AnyElement, AppContext, AsyncWindowContext, ClickEvent, EventEmitter, FocusHandle,
  26    FocusableView, ListAlignment, ListState, Model, Render, Task, View, WeakView,
  27};
  28use language::{language_settings::SoftWrap, LanguageRegistry};
  29use markdown::{Markdown, MarkdownStyle};
  30use open_ai::{FunctionContent, ToolCall, ToolCallContent};
  31use saved_conversation::{SavedAssistantMessagePart, SavedChatMessage, SavedConversation};
  32use saved_conversations::SavedConversations;
  33use semantic_index::{CloudEmbeddingProvider, ProjectIndex, ProjectIndexDebugView, SemanticIndex};
  34use serde::{Deserialize, Serialize};
  35use settings::Settings;
  36use std::sync::Arc;
  37use tools::{AnnotationTool, CreateBufferTool, ProjectIndexTool};
  38use ui::{ActiveFileButton, Composer, ProjectIndexButton};
  39use util::paths::CONVERSATIONS_DIR;
  40use util::{maybe, paths::EMBEDDINGS_DIR, ResultExt};
  41use workspace::{
  42    dock::{DockPosition, Panel, PanelEvent},
  43    Workspace,
  44};
  45
  46pub use assistant_settings::AssistantSettings;
  47
  48const MAX_COMPLETION_CALLS_PER_SUBMISSION: usize = 5;
  49
  50#[derive(Eq, PartialEq, Copy, Clone, Deserialize)]
  51pub struct Submit(SubmitMode);
  52
  53/// There are multiple different ways to submit a model request, represented by this enum.
  54#[derive(Eq, PartialEq, Copy, Clone, Deserialize)]
  55pub enum SubmitMode {
  56    /// Only include the conversation.
  57    Simple,
  58    /// Send the current file as context.
  59    CurrentFile,
  60    /// Search the codebase and send relevant excerpts.
  61    Codebase,
  62}
  63
  64gpui::actions!(assistant2, [Cancel, ToggleFocus, DebugProjectIndex,]);
  65gpui::impl_actions!(assistant2, [Submit]);
  66
  67pub fn init(client: Arc<Client>, cx: &mut AppContext) {
  68    AssistantSettings::register(cx);
  69
  70    cx.spawn(|mut cx| {
  71        let client = client.clone();
  72        async move {
  73            let embedding_provider = CloudEmbeddingProvider::new(client.clone());
  74            let semantic_index = SemanticIndex::new(
  75                EMBEDDINGS_DIR.join("semantic-index-db.0.mdb"),
  76                Arc::new(embedding_provider),
  77                &mut cx,
  78            )
  79            .await?;
  80            cx.update(|cx| cx.set_global(semantic_index))
  81        }
  82    })
  83    .detach();
  84
  85    cx.set_global(CompletionProvider::new(CloudCompletionProvider::new(
  86        client,
  87    )));
  88
  89    cx.observe_new_views(
  90        |workspace: &mut Workspace, _cx: &mut ViewContext<Workspace>| {
  91            workspace.register_action(|workspace, _: &ToggleFocus, cx| {
  92                workspace.toggle_panel_focus::<AssistantPanel>(cx);
  93            });
  94            workspace.register_action(|workspace, _: &DebugProjectIndex, cx| {
  95                if let Some(panel) = workspace.panel::<AssistantPanel>(cx) {
  96                    let index = panel.read(cx).chat.read(cx).project_index.clone();
  97                    let view = cx.new_view(|cx| ProjectIndexDebugView::new(index, cx));
  98                    workspace.add_item_to_center(Box::new(view), cx);
  99                }
 100            });
 101        },
 102    )
 103    .detach();
 104}
 105
 106pub fn enabled(cx: &AppContext) -> bool {
 107    cx.is_staff()
 108}
 109
 110pub struct AssistantPanel {
 111    chat: View<AssistantChat>,
 112    width: Option<Pixels>,
 113}
 114
 115impl AssistantPanel {
 116    pub fn load(
 117        workspace: WeakView<Workspace>,
 118        cx: AsyncWindowContext,
 119    ) -> Task<Result<View<Self>>> {
 120        cx.spawn(|mut cx| async move {
 121            let (app_state, project) = workspace.update(&mut cx, |workspace, _| {
 122                (workspace.app_state().clone(), workspace.project().clone())
 123            })?;
 124
 125            cx.new_view(|cx| {
 126                let project_index = cx.update_global(|semantic_index: &mut SemanticIndex, cx| {
 127                    semantic_index.project_index(project.clone(), cx)
 128                });
 129
 130                let mut tool_registry = ToolRegistry::new();
 131                tool_registry
 132                    .register(ProjectIndexTool::new(project_index.clone()))
 133                    .unwrap();
 134                tool_registry
 135                    .register(CreateBufferTool::new(workspace.clone(), project.clone()))
 136                    .unwrap();
 137                tool_registry
 138                    .register(AnnotationTool::new(workspace.clone(), project.clone()))
 139                    .unwrap();
 140
 141                let mut attachment_registry = AttachmentRegistry::new();
 142                attachment_registry
 143                    .register(ActiveEditorAttachmentTool::new(workspace.clone(), cx));
 144
 145                Self::new(
 146                    project.read(cx).fs().clone(),
 147                    app_state.languages.clone(),
 148                    Arc::new(tool_registry),
 149                    Arc::new(attachment_registry),
 150                    app_state.user_store.clone(),
 151                    project_index,
 152                    workspace,
 153                    cx,
 154                )
 155            })
 156        })
 157    }
 158
 159    #[allow(clippy::too_many_arguments)]
 160    pub fn new(
 161        fs: Arc<dyn Fs>,
 162        language_registry: Arc<LanguageRegistry>,
 163        tool_registry: Arc<ToolRegistry>,
 164        attachment_registry: Arc<AttachmentRegistry>,
 165        user_store: Model<UserStore>,
 166        project_index: Model<ProjectIndex>,
 167        workspace: WeakView<Workspace>,
 168        cx: &mut ViewContext<Self>,
 169    ) -> Self {
 170        let chat = cx.new_view(|cx| {
 171            AssistantChat::new(
 172                fs,
 173                language_registry,
 174                tool_registry.clone(),
 175                attachment_registry,
 176                user_store,
 177                project_index,
 178                workspace,
 179                cx,
 180            )
 181        });
 182
 183        Self { width: None, chat }
 184    }
 185}
 186
 187impl Render for AssistantPanel {
 188    fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
 189        div()
 190            .size_full()
 191            .v_flex()
 192            .bg(cx.theme().colors().panel_background)
 193            .child(self.chat.clone())
 194    }
 195}
 196
 197impl Panel for AssistantPanel {
 198    fn persistent_name() -> &'static str {
 199        "AssistantPanelv2"
 200    }
 201
 202    fn position(&self, _cx: &WindowContext) -> workspace::dock::DockPosition {
 203        // todo!("Add a setting / use assistant settings")
 204        DockPosition::Right
 205    }
 206
 207    fn position_is_valid(&self, position: workspace::dock::DockPosition) -> bool {
 208        matches!(position, DockPosition::Right)
 209    }
 210
 211    fn set_position(&mut self, _: workspace::dock::DockPosition, _: &mut ViewContext<Self>) {
 212        // Do nothing until we have a setting for this
 213    }
 214
 215    fn size(&self, _cx: &WindowContext) -> Pixels {
 216        self.width.unwrap_or(px(400.))
 217    }
 218
 219    fn set_size(&mut self, size: Option<Pixels>, cx: &mut ViewContext<Self>) {
 220        self.width = size;
 221        cx.notify();
 222    }
 223
 224    fn icon(&self, _cx: &WindowContext) -> Option<::ui::IconName> {
 225        Some(IconName::ZedAssistant)
 226    }
 227
 228    fn icon_tooltip(&self, _: &WindowContext) -> Option<&'static str> {
 229        Some("Assistant Panel ✨")
 230    }
 231
 232    fn toggle_action(&self) -> Box<dyn gpui::Action> {
 233        Box::new(ToggleFocus)
 234    }
 235}
 236
 237impl EventEmitter<PanelEvent> for AssistantPanel {}
 238
 239impl FocusableView for AssistantPanel {
 240    fn focus_handle(&self, cx: &AppContext) -> FocusHandle {
 241        self.chat.read(cx).composer_editor.read(cx).focus_handle(cx)
 242    }
 243}
 244
 245pub struct AssistantChat {
 246    model: String,
 247    messages: Vec<ChatMessage>,
 248    list_state: ListState,
 249    fs: Arc<dyn Fs>,
 250    language_registry: Arc<LanguageRegistry>,
 251    composer_editor: View<Editor>,
 252    saved_conversations: View<SavedConversations>,
 253    saved_conversations_open: bool,
 254    project_index_button: View<ProjectIndexButton>,
 255    active_file_button: Option<View<ActiveFileButton>>,
 256    user_store: Model<UserStore>,
 257    next_message_id: MessageId,
 258    collapsed_messages: HashMap<MessageId, bool>,
 259    editing_message: Option<EditingMessage>,
 260    pending_completion: Option<Task<()>>,
 261    tool_registry: Arc<ToolRegistry>,
 262    attachment_registry: Arc<AttachmentRegistry>,
 263    project_index: Model<ProjectIndex>,
 264    markdown_style: MarkdownStyle,
 265}
 266
 267struct EditingMessage {
 268    id: MessageId,
 269    body: View<Editor>,
 270}
 271
 272impl AssistantChat {
 273    #[allow(clippy::too_many_arguments)]
 274    fn new(
 275        fs: Arc<dyn Fs>,
 276        language_registry: Arc<LanguageRegistry>,
 277        tool_registry: Arc<ToolRegistry>,
 278        attachment_registry: Arc<AttachmentRegistry>,
 279        user_store: Model<UserStore>,
 280        project_index: Model<ProjectIndex>,
 281        workspace: WeakView<Workspace>,
 282        cx: &mut ViewContext<Self>,
 283    ) -> Self {
 284        let model = CompletionProvider::get(cx).default_model();
 285        let view = cx.view().downgrade();
 286        let list_state = ListState::new(
 287            0,
 288            ListAlignment::Bottom,
 289            px(1024.),
 290            move |ix, cx: &mut WindowContext| {
 291                view.update(cx, |this, cx| this.render_message(ix, cx))
 292                    .unwrap()
 293            },
 294        );
 295
 296        let project_index_button = cx.new_view(|cx| {
 297            ProjectIndexButton::new(project_index.clone(), tool_registry.clone(), cx)
 298        });
 299
 300        let active_file_button = match workspace.upgrade() {
 301            Some(workspace) => {
 302                Some(cx.new_view(
 303                    |cx| ActiveFileButton::new(attachment_registry.clone(), workspace, cx), //
 304                ))
 305            }
 306            _ => None,
 307        };
 308
 309        let saved_conversations = cx.new_view(|cx| SavedConversations::new(cx));
 310        cx.spawn({
 311            let fs = fs.clone();
 312            let saved_conversations = saved_conversations.downgrade();
 313            |_assistant_chat, mut cx| async move {
 314                let saved_conversation_metadata = SavedConversationMetadata::list(fs).await?;
 315
 316                cx.update(|cx| {
 317                    saved_conversations.update(cx, |this, cx| {
 318                        this.init(saved_conversation_metadata, cx);
 319                    })
 320                })??;
 321
 322                anyhow::Ok(())
 323            }
 324        })
 325        .detach_and_log_err(cx);
 326
 327        Self {
 328            model,
 329            messages: Vec::new(),
 330            composer_editor: cx.new_view(|cx| {
 331                let mut editor = Editor::auto_height(80, cx);
 332                editor.set_soft_wrap_mode(SoftWrap::EditorWidth, cx);
 333                editor.set_placeholder_text("Send a message…", cx);
 334                editor
 335            }),
 336            saved_conversations,
 337            saved_conversations_open: false,
 338            list_state,
 339            user_store,
 340            fs,
 341            language_registry,
 342            project_index_button,
 343            active_file_button,
 344            project_index,
 345            next_message_id: MessageId(0),
 346            editing_message: None,
 347            collapsed_messages: HashMap::default(),
 348            pending_completion: None,
 349            attachment_registry,
 350            tool_registry,
 351            markdown_style: MarkdownStyle {
 352                code_block: gpui::TextStyleRefinement {
 353                    font_family: Some("Zed Mono".into()),
 354                    color: Some(cx.theme().colors().editor_foreground),
 355                    background_color: Some(cx.theme().colors().editor_background),
 356                    ..Default::default()
 357                },
 358                inline_code: gpui::TextStyleRefinement {
 359                    font_family: Some("Zed Mono".into()),
 360                    // @nate: Could we add inline-code specific styles to the theme?
 361                    color: Some(cx.theme().colors().editor_foreground),
 362                    background_color: Some(cx.theme().colors().editor_background),
 363                    ..Default::default()
 364                },
 365                rule_color: Color::Muted.color(cx),
 366                block_quote_border_color: Color::Muted.color(cx),
 367                block_quote: gpui::TextStyleRefinement {
 368                    color: Some(Color::Muted.color(cx)),
 369                    ..Default::default()
 370                },
 371                link: gpui::TextStyleRefinement {
 372                    color: Some(Color::Accent.color(cx)),
 373                    underline: Some(gpui::UnderlineStyle {
 374                        thickness: px(1.),
 375                        color: Some(Color::Accent.color(cx)),
 376                        wavy: false,
 377                    }),
 378                    ..Default::default()
 379                },
 380                syntax: cx.theme().syntax().clone(),
 381                selection_background_color: {
 382                    let mut selection = cx.theme().players().local().selection;
 383                    selection.fade_out(0.7);
 384                    selection
 385                },
 386            },
 387        }
 388    }
 389
 390    fn message_for_id(&self, id: MessageId) -> Option<&ChatMessage> {
 391        self.messages.iter().find(|message| match message {
 392            ChatMessage::User(message) => message.id == id,
 393            ChatMessage::Assistant(message) => message.id == id,
 394        })
 395    }
 396
 397    fn toggle_saved_conversations(&mut self) {
 398        self.saved_conversations_open = !self.saved_conversations_open;
 399    }
 400
 401    fn cancel(&mut self, _: &Cancel, cx: &mut ViewContext<Self>) {
 402        // If we're currently editing a message, cancel the edit.
 403        if self.editing_message.take().is_some() {
 404            cx.notify();
 405            return;
 406        }
 407
 408        if self.pending_completion.take().is_some() {
 409            if let Some(ChatMessage::Assistant(grouping)) = self.messages.last() {
 410                if grouping.messages.is_empty() {
 411                    self.pop_message(cx);
 412                }
 413            }
 414            return;
 415        }
 416
 417        cx.propagate();
 418    }
 419
 420    fn submit(&mut self, Submit(mode): &Submit, cx: &mut ViewContext<Self>) {
 421        if self.composer_editor.focus_handle(cx).is_focused(cx) {
 422            // Don't allow multiple concurrent completions.
 423            if self.pending_completion.is_some() {
 424                cx.propagate();
 425                return;
 426            }
 427
 428            let message = self.composer_editor.update(cx, |composer_editor, cx| {
 429                let text = composer_editor.text(cx);
 430                let id = self.next_message_id.post_inc();
 431                let body = cx.new_view(|cx| {
 432                    Markdown::new(
 433                        text,
 434                        self.markdown_style.clone(),
 435                        self.language_registry.clone(),
 436                        cx,
 437                    )
 438                });
 439                composer_editor.clear(cx);
 440
 441                ChatMessage::User(UserMessage {
 442                    id,
 443                    body,
 444                    attachments: Vec::new(),
 445                })
 446            });
 447            self.push_message(message, cx);
 448        } else if let Some(editing_message) = self.editing_message.as_ref() {
 449            let focus_handle = editing_message.body.focus_handle(cx);
 450            if focus_handle.contains_focused(cx) {
 451                if let Some(ChatMessage::User(user_message)) =
 452                    self.message_for_id(editing_message.id)
 453                {
 454                    user_message.body.update(cx, |body, cx| {
 455                        body.reset(editing_message.body.read(cx).text(cx), cx);
 456                    });
 457                }
 458
 459                self.truncate_messages(editing_message.id, cx);
 460
 461                self.pending_completion.take();
 462                self.composer_editor.focus_handle(cx).focus(cx);
 463                self.editing_message.take();
 464            } else {
 465                log::error!("unexpected state: no user message editor is focused.");
 466                return;
 467            }
 468        } else {
 469            log::error!("unexpected state: no user message editor is focused.");
 470            return;
 471        }
 472
 473        let mode = *mode;
 474        self.pending_completion = Some(cx.spawn(move |this, mut cx| async move {
 475            let attachments_task = this.update(&mut cx, |this, cx| {
 476                let attachment_registry = this.attachment_registry.clone();
 477                attachment_registry.call_all_attachment_tools(cx)
 478            });
 479
 480            let attachments = maybe!(async {
 481                let attachments_task = attachments_task?;
 482                let attachments = attachments_task.await?;
 483
 484                anyhow::Ok(attachments)
 485            })
 486            .await
 487            .log_err()
 488            .unwrap_or_default();
 489
 490            // Set the attachments to the _last_ user message
 491            this.update(&mut cx, |this, _cx| {
 492                if let Some(ChatMessage::User(message)) = this.messages.last_mut() {
 493                    message.attachments = attachments;
 494                }
 495            })
 496            .log_err();
 497
 498            Self::request_completion(
 499                this.clone(),
 500                mode,
 501                MAX_COMPLETION_CALLS_PER_SUBMISSION,
 502                &mut cx,
 503            )
 504            .await
 505            .log_err();
 506
 507            this.update(&mut cx, |this, _cx| {
 508                this.pending_completion = None;
 509            })
 510            .context("Failed to push new user message")
 511            .log_err();
 512        }));
 513    }
 514
 515    async fn request_completion(
 516        this: WeakView<Self>,
 517        mode: SubmitMode,
 518        limit: usize,
 519        cx: &mut AsyncWindowContext,
 520    ) -> Result<()> {
 521        let mut call_count = 0;
 522        loop {
 523            let complete = async {
 524                let (tool_definitions, model_name, messages) = this.update(cx, |this, cx| {
 525                    this.push_new_assistant_message(cx);
 526
 527                    let definitions = if call_count < limit
 528                        && matches!(mode, SubmitMode::Codebase | SubmitMode::Simple)
 529                    {
 530                        this.tool_registry.definitions()
 531                    } else {
 532                        Vec::new()
 533                    };
 534                    call_count += 1;
 535
 536                    (
 537                        definitions,
 538                        this.model.clone(),
 539                        this.completion_messages(cx),
 540                    )
 541                })?;
 542
 543                let messages = messages.await?;
 544
 545                let completion = cx.update(|cx| {
 546                    CompletionProvider::get(cx).complete(
 547                        model_name,
 548                        messages,
 549                        Vec::new(),
 550                        1.0,
 551                        tool_definitions,
 552                    )
 553                });
 554
 555                let mut stream = completion?.await?;
 556                while let Some(delta) = stream.next().await {
 557                    let delta = delta?;
 558                    this.update(cx, |this, cx| {
 559                        if let Some(ChatMessage::Assistant(AssistantMessage { messages, .. })) =
 560                            this.messages.last_mut()
 561                        {
 562                            if messages.is_empty() {
 563                                messages.push(AssistantMessagePart {
 564                                    body: cx.new_view(|cx| {
 565                                        Markdown::new(
 566                                            "".into(),
 567                                            this.markdown_style.clone(),
 568                                            this.language_registry.clone(),
 569                                            cx,
 570                                        )
 571                                    }),
 572                                    tool_calls: Vec::new(),
 573                                })
 574                            }
 575
 576                            let message = messages.last_mut().unwrap();
 577
 578                            if let Some(content) = &delta.content {
 579                                message
 580                                    .body
 581                                    .update(cx, |message, cx| message.append(&content, cx));
 582                            }
 583
 584                            for tool_call_delta in delta.tool_calls {
 585                                let index = tool_call_delta.index as usize;
 586                                if index >= message.tool_calls.len() {
 587                                    message.tool_calls.resize_with(index + 1, Default::default);
 588                                }
 589                                let tool_call = &mut message.tool_calls[index];
 590
 591                                if let Some(id) = &tool_call_delta.id {
 592                                    tool_call.id.push_str(id);
 593                                }
 594
 595                                match tool_call_delta.variant {
 596                                    Some(proto::tool_call_delta::Variant::Function(
 597                                        tool_call_delta,
 598                                    )) => {
 599                                        this.tool_registry.update_tool_call(
 600                                            tool_call,
 601                                            tool_call_delta.name.as_deref(),
 602                                            tool_call_delta.arguments.as_deref(),
 603                                            cx,
 604                                        );
 605                                    }
 606                                    None => {}
 607                                }
 608                            }
 609
 610                            cx.notify();
 611                        } else {
 612                            unreachable!()
 613                        }
 614                    })?;
 615                }
 616
 617                anyhow::Ok(())
 618            }
 619            .await;
 620
 621            let mut tool_tasks = Vec::new();
 622            this.update(cx, |this, cx| {
 623                if let Some(ChatMessage::Assistant(AssistantMessage {
 624                    error: message_error,
 625                    messages,
 626                    ..
 627                })) = this.messages.last_mut()
 628                {
 629                    if let Err(error) = complete {
 630                        message_error.replace(SharedString::from(error.to_string()));
 631                        cx.notify();
 632                    } else {
 633                        if let Some(current_message) = messages.last_mut() {
 634                            for tool_call in current_message.tool_calls.iter_mut() {
 635                                tool_tasks
 636                                    .extend(this.tool_registry.execute_tool_call(tool_call, cx));
 637                            }
 638                        }
 639                    }
 640                }
 641            })?;
 642
 643            // This ends recursion on calling for responses after tools
 644            if tool_tasks.is_empty() {
 645                return Ok(());
 646            }
 647
 648            join_all(tool_tasks.into_iter()).await;
 649        }
 650    }
 651
 652    fn push_new_assistant_message(&mut self, cx: &mut ViewContext<Self>) {
 653        // If the last message is a grouped assistant message, add to the grouped message
 654        if let Some(ChatMessage::Assistant(AssistantMessage { messages, .. })) =
 655            self.messages.last_mut()
 656        {
 657            messages.push(AssistantMessagePart {
 658                body: cx.new_view(|cx| {
 659                    Markdown::new(
 660                        "".into(),
 661                        self.markdown_style.clone(),
 662                        self.language_registry.clone(),
 663                        cx,
 664                    )
 665                }),
 666                tool_calls: Vec::new(),
 667            });
 668            return;
 669        }
 670
 671        let message = ChatMessage::Assistant(AssistantMessage {
 672            id: self.next_message_id.post_inc(),
 673            messages: vec![AssistantMessagePart {
 674                body: cx.new_view(|cx| {
 675                    Markdown::new(
 676                        "".into(),
 677                        self.markdown_style.clone(),
 678                        self.language_registry.clone(),
 679                        cx,
 680                    )
 681                }),
 682                tool_calls: Vec::new(),
 683            }],
 684            error: None,
 685        });
 686        self.push_message(message, cx);
 687    }
 688
 689    fn push_message(&mut self, message: ChatMessage, cx: &mut ViewContext<Self>) {
 690        let old_len = self.messages.len();
 691        let focus_handle = Some(message.focus_handle(cx));
 692        self.messages.push(message);
 693        self.list_state
 694            .splice_focusable(old_len..old_len, focus_handle);
 695        cx.notify();
 696    }
 697
 698    fn pop_message(&mut self, cx: &mut ViewContext<Self>) {
 699        if self.messages.is_empty() {
 700            return;
 701        }
 702
 703        self.messages.pop();
 704        self.list_state
 705            .splice(self.messages.len()..self.messages.len() + 1, 0);
 706        cx.notify();
 707    }
 708
 709    fn truncate_messages(&mut self, last_message_id: MessageId, cx: &mut ViewContext<Self>) {
 710        if let Some(index) = self.messages.iter().position(|message| match message {
 711            ChatMessage::User(message) => message.id == last_message_id,
 712            ChatMessage::Assistant(message) => message.id == last_message_id,
 713        }) {
 714            self.list_state.splice(index + 1..self.messages.len(), 0);
 715            self.messages.truncate(index + 1);
 716            cx.notify();
 717        }
 718    }
 719
 720    fn is_message_collapsed(&self, id: &MessageId) -> bool {
 721        self.collapsed_messages.get(id).copied().unwrap_or_default()
 722    }
 723
 724    fn toggle_message_collapsed(&mut self, id: MessageId) {
 725        let entry = self.collapsed_messages.entry(id).or_insert(false);
 726        *entry = !*entry;
 727    }
 728
 729    fn reset(&mut self) {
 730        self.messages.clear();
 731        self.list_state.reset(0);
 732        self.editing_message.take();
 733        self.collapsed_messages.clear();
 734    }
 735
 736    fn new_conversation(&mut self, cx: &mut ViewContext<Self>) {
 737        let messages = std::mem::take(&mut self.messages)
 738            .into_iter()
 739            .map(|message| self.serialize_message(message, cx))
 740            .collect::<Vec<_>>();
 741
 742        self.reset();
 743
 744        let title = messages
 745            .first()
 746            .map(|message| match message {
 747                SavedChatMessage::User { body, .. } => body.clone(),
 748                SavedChatMessage::Assistant { messages, .. } => messages
 749                    .first()
 750                    .map(|message| message.body.to_string())
 751                    .unwrap_or_default(),
 752            })
 753            .unwrap_or_else(|| "A conversation with the assistant.".to_string());
 754
 755        let saved_conversation = SavedConversation {
 756            version: "0.3.0".to_string(),
 757            title,
 758            messages,
 759        };
 760
 761        let discriminant = 1;
 762
 763        let path = CONVERSATIONS_DIR.join(&format!(
 764            "{title} - {discriminant}.zed.{version}.json",
 765            title = saved_conversation.title,
 766            version = saved_conversation.version
 767        ));
 768
 769        cx.spawn({
 770            let fs = self.fs.clone();
 771            |_this, _cx| async move {
 772                fs.create_dir(CONVERSATIONS_DIR.as_ref()).await?;
 773                fs.atomic_write(path, serde_json::to_string(&saved_conversation)?)
 774                    .await?;
 775
 776                anyhow::Ok(())
 777            }
 778        })
 779        .detach_and_log_err(cx);
 780    }
 781
 782    fn render_error(
 783        &self,
 784        error: Option<SharedString>,
 785        _ix: usize,
 786        cx: &mut ViewContext<Self>,
 787    ) -> AnyElement {
 788        let theme = cx.theme();
 789
 790        if let Some(error) = error {
 791            div()
 792                .py_1()
 793                .px_2()
 794                .mx_neg_1()
 795                .rounded_md()
 796                .border_1()
 797                .border_color(theme.status().error_border)
 798                // .bg(theme.status().error_background)
 799                .text_color(theme.status().error)
 800                .child(error.clone())
 801                .into_any_element()
 802        } else {
 803            div().into_any_element()
 804        }
 805    }
 806
 807    fn render_message(&self, ix: usize, cx: &mut ViewContext<Self>) -> AnyElement {
 808        let is_first = ix == 0;
 809        let is_last = ix == self.messages.len().saturating_sub(1);
 810
 811        let padding = Spacing::Large.rems(cx);
 812
 813        // Whenever there's a run of assistant messages, group as one Assistant UI element
 814
 815        match &self.messages[ix] {
 816            ChatMessage::User(UserMessage {
 817                id,
 818                body,
 819                attachments,
 820            }) => div()
 821                .id(SharedString::from(format!("message-{}-container", id.0)))
 822                .when(is_first, |this| this.pt(padding))
 823                .map(|element| {
 824                    if let Some(editing_message) = self.editing_message.as_ref() {
 825                        if editing_message.id == *id {
 826                            return element.child(Composer::new(
 827                                editing_message.body.clone(),
 828                                self.project_index_button.clone(),
 829                                self.active_file_button.clone(),
 830                                crate::ui::ModelSelector::new(
 831                                    cx.view().downgrade(),
 832                                    self.model.clone(),
 833                                )
 834                                .into_any_element(),
 835                            ));
 836                        }
 837                    }
 838
 839                    element
 840                        .on_click(cx.listener({
 841                            let id = *id;
 842                            let body = body.clone();
 843                            move |assistant_chat, event: &ClickEvent, cx| {
 844                                if event.up.click_count == 2 {
 845                                    let body = cx.new_view(|cx| {
 846                                        let mut editor = Editor::auto_height(80, cx);
 847                                        let source = Arc::from(body.read(cx).source());
 848                                        editor.set_text(source, cx);
 849                                        editor.set_soft_wrap_mode(SoftWrap::EditorWidth, cx);
 850                                        editor
 851                                    });
 852                                    assistant_chat.editing_message = Some(EditingMessage {
 853                                        id,
 854                                        body: body.clone(),
 855                                    });
 856                                    body.focus_handle(cx).focus(cx);
 857                                }
 858                            }
 859                        }))
 860                        .child(
 861                            crate::ui::ChatMessage::new(
 862                                *id,
 863                                UserOrAssistant::User(self.user_store.read(cx).current_user()),
 864                                // todo!(): clean up the vec usage
 865                                vec![
 866                                    body.clone().into_any_element(),
 867                                    h_flex()
 868                                        .gap_2()
 869                                        .children(
 870                                            attachments
 871                                                .iter()
 872                                                .map(|attachment| attachment.view.clone()),
 873                                        )
 874                                        .into_any_element(),
 875                                ],
 876                                self.is_message_collapsed(id),
 877                                Box::new(cx.listener({
 878                                    let id = *id;
 879                                    move |assistant_chat, _event, _cx| {
 880                                        assistant_chat.toggle_message_collapsed(id)
 881                                    }
 882                                })),
 883                            )
 884                            // TODO: Wire up selections.
 885                            .selected(is_last),
 886                        )
 887                })
 888                .into_any(),
 889            ChatMessage::Assistant(AssistantMessage {
 890                id,
 891                messages,
 892                error,
 893                ..
 894            }) => {
 895                let mut message_elements = Vec::new();
 896
 897                for message in messages {
 898                    if !message.body.read(cx).source().is_empty() {
 899                        message_elements.push(div().child(message.body.clone()).into_any())
 900                    }
 901
 902                    let tools = message
 903                        .tool_calls
 904                        .iter()
 905                        .filter_map(|tool_call| self.tool_registry.render_tool_call(tool_call, cx))
 906                        .collect::<Vec<AnyElement>>();
 907
 908                    if !tools.is_empty() {
 909                        message_elements.push(div().children(tools).into_any())
 910                    }
 911                }
 912
 913                if message_elements.is_empty() {
 914                    message_elements.push(::ui::Label::new("Researching...").into_any_element())
 915                }
 916
 917                div()
 918                    .when(is_first, |this| this.pt(padding))
 919                    .child(
 920                        crate::ui::ChatMessage::new(
 921                            *id,
 922                            UserOrAssistant::Assistant,
 923                            message_elements,
 924                            self.is_message_collapsed(id),
 925                            Box::new(cx.listener({
 926                                let id = *id;
 927                                move |assistant_chat, _event, _cx| {
 928                                    assistant_chat.toggle_message_collapsed(id)
 929                                }
 930                            })),
 931                        )
 932                        // TODO: Wire up selections.
 933                        .selected(is_last),
 934                    )
 935                    .child(self.render_error(error.clone(), ix, cx))
 936                    .into_any()
 937            }
 938        }
 939    }
 940
 941    fn completion_messages(&self, cx: &mut WindowContext) -> Task<Result<Vec<CompletionMessage>>> {
 942        let project_index = self.project_index.read(cx);
 943        let project = project_index.project();
 944        let fs = project_index.fs();
 945
 946        let mut project_context = ProjectContext::new(project, fs);
 947        let mut completion_messages = Vec::new();
 948
 949        for message in &self.messages {
 950            match message {
 951                ChatMessage::User(UserMessage {
 952                    body, attachments, ..
 953                }) => {
 954                    for attachment in attachments {
 955                        if let Some(content) = attachment.generate(&mut project_context, cx) {
 956                            completion_messages.push(CompletionMessage::System { content });
 957                        }
 958                    }
 959
 960                    // Show user's message last so that the assistant is grounded in the user's request
 961                    completion_messages.push(CompletionMessage::User {
 962                        content: body.read(cx).source().to_string(),
 963                    });
 964                }
 965                ChatMessage::Assistant(AssistantMessage { messages, .. }) => {
 966                    for message in messages {
 967                        let body = message.body.clone();
 968
 969                        if body.read(cx).source().is_empty() && message.tool_calls.is_empty() {
 970                            continue;
 971                        }
 972
 973                        let tool_calls_from_assistant = message
 974                            .tool_calls
 975                            .iter()
 976                            .map(|tool_call| ToolCall {
 977                                content: ToolCallContent::Function {
 978                                    function: FunctionContent {
 979                                        name: tool_call.name.clone(),
 980                                        arguments: tool_call.arguments.clone(),
 981                                    },
 982                                },
 983                                id: tool_call.id.clone(),
 984                            })
 985                            .collect();
 986
 987                        completion_messages.push(CompletionMessage::Assistant {
 988                            content: Some(body.read(cx).source().to_string()),
 989                            tool_calls: tool_calls_from_assistant,
 990                        });
 991
 992                        for tool_call in &message.tool_calls {
 993                            // Every tool call _must_ have a result by ID, otherwise OpenAI will error.
 994                            let content = self.tool_registry.content_for_tool_call(
 995                                tool_call,
 996                                &mut project_context,
 997                                cx,
 998                            );
 999                            completion_messages.push(CompletionMessage::Tool {
1000                                content,
1001                                tool_call_id: tool_call.id.clone(),
1002                            });
1003                        }
1004                    }
1005                }
1006            }
1007        }
1008
1009        let system_message = project_context.generate_system_message(cx);
1010
1011        cx.background_executor().spawn(async move {
1012            let content = system_message.await?;
1013            completion_messages.insert(0, CompletionMessage::System { content });
1014            Ok(completion_messages)
1015        })
1016    }
1017
1018    fn serialize_message(
1019        &self,
1020        message: ChatMessage,
1021        cx: &mut ViewContext<AssistantChat>,
1022    ) -> SavedChatMessage {
1023        match message {
1024            ChatMessage::User(message) => SavedChatMessage::User {
1025                id: message.id,
1026                body: message.body.read(cx).source().into(),
1027                attachments: message
1028                    .attachments
1029                    .iter()
1030                    .map(|attachment| {
1031                        self.attachment_registry
1032                            .serialize_user_attachment(attachment)
1033                    })
1034                    .collect(),
1035            },
1036            ChatMessage::Assistant(message) => SavedChatMessage::Assistant {
1037                id: message.id,
1038                error: message.error,
1039                messages: message
1040                    .messages
1041                    .iter()
1042                    .map(|message| SavedAssistantMessagePart {
1043                        body: message.body.read(cx).source().to_string().into(),
1044                        tool_calls: message
1045                            .tool_calls
1046                            .iter()
1047                            .filter_map(|tool_call| {
1048                                self.tool_registry
1049                                    .serialize_tool_call(tool_call, cx)
1050                                    .log_err()
1051                            })
1052                            .collect(),
1053                    })
1054                    .collect(),
1055            },
1056        }
1057    }
1058}
1059
1060impl Render for AssistantChat {
1061    fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
1062        let header_height = Spacing::Small.rems(cx) * 2.0 + ButtonSize::Default.rems();
1063
1064        div()
1065            .relative()
1066            .flex_1()
1067            .v_flex()
1068            .key_context("AssistantChat")
1069            .on_action(cx.listener(Self::submit))
1070            .on_action(cx.listener(Self::cancel))
1071            .text_color(Color::Default.color(cx))
1072            .child(list(self.list_state.clone()).flex_1().pt(header_height))
1073            .child(
1074                h_flex()
1075                    .absolute()
1076                    .top_0()
1077                    .justify_between()
1078                    .w_full()
1079                    .h(header_height)
1080                    .p(Spacing::Small.rems(cx))
1081                    .child(
1082                        IconButton::new(
1083                            "toggle-saved-conversations",
1084                            if self.saved_conversations_open {
1085                                IconName::ChevronRight
1086                            } else {
1087                                IconName::ChevronLeft
1088                            },
1089                        )
1090                        .on_click(cx.listener(|this, _event, _cx| {
1091                            this.toggle_saved_conversations();
1092                        }))
1093                        .tooltip(move |cx| Tooltip::text("Switch Conversations", cx)),
1094                    )
1095                    .child(
1096                        h_flex()
1097                            .gap(Spacing::Large.rems(cx))
1098                            .child(
1099                                IconButton::new("new-conversation", IconName::Plus)
1100                                    .on_click(cx.listener(move |this, _event, cx| {
1101                                        this.new_conversation(cx);
1102                                    }))
1103                                    .tooltip(move |cx| Tooltip::text("New Conversation", cx)),
1104                            )
1105                            .child(
1106                                IconButton::new("assistant-menu", IconName::Menu)
1107                                    .disabled(true)
1108                                    .tooltip(move |cx| {
1109                                        Tooltip::text(
1110                                            "Coming soon – Assistant settings & controls",
1111                                            cx,
1112                                        )
1113                                    }),
1114                            ),
1115                    ),
1116            )
1117            .when(self.saved_conversations_open, |element| {
1118                element.child(
1119                    h_flex()
1120                        .absolute()
1121                        .top(header_height)
1122                        .w_full()
1123                        .child(self.saved_conversations.clone()),
1124                )
1125            })
1126            .child(Composer::new(
1127                self.composer_editor.clone(),
1128                self.project_index_button.clone(),
1129                self.active_file_button.clone(),
1130                crate::ui::ModelSelector::new(cx.view().downgrade(), self.model.clone())
1131                    .into_any_element(),
1132            ))
1133    }
1134}
1135
1136#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy, Serialize, Deserialize)]
1137pub struct MessageId(usize);
1138
1139impl MessageId {
1140    fn post_inc(&mut self) -> Self {
1141        let id = *self;
1142        self.0 += 1;
1143        id
1144    }
1145}
1146
1147enum ChatMessage {
1148    User(UserMessage),
1149    Assistant(AssistantMessage),
1150}
1151
1152impl ChatMessage {
1153    fn focus_handle(&self, cx: &AppContext) -> Option<FocusHandle> {
1154        match self {
1155            ChatMessage::User(message) => Some(message.body.focus_handle(cx)),
1156            ChatMessage::Assistant(_) => None,
1157        }
1158    }
1159}
1160
1161struct UserMessage {
1162    pub id: MessageId,
1163    pub body: View<Markdown>,
1164    pub attachments: Vec<UserAttachment>,
1165}
1166
1167struct AssistantMessagePart {
1168    pub body: View<Markdown>,
1169    pub tool_calls: Vec<ToolFunctionCall>,
1170}
1171
1172struct AssistantMessage {
1173    pub id: MessageId,
1174    pub messages: Vec<AssistantMessagePart>,
1175    pub error: Option<SharedString>,
1176}