assistant2.rs

   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 open_ai::{FunctionContent, ToolCall, ToolCallContent};
  30use rich_text::RichText;
  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}
 265
 266struct EditingMessage {
 267    id: MessageId,
 268    old_body: Arc<str>,
 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        }
 352    }
 353
 354    fn editing_message_id(&self) -> Option<MessageId> {
 355        self.editing_message.as_ref().map(|message| message.id)
 356    }
 357
 358    fn focused_message_id(&self, cx: &WindowContext) -> Option<MessageId> {
 359        self.messages.iter().find_map(|message| match message {
 360            ChatMessage::User(message) => message
 361                .body
 362                .focus_handle(cx)
 363                .contains_focused(cx)
 364                .then_some(message.id),
 365            ChatMessage::Assistant(_) => None,
 366        })
 367    }
 368
 369    fn toggle_saved_conversations(&mut self) {
 370        self.saved_conversations_open = !self.saved_conversations_open;
 371    }
 372
 373    fn cancel(&mut self, _: &Cancel, cx: &mut ViewContext<Self>) {
 374        // If we're currently editing a message, cancel the edit.
 375        if let Some(editing_message) = self.editing_message.take() {
 376            editing_message
 377                .body
 378                .update(cx, |body, cx| body.set_text(editing_message.old_body, cx));
 379            return;
 380        }
 381
 382        if self.pending_completion.take().is_some() {
 383            if let Some(ChatMessage::Assistant(grouping)) = self.messages.last() {
 384                if grouping.messages.is_empty() {
 385                    self.pop_message(cx);
 386                }
 387            }
 388            return;
 389        }
 390
 391        cx.propagate();
 392    }
 393
 394    fn submit(&mut self, Submit(mode): &Submit, cx: &mut ViewContext<Self>) {
 395        if let Some(focused_message_id) = self.focused_message_id(cx) {
 396            self.truncate_messages(focused_message_id, cx);
 397            self.pending_completion.take();
 398            self.composer_editor.focus_handle(cx).focus(cx);
 399            if self.editing_message_id() == Some(focused_message_id) {
 400                self.editing_message.take();
 401            }
 402        } else if self.composer_editor.focus_handle(cx).is_focused(cx) {
 403            // Don't allow multiple concurrent completions.
 404            if self.pending_completion.is_some() {
 405                cx.propagate();
 406                return;
 407            }
 408
 409            let message = self.composer_editor.update(cx, |composer_editor, cx| {
 410                let text = composer_editor.text(cx);
 411                let id = self.next_message_id.post_inc();
 412                let body = cx.new_view(|cx| {
 413                    let mut editor = Editor::auto_height(80, cx);
 414                    editor.set_text(text, cx);
 415                    editor.set_soft_wrap_mode(SoftWrap::EditorWidth, cx);
 416                    editor
 417                });
 418                composer_editor.clear(cx);
 419
 420                ChatMessage::User(UserMessage {
 421                    id,
 422                    body,
 423                    attachments: Vec::new(),
 424                })
 425            });
 426            self.push_message(message, cx);
 427        } else {
 428            log::error!("unexpected state: no user message editor is focused.");
 429            return;
 430        }
 431
 432        let mode = *mode;
 433        self.pending_completion = Some(cx.spawn(move |this, mut cx| async move {
 434            let attachments_task = this.update(&mut cx, |this, cx| {
 435                let attachment_registry = this.attachment_registry.clone();
 436                attachment_registry.call_all_attachment_tools(cx)
 437            });
 438
 439            let attachments = maybe!(async {
 440                let attachments_task = attachments_task?;
 441                let attachments = attachments_task.await?;
 442
 443                anyhow::Ok(attachments)
 444            })
 445            .await
 446            .log_err()
 447            .unwrap_or_default();
 448
 449            // Set the attachments to the _last_ user message
 450            this.update(&mut cx, |this, _cx| {
 451                if let Some(ChatMessage::User(message)) = this.messages.last_mut() {
 452                    message.attachments = attachments;
 453                }
 454            })
 455            .log_err();
 456
 457            Self::request_completion(
 458                this.clone(),
 459                mode,
 460                MAX_COMPLETION_CALLS_PER_SUBMISSION,
 461                &mut cx,
 462            )
 463            .await
 464            .log_err();
 465
 466            this.update(&mut cx, |this, _cx| {
 467                this.pending_completion = None;
 468            })
 469            .context("Failed to push new user message")
 470            .log_err();
 471        }));
 472    }
 473
 474    async fn request_completion(
 475        this: WeakView<Self>,
 476        mode: SubmitMode,
 477        limit: usize,
 478        cx: &mut AsyncWindowContext,
 479    ) -> Result<()> {
 480        let mut call_count = 0;
 481        loop {
 482            let complete = async {
 483                let (tool_definitions, model_name, messages) = this.update(cx, |this, cx| {
 484                    this.push_new_assistant_message(cx);
 485
 486                    let definitions = if call_count < limit
 487                        && matches!(mode, SubmitMode::Codebase | SubmitMode::Simple)
 488                    {
 489                        this.tool_registry.definitions()
 490                    } else {
 491                        Vec::new()
 492                    };
 493                    call_count += 1;
 494
 495                    (
 496                        definitions,
 497                        this.model.clone(),
 498                        this.completion_messages(cx),
 499                    )
 500                })?;
 501
 502                let messages = messages.await?;
 503
 504                let completion = cx.update(|cx| {
 505                    CompletionProvider::get(cx).complete(
 506                        model_name,
 507                        messages,
 508                        Vec::new(),
 509                        1.0,
 510                        tool_definitions,
 511                    )
 512                });
 513
 514                let mut stream = completion?.await?;
 515                let mut body = String::new();
 516                while let Some(delta) = stream.next().await {
 517                    let delta = delta?;
 518                    this.update(cx, |this, cx| {
 519                        if let Some(ChatMessage::Assistant(AssistantMessage { messages, .. })) =
 520                            this.messages.last_mut()
 521                        {
 522                            if messages.is_empty() {
 523                                messages.push(AssistantMessagePart {
 524                                    body: RichText::default(),
 525                                    tool_calls: Vec::new(),
 526                                })
 527                            }
 528
 529                            let message = messages.last_mut().unwrap();
 530
 531                            if let Some(content) = &delta.content {
 532                                body.push_str(content);
 533                            }
 534
 535                            for tool_call_delta in delta.tool_calls {
 536                                let index = tool_call_delta.index as usize;
 537                                if index >= message.tool_calls.len() {
 538                                    message.tool_calls.resize_with(index + 1, Default::default);
 539                                }
 540                                let tool_call = &mut message.tool_calls[index];
 541
 542                                if let Some(id) = &tool_call_delta.id {
 543                                    tool_call.id.push_str(id);
 544                                }
 545
 546                                match tool_call_delta.variant {
 547                                    Some(proto::tool_call_delta::Variant::Function(
 548                                        tool_call_delta,
 549                                    )) => {
 550                                        this.tool_registry.update_tool_call(
 551                                            tool_call,
 552                                            tool_call_delta.name.as_deref(),
 553                                            tool_call_delta.arguments.as_deref(),
 554                                            cx,
 555                                        );
 556                                    }
 557                                    None => {}
 558                                }
 559                            }
 560
 561                            message.body =
 562                                RichText::new(body.clone(), &[], &this.language_registry);
 563                            cx.notify();
 564                        } else {
 565                            unreachable!()
 566                        }
 567                    })?;
 568                }
 569
 570                anyhow::Ok(())
 571            }
 572            .await;
 573
 574            let mut tool_tasks = Vec::new();
 575            this.update(cx, |this, cx| {
 576                if let Some(ChatMessage::Assistant(AssistantMessage {
 577                    error: message_error,
 578                    messages,
 579                    ..
 580                })) = this.messages.last_mut()
 581                {
 582                    if let Err(error) = complete {
 583                        message_error.replace(SharedString::from(error.to_string()));
 584                        cx.notify();
 585                    } else {
 586                        if let Some(current_message) = messages.last_mut() {
 587                            for tool_call in current_message.tool_calls.iter_mut() {
 588                                tool_tasks
 589                                    .extend(this.tool_registry.execute_tool_call(tool_call, cx));
 590                            }
 591                        }
 592                    }
 593                }
 594            })?;
 595
 596            // This ends recursion on calling for responses after tools
 597            if tool_tasks.is_empty() {
 598                return Ok(());
 599            }
 600
 601            join_all(tool_tasks.into_iter()).await;
 602        }
 603    }
 604
 605    fn push_new_assistant_message(&mut self, cx: &mut ViewContext<Self>) {
 606        // If the last message is a grouped assistant message, add to the grouped message
 607        if let Some(ChatMessage::Assistant(AssistantMessage { messages, .. })) =
 608            self.messages.last_mut()
 609        {
 610            messages.push(AssistantMessagePart {
 611                body: RichText::default(),
 612                tool_calls: Vec::new(),
 613            });
 614            return;
 615        }
 616
 617        let message = ChatMessage::Assistant(AssistantMessage {
 618            id: self.next_message_id.post_inc(),
 619            messages: vec![AssistantMessagePart {
 620                body: RichText::default(),
 621                tool_calls: Vec::new(),
 622            }],
 623            error: None,
 624        });
 625        self.push_message(message, cx);
 626    }
 627
 628    fn push_message(&mut self, message: ChatMessage, cx: &mut ViewContext<Self>) {
 629        let old_len = self.messages.len();
 630        let focus_handle = Some(message.focus_handle(cx));
 631        self.messages.push(message);
 632        self.list_state
 633            .splice_focusable(old_len..old_len, focus_handle);
 634        cx.notify();
 635    }
 636
 637    fn pop_message(&mut self, cx: &mut ViewContext<Self>) {
 638        if self.messages.is_empty() {
 639            return;
 640        }
 641
 642        self.messages.pop();
 643        self.list_state
 644            .splice(self.messages.len()..self.messages.len() + 1, 0);
 645        cx.notify();
 646    }
 647
 648    fn truncate_messages(&mut self, last_message_id: MessageId, cx: &mut ViewContext<Self>) {
 649        if let Some(index) = self.messages.iter().position(|message| match message {
 650            ChatMessage::User(message) => message.id == last_message_id,
 651            ChatMessage::Assistant(message) => message.id == last_message_id,
 652        }) {
 653            self.list_state.splice(index + 1..self.messages.len(), 0);
 654            self.messages.truncate(index + 1);
 655            cx.notify();
 656        }
 657    }
 658
 659    fn is_message_collapsed(&self, id: &MessageId) -> bool {
 660        self.collapsed_messages.get(id).copied().unwrap_or_default()
 661    }
 662
 663    fn toggle_message_collapsed(&mut self, id: MessageId) {
 664        let entry = self.collapsed_messages.entry(id).or_insert(false);
 665        *entry = !*entry;
 666    }
 667
 668    fn reset(&mut self) {
 669        self.messages.clear();
 670        self.list_state.reset(0);
 671        self.editing_message.take();
 672        self.collapsed_messages.clear();
 673    }
 674
 675    fn new_conversation(&mut self, cx: &mut ViewContext<Self>) {
 676        let messages = std::mem::take(&mut self.messages)
 677            .into_iter()
 678            .map(|message| self.serialize_message(message, cx))
 679            .collect::<Vec<_>>();
 680
 681        self.reset();
 682
 683        let title = messages
 684            .first()
 685            .map(|message| match message {
 686                SavedChatMessage::User { body, .. } => body.clone(),
 687                SavedChatMessage::Assistant { messages, .. } => messages
 688                    .first()
 689                    .map(|message| message.body.to_string())
 690                    .unwrap_or_default(),
 691            })
 692            .unwrap_or_else(|| "A conversation with the assistant.".to_string());
 693
 694        let saved_conversation = SavedConversation {
 695            version: "0.3.0".to_string(),
 696            title,
 697            messages,
 698        };
 699
 700        let discriminant = 1;
 701
 702        let path = CONVERSATIONS_DIR.join(&format!(
 703            "{title} - {discriminant}.zed.{version}.json",
 704            title = saved_conversation.title,
 705            version = saved_conversation.version
 706        ));
 707
 708        cx.spawn({
 709            let fs = self.fs.clone();
 710            |_this, _cx| async move {
 711                fs.create_dir(CONVERSATIONS_DIR.as_ref()).await?;
 712                fs.atomic_write(path, serde_json::to_string(&saved_conversation)?)
 713                    .await?;
 714
 715                anyhow::Ok(())
 716            }
 717        })
 718        .detach_and_log_err(cx);
 719    }
 720
 721    fn render_error(
 722        &self,
 723        error: Option<SharedString>,
 724        _ix: usize,
 725        cx: &mut ViewContext<Self>,
 726    ) -> AnyElement {
 727        let theme = cx.theme();
 728
 729        if let Some(error) = error {
 730            div()
 731                .py_1()
 732                .px_2()
 733                .mx_neg_1()
 734                .rounded_md()
 735                .border_1()
 736                .border_color(theme.status().error_border)
 737                // .bg(theme.status().error_background)
 738                .text_color(theme.status().error)
 739                .child(error.clone())
 740                .into_any_element()
 741        } else {
 742            div().into_any_element()
 743        }
 744    }
 745
 746    fn render_message(&self, ix: usize, cx: &mut ViewContext<Self>) -> AnyElement {
 747        let is_first = ix == 0;
 748        let is_last = ix == self.messages.len().saturating_sub(1);
 749
 750        let padding = Spacing::Large.rems(cx);
 751
 752        // Whenever there's a run of assistant messages, group as one Assistant UI element
 753
 754        match &self.messages[ix] {
 755            ChatMessage::User(UserMessage {
 756                id,
 757                body,
 758                attachments,
 759            }) => div()
 760                .id(SharedString::from(format!("message-{}-container", id.0)))
 761                .when(is_first, |this| this.pt(padding))
 762                .map(|element| {
 763                    if self.editing_message_id() == Some(*id) {
 764                        element.child(Composer::new(
 765                            body.clone(),
 766                            self.project_index_button.clone(),
 767                            self.active_file_button.clone(),
 768                            crate::ui::ModelSelector::new(
 769                                cx.view().downgrade(),
 770                                self.model.clone(),
 771                            )
 772                            .into_any_element(),
 773                        ))
 774                    } else {
 775                        element
 776                            .on_click(cx.listener({
 777                                let id = *id;
 778                                let body = body.clone();
 779                                move |assistant_chat, event: &ClickEvent, cx| {
 780                                    if event.up.click_count == 2 {
 781                                        assistant_chat.editing_message = Some(EditingMessage {
 782                                            id,
 783                                            body: body.clone(),
 784                                            old_body: body.read(cx).text(cx).into(),
 785                                        });
 786                                        body.focus_handle(cx).focus(cx);
 787                                    }
 788                                }
 789                            }))
 790                            .child(
 791                                crate::ui::ChatMessage::new(
 792                                    *id,
 793                                    UserOrAssistant::User(self.user_store.read(cx).current_user()),
 794                                    // todo!(): clean up the vec usage
 795                                    vec![
 796                                        RichText::new(
 797                                            body.read(cx).text(cx),
 798                                            &[],
 799                                            &self.language_registry,
 800                                        )
 801                                        .element(ElementId::from(id.0), cx),
 802                                        h_flex()
 803                                            .gap_2()
 804                                            .children(
 805                                                attachments
 806                                                    .iter()
 807                                                    .map(|attachment| attachment.view.clone()),
 808                                            )
 809                                            .into_any_element(),
 810                                    ],
 811                                    self.is_message_collapsed(id),
 812                                    Box::new(cx.listener({
 813                                        let id = *id;
 814                                        move |assistant_chat, _event, _cx| {
 815                                            assistant_chat.toggle_message_collapsed(id)
 816                                        }
 817                                    })),
 818                                )
 819                                // TODO: Wire up selections.
 820                                .selected(is_last),
 821                            )
 822                    }
 823                })
 824                .into_any(),
 825            ChatMessage::Assistant(AssistantMessage {
 826                id,
 827                messages,
 828                error,
 829                ..
 830            }) => {
 831                let mut message_elements = Vec::new();
 832
 833                for message in messages {
 834                    if !message.body.text.is_empty() {
 835                        message_elements.push(
 836                            div()
 837                                // todo!(): The element Id will need to be a combo of the base ID and the index within the grouping
 838                                .child(message.body.element(ElementId::from(id.0), cx))
 839                                .into_any_element(),
 840                        )
 841                    }
 842
 843                    let tools = message
 844                        .tool_calls
 845                        .iter()
 846                        .filter_map(|tool_call| self.tool_registry.render_tool_call(tool_call, cx))
 847                        .collect::<Vec<AnyElement>>();
 848
 849                    if !tools.is_empty() {
 850                        message_elements.push(div().children(tools).into_any_element())
 851                    }
 852                }
 853
 854                if message_elements.is_empty() {
 855                    message_elements.push(::ui::Label::new("Researching...").into_any_element())
 856                }
 857
 858                div()
 859                    .when(is_first, |this| this.pt(padding))
 860                    .child(
 861                        crate::ui::ChatMessage::new(
 862                            *id,
 863                            UserOrAssistant::Assistant,
 864                            message_elements,
 865                            self.is_message_collapsed(id),
 866                            Box::new(cx.listener({
 867                                let id = *id;
 868                                move |assistant_chat, _event, _cx| {
 869                                    assistant_chat.toggle_message_collapsed(id)
 870                                }
 871                            })),
 872                        )
 873                        // TODO: Wire up selections.
 874                        .selected(is_last),
 875                    )
 876                    .child(self.render_error(error.clone(), ix, cx))
 877                    .into_any()
 878            }
 879        }
 880    }
 881
 882    fn completion_messages(&self, cx: &mut WindowContext) -> Task<Result<Vec<CompletionMessage>>> {
 883        let project_index = self.project_index.read(cx);
 884        let project = project_index.project();
 885        let fs = project_index.fs();
 886
 887        let mut project_context = ProjectContext::new(project, fs);
 888        let mut completion_messages = Vec::new();
 889
 890        for message in &self.messages {
 891            match message {
 892                ChatMessage::User(UserMessage {
 893                    body, attachments, ..
 894                }) => {
 895                    for attachment in attachments {
 896                        if let Some(content) = attachment.generate(&mut project_context, cx) {
 897                            completion_messages.push(CompletionMessage::System { content });
 898                        }
 899                    }
 900
 901                    // Show user's message last so that the assistant is grounded in the user's request
 902                    completion_messages.push(CompletionMessage::User {
 903                        content: body.read(cx).text(cx),
 904                    });
 905                }
 906                ChatMessage::Assistant(AssistantMessage { messages, .. }) => {
 907                    for message in messages {
 908                        let body = message.body.clone();
 909
 910                        if body.text.is_empty() && message.tool_calls.is_empty() {
 911                            continue;
 912                        }
 913
 914                        let tool_calls_from_assistant = message
 915                            .tool_calls
 916                            .iter()
 917                            .map(|tool_call| ToolCall {
 918                                content: ToolCallContent::Function {
 919                                    function: FunctionContent {
 920                                        name: tool_call.name.clone(),
 921                                        arguments: tool_call.arguments.clone(),
 922                                    },
 923                                },
 924                                id: tool_call.id.clone(),
 925                            })
 926                            .collect();
 927
 928                        completion_messages.push(CompletionMessage::Assistant {
 929                            content: Some(body.text.to_string()),
 930                            tool_calls: tool_calls_from_assistant,
 931                        });
 932
 933                        for tool_call in &message.tool_calls {
 934                            // Every tool call _must_ have a result by ID, otherwise OpenAI will error.
 935                            let content = self.tool_registry.content_for_tool_call(
 936                                tool_call,
 937                                &mut project_context,
 938                                cx,
 939                            );
 940                            completion_messages.push(CompletionMessage::Tool {
 941                                content,
 942                                tool_call_id: tool_call.id.clone(),
 943                            });
 944                        }
 945                    }
 946                }
 947            }
 948        }
 949
 950        let system_message = project_context.generate_system_message(cx);
 951
 952        cx.background_executor().spawn(async move {
 953            let content = system_message.await?;
 954            completion_messages.insert(0, CompletionMessage::System { content });
 955            Ok(completion_messages)
 956        })
 957    }
 958
 959    fn serialize_message(
 960        &self,
 961        message: ChatMessage,
 962        cx: &mut ViewContext<AssistantChat>,
 963    ) -> SavedChatMessage {
 964        match message {
 965            ChatMessage::User(message) => SavedChatMessage::User {
 966                id: message.id,
 967                body: message.body.read(cx).text(cx),
 968                attachments: message
 969                    .attachments
 970                    .iter()
 971                    .map(|attachment| {
 972                        self.attachment_registry
 973                            .serialize_user_attachment(attachment)
 974                    })
 975                    .collect(),
 976            },
 977            ChatMessage::Assistant(message) => SavedChatMessage::Assistant {
 978                id: message.id,
 979                error: message.error,
 980                messages: message
 981                    .messages
 982                    .iter()
 983                    .map(|message| SavedAssistantMessagePart {
 984                        body: message.body.text.clone(),
 985                        tool_calls: message
 986                            .tool_calls
 987                            .iter()
 988                            .filter_map(|tool_call| {
 989                                self.tool_registry
 990                                    .serialize_tool_call(tool_call, cx)
 991                                    .log_err()
 992                            })
 993                            .collect(),
 994                    })
 995                    .collect(),
 996            },
 997        }
 998    }
 999}
1000
1001impl Render for AssistantChat {
1002    fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
1003        let header_height = Spacing::Small.rems(cx) * 2.0 + ButtonSize::Default.rems();
1004
1005        div()
1006            .relative()
1007            .flex_1()
1008            .v_flex()
1009            .key_context("AssistantChat")
1010            .on_action(cx.listener(Self::submit))
1011            .on_action(cx.listener(Self::cancel))
1012            .text_color(Color::Default.color(cx))
1013            .child(list(self.list_state.clone()).flex_1().pt(header_height))
1014            .child(
1015                h_flex()
1016                    .absolute()
1017                    .top_0()
1018                    .justify_between()
1019                    .w_full()
1020                    .h(header_height)
1021                    .p(Spacing::Small.rems(cx))
1022                    .child(
1023                        IconButton::new(
1024                            "toggle-saved-conversations",
1025                            if self.saved_conversations_open {
1026                                IconName::ChevronRight
1027                            } else {
1028                                IconName::ChevronLeft
1029                            },
1030                        )
1031                        .on_click(cx.listener(|this, _event, _cx| {
1032                            this.toggle_saved_conversations();
1033                        }))
1034                        .tooltip(move |cx| Tooltip::text("Switch Conversations", cx)),
1035                    )
1036                    .child(
1037                        h_flex()
1038                            .gap(Spacing::Large.rems(cx))
1039                            .child(
1040                                IconButton::new("new-conversation", IconName::Plus)
1041                                    .on_click(cx.listener(move |this, _event, cx| {
1042                                        this.new_conversation(cx);
1043                                    }))
1044                                    .tooltip(move |cx| Tooltip::text("New Conversation", cx)),
1045                            )
1046                            .child(
1047                                IconButton::new("assistant-menu", IconName::Menu)
1048                                    .disabled(true)
1049                                    .tooltip(move |cx| {
1050                                        Tooltip::text(
1051                                            "Coming soon – Assistant settings & controls",
1052                                            cx,
1053                                        )
1054                                    }),
1055                            ),
1056                    ),
1057            )
1058            .when(self.saved_conversations_open, |element| {
1059                element.child(
1060                    h_flex()
1061                        .absolute()
1062                        .top(header_height)
1063                        .w_full()
1064                        .child(self.saved_conversations.clone()),
1065                )
1066            })
1067            .child(Composer::new(
1068                self.composer_editor.clone(),
1069                self.project_index_button.clone(),
1070                self.active_file_button.clone(),
1071                crate::ui::ModelSelector::new(cx.view().downgrade(), self.model.clone())
1072                    .into_any_element(),
1073            ))
1074    }
1075}
1076
1077#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy, Serialize, Deserialize)]
1078pub struct MessageId(usize);
1079
1080impl MessageId {
1081    fn post_inc(&mut self) -> Self {
1082        let id = *self;
1083        self.0 += 1;
1084        id
1085    }
1086}
1087
1088enum ChatMessage {
1089    User(UserMessage),
1090    Assistant(AssistantMessage),
1091}
1092
1093impl ChatMessage {
1094    fn focus_handle(&self, cx: &AppContext) -> Option<FocusHandle> {
1095        match self {
1096            ChatMessage::User(UserMessage { body, .. }) => Some(body.focus_handle(cx)),
1097            ChatMessage::Assistant(_) => None,
1098        }
1099    }
1100}
1101
1102struct UserMessage {
1103    pub id: MessageId,
1104    pub body: View<Editor>,
1105    pub attachments: Vec<UserAttachment>,
1106}
1107
1108struct AssistantMessagePart {
1109    pub body: RichText,
1110    pub tool_calls: Vec<ToolFunctionCall>,
1111}
1112
1113struct AssistantMessage {
1114    pub id: MessageId,
1115    pub messages: Vec<AssistantMessagePart>,
1116    pub error: Option<SharedString>,
1117}