Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

assistant2: Stream in completion text (#21182)

Browse source 
This PR makes it so that the completion text streams into the message
list rather than being buffered until the end.

Release Notes:

- N/A
This commit is contained in:
Marshall Bowers 2024-11-25 16:13:27 -05:00 • committed by GitHub
parent 91a565f5fa
commit 9ee1aba80a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 75 additions and 51 deletions
Download patch file Download diff file Expand all files Collapse all files

51
crates/assistant2/src/message_editor.rs
View file

@ -1,14 +1,11 @@
use editor::{Editor, EditorElement, EditorStyle};
use futures::StreamExt;
use gpui::{AppContext, Model, TextStyle, View};
use language_model::{
LanguageModelCompletionEvent, LanguageModelRegistry, LanguageModelRequest,
LanguageModelRequestMessage, MessageContent, Role, StopReason,
LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage, MessageContent, Role,
};
use settings::Settings;
use theme::ThemeSettings;
use ui::{prelude::*, ButtonLike, ElevationIndex, KeyBinding};
use util::ResultExt;
use crate::thread::{self, Thread};
use crate::Chat;
@ -71,50 +68,8 @@ impl MessageEditor {
editor.clear(cx);
});
let task = cx.spawn(|this, mut cx| async move {
let stream = model.stream_completion(request, &cx);
let stream_completion = async {
let mut events = stream.await?;
let mut stop_reason = StopReason::EndTurn;
let mut text = String::new();
while let Some(event) = events.next().await {
let event = event?;
match event {
LanguageModelCompletionEvent::StartMessage { .. } => {}
LanguageModelCompletionEvent::Stop(reason) => {
stop_reason = reason;
}
LanguageModelCompletionEvent::Text(chunk) => {
text.push_str(&chunk);
}
LanguageModelCompletionEvent::ToolUse(_tool_use) => {}
}
smol::future::yield_now().await;
}
anyhow::Ok((stop_reason, text))
};
let result = stream_completion.await;
this.update(&mut cx, |this, cx| {
if let Some((_stop_reason, text)) = result.log_err() {
this.thread.update(cx, |thread, _cx| {
thread.messages.push(thread::Message {
role: Role::Assistant,
text,
});
});
}
})
.ok();
});
self.thread.update(cx, |thread, _cx| {
thread.pending_completion_tasks.push(task);
self.thread.update(cx, |thread, cx| {
thread.stream_completion(request, model, cx)
});
None