Pins rig to the newest published crates.io release (0.41.0) instead of the git main branch, and adapts swear_cleanup's revise.rs to that release's API (OneOrMany::first() returns T directly, raw_completion folded into CompletionResponse::raw_response). The research agent (gemma4:26b) drives rig's AgentRunner tool-calling loop with two lean #[rig::tool_macro] tools — a DuckDuckGo HTML search and a page-text fetcher — to gather and cross-check findings. A second agent (gemma4-e4b) turns those findings into a structured report; the smaller/faster model suffices there since it's reformatting already- digested notes rather than doing multi-step research reasoning.
60 lines
2.3 KiB
Rust
60 lines
2.3 KiB
Rust
use anyhow;
|
|
use rig::client::{CompletionClient, Nothing};
|
|
use rig::providers::llamafile::LlamafileExt;
|
|
use rig::providers::openai::GenericCompletionModel;
|
|
use rig::providers::{llamafile, ollama};
|
|
use crate::swear_cleanup::{revise, server};
|
|
|
|
/// Respects RUST_LOG if the shell sets one (e.g. `RUST_LOG=debug cargo run`),
|
|
/// otherwise defaults to "info" — the level Rig's own completion spans use.
|
|
/// `with_span_events(CLOSE)` is the part that actually makes anything print:
|
|
/// Rig records fields (model, token usage, ...) onto the span itself rather
|
|
/// than emitting log events, so without this, fmt's default event-only
|
|
/// logging shows nothing even though tracing is "on".
|
|
/// Logs go to stderr, not stdout — keeps stdout reserved for the actual
|
|
/// result (the final `println!` below), so it stays pipeable/parseable
|
|
/// without log lines mixed in.
|
|
fn initialize_observability() {
|
|
tracing_subscriber::fmt()
|
|
.with_env_filter(
|
|
tracing_subscriber::EnvFilter::try_from_default_env()
|
|
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("debug")),
|
|
)
|
|
.with_span_events(tracing_subscriber::fmt::format::FmtSpan::CLOSE)
|
|
.with_writer(std::io::stderr)
|
|
.init();
|
|
}
|
|
|
|
/// Execute the main functionality of this demo.
|
|
pub(crate) async fn run() -> anyhow::Result<()> {
|
|
initialize_observability();
|
|
|
|
// Independent setup steps (talk to unrelated backends, no data dependency) — run concurrently.
|
|
let (gemma, ()) = tokio::try_join!(wire_gemma_client(), server::ensure_running())?;
|
|
|
|
let shieldstral = wire_shieldstral().await?;
|
|
|
|
let result = revise::generate_below_threshold(&gemma, &shieldstral).await?;
|
|
|
|
tracing::info!("Final score (score={:.6})", result.score);
|
|
tracing::info!("Final text ({})", result.text);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
async fn wire_gemma_client() -> anyhow::Result<ollama::CompletionModel> {
|
|
let gemma_client = ollama::Client::new(Nothing)?;
|
|
|
|
let gemma = gemma_client.completion_model("gemma4-e4b:latest");
|
|
|
|
Ok(gemma)
|
|
}
|
|
|
|
async fn wire_shieldstral() -> anyhow::Result<GenericCompletionModel<LlamafileExt>> {
|
|
let client = llamafile::Client::from_url(&server::url())?;
|
|
|
|
// Name doesn't matter here, server just uses whatever is running on it.
|
|
let shieldstral = client.completion_model("shieldstral");
|
|
|
|
Ok(shieldstral)
|
|
}
|