use crate::context::attach_context_to_message;use crate::context_store::ContextStore;use crate::inline_prompt_editor::CodegenStatus;use crate::{    prompts::PromptBuilder,    streaming_diff::{CharOperation, LineDiff, LineOperation, StreamingDiff},};use anyhow::{Context as _, Result};use client::telemetry::Telemetry;use collections::HashSet;use editor::{Anchor, AnchorRangeExt, MultiBuffer, MultiBufferSnapshot, ToOffset as _, ToPoint};use futures::{channel::mpsc, future::LocalBoxFuture, join, SinkExt, Stream, StreamExt};use gpui::{AppContext, Context as _, EventEmitter, Model, ModelContext, Subscription, Task};use language::{Buffer, IndentKind, Point, TransactionId};use language_model::{    LanguageModel, LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage,    LanguageModelTextStream, Role,};use language_models::report_assistant_event;use multi_buffer::MultiBufferRow;use parking_lot::Mutex;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 telemetry_events::{AssistantEvent, AssistantKind, AssistantPhase}; pub struct BufferCodegen {    alternatives: Vec<Model<CodegenAlternative>>,    pub active_alternative: usize,    seen_alternatives: HashSet<usize>,    subscriptions: Vec<Subscription>,    buffer: Model<MultiBuffer>,    range: Range<Anchor>,    initial_transaction_id: Option<TransactionId>,    context_store: Model<ContextStore>,    telemetry: Arc<Telemetry>,    builder: Arc<PromptBuilder>,    pub is_insertion: bool,} impl BufferCodegen {    pub fn new(        buffer: Model<MultiBuffer>,        range: Range<Anchor>,        initial_transaction_id: Option<TransactionId>,        context_store: Model<ContextStore>,        telemetry: Arc<Telemetry>,        builder: Arc<PromptBuilder>,        cx: &mut ModelContext<Self>,    ) -> Self {        let codegen = cx.new_model(|cx| {            CodegenAlternative::new(                buffer.clone(),                range.clone(),                false,                Some(context_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,            telemetry,            builder,        };        this.activate(0, cx);        this    }     fn subscribe_to_alternative(&mut self, cx: &mut ModelContext<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) -> &Model<CodegenAlternative> {        &self.alternatives[self.active_alternative]    }     pub fn status<'a>(&self, cx: &'a AppContext) -> &'a CodegenStatus {        &self.active_alternative().read(cx).status    }     pub fn alternative_count(&self, cx: &AppContext) -> usize {        LanguageModelRegistry::read_global(cx)            .inline_alternative_models()            .len()            + 1    }     pub fn cycle_prev(&mut self, cx: &mut ModelContext<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 ModelContext<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 ModelContext<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, user_prompt: String, cx: &mut ModelContext<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_model(|cx| {                CodegenAlternative::new(                    self.buffer.clone(),                    self.range.clone(),                    false,                    Some(self.context_store.clone()),                    Some(self.telemetry.clone()),                    self.builder.clone(),                    cx,                )            }));        }         let primary_model = LanguageModelRegistry::read_global(cx)            .active_model()            .context("no active model")?;         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 ModelContext<Self>) {        for codegen in &self.alternatives {            codegen.update(cx, |codegen, cx| codegen.stop(cx));        }    }     pub fn undo(&mut self, cx: &mut ModelContext<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: &AppContext) -> Model<MultiBuffer> {        self.active_alternative().read(cx).buffer.clone()    }     pub fn old_buffer(&self, cx: &AppContext) -> Model<Buffer> {        self.active_alternative().read(cx).old_buffer.clone()    }     pub fn snapshot(&self, cx: &AppContext) -> MultiBufferSnapshot {        self.active_alternative().read(cx).snapshot.clone()    }     pub fn edit_position(&self, cx: &AppContext) -> Option<Anchor> {        self.active_alternative().read(cx).edit_position    }     pub fn diff<'a>(&self, cx: &'a AppContext) -> &'a Diff {        &self.active_alternative().read(cx).diff    }     pub fn last_equal_ranges<'a>(&self, cx: &'a AppContext) -> &'a [Range<Anchor>] {        self.active_alternative().read(cx).last_equal_ranges()    }} impl EventEmitter<CodegenEvent> for BufferCodegen {} pub struct CodegenAlternative {    buffer: Model<MultiBuffer>,    old_buffer: Model<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<Model<ContextStore>>,    telemetry: Option<Arc<Telemetry>>,    _subscription: gpui::Subscription,    builder: Arc<PromptBuilder>,    active: bool,    edits: Vec<(Range<Anchor>, String)>,    line_operations: Vec<LineOperation>,    request: Option<LanguageModelRequest>,    elapsed_time: Option<f64>,    completion: Option<String>,    pub message_id: Option<String>,} impl EventEmitter<CodegenEvent> for CodegenAlternative {} impl CodegenAlternative {    pub fn new(        buffer: Model<MultiBuffer>,        range: Range<Anchor>,        active: bool,        context_store: Option<Model<ContextStore>>,        telemetry: Option<Arc<Telemetry>>,        builder: Arc<PromptBuilder>,        cx: &mut ModelContext<Self>,    ) -> Self {        let snapshot = buffer.read(cx).snapshot(cx);         let (old_excerpt, _) = snapshot            .range_to_buffer_ranges(range.clone())            .pop()            .unwrap();        let old_buffer = cx.new_model(|cx| {            let text = old_excerpt.buffer().as_rope().clone();            let line_ending = old_excerpt.buffer().line_ending();            let language = old_excerpt.buffer().language().cloned();            let language_registry = buffer                .read(cx)                .buffer(old_excerpt.buffer_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,            telemetry,            _subscription: cx.subscribe(&buffer, Self::handle_buffer_event),            builder,            active,            edits: Vec::new(),            line_operations: Vec::new(),            range,            request: None,            elapsed_time: None,            completion: None,        }    }     pub fn set_active(&mut self, active: bool, cx: &mut ModelContext<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: Model<MultiBuffer>,        event: &multi_buffer::Event,        cx: &mut ModelContext<Self>,    ) {        if let multi_buffer::Event::TransactionUndone { transaction_id } = event {            if 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 ModelContext<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(user_prompt, cx)?;                self.request = Some(request.clone());                 cx.spawn(|_, cx| async move { model.stream_completion_text(request, &cx).await })                    .boxed_local()            };        self.handle_stream(telemetry_id, provider_id.to_string(), api_key, stream, cx);        Ok(())    }     fn build_request(        &self,        user_prompt: String,        cx: &mut AppContext,    ) -> Result<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 {                return Err(anyhow::anyhow!("invalid transformation range"));            }        } else {            return Err(anyhow::anyhow!("invalid transformation range"));        };         let prompt = self            .builder            .generate_inline_transformation_prompt(user_prompt, language_name, buffer, range)            .map_err(|e| anyhow::anyhow!("Failed to generate content prompt: {}", e))?;         let mut request_message = LanguageModelRequestMessage {            role: Role::User,            content: Vec::new(),            cache: false,        };         if let Some(context_store) = &self.context_store {            let context = context_store.update(cx, |this, _cx| this.context().clone());            attach_context_to_message(&mut request_message, context);        }         request_message.content.push(prompt.into());         Ok(LanguageModelRequest {            tools: Vec::new(),            stop: Vec::new(),            temperature: None,            messages: vec![request_message],        })    }     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 ModelContext<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            .suggested_indents(selection_start.row..=selection_start.row, cx)            .into_values()            .next()            .unwrap_or_else(|| snapshot.indent_size_for_line(MultiBufferRow(selection_start.row)));         // If the first line in the selection does not have indentation, check the following lines        if suggested_line_indent.len == 0 && suggested_line_indent.kind == IndentKind::Space {            for row in selection_start.row..=self.range.end.to_point(&snapshot).row {                let line_indent = snapshot.indent_size_for_line(MultiBufferRow(row));                // Prefer tabs if a line in the selection uses tabs as indentation                if line_indent.kind == IndentKind::Tab {                    suggested_line_indent.kind = IndentKind::Tab;                    break;                }            }        }         let http_client = cx.http_client().clone();        let telemetry = self.telemetry.clone();        let language_name = {            let multibuffer = self.buffer.read(cx);            let snapshot = multibuffer.snapshot(cx);            let ranges = snapshot.range_to_buffer_ranges(self.range.clone());            ranges                .first()                .and_then(|(excerpt, _)| excerpt.buffer().language())                .map(|language| language.name())        };         self.diff = Diff::default();        self.status = CodegenStatus::Pending;        let mut edit_start = self.range.start.to_offset(&snapshot);        let completion = Arc::new(Mutex::new(String::new()));        let completion_clone = completion.clone();         self.generation = cx.spawn(|codegen, mut cx| {            async move {                let stream = stream.await;                let message_id = stream                    .as_ref()                    .ok()                    .and_then(|stream| stream.message_id.clone());                let generate = async {                    let (mut diff_tx, mut diff_rx) = mpsc::channel(1);                    let executor = cx.background_executor().clone();

Lines 1–500 of 1479. Load more