1184 lines · 42.9 KB · Rust
RawDownload
mod assistant_settings;mod attachments;mod completion_provider;mod saved_conversation;mod saved_conversations;mod tools;pub mod ui; use crate::saved_conversation::SavedConversationMetadata;use crate::ui::UserOrAssistant;use ::ui::{div, prelude::*, Color, Tooltip, ViewContext};use anyhow::{Context, Result};use assistant_tooling::{    AttachmentRegistry, ProjectContext, ToolFunctionCall, ToolRegistry, UserAttachment,};use attachments::ActiveEditorAttachmentTool;use client::{proto, Client, UserStore};use collections::HashMap;use completion_provider::*;use editor::Editor;use feature_flags::FeatureFlagAppExt as _;use file_icons::FileIcons;use fs::Fs;use futures::{future::join_all, StreamExt};use gpui::{    list, AnyElement, AppContext, AsyncWindowContext, ClickEvent, EventEmitter, FocusHandle,    FocusableView, ListAlignment, ListState, Model, ReadGlobal, Render, Task, UpdateGlobal, View,    WeakView,};use language::{language_settings::SoftWrap, LanguageRegistry};use markdown::{Markdown, MarkdownStyle};use open_ai::{FunctionContent, ToolCall, ToolCallContent};use saved_conversation::{SavedAssistantMessagePart, SavedChatMessage, SavedConversation};use saved_conversations::SavedConversations;use semantic_index::{CloudEmbeddingProvider, ProjectIndex, ProjectIndexDebugView, SemanticIndex};use serde::{Deserialize, Serialize};use settings::Settings;use std::sync::Arc;use tools::{AnnotationTool, CreateBufferTool, ProjectIndexTool};use ui::{ActiveFileButton, Composer, ProjectIndexButton};use util::paths::CONVERSATIONS_DIR;use util::{maybe, paths::EMBEDDINGS_DIR, ResultExt};use workspace::{    dock::{DockPosition, Panel, PanelEvent},    Workspace,}; pub use assistant_settings::AssistantSettings; const MAX_COMPLETION_CALLS_PER_SUBMISSION: usize = 5; #[derive(Eq, PartialEq, Copy, Clone, Deserialize)]pub struct Submit(SubmitMode); /// There are multiple different ways to submit a model request, represented by this enum.#[derive(Eq, PartialEq, Copy, Clone, Deserialize)]pub enum SubmitMode {    /// Only include the conversation.    Simple,    /// Send the current file as context.    CurrentFile,    /// Search the codebase and send relevant excerpts.    Codebase,} gpui::actions!(assistant2, [Cancel, ToggleFocus, DebugProjectIndex,]);gpui::impl_actions!(assistant2, [Submit]); pub fn init(client: Arc<Client>, cx: &mut AppContext) {    AssistantSettings::register(cx);     cx.spawn(|mut cx| {        let client = client.clone();        async move {            let embedding_provider = CloudEmbeddingProvider::new(client.clone());            let semantic_index = SemanticIndex::new(                EMBEDDINGS_DIR.join("semantic-index-db.0.mdb"),                Arc::new(embedding_provider),                &mut cx,            )            .await?;            cx.update(|cx| cx.set_global(semantic_index))        }    })    .detach();     cx.set_global(CompletionProvider::new(CloudCompletionProvider::new(        client,    )));     cx.observe_new_views(        |workspace: &mut Workspace, _cx: &mut ViewContext<Workspace>| {            workspace.register_action(|workspace, _: &ToggleFocus, cx| {                workspace.toggle_panel_focus::<AssistantPanel>(cx);            });            workspace.register_action(|workspace, _: &DebugProjectIndex, cx| {                if let Some(panel) = workspace.panel::<AssistantPanel>(cx) {                    let index = panel.read(cx).chat.read(cx).project_index.clone();                    let view = cx.new_view(|cx| ProjectIndexDebugView::new(index, cx));                    workspace.add_item_to_center(Box::new(view), cx);                }            });        },    )    .detach();} pub fn enabled(cx: &AppContext) -> bool {    cx.is_staff()} pub struct AssistantPanel {    chat: View<AssistantChat>,    width: Option<Pixels>,} impl AssistantPanel {    pub fn load(        workspace: WeakView<Workspace>,        cx: AsyncWindowContext,    ) -> Task<Result<View<Self>>> {        cx.spawn(|mut cx| async move {            let (app_state, project) = workspace.update(&mut cx, |workspace, _| {                (workspace.app_state().clone(), workspace.project().clone())            })?;             cx.new_view(|cx| {                let project_index = SemanticIndex::update_global(cx, |semantic_index, cx| {                    semantic_index.project_index(project.clone(), cx)                });                 // Used in tools to render file icons                cx.observe_global::<FileIcons>(|_, cx| {                    cx.notify();                })                .detach();                 let mut tool_registry = ToolRegistry::new();                tool_registry                    .register(ProjectIndexTool::new(project_index.clone()))                    .unwrap();                tool_registry                    .register(CreateBufferTool::new(workspace.clone(), project.clone()))                    .unwrap();                tool_registry                    .register(AnnotationTool::new(workspace.clone(), project.clone()))                    .unwrap();                 let mut attachment_registry = AttachmentRegistry::new();                attachment_registry                    .register(ActiveEditorAttachmentTool::new(workspace.clone(), cx));                 Self::new(                    project.read(cx).fs().clone(),                    app_state.languages.clone(),                    Arc::new(tool_registry),                    Arc::new(attachment_registry),                    app_state.user_store.clone(),                    project_index,                    workspace,                    cx,                )            })        })    }     #[allow(clippy::too_many_arguments)]    pub fn new(        fs: Arc<dyn Fs>,        language_registry: Arc<LanguageRegistry>,        tool_registry: Arc<ToolRegistry>,        attachment_registry: Arc<AttachmentRegistry>,        user_store: Model<UserStore>,        project_index: Model<ProjectIndex>,        workspace: WeakView<Workspace>,        cx: &mut ViewContext<Self>,    ) -> Self {        let chat = cx.new_view(|cx| {            AssistantChat::new(                fs,                language_registry,                tool_registry.clone(),                attachment_registry,                user_store,                project_index,                workspace,                cx,            )        });         Self { width: None, chat }    }} impl Render for AssistantPanel {    fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {        div()            .size_full()            .v_flex()            .bg(cx.theme().colors().panel_background)            .child(self.chat.clone())    }} impl Panel for AssistantPanel {    fn persistent_name() -> &'static str {        "AssistantPanelv2"    }     fn position(&self, _cx: &WindowContext) -> workspace::dock::DockPosition {        // todo!("Add a setting / use assistant settings")        DockPosition::Right    }     fn position_is_valid(&self, position: workspace::dock::DockPosition) -> bool {        matches!(position, DockPosition::Right)    }     fn set_position(&mut self, _: workspace::dock::DockPosition, _: &mut ViewContext<Self>) {        // Do nothing until we have a setting for this    }     fn size(&self, _cx: &WindowContext) -> Pixels {        self.width.unwrap_or(px(400.))    }     fn set_size(&mut self, size: Option<Pixels>, cx: &mut ViewContext<Self>) {        self.width = size;        cx.notify();    }     fn icon(&self, _cx: &WindowContext) -> Option<::ui::IconName> {        Some(IconName::ZedAssistant)    }     fn icon_tooltip(&self, _: &WindowContext) -> Option<&'static str> {        Some("Assistant Panel ✨")    }     fn toggle_action(&self) -> Box<dyn gpui::Action> {        Box::new(ToggleFocus)    }} impl EventEmitter<PanelEvent> for AssistantPanel {} impl FocusableView for AssistantPanel {    fn focus_handle(&self, cx: &AppContext) -> FocusHandle {        self.chat.read(cx).composer_editor.read(cx).focus_handle(cx)    }} pub struct AssistantChat {    model: String,    messages: Vec<ChatMessage>,    list_state: ListState,    fs: Arc<dyn Fs>,    language_registry: Arc<LanguageRegistry>,    composer_editor: View<Editor>,    saved_conversations: View<SavedConversations>,    saved_conversations_open: bool,    project_index_button: View<ProjectIndexButton>,    active_file_button: Option<View<ActiveFileButton>>,    user_store: Model<UserStore>,    next_message_id: MessageId,    collapsed_messages: HashMap<MessageId, bool>,    editing_message: Option<EditingMessage>,    pending_completion: Option<Task<()>>,    tool_registry: Arc<ToolRegistry>,    attachment_registry: Arc<AttachmentRegistry>,    project_index: Model<ProjectIndex>,    markdown_style: MarkdownStyle,} struct EditingMessage {    id: MessageId,    body: View<Editor>,} impl AssistantChat {    #[allow(clippy::too_many_arguments)]    fn new(        fs: Arc<dyn Fs>,        language_registry: Arc<LanguageRegistry>,        tool_registry: Arc<ToolRegistry>,        attachment_registry: Arc<AttachmentRegistry>,        user_store: Model<UserStore>,        project_index: Model<ProjectIndex>,        workspace: WeakView<Workspace>,        cx: &mut ViewContext<Self>,    ) -> Self {        let model = CompletionProvider::global(cx).default_model();        let view = cx.view().downgrade();        let list_state = ListState::new(            0,            ListAlignment::Bottom,            px(1024.),            move |ix, cx: &mut WindowContext| {                view.update(cx, |this, cx| this.render_message(ix, cx))                    .unwrap()            },        );         let project_index_button = cx.new_view(|cx| {            ProjectIndexButton::new(project_index.clone(), tool_registry.clone(), cx)        });         let active_file_button = match workspace.upgrade() {            Some(workspace) => {                Some(cx.new_view(                    |cx| ActiveFileButton::new(attachment_registry.clone(), workspace, cx), //                ))            }            _ => None,        };         let saved_conversations = cx.new_view(|cx| SavedConversations::new(cx));        cx.spawn({            let fs = fs.clone();            let saved_conversations = saved_conversations.downgrade();            |_assistant_chat, mut cx| async move {                let saved_conversation_metadata = SavedConversationMetadata::list(fs).await?;                 cx.update(|cx| {                    saved_conversations.update(cx, |this, cx| {                        this.init(saved_conversation_metadata, cx);                    })                })??;                 anyhow::Ok(())            }        })        .detach_and_log_err(cx);         Self {            model,            messages: Vec::new(),            composer_editor: cx.new_view(|cx| {                let mut editor = Editor::auto_height(80, cx);                editor.set_soft_wrap_mode(SoftWrap::EditorWidth, cx);                editor.set_placeholder_text("Send a message…", cx);                editor            }),            saved_conversations,            saved_conversations_open: false,            list_state,            user_store,            fs,            language_registry,            project_index_button,            active_file_button,            project_index,            next_message_id: MessageId(0),            editing_message: None,            collapsed_messages: HashMap::default(),            pending_completion: None,            attachment_registry,            tool_registry,            markdown_style: MarkdownStyle {                code_block: gpui::TextStyleRefinement {                    font_family: Some("Zed Mono".into()),                    color: Some(cx.theme().colors().editor_foreground),                    background_color: Some(cx.theme().colors().editor_background),                    ..Default::default()                },                inline_code: gpui::TextStyleRefinement {                    font_family: Some("Zed Mono".into()),                    // @nate: Could we add inline-code specific styles to the theme?                    color: Some(cx.theme().colors().editor_foreground),                    background_color: Some(cx.theme().colors().editor_background),                    ..Default::default()                },                rule_color: Color::Muted.color(cx),                block_quote_border_color: Color::Muted.color(cx),                block_quote: gpui::TextStyleRefinement {                    color: Some(Color::Muted.color(cx)),                    ..Default::default()                },                link: gpui::TextStyleRefinement {                    color: Some(Color::Accent.color(cx)),                    underline: Some(gpui::UnderlineStyle {                        thickness: px(1.),                        color: Some(Color::Accent.color(cx)),                        wavy: false,                    }),                    ..Default::default()                },                syntax: cx.theme().syntax().clone(),                selection_background_color: {                    let mut selection = cx.theme().players().local().selection;                    selection.fade_out(0.7);                    selection                },            },        }    }     fn message_for_id(&self, id: MessageId) -> Option<&ChatMessage> {        self.messages.iter().find(|message| match message {            ChatMessage::User(message) => message.id == id,            ChatMessage::Assistant(message) => message.id == id,        })    }     fn toggle_saved_conversations(&mut self) {        self.saved_conversations_open = !self.saved_conversations_open;    }     fn cancel(&mut self, _: &Cancel, cx: &mut ViewContext<Self>) {        // If we're currently editing a message, cancel the edit.        if self.editing_message.take().is_some() {            cx.notify();            return;        }         if self.pending_completion.take().is_some() {            if let Some(ChatMessage::Assistant(grouping)) = self.messages.last() {                if grouping.messages.is_empty() {                    self.pop_message(cx);                }            }            return;        }         cx.propagate();    }     fn submit(&mut self, Submit(mode): &Submit, cx: &mut ViewContext<Self>) {        if self.composer_editor.focus_handle(cx).is_focused(cx) {            // Don't allow multiple concurrent completions.            if self.pending_completion.is_some() {                cx.propagate();                return;            }             let message = self.composer_editor.update(cx, |composer_editor, cx| {                let text = composer_editor.text(cx);                let id = self.next_message_id.post_inc();                let body = cx.new_view(|cx| {                    Markdown::new(                        text,                        self.markdown_style.clone(),                        Some(self.language_registry.clone()),                        cx,                    )                });                composer_editor.clear(cx);                 ChatMessage::User(UserMessage {                    id,                    body,                    attachments: Vec::new(),                })            });            self.push_message(message, cx);        } else if let Some(editing_message) = self.editing_message.as_ref() {            let focus_handle = editing_message.body.focus_handle(cx);            if focus_handle.contains_focused(cx) {                if let Some(ChatMessage::User(user_message)) =                    self.message_for_id(editing_message.id)                {                    user_message.body.update(cx, |body, cx| {                        body.reset(editing_message.body.read(cx).text(cx), cx);                    });                }                 self.truncate_messages(editing_message.id, cx);                 self.pending_completion.take();                self.composer_editor.focus_handle(cx).focus(cx);                self.editing_message.take();            } else {                log::error!("unexpected state: no user message editor is focused.");                return;            }        } else {            log::error!("unexpected state: no user message editor is focused.");            return;        }         let mode = *mode;        self.pending_completion = Some(cx.spawn(move |this, mut cx| async move {            let attachments_task = this.update(&mut cx, |this, cx| {                let attachment_registry = this.attachment_registry.clone();                attachment_registry.call_all_attachment_tools(cx)            });             let attachments = maybe!(async {                let attachments_task = attachments_task?;                let attachments = attachments_task.await?;                 anyhow::Ok(attachments)            })            .await            .log_err()            .unwrap_or_default();             // Set the attachments to the _last_ user message            this.update(&mut cx, |this, _cx| {                if let Some(ChatMessage::User(message)) = this.messages.last_mut() {

Lines 1–500 of 1184.

Keyboard shortcuts

General

Show keyboard shortcuts
?
Focus the current search
/

Issues and pull requests

Open filters
f

Issues, pull requests and Inbox

Select next item
j↓
Select previous item
k↑
Open selected item
Enter
Dismiss filters or clear selection
Esc

Repository code

Go to file
t