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