4f78165ee8
Release Notes: - Improved latency when the agent starts streaming edits. --------- Co-authored-by: Ben Brandt <benjamin.j.brandt@gmail.com>
1327 lines
46 KiB
Rust
1327 lines
46 KiB
Rust
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;
|
|
|
|
#[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,
|
|
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, 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, 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_, 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 Some(resolved_old_text) = resolved_old_text else {
|
|
output_events
|
|
.unbounded_send(EditAgentOutputEvent::UnresolvedEditRange)
|
|
.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, Option<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 old_range = matcher.finish();
|
|
old_range_tx.send(old_range.clone())?;
|
|
if let Some(old_range) = old_range {
|
|
let line_indent =
|
|
LineIndent::from_iter(matcher.query_lines().first().unwrap().chars());
|
|
Ok((
|
|
edit_events,
|
|
Some(ResolvedOldText {
|
|
range: old_range,
|
|
indent: line_indent,
|
|
}),
|
|
))
|
|
} else {
|
|
Ok((edit_events, None))
|
|
}
|
|
});
|
|
|
|
(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;
|
|
let mut new_text_chunks =
|
|
Self::reindent_new_text_chunks(indent_delta, &mut edit_events);
|
|
let mut done = false;
|
|
while !done {
|
|
let char_operations = if let Some(new_text_chunk) = new_text_chunks.next().await {
|
|
diff.push_new(&new_text_chunk?)
|
|
} else {
|
|
done = true;
|
|
mem::take(&mut diff).finish()
|
|
};
|
|
|
|
for op in char_operations {
|
|
match op {
|
|
CharOperation::Insert { text } => {
|
|
let edit_start = snapshot.anchor_after(edit_start);
|
|
edits_tx.unbounded_send((edit_start..edit_start, Arc::from(text)))?;
|
|
}
|
|
CharOperation::Delete { bytes } => {
|
|
let edit_end = edit_start + bytes;
|
|
let edit_range =
|
|
snapshot.anchor_after(edit_start)..snapshot.anchor_before(edit_end);
|
|
edit_start = edit_end;
|
|
edits_tx.unbounded_send((edit_range, Arc::from("")))?;
|
|
}
|
|
CharOperation::Keep { bytes } => edit_start += bytes,
|
|
}
|
|
}
|
|
}
|
|
|
|
drop(new_text_chunks);
|
|
anyhow::Ok(edit_events)
|
|
});
|
|
|
|
(compute_edits, edits_rx)
|
|
}
|
|
|
|
fn reindent_new_text_chunks(
|
|
delta: IndentDelta,
|
|
mut stream: impl Unpin + Stream<Item = Result<EditParserEvent>>,
|
|
) -> impl Stream<Item = Result<String>> {
|
|
let mut buffer = String::new();
|
|
let mut in_leading_whitespace = true;
|
|
let mut done = false;
|
|
futures::stream::poll_fn(move |cx| {
|
|
while !done {
|
|
let (chunk, is_last_chunk) = match stream.poll_next_unpin(cx) {
|
|
Poll::Ready(Some(Ok(EditParserEvent::NewTextChunk { chunk, done }))) => {
|
|
(chunk, done)
|
|
}
|
|
Poll::Ready(Some(Err(err))) => return Poll::Ready(Some(Err(err))),
|
|
Poll::Pending => return Poll::Pending,
|
|
_ => return Poll::Ready(None),
|
|
};
|
|
|
|
buffer.push_str(&chunk);
|
|
|
|
let mut indented_new_text = String::new();
|
|
let mut start_ix = 0;
|
|
let mut newlines = buffer.match_indices('\n').peekable();
|
|
loop {
|
|
let (line_end, is_pending_line) = match newlines.next() {
|
|
Some((ix, _)) => (ix, false),
|
|
None => (buffer.len(), true),
|
|
};
|
|
let line = &buffer[start_ix..line_end];
|
|
|
|
if in_leading_whitespace {
|
|
if let Some(non_whitespace_ix) = line.find(|c| delta.character() != c) {
|
|
// We found a non-whitespace character, adjust
|
|
// indentation based on the delta.
|
|
let new_indent_len =
|
|
cmp::max(0, non_whitespace_ix as isize + delta.len()) as usize;
|
|
indented_new_text
|
|
.extend(iter::repeat(delta.character()).take(new_indent_len));
|
|
indented_new_text.push_str(&line[non_whitespace_ix..]);
|
|
in_leading_whitespace = false;
|
|
} else if is_pending_line {
|
|
// We're still in leading whitespace and this line is incomplete.
|
|
// Stop processing until we receive more input.
|
|
break;
|
|
} else {
|
|
// This line is entirely whitespace. Push it without indentation.
|
|
indented_new_text.push_str(line);
|
|
}
|
|
} else {
|
|
indented_new_text.push_str(line);
|
|
}
|
|
|
|
if is_pending_line {
|
|
start_ix = line_end;
|
|
break;
|
|
} else {
|
|
in_leading_whitespace = true;
|
|
indented_new_text.push('\n');
|
|
start_ix = line_end + 1;
|
|
}
|
|
}
|
|
buffer.replace_range(..start_ix, "");
|
|
|
|
// This was the last chunk, push all the buffered content as-is.
|
|
if is_last_chunk {
|
|
indented_new_text.push_str(&buffer);
|
|
buffer.clear();
|
|
done = true;
|
|
}
|
|
|
|
if !indented_new_text.is_empty() {
|
|
return Poll::Ready(Some(Ok(indented_new_text)));
|
|
}
|
|
}
|
|
|
|
Poll::Ready(None)
|
|
})
|
|
}
|
|
|
|
async fn request(
|
|
&self,
|
|
mut conversation: LanguageModelRequest,
|
|
prompt: String,
|
|
cx: &mut AsyncApp,
|
|
) -> Result<BoxStream<'static, Result<String, LanguageModelCompletionError>>> {
|
|
let mut messages_iter = conversation.messages.iter_mut();
|
|
if let Some(last_message) = messages_iter.next_back() {
|
|
if last_message.role == Role::Assistant {
|
|
let old_content_len = last_message.content.len();
|
|
last_message
|
|
.content
|
|
.retain(|content| !matches!(content, MessageContent::ToolUse(_)));
|
|
let new_content_len = last_message.content.len();
|
|
|
|
// We just removed pending tool uses from the content of the
|
|
// last message, so it doesn't make sense to cache it anymore
|
|
// (e.g., the message will look very different on the next
|
|
// request). Thus, we move the flag to the message prior to it,
|
|
// as it will still be a valid prefix of the conversation.
|
|
if old_content_len != new_content_len && last_message.cache {
|
|
if let Some(prev_message) = messages_iter.next_back() {
|
|
last_message.cache = false;
|
|
prev_message.cache = true;
|
|
}
|
|
}
|
|
|
|
if last_message.content.is_empty() {
|
|
conversation.messages.pop();
|
|
}
|
|
} else {
|
|
debug_panic!(
|
|
"Last message must be an Assistant tool calling! Got {:?}",
|
|
last_message.content
|
|
);
|
|
}
|
|
}
|
|
|
|
conversation.messages.push(LanguageModelRequestMessage {
|
|
role: Role::User,
|
|
content: vec![MessageContent::Text(prompt)],
|
|
cache: false,
|
|
});
|
|
|
|
// Include tools in the request so that we can take advantage of
|
|
// caching when ToolChoice::None is supported.
|
|
let mut tool_choice = None;
|
|
let mut tools = Vec::new();
|
|
if !conversation.tools.is_empty()
|
|
&& self
|
|
.model
|
|
.supports_tool_choice(LanguageModelToolChoice::None)
|
|
{
|
|
tool_choice = Some(LanguageModelToolChoice::None);
|
|
tools = conversation.tools.clone();
|
|
}
|
|
|
|
let request = LanguageModelRequest {
|
|
thread_id: conversation.thread_id,
|
|
prompt_id: conversation.prompt_id,
|
|
mode: conversation.mode,
|
|
messages: conversation.messages,
|
|
tool_choice,
|
|
tools,
|
|
stop: Vec::new(),
|
|
temperature: None,
|
|
};
|
|
|
|
Ok(self.model.stream_completion_text(request, cx).await?.stream)
|
|
}
|
|
}
|
|
|
|
struct ResolvedOldText {
|
|
range: Range<usize>,
|
|
indent: LineIndent,
|
|
}
|
|
|
|
#[derive(Copy, Clone, Debug)]
|
|
enum IndentDelta {
|
|
Spaces(isize),
|
|
Tabs(isize),
|
|
}
|
|
|
|
impl IndentDelta {
|
|
fn character(&self) -> char {
|
|
match self {
|
|
IndentDelta::Spaces(_) => ' ',
|
|
IndentDelta::Tabs(_) => '\t',
|
|
}
|
|
}
|
|
|
|
fn len(&self) -> isize {
|
|
match self {
|
|
IndentDelta::Spaces(n) => *n,
|
|
IndentDelta::Tabs(n) => *n,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use fs::FakeFs;
|
|
use futures::stream;
|
|
use gpui::{AppContext, TestAppContext};
|
|
use indoc::indoc;
|
|
use language_model::fake_provider::FakeLanguageModel;
|
|
use project::{AgentLocation, Project};
|
|
use rand::prelude::*;
|
|
use rand::rngs::StdRng;
|
|
use std::cmp;
|
|
|
|
#[gpui::test(iterations = 100)]
|
|
async fn test_empty_old_text(cx: &mut TestAppContext, mut rng: StdRng) {
|
|
let agent = init_test(cx).await;
|
|
let buffer = cx.new(|cx| {
|
|
Buffer::local(
|
|
indoc! {"
|
|
abc
|
|
def
|
|
ghi
|
|
"},
|
|
cx,
|
|
)
|
|
});
|
|
let (apply, _events) = agent.edit(
|
|
buffer.clone(),
|
|
String::new(),
|
|
&LanguageModelRequest::default(),
|
|
&mut cx.to_async(),
|
|
);
|
|
cx.run_until_parked();
|
|
|
|
simulate_llm_output(
|
|
&agent,
|
|
indoc! {"
|
|
<old_text></old_text>
|
|
<new_text>jkl</new_text>
|
|
<old_text>def</old_text>
|
|
<new_text>DEF</new_text>
|
|
"},
|
|
&mut rng,
|
|
cx,
|
|
);
|
|
apply.await.unwrap();
|
|
|
|
pretty_assertions::assert_eq!(
|
|
buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
|
|
indoc! {"
|
|
abc
|
|
DEF
|
|
ghi
|
|
"}
|
|
);
|
|
}
|
|
|
|
#[gpui::test(iterations = 100)]
|
|
async fn test_indentation(cx: &mut TestAppContext, mut rng: StdRng) {
|
|
let agent = init_test(cx).await;
|
|
let buffer = cx.new(|cx| {
|
|
Buffer::local(
|
|
indoc! {"
|
|
lorem
|
|
ipsum
|
|
dolor
|
|
sit
|
|
"},
|
|
cx,
|
|
)
|
|
});
|
|
let (apply, _events) = agent.edit(
|
|
buffer.clone(),
|
|
String::new(),
|
|
&LanguageModelRequest::default(),
|
|
&mut cx.to_async(),
|
|
);
|
|
cx.run_until_parked();
|
|
|
|
simulate_llm_output(
|
|
&agent,
|
|
indoc! {"
|
|
<old_text>
|
|
ipsum
|
|
dolor
|
|
sit
|
|
</old_text>
|
|
<new_text>
|
|
ipsum
|
|
dolor
|
|
sit
|
|
amet
|
|
</new_text>
|
|
"},
|
|
&mut rng,
|
|
cx,
|
|
);
|
|
apply.await.unwrap();
|
|
|
|
pretty_assertions::assert_eq!(
|
|
buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
|
|
indoc! {"
|
|
lorem
|
|
ipsum
|
|
dolor
|
|
sit
|
|
amet
|
|
"}
|
|
);
|
|
}
|
|
|
|
#[gpui::test(iterations = 100)]
|
|
async fn test_dependent_edits(cx: &mut TestAppContext, mut rng: StdRng) {
|
|
let agent = init_test(cx).await;
|
|
let buffer = cx.new(|cx| Buffer::local("abc\ndef\nghi", cx));
|
|
let (apply, _events) = agent.edit(
|
|
buffer.clone(),
|
|
String::new(),
|
|
&LanguageModelRequest::default(),
|
|
&mut cx.to_async(),
|
|
);
|
|
cx.run_until_parked();
|
|
|
|
simulate_llm_output(
|
|
&agent,
|
|
indoc! {"
|
|
<old_text>
|
|
def
|
|
</old_text>
|
|
<new_text>
|
|
DEF
|
|
</new_text>
|
|
|
|
<old_text>
|
|
DEF
|
|
</old_text>
|
|
<new_text>
|
|
DeF
|
|
</new_text>
|
|
"},
|
|
&mut rng,
|
|
cx,
|
|
);
|
|
apply.await.unwrap();
|
|
|
|
assert_eq!(
|
|
buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
|
|
"abc\nDeF\nghi"
|
|
);
|
|
}
|
|
|
|
#[gpui::test(iterations = 100)]
|
|
async fn test_old_text_hallucination(cx: &mut TestAppContext, mut rng: StdRng) {
|
|
let agent = init_test(cx).await;
|
|
let buffer = cx.new(|cx| Buffer::local("abc\ndef\nghi", cx));
|
|
let (apply, _events) = agent.edit(
|
|
buffer.clone(),
|
|
String::new(),
|
|
&LanguageModelRequest::default(),
|
|
&mut cx.to_async(),
|
|
);
|
|
cx.run_until_parked();
|
|
|
|
simulate_llm_output(
|
|
&agent,
|
|
indoc! {"
|
|
<old_text>
|
|
jkl
|
|
</old_text>
|
|
<new_text>
|
|
mno
|
|
</new_text>
|
|
|
|
<old_text>
|
|
abc
|
|
</old_text>
|
|
<new_text>
|
|
ABC
|
|
</new_text>
|
|
"},
|
|
&mut rng,
|
|
cx,
|
|
);
|
|
apply.await.unwrap();
|
|
|
|
assert_eq!(
|
|
buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
|
|
"ABC\ndef\nghi"
|
|
);
|
|
}
|
|
|
|
#[gpui::test]
|
|
async fn test_edit_events(cx: &mut TestAppContext) {
|
|
let agent = init_test(cx).await;
|
|
let model = agent.model.as_fake();
|
|
let project = agent
|
|
.action_log
|
|
.read_with(cx, |log, _| log.project().clone());
|
|
let buffer = cx.new(|cx| Buffer::local("abc\ndef\nghi\njkl", cx));
|
|
|
|
let mut async_cx = cx.to_async();
|
|
let (apply, mut events) = agent.edit(
|
|
buffer.clone(),
|
|
String::new(),
|
|
&LanguageModelRequest::default(),
|
|
&mut async_cx,
|
|
);
|
|
cx.run_until_parked();
|
|
|
|
model.stream_last_completion_response("<old_text>a");
|
|
cx.run_until_parked();
|
|
assert_eq!(drain_events(&mut events), vec![]);
|
|
assert_eq!(
|
|
buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
|
|
"abc\ndef\nghi\njkl"
|
|
);
|
|
assert_eq!(
|
|
project.read_with(cx, |project, _| project.agent_location()),
|
|
None
|
|
);
|
|
|
|
model.stream_last_completion_response("bc</old_text>");
|
|
cx.run_until_parked();
|
|
assert_eq!(
|
|
drain_events(&mut events),
|
|
vec![EditAgentOutputEvent::ResolvingEditRange(buffer.read_with(
|
|
cx,
|
|
|buffer, _| buffer.anchor_before(Point::new(0, 0))
|
|
..buffer.anchor_before(Point::new(0, 3))
|
|
))]
|
|
);
|
|
assert_eq!(
|
|
buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
|
|
"abc\ndef\nghi\njkl"
|
|
);
|
|
assert_eq!(
|
|
project.read_with(cx, |project, _| project.agent_location()),
|
|
Some(AgentLocation {
|
|
buffer: buffer.downgrade(),
|
|
position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(0, 3)))
|
|
})
|
|
);
|
|
|
|
model.stream_last_completion_response("<new_text>abX");
|
|
cx.run_until_parked();
|
|
assert_eq!(drain_events(&mut events), [EditAgentOutputEvent::Edited]);
|
|
assert_eq!(
|
|
buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
|
|
"abXc\ndef\nghi\njkl"
|
|
);
|
|
assert_eq!(
|
|
project.read_with(cx, |project, _| project.agent_location()),
|
|
Some(AgentLocation {
|
|
buffer: buffer.downgrade(),
|
|
position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(0, 3)))
|
|
})
|
|
);
|
|
|
|
model.stream_last_completion_response("cY");
|
|
cx.run_until_parked();
|
|
assert_eq!(drain_events(&mut events), [EditAgentOutputEvent::Edited]);
|
|
assert_eq!(
|
|
buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
|
|
"abXcY\ndef\nghi\njkl"
|
|
);
|
|
assert_eq!(
|
|
project.read_with(cx, |project, _| project.agent_location()),
|
|
Some(AgentLocation {
|
|
buffer: buffer.downgrade(),
|
|
position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(0, 5)))
|
|
})
|
|
);
|
|
|
|
model.stream_last_completion_response("</new_text>");
|
|
model.stream_last_completion_response("<old_text>hall");
|
|
cx.run_until_parked();
|
|
assert_eq!(drain_events(&mut events), vec![]);
|
|
assert_eq!(
|
|
buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
|
|
"abXcY\ndef\nghi\njkl"
|
|
);
|
|
assert_eq!(
|
|
project.read_with(cx, |project, _| project.agent_location()),
|
|
Some(AgentLocation {
|
|
buffer: buffer.downgrade(),
|
|
position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(0, 5)))
|
|
})
|
|
);
|
|
|
|
model.stream_last_completion_response("ucinated old</old_text>");
|
|
model.stream_last_completion_response("<new_text>");
|
|
cx.run_until_parked();
|
|
assert_eq!(
|
|
drain_events(&mut events),
|
|
vec![EditAgentOutputEvent::UnresolvedEditRange]
|
|
);
|
|
assert_eq!(
|
|
buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
|
|
"abXcY\ndef\nghi\njkl"
|
|
);
|
|
assert_eq!(
|
|
project.read_with(cx, |project, _| project.agent_location()),
|
|
Some(AgentLocation {
|
|
buffer: buffer.downgrade(),
|
|
position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(0, 5)))
|
|
})
|
|
);
|
|
|
|
model.stream_last_completion_response("hallucinated new</new_");
|
|
model.stream_last_completion_response("text>");
|
|
cx.run_until_parked();
|
|
assert_eq!(drain_events(&mut events), vec![]);
|
|
assert_eq!(
|
|
buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
|
|
"abXcY\ndef\nghi\njkl"
|
|
);
|
|
assert_eq!(
|
|
project.read_with(cx, |project, _| project.agent_location()),
|
|
Some(AgentLocation {
|
|
buffer: buffer.downgrade(),
|
|
position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(0, 5)))
|
|
})
|
|
);
|
|
|
|
model.stream_last_completion_response("<old_text>\nghi\nj");
|
|
cx.run_until_parked();
|
|
assert_eq!(
|
|
drain_events(&mut events),
|
|
vec![EditAgentOutputEvent::ResolvingEditRange(buffer.read_with(
|
|
cx,
|
|
|buffer, _| buffer.anchor_before(Point::new(2, 0))
|
|
..buffer.anchor_before(Point::new(2, 3))
|
|
))]
|
|
);
|
|
assert_eq!(
|
|
buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
|
|
"abXcY\ndef\nghi\njkl"
|
|
);
|
|
assert_eq!(
|
|
project.read_with(cx, |project, _| project.agent_location()),
|
|
Some(AgentLocation {
|
|
buffer: buffer.downgrade(),
|
|
position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(2, 3)))
|
|
})
|
|
);
|
|
|
|
model.stream_last_completion_response("kl</old_text>");
|
|
model.stream_last_completion_response("<new_text>");
|
|
cx.run_until_parked();
|
|
assert_eq!(
|
|
drain_events(&mut events),
|
|
vec![EditAgentOutputEvent::ResolvingEditRange(buffer.read_with(
|
|
cx,
|
|
|buffer, _| buffer.anchor_before(Point::new(2, 0))
|
|
..buffer.anchor_before(Point::new(3, 3))
|
|
))]
|
|
);
|
|
assert_eq!(
|
|
buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
|
|
"abXcY\ndef\nghi\njkl"
|
|
);
|
|
assert_eq!(
|
|
project.read_with(cx, |project, _| project.agent_location()),
|
|
Some(AgentLocation {
|
|
buffer: buffer.downgrade(),
|
|
position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(3, 3)))
|
|
})
|
|
);
|
|
|
|
model.stream_last_completion_response("GHI</new_text>");
|
|
cx.run_until_parked();
|
|
assert_eq!(
|
|
drain_events(&mut events),
|
|
vec![EditAgentOutputEvent::Edited]
|
|
);
|
|
assert_eq!(
|
|
buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
|
|
"abXcY\ndef\nGHI"
|
|
);
|
|
assert_eq!(
|
|
project.read_with(cx, |project, _| project.agent_location()),
|
|
Some(AgentLocation {
|
|
buffer: buffer.downgrade(),
|
|
position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(2, 3)))
|
|
})
|
|
);
|
|
|
|
model.end_last_completion_stream();
|
|
apply.await.unwrap();
|
|
assert_eq!(
|
|
buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
|
|
"abXcY\ndef\nGHI"
|
|
);
|
|
assert_eq!(drain_events(&mut events), vec![]);
|
|
assert_eq!(
|
|
project.read_with(cx, |project, _| project.agent_location()),
|
|
Some(AgentLocation {
|
|
buffer: buffer.downgrade(),
|
|
position: buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(2, 3)))
|
|
})
|
|
);
|
|
}
|
|
|
|
#[gpui::test]
|
|
async fn test_overwrite_events(cx: &mut TestAppContext) {
|
|
let agent = init_test(cx).await;
|
|
let project = agent
|
|
.action_log
|
|
.read_with(cx, |log, _| log.project().clone());
|
|
let buffer = cx.new(|cx| Buffer::local("abc\ndef\nghi", cx));
|
|
let (chunks_tx, chunks_rx) = mpsc::unbounded();
|
|
let (apply, mut events) = agent.overwrite_with_chunks(
|
|
buffer.clone(),
|
|
chunks_rx.map(|chunk: &str| Ok(chunk.to_string())),
|
|
&mut cx.to_async(),
|
|
);
|
|
|
|
cx.run_until_parked();
|
|
assert_eq!(
|
|
drain_events(&mut events),
|
|
vec![EditAgentOutputEvent::Edited]
|
|
);
|
|
assert_eq!(
|
|
buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
|
|
""
|
|
);
|
|
assert_eq!(
|
|
project.read_with(cx, |project, _| project.agent_location()),
|
|
Some(AgentLocation {
|
|
buffer: buffer.downgrade(),
|
|
position: language::Anchor::MAX
|
|
})
|
|
);
|
|
|
|
chunks_tx.unbounded_send("```\njkl\n").unwrap();
|
|
cx.run_until_parked();
|
|
assert_eq!(
|
|
drain_events(&mut events),
|
|
vec![EditAgentOutputEvent::Edited]
|
|
);
|
|
assert_eq!(
|
|
buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
|
|
"jkl"
|
|
);
|
|
assert_eq!(
|
|
project.read_with(cx, |project, _| project.agent_location()),
|
|
Some(AgentLocation {
|
|
buffer: buffer.downgrade(),
|
|
position: language::Anchor::MAX
|
|
})
|
|
);
|
|
|
|
chunks_tx.unbounded_send("mno\n").unwrap();
|
|
cx.run_until_parked();
|
|
assert_eq!(
|
|
drain_events(&mut events),
|
|
vec![EditAgentOutputEvent::Edited]
|
|
);
|
|
assert_eq!(
|
|
buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
|
|
"jkl\nmno"
|
|
);
|
|
assert_eq!(
|
|
project.read_with(cx, |project, _| project.agent_location()),
|
|
Some(AgentLocation {
|
|
buffer: buffer.downgrade(),
|
|
position: language::Anchor::MAX
|
|
})
|
|
);
|
|
|
|
chunks_tx.unbounded_send("pqr\n```").unwrap();
|
|
cx.run_until_parked();
|
|
assert_eq!(
|
|
drain_events(&mut events),
|
|
vec![EditAgentOutputEvent::Edited]
|
|
);
|
|
assert_eq!(
|
|
buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
|
|
"jkl\nmno\npqr"
|
|
);
|
|
assert_eq!(
|
|
project.read_with(cx, |project, _| project.agent_location()),
|
|
Some(AgentLocation {
|
|
buffer: buffer.downgrade(),
|
|
position: language::Anchor::MAX
|
|
})
|
|
);
|
|
|
|
drop(chunks_tx);
|
|
apply.await.unwrap();
|
|
assert_eq!(
|
|
buffer.read_with(cx, |buffer, _| buffer.snapshot().text()),
|
|
"jkl\nmno\npqr"
|
|
);
|
|
assert_eq!(drain_events(&mut events), vec![]);
|
|
assert_eq!(
|
|
project.read_with(cx, |project, _| project.agent_location()),
|
|
Some(AgentLocation {
|
|
buffer: buffer.downgrade(),
|
|
position: language::Anchor::MAX
|
|
})
|
|
);
|
|
}
|
|
|
|
#[gpui::test(iterations = 100)]
|
|
async fn test_indent_new_text_chunks(mut rng: StdRng) {
|
|
let chunks = to_random_chunks(&mut rng, " abc\n def\n ghi");
|
|
let new_text_chunks = stream::iter(chunks.iter().enumerate().map(|(index, chunk)| {
|
|
Ok(EditParserEvent::NewTextChunk {
|
|
chunk: chunk.clone(),
|
|
done: index == chunks.len() - 1,
|
|
})
|
|
}));
|
|
let indented_chunks =
|
|
EditAgent::reindent_new_text_chunks(IndentDelta::Spaces(2), new_text_chunks)
|
|
.collect::<Vec<_>>()
|
|
.await;
|
|
let new_text = indented_chunks
|
|
.into_iter()
|
|
.collect::<Result<String>>()
|
|
.unwrap();
|
|
assert_eq!(new_text, " abc\n def\n ghi");
|
|
}
|
|
|
|
#[gpui::test(iterations = 100)]
|
|
async fn test_outdent_new_text_chunks(mut rng: StdRng) {
|
|
let chunks = to_random_chunks(&mut rng, "\t\t\t\tabc\n\t\tdef\n\t\t\t\t\t\tghi");
|
|
let new_text_chunks = stream::iter(chunks.iter().enumerate().map(|(index, chunk)| {
|
|
Ok(EditParserEvent::NewTextChunk {
|
|
chunk: chunk.clone(),
|
|
done: index == chunks.len() - 1,
|
|
})
|
|
}));
|
|
let indented_chunks =
|
|
EditAgent::reindent_new_text_chunks(IndentDelta::Tabs(-2), new_text_chunks)
|
|
.collect::<Vec<_>>()
|
|
.await;
|
|
let new_text = indented_chunks
|
|
.into_iter()
|
|
.collect::<Result<String>>()
|
|
.unwrap();
|
|
assert_eq!(new_text, "\t\tabc\ndef\n\t\t\t\tghi");
|
|
}
|
|
|
|
#[gpui::test(iterations = 100)]
|
|
async fn test_random_indents(mut rng: StdRng) {
|
|
let len = rng.gen_range(1..=100);
|
|
let new_text = util::RandomCharIter::new(&mut rng)
|
|
.with_simple_text()
|
|
.take(len)
|
|
.collect::<String>();
|
|
let new_text = new_text
|
|
.split('\n')
|
|
.map(|line| format!("{}{}", " ".repeat(rng.gen_range(0..=8)), line))
|
|
.collect::<Vec<_>>()
|
|
.join("\n");
|
|
let delta = IndentDelta::Spaces(rng.gen_range(-4..=4));
|
|
|
|
let chunks = to_random_chunks(&mut rng, &new_text);
|
|
let new_text_chunks = stream::iter(chunks.iter().enumerate().map(|(index, chunk)| {
|
|
Ok(EditParserEvent::NewTextChunk {
|
|
chunk: chunk.clone(),
|
|
done: index == chunks.len() - 1,
|
|
})
|
|
}));
|
|
let reindented_chunks = EditAgent::reindent_new_text_chunks(delta, new_text_chunks)
|
|
.collect::<Vec<_>>()
|
|
.await;
|
|
let actual_reindented_text = reindented_chunks
|
|
.into_iter()
|
|
.collect::<Result<String>>()
|
|
.unwrap();
|
|
let expected_reindented_text = new_text
|
|
.split('\n')
|
|
.map(|line| {
|
|
if let Some(ix) = line.find(|c| c != ' ') {
|
|
let new_indent = cmp::max(0, ix as isize + delta.len()) as usize;
|
|
format!("{}{}", " ".repeat(new_indent), &line[ix..])
|
|
} else {
|
|
line.to_string()
|
|
}
|
|
})
|
|
.collect::<Vec<_>>()
|
|
.join("\n");
|
|
assert_eq!(actual_reindented_text, expected_reindented_text);
|
|
}
|
|
|
|
fn to_random_chunks(rng: &mut StdRng, input: &str) -> Vec<String> {
|
|
let chunk_count = rng.gen_range(1..=cmp::min(input.len(), 50));
|
|
let mut chunk_indices = (0..input.len()).choose_multiple(rng, chunk_count);
|
|
chunk_indices.sort();
|
|
chunk_indices.push(input.len());
|
|
|
|
let mut chunks = Vec::new();
|
|
let mut last_ix = 0;
|
|
for chunk_ix in chunk_indices {
|
|
chunks.push(input[last_ix..chunk_ix].to_string());
|
|
last_ix = chunk_ix;
|
|
}
|
|
chunks
|
|
}
|
|
|
|
fn simulate_llm_output(
|
|
agent: &EditAgent,
|
|
output: &str,
|
|
rng: &mut StdRng,
|
|
cx: &mut TestAppContext,
|
|
) {
|
|
let executor = cx.executor();
|
|
let chunks = to_random_chunks(rng, output);
|
|
let model = agent.model.clone();
|
|
cx.background_spawn(async move {
|
|
for chunk in chunks {
|
|
executor.simulate_random_delay().await;
|
|
model.as_fake().stream_last_completion_response(chunk);
|
|
}
|
|
model.as_fake().end_last_completion_stream();
|
|
})
|
|
.detach();
|
|
}
|
|
|
|
async fn init_test(cx: &mut TestAppContext) -> EditAgent {
|
|
cx.update(settings::init);
|
|
cx.update(Project::init_settings);
|
|
let project = Project::test(FakeFs::new(cx.executor()), [], cx).await;
|
|
let model = Arc::new(FakeLanguageModel::default());
|
|
let action_log = cx.new(|_| ActionLog::new(project.clone()));
|
|
EditAgent::new(model, project, action_log, Templates::new())
|
|
}
|
|
|
|
fn drain_events(
|
|
stream: &mut UnboundedReceiver<EditAgentOutputEvent>,
|
|
) -> Vec<EditAgentOutputEvent> {
|
|
let mut events = Vec::new();
|
|
while let Ok(Some(event)) = stream.try_next() {
|
|
events.push(event);
|
|
}
|
|
events
|
|
}
|
|
}
|