mod create_file_parser;mod edit_parser;#[cfg(test)]mod evals;pub mod streaming_fuzzy_matcher; use crate::{Template, Templates};use action_log::ActionLog;use anyhow::Result;use cloud_llm_client::CompletionIntent;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::{    LanguageModel, LanguageModelCompletionError, LanguageModelRequest, LanguageModelRequestMessage,    LanguageModelToolChoice, MessageContent, Role,};use project::{AgentLocation, Project};use schemars::JsonSchema;use serde::{Deserialize, Serialize};use std::{cmp, iter, 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,} impl EditAgent {    pub fn new(        model: Arc<dyn LanguageModel>,        project: Entity<Project>,        action_log: Entity<ActionLog>,        templates: Arc<Templates>,        edit_format: EditFormat,    ) -> Self {        EditAgent {            model,            project,            action_log,            templates,            edit_format,        }    }     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<()> {        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);            });            self.project.update(cx, |project, cx| {                project.set_agent_location(                    Some(AgentLocation {                        buffer: buffer.downgrade(),                        position: language::Anchor::max_for_buffer(buffer.read(cx).remote_id()),                    }),                    cx,                )            });            output_events_tx                .unbounded_send(EditAgentOutputEvent::Edited(                    Anchor::min_max_range_for_buffer(buffer.read(cx).remote_id()),                ))                .ok();        });         while let Some(event) = parse_rx.next().await {            match event? {                CreateFileParserEvent::NewTextChunk { chunk } => {                    let buffer_id = cx.update(|cx| {                        buffer.update(cx, |buffer, cx| buffer.append(chunk, cx));                        self.action_log                            .update(cx, |log, cx| log.buffer_edited(buffer.clone(), cx));                        self.project.update(cx, |project, cx| {                            project.set_agent_location(                                Some(AgentLocation {                                    buffer: buffer.downgrade(),                                    position: language::Anchor::max_for_buffer(                                        buffer.read(cx).remote_id(),                                    ),                                }),                                cx,                            )                        });                        buffer.read(cx).remote_id()                    });                    output_events_tx                        .unbounded_send(EditAgentOutputEvent::Edited(                            Anchor::min_max_range_for_buffer(buffer_id),                        ))                        .ok();                }            }        }         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);                    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));                    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>>,    {        let (mut old_range_tx, old_range_rx) = watch::channel(None);        let task = cx.background_spawn(async move {            let mut matcher = StreamingFuzzyMatcher::new(snapshot);            while let Some(edit_event) = edit_events.next().await {                let EditParserEvent::OldTextChunk {                    chunk,                    done,                    line_hint,                } = edit_event?                else {                    break;                };                 old_range_tx.send(matcher.push(&chunk, line_hint))?;                if done {                    break;                }            }             let matches = matcher.finish();            let best_match = matcher.select_best_match();             old_range_tx.send(best_match.clone())?;             let indent = LineIndent::from_iter(                matcher                    .query_lines()

Lines 1–500 of 1505. Load more