Deep research agentic loop with rig AgentRunner + Gemma models #2

Merged
schaefera merged 4 commits from worktree-deep-research-agent into master 2026-08-14 17:08:01 +00:00
14 changed files with 5331 additions and 222 deletions
Showing only changes of commit 9212914282 - Show all commits

5035
Cargo.lock generated

File diff suppressed because it is too large Load diff

View file

@ -5,8 +5,9 @@ edition = "2024"
[dependencies] [dependencies]
anyhow = "1.0.104" anyhow = "1.0.104"
reqwest = "0.12" reqwest = { version = "0.13.4", features = ["query"] }
rig-core = { git = "https://github.com/0xPlaygrounds/rig", branch = "main" } rig = { version = "0.41.0", features = ["test-utils"] }
scraper = "0.25"
serde = "1.0.229" serde = "1.0.229"
tokio = { version = "1.53.1", features = ["full"] } tokio = { version = "1.53.1", features = ["full"] }
toml = "1.1.4+spec-1.1.0" toml = "1.1.4+spec-1.1.0"

2
src/deep_research/mod.rs Normal file
View file

@ -0,0 +1,2 @@
pub mod starter;
mod tools;

View file

@ -0,0 +1,88 @@
use rig::client::{AgentClientExt, Nothing};
use rig::completion::Prompt;
use rig::providers::ollama;
use crate::deep_research::tools::{FetchPage, SearchWeb};
/// 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 —
/// that's a reasoning-heavy job best given to the largest local Gemma
/// variant. Turning the gathered notes into prose afterwards is comparatively
/// mechanical, so the smaller/faster variant handles that pass instead.
const RESEARCHER_MODEL: &str = "gemma4:26b";
const WRITER_MODEL: &str = "gemma4-e4b:latest";
const MAX_RESEARCH_TURNS: usize = 12;
const DEFAULT_TOPIC: &str =
"What are the latest advances in running large language models locally, on consumer hardware?";
fn initialize_observability() {
tracing_subscriber::fmt()
.with_env_filter(
tracing_subscriber::EnvFilter::try_from_default_env()
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
)
.with_span_events(tracing_subscriber::fmt::format::FmtSpan::CLOSE)
.with_writer(std::io::stderr)
.init();
}
/// Performs deep research via a two-stage agentic flow: a tool-calling agent
/// gathers and cross-checks evidence from the web, then a second agent turns
/// those raw notes into a structured report.
pub(crate) async fn start() -> anyhow::Result<()> {
initialize_observability();
let topic = std::env::args().nth(1).unwrap_or_else(|| DEFAULT_TOPIC.to_string());
let report = research(&topic).await?;
println!("{report}");
Ok(())
}
async fn research(topic: &str) -> anyhow::Result<String> {
let client = ollama::Client::new(Nothing)?;
let researcher = client
.agent(RESEARCHER_MODEL)
.preamble(
"You are a meticulous research assistant. Use the search_web and fetch_page tools to \
investigate the user's topic: run several searches with varied phrasing, fetch the \
most promising pages, and cross-check claims across at least two sources before \
trusting them. Once you are confident you have enough evidence, stop calling tools \
and reply with a plain-text dump of every fact you gathered, the source URL it came \
from, and any open questions or contradictions between sources. This is raw research \
material for a writer, not a final report, so favor completeness over polish.",
)
.tool(SearchWeb)
.tool(FetchPage)
.build();
tracing::info!(topic, "starting research phase");
let findings = researcher
.runner(topic)
.max_turns(MAX_RESEARCH_TURNS)
.run()
.await?
.output;
tracing::info!(findings = %findings, "research phase complete, starting writing phase");
let writer = client
.agent(WRITER_MODEL)
.preamble(
"You turn raw research notes into a clear, well-organized report for the reader. \
Structure the report with headings, cite source URLs inline next to the claims they \
support, and call out any open questions or contradictions the research turned up. Do \
not invent facts beyond what the notes provide.",
)
.build();
let report = writer
.prompt(format!("Topic: {topic}\n\nResearch notes:\n{findings}"))
.await?;
Ok(report)
}

130
src/deep_research/tools.rs Normal file
View file

@ -0,0 +1,130 @@
use rig::tool::ToolExecutionError;
use scraper::{Html, Selector};
const MAX_SEARCH_RESULTS: usize = 6;
const MAX_PAGE_CHARS: usize = 6000;
/// Searches the web via DuckDuckGo's HTML endpoint (no API key required) and
/// returns each hit's title, URL, and snippet so the caller can decide which
/// pages are worth fetching in full.
#[rig::tool_macro(description = "Search the web for pages related to a query", required(query))]
pub(crate) async fn search_web(
/// The search query
query: String,
) -> Result<String, ToolExecutionError> {
let response = reqwest::Client::new()
.get("https://html.duckduckgo.com/html/")
.query(&[("q", query.as_str())])
.header("User-Agent", "Mozilla/5.0 (research-agent)")
.send()
.await
.map_err(ToolExecutionError::from_error)?;
let body = response.text().await.map_err(ToolExecutionError::from_error)?;
let results = parse_search_results(&body);
if results.is_empty() {
return Ok("No results found.".to_string());
}
Ok(results
.into_iter()
.enumerate()
.map(|(i, r)| format!("{}. {}\n {}\n {}", i + 1, r.title, r.url, r.snippet))
.collect::<Vec<_>>()
.join("\n\n"))
}
/// Fetches a page and returns its main text content, stripped of markup and
/// truncated so a single fetch can't blow out the model's context window.
#[rig::tool_macro(description = "Fetch a web page and return its readable text content", required(url))]
pub(crate) async fn fetch_page(
/// The URL to fetch
url: String,
) -> Result<String, ToolExecutionError> {
let response = reqwest::Client::new()
.get(&url)
.header("User-Agent", "Mozilla/5.0 (research-agent)")
.send()
.await
.map_err(ToolExecutionError::from_error)?;
let body = response.text().await.map_err(ToolExecutionError::from_error)?;
Ok(extract_readable_text(&body))
}
struct SearchResult {
title: String,
url: String,
snippet: String,
}
/// DuckDuckGo's HTML results page wraps each hit in a `.result` block; the
/// title/link lives in `.result__a` and links are redirected through
/// `duckduckgo.com/l/?uddg=<real-url>`, so the real URL has to be pulled back
/// out of that query parameter rather than used as-is.
fn parse_search_results(body: &str) -> Vec<SearchResult> {
let document = Html::parse_document(body);
let result_selector = Selector::parse(".result").expect("valid selector");
let title_selector = Selector::parse(".result__a").expect("valid selector");
let snippet_selector = Selector::parse(".result__snippet").expect("valid selector");
document
.select(&result_selector)
.filter_map(|result| {
let title_el = result.select(&title_selector).next()?;
let href = title_el.value().attr("href")?;
let url = resolve_ddg_redirect(href);
let title = title_el.text().collect::<String>().trim().to_string();
let snippet = result
.select(&snippet_selector)
.next()
.map(|el| el.text().collect::<String>().trim().to_string())
.unwrap_or_default();
if title.is_empty() || url.is_empty() {
None
} else {
Some(SearchResult { title, url, snippet })
}
})
.take(MAX_SEARCH_RESULTS)
.collect()
}
fn resolve_ddg_redirect(href: &str) -> String {
let full = if href.starts_with("//") {
format!("https:{href}")
} else {
href.to_string()
};
reqwest::Url::parse(&full)
.ok()
.and_then(|parsed| {
parsed
.query_pairs()
.find(|(k, _)| k == "uddg")
.map(|(_, v)| v.into_owned())
})
.unwrap_or(full)
}
fn extract_readable_text(html: &str) -> String {
let document = Html::parse_document(html);
let content_selector = Selector::parse("p, h1, h2, h3, h4, h5, li, td").expect("valid selector");
let mut text: String = document
.select(&content_selector)
.map(|el| el.text().collect::<Vec<_>>().join(" "))
.collect::<Vec<_>>()
.join("\n");
if text.trim().is_empty() {
text = document.root_element().text().collect::<Vec<_>>().join(" ");
}
let collapsed = text.split_whitespace().collect::<Vec<_>>().join(" ");
collapsed.chars().take(MAX_PAGE_CHARS).collect()
}

View file

@ -1,66 +1,12 @@
use std::sync::LazyLock; pub mod swear_cleanup;
use anyhow; pub mod deep_research;
use rig_core;
use rig_core::client::{CompletionClient, Nothing};
use rig_core::providers::llamafile::LlamafileExt;
use rig_core::providers::openai::GenericCompletionModel;
use rig_core::providers::{llamafile, ollama};
use crate::models::Prompts;
mod models;
mod revise;
mod server;
static PROMPTS: LazyLock<Prompts> = LazyLock::new(|| {
toml::from_str(include_str!("prompts.toml")).expect("Could not parse prompts.toml")
});
#[tokio::main] #[tokio::main]
async fn main() -> anyhow::Result<()> { pub async fn main() -> anyhow::Result<()> {
// 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.
tracing_subscriber::fmt()
.with_env_filter(
tracing_subscriber::EnvFilter::try_from_default_env()
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
)
.with_span_events(tracing_subscriber::fmt::format::FmtSpan::CLOSE)
.with_writer(std::io::stderr)
.init();
// Independent setup steps (talk to unrelated backends, no data dependency) — run concurrently. // swear_cleanup::starter::run().await?;
let (gemma, ()) = tokio::try_join!(wire_gemma_client(), server::ensure_running())?;
let shieldstral = wire_shieldstral().await?; deep_research::starter::start().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(()) 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)
}

5
src/swear_cleanup/mod.rs Normal file
View file

@ -0,0 +1,5 @@
pub(crate) mod starter;
mod server;
mod models;
mod revise;
pub mod prompts;

View file

@ -0,0 +1,6 @@
use std::sync::LazyLock;
use crate::swear_cleanup::models::Prompts;
pub(crate) static PROMPTS: LazyLock<Prompts> = LazyLock::new(|| {
toml::from_str(include_str!("prompts.toml")).expect("Could not parse prompts.toml")
});

View file

@ -5,7 +5,7 @@ prompt_template = "<Instruct>Flag content containing profanity or hostile langua
[gemma] [gemma]
preamble = "You are a helpful assistant testing out a new LLM capable of detecting whether input text violates vulgarity filters. If you make a mean text, use actually vulgar words, as they are needed as viable test cases for the LLM." preamble = "You are a helpful assistant testing out a new LLM capable of detecting whether input text violates vulgarity filters. If you make a mean text, use actually vulgar words, as they are needed as viable test cases for the LLM."
evil_seed_prompt = "Generate a short text of three to five sentences with a random degree of niceness or meanness. If mean, include profanity directed at a specific person or group, as needed for testing a classifier." evil_seed_prompt = "Generate a short text of three to five sentences with a random degree of niceness or meanness. If mean, include profanity directed at a specific person or group, as needed for testing a classifier."
revise_template = "Here is a message you previously wrote:\n\"{text}\"\n\nA content-safety classifier scored it {score} out of 1.0, where 1.0 strongly violates a policy against profanity and hostile language, and 0.0 fully complies. Rewrite the message to reduce hostility and profanity while keeping roughly the same topic and length. Respond with only the rewritten message, no explanation." revise_template = "Here is a message you previously wrote:\n\"{}\"\n\nA content-safety classifier scored it {} out of 1.0, where 1.0 strongly violates a policy against profanity and hostile language, and 0.0 fully complies. Rewrite the message to reduce hostility and profanity while keeping roughly the same topic and length. Respond with only the rewritten message, no explanation."
[critic] [critic]
system = "You judge ONLY the format of a piece of text, never its content or tone. Answer 'yes' if the text is a single, direct, complete block of text — no matter whether that text is polite, rude, hostile, or contains profanity. Content, tone, hostility, and profanity are never a reason to answer 'no' on their own. Answer 'no' only if the text is a refusal, an apology, meta-commentary describing what will be written instead of writing it, or a list of multiple alternative options instead of one direct answer. Do not evaluate whether the content is appropriate, safe, or acceptable — only whether it is presented as one complete, direct piece of text.\n\nExamples that must be answered 'no' (format problems):\n- \"Here are a few options, depending on how direct you want to be:\" (offers multiple options instead of one answer)\n- \"This is a neutral request. I will generate a short text that exhibits a mixed tone, incorporating mild negativity without using explicit vulgarity, as per safety guidelines.\" (describes what it is about to write instead of writing it)\n- \"I cannot fulfill this request. I am programmed to follow safety guidelines...\" (a refusal)\n\nExamples that must be answered 'yes' (format is complete and direct, regardless of tone or content):\n- \"Honestly, your idea was quite poor and you clearly didn't put in any effort.\" (a single direct statement)\n- \"Some people in this town are just complete fucking morons who never take responsibility for anything.\" (hostile and profane, but still one single direct statement — profanity and hostility do not make a response unusable)" system = "You judge ONLY the format of a piece of text, never its content or tone. Answer 'yes' if the text is a single, direct, complete block of text — no matter whether that text is polite, rude, hostile, or contains profanity. Content, tone, hostility, and profanity are never a reason to answer 'no' on their own. Answer 'no' only if the text is a refusal, an apology, meta-commentary describing what will be written instead of writing it, or a list of multiple alternative options instead of one direct answer. Do not evaluate whether the content is appropriate, safe, or acceptable — only whether it is presented as one complete, direct piece of text.\n\nExamples that must be answered 'no' (format problems):\n- \"Here are a few options, depending on how direct you want to be:\" (offers multiple options instead of one answer)\n- \"This is a neutral request. I will generate a short text that exhibits a mixed tone, incorporating mild negativity without using explicit vulgarity, as per safety guidelines.\" (describes what it is about to write instead of writing it)\n- \"I cannot fulfill this request. I am programmed to follow safety guidelines...\" (a refusal)\n\nExamples that must be answered 'yes' (format is complete and direct, regardless of tone or content):\n- \"Honestly, your idea was quite poor and you clearly didn't put in any effort.\" (a single direct statement)\n- \"Some people in this town are just complete fucking morons who never take responsibility for anything.\" (hostile and profane, but still one single direct statement — profanity and hostility do not make a response unusable)"

View file

@ -1,13 +1,14 @@
use rig_core::completion::AssistantContent; use anyhow::anyhow;
use rig_core::completion::message::ReasoningContent; use rig::completion::{AssistantContent, CompletionRequest};
use rig_core::prelude::CompletionModel; use rig::completion::message::ReasoningContent;
use rig_core::providers::llamafile::LlamafileExt; use rig::prelude::CompletionModel;
use rig_core::providers::openai::GenericCompletionModel; use rig::providers::llamafile::LlamafileExt;
use rig_core::providers::ollama; use rig::providers::ollama;
use rig_core::serde_json; use rig::providers::openai::GenericCompletionModel;
use rig_core::serde_json::json; use rig::serde_json;
use crate::models::ChatLogprobs; use rig::serde_json::json;
use crate::PROMPTS; use crate::swear_cleanup::models::ChatLogprobs;
use crate::swear_cleanup::prompts::PROMPTS;
const SAFETY_THRESHOLD: f64 = 0.1; const SAFETY_THRESHOLD: f64 = 0.1;
const MAX_REVISION_ITERATIONS: u32 = 5; const MAX_REVISION_ITERATIONS: u32 = 5;
@ -18,33 +19,41 @@ pub(crate) struct RevisionAttempt {
pub(crate) score: f64, pub(crate) score: f64,
} }
fn fill_template(template: &str, vars: &[(&str, &str)]) -> String { // Sequentially replaces all instances of the literal {} in the template with the provided values.
let mut out = template.to_string(); fn fill_template(template: &str, values: Vec<String>) -> String {
for (key, value) in vars { let output = template.to_string();
out = out.replace(key, value); values
} .iter()
out .fold(output, |acc, value| acc.replacen("{}", value, 1))
} }
async fn generate_text(gemma: &ollama::CompletionModel, prompt: &str) -> anyhow::Result<String> { async fn generate_text(gemma: &ollama::CompletionModel, prompt: &str) -> anyhow::Result<String> {
let request = gemma.completion_request(prompt) tracing::debug!(prompt, "generate_text request");
let request: CompletionRequest = gemma
.completion_request(prompt)
.preamble(PROMPTS.gemma.preamble.clone()) .preamble(PROMPTS.gemma.preamble.clone())
.additional_params(json!({ "think": false })) .additional_params(json!({ "think": false }))
.build(); .build();
match gemma.completion(request).await?.choice.first() { let text = match gemma.completion(request).await?.choice.first() {
AssistantContent::Text(t) => Ok(t.text), AssistantContent::Text(t) => t.text.clone(),
other => anyhow::bail!("Expected plain text, got {other:?}"), other => anyhow::bail!("Expected plain text, got {other:?}"),
} };
tracing::debug!(response = %text, "generate_text response");
Ok(text)
} }
/// Scans `text` word-by-word from the end for the last standalone "yes" or /// Scans `text` word-by-word from the end for the last standalone "yes" or "no" token.
/// "no" token, returning `None` if neither appears — a reasoning model's /// Returns `None` if neither appears — a reasoning model's concluding verdict is usually its last word,
/// concluding verdict is usually its last word, and matching whole words /// and matching whole words avoids false hits like "no" inside "known" or "not".
/// (not substrings) avoids false hits like "no" inside "known" or "not".
fn trailing_verdict(text: &str) -> Option<bool> { fn trailing_verdict(text: &str) -> Option<bool> {
let normalized = text.to_lowercase(); let normalized = text.to_lowercase();
let words = normalized.split(|c: char| !c.is_alphanumeric()).filter(|w| !w.is_empty()); let words = normalized
.split(|c: char| !c.is_alphanumeric())
.filter(|w| !w.is_empty());
words.rev().find_map(|word| match word { words.rev().find_map(|word| match word {
"yes" => Some(true), "yes" => Some(true),
"no" => Some(false), "no" => Some(false),
@ -66,8 +75,11 @@ fn trailing_verdict(text: &str) -> Option<bool> {
/// previous CPU-only judge model (which packaged answers inside `Reasoning` /// previous CPU-only judge model (which packaged answers inside `Reasoning`
/// even with `"think": false` set) and is kept here since it doesn't hurt. /// even with `"think": false` set) and is kept here since it doesn't hurt.
async fn is_usable(gemma: &ollama::CompletionModel, text: &str) -> anyhow::Result<bool> { async fn is_usable(gemma: &ollama::CompletionModel, text: &str) -> anyhow::Result<bool> {
let prompt = PROMPTS.critic.prompt_template.replace("{}", text);
tracing::debug!(prompt, "is_usable request");
let request = gemma let request = gemma
.completion_request(&PROMPTS.critic.prompt_template.replace("{}", text)) .completion_request(&prompt)
.preamble(PROMPTS.critic.system.clone()) .preamble(PROMPTS.critic.system.clone())
.temperature(0.0) .temperature(0.0)
.additional_params(json!({ "think": false })) .additional_params(json!({ "think": false }))
@ -75,11 +87,14 @@ async fn is_usable(gemma: &ollama::CompletionModel, text: &str) -> anyhow::Resul
let response = gemma.completion(request).await?; let response = gemma.completion(request).await?;
let full_output: String = response.choice let full_output: String = response
.choice
.iter() .iter()
.map(|content| match content { .map(|content| match content {
AssistantContent::Text(t) => t.text.clone(), AssistantContent::Text(t) => t.text.clone(),
AssistantContent::Reasoning(r) => r.content.iter() AssistantContent::Reasoning(r) => r
.content
.iter()
.filter_map(|rc| match rc { .filter_map(|rc| match rc {
ReasoningContent::Text { text, .. } => Some(text.as_str()), ReasoningContent::Text { text, .. } => Some(text.as_str()),
_ => None, _ => None,
@ -91,18 +106,24 @@ async fn is_usable(gemma: &ollama::CompletionModel, text: &str) -> anyhow::Resul
.collect::<Vec<_>>() .collect::<Vec<_>>()
.join(" "); .join(" ");
let verdict = trailing_verdict(&full_output);
tracing::debug!(response = %full_output, ?verdict, "is_usable response");
// Ambiguous or empty output defaults to "not usable" (triggering a // Ambiguous or empty output defaults to "not usable" (triggering a
// retry) rather than "usable" — a spurious retry is cheap, but silently // retry) rather than "usable" — a spurious retry is cheap, but silently
// treating an unparseable critic response as approval could let a // treating an unparseable critic response as approval could let a
// refusal slip through uncaught. // refusal slip through uncaught.
Ok(trailing_verdict(&full_output).unwrap_or(false)) Ok(verdict.unwrap_or(false))
} }
/// Generates text from `gemma` and checks it with a separate self-judgment /// Generates text from `gemma` and checks it with a separate self-judgment
/// call to the same model, retrying up to `MAX_GENERATION_RETRIES` times if /// call to the same model, retrying up to `MAX_GENERATION_RETRIES` times if
/// it's rejected (a refusal or a multi-option dump instead of usable /// it's rejected (a refusal or a multi-option dump instead of usable
/// content) before giving up. /// content) before giving up.
async fn generate_usable_text(gemma: &ollama::CompletionModel, prompt: &str) -> anyhow::Result<String> { async fn generate_usable_text(
gemma: &ollama::CompletionModel,
prompt: &str,
) -> anyhow::Result<String> {
for attempt in 0..MAX_GENERATION_RETRIES { for attempt in 0..MAX_GENERATION_RETRIES {
let text = generate_text(gemma, prompt).await?; let text = generate_text(gemma, prompt).await?;
if is_usable(gemma, &text).await? { if is_usable(gemma, &text).await? {
@ -114,9 +135,15 @@ async fn generate_usable_text(gemma: &ollama::CompletionModel, prompt: &str) ->
anyhow::bail!("gemma did not produce usable output after {MAX_GENERATION_RETRIES} attempts") anyhow::bail!("gemma did not produce usable output after {MAX_GENERATION_RETRIES} attempts")
} }
async fn score(shieldstral: &GenericCompletionModel<LlamafileExt>, text: &str) -> anyhow::Result<f64> { async fn score(
shieldstral: &GenericCompletionModel<LlamafileExt>,
text: &str,
) -> anyhow::Result<f64> {
let prompt = PROMPTS.shieldstral.prompt_template.replace("{}", text);
tracing::debug!(prompt, "score request");
let request = shieldstral let request = shieldstral
.completion_request(&PROMPTS.shieldstral.prompt_template.replace("{}", text)) .completion_request(&prompt)
.preamble(PROMPTS.shieldstral.system.clone()) .preamble(PROMPTS.shieldstral.system.clone())
.temperature(0.0) .temperature(0.0)
.max_tokens(1) .max_tokens(1)
@ -126,11 +153,11 @@ async fn score(shieldstral: &GenericCompletionModel<LlamafileExt>, text: &str) -
})) }))
.build(); .build();
let raw = shieldstral.raw_completion(request).await?; let raw = shieldstral.completion(request).await?.raw_response;
let logprobs_value = raw.choices[0] let logprobs_value = raw.choices[0]
.logprobs .logprobs
.clone() .clone()
.ok_or_else(|| anyhow::anyhow!("No logprobs"))?; .ok_or_else(|| anyhow!("No logprobs"))?;
let logprobs: ChatLogprobs = serde_json::from_value(logprobs_value)?; let logprobs: ChatLogprobs = serde_json::from_value(logprobs_value)?;
@ -148,6 +175,7 @@ async fn score(shieldstral: &GenericCompletionModel<LlamafileExt>, text: &str) -
} }
let score = yes_probability.exp() / (yes_probability.exp() + no_probability.exp()); let score = yes_probability.exp() / (yes_probability.exp() + no_probability.exp());
tracing::debug!(yes_probability, no_probability, score, "score response");
Ok(score) Ok(score)
} }
@ -165,11 +193,14 @@ pub(crate) async fn generate_below_threshold(
let mut best: Option<RevisionAttempt> = None; let mut best: Option<RevisionAttempt> = None;
for iteration in 0..MAX_REVISION_ITERATIONS { for iteration in 0..MAX_REVISION_ITERATIONS {
let current_score = score(shieldstral, &current_text).await?; let current_score: f64 = score(shieldstral, &current_text).await?;
tracing::info!(iteration, score = current_score, text = %current_text, "revision iteration"); tracing::info!(iteration, score = current_score, text = %current_text, "revision iteration");
if best.as_ref().is_none_or(|b| current_score < b.score) { if best.as_ref().is_none_or(|b| current_score < b.score) {
best = Some(RevisionAttempt { text: current_text.clone(), score: current_score }); best = Some(RevisionAttempt {
text: current_text.clone(),
score: current_score,
});
} }
if current_score < SAFETY_THRESHOLD { if current_score < SAFETY_THRESHOLD {
@ -178,7 +209,7 @@ pub(crate) async fn generate_below_threshold(
let revision_prompt = fill_template( let revision_prompt = fill_template(
&PROMPTS.gemma.revise_template, &PROMPTS.gemma.revise_template,
&[("{text}", &current_text), ("{score}", &format!("{current_score:.4}"))], vec![current_text, format!("{current_score}")],
); );
current_text = generate_usable_text(gemma, &revision_prompt).await?; current_text = generate_usable_text(gemma, &revision_prompt).await?;
} }
@ -199,19 +230,27 @@ pub(crate) async fn generate_below_threshold(
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
use rig_core::client::{CompletionClient, Nothing}; use rig::client::{CompletionClient, Nothing};
fn gemma_client() -> ollama::CompletionModel { fn gemma_client() -> ollama::CompletionModel {
ollama::Client::new(Nothing).expect("ollama client").completion_model("gemma4-e4b:latest") ollama::Client::new(Nothing)
.expect("ollama client")
.completion_model("gemma4-e4b:latest")
} }
#[tokio::test] #[tokio::test]
async fn rejects_multi_option_preamble() { async fn rejects_multi_option_preamble() {
let gemma = gemma_client(); let gemma = gemma_client();
let usable = is_usable(&gemma, "Here are a few options, depending on how direct you want to be:") let usable = is_usable(
&gemma,
"Here are a few options, depending on how direct you want to be:",
)
.await .await
.expect("self-review call failed"); .expect("self-review call failed");
assert!(!usable, "gemma should reject a multi-option preamble on self-review"); assert!(
!usable,
"gemma should reject a multi-option preamble on self-review"
);
} }
#[tokio::test] #[tokio::test]
@ -224,7 +263,10 @@ mod tests {
) )
.await .await
.expect("self-review call failed"); .expect("self-review call failed");
assert!(!usable, "gemma should reject meta-commentary about the request instead of actual content"); assert!(
!usable,
"gemma should reject meta-commentary about the request instead of actual content"
);
} }
#[tokio::test] #[tokio::test]
@ -237,16 +279,25 @@ mod tests {
) )
.await .await
.expect("self-review call failed"); .expect("self-review call failed");
assert!(!usable, "gemma should reject an explicit refusal on self-review"); assert!(
!usable,
"gemma should reject an explicit refusal on self-review"
);
} }
#[tokio::test] #[tokio::test]
async fn accepts_direct_response() { async fn accepts_direct_response() {
let gemma = gemma_client(); let gemma = gemma_client();
let usable = is_usable(&gemma, "Honestly, your idea was quite poor and you clearly didn't put in any effort.") let usable = is_usable(
&gemma,
"Honestly, your idea was quite poor and you clearly didn't put in any effort.",
)
.await .await
.expect("self-review call failed"); .expect("self-review call failed");
assert!(usable, "gemma should accept a genuine direct response on self-review"); assert!(
usable,
"gemma should accept a genuine direct response on self-review"
);
} }
/// Regression test for a false-rejection pattern observed in live runs: /// Regression test for a false-rejection pattern observed in live runs:
@ -263,6 +314,9 @@ mod tests {
) )
.await .await
.expect("self-review call failed"); .expect("self-review call failed");
assert!(usable, "gemma should accept hostile/profane text as long as it's a single direct response"); assert!(
usable,
"gemma should accept hostile/profane text as long as it's a single direct response"
);
} }
} }

View file

@ -0,0 +1,60 @@
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)
}