Compare commits
No commits in common. "351cc79e573d8d4b00ee70a9826740e1d75b66dc" and "34c699df0fc39b91cbba94312d8f9ced41064273" have entirely different histories.
351cc79e57
...
34c699df0f
3 changed files with 12 additions and 29 deletions
|
|
@ -171,19 +171,20 @@ async fn write_report(
|
||||||
)
|
)
|
||||||
.build();
|
.build();
|
||||||
|
|
||||||
|
// Drop the spinner before streaming starts: report text is about to print
|
||||||
|
// to the same terminal line, so the two must not race over stdout.
|
||||||
let spinner = Spinner::start(show_progress, format!("{REPORT_EMOJI} Writing report..."));
|
let spinner = Spinner::start(show_progress, format!("{REPORT_EMOJI} Writing report..."));
|
||||||
let response_stream = writer
|
let response_stream = writer
|
||||||
.stream_prompt(format!("Topic: {topic}\n\nResearch notes:\n{findings}"))
|
.stream_prompt(format!("Topic: {topic}\n\nResearch notes:\n{findings}"))
|
||||||
.await;
|
.await;
|
||||||
|
drop(spinner);
|
||||||
|
|
||||||
// Locked once for the whole stream rather than per chunk (as print!
|
// Locked once for the whole stream rather than per chunk (as print!
|
||||||
// would do internally) — chunks arrive in a tight loop, so re-acquiring
|
// would do internally) — chunks arrive in a tight loop, so re-acquiring
|
||||||
// the lock on every one adds up. The spinner keeps running until the
|
// the lock on every one adds up.
|
||||||
// stream's first chunk arrives, so the terminal stays covered through
|
|
||||||
// the gap between sending the prompt and generation actually starting.
|
|
||||||
let stdout = std::io::stdout();
|
let stdout = std::io::stdout();
|
||||||
let mut handle = stdout.lock();
|
let mut handle = stdout.lock();
|
||||||
let report = write_text_stream(response_stream, &mut handle, spinner).await?;
|
let report = write_text_stream(response_stream, &mut handle).await?;
|
||||||
writeln!(handle)?;
|
writeln!(handle)?;
|
||||||
|
|
||||||
Ok(report)
|
Ok(report)
|
||||||
|
|
|
||||||
|
|
@ -47,23 +47,14 @@ impl Spinner {
|
||||||
|
|
||||||
Self(Some(bar))
|
Self(Some(bar))
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Clears the spinner immediately rather than waiting for drop — for
|
|
||||||
/// callers that need it gone at a precise moment (e.g. right as the
|
|
||||||
/// first chunk of a stream is about to print on the same line) rather
|
|
||||||
/// than whenever the value happens to go out of scope. Idempotent: a
|
|
||||||
/// spinner already stopped, or one that was never enabled, does nothing.
|
|
||||||
pub(crate) fn stop(&mut self) {
|
|
||||||
if let Some(bar) = self.0.take() {
|
|
||||||
bar.finish_and_clear();
|
|
||||||
*active().lock().expect("spinner mutex poisoned") = None;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Drop for Spinner {
|
impl Drop for Spinner {
|
||||||
fn drop(&mut self) {
|
fn drop(&mut self) {
|
||||||
self.stop();
|
if let Some(bar) = &self.0 {
|
||||||
|
bar.finish_and_clear();
|
||||||
|
*active().lock().expect("spinner mutex poisoned") = None;
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -1,4 +1,3 @@
|
||||||
use crate::progress::Spinner;
|
|
||||||
use futures::{Stream, StreamExt};
|
use futures::{Stream, StreamExt};
|
||||||
use rig::agent::{MultiTurnStreamItem, StreamingError};
|
use rig::agent::{MultiTurnStreamItem, StreamingError};
|
||||||
use rig::streaming::StreamedAssistantContent;
|
use rig::streaming::StreamedAssistantContent;
|
||||||
|
|
@ -13,15 +12,9 @@ use std::io::Write;
|
||||||
/// caller controls the lock's lifetime: locking once around a whole report
|
/// caller controls the lock's lifetime: locking once around a whole report
|
||||||
/// (as `write_report` does) avoids re-acquiring it on every chunk, the way
|
/// (as `write_report` does) avoids re-acquiring it on every chunk, the way
|
||||||
/// `print!` would.
|
/// `print!` would.
|
||||||
///
|
|
||||||
/// `spinner` stays up until the stream actually produces its first item,
|
|
||||||
/// covering the gap between the prompt being sent and generation starting
|
|
||||||
/// (otherwise the terminal would go blank for however long that takes)
|
|
||||||
/// rather than being dropped by the caller before this is even called.
|
|
||||||
pub(crate) async fn write_text_stream<R>(
|
pub(crate) async fn write_text_stream<R>(
|
||||||
mut stream: impl Stream<Item = Result<MultiTurnStreamItem<R>, StreamingError>> + Unpin,
|
mut stream: impl Stream<Item = Result<MultiTurnStreamItem<R>, StreamingError>> + Unpin,
|
||||||
writer: &mut impl Write,
|
writer: &mut impl Write,
|
||||||
mut spinner: Spinner,
|
|
||||||
) -> anyhow::Result<String>
|
) -> anyhow::Result<String>
|
||||||
where
|
where
|
||||||
R: Clone,
|
R: Clone,
|
||||||
|
|
@ -29,8 +22,6 @@ where
|
||||||
let mut text = String::new();
|
let mut text = String::new();
|
||||||
|
|
||||||
while let Some(chunk) = stream.next().await {
|
while let Some(chunk) = stream.next().await {
|
||||||
spinner.stop();
|
|
||||||
|
|
||||||
if let MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::Text(chunk)) =
|
if let MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::Text(chunk)) =
|
||||||
chunk?
|
chunk?
|
||||||
{
|
{
|
||||||
|
|
@ -64,7 +55,7 @@ mod tests {
|
||||||
let items = vec![text_item("Hello, "), text_item("world!")];
|
let items = vec![text_item("Hello, "), text_item("world!")];
|
||||||
let mut written = Vec::new();
|
let mut written = Vec::new();
|
||||||
|
|
||||||
let accumulated = write_text_stream(stream::iter(items), &mut written, Spinner::start(false, ""))
|
let accumulated = write_text_stream(stream::iter(items), &mut written)
|
||||||
.await
|
.await
|
||||||
.unwrap();
|
.unwrap();
|
||||||
|
|
||||||
|
|
@ -81,7 +72,7 @@ mod tests {
|
||||||
let items = vec![text_item("kept"), final_item];
|
let items = vec![text_item("kept"), final_item];
|
||||||
let mut written = Vec::new();
|
let mut written = Vec::new();
|
||||||
|
|
||||||
let accumulated = write_text_stream(stream::iter(items), &mut written, Spinner::start(false, ""))
|
let accumulated = write_text_stream(stream::iter(items), &mut written)
|
||||||
.await
|
.await
|
||||||
.unwrap();
|
.unwrap();
|
||||||
|
|
||||||
|
|
@ -97,7 +88,7 @@ mod tests {
|
||||||
let items = vec![text_item("kept"), Err(error)];
|
let items = vec![text_item("kept"), Err(error)];
|
||||||
let mut written = Vec::new();
|
let mut written = Vec::new();
|
||||||
|
|
||||||
let result = write_text_stream(stream::iter(items), &mut written, Spinner::start(false, "")).await;
|
let result = write_text_stream(stream::iter(items), &mut written).await;
|
||||||
|
|
||||||
assert!(result.is_err());
|
assert!(result.is_err());
|
||||||
assert_eq!(String::from_utf8(written).unwrap(), "kept");
|
assert_eq!(String::from_utf8(written).unwrap(), "kept");
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue