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