1527 lines · 52.1 KB · Rust
RawDownload
mod create_file_parser;mod edit_parser;#[cfg(all(test, feature = "unit-eval"))]mod evals;pub mod reindent;pub mod streaming_fuzzy_matcher; use crate::{Template, Templates};use action_log::ActionLog;use anyhow::Result;use create_file_parser::{CreateFileParser, CreateFileParserEvent};pub use edit_parser::EditFormat;use edit_parser::{EditParser, EditParserEvent, EditParserMetrics};use futures::{    Stream, StreamExt,    channel::mpsc::{self, UnboundedReceiver},    pin_mut,    stream::BoxStream,};use gpui::{AppContext, AsyncApp, Entity, Task};use language::{Anchor, Buffer, BufferSnapshot, LineIndent, Point, TextBufferSnapshot};use language_model::{    CompletionIntent, LanguageModel, LanguageModelCompletionError, LanguageModelRequest,    LanguageModelRequestMessage, LanguageModelToolChoice, MessageContent, Role,};use project::{AgentLocation, Project};use reindent::{IndentDelta, Reindenter};use schemars::JsonSchema;use serde::{Deserialize, Serialize};use std::{mem, ops::Range, pin::Pin, sync::Arc, task::Poll};use streaming_diff::{CharOperation, StreamingDiff};use streaming_fuzzy_matcher::StreamingFuzzyMatcher; #[derive(Serialize)]struct CreateFilePromptTemplate {    path: Option<String>,    edit_description: String,} impl Template for CreateFilePromptTemplate {    const TEMPLATE_NAME: &'static str = "create_file_prompt.hbs";} #[derive(Serialize)]struct EditFileXmlPromptTemplate {    path: Option<String>,    edit_description: String,} impl Template for EditFileXmlPromptTemplate {    const TEMPLATE_NAME: &'static str = "edit_file_prompt_xml.hbs";} #[derive(Serialize)]struct EditFileDiffFencedPromptTemplate {    path: Option<String>,    edit_description: String,} impl Template for EditFileDiffFencedPromptTemplate {    const TEMPLATE_NAME: &'static str = "edit_file_prompt_diff_fenced.hbs";} #[derive(Clone, Debug, PartialEq, Eq)]pub enum EditAgentOutputEvent {    ResolvingEditRange(Range<Anchor>),    UnresolvedEditRange,    AmbiguousEditRange(Vec<Range<usize>>),    Edited(Range<Anchor>),} #[derive(Clone, Debug, Serialize, Deserialize, JsonSchema)]pub struct EditAgentOutput {    pub raw_edits: String,    pub parser_metrics: EditParserMetrics,} #[derive(Clone)]pub struct EditAgent {    model: Arc<dyn LanguageModel>,    action_log: Entity<ActionLog>,    project: Entity<Project>,    templates: Arc<Templates>,    edit_format: EditFormat,    thinking_allowed: bool,    update_agent_location: bool,} impl EditAgent {    pub fn new(        model: Arc<dyn LanguageModel>,        project: Entity<Project>,        action_log: Entity<ActionLog>,        templates: Arc<Templates>,        edit_format: EditFormat,        allow_thinking: bool,        update_agent_location: bool,    ) -> Self {        EditAgent {            model,            project,            action_log,            templates,            edit_format,            thinking_allowed: allow_thinking,            update_agent_location,        }    }     pub fn overwrite(        &self,        buffer: Entity<Buffer>,        edit_description: String,        conversation: &LanguageModelRequest,        cx: &mut AsyncApp,    ) -> (        Task<Result<EditAgentOutput>>,        mpsc::UnboundedReceiver<EditAgentOutputEvent>,    ) {        let this = self.clone();        let (events_tx, events_rx) = mpsc::unbounded();        let conversation = conversation.clone();        let output = cx.spawn(async move |cx| {            let snapshot = buffer.read_with(cx, |buffer, _| buffer.snapshot());            let path = cx.update(|cx| snapshot.resolve_file_path(true, cx));            let prompt = CreateFilePromptTemplate {                path,                edit_description,            }            .render(&this.templates)?;            let new_chunks = this                .request(conversation, CompletionIntent::CreateFile, prompt, cx)                .await?;             let (output, mut inner_events) = this.overwrite_with_chunks(buffer, new_chunks, cx);            while let Some(event) = inner_events.next().await {                events_tx.unbounded_send(event).ok();            }            output.await        });        (output, events_rx)    }     fn overwrite_with_chunks(        &self,        buffer: Entity<Buffer>,        edit_chunks: impl 'static + Send + Stream<Item = Result<String, LanguageModelCompletionError>>,        cx: &mut AsyncApp,    ) -> (        Task<Result<EditAgentOutput>>,        mpsc::UnboundedReceiver<EditAgentOutputEvent>,    ) {        let (output_events_tx, output_events_rx) = mpsc::unbounded();        let (parse_task, parse_rx) = Self::parse_create_file_chunks(edit_chunks, cx);        let this = self.clone();        let task = cx.spawn(async move |cx| {            this.action_log                .update(cx, |log, cx| log.buffer_created(buffer.clone(), cx));            this.overwrite_with_chunks_internal(buffer, parse_rx, output_events_tx, cx)                .await?;            parse_task.await        });        (task, output_events_rx)    }     async fn overwrite_with_chunks_internal(        &self,        buffer: Entity<Buffer>,        mut parse_rx: UnboundedReceiver<Result<CreateFileParserEvent>>,        output_events_tx: mpsc::UnboundedSender<EditAgentOutputEvent>,        cx: &mut AsyncApp,    ) -> Result<()> {        let buffer_id = cx.update(|cx| {            let buffer_id = buffer.read(cx).remote_id();            if self.update_agent_location {                self.project.update(cx, |project, cx| {                    project.set_agent_location(                        Some(AgentLocation {                            buffer: buffer.downgrade(),                            position: language::Anchor::min_for_buffer(buffer_id),                        }),                        cx,                    )                });            }            buffer_id        });         let send_edit_event = || {            output_events_tx                .unbounded_send(EditAgentOutputEvent::Edited(                    Anchor::min_max_range_for_buffer(buffer_id),                ))                .ok()        };        let set_agent_location = |cx: &mut _| {            if self.update_agent_location {                self.project.update(cx, |project, cx| {                    project.set_agent_location(                        Some(AgentLocation {                            buffer: buffer.downgrade(),                            position: language::Anchor::max_for_buffer(buffer_id),                        }),                        cx,                    )                })            }        };        let mut first_chunk = true;        while let Some(event) = parse_rx.next().await {            match event? {                CreateFileParserEvent::NewTextChunk { chunk } => {                    cx.update(|cx| {                        buffer.update(cx, |buffer, cx| {                            if mem::take(&mut first_chunk) {                                buffer.set_text(chunk, cx)                            } else {                                buffer.append(chunk, cx)                            }                        });                        self.action_log                            .update(cx, |log, cx| log.buffer_edited(buffer.clone(), cx));                        set_agent_location(cx);                    });                    send_edit_event();                }            }        }         if first_chunk {            cx.update(|cx| {                buffer.update(cx, |buffer, cx| buffer.set_text("", cx));                self.action_log                    .update(cx, |log, cx| log.buffer_edited(buffer.clone(), cx));                set_agent_location(cx);            });            send_edit_event();        }         Ok(())    }     pub fn edit(        &self,        buffer: Entity<Buffer>,        edit_description: String,        conversation: &LanguageModelRequest,        cx: &mut AsyncApp,    ) -> (        Task<Result<EditAgentOutput>>,        mpsc::UnboundedReceiver<EditAgentOutputEvent>,    ) {        let this = self.clone();        let (events_tx, events_rx) = mpsc::unbounded();        let conversation = conversation.clone();        let edit_format = self.edit_format;        let output = cx.spawn(async move |cx| {            let snapshot = buffer.read_with(cx, |buffer, _| buffer.snapshot());            let path = cx.update(|cx| snapshot.resolve_file_path(true, cx));            let prompt = match edit_format {                EditFormat::XmlTags => EditFileXmlPromptTemplate {                    path,                    edit_description,                }                .render(&this.templates)?,                EditFormat::DiffFenced => EditFileDiffFencedPromptTemplate {                    path,                    edit_description,                }                .render(&this.templates)?,            };             let edit_chunks = this                .request(conversation, CompletionIntent::EditFile, prompt, cx)                .await?;            this.apply_edit_chunks(buffer, edit_chunks, events_tx, cx)                .await        });        (output, events_rx)    }     async fn apply_edit_chunks(        &self,        buffer: Entity<Buffer>,        edit_chunks: impl 'static + Send + Stream<Item = Result<String, LanguageModelCompletionError>>,        output_events: mpsc::UnboundedSender<EditAgentOutputEvent>,        cx: &mut AsyncApp,    ) -> Result<EditAgentOutput> {        self.action_log            .update(cx, |log, cx| log.buffer_read(buffer.clone(), cx));         let (output, edit_events) = Self::parse_edit_chunks(edit_chunks, self.edit_format, cx);        let mut edit_events = edit_events.peekable();        while let Some(edit_event) = Pin::new(&mut edit_events).peek().await {            // Skip events until we're at the start of a new edit.            let Ok(EditParserEvent::OldTextChunk { .. }) = edit_event else {                edit_events.next().await.unwrap()?;                continue;            };             let snapshot = buffer.read_with(cx, |buffer, _| buffer.snapshot());             // Resolve the old text in the background, updating the agent            // location as we keep refining which range it corresponds to.            let (resolve_old_text, mut old_range) =                Self::resolve_old_text(snapshot.text.clone(), edit_events, cx);            while let Ok(old_range) = old_range.recv().await {                if let Some(old_range) = old_range {                    let old_range = snapshot.anchor_before(old_range.start)                        ..snapshot.anchor_before(old_range.end);                    if self.update_agent_location {                        self.project.update(cx, |project, cx| {                            project.set_agent_location(                                Some(AgentLocation {                                    buffer: buffer.downgrade(),                                    position: old_range.end,                                }),                                cx,                            );                        });                    }                    output_events                        .unbounded_send(EditAgentOutputEvent::ResolvingEditRange(old_range))                        .ok();                }            }             let (edit_events_, mut resolved_old_text) = resolve_old_text.await?;            edit_events = edit_events_;             // If we can't resolve the old text, restart the loop waiting for a            // new edit (or for the stream to end).            let resolved_old_text = match resolved_old_text.len() {                1 => resolved_old_text.pop().unwrap(),                0 => {                    output_events                        .unbounded_send(EditAgentOutputEvent::UnresolvedEditRange)                        .ok();                    continue;                }                _ => {                    let ranges = resolved_old_text                        .into_iter()                        .map(|text| {                            let start_line =                                (snapshot.offset_to_point(text.range.start).row + 1) as usize;                            let end_line =                                (snapshot.offset_to_point(text.range.end).row + 1) as usize;                            start_line..end_line                        })                        .collect();                    output_events                        .unbounded_send(EditAgentOutputEvent::AmbiguousEditRange(ranges))                        .ok();                    continue;                }            };             // Compute edits in the background and apply them as they become            // available.            let (compute_edits, edits) =                Self::compute_edits(snapshot, resolved_old_text, edit_events, cx);            let mut edits = edits.ready_chunks(32);            while let Some(edits) = edits.next().await {                if edits.is_empty() {                    continue;                }                 // Edit the buffer and report edits to the action log as part of the                // same effect cycle, otherwise the edit will be reported as if the                // user made it.                let (min_edit_start, max_edit_end) = cx.update(|cx| {                    let (min_edit_start, max_edit_end) = buffer.update(cx, |buffer, cx| {                        buffer.edit(edits.iter().cloned(), None, cx);                        let max_edit_end = buffer                            .summaries_for_anchors::<Point, _>(                                edits.iter().map(|(range, _)| range.end),                            )                            .max()                            .unwrap();                        let min_edit_start = buffer                            .summaries_for_anchors::<Point, _>(                                edits.iter().map(|(range, _)| range.start),                            )                            .min()                            .unwrap();                        (                            buffer.anchor_after(min_edit_start),                            buffer.anchor_before(max_edit_end),                        )                    });                    self.action_log                        .update(cx, |log, cx| log.buffer_edited(buffer.clone(), cx));                    if self.update_agent_location {                        self.project.update(cx, |project, cx| {                            project.set_agent_location(                                Some(AgentLocation {                                    buffer: buffer.downgrade(),                                    position: max_edit_end,                                }),                                cx,                            );                        });                    }                    (min_edit_start, max_edit_end)                });                output_events                    .unbounded_send(EditAgentOutputEvent::Edited(min_edit_start..max_edit_end))                    .ok();            }             edit_events = compute_edits.await?;        }         output.await    }     fn parse_edit_chunks(        chunks: impl 'static + Send + Stream<Item = Result<String, LanguageModelCompletionError>>,        edit_format: EditFormat,        cx: &mut AsyncApp,    ) -> (        Task<Result<EditAgentOutput>>,        UnboundedReceiver<Result<EditParserEvent>>,    ) {        let (tx, rx) = mpsc::unbounded();        let output = cx.background_spawn(async move {            pin_mut!(chunks);             let mut parser = EditParser::new(edit_format);            let mut raw_edits = String::new();            while let Some(chunk) = chunks.next().await {                match chunk {                    Ok(chunk) => {                        raw_edits.push_str(&chunk);                        for event in parser.push(&chunk) {                            tx.unbounded_send(Ok(event))?;                        }                    }                    Err(error) => {                        tx.unbounded_send(Err(error.into()))?;                    }                }            }            Ok(EditAgentOutput {                raw_edits,                parser_metrics: parser.finish(),            })        });        (output, rx)    }     fn parse_create_file_chunks(        chunks: impl 'static + Send + Stream<Item = Result<String, LanguageModelCompletionError>>,        cx: &mut AsyncApp,    ) -> (        Task<Result<EditAgentOutput>>,        UnboundedReceiver<Result<CreateFileParserEvent>>,    ) {        let (tx, rx) = mpsc::unbounded();        let output = cx.background_spawn(async move {            pin_mut!(chunks);             let mut parser = CreateFileParser::new();            let mut raw_edits = String::new();            while let Some(chunk) = chunks.next().await {                match chunk {                    Ok(chunk) => {                        raw_edits.push_str(&chunk);                        for event in parser.push(Some(&chunk)) {                            tx.unbounded_send(Ok(event))?;                        }                    }                    Err(error) => {                        tx.unbounded_send(Err(error.into()))?;                    }                }            }            // Send final events with None to indicate completion            for event in parser.push(None) {                tx.unbounded_send(Ok(event))?;            }            Ok(EditAgentOutput {                raw_edits,                parser_metrics: EditParserMetrics::default(),            })        });        (output, rx)    }     fn resolve_old_text<T>(        snapshot: TextBufferSnapshot,        mut edit_events: T,        cx: &mut AsyncApp,    ) -> (        Task<Result<(T, Vec<ResolvedOldText>)>>,        watch::Receiver<Option<Range<usize>>>,    )    where        T: 'static + Send + Unpin + Stream<Item = Result<EditParserEvent>>,

Lines 1–500 of 1527.