ec15893c9c
All checks were successful
CI / test (pull_request) Successful in 1m52s
Adds --doc (repeatable, file or directory) so the researcher can draw on user-supplied documents alongside the web: documents.rs resolves paths into embeddable text, retrieval.rs embeds them with a dedicated embedding model (nomic-embed-text, separate from the chat models used elsewhere) into an in-memory vector index and retrieves the excerpts most relevant to the topic once up front, and researcher.rs folds those excerpts into the researcher's task under the same footnote-citation scheme already used for web sources. The embedding and retrieval phases show progress the same way every other phase does — a spinner while embedding, a summary line once excerpts are retrieved, tracing spans for -l mode. Verified against a live Ollama nomic-embed-text pull and a real research round: a planted fact sheet was correctly ranked as the most relevant of several embedded documents and appeared in the researcher's task before its first turn.
80 lines
3.2 KiB
Rust
80 lines
3.2 KiB
Rust
use std::sync::{Mutex, OnceLock};
|
|
use std::time::Duration;
|
|
|
|
/// Emoji shown on the spinner line for whichever activity is currently
|
|
/// running, so the different phases (and the tool calls within the
|
|
/// research phase) stay visually distinct at a glance.
|
|
pub(crate) const RESEARCH_EMOJI: &str = "🔎";
|
|
pub(crate) const SEARCH_EMOJI: &str = "🌐";
|
|
pub(crate) const FETCH_EMOJI: &str = "📄";
|
|
pub(crate) const REVIEW_EMOJI: &str = "🧐";
|
|
pub(crate) const REJECTED_EMOJI: &str = "❌";
|
|
pub(crate) const REPORT_EMOJI: &str = "✍️";
|
|
pub(crate) const SUMMARIZE_EMOJI: &str = "🧩";
|
|
pub(crate) const EMBED_EMOJI: &str = "📚";
|
|
|
|
/// The spinner currently on screen, if any — set by `Spinner::start` and
|
|
/// cleared on drop. Tool implementations don't otherwise have a handle to
|
|
/// the active spinner (they're plain functions invoked by the model, not
|
|
/// passed one down through the tool-calling loop), so `set_activity` lets
|
|
/// them reach it here instead to reflect what they're doing — e.g. which
|
|
/// page they're fetching — on the same line.
|
|
static ACTIVE: OnceLock<Mutex<Option<indicatif::ProgressBar>>> = OnceLock::new();
|
|
|
|
fn active() -> &'static Mutex<Option<indicatif::ProgressBar>> {
|
|
ACTIVE.get_or_init(|| Mutex::new(None))
|
|
}
|
|
|
|
/// A terminal spinner for a research phase, shown only when logging is off —
|
|
/// with logging on, the trace output already tells the user something is
|
|
/// happening, and interleaving both would just be noisy. Clearing on drop
|
|
/// means call sites don't need an explicit "stop" at every early return.
|
|
pub(crate) struct Spinner(Option<indicatif::ProgressBar>);
|
|
|
|
impl Spinner {
|
|
pub(crate) fn start(enabled: bool, message: impl Into<String>) -> Self {
|
|
if !enabled {
|
|
return Self(None);
|
|
}
|
|
|
|
let bar = indicatif::ProgressBar::new_spinner();
|
|
bar.enable_steady_tick(Duration::from_millis(100));
|
|
bar.set_style(
|
|
indicatif::ProgressStyle::with_template("{spinner:.cyan} {msg}")
|
|
.expect("static template is valid"),
|
|
);
|
|
bar.set_message(message.into());
|
|
|
|
*active().lock().expect("spinner mutex poisoned") = Some(bar.clone());
|
|
|
|
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 {
|
|
fn drop(&mut self) {
|
|
self.stop();
|
|
}
|
|
}
|
|
|
|
/// Updates the message of whatever spinner is currently running, if any.
|
|
/// A no-op when progress display is off (no spinner was ever started, so
|
|
/// `active()` stays empty) or between phases (the previous `Spinner` has
|
|
/// already dropped and cleared it).
|
|
pub(crate) fn set_activity(message: impl Into<String>) {
|
|
if let Some(bar) = active().lock().expect("spinner mutex poisoned").as_ref() {
|
|
bar.set_message(message.into());
|
|
}
|
|
}
|