diff --git a/deep_research/src/core.rs b/deep_research/src/core.rs index 50ae6e4..3711c6d 100644 --- a/deep_research/src/core.rs +++ b/deep_research/src/core.rs @@ -1,12 +1,11 @@ use crate::progress::Spinner; use crate::review::{self, Review}; +use crate::stream::write_text_stream; use crate::tools::{FetchPage, SearchWeb}; use clap::Parser; -use futures::StreamExt; -use rig::agent::MultiTurnStreamItem; use rig::client::{AgentClientExt, Nothing}; use rig::providers::ollama; -use rig::streaming::{StreamedAssistantContent, StreamingPrompt}; +use rig::streaming::StreamingPrompt; use std::io::Write; /// The tool-calling research loop needs to reliably decide what to search @@ -171,34 +170,17 @@ async fn write_report( // 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, "Writing report..."); - let mut response_stream = writer + let response_stream = writer .stream_prompt(format!("Topic: {topic}\n\nResearch notes:\n{findings}")) .await; drop(spinner); - let mut report = String::new(); // Locked once for the whole stream rather than per chunk (as print! // would do internally) — chunks arrive in a tight loop, so re-acquiring // the lock on every one adds up. let stdout = std::io::stdout(); let mut handle = stdout.lock(); - - while let Some(chunk) = response_stream.next().await { - match chunk? { - MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::Text(text)) => { - report.push_str(&text.text); - // Terminal stdout is line-buffered, so a flush is needed here — - // otherwise a chunk without a trailing newline sits in the - // buffer instead of appearing as it streams in. - write!(handle, "{}", text.text)?; - handle.flush()?; - } - _ => continue, - } - } - // `handle` still holds stdout's lock here, so this must go through it - // rather than `println!` — reacquiring the same lock from this thread - // would deadlock. + let report = write_text_stream(response_stream, &mut handle).await?; writeln!(handle)?; Ok(report) diff --git a/deep_research/src/main.rs b/deep_research/src/main.rs index 6acaf89..1d1b553 100644 --- a/deep_research/src/main.rs +++ b/deep_research/src/main.rs @@ -3,6 +3,7 @@ use clap::Parser; mod progress; mod review; mod core; +mod stream; mod tools; pub(crate) const DEFAULT_TOPIC: &str = diff --git a/deep_research/src/stream.rs b/deep_research/src/stream.rs new file mode 100644 index 0000000..36d71cb --- /dev/null +++ b/deep_research/src/stream.rs @@ -0,0 +1,96 @@ +use futures::{Stream, StreamExt}; +use rig::agent::{MultiTurnStreamItem, StreamingError}; +use rig::streaming::StreamedAssistantContent; +use std::io::Write; + +/// Drains a multi-turn prompt stream, writing each text chunk to `writer` as +/// it arrives and returning the full accumulated text. Non-text items (tool +/// calls, reasoning, completion-call metadata, ...) are ignored — this is +/// only concerned with the assistant's prose. +/// +/// `writer` is taken by reference rather than locked internally so the +/// 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 +/// `print!` would. +pub(crate) async fn write_text_stream( + mut stream: impl Stream, StreamingError>> + Unpin, + writer: &mut impl Write, +) -> anyhow::Result +where + R: Clone, +{ + let mut text = String::new(); + + while let Some(chunk) = stream.next().await { + if let MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::Text(chunk)) = + chunk? + { + write!(writer, "{}", chunk.text)?; + writer.flush()?; + text.push_str(&chunk.text); + } + } + + Ok(text) +} + +#[cfg(test)] +mod tests { + use super::*; + use futures::stream; + use rig::completion::CompletionError; + use rig::message::Text; + + #[derive(Clone)] + struct DummyResponse; + + fn text_item(text: &str) -> Result, StreamingError> { + Ok(MultiTurnStreamItem::StreamAssistantItem( + StreamedAssistantContent::Text(Text::new(text)), + )) + } + + #[tokio::test] + async fn writes_and_accumulates_text_chunks() { + let items = vec![text_item("Hello, "), text_item("world!")]; + let mut written = Vec::new(); + + let accumulated = write_text_stream(stream::iter(items), &mut written) + .await + .unwrap(); + + assert_eq!(accumulated, "Hello, world!"); + assert_eq!(String::from_utf8(written).unwrap(), "Hello, world!"); + } + + #[tokio::test] + async fn ignores_non_text_items() { + let final_item = Ok(MultiTurnStreamItem::final_response( + rig::OneOrMany::one(rig::message::AssistantContent::text("ignored")), + rig::completion::Usage::new(), + )); + let items = vec![text_item("kept"), final_item]; + let mut written = Vec::new(); + + let accumulated = write_text_stream(stream::iter(items), &mut written) + .await + .unwrap(); + + assert_eq!(accumulated, "kept"); + assert_eq!(String::from_utf8(written).unwrap(), "kept"); + } + + #[tokio::test] + async fn propagates_stream_errors() { + let error = StreamingError::Completion(CompletionError::RequestError(Box::new( + std::io::Error::other("boom"), + ))); + let items = vec![text_item("kept"), Err(error)]; + let mut written = Vec::new(); + + let result = write_text_stream(stream::iter(items), &mut written).await; + + assert!(result.is_err()); + assert_eq!(String::from_utf8(written).unwrap(), "kept"); + } +}