diff --git a/Cargo.lock b/Cargo.lock index 82f2169..15dba5d 100644 --- a/Cargo.lock +++ b/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", - "mac", - "utf-8", + "new_debug_unreachable", ] [[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", diff --git a/Cargo.toml b/Cargo.toml index ccf4643..844d787 100644 --- a/Cargo.toml +++ b/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"] } diff --git a/deep_research/Cargo.toml b/deep_research/Cargo.toml index 016c13b..9c4e5b8 100644 --- a/deep_research/Cargo.toml +++ b/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 } diff --git a/deep_research/src/starter.rs b/deep_research/src/core.rs similarity index 88% rename from deep_research/src/starter.rs rename to deep_research/src/core.rs index 229c465..3711c6d 100644 --- a/deep_research/src/starter.rs +++ b/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, + pub(crate) topic: Option, /// 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, + pub(crate) log_level: Option, } /// 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 { +pub(crate) async fn research(topic: &str, show_progress: bool) -> anyhow::Result { 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 - .prompt(format!("Topic: {topic}\n\nResearch notes:\n{findings}")) - .await?; + let response_stream = writer + .stream_prompt(format!("Topic: {topic}\n\nResearch notes:\n{findings}")) + .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) } diff --git a/deep_research/src/main.rs b/deep_research/src/main.rs index f139ad4..1d1b553 100644 --- a/deep_research/src/main.rs +++ b/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(()) } diff --git a/deep_research/src/review.rs b/deep_research/src/review.rs index e6e2296..54211b2 100644 --- a/deep_research/src/review.rs +++ b/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; 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"); + } +}