1518 lines · 55.7 KB · Rust
RawDownload
use crate::inline_prompt_editor::CodegenStatus;use agent::{    ContextStore,    context::{ContextLoadResult, load_context},};use agent_settings::AgentSettings;use anyhow::{Context as _, Result};use client::telemetry::Telemetry;use cloud_llm_client::CompletionIntent;use collections::HashSet;use editor::{Anchor, AnchorRangeExt, MultiBuffer, MultiBufferSnapshot, ToOffset as _, ToPoint};use futures::{    SinkExt, Stream, StreamExt, TryStreamExt as _, channel::mpsc, future::LocalBoxFuture, join,};use gpui::{App, AppContext as _, Context, Entity, EventEmitter, Subscription, Task, WeakEntity};use language::{Buffer, IndentKind, Point, TransactionId, line_diff};use language_model::{    LanguageModel, LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage,    LanguageModelTextStream, Role, report_assistant_event,};use multi_buffer::MultiBufferRow;use parking_lot::Mutex;use project::Project;use prompt_store::{PromptBuilder, PromptStore};use rope::Rope;use smol::future::FutureExt;use std::{    cmp,    future::Future,    iter,    ops::{Range, RangeInclusive},    pin::Pin,    sync::Arc,    task::{self, Poll},    time::Instant,};use streaming_diff::{CharOperation, LineDiff, LineOperation, StreamingDiff};use telemetry_events::{AssistantEventData, AssistantKind, AssistantPhase}; pub struct BufferCodegen {    alternatives: Vec<Entity<CodegenAlternative>>,    pub active_alternative: usize,    seen_alternatives: HashSet<usize>,    subscriptions: Vec<Subscription>,    buffer: Entity<MultiBuffer>,    range: Range<Anchor>,    initial_transaction_id: Option<TransactionId>,    context_store: Entity<ContextStore>,    project: WeakEntity<Project>,    prompt_store: Option<Entity<PromptStore>>,    telemetry: Arc<Telemetry>,    builder: Arc<PromptBuilder>,    pub is_insertion: bool,} impl BufferCodegen {    pub fn new(        buffer: Entity<MultiBuffer>,        range: Range<Anchor>,        initial_transaction_id: Option<TransactionId>,        context_store: Entity<ContextStore>,        project: WeakEntity<Project>,        prompt_store: Option<Entity<PromptStore>>,        telemetry: Arc<Telemetry>,        builder: Arc<PromptBuilder>,        cx: &mut Context<Self>,    ) -> Self {        let codegen = cx.new(|cx| {            CodegenAlternative::new(                buffer.clone(),                range.clone(),                false,                Some(context_store.clone()),                project.clone(),                prompt_store.clone(),                Some(telemetry.clone()),                builder.clone(),                cx,            )        });        let mut this = Self {            is_insertion: range.to_offset(&buffer.read(cx).snapshot(cx)).is_empty(),            alternatives: vec![codegen],            active_alternative: 0,            seen_alternatives: HashSet::default(),            subscriptions: Vec::new(),            buffer,            range,            initial_transaction_id,            context_store,            project,            prompt_store,            telemetry,            builder,        };        this.activate(0, cx);        this    }     fn subscribe_to_alternative(&mut self, cx: &mut Context<Self>) {        let codegen = self.active_alternative().clone();        self.subscriptions.clear();        self.subscriptions            .push(cx.observe(&codegen, |_, _, cx| cx.notify()));        self.subscriptions            .push(cx.subscribe(&codegen, |_, _, event, cx| cx.emit(*event)));    }     pub fn active_alternative(&self) -> &Entity<CodegenAlternative> {        &self.alternatives[self.active_alternative]    }     pub fn status<'a>(&self, cx: &'a App) -> &'a CodegenStatus {        &self.active_alternative().read(cx).status    }     pub fn alternative_count(&self, cx: &App) -> usize {        LanguageModelRegistry::read_global(cx)            .inline_alternative_models()            .len()            + 1    }     pub fn cycle_prev(&mut self, cx: &mut Context<Self>) {        let next_active_ix = if self.active_alternative == 0 {            self.alternatives.len() - 1        } else {            self.active_alternative - 1        };        self.activate(next_active_ix, cx);    }     pub fn cycle_next(&mut self, cx: &mut Context<Self>) {        let next_active_ix = (self.active_alternative + 1) % self.alternatives.len();        self.activate(next_active_ix, cx);    }     fn activate(&mut self, index: usize, cx: &mut Context<Self>) {        self.active_alternative()            .update(cx, |codegen, cx| codegen.set_active(false, cx));        self.seen_alternatives.insert(index);        self.active_alternative = index;        self.active_alternative()            .update(cx, |codegen, cx| codegen.set_active(true, cx));        self.subscribe_to_alternative(cx);        cx.notify();    }     pub fn start(        &mut self,        primary_model: Arc<dyn LanguageModel>,        user_prompt: String,        cx: &mut Context<Self>,    ) -> Result<()> {        let alternative_models = LanguageModelRegistry::read_global(cx)            .inline_alternative_models()            .to_vec();         self.active_alternative()            .update(cx, |alternative, cx| alternative.undo(cx));        self.activate(0, cx);        self.alternatives.truncate(1);         for _ in 0..alternative_models.len() {            self.alternatives.push(cx.new(|cx| {                CodegenAlternative::new(                    self.buffer.clone(),                    self.range.clone(),                    false,                    Some(self.context_store.clone()),                    self.project.clone(),                    self.prompt_store.clone(),                    Some(self.telemetry.clone()),                    self.builder.clone(),                    cx,                )            }));        }         for (model, alternative) in iter::once(primary_model)            .chain(alternative_models)            .zip(&self.alternatives)        {            alternative.update(cx, |alternative, cx| {                alternative.start(user_prompt.clone(), model.clone(), cx)            })?;        }         Ok(())    }     pub fn stop(&mut self, cx: &mut Context<Self>) {        for codegen in &self.alternatives {            codegen.update(cx, |codegen, cx| codegen.stop(cx));        }    }     pub fn undo(&mut self, cx: &mut Context<Self>) {        self.active_alternative()            .update(cx, |codegen, cx| codegen.undo(cx));         self.buffer.update(cx, |buffer, cx| {            if let Some(transaction_id) = self.initial_transaction_id.take() {                buffer.undo_transaction(transaction_id, cx);                buffer.refresh_preview(cx);            }        });    }     pub fn buffer(&self, cx: &App) -> Entity<MultiBuffer> {        self.active_alternative().read(cx).buffer.clone()    }     pub fn old_buffer(&self, cx: &App) -> Entity<Buffer> {        self.active_alternative().read(cx).old_buffer.clone()    }     pub fn snapshot(&self, cx: &App) -> MultiBufferSnapshot {        self.active_alternative().read(cx).snapshot.clone()    }     pub fn edit_position(&self, cx: &App) -> Option<Anchor> {        self.active_alternative().read(cx).edit_position    }     pub fn diff<'a>(&self, cx: &'a App) -> &'a Diff {        &self.active_alternative().read(cx).diff    }     pub fn last_equal_ranges<'a>(&self, cx: &'a App) -> &'a [Range<Anchor>] {        self.active_alternative().read(cx).last_equal_ranges()    }} impl EventEmitter<CodegenEvent> for BufferCodegen {} pub struct CodegenAlternative {    buffer: Entity<MultiBuffer>,    old_buffer: Entity<Buffer>,    snapshot: MultiBufferSnapshot,    edit_position: Option<Anchor>,    range: Range<Anchor>,    last_equal_ranges: Vec<Range<Anchor>>,    transformation_transaction_id: Option<TransactionId>,    status: CodegenStatus,    generation: Task<()>,    diff: Diff,    context_store: Option<Entity<ContextStore>>,    project: WeakEntity<Project>,    prompt_store: Option<Entity<PromptStore>>,    telemetry: Option<Arc<Telemetry>>,    _subscription: gpui::Subscription,    builder: Arc<PromptBuilder>,    active: bool,    edits: Vec<(Range<Anchor>, String)>,    line_operations: Vec<LineOperation>,    elapsed_time: Option<f64>,    completion: Option<String>,    pub message_id: Option<String>,} impl EventEmitter<CodegenEvent> for CodegenAlternative {} impl CodegenAlternative {    pub fn new(        buffer: Entity<MultiBuffer>,        range: Range<Anchor>,        active: bool,        context_store: Option<Entity<ContextStore>>,        project: WeakEntity<Project>,        prompt_store: Option<Entity<PromptStore>>,        telemetry: Option<Arc<Telemetry>>,        builder: Arc<PromptBuilder>,        cx: &mut Context<Self>,    ) -> Self {        let snapshot = buffer.read(cx).snapshot(cx);         let (old_buffer, _, _) = snapshot            .range_to_buffer_ranges(range.clone())            .pop()            .unwrap();        let old_buffer = cx.new(|cx| {            let text = old_buffer.as_rope().clone();            let line_ending = old_buffer.line_ending();            let language = old_buffer.language().cloned();            let language_registry = buffer                .read(cx)                .buffer(old_buffer.remote_id())                .unwrap()                .read(cx)                .language_registry();             let mut buffer = Buffer::local_normalized(text, line_ending, cx);            buffer.set_language(language, cx);            if let Some(language_registry) = language_registry {                buffer.set_language_registry(language_registry)            }            buffer        });         Self {            buffer: buffer.clone(),            old_buffer,            edit_position: None,            message_id: None,            snapshot,            last_equal_ranges: Default::default(),            transformation_transaction_id: None,            status: CodegenStatus::Idle,            generation: Task::ready(()),            diff: Diff::default(),            context_store,            project,            prompt_store,            telemetry,            _subscription: cx.subscribe(&buffer, Self::handle_buffer_event),            builder,            active,            edits: Vec::new(),            line_operations: Vec::new(),            range,            elapsed_time: None,            completion: None,        }    }     pub fn set_active(&mut self, active: bool, cx: &mut Context<Self>) {        if active != self.active {            self.active = active;             if self.active {                let edits = self.edits.clone();                self.apply_edits(edits, cx);                if matches!(self.status, CodegenStatus::Pending) {                    let line_operations = self.line_operations.clone();                    self.reapply_line_based_diff(line_operations, cx);                } else {                    self.reapply_batch_diff(cx).detach();                }            } else if let Some(transaction_id) = self.transformation_transaction_id.take() {                self.buffer.update(cx, |buffer, cx| {                    buffer.undo_transaction(transaction_id, cx);                    buffer.forget_transaction(transaction_id, cx);                });            }        }    }     fn handle_buffer_event(        &mut self,        _buffer: Entity<MultiBuffer>,        event: &multi_buffer::Event,        cx: &mut Context<Self>,    ) {        if let multi_buffer::Event::TransactionUndone { transaction_id } = event            && self.transformation_transaction_id == Some(*transaction_id)        {            self.transformation_transaction_id = None;            self.generation = Task::ready(());            cx.emit(CodegenEvent::Undone);        }    }     pub fn last_equal_ranges(&self) -> &[Range<Anchor>] {        &self.last_equal_ranges    }     pub fn start(        &mut self,        user_prompt: String,        model: Arc<dyn LanguageModel>,        cx: &mut Context<Self>,    ) -> Result<()> {        if let Some(transformation_transaction_id) = self.transformation_transaction_id.take() {            self.buffer.update(cx, |buffer, cx| {                buffer.undo_transaction(transformation_transaction_id, cx);            });        }         self.edit_position = Some(self.range.start.bias_right(&self.snapshot));         let api_key = model.api_key(cx);        let telemetry_id = model.telemetry_id();        let provider_id = model.provider_id();        let stream: LocalBoxFuture<Result<LanguageModelTextStream>> =            if user_prompt.trim().to_lowercase() == "delete" {                async { Ok(LanguageModelTextStream::default()) }.boxed_local()            } else {                let request = self.build_request(&model, user_prompt, cx)?;                cx.spawn(async move |_, cx| {                    Ok(model.stream_completion_text(request.await, cx).await?)                })                .boxed_local()            };        self.handle_stream(telemetry_id, provider_id.to_string(), api_key, stream, cx);        Ok(())    }     fn build_request(        &self,        model: &Arc<dyn LanguageModel>,        user_prompt: String,        cx: &mut App,    ) -> Result<Task<LanguageModelRequest>> {        let buffer = self.buffer.read(cx).snapshot(cx);        let language = buffer.language_at(self.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        };         let language_name = language_name.as_ref();        let start = buffer.point_to_buffer_offset(self.range.start);        let end = buffer.point_to_buffer_offset(self.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 {                anyhow::bail!("invalid transformation range");            }        } else {            anyhow::bail!("invalid transformation range");        };         let prompt = self            .builder            .generate_inline_transformation_prompt(user_prompt, language_name, buffer, range)            .context("generating content prompt")?;         let context_task = self.context_store.as_ref().map(|context_store| {            if let Some(project) = self.project.upgrade() {                let context = context_store                    .read(cx)                    .context()                    .cloned()                    .collect::<Vec<_>>();                load_context(context, &project, &self.prompt_store, cx)            } else {                Task::ready(ContextLoadResult::default())            }        });         let temperature = AgentSettings::temperature_for_model(model, cx);         Ok(cx.spawn(async move |_cx| {            let mut request_message = LanguageModelRequestMessage {                role: Role::User,                content: Vec::new(),                cache: false,            };             if let Some(context_task) = context_task {                context_task                    .await                    .loaded_context                    .add_to_request_message(&mut request_message);            }             request_message.content.push(prompt.into());             LanguageModelRequest {                thread_id: None,                prompt_id: None,                intent: Some(CompletionIntent::InlineAssist),                mode: None,                tools: Vec::new(),                tool_choice: None,                stop: Vec::new(),                temperature,                messages: vec![request_message],                thinking_allowed: false,            }        }))    }     pub fn handle_stream(        &mut self,        model_telemetry_id: String,        model_provider_id: String,        model_api_key: Option<String>,        stream: impl 'static + Future<Output = Result<LanguageModelTextStream>>,        cx: &mut Context<Self>,    ) {        let start_time = Instant::now();        let snapshot = self.snapshot.clone();        let selected_text = snapshot            .text_for_range(self.range.start..self.range.end)            .collect::<Rope>();         let selection_start = self.range.start.to_point(&snapshot);         // Start with the indentation of the first line in the selection        let mut suggested_line_indent = snapshot

Lines 1–500 of 1518.