assistant2.rs

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