1961 lines · 73.8 KB · Rust
RawDownload
use crate::{context::LoadedContext, inline_prompt_editor::CodegenStatus};use agent_settings::AgentSettings;use anyhow::{Context as _, Result};use uuid::Uuid; 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, Shared},    join,    stream::BoxStream,};use gpui::{App, AppContext as _, AsyncApp, Context, Entity, EventEmitter, Subscription, Task};use language::{Buffer, IndentKind, LanguageName, Point, TransactionId, line_diff};use language_model::{    LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent,    LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage,    LanguageModelRequestTool, LanguageModelTextStream, LanguageModelToolChoice,    LanguageModelToolUse, Role, TokenUsage,};use multi_buffer::MultiBufferRow;use parking_lot::Mutex;use prompt_store::PromptBuilder;use rope::Rope;use schemars::JsonSchema;use serde::{Deserialize, Serialize};use settings::Settings as _;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 this tool when you cannot or should not make a rewrite. This includes:/// - The user's request is unclear, ambiguous, or nonsensical/// - The requested change cannot be made by only editing the <rewrite_this> section#[derive(Debug, Serialize, Deserialize, JsonSchema)]pub struct FailureMessageInput {    /// A brief message to the user explaining why you're unable to fulfill the request or to ask a question about the request.    #[serde(default)]    pub message: String,} /// Replaces text in <rewrite_this></rewrite_this> tags with your replacement_text./// Only use this tool when you are confident you understand the user's request and can fulfill it/// by editing the marked section.#[derive(Debug, Serialize, Deserialize, JsonSchema)]pub struct RewriteSectionInput {    /// The text to replace the section with.    #[serde(default)]    pub replacement_text: String,} 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>,    builder: Arc<PromptBuilder>,    pub is_insertion: bool,    session_id: Uuid,} pub const REWRITE_SECTION_TOOL_NAME: &str = "rewrite_section";pub const FAILURE_MESSAGE_TOOL_NAME: &str = "failure_message"; impl BufferCodegen {    pub fn new(        buffer: Entity<MultiBuffer>,        range: Range<Anchor>,        initial_transaction_id: Option<TransactionId>,        session_id: Uuid,        builder: Arc<PromptBuilder>,        cx: &mut Context<Self>,    ) -> Self {        let codegen = cx.new(|cx| {            CodegenAlternative::new(                buffer.clone(),                range.clone(),                false,                builder.clone(),                session_id,                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,            builder,            session_id,        };        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_completion(&self, cx: &App) -> Option<String> {        self.active_alternative().read(cx).current_completion()    }     pub fn active_alternative(&self) -> &Entity<CodegenAlternative> {        &self.alternatives[self.active_alternative]    }     pub fn language_name(&self, cx: &App) -> Option<LanguageName> {        self.active_alternative().read(cx).language_name(cx)    }     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,        context_task: Shared<Task<Option<LoadedContext>>>,        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,                    self.builder.clone(),                    self.session_id,                    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(), context_task.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()    }     pub fn selected_text<'a>(&self, cx: &'a App) -> Option<&'a str> {        self.active_alternative().read(cx).selected_text()    }     pub fn session_id(&self) -> Uuid {        self.session_id    }} 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,    _subscription: gpui::Subscription,    builder: Arc<PromptBuilder>,    active: bool,    edits: Vec<(Range<Anchor>, String)>,    line_operations: Vec<LineOperation>,    elapsed_time: Option<f64>,    completion: Option<String>,    selected_text: Option<String>,    pub message_id: Option<String>,    session_id: Uuid,    pub description: Option<String>,    pub failure: Option<String>,} impl EventEmitter<CodegenEvent> for CodegenAlternative {} impl CodegenAlternative {    pub fn new(        buffer: Entity<MultiBuffer>,        range: Range<Anchor>,        active: bool,        builder: Arc<PromptBuilder>,        session_id: Uuid,        cx: &mut Context<Self>,    ) -> Self {        let snapshot = buffer.read(cx).snapshot(cx);         let (old_buffer, _, _) = snapshot            .range_to_buffer_ranges(range.start..=range.end)            .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(),            builder,            active: active,            edits: Vec::new(),            line_operations: Vec::new(),            range,            elapsed_time: None,            completion: None,            selected_text: None,            session_id,            description: None,            failure: None,            _subscription: cx.subscribe(&buffer, Self::handle_buffer_event),        }    }     pub fn language_name(&self, cx: &App) -> Option<LanguageName> {        self.old_buffer            .read(cx)            .language()            .map(|language| language.name())    }     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 use_streaming_tools(model: &dyn LanguageModel, cx: &App) -> bool {        model.supports_streaming_tools()            && AgentSettings::get_global(cx).inline_assistant_use_streaming_tools    }     pub fn start(        &mut self,        user_prompt: String,        context_task: Shared<Task<Option<LoadedContext>>>,        model: Arc<dyn LanguageModel>,        cx: &mut Context<Self>,    ) -> Result<()> {        // Clear the model explanation since the user has started a new generation.        self.description = None;         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));         if Self::use_streaming_tools(model.as_ref(), cx) {            let request = self.build_request(&model, user_prompt, context_task, cx)?;            let completion_events = cx.spawn({                let model = model.clone();                async move |_, cx| model.stream_completion(request.await, cx).await            });            self.generation = self.handle_completion(model, completion_events, cx);        } else {            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, context_task, cx)?;                    cx.spawn({                        let model = model.clone();                        async move |_, cx| {                            Ok(model.stream_completion_text(request.await, cx).await?)                        }                    })                    .boxed_local()                };            self.generation =                self.handle_stream(model, /* strip_invalid_spans: */ true, stream, cx);        }         Ok(())    }     fn build_request_tools(        &self,        model: &Arc<dyn LanguageModel>,        user_prompt: String,        context_task: Shared<Task<Option<LoadedContext>>>,        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 system_prompt = self            .builder            .generate_inline_transformation_prompt_tools(                language_name,                buffer,                range.start.0..range.end.0,            )            .context("generating content prompt")?;         let temperature = AgentSettings::temperature_for_model(model, cx);         let tool_input_format = model.tool_input_format();        let tool_choice = model            .supports_tool_choice(LanguageModelToolChoice::Any)            .then_some(LanguageModelToolChoice::Any);

Lines 1–500 of 1961.