1536 lines · 56.4 KB · Rust
RawDownload
use crate::{    prompts::generate_content_prompt, AssistantPanel, CompletionProvider, Hunk,    LanguageModelRequest, LanguageModelRequestMessage, Role, StreamingDiff,};use anyhow::Result;use client::telemetry::Telemetry;use collections::{hash_map, HashMap, HashSet, VecDeque};use editor::{    actions::{MoveDown, MoveUp},    display_map::{        BlockContext, BlockDisposition, BlockId, BlockProperties, BlockStyle, RenderBlock,    },    scroll::{Autoscroll, AutoscrollStrategy},    Anchor, Editor, EditorElement, EditorEvent, EditorStyle, GutterDimensions, MultiBuffer,    MultiBufferSnapshot, ToOffset, ToPoint,};use futures::{channel::mpsc, SinkExt, Stream, StreamExt};use gpui::{    AnyWindowHandle, AppContext, EventEmitter, FocusHandle, FocusableView, FontStyle, FontWeight,    Global, HighlightStyle, Model, ModelContext, Subscription, Task, TextStyle, UpdateGlobal, View,    ViewContext, WeakView, WhiteSpace, WindowContext,};use language::{Point, TransactionId};use multi_buffer::MultiBufferRow;use parking_lot::Mutex;use rope::Rope;use settings::Settings;use std::{cmp, future, ops::Range, sync::Arc, time::Instant};use theme::ThemeSettings;use ui::{prelude::*, Tooltip};use workspace::{notifications::NotificationId, Toast, Workspace}; pub fn init(telemetry: Arc<Telemetry>, cx: &mut AppContext) {    cx.set_global(InlineAssistant::new(telemetry));} const PROMPT_HISTORY_MAX_LEN: usize = 20; pub struct InlineAssistant {    next_assist_id: InlineAssistId,    pending_assists: HashMap<InlineAssistId, PendingInlineAssist>,    pending_assist_ids_by_editor: HashMap<WeakView<Editor>, EditorPendingAssists>,    prompt_history: VecDeque<String>,    telemetry: Option<Arc<Telemetry>>,} struct EditorPendingAssists {    window: AnyWindowHandle,    assist_ids: Vec<InlineAssistId>,} impl Global for InlineAssistant {} impl InlineAssistant {    pub fn new(telemetry: Arc<Telemetry>) -> Self {        Self {            next_assist_id: InlineAssistId::default(),            pending_assists: HashMap::default(),            pending_assist_ids_by_editor: HashMap::default(),            prompt_history: VecDeque::default(),            telemetry: Some(telemetry),        }    }     pub fn assist(        &mut self,        editor: &View<Editor>,        workspace: Option<WeakView<Workspace>>,        include_context: bool,        cx: &mut WindowContext,    ) {        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(MultiBufferRow(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 = self.next_assist_id.post_inc();        let codegen = cx.new_model(|cx| {            Codegen::new(                editor.read(cx).buffer().clone(),                codegen_kind,                self.telemetry.clone(),                cx,            )        });         let gutter_dimensions = Arc::new(Mutex::new(GutterDimensions::default()));        let inline_assist_editor = cx.new_view(|cx| {            InlineAssistEditor::new(                inline_assist_id,                gutter_dimensions.clone(),                self.prompt_history.clone(),                codegen.clone(),                cx,            )        });        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::Sticky,                    position: snapshot.anchor_before(Point::new(point_selection.head().row, 0)),                    height: inline_assist_editor.read(cx).height_in_lines,                    render: build_inline_assist_editor_renderer(                        &inline_assist_editor,                        gutter_dimensions,                    ),                    disposition: if selection.reversed {                        BlockDisposition::Above                    } else {                        BlockDisposition::Below                    },                }],                Some(Autoscroll::Strategy(AutoscrollStrategy::Newest)),                cx,            )[0]        });         self.pending_assists.insert(            inline_assist_id,            PendingInlineAssist {                include_context,                editor: editor.downgrade(),                inline_assist_editor: Some((block_id, inline_assist_editor.clone())),                codegen: codegen.clone(),                workspace,                _subscriptions: vec![                    cx.subscribe(&inline_assist_editor, |inline_assist_editor, event, cx| {                        InlineAssistant::update_global(cx, |this, cx| {                            this.handle_inline_assistant_event(inline_assist_editor, event, cx)                        })                    }),                    cx.subscribe(editor, {                        let inline_assist_editor = inline_assist_editor.downgrade();                        move |editor, event, cx| {                            if let Some(inline_assist_editor) = inline_assist_editor.upgrade() {                                if let EditorEvent::SelectionsChanged { local } = event {                                    if *local                                        && inline_assist_editor                                            .focus_handle(cx)                                            .contains_focused(cx)                                    {                                        cx.focus_view(&editor);                                    }                                }                            }                        }                    }),                    cx.observe(&codegen, {                        let editor = editor.downgrade();                        move |_, cx| {                            if let Some(editor) = editor.upgrade() {                                InlineAssistant::update_global(cx, |this, cx| {                                    this.update_highlights_for_editor(&editor, cx);                                })                            }                        }                    }),                    cx.subscribe(&codegen, move |codegen, event, cx| {                        InlineAssistant::update_global(cx, |this, cx| match event {                            CodegenEvent::Undone => {                                this.finish_inline_assist(inline_assist_id, false, cx)                            }                            CodegenEvent::Finished => {                                let pending_assist = if let Some(pending_assist) =                                    this.pending_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_assist_editor.is_none() {                                        if let Some(workspace) = pending_assist                                            .workspace                                            .as_ref()                                            .and_then(|workspace| workspace.upgrade())                                        {                                            workspace.update(cx, |workspace, cx| {                                                struct InlineAssistantError;                                                 let id = NotificationId::identified::<                                                    InlineAssistantError,                                                >(                                                    inline_assist_id.0                                                );                                                 workspace.show_toast(Toast::new(id, error), cx);                                            })                                        }                                         this.finish_inline_assist(inline_assist_id, false, cx);                                    }                                } else {                                    this.finish_inline_assist(inline_assist_id, false, cx);                                }                            }                        })                    }),                ],            },        );         self.pending_assist_ids_by_editor            .entry(editor.downgrade())            .or_insert_with(|| EditorPendingAssists {                window: cx.window_handle(),                assist_ids: Vec::new(),            })            .assist_ids            .push(inline_assist_id);        self.update_highlights_for_editor(editor, cx);    }     fn handle_inline_assistant_event(        &mut self,        inline_assist_editor: View<InlineAssistEditor>,        event: &InlineAssistEditorEvent,        cx: &mut WindowContext,    ) {        let assist_id = inline_assist_editor.read(cx).id;        match event {            InlineAssistEditorEvent::Confirmed { prompt } => {                self.confirm_inline_assist(assist_id, prompt, cx);            }            InlineAssistEditorEvent::Canceled => {                self.finish_inline_assist(assist_id, true, cx);            }            InlineAssistEditorEvent::Dismissed => {                self.hide_inline_assist(assist_id, cx);            }            InlineAssistEditorEvent::Resized { height_in_lines } => {                self.resize_inline_assist(assist_id, *height_in_lines, cx);            }        }    }     pub fn cancel_last_inline_assist(&mut self, cx: &mut WindowContext) -> bool {        for (editor, pending_assists) in &self.pending_assist_ids_by_editor {            if pending_assists.window == cx.window_handle() {                if let Some(editor) = editor.upgrade() {                    if editor.read(cx).is_focused(cx) {                        if let Some(assist_id) = pending_assists.assist_ids.last().copied() {                            self.finish_inline_assist(assist_id, true, cx);                            return true;                        }                    }                }            }        }        false    }     fn finish_inline_assist(        &mut self,        assist_id: InlineAssistId,        undo: bool,        cx: &mut WindowContext,    ) {        self.hide_inline_assist(assist_id, cx);         if let Some(pending_assist) = self.pending_assists.remove(&assist_id) {            if let hash_map::Entry::Occupied(mut entry) = self                .pending_assist_ids_by_editor                .entry(pending_assist.editor.clone())            {                entry.get_mut().assist_ids.retain(|id| *id != assist_id);                if entry.get().assist_ids.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: InlineAssistId, cx: &mut WindowContext) {        if let Some(pending_assist) = self.pending_assists.get_mut(&assist_id) {            if let Some(editor) = pending_assist.editor.upgrade() {                if let Some((block_id, inline_assist_editor)) =                    pending_assist.inline_assist_editor.take()                {                    editor.update(cx, |editor, cx| {                        editor.remove_blocks(HashSet::from_iter([block_id]), None, cx);                        if inline_assist_editor.focus_handle(cx).contains_focused(cx) {                            editor.focus(cx);                        }                    });                }            }        }    }     fn resize_inline_assist(        &mut self,        assist_id: InlineAssistId,        height_in_lines: u8,        cx: &mut WindowContext,    ) {        if let Some(pending_assist) = self.pending_assists.get_mut(&assist_id) {            if let Some(editor) = pending_assist.editor.upgrade() {                if let Some((block_id, inline_assist_editor)) =                    pending_assist.inline_assist_editor.as_ref()                {                    let gutter_dimensions = inline_assist_editor.read(cx).gutter_dimensions.clone();                    let mut new_blocks = HashMap::default();                    new_blocks.insert(                        *block_id,                        (                            Some(height_in_lines),                            build_inline_assist_editor_renderer(                                inline_assist_editor,                                gutter_dimensions,                            ),                        ),                    );                    editor.update(cx, |editor, cx| {                        editor                            .display_map                            .update(cx, |map, cx| map.replace_blocks(new_blocks, cx))                    });                }            }        }    }     fn confirm_inline_assist(        &mut self,        assist_id: InlineAssistId,        user_prompt: &str,        cx: &mut WindowContext,    ) {        let pending_assist = if let Some(pending_assist) = self.pending_assists.get_mut(&assist_id)        {            pending_assist        } else {            return;        };         let context = if pending_assist.include_context {            pending_assist.workspace.as_ref().and_then(|workspace| {                let workspace = workspace.upgrade()?.read(cx);                let assistant_panel = workspace.panel::<AssistantPanel>(cx)?;                assistant_panel.read(cx).active_context(cx)            })        } else {            None        };         let editor = if let Some(editor) = pending_assist.editor.upgrade() {            editor        } else {            return;        };         let project_name = pending_assist.workspace.as_ref().and_then(|workspace| {            let workspace = workspace.upgrade()?;            Some(                workspace                    .read(cx)                    .project()                    .read(cx)                    .worktree_root_names(cx)                    .collect::<Vec<&str>>()                    .join("/"),            )        });         self.prompt_history.retain(|prompt| prompt != user_prompt);        self.prompt_history.push_back(user_prompt.into());        if self.prompt_history.len() > PROMPT_HISTORY_MAX_LEN {            self.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(assist_id, false, cx);                return;            }        } else {            self.finish_inline_assist(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.as_ref() == "Markdown" {                1.0            } else {                0.5            }        } else {            1.0        };         let user_prompt = user_prompt.to_string();         let prompt = cx.background_executor().spawn(async move {            let language_name = language_name.as_deref();            generate_content_prompt(user_prompt, language_name, buffer, range, project_name)        });         let mut messages = Vec::new();        if let Some(context) = context {            let request = context.read(cx).to_completion_request(cx);            messages = request.messages;        }        let model = CompletionProvider::global(cx).model();         cx.spawn(|mut cx| async move {            let prompt = prompt.await?;             messages.push(LanguageModelRequestMessage {                role: Role::User,                content: prompt,            });             let request = LanguageModelRequest {                model,                messages,                stop: vec!["|END|>".to_string()],                temperature,            };             codegen.update(&mut cx, |codegen, cx| codegen.start(request, cx))?;            anyhow::Ok(())        })        .detach_and_log_err(cx);    }     fn update_highlights_for_editor(&self, editor: &View<Editor>, cx: &mut WindowContext) {        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_assist_ids_by_editor            .get(&editor.downgrade())            .map_or(&empty_inline_assist_ids, |pending_assists| {                &pending_assists.assist_ids            });

Lines 1–500 of 1536.