Stream the report to the terminal as the writer generates it #5
7 changed files with 201 additions and 76 deletions
94
Cargo.lock
generated
94
Cargo.lock
generated
|
|
@ -995,6 +995,18 @@ dependencies = [
|
||||||
"windows-sys 0.59.0",
|
"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]]
|
[[package]]
|
||||||
name = "const-random"
|
name = "const-random"
|
||||||
version = "0.1.18"
|
version = "0.1.18"
|
||||||
|
|
@ -1154,9 +1166,9 @@ dependencies = [
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "cssparser"
|
name = "cssparser"
|
||||||
version = "0.36.0"
|
version = "0.37.0"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "dae61cf9c0abb83bd659dab65b7e4e38d8236824c85f0f804f173567bda257d2"
|
checksum = "8c9cdaae01d5ed7882b04d795e7f752f46ff52d2fa3b50a20d28c464510bba98"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"cssparser-macros",
|
"cssparser-macros",
|
||||||
"dtoa-short",
|
"dtoa-short",
|
||||||
|
|
@ -1167,9 +1179,9 @@ dependencies = [
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "cssparser-macros"
|
name = "cssparser-macros"
|
||||||
version = "0.6.1"
|
version = "0.7.0"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "13b588ba4ac1a99f7f2964d24b3d896ddc6bf847ee3855dbd4366f058cfcd331"
|
checksum = "10a2a99df6e410a8ff4245aa2006499ea662245f967cc7c0a38c83ef8eb44dbf"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"quote",
|
"quote",
|
||||||
"syn 2.0.119",
|
"syn 2.0.119",
|
||||||
|
|
@ -1922,7 +1934,8 @@ dependencies = [
|
||||||
"anyhow",
|
"anyhow",
|
||||||
"chrono",
|
"chrono",
|
||||||
"clap",
|
"clap",
|
||||||
"indicatif",
|
"futures",
|
||||||
|
"indicatif 0.18.6",
|
||||||
"reqwest 0.13.4",
|
"reqwest 0.13.4",
|
||||||
"rig",
|
"rig",
|
||||||
"schemars 1.2.2",
|
"schemars 1.2.2",
|
||||||
|
|
@ -2117,9 +2130,9 @@ checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555"
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "ego-tree"
|
name = "ego-tree"
|
||||||
version = "0.10.0"
|
version = "0.11.0"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "b2972feb8dffe7bc8c5463b1dacda1b0dfbed3710e50f977d965429692d74cd8"
|
checksum = "b04dc5a38e4f151a79d9f2451ae6037fb6eaf5cba34771f44781f80e508498e3"
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "either"
|
name = "either"
|
||||||
|
|
@ -2381,16 +2394,6 @@ version = "2.0.0"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "e6d5a32815ae3f33302d95fdcb2ce17862f8c65363dcfd29360480ba1001fc9c"
|
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]]
|
[[package]]
|
||||||
name = "futures"
|
name = "futures"
|
||||||
version = "0.3.34"
|
version = "0.3.34"
|
||||||
|
|
@ -2685,7 +2688,7 @@ checksum = "629d8f3bbeda9d148036d6b0de0a3ab947abd08ce90626327fc3547a49d59d97"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"dirs",
|
"dirs",
|
||||||
"http",
|
"http",
|
||||||
"indicatif",
|
"indicatif 0.17.11",
|
||||||
"libc",
|
"libc",
|
||||||
"log",
|
"log",
|
||||||
"rand 0.9.5",
|
"rand 0.9.5",
|
||||||
|
|
@ -2699,9 +2702,9 @@ dependencies = [
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "html5ever"
|
name = "html5ever"
|
||||||
version = "0.36.1"
|
version = "0.39.0"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "6452c4751a24e1b99c3260d505eaeee76a050573e61f30ac2c924ddc7236f01e"
|
checksum = "46a1761807faccc9a19e86944bbf40610014066306f96edcdedc2fb714bcb7b8"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"log",
|
"log",
|
||||||
"markup5ever",
|
"markup5ever",
|
||||||
|
|
@ -3027,13 +3030,26 @@ version = "0.17.11"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "183b3088984b400f4cfac3620d5e076c84da5364016b4f49473de574b2586235"
|
checksum = "183b3088984b400f4cfac3620d5e076c84da5364016b4f49473de574b2586235"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"console",
|
"console 0.15.11",
|
||||||
"number_prefix",
|
"number_prefix",
|
||||||
"portable-atomic",
|
"portable-atomic",
|
||||||
"unicode-width",
|
"unicode-width",
|
||||||
"web-time",
|
"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]]
|
[[package]]
|
||||||
name = "interpolate_name"
|
name = "interpolate_name"
|
||||||
version = "0.2.4"
|
version = "0.2.4"
|
||||||
|
|
@ -4073,12 +4089,6 @@ dependencies = [
|
||||||
"twox-hash",
|
"twox-hash",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
|
||||||
name = "mac"
|
|
||||||
version = "0.1.1"
|
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
||||||
checksum = "c41e0c4fef86961ac6d6f8a82609f55f31b05e4fce149ac5710e439df7619ba4"
|
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "macro_rules_attribute"
|
name = "macro_rules_attribute"
|
||||||
version = "0.2.3"
|
version = "0.2.3"
|
||||||
|
|
@ -4097,9 +4107,9 @@ checksum = "fc04a4c58212d57930a24bf47d3fa87485264a3a054e9c10e042eb373573ad3c"
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "markup5ever"
|
name = "markup5ever"
|
||||||
version = "0.36.1"
|
version = "0.39.0"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "6c3294c4d74d0742910f8c7b466f44dda9eb2d5742c1e430138df290a1e8451c"
|
checksum = "7122d987ec5f704ee56f6e5b41a7d93722e9aae27ae07cafa4036c4d3f9757de"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"log",
|
"log",
|
||||||
"tendril",
|
"tendril",
|
||||||
|
|
@ -5790,9 +5800,9 @@ checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49"
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "scraper"
|
name = "scraper"
|
||||||
version = "0.25.0"
|
version = "0.27.0"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "93cecd86d6259499c844440546d02f55f3e17bd286e529e48d1f9f67e92315cb"
|
checksum = "bdd0be4d296f048bfb06dd01bbc80ef789ddd2e55583e8d2e6b804942abfabc2"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"cssparser",
|
"cssparser",
|
||||||
"ego-tree",
|
"ego-tree",
|
||||||
|
|
@ -5828,9 +5838,9 @@ dependencies = [
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "selectors"
|
name = "selectors"
|
||||||
version = "0.33.0"
|
version = "0.38.0"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "feef350c36147532e1b79ea5c1f3791373e61cbd9a6a2615413b3807bb164fb7"
|
checksum = "8adfa1c298912827b8a28b223b3b874357397ae706e6190acd9bf28cee99114d"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"bitflags 2.13.1",
|
"bitflags 2.13.1",
|
||||||
"cssparser",
|
"cssparser",
|
||||||
|
|
@ -6414,13 +6424,11 @@ dependencies = [
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "tendril"
|
name = "tendril"
|
||||||
version = "0.4.3"
|
version = "0.5.1"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "d24a120c5fc464a3458240ee02c299ebcb9d67b5249c8848b09d639dca8d7bb0"
|
checksum = "5fed54709c5b3a53d09bb1c113ea4f5ceafd1e772ddcb0030a82e1d56c087b08"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"futf",
|
"new_debug_unreachable",
|
||||||
"mac",
|
|
||||||
"utf-8",
|
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
|
|
@ -6920,6 +6928,12 @@ version = "0.1.1"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e"
|
checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e"
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "unit-prefix"
|
||||||
|
version = "0.5.2"
|
||||||
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
|
checksum = "81e544489bf3d8ef66c953931f56617f423cd4b5494be343d9b9d3dda037b9a3"
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "unsafe-libyaml"
|
name = "unsafe-libyaml"
|
||||||
version = "0.2.11"
|
version = "0.2.11"
|
||||||
|
|
@ -6989,9 +7003,9 @@ checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821"
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "uuid"
|
name = "uuid"
|
||||||
version = "1.24.0"
|
version = "1.24.1"
|
||||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||||
checksum = "bf3923a6f5c4c6382e0b653c4117f48d631ea17f38ed86e2a828e6f7412f5239"
|
checksum = "2cefc03fd367c0c6d4305de1b312cf00248c4114f4a0418ce6a6af769e3b0bd9"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"getrandom 0.4.3",
|
"getrandom 0.4.3",
|
||||||
"js-sys",
|
"js-sys",
|
||||||
|
|
|
||||||
|
|
@ -4,6 +4,7 @@ members = ["deep_research", "swear_cleanup"]
|
||||||
|
|
||||||
[workspace.dependencies]
|
[workspace.dependencies]
|
||||||
anyhow = "1.0.104"
|
anyhow = "1.0.104"
|
||||||
|
futures = "0.3"
|
||||||
reqwest = "0.13.4"
|
reqwest = "0.13.4"
|
||||||
rig = "0.41.0"
|
rig = "0.41.0"
|
||||||
serde = { version = "1.0.229", features = ["derive"] }
|
serde = { version = "1.0.229", features = ["derive"] }
|
||||||
|
|
|
||||||
|
|
@ -7,11 +7,12 @@ edition = "2024"
|
||||||
anyhow = { workspace = true }
|
anyhow = { workspace = true }
|
||||||
chrono = "0.4.45"
|
chrono = "0.4.45"
|
||||||
clap = { version = "4", features = ["derive"] }
|
clap = { version = "4", features = ["derive"] }
|
||||||
indicatif = "0.17"
|
futures = { workspace = true }
|
||||||
|
indicatif = "0.18.6"
|
||||||
reqwest = { workspace = true, features = ["query"] }
|
reqwest = { workspace = true, features = ["query"] }
|
||||||
rig = { workspace = true }
|
rig = { workspace = true }
|
||||||
schemars = "1"
|
schemars = "1"
|
||||||
scraper = "0.25"
|
scraper = "0.27"
|
||||||
serde = { workspace = true }
|
serde = { workspace = true }
|
||||||
tokio = { workspace = true }
|
tokio = { workspace = true }
|
||||||
tracing = { workspace = true }
|
tracing = { workspace = true }
|
||||||
|
|
|
||||||
|
|
@ -1,10 +1,12 @@
|
||||||
use crate::progress::Spinner;
|
use crate::progress::Spinner;
|
||||||
use crate::review::{self, Review};
|
use crate::review::{self, Review};
|
||||||
|
use crate::stream::write_text_stream;
|
||||||
use crate::tools::{FetchPage, SearchWeb};
|
use crate::tools::{FetchPage, SearchWeb};
|
||||||
use clap::Parser;
|
use clap::Parser;
|
||||||
use rig::client::{AgentClientExt, Nothing};
|
use rig::client::{AgentClientExt, Nothing};
|
||||||
use rig::completion::Prompt;
|
|
||||||
use rig::providers::ollama;
|
use rig::providers::ollama;
|
||||||
|
use rig::streaming::StreamingPrompt;
|
||||||
|
use std::io::Write;
|
||||||
|
|
||||||
/// The tool-calling research loop needs to reliably decide what to search
|
/// 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 —
|
/// 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.
|
/// the reviewer can never be satisfied with.
|
||||||
const MAX_RESEARCH_ROUNDS: usize = 3;
|
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
|
/// 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
|
/// gathers and cross-checks web evidence, a reviewer agent gates it, and a
|
||||||
/// writer agent turns approved findings into a structured report.
|
/// writer agent turns approved findings into a structured report.
|
||||||
#[derive(Parser)]
|
#[derive(Parser)]
|
||||||
#[command(name = "doubleo7-research", version, about)]
|
#[command(name = "doubleo7-research", version, about)]
|
||||||
struct Cli {
|
pub(crate) struct Cli {
|
||||||
/// Research topic to investigate
|
/// 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
|
/// Emit logs at this level (off by default; passing this also enables a
|
||||||
/// progress spinner to switch off, since the logs already show progress)
|
/// progress spinner to switch off, since the logs already show progress)
|
||||||
#[arg(short = 'l', long, value_name = "LEVEL")]
|
#[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)
|
/// Only initializes a subscriber (and thus produces any log output at all)
|
||||||
/// when the caller opted in via `--log-level` — otherwise tracing's macros
|
/// when the caller opted in via `--log-level` — otherwise tracing's macros
|
||||||
/// are no-ops, leaving the terminal clean for the spinner.
|
/// 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()
|
tracing_subscriber::fmt()
|
||||||
.with_env_filter(
|
.with_env_filter(
|
||||||
tracing_subscriber::EnvFilter::try_from_default_env()
|
tracing_subscriber::EnvFilter::try_from_default_env()
|
||||||
|
|
@ -51,31 +51,11 @@ fn initialize_observability(log_level: tracing::Level) {
|
||||||
.init();
|
.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*,
|
/// 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
|
/// 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
|
/// reviewer's feedback folded in until it approves or the round budget runs
|
||||||
/// out, then writing the report from whatever the last pass produced.
|
/// 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 client = ollama::Client::new(Nothing)?;
|
||||||
|
|
||||||
let mut findings = String::new();
|
let mut findings = String::new();
|
||||||
|
|
@ -187,11 +167,21 @@ 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, "Writing report...");
|
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);
|
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)
|
Ok(report)
|
||||||
}
|
}
|
||||||
|
|
@ -1,9 +1,32 @@
|
||||||
|
use clap::Parser;
|
||||||
|
|
||||||
mod progress;
|
mod progress;
|
||||||
mod review;
|
mod review;
|
||||||
mod starter;
|
mod core;
|
||||||
|
mod stream;
|
||||||
mod tools;
|
mod tools;
|
||||||
|
|
||||||
|
pub(crate) const DEFAULT_TOPIC: &str =
|
||||||
|
"What are the latest advances in running large language models locally, on consumer hardware?";
|
||||||
|
|
||||||
#[tokio::main]
|
#[tokio::main]
|
||||||
async fn main() -> anyhow::Result<()> {
|
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(())
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -1,6 +1,6 @@
|
||||||
use rig::client::AgentClientExt;
|
use rig::client::AgentClientExt;
|
||||||
use rig::providers::ollama;
|
use rig::providers::ollama;
|
||||||
use schemars::JsonSchema;
|
use rig::schemars::JsonSchema;
|
||||||
use serde::{Deserialize, Serialize};
|
use serde::{Deserialize, Serialize};
|
||||||
|
|
||||||
use crate::progress::Spinner;
|
use crate::progress::Spinner;
|
||||||
|
|
|
||||||
96
deep_research/src/stream.rs
Normal file
96
deep_research/src/stream.rs
Normal file
|
|
@ -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");
|
||||||
|
}
|
||||||
|
}
|
||||||
Loading…
Reference in a new issue