Merge pull request 'Stream the report to the terminal as the writer generates it' (#5) from worktree-deep-research-streaming into master

Reviewed-on: #5
Reviewed-by: Austin Schaefer <austin.schaefer@mailo.eu>

Cargo.lock

@@ -995,6 +995,18 @@ dependencies = [
  "windows-sys 0.59.0",
 ]
 
+[[package]]
+name = "console"
+version = "0.16.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4fe5f465a4f6fee88fad41b85d990f84c835335e85b5d9e6e63e0d06d28cba7c"
+dependencies = [
+ "encode_unicode",
+ "libc",
+ "unicode-width",
+ "windows-sys 0.61.2",
+]
+
 [[package]]
 name = "const-random"
 version = "0.1.18"
@@ -1154,9 +1166,9 @@ dependencies = [
 
 [[package]]
 name = "cssparser"
-version = "0.36.0"
+version = "0.37.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "dae61cf9c0abb83bd659dab65b7e4e38d8236824c85f0f804f173567bda257d2"
+checksum = "8c9cdaae01d5ed7882b04d795e7f752f46ff52d2fa3b50a20d28c464510bba98"
 dependencies = [
  "cssparser-macros",
  "dtoa-short",
@@ -1167,9 +1179,9 @@ dependencies = [
 
 [[package]]
 name = "cssparser-macros"
-version = "0.6.1"
+version = "0.7.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "13b588ba4ac1a99f7f2964d24b3d896ddc6bf847ee3855dbd4366f058cfcd331"
+checksum = "10a2a99df6e410a8ff4245aa2006499ea662245f967cc7c0a38c83ef8eb44dbf"
 dependencies = [
  "quote",
  "syn 2.0.119",
@@ -1922,7 +1934,8 @@ dependencies = [
  "anyhow",
  "chrono",
  "clap",
- "indicatif",
+ "futures",
+ "indicatif 0.18.6",
  "reqwest 0.13.4",
  "rig",
  "schemars 1.2.2",
@@ -2117,9 +2130,9 @@ checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555"
 
 [[package]]
 name = "ego-tree"
-version = "0.10.0"
+version = "0.11.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b2972feb8dffe7bc8c5463b1dacda1b0dfbed3710e50f977d965429692d74cd8"
+checksum = "b04dc5a38e4f151a79d9f2451ae6037fb6eaf5cba34771f44781f80e508498e3"
 
 [[package]]
 name = "either"
@@ -2381,16 +2394,6 @@ version = "2.0.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e6d5a32815ae3f33302d95fdcb2ce17862f8c65363dcfd29360480ba1001fc9c"
 
-[[package]]
-name = "futf"
-version = "0.1.5"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "df420e2e84819663797d1ec6544b13c5be84629e7bb00dc960d6917db2987843"
-dependencies = [
- "mac",
- "new_debug_unreachable",
-]
-
 [[package]]
 name = "futures"
 version = "0.3.34"
@@ -2685,7 +2688,7 @@ checksum = "629d8f3bbeda9d148036d6b0de0a3ab947abd08ce90626327fc3547a49d59d97"
 dependencies = [
  "dirs",
  "http",
- "indicatif",
+ "indicatif 0.17.11",
  "libc",
  "log",
  "rand 0.9.5",
@@ -2699,9 +2702,9 @@ dependencies = [
 
 [[package]]
 name = "html5ever"
-version = "0.36.1"
+version = "0.39.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6452c4751a24e1b99c3260d505eaeee76a050573e61f30ac2c924ddc7236f01e"
+checksum = "46a1761807faccc9a19e86944bbf40610014066306f96edcdedc2fb714bcb7b8"
 dependencies = [
  "log",
  "markup5ever",
@@ -3027,13 +3030,26 @@ version = "0.17.11"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "183b3088984b400f4cfac3620d5e076c84da5364016b4f49473de574b2586235"
 dependencies = [
- "console",
+ "console 0.15.11",
  "number_prefix",
  "portable-atomic",
  "unicode-width",
  "web-time",
 ]
 
+[[package]]
+name = "indicatif"
+version = "0.18.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9433806cd6b4ec1aba79c021c7e4c58fb4c3b9977c085062e611ac929998fb0c"
+dependencies = [
+ "console 0.16.4",
+ "portable-atomic",
+ "unicode-width",
+ "unit-prefix",
+ "web-time",
+]
+
 [[package]]
 name = "interpolate_name"
 version = "0.2.4"
@@ -4073,12 +4089,6 @@ dependencies = [
  "twox-hash",
 ]
 
-[[package]]
-name = "mac"
-version = "0.1.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c41e0c4fef86961ac6d6f8a82609f55f31b05e4fce149ac5710e439df7619ba4"
-
 [[package]]
 name = "macro_rules_attribute"
 version = "0.2.3"
@@ -4097,9 +4107,9 @@ checksum = "fc04a4c58212d57930a24bf47d3fa87485264a3a054e9c10e042eb373573ad3c"
 
 [[package]]
 name = "markup5ever"
-version = "0.36.1"
+version = "0.39.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6c3294c4d74d0742910f8c7b466f44dda9eb2d5742c1e430138df290a1e8451c"
+checksum = "7122d987ec5f704ee56f6e5b41a7d93722e9aae27ae07cafa4036c4d3f9757de"
 dependencies = [
  "log",
  "tendril",
@@ -5790,9 +5800,9 @@ checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49"
 
 [[package]]
 name = "scraper"
-version = "0.25.0"
+version = "0.27.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "93cecd86d6259499c844440546d02f55f3e17bd286e529e48d1f9f67e92315cb"
+checksum = "bdd0be4d296f048bfb06dd01bbc80ef789ddd2e55583e8d2e6b804942abfabc2"
 dependencies = [
  "cssparser",
  "ego-tree",
@@ -5828,9 +5838,9 @@ dependencies = [
 
 [[package]]
 name = "selectors"
-version = "0.33.0"
+version = "0.38.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "feef350c36147532e1b79ea5c1f3791373e61cbd9a6a2615413b3807bb164fb7"
+checksum = "8adfa1c298912827b8a28b223b3b874357397ae706e6190acd9bf28cee99114d"
 dependencies = [
  "bitflags 2.13.1",
  "cssparser",
@@ -6414,13 +6424,11 @@ dependencies = [
 
 [[package]]
 name = "tendril"
-version = "0.4.3"
+version = "0.5.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d24a120c5fc464a3458240ee02c299ebcb9d67b5249c8848b09d639dca8d7bb0"
+checksum = "5fed54709c5b3a53d09bb1c113ea4f5ceafd1e772ddcb0030a82e1d56c087b08"
 dependencies = [
- "futf",
+ "new_debug_unreachable",
- "mac",
- "utf-8",
 ]
 
 [[package]]
@@ -6920,6 +6928,12 @@ version = "0.1.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e"
 
+[[package]]
+name = "unit-prefix"
+version = "0.5.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "81e544489bf3d8ef66c953931f56617f423cd4b5494be343d9b9d3dda037b9a3"
+
 [[package]]
 name = "unsafe-libyaml"
 version = "0.2.11"
@@ -6989,9 +7003,9 @@ checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821"
 
 [[package]]
 name = "uuid"
-version = "1.24.0"
+version = "1.24.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "bf3923a6f5c4c6382e0b653c4117f48d631ea17f38ed86e2a828e6f7412f5239"
+checksum = "2cefc03fd367c0c6d4305de1b312cf00248c4114f4a0418ce6a6af769e3b0bd9"
 dependencies = [
  "getrandom 0.4.3",
  "js-sys",

Cargo.toml

@@ -4,6 +4,7 @@ members = ["deep_research", "swear_cleanup"]
 
 [workspace.dependencies]
 anyhow = "1.0.104"
+futures = "0.3"
 reqwest = "0.13.4"
 rig = "0.41.0"
 serde = { version = "1.0.229", features = ["derive"] }

deep_research/Cargo.toml

@@ -7,11 +7,12 @@ edition = "2024"
 anyhow = { workspace = true }
 chrono = "0.4.45"
 clap = { version = "4", features = ["derive"] }
-indicatif = "0.17"
+futures = { workspace = true }
+indicatif = "0.18.6"
 reqwest = { workspace = true, features = ["query"] }
 rig = { workspace = true }
 schemars = "1"
-scraper = "0.25"
+scraper = "0.27"
 serde = { workspace = true }
 tokio = { workspace = true }
 tracing = { workspace = true }

deep_research/src/core.rs

@@ -1,10 +1,12 @@
 use crate::progress::Spinner;
 use crate::review::{self, Review};
+use crate::stream::write_text_stream;
 use crate::tools::{FetchPage, SearchWeb};
 use clap::Parser;
 use rig::client::{AgentClientExt, Nothing};
-use rig::completion::Prompt;
 use rig::providers::ollama;
+use rig::streaming::StreamingPrompt;
+use std::io::Write;
 
 /// The tool-calling research loop needs to reliably decide what to search
 /// for, when a page is worth fetching, and when it has enough evidence —
@@ -19,28 +21,26 @@ const MAX_RESEARCH_TURNS: usize = 12;
 /// the reviewer can never be satisfied with.
 const MAX_RESEARCH_ROUNDS: usize = 3;
 
-const DEFAULT_TOPIC: &str =
-    "What are the latest advances in running large language models locally, on consumer hardware?";
 
 /// Deep research agentic loop over local Gemma models: a tool-calling agent
 /// gathers and cross-checks web evidence, a reviewer agent gates it, and a
 /// writer agent turns approved findings into a structured report.
 #[derive(Parser)]
 #[command(name = "doubleo7-research", version, about)]
-struct Cli {
+pub(crate) struct Cli {
     /// Research topic to investigate
-    topic: Option<String>,
+    pub(crate) topic: Option<String>,
 
     /// Emit logs at this level (off by default; passing this also enables a
     /// progress spinner to switch off, since the logs already show progress)
     #[arg(short = 'l', long, value_name = "LEVEL")]
-    log_level: Option<tracing::Level>,
+    pub(crate) log_level: Option<tracing::Level>,
 }
 
 /// Only initializes a subscriber (and thus produces any log output at all)
 /// when the caller opted in via `--log-level` — otherwise tracing's macros
 /// are no-ops, leaving the terminal clean for the spinner.
-fn initialize_observability(log_level: tracing::Level) {
+pub(crate) fn initialize_observability(log_level: tracing::Level) {
     tracing_subscriber::fmt()
         .with_env_filter(
             tracing_subscriber::EnvFilter::try_from_default_env()
@@ -51,31 +51,11 @@ fn initialize_observability(log_level: tracing::Level) {
         .init();
 }
 
-pub(crate) async fn start() -> anyhow::Result<()> {
-    let cli = Cli::parse();
-
-    let show_progress = match cli.log_level {
-        Some(level) => {
-            initialize_observability(level);
-            false
-        }
-        None => true,
-    };
-
-    let topic = cli.topic.unwrap_or_else(|| DEFAULT_TOPIC.to_string());
-
-    let report = research(&topic, show_progress).await?;
-
-    println!("{report}");
-
-    Ok(())
-}
-
 /// The least-agentic shape that fits: a plain Rust loop putting *this code*,
 /// not a model, in charge of when to stop — re-running research with the
 /// reviewer's feedback folded in until it approves or the round budget runs
 /// out, then writing the report from whatever the last pass produced.
-async fn research(topic: &str, show_progress: bool) -> anyhow::Result<String> {
+pub(crate) async fn research(topic: &str, show_progress: bool) -> anyhow::Result<String> {
     let client = ollama::Client::new(Nothing)?;
 
     let mut findings = String::new();
@@ -187,11 +167,21 @@ async fn write_report(
         )
         .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, "Writing report...");
-    let report = writer
+    let response_stream = writer
-        .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!
+    // 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();
+    let report = write_text_stream(response_stream, &mut handle).await?;
+    writeln!(handle)?;
+
     Ok(report)
 }

deep_research/src/main.rs

@@ -1,9 +1,32 @@
+use clap::Parser;
+
 mod progress;
 mod review;
-mod starter;
+mod core;
+mod stream;
 mod tools;
 
+pub(crate) const DEFAULT_TOPIC: &str =
+    "What are the latest advances in running large language models locally, on consumer hardware?";
+
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
-    starter::start().await
+    let cli = core::Cli::parse();
+
+    let show_progress = match cli.log_level {
+        Some(level) => {
+            core::initialize_observability(level);
+            false
+        }
+        None => true,
+    };
+
+    let topic = cli.topic.unwrap_or_else(|| DEFAULT_TOPIC.to_string());
+
+    // The report streams to stdout as the writer generates it, so nothing
+    // left to print here — the return value only matters to callers that
+    // embed `research` rather than running it as this binary.
+    core::research(&topic, show_progress).await?;
+
+    Ok(())
 }

deep_research/src/review.rs

@@ -1,6 +1,6 @@
 use rig::client::AgentClientExt;
 use rig::providers::ollama;
-use schemars::JsonSchema;
+use rig::schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 
 use crate::progress::Spinner;

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<R>(
+    mut stream: impl Stream<Item = Result<MultiTurnStreamItem<R>, StreamingError>> + Unpin,
+    writer: &mut impl Write,
+) -> anyhow::Result<String>
+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<MultiTurnStreamItem<DummyResponse>, 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");
+    }
+}