mod create_file_parser;mod edit_parser;#[cfg(test)]mod evals;mod streaming_fuzzy_matcher; use crate::{Template, Templates};use anyhow::Result;use assistant_tool::ActionLog;use create_file_parser::{CreateFileParser, CreateFileParserEvent};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, path::PathBuf, pin::Pin, sync::Arc, task::Poll};use streaming_diff::{CharOperation, StreamingDiff};use streaming_fuzzy_matcher::StreamingFuzzyMatcher;use util::debug_panic;use zed_llm_client::CompletionIntent; #[derive(Serialize)]struct CreateFilePromptTemplate {    path: Option<PathBuf>,    edit_description: String,} impl Template for CreateFilePromptTemplate {    const TEMPLATE_NAME: &'static str = "create_file_prompt.hbs";} #[derive(Serialize)]struct EditFilePromptTemplate {    path: Option<PathBuf>,    edit_description: String,} impl Template for EditFilePromptTemplate {    const TEMPLATE_NAME: &'static str = "edit_file_prompt.hbs";} #[derive(Clone, Debug, PartialEq, Eq)]pub enum EditAgentOutputEvent {    ResolvingEditRange(Range<Anchor>),    UnresolvedEditRange,    AmbiguousEditRange(Vec<Range<usize>>),    Edited,} #[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>,} impl EditAgent {    pub fn new(        model: Arc<dyn LanguageModel>,        project: Entity<Project>,        action_log: Entity<ActionLog>,        templates: Arc<Templates>,    ) -> Self {        EditAgent {            model,            project,            action_log,            templates,        }    }     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(cx, true))?;            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,                    }),                    cx,                )            });            output_events_tx                .unbounded_send(EditAgentOutputEvent::Edited)                .ok();        })?;         while let Some(event) = parse_rx.next().await {            match event? {                CreateFileParserEvent::NewTextChunk { chunk } => {                    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,                                }),                                cx,                            )                        });                    })?;                    output_events_tx                        .unbounded_send(EditAgentOutputEvent::Edited)                        .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 output = cx.spawn(async move |cx| {            let snapshot = buffer.read_with(cx, |buffer, _| buffer.snapshot())?;            let path = cx.update(|cx| snapshot.resolve_file_path(cx, true))?;            let prompt = EditFilePromptTemplate {                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, 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| text.range)                        .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.                cx.update(|cx| {                    let 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();                        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,                        );                    });                })?;                output_events                    .unbounded_send(EditAgentOutputEvent::Edited)                    .ok();            }             edit_events = compute_edits.await?;        }         output.await    }     fn parse_edit_chunks(        chunks: impl 'static + Send + Stream<Item = Result<String, LanguageModelCompletionError>>,        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();            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>)>>,        async_watch::Receiver<Option<Range<usize>>>,    )    where        T: 'static + Send + Unpin + Stream<Item = Result<EditParserEvent>>,    {        let (old_range_tx, old_range_rx) = async_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 } = edit_event? else {                    break;                };                 old_range_tx.send(matcher.push(&chunk))?;                if done {                    break;                }            }             let matches = matcher.finish();             let old_range = if matches.len() == 1 {                matches.first()            } else {                // No matches or multiple ambiguous matches                None            };            old_range_tx.send(old_range.cloned())?;             let indent = LineIndent::from_iter(                matcher                    .query_lines()                    .first()                    .unwrap_or(&String::new())                    .chars(),            );            let resolved_old_texts = matches                .into_iter()                .map(|range| ResolvedOldText { range, indent })                .collect::<Vec<_>>();             Ok((edit_events, resolved_old_texts))        });         (task, old_range_rx)    }     fn compute_edits<T>(        snapshot: BufferSnapshot,        resolved_old_text: ResolvedOldText,        mut edit_events: T,        cx: &mut AsyncApp,    ) -> (        Task<Result<T>>,        UnboundedReceiver<(Range<Anchor>, Arc<str>)>,    )    where        T: 'static + Send + Unpin + Stream<Item = Result<EditParserEvent>>,    {        let (edits_tx, edits_rx) = mpsc::unbounded();        let compute_edits = cx.background_spawn(async move {            let buffer_start_indent = snapshot                .line_indent_for_row(snapshot.offset_to_point(resolved_old_text.range.start).row);            let indent_delta = if buffer_start_indent.tabs > 0 {                IndentDelta::Tabs(                    buffer_start_indent.tabs as isize - resolved_old_text.indent.tabs as isize,                )            } else {                IndentDelta::Spaces(                    buffer_start_indent.spaces as isize - resolved_old_text.indent.spaces as isize,                )            };             let old_text = snapshot                .text_for_range(resolved_old_text.range.clone())                .collect::<String>();            let mut diff = StreamingDiff::new(old_text);            let mut edit_start = resolved_old_text.range.start;

Lines 1–500 of 1426. Load more