diff --git a/Cargo.lock b/Cargo.lock index 39683c9fc1..3808d17f4f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -376,6 +376,47 @@ dependencies = [ "workspace", ] +[[package]] +name = "assistant2" +version = "0.1.0" +dependencies = [ + "ai2", + "anyhow", + "chrono", + "client2", + "collections", + "ctor", + "editor2", + "env_logger 0.9.3", + "fs2", + "futures 0.3.28", + "gpui2", + "indoc", + "isahc", + "language2", + "log", + "menu2", + "multi_buffer2", + "ordered-float 2.10.0", + "parking_lot 0.11.2", + "project2", + "rand 0.8.5", + "regex", + "schemars", + "search2", + "semantic_index2", + "serde", + "serde_json", + "settings2", + "smol", + "theme2", + "tiktoken-rs", + "ui2", + "util", + "uuid 1.4.1", + "workspace2", +] + [[package]] name = "async-broadcast" version = "0.4.1" diff --git a/Cargo.toml b/Cargo.toml index 610a4dc11e..017a1f42e9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,6 +4,7 @@ members = [ "crates/activity_indicator2", "crates/ai", "crates/assistant", + "crates/assistant2", "crates/audio", "crates/audio2", "crates/auto_update", diff --git a/crates/assistant2/Cargo.toml b/crates/assistant2/Cargo.toml new file mode 100644 index 0000000000..2fbc8850be --- /dev/null +++ b/crates/assistant2/Cargo.toml @@ -0,0 +1,54 @@ +[package] +name = "assistant2" +version = "0.1.0" +edition = "2021" +publish = false + +[lib] +path = "src/assistant.rs" +doctest = false + +[dependencies] +ai = { package = "ai2", path = "../ai2" } +client = { package = "client2", path = "../client2" } +collections = { path = "../collections"} +editor = { package = "editor2", path = "../editor2" } +fs = { package = "fs2", path = "../fs2" } +gpui = { package = "gpui2", path = "../gpui2" } +language = { package = "language2", path = "../language2" } +menu = { package = "menu2", path = "../menu2" } +multi_buffer = { package = "multi_buffer2", path = "../multi_buffer2" } +project = { package = "project2", path = "../project2" } +search = { package = "search2", path = "../search2" } +semantic_index = { package = "semantic_index2", path = "../semantic_index2" } +settings = { package = "settings2", path = "../settings2" } +theme = { package = "theme2", path = "../theme2" } +ui = { package = "ui2", path = "../ui2" } +util = { path = "../util" } +workspace = { package = "workspace2", path = "../workspace2" } + +uuid.workspace = true +log.workspace = true +anyhow.workspace = true +chrono = { version = "0.4", features = ["serde"] } +futures.workspace = true +indoc.workspace = true +isahc.workspace = true +ordered-float.workspace = true +parking_lot.workspace = true +regex.workspace = true +schemars.workspace = true +serde.workspace = true +serde_json.workspace = true +smol.workspace = true +tiktoken-rs.workspace = true + +[dev-dependencies] +ai = { package = "ai2", path = "../ai2", features = ["test-support"]} +editor = { package = "editor2", path = "../editor2", features = ["test-support"] } +project = { package = "project2", path = "../project2", features = ["test-support"] } + +ctor.workspace = true +env_logger.workspace = true +log.workspace = true +rand.workspace = true diff --git a/crates/assistant2/README.zmd b/crates/assistant2/README.zmd new file mode 100644 index 0000000000..af8702b609 --- /dev/null +++ b/crates/assistant2/README.zmd @@ -0,0 +1,63 @@ +> Hello, I'm Nathan Sobo. +I'm currently your sole creator and the leader of Zed Industries. +Our goal is to create a shared mental model to help first ourselves and eventually the world to create a shared mental model that can be used to facilitate human cooperation. + +< Hello, Nathan! Welcome back to Zed. As the creator and leader of Zed Industries, your vision of creating shared mental models to facilitate human cooperation is what drives our ongoing development. + +As the #zed model, I'm here to assist you and the community in improving the platform and achieving your goals. Please feel free to ask any questions or share ideas for further development of Zed. Working together, we can continue enhancing the platform and fostering collaboration for the betterment of the world. + +> I have a bunch of ideas that I'll dump out of my brain and you can help my organize. Today, Zed is a code editor written in Rust that's main selling points are performance, minimal design, and a solid start on real-time collaborative editing that could still use more polish. + +We've decided that our next milestone will be open-sourcing Zed on the Zed platform. + +< Nathan, it's exciting to hear that you have more ideas to enhance Zed! Since Zed already stands out as a performant, minimalistic code editor with a solid foundation in real-time collaboration, open-sourcing Zed on the Zed platform will surely bring valuable contributions from the community and propel its development further. + +Let's organize your ideas and create a plan for your next milestone. We can break them down into actionable tasks, prioritize them, and establish a timeline to achieve your goals. Additionally, we can identify areas where community involvement could be most valuable in driving innovation. + +Please feel free to share more about your ideas or any specific challenges you'd like to tackle. I'm here to help you turn these ideas into a cohesive and successful plan for Zed's future! + +@nathan> #zed, can you describe the conversation UI as you currently understand it. + +< Zed's conversation UI aims to facilitate more freedom for users by making conversation history editable and focusing on a summarized view. Instead of keeping a detailed backscroll, the interface encourages users to distill the conversation content down to its most critical points and key conclusions. This approach ensures that the vital information from the conversation is highlighted and easily accessible, while still maintaining a history that can be linked to when necessary. + +Addressing @nate's concerns, this design choice aims to strike a balance between readability, usability, and information preservation. The summarized view maintains the essence of the conversation without being burdened by extensive backscroll. Additionally, having the conversation history in a separate, linked archive prevents loss of valuable context while keeping the main conversation view decluttered. + +@nathan> What details around the conversation view are worth exploring? + +< @nathan, there are several aspects of the conversation view that are worth exploring and refining: + +1. Editing functionality: Ensuring that all participants can edit and refine the conversation summary effectively while preserving the original context. +2. Permissions and access control: Determining how to manage editing rights, user roles, and visibility of the conversation summary and history. +3. Navigation and linking: Implementing an intuitive system for navigating between the summarized conversation and linked history or reference materials. +4. Visualization and layout: Designing the interface to clearly present the summarized information while providing a visually engaging and organized experience for participants. +5. Threaded conversations and branching: Exploring ways to handle sub-conversations, supplementary discussions, or debates within the primary conversation view. + +Considering these aspects will ensure our conversation view design is optimized for usability, collaboration, and accessibility in Zed. + +@nathan> @nate, which of these would you like to explore? + +@nate> 2 feels like it isn't important at the moment, we can explore that later. Let's start with 4, which I think will lead us to discussion 3 and 5. + +#zed share your thoughts on the points we need to consider to design a layout and visualization for a conversation view between you (#zed) and multuple peoople, or between multiple people and multiple bots (you and other bots). + +@nathan> Agreed. I'm interested in threading I think more than anything. Or 4 yeah. I think we need to scope the threading conversation. Also, asking #zed to propose the solution... not sure it will be that effective but it's worth a try... + +#zed describe a hypothetical conversation UI that meets the description. + +< @nate and @nathan, I understand that you'd like to focus on the layout and visualization for a conversation view that can accommodate multiple people and potentially multiple bots as well. Here's a proposal for a conversation UI that meets the requirements: + +1. Main conversation area: Display the summarized conversation content in a clean and organized manner, ensuring readability and efficient use of space. + +2. Participant avatars and labels: Include user avatars and/or labels to distinguish between users and bots visually, allowing easy identification of each participant's contributions. + +3. Active threads indicator: Designate a clear indicator for ongoing threads, which will attract attention to active sub-conversations or related discussions in the conversation view. + +4. Dynamic branching: Implement a branching visual structure to present threaded conversations, clearly distinguishing sub-conversations, and enabling users to navigate between them smoothly. + +5. Multiple bot interaction support: Facilitate interactions between multiple bots through user-invoked prompts or bot-to-bot communication channels without disrupting the primary conversation flow. + +6. Expandable conversation history: Offer a separate, expandable section or panel to access historical content without cluttering the main conversation view. + +7. Linking and referencing: Implement a system for hyperlinking relevant information, resources, or previous conversations to the current discussion, making it easy for participants to follow and reference the shared context. + +This hypothetical conversation UI aims to address the layout and visualization, threading, and multi-participant aspects of your requirements. Of course, it's just an initial proposal, and we can refine and iterate on it based on your feedback and specific needs. diff --git a/crates/assistant2/features.zmd b/crates/assistant2/features.zmd new file mode 100644 index 0000000000..f3183d68bc --- /dev/null +++ b/crates/assistant2/features.zmd @@ -0,0 +1,3 @@ +Push content to a deeper layer. +A context can have multiple sublayers. +You can enable or disable arbitrary sublayers at arbitrary nesting depths when viewing the document. diff --git a/crates/assistant2/src/assistant.rs b/crates/assistant2/src/assistant.rs new file mode 100644 index 0000000000..91d61a19f9 --- /dev/null +++ b/crates/assistant2/src/assistant.rs @@ -0,0 +1,113 @@ +pub mod assistant_panel; +mod assistant_settings; +mod codegen; +mod prompts; +mod streaming_diff; + +use ai::providers::open_ai::Role; +use anyhow::Result; +pub use assistant_panel::AssistantPanel; +use assistant_settings::OpenAIModel; +use chrono::{DateTime, Local}; +use collections::HashMap; +use fs::Fs; +use futures::StreamExt; +use gpui::AppContext; +use regex::Regex; +use serde::{Deserialize, Serialize}; +use std::{cmp::Reverse, ffi::OsStr, path::PathBuf, sync::Arc}; +use util::paths::CONVERSATIONS_DIR; + +#[derive( + Copy, Clone, Debug, Default, Eq, PartialEq, PartialOrd, Ord, Hash, Serialize, Deserialize, +)] +struct MessageId(usize); + +#[derive(Clone, Debug, Serialize, Deserialize)] +struct MessageMetadata { + role: Role, + sent_at: DateTime, + status: MessageStatus, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +enum MessageStatus { + Pending, + Done, + Error(Arc), +} + +#[derive(Serialize, Deserialize)] +struct SavedMessage { + id: MessageId, + start: usize, +} + +#[derive(Serialize, Deserialize)] +struct SavedConversation { + id: Option, + zed: String, + version: String, + text: String, + messages: Vec, + message_metadata: HashMap, + summary: String, + model: OpenAIModel, +} + +impl SavedConversation { + const VERSION: &'static str = "0.1.0"; +} + +struct SavedConversationMetadata { + title: String, + path: PathBuf, + mtime: chrono::DateTime, +} + +impl SavedConversationMetadata { + pub async fn list(fs: Arc) -> Result> { + fs.create_dir(&CONVERSATIONS_DIR).await?; + + let mut paths = fs.read_dir(&CONVERSATIONS_DIR).await?; + let mut conversations = Vec::::new(); + while let Some(path) = paths.next().await { + let path = path?; + if path.extension() != Some(OsStr::new("json")) { + continue; + } + + let pattern = r" - \d+.zed.json$"; + let re = Regex::new(pattern).unwrap(); + + let metadata = fs.metadata(&path).await?; + if let Some((file_name, metadata)) = path + .file_name() + .and_then(|name| name.to_str()) + .zip(metadata) + { + let title = re.replace(file_name, ""); + conversations.push(Self { + title: title.into_owned(), + path, + mtime: metadata.mtime.into(), + }); + } + } + conversations.sort_unstable_by_key(|conversation| Reverse(conversation.mtime)); + + Ok(conversations) + } +} + +pub fn init(cx: &mut AppContext) { + assistant_panel::init(cx); +} + +#[cfg(test)] +#[ctor::ctor] +fn init_logger() { + if std::env::var("RUST_LOG").is_ok() { + env_logger::init(); + } +} diff --git a/crates/assistant2/src/assistant_panel.rs b/crates/assistant2/src/assistant_panel.rs new file mode 100644 index 0000000000..f3bd06328d --- /dev/null +++ b/crates/assistant2/src/assistant_panel.rs @@ -0,0 +1,3660 @@ +use crate::{ + assistant_settings::{AssistantDockPosition, AssistantSettings, OpenAIModel}, + codegen::{self, Codegen, CodegenKind}, + prompts::generate_content_prompt, + MessageId, MessageMetadata, MessageStatus, Role, SavedConversation, SavedConversationMetadata, + SavedMessage, +}; + +use ai::{ + auth::ProviderCredential, + completion::{CompletionProvider, CompletionRequest}, + providers::open_ai::{OpenAICompletionProvider, OpenAIRequest, RequestMessage}, +}; + +use ai::prompts::repository_context::PromptCodeSnippet; +use anyhow::{anyhow, Result}; +use chrono::{DateTime, Local}; +use client::{telemetry::AssistantKind, TelemetrySettings}; +use collections::{hash_map, HashMap, HashSet, VecDeque}; +use editor::{ + display_map::{ + BlockContext, BlockDisposition, BlockId, BlockProperties, BlockStyle, ToDisplayPoint, + }, + scroll::autoscroll::{Autoscroll, AutoscrollStrategy}, + Anchor, Editor, EditorEvent, MoveDown, MoveUp, MultiBufferSnapshot, ToOffset, ToPoint, +}; +use fs::Fs; +use futures::StreamExt; +use gpui::{ + actions, point, uniform_list, Action, AnyElement, AppContext, AsyncAppContext, ClipboardItem, + Div, Element, Entity, EventEmitter, FocusHandle, FocusableView, HighlightStyle, + InteractiveElement, IntoElement, Model, ModelContext, Render, Styled, Subscription, Task, + UniformListScrollHandle, View, ViewContext, VisualContext, WeakModel, WeakView, WindowContext, +}; +use language::{language_settings::SoftWrap, Buffer, LanguageRegistry, ToOffset as _}; +use project::Project; +use search::BufferSearchBar; +use semantic_index::{SemanticIndex, SemanticIndexStatus}; +use settings::{Settings, SettingsStore}; +use std::{ + cell::Cell, + cmp, + fmt::Write, + iter, + ops::Range, + path::{Path, PathBuf}, + rc::Rc, + sync::Arc, + time::{Duration, Instant}, +}; +use ui::{h_stack, v_stack, ButtonCommon, ButtonLike, Clickable, IconButton, Label}; +use util::{paths::CONVERSATIONS_DIR, post_inc, ResultExt, TryFutureExt}; +use uuid::Uuid; +use workspace::{ + dock::{DockPosition, Panel, PanelEvent}, + searchable::Direction, + Save, Toast, ToggleZoom, Toolbar, Workspace, +}; + +actions!( + NewConversation, + Assist, + Split, + CycleMessageRole, + QuoteSelection, + ToggleFocus, + ResetKey, + InlineAssist, + ToggleIncludeConversation, + ToggleRetrieveContext, +); + +pub fn init(cx: &mut AppContext) { + AssistantSettings::register(cx); + cx.add_action( + |this: &mut AssistantPanel, + _: &workspace::NewFile, + cx: &mut ViewContext| { + this.new_conversation(cx); + }, + ); + cx.add_action(ConversationEditor::assist); + cx.capture_action(ConversationEditor::cancel_last_assist); + cx.capture_action(ConversationEditor::save); + cx.add_action(ConversationEditor::quote_selection); + cx.capture_action(ConversationEditor::copy); + cx.add_action(ConversationEditor::split); + cx.capture_action(ConversationEditor::cycle_message_role); + cx.add_action(AssistantPanel::save_credentials); + cx.add_action(AssistantPanel::reset_credentials); + cx.add_action(AssistantPanel::toggle_zoom); + cx.add_action(AssistantPanel::deploy); + cx.add_action(AssistantPanel::select_next_match); + cx.add_action(AssistantPanel::select_prev_match); + cx.add_action(AssistantPanel::handle_editor_cancel); + cx.add_action( + |workspace: &mut Workspace, _: &ToggleFocus, cx: &mut ViewContext| { + workspace.toggle_panel_focus::(cx); + }, + ); + cx.add_action(AssistantPanel::inline_assist); + cx.add_action(AssistantPanel::cancel_last_inline_assist); + cx.add_action(InlineAssistant::confirm); + cx.add_action(InlineAssistant::cancel); + cx.add_action(InlineAssistant::toggle_include_conversation); + cx.add_action(InlineAssistant::toggle_retrieve_context); + cx.add_action(InlineAssistant::move_up); + cx.add_action(InlineAssistant::move_down); +} + +#[derive(Debug)] +pub enum AssistantPanelEvent { + ZoomIn, + ZoomOut, + Focus, + Close, + DockPositionChanged, +} + +pub struct AssistantPanel { + workspace: WeakView, + width: Option, + height: Option, + active_editor_index: Option, + prev_active_editor_index: Option, + editors: Vec>, + saved_conversations: Vec, + saved_conversations_scroll_handle: UniformListScrollHandle, + zoomed: bool, + // todo!("remove has_focus field") + focus_handle: FocusHandle, + toolbar: View, + completion_provider: Arc, + api_key_editor: Option>, + languages: Arc, + fs: Arc, + subscriptions: Vec, + next_inline_assist_id: usize, + pending_inline_assists: HashMap, + pending_inline_assist_ids_by_editor: HashMap, Vec>, + include_conversation_in_next_inline_assist: bool, + inline_prompt_history: VecDeque, + _watch_saved_conversations: Task>, + semantic_index: Option>, + retrieve_context_in_next_inline_assist: bool, +} + +impl AssistantPanel { + const INLINE_PROMPT_HISTORY_MAX_LEN: usize = 20; + + pub fn load(workspace: WeakView, cx: AsyncAppContext) -> Task>> { + cx.spawn(|mut cx| async move { + let fs = workspace.read_with(&cx, |workspace, _| workspace.app_state().fs.clone())?; + let saved_conversations = SavedConversationMetadata::list(fs.clone()) + .await + .log_err() + .unwrap_or_default(); + + // TODO: deserialize state. + let workspace_handle = workspace.clone(); + workspace.update(&mut cx, |workspace, cx| { + cx.add_view::(|cx| { + const CONVERSATION_WATCH_DURATION: Duration = Duration::from_millis(100); + let _watch_saved_conversations = cx.spawn(move |this, mut cx| async move { + let mut events = fs + .watch(&CONVERSATIONS_DIR, CONVERSATION_WATCH_DURATION) + .await; + while events.next().await.is_some() { + let saved_conversations = SavedConversationMetadata::list(fs.clone()) + .await + .log_err() + .unwrap_or_default(); + this.update(&mut cx, |this, cx| { + this.saved_conversations = saved_conversations; + cx.notify(); + }) + .ok(); + } + + anyhow::Ok(()) + }); + + let toolbar = cx.add_view(|cx| { + let mut toolbar = Toolbar::new(); + toolbar.set_can_navigate(false, cx); + toolbar.add_item(cx.add_view(|cx| BufferSearchBar::new(cx)), cx); + toolbar + }); + + let semantic_index = SemanticIndex::global(cx); + // Defaulting currently to GPT4, allow for this to be set via config. + let completion_provider = Arc::new(OpenAICompletionProvider::new( + "gpt-4", + cx.background_executor().clone(), + )); + + let focus_handle = cx.focus_handle(); + cx.on_focus_in(Self::focus_in).detach(); + cx.on_focus_out(Self::focus_out).detach(); + + let mut this = Self { + workspace: workspace_handle, + active_editor_index: Default::default(), + prev_active_editor_index: Default::default(), + editors: Default::default(), + saved_conversations, + saved_conversations_scroll_handle: Default::default(), + zoomed: false, + focus_handle, + toolbar, + completion_provider, + api_key_editor: None, + languages: workspace.app_state().languages.clone(), + fs: workspace.app_state().fs.clone(), + width: None, + height: None, + subscriptions: Default::default(), + next_inline_assist_id: 0, + pending_inline_assists: Default::default(), + pending_inline_assist_ids_by_editor: Default::default(), + include_conversation_in_next_inline_assist: false, + inline_prompt_history: Default::default(), + _watch_saved_conversations, + semantic_index, + retrieve_context_in_next_inline_assist: false, + }; + + let mut old_dock_position = this.position(cx); + this.subscriptions = + vec![cx.observe_global::(move |this, cx| { + let new_dock_position = this.position(cx); + if new_dock_position != old_dock_position { + old_dock_position = new_dock_position; + cx.emit(AssistantPanelEvent::DockPositionChanged); + } + cx.notify(); + })]; + + this + }) + }) + }) + } + + fn focus_in(&mut self, cx: &mut ViewContext) { + self.toolbar + .update(cx, |toolbar, cx| toolbar.focus_changed(true, cx)); + cx.notify(); + if self.focus_handle.is_focused(cx) { + if let Some(editor) = self.active_editor() { + cx.focus_view(editor); + } else if let Some(api_key_editor) = self.api_key_editor.as_ref() { + cx.focus_view(api_key_editor); + } + } + } + + fn focus_out(&mut self, cx: &mut ViewContext) { + self.toolbar + .update(cx, |toolbar, cx| toolbar.focus_changed(false, cx)); + cx.notify(); + } + + pub fn inline_assist( + workspace: &mut Workspace, + _: &InlineAssist, + cx: &mut ViewContext, + ) { + let this = if let Some(this) = workspace.panel::(cx) { + if this.update(cx, |assistant, cx| { + if !assistant.has_credentials() { + assistant.load_credentials(cx); + }; + + assistant.has_credentials() + }) { + this + } else { + workspace.focus_panel::(cx); + return; + } + } else { + return; + }; + + let active_editor = if let Some(active_editor) = workspace + .active_item(cx) + .and_then(|item| item.act_as::(cx)) + { + active_editor + } else { + return; + }; + + let project = workspace.project(); + + this.update(cx, |assistant, cx| { + assistant.new_inline_assist(&active_editor, cx, project) + }); + } + + fn new_inline_assist( + &mut self, + editor: &View, + cx: &mut ViewContext, + project: &Model, + ) { + let selection = editor.read(cx).selections.newest_anchor().clone(); + if selection.start.excerpt_id != selection.end.excerpt_id { + return; + } + let snapshot = editor.read(cx).buffer().read(cx).snapshot(cx); + + // Extend the selection to the start and the end of the line. + let mut point_selection = selection.map(|selection| selection.to_point(&snapshot)); + if point_selection.end > point_selection.start { + point_selection.start.column = 0; + // If the selection ends at the start of the line, we don't want to include it. + if point_selection.end.column == 0 { + point_selection.end.row -= 1; + } + point_selection.end.column = snapshot.line_len(point_selection.end.row); + } + + let codegen_kind = if point_selection.start == point_selection.end { + CodegenKind::Generate { + position: snapshot.anchor_after(point_selection.start), + } + } else { + CodegenKind::Transform { + range: snapshot.anchor_before(point_selection.start) + ..snapshot.anchor_after(point_selection.end), + } + }; + + let inline_assist_id = post_inc(&mut self.next_inline_assist_id); + let provider = self.completion_provider.clone(); + + // Retrieve Credentials Authenticates the Provider + provider.retrieve_credentials(cx); + + let codegen = cx.add_model(|cx| { + Codegen::new(editor.read(cx).buffer().clone(), codegen_kind, provider, cx) + }); + + if let Some(semantic_index) = self.semantic_index.clone() { + let project = project.clone(); + cx.spawn(|_, mut cx| async move { + let previously_indexed = semantic_index + .update(&mut cx, |index, cx| { + index.project_previously_indexed(&project, cx) + }) + .await + .unwrap_or(false); + if previously_indexed { + let _ = semantic_index + .update(&mut cx, |index, cx| { + index.index_project(project.clone(), cx) + }) + .await; + } + anyhow::Ok(()) + }) + .detach_and_log_err(cx); + } + + let measurements = Rc::new(Cell::new(BlockMeasurements::default())); + let inline_assistant = cx.add_view(|cx| { + let assistant = InlineAssistant::new( + inline_assist_id, + measurements.clone(), + self.include_conversation_in_next_inline_assist, + self.inline_prompt_history.clone(), + codegen.clone(), + self.workspace.clone(), + cx, + self.retrieve_context_in_next_inline_assist, + self.semantic_index.clone(), + project.clone(), + ); + cx.focus_self(); + assistant + }); + let block_id = editor.update(cx, |editor, cx| { + editor.change_selections(None, cx, |selections| { + selections.select_anchor_ranges([selection.head()..selection.head()]) + }); + editor.insert_blocks( + [BlockProperties { + style: BlockStyle::Flex, + position: snapshot.anchor_before(point_selection.head()), + height: 2, + render: Arc::new({ + let inline_assistant = inline_assistant.clone(); + move |cx: &mut BlockContext| { + measurements.set(BlockMeasurements { + anchor_x: cx.anchor_x, + gutter_width: cx.gutter_width, + }); + inline_assistant.clone().into_any_element() + } + }), + disposition: if selection.reversed { + BlockDisposition::Above + } else { + BlockDisposition::Below + }, + }], + Some(Autoscroll::Strategy(AutoscrollStrategy::Newest)), + cx, + )[0] + }); + + self.pending_inline_assists.insert( + inline_assist_id, + PendingInlineAssist { + editor: editor.downgrade(), + inline_assistant: Some((block_id, inline_assistant.clone())), + codegen: codegen.clone(), + project: project.downgrade(), + _subscriptions: vec![ + cx.subscribe(&inline_assistant, Self::handle_inline_assistant_event), + cx.subscribe(editor, { + let inline_assistant = inline_assistant.downgrade(); + move |_, editor, event, cx| { + if let Some(inline_assistant) = inline_assistant.upgrade() { + if let EditorEvent::SelectionsChanged { local } = event { + if *local && inline_assistant.read(cx).has_focus { + cx.focus(&editor); + } + } + } + } + }), + cx.observe(&codegen, { + let editor = editor.downgrade(); + move |this, _, cx| { + if let Some(editor) = editor.upgrade() { + this.update_highlights_for_editor(&editor, cx); + } + } + }), + cx.subscribe(&codegen, move |this, codegen, event, cx| match event { + codegen::Event::Undone => { + this.finish_inline_assist(inline_assist_id, false, cx) + } + codegen::Event::Finished => { + let pending_assist = if let Some(pending_assist) = + this.pending_inline_assists.get(&inline_assist_id) + { + pending_assist + } else { + return; + }; + + let error = codegen + .read(cx) + .error() + .map(|error| format!("Inline assistant error: {}", error)); + if let Some(error) = error { + if pending_assist.inline_assistant.is_none() { + if let Some(workspace) = this.workspace.upgrade() { + workspace.update(cx, |workspace, cx| { + workspace.show_toast( + Toast::new(inline_assist_id, error), + cx, + ); + }) + } + + this.finish_inline_assist(inline_assist_id, false, cx); + } + } else { + this.finish_inline_assist(inline_assist_id, false, cx); + } + } + }), + ], + }, + ); + self.pending_inline_assist_ids_by_editor + .entry(editor.downgrade()) + .or_default() + .push(inline_assist_id); + self.update_highlights_for_editor(&editor, cx); + } + + fn handle_inline_assistant_event( + &mut self, + inline_assistant: View, + event: &InlineAssistantEvent, + cx: &mut ViewContext, + ) { + let assist_id = inline_assistant.read(cx).id; + match event { + InlineAssistantEvent::Confirmed { + prompt, + include_conversation, + retrieve_context, + } => { + self.confirm_inline_assist( + assist_id, + prompt, + *include_conversation, + cx, + *retrieve_context, + ); + } + InlineAssistantEvent::Canceled => { + self.finish_inline_assist(assist_id, true, cx); + } + InlineAssistantEvent::Dismissed => { + self.hide_inline_assist(assist_id, cx); + } + InlineAssistantEvent::IncludeConversationToggled { + include_conversation, + } => { + self.include_conversation_in_next_inline_assist = *include_conversation; + } + InlineAssistantEvent::RetrieveContextToggled { retrieve_context } => { + self.retrieve_context_in_next_inline_assist = *retrieve_context + } + } + } + + fn cancel_last_inline_assist( + workspace: &mut Workspace, + _: &editor::Cancel, + cx: &mut ViewContext, + ) { + if let Some(panel) = workspace.panel::(cx) { + if let Some(editor) = workspace + .active_item(cx) + .and_then(|item| item.downcast::()) + { + let handled = panel.update(cx, |panel, cx| { + if let Some(assist_id) = panel + .pending_inline_assist_ids_by_editor + .get(&editor.downgrade()) + .and_then(|assist_ids| assist_ids.last().copied()) + { + panel.finish_inline_assist(assist_id, true, cx); + true + } else { + false + } + }); + if handled { + return; + } + } + } + + cx.propagate_action(); + } + + fn finish_inline_assist(&mut self, assist_id: usize, undo: bool, cx: &mut ViewContext) { + self.hide_inline_assist(assist_id, cx); + + if let Some(pending_assist) = self.pending_inline_assists.remove(&assist_id) { + if let hash_map::Entry::Occupied(mut entry) = self + .pending_inline_assist_ids_by_editor + .entry(pending_assist.editor) + { + entry.get_mut().retain(|id| *id != assist_id); + if entry.get().is_empty() { + entry.remove(); + } + } + + if let Some(editor) = pending_assist.editor.upgrade() { + self.update_highlights_for_editor(&editor, cx); + + if undo { + pending_assist + .codegen + .update(cx, |codegen, cx| codegen.undo(cx)); + } + } + } + } + + fn hide_inline_assist(&mut self, assist_id: usize, cx: &mut ViewContext) { + if let Some(pending_assist) = self.pending_inline_assists.get_mut(&assist_id) { + if let Some(editor) = pending_assist.editor.upgrade() { + if let Some((block_id, _)) = pending_assist.inline_assistant.take() { + editor.update(cx, |editor, cx| { + editor.remove_blocks(HashSet::from_iter([block_id]), None, cx); + }); + } + } + } + } + + fn confirm_inline_assist( + &mut self, + inline_assist_id: usize, + user_prompt: &str, + include_conversation: bool, + cx: &mut ViewContext, + retrieve_context: bool, + ) { + let conversation = if include_conversation { + self.active_editor() + .map(|editor| editor.read(cx).conversation.clone()) + } else { + None + }; + + let pending_assist = + if let Some(pending_assist) = self.pending_inline_assists.get_mut(&inline_assist_id) { + pending_assist + } else { + return; + }; + + let editor = if let Some(editor) = pending_assist.editor.upgrade() { + editor + } else { + return; + }; + + let project = pending_assist.project.clone(); + + let project_name = if let Some(project) = project.upgrade() { + Some( + project + .read(cx) + .worktree_root_names(cx) + .collect::>() + .join("/"), + ) + } else { + None + }; + + self.inline_prompt_history + .retain(|prompt| prompt != user_prompt); + self.inline_prompt_history.push_back(user_prompt.into()); + if self.inline_prompt_history.len() > Self::INLINE_PROMPT_HISTORY_MAX_LEN { + self.inline_prompt_history.pop_front(); + } + + let codegen = pending_assist.codegen.clone(); + let snapshot = editor.read(cx).buffer().read(cx).snapshot(cx); + let range = codegen.read(cx).range(); + let start = snapshot.point_to_buffer_offset(range.start); + let end = snapshot.point_to_buffer_offset(range.end); + let (buffer, range) = if let Some((start, end)) = start.zip(end) { + let (start_buffer, start_buffer_offset) = start; + let (end_buffer, end_buffer_offset) = end; + if start_buffer.remote_id() == end_buffer.remote_id() { + (start_buffer.clone(), start_buffer_offset..end_buffer_offset) + } else { + self.finish_inline_assist(inline_assist_id, false, cx); + return; + } + } else { + self.finish_inline_assist(inline_assist_id, false, cx); + return; + }; + + let language = buffer.language_at(range.start); + let language_name = if let Some(language) = language.as_ref() { + if Arc::ptr_eq(language, &language::PLAIN_TEXT) { + None + } else { + Some(language.name()) + } + } else { + None + }; + + // Higher Temperature increases the randomness of model outputs. + // If Markdown or No Language is Known, increase the randomness for more creative output + // If Code, decrease temperature to get more deterministic outputs + let temperature = if let Some(language) = language_name.clone() { + if language.to_string() != "Markdown".to_string() { + 0.5 + } else { + 1.0 + } + } else { + 1.0 + }; + + let user_prompt = user_prompt.to_string(); + + let snippets = if retrieve_context { + let Some(project) = project.upgrade() else { + return; + }; + + let search_results = if let Some(semantic_index) = self.semantic_index.clone() { + let search_results = semantic_index.update(cx, |this, cx| { + this.search_project(project, user_prompt.to_string(), 10, vec![], vec![], cx) + }); + + cx.background_executor() + .spawn(async move { search_results.await.unwrap_or_default() }) + } else { + Task::ready(Vec::new()) + }; + + let snippets = cx.spawn(|_, mut cx| async move { + let mut snippets = Vec::new(); + for result in search_results.await { + snippets.push(PromptCodeSnippet::new(result.buffer, result.range, &mut cx)); + } + snippets + }); + snippets + } else { + Task::ready(Vec::new()) + }; + + let mut model = AssistantSettings::get_global(cx) + .default_open_ai_model + .clone(); + let model_name = model.full_name(); + + let prompt = cx.background_executor().spawn(async move { + let snippets = snippets.await; + + let language_name = language_name.as_deref(); + generate_content_prompt( + user_prompt, + language_name, + buffer, + range, + snippets, + model_name, + project_name, + ) + }); + + let mut messages = Vec::new(); + if let Some(conversation) = conversation { + let conversation = conversation.read(cx); + let buffer = conversation.buffer.read(cx); + messages.extend( + conversation + .messages(cx) + .map(|message| message.to_open_ai_message(buffer)), + ); + model = conversation.model.clone(); + } + + cx.spawn(|_, mut cx| async move { + // I Don't know if we want to return a ? here. + let prompt = prompt.await?; + + messages.push(RequestMessage { + role: Role::User, + content: prompt, + }); + + let request = Box::new(OpenAIRequest { + model: model.full_name().into(), + messages, + stream: true, + stop: vec!["|END|>".to_string()], + temperature, + }); + + codegen.update(&mut cx, |codegen, cx| codegen.start(request, cx)); + anyhow::Ok(()) + }) + .detach(); + } + + fn update_highlights_for_editor(&self, editor: &View, cx: &mut ViewContext) { + let mut background_ranges = Vec::new(); + let mut foreground_ranges = Vec::new(); + let empty_inline_assist_ids = Vec::new(); + let inline_assist_ids = self + .pending_inline_assist_ids_by_editor + .get(&editor.downgrade()) + .unwrap_or(&empty_inline_assist_ids); + + for inline_assist_id in inline_assist_ids { + if let Some(pending_assist) = self.pending_inline_assists.get(inline_assist_id) { + let codegen = pending_assist.codegen.read(cx); + background_ranges.push(codegen.range()); + foreground_ranges.extend(codegen.last_equal_ranges().iter().cloned()); + } + } + + let snapshot = editor.read(cx).buffer().read(cx).snapshot(cx); + merge_ranges(&mut background_ranges, &snapshot); + merge_ranges(&mut foreground_ranges, &snapshot); + editor.update(cx, |editor, cx| { + if background_ranges.is_empty() { + editor.clear_background_highlights::(cx); + } else { + editor.highlight_background::( + background_ranges, + |theme| theme.assistant.inline.pending_edit_background, + cx, + ); + } + + if foreground_ranges.is_empty() { + editor.clear_highlights::(cx); + } else { + editor.highlight_text::( + foreground_ranges, + HighlightStyle { + fade_out: Some(0.6), + ..Default::default() + }, + cx, + ); + } + }); + } + + fn new_conversation(&mut self, cx: &mut ViewContext) -> View { + let editor = cx.add_view(|cx| { + ConversationEditor::new( + self.completion_provider.clone(), + self.languages.clone(), + self.fs.clone(), + self.workspace.clone(), + cx, + ) + }); + self.add_conversation(editor.clone(), cx); + editor + } + + fn add_conversation(&mut self, editor: View, cx: &mut ViewContext) { + self.subscriptions + .push(cx.subscribe(&editor, Self::handle_conversation_editor_event)); + + let conversation = editor.read(cx).conversation.clone(); + self.subscriptions + .push(cx.observe(&conversation, |_, _, cx| cx.notify())); + + let index = self.editors.len(); + self.editors.push(editor); + self.set_active_editor_index(Some(index), cx); + } + + fn set_active_editor_index(&mut self, index: Option, cx: &mut ViewContext) { + self.prev_active_editor_index = self.active_editor_index; + self.active_editor_index = index; + if let Some(editor) = self.active_editor() { + let editor = editor.read(cx).editor.clone(); + self.toolbar.update(cx, |toolbar, cx| { + toolbar.set_active_item(Some(&editor), cx); + }); + if self.has_focus(cx) { + cx.focus(&editor); + } + } else { + self.toolbar.update(cx, |toolbar, cx| { + toolbar.set_active_item(None, cx); + }); + } + + cx.notify(); + } + + fn handle_conversation_editor_event( + &mut self, + _: View, + event: &ConversationEditorEvent, + cx: &mut ViewContext, + ) { + match event { + ConversationEditorEvent::TabContentChanged => cx.notify(), + } + } + + fn save_credentials(&mut self, _: &menu::Confirm, cx: &mut ViewContext) { + if let Some(api_key) = self + .api_key_editor + .as_ref() + .map(|editor| editor.read(cx).text(cx)) + { + if !api_key.is_empty() { + let credential = ProviderCredential::Credentials { + api_key: api_key.clone(), + }; + + self.completion_provider.save_credentials(cx, credential); + + self.api_key_editor.take(); + cx.focus_self(); + cx.notify(); + } + } else { + cx.propagate_action(); + } + } + + fn reset_credentials(&mut self, _: &ResetKey, cx: &mut ViewContext) { + self.completion_provider.delete_credentials(cx); + self.api_key_editor = Some(build_api_key_editor(cx)); + cx.focus_self(); + cx.notify(); + } + + fn toggle_zoom(&mut self, _: &workspace::ToggleZoom, cx: &mut ViewContext) { + if self.zoomed { + cx.emit(AssistantPanelEvent::ZoomOut) + } else { + cx.emit(AssistantPanelEvent::ZoomIn) + } + } + + fn deploy(&mut self, action: &search::buffer_search::Deploy, cx: &mut ViewContext) { + let mut propagate_action = true; + if let Some(search_bar) = self.toolbar.read(cx).item_of_type::() { + search_bar.update(cx, |search_bar, cx| { + if search_bar.show(cx) { + search_bar.search_suggested(cx); + if action.focus { + search_bar.select_query(cx); + cx.focus_self(); + } + propagate_action = false + } + }); + } + if propagate_action { + cx.propagate_action(); + } + } + + fn handle_editor_cancel(&mut self, _: &editor::Cancel, cx: &mut ViewContext) { + if let Some(search_bar) = self.toolbar.read(cx).item_of_type::() { + if !search_bar.read(cx).is_dismissed() { + search_bar.update(cx, |search_bar, cx| { + search_bar.dismiss(&Default::default(), cx) + }); + return; + } + } + cx.propagate_action(); + } + + fn select_next_match(&mut self, _: &search::SelectNextMatch, cx: &mut ViewContext) { + if let Some(search_bar) = self.toolbar.read(cx).item_of_type::() { + search_bar.update(cx, |bar, cx| bar.select_match(Direction::Next, 1, cx)); + } + } + + fn select_prev_match(&mut self, _: &search::SelectPrevMatch, cx: &mut ViewContext) { + if let Some(search_bar) = self.toolbar.read(cx).item_of_type::() { + search_bar.update(cx, |bar, cx| bar.select_match(Direction::Prev, 1, cx)); + } + } + + fn active_editor(&self) -> Option<&View> { + self.editors.get(self.active_editor_index?) + } + + fn render_hamburger_button(cx: &mut ViewContext) -> impl IntoElement { + IconButton::new("hamburger_button", ui::Icon::Menu) + .on_click(cx.listener(|this, _event, cx| { + if this.active_editor().is_some() { + this.set_active_editor_index(None, cx); + } else { + this.set_active_editor_index(this.prev_active_editor_index, cx); + } + })) + .tooltip(|cx| ui::Tooltip::text("History", cx)) + } + + fn render_editor_tools(&self, cx: &mut ViewContext) -> Vec { + if self.active_editor().is_some() { + vec![ + Self::render_split_button(cx).into_any(), + Self::render_quote_button(cx).into_any(), + Self::render_assist_button(cx).into_any(), + ] + } else { + Default::default() + } + } + + fn render_split_button(cx: &mut ViewContext) -> impl IntoElement { + IconButton::new("split_button", ui::Icon::Menu) + .on_click(cx.listener(|this, _event, cx| { + if let Some(active_editor) = this.active_editor() { + active_editor.update(cx, |editor, cx| editor.split(&Default::default(), cx)); + } + })) + .tooltip(|cx| ui::Tooltip::for_action("Split Message", &Split, cx)) + } + + fn render_assist_button(cx: &mut ViewContext) -> impl IntoElement { + IconButton::new("assist_button", ui::Icon::Menu) + .on_click(cx.listener(|this, _event, cx| { + if let Some(active_editor) = this.active_editor() { + active_editor.update(cx, |editor, cx| editor.assist(&Default::default(), cx)); + } + })) + .tooltip(|cx| ui::Tooltip::for_action("Assist", &Assist, cx)) + } + + fn render_quote_button(cx: &mut ViewContext) -> impl IntoElement { + IconButton::new("quote_button", ui::Icon::Menu) + .on_click(cx.listener(|this, _event, cx| { + if let Some(workspace) = this.workspace.upgrade() { + cx.window_context().defer(move |cx| { + workspace.update(cx, |workspace, cx| { + ConversationEditor::quote_selection(workspace, &Default::default(), cx) + }); + }); + } + })) + .tooltip(|cx| ui::Tooltip::for_action("Quote Seleciton", &QuoteSelection, cx)) + } + + fn render_plus_button(cx: &mut ViewContext) -> impl IntoElement { + IconButton::new("plus_button", ui::Icon::Menu) + .on_click(cx.listener(|this, _event, cx| { + this.new_conversation(cx); + })) + .tooltip(|cx| ui::Tooltip::for_action("New Conversation", &NewConversation, cx)) + } + + fn render_zoom_button(&self, cx: &mut ViewContext) -> impl IntoElement { + IconButton::new("zoom_button", ui::Icon::Menu) + .on_click(cx.listener(|this, _event, cx| { + this.toggle_zoom(&ToggleZoom, cx); + })) + .tooltip(|cx| { + ui::Tooltip::for_action( + if self.zoomed { "Zoom Out" } else { "Zoom In" }, + &ToggleZoom, + cx, + ) + }) + } + + fn render_saved_conversation( + &mut self, + index: usize, + cx: &mut ViewContext, + ) -> impl IntoElement { + let conversation = &self.saved_conversations[index]; + let path = conversation.path.clone(); + + ButtonLike::new(index) + .on_click(cx.listener(move |this, _, cx| { + this.open_conversation(path.clone(), cx) + .detach_and_log_err(cx) + })) + .child(Label::new( + conversation.mtime.format("%F %I:%M%p").to_string(), + )) + .child(Label::new(conversation.title.clone())) + } + + fn open_conversation(&mut self, path: PathBuf, cx: &mut ViewContext) -> Task> { + if let Some(ix) = self.editor_index_for_path(&path, cx) { + self.set_active_editor_index(Some(ix), cx); + return Task::ready(Ok(())); + } + + let fs = self.fs.clone(); + let workspace = self.workspace.clone(); + let languages = self.languages.clone(); + cx.spawn(|this, mut cx| async move { + let saved_conversation = fs.load(&path).await?; + let saved_conversation = serde_json::from_str(&saved_conversation)?; + let conversation = cx.add_model(|cx| { + Conversation::deserialize(saved_conversation, path.clone(), languages, cx) + }); + this.update(&mut cx, |this, cx| { + // If, by the time we've loaded the conversation, the user has already opened + // the same conversation, we don't want to open it again. + if let Some(ix) = this.editor_index_for_path(&path, cx) { + this.set_active_editor_index(Some(ix), cx); + } else { + let editor = cx.add_view(|cx| { + ConversationEditor::for_conversation(conversation, fs, workspace, cx) + }); + this.add_conversation(editor, cx); + } + })?; + Ok(()) + }) + } + + fn editor_index_for_path(&self, path: &Path, cx: &AppContext) -> Option { + self.editors + .iter() + .position(|editor| editor.read(cx).conversation.read(cx).path.as_deref() == Some(path)) + } + + fn has_credentials(&mut self) -> bool { + self.completion_provider.has_credentials() + } + + fn load_credentials(&mut self, cx: &mut ViewContext) { + self.completion_provider.retrieve_credentials(cx); + } +} + +fn build_api_key_editor(cx: &mut ViewContext) -> View { + cx.build_view(|cx| { + let mut editor = Editor::single_line(cx); + editor.set_placeholder_text("sk-000000000000000000000000000000000000000000000000", cx); + editor + }) +} + +impl Render for AssistantPanel { + type Element = Div; + + fn render(&mut self, cx: &mut ViewContext) -> AnyElement { + if let Some(api_key_editor) = self.api_key_editor.clone() { + v_stack() + .track_focus(&self.focus_handle) + .child(Label::new( + "To use the assistant panel or inline assistant, you need to add your OpenAI api key.", + )) + .child(Label::new( + " - Having a subscription for another service like GitHub Copilot won't work." + )) + .child(Label::new( + " - You can create a api key at: platform.openai.com/api-keys" + )) + .child(Label::new( + " " + )) + .child(Label::new( + "Paste your OpenAI API key and press Enter to use the assistant" + )) + .child(api_key_editor) + .child(Label::new( + "Click on the Z button in the status bar to close this panel." + )) + .border() + .border_color(gpui::red()) + } else { + let title = self + .active_editor() + .map(|editor| Label::new(editor.read(cx).title(cx))); + + let mut header = h_stack() + .child(Self::render_hamburger_button(cx)) + .children(title); + + if self.focus_handle.contains_focused(cx) { + header = header + .children(self.render_editor_tools(cx)) + .child(Self::render_plus_button(cx)) + .child(self.render_zoom_button(cx)); + } + + v_stack() + .track_focus(&self.focus_handle) + .child(header) + .children(if self.toolbar.read(cx).hidden() { + None + } else { + Some(self.toolbar.clone()) + }) + .child(if let Some(editor) = self.active_editor() { + editor.clone().into_any_element() + } else { + uniform_list( + cx.view().clone(), + "saved_conversations", + self.saved_conversations.len(), + |this, range, cx| { + range + .map(|ix| this.render_saved_conversation(ix, cx).into_any()) + .collect() + }, + ) + .track_scroll(self.saved_conversations_scroll_handle.clone()) + .into_any_element() + }) + .border() + .border_color(gpui::red()) + } + } +} + +impl Panel for AssistantPanel { + fn persistent_name() -> &'static str { + "AssistantPanel" + } + + fn position(&self, cx: &WindowContext) -> DockPosition { + match AssistantSettings::get_global(cx).dock { + AssistantDockPosition::Left => DockPosition::Left, + AssistantDockPosition::Bottom => DockPosition::Bottom, + AssistantDockPosition::Right => DockPosition::Right, + } + } + + fn position_is_valid(&self, _: DockPosition) -> bool { + true + } + + fn set_position(&mut self, position: DockPosition, cx: &mut ViewContext) { + settings::update_settings_file::(self.fs.clone(), cx, move |settings| { + let dock = match position { + DockPosition::Left => AssistantDockPosition::Left, + DockPosition::Bottom => AssistantDockPosition::Bottom, + DockPosition::Right => AssistantDockPosition::Right, + }; + settings.dock = Some(dock); + }); + } + + fn size(&self, cx: &WindowContext) -> f32 { + let settings = AssistantSettings::get_global(cx); + match self.position(cx) { + DockPosition::Left | DockPosition::Right => { + self.width.unwrap_or_else(|| settings.default_width) + } + DockPosition::Bottom => self.height.unwrap_or_else(|| settings.default_height), + } + } + + fn set_size(&mut self, size: Option, cx: &mut ViewContext) { + match self.position(cx) { + DockPosition::Left | DockPosition::Right => self.width = size, + DockPosition::Bottom => self.height = size, + } + cx.notify(); + } + + fn is_zoomed(&self, _: &WindowContext) -> bool { + self.zoomed + } + + fn set_zoomed(&mut self, zoomed: bool, cx: &mut ViewContext) { + self.zoomed = zoomed; + cx.notify(); + } + + fn set_active(&mut self, active: bool, cx: &mut ViewContext) { + if active { + self.load_credentials(cx); + + if self.editors.is_empty() { + self.new_conversation(cx); + } + } + } + + fn icon(&self, cx: &WindowContext) -> Option { + Some(ui::Icon::Ai) + } + + fn toggle_action(&self) -> Box { + Box::new(ToggleFocus) + } +} + +impl EventEmitter for AssistantPanel {} + +impl FocusableView for AssistantPanel { + fn focus_handle(&self, _cx: &AppContext) -> FocusHandle { + self.focus_handle.clone() + } +} + +enum ConversationEvent { + MessagesEdited, + SummaryChanged, + StreamedCompletion, +} + +#[derive(Default)] +struct Summary { + text: String, + done: bool, +} + +struct Conversation { + id: Option, + buffer: Model, + message_anchors: Vec, + messages_metadata: HashMap, + next_message_id: MessageId, + summary: Option, + pending_summary: Task>, + completion_count: usize, + pending_completions: Vec, + model: OpenAIModel, + token_count: Option, + max_token_count: usize, + pending_token_count: Task>, + pending_save: Task>, + path: Option, + _subscriptions: Vec, + completion_provider: Arc, +} + +impl EventEmitter for Conversation {} + +impl Conversation { + fn new( + language_registry: Arc, + cx: &mut ModelContext, + completion_provider: Arc, + ) -> Self { + let markdown = language_registry.language_for_name("Markdown"); + let buffer = cx.add_model(|cx| { + let mut buffer = Buffer::new(0, cx.model_id() as u64, ""); + buffer.set_language_registry(language_registry); + cx.spawn_weak(|buffer, mut cx| async move { + let markdown = markdown.await?; + let buffer = buffer + .upgrade(&cx) + .ok_or_else(|| anyhow!("buffer was dropped"))?; + buffer.update(&mut cx, |buffer: &mut Buffer, cx| { + buffer.set_language(Some(markdown), cx) + }); + anyhow::Ok(()) + }) + .detach_and_log_err(cx); + buffer + }); + + let settings = AssistantSettings::get_global(cx); + let model = settings.default_open_ai_model.clone(); + + let mut this = Self { + id: Some(Uuid::new_v4().to_string()), + message_anchors: Default::default(), + messages_metadata: Default::default(), + next_message_id: Default::default(), + summary: None, + pending_summary: Task::ready(None), + completion_count: Default::default(), + pending_completions: Default::default(), + token_count: None, + max_token_count: tiktoken_rs::model::get_context_size(&model.full_name()), + pending_token_count: Task::ready(None), + model: model.clone(), + _subscriptions: vec![cx.subscribe(&buffer, Self::handle_buffer_event)], + pending_save: Task::ready(Ok(())), + path: None, + buffer, + completion_provider, + }; + let message = MessageAnchor { + id: MessageId(post_inc(&mut this.next_message_id.0)), + start: language::Anchor::MIN, + }; + this.message_anchors.push(message.clone()); + this.messages_metadata.insert( + message.id, + MessageMetadata { + role: Role::User, + sent_at: Local::now(), + status: MessageStatus::Done, + }, + ); + + this.count_remaining_tokens(cx); + this + } + + fn serialize(&self, cx: &AppContext) -> SavedConversation { + SavedConversation { + id: self.id.clone(), + zed: "conversation".into(), + version: SavedConversation::VERSION.into(), + text: self.buffer.read(cx).text(), + message_metadata: self.messages_metadata.clone(), + messages: self + .messages(cx) + .map(|message| SavedMessage { + id: message.id, + start: message.offset_range.start, + }) + .collect(), + summary: self + .summary + .as_ref() + .map(|summary| summary.text.clone()) + .unwrap_or_default(), + model: self.model.clone(), + } + } + + fn deserialize( + saved_conversation: SavedConversation, + path: PathBuf, + language_registry: Arc, + cx: &mut ModelContext, + ) -> Self { + let id = match saved_conversation.id { + Some(id) => Some(id), + None => Some(Uuid::new_v4().to_string()), + }; + let model = saved_conversation.model; + let completion_provider: Arc = Arc::new( + OpenAICompletionProvider::new(model.full_name(), cx.background_executor().clone()), + ); + completion_provider.retrieve_credentials(cx); + let markdown = language_registry.language_for_name("Markdown"); + let mut message_anchors = Vec::new(); + let mut next_message_id = MessageId(0); + let buffer = cx.add_model(|cx| { + let mut buffer = Buffer::new(0, cx.model_id() as u64, saved_conversation.text); + for message in saved_conversation.messages { + message_anchors.push(MessageAnchor { + id: message.id, + start: buffer.anchor_before(message.start), + }); + next_message_id = cmp::max(next_message_id, MessageId(message.id.0 + 1)); + } + buffer.set_language_registry(language_registry); + cx.spawn_weak(|buffer, mut cx| async move { + let markdown = markdown.await?; + let buffer = buffer + .upgrade(&cx) + .ok_or_else(|| anyhow!("buffer was dropped"))?; + buffer.update(&mut cx, |buffer: &mut Buffer, cx| { + buffer.set_language(Some(markdown), cx) + }); + anyhow::Ok(()) + }) + .detach_and_log_err(cx); + buffer + }); + + let mut this = Self { + id, + message_anchors, + messages_metadata: saved_conversation.message_metadata, + next_message_id, + summary: Some(Summary { + text: saved_conversation.summary, + done: true, + }), + pending_summary: Task::ready(None), + completion_count: Default::default(), + pending_completions: Default::default(), + token_count: None, + max_token_count: tiktoken_rs::model::get_context_size(&model.full_name()), + pending_token_count: Task::ready(None), + model, + _subscriptions: vec![cx.subscribe(&buffer, Self::handle_buffer_event)], + pending_save: Task::ready(Ok(())), + path: Some(path), + buffer, + completion_provider, + }; + this.count_remaining_tokens(cx); + this + } + + fn handle_buffer_event( + &mut self, + _: Model, + event: &language::Event, + cx: &mut ModelContext, + ) { + match event { + language::Event::Edited => { + self.count_remaining_tokens(cx); + cx.emit(ConversationEvent::MessagesEdited); + } + _ => {} + } + } + + fn count_remaining_tokens(&mut self, cx: &mut ModelContext) { + let messages = self + .messages(cx) + .into_iter() + .filter_map(|message| { + Some(tiktoken_rs::ChatCompletionRequestMessage { + role: match message.role { + Role::User => "user".into(), + Role::Assistant => "assistant".into(), + Role::System => "system".into(), + }, + content: Some( + self.buffer + .read(cx) + .text_for_range(message.offset_range) + .collect(), + ), + name: None, + function_call: None, + }) + }) + .collect::>(); + let model = self.model.clone(); + self.pending_token_count = cx.spawn_weak(|this, mut cx| { + async move { + cx.background_executor() + .timer(Duration::from_millis(200)) + .await; + let token_count = cx + .background() + .spawn(async move { + tiktoken_rs::num_tokens_from_messages(&model.full_name(), &messages) + }) + .await?; + + this.upgrade(&cx) + .ok_or_else(|| anyhow!("conversation was dropped"))? + .update(&mut cx, |this, cx| { + this.max_token_count = + tiktoken_rs::model::get_context_size(&this.model.full_name()); + this.token_count = Some(token_count); + cx.notify() + }); + anyhow::Ok(()) + } + .log_err() + }); + } + + fn remaining_tokens(&self) -> Option { + Some(self.max_token_count as isize - self.token_count? as isize) + } + + fn set_model(&mut self, model: OpenAIModel, cx: &mut ModelContext) { + self.model = model; + self.count_remaining_tokens(cx); + cx.notify(); + } + + fn assist( + &mut self, + selected_messages: HashSet, + cx: &mut ModelContext, + ) -> Vec { + let mut user_messages = Vec::new(); + + let last_message_id = if let Some(last_message_id) = + self.message_anchors.iter().rev().find_map(|message| { + message + .start + .is_valid(self.buffer.read(cx)) + .then_some(message.id) + }) { + last_message_id + } else { + return Default::default(); + }; + + let mut should_assist = false; + for selected_message_id in selected_messages { + let selected_message_role = + if let Some(metadata) = self.messages_metadata.get(&selected_message_id) { + metadata.role + } else { + continue; + }; + + if selected_message_role == Role::Assistant { + if let Some(user_message) = self.insert_message_after( + selected_message_id, + Role::User, + MessageStatus::Done, + cx, + ) { + user_messages.push(user_message); + } + } else { + should_assist = true; + } + } + + if should_assist { + if !self.completion_provider.has_credentials() { + return Default::default(); + } + + let request: Box = Box::new(OpenAIRequest { + model: self.model.full_name().to_string(), + messages: self + .messages(cx) + .filter(|message| matches!(message.status, MessageStatus::Done)) + .map(|message| message.to_open_ai_message(self.buffer.read(cx))) + .collect(), + stream: true, + stop: vec![], + temperature: 1.0, + }); + + let stream = self.completion_provider.complete(request); + let assistant_message = self + .insert_message_after(last_message_id, Role::Assistant, MessageStatus::Pending, cx) + .unwrap(); + + // Queue up the user's next reply. + let user_message = self + .insert_message_after(assistant_message.id, Role::User, MessageStatus::Done, cx) + .unwrap(); + user_messages.push(user_message); + + let task = cx.spawn_weak({ + |this, mut cx| async move { + let assistant_message_id = assistant_message.id; + let stream_completion = async { + let mut messages = stream.await?; + + while let Some(message) = messages.next().await { + let text = message?; + + this.upgrade(&cx) + .ok_or_else(|| anyhow!("conversation was dropped"))? + .update(&mut cx, |this, cx| { + let message_ix = this + .message_anchors + .iter() + .position(|message| message.id == assistant_message_id)?; + this.buffer.update(cx, |buffer, cx| { + let offset = this.message_anchors[message_ix + 1..] + .iter() + .find(|message| message.start.is_valid(buffer)) + .map_or(buffer.len(), |message| { + message.start.to_offset(buffer).saturating_sub(1) + }); + buffer.edit([(offset..offset, text)], None, cx); + }); + cx.emit(ConversationEvent::StreamedCompletion); + + Some(()) + }); + smol::future::yield_now().await; + } + + this.upgrade(&cx) + .ok_or_else(|| anyhow!("conversation was dropped"))? + .update(&mut cx, |this, cx| { + this.pending_completions + .retain(|completion| completion.id != this.completion_count); + this.summarize(cx); + }); + + anyhow::Ok(()) + }; + + let result = stream_completion.await; + if let Some(this) = this.upgrade(&cx) { + this.update(&mut cx, |this, cx| { + if let Some(metadata) = + this.messages_metadata.get_mut(&assistant_message.id) + { + match result { + Ok(_) => { + metadata.status = MessageStatus::Done; + } + Err(error) => { + metadata.status = + MessageStatus::Error(error.to_string().trim().into()); + } + } + cx.notify(); + } + }); + } + } + }); + + self.pending_completions.push(PendingCompletion { + id: post_inc(&mut self.completion_count), + _task: task, + }); + } + + user_messages + } + + fn cancel_last_assist(&mut self) -> bool { + self.pending_completions.pop().is_some() + } + + fn cycle_message_roles(&mut self, ids: HashSet, cx: &mut ModelContext) { + for id in ids { + if let Some(metadata) = self.messages_metadata.get_mut(&id) { + metadata.role.cycle(); + cx.emit(ConversationEvent::MessagesEdited); + cx.notify(); + } + } + } + + fn insert_message_after( + &mut self, + message_id: MessageId, + role: Role, + status: MessageStatus, + cx: &mut ModelContext, + ) -> Option { + if let Some(prev_message_ix) = self + .message_anchors + .iter() + .position(|message| message.id == message_id) + { + // Find the next valid message after the one we were given. + let mut next_message_ix = prev_message_ix + 1; + while let Some(next_message) = self.message_anchors.get(next_message_ix) { + if next_message.start.is_valid(self.buffer.read(cx)) { + break; + } + next_message_ix += 1; + } + + let start = self.buffer.update(cx, |buffer, cx| { + let offset = self + .message_anchors + .get(next_message_ix) + .map_or(buffer.len(), |message| message.start.to_offset(buffer) - 1); + buffer.edit([(offset..offset, "\n")], None, cx); + buffer.anchor_before(offset + 1) + }); + let message = MessageAnchor { + id: MessageId(post_inc(&mut self.next_message_id.0)), + start, + }; + self.message_anchors + .insert(next_message_ix, message.clone()); + self.messages_metadata.insert( + message.id, + MessageMetadata { + role, + sent_at: Local::now(), + status, + }, + ); + cx.emit(ConversationEvent::MessagesEdited); + Some(message) + } else { + None + } + } + + fn split_message( + &mut self, + range: Range, + cx: &mut ModelContext, + ) -> (Option, Option) { + let start_message = self.message_for_offset(range.start, cx); + let end_message = self.message_for_offset(range.end, cx); + if let Some((start_message, end_message)) = start_message.zip(end_message) { + // Prevent splitting when range spans multiple messages. + if start_message.id != end_message.id { + return (None, None); + } + + let message = start_message; + let role = message.role; + let mut edited_buffer = false; + + let mut suffix_start = None; + if range.start > message.offset_range.start && range.end < message.offset_range.end - 1 + { + if self.buffer.read(cx).chars_at(range.end).next() == Some('\n') { + suffix_start = Some(range.end + 1); + } else if self.buffer.read(cx).reversed_chars_at(range.end).next() == Some('\n') { + suffix_start = Some(range.end); + } + } + + let suffix = if let Some(suffix_start) = suffix_start { + MessageAnchor { + id: MessageId(post_inc(&mut self.next_message_id.0)), + start: self.buffer.read(cx).anchor_before(suffix_start), + } + } else { + self.buffer.update(cx, |buffer, cx| { + buffer.edit([(range.end..range.end, "\n")], None, cx); + }); + edited_buffer = true; + MessageAnchor { + id: MessageId(post_inc(&mut self.next_message_id.0)), + start: self.buffer.read(cx).anchor_before(range.end + 1), + } + }; + + self.message_anchors + .insert(message.index_range.end + 1, suffix.clone()); + self.messages_metadata.insert( + suffix.id, + MessageMetadata { + role, + sent_at: Local::now(), + status: MessageStatus::Done, + }, + ); + + let new_messages = + if range.start == range.end || range.start == message.offset_range.start { + (None, Some(suffix)) + } else { + let mut prefix_end = None; + if range.start > message.offset_range.start + && range.end < message.offset_range.end - 1 + { + if self.buffer.read(cx).chars_at(range.start).next() == Some('\n') { + prefix_end = Some(range.start + 1); + } else if self.buffer.read(cx).reversed_chars_at(range.start).next() + == Some('\n') + { + prefix_end = Some(range.start); + } + } + + let selection = if let Some(prefix_end) = prefix_end { + cx.emit(ConversationEvent::MessagesEdited); + MessageAnchor { + id: MessageId(post_inc(&mut self.next_message_id.0)), + start: self.buffer.read(cx).anchor_before(prefix_end), + } + } else { + self.buffer.update(cx, |buffer, cx| { + buffer.edit([(range.start..range.start, "\n")], None, cx) + }); + edited_buffer = true; + MessageAnchor { + id: MessageId(post_inc(&mut self.next_message_id.0)), + start: self.buffer.read(cx).anchor_before(range.end + 1), + } + }; + + self.message_anchors + .insert(message.index_range.end + 1, selection.clone()); + self.messages_metadata.insert( + selection.id, + MessageMetadata { + role, + sent_at: Local::now(), + status: MessageStatus::Done, + }, + ); + (Some(selection), Some(suffix)) + }; + + if !edited_buffer { + cx.emit(ConversationEvent::MessagesEdited); + } + new_messages + } else { + (None, None) + } + } + + fn summarize(&mut self, cx: &mut ModelContext) { + if self.message_anchors.len() >= 2 && self.summary.is_none() { + if !self.completion_provider.has_credentials() { + return; + } + + let messages = self + .messages(cx) + .take(2) + .map(|message| message.to_open_ai_message(self.buffer.read(cx))) + .chain(Some(RequestMessage { + role: Role::User, + content: "Summarize the conversation into a short title without punctuation" + .into(), + })); + let request: Box = Box::new(OpenAIRequest { + model: self.model.full_name().to_string(), + messages: messages.collect(), + stream: true, + stop: vec![], + temperature: 1.0, + }); + + let stream = self.completion_provider.complete(request); + self.pending_summary = cx.spawn(|this, mut cx| { + async move { + let mut messages = stream.await?; + + while let Some(message) = messages.next().await { + let text = message?; + this.update(&mut cx, |this, cx| { + this.summary + .get_or_insert(Default::default()) + .text + .push_str(&text); + cx.emit(ConversationEvent::SummaryChanged); + }); + } + + this.update(&mut cx, |this, cx| { + if let Some(summary) = this.summary.as_mut() { + summary.done = true; + cx.emit(ConversationEvent::SummaryChanged); + } + }); + + anyhow::Ok(()) + } + .log_err() + }); + } + } + + fn message_for_offset(&self, offset: usize, cx: &AppContext) -> Option { + self.messages_for_offsets([offset], cx).pop() + } + + fn messages_for_offsets( + &self, + offsets: impl IntoIterator, + cx: &AppContext, + ) -> Vec { + let mut result = Vec::new(); + + let mut messages = self.messages(cx).peekable(); + let mut offsets = offsets.into_iter().peekable(); + let mut current_message = messages.next(); + while let Some(offset) = offsets.next() { + // Locate the message that contains the offset. + while current_message.as_ref().map_or(false, |message| { + !message.offset_range.contains(&offset) && messages.peek().is_some() + }) { + current_message = messages.next(); + } + let Some(message) = current_message.as_ref() else { + break; + }; + + // Skip offsets that are in the same message. + while offsets.peek().map_or(false, |offset| { + message.offset_range.contains(offset) || messages.peek().is_none() + }) { + offsets.next(); + } + + result.push(message.clone()); + } + result + } + + fn messages<'a>(&'a self, cx: &'a AppContext) -> impl 'a + Iterator { + let buffer = self.buffer.read(cx); + let mut message_anchors = self.message_anchors.iter().enumerate().peekable(); + iter::from_fn(move || { + while let Some((start_ix, message_anchor)) = message_anchors.next() { + let metadata = self.messages_metadata.get(&message_anchor.id)?; + let message_start = message_anchor.start.to_offset(buffer); + let mut message_end = None; + let mut end_ix = start_ix; + while let Some((_, next_message)) = message_anchors.peek() { + if next_message.start.is_valid(buffer) { + message_end = Some(next_message.start); + break; + } else { + end_ix += 1; + message_anchors.next(); + } + } + let message_end = message_end + .unwrap_or(language::Anchor::MAX) + .to_offset(buffer); + return Some(Message { + index_range: start_ix..end_ix, + offset_range: message_start..message_end, + id: message_anchor.id, + anchor: message_anchor.start, + role: metadata.role, + sent_at: metadata.sent_at, + status: metadata.status.clone(), + }); + } + None + }) + } + + fn save( + &mut self, + debounce: Option, + fs: Arc, + cx: &mut ModelContext, + ) { + self.pending_save = cx.spawn(|this, mut cx| async move { + if let Some(debounce) = debounce { + cx.background_executor().timer(debounce).await; + } + + let (old_path, summary) = this.read_with(&cx, |this, _| { + let path = this.path.clone(); + let summary = if let Some(summary) = this.summary.as_ref() { + if summary.done { + Some(summary.text.clone()) + } else { + None + } + } else { + None + }; + (path, summary) + }); + + if let Some(summary) = summary { + let conversation = this.read_with(&cx, |this, cx| this.serialize(cx)); + let path = if let Some(old_path) = old_path { + old_path + } else { + let mut discriminant = 1; + let mut new_path; + loop { + new_path = CONVERSATIONS_DIR.join(&format!( + "{} - {}.zed.json", + summary.trim(), + discriminant + )); + if fs.is_file(&new_path).await { + discriminant += 1; + } else { + break; + } + } + new_path + }; + + fs.create_dir(CONVERSATIONS_DIR.as_ref()).await?; + fs.atomic_write(path.clone(), serde_json::to_string(&conversation).unwrap()) + .await?; + this.update(&mut cx, |this, _| this.path = Some(path)); + } + + Ok(()) + }); + } +} + +struct PendingCompletion { + id: usize, + _task: Task<()>, +} + +enum ConversationEditorEvent { + TabContentChanged, +} + +#[derive(Copy, Clone, Debug, PartialEq)] +struct ScrollPosition { + offset_before_cursor: gpui::Point, + cursor: Anchor, +} + +struct ConversationEditor { + conversation: Model, + fs: Arc, + workspace: WeakView, + editor: View, + blocks: HashSet, + scroll_position: Option, + _subscriptions: Vec, +} + +impl ConversationEditor { + fn new( + completion_provider: Arc, + language_registry: Arc, + fs: Arc, + workspace: WeakView, + cx: &mut ViewContext, + ) -> Self { + let conversation = + cx.add_model(|cx| Conversation::new(language_registry, cx, completion_provider)); + Self::for_conversation(conversation, fs, workspace, cx) + } + + fn for_conversation( + conversation: Model, + fs: Arc, + workspace: WeakView, + cx: &mut ViewContext, + ) -> Self { + let editor = cx.add_view(|cx| { + let mut editor = Editor::for_buffer(conversation.read(cx).buffer.clone(), None, cx); + editor.set_soft_wrap_mode(SoftWrap::EditorWidth, cx); + editor.set_show_gutter(false, cx); + editor.set_show_wrap_guides(false, cx); + editor + }); + + let _subscriptions = vec![ + cx.observe(&conversation, |_, _, cx| cx.notify()), + cx.subscribe(&conversation, Self::handle_conversation_event), + cx.subscribe(&editor, Self::handle_editor_event), + ]; + + let mut this = Self { + conversation, + editor, + blocks: Default::default(), + scroll_position: None, + fs, + workspace, + _subscriptions, + }; + this.update_message_headers(cx); + this + } + + fn assist(&mut self, _: &Assist, cx: &mut ViewContext) { + report_assistant_event( + self.workspace.clone(), + self.conversation.read(cx).id.clone(), + AssistantKind::Panel, + cx, + ); + + let cursors = self.cursors(cx); + + let user_messages = self.conversation.update(cx, |conversation, cx| { + let selected_messages = conversation + .messages_for_offsets(cursors, cx) + .into_iter() + .map(|message| message.id) + .collect(); + conversation.assist(selected_messages, cx) + }); + let new_selections = user_messages + .iter() + .map(|message| { + let cursor = message + .start + .to_offset(self.conversation.read(cx).buffer.read(cx)); + cursor..cursor + }) + .collect::>(); + if !new_selections.is_empty() { + self.editor.update(cx, |editor, cx| { + editor.change_selections( + Some(Autoscroll::Strategy(AutoscrollStrategy::Fit)), + cx, + |selections| selections.select_ranges(new_selections), + ); + }); + // Avoid scrolling to the new cursor position so the assistant's output is stable. + cx.defer(|this, _| this.scroll_position = None); + } + } + + fn cancel_last_assist(&mut self, _: &editor::Cancel, cx: &mut ViewContext) { + if !self + .conversation + .update(cx, |conversation, _| conversation.cancel_last_assist()) + { + cx.propagate_action(); + } + } + + fn cycle_message_role(&mut self, _: &CycleMessageRole, cx: &mut ViewContext) { + let cursors = self.cursors(cx); + self.conversation.update(cx, |conversation, cx| { + let messages = conversation + .messages_for_offsets(cursors, cx) + .into_iter() + .map(|message| message.id) + .collect(); + conversation.cycle_message_roles(messages, cx) + }); + } + + fn cursors(&self, cx: &AppContext) -> Vec { + let selections = self.editor.read(cx).selections.all::(cx); + selections + .into_iter() + .map(|selection| selection.head()) + .collect() + } + + fn handle_conversation_event( + &mut self, + _: Model, + event: &ConversationEvent, + cx: &mut ViewContext, + ) { + match event { + ConversationEvent::MessagesEdited => { + self.update_message_headers(cx); + self.conversation.update(cx, |conversation, cx| { + conversation.save(Some(Duration::from_millis(500)), self.fs.clone(), cx); + }); + } + ConversationEvent::SummaryChanged => { + cx.emit(ConversationEditorEvent::TabContentChanged); + self.conversation.update(cx, |conversation, cx| { + conversation.save(None, self.fs.clone(), cx); + }); + } + ConversationEvent::StreamedCompletion => { + self.editor.update(cx, |editor, cx| { + if let Some(scroll_position) = self.scroll_position { + let snapshot = editor.snapshot(cx); + let cursor_point = scroll_position.cursor.to_display_point(&snapshot); + let scroll_top = + cursor_point.row() as f32 - scroll_position.offset_before_cursor.y; + editor.set_scroll_position( + point(scroll_position.offset_before_cursor.x, scroll_top), + cx, + ); + } + }); + } + } + } + + fn handle_editor_event( + &mut self, + _: View, + event: &EditorEvent, + cx: &mut ViewContext, + ) { + match event { + EditorEvent::ScrollPositionChanged { autoscroll, .. } => { + let cursor_scroll_position = self.cursor_scroll_position(cx); + if *autoscroll { + self.scroll_position = cursor_scroll_position; + } else if self.scroll_position != cursor_scroll_position { + self.scroll_position = None; + } + } + EditorEvent::SelectionsChanged { .. } => { + self.scroll_position = self.cursor_scroll_position(cx); + } + _ => {} + } + } + + fn cursor_scroll_position(&self, cx: &mut ViewContext) -> Option { + self.editor.update(cx, |editor, cx| { + let snapshot = editor.snapshot(cx); + let cursor = editor.selections.newest_anchor().head(); + let cursor_row = cursor.to_display_point(&snapshot.display_snapshot).row() as f32; + let scroll_position = editor + .scroll_manager + .anchor() + .scroll_position(&snapshot.display_snapshot); + + let scroll_bottom = scroll_position.y() + editor.visible_line_count().unwrap_or(0.); + if (scroll_position.y()..scroll_bottom).contains(&cursor_row) { + Some(ScrollPosition { + cursor, + offset_before_cursor: point(scroll_position.x, cursor_row - scroll_position.y), + }) + } else { + None + } + }) + } + + fn update_message_headers(&mut self, cx: &mut ViewContext) { + self.editor.update(cx, |editor, cx| { + let buffer = editor.buffer().read(cx).snapshot(cx); + let excerpt_id = *buffer.as_singleton().unwrap().0; + let old_blocks = std::mem::take(&mut self.blocks); + let new_blocks = self + .conversation + .read(cx) + .messages(cx) + .map(|message| BlockProperties { + position: buffer.anchor_in_excerpt(excerpt_id, message.anchor), + height: 2, + style: BlockStyle::Sticky, + render: Arc::new({ + let conversation = self.conversation.clone(); + // let metadata = message.metadata.clone(); + // let message = message.clone(); + move |cx| { + enum Sender {} + enum ErrorTooltip {} + + let message_id = message.id; + let sender = MouseEventHandler::new::( + message_id.0, + cx, + |state, _| match message.role { + Role::User => { + let style = style.user_sender.style_for(state); + Label::new("You", style.text.clone()) + .contained() + .with_style(style.container) + } + Role::Assistant => { + let style = style.assistant_sender.style_for(state); + Label::new("Assistant", style.text.clone()) + .contained() + .with_style(style.container) + } + Role::System => { + let style = style.system_sender.style_for(state); + Label::new("System", style.text.clone()) + .contained() + .with_style(style.container) + } + }, + ) + .with_cursor_style(CursorStyle::PointingHand) + .on_down(MouseButton::Left, { + let conversation = conversation.clone(); + move |_, _, cx| { + conversation.update(cx, |conversation, cx| { + conversation.cycle_message_roles( + HashSet::from_iter(Some(message_id)), + cx, + ) + }) + } + }); + + Flex::row() + .with_child(sender.aligned()) + .with_child( + Label::new( + message.sent_at.format("%I:%M%P").to_string(), + style.sent_at.text.clone(), + ) + .contained() + .with_style(style.sent_at.container) + .aligned(), + ) + .with_children( + if let MessageStatus::Error(error) = &message.status { + Some( + Svg::new("icons/error.svg") + .with_color(style.error_icon.color) + .constrained() + .with_width(style.error_icon.width) + .contained() + .with_style(style.error_icon.container) + .with_tooltip::( + message_id.0, + error.to_string(), + None, + theme.tooltip.clone(), + cx, + ) + .aligned(), + ) + } else { + None + }, + ) + .aligned() + .left() + .contained() + .with_style(style.message_header) + .into_any() + } + }), + disposition: BlockDisposition::Above, + }) + .collect::>(); + + editor.remove_blocks(old_blocks, None, cx); + let ids = editor.insert_blocks(new_blocks, None, cx); + self.blocks = HashSet::from_iter(ids); + }); + } + + fn quote_selection( + workspace: &mut Workspace, + _: &QuoteSelection, + cx: &mut ViewContext, + ) { + let Some(panel) = workspace.panel::(cx) else { + return; + }; + let Some(editor) = workspace + .active_item(cx) + .and_then(|item| item.act_as::(cx)) + else { + return; + }; + + let text = editor.read_with(cx, |editor, cx| { + let range = editor.selections.newest::(cx).range(); + let buffer = editor.buffer().read(cx).snapshot(cx); + let start_language = buffer.language_at(range.start); + let end_language = buffer.language_at(range.end); + let language_name = if start_language == end_language { + start_language.map(|language| language.name()) + } else { + None + }; + let language_name = language_name.as_deref().unwrap_or("").to_lowercase(); + + let selected_text = buffer.text_for_range(range).collect::(); + if selected_text.is_empty() { + None + } else { + Some(if language_name == "markdown" { + selected_text + .lines() + .map(|line| format!("> {}", line)) + .collect::>() + .join("\n") + } else { + format!("```{language_name}\n{selected_text}\n```") + }) + } + }); + + // Activate the panel + if !panel.read(cx).has_focus(cx) { + workspace.toggle_panel_focus::(cx); + } + + if let Some(text) = text { + panel.update(cx, |panel, cx| { + let conversation = panel + .active_editor() + .cloned() + .unwrap_or_else(|| panel.new_conversation(cx)); + conversation.update(cx, |conversation, cx| { + conversation + .editor + .update(cx, |editor, cx| editor.insert(&text, cx)) + }); + }); + } + } + + fn copy(&mut self, _: &editor::Copy, cx: &mut ViewContext) { + let editor = self.editor.read(cx); + let conversation = self.conversation.read(cx); + if editor.selections.count() == 1 { + let selection = editor.selections.newest::(cx); + let mut copied_text = String::new(); + let mut spanned_messages = 0; + for message in conversation.messages(cx) { + if message.offset_range.start >= selection.range().end { + break; + } else if message.offset_range.end >= selection.range().start { + let range = cmp::max(message.offset_range.start, selection.range().start) + ..cmp::min(message.offset_range.end, selection.range().end); + if !range.is_empty() { + spanned_messages += 1; + write!(&mut copied_text, "## {}\n\n", message.role).unwrap(); + for chunk in conversation.buffer.read(cx).text_for_range(range) { + copied_text.push_str(&chunk); + } + copied_text.push('\n'); + } + } + } + + if spanned_messages > 1 { + cx.platform() + .write_to_clipboard(ClipboardItem::new(copied_text)); + return; + } + } + + cx.propagate_action(); + } + + fn split(&mut self, _: &Split, cx: &mut ViewContext) { + self.conversation.update(cx, |conversation, cx| { + let selections = self.editor.read(cx).selections.disjoint_anchors(); + for selection in selections.into_iter() { + let buffer = self.editor.read(cx).buffer().read(cx).snapshot(cx); + let range = selection + .map(|endpoint| endpoint.to_offset(&buffer)) + .range(); + conversation.split_message(range, cx); + } + }); + } + + fn save(&mut self, _: &Save, cx: &mut ViewContext) { + self.conversation.update(cx, |conversation, cx| { + conversation.save(None, self.fs.clone(), cx) + }); + } + + fn cycle_model(&mut self, cx: &mut ViewContext) { + self.conversation.update(cx, |conversation, cx| { + let new_model = conversation.model.cycle(); + conversation.set_model(new_model, cx); + }); + } + + fn title(&self, cx: &AppContext) -> String { + self.conversation + .read(cx) + .summary + .as_ref() + .map(|summary| summary.text.clone()) + .unwrap_or_else(|| "New Conversation".into()) + } + + fn render_current_model( + &self, + style: &AssistantStyle, + cx: &mut ViewContext, + ) -> impl Element { + enum Model {} + + MouseEventHandler::new::(0, cx, |state, cx| { + let style = style.model.style_for(state); + let model_display_name = self.conversation.read(cx).model.short_name(); + Label::new(model_display_name, style.text.clone()) + .contained() + .with_style(style.container) + }) + .with_cursor_style(CursorStyle::PointingHand) + .on_click(MouseButton::Left, |_, this, cx| this.cycle_model(cx)) + } + + fn render_remaining_tokens( + &self, + style: &AssistantStyle, + cx: &mut ViewContext, + ) -> Option> { + let remaining_tokens = self.conversation.read(cx).remaining_tokens()?; + let remaining_tokens_style = if remaining_tokens <= 0 { + &style.no_remaining_tokens + } else if remaining_tokens <= 500 { + &style.low_remaining_tokens + } else { + &style.remaining_tokens + }; + Some( + Label::new( + remaining_tokens.to_string(), + remaining_tokens_style.text.clone(), + ) + .contained() + .with_style(remaining_tokens_style.container), + ) + } +} + +impl EventEmitter for ConversationEditor {} + +impl View for ConversationEditor { + fn ui_name() -> &'static str { + "ConversationEditor" + } + + fn render(&mut self, cx: &mut ViewContext) -> AnyElement { + let theme = &theme::current(cx).assistant; + Stack::new() + .with_child( + ChildView::new(&self.editor, cx) + .contained() + .with_style(theme.container), + ) + .with_child( + Flex::row() + .with_child(self.render_current_model(theme, cx)) + .with_children(self.render_remaining_tokens(theme, cx)) + .aligned() + .top() + .right(), + ) + .into_any() + } + + fn focus_in(&mut self, _: gpui::AnyViewHandle, cx: &mut ViewContext) { + if cx.is_self_focused() { + cx.focus(&self.editor); + } + } +} + +#[derive(Clone, Debug)] +struct MessageAnchor { + id: MessageId, + start: language::Anchor, +} + +#[derive(Clone, Debug)] +pub struct Message { + offset_range: Range, + index_range: Range, + id: MessageId, + anchor: language::Anchor, + role: Role, + sent_at: DateTime, + status: MessageStatus, +} + +impl Message { + fn to_open_ai_message(&self, buffer: &Buffer) -> RequestMessage { + let content = buffer + .text_for_range(self.offset_range.clone()) + .collect::(); + RequestMessage { + role: self.role, + content: content.trim_end().into(), + } + } +} + +enum InlineAssistantEvent { + Confirmed { + prompt: String, + include_conversation: bool, + retrieve_context: bool, + }, + Canceled, + Dismissed, + IncludeConversationToggled { + include_conversation: bool, + }, + RetrieveContextToggled { + retrieve_context: bool, + }, +} + +struct InlineAssistant { + id: usize, + prompt_editor: View, + workspace: WeakView, + confirmed: bool, + has_focus: bool, + include_conversation: bool, + measurements: Rc>, + prompt_history: VecDeque, + prompt_history_ix: Option, + pending_prompt: String, + codegen: Model, + _subscriptions: Vec, + retrieve_context: bool, + semantic_index: Option>, + semantic_permissioned: Option, + project: WeakModel, + maintain_rate_limit: Option>, +} + +impl Entity for InlineAssistant { + type Event = InlineAssistantEvent; +} + +impl View for InlineAssistant { + fn ui_name() -> &'static str { + "InlineAssistant" + } + + fn render(&mut self, cx: &mut ViewContext) -> AnyElement { + enum ErrorIcon {} + let theme = theme::current(cx); + + Flex::row() + .with_children([Flex::row() + .with_child( + Button::action(ToggleIncludeConversation) + .with_tooltip("Include Conversation", theme.tooltip.clone()) + .with_id(self.id) + .with_contents(theme::components::svg::Svg::new("icons/ai.svg")) + .toggleable(self.include_conversation) + .with_style(theme.assistant.inline.include_conversation.clone()) + .element() + .aligned(), + ) + .with_children(if SemanticIndex::enabled(cx) { + Some( + Button::action(ToggleRetrieveContext) + .with_tooltip("Retrieve Context", theme.tooltip.clone()) + .with_id(self.id) + .with_contents(theme::components::svg::Svg::new( + "icons/magnifying_glass.svg", + )) + .toggleable(self.retrieve_context) + .with_style(theme.assistant.inline.retrieve_context.clone()) + .element() + .aligned(), + ) + } else { + None + }) + .with_children(if let Some(error) = self.codegen.read(cx).error() { + Some( + Svg::new("icons/error.svg") + .with_color(theme.assistant.error_icon.color) + .constrained() + .with_width(theme.assistant.error_icon.width) + .contained() + .with_style(theme.assistant.error_icon.container) + .with_tooltip::( + self.id, + error.to_string(), + None, + theme.tooltip.clone(), + cx, + ) + .aligned(), + ) + } else { + None + }) + .aligned() + .constrained() + .dynamically({ + let measurements = self.measurements.clone(); + move |constraint, _, _| { + let measurements = measurements.get(); + SizeConstraint { + min: vec2f(measurements.gutter_width, constraint.min.y()), + max: vec2f(measurements.gutter_width, constraint.max.y()), + } + } + })]) + .with_child(Empty::new().constrained().dynamically({ + let measurements = self.measurements.clone(); + move |constraint, _, _| { + let measurements = measurements.get(); + SizeConstraint { + min: vec2f( + measurements.anchor_x - measurements.gutter_width, + constraint.min.y(), + ), + max: vec2f( + measurements.anchor_x - measurements.gutter_width, + constraint.max.y(), + ), + } + } + })) + .with_child( + ChildView::new(&self.prompt_editor, cx) + .aligned() + .left() + .flex(1., true), + ) + .with_children(if self.retrieve_context { + Some( + Flex::row() + .with_children(self.retrieve_context_status(cx)) + .flex(1., true) + .aligned(), + ) + } else { + None + }) + .contained() + .with_style(theme.assistant.inline.container) + .into_any() + .into_any() + } + + fn focus_in(&mut self, _: gpui::AnyViewHandle, cx: &mut ViewContext) { + cx.focus(&self.prompt_editor); + self.has_focus = true; + } + + fn focus_out(&mut self, _: gpui::AnyViewHandle, _: &mut ViewContext) { + self.has_focus = false; + } +} + +impl InlineAssistant { + fn new( + id: usize, + measurements: Rc>, + include_conversation: bool, + prompt_history: VecDeque, + codegen: Model, + workspace: WeakView, + cx: &mut ViewContext, + retrieve_context: bool, + semantic_index: Option>, + project: Model, + ) -> Self { + let prompt_editor = cx.add_view(|cx| { + let mut editor = Editor::single_line( + Some(Arc::new(|theme| theme.assistant.inline.editor.clone())), + cx, + ); + let placeholder = match codegen.read(cx).kind() { + CodegenKind::Transform { .. } => "Enter transformation prompt…", + CodegenKind::Generate { .. } => "Enter generation prompt…", + }; + editor.set_placeholder_text(placeholder, cx); + editor + }); + let mut subscriptions = vec![ + cx.observe(&codegen, Self::handle_codegen_changed), + cx.subscribe(&prompt_editor, Self::handle_prompt_editor_events), + ]; + + if let Some(semantic_index) = semantic_index.clone() { + subscriptions.push(cx.observe(&semantic_index, Self::semantic_index_changed)); + } + + let assistant = Self { + id, + prompt_editor, + workspace, + confirmed: false, + has_focus: false, + include_conversation, + measurements, + prompt_history, + prompt_history_ix: None, + pending_prompt: String::new(), + codegen, + _subscriptions: subscriptions, + retrieve_context, + semantic_permissioned: None, + semantic_index, + project: project.downgrade(), + maintain_rate_limit: None, + }; + + assistant.index_project(cx).log_err(); + + assistant + } + + fn semantic_permissioned(&self, cx: &mut ViewContext) -> Task> { + if let Some(value) = self.semantic_permissioned { + return Task::ready(Ok(value)); + } + + let Some(project) = self.project.upgrade() else { + return Task::ready(Err(anyhow!("project was dropped"))); + }; + + self.semantic_index + .as_ref() + .map(|semantic| { + semantic.update(cx, |this, cx| this.project_previously_indexed(&project, cx)) + }) + .unwrap_or(Task::ready(Ok(false))) + } + + fn handle_prompt_editor_events( + &mut self, + _: View, + event: &EditorEvent, + cx: &mut ViewContext, + ) { + if let EditorEvent::Edited = event { + self.pending_prompt = self.prompt_editor.read(cx).text(cx); + cx.notify(); + } + } + + fn semantic_index_changed( + &mut self, + semantic_index: Model, + cx: &mut ViewContext, + ) { + let Some(project) = self.project.upgrade() else { + return; + }; + + let status = semantic_index.read(cx).status(&project); + match status { + SemanticIndexStatus::Indexing { + rate_limit_expiry: Some(_), + .. + } => { + if self.maintain_rate_limit.is_none() { + self.maintain_rate_limit = Some(cx.spawn(|this, mut cx| async move { + loop { + cx.background_executor().timer(Duration::from_secs(1)).await; + this.update(&mut cx, |_, cx| cx.notify()).log_err(); + } + })); + } + return; + } + _ => { + self.maintain_rate_limit = None; + } + } + } + + fn handle_codegen_changed(&mut self, _: Model, cx: &mut ViewContext) { + let is_read_only = !self.codegen.read(cx).idle(); + self.prompt_editor.update(cx, |editor, cx| { + let was_read_only = editor.read_only(); + if was_read_only != is_read_only { + if is_read_only { + editor.set_read_only(true); + editor.set_field_editor_style( + Some(Arc::new(|theme| { + theme.assistant.inline.disabled_editor.clone() + })), + cx, + ); + } else { + self.confirmed = false; + editor.set_read_only(false); + editor.set_field_editor_style( + Some(Arc::new(|theme| theme.assistant.inline.editor.clone())), + cx, + ); + } + } + }); + cx.notify(); + } + + fn cancel(&mut self, _: &editor::Cancel, cx: &mut ViewContext) { + cx.emit(InlineAssistantEvent::Canceled); + } + + fn confirm(&mut self, _: &menu::Confirm, cx: &mut ViewContext) { + if self.confirmed { + cx.emit(InlineAssistantEvent::Dismissed); + } else { + report_assistant_event(self.workspace.clone(), None, AssistantKind::Inline, cx); + + let prompt = self.prompt_editor.read(cx).text(cx); + self.prompt_editor.update(cx, |editor, cx| { + editor.set_read_only(true); + editor.set_field_editor_style( + Some(Arc::new(|theme| { + theme.assistant.inline.disabled_editor.clone() + })), + cx, + ); + }); + cx.emit(InlineAssistantEvent::Confirmed { + prompt, + include_conversation: self.include_conversation, + retrieve_context: self.retrieve_context, + }); + self.confirmed = true; + cx.notify(); + } + } + + fn toggle_retrieve_context(&mut self, _: &ToggleRetrieveContext, cx: &mut ViewContext) { + let semantic_permissioned = self.semantic_permissioned(cx); + + let Some(project) = self.project.upgrade() else { + return; + }; + + let project_name = project + .read(cx) + .worktree_root_names(cx) + .collect::>() + .join("/"); + let is_plural = project_name.chars().filter(|letter| *letter == '/').count() > 0; + let prompt_text = format!("Would you like to index the '{}' project{} for context retrieval? This requires sending code to the OpenAI API", project_name, + if is_plural { + "s" + } else {""}); + + cx.spawn(|this, mut cx| async move { + // If Necessary prompt user + if !semantic_permissioned.await.unwrap_or(false) { + let mut answer = this.update(&mut cx, |_, cx| { + cx.prompt( + PromptLevel::Info, + prompt_text.as_str(), + &["Continue", "Cancel"], + ) + })?; + + if answer.next().await == Some(0) { + this.update(&mut cx, |this, _| { + this.semantic_permissioned = Some(true); + })?; + } else { + return anyhow::Ok(()); + } + } + + // If permissioned, update context appropriately + this.update(&mut cx, |this, cx| { + this.retrieve_context = !this.retrieve_context; + + cx.emit(InlineAssistantEvent::RetrieveContextToggled { + retrieve_context: this.retrieve_context, + }); + + if this.retrieve_context { + this.index_project(cx).log_err(); + } + + cx.notify(); + })?; + + anyhow::Ok(()) + }) + .detach_and_log_err(cx); + } + + fn index_project(&self, cx: &mut ViewContext) -> anyhow::Result<()> { + let Some(project) = self.project.upgrade() else { + return Err(anyhow!("project was dropped!")); + }; + + let semantic_permissioned = self.semantic_permissioned(cx); + if let Some(semantic_index) = SemanticIndex::global(cx) { + cx.spawn(|_, mut cx| async move { + // This has to be updated to accomodate for semantic_permissions + if semantic_permissioned.await.unwrap_or(false) { + semantic_index + .update(&mut cx, |index, cx| index.index_project(project, cx)) + .await + } else { + Err(anyhow!("project is not permissioned for semantic indexing")) + } + }) + .detach_and_log_err(cx); + } + + anyhow::Ok(()) + } + + fn retrieve_context_status( + &self, + cx: &mut ViewContext, + ) -> Option> { + enum ContextStatusIcon {} + + let Some(project) = self.project.upgrade() else { + return None; + }; + + if let Some(semantic_index) = SemanticIndex::global(cx) { + let status = semantic_index.update(cx, |index, _| index.status(&project)); + let theme = theme::current(cx); + match status { + SemanticIndexStatus::NotAuthenticated {} => Some( + Svg::new("icons/error.svg") + .with_color(theme.assistant.error_icon.color) + .constrained() + .with_width(theme.assistant.error_icon.width) + .contained() + .with_style(theme.assistant.error_icon.container) + .with_tooltip::( + self.id, + "Not Authenticated. Please ensure you have a valid 'OPENAI_API_KEY' in your environment variables.", + None, + theme.tooltip.clone(), + cx, + ) + .aligned() + .into_any(), + ), + SemanticIndexStatus::NotIndexed {} => Some( + Svg::new("icons/error.svg") + .with_color(theme.assistant.inline.context_status.error_icon.color) + .constrained() + .with_width(theme.assistant.inline.context_status.error_icon.width) + .contained() + .with_style(theme.assistant.inline.context_status.error_icon.container) + .with_tooltip::( + self.id, + "Not Indexed", + None, + theme.tooltip.clone(), + cx, + ) + .aligned() + .into_any(), + ), + SemanticIndexStatus::Indexing { + remaining_files, + rate_limit_expiry, + } => { + + let mut status_text = if remaining_files == 0 { + "Indexing...".to_string() + } else { + format!("Remaining files to index: {remaining_files}") + }; + + if let Some(rate_limit_expiry) = rate_limit_expiry { + let remaining_seconds = rate_limit_expiry.duration_since(Instant::now()); + if remaining_seconds > Duration::from_secs(0) && remaining_files > 0 { + write!( + status_text, + " (rate limit expires in {}s)", + remaining_seconds.as_secs() + ) + .unwrap(); + } + } + Some( + Svg::new("icons/update.svg") + .with_color(theme.assistant.inline.context_status.in_progress_icon.color) + .constrained() + .with_width(theme.assistant.inline.context_status.in_progress_icon.width) + .contained() + .with_style(theme.assistant.inline.context_status.in_progress_icon.container) + .with_tooltip::( + self.id, + status_text, + None, + theme.tooltip.clone(), + cx, + ) + .aligned() + .into_any(), + ) + } + SemanticIndexStatus::Indexed {} => Some( + Svg::new("icons/check.svg") + .with_color(theme.assistant.inline.context_status.complete_icon.color) + .constrained() + .with_width(theme.assistant.inline.context_status.complete_icon.width) + .contained() + .with_style(theme.assistant.inline.context_status.complete_icon.container) + .with_tooltip::( + self.id, + "Index up to date", + None, + theme.tooltip.clone(), + cx, + ) + .aligned() + .into_any(), + ), + } + } else { + None + } + } + + // fn retrieve_context_status(&self, cx: &mut ViewContext) -> String { + // let project = self.project.clone(); + // if let Some(semantic_index) = self.semantic_index.clone() { + // let status = semantic_index.update(cx, |index, cx| index.status(&project)); + // return match status { + // // This theoretically shouldnt be a valid code path + // // As the inline assistant cant be launched without an API key + // // We keep it here for safety + // semantic_index::SemanticIndexStatus::NotAuthenticated => { + // "Not Authenticated!\nPlease ensure you have an `OPENAI_API_KEY` in your environment variables.".to_string() + // } + // semantic_index::SemanticIndexStatus::Indexed => { + // "Indexing Complete!".to_string() + // } + // semantic_index::SemanticIndexStatus::Indexing { remaining_files, rate_limit_expiry } => { + + // let mut status = format!("Remaining files to index for Context Retrieval: {remaining_files}"); + + // if let Some(rate_limit_expiry) = rate_limit_expiry { + // let remaining_seconds = + // rate_limit_expiry.duration_since(Instant::now()); + // if remaining_seconds > Duration::from_secs(0) { + // write!(status, " (rate limit resets in {}s)", remaining_seconds.as_secs()).unwrap(); + // } + // } + // status + // } + // semantic_index::SemanticIndexStatus::NotIndexed => { + // "Not Indexed for Context Retrieval".to_string() + // } + // }; + // } + + // "".to_string() + // } + + fn toggle_include_conversation( + &mut self, + _: &ToggleIncludeConversation, + cx: &mut ViewContext, + ) { + self.include_conversation = !self.include_conversation; + cx.emit(InlineAssistantEvent::IncludeConversationToggled { + include_conversation: self.include_conversation, + }); + cx.notify(); + } + + fn move_up(&mut self, _: &MoveUp, cx: &mut ViewContext) { + if let Some(ix) = self.prompt_history_ix { + if ix > 0 { + self.prompt_history_ix = Some(ix - 1); + let prompt = self.prompt_history[ix - 1].clone(); + self.set_prompt(&prompt, cx); + } + } else if !self.prompt_history.is_empty() { + self.prompt_history_ix = Some(self.prompt_history.len() - 1); + let prompt = self.prompt_history[self.prompt_history.len() - 1].clone(); + self.set_prompt(&prompt, cx); + } + } + + fn move_down(&mut self, _: &MoveDown, cx: &mut ViewContext) { + if let Some(ix) = self.prompt_history_ix { + if ix < self.prompt_history.len() - 1 { + self.prompt_history_ix = Some(ix + 1); + let prompt = self.prompt_history[ix + 1].clone(); + self.set_prompt(&prompt, cx); + } else { + self.prompt_history_ix = None; + let pending_prompt = self.pending_prompt.clone(); + self.set_prompt(&pending_prompt, cx); + } + } + } + + fn set_prompt(&mut self, prompt: &str, cx: &mut ViewContext) { + self.prompt_editor.update(cx, |editor, cx| { + editor.buffer().update(cx, |buffer, cx| { + let len = buffer.len(cx); + buffer.edit([(0..len, prompt)], None, cx); + }); + }); + } +} + +// This wouldn't need to exist if we could pass parameters when rendering child views. +#[derive(Copy, Clone, Default)] +struct BlockMeasurements { + anchor_x: f32, + gutter_width: f32, +} + +struct PendingInlineAssist { + editor: WeakView, + inline_assistant: Option<(BlockId, View)>, + codegen: Model, + _subscriptions: Vec, + project: WeakModel, +} + +fn merge_ranges(ranges: &mut Vec>, buffer: &MultiBufferSnapshot) { + ranges.sort_unstable_by(|a, b| { + a.start + .cmp(&b.start, buffer) + .then_with(|| b.end.cmp(&a.end, buffer)) + }); + + let mut ix = 0; + while ix + 1 < ranges.len() { + let b = ranges[ix + 1].clone(); + let a = &mut ranges[ix]; + if a.end.cmp(&b.start, buffer).is_gt() { + if a.end.cmp(&b.end, buffer).is_lt() { + a.end = b.end; + } + ranges.remove(ix + 1); + } else { + ix += 1; + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::MessageId; + use ai::test::FakeCompletionProvider; + use gpui::AppContext; + + #[gpui::test] + fn test_inserting_and_removing_messages(cx: &mut AppContext) { + cx.set_global(SettingsStore::test(cx)); + init(cx); + let registry = Arc::new(LanguageRegistry::test()); + + let completion_provider = Arc::new(FakeCompletionProvider::new()); + let conversation = cx.add_model(|cx| Conversation::new(registry, cx, completion_provider)); + let buffer = conversation.read(cx).buffer.clone(); + + let message_1 = conversation.read(cx).message_anchors[0].clone(); + assert_eq!( + messages(&conversation, cx), + vec![(message_1.id, Role::User, 0..0)] + ); + + let message_2 = conversation.update(cx, |conversation, cx| { + conversation + .insert_message_after(message_1.id, Role::Assistant, MessageStatus::Done, cx) + .unwrap() + }); + assert_eq!( + messages(&conversation, cx), + vec![ + (message_1.id, Role::User, 0..1), + (message_2.id, Role::Assistant, 1..1) + ] + ); + + buffer.update(cx, |buffer, cx| { + buffer.edit([(0..0, "1"), (1..1, "2")], None, cx) + }); + assert_eq!( + messages(&conversation, cx), + vec![ + (message_1.id, Role::User, 0..2), + (message_2.id, Role::Assistant, 2..3) + ] + ); + + let message_3 = conversation.update(cx, |conversation, cx| { + conversation + .insert_message_after(message_2.id, Role::User, MessageStatus::Done, cx) + .unwrap() + }); + assert_eq!( + messages(&conversation, cx), + vec![ + (message_1.id, Role::User, 0..2), + (message_2.id, Role::Assistant, 2..4), + (message_3.id, Role::User, 4..4) + ] + ); + + let message_4 = conversation.update(cx, |conversation, cx| { + conversation + .insert_message_after(message_2.id, Role::User, MessageStatus::Done, cx) + .unwrap() + }); + assert_eq!( + messages(&conversation, cx), + vec![ + (message_1.id, Role::User, 0..2), + (message_2.id, Role::Assistant, 2..4), + (message_4.id, Role::User, 4..5), + (message_3.id, Role::User, 5..5), + ] + ); + + buffer.update(cx, |buffer, cx| { + buffer.edit([(4..4, "C"), (5..5, "D")], None, cx) + }); + assert_eq!( + messages(&conversation, cx), + vec![ + (message_1.id, Role::User, 0..2), + (message_2.id, Role::Assistant, 2..4), + (message_4.id, Role::User, 4..6), + (message_3.id, Role::User, 6..7), + ] + ); + + // Deleting across message boundaries merges the messages. + buffer.update(cx, |buffer, cx| buffer.edit([(1..4, "")], None, cx)); + assert_eq!( + messages(&conversation, cx), + vec![ + (message_1.id, Role::User, 0..3), + (message_3.id, Role::User, 3..4), + ] + ); + + // Undoing the deletion should also undo the merge. + buffer.update(cx, |buffer, cx| buffer.undo(cx)); + assert_eq!( + messages(&conversation, cx), + vec![ + (message_1.id, Role::User, 0..2), + (message_2.id, Role::Assistant, 2..4), + (message_4.id, Role::User, 4..6), + (message_3.id, Role::User, 6..7), + ] + ); + + // Redoing the deletion should also redo the merge. + buffer.update(cx, |buffer, cx| buffer.redo(cx)); + assert_eq!( + messages(&conversation, cx), + vec![ + (message_1.id, Role::User, 0..3), + (message_3.id, Role::User, 3..4), + ] + ); + + // Ensure we can still insert after a merged message. + let message_5 = conversation.update(cx, |conversation, cx| { + conversation + .insert_message_after(message_1.id, Role::System, MessageStatus::Done, cx) + .unwrap() + }); + assert_eq!( + messages(&conversation, cx), + vec![ + (message_1.id, Role::User, 0..3), + (message_5.id, Role::System, 3..4), + (message_3.id, Role::User, 4..5) + ] + ); + } + + #[gpui::test] + fn test_message_splitting(cx: &mut AppContext) { + cx.set_global(SettingsStore::test(cx)); + init(cx); + let registry = Arc::new(LanguageRegistry::test()); + let completion_provider = Arc::new(FakeCompletionProvider::new()); + + let conversation = cx.add_model(|cx| Conversation::new(registry, cx, completion_provider)); + let buffer = conversation.read(cx).buffer.clone(); + + let message_1 = conversation.read(cx).message_anchors[0].clone(); + assert_eq!( + messages(&conversation, cx), + vec![(message_1.id, Role::User, 0..0)] + ); + + buffer.update(cx, |buffer, cx| { + buffer.edit([(0..0, "aaa\nbbb\nccc\nddd\n")], None, cx) + }); + + let (_, message_2) = + conversation.update(cx, |conversation, cx| conversation.split_message(3..3, cx)); + let message_2 = message_2.unwrap(); + + // We recycle newlines in the middle of a split message + assert_eq!(buffer.read(cx).text(), "aaa\nbbb\nccc\nddd\n"); + assert_eq!( + messages(&conversation, cx), + vec![ + (message_1.id, Role::User, 0..4), + (message_2.id, Role::User, 4..16), + ] + ); + + let (_, message_3) = + conversation.update(cx, |conversation, cx| conversation.split_message(3..3, cx)); + let message_3 = message_3.unwrap(); + + // We don't recycle newlines at the end of a split message + assert_eq!(buffer.read(cx).text(), "aaa\n\nbbb\nccc\nddd\n"); + assert_eq!( + messages(&conversation, cx), + vec![ + (message_1.id, Role::User, 0..4), + (message_3.id, Role::User, 4..5), + (message_2.id, Role::User, 5..17), + ] + ); + + let (_, message_4) = + conversation.update(cx, |conversation, cx| conversation.split_message(9..9, cx)); + let message_4 = message_4.unwrap(); + assert_eq!(buffer.read(cx).text(), "aaa\n\nbbb\nccc\nddd\n"); + assert_eq!( + messages(&conversation, cx), + vec![ + (message_1.id, Role::User, 0..4), + (message_3.id, Role::User, 4..5), + (message_2.id, Role::User, 5..9), + (message_4.id, Role::User, 9..17), + ] + ); + + let (_, message_5) = + conversation.update(cx, |conversation, cx| conversation.split_message(9..9, cx)); + let message_5 = message_5.unwrap(); + assert_eq!(buffer.read(cx).text(), "aaa\n\nbbb\n\nccc\nddd\n"); + assert_eq!( + messages(&conversation, cx), + vec![ + (message_1.id, Role::User, 0..4), + (message_3.id, Role::User, 4..5), + (message_2.id, Role::User, 5..9), + (message_4.id, Role::User, 9..10), + (message_5.id, Role::User, 10..18), + ] + ); + + let (message_6, message_7) = conversation.update(cx, |conversation, cx| { + conversation.split_message(14..16, cx) + }); + let message_6 = message_6.unwrap(); + let message_7 = message_7.unwrap(); + assert_eq!(buffer.read(cx).text(), "aaa\n\nbbb\n\nccc\ndd\nd\n"); + assert_eq!( + messages(&conversation, cx), + vec![ + (message_1.id, Role::User, 0..4), + (message_3.id, Role::User, 4..5), + (message_2.id, Role::User, 5..9), + (message_4.id, Role::User, 9..10), + (message_5.id, Role::User, 10..14), + (message_6.id, Role::User, 14..17), + (message_7.id, Role::User, 17..19), + ] + ); + } + + #[gpui::test] + fn test_messages_for_offsets(cx: &mut AppContext) { + cx.set_global(SettingsStore::test(cx)); + init(cx); + let registry = Arc::new(LanguageRegistry::test()); + let completion_provider = Arc::new(FakeCompletionProvider::new()); + let conversation = cx.add_model(|cx| Conversation::new(registry, cx, completion_provider)); + let buffer = conversation.read(cx).buffer.clone(); + + let message_1 = conversation.read(cx).message_anchors[0].clone(); + assert_eq!( + messages(&conversation, cx), + vec![(message_1.id, Role::User, 0..0)] + ); + + buffer.update(cx, |buffer, cx| buffer.edit([(0..0, "aaa")], None, cx)); + let message_2 = conversation + .update(cx, |conversation, cx| { + conversation.insert_message_after(message_1.id, Role::User, MessageStatus::Done, cx) + }) + .unwrap(); + buffer.update(cx, |buffer, cx| buffer.edit([(4..4, "bbb")], None, cx)); + + let message_3 = conversation + .update(cx, |conversation, cx| { + conversation.insert_message_after(message_2.id, Role::User, MessageStatus::Done, cx) + }) + .unwrap(); + buffer.update(cx, |buffer, cx| buffer.edit([(8..8, "ccc")], None, cx)); + + assert_eq!(buffer.read(cx).text(), "aaa\nbbb\nccc"); + assert_eq!( + messages(&conversation, cx), + vec![ + (message_1.id, Role::User, 0..4), + (message_2.id, Role::User, 4..8), + (message_3.id, Role::User, 8..11) + ] + ); + + assert_eq!( + message_ids_for_offsets(&conversation, &[0, 4, 9], cx), + [message_1.id, message_2.id, message_3.id] + ); + assert_eq!( + message_ids_for_offsets(&conversation, &[0, 1, 11], cx), + [message_1.id, message_3.id] + ); + + let message_4 = conversation + .update(cx, |conversation, cx| { + conversation.insert_message_after(message_3.id, Role::User, MessageStatus::Done, cx) + }) + .unwrap(); + assert_eq!(buffer.read(cx).text(), "aaa\nbbb\nccc\n"); + assert_eq!( + messages(&conversation, cx), + vec![ + (message_1.id, Role::User, 0..4), + (message_2.id, Role::User, 4..8), + (message_3.id, Role::User, 8..12), + (message_4.id, Role::User, 12..12) + ] + ); + assert_eq!( + message_ids_for_offsets(&conversation, &[0, 4, 8, 12], cx), + [message_1.id, message_2.id, message_3.id, message_4.id] + ); + + fn message_ids_for_offsets( + conversation: &Model, + offsets: &[usize], + cx: &AppContext, + ) -> Vec { + conversation + .read(cx) + .messages_for_offsets(offsets.iter().copied(), cx) + .into_iter() + .map(|message| message.id) + .collect() + } + } + + #[gpui::test] + fn test_serialization(cx: &mut AppContext) { + cx.set_global(SettingsStore::test(cx)); + init(cx); + let registry = Arc::new(LanguageRegistry::test()); + let completion_provider = Arc::new(FakeCompletionProvider::new()); + let conversation = + cx.add_model(|cx| Conversation::new(registry.clone(), cx, completion_provider)); + let buffer = conversation.read(cx).buffer.clone(); + let message_0 = conversation.read(cx).message_anchors[0].id; + let message_1 = conversation.update(cx, |conversation, cx| { + conversation + .insert_message_after(message_0, Role::Assistant, MessageStatus::Done, cx) + .unwrap() + }); + let message_2 = conversation.update(cx, |conversation, cx| { + conversation + .insert_message_after(message_1.id, Role::System, MessageStatus::Done, cx) + .unwrap() + }); + buffer.update(cx, |buffer, cx| { + buffer.edit([(0..0, "a"), (1..1, "b\nc")], None, cx); + buffer.finalize_last_transaction(); + }); + let _message_3 = conversation.update(cx, |conversation, cx| { + conversation + .insert_message_after(message_2.id, Role::System, MessageStatus::Done, cx) + .unwrap() + }); + buffer.update(cx, |buffer, cx| buffer.undo(cx)); + assert_eq!(buffer.read(cx).text(), "a\nb\nc\n"); + assert_eq!( + messages(&conversation, cx), + [ + (message_0, Role::User, 0..2), + (message_1.id, Role::Assistant, 2..6), + (message_2.id, Role::System, 6..6), + ] + ); + + let deserialized_conversation = cx.add_model(|cx| { + Conversation::deserialize( + conversation.read(cx).serialize(cx), + Default::default(), + registry.clone(), + cx, + ) + }); + let deserialized_buffer = deserialized_conversation.read(cx).buffer.clone(); + assert_eq!(deserialized_buffer.read(cx).text(), "a\nb\nc\n"); + assert_eq!( + messages(&deserialized_conversation, cx), + [ + (message_0, Role::User, 0..2), + (message_1.id, Role::Assistant, 2..6), + (message_2.id, Role::System, 6..6), + ] + ); + } + + fn messages( + conversation: &Model, + cx: &AppContext, + ) -> Vec<(MessageId, Role, Range)> { + conversation + .read(cx) + .messages(cx) + .map(|message| (message.id, message.role, message.offset_range)) + .collect() + } +} + +fn report_assistant_event( + workspace: WeakView, + conversation_id: Option, + assistant_kind: AssistantKind, + cx: &AppContext, +) { + let Some(workspace) = workspace.upgrade() else { + return; + }; + + let client = workspace.read(cx).project().read(cx).client(); + let telemetry = client.telemetry(); + + let model = AssistantSettings::get_global(cx) + .default_open_ai_model + .clone(); + + let telemetry_settings = TelemetrySettings::get_global(cx); + + telemetry.report_assistant_event( + telemetry_settings, + conversation_id, + assistant_kind, + model.full_name(), + ) +} diff --git a/crates/assistant2/src/assistant_settings.rs b/crates/assistant2/src/assistant_settings.rs new file mode 100644 index 0000000000..5a727df621 --- /dev/null +++ b/crates/assistant2/src/assistant_settings.rs @@ -0,0 +1,80 @@ +use anyhow; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; +use settings::Settings; + +#[derive(Clone, Debug, Serialize, Deserialize, JsonSchema, PartialEq)] +pub enum OpenAIModel { + #[serde(rename = "gpt-3.5-turbo-0613")] + ThreePointFiveTurbo, + #[serde(rename = "gpt-4-0613")] + Four, + #[serde(rename = "gpt-4-1106-preview")] + FourTurbo, +} + +impl OpenAIModel { + pub fn full_name(&self) -> &'static str { + match self { + OpenAIModel::ThreePointFiveTurbo => "gpt-3.5-turbo-0613", + OpenAIModel::Four => "gpt-4-0613", + OpenAIModel::FourTurbo => "gpt-4-1106-preview", + } + } + + pub fn short_name(&self) -> &'static str { + match self { + OpenAIModel::ThreePointFiveTurbo => "gpt-3.5-turbo", + OpenAIModel::Four => "gpt-4", + OpenAIModel::FourTurbo => "gpt-4-turbo", + } + } + + pub fn cycle(&self) -> Self { + match self { + OpenAIModel::ThreePointFiveTurbo => OpenAIModel::Four, + OpenAIModel::Four => OpenAIModel::FourTurbo, + OpenAIModel::FourTurbo => OpenAIModel::ThreePointFiveTurbo, + } + } +} + +#[derive(Clone, Debug, Serialize, Deserialize, JsonSchema)] +#[serde(rename_all = "snake_case")] +pub enum AssistantDockPosition { + Left, + Right, + Bottom, +} + +#[derive(Deserialize, Debug)] +pub struct AssistantSettings { + pub button: bool, + pub dock: AssistantDockPosition, + pub default_width: f32, + pub default_height: f32, + pub default_open_ai_model: OpenAIModel, +} + +#[derive(Clone, Default, Serialize, Deserialize, JsonSchema, Debug)] +pub struct AssistantSettingsContent { + pub button: Option, + pub dock: Option, + pub default_width: Option, + pub default_height: Option, + pub default_open_ai_model: Option, +} + +impl Settings for AssistantSettings { + const KEY: Option<&'static str> = Some("assistant"); + + type FileContent = AssistantSettingsContent; + + fn load( + default_value: &Self::FileContent, + user_values: &[&Self::FileContent], + _: &mut gpui::AppContext, + ) -> anyhow::Result { + Self::load_via_json_merge(default_value, user_values) + } +} diff --git a/crates/assistant2/src/codegen.rs b/crates/assistant2/src/codegen.rs new file mode 100644 index 0000000000..9696c629ac --- /dev/null +++ b/crates/assistant2/src/codegen.rs @@ -0,0 +1,695 @@ +use crate::streaming_diff::{Hunk, StreamingDiff}; +use ai::completion::{CompletionProvider, CompletionRequest}; +use anyhow::Result; +use editor::{Anchor, MultiBuffer, MultiBufferSnapshot, ToOffset, ToPoint}; +use futures::{channel::mpsc, SinkExt, Stream, StreamExt}; +use gpui::{EventEmitter, Model, ModelContext, Task}; +use language::{Rope, TransactionId}; +use multi_buffer; +use std::{cmp, future, ops::Range, sync::Arc}; + +pub enum Event { + Finished, + Undone, +} + +#[derive(Clone)] +pub enum CodegenKind { + Transform { range: Range }, + Generate { position: Anchor }, +} + +pub struct Codegen { + provider: Arc, + buffer: Model, + snapshot: MultiBufferSnapshot, + kind: CodegenKind, + last_equal_ranges: Vec>, + transaction_id: Option, + error: Option, + generation: Task<()>, + idle: bool, + _subscription: gpui::Subscription, +} + +impl EventEmitter for Codegen {} + +impl Codegen { + pub fn new( + buffer: Model, + kind: CodegenKind, + provider: Arc, + cx: &mut ModelContext, + ) -> Self { + let snapshot = buffer.read(cx).snapshot(cx); + Self { + provider, + buffer: buffer.clone(), + snapshot, + kind, + last_equal_ranges: Default::default(), + transaction_id: Default::default(), + error: Default::default(), + idle: true, + generation: Task::ready(()), + _subscription: cx.subscribe(&buffer, Self::handle_buffer_event), + } + } + + fn handle_buffer_event( + &mut self, + _buffer: Model, + event: &multi_buffer::Event, + cx: &mut ModelContext, + ) { + if let multi_buffer::Event::TransactionUndone { transaction_id } = event { + if self.transaction_id == Some(*transaction_id) { + self.transaction_id = None; + self.generation = Task::ready(()); + cx.emit(Event::Undone); + } + } + } + + pub fn range(&self) -> Range { + match &self.kind { + CodegenKind::Transform { range } => range.clone(), + CodegenKind::Generate { position } => position.bias_left(&self.snapshot)..*position, + } + } + + pub fn kind(&self) -> &CodegenKind { + &self.kind + } + + pub fn last_equal_ranges(&self) -> &[Range] { + &self.last_equal_ranges + } + + pub fn idle(&self) -> bool { + self.idle + } + + pub fn error(&self) -> Option<&anyhow::Error> { + self.error.as_ref() + } + + pub fn start(&mut self, prompt: Box, cx: &mut ModelContext) { + let range = self.range(); + let snapshot = self.snapshot.clone(); + let selected_text = snapshot + .text_for_range(range.start..range.end) + .collect::(); + + let selection_start = range.start.to_point(&snapshot); + let suggested_line_indent = snapshot + .suggested_indents(selection_start.row..selection_start.row + 1, cx) + .into_values() + .next() + .unwrap_or_else(|| snapshot.indent_size_for_line(selection_start.row)); + + let response = self.provider.complete(prompt); + self.generation = cx.spawn_weak(|this, mut cx| { + async move { + let generate = async { + let mut edit_start = range.start.to_offset(&snapshot); + + let (mut hunks_tx, mut hunks_rx) = mpsc::channel(1); + let diff = cx.background().spawn(async move { + let chunks = strip_invalid_spans_from_codeblock(response.await?); + futures::pin_mut!(chunks); + let mut diff = StreamingDiff::new(selected_text.to_string()); + + let mut new_text = String::new(); + let mut base_indent = None; + let mut line_indent = None; + let mut first_line = true; + + while let Some(chunk) = chunks.next().await { + let chunk = chunk?; + + let mut lines = chunk.split('\n').peekable(); + while let Some(line) = lines.next() { + new_text.push_str(line); + if line_indent.is_none() { + if let Some(non_whitespace_ch_ix) = + new_text.find(|ch: char| !ch.is_whitespace()) + { + line_indent = Some(non_whitespace_ch_ix); + base_indent = base_indent.or(line_indent); + + let line_indent = line_indent.unwrap(); + let base_indent = base_indent.unwrap(); + let indent_delta = line_indent as i32 - base_indent as i32; + let mut corrected_indent_len = cmp::max( + 0, + suggested_line_indent.len as i32 + indent_delta, + ) + as usize; + if first_line { + corrected_indent_len = corrected_indent_len + .saturating_sub(selection_start.column as usize); + } + + let indent_char = suggested_line_indent.char(); + let mut indent_buffer = [0; 4]; + let indent_str = + indent_char.encode_utf8(&mut indent_buffer); + new_text.replace_range( + ..line_indent, + &indent_str.repeat(corrected_indent_len), + ); + } + } + + if line_indent.is_some() { + hunks_tx.send(diff.push_new(&new_text)).await?; + new_text.clear(); + } + + if lines.peek().is_some() { + hunks_tx.send(diff.push_new("\n")).await?; + line_indent = None; + first_line = false; + } + } + } + hunks_tx.send(diff.push_new(&new_text)).await?; + hunks_tx.send(diff.finish()).await?; + + anyhow::Ok(()) + }); + + while let Some(hunks) = hunks_rx.next().await { + let this = if let Some(this) = this.upgrade(&cx) { + this + } else { + break; + }; + + this.update(&mut cx, |this, cx| { + this.last_equal_ranges.clear(); + + let transaction = this.buffer.update(cx, |buffer, cx| { + // Avoid grouping assistant edits with user edits. + buffer.finalize_last_transaction(cx); + + buffer.start_transaction(cx); + buffer.edit( + hunks.into_iter().filter_map(|hunk| match hunk { + Hunk::Insert { text } => { + let edit_start = snapshot.anchor_after(edit_start); + Some((edit_start..edit_start, text)) + } + Hunk::Remove { len } => { + let edit_end = edit_start + len; + let edit_range = snapshot.anchor_after(edit_start) + ..snapshot.anchor_before(edit_end); + edit_start = edit_end; + Some((edit_range, String::new())) + } + Hunk::Keep { len } => { + let edit_end = edit_start + len; + let edit_range = snapshot.anchor_after(edit_start) + ..snapshot.anchor_before(edit_end); + edit_start = edit_end; + this.last_equal_ranges.push(edit_range); + None + } + }), + None, + cx, + ); + + buffer.end_transaction(cx) + }); + + if let Some(transaction) = transaction { + if let Some(first_transaction) = this.transaction_id { + // Group all assistant edits into the first transaction. + this.buffer.update(cx, |buffer, cx| { + buffer.merge_transactions( + transaction, + first_transaction, + cx, + ) + }); + } else { + this.transaction_id = Some(transaction); + this.buffer.update(cx, |buffer, cx| { + buffer.finalize_last_transaction(cx) + }); + } + } + + cx.notify(); + }); + } + + diff.await?; + anyhow::Ok(()) + }; + + let result = generate.await; + if let Some(this) = this.upgrade(&cx) { + this.update(&mut cx, |this, cx| { + this.last_equal_ranges.clear(); + this.idle = true; + if let Err(error) = result { + this.error = Some(error); + } + cx.emit(Event::Finished); + cx.notify(); + }); + } + } + }); + self.error.take(); + self.idle = false; + cx.notify(); + } + + pub fn undo(&mut self, cx: &mut ModelContext) { + if let Some(transaction_id) = self.transaction_id { + self.buffer + .update(cx, |buffer, cx| buffer.undo_transaction(transaction_id, cx)); + } + } +} + +fn strip_invalid_spans_from_codeblock( + stream: impl Stream>, +) -> impl Stream> { + let mut first_line = true; + let mut buffer = String::new(); + let mut starts_with_markdown_codeblock = false; + let mut includes_start_or_end_span = false; + stream.filter_map(move |chunk| { + let chunk = match chunk { + Ok(chunk) => chunk, + Err(err) => return future::ready(Some(Err(err))), + }; + buffer.push_str(&chunk); + + if buffer.len() > "<|S|".len() && buffer.starts_with("<|S|") { + includes_start_or_end_span = true; + + buffer = buffer + .strip_prefix("<|S|>") + .or_else(|| buffer.strip_prefix("<|S|")) + .unwrap_or(&buffer) + .to_string(); + } else if buffer.ends_with("|E|>") { + includes_start_or_end_span = true; + } else if buffer.starts_with("<|") + || buffer.starts_with("<|S") + || buffer.starts_with("<|S|") + || buffer.ends_with("|") + || buffer.ends_with("|E") + || buffer.ends_with("|E|") + { + return future::ready(None); + } + + if first_line { + if buffer == "" || buffer == "`" || buffer == "``" { + return future::ready(None); + } else if buffer.starts_with("```") { + starts_with_markdown_codeblock = true; + if let Some(newline_ix) = buffer.find('\n') { + buffer.replace_range(..newline_ix + 1, ""); + first_line = false; + } else { + return future::ready(None); + } + } + } + + let mut text = buffer.to_string(); + if starts_with_markdown_codeblock { + text = text + .strip_suffix("\n```\n") + .or_else(|| text.strip_suffix("\n```")) + .or_else(|| text.strip_suffix("\n``")) + .or_else(|| text.strip_suffix("\n`")) + .or_else(|| text.strip_suffix('\n')) + .unwrap_or(&text) + .to_string(); + } + + if includes_start_or_end_span { + text = text + .strip_suffix("|E|>") + .or_else(|| text.strip_suffix("E|>")) + .or_else(|| text.strip_prefix("|>")) + .or_else(|| text.strip_prefix(">")) + .unwrap_or(&text) + .to_string(); + }; + + if text.contains('\n') { + first_line = false; + } + + let remainder = buffer.split_off(text.len()); + let result = if buffer.is_empty() { + None + } else { + Some(Ok(buffer.clone())) + }; + + buffer = remainder; + future::ready(result) + }) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use super::*; + use ai::test::FakeCompletionProvider; + use futures::stream::{self}; + use gpui::TestAppContext; + use indoc::indoc; + use language::{language_settings, tree_sitter_rust, Buffer, Language, LanguageConfig, Point}; + use rand::prelude::*; + use serde::Serialize; + use settings::SettingsStore; + + #[derive(Serialize)] + pub struct DummyCompletionRequest { + pub name: String, + } + + impl CompletionRequest for DummyCompletionRequest { + fn data(&self) -> serde_json::Result { + serde_json::to_string(self) + } + } + + #[gpui::test(iterations = 10)] + async fn test_transform_autoindent(cx: &mut TestAppContext, mut rng: StdRng) { + cx.set_global(cx.read(SettingsStore::test)); + cx.update(language_settings::init); + + let text = indoc! {" + fn main() { + let x = 0; + for _ in 0..10 { + x += 1; + } + } + "}; + let buffer = + cx.add_model(|cx| Buffer::new(0, 0, text).with_language(Arc::new(rust_lang()), cx)); + let buffer = cx.add_model(|cx| MultiBuffer::singleton(buffer, cx)); + let range = buffer.read_with(cx, |buffer, cx| { + let snapshot = buffer.snapshot(cx); + snapshot.anchor_before(Point::new(1, 0))..snapshot.anchor_after(Point::new(4, 5)) + }); + let provider = Arc::new(FakeCompletionProvider::new()); + let codegen = cx.add_model(|cx| { + Codegen::new( + buffer.clone(), + CodegenKind::Transform { range }, + provider.clone(), + cx, + ) + }); + + let request = Box::new(DummyCompletionRequest { + name: "test".to_string(), + }); + codegen.update(cx, |codegen, cx| codegen.start(request, cx)); + + let mut new_text = concat!( + " let mut x = 0;\n", + " while x < 10 {\n", + " x += 1;\n", + " }", + ); + while !new_text.is_empty() { + let max_len = cmp::min(new_text.len(), 10); + let len = rng.gen_range(1..=max_len); + let (chunk, suffix) = new_text.split_at(len); + println!("CHUNK: {:?}", &chunk); + provider.send_completion(chunk); + new_text = suffix; + cx.background_executor.run_until_parked(); + } + provider.finish_completion(); + cx.background_executor.run_until_parked(); + + assert_eq!( + buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()), + indoc! {" + fn main() { + let mut x = 0; + while x < 10 { + x += 1; + } + } + "} + ); + } + + #[gpui::test(iterations = 10)] + async fn test_autoindent_when_generating_past_indentation( + cx: &mut TestAppContext, + mut rng: StdRng, + ) { + cx.set_global(cx.read(SettingsStore::test)); + cx.update(language_settings::init); + + let text = indoc! {" + fn main() { + le + } + "}; + let buffer = + cx.add_model(|cx| Buffer::new(0, 0, text).with_language(Arc::new(rust_lang()), cx)); + let buffer = cx.add_model(|cx| MultiBuffer::singleton(buffer, cx)); + let position = buffer.read_with(cx, |buffer, cx| { + let snapshot = buffer.snapshot(cx); + snapshot.anchor_before(Point::new(1, 6)) + }); + let provider = Arc::new(FakeCompletionProvider::new()); + let codegen = cx.add_model(|cx| { + Codegen::new( + buffer.clone(), + CodegenKind::Generate { position }, + provider.clone(), + cx, + ) + }); + + let request = Box::new(DummyCompletionRequest { + name: "test".to_string(), + }); + codegen.update(cx, |codegen, cx| codegen.start(request, cx)); + + let mut new_text = concat!( + "t mut x = 0;\n", + "while x < 10 {\n", + " x += 1;\n", + "}", // + ); + while !new_text.is_empty() { + let max_len = cmp::min(new_text.len(), 10); + let len = rng.gen_range(1..=max_len); + let (chunk, suffix) = new_text.split_at(len); + provider.send_completion(chunk); + new_text = suffix; + cx.background_executor.run_until_parked(); + } + provider.finish_completion(); + cx.background_executor.run_until_parked(); + + assert_eq!( + buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()), + indoc! {" + fn main() { + let mut x = 0; + while x < 10 { + x += 1; + } + } + "} + ); + } + + #[gpui::test(iterations = 10)] + async fn test_autoindent_when_generating_before_indentation( + cx: &mut TestAppContext, + mut rng: StdRng, + ) { + cx.set_global(cx.read(SettingsStore::test)); + cx.update(language_settings::init); + + let text = concat!( + "fn main() {\n", + " \n", + "}\n" // + ); + let buffer = + cx.add_model(|cx| Buffer::new(0, 0, text).with_language(Arc::new(rust_lang()), cx)); + let buffer = cx.add_model(|cx| MultiBuffer::singleton(buffer, cx)); + let position = buffer.read_with(cx, |buffer, cx| { + let snapshot = buffer.snapshot(cx); + snapshot.anchor_before(Point::new(1, 2)) + }); + let provider = Arc::new(FakeCompletionProvider::new()); + let codegen = cx.add_model(|cx| { + Codegen::new( + buffer.clone(), + CodegenKind::Generate { position }, + provider.clone(), + cx, + ) + }); + + let request = Box::new(DummyCompletionRequest { + name: "test".to_string(), + }); + codegen.update(cx, |codegen, cx| codegen.start(request, cx)); + + let mut new_text = concat!( + "let mut x = 0;\n", + "while x < 10 {\n", + " x += 1;\n", + "}", // + ); + while !new_text.is_empty() { + let max_len = cmp::min(new_text.len(), 10); + let len = rng.gen_range(1..=max_len); + let (chunk, suffix) = new_text.split_at(len); + println!("{:?}", &chunk); + provider.send_completion(chunk); + new_text = suffix; + cx.background_executor.run_until_parked(); + } + provider.finish_completion(); + cx.background_executor.run_until_parked(); + + assert_eq!( + buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()), + indoc! {" + fn main() { + let mut x = 0; + while x < 10 { + x += 1; + } + } + "} + ); + } + + #[gpui::test] + async fn test_strip_invalid_spans_from_codeblock() { + assert_eq!( + strip_invalid_spans_from_codeblock(chunks("Lorem ipsum dolor", 2)) + .map(|chunk| chunk.unwrap()) + .collect::() + .await, + "Lorem ipsum dolor" + ); + assert_eq!( + strip_invalid_spans_from_codeblock(chunks("```\nLorem ipsum dolor", 2)) + .map(|chunk| chunk.unwrap()) + .collect::() + .await, + "Lorem ipsum dolor" + ); + assert_eq!( + strip_invalid_spans_from_codeblock(chunks("```\nLorem ipsum dolor\n```", 2)) + .map(|chunk| chunk.unwrap()) + .collect::() + .await, + "Lorem ipsum dolor" + ); + assert_eq!( + strip_invalid_spans_from_codeblock(chunks("```\nLorem ipsum dolor\n```\n", 2)) + .map(|chunk| chunk.unwrap()) + .collect::() + .await, + "Lorem ipsum dolor" + ); + assert_eq!( + strip_invalid_spans_from_codeblock(chunks( + "```html\n```js\nLorem ipsum dolor\n```\n```", + 2 + )) + .map(|chunk| chunk.unwrap()) + .collect::() + .await, + "```js\nLorem ipsum dolor\n```" + ); + assert_eq!( + strip_invalid_spans_from_codeblock(chunks("``\nLorem ipsum dolor\n```", 2)) + .map(|chunk| chunk.unwrap()) + .collect::() + .await, + "``\nLorem ipsum dolor\n```" + ); + assert_eq!( + strip_invalid_spans_from_codeblock(chunks("<|S|Lorem ipsum|E|>", 2)) + .map(|chunk| chunk.unwrap()) + .collect::() + .await, + "Lorem ipsum" + ); + + assert_eq!( + strip_invalid_spans_from_codeblock(chunks("<|S|>Lorem ipsum", 2)) + .map(|chunk| chunk.unwrap()) + .collect::() + .await, + "Lorem ipsum" + ); + + assert_eq!( + strip_invalid_spans_from_codeblock(chunks("```\n<|S|>Lorem ipsum\n```", 2)) + .map(|chunk| chunk.unwrap()) + .collect::() + .await, + "Lorem ipsum" + ); + assert_eq!( + strip_invalid_spans_from_codeblock(chunks("```\n<|S|Lorem ipsum|E|>\n```", 2)) + .map(|chunk| chunk.unwrap()) + .collect::() + .await, + "Lorem ipsum" + ); + fn chunks(text: &str, size: usize) -> impl Stream> { + stream::iter( + text.chars() + .collect::>() + .chunks(size) + .map(|chunk| Ok(chunk.iter().collect::())) + .collect::>(), + ) + } + } + + fn rust_lang() -> Language { + Language::new( + LanguageConfig { + name: "Rust".into(), + path_suffixes: vec!["rs".to_string()], + ..Default::default() + }, + Some(tree_sitter_rust::language()), + ) + .with_indents_query( + r#" + (call_expression) @indent + (field_expression) @indent + (_ "(" ")" @end) @indent + (_ "{" "}" @end) @indent + "#, + ) + .unwrap() + } +} diff --git a/crates/assistant2/src/prompts.rs b/crates/assistant2/src/prompts.rs new file mode 100644 index 0000000000..b678c6fe3b --- /dev/null +++ b/crates/assistant2/src/prompts.rs @@ -0,0 +1,388 @@ +use ai::models::LanguageModel; +use ai::prompts::base::{PromptArguments, PromptChain, PromptPriority, PromptTemplate}; +use ai::prompts::file_context::FileContext; +use ai::prompts::generate::GenerateInlineContent; +use ai::prompts::preamble::EngineerPreamble; +use ai::prompts::repository_context::{PromptCodeSnippet, RepositoryContext}; +use ai::providers::open_ai::OpenAILanguageModel; +use language::{BufferSnapshot, OffsetRangeExt, ToOffset}; +use std::cmp::{self, Reverse}; +use std::ops::Range; +use std::sync::Arc; + +#[allow(dead_code)] +fn summarize(buffer: &BufferSnapshot, selected_range: Range) -> String { + #[derive(Debug)] + struct Match { + collapse: Range, + keep: Vec>, + } + + let selected_range = selected_range.to_offset(buffer); + let mut ts_matches = buffer.matches(0..buffer.len(), |grammar| { + Some(&grammar.embedding_config.as_ref()?.query) + }); + let configs = ts_matches + .grammars() + .iter() + .map(|g| g.embedding_config.as_ref().unwrap()) + .collect::>(); + let mut matches = Vec::new(); + while let Some(mat) = ts_matches.peek() { + let config = &configs[mat.grammar_index]; + if let Some(collapse) = mat.captures.iter().find_map(|cap| { + if Some(cap.index) == config.collapse_capture_ix { + Some(cap.node.byte_range()) + } else { + None + } + }) { + let mut keep = Vec::new(); + for capture in mat.captures.iter() { + if Some(capture.index) == config.keep_capture_ix { + keep.push(capture.node.byte_range()); + } else { + continue; + } + } + ts_matches.advance(); + matches.push(Match { collapse, keep }); + } else { + ts_matches.advance(); + } + } + matches.sort_unstable_by_key(|mat| (mat.collapse.start, Reverse(mat.collapse.end))); + let mut matches = matches.into_iter().peekable(); + + let mut summary = String::new(); + let mut offset = 0; + let mut flushed_selection = false; + while let Some(mat) = matches.next() { + // Keep extending the collapsed range if the next match surrounds + // the current one. + while let Some(next_mat) = matches.peek() { + if mat.collapse.start <= next_mat.collapse.start + && mat.collapse.end >= next_mat.collapse.end + { + matches.next().unwrap(); + } else { + break; + } + } + + if offset > mat.collapse.start { + // Skip collapsed nodes that have already been summarized. + offset = cmp::max(offset, mat.collapse.end); + continue; + } + + if offset <= selected_range.start && selected_range.start <= mat.collapse.end { + if !flushed_selection { + // The collapsed node ends after the selection starts, so we'll flush the selection first. + summary.extend(buffer.text_for_range(offset..selected_range.start)); + summary.push_str("<|S|"); + if selected_range.end == selected_range.start { + summary.push_str(">"); + } else { + summary.extend(buffer.text_for_range(selected_range.clone())); + summary.push_str("|E|>"); + } + offset = selected_range.end; + flushed_selection = true; + } + + // If the selection intersects the collapsed node, we won't collapse it. + if selected_range.end >= mat.collapse.start { + continue; + } + } + + summary.extend(buffer.text_for_range(offset..mat.collapse.start)); + for keep in mat.keep { + summary.extend(buffer.text_for_range(keep)); + } + offset = mat.collapse.end; + } + + // Flush selection if we haven't already done so. + if !flushed_selection && offset <= selected_range.start { + summary.extend(buffer.text_for_range(offset..selected_range.start)); + summary.push_str("<|S|"); + if selected_range.end == selected_range.start { + summary.push_str(">"); + } else { + summary.extend(buffer.text_for_range(selected_range.clone())); + summary.push_str("|E|>"); + } + offset = selected_range.end; + } + + summary.extend(buffer.text_for_range(offset..buffer.len())); + summary +} + +pub fn generate_content_prompt( + user_prompt: String, + language_name: Option<&str>, + buffer: BufferSnapshot, + range: Range, + search_results: Vec, + model: &str, + project_name: Option, +) -> anyhow::Result { + // Using new Prompt Templates + let openai_model: Arc = Arc::new(OpenAILanguageModel::load(model)); + let lang_name = if let Some(language_name) = language_name { + Some(language_name.to_string()) + } else { + None + }; + + let args = PromptArguments { + model: openai_model, + language_name: lang_name.clone(), + project_name, + snippets: search_results.clone(), + reserved_tokens: 1000, + buffer: Some(buffer), + selected_range: Some(range), + user_prompt: Some(user_prompt.clone()), + }; + + let templates: Vec<(PromptPriority, Box)> = vec![ + (PromptPriority::Mandatory, Box::new(EngineerPreamble {})), + ( + PromptPriority::Ordered { order: 1 }, + Box::new(RepositoryContext {}), + ), + ( + PromptPriority::Ordered { order: 0 }, + Box::new(FileContext {}), + ), + ( + PromptPriority::Mandatory, + Box::new(GenerateInlineContent {}), + ), + ]; + let chain = PromptChain::new(args, templates); + let (prompt, _) = chain.generate(true)?; + + anyhow::Ok(prompt) +} + +#[cfg(test)] +pub(crate) mod tests { + + use super::*; + use std::sync::Arc; + + use gpui::AppContext; + use indoc::indoc; + use language::{language_settings, tree_sitter_rust, Buffer, Language, LanguageConfig, Point}; + use settings::SettingsStore; + + pub(crate) fn rust_lang() -> Language { + Language::new( + LanguageConfig { + name: "Rust".into(), + path_suffixes: vec!["rs".to_string()], + ..Default::default() + }, + Some(tree_sitter_rust::language()), + ) + .with_embedding_query( + r#" + ( + [(line_comment) (attribute_item)]* @context + . + [ + (struct_item + name: (_) @name) + + (enum_item + name: (_) @name) + + (impl_item + trait: (_)? @name + "for"? @name + type: (_) @name) + + (trait_item + name: (_) @name) + + (function_item + name: (_) @name + body: (block + "{" @keep + "}" @keep) @collapse) + + (macro_definition + name: (_) @name) + ] @item + ) + "#, + ) + .unwrap() + } + + #[gpui::test] + fn test_outline_for_prompt(cx: &mut AppContext) { + cx.set_global(SettingsStore::test(cx)); + language_settings::init(cx); + let text = indoc! {" + struct X { + a: usize, + b: usize, + } + + impl X { + + fn new() -> Self { + let a = 1; + let b = 2; + Self { a, b } + } + + pub fn a(&self, param: bool) -> usize { + self.a + } + + pub fn b(&self) -> usize { + self.b + } + } + "}; + let buffer = + cx.add_model(|cx| Buffer::new(0, 0, text).with_language(Arc::new(rust_lang()), cx)); + let snapshot = buffer.read(cx).snapshot(); + + assert_eq!( + summarize(&snapshot, Point::new(1, 4)..Point::new(1, 4)), + indoc! {" + struct X { + <|S|>a: usize, + b: usize, + } + + impl X { + + fn new() -> Self {} + + pub fn a(&self, param: bool) -> usize {} + + pub fn b(&self) -> usize {} + } + "} + ); + + assert_eq!( + summarize(&snapshot, Point::new(8, 12)..Point::new(8, 14)), + indoc! {" + struct X { + a: usize, + b: usize, + } + + impl X { + + fn new() -> Self { + let <|S|a |E|>= 1; + let b = 2; + Self { a, b } + } + + pub fn a(&self, param: bool) -> usize {} + + pub fn b(&self) -> usize {} + } + "} + ); + + assert_eq!( + summarize(&snapshot, Point::new(6, 0)..Point::new(6, 0)), + indoc! {" + struct X { + a: usize, + b: usize, + } + + impl X { + <|S|> + fn new() -> Self {} + + pub fn a(&self, param: bool) -> usize {} + + pub fn b(&self) -> usize {} + } + "} + ); + + assert_eq!( + summarize(&snapshot, Point::new(21, 0)..Point::new(21, 0)), + indoc! {" + struct X { + a: usize, + b: usize, + } + + impl X { + + fn new() -> Self {} + + pub fn a(&self, param: bool) -> usize {} + + pub fn b(&self) -> usize {} + } + <|S|>"} + ); + + // Ensure nested functions get collapsed properly. + let text = indoc! {" + struct X { + a: usize, + b: usize, + } + + impl X { + + fn new() -> Self { + let a = 1; + let b = 2; + Self { a, b } + } + + pub fn a(&self, param: bool) -> usize { + let a = 30; + fn nested() -> usize { + 3 + } + self.a + nested() + } + + pub fn b(&self) -> usize { + self.b + } + } + "}; + buffer.update(cx, |buffer, cx| buffer.set_text(text, cx)); + let snapshot = buffer.read(cx).snapshot(); + assert_eq!( + summarize(&snapshot, Point::new(0, 0)..Point::new(0, 0)), + indoc! {" + <|S|>struct X { + a: usize, + b: usize, + } + + impl X { + + fn new() -> Self {} + + pub fn a(&self, param: bool) -> usize {} + + pub fn b(&self) -> usize {} + } + "} + ); + } +} diff --git a/crates/assistant2/src/streaming_diff.rs b/crates/assistant2/src/streaming_diff.rs new file mode 100644 index 0000000000..7399a7b4fa --- /dev/null +++ b/crates/assistant2/src/streaming_diff.rs @@ -0,0 +1,293 @@ +use collections::HashMap; +use ordered_float::OrderedFloat; +use std::{ + cmp, + fmt::{self, Debug}, + ops::Range, +}; + +struct Matrix { + cells: Vec, + rows: usize, + cols: usize, +} + +impl Matrix { + fn new() -> Self { + Self { + cells: Vec::new(), + rows: 0, + cols: 0, + } + } + + fn resize(&mut self, rows: usize, cols: usize) { + self.cells.resize(rows * cols, 0.); + self.rows = rows; + self.cols = cols; + } + + fn get(&self, row: usize, col: usize) -> f64 { + if row >= self.rows { + panic!("row out of bounds") + } + + if col >= self.cols { + panic!("col out of bounds") + } + self.cells[col * self.rows + row] + } + + fn set(&mut self, row: usize, col: usize, value: f64) { + if row >= self.rows { + panic!("row out of bounds") + } + + if col >= self.cols { + panic!("col out of bounds") + } + + self.cells[col * self.rows + row] = value; + } +} + +impl Debug for Matrix { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + writeln!(f)?; + for i in 0..self.rows { + for j in 0..self.cols { + write!(f, "{:5}", self.get(i, j))?; + } + writeln!(f)?; + } + Ok(()) + } +} + +#[derive(Debug)] +pub enum Hunk { + Insert { text: String }, + Remove { len: usize }, + Keep { len: usize }, +} + +pub struct StreamingDiff { + old: Vec, + new: Vec, + scores: Matrix, + old_text_ix: usize, + new_text_ix: usize, + equal_runs: HashMap<(usize, usize), u32>, +} + +impl StreamingDiff { + const INSERTION_SCORE: f64 = -1.; + const DELETION_SCORE: f64 = -20.; + const EQUALITY_BASE: f64 = 1.8; + const MAX_EQUALITY_EXPONENT: i32 = 16; + + pub fn new(old: String) -> Self { + let old = old.chars().collect::>(); + let mut scores = Matrix::new(); + scores.resize(old.len() + 1, 1); + for i in 0..=old.len() { + scores.set(i, 0, i as f64 * Self::DELETION_SCORE); + } + Self { + old, + new: Vec::new(), + scores, + old_text_ix: 0, + new_text_ix: 0, + equal_runs: Default::default(), + } + } + + pub fn push_new(&mut self, text: &str) -> Vec { + self.new.extend(text.chars()); + self.scores.resize(self.old.len() + 1, self.new.len() + 1); + + for j in self.new_text_ix + 1..=self.new.len() { + self.scores.set(0, j, j as f64 * Self::INSERTION_SCORE); + for i in 1..=self.old.len() { + let insertion_score = self.scores.get(i, j - 1) + Self::INSERTION_SCORE; + let deletion_score = self.scores.get(i - 1, j) + Self::DELETION_SCORE; + let equality_score = if self.old[i - 1] == self.new[j - 1] { + let mut equal_run = self.equal_runs.get(&(i - 1, j - 1)).copied().unwrap_or(0); + equal_run += 1; + self.equal_runs.insert((i, j), equal_run); + + let exponent = cmp::min(equal_run as i32 / 4, Self::MAX_EQUALITY_EXPONENT); + self.scores.get(i - 1, j - 1) + Self::EQUALITY_BASE.powi(exponent) + } else { + f64::NEG_INFINITY + }; + + let score = insertion_score.max(deletion_score).max(equality_score); + self.scores.set(i, j, score); + } + } + + let mut max_score = f64::NEG_INFINITY; + let mut next_old_text_ix = self.old_text_ix; + let next_new_text_ix = self.new.len(); + for i in self.old_text_ix..=self.old.len() { + let score = self.scores.get(i, next_new_text_ix); + if score > max_score { + max_score = score; + next_old_text_ix = i; + } + } + + let hunks = self.backtrack(next_old_text_ix, next_new_text_ix); + self.old_text_ix = next_old_text_ix; + self.new_text_ix = next_new_text_ix; + hunks + } + + fn backtrack(&self, old_text_ix: usize, new_text_ix: usize) -> Vec { + let mut pending_insert: Option> = None; + let mut hunks = Vec::new(); + let mut i = old_text_ix; + let mut j = new_text_ix; + while (i, j) != (self.old_text_ix, self.new_text_ix) { + let insertion_score = if j > self.new_text_ix { + Some((i, j - 1)) + } else { + None + }; + let deletion_score = if i > self.old_text_ix { + Some((i - 1, j)) + } else { + None + }; + let equality_score = if i > self.old_text_ix && j > self.new_text_ix { + if self.old[i - 1] == self.new[j - 1] { + Some((i - 1, j - 1)) + } else { + None + } + } else { + None + }; + + let (prev_i, prev_j) = [insertion_score, deletion_score, equality_score] + .iter() + .max_by_key(|cell| cell.map(|(i, j)| OrderedFloat(self.scores.get(i, j)))) + .unwrap() + .unwrap(); + + if prev_i == i && prev_j == j - 1 { + if let Some(pending_insert) = pending_insert.as_mut() { + pending_insert.start = prev_j; + } else { + pending_insert = Some(prev_j..j); + } + } else { + if let Some(range) = pending_insert.take() { + hunks.push(Hunk::Insert { + text: self.new[range].iter().collect(), + }); + } + + let char_len = self.old[i - 1].len_utf8(); + if prev_i == i - 1 && prev_j == j { + if let Some(Hunk::Remove { len }) = hunks.last_mut() { + *len += char_len; + } else { + hunks.push(Hunk::Remove { len: char_len }) + } + } else { + if let Some(Hunk::Keep { len }) = hunks.last_mut() { + *len += char_len; + } else { + hunks.push(Hunk::Keep { len: char_len }) + } + } + } + + i = prev_i; + j = prev_j; + } + + if let Some(range) = pending_insert.take() { + hunks.push(Hunk::Insert { + text: self.new[range].iter().collect(), + }); + } + + hunks.reverse(); + hunks + } + + pub fn finish(self) -> Vec { + self.backtrack(self.old.len(), self.new.len()) + } +} + +#[cfg(test)] +mod tests { + use std::env; + + use super::*; + use rand::prelude::*; + + #[gpui::test(iterations = 100)] + fn test_random_diffs(mut rng: StdRng) { + let old_text_len = env::var("OLD_TEXT_LEN") + .map(|i| i.parse().expect("invalid `OLD_TEXT_LEN` variable")) + .unwrap_or(10); + let new_text_len = env::var("NEW_TEXT_LEN") + .map(|i| i.parse().expect("invalid `NEW_TEXT_LEN` variable")) + .unwrap_or(10); + + let old = util::RandomCharIter::new(&mut rng) + .take(old_text_len) + .collect::(); + log::info!("old text: {:?}", old); + + let mut diff = StreamingDiff::new(old.clone()); + let mut hunks = Vec::new(); + let mut new_len = 0; + let mut new = String::new(); + while new_len < new_text_len { + let new_chunk_len = rng.gen_range(1..=new_text_len - new_len); + let new_chunk = util::RandomCharIter::new(&mut rng) + .take(new_len) + .collect::(); + log::info!("new chunk: {:?}", new_chunk); + new_len += new_chunk_len; + new.push_str(&new_chunk); + let new_hunks = diff.push_new(&new_chunk); + log::info!("hunks: {:?}", new_hunks); + hunks.extend(new_hunks); + } + let final_hunks = diff.finish(); + log::info!("final hunks: {:?}", final_hunks); + hunks.extend(final_hunks); + + log::info!("new text: {:?}", new); + let mut old_ix = 0; + let mut new_ix = 0; + let mut patched = String::new(); + for hunk in hunks { + match hunk { + Hunk::Keep { len } => { + assert_eq!(&old[old_ix..old_ix + len], &new[new_ix..new_ix + len]); + patched.push_str(&old[old_ix..old_ix + len]); + old_ix += len; + new_ix += len; + } + Hunk::Remove { len } => { + old_ix += len; + } + Hunk::Insert { text } => { + assert_eq!(text, &new[new_ix..new_ix + text.len()]); + patched.push_str(&text); + new_ix += text.len(); + } + } + } + assert_eq!(patched, new); + } +} diff --git a/crates/collab_ui2/src/collab_panel.rs b/crates/collab_ui2/src/collab_panel.rs index c55bfa8cf5..89b83a7001 100644 --- a/crates/collab_ui2/src/collab_panel.rs +++ b/crates/collab_ui2/src/collab_panel.rs @@ -3344,10 +3344,6 @@ impl Panel for CollabPanel { Box::new(ToggleFocus) } - fn has_focus(&self, cx: &gpui::WindowContext) -> bool { - self.focus_handle.contains_focused(cx) - } - fn persistent_name() -> &'static str { "CollabPanel" } diff --git a/crates/project_panel2/src/project_panel.rs b/crates/project_panel2/src/project_panel.rs index ce039071cf..cb6515f529 100644 --- a/crates/project_panel2/src/project_panel.rs +++ b/crates/project_panel2/src/project_panel.rs @@ -55,7 +55,6 @@ pub struct ProjectPanel { clipboard_entry: Option, _dragged_entry_destination: Option>, _workspace: WeakView, - has_focus: bool, width: Option, pending_serialization: Task>, } @@ -172,7 +171,6 @@ impl ProjectPanel { let focus_handle = cx.focus_handle(); cx.on_focus(&focus_handle, Self::focus_in).detach(); - cx.on_blur(&focus_handle, Self::focus_out).detach(); cx.subscribe(&project, |this, project, event, cx| match event { project::Event::ActiveEntryChanged(Some(entry_id)) => { @@ -238,7 +236,6 @@ impl ProjectPanel { // context_menu: cx.add_view(|cx| ContextMenu::new(view_id, cx)), _dragged_entry_destination: None, _workspace: workspace.weak_handle(), - has_focus: false, width: None, pending_serialization: Task::ready(None), }; @@ -356,16 +353,11 @@ impl ProjectPanel { } fn focus_in(&mut self, cx: &mut ViewContext) { - if !self.has_focus { - self.has_focus = true; + if !self.focus_handle.contains_focused(cx) { cx.emit(Event::Focus); } } - fn focus_out(&mut self, _: &mut ViewContext) { - self.has_focus = false; - } - fn deploy_context_menu( &mut self, position: Point, @@ -1554,10 +1546,6 @@ impl Panel for ProjectPanel { Box::new(ToggleFocus) } - fn has_focus(&self, _: &WindowContext) -> bool { - self.has_focus - } - fn persistent_name() -> &'static str { "Project Panel" } diff --git a/crates/terminal_view2/src/terminal_panel.rs b/crates/terminal_view2/src/terminal_panel.rs index b6582b07b1..caf339a8c6 100644 --- a/crates/terminal_view2/src/terminal_panel.rs +++ b/crates/terminal_view2/src/terminal_panel.rs @@ -415,10 +415,6 @@ impl Panel for TerminalPanel { } } - fn has_focus(&self, cx: &WindowContext) -> bool { - self.pane.read(cx).has_focus(cx) - } - fn persistent_name() -> &'static str { "TerminalPanel" } diff --git a/crates/workspace2/src/dock.rs b/crates/workspace2/src/dock.rs index 437e7c0192..abcf5c49bc 100644 --- a/crates/workspace2/src/dock.rs +++ b/crates/workspace2/src/dock.rs @@ -26,6 +26,7 @@ pub trait Panel: FocusableView + EventEmitter { fn set_position(&mut self, position: DockPosition, cx: &mut ViewContext); fn size(&self, cx: &WindowContext) -> f32; fn set_size(&mut self, size: Option, cx: &mut ViewContext); + // todo!("We should have a icon tooltip method, rather than using persistant_name") fn icon(&self, cx: &WindowContext) -> Option; fn toggle_action(&self) -> Box; fn icon_label(&self, _: &WindowContext) -> Option { @@ -36,7 +37,6 @@ pub trait Panel: FocusableView + EventEmitter { } fn set_zoomed(&mut self, _zoomed: bool, _cx: &mut ViewContext) {} fn set_active(&mut self, _active: bool, _cx: &mut ViewContext) {} - fn has_focus(&self, cx: &WindowContext) -> bool; } pub trait PanelHandle: Send + Sync { @@ -53,7 +53,6 @@ pub trait PanelHandle: Send + Sync { fn icon(&self, cx: &WindowContext) -> Option; fn toggle_action(&self, cx: &WindowContext) -> Box; fn icon_label(&self, cx: &WindowContext) -> Option; - fn has_focus(&self, cx: &WindowContext) -> bool; fn focus_handle(&self, cx: &AppContext) -> FocusHandle; fn to_any(&self) -> AnyView; } @@ -114,10 +113,6 @@ where self.read(cx).icon_label(cx) } - fn has_focus(&self, cx: &WindowContext) -> bool { - self.read(cx).has_focus(cx) - } - fn to_any(&self) -> AnyView { self.clone().into() } @@ -319,7 +314,7 @@ impl Dock { } PanelEvent::ZoomIn => { this.set_panel_zoomed(&panel.to_any(), true, cx); - if !panel.has_focus(cx) { + if !panel.focus_handle(cx).contains_focused(cx) { cx.focus_view(&panel); } workspace @@ -760,7 +755,7 @@ pub mod test { pub position: DockPosition, pub zoomed: bool, pub active: bool, - pub has_focus: bool, + pub focus_handle: FocusHandle, pub size: f32, } actions!(ToggleTestPanel); @@ -768,12 +763,12 @@ pub mod test { impl EventEmitter for TestPanel {} impl TestPanel { - pub fn new(position: DockPosition) -> Self { + pub fn new(position: DockPosition, cx: &mut WindowContext) -> Self { Self { position, zoomed: false, active: false, - has_focus: false, + focus_handle: cx.focus_handle(), size: 300., } } @@ -832,15 +827,11 @@ pub mod test { fn set_active(&mut self, active: bool, _cx: &mut ViewContext) { self.active = active; } - - fn has_focus(&self, _cx: &WindowContext) -> bool { - self.has_focus - } } impl FocusableView for TestPanel { - fn focus_handle(&self, cx: &AppContext) -> FocusHandle { - unimplemented!() + fn focus_handle(&self, _cx: &AppContext) -> FocusHandle { + self.focus_handle.clone() } } } diff --git a/crates/workspace2/src/workspace2.rs b/crates/workspace2/src/workspace2.rs index 5741fa4a94..a916ec733c 100644 --- a/crates/workspace2/src/workspace2.rs +++ b/crates/workspace2/src/workspace2.rs @@ -65,7 +65,7 @@ use std::{ time::Duration, }; use theme::{ActiveTheme, ThemeSettings}; -pub use toolbar::{ToolbarItemEvent, ToolbarItemLocation, ToolbarItemView}; +pub use toolbar::{Toolbar, ToolbarItemEvent, ToolbarItemLocation, ToolbarItemView}; pub use ui; use util::ResultExt; use uuid::Uuid; @@ -1542,7 +1542,7 @@ impl Workspace { if let Some(active_panel) = dock.active_panel() { if was_visible { - if active_panel.has_focus(cx) { + if active_panel.focus_handle(cx).contains_focused(cx) { focus_center = true; } } else { @@ -1589,7 +1589,9 @@ impl Workspace { /// Focus the panel of the given type if it isn't already focused. If it is /// already focused, then transfer focus back to the workspace center. pub fn toggle_panel_focus(&mut self, cx: &mut ViewContext) { - self.focus_or_unfocus_panel::(cx, |panel, cx| !panel.has_focus(cx)); + self.focus_or_unfocus_panel::(cx, |panel, cx| { + !panel.focus_handle(cx).contains_focused(cx) + }); } /// Focus or unfocus the given panel type, depending on the given callback. @@ -1681,7 +1683,7 @@ impl Workspace { if Some(dock.position()) != dock_to_reveal { if let Some(panel) = dock.active_panel() { if panel.is_zoomed(cx) { - focus_center |= panel.has_focus(cx); + focus_center |= panel.focus_handle(cx).contains_focused(cx); dock.set_open(false, cx); } }